训练每小时崩一次？豆包大模型港大团队为脆皮万卡训练提效

机器之心发布

机器之心编辑部

伴随大模型迭代速度越来越快，训练集群规模越来越大，高频率的软硬件故障已经成为阻碍训练效率进一步提高的痛点，检查点（Checkpoint）系统在训练过程中负责状态的存储和恢复，已经成为克服训练故障、保障训练进度和提高训练效率的关键。

近日，字节跳动豆包大模型团队与香港大学联合提出了ByteCheckpoint。这是一个PyTorch原生，兼容多个训练框架，支持Checkpoint的高效读写和自动重新切分的大模型Checkpointing系统，相比现有方法有显著性能提升和易用性优势。本文介绍了大模型训练提效中Checkpoint方向面临的挑战，总结ByteCheckpoint的解决思路、系统设计、I/O性能优化技术，以及在存储性能和读取性能测试的实验结果。

Meta官方最近披露了在16384块H10080GB训练集群上进行Llama3405B训练的故障率——短短54天，发生419次中断，平均每三小时崩溃一次，引来不少从业者关注。

正如业内一句常言，大型训练系统唯一确定的，便是软硬件故障。随着训练规模与模型大小的日益增长，克服软硬件故障，提高训练效率成为大模型迭代的重要影响要素。

Checkpoint已成为训练提效关键。在Llama训练报告中，技术团队提到，为了对抗高故障率，需要在训练过程中频繁地进行Checkpoint，保存训练中的模型、优化器、数据读取器状态，减少训练进度损失。

字节跳动豆包大模型团队与港大近期公开了成果——ByteCheckpoint，一个PyTorch原生，兼容多个训练框架，支持Checkpoint的高效读写和自动重新切分的大模型Checkpointing系统。

与基线方法相比，ByteCheckpoint在Checkpoint保存上性能提升高达529.22倍，在加载上，性能提升高达3.51倍。极简的用户接口和Checkpoint自动重新切分功能，显著降低了用户上手和使用成本，提高了系统的易用性。

目前论文成果已对外公开。

ByteCheckpoint:AUnifiedCheckpointingSystemforLLMDevelopment

论文链接：https://team.doubao.com/zh/publication/bytecheckpoint-a-unified-checkpointing-system-for-llm-development?view_from=research

Checkpoint技术在大模型训练中的技术挑战

当前Checkpoint相关技术在支持大模型训练提效中，共面临四个方面挑战：

现有系统设计存在缺陷，显著增加训练额外I/O开销

在训练工业级别的大语言模型(LLM)的过程中，训练状态需要通过检查点技术(Checkpointing)进行保存和持久化。通常情况下，一个Checkpoint包括5个部分(模型，优化器，数据读取器，随机数和用户自定义配置)。这一过程往往会给训练带来分钟级别的阻塞，严重影响训练效率。

在使用远程持久化存储系统的大规模训练场景下，现有的Checkpointing系统没有充分利用Checkpoint保存过程中GPU到CPU内存拷贝(D2H复制)，序列化，本地存盘，上传到存储系统各个阶段的执行独立性。

此外，不同训练进程共同分担Checkpoint存取任务的并行处理潜力也没有被充分发掘。这些系统设计上的不足增加了Checkpoint训练带来的额外I/O开销。

Checkpoint重新切分困难，手动切分脚本开发维护开销过高

在LLM的不同训练阶段(预训练到SFT或者RLHF)以及不同任务(从训练任务拉取不同阶段的Checkpoint进行执行自动评估)之间进行Checkpoint迁移时，通常需要对保存在持久化存储系统中的Checkpoint进行重新切分(CheckpointResharding)，训练每小时崩一次？豆包大模型港大团队为脆皮万卡训练提效以适应下游任务的新并行度配置以及可用GPU资源的配额。

