DeepSpeed v0.18.0 更新详解：DeepCompile、SuperOffload 与 ZeRO 优化全面升级

DeepSpeed 发布了 v0.18.0 版本，本次更新在分布式训练性能优化、内存管理、日志系统、模型检查点等方面进行了大量改进和修复，重点增强了DeepCompile、SuperOffload以及ZeRO 优化器的功能。以下为本次版本的详细更新内容：

1. DeepCompile 功能优化与修复

本次更新中，DeepCompile 进行了多项改进，提升了在大型模型与复杂计算图中的表现：

• •ZeRO-3 优化：增强了 allgather 操作在不均匀分片情况下的鲁棒性，并修复了性能分析问题。

• •VLA 类模型支持优化：修复了应用 DeepCompile 到 VLA-like 模型时的错误。

• •激活释放大小可配置：用户可根据需求配置释放激活的大小，提高显存利用率。

• •张量别名指定：在 C++ 算子 schema 中明确张量别名关系，减少内存冲突。

• •allgather 与 downcast 融合：提升通信效率与类型转换性能。

• •启用状态检测：当 DeepCompile 已启用但未激活时，系统可以正确处理此状态。

• •IPG bucket 清理修复：避免因 bucket 清理不当导致的训练过程异常。

• •最小切割重物化分区策略：在联合计算图分区时采用 min_cut_rematerialization 策略，提高计算与显存的平衡性。

• •测试增强：在测试失败时显示不匹配的数值，便于定位问题。

SuperOffload 在本次版本正式发布，主要应用于大模型训练的显存与存储优化：

• •功能发布：支持将部分张量及中间结果离线到 NVMe 等存储设备，提高大模型训练的可扩展性。

• •初始化支持：superoffload 文件夹增加初始化文件，保证模块加载正常。

• •博客文档更新：增加了 SuperOffload 的技术文档与中文版本，帮助开发者快速上手。

• •NVMe offload 检查点写入优化：修复了在 NVMe offload 张量合并过程中可能发生的保存竞争问题。

ZeRO 是 DeepSpeed 中用于高效分布式训练的核心技术，本次更新中涉及多个优化：

• •fp16 溢出广播：在 ZeRO-1 中支持 fp16 溢出状态的同步。

• •checkpoint 加载路径优化：当 stage=0 时禁用 ZeRO 检查点加载路径，避免冗余操作。

• •多子组通用检查点修复：解决了 stage3 多子组情况下加载通用检查点的问题。

• •跨机器通用检查点加载修复：修复了在多机环境下加载通用检查点的错误。

• •world size 扩展支持：修复了 stage3 在 world size 扩展情况下无法加载通用检查点的错误。

• •ZeRO-0 与自动混合精度集成：在 ZeRO-0 下支持与 torch.autocast 联合使用梯度缩放器。

除了核心模块外，本次更新还对日志、工具方法、文档等进行了完善：

• •日志系统优化：在设置日志级别时，同时调整各 logger handler 的级别，确保输出一致性。

• •设备命名优化：将current_device()更改为current_device_name()，提升可读性。

• •Leaf 模块接口改进

• • • 支持通过配置启用

  • • 放宽匹配条件

  • • 补充文档说明

  • • 简化 hook 实现

• •wall_clock_breakdown 改进：启用后始终记录统计数据，便于性能分析。

• •print_dist 工具：新增分布式打印工具方法，便于调试分布式信息。

• •TiledMLP moe 支持：在 TiledMLP 中增加 mixture-of-experts 支持。

• •文档更新：进一步明确了 Leaf 模块相关配置说明。

• •环境支持优化：在.gitignore中加入 venv 文件夹，方便虚拟环境管理。

• •邮箱地址更新：更新了系统中的邮件地址信息。

总结
DeepSpeed v0.18.0 的更新着重优化了 DeepCompile 与 SuperOffload 两大核心模块，同时在 ZeRO 优化器的稳定性与功能扩展方面也带来了显著提升。新增的工具和文档更新进一步降低了分布式训练的使用门槛，为开发者在超大规模模型训练中的性能、稳定性和可扩展性提供了更强的支持。

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