第7章：PyTorch Profiling 与 Tuning

CUDA kernel 写得再好，框架调度不对也跑不快。这一章把 Ch5-6 的”硬件视角”和上层 PyTorch 框架层串起来——从 profiler 看 timeline、torch.compile 工作机制、CUDA Graph 决策、AMP 精度选型、内存布局调优，到 DataLoader 系列参数——给出一份可直接照着改的 PyTorch 性能调优决策树。

📑 目录

1. PyTorch Profiler：框架层的时间线
2. torch.compile 工作机制
3. CUDA Graph：什么时候开、什么时候不开
4. 精度决策树：TF32 / BF16 / FP16 / FP8
5. AMP 默认策略与坑
6. Memory snapshot 与 OOM 排查
7. channels_last 内存布局
8. DataLoader 调参清单
9. 一份 PyTorch 调优决策树
自我检验清单
参考资料

1. PyTorch Profiler：框架层的时间线

1.1 PyTorch Profiler vs Nsight Systems

工具	视角	强项
Nsight Systems	OS / 硬件	看 CPU/GPU/通信全貌
PyTorch Profiler	框架层	看 op 级 / module 级时间分布

两者互补：通常先用 PyTorch Profiler 找出最慢的 op，再用 Nsight Systems 看 op 之间的间隙是不是被 CPU/IO 占了。

1.2 基本用法

from torch.profiler import profile, record_function, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step, batch in enumerate(loader):
        with record_function("training_step"):
            train_step(batch)
        prof.step()

输出可以直接打开 TensorBoard 看 timeline。

1.3 该看什么

打开 TensorBoard Profiler 后按这个顺序看：

1. Overview 页

GPU Kernel time / CPU time / Other time 比例
“GPU Idle Time” → 严重时说明 GPU 在等
推荐操作：直接列出”top 10 GPU ops by cuda time”——找最大头

2. Operator 页

每个 op 的累计时间、调用次数
self CPU time vs cumulative CPU time（包含子 op）
shape 列：检查每个 op 实际的输入 shape

3. Trace 页

时间线视图——可以放大看 step 内每一段
CPU 行 + CUDA stream 行同时显示
通过 NVTX 标签可以加自定义层次

1.4 实战经验

只 profile 几个 step：activity 设 3-5 步就够，profile 本身有开销
配合 NVTX：在自己代码里加 torch.cuda.nvtx.range_push("xxx")，让时间线更可读
memory profiling：profile_memory=True 能看每个 op 分配 / 释放多少显存——OOM 排查必备

2. torch.compile 工作机制

2.1 三件套：Dynamo + AOTAutograd + Inductor

torch.compile(model) 背后是三个组件：

Python 代码
    │
    ▼
┌─────────────────────────────────────┐
│ Dynamo                              │
│ - 在 Python bytecode 层面 trace    │
│ - 把 model 切成 "fx graph" 子段     │
│ - 遇到 Python 控制流就 graph break  │
└──────────────┬──────────────────────┘
               │
               ▼
┌─────────────────────────────────────┐
│ AOTAutograd                         │
│ - 把 forward graph 自动加上 backward│
│ - 输出 forward + backward 的整图    │
└──────────────┬──────────────────────┘
               │
               ▼
┌─────────────────────────────────────┐
│ Inductor                            │
│ - 后端编译器：把 fx graph → Triton  │
│ - 自动做 fusion / vectorize / tiling│
│ - 生成 CUDA / CPU / ROCm kernel    │
└─────────────────────────────────────┘

2.2 三种编译模式

model = torch.compile(model, mode="default")        # 平衡
model = torch.compile(model, mode="reduce-overhead") # 用 CUDA Graph
model = torch.compile(model, mode="max-autotune")   # 极致优化（编译慢）

模式	编译时间	运行性能	推荐场景
default	几十秒	1.2-1.5x	大多数训练
reduce-overhead	1-2 分钟	1.5-2x	推理 / 小 batch
max-autotune	10-30 分钟	2-3x	长跑生产任务

