大模型算力中被忽视的I/O问题

中国计算机学会

摘要—大模型时代的到来意味着模型的参数量和训练数据量飞速增长，算力的快速提升导致大模型算力系统的瓶颈逐渐向I/O转移。本文介绍了大模型算力系统中的I/O来源、特性以及优化方案。

林清音（中山大学）
陈志广（中山大学）

关键词 ：大模型算力 I/O 瓶颈 数据加载 保存检查点

  

大模型时代高算力需求

        人工智能（AI）大模型时代已经到来，国内外已经掀起了一股大模型热潮。OpenAI率先推出GPT，随后阿里达摩院的通义M6、谷歌的GLaM、百度的“文心一言”等大模型也相继被推出。各大互联网公司已经对AI大模型投入了深度研究，学术界更是如此。从人工智能发展历程来看，大模型是不可阻挡的趋势，模型的参数量以及训练数据量都在持续增加。2020年5月发布的GPT-3模型，其参数量已经达到1750亿，预训练数据量更是达到45 TB。从人工智能目前的发展状态来看，更好的模型应用效果意味着更大的参数量和更大的训练数据集，而GPU等硬件加速器的发展进一步推动了这种趋势。

        出于数据安全或是核心竞争力等因素的考虑，各个国家、各大互联网厂商都在不断尝试推出属于自己的大模型，这种爆发性增长的需求对大模型算力系统带来了各种挑战，算力资源首当其冲。为了更好地构建支持大模型的算力系统，学业界和工业界做了许多努力。例如不断推出性能更高的硬件加速器，英伟达推出了高性能的A100、H100 GPU，AMD推出了MI300，我国厂商也推出了寒武纪思元370、海光8100等AI加速器。在大模型的高算力要求下，构建万卡规模的算力集群，就必须设计新型的网络拓扑以支持万卡级别的高速互连。此外，单机的模型训练框架早已无法满足需求，各种分布式训练框架被陆续提出并不断优化，以高效利用算力资源。

大模型算力系统中的I/O瓶颈

       完整的大模型算力系统不仅需要高性能算力、调度能力以及高效的训练框架，还需要高性能的存储子系统，现有研究在算力资源不足的巨大压力下忽略了算力系统存储模块面临的挑战。参数量和训练数据集规模的快速增长，意味着存储子系统需要应对更多的I/O请求，高效的算力系统离不开高性能的I/O模块。

       大规模算力系统以机器学习系统为基础，包含许多模块，需要服务多种机器学习流程。斯卡利（Sculley）等人[1]将机器学习系统分解为多个模块，如图1所示。能够提供服务的机器学习流程包括：数据收集（data collection）、模型训练（model training）、数据验证（data verification）、模型推理服务（model serving）等。其中，模型训练需要较高的算力和时间成本，其效率受到的关注最为广泛。随着模型规模的增大，提高模型训练性能成为重中之重。然而，高性能硬件加速器的出现，算法和框架的不断优化，使得机器学习系统的计算能力和计算速度迅速提高，其性能瓶颈逐渐转移到I/O。根据这种趋势，数据加载的效率将取代算力成为机器学习系统性能的决定性因素。相关研究表明[2]，当训练较大规模的ResNet-50模型时，I/O开销占据了运行时开销的85%。特别是在分布式机器学习训练系统中，计算的并行可以简单且高效地进行扩展，而数据却需要从分布式的文件系统加载到本地内存，使I/O瓶颈问题越发突出。相关研究表明[3]，当训练任务扩展到9216个进程时，有近90%的时间在执行I/O任务。

       要提高机器学习系统的I/O性能，就要清楚模型训练过程的I/O来源。模型训练过程包括数据加载、数据预处理、模型计算、计算损失、参数更新、保存检查点和模型参数等阶段。其中，数据加载和保存检查点涉及大量I/O操作。数据加载是不可避免的阶段，后续的一系列步骤都必须等待数据加载完成才能进行。当I/O速度无法匹配计算速度时，训练过程将花费大量的时间等待I/O任务的完成。检查点的保存可根据不同的需求适当增加或减少，减少检查点即减少了I/O，减少总体训练时间，但当程序意外崩溃时恢复训练进度则需要更多时间。此外，模型的增大意味着每次保存检查点都需要将更多的参数写入持久存储，也意味着减少检查点将以更长的计算时间为代价。总之，训练数据集规模的快速增长，模型参数量的不断增加，进一步加重了I/O系统的任务，设计高效的I/O子系统迫在眉睫。

       由于模型训练过程中的主要I/O来源是数据加载和保存检查点，据此可以分析I/O请求类型。数据加载过程中的I/O请求是读密集型的，保存检查点的I/O请求则是写密集型的，机器学习任务总体的I/O类型需要结合这两部分进行判断。阿纳柏·保罗（Arnab K. Paul）等人[4]对不同科学领域的机器学习任务进行了测试，包括生物学、计算机科学、材料学、化学四个领域。如表1所示，读密集型任务是科学领域机器学习的主要任务类型，占比高达82%。此外，他们还对不同规模的机器学习任务进行了测试，实验数据表明，随着机器学习任务规模的增大，读密集型任务占比也越来越高。随着训练数据集规模的增大，单机的文件系统早已无法满足需求，大规模的训练任务往往从并行文件系统中加载数据，例如BeeGFS以及IBM的通用并行文件系统（General Parallel FileSystem，GPFS）等。尽管这些文件系统能够实现并行I/O，但它们并不是为读密集型I/O请求而设计的，相比之下，高性能集群上的burst buffer更适合服务读密集型I/O请求。文献[4]证明了今后的机器学习任务将继续呈现读密集趋势，因此burst buffer可以作为优化机器学习系统I/O性能的有力手段。