现有Checkpointing系统[1,2,3,4]都假设存储和加载时，并行度配置和GPU资源保持不变，无法处理Checkpoint重新切分的需求。工业界目前常见的解决办法是——为不同模型定制Checkpoint合并或者重新切分脚本。这种方法带来了大量开发与维护开销，可扩展性较差。

不同的训练框架Checkpoint模块割裂，为Checkpoint统一管理和性能优化带来挑战

在工业界的训练平台上，工程师与科学家往往会根据任务特性，选择合适框架(Megatron-LM[5],FSDP[6],DeepSpeed[7],veScale[8,9])进行训练，并保存Checkpoint到存储系统。然而，这些不同的训练框架都具有自己独立的Checkpoint格式以及读写模块。不同训练框架的Checkpoint模块设计不尽相同，为底层系统进行统一的Checkpoint管理以及性能优化带来了挑战。

分布式训练系统的用户面临多重困扰

从训练系统的用户（AI研究科学家或工程师）的角度出发，用户使用分布式训练系统时，在Checkpoint方向往往会被三个问题困扰：

1）如何高效地存储Checkpoint，在不影响训练效率的情况下保存Checkpoint。

2）如何重新切分Checkpoint，对于在一个并行度下存储的Checkpoint，根据新的并行度正确读入。

3）如何把训练得到的产物上传到云存储系统上（HDFS，S3等），手动管理多个存储系统，对用户来说学习和使用成本较高。

针对上述问题，字节跳动豆包大模型团队和香港大学吴川教授实验室联合推出了ByteCheckpoint。

ByteCheckpoint是一个多训练框架统一，支持多存储后端，具备自动Checkpoint重新切分能力的高性能分布式Checkpointing系统。ByteCheckpoint提供了简单易用的用户接口，实现了大量I/O性能优化技术提高了存储和读取Checkpoint性能，并支持Checkpoint在不同并行度配置的任务中的灵活迁移。

系统设计

存储架构

ByteCheckpoint采用了元数据/张量数据分离的存储架构，实现了Checkpoint管理与训练框架和并行度的解耦合。

不同训练框架中的模型以及优化器的张量切片(TensorShard)存储在storage文件中，元信息(TensorMeta,ShardMeta,ByteMeta)存储到全局唯一的metadata文件中。

当使用不同的并行度配置读取Checkpoint时，如下图所示，每个训练进程只需要根据当前的并行度设置查询元信息，便能够获取进程所需要张量的存储位置，再根据位置直接读取，实现自动Checkpoint重新切分。

巧解不规则张量切分

不同训练框架在运行时，往往会把模型或者优化器中张量的形状摊平(Flatten)成一维，从而提高集合通信性能。这种摊平操作给Checkpoint存储带来了不规则张量切分(IrregularTensorSharding)的挑战。

如下图所示，在Megatron-LM(由NVIDIA研发的分布式大模型训练框架)和veScale(由字节跳动研发的PyTorch原生分布式大模型训练框架)中，模型参数对应的优化器状态会被展平为一维后合并，再根据数据并行度切分。这导致张量被不规则地切分到不同进程之中，张量切片的元信息无法使用偏移量和长度元组来表示，给存储和读取带来困难。

不规则张量切分的问题在FSDP框架中也同样存在。

为消除不规则切分的张量切片，FSDP框架在存储Checkpoint之前会在所有进程上对一维张量切片进行all-gather集合通信以及D2H复制操作，以获取完整不规则切分的张量。这种方案带来了极大的通信和频繁的GPU-CPU同步开销，严重影响了Checkpoint存储的性能。

针对这个问题，ByteCheckpoint提出了异步张量合并(AsynchronousTensorMerging)技术。

ByteCheckpoint首先找出不同进程中被不规则切分的张量，之后采用异步的P2P通信，把这些不规则的张量分配到不同进程上进行合并。所有针对这些不规则张量的P2P通信等待（Wait)以及张量D2H复制操作被推迟到他们即将进入序列化阶段的时候，从而消除了频繁的同步开销，也增加了通信与其他Checkpoint存储流程的执行重叠度。

