DeepSpeed 集成 DataStates-LLM：大模型惰性异步检查点（Async Checkpointing）实战指南

DeepSpeed 集成 DataStates-LLM大模型惰性异步检查点Async Checkpointing实战指南【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本篇技术指南讲解如何将 DataStates-LLM 的惰性异步检查点lazy asynchronous checkpointing机制接入 DeepSpeed 训练框架从而把频繁保存检查点带来的 I/O 阻塞与训练进程解耦显著降低保存开销。读完本文你将掌握 DataStates-LLM 在 DeepSpeed 中的启用前提、deepspeed_config.json配置方式、训练脚本中的检查点频率控制以及当前实现边界与源码级工作原理。为什么 LLM 训练需要异步检查点大语言模型训练周期长、易中断检查点checkpoint保存是保障容错恢复的核心手段。但模型规模越大检查点越大同步写入持久化存储如并行文件系统、对象存储造成的 I/O 阻塞越明显甚至拖慢整个训练进度。DataStates-LLM 正是针对这一痛点提出的方案它利用前向/反向传播过程中模型参数与优化器状态在内存中的不可变性immutability将数据搬运到宿主机内存乃至更深层存储的步骤做成非阻塞non-blocking的多级异步传输从而最小化对训练进程的干扰。据其论文DataStates-LLM: Lazy Asynchronous Checkpointing for Large Language ModelsarXiv:2406.10707报告相较传统同步保存方式该方法可带来最高约 48 倍的检查点保存加速与约 2.2 倍的端到端训练加速。上述数字为论文所述实验结果具体收益取决于模型规模、存储带宽与检查点频率。本仓库中该特性对应一个专门模块 deepspeed/datastates并在 deepspeed/datastates/README.md 中说明此特性默认不开启需在 ds_config.json 中配置并另行安装 DataStates-LLM 检查点库完整教程即指向本文对应的文档 docs/_tutorials/datastates-async-checkpointing.md。DeepSpeed 中的检查点引擎抽象与 DataStates 接入点要理解 DataStates 如何工作先要了解 DeepSpeed 的检查点引擎Checkpoint Engine抽象。DeepSpeed 将检查点的写入、读取、提交统一抽象为引擎接口并支持按配置自动选择实现。引擎选择链路创建引擎的入口是create_checkpoint_engine()位于 deepspeed/runtime/checkpoint_engine/utils.py其选择顺序为若配置了checkpointwriter 相关选项则创建DecoupledCheckpointEngine或FastCheckpointEngine用于数据并行 checkpoint 优化若nebula_config.enabled为真则尝试创建NebulaCheckpointEngine若datastates_config.enabled为真则尝试创建DataStatesCheckpointEngineutils.py以上均不命中或依赖缺失时回退到默认的TorchCheckpointEngine。各引擎统一从 deepspeed/runtime/checkpoint_engine/init.py 导出并在 engine.py 的_configure_checkpointing()约 engine.py#L1645-L1652中完成装配。DataStatesCheckpointEngine 包装实现DataStatesCheckpointEngine 是 DeepSpeed 官方提供的适配器其关键行为延迟导入第三方库构造函数在运行时from datastates import CheckpointEnginedatastates_checkpoint_engine.py#L20-L26若未安装会抛出RuntimeError(Please install DataStates ...)解耦式引擎is_decoupled()返回Truedatastates_checkpoint_engine.py#L55-L56意味着save()只负责把状态交给 DataStates 内部异步管线真正的落盘提交commit()调用self.ckpt_engine.wait(persistTrue)可延后触发实现先记账、后刷盘生命周期管理__del__中调用cleanup()cleanup()会先提交未完成的提交记录再等待 DataStates 引擎持久化完成datastates_checkpoint_engine.py#L49-L53避免退出时丢数据。从 engine.py 的使用方式看普通同步引擎在save_checkpoint()末尾立即commit()约 engine.py#L4932-L4934而解耦引擎则通过is_decoupled()分支在梯度累积边界统一执行_commit_decoupled_checkpoint()约 engine.py#L3547-L3548把多份待提交数据合并刷盘。模型/优化器状态经由checkpoint_engine.save(...)写入如 engine.py#L5246-L5246 与 engine.py#L5360-L5361从而实现整条保存链路对训练主循环的异步化。接入前的准备工作Prerequisites在将 DataStates-LLM 与 DeepSpeed 集成前需要满足两个前置条件安装 DeepSpeed环境中应已安装 DeepSpeed。若尚未安装可参考仓库内的 DeepSpeed 快速开始指南 完成安装获取 DataStates-LLM 库DataStates-LLM 是独立第三方开源项目需从其官方 GitHub 仓库DataStates/datastates-llm获取源码并按其自带的安装说明完成安装。DeepSpeed 只在运行时按需导入它见上文DataStatesCheckpointEngine的惰性import因此必须先安装否则将触发导入错误或自动回退到TorchCheckpointEngine。配置 DeepSpeed启用datastates_ckpt要启用 DataStates-LLM 的异步检查点需要在deepspeed_config.json即传给deepspeed.initialize的配置中增加datastates_ckpt一节。以下为示例配置{ // ... 其他 DeepSpeed 配置选项 datastates_ckpt: { host_cache_size: 16 } }注意原 JSON 语法并不支持//行注释实际使用请移除注释。该配置会被 DeepSpeedDataStatesConfig 解析。从源码看其判定规则为只要配置字典中存在datastates_ckpt键且其值不为False即视为启用param_dict.get(datastates_ckpt, False) is not False并将该子配置深拷贝保留后整体传给 DataStates 引擎。也就是说仅写入一个空对象datastates_ckpt: {}也能触发启用分支但实际行为仍取决于后续的host_cache_size等参数。配置参数host_cache_sizehost_cache_size预留的钉页宿主机内存pinned host memory容量单位为吉字节GB用于异步检查点数据的刷写缓冲。该值越大可暂存在 CPU 端等待异步搬运/持久化的检查点数据越多训练侧的阻塞越小但它直接占用系统内存应依据机器内存总量与单次模型检查点大小综合设定例如示例中的 16 表示预留 16 GB。该参数在 DeepSpeed 侧不做额外校验实际由 DataStates 库读取并生效。在训练脚本中触发按迭代保存启用datastates_ckpt后检查点的捕获频率由训练脚本侧的迭代间隔参数控制即指定每经过多少个迭代iteration后捕获一次检查点典型做法是在训练入口脚本中通过命令行参数--save-interval传入该间隔。其原理是检查点捕获仍由训练脚本显式驱动例如在for训练循环中每N步调用一次引擎的save_checkpoint区别仅在于保存动作已交给 DataStates 异步化调用方只需登记该次保存真正的落盘在后台完成。仓库测试中对--save-interval的典型用法可见 tests/model/Megatron_GPT2/run_checkpoint_test.py其中以--save {checkpoint_folder} --save-interval {checkpoint_interval}形式把间隔透传给模型训练脚本用于验证保存 N 步 → 中断 → 恢复 → 对比结果一致的检查点闭环。推荐实践save-interval与host_cache_size需协同调整——间隔越短、单次状态越大越需要充足的 host cache 缓冲否则异步管道来不及消化时会重新形成排队阻塞启用 DataStates 后由于保存是解耦的不要依赖save_checkpoint返回即代表数据已落盘正常的提交/等待流程由 DeepSpeed 引擎在梯度累积边界或cleanup()时统一触发。当前限制与后续演进方向根据官方教程 datastates-async-checkpointing.md 的说明集成 DataStates-LLM 目前存在以下边界选型与排障时需特别留意仅支持 NVIDIA GPU 的 CUDA 运行时DataStates-LLM 当前只适配 CUDA 运行环境在 AMD ROCm、Intel XPU、昇腾 NPU、寒武纪 MLU 等非 CUDA 后端上无法使用仅验证过 ZeRO stage-1 且不向其他层级卸载offload即配置为 ZeRO 第一阶段、且不启用 CPU/NVMe offload 的组合启用 ZeRO stage-2/3 或 offload 的组合尚未得到验证需自行评估风险检查点布局接近但不等同于 safetensorsDataStates 产生的检查点文件布局与 Hugging Face 的 safetensors 格式相近但 DeepSpeed 在重启恢复时仍需依赖 pickle 序列化的对象因此目前尚不能与 safetensors 库完全互操作暂不支持 universal / elastic checkpointingDeepSpeed 提供的通用检查点universal checkpoint与弹性检查点elastic checkpoint能力尚无法与 DataStates 引擎叠加使用。仓库中的相关能力实现可分别在 deepspeed/checkpoint 与 deepspeed/elasticity 下查看作为了解其差异的背景。故障排查与支持渠道接入过程中最常见的两类失败都源于第三方库缺失且两类报错位置不同可用于快速定位初始化阶段硬失败若配置已启用但库缺失DataStatesCheckpointEngine构造函数会抛出RuntimeError(Please install DataStates from https://github.com/DataStates/datastates-llm.)训练直接中断提示信息即安装指引引擎装配阶段软回退create_checkpoint_engine()中的try/except会捕获ImportError打印No datastates engine found! ... Will fall back to torch.save.日志并静默回退到TorchCheckpointEngineutils.py#L42-L46此时训练仍可运行但并未真正获得异步加速。因此建议训练启动时留意日志确认使用的是 DataStates 引擎而非 torch 回退引擎。关于 DataStates-LLM 的疑问、Issue 与特性请求官方建议直接提交到 DataStates-LLM 的 GitHub 仓库获取维护者支持而 DeepSpeed 侧适配层的源码、导出与装配关系可在 deepspeed/datastates、deepspeed/runtime/checkpoint_engine 与 engine.py 中进一步追踪验证。小结DataStates-LLM 与 DeepSpeed 的集成本质上是把保存检查点从训练关键路径上摘除通过datastates_ckpt配置段启用解耦式检查点引擎配合host_cache_size预留的宿主机钉页缓冲并借助训练脚本的--save-interval控制捕获节奏即可在不改变原有save_checkpoint/load_checkpoint使用习惯的前提下获得异步化收益。当然其 CUDA-only、ZeRO stage-1-only 的当前边界决定了它更适合作为特定场景下的性能优化选项而非通用默认方案。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考