2.3 Graph break：torch.compile 的最大杀手

Dynamo 遇到不可 trace 的代码会 graph break——把图切成多段，每段单独编译，段之间退到 eager mode：

触发 graph break 的常见操作	替代方案
`if x.item() > 0`	用 mask + select
`print(x)`	删掉，profiling 用 NVTX
自定义 `torch.autograd.Function`（旧 API）	用 `@torch.compile(fullgraph=True)` 强制报错定位
`.tolist()` / `.cpu()` 中间出现	推到模型外

调试 graph break：

import torch._dynamo as dynamo
dynamo.config.verbose = True
torch.compile(model)(input)
# log 会列出所有 graph break 的位置和原因

2.4 一个真实陷阱：dynamic shape

如果输入 shape 经常变（比如不同长度的序列），torch.compile 会反复重新编译——每个新 shape 一次。两条解决路径：

Padding 到固定 shape：丢一点效率换稳定性
torch.compile(model, dynamic=True)：让编译器生成支持 dynamic 的 kernel（性能略差但稳定）

3. CUDA Graph：什么时候开、什么时候不开

3.1 CUDA Graph 是什么

CUDA Graph 把一段固定的 GPU 操作序列录制成一张图，每次 launch 这张图只需要一次系统调用——把 N 个 kernel launch 的 CPU 开销摊薄到 1 个。

适合：

训练 / 推理的核心 step（每次 step 操作序列固定）
batch size 固定、shape 固定的负载

3.2 怎么用 CUDA Graph

PyTorch 里有几种方式：

1. torch.compile(mode="reduce-overhead")：自动开 CUDA Graph

2. 手动 record + replay：

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    output = model(static_input)
    loss = criterion(output, static_target)
    loss.backward()

# 后续每个 step 只需 replay
for batch in loader:
    static_input.copy_(batch)
    g.replay()
    optimizer.step()

3. torch.cuda.make_graphed_callables：自动版本

3.3 什么时候不开

CUDA Graph 不适合：

❌ Shape / control flow 经常变的负载（每次重 record 反而慢）
❌ kernel 单次很大的负载（launch 开销占比小）
❌ debug 阶段（graph 内部不能 print）

🌟 经验：CUDA Graph 适合单步 kernel 数 > 50 且每个 kernel < 100 µs的场景——典型是小 batch size 推理。

4. 精度决策树：TF32 / BF16 / FP16 / FP8

4.1 五种精度的对照

精度	指数 / 尾数位	数值范围	精度损失	Tensor Core 加速
FP32	8 / 23	极宽	无	1x（基线）
TF32	8 / 10	同 FP32	中	2x
BF16	8 / 7	同 FP32	大	4x
FP16	5 / 10	较窄	小	4x
FP8	4-5 / 2-3	极窄	巨大	8x（H100）

4.2 决策树

你在做什么？
├─ 训练
│  ├─ 默认：BF16 + FP32 master weight (AMP)
│  ├─ 不稳定（loss spike）：FP16 + FP32 master weight + GradScaler
│  ├─ H100+ 上的大模型：FP8（用 transformer-engine）
│  └─ 全精度调试：FP32
└─ 推理
   ├─ 默认：BF16 / FP16
   ├─ 极致延迟：FP8 (H100+)
   ├─ 量化部署：INT8 / INT4
   └─ 高精度业务（金融 / 医疗）：FP32 / TF32

4.3 BF16 vs FP16 的选择

BF16 数值范围和 FP32 一致——几乎不会 overflow / underflow——适合大部分训练
FP16 数值范围窄——可能 overflow，但精度位多——适合特定 inference / 老硬件
2024+ 默认选 BF16

4.4 FP8 的工程现实

H100 引入 FP8 (E4M3 + E5M2 两种格式)。开 FP8 的方法：

from transformer_engine.pytorch import Linear, fp8_autocast

with fp8_autocast():
    output = my_linear(input)