系统架构

下图展示了ByteCheckpoint的系统架构：

API层为不同训练框架提供了简单，易用且统一的读取和写入(Save）和读取（Load）接口。

Planner层会根据存取对象为不同训练进程生成存取方案，交由Execution层执行实际的I/O任务。

Execution层执行I/O任务并与Storage层进行交互，利用各种I/O优化技术进行高性能的Checkpoint存取。

Storage层管理不同的存储后端，并在I/O任务过程中根据不同存储后端进行相应的优化。

分层设计增强了系统的可扩展性，以便未来支持更多的训练框架和存储后端。

API用例

ByteCheckpoint的API用例如下：

ByteCheckpoint提供了极简API，降低了用户上手的成本。用户在存储和读取Checkpoint时，只需要调用存储和加载函数，传入需要存储和读取的内容，文件系统路径和各种性能优化选项。

I/O性能优化技术

Checkpoint存储优化

流水线执行

如下图所示，ByteCheckpoint设计了全异步的存储流水线（SavePipeline），将Checkpoint存储的不同阶段（P2P张量传输，D2H复制，序列化，保存本地和上传文件系统）进行拆分，实现高效的流水线执行。

避免内存重复分配

在D2H复制过程，ByteCheckpoint采用固定内存池（PinnedMemoryPool），减少了内存反复分配的时间开销。

除此之外，为了降低高频存储场景中因为同步等待固定内存池回收而带来的额外时间开销，ByteCheckpoint在固定内存池的基础上加入了Ping-Pongbuffering的机制。两个独立的内存池交替扮演着读写buffer的角色，与GPU和执行后续I/O操作的I/Oworkers进行交互，进一步提升存储效率。

负载均衡

在数据并行(Data-ParallelorDP)训练中，模型在不同的数据并行进程组（DPGroup）之间是冗余的，ByteCheckpoint采用了负载均衡算法把冗余的模型张量均匀分配到不同进程组中进行存储，有效地提高了Checkpoint存储效率。

Checkpoint读取优化

零冗余加载

如图所示，在改变并行度读取Checkpoint时，新的训练进程可能只需要从原来的张量切片中读取其中的一部分。

ByteCheckpoint采用按需部分文件读取（PartialFileReading）技术，直接从远程存储中读取需要的文件片段，避免下载和读取不必要的数据。

在数据并行(Data-ParallelorDP)训练中，模型在不同的数据并行进程组（DPGroup）之间是冗余的，不同进程组会重复读取同一个张量切片。在大规模训练的场景下，不同进程组同时发给远程持久化存储系统（比如HDFS）大量请求，会给存储系统带来巨大压力。

为了消除重复数据读取，减少训练进程发给HDFS的请求，优化加载的性能，ByteCheckpoint把相同的张量切片读取任务均匀分配到不同进程上，并在对远程文件进行读取的同时，利用GPU之间闲置的带宽进行张量切片传输。

实验结果

实验配置

团队使用DenseGPT与SparseGPT模型(基于GPT-3[10]结构实现)，在不同模型参数量，不同训练框架和不同规模的训练任务中评估了ByteCheckpoint的Checkpoint存取正确性、存储性能和读取性能。更多实验配置和正确性测试细节请移步完整论文。

存储性能测试

在存储性能测试中，团队比较了不同模型规模和训练框架，在训练过程中每50或者100步存一次Checkpoint，Bytecheckpoint和基线（Baseline）方法给训练带来的总的阻塞时间(Checkpointstalls)。

得益于对写入性能的深度优化，ByteCheckpoint在各类实验场景中均取得了很高的表现，在576卡SparseGPT110B-Megatron-LM训练任务中相比基线存储方法取得了66.65~74.55倍的性能提升，在256卡DenseGPT10B-FSDP训练任务中甚至能达到529.22倍的性能提升。

读取性能测试

在读取性能测试中，团队比较不同方法根据下游任务并行度读取Checkpoint的加载时间。ByteCheckpoint相比基线方法取得了1.55～3.37倍的性能提升。

团队观察到ByteCheckpoint相对于Megatron-LM基线方法的性能提升更为显著。这是因为Megatron-LM在读取Checkpoint到新的并行度配置之前，需要运行离线的脚本对分布式Checkpoint进行重新分片。相比之下，ByteCheckpoint能够直接进行自动Checkpoint重新切分，无需运行离线脚本，高效完成读取。

最后，关于ByteCheckpoint的未来规划，团队希望从两个方面着手：

其一，实现支持超大规模GPU集群训练任务高效Checkpointing的长远目标。

其二，实现大模型训练全生命周期的Checkpoint管理，支持全场景的Checkpoint，从预训练（Pre-Training)，到监督微调（SFT），再到强化学习（RLHF）和评估(Evaluation)等场景。