面向大模型训练的I/O优化

       除了了解机器学习任务的I/O类型之外，我们还可以针对机器学习训练过程的特点，分析机器学习系统在数据加载和保存检查点过程的I/O特性。

数据加载

       机器学习的数据集主要由图片、文本等小文件组成，而加载数据集是机器学习过程的主要I/O来源之一。许多工作[2~6]提出机器学习过程中存在大量小文件访问，并分析了这种I/O特征对模型训练效率带来的负面影响。阿纳柏·保罗等人[4]对科学领域机器学习任务中的文件访问大小进行了测试，其中99%的读写请求小于10 MB。对于大规模的机器学习任务，其输入的数据集总量可以达到TB级，因此往往采用分布式文件系统存储数据，大量的小文件访问将给文件系统带来多种压力。首先，频繁的小文件访问会加重共享文件系统内部的读写竞争，现有的框架大多将I/O与计算重叠以减少等待时间，然而随着计算速度的加快和数据集规模的增大，这种方法已无法缓解I/O瓶颈。其次，过多小文件存储会对文件系统的元数据服务器带来巨大的压力，最终导致文件系统无法工作，这表明传统的文件系统难以应对机器学习训练场景，一些工作[5]采用分布式应用对象存储（Distributed Application Object Storage，DAOS）技术解决元数据服务器瓶颈问题。此外，频繁的小文件访问给文件系统带来了大量的查询请求，而分布式文件系统的查询需要客户端与NameNode和DataNode进行通信，大量的查询请求导致频繁的DataNode通信，降低了I/O性能。相关工作[6]分析了小文件之间的关联性，并将可能同时访问的小文件合并，以缓解分布式文件系统的压力。

       大量工作[2, 6~12]都已指出机器学习过程存在对文件系统的大量随机访问，严重降低了系统性能。梯度下降算法是机器学习中常见的优化算法，当前几乎每一个机器学习库都含有梯度下降算法，其中，随机梯度下降（Stochastic Gradient Descent，SGD）算法由于优化效果较好被广泛使用。随机梯度下降每次从训练集中随机选择样本进行计算并更新模型参数，直至模型收敛，在此过程中，随机性是至关重要的，关系到模型是否收敛。对数据的随机访问意味着I/O请求无序且难以预测，这对持久存储设备非常不友好，而分布式文件系统在此基础上还会引入网络传输时间，进一步造成了I/O瓶颈问题。

       分批（batch）训练是机器学习的重要特征。机器学习的训练包含多个时期（epoch），每个时期内部包含多个迭代（iteration），每个迭代负责对某个批次的数据进行计算。也就是说，训练数据集在开始训练之前被划分为多个批次，分发给不同的迭代进行多次计算。批次内部的数据会被随机访问，且在多次计算中每次访问顺序都不同，但由于划分好的批次不会再变更，每个批次包含的数据是固定的。PyTorch和TensorFlow是常用的两种机器学习框架。PyTorch将数据加载模块与训练计算分隔开，可以启动多个进程并行读取数据，然后将训练所需要的批次数据重复地分发到共享内存队列，加快数据读取并减少对持久存储的访问。TensorFlow将随机采样限制在固定长度的缓存中避免大范围的随机，并将数据预取到该缓存中供后续访问。此外，也有相关工作根据分批训练这一特性进行了I/O优化。由于批次所包含的数据是固定的，当批次内某个数据被访问时其他数据也即将被访问。PHDFS[6]利用了分批训练中文件的关联性，将多个小文件合并成较大的文件，减少随机访问。尽管每个时期对批次内数据随机访问的顺序不同，但数据访问仍然具有重复性，因此可以将上个时期加载的数据缓存起来供下个时期使用，但每个时期涉及的数据量非常大，计算节点的本地缓存无法容纳。杨智傑（Chih-Chieh Yang）等人[10]将所有计算节点的本地缓存聚合成分布式缓存，大大减少了I/O请求。

       在现实世界中，越来越多的模型基于相同的数据集训练而成，同一公司的不同AI应用也往往基于该公司收集的相同的基础数据训练出不同的功能。然而，在多个机器学习作业之间高效地共享存储在云中的数据仍然具有挑战性。由于训练数据集的规模越来越大，具有相同数据集的机器学习作业很可能持有大量的重复数据缓存，这对本就昂贵的缓存空间十分不利。数据缓存共享是一个直观的解决方案。然而，在同一时间段内读取同一组数据集进行多种机器学习任务，仍需要谨慎处理数据缓存的共享，因为不恰当的任务调度会导致频繁的缓存驱逐和颠簸，降低缓存利用率，最终导致整个机器学习任务训练性能下降。文献[11]提出了一种基于数据缓存的机器学习任务调度策略，该策略考虑了数据集的缓存状态，以较高的优先级运行已缓存数据集的任务，最终降低了总体的训练时间。

       总而言之，机器学习训练数据加载过程中存在的主要特征为大量的随机小文件访问，针对这一特点，现有的I/O优化方案主要包括内存缓存和优化数据访问模式。然而，从整体来看，大规模算力系统也是大规模的机器学习系统，所用到的硬件资源、硬件架构、网络拓扑等都更为复杂，针对算力系统中的I/O瓶颈问题也不能局限于文件系统本身，需要综合整体的资源架构进行考虑。首先，针对海量的训练数据，需要采用大规模的存储系统。以高性能文件系统（High Performance File System，HPFS）为代表的分布式文件系统已被广泛应用，然而，集中式的元数据管理使文件系统规模受限，难以快速扩展，且往往需要将多个小文件合并成大文件后进行存储，因而在读取数据时引入了额外的流程开销。其次，出于成本的考量，对象存储更适合存算分离的场景。与分布式文件系统相比，对象存储拥有良好的扩展性，且具有多种存储介质分级存储的能力，在存储海量数据方面能够达到比分布式文件系统更低的成本。然而，尽管对象存储采用分布式键值（Key-Value，KV）存储等方法能够解决小文件的规模扩展问题，对象存储在I/O路径上却更为复杂，频繁地在小文件上读仍然具有不可忽略的延迟开销。综合看来，仅依赖分布式存储系统解决算力系统的I/O瓶颈问题较为困难。现有的分布式文件系统和对象存储都有各自难以克服的短板，要构建高性能的大规模算力系统，就必须构建与计算框架相匹配的存储系统。例如，根据数据存储格式设计具有良好扩展性且I/O路径较短的索引系统，或是利用数据集只读的特点，将数据放置到离计算节点更近的地方，构建高可扩展的缓存系统，同时根据模型训练过程的计算特点设计相对应的缓存替换算法。

