DeepSpeed v0.18.9 正式发布!AutoTP 全面增强、ZeRO 与 SuperOffload 深度优化,大模型训练再升级



在大模型分布式训练领域,DeepSpeed作为业界主流的优化框架,持续通过版本迭代为开发者提供更高效、更稳定、更兼容的训练能力。近日,DeepSpeed正式推出v0.18.9版本,本次更新聚焦于AutoTP自动张量并行能力增强、ZeRO优化器适配、SuperOffload内存卸载优化、PyTorch兼容性修复、环境变量适配完善以及CI流程与文档更新等多个核心方向,同时完成了AutoSP功能的合并入库,进一步完善了序列并行相关能力,修复了多项影响训练稳定性与环境适配的问题,为大模型训练提供了更可靠的底层支撑。
本次版本更新共计16次提交、56个文件变更,由16位社区贡献者共同完成,代码层面实现3891行新增与321行删除,覆盖功能新增、缺陷修复、兼容性优化、文档更新、CI流程调整等全维度内容,本文将基于官方更新日志,对DeepSpeed v0.18.9的所有变更进行全面、详细的解析,帮助开发者快速掌握版本核心升级点,顺利完成版本升级与业务适配。
一、环境适配与硬件检测优化
DeepSpeed v0.18.9针对环境变量适配、GPU硬件检测逻辑进行了多项优化,解决了不同硬件平台、自定义环境路径下的适配问题,提升了框架的环境兼容性与部署灵活性。
1. 支持$TRITON_HOME环境变量
本次更新新增对TRITONHOME环境变量的适配,框架会自动识别该环境变量配置的路径,解决了Triton相关组件在自定义安装路径下无法被DeepSpeed正确加载的问题。在过往版本中,DeepSpeed仅默认读取系统默认路径下的Triton相关文件,当用户通过自定义路径安装Triton时,会出现路径查找失败、组件调用异常等问题,而v0.18.9通过适配TRITON_HOME环境变量的适配,框架会自动识别该环境变量配置的路径,解决了Triton相关组件在自定义安装路径下无法被DeepSpeed正确加载的问题。在过往版本中,DeepSpeed仅默认读取系统默认路径下的Triton相关文件,当用户通过自定义路径安装Triton时,会出现路径查找失败、组件调用异常等问题,而v0.18.9通过适配TRITONHOME环境变量的适配,框架会自动识别该环境变量配置的路径,解决了Triton相关组件在自定义安装路径下无法被DeepSpeed正确加载的问题。在过往版本中,DeepSpeed仅默认读取系统默认路径下的Triton相关文件,当用户通过自定义路径安装Triton时,会出现路径查找失败、组件调用异常等问题,而v0.18.9通过适配TRITON_HOME环境变量,允许用户自由指定Triton的安装目录,大幅提升了框架在自定义部署环境下的适配能力,尤其适用于多版本Triton共存、离线部署等复杂场景。
2. ROCm GPU架构检测移除不必要的shell=True
针对AMD ROCm平台的GPU架构检测逻辑,本次更新移除了代码中不必要的shell=True参数。在Python的subprocess调用中,shell=True会带来潜在的安全风险与执行效率问题,同时可能引发命令解析异常。DeepSpeed v0.18.9优化了ROCm GPU架构检测的代码实现,采用无shell模式执行硬件检测命令,既提升了检测逻辑的安全性,又避免了因shell环境差异导致的GPU架构识别失败问题,保障了AMD GPU平台上DeepSpeed的稳定运行。
3. 支持$DS_IGNORE_CUDA_DETECTION环境变量跳过本地GPU检测
新增$DS_IGNORE_CUDA_DETECTION环境变量控制逻辑,当用户设置该环境变量时,DeepSpeed会跳过本地CUDA GPU的检测流程。该优化主要针对无GPU环境、远程训练节点、容器化部署等场景,在这些场景下,本地GPU检测不仅无意义,还可能因检测失败导致框架初始化异常。通过该环境变量,开发者可以灵活控制GPU检测行为,避免非必要的硬件检测流程,提升框架在纯CPU环境、云服务器远程调度场景下的初始化效率。
4. 修复Triton自动调优缓存的NFS路径判断问题
修复了Triton自动调优缓存模块中,is_nfs_path函数对不存在路径的处理异常问题。在过往版本中,当判断的NFS路径不存在时,该函数会抛出异常,导致Triton自动调优流程中断,影响模型训练的性能调优。v0.18.9优化了路径存在性校验逻辑,当目标路径不存在时,函数会返回合理的默认值,不会中断调优流程,保障了Triton自动调优功能在任意路径配置下的稳定运行,提升了大模型训练的性能调优可靠性。
二、AutoTP自动张量并行核心功能增强
AutoTP作为DeepSpeed实现自动张量并行的核心模块,是大模型分布式训练的关键能力,本次v0.18.9版本对AutoTP进行了多项功能升级,大幅提升了其兼容性、易用性与模型适配范围。
1. 新增AutoTP通用Checkpoint功能
本次更新为AutoTP新增通用Checkpoint(通用 checkpoint)功能,在checkpoint/constants.py中新增UNIVERSAL_CHECKPOINT_VERSION_KEY、UNIVERSAL_CHECKPOINT_VERSION_VALUE、DS_AUTOTP_UC_META等常量定义,同时在universal_checkpoint.py中实现了AutoTP通用 checkpoint 的加载与解析逻辑。
该功能支持AutoTP模式下的模型权重、并行元信息的统一存储与恢复,新增_get_param_uc_restore_meta函数用于获取AutoTP通用 checkpoint 的参数元数据,新增_resolve_autotp_partition函数实现AutoTP并行参数的分区解析与恢复,支持分区维度、逻辑形状、子参数形状、子参数大小等元信息的读取与处理,同时支持复制型参数的特殊处理。
通用Checkpoint功能解决了AutoTP模式下模型 checkpoint 兼容性差、跨节点恢复困难的问题,支持不同并行配置下的模型权重迁移,尤其适用于大模型训练中断恢复、多节点并行训练切换、模型微调等场景,大幅提升了AutoTP模式下训练任务的可迁移性与可靠性。
2. 支持HuggingFace官方tp_plan配置
新增对HuggingFace Transformers库中tp_plan的原生支持,AutoTP可自动识别并使用Llama、Qwen、Gemma2等模型内置的base_model_tp_plan配置,无需开发者手动设置preset_model或partition_config参数。该优化实现了DeepSpeed AutoTP与HuggingFace模型生态的无缝衔接,开发者无需额外编写并行分区规则,直接使用HuggingFace官方模型即可开启AutoTP并行训练,支持colwise与rowwise两种核心分区类型,大幅降低了AutoTP的使用门槛。
同时,官方在huggingface-tp/README.md中更新了相关文档,详细说明HuggingFace tp_plan的适配逻辑与使用方法,明确了AutoTP对自定义分区的支持方式,让开发者可以快速上手基于HuggingFace模型的自动张量并行训练。
三、PyTorch兼容性与训练稳定性修复
DeepSpeed的核心能力依赖PyTorch生态,本次版本针对低版本PyTorch的兼容性、训练过程中的内存与梯度更新问题进行了关键修复,保障了不同PyTorch版本下训练任务的稳定运行。
1. 修复PyTorch 2.4以下版本torch.amp.custom_fwd兼容性问题
修复了torch.amp.custom_fwd接口在PyTorch 2.4以下版本的向后兼容问题。在过往版本中,DeepSpeed的自动混合精度训练模块调用custom_fwd接口时,会因PyTorch版本差异出现接口不兼容、训练中断等异常,v0.18.9通过接口适配与逻辑兼容处理,确保在PyTorch 2.4以下版本中,自动混合精度训练功能可以正常运行,无需开发者升级PyTorch版本,提升了框架对老旧环境的适配能力。
2. SuperOffload模块修复多分组共享CPU缓冲区梯度更新问题
修复了SuperOffload模块中,多分组梯度更新使用共享CPU缓冲区时的数据丢失问题。SuperOffload是DeepSpeed实现内存卸载、降低GPU显存占用的核心功能,在大模型训练中,多分组梯度更新场景下,共享CPU缓冲区会导致梯度数据覆盖、更新失效等问题,严重影响训练收敛效果。v0.18.9优化了缓冲区管理逻辑,保障多分组梯度更新时的数据独立性,确保梯度更新结果正确保留,提升了SuperOffload模式下大模型训练的稳定性与收敛效果。
四、ZeRO优化与Muon优化器适配扩展
ZeRO作为DeepSpeed的核心内存优化技术,本次更新对ZeRO Stage 3进行了优化,同时扩展了Muon优化器的支持范围,提升了分布式训练的内存效率与优化器兼容性。
1. 扩展Muon优化器对ZeRO Stage 3的支持
新增Muon优化器在ZeRO Stage 3模式下的完整支持,解决了过往版本中Muon优化器与ZeRO Stage 3不兼容、参数分区异常、优化器状态同步失败等问题。ZeRO Stage 3是DeepSpeed显存优化的最高级别,可将模型参数、梯度、优化器状态全部分区存储,大幅降低单卡显存占用,而Muon优化器作为新型优化器,在大模型训练中具备收敛速度快、稳定性强的优势,本次适配让开发者可以在ZeRO Stage 3模式下使用Muon优化器,兼顾显存效率与训练性能。
2. ZeRO相关底层逻辑优化
配合Muon优化器适配,DeepSpeed v0.18.9对ZeRO Stage 3的参数分区、状态同步、内存管理逻辑进行了微调,确保优化器状态与模型分区参数的正确映射,避免了显存泄漏、状态同步延迟等问题,进一步提升了ZeRO Stage 3模式下的训练稳定性。
五、AutoSP功能合并入库
本次版本最大的功能升级之一,是将AutoSP(自动序列并行)功能正式合并入DeepSpeed主分支,完成了AutoSP模块的全量集成。AutoSP是DeepSpeed针对长序列大模型训练推出的自动序列并行技术,基于torch.compile与FX图优化实现序列维度的自动分片,无需手动修改模型代码,即可实现序列并行训练,大幅提升长上下文大模型的训练效率。
1. AutoSP核心测试用例新增
在tests/unit/v1/compile/目录下新增大量AutoSP相关测试用例,覆盖SDPA节点检测、输入ID/标签ID/位置ID节点识别、符号序列长度节点查找、分片偏移量生成、符号切片索引创建、张量分片节点替换等核心功能测试,测试序列长度覆盖64、128、256等常用场景,确保AutoSP功能在不同序列长度下的稳定性。
2. AutoSP与Ulysses SP对齐验证
新增compare_sp_loss测试函数,实现AutoSP编译模型与Ulysses SP基准模型的损失值对齐验证,在分布式环境下对比两种序列并行方案的训练损失,确保AutoSP的计算精度与Ulysses SP一致,误差控制在RTOL=0.1、ATOL=0.01范围内,保障自动序列并行的训练精度。
3. AutoSP图优化逻辑完善
新增create_gm_nodes函数实现模型FX图捕获,通过自定义torch.compile后端捕获模型计算图,确保AutoSP可以正确识别序列维度与输入节点;新增find_sym_seq_node函数实现符号序列长度节点的查找,解决了torch 2.9 bf16模式下符号节点顺序异常的问题,保障AutoSP在不同PyTorch版本与精度模式下的正常运行。
六、CI流程、文档与版权信息更新
1. CI工作流优化
对.github/workflows/aws-torch-latest-full.yml进行调整,优化AWS L40S GPU集群的CI测试流程,新增容量 fallback 逻辑,当L40S GPU资源不足时,自动切换至8xA100节点执行测试;优化提交检测逻辑,简化GitHub Actions的命令执行格式,统一标点符号规范;新增GPU架构自动检测步骤,通过Python代码实时获取GPU计算能力与设备数量,动态设置TORCH_CUDA_ARCH_LIST环境变量,提升CI测试的兼容性与灵活性。
同时新增测试 fallback 逻辑,当部分测试用例执行失败时,自动切换至完整测试模式,确保测试覆盖率,保障版本发布的稳定性。
2. 开发文档更新
对AGENTS.md与CLAUDE.md文档进行优化,在CI开发规范中新增pre-commit命令说明,明确要求开发者提交代码前执行pre-commit run --files <changed_files>命令,仅检查修改文件而非全量代码库,提升开发效率;移除文档中Microsoft Corporation版权声明,统一版权规范;更新pre-commit配置说明,明确代码格式化、静态检查的执行要求。
3. 项目README与新闻更新
更新README.md的Latest News板块,新增两项重要资讯:一是DeepSpeed团队在ASPLOS 2026发布大模型系统构建教程;二是SuperOffload相关工作获得ASPLOS 2026最佳论文荣誉提名,同步更新相关链接,提升项目社区影响力。
同时新增ASPLOS 2026最佳论文奖相关新闻条目,完善项目荣誉展示,增强社区开发者对DeepSpeed技术价值的认可。
4. 版本号更新
将version.txt文件中的版本号从0.18.8更新至0.18.9,完成正式版本的标识,方便开发者通过版本号区分迭代内容,顺利进行版本升级。
七、版本更新总结与升级建议
DeepSpeed v0.18.9作为一次迭代型版本更新,聚焦于功能增强、兼容性修复、稳定性提升、生态适配四大核心方向,所有变更均围绕大模型分布式训练的实际痛点展开,无破坏性变更,适合所有使用DeepSpeed的开发者升级。
本次更新的核心价值体现在:
- AutoTP能力全面升级,通用Checkpoint与HuggingFace tp_plan支持,大幅降低自动张量并行使用门槛;
- 硬件与环境适配优化,支持自定义环境变量、多平台GPU检测,提升框架部署灵活性;
- 训练稳定性修复,解决PyTorch兼容、SuperOffload梯度更新、Triton路径异常等关键问题;
- AutoSP正式合并,长序列自动并行能力落地,适配长上下文大模型训练;
- ZeRO与优化器适配扩展,支持Muon优化器在ZeRO Stage 3模式下运行;
- CI与文档完善,提升开发协作效率与项目规范性。
升级建议:
- 正在使用DeepSpeed v0.18.x系列版本的开发者,可直接升级至v0.18.9,无需修改业务代码;
- 使用HuggingFace模型进行张量并行训练的开发者,升级后可直接使用官方tp_plan,简化配置;
- 长序列大模型训练场景,可尝试AutoSP功能,提升序列并行训练效率;
- AMD ROCm平台、自定义Triton路径、无GPU环境的开发者,升级后可解决环境适配问题;
- 使用PyTorch 2.4以下版本的开发者,升级后可修复混合精度训练接口兼容问题。
结语
代码地址:github.com/deepspeedai/DeepSpeed
DeepSpeed v0.18.9版本通过社区协同开发,完成了多项实用功能升级与缺陷修复,进一步巩固了其在大模型分布式训练领域的领先地位。AutoTP的生态适配、AutoSP的正式落地、ZeRO与优化器的深度优化,让大模型训练更高效、更简单、更稳定;环境适配与兼容性修复,保障了框架在不同硬件、不同软件版本下的可靠运行。
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。
更多推荐


所有评论(0)