DeepSpeed项目在AMD GPU Mi200上启用ZeRO++的实践指南

背景介绍

DeepSpeed是微软开发的一个深度学习优化库，它通过多种技术手段来加速大规模模型的训练。其中ZeRO++是DeepSpeed的一项重要优化技术，能够显著减少内存占用并提高训练效率。然而，在AMD GPU（特别是MI200系列）上启用ZeRO++时，开发者可能会遇到一些技术挑战。

问题现象

当尝试在AMD MI200 GPU上启用ZeRO++时，系统会抛出异常，提示无法加载quantizer.so共享对象文件。错误信息表明系统尝试加载CUDA量化器模块失败，这通常是由于底层编译问题导致的。

根本原因分析

经过深入调查，发现这个问题主要有两个层面的原因：

1. 运行时缓存问题：系统尝试加载之前编译失败的缓存文件，导致运行时错误。

2. Torch版本兼容性问题：在Torch 2.3.0版本中，缺少对HIP(AMD GPU计算平台)的cublasGemmAlgo_t类型的支持，而这个类型是ZeRO++量化功能所必需的。这个问题在Torch 2.3.1版本中得到了修复。

解决方案

针对上述问题，我们推荐以下解决方案：

1. 清理编译缓存： 删除~/.cache/torch_extensions目录下的所有内容，强制系统重新编译相关组件。

2. 升级Torch版本： 将PyTorch升级到2.3.1或更高版本，以获得完整的HIP支持。

3. 正确设置环境变量： 在AMD GPU上编译时，确保设置正确的环境变量：

   DS_BUILD_QUANTIZER=1
   

技术细节

在AMD GPU上使用ZeRO++时，系统会尝试编译特定的量化器模块。这个模块依赖于HIP版本的BLAS库函数。Torch 2.3.0版本中缺少对cublasGemmAlgo_t类型的HIP映射，导致编译失败。这个类型用于指定矩阵乘法算法的选择策略。

在Torch 2.3.1中，开发团队添加了从CUDA到HIP的类型映射，使得cublasGemmAlgo_t能够正确映射到AMD平台上的hipblasGemmAlgo_t，从而解决了兼容性问题。

最佳实践建议

1. 在AMD GPU平台上使用DeepSpeed时，建议始终使用最新稳定版的PyTorch。

2. 遇到类似编译问题时，可以尝试以下步骤：

   • 检查错误日志中的具体编译错误
   • 清理编译缓存
   • 确认所有必要的环境变量已正确设置
   • 查阅相关组件的版本兼容性矩阵

3. 对于生产环境，建议先在测试环境中验证ZeRO++的功能，确保所有组件都能正常工作后再部署到生产环境。

总结

在AMD MI200系列GPU上使用DeepSpeed的ZeRO++功能是完全可行的，但需要注意PyTorch版本的兼容性问题。通过升级到Torch 2.3.1或更高版本，并正确配置编译环境，开发者可以充分利用ZeRO++带来的性能优势。这一案例也提醒我们，在异构计算环境中，保持软件栈的版本兼容性至关重要。