保存检查点

       除了加载大数据集之外，大模型训练出于调试、迁移学习、故障恢复等目的，需要频繁地保存检查点，GPU需等待检查点保存完毕之后才能继续训练，因此保存检查点的I/O效率对整个训练过程至关重要。保存检查点的I/O过程是将大量模型参数写入持久存储，其I/O特性较为直观，即大量的顺序写入。尽管现有的存储介质已经能够支持较高性能的顺序写入，但大模型参数规模的增大使I/O开销和存储开销也快速增长，因此，探究对保存检查点的I/O优化方案仍然不可忽视。

       现有优化方案大多采用异步检查点的方式，将需要持久化的数据传输到更快的存储介质上，再异步地将数据持久化到传输速度较慢但容量更大的存储介质上，能够大大减少阻塞GPU计算的时间。Chien等人[13]使用突发缓冲区实现异步检查点机制，将保存检查点的性能提高了2.6倍。Checkfreq[14]借助CPU内存实现了异步检查点，将GPU计算与检查点保存流水线化，并通过算法自适应调整检查点保存频率，能够以最小的成本实现最频繁的检查点保存。Deepfreeze[15]引入了细粒度的异步序列化和I/O技术，并将负载分配给所有参与进程，大大降低了保存检查点的开销。

       除了异步检查点之外，另一种优化方案是减少保存检查点时需要持久化的数据量。算子融合[16]和量化[17]能够压缩模型参数，但会牺牲一定的精确度。Check-N-Run[17]跟踪了特定模型训练过程的参数更新情况，发现该模型在训练过程中仅更新一小部分参数，因此采用差异检查点的方案，首先将全部参数持久化作为基准检查点，然后仅保存与基准检查点有差异的参数，大大降低了保存检查点的I/O开销和存储开销。

       总之，保存检查点是大模型训练中不可忽视的一环，为了节省故障恢复重计算时间以及方便模型调试等，需要频繁保存检查点，存储大量模型参数带来了巨大的I/O和存储开销。现有的模型训练框架往往直接将模型参数保存为单个大文件，尽管利用了大量顺序写的特点，但忽略了检查点存储系统的可扩展性和存储开销。如何跟踪模型训练过程中参数的更新情况，使保存检查点时需要持久化的数据量减少，成为降低保存检查点成本的关键因素。此外，如何结合模型参数更新和层次化存储的特点设计高性能、可扩展的存储系统，也是提升大模型训练I/O性能的重要途径。