团队介绍

字节跳动豆包大模型团队成立于2023年，致力于开发业界最先进的AI大模型技术，成为世界一流的研究团队，为科技和社会发展作出贡献。

目前，团队正在持续吸引优秀人才加入，硬核、开放且充满创新精神是团队氛围关键词，团队致力于创造一个积极向上的工作环境，鼓励团队成员不断学习和成长，不畏挑战，追求卓越。

希望与具备创新精神、责任心的技术人才一起，推进大模型训练提效工作取得更多进展和成果。

参考文献

[1]Mohan,Jayashree,AmarPhanishayee,andVijayChidambaram."{CheckFreq}:Frequent,{Fine-Grained}{DNN}Checkpointing."19thUSENIXConferenceonFileandStorageTechnologies(FAST21).2021.

[2]Eisenman,Assaf,etal."{Check-N-Run}:ACheckpointingsystemfortrainingdeeplearningrecommendationmodels."19thUSENIXSymposiumonNetworkedSystemsDesignandImplementation(NSDI22).2022.

[3]Wang,Zhuang,etal."Gemini:Fastfailurerecoveryin分布式trainingwithin-memoryCheckpoints."Proceedingsofthe29thSymposiumonOperatingSystemsPrinciples.2023.

[4]Gupta,Tanmaey,etal."Just-In-TimeCheckpointing:LowCostErrorRecoveryfromDeepLearningTrainingFailures."ProceedingsoftheNineteenthEuropeanConferenceonComputerSystems.2024.

[5]Shoeybi,Mohammad,etal."Megatron-lm:Trainingmulti-billionparameterlanguagemodelsusingmodelparallelism."arXivpreprintarXiv:1909.08053(2019).

[6]Zhao,Yanli,etal."Pytorchfsdp:experiencesonscalingfullyshardeddataparallel."arXivpreprintarXiv:2304.11277(2023).

[7]Rasley,Jeff,etal."Deepspeed:Systemoptimizationsenabletrainingdeeplearningmodelswithover100billionparameters."Proceedingsofthe26thACMSIGKDDInternationalConferenceonKnowledgeDiscovery&DataMining.2020.

[8]Jiang,Ziheng,etal."{MegaScale}:Scalinglargelanguagemodeltrainingtomorethan10,000{GPUs}."21stUSENIXSymposiumonNetworkedSystemsDesignandImplementation(NSDI24).2024.

[9]veScale:APyTorchNativeLLMTrainingFrameworkhttps://github.com/volcengine/veScale

[10]Brown,Tom,etal."Languagemodelsarefew-shotlearners."Advancesinneuralinformationprocessingsystems33(2020):1877-1901.

https://jobs.bytedance.com/campus/position?keywords=seed foundation&category=&location=&project=&type=&job_hot_flag=¤t=1&limit=10&functionCategory=&tag=&spread=FBUF844