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。

• •DeepSpeedZeroOptimizer子类初始化方法支持**kwargs，提升了参数传递的灵活性。

• • 移除了在 embedding 和 lm_head 层使用 Muon 优化器的设置。

• • Stage 3 SuperOffload 优化器初始化参数简化，通过**kwargs统一传递，支持cpuadam_cores_perc配置。

• •update_lr方法支持标量与Tensor 类型学习率，可根据现有param_group['lr']类型自动匹配。

• • 当当前学习率为Tensor时，会保持原有的张量设备与类型。

• • 新增preserves_storage_sharing()接口，用于判断检查点引擎是否需要克隆张量，以支持存储共享。

• • 在 MoE 层、优化器状态以及模型保存流程中根据该机制动态调整张量处理方式，避免无谓的内存复制。

• • 增加了ZenFlowSelectiveAdamW_stage3，并在 Zero 优化器中集成 ZenFlow Stage 3 功能。

• • 对优化器初始化流程进行了调整，以便更好地集成 ZenFlow 特性。

• • 改进了register_with_transformers方法：

• • • 支持直接传入已加载的 HF 模型对象。

  • • 增加对flash_attention_2、flash_attention_3以及sdpa作为注意力实现的支持。

• • 在数据加载器中增加形状验证，确保张量的seqlen与input_ids一致，否则抛出异常，提升数据一致性保障。

• • 修复multi_tensor_apply在size超过INT_MAX时出现的非法内存访问问题（int改为int64_t）。

• • 恢复 XPUxpu-max1100CI 流程，移除部分冗余测试以提升运行效率。

• •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助力您的未来发展。