3. 基准测试工具：PyTorch Profiler、TensorBoard、NVIDIA DCGM、DeepSpeed Profiler 的安装与基本使用
做AI训练吞吐压测，说白了就是跟时间赛跑。你得知道每一毫秒花在哪了，GPU到底有没有偷懒。我见过太多团队，模型能跑起来就觉得万事大吉，结果一压测才发现，GPU利用率不到30%。嗯，这就是我们今天要解决的问题。
这一章，我带你装好四把趁手的兵器：PyTorch Profiler、TensorBoard、NVIDIA DCGM、DeepSpeed Profiler。它们各有各的脾气，但组合起来，能让你把训练过程看得透透的。
3.1 PyTorch Profiler：最亲民的性能放大镜
PyTorch Profiler 是 PyTorch 官方自带的性能分析工具。我个人习惯把它当作第一道防线——先用它跑一遍，看看有没有明显的瓶颈。
安装
从 PyTorch 1.8 开始，Profiler 就内置在 torch.profiler 模块里了。你只要装好 PyTorch 就行，不需要额外折腾。
# 确认版本
import torch
print(torch.__version__) # 建议 >= 1.8
# 如果版本太老，升级一下
pip install torch --upgrade基本使用
我写了个最简示例，你直接复制就能跑。注意看注释，每一行都有讲究。
import torch
import torch.nn as nn
from torch.profiler import profile, record_function, ProfilerActivity
# 造个简单的模型
model = nn.Linear(100, 10).cuda()
input_data = torch.randn(32, 100).cuda()
# 开始分析
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
profile_memory=True,
with_stack=True
) as prof:
with record_function("model_inference"):
output = model(input_data)
loss = output.sum()
loss.backward()
# 打印结果
print(prof.key_averages().table(
sort_by="cuda_time_total",
row_limit=10
))profile_memory=True。我在项目中遇到过好几次，模型跑得慢不是因为计算，而是因为显存频繁分配释放。这个参数能帮你揪出内存碎片问题。
关键参数解读
| 参数 | 作用 | 我建议的设置 |
|---|---|---|
| activities | 指定监控CPU还是CUDA | 两个都开，别省 |
| record_shapes | 记录张量形状 | True，方便定位形状不匹配 |
| profile_memory | 记录内存使用 | True，排查显存泄漏必备 |
| with_stack | 记录调用栈 | 调试时开，正式压测关掉 |
3.2 TensorBoard：把枯燥的数字变成图
光看表格不过瘾，对吧？TensorBoard 能把 Profiler 的数据可视化。你想想看，一条曲线比一百个数字直观多了。
安装
pip install tensorboard导出 Profiler 数据到 TensorBoard
from torch.profiler import tensorboard_trace_handler
# 在 profiler 中指定输出路径
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
on_trace_ready=tensorboard_trace_handler("./logs")
) as prof:
# 你的训练代码
pass
# 启动 TensorBoard
# tensorboard --logdir=./logs./logs/run_20241001。
3.3 NVIDIA DCGM：GPU 硬件的体检医生
PyTorch Profiler 看的是软件层面，但 GPU 硬件本身有没有问题？温度是不是过高？功耗有没有被限制？这时候就得请出 DCGM 了。
安装
# Ubuntu/Debian
sudo apt-get install datacenter-gpu-manager
# 或者用容器
docker run --gpus all -d --name dcgm nvidia/dcgm:latest
# 启动服务
sudo systemctl start nvidia-dcgm基本使用
DCGM 提供了命令行工具 dcgmi，我最常用的是监控模式：
# 实时监控，每秒刷新
dcgmi monitor -e 1001,1002,1003,1004
# 参数说明：
# 1001 - GPU 利用率
# 1002 - 显存利用率
# 1003 - 温度
# 1004 - 功耗你还可以用 Python 调用 DCGM 的 API：
import pydcgm
from dcgm_structs import dcgmException
# 连接到 DCGM 服务
handle = pydcgm.DcgmHandle()
group = handle.GetGroupWithAllDevices()
# 获取 GPU 利用率
for gpu_id in range(group.GetInfo()['entityCount']):
gpu = group.GetEntityGroupHandle(gpu_id)
stats = gpu.GetStats()
print(f"GPU {gpu_id}: 利用率 {stats.gpuUtil}%")3.4 DeepSpeed Profiler：分布式训练的专属工具
当你开始用多卡训练，事情就复杂了。通信开销、负载不均、梯度同步延迟……DeepSpeed Profiler 就是为这个场景设计的。
安装
pip install deepspeed
# 验证安装
ds_report基本使用
DeepSpeed Profiler 不需要改你的训练代码，只要在启动命令里加个参数：
# 单节点多卡
deepspeed --num_gpus=8 train.py \
--deepspeed_config ds_config.json \
--profile
# 多节点
deepspeed --num_gpus=8 --num_nodes=2 \
train.py --deepspeed_config ds_config.json --profile它会自动生成一个 profile_results.json，里面包含了：
- 前向/反向传播时间分布
- 通信耗时（All-Reduce、All-Gather 等）
- 各 GPU 负载均衡情况
- ZeRO 优化阶段的内存节省
3.5 知识体系总览
下面这张图，是我自己总结的。四个工具各司其职，覆盖了从单卡到多卡、从软件到硬件的全链路。
3.6 我的实战建议
工具装好了，怎么用才高效？我分享几个经验：
- 先单卡后多卡：别一上来就搞分布式。先用 PyTorch Profiler 把单卡跑透，再上 DeepSpeed Profiler 看多卡通信。
- DCGM 常驻后台：我习惯在压测全程开着 DCGM 监控。有一次发现 GPU 温度飙到 92°C，赶紧停了任务检查散热，不然卡就烧了。
- TensorBoard 做对比：每次改完参数，把日志存到不同目录。用 TensorBoard 对比曲线，一眼就能看出哪个改动有效。
- 别贪多：刚开始压测，先盯住 GPU 利用率和吞吐量这两个指标。其他的等遇到瓶颈再细看。
好了，工具都装好了，也跑通了。下一章我们就要动真格的了——用这些工具去压测一个真实的训练任务。到时候你会发现，前面这些准备工作，每一分钟都值。