性能调优指南

先测量，别瞎猜

性能调优最容易犯的错：凭感觉猜瓶颈。有人一慢就换显卡，结果瓶颈在数据加载。正确的流程只有三步：测量、定位、优化，然后回到第一步验证。

测量工具有三件：

• torch.profiler（§32 讲过）：看每个算子的耗时分布，定位计算瓶颈；
• torch.utils.benchmark：精确对比两段代码谁快；
• 内存快照：查显存都花在哪了。

定位计算瓶颈用 torch.profiler（§32）：跑几十步训练，看表格里哪个算子耗时最长。经验上，Transformer 模型里 attention 和线性层应该占大头；如果看到大量 to、copy_ 或者数据传输类算子排前列，说明问题在数据搬运而不是计算，先查 DataLoader 和 dtype 是否一致。

基准测试的正确姿势

torch.utils.benchmark.Timer 用法和 Python 的 timeit 类似，但有两个关键改进：自动预热、自动同步 CUDA。

import torch.utils.benchmark as benchmark

def batched_dot_mul_sum(a, b):
    return a.mul(b).sum(-1)

def batched_dot_bmm(a, b):
    a = a.reshape(-1, 1, a.shape[-1])
    b = b.reshape(-1, b.shape[-1], 1)
    return torch.bmm(a, b).flatten(-3)

x = torch.randn(10000, 64, device="cuda")

t0 = benchmark.Timer(stmt="batched_dot_mul_sum(x, x)", globals={"x": x})
t1 = benchmark.Timer(stmt="batched_dot_bmm(x, x)", globals={"x": x})

print(t0.blocked_autorange())
print(t1.blocked_autorange())

blocked_autorange 自动决定跑多少轮，直到测够时间，返回中位数等统计量。多组结果用 Compare 汇总成表格，按输入形状批量对比：

results = []
for n in [64, 1024, 10000]:
    x = torch.randn(n, 64)
    for name, fn in [("mul_sum", batched_dot_mul_sum), ("bmm", batched_dot_bmm)]:
        results.append(benchmark.Timer(
            stmt="fn(x, x)", globals={"x": x, "fn": fn},
            label="batched_dot", sub_label=f"[{n}, 64]", description=name,
        ).blocked_autorange())

benchmark.Compare(results).print()

两个坑记住。一是必须预热：cuBLAS 首次调用要加载库，不预热直接测会虚高。二是 GPU 上要用 benchmark 模块而不是裸 timeit——后者不同步 CUDA，测的是「启动内核」的时间，不是真实耗时。

通用优化清单

下面这些改动小、收益稳，按性价比排：

1. 数据加载异步化：DataLoader(num_workers=4, pin_memory=True)。默认 num_workers=0 是同步加载，训练要干等数据。多开几个子进程，再开 pin_memory 加速拷贝到 GPU。

2. 验证/推理关梯度：包一层 torch.no_grad()，省掉中间缓冲的保存，又快又省显存。

3. BatchNorm 前面的卷积去掉 bias：nn.Conv2d(..., bias=False)。BatchNorm 第一步就减均值，bias 会被抵消，纯属白算。

4. 梯度清零用 set_to_none：optimizer.zero_grad(set_to_none=True)，比置零少一次内存写，反向时用赋值代替累加。

5. 算子融合：逐点运算（加法、激活函数这类）用 torch.compile 一键融合，少启动内核、少读写内存。第 46、47 章讲过。注意输入形状多变时缓存会反复重建，推理服务里尤其明显。

6. 关闭调试 API：detect_anomaly、profiler、gradcheck 这些只该在调试时开，训练时开着都是纯开销。

GPU 专项优化

**混合精度（AMP）**是最划算的 GPU 加速。Tensor Core 专吃低精度矩阵乘，第 54 章讲过。记得把张量维度凑成 8 的倍数，方便 Tensor Core 对齐。

CUDA Graph 把一串内核启动变成一张图，省掉 CPU 反复启动内核的开销。torch.compile(model, mode="reduce-overhead") 就能启用，代价是显存占用变高。

cuDNN 自动调优器：卷积网络训练前加一行：

torch.backends.cudnn.benchmark = True

它会先跑个小基准，挑最快的卷积算法。输入尺寸变化大的场景别开，选算法本身有开销。

避免 CPU-GPU 同步。print(cuda_tensor)、.item()、.cpu() 都会让 CPU 停下来等 GPU，训练循环里尽量少用。损失打印用 .item() 是少数可接受的例外，但也别太频繁。

直接在目标设备建张量：torch.rand(1000, device="cuda") 而不是 torch.rand(1000).cuda()，少一次拷贝。

变长输入预分配：NLP 或语音场景里，短批次后跟着长批次，会让缓存分配器反复分配、碎片化，严重时直接 OOM。先用最大序列长度跑一次前向和反向（不更新参数），把最大缓冲区预分配好，再开始正式训练。

显存分析：Mosaic

调优不只看速度，还要看显存。用官方内存快照 API 抓一份快照，交给 Mosaic 工具做分类分析：

torch.cuda.memory._record_memory_history(max_entries=100000)
# ... 训练代码 ...
snapshot = torch.cuda.memory._snapshot()

快照里存着每一次分配的记录。用 mosaic_get_memory_profile --snapshot <文件> --profile categories 生成 HTML 报告，把内存按激活、梯度、优化器状态、参数分好类；mosaic_get_memory_usage_peak 则能看出峰值是哪段代码分配的。

官方教程有个很直观的案例：GPT-2 训练开启激活检查点（activation checkpointing）后，峰值显存从 4.62 GB 降到 2.55 GB，其中激活内存从 2.93 GB 降到 873 MB，省了 71%。激活检查点不存中间激活，反向时重算，用算力换显存：

model.gradient_checkpointing_enable()   # 部分模型直接支持

分布式优化

多卡训练还有几个专项技巧。

梯度累加跳过通信：用 DDP 做梯度累加时，前 N-1 次反向不需要 all-reduce，包在 no_sync() 里，最后一次再同步：

for i, batch in enumerate(dataloader):
    if i % accum_steps != accum_steps - 1:
        with model.no_sync():          # 跳过梯度同步
            loss = model(batch).loss
            loss.backward()
    else:
        loss = model(batch).loss
        loss.backward()
        optim.step()
        optim.zero_grad()

负载均衡：序列数据长短不一，一个 batch 里某张卡分到的序列特别长，所有卡都得等它。解法是按序列长度分桶打乱（bucket），或者按 token 数大致恒定来组 batch。

CPU 侧还有 OMP_NUM_THREADS 线程数、numactl 内存绑定、jemalloc 内存分配器这些招。多数情况默认就够，遇到 CPU 瓶颈再研究。

最后提醒一句：优化完记得测端到端时间，而不是只看单算子。有时候省了 30% 的算子耗时，因为瓶颈在别处，整体只快了 3%。

小结

调优心法一句话：测量先行，逐项验证。优先级建议：数据加载和 no_grad 这类基础项先查，然后 AMP 加 torch.compile 白拿加速，显存紧张上激活检查点，多卡再看通信。到这里，分布式和性能的拼图齐了。最后一章，把 60 章串成一张学习地图。