总结

       大模型训练参数量和样本量的快速增长为计算系统带来了挑战。在算力瓶颈引起广泛关注的当下，存储系统面临的I/O挑战往往被忽略。训练数据加载和模型检查点保存是模型训练过程中的关键步骤，也是涉及大量I/O操作的过程。当存储系统无法提供与计算速度相匹配的I/O性能时，I/O将成为大模型算力系统中的瓶颈之一。从I/O特性出发，数据加载存在对大量小文件随机读取的情况，检查点保存则是对大量参数的序列化和顺序持久化。针对大模型训练系统的I/O需求和I/O特性，为其提供高性能、可扩展的存储系统，将大大提升大模型算力系统的总体效率。  ■

参考文献：

[1] Sculley D, Holt G, Golovin D, et al. Hidden technical debt in machine learning systems[J]. Advances in Neural Information Processing Systems, 2015: 2494-2502.

[2] Dryden N, Böhringer R, Ben-Nun T, et al. Clairvoyant prefetching for distributed machine learning I/O[C]//Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 2021: 1-15.

[3] Pumma S, Si M, Feng W C, et al. Scalable deep learning via I/O analysis and optimization[J]. ACM Transactions on Parallel Computing (TOPC), 2019, 6(2): 1-34.

[4] Paul A K, Karimi A M, Wang F. Characterizing machine learning i/o workloads on leadership scale hpc systems[C]// 2021 29th International Symposium on Modeling, Analysis, and Simulation of Computer and Telecommunication Systems (MASCOTS). IEEE, 2021: 1-8.

