如何用 DS_SDMA_ALLGATHER 在 AMD MI300 上加速 DeepSpeed ZeRO-3 AllGather【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed在 AMD MI300 上用 ZeRO-3 训练时RCCL 的 AllGather kernel 和 GEMM/attention 计算跑在同一组计算单元CU上即使 prefetch 时间重叠调得再好重叠窗口内的有效计算吞吐也可能低于峰值。DeepSpeed 为此提供了一条 SDMA 快速路径只要设置一个环境变量DS_SDMA_ALLGATHER1WORLD 进程组的all_gather_into_tensor就会改走 MI300 的 System DMASDMA拷贝引擎让集合通信和计算跑在物理上分离的硬件上从而在不修改ds_config、不改 ZeRO-3 源码的前提下拿到 step-time 收益。适用前提缺一不可AMD/ROCm 运行时ROCm 的mori包可 import这是外部依赖DeepSpeed 本身不带GPU 为含独立 SDMA 拷贝引擎的 MI300X / MI325X / MI355X节点内 SDMA 拷贝官方复现脚本按单节点 8 卡deepspeed --num_gpus 8组织。如果 mori 装不上或初始化失败不会报错中断而是静默回退到 RCCL/NCCL训练照常进行——这一点在验证环节会再检查。开启方式只靠环境变量没有 ds_config 开关SDMA 路径完全由环境变量控制与ds_config无关。相关变量如下变量作用DS_SDMA_ALLGATHER1唯一开关。必须显式设置才启用 SDMA 快速路径即使已安装 mori默认也是关闭的。设置后 DeepSpeed 会自动导出MORI_ENABLE_SDMA1让 mori 分配 SDMA kernel 需要的 uncached transit buffer。DS_SDMA_ALLGATHER_MAX_NUMELNtransit buffer 大小单位元素数默认 64M每 rank 输入 256 MiB8 rank 时输出约 2 GiB。超过该大小的 allgather 调用自动回退 RCCL/NCCL。MORI_ENABLE_SDMA1mori 自身的开关由 DeepSpeed 在你设置DS_SDMA_ALLGATHER1时自动设置一般不需要手动导出。两个容易忽略的路由细节只有WORLD 进程组的all_gather_into_tensor会走 SDMA如果你用非 WORLD 的数据并行组初始化 ZeRO-3例如设置了sub_group_size产生二级 zero-param 组那些子组的 allgather 会自动走 RCCL/NCCL。回退条件未设置DS_SDMA_ALLGATHER1、非 WORLD 组、单 rank shard 超过 transit buffer、dtype 不在支持列表、或 mori 初始化失败——任一条件命中即回退不需要人工干预。用示例脚本跑一遍 A/B 对比仓库里examples/sdma_allgather/提供了两条完整的复现路径on/off 两组脚本共用同一份 ds_config_zero3.jsonZeRO-3 bf16 DeepSpeed 默认 bucket 配置overlap_comm: true唯一区别就是是否导出DS_SDMA_ALLGATHER1。这些命令只启动训练进程不会修改环境GPT demo 用合成数据、无需下载模型权重Qwen3 demo 通过from_config随机初始化同样不需要权重文件。cd examples/sdma_allgather # Demo 1 — GPT-7B 级别约 1 分钟量级无需下载 bash run_gpt_sdma_off.sh # 基线DS_SDMA_ALLGATHER 未设置走 RCCL bash run_gpt_sdma_on.sh # DS_SDMA_ALLGATHER1走 SDMA 快速路径 # Demo 2 — Qwen3-32B几分钟量级随机初始化无需权重 bash run_qwen3_sdma_off.sh # 基线 bash run_qwen3_sdma_on.sh # SDMA 开启Qwen3 脚本额外导出PYTORCH_HIP_ALLOC_CONFexpandable_segments:True、PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True和TORCH_NCCL_ENABLE_MONITORING0用于缓解 32B 模型长生命周期 tensor 下的 HIP 分配器碎片这两个 A/B 脚本共用所以对比仍然公平。脚本支持通过环境变量覆盖运行参数SEQ_LEN、BATCH_SIZE、NUM_STEPS、WARMUP_STEPS、NUM_GPUS、MODEL、DS_CONFIG例如只想跑 20 步冒烟测试可以NUM_STEPS20 bash run_gpt_sdma_on.sh。验证确认走的是 SDMA并量化收益1. 单元级验证确认后端确实切换了examples/sdma_allgather/test_sdma_allgather_zero3.py 按 ZeRO-3_all_gather_dtype的调用方式flat 输出、每 rank shard 输入、async_opTrue对 bfloat16/float16/float32 做正确性校验并测带宽。由于 SDMA 路径必须显式 opt-in跑测试时同样要带上开关cd examples/sdma_allgather DS_SDMA_ALLGATHER1 deepspeed --num_gpus 8 test_sdma_allgather_zero3.py # 可选调整每 rank 元素数与测量轮数 DS_SDMA_ALLGATHER1 deepspeed --num_gpus 8 test_sdma_allgather_zero3.py --partition_sz 4194304 --iterations 50结果判读rank 0 输出结构来自测试代码头部backend一行显示SDMA (mori)说明快速路径已生效显示RCCL/NCCL (mori unavailable or disabled)说明 mori 不可用或开关没设上此时测得的只是 RCCL 基线。正确性部分每个 dtype 输出PASSED/FAILED末尾汇总Result: All correctness tests PASSED。带宽部分按 dtype 输出 data_MB / avg_ms / min_ms / max_ms / algo_BWGB/s表格。文档没有给出必须达到的固定带宽值on/off 各跑一遍对比同一张表即可。在 mori 后端实现 中init_distributed()时 rank 0 还会打一行日志启用成功为SDMA allgather enabled via mori_cpp.AllGatherIntoTensor (max_numel...)初始化失败则为SDMA allgather unavailable (...); using RCCL/NCCL allgather。训练日志里看到后一行就说明本次运行实际走的是 RCCL 回退路径。2. 训练级验证step-time 与 loss 数值稳定性README 给出了 8x MI300X 上的已验证结果以下为文档中的实测数据你的工作负载数值会不同GPT-7B 级别Qwen3-32B配置seq 2048 / micro batch 1wikitext-2-raw-v1seq 1024 / micro batch 1wikitext-103-raw-v110%测量100 步 10 步 warmup100 步 10 步 warmupSDMA off (RCCL)697.7 ms/step1402.5 ms/stepSDMA on622.0 ms/step1263.2 ms/step收益10.85 %9.93 %rank 0 峰值显存12.12 GBoff/on 不变96.45 GBoff/on 不变数值稳定性方面两个 demo 各跑了 2000 步同一 wikitext shuffleoff 与 on 对照分桶平均 |off − on| loss 差在 GPT 上 ≤ 0.026、Qwen3 上 ≤ 0.048文档结论是 SDMA 路径没有引入数值漂移。对应的 loss 曲线图见 loss_gpt_2k.png 和 loss_qwen3_2k.png。如果你要自己下结论先跑 off 脚本记录每步耗时再跑 on 脚本同一份ds_config、同一份数据、同一 shuffle对比平均 ms/step 即可。README 提醒收益是 workload-dependent 的——当 allgather 无法与计算重叠时例如 payload 非常小或overlap_commfalse收益会缩水甚至反转所以不要把上面的 10% 当作固定预期。常见不生效的原因按文档给出的回退条件逐条核对没有设DS_SDMA_ALLGATHER1最常见以为装了 mori 就自动生效实际默认关闭mori 不可 import 或 shmem 初始化失败——看 rank 0 的SDMA allgather unavailable日志你的 allgather 发生在非 WORLD 组sub_group_size等配置产生的子组永远走 RCCL单次调用的 per-rank 元素数超过 transit buffer默认 64M 元素可用DS_SDMA_ALLGATHER_MAX_NUMEL调大后重试运行时不是 AMD/ROCm。核对完如果仍然走 RCCL 且训练本身正常说明触发的是上面某个安全回退条件而不是故障。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考