DeepSpeed v0.18.1 发布：新增 DataStates 异步检查点引擎，优化 Zero 优化器与存储共享机制

版本发布时间

2025年10月23日

主要更新内容

DeepSpeed v0.18.1 在性能优化、功能扩展与稳定性方面进行了多项改进，尤其在优化器灵活性、学习率控制、检查点引擎以及 HF Accelerate 集成等方面有显著提升。

核心更新一：新增 DataStates-LLM 异步检查点引擎

• • 引入 DataStatesCheckpointEngine，支持异步检查点保存与加载。
• • 可通过 ds_config.json 启用： .

{
    ... 其它配置,
    "datastates_ckpt": {
        "host_cache_size": 16
    }
}

• • 当启用 DataStates 引擎后，可在保存模型、优化器和 MoE 层检查点时根据 preserves_storage_sharing() 判断是否克隆张量，有效减少内存占用。
• • 如果未安装 datastates-llm 库，会自动回退至 TorchCheckpointEngine。

核心更新二：Zero 优化器增强

• • DeepSpeedZeroOptimizer 子类初始化方法支持 **kwargs，提升了参数传递的灵活性。
• • 移除了在 embedding 和 lm_head 层使用 Muon 优化器的设置。
• • Stage 3 SuperOffload 优化器初始化参数简化，通过 **kwargs 统一传递，支持 cpuadam_cores_perc 配置。

核心更新三：学习率支持张量类型

• • update_lr 方法支持标量与 Tensor 类型学习率，可根据现有 param_group['lr'] 类型自动匹配。
• • 当当前学习率为 Tensor 时，会保持原有的张量设备与类型。

核心更新四：存储共享机制优化

• • 新增 preserves_storage_sharing() 接口，用于判断检查点引擎是否需要克隆张量，以支持存储共享。
• • 在 MoE 层、优化器状态以及模型保存流程中根据该机制动态调整张量处理方式，避免无谓的内存复制。

核心更新五：ZenFlow Stage 3 支持

• • 增加了 ZenFlowSelectiveAdamW_stage3，并在 Zero 优化器中集成 ZenFlow Stage 3 功能。
• • 对优化器初始化流程进行了调整，以便更好地集成 ZenFlow 特性。

核心更新六：Ulysses HF Accelerate 集成

• • 改进了 register_with_transformers 方法：
  • • 支持直接传入已加载的 HF 模型对象。
  • • 增加对 flash_attention_2、flash_attention_3 以及 sdpa 作为注意力实现的支持。
• • 在数据加载器中增加形状验证，确保张量的 seqlen 与 input_ids 一致，否则抛出异常，提升数据一致性保障。

核心更新七：内存安全与CI修复

• • 修复 multi_tensor_apply 在 size 超过 INT_MAX 时出现的非法内存访问问题（int 改为 int64_t）。
• • 恢复 XPU xpu-max1100 CI 流程，移除部分冗余测试以提升运行效率。

主要文件改动

• • deepspeed/datastates/ 新增 DataStates 配置与引擎实现文件。
• • deepspeed/runtime/checkpoint_engine/ 新增 datastates_checkpoint_engine.py，并在 utils.py 中集成引擎选择逻辑。
• • deepspeed/runtime/engine.py、deepspeed/runtime/pipe/module.py、deepspeed/runtime/config.py 等文件统一集成 preserves_storage_sharing() 判断逻辑。
• • deepspeed/runtime/lr_schedules.py 新增 Tensor 学习率支持。
• • deepspeed/ops/adam/zenflow_torch_adam.py 增加 Stage 3 相关优化器类。
• • deepspeed/runtime/superoffload/superoffload_stage3.py 优化初始化参数处理。

总结

代码地址：github.com/deepspeedai/DeepSpeed

DeepSpeed v0.18.1 在分布式训练的可扩展性与灵活性方面再次升级，尤其是 异步检查点引擎 DataStates 的引入，为大模型的高效训练与部署提供了强有力的支持。与此同时，针对优化器、学习率控制以及 HF 模型集成的改进，也使该框架更易于与最新的大模型技术栈适配。

我们相信人工智能为普通人提供了一种“增强工具”，并致力于分享全方位的AI知识。在这里，您可以找到最新的AI科普文章、工具评测、提升效率的秘籍以及行业洞察。 欢迎关注“福大大架构师每日一题”，发消息可获得面试资料，让AI助力您的未来发展