实测在大模型训练上能拿到 1.3-1.6x 加速 vs BF16——但需要 H100、需要 transformer-engine、需要 calibration。

5. AMP 默认策略与坑

5.1 AMP 是什么

AMP (Automatic Mixed Precision) 让一段计算自动用低精度跑：

from torch import autocast, GradScaler

scaler = GradScaler()
for batch in loader:
    with autocast(device_type="cuda", dtype=torch.bfloat16):
        output = model(batch)
        loss = criterion(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

autocast 自动决定每个 op 用什么精度。

5.2 AMP 内部的精度规则

PyTorch AMP 把 op 分三类：

Always FP32：torch.softmax / log / exp / pow 等数值敏感的
Always BF16/FP16：matmul / conv 等吞吐密集的
Match input：随上下游而定

工程含义：不要担心 AMP 把所有东西都 cast 到 FP16——它会自动保留数值敏感 op 的精度。

5.3 三个常见的 AMP 坑

坑 1：自定义 op 不在 AMP 列表

如果你写了 custom CUDA op，AMP 默认会按”match input”处理——可能导致它在 BF16 输入上跑——而你的 kernel 没考虑 BF16。

对策：

@torch.amp.custom_fwd(cast_inputs=torch.float32)
def my_op_forward(...):
    ...

坑 2：Loss 是 NaN

FP16 的 overflow 是常见原因。诊断：

看 loss 是不是几个 step 后才变 NaN
看 GradScaler 的 scale 值（连续报告 inf 时会自动减半）
解决方法：换 BF16

坑 3：不需要 GradScaler 的 BF16

GradScaler 是 FP16 的伴生工具，BF16 范围足够大不需要。用 BF16 时不要用 GradScaler——会引入 host-device sync。

6. Memory snapshot 与 OOM 排查

6.1 PyTorch 显存模型

PyTorch 有自己的显存分配器（caching allocator）——不是直接调 cudaMalloc。这意味着：

nvidia-smi 看到的”显存使用”≠ 实际程序真正用的
PyTorch 会保留释放的显存以备下次用

6.2 Memory snapshot 用法

torch.cuda.memory._record_memory_history(max_entries=10000)

# 跑你的训练 / 推理代码
train()

# 导出 snapshot
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

然后用 PyTorch 提供的可视化工具打开 pickle——能看到：

每个 tensor 的分配 / 释放时间线
哪些 tensor 一直没被释放（潜在泄漏）
peak memory 时谁占用最大

6.3 OOM 的三类常见根因

根因	表现	排查方法
Tensor 泄漏	显存随 step 数增长	snapshot 看哪些 tensor 长期不释放
中间激活值过大	单 step 内 peak 超	看 forward 阶段哪个 op 输出最大
优化器状态过大	模型和优化器各占一半	用 ZeRO / FSDP 分摊

6.4 减少显存的常见手段

按”好用 → 改动大”排序：

gradient checkpointing：用 30% 计算换 50%+ 激活值显存
AMP：BF16 直接砍一半激活值显存
torch.compile(mode="reduce-overhead")：编译时做 memory planning
fused optimizer：optim.AdamW(fused=True) 减少中间 buffer
batch size halving：最朴素但最 reliable
FSDP / ZeRO：分布式分摊优化器状态和模型参数

7. channels_last 内存布局

7.1 NCHW vs NHWC

PyTorch 默认用 NCHW（通道在前），但 cuDNN / Tensor Core 在 NHWC 上更快——尤其卷积。

# 改成 channels_last (NHWC)
model = model.to(memory_format=torch.channels_last)
input = input.to(memory_format=torch.channels_last)

# forward 走 channels_last 路径
output = model(input)

7.2 性能收益

卷积训练 / 推理：通常 1.3-2x 加速
ResNet-50 / EfficientNet 这种典型 CV 模型：直接开 channels_last
对 transformer：基本没用（没有空间维度）

7.3 一个常见坑

并不是所有 op 都支持 channels_last——如果中间某个 op 不支持，PyTorch 会自动 cast 回 NCHW，然后输出又 cast 回 NHWC——两次 layout 转换比不开还慢。

调试：用 PyTorch Profiler 看是不是有大量 _to_copy op——那就是 layout 转换的体现。

8. DataLoader 调参清单

8.1 关键参数

loader = DataLoader(
    dataset,
    batch_size=128,
    num_workers=8,           # CPU worker 数
    pin_memory=True,         # 用 pinned memory 加速 H2D
    prefetch_factor=2,       # 每 worker 预读 2 个 batch
    persistent_workers=True, # epoch 之间 worker 不重启
    drop_last=True,          # 避免最后一个不完整 batch
)

8.2 num_workers 的经验值

CPU 数 ≥ GPU 数 × 4：num_workers = CPU 数
CPU 数有限：num_workers = max(2, CPU/2)
从 0 开始：先验证模型本身正确，再开 workers

8.3 pin_memory + non_blocking

input = batch.to(device, non_blocking=True)  # H2D 异步

pin_memory=True 配合 non_blocking=True 能让 H2D 拷贝异步进行——和上一 step 的 GPU 计算重叠。

8.4 prefetch_factor

每个 worker 预先读 prefetch_factor 个 batch 准备着。太小（=1）→ workers 经常 idle；太大（>4）→ 内存占用大。默认 2 通常够。

9. 一份 PyTorch 调优决策树

把上面所有内容压缩成一份可直接照着改的决策树：

你的模型 / 训练 / 推理慢吗？
│
├─ 先用 PyTorch Profiler 看
│   ├─ GPU Idle Time 高 → DataLoader / IO / CPU 是瓶颈
│   │   ├─ num_workers ↑ + pin_memory + prefetch_factor
│   │   ├─ 数据预处理移到 GPU（torchvision.transforms.v2）
│   │   └─ 文件 IO 调优（NVMe / RDMA / shm）
│   │
│   └─ GPU 利用率高但慢
│       ├─ 用 Nsight Systems 看 op 间隙
│       └─ 用 Nsight Compute 看具体 kernel
│
├─ 模型本身可优化
│   ├─ 默认开 AMP（BF16）
│   ├─ 默认开 channels_last（CV 模型）
│   ├─ 训练：torch.compile(mode="default")
│   ├─ 推理：torch.compile(mode="reduce-overhead") 或 CUDA Graph
│   └─ 长跑生产：torch.compile(mode="max-autotune")
│
├─ 显存压力
│   ├─ AMP / FP16 / BF16
│   ├─ Gradient checkpointing
│   ├─ Fused optimizer
│   ├─ 减 batch size
│   └─ 上 FSDP / DeepSpeed
│
└─ 还慢
    └─ 自己写 kernel（Ch6 + Triton）

🎯 自我检验清单

PyTorch Profiler 和 Nsight Systems 在调优工作中各自的角色是什么？
torch.compile 的三件套（Dynamo / AOTAutograd / Inductor）各自做什么？什么是 graph break？
BF16 和 FP16 的区别是什么？在训练中默认选哪个？为什么？
AMP 内部的”Always FP32 / Always BF16 / Match input”三类 op 划分背后的工程理由是什么？
DataLoader 慢的几种典型表现是什么？num_workers / pin_memory / prefetch_factor 各自调什么？

📚 参考资料

PyTorch Profiler Tutorial：pytorch.org/tutorials/recipes/recipes/profiler_recipe.html
torch.compile 文档：pytorch.org/docs/stable/torch.compiler.html
AMP 文档：pytorch.org/docs/stable/amp.html
CUDA Graph in PyTorch：pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/
NVIDIA Transformer Engine（FP8）：github.com/NVIDIA/TransformerEngine
本系列模块一第 4 章：PyTorch 框架基础

下一章预告：Ch8 把视角从”单卡 / 单机”扩到”万卡集群”——3D 并行、ZeRO/FSDP、MoE、Sequence Parallelism、容错——给一份完整的超大规模分布式训练运维指南。

搜索