ByteCheckpoint: A Unified Checkpointing System for Large Foundation Model Development
ByteCheckpoint: A Unified Checkpointing System for Large Foundation Model Development
发表时间: 2024-07 · arXiv:2407.20143
原文: https://arxiv.org/abs/2407.20143
Borui Wan1,∗ Mingji Han2,∗ Yiyao Sheng2 Yanghua Peng2 Haibin Lin2 Mofan Zhang2 Zhichao Lai2 Menghan Yu2 Junda Zhang2 Zuquan Song2 Xin Liu2 Chuan Wu1
1The University of Hong Kong 2ByteDance
速读
一句话结论 针对大型基础模型开发中频繁的并行配置变更与多框架兼容难题,本文提出了一个与并行策略无关的统一检查点系统 ByteCheckpoint,通过解耦元数据与张量数据并引入全栈异步 I/O 优化,在万卡规模下将检查点保存停顿最多缩短 161.50 倍,并实现了加载时的自动高效重新分片。
要解决什么问题 大型基础模型在预训练、微调及评估等阶段中,分配的 GPU 资源和并行配置会频繁变化。原有检查点机制将数据与特定并行策略深度耦合,跨阶段转移时必须依赖离线脚本重新分片。这不仅耗费人力维护复杂脚本,还导致目标任务在等待期间白白浪费昂贵的计算资源。此外,生产环境混合使用 Megatron-LM、FSDP 等框架及 HDFS 等存储后端,现有系统往往只针对单一框架设计,缺乏统一支持。更致命的是,当模型达千亿参数、集群扩展至万卡时,传统系统遭遇 I/O 与可扩展性瓶颈:朴素的集体通信屏障会导致初始化延迟甚至显存溢出,海量分布式读写请求会瞬间压垮存储主节点导致元数据超时,最终引发训练崩溃。
怎么做的 核心思路是将检查点表示形式与运行时并行策略彻底解耦,通过分层架构屏蔽底层差异。系统由四个部件构成:统一应用程序接口提供入口;规划器对接不同框架,提取张量分片规范并生成执行计划;执行引擎在各节点异步执行计划;存储 I/O 层封装后端读写逻辑。为实现加载时自动重新分片,方法引入与并行无关的张量分片元数据定义,核心不变量为索引元组:$$ ShardMeta = (fqn, nD\_offsets, nD\_lengths) $$ 其中 $fqn$ 是张量完全限定名,$nD\_offsets$ 和 $nD\_lengths$ 记录本地分片在全局多维形状中的偏移与长度。所有元数据整合到全局文件,建立逻辑分片到物理字节偏移的映射。并行配置改变时,节点查询该文件即可精准拉取数据。针对 ZeRO 优化中被展平且不规则切分的张量,系统采用张量分解策略,将其拆分为多个规则的索引元组,巧妙绕开了保存时昂贵的全局收集通信开销。在执行层面,系统构建完全异步流水线,将文件读取、反序列化、主机到设备拷贝及节点间通信重叠执行。为消除数据并行组间的冗余读取,规划器利用 Worst-Fit 算法在组间均匀分配读取任务,随后通过空闲节点间带宽将数据广播给对等节点。针对万卡规模挑战,系统摒弃传统 NCCL 通信后端,改用 gRPC 框架并构建基于树状层次的通信拓扑下发执行计划,彻底消除规划阶段的显存占用和超时问题;同时在存储侧引入无状态代理组件 NNProxy,专门应对海量元数据查询请求。
效果如何 实验在 32 到 8960 张 NVIDIA H800 和 A100 GPU 上展开,采用 InfiniBand 网络和 HDFS 存储。测试模型包括视频生成 vDiT 模型以及参数量达 4050 亿的文本生成 tGPT 模型。对比基线为代表 PyTorch 原生分布式检查点路线的 DCP,以及代表 Megatron 生态路线的 MCP。量化结果显示,在不同规模工作负载下,ByteCheckpoint 显著降低了运行时检查点停顿,相比基线减少 13.06 倍到 161.50 倍,在 8960 卡规模下平均停顿时间控制在 600 毫秒以内。在端到端保存和加载时间上,该系统分别实现平均 6.05 倍和 3.88 倍加速,最高加速比达 9.96 倍和 8.80 倍。对于耗时严重的重新分片场景,不规则张量分解技术将阻塞时间从 5.03 秒断崖式降至 0.20 秒。作者也坦承了方法的局限与代价:将不规则张量分解会轻微增加全局元数据文件大小,并在加载时引入额外查找拼接步骤;此外,数据加载器状态(如累积的令牌缓冲区)的收集仍会带来一定延迟,其耗时高度依赖于缓冲区大小,这也是未来需要进一步优化数据加载器设计的方向。
A1 主要贡献
本文针对大型基础模型(LFMs)的开发过程中,检查点保存以保留训练状态至关重要,用于在各种故障或GPU资源和并行配置变化时恢复训练。此外,保存的检查点被分发到评估任务或在不同训练阶段之间转移(如从预训练到后训练)。所有这些场景都需要将分布式检查点从一种并行配置重新分片到另一种。在生产环境中,不同的LFMs使用各种框架和存储后端进行训练,取决于模型规模和训练规模。需要一个高性能的检查点系统来实现整个LFM开发生命周期中高效的检查点管理。
本文引入ByteCheckpoint,一个工业级检查点系统,用于大规模LFM训练。ByteCheckpoint的特点包括:一个与并行无关的检查点表示形式,支持高效的加载时检查点重新分片;一个通用的检查点保存/加载工作流,以适应多种训练框架并支持不同的存储后端;全栈优化以确保高I/O效率和可扩展性;一套监控工具以简化大规模性能分析和瓶颈检测。与现有的开源检查点系统【52,Megatron Distributed Checkpoint, 2024, 无会议, https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/dist_checkpointing.html】【58,Getting started with Distributed Checkpoint (DCP), 2023, 无会议, https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html】相比,ByteCheckpoint显著减少了运行时检查点停顿,平均减少54.20×。在保存和加载时间上,ByteCheckpoint分别实现了高达9.96×和8.80×的改进。
核心问题包括:LFM训练的复杂性、多阶段开发、资源密集型特性以及大规模(例如DeepSeek-V3有6710亿参数,在14.8万亿token上预训练,使用多达12288个GPU)。检查点需要用于训练恢复、并发评估和跨阶段转移,但面临重新分片挑战、框架和存储后端的多样性以及I/O效率问题。现有的系统假设一致的并行性,无法处理重新分片,或仅支持特定框架,且I/O性能和可扩展性不足。
研究目标是设计高效统一的检查点管理系统,支持整个LFM开发生命周期,提供高效重新分片、通用的工作流和高性能I/O。
创新点包括:与并行无关的检查点表示形式,通过分离元数据和数值并整合到全局文件中,实现高效加载时重新分片;通用的工作流,通过为每个框架定制规划器生成统一计划,然后由无关引擎执行;全栈优化,如平衡计划生成、异步管道和监控工具。ByteCheckpoint部署在具有数万个GPU的工业AI平台上,支持各种LFM任务,与基线相比,检查点停顿减少12.13×到161.50×,保存和加载过程平均加速6.05×和3.88×。
A3 背景知识/关键Observation/设计原则
LFM开发生产管道。 如图1所示,LFM的开发包括预训练和后训练阶段。在初始预训练阶段,LFM在从多个来源收集的大量数据上迭代训练,以吸收世界知识。随后,持续预训练用于增强基础模型的能力。例如,大型语言模型(LLMs)的预训练通常涉及长上下文持续训练,以逐步增加LLMs支持的上下文长度。后训练用于将预训练模型与人类反馈对齐或增强模型的推理能力【18, Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, 2025, arXiv, arXiv:2501.12948】【57, Introducing OpenAI o1, 2024, 无会议, https://openai.com/o1/】。各种任务特定的标记数据集(例如多语言、代码、数学、推理等)参与微调LFM,随后是强化学习,通常包括奖励建模然后执行Proximal Policy Optimization【45, Proximal policy optimization algorithms, 2017, arXiv, arXiv:1707.06347】(PPO),或直接进行Direct Preference Optimization【41, Direct preference optimization: Your language model is secretly a reward model, 2024, NeurIPS, 无URL】(DPO)。由于这些数据集的规模减小,后训练中涉及的GPU相对较少。自动评估【14, Check-N-Run: A checkpointing system for training deep learning recommendation models, 2022, NSDI, 无URL】【22, Characterization of large language model development in the datacenter, 2024, NSDI, 无URL】被定期触发,以获取中间模型检查点并使用多样化标准评估质量。
检查点保存。 LFM训练作业的训练状态包括GPU和CPU状态。GPU状态是LFM模型中的可学习参数和优化器信息(例如模型的float32精度副本及其在Adam【26, Adam: A method for stochastic optimization, 2014, arXiv, arXiv:1412.6980】中的动量和方差)。在最先进的并行训练【28, Reducing activation recomputation in large transformer models, 2023, MLSys, 无URL】中,这些状态被分片并放置在多个GPU上。CPU状态包括数据加载器模块、随机数生成器(RNG)状态、全局训练步数和学习率调度器,所有存储在CPU内存中。我们的数据加载器模块包含一个令牌缓冲区,用于缓存从数据源读取的变长输入样本;当累积令牌数量达到上下文窗口大小【13, The Llama 3 herd of models, 2024, arXiv, arXiv:2407.21783】【25, MegaScale: Scaling large language model training to more than 10,000 GPUs, 2024, NSDI, 无URL】【69, GLM-130b: An open bilingual pre-trained model, 2022, arXiv, arXiv:2210.02414】【71, OPT: Open pre-trained transformer language models, 2022, arXiv, arXiv:2205.01068】时,数据加载器将所有缓存样本组装成一个批次(微批次)。由于GPU和CPU内存的易失性,这些训练状态应定期保存到持久存储中,以容忍任何故障并为评估任务准备。
存储后端。 在生产环境中,使用独立的分布式文件系统(例如用于Llama 3.1训练【13, The Llama 3 herd of models, 2024, arXiv, arXiv:2407.21783】的Tectonic【39, Facebook’s tectonic filesystem: Efficiency from exascale, 2021, FAST, 无URL】)来存储检查点【14, Check-N-Run: A checkpointing system for training deep learning recommendation models, 2022, NSDI, 无URL】【20, Unicron: Economizing self-healing llm training at scale, 2023, arXiv, arXiv:2401.00134】用于正式任务。鉴于训练期间各种硬件故障和软件错误不可避免【13, The Llama 3 herd of models, 2024, arXiv, arXiv:2407.21783】【22, Characterization of large language model development in the datacenter, 2024, NSDI, 无URL】【25, MegaScale: Scaling large language model training to more than 10,000 GPUs, 2024, NSDI, 无URL】,在不同全局训练步存储检查点是必要的,以保护训练。分布式文件系统(例如HDFS、NAS)提供足够的存储容量来容纳大型模型的多个检查点。
LFM开发生命周期中的检查点重新分片场景。 在LFM开发的生命周期中,由于不同场景下并行性的变化(图2),检查点重新分片始终是必需的:(1) 训练恢复。分配给LFM训练作业的GPU配额可能因移除故障机器【22, Characterization of large language model development in the datacenter, 2024, NSDI, 无URL】【25, MegaScale: Scaling large language model training to more than 10,000 GPUs, 2024, NSDI, 无URL】、添加从完成任务释放的新机器或运行在潮汐资源【12, Parcae: Proactive, Liveput-Optimized DNN training on preemptible instances, 2024, NSDI, 无URL】【29, Lyra: Elastic scheduling for deep learning clusters, 2023, EuroSys, 无URL】【68, AntMan: Dynamic scaling on GPU clusters for deep learning, 2020, OSDI, 无URL】【72, QSync: Quantization-minimized synchronous distributed training across hybrid devices, 2024, arXiv, arXiv:2407.02327】而变化。为了最大化资源利用,通常需要调整训练并行性以响应资源变化。此外,在大规模预训练开始时,AI工程师通常需要实验各种模型配置和优化技术【25, MegaScale: Scaling large language model training to more than 10,000 GPUs, 2024, NSDI, 无URL】,因为小规模剖析或模拟的结论并不总是转化为大规模训练中的最佳性能。在此阶段,调整并行配置是常见的,导致频繁的训练恢复。而且,长上下文训练改变上下文长度,这也需要GPU配额或并行调整。如图2示例所示,最初存储8个分布式检查点文件,然后在恢复时加载到6个训练工作者中。(2) 跨阶段转移。进入后训练时,由于后者训练数据减少,涉及的GPU数量通常减少。从预训练保存的检查点经常被重新分片以与每个后训练任务的具体工作负载对齐。如图2所示,后训练阶段的微调任务仅涉及4个GPU,因此检查点相应地被重新分片。(3) 评估。两个阶段的评估任务需要加载模型检查点并在单独资源上进行推理;它们的并行性需要调整以与用于特定数据集的分配GPU对齐。图2显示了一个4-GPU评估任务从预训练中重新分片模型检查点。我们收集了过去六个月在我们AI平台上检查点重新分片需求次数(训练文本、图像和视频生成模型),识别出1870个预训练恢复实例、13080个跨阶段重新配置实例和19844个运行评估任务实例。
加载时检查点重新分片。 作为LFM开发生命周期中的关键步骤,检查点重新分片的效率对于最小化额外开销至关重要。在我们的AI平台中,先前的常见实践是开发离线检查点重新分片脚本,并在每次出现新检查点重新分片场景时适应脚本。这种方法效率低下且劳动密集(详见附录A)。除了开发成本外,运行重新分片脚本导致GPU时间和资源的巨大浪费。表1展示了各种场景下执行重新分片作业的成本。在训练恢复或新评估任务开始之前,必须提前提交执行重新分片脚本的独立作业。这些作业从存储系统下载检查点,将分布式检查点重新分片到给定的并行配置,并将新检查点上传回存储系统。目标训练或评估作业直到重新分片作业完成才能执行,导致延长挂起时间。而且,由于重新分片脚本创建的检查点与特定并行性耦合,它们无法自由重用,从而增加存储开销。
表1: 不同场景下执行离线重新分片作业的平均完成时间。
表2: 我们平台上使用的三大训练框架。
多种框架和存储后端。 在我们的AI平台上,使用了广泛的训练框架,例如Megatron-LM【47, Megatron-LM: Training multi-billion parameter language models using model parallelism, 2019, arXiv, arXiv:1909.08053】、DDP【31, Pytorch distributed: Experiences on accelerating data parallel training, 2020, arXiv, arXiv:2006.15704】、FSDP【73, Pytorch FSDP: Experiences on scaling fully sharded data parallel, 2023, arXiv, arXiv:2304.11277】、veScale【25, MegaScale: Scaling large language model training to more than 10,000 GPUs, 2024, NSDI, 无URL】【63, veScale: A PyTorch Native LLM Training Framework, 2024, 无会议, https://github.com/volcengine/veScale】等。表2基于六个月的跟踪分析列出了平台上三大首选训练框架。用户通常采用Megatron-LM【47, Megatron-LM: Training multi-billion parameter language models using model parallelism, 2019, arXiv, arXiv:1909.08053】训练大型语言基础模型。FSDP【73, Pytorch FSDP: Experiences on scaling fully sharded data parallel, 2023, arXiv, arXiv:2304.11277】用于涉及文本到视频或文本到语音模型的训练任务,DDP【31, Pytorch distributed: Experiences on accelerating data parallel training, 2020, arXiv, arXiv:2006.15704】通常用于训练多模态基础模型的图像编码器组件或常规算法测试。此外,用户可以根据场景选择各种存储后端用于检查点持久化,从调试的本地磁盘到正式训练任务的HDFS或NAS。每个训练框架都有自己的检查点模块、文件格式和保存/加载逻辑的特定实现。然而,这些模块缺乏生产所需的关键特性,例如加载时重新分片、异步检查点和远程持久存储支持。为每个框架的检查点模块集成这些特性并定制优化实现需要重复的工程努力。这导致跨训练框架和存储后端的检查点接口不一致,复杂化代码库。而且,不同框架的多样检查点文件格式的维护增加了跨训练阶段实现检查点转移逻辑和部署模型用于评估和推理任务的复杂性。因此,提供针对不同训练框架和存储后端的通用工作流至关重要。
高效且可扩展的I/O性能。 主流LFM具有庞大的模型规模,达到数千亿参数【13, The Llama 3 herd of models, 2024, arXiv, arXiv:2407.21783】。因此,训练状态的大小也显著增加,对检查点保存和加载施加了大量开销。通过分析我们先前的LFM训练作业,我们观察到将训练在4096个GPU上的GPT 175B模型的检查点保存到HDFS的平均端到端时间可能为200秒。这一持续时间大大超过单个训练迭代所需的时间。尽管这种耗时过程可以通过采用异步检查点【25, MegaScale: Scaling large language model training to more than 10,000 GPUs, 2024, NSDI, 无URL】【34, DataStates-LLM: Lazy asynchronous checkpointing for large language models, 2024, arXiv, arXiv:2406.10707】【35, CheckFreq: Frequent, fine-Grained DNN checkpointing, 2021, FAST, 无URL】【65, Reliable and efficient in-memory fault tolerance of large language model pretraining, 2023, arXiv, arXiv:2310.12670】部分从模型训练的关键路径中移除,但加速端到端检查点保存对于最小化大规模训练中不可避免的频繁故障【25, MegaScale: Scaling large language model training to more than 10,000 GPUs, 2024, NSDI, 无URL】引起的训练进度损失【66, Gemini: Fast failure recovery in distributed training with in-memory checkpoints, 2023, SOSP, 无URL】仍然至关重要。如图3所示,尽管检查点与训练重叠,但其快速完成允许在故障发生前存储更多中间检查点,从而从更近的状态恢复并改善ETTR。此外,评估任务在训练期间被触发,中间检查点被定期拉取用于这些任务。更快的检查点保存确保它们的及时执行,减少由于在远程持久存储中准备这些检查点的阻塞时间。
检查点系统的可扩展性。 扩展检查点系统同时保持高I/O性能是非平凡的。小规模设置中难以检测的次优和风险设计可能导致大规模训练中的严重性能瓶颈或甚至灾难性作业故障。例如,从训练集群到存储系统(例如HDFS)的检查点文件的大量读/写请求可能过载主节点,导致文件元数据操作延迟。此外,集体通信的朴素实现,例如完整性检查屏障,引入了大量的初始化和同步开销。这些开销甚至可能导致通信超时,最终导致整个训练作业失败。而且,随着训练规模扩大,在多个机器上的训练和I/O工作者之间高效分析系统性能和检测错误变得越来越具有挑战性。
现有检查点系统。 现有的检查点系统,例如CheckFreq【35, CheckFreq: Frequent, fine-Grained DNN checkpointing, 2021, FAST, 无URL】、Check-N-Run【14, Check-N-Run: A checkpointing system for training deep learning recommendation models, 2022, NSDI, 无URL】和Gemini【66, Gemini: Fast failure recovery in distributed training with in-memory checkpoints, 2023, SOSP, 无URL】,假设一致的并行性,并不解决检查点重新分片的需求。DCP【58, Getting started with Distributed Checkpoint (DCP), 2023, 无会议, https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html】和MCP【52, Dist checkpointing package, 2024, 无会议, https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/dist_checkpointing.html】包含检查点重新分片能力,但受限于支持的并行策略和训练框架、I/O性能和可扩展性。相关工作的全面讨论在附录F中。在ByteCheckpoint中,我们设计了分离的存储表示形式用于高效加载时重新分片,提出了通用保存/加载工作流以支持多种框架和存储后端,集成了全栈优化以提升I/O性能,并分享了扩展检查点系统以支持真实世界LFM训练的经验。
ByteCheckpoint的设计原则。 基于上述观察,我们根据以下关键原则设计ByteCheckpoint:(1) 解耦。检查点表示形式独立于特定运行时并行性。训练框架和存储后端的接口与核心执行引擎分离,确保强大的可扩展性。(2) 用户友好性。API应简洁,使AI研究人员和工程师无缝集成到他们的代码和运行时环境中。
A2 方法细节
统一架构概述。 ByteCheckpoint的架构如图4所示。每个组件详细介绍如下:API。ByteCheckpoint的API(bytecheckpoint.save和bytecheckpoint.load)为各种训练框架的用户代码提供统一的入口点。例如,要保存检查点,用户首先准备相应的训练状态、检查点路径、框架名称和性能相关选项,然后调用bytecheckpoint.save。这个高层入口点抽象了底层系统复杂性,例如分片规范、保存/重新分片计划生成和I/O操作。规划器。规划器作为训练框架的接口。它从API层接收参数(训练状态、检查点路径等),基于工作者的rank和框架特定分片规范(如Megatron ShardedTensor或FSDP DTensor)为每个张量分片创建ShardMeta(第3.2节),并确定每个工作者的保存/加载张量和其他状态。每个工作者利用规划器最初创建本地计划,随后协作创建全局计划。我们为每个训练框架实现了定制的规划器,以从这些规范中提取信息并生成计划。执行引擎。引擎在每个训练工作者上运行,当调用相应API时执行规划器生成的保存/加载计划。它分析给定的检查点路径以确定适当的存储后端,然后与存储I/O层交互以执行计划中指定的I/O任务。存储I/O。与规划器层的设计类似,存储I/O层封装不同的存储后端并管理后端特定的读/写操作和优化。引擎层和存储I/O层之间的接口在不同存储后端中保持统一,便于无缝集成新存储后端。ByteCheckpoint支持多种存储选项,包括内存中检查点存储【66, Gemini: Fast failure recovery in distributed training with in-memory checkpoints, 2023, SOSP, 无URL】、本地磁盘存储和远程存储系统。图5展示了ByteCheckpoint的典型用例。最初,用户定义一个字典指定要保存或加载的状态,包括模型、优化器、数据加载器和其他状态。在训练恢复、阶段转移或评估开始时,调用bytecheckpoint.load来加载保存的检查点。当并行性变化时,检查点重新分片在加载期间自动发生。在训练期间,定期调用bytecheckpoint.save来保存检查点。
模型和优化器状态的解耦检查点表示。 为了实现与并行无关的检查点保存以启用自动加载时重新分片,我们开发了分布式训练中模型和优化器张量的正式规范。我们的设计符合PyTorch Distributed库【31, Pytorch distributed: Experiences on accelerating data parallel training, 2020, arXiv, arXiv:2006.15704】中的Distributed Tensor和Checkpoint概念。每个张量由“完全限定名称”(FQN)唯一标识,并具有全局形状,表示其分片前的原始形状。张量可以跨rank(即训练工作者)分片或复制。对于分片张量,工作者持有的特定分片由三个因素决定:应用于张量的并行性(例如张量并行)、分片维度以及相应并行组中的组rank。基于张量的FQN、全局形状和分片规范,ByteCheckpoint为每个rank创建张量分片元数据(ShardMeta),用于检查点存储中表示张量分片,独立于并行性。更精确地说,张量分片的ShardMeta是一个索引元组(fqn, nD_offsets, nD_lengths),其中nD_offsets和nD_lengths指示本地分片在其全局形状的多维轴上的偏移和长度。
数据加载器和其他状态。 数据加载器状态可分为两类:复制状态和分片状态。复制状态包括数据读取工作者的数量、源数据集路径和采样比率,并在不同rank的所有I/O工作者(子进程)中相同。分片状态对每个I/O工作者独特,包括令牌缓冲区和不同数据源的数据检索偏移。在ByteCheckpoint中,分片状态保存在单个文件中,而复制状态仅由全局rank为0的训练工作者保存。这减少了要保存的数据加载器状态的总体大小,并促进当并行配置变化时的数据加载器重新分片,因为将数据状态分离到不同的文件中简化了根据新并行性的状态合并和重新分配。对于其他额外状态,如RNG状态,我们在将它们转储到存储之前打包并序列化为一个紧凑的字节对象。
检查点表示形式。 图6展示了ByteCheckpoint中的检查点表示形式。分布式检查点包括一个全局元数据文件和多个存储文件。每个rank生成三个不同的文件:模型状态文件、优化器状态文件和额外状态文件。数据加载器状态文件仅由除DP度外的所有并行度rank为0的训练工作者生成。对于模型和优化器状态中的张量分片,它们的元数据包括三部分:BasicMeta,记录单个张量分片的基本信息,如步幅和设备,对恢复运行时状态至关重要;ShardMeta,如前所述,记录分片在完整张量中的相对位置信息;ByteMeta,指定每个张量分片在存储文件中的字节起始偏移和长度。所有张量元数据整合到全局元数据文件中,并基于每个张量分片的ShardMeta、BasicMeta和ByteMeta,建立保存张量分片与存储文件之间的映射,称为TensorShardtoBasicByteMap,确保准确的数据检索。全局元数据文件还包括LoaderShardtoByteMap,记录每个数据加载器中分片状态的文件索引信息。所有存储文件和全局元数据文件存储在检查点路径指定的存储后端中。
分解不规则张量。 我们将不规则张量定义为在分片前展平后无法重塑回原始维度的张量。这些张量通常在应用Zero Redundancy Optimizer(ZeRO)【42, ZERO: Memory optimizations toward training trillion parameter models, 2020, SC, 无URL】时出现,如在Megatron-LM ZeRO2和FSDP ZeRO3【73, Pytorch FSDP: Experiences on scaling fully sharded data parallel, 2023, arXiv, arXiv:2304.11277】中实现。在这些实现中,DP组内的优化器状态被展平、连接并分片。因此,结果1D张量切片通常无法使用n维形状和偏移直接表示。图7展示了不规则张量分片的示例:与张量A和C不同,具有原始形状(3,2)的张量B被均匀分割成两个分片。每个分片包含三个元素,无法直接表示为二维形状与相应偏移。解决不规则张量分片挑战的一个直观方法是在保存检查点前将所有张量分片合并成完整张量,从而简化ShardMeta的生成。例如,为了消除DCP【58, Getting started with Distributed Checkpoint (DCP), 2023, 无会议, https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html】中的潜在不规则张量,FSDP执行同步all-gather通信操作,与每个张量分片的D2H复制操作交错,而不管分片是否不规则分片。然而,这种方法会产生显著的通信开销,并需要GPU和CPU之间的频繁同步,大大阻碍效率。为了缓解与合并张量分片相关的开销,ByteCheckpoint采用张量分解策略来管理不规则张量分片。具体来说,ByteCheckpoint将不规则张量分解成一系列规则张量,每个用索引元组表示。例如,考虑图7中rank 0上具有三个元素的不规则张量B分片。ByteCheckpoint将此分片分解成两个可由nD_offsets和nD_lengths直接表示的规则张量。这种方法允许ByteCheckpoint使用多个ShardMeta条目表示单个不规则张量分片。尽管这种分解略微增加了元数据大小并在加载过程中添加了步骤,因为重建目标张量可能需要查找不规则分片的多个较小段,但它避免了张量分片的昂贵通信,而没有额外的保存阻塞时间。
检查点重新分片工作流。 ByteCheckpoint实现了通用的保存和加载工作流,在加载过程中自动重新分片检查点。利用引擎层、规划器层和存储I/O层之间的独立性,我们可以在不同训练框架和存储后端执行一致的保存和加载步骤。我们以张量重新分片为例(如图8),强调我们的检查点表示形式如何启用灵活的加载时重新分片。没有重新分片的检查点保存和加载过程遵循类似程序。步骤1。为了启动加载时检查点重新分片,每个rank调用bytecheckpoint.load() API,指定检查点路径以及要恢复的模型/优化器。然后,所有rank从路径加载全局元数据文件。步骤2。对于给定模型/优化器中的每个张量分片,每个rank查询全局元数据文件中的TensorShardToBasicByteMap,识别保存张量分片与新分片规范之间的匹配段。在识别这些匹配后,规划器构建本地加载计划,包括针对特定分片的BasicMeta和ByteMeta。此识别机制如图8底部所示。步骤3。协调规划器,通常位于rank 0,启动gather操作以聚合所有rank的加载计划。然后,它通过应用冗余消除优化(第4.1节)优化每个本地计划,以分布张量分片加载工作负载以减少完成时间。步骤4。协调器启动scatter操作以分布最终加载计划。每个rank然后从协调器接收其最终加载计划。步骤5。每个rank上的执行引擎基于检查点路径选择存储I/O层中的存储后端包装器,然后执行加载管道(第4.2节)。步骤6。在加载管道完成后,每个rank利用优化的异步集体屏障原语确保分布式加载的原子性。进一步细节可在附录B中找到。
数据加载器重新分片。 与图8中的张量重新分片类似,ByteCheckpoint通过查询LoaderShardtoByteMap重新分片数据加载器。数据加载器重新分片的示意图如图9所示。当并行配置变化时,对于数据加载器检查点,数据加载器检查点中的常见项可以直接加载,而独特项如累积令牌缓冲区和数据检索偏移需要重新分片。具体来说,当DP度大小保持不变而其他并行度改变时(图9示例中TP度变化),令牌缓冲区应复制到目标工作者以实现位精确恢复。当DP度大小变化时,令牌缓冲区必须相应地拆分或合并,以确保恢复的数据加载器不丢弃缓存数据且不重新训练已采样和输入的数据。多亏了为数据加载器模块设计的分离存储策略,ByteCheckpoint可以精确识别需要重新分片的独特项并高效处理它们。
性能优化概述。 我们现在详细说明ByteCheckpoint的性能优化技术,重点是最小化与检查点保存和加载相关的开销。
平衡保存工作负载。 在采用数据并行(DP)的训练场景中,模型状态在所有DP组中复制,导致模型状态重复。现有的检查点系统【52, Dist checkpointing package, 2024, 无会议, https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/dist_checkpointing.html】【58, Getting started with Distributed Checkpoint (DCP), 2023, 无会议, https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html】通过指定第一个DP组保存所有模型状态来解决此问题。然而,这种方法导致工作负载不平衡,可能使第一个DP组中的训练工作者成为落后者。为了解决这一挑战,我们在规划过程中实现了工作负载平衡去重机制,利用Worst-Fit算法。具体来说,协调规划器基于每个张量分片的大小分布保存工作负载,将当前张量分片分配给累积张量分片大小最小的rank。这种方法确保了跨rank的保存工作负载更公平分布。这平衡了所有工作者的工作负载,提高了保存效率。
消除冗余加载。 当加载检查点到包含数据并行的并行配置时,ByteCheckpoint通过消除DP组间的重复张量读取来优化过程,有效地将存储文件读取与张量传输结合。一个关键观察是,在文件读取过程中,可以利用空闲的GPU间带宽并发传输加载的张量到对等GPU。如图10所示,在规划阶段,DP组间的张量读取工作负载均匀分布在训练工作者中,从而避免读取重复。随后,启动I/O线程读取分配的张量分片。同时,在主线程中,首先将读取到CPU内存的分片复制到GPU内存,然后使用all-to-all集体通信传输到需要它们的其他工作者。
计划和元数据缓存。 在大规模训练中,规划过程的执行可能引入显著的通信开销,特别是当检查点保存频繁发生时。例如,为分布在8960个GPU上的405B transformer模型规划保存过程需要62秒。然而,我们观察到保存计划和全局元数据文件,虽然与特定并行性耦合,但在单个训练会话中保持不变。这允许缓存策略减少规划时间和后续检查点操作的相关开销。我们引入了计划元数据缓存,将规划转化为一次性成本。一旦首次建立,保存计划和全局元数据文件被缓存以供未来重用,消除重复规划。
完全异步引擎管道。 ByteCheckpoint引擎通过管道化优化检查点保存和加载(重新分片)期间的操作执行。以加载为例(如图10),我们为每个张量分片管道化文件读取、反序列化、主机到设备(H2D)复制和GPU间通信,从而实现更高效率。读取操作从存储系统下载包含所需张量的文件并放置到共享内存(例如dev/shm目录)。反序列化操作从共享内存反序列化张量。ByteCheckpoint使用多个线程并行文件下载和反序列化。H2D复制操作将张量从CPU内存传输到GPU,而All2All促进DP组内的张量传输。对于保存,我们实现了对称的完全异步管道,包括D2H复制、序列化和文件上传操作。为了缓解D2H复制对训练的性能影响,我们采用固定CPU内存池结合Ping-Pong缓冲机制来加速此操作。我们运行多个并行进程来序列化张量并将文件转储到共享内存。上传线程主动监控并在转储阶段完成后启动文件上传。
高性能读/写。 我们通过针对常用存储后端的优化I/O使用来优化检查点的读写吞吐量。例如,虽然HDFS并非主要设计用于随机数据访问,但它通过其SDK提供一些随机读取能力,允许应用程序访问特定文件偏移并从这些位置检索数据。我们利用此特性并启用单个文件的多线程读取,大大加速从HDFS下载检查点文件。在我们的生产平台中,单文件读取速度从400 MB/s提高到2-3 GB/s,所有测试都在配备超过100个CPU核心、TB级内存和200Gbps NIC的H800服务器上进行。对于文件写入,HDFS的仅追加写入使得基于偏移将单个文件分割成多个部分进行多线程写入不切实际。为了克服这一限制,我们将目标文件分割成几个固定大小的子文件,并使用多线程并发将它们写入HDFS。上传完成后,我们执行元数据级连接以无缝将子文件合并回单个实体,确保存储数据块的完整性。在没有网络拥塞的情况下,单文件上传速度可达到3 GB/s,远远超过单个HDFS客户端的平均读/写吞吐量(低于100 MB/s【48, HDFS Scalability: The limits to growth, 2010, USENIX Magazine, 无URL】)。
预取数据加载器状态。 每个数据加载器采用多个子进程,称为读取工作者,来处理数据加载和预处理。当启动检查点保存时,主读取工作者向所有其他读取工作者发出信号以准备它们的状态。训练过程暂停直到收集所有状态以确保准确性,因为主进程中的任何更新都会导致工作者状态变化。此阻塞期的持续时间取决于工作者数量和累积令牌缓冲区的大小。例如,当使用配置有4个工作者的数据加载器且总状态大小约为1GB时,状态收集过程通常需要约8秒。为了缓解此开销,我们采用预取。根据预设的检查点频率,每个读取工作者在检查点保存前的训练步骤中准备其状态并将其放入状态队列。在检查点保存步骤中,主读取工作者通过队列轮询立即收集这些准备好的状态。此优化允许ByteCheckpoint实现近零数据加载器状态收集延迟。
大规模检查点概述。 本节深入探讨了为大规模LFM训练优化检查点保存的方法。
高吞吐量、可扩展存储系统。 HDFS是ByteCheckpoint的主要存储后端,我们实现了几个优化来增强大规模检查点保存的吞吐量和可扩展性。我们用C++重写了所有HDFS组件,包括NameNode、DataNode和SDK【15, Hadoop Distributed File System, 2024, 无会议, https://hadoop.apache.org/docs/current/hadoop-project-dist/hadoop-hdfs/HdfsDesign.html】,与原始基于Java的实现相比,性能有效翻倍。此外,我们通过引入名为NNProxy【1, HDFS Federation Solution: NameNodeProxy, 2016, 无会议, https://github.com/bytedance/nnproxy】的新组件优化了其架构。NNProxy作为HDFS NameNode的无状态RPC代理,促进NameNode的大规模联邦部署,同时保持最小查询延迟。NNProxy设计解决了HDFS NameNode元数据请求的QPS瓶颈,这一挑战因分布式检查点的大量而加剧。NNProxy还提供额外特性,如认证、速率限制和元数据查询缓存。这些能力启用生产环境中检查点的更细粒度管理。我们优化的HDFS实现了数百PB的巨大容量、10 TB/s读/写带宽和约100,000元数据QPS。
检查点冷却策略。 我们使用SSD和HDD存储服务器的组合实现了两级热冷存储架构。我们的关键观察是,新存储的检查点文件通常在创建后不久被评估任务下载。然而,在没有训练异常(如损失峰值)的情况下,这些文件的访问频率在被评估任务下载后显著下降。尽管如此,所有检查点文件必须保留以用于追溯。为了高效管理存储空间同时确保数据可用性,我们开发了数据冷却机制,将数据从SSD迁移到HDD存储,从而确保高性能热存储始终有足够空间用于操作。具体来说,我们基于最后修改时间冷却超过保留阈值的所有文件。然后,通过纯元数据操作将原始文件路径重新映射到新的HDD存储位置。此策略适用于HDFS目录,并保留冷却文件的原始访问路径,提供无缝用户体验。
集体通信。 集体通信(例如scatter、gather、barrier)在ByteCheckpoint的工作流(第3.3节)中至关重要,特别是用于规划和检查点完整性保证。我们最初使用NCCL【54, Nvidia collective communications library (nccl), 2024, 无会议, https://developer.nvidia.com/nccl】作为通信后端来在协调器执行gather和scatter操作。然而,当将预训练任务扩展到8960个GPU时,我们观察到NCCL在ByteCheckpoint保存工作流的规划阶段需要长时间懒惰构建通信通道和分配GPU内存。在某些情况下,它变得无响应或导致CUDA内存不足(OOM)错误,因为scatter或gather操作需要与每个GPU建立点对点通信。这些GPU OOM问题和长初始化时间在小规模试验中不明显,但在更大规模时变得显著。为了增强规划期间的通信稳定性,我们使用gRPC框架【17, grpc: A high performance, open source universal rpc framework get started!, 2024, 无会议, https://grpc.io/】重新实现了过程,这消除了规划期间的GPU内存使用。然而,当将训练扩展到数万个GPU时,集中的gather和scatter操作继续对协调器施加显著负担,导致通信失败。我们通过实现基于树的层次通信拓扑进一步改进了其稳定性。单个机器上的训练工作者组织成一级子树,本地rank 0的工作者指定为根。对于机器间通信,我们迭代分组多个机器,将每个组中全局rank最低的工作者指定为根。此过程继续直到所有工作者集成到收敛于全局根(即协调器)的层次结构中。在大规模3D并行训练场景中,自然形成TP-DP-PP通信树,移除额外连接。
监控和分析。 在训练期间,ByteCheckpoint持续收集关键性能测量并可视化它们,用于实时性能监控和分析。这种方法启用快速检测检查点问题,例如低读/写吞吐量、落后者和上传失败及重试。数据收集。我们基于Python的上下文管理器和装饰器语法设计了用户友好的指标系统,以灵活监控关键过程。它自动捕获每个操作的持续时间和I/O大小,以及相关元数据,如每个工作者的rank、文件路径和当前步骤。所有收集的指标通过后台消息队列传输到远程数据库。可视化。ByteCheckpoint为用户提供与不同检查点阶段相关的所有rank的全面拓扑性能概述,以及任何特定rank的详细持续时间分解。图11展示了3D并行训练拓扑中检查点保存时间的示例热图可视化。可视化突出显示rank 0、4、8和12经历最长的保存时间,因为它们的检查点包括数据加载器状态。额外的指标可视化也可用于详细分析,例如阻塞时间和细粒度阶段如规划和D2H复制。这些可视化使用户能够轻松定位系统开发和生产环境中的落后节点或阶段。例如,如果某些节点遇到网络问题,HDFS传输阶段的上传或下载时间增加将显而易见。而且,通过热图拓扑概述,可以访问每个rank的检查点过程的详细时间线分解,使用户能够彻底评估系统中的所有优化。图12展示了特定rank的每个检查点保存阶段的执行细节。存储侧监控。在存储客户端侧,我们监控每个原子读/写操作的延迟和I/O大小在I/O块级别。所有收集的指标定期通过消息队列传输到ClickHouse服务,在那里数据被聚合和分析。意外的高延迟或低带宽触发工程师的警报以进一步调查。在存储集群侧,我们的主要焦点是整体性能指标,包括元数据请求QPS、集群级读/写吞吐量和存储容量利用率。这个全面监控系统在识别问题如慢检查点读写,甚至由于存储容量耗尽导致的系统不可用性中发挥关键作用,促进实施预防措施以提升系统可靠性和效率。
A4 实验环境
数据集名称、规模及用途。 采用两种基于transformer【62, Attention is all you need, 2017, NeurIPS, 无URL】的结构:DiT【40, Scalable diffusion models with transformers, 2023, ICCV, 无URL】和GPT-3【9, Language models are few-shot learners, 2020, NeurIPS, 无URL】,实现vDiT和tGPT模型。vDiT用于视频生成微调,tGPT用于文本生成训练。训练超过500步,每100步保存检查点。数据集规模未具体量化,但用于预训练和后训练阶段的评估。
模型架构关键参数。 vDiT:基于DiT,FSDP框架。tGPT:基于GPT-3,模型大小包括13B、30B、70B和更大变体如175B和405B。并行配置详见表3,包括源和目标GPU数量、TP、DP、PP度。
硬件配置。 NVIDIA A100 80GB GPU用于vDiT微调,NVIDIA H800 80GB GPU用于tGPT训练和大规模部署。GPU数量从32到8960。机器通过InfiniBand互连。H800服务器配备超过100个CPU核心、TB级内存和200Gbps NIC。
软件配置。 基于PyTorch Distributed【31, Pytorch distributed: Experiences on accelerating data parallel training, 2020, arXiv, arXiv:2006.15704】,框架包括Megatron-LM【47, Megatron-LM: Training multi-billion parameter language models using model parallelism, 2019, arXiv, arXiv:1909.08053】、FSDP【73, Pytorch FSDP: Experiences on scaling fully sharded data parallel, 2023, arXiv, arXiv:2304.11277】。存储后端:HDFS。实现约20,000行Python代码,基于DCP【58, Getting started with Distributed Checkpoint (DCP), 2023, 无会议, https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html】(commit hash: 80c07df)。
A4 实验结果
检查点效率实验。 比较ByteCheckpoint与基线DCP【58】和MCP【52】在不同工作负载下的性能。工作负载:vDiT在A100 GPU上微调,tGPT在H800 GPU上训练(表3配置)。保存性能:ByteCheckpoint将检查点停顿减少13.06×到161.50×,端到端保存时间平均加速6.05×,在更大规模下加速更多(表4)。加载和重新分片性能:标准加载改进1.41×到8.80×,重新分片平均加速3.64×。包括CPU状态(如数据加载器)增加时间,由于令牌缓冲区大小(表4)。端到端ETTR:改进1.16×到1.29×(表4)。
微基准实验。 保存优化:异步管道、工作负载平衡和计划缓存将tGPT 13B和30B的保存时间从48.3s减少到19.27s(表5)。加载优化:异步管道和读取-通信重叠将加载时间缩短1.53×和1.58×(表6)。重新分片优化:不规则张量分解将阻塞时间从FSDP的5.03s减少到0.20s,加速25.15×(表7)。
正确性验证实验。 tGPT 13B的归一化训练损失曲线在PP、TP、DP和混合重新分片后平滑匹配(图13、图16)。位精确恢复:175B模型的损失曲线在恢复后相同(图14)。数据加载器:归一化样本长度曲线在多次重启后相同(图17)。
真实LFM生产中的检查点。 在H800 GPU上训练图像和文本生成LFM,规模到8960 GPU,平均停顿<600ms,端到端几秒内完成(表8)。解决了数据加载器落后者和HDFS元数据瓶颈问题。
A5 结论
ByteCheckpoint是一个生产级检查点系统,用于LFM开发。它倡导统一的检查点表示形式和工作流架构,支持高效加载时重新分片,并兼容多种训练框架和存储后端。它实现了全栈I/O性能和可扩展性优化,并集成了高效的性能监控和分析工具,用于大规模LFM开发。与最先进的基线【52】【58】相比,ByteCheckpoint将检查点停顿减少高达161.50×,端到端检查点完成时间缩短高达9.96×。此外,加载(重新分片)过程平均加速3.88×。本文相信该系统不仅为构建真实世界LFM开发检查点系统的人员提供宝贵实践经验,还为社区未来的研究提供深刻洞见。未来工作包括设计更高效的数据加载器。
A6 附录
维护重新分片脚本的更多背景。 离线重新分片与定制脚本的示例如图15所示。定制离线脚本是劳动密集的。对于GPU状态重新分片,离线脚本必须覆盖LFM及其优化器中的所有不同组件,并适应每个组件在不同并行策略下的多样行为。例如,在张量并行(TP)下,注意和MLP块中的GEMM操作者沿不同维度分片,而其他操作者如LayerNorm【8, Layer normalization, 2016, arXiv, arXiv:1607.06450】在GPU间复制。当采用混合3D并行【36, Efficient large-scale language model training on gpu clusters using MegatronLM, 2021, SC, 无URL】时,分布式优化器中一层(模块)的TP分片张量首先展平然后合并,然后根据指定的数据并行(DP)度分片。离线脚本必须实现与模型(优化器)组件和并行策略组合紧密耦合的重新分片逻辑。此外,特殊算法优化技术如GQA【23, GQA: A new dataset for real-world visual reasoning and compositional question answering, 2019, CVPR, 无URL】和MLA【75, DeepSeek-Coder-V2: Breaking the Barrier of Closed-Source Models in Code Intelligence, 2024, arXiv, arXiv:2406.11931】改变了某些操作者(例如注意块中的查询-键-值投影GEMM操作者)的张量布局,需要相应的重新分片支持。为了处理我们生产环境中的各种情况,我们最大的脚本甚至包括3193行Python代码。这种复杂性导致开发和维护的显著工程努力。
高效完整性保证。 完整的检查点由不同工作者存储的多个文件组成。任何单个工作者的失败都可能破坏整个检查点。为了防止此类问题,屏障机制对于实现所有训练工作者间的原子保存/加载操作至关重要。训练框架如Megatron-LM【47】中的检查点模块依赖torch.distributed中的屏障函数来执行完整性检查。这种方法同步训练工作者以确保所有检查点保存/加载操作完成。我们观察到,当将训练扩展到涉及约10,000个GPU时,这种行为每次导致约20秒的停顿。为了解决这种低效,我们使用上述方法(gRPC与基于树的通信拓扑)重新实现了屏障函数,并异步执行完整性检查,有效消除阻塞时间。我们还在ByteCheckpoint的I/O工作者中集成了上传/下载重试机制,并集成了失败日志,记录未能完成检查点任务的工作者中检查点保存/加载管道的精确失败阶段。
平均ETTR计算。 训练期间发生的故障引入进度损失和最后检查点加载开销。假设故障在每个检查点间隔内均匀分布【66】,最佳情况是故障刚好在检查点端到端保存过程完成后发生,而最坏情况是刚好之前。给定每迭代训练时间$T_{iter}$、检查点间隔$N$、端到端检查点保存时间$T_{save}$和加载(重新分片)时间$T_{load}$,我们推导出平均浪费时间$T_{wasted}$为:
$T_{wasted} = \frac{1}{2} (N \cdot T_{iter} + T_{save}) + T_{load}$
因此,平均ETTR为:
$ETTR = \frac{N \cdot T_{iter}}{N \cdot T_{iter} + T_{wasted}}$
表4中呈现的端到端ETTR结果是标准加载和重新分片设置的平均值。
检查点开销分解。 我们将rank 0的端到端检查点保存时间($T_{save}$)分解成几个阶段,并调查每个部分的开销。结果如表9所示,其中$T_{FirstPlan}$表示初始规划成本,而$T_{CachePlan}$表示后续检查点操作的缓存。我们发现,随着训练规模扩大,规划的通信开销增加,导致显著停顿。多亏了缓存策略(第4.1节),它成为每个(恢复)训练会话的一次性开销。此外,采用固定内存CPU池使D2H的阻塞时间几乎可以忽略,而我们的异步引擎管道重叠了序列化、共享内存转储和HDFS上传的执行时间,减少了端到端时间。而且,负载平衡机制利用每个DP组内的并行上传能力,并在更大训练规模下实现更多性能提升(例如,4800 GPU的模型状态上传速度比2400 GPU快3.03×)。
表9: rank 0检查点保存过程的详细开销分解。
更多重新分片正确性实验。 DP和混合重新分片的归一化损失曲线如图16所示。在DP和混合重新分片的情况下,由于我们也增加了全局批次大小,重新分片后的损失曲线下降更快。我们进一步展示了ByteCheckpoint在训练恢复时实现数据加载器状态的位精确对齐能力。由于RNG状态是固定的,正确恢复应产生相同的数据采样轨迹。因此,我们使用归一化数据样本长度曲线(图17)进行评估。如突出显示,多次重启前后归一化数据样本长度相同。
相关工作。 检查点框架。一些工业举措专注于为深度学习开发检查点系统。在DCP【58】之前,PyTorch提供了torch.save和torch.load API用于本地检查点管理,没有原生重新分片支持。DCP【58】为FSDP引入了重新分片能力,但缺乏对并行策略如TP和PP的支持。DeepSpeed-UCP【32, Universal checkpointing: Efficient and flexible checkpointing for large scale distributed training, 2024, arXiv, arXiv:2406.18820】为DeepSpeed检查点提供统一格式,并通过离线脚本提供重新分片能力。Megatron MCP【52】基于DCP【58】的工作流,并将存储选项扩展到格式如Zarr【60, Zarr: chunked, compressed, n-dimensional arrays, 2024, 无会议, https://zarr.dev/】。所有这些框架对各种并行策略和训练框架的支持有限,且它们的I/O性能在大规模训练中扩展性不佳。检查点优化。几个工作【10, Project adam: Building an efficient and scalable deep learning training system, 2014, OSDI, 无URL】【37, Deepfreeze: Towards scalable asynchronous checkpointing of deep learning models, 2020, CCGRID, 无URL】从不同角度调查了减少检查点成本。Check-N-Run【14】,专为推荐模型设计,采用差异检查点仅存储模型的修改部分,以及量化以减少检查点大小。CheckFreq【35】将快照和保存操作与计算管道化以最小化检查点停顿,并引入在线算法调整检查点频率以进一步降低成本。Gemini【66】倡导内存中检查点与机器间备份用于快速恢复,将检查点通信与训练流量交织以启用每迭代频繁检查点。JIT-Checkpointing【19, JustIn-Time Checkpointing: Low cost error recovery from deep learning training failures, 2024, EuroSys, 无URL】采用即时检查点用于低成本错误恢复。在工业用例中,将检查点存储在单独的持久存储中对于各种任务如自动评估【22】、超参数调优和模型调试至关重要。ServerlessLLM【16, ServerlessLLM: Low-Latency serverless inference for large language models, 2024, OSDI, 无URL】专注于加载优化的需求,这在无服务器推理场景中至关重要。它提出基于块的多级加载管道来加速检查点加载。与现有解决方案不同,ByteCheckpoint倡导优化的I/O性能和灵活的加载时检查点重新分片,支持通用LFM开发。检查点表示形式。PyTorch的torch.save/load功能依赖pickle进行序列化和反序列化。此格式缺乏关键张量分片元数据,如全局形状信息,阻止自动重新分片。DCP【58】通过引入分离格式分离元数据与张量数据来解决此限制。此元数据包括全局形状和偏移细节,在DCP中启用自动重新分片。Tenplex【64, Tenplex: Dynamic parallelism for deep learning using parallelizable tensor collections, 2024, SOSP, 无URL】引入Parallelizable Tensor Collection(PTC)以灵活表示和转换跨多样并行配置的张量状态。为了构建PyTorch原生检查点系统,ByteCheckpoint采用DCP的表示形式,并包含必要的适应以处理不规则张量分片。基于阵列的存储系统如Zarr【4, Zarr, 2024, 无会议, https://zarr.dev/】和TensorStore【3, TensorStore, 2024, 无会议, https://google.github.io/tensorstore/】允许张量保存为单个阵列,支持并发读写。MCP【52】使用Zarr格式支持分布式检查点。对于安全高效的张量存储,Safetensors【2, Safetensors, 2024, 无会议, https://huggingface.co/docs/safetensors/index】是一个文件格式,用于安全保存张量并高效加载它们。为了提高与Hugging Face开源生态系统的兼容性,ByteCheckpoint集成了导出检查点到Safetensors格式的功能。LFM的存储系统。大规模检查点需要健壮的存储系统。LLaMA 3.1报告【13】突出Tectonic【39】,Meta的通用文件系统,作为预训练期间存储检查点的骨干。该系统的一个关键挑战是管理高频检查点写入。类似地,DeepSeek【11, Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024, 无会议, 无URL】【18】【33, Deepseek-v3 technical report, 2024, arXiv, arXiv:2412.19437】开发的AI-HPC系统FireFlyer【6, Fire-Flyer AI-HPC: A cost-effective software-hardware co-design for deep learning, 2024, arXiv, arXiv:2408.14158】利用自定义构建的3FS分布式文件系统【50, Fire-flyer file system, 2025, 无会议, https://github.com/deepseek-ai/3FS】,类似于其他并行文件系统如BeeGFS【21, An introduction to BeeGFS, 2014, 无会议, 无URL】,来管理检查点存储。弹性训练系统。一些正交努力致力于增强DL训练作业的弹性【7, Varuna: scalable, low-cost training of massive deep learning models, 2022, EuroSys, 无URL】【12】【24, Oobleck: Resilient distributed training of large models using pipeline templates, 2023, SOSP, 无URL】【30, Easyscale: Elastic training with consistent accuracy and improved utilization on gpus, 2023, SC, 无URL】【61, Bamboo: Making preemptible instances resilient for affordable training of large DNNs, 2023, NSDI, 无URL】【74, Swift: Expedited failure recovery for large-scale dnn training, 2023, PPoPP, 无URL】。例如,Varuna引入作业变形以在管道和数据并行度中重新配置训练作业,而不改变超参数。Bamboo【61】利用跨阶段冗余计算到管道中,用于在spot实例上的弹性训练。Oobleck【24】通过预定义模板实现管道实例化机制以容忍不同管道中的并发故障。这些工作限于特定并行策略(例如无ZeRO的DP),而ByteCheckpoint不假设任何特定并行,支持真实世界生产。采用我们的检查点表示形式用于训练状态管理可以显著增强弹性训练系统的灵活性。
💬 评论讨论
欢迎在这里分享您的想法和见解!