如何用 DS_SDMA_ALLGATHER 在 AMD MI300 上加速 DeepSpeed ZeRO-3 AllGather？

如何用 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_ALLGATHER=1，WORLD 进程组的all_gather_into_tensor就会改走 MI300 的 System DMA（SDMA）拷贝引擎，让集合通信和计算跑在物理上分离的硬件上，从而在不修改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无关。相关变量如下：

两个容易忽略的路由细节：

• 只有WORLD 进程组的all_gather_into_tensor会走 SDMA；如果你用非 WORLD 的数据并行组初始化 ZeRO-3（例如设置了sub_group_size产生二级 zero-param 组），那些子组的 allgather 会自动走 RCCL/NCCL。
• 回退条件：未设置DS_SDMA_ALLGATHER=1、非 WORLD 组、单 rank shard 超过 transit buffer、dtype 不在支持列表、或 mori 初始化失败——任一条件命中即回退，不需要人工干预。

用示例脚本跑一遍 A/B 对比

仓库里examples/sdma_allgather/提供了两条完整的复现路径，on/off 两组脚本共用同一份 ds_config_zero3.json（ZeRO-3 + bf16 + DeepSpeed 默认 bucket 配置，overlap_comm: true），唯一区别就是是否导出DS_SDMA_ALLGATHER=1。这些命令只启动训练进程，不会修改环境；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_ALLGATHER=1，走 SDMA 快速路径 # Demo 2 — Qwen3-32B，几分钟量级，随机初始化（无需权重） bash run_qwen3_sdma_off.sh # 基线 bash run_qwen3_sdma_on.sh # SDMA 开启

Qwen3 脚本额外导出PYTORCH_HIP_ALLOC_CONF=expandable_segments:True、PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True和TORCH_NCCL_ENABLE_MONITORING=0，用于缓解 32B 模型长生命周期 tensor 下的 HIP 分配器碎片，这两个 A/B 脚本共用，所以对比仍然公平。

脚本支持通过环境变量覆盖运行参数：SEQ_LEN、BATCH_SIZE、NUM_STEPS、WARMUP_STEPS、NUM_GPUS、MODEL、DS_CONFIG，例如只想跑 20 步冒烟测试可以NUM_STEPS=20 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_op=True）对 bfloat16/float16/float32 做正确性校验并测带宽。由于 SDMA 路径必须显式 opt-in，跑测试时同样要带上开关：

cd examples/sdma_allgather DS_SDMA_ALLGATHER=1 deepspeed --num_gpus 8 test_sdma_allgather_zero3.py # 可选：调整每 rank 元素数与测量轮数 DS_SDMA_ALLGATHER=1 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_BW（GB/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 上的已验证结果（以下为文档中的实测数据，你的工作负载数值会不同）：

数值稳定性方面，两个 demo 各跑了 2000 步（同一 wikitext shuffle，off 与 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_comm=false），收益会缩水甚至反转，所以不要把上面的 +10% 当作固定预期。

常见不生效的原因

按文档给出的回退条件逐条核对：

1. 没有设DS_SDMA_ALLGATHER=1（最常见：以为装了 mori 就自动生效，实际默认关闭）；
2. mori 不可 import 或 shmem 初始化失败——看 rank 0 的SDMA allgather unavailable日志；
3. 你的 allgather 发生在非 WORLD 组（sub_group_size等配置产生的子组永远走 RCCL）；
4. 单次调用的 per-rank 元素数超过 transit buffer（默认 64M 元素），可用DS_SDMA_ALLGATHER_MAX_NUMEL调大后重试；
5. 运行时不是 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