[5] Párraga E, León B, Bond R, et al. Analyzing the I/O Patterns of Deep Learning Applications[C]// Cloud Computing, Big Data & Emerging Topics: 9th Conference. Springer International Publishing, 2021: 3-16.

[6] Zhu Z, Tan L, Li Y, et al. PHDFS: Optimizing I/O performance of HDFS in deep learning cloud computing platform[J]. Journal of Systems Architecture, 2020, 109: 101810.

[7] Zhu Y, Chowdhury F, Fu H, et al. Entropy-aware I/O pipelining for large-scale deep learning on HPC systems[C]//2018 IEEE 26th International Symposium on Modeling, Analysis, and Simulation of Computer and Telecommunication Systems (MASCOTS). IEEE, 2018: 145-156.

[8] Aizman A, Maltby G, Breuel T. High performance I/O for large scale deep learning[C]//2019 IEEE International Conference on Big Data (Big Data). IEEE, 2019: 5965-5967.

[9] Kakaraparthy A, Venkatesh A, Phanishayee A, et al. The case for unifying data loading in machine learning clusters[C]//11th USENIX Workshop on Hot Topics in Cloud Computing (HotCloud 19). 2019.

[10] Yang C C, Cong G. Accelerating data loading in deep neural network training[C]//2019 IEEE 26th International Conference on High Performance Computing, Data, and Analytics (HiPC). IEEE, 2019: 235-245.

[11] Gu R, Zhang K, Xu Z, et al. Fluid: Dataset abstraction and elastic acceleration for cloud-native deep learning training jobs[C]//2022 IEEE 38th International Conference on Data Engineering (ICDE). IEEE, 2022: 2182-2195.

[12] Li S, Kumar A. Towards an optimized GROUP by abstraction for large-scale machine learning[J]. Proceedings of the VLDB Endowment, 2021, 14(11): 2327-2340.

[13] Chien S W D, Markidis S, Sishtla C P, et al. Characterizing deep-learning I/O workloads in TensorFlow[C]//2018 IEEE/ACM 3rd International Workshop on Parallel Data Storage & Data Intensive Scalable Computing Systems (PDSW-DISCS). IEEE, 2018: 54-63.

[14] Mohan J, Phanishayee A, Chidambaram V. CheckFreq: Frequent,Fine-Grained DNN Checkpointing[C]//19th USENIX Conference on File and Storage Technologies (FAST 21). 2021: 203-216.

[15] Nicolae B, Li J, Wozniak J M, et al. Deepfreeze: Towards scalable asynchronous checkpointing of deep learning models[C]//2020 20th IEEE/ACM International Symposium on Cluster, Cloud and Internet Computing (CCGRID). IEEE, 2020: 172-181.

[16] He H, Yu S. Transcending Runtime-Memory Tradeoffs in Checkpointing by being Fusion Aware[J]. Proceedings of Machine Learning and Systems, 2023, 5.

[17] Eisenman A, Matam K K, Ingram S, et al. Check-N-Run: A checkpointing system for training deep learning recommendation models[C]//19th USENIX Symposium on Networked Systems Design and Implementation (NSDI 22). 2022: 929-943.

版权声明：中国计算机学会(CCF)拥有《中国计算机学会通讯》(CCCF)所刊登内容的所有版权，未经CCF允许，不得转载本刊文字及照片，否则被视为侵权。对于侵权行为，CCF将追究其法律责任。