3. 压测工具选型

说到压测工具选型，我见过不少团队一上来就选最火的工具，结果跑起来才发现水土不服。其实说白了，没有万能工具，只有最合适的组合。今天我就把几个主流工具掰开揉碎讲清楚，帮你少走弯路。

3.1 主流工具速览

目前业界常用的GPU集群压测工具，我按使用场景分成了四类：

工具	核心定位	我推荐的使用场景
NCCL Tests	通信原语基准测试	网络带宽、延迟、拓扑验证
DeepSpeed	分布式训练框架	ZeRO优化、大模型训练压测
Megatron-LM	模型并行训练框架	张量/流水线并行压测
MLPerf	行业标准基准	横向对比、合规性验证
自定义脚本	灵活定制	特定场景、故障注入

3.2 NCCL Tests：通信层面的“听诊器”

我个人习惯，拿到一个新集群，第一件事就是跑NCCL Tests。为什么？因为它能直接暴露网络问题。

NCCL Tests主要测这几个东西：

• all_reduce：最常用的集合通信操作
• all_gather：数据收集性能
• reduce_scatter：数据分发性能
• broadcast：广播延迟

举个例子，跑个8卡all_reduce测试：

# 单节点8卡测试
mpirun -np 8 \
  -host localhost:8 \
  /path/to/nccl-tests/build/all_reduce_perf \
  -b 8 -e 128M -f 2 -g 1

# 跨节点测试（4节点，每节点8卡）
mpirun -np 32 \
  -host node1:8,node2:8,node3:8,node4:8 \
  /path/to/nccl-tests/build/all_reduce_perf \
  -b 8 -e 128M -f 2 -g 1

我在项目中遇到过，某次跨节点测试带宽只有理论值的60%。排查了半天，发现是网卡中断绑定没做好。嗯，这种问题NCCL Tests一跑就现原形。

3.3 DeepSpeed：ZeRO优化的“试金石”

DeepSpeed是微软开源的分布式训练框架，它的压测价值在于验证ZeRO系列优化是否生效。

我个人建议重点关注：

• ZeRO-1：优化器状态分片，适合小规模集群
• ZeRO-2：梯度分片，中等规模集群首选
• ZeRO-3：参数分片，大模型必备

跑个简单的压测脚本：

# deepspeed 压测示例
deepspeed --num_gpus=8 \
  --num_nodes=4 \
  train.py \
  --model_name -large \
  --batch_size 16 \
  --zero_stage 3 \
  --gradient_accumulation_steps 4

3.4 Megatron-LM：模型并行的“压力锅”

Megatron是NVIDIA的模型并行框架，专门压测张量并行和流水线并行。说实话，这个工具对网络要求极高。

核心压测维度：

• 张量并行（TP）：单节点内GPU间通信，依赖NVLink
• 流水线并行（PP）：跨节点通信，依赖RDMA网络
• 数据并行（DP）：梯度同步，依赖all_reduce

一个典型的压测命令：

# Megatron 压测示例
python pretrain_ \
  --tensor-model-parallel-size 4 \
  --pipeline-model-parallel-size 4 \
  --num-layers 40 \
  --hidden-size 5120 \
  --num-attention-heads 32 \
  --micro-batch-size 4 \
  --global-batch-size 1024

3.5 MLPerf：行业标准的“裁判员”

MLPerf是MLCommons推出的行业基准。如果你需要向客户或领导证明集群性能，MLPerf是最有说服力的。

它包含多个基准：

• 训练（Training）：BERT、ResNet、等模型
• 推理（Inference）：在线/离线场景
• 存储（Storage）：数据加载性能

3.6 自定义脚本：最后的“杀手锏”

当现成工具满足不了需求时，就得自己动手了。我经常写自定义脚本做这几件事：

• 故障注入：模拟网卡故障、GPU掉卡
• 混合负载：同时跑训练和推理
• 长时间稳定性：跑72小时以上看有没有内存泄漏

一个简单的自定义压测脚本框架：

import torch
import torch.distributed as dist
import time

def custom_benchmark():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    
    # 自定义通信模式
    tensor = torch.randn(1024, 1024).cuda()
    
    # 模拟混合精度训练
    for i in range(1000):
        # all_reduce
        dist.all_reduce(tensor)
        
        # 模拟计算
        torch.mm(tensor, tensor)
        
        if rank == 0 and i % 100 == 0:
            print(f"Iter {i}: bandwidth = {measure_bandwidth()}")
    
    dist.destroy_process_group()

3.7 选型策略：场景决定工具

说了这么多，到底怎么选？我总结了一个简单的决策树：

3.8 避坑指南

最后分享几个我踩过的坑：

嗯，工具选型就聊到这儿。记住一个原则：先通信，再训练，最后基准。按这个顺序来，基本不会出大问题。

---