CANN-昇腾NPU分布式训练-8卡到64卡怎么线性扩展
8 卡训练 Llama2-7B 的吞吐约 1800 tokens/s/p。64 卡应该是 8 卡的 8 倍吗？实际上只能到 6-7 倍。缺失的 1-2 倍被通信开销吃了。这篇分析昇腾NPU上分布式训练的扩展效率。
扩展效率定义
扩展效率 = 实际加速比 / 理论加速比 8 卡 → 64 卡： 理论加速比 = 8× 实际加速比 = 6.5× 扩展效率 = 6.5 / 8 = 81%80% 以上的扩展效率算优秀。低于 70% 需要排查通信瓶颈。
通信瓶颈分析
8 卡 TP 的通信：每层 2 次 All-Reduce，32 层 64 次。单机 HCCS 带宽 200GB/s，单次 All-Reduce 约 0.05ms，64 次 3.2ms。总步长约 50ms，通信占 6.4%。
64 卡（8 机 × 8 卡）的通信：除了机内 All-Reduce，还有跨机 All-Reduce。跨机走 RoCE 100GbE，单次 All-Reduce 约 1.8ms。32 层 64 次，总通信 115ms。
| 配置 | 计算时间 | 通信时间 | 通信占比 | 扩展效率 |
|---|---|---|---|---|
| 8 卡（单机） | 47ms | 3.2ms | 6.4% | 100% |
| 16 卡（2 机） | 24ms | 12ms | 33% | 82% |
| 32 卡（4 机） | 12ms | 35ms | 74% | 65% |
| 64 卡（8 机） | 6ms | 115ms | 95% | 42% |
64 卡的通信占比 95%，计算只占 5%。扩展效率只有 42%——每加一张卡多出的算力都被通信抵消了。
优化策略
策略 1：DP + TP 混合
不要把 64 卡全用来做 TP。4 卡做 TP（单机内），16 组 DP（跨机）：
每组：4 卡 TP，HCCS 通信，0.05ms/All-Reduce 跨组：DP，只需要梯度 All-Reduce（每步 1 次，不在每层） 每步通信： TP 通信：64 次 × 0.05ms = 3.2ms（机内） DP 通信：1 次 × 1.8ms = 1.8ms（跨机） 总通信：5ms64 卡扩展效率提升到 85%。
策略 2：通信计算重叠
MC2 的通算融合让 All-Reduce 跟计算重叠。每组 TP 的 64 次 All-Reduce 跟 forward/backward 的计算并行，通信时间从 3.2ms 降到接近 0。
策略 3：梯度压缩
跨机 DP 的梯度 All-Reduce 可以用 8bit 压缩：
fromtorch_npu.npuimportcompress_allreduce# 代替标准的 DDP All-Reducecompress_allreduce(model.parameters(),compression="8bit")8bit 压缩把通信量减半，精度损失约 0.1-0.3%。在 RoCE 带宽受限时效果显著。
实测扩展效率
Atlas 800I A2，Llama2-7B，bf16：
| 配置 | 吞吐 (tokens/s) | 加速比 | 扩展效率 |
|---|---|---|---|
| 8 卡 TP | 1,800 | 1× | 100% |
| 16 卡 (4TP+4DP) | 3,500 | 1.94× | 97% |
| 32 卡 (4TP+8DP) | 6,600 | 3.67× | 92% |
| 64 卡 (4TP+16DP) | 12,400 | 6.89× | 86% |
DP+TP 混合的扩展效率远好于纯 TP。64 卡时 86% vs 42%。
超大集群的建议
1000+ 卡的训练集群：
- TP=4 或 8（单机内，HCCS 通信）
- DP 按需扩展（跨机，RoCE 通信 + 梯度压缩）
- PP 不推荐（昇腾NPU的 HCCS 带宽高，不需要用 PP 避免通信）
- 梯度累积（减少 DP 通信频率）
分布式训练的扩展效率取决于通信策略。纯 TP 的扩展性差，DP+TP 混合是昇腾NPU上的最优解。通信计算重叠和梯度压缩是锦上添花。仓库在这里：
https://atomgit.com/cann/hccl