压测工具选型：nccl-tests、DeepSpeed、Megatron-LM 对比
说到压测工具选型，我见过不少团队一上来就纠结。其实没那么复杂。你想想看，这三个工具解决的问题层次完全不同。我习惯把它们分成三类：底层通信测试、框架级压测、模型级压测。
嗯，咱们一个一个聊。
1. nccl-tests：最底层的通信基准
nccl-tests 是 NVIDIA 官方出的。说白了，它只测一件事：GPU 之间的带宽和延迟。
它不关心你跑什么模型，也不管你用哪个框架。它只管发数据、收数据、算时间。
核心价值：验证集群的物理通信能力是否达标。
我在项目中遇到过一件事。有一次新集群上线，跑 DeepSpeed 死活上不去吞吐。折腾了两天，最后用 nccl-tests 一测，发现跨节点带宽只有理论值的 30%。原来是 IB 网卡配置错了。你看，要是早点跑这个，能省多少时间。
用法很简单：
# 单节点 8 卡 allreduce 测试
./build/all_reduce_perf -b 8 -e 128M -f 2 -g 8
# 跨节点测试（需要 mpirun）
mpirun -np 16 -hostfile hosts ./build/all_reduce_perf -b 8 -e 128M -f 2 -g 8输出结果里，我主要看两个指标：algbw（算法带宽）和 busbw（总线带宽）。如果 busbw 低于理论值的 80%，那肯定有问题。
我的习惯：每次集群变更后，先跑一轮 nccl-tests 全家桶。包括 allreduce、allgather、reducescatter。全部通过再上框架。
2. DeepSpeed：微软出品的框架级压测
DeepSpeed 的压测思路不一样。它不测裸带宽，而是测「在 ZeRO 优化下，模型能跑多快」。
我个人觉得，DeepSpeed 的压测工具更适合做「容量规划」。比如你想知道：用 ZeRO-3 跑 175B 的模型，需要多少节点？每节点能塞多少参数？
它内置了一个 benchmark 脚本：
# 在 DeepSpeed 目录下
python benchmark.py --model bert-large --dtype fp16 --num_gpus 8输出会告诉你：
- 每 GPU 的吞吐（tokens/sec）
- 显存占用
- 通信开销占比
我曾经用这个工具帮一个客户做选型。他们纠结用 8 卡 A100 还是 16 卡 A10G。跑完 benchmark 发现，A10G 虽然卡多，但通信效率低，最终吞吐反而不如 8 卡 A100。嗯，数据说话。
注意：DeepSpeed 的 benchmark 默认用 HuggingFace 模型。如果你用自定义模型，需要自己写一个简单的 forward/backward 脚本。别偷懒，我见过有人直接拿默认结果去汇报，结果跟实际差 30%。
3. Megatron-LM：模型级压测的标杆
Megatron-LM 是 NVIDIA 的。它做压测的方式最「重」——直接跑一个完整的 GPT 或 BERT 训练流程。
为什么说它是标杆？因为它同时压了三个维度：
- 计算：矩阵乘法、注意力机制
- 通信：张量并行、流水线并行
- 显存：激活值、优化器状态
它的压测脚本长这样：
# 跑一个 1.3B 参数的 GPT 模型
bash examples/gpt3/train_gpt3_1.3b_distributed.sh \
--tensor-model-parallel-size 2 \
--pipeline-model-parallel-size 2 \
--num-layers 24 \
--hidden-size 2048 \
--num-attention-heads 32输出结果里，我最关注 tokens per second per GPU。这个指标能直接反映集群的「有效算力利用率」。
我记得有一次，用 Megatron 压测发现吞吐波动很大。排查了半天，发现是 NVLink 带宽被其他进程抢占了。嗯，这种问题 nccl-tests 测不出来，DeepSpeed 也测不出来。只有跑完整模型才能暴露。
4. 三者的对比总结
| 维度 | nccl-tests | DeepSpeed | Megatron-LM |
|---|---|---|---|
| 测试层级 | 通信层 | 框架层 | 模型层 |
| 测试内容 | 带宽、延迟 | ZeRO 优化下的吞吐 | 完整训练流程 |
| 耗时 | 几分钟 | 十几分钟 | 几小时 |
| 适用场景 | 集群验收、故障排查 | 容量规划、框架选型 | 性能调优、上线前验证 |
| 依赖 | 仅 NCCL | DeepSpeed + PyTorch | Megatron + PyTorch + Apex |
5. 我的选型建议
别想着只用一种工具。我习惯分三步走：
- 第一步：nccl-tests 快速验证。确保物理链路没问题。这一步不过，后面都是白搭。
- 第二步：DeepSpeed benchmark 做容量估算。确定能跑多大的模型、用多少资源。
- 第三步：Megatron-LM 跑完整压测。模拟真实训练场景，暴露隐藏问题。
避坑指南：我曾经跳过第一步，直接上 Megatron。结果跑了一天发现吞吐不对，回头一查是网卡固件版本不匹配。白白浪费一天算力。所以，别嫌麻烦，按顺序来。
6. 核心逻辑图
下面这张图展示了三个工具的定位和关系。我建议你保存下来，以后做方案时对照着看。
你看，这三层是递进关系。底层不过，上层跑得再花哨也没用。我见过太多人一上来就调 Megatron 的参数，结果发现是网线没插好。嗯，别做那种事。
好了，工具选型就聊到这儿。记住一个原则：先通后快，先简后繁。下次你拿到一个新集群，按这个顺序来，不会错。