PyTorch 训练框架内核深度解析

第21章 Profiler 与性能调优

作者 杨艺韬 · 14,156 字 · 发布于 · 更新于

"An experimental config for Kineto features. Please note that backward compatibility is not guaranteed."

—— torch/csrc/profiler/python/init.cpp:384 对 _ExperimentalConfig 的说明；profiler 最有用的那些开关，恰恰都在这行警告后面

本章要点

  • torch.profiler.profile 是 PyTorch 内置 profiler，能同时记录 CPU 函数调用、CUDA kernel 执行、显存分配、autograd 事件
  • 底层是 Kineto + CUPTI：Kineto 是 PyTorch 自家的库（third_party/kineto 指向 github.com/pytorch/kineto），CUPTI 才是 NVIDIA 的 Profiling Tools Interface
  • 三种主要输出：text summary（key_averages()）、chrome trace（拖进 chrome://tracing 看时间线）、TensorBoard 插件（更友好的可视化）
  • Memory Snapshot：torch.cuda.memory._record_memory_history() 捕获每次 alloc/free 的栈，配 pytorch.org/memory_viz 可视化（第 4 章 §4.11）
  • Schedule API 让 profiler 在多 step 训练里"录制几步、跳过几步、重复"，避免 trace 体积失控
  • record_function 是用户级注解：在自定义代码段插入 marker，让 trace 有可读名字

21.1 训练性能问题的诊断思路

训练性能问题大概分四类：

症状 可能原因
GPU 利用率低（< 80%） DataLoader 跟不上 / 通信瓶颈 / 主机端瓶颈
GPU 利用率高但吞吐低 算子选错 dtype / kernel 调用频繁（调度开销）
OOM 显存爆 activation 太多 / 碎片化 / state 没释放
训练突然慢 某 step 触发隐式同步（如 .item()）、CUDA Graph 失效

诊断这四类问题的标配工具是 profiler。

21.2 一个最简 profiler 用法

from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for i, batch in enumerate(loader):
        loss = model(batch).sum()
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        if i >= 4:
            break

print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=20))
prof.export_chrome_trace('trace.json')

输出表格大致如下（只节选几列，数字为示意；完整表还有 Self CPU %、Self CUDA、profile_memory=True 时的 CPU Mem / CUDA Mem、# of Calls 等列）：

---------------------------  ------------  -------------  ----------
Name                         CUDA total    CUDA time avg  CPU total
---------------------------  ------------  -------------  ----------
aten::mm                     1.234s        12.34ms        0.123s
aten::add                    0.567s        5.67ms         0.056s
aten::layer_norm             0.234s        2.34ms         0.023s
...

sort_by='cuda_time_total' 在 v2.11 里仍可用，排序时会被换成 device_time_total（torch/autograd/profiler_util.py 里 sort_by.replace("cuda", "device")）。

关键开关：

  • record_shapes=True：记录每次算子调用的输入 shape，方便对比"形状相同的 op 加起来多少时间"
  • profile_memory=True：记录每次 alloc/free 的字节数
  • with_stack=True：附带 Python 栈，让你知道某算子是从哪行代码调的（开销更大、慢一些）
  • with_flops=True：按公式估算 matmul / conv 等少数算子的 FLOPs，表里多出一列 Total KFLOPs（单位随量级自动换成 MFLOPs / GFLOPs 等）

注意 prof.step_num 只在你显式调 prof.step() 时才递增（torch/profiler/profiler.py:891），上面这种没有 schedule 的用法里它一直是 0 —— 想按步数停就用普通计数器。要按步录制见 §21.5。

21.3 chrome trace：时间线可视化

prof.export_chrome_trace('trace.json') 输出一个 JSON 文件。把它拖进 chrome://tracing/ 或者 https://ui.perfetto.dev/，能看到：

CPU thread:  [Python] [aten::mm] [aten::relu] [Python] [aten::add] ...
                ↓ launch       ↓ launch
CUDA stream: ........ [mm kernel] [relu kernel] ........ [add kernel]

横轴是时间，纵轴是不同 thread / stream。每个矩形是一个事件，鼠标悬停显示详细信息（duration / args）。

诊断思路：

  • GPU stream 上有大段空白 → CPU 端追不上 GPU，可能是 DataLoader / Python 慢
  • CPU 上有 dispatcher 密集小事件 → 调用太多小算子，考虑 fuse / torch.compile
  • kernel 之间有间隙 → cudaStreamSynchronize 等同步操作（找隐式同步的 .item() / .cpu()）
  • memcpy H2D / D2H 异常多 → 数据没 pin / 频繁跨 device

21.4 Kineto：底层数据采集

这里要澄清一个常见误解：Kineto 是 PyTorch 自家的项目（.gitmodules 里 third_party/kineto 指向 https://github.com/pytorch/kineto），不是 NVIDIA 的。它是一层跨后端的性能数据采集抽象；真正属于 NVIDIA 的是它调用的 CUPTI（CUDA Profiling Tools Interface）。Kineto 经由 CUPTI 拿到：

  • 每个 CUDA kernel 的 launch / completion 时间
  • memcpy 的 H2D / D2H 时间
  • NVTX range 标记
  • GPU 上 SM utilization、DRAM bandwidth 等硬件计数器

PyTorch 的 torch/csrc/autograd/profiler_kineto.cpp（1183 行）把 Kineto 数据接进 PyTorch 的事件 timeline，与 CPU 端的 dispatcher 调用、autograd 事件合并成统一 trace。Kineto 这层抽象也是 ROCm / XPU 等后端能复用同一套 profiler API 的原因（§21.35）。

开销不可忽略：开 profiler 后每个算子多一次 RecordFunction 回调（§21.12），CUPTI 那边还要额外缓冲 GPU 事件。具体拖慢多少取决于算子密度——小算子越多、相对开销越大。结论是一样的：生产代码不要长开 profiler，用 schedule 录几步就停。

21.5 Schedule API：录几步停几步

朴素 profile 整个 epoch，trace 文件会大到根本打不开。Schedule API：

from torch.profiler import schedule

# 注意 schedule 是 keyword-only, wait / warmup / active 三个必填
my_schedule = schedule(
    skip_first=1,          # 最前面跳过 1 步
    wait=1,                # 然后等 1 步不录
    warmup=1,              # 预热 1 步 (开采集但丢弃)
    active=3,              # 真正录 3 步
    repeat=2,              # 循环 2 次 (默认 0 = 一直循环到 profile 结束)
)

with profile(
    activities=[...],
    schedule=my_schedule,
    on_trace_ready=lambda p: p.export_chrome_trace(f'/tmp/trace_{p.step_num}.json'),
) as prof:
    for step, batch in enumerate(loader):
        # 训练循环
        prof.step()    # 必须每步调一次让 schedule 推进

prof.step() 推进 profiler 状态机。ProfilerAction 只有四个取值（torch/profiler/profiler.py:490）：NONE / WARMUP / RECORD / RECORD_AND_SAVE；skip_first 与 wait 都映射到 NONE，一个 cycle 的最后一步是 RECORD_AND_SAVE，在它切走时触发 on_trace_ready。另有一个 skip_first_wait 参数：非零时让第一个 cycle 跳过 wait 阶段（想 profile 早点开始但后续 cycle 间隔大时有用）。这种"采样式"录制让 trace 体积可控。

21.6 TensorBoard 插件

PyTorch 提供 TensorBoard 集成（独立包 torch_tb_profiler）。先用 on_trace_ready=torch.profiler.tensorboard_trace_handler('/tmp/profiler_output') 把每个 cycle 的 trace 写成 <worker_name>.<时间戳>.pt.trace.json，再：

pip install torch_tb_profiler
tensorboard --logdir=/tmp/profiler_output

打开浏览器看 GPU Kernel View / Distributed View / Memory View / Trace View。比 chrome trace 多一层自动分析 ——"GPU idle 时间占比"、"top kernels"、"memory timeline"。不过这个插件在 PyPI 上的最后一版是 0.4.3（2023-10），之后没有再发布；新环境里装不上或显示异常时，直接用 Perfetto 打开 chrome trace 更稳妥。

21.7 Memory Snapshot：显存调试利器

第 4 章 §4.11 提过：

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

# 跑训练几个 step
for step, batch in enumerate(loader):
    train_step(batch)
    if step >= 5: break

torch.cuda.memory._dump_snapshot('mem.pickle')
torch.cuda.memory._record_memory_history(enabled=None)

把 mem.pickle 拖进 https://pytorch.org/memory_viz（PyTorch 官方 viewer），能看到：

  • 每次 alloc 的时间、大小、调用栈
  • 每次 free 的时间
  • 显存随时间的曲线
  • 哪些 alloc 长期占用、哪些短期
  • 碎片化情况

诊断 OOM / 显存膨胀的金钥匙。第 4 章 §4.13 有完整诊断流程。

21.8 record_function：用户级注解

profiler 默认只看到 ATen 算子。如果你想给"自定义代码段"加 marker：

from torch.profiler import record_function

with record_function("my_data_preprocess"):
    data = preprocess(raw)

with record_function("my_attention_block"):
    out = my_attention(data)

record_function 在 trace 里产生一个父事件，包住下面所有子事件。在 chrome trace / TB 里能看到自定义名字，让 trace 易读得多。对复杂 pipeline 强烈建议加 record_function 标注关键阶段。

NVTX range 是面向 NVIDIA 工具的对应物（在 Nsight Systems 里能看到）：

torch.cuda.nvtx.range_push("my_section")
do_work()
torch.cuda.nvtx.range_pop()

21.9 Distributed View：DDP/FSDP 的通信分析

profiler 在多 rank 训练时每个 rank 各自录一个 trace。把所有 rank 的 trace 文件夹喂给 TensorBoard，开 Distributed View 能看到：

  • 每个 rank 的 timeline 对比
  • collective 通信的时间分布
  • 各 rank 在 AllReduce 上等的时间（"等最慢的 rank"现象）
  • communication / computation overlap 比例

DDP / FSDP 调优靠这个 view 找瓶颈。"哪个 rank 慢"是分布式诊断第一问，distributed view 直接告诉你。

21.10 _ExperimentalConfig：高级 profiler 选项

torch/csrc/profiler/python/init.cpp:368 注册的 _ExperimentalConfig 暴露 12 个高级开关：

字段 用途
profiler_metrics CUPTI 硬件计数器列表（如 sm__cycles_active.avg.pct_of_peak_sustained_elapsed），让 profiler 进 CUPTI profiling mode 拿 GPU 占用率 / DRAM 带宽等
profiler_measure_per_kernel 硬件指标按 kernel 测还是整段测
verbose trace 文件里是否带 Call stack 字段
performance_events 进一步定制要测的事件
enable_cuda_sync_events 把 CUDA stream / event 同步操作也放进 trace（默认关）—— 调试隐式同步用
adjust_profiler_step 让 profiler.step 时长与父 Python 事件对齐
disable_external_correlation 关闭 CPU-GPU 事件关联
profile_all_threads 在所有线程上挂 callback（默认只挂在开启 profiler 的那个线程上，torch/csrc/profiler/orchestration/observer.cpp 的 pushGlobalCallbacks()）
capture_overload_names trace 里出现 aten::add.Tensor 而不只是 aten::add
record_python_gc_info 把 Python GC 事件加入 trace（诊断 GC 卡顿）
expose_kineto_event_metadata 暴露 Kineto 事件原始 metadata 到 Python 端
custom_profiler_config 自定义 backend（如某厂商自家硬件）的配置字符串 
from torch.profiler import _ExperimentalConfig, profile

config = _ExperimentalConfig(
    enable_cuda_sync_events=True,         # 看隐式同步
    capture_overload_names=True,           # 区分 aten::add.Tensor / .Scalar
    record_python_gc_info=True,             # 找 GC pause
)

with profile(activities=[...], experimental_config=config) as prof:
    train_step()

诊断高级问题（如训练偶尔卡顿）时这些开关价值很高 —— 能看到默认 trace 里被隐藏的 GC pause、stream 同步、overload 切换等"非典型"事件。

21.11 chrome trace 的 JSON schema

prof.export_chrome_trace('trace.json') 输出符合 Trace Event Format（Google Chrome / Perfetto 共用）的 JSON：

{
  "traceEvents": [
    {"name": "aten::add", "cat": "cpu_op", "ph": "X", "ts": 12345, "dur": 8,
     "tid": 0, "pid": 12345,
     "args": {"Input Dims": [[4, 4], [4, 4], []], "Input type": ["float", "float", "Scalar"]}},
    {"name": "ProfilerStep#1", "ph": "X", "ts": 100, "dur": 50000, "tid": 0},
    {"name": "Memcpy HtoD", "ph": "X", "ts": 200, "dur": 30,
     "tid": 99, "args": {"bytes": 4096}}
  ]
}

每条事件的核心字段：name / ph（X = duration、B/E = begin/end pair、i = instant）/ ts（开始时间，微秒）/ dur（持续时长）/ pid 与 tid（进程 / 线程，CUDA stream 也用 tid 表达）/ args（额外属性）；PyTorch 还会写 cat 区分事件类别（cpu_op、python_function、user_annotation 等）。record_shapes=True 时输入形状与 dtype 记在 args 的 Input Dims / Input type 里（本机 CPU 实跑所见）。

tid / pid 的具体取值由 Kineto 决定、随版本变化，别去背魔数。可靠的做法是看 trace 里的 metadata 事件：Chrome Trace Format 用 {"name": "thread_name", "ph": "M", ...} 与 "process_name" 这类 ph: "M"（metadata）事件给每条泳道命名，CPU 线程会被命名成 thread 12345 (python) 这样的可读名字（本机 CPU 实跑所见），GPU stream 泳道同样有 Kineto 写入的 thread_name。

# 先看有哪些泳道
python -c "
import json; d=json.load(open('trace.json'))
for e in d['traceEvents']:
    if e.get('ph')=='M' and e.get('name') in ('thread_name','process_name'):
        print(e['pid'], e.get('tid'), e['args'])
" | sort -u

拿到名字与 tid 的对应关系，再按 tid grep 具体事件 —— 比在 viewer 里滚动快得多。

21.12 record_function 的 C++ 实现

Python 端 with record_function("foo"): 进 C++ 是 at::RecordFunction 类（aten/src/ATen/record_function.h）。它的工作：

  1. 构造时调注册的 callbacks（profiler 会在这里挂钩，记录 begin 时间）
  2. 析构时调注册的"end" callbacks（profiler 记录 end 时间）
  3. profiler 的线程状态挂在 c10::ThreadLocalDebugInfo 的 PROFILER_STATE 槽上（torch/csrc/profiler/orchestration/observer.cpp），callback 据此找到当前线程该写入的事件缓冲

dispatcher 主路径上每个算子调用都可能触发一次 RecordFunction。实现在 aten/src/ATen/core/dispatch/Dispatcher.h：Dispatcher::call 先取 at::StepCallbacks，只有非空时才走 callWithDispatchKeySlowPath（:710），在那里构造 RecordFunction guard 并调 runRecordFunction（:741 / :753 / :756）。这就是 profiler 能"看到所有算子"的实现。

关键是这个快慢路径分离：没注册 callback 时，call 连 RecordFunction 对象都不构造，代价只有一次"callbacks 是否为空"的检查；开了 profiler 才落到慢路径。所以 profiler 的 overhead 与算子数量成正比 —— 小算子越密集，相对开销越大，这也解释了为什么"很多小 op 的模型"开 profiler 后掉速最明显。

runRecordFunction 有两个重载：一个带 inputs（callback 声明了 needsInputs(true) 时走它；typed 路径里由 :725 的 if (guard.needsInputs()) 选择，boxed 路径 callBoxed 里是 :876 起的三元表达式），一个不带。record_shapes=True 走的就是前者，所以它比默认更慢。

21.13 CUPTI 硬件计数器：从 kernel 时长走到 SM 利用率

profile() 默认只记 kernel 时长。要看更深的硬件指标（如 SM 占用率、DRAM 带宽），用 _ExperimentalConfig.profiler_metrics：

from torch.profiler import _ExperimentalConfig

config = _ExperimentalConfig(
    profiler_metrics=[
        "sm__cycles_active.avg.pct_of_peak_sustained_elapsed",       # SM 活跃率
        "smsp__sass_thread_inst_executed_op_fadd_pred_on.sum",        # FP add 总数
        "dram__bytes.sum",                                              # DRAM 读写字节
        "lts__t_sectors.sum.per_second",                                # L2 cache traffic
    ],
    profiler_measure_per_kernel=True,
)

CUPTI（CUDA Profiling Tools Interface）暴露了上千个 metric，分几大类：

类别 代表 metric 看什么
Compute sm__cycles_active, smsp__inst_executed SM 是否在算 vs idle
Memory dram__bytes, lts__t_sectors DRAM/L2 traffic
Latency gpu__time_active.avg, sm__warps_active warp 占用率（latency-hiding 程度）
Tensor Core sm__inst_executed_pipe_tensor TC 是否被用

实战分析路径：

flowchart TB
    Slow[kernel 执行慢]
    Slow --> Q1{SM 利用率 > 80%?}
    Q1 -->|否| MB[memory bound:<br/>DRAM 带宽是否打满?]
    Q1 -->|是| CB[compute bound:<br/>是否用 Tensor Core?]
    MB -->|是| OK1[已极限, 看能否减少 traffic<br/>fuse / 量化]
    MB -->|否| Latency[latency bound:<br/>warp occupancy 太低]
    Latency --> Block[调 block size / num_warps]
    CB -->|否| Force[force Tensor Core<br/>用对齐 shape / fp16 mm]
    CB -->|是| Try[已极限, 考虑算法改进]
    style MB fill:#fef3c7
    style CB fill:#dcfce7

一个典型的排查形态：某个 GEMM kernel 看着比同尺寸的别的 GEMM 慢很多，metric 显示 Tensor Core 相关计数很低 → 去看 input shape，发现某一维不是 Tensor Core 要求的对齐倍数（各代要求不同，fp16 常见是 8 或 16 的倍数），kernel 只好回落到非 TC 路径。padding 到对齐尺寸后 TC 计数上来、kernel 变快。

理解 CUPTI 让你从"猜"过渡到"测"：每个性能问题都有具体 metric 暴露根因。metric 名字与语义随 CUDA 版本变化，用 ncu --query-metrics 查当前环境支持哪些，别抄固定列表。Nsight Compute / nsys 是更专业的 viewer，PyTorch profiler 内置的够覆盖日常调优。

21.14 profile × torch.compile：看 fused kernel 的真面目

torch.compile 让多个 ATen 算子 fuse 成一个 Triton kernel。profile 时这种 kernel 名字看起来像 triton_poi_fused_add_mul_relu_0——通过名字就能看出 fused 了哪些 op（默认按原始 aten 算子名排序拼接，见 torch/_inductor/utils.py 的 get_fused_kernel_name）。

@torch.compile
def f(x, w):
    return torch.relu(x @ w + 1.0).mul(2.0)

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    f(input, weight)

print(prof.key_averages().table(sort_by='cuda_time_total'))

输出（精简）：

Name                                CUDA total
triton_poi_fused_add_mul_relu_0     0.234 ms      <- add / relu / mul 三个 op 融进 1 个 Triton kernel
aten::mm                            0.523 ms      <- matmul 仍交给 cuBLAS（extern kernel）

对比不开 compile：

Name              CUDA total
aten::mm          0.523 ms
aten::add         0.078 ms
aten::relu        0.156 ms
aten::mul         0.043 ms

一眼看出 fusion 有没有发生：原本 4 个独立 kernel（各自一次 HBM 读写）变成 2 个 —— 不开 max_autotune 时 Inductor 默认把 mm 交给 aten 的 extern kernel（torch/_inductor/kernel/mm.py 的 tuned_mm），三个 elementwise 融成 1 个，中间结果留在寄存器 / shared memory 里不落 HBM —— 收益大小取决于被省掉的那几趟显存往返占原时间多少，memory-bound 的 elementwise 链收益最大，GEMM 主导时收益有限。

如果某个预期会 fuse 的 op 没 fuse（trace 里仍是 aten::xxx 而非 triton_yyy），说明 fusion pattern miss → 看 output_code.py 找原因（§15.11.2）。

调优路径：

  1. profile 看每个 fused kernel 时长 + 内含 op 名
  2. 找耗时大头 kernel
  3. 看 Inductor output_code.py 对应 Triton 源码
  4. 决定是否能进一步优化（如改 dtype、加 mark_dynamic、调 max_autotune）

这套流程让 compile 调优有迹可循。没有 profile 时 compile 是黑盒，开 profile 后每个决策都有数据支撑。

21.15 配合 py-spy：Python 端火焰图

profile 看 PyTorch 内部，但Python 端的瓶颈（如 Python 循环慢、numpy 操作多）profile 看不到。互补工具是 py-spy：

# 在训练进程跑起来后，单独终端
pip install py-spy
py-spy record --pid <pid> -o flame.svg --duration 60

输出 SVG 火焰图：每个矩形宽度 = 该函数 CPU 占用时间。互补 PyTorch profiler 的视角：

工具 看什么 不看什么
torch.profiler ATen 算子 / CUDA kernel / Python 调用栈关联 纯 Python 循环细节
py-spy Python 函数级 CPU 时间分布 GPU kernel
nsys 硬件级 timeline 高层 ATen 语义

典型场景：GPU 利用率上不去，torch.profiler 看 GPU stream 上没有明显异常、但 kernel 之间大段空白 —— 说明时间花在 profiler 看不见的地方。py-spy 一挂，发现某个 Python 循环里在做大量 numpy 小操作。改成 batched 的 torch op，CPU 就不再是瓶颈。

py-spy 优势：

  • 不需要修改代码：直接 attach 到 running process
  • 采样式：按固定频率读取目标进程的栈，不插桩，对被测进程影响很小
  • 支持子进程：py-spy record --pid <pid> --subprocesses 自动跟 fork 出的 worker

生产建议把 py-spy 加进训练 toolkit，与 PyTorch profiler 配合用。一个看"PyTorch 内部"、一个看"Python 端"，共同覆盖完整调优空间。

21.16 真实案例：FSDP all_gather × forward overlap 调优

第 18 章 §18.4 讲了 FSDP 用 prefetch + 独立 stream 让 AllGather 与 forward 计算重叠。怎么验证 overlap 真的发生了？用 profiler。

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for _ in range(3):
        out = fsdp_model(input)
        loss = out.sum()
        loss.backward()
        prof.step()

prof.export_chrome_trace('fsdp_trace.json')

打开 chrome trace，找到几条关键 stream（下面括号里的泳道名只是示意，实际名字以 trace 的 thread_name 元数据为准，见 §21.11）：

  • default stream (CUDA:0:0)：跑 forward / backward 计算 kernel
  • comm stream (CUDA:0:1)：跑 NCCL all_gather / reduce_scatter
  • CPU thread：dispatcher 调用

理想情况：

default stream:   [layer0 fwd] [layer1 fwd] [layer2 fwd] ...
comm stream:      [ag layer1 ] [ag layer2 ] [ag layer3 ] ...
                       ↑ 与 layer0 同时跑   ↑ 与 layer1 同时跑

如果看到 [ag layerN] 与 [layer(N-1) fwd] 并行、相邻两层 forward 之间几乎没有空隙 → overlap 成功。

排查形态 1：每层 forward 前都有一小段空隙 → prefetch 没赶上。FSDP-1 侧检查 forward_prefetch（默认 False）与 backward_prefetch（默认 BackwardPrefetch.BACKWARD_PRE）；FSDP-2 侧则用 set_modules_to_forward_prefetch / set_modules_to_backward_prefetch 显式指定 prefetch 目标（第 18 章 §18.20）。注意没有 prefetch_policy / prefetch_size 这两个参数，网上流传的这类写法是错的。

排查形态 2：comm stream 上某次 all_gather 明显长于同伴 → 通常是 wrap 粒度不对，某个 unit 的参数量远大于其他，或者同一批参数被反复 unshard。回到第 18 章 §18.13 调 wrap 粒度。

这种"profile 看 stream 关系"是 FSDP 调优最直接的工具。没看 trace 之前的 FSDP 调优都是猜测。

21.17 真实案例：Transformer Attention 性能优化

最经典的 profile-driven 优化案例：transformer attention 从"naive 实现"到"FlashAttention"的演进。每一步都用 profile 数据驱动。

baseline (naive attention)：

def attention(q, k, v):
    scores = q @ k.transpose(-2, -1) / sqrt(d)    # [B, H, L, L]
    attn = softmax(scores, dim=-1)
    return attn @ v

profile 会给出这样的形状（B=4, H=32, L=2048, d=128）：

Name                    说明
aten::matmul            Q @ K^T          <- 大头
aten::div               /sqrt(d)         <- 纯 elementwise, 但要完整读写一遍 scores
aten::softmax           exp + reduction + div
aten::matmul            @ V              <- 另一个大头

真正致命的不是时间而是显存：scores 的形状是 [4, 32, 2048, 2048]，fp16 下就是 4×32×2048×2048×2 字节 = 1 GB，而且这个中间张量必须完整实例化（softmax 要沿最后一维归一化）。

优化 1：torch.compile 自动 fuse

Inductor 的 pattern matcher（torch/_inductor/fx_passes/fuse_attention.py，如 _sfdp_pattern_1）能认出 matmul → div → softmax → matmul 这类朴素写法，直接替换成 scaled_dot_product_attention（本机 CPU 实跑上面这段 attention，counters["inductor"]["fuse_attention"] 计数为 1）。写法不匹配这些模板时，compile 才只是把 div 与 softmax 融进一趟、省掉一次对 scores 的完整读写，1 GB 的 scores 张量仍然要实例化。

优化 2：FlashAttention（手写算子）

FlashAttention 把 attention 做成"分 block 的 streaming 计算"：

for block_K, block_V in chunks(K, V):
    scores_block = Q @ block_K.T
    softmax 一部分
    output += softmax_part @ block_V

中间 scores 永不完整实例化，显存占用从与 L² 成正比降到与 block 大小成正比。省下的不只是显存，还有那 1 GB 数据反复过 HBM 的时间 —— 这才是 FlashAttention 提速的主因。

PyTorch 用 F.scaled_dot_product_attention 自动选 backend：

out = F.scaled_dot_product_attention(q, k, v)
# 内部按输入 dtype / shape / mask 情况在几个 backend 之间选
# (flash attention / memory-efficient / cuDNN / math fallback 等)

顺带一提：scaled_dot_product_attention 本身在 autocast 的 AT_FORALL_LOWER_PRECISION_FP 白名单里（第 20 章 §20.6），amp 下会自动走低精度。

profile 一对比，三种实现的 kernel 数量、名字与显存占用差异一目了然。这种"测量-优化-再测量"的循环是性能工程的标准动作。至于具体快多少，取决于 B / H / L / d 与硬件，必须自己测。

理解这个案例让你看到 PyTorch 内置 SDPA 不是"魔法 API"——它是几年研究 + 工程优化的结果，profile 是验证每个优化真的有效的关键工具。

21.18 memory_viz 实战阅读

memory_viz 是 PyTorch 官方显存可视化（pytorch.org/memory_viz）。把 dump 出的 pickle 拖进去看到的核心视图：

Active Memory Timeline（横轴时间、纵轴活跃显存）：

显存
 ↑
80GB ─────────────────────────────────────
70GB        ╱╲          ╱╲         ╱╲
60GB       ╱  ╲        ╱  ╲       ╱  ╲
50GB ─────╱    ╲──────╱    ╲─────╱    ╲──── (训练 step 边界)
40GB ────fwd───bwd───fwd───bwd───fwd───bwd
30GB ────                                 (常驻 weight + optim state)
                                           → 时间

每个尖峰对应"forward 累积 activation"，回落对应"backward 释放 activation"。常驻部分是 weight + optim state。

实战诊断：

  • 尖峰持续上涨（每 step 比上一 step 高几 GB）：activation 没释放、内存泄漏。检查代码里有没有把 tensor 加到 list / 缓存里
  • 常驻部分在第 1 步之后抬高一截：多半是 optimizer state 在第一次 step() 时才分配，属正常；之后还在逐步上涨才是问题
  • 怀疑碎片化：这张图画的是活跃（allocated）显存，看不出碎片；要切到 Active Cached Segment Timeline 或下面的 Allocator State History，看 segment 总量与块之间的空洞
  • 常驻部分大：weight 太大 → 用 FSDP / 量化降低；optim state 大 → 用 8-bit Adam

时间线上每一块色带就是一次 alloc，鼠标悬停会显示它的大小和分配时的 Python / C++ 调用栈 —— 直接定位到哪个文件、哪一行。某块特别宽的色带若落在 attention 的 softmax 附近，多半就是 naive attention 的 scores，应换 FlashAttention。

Allocator State History：

按事件顺序回放 allocator 的每一次 segment 分配、alloc、free：左侧是事件列表，右侧画出选中时刻各 segment 里块的占用情况，下方显示该事件的调用栈。看碎片化、看哪些 segment 被长寿小块钉住，用的就是这个视图。

（v2.11 的 viewer 源码在 torch/utils/viz/MemoryViz.js，下拉框里共有 Active Memory Timeline / Allocator State History / Active Cached Segment Timeline / Allocator Settings 四项。）

理解 memory_viz 这几个视图让你在 OOM 时快速定位根因。第 4 章 §4.11 给了机制，本章给"怎么读"。两章配合让显存调试不再玄学。

21.19 .item() / .cpu() / .tolist()：隐式同步陷阱

trace 里看到 cudaStreamSynchronize 频繁出现 → 几乎肯定是隐式同步。常见触发点：

loss = compute_loss(out, target)

# 隐式同步！
print(loss)                    # str() 触发 .item()
loss_value = loss.item()       # 显式触发
losses.append(loss.item())     # 训练循环里累积 metric

# 也是隐式同步
if loss < 1e-4: break          # bool() 触发 .item()
preds = out.cpu().numpy()      # GPU → CPU + Python conversion

每次 .item() 让 CPU 等 GPU stream 完成 → 打破 async pipeline。在 hot loop 里反复调用，吞吐会明显下降。

修复模式：

# 错误：每 step 同步
losses = []
for batch in loader:
    loss = train_step(batch)
    losses.append(loss.item())   # ← 每 step 同步

# 正确：把 metric 累积在 GPU 上, epoch 末再 .item()
loss_sum_gpu = torch.zeros(1, device='cuda')
for batch in loader:
    loss = train_step(batch)
    loss_sum_gpu += loss.detach()
mean_loss = (loss_sum_gpu / len(loader)).item()   # 只同步 1 次

profile 验证：修复前 trace 里每 step 之间夹着 cudaStreamSynchronize、GPU stream 有空白；修复后这些事件消失、空白合拢。收益大小取决于原来每步同步几次、单次要等多久。

判断规则很简单：任何要把 GPU 上的"数值"变成 Python 值的操作都会同步；只读"元信息"的不会。

  • tensor.numel() / tensor.shape / tensor.dtype / tensor.device：不同步（这些存在 TensorImpl 里，不碰显存）
  • tensor.item() / float(tensor) / int(tensor)：同步
  • if tensor: / while tensor > 0:（隐式调 bool()）：同步
  • tensor.tolist() / tensor.numpy() / tensor.cpu()：同步
  • assert tensor.allclose(...)：allclose 返回 bool，同步
  • print(tensor)：要格式化数值，同步

"避免在 hot loop 里调 .item() / .cpu()" 是 PyTorch 训练性能 No.1 经验法则。理解 trace 上的 cudaStreamSynchronize 就是这套法则的可视化体现。

21.20 NVTX range + Nsight Systems 实战

torch.cuda.nvtx 比 record_function 更专业（在 Nsight Systems 看到）：

import torch.cuda.nvtx as nvtx

for step, batch in enumerate(loader):
    nvtx.range_push(f"step_{step}")          # 嵌套层 1
    nvtx.range_push("forward")                # 嵌套层 2
    out = model(batch)
    nvtx.range_pop()                           # forward end

    nvtx.range_push("backward")
    out.sum().backward()
    nvtx.range_pop()                           # backward end

    nvtx.range_push("optimizer")
    optimizer.step()
    nvtx.range_pop()                           # optimizer end

    nvtx.range_pop()                           # step end

用 nsys 录制：

nsys profile --trace=cuda,nvtx,osrt -o report python train.py
# 新版本产物是 report.nsys-rep, 老版本是 .qdrep; 用 nsys-ui 打开

Nsight Systems 比 chrome trace 强大得多：

  • 硬件级 timeline：含 GPU SM 占用、DRAM 带宽、PCIe traffic 实时曲线（GPU metrics 采样要在 nsys profile 时另开选项）
  • NVTX range 层级：你定义的 step / forward / backward 自动嵌套展开
  • stream 关系：跨 stream 依赖以箭头显示
  • 分析报告：自动生成"top kernel by occupancy / latency"等

生产调优 LLM 训练时 nsys 是终极工具。chrome trace 适合快速看一眼、nsys 适合深入分析。

实战工作流：

  1. 训练异常 → torch.profiler 录 trace 看大概问题
  2. 怀疑 GPU 端瓶颈 → nsys 录详细 trace 看 SM/DRAM 利用率
  3. 怀疑 Python 端 → py-spy 看火焰图

三个工具组合用，能覆盖大多数性能问题。每个都精通需要时间，但理解各自定位让你在不同场景选对工具。

21.21 profile 数据后处理：自定义聚合分析

prof.events() 返回所有事件列表，可以自己写分析逻辑：

from collections import defaultdict
from torch.autograd import DeviceType

events = prof.events()

# 按 kernel 名聚合 GPU 时间（device 侧事件的 name 是 kernel 名）
kernel_total = defaultdict(int)
for e in events:
    if e.device_type == DeviceType.CUDA:
        kernel_total[e.name] += e.device_time_total   # cuda_time_total 已 deprecated

# 按 (op, shape) 聚合：相同 op 不同 shape 各自时间（需 record_shapes=True）
# 注意 CPU op 的 device_time_total 含子 op 的 kernel，嵌套 op 会重复计入；
# 只看"自身"时间用 self_device_time_total
op_shape_total = defaultdict(int)
for e in events:
    if e.device_type == DeviceType.CPU:
        key = (e.name, str(e.input_shapes))
        op_shape_total[key] += e.self_device_time_total

按 shape 聚合其实也有现成的：prof.key_averages(group_by_input_shape=True) 会多出一列 Input Shapes（torch/profiler/profiler.py 的 key_averages，另有 group_by_stack_n / group_by_overload_name）。自己写循环的价值在于内置分组给不了的维度。

实战聚合：

  • Top kernel by total time：哪个 op 占总时间最多
  • Top kernel by call count：哪个 op 调用次数最多（可能小但累积大）
  • Per-shape analysis：发现"某个 shape 的 mm 特别慢"
  • Time per layer：用 record_function 标 layer 名后聚合

这种"自定义分析"在大规模训练里价值很高。预定义 view（chrome trace / TB 插件 / key_averages）能解决大部分问题，剩下的需要自己写脚本聚合。

例子：某个 GEMM 在总时间里占比很高，但它分散在几百次调用里，viewer 里一个个看没意义。按 (op name, input_shapes) 聚合之后，立刻能看出某一类 shape 明显偏慢 —— 通常是没对齐到 Tensor Core 要求的倍数（§21.13）。这种按维度聚合的分析是 profile 的高级用法。

21.22 H2D / D2H 拷贝优化案例

profile 看到 Memcpy HtoD 占了总时间的一大块 → 多半是数据传输方式有问题。诊断 + 修复流程（下面两段输出是示意，数字随硬件与 batch 而变）：

症状：

Name                         CUDA total
Memcpy HtoD (Pageable→Device) 250 ms     <- 占整 step 30%
aten::add                     50 ms
aten::mm                      300 ms

Pageable→Device 是不走 pinned memory 的慢路径：驱动要先把数据拷进内部的 pinned 暂存区再 DMA，这次拷贝对 host 是同步的。

诊断步骤：

  1. 检查 DataLoader 配置：pin_memory=True？没开就立刻开
  2. 检查代码：tensor.cuda() vs tensor.cuda(non_blocking=True)？后者必须配 pin_memory 才有效
  3. 检查 source dtype：CPU tensor dtype 与 GPU 一致吗？不一致触发 implicit cast → 在 CPU 端做 cast 后才能 DMA
  4. 检查 copy 频率：是否每 step 多次 .cuda() 同一份数据？应该只 cuda 一次然后复用

修复案例：

# 错误：每 step pageable copy
for batch in loader:               # batch 在 CPU pageable memory
    batch = batch.cuda()           # 触发 pageable → device, 慢路径
    out = model(batch)

# 修复 1: pin_memory + non_blocking
loader = DataLoader(dataset, ..., pin_memory=True)
for batch in loader:
    batch = batch.cuda(non_blocking=True)    # pinned → device, 异步
    out = model(batch)

# 修复 2: 预分配 GPU buffer, 每步只 copy_ 新数据 (batch_cpu 仍需是 pinned 才异步)
buffer_gpu = torch.empty(shape, device='cuda')
for batch_cpu in loader:
    buffer_gpu.copy_(batch_cpu, non_blocking=True)   # 只填新数据
    out = model(buffer_gpu)

profile 修复后：

Name                       CUDA total
Memcpy HtoD (Pinned→Device) 50 ms     <- 快路径, async, 与计算 overlap
aten::add                    50 ms
aten::mm                     300 ms

pinned 走的是可异步的快路径，H2D 能与计算 overlap，这段拷贝时间大部分被隐藏掉 —— 具体省多少要在自己的机器上量。这种"profile 看到问题 → 立刻修复 → profile 验证"是性能调优最爽的循环。

21.23 长跑训练的滚动 profile + 异常检测

生产 LLM 训练跑几天到几周，偶发慢 step 不易复现。解法：rolling profile + 异常检测。

import time
from collections import deque

step_times = []
profile_buffer = deque(maxlen=5)    # 只留最近几次 sample

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=100, warmup=2, active=3, repeat=10000),
    on_trace_ready=lambda p: profile_buffer.append((time.time(), p.events())),
) as prof:
    for step, batch in enumerate(loader):
        t0 = time.time()
        train_step(batch)
        torch.cuda.synchronize()     # CUDA 是异步的, 不同步就只量到了 launch 时间
        step_time = time.time() - t0
        step_times.append(step_time)
        prof.step()

        # 异常检测
        if step > 100:
            mean = sum(step_times[-100:]) / 100
            if step_time > mean * 2:
                # 这步比近 100 步均值慢 2x → 触发详细 dump
                dump_full_profile(prof, batch_id=step)
                send_alert(f"slow step at step {step}: {step_time:.2f}s vs mean {mean:.2f}s")

整套机制：

  • 大约每 100 步 sample 一次 profile（schedule 控制：wait 100 + warmup 2 + active 3，105 步一个周期）
  • 保留最近几次 sample 在内存（不上传）
  • 异常时立刻 dump + 上报
  • 正常时 sample 数据可丢弃 / 抽样上传

这类机制最常抓到的根因是 IO 抖动：网络文件系统偶发变慢，DataLoader 的 worker 卡在读数据上，trace 里表现为 GPU stream 一大段空白而 CPU 侧某个线程一直在等。没有这套采样 + 异常触发，这种"十小时一次"的问题基本无从复现。

理解这套 monitoring 模式让你能在生产主动发现问题，而非等到训练失败再回头查。是大模型 SRE 的标配能力。

21.24 Inductor output_code + profile 联合调试

profile 看到某 fused kernel 慢，怎么深入？联合 TORCH_COMPILE_DEBUG=1（§15.11.2）拿到 output_code，对照 profile 看每个 kernel 的实际行为。

工作流：

import os
os.environ['TORCH_COMPILE_DEBUG'] = '1'

@torch.compile
def f(x, w):
    return torch.nn.functional.gelu(x @ w)

with profile(activities=[ProfilerActivity.CUDA]) as prof:
    f(input, weight)

# 1. profile 里看到某个 triton_poi_fused_gelu_0 出乎意料地慢 (x @ w 本身默认走 cuBLAS)
# 2. 看 torch_compile_debug/.../ 下对应的 Triton 源码
# 3. 看 output_code.py 里 wrapper 是怎么调它的 (grid / 参数)

两类常见诊断结论：

并行度不足 —— output_code 里看到 grid 只有几十上百个 block，远小于 GPU 的 SM 数量，大部分 SM 空着。这时开 max_autotune 让 Inductor 试更小的 block size、把 grid 撑大。

dtype 没降下来 —— Triton 源码里全是 tl.float32，说明 autocast 没覆盖到这段（比如 forward 根本没包在 autocast 里，或者这个算子归在 fp32 policy，第 20 章 §20.6）。

理解这种"profile + output_code"联合调试让你能深入到 Inductor 内部优化。两层信息缺一不可：profile 给"哪个 kernel 慢"、output_code 给"为什么慢"。

21.25 step time 异常检测的"指纹"

每个训练任务都有特定 step time 指纹（mean / variance / outlier 模式）。学会识别"指纹异常"：

graph TB
    subgraph Healthy[正常]
        H1[每 step 5.0 ± 0.1 s<br/>偶尔 5.5 s 是正常]
    end
    subgraph Pattern1[周期性突变]
        P1["每 100 step 出 5.5s"]
        P1 --> R1[Eval / ckpt 周期]
    end
    subgraph Pattern2[逐渐变慢]
        P2[step time 从 5s 涨到 8s]
        P2 --> R2[memory leak / 显存碎片化]
    end
    subgraph Pattern3[随机 spike]
        P3[偶尔 30s spike]
        P3 --> R3[NCCL straggler / IO hang]
    end
    subgraph Pattern4[loss 爆炸 + 慢]
        P4[loss spike + step time 长]
        P4 --> R4[GradScaler 减半重试 / inf 出现]
    end
    style R1 fill:#dcfce7
    style R2 fill:#fef3c7
    style R3 fill:#fee2e2
    style R4 fill:#fecaca

每种 pattern 对应不同根因，profile 能验证：

  • Pattern1：trace 显示 100 step 一次有 ckpt write 阻塞（async ckpt 没开）
  • Pattern2：memory_viz 显示常驻显存涨 1GB/step → 找代码里的 leak
  • Pattern3：distributed view 显示某 rank trace 在 AllReduce 上等 25s
  • Pattern4：trace 含 GradScaler.skip 事件、inf check 触发

把这个 pattern 表存为 SOP，新人遇到性能问题对照查 → 节省时间。"性能问题不是无规律的，每种症状有对应根因"。

21.26 多 rank trace 时序对齐：distributed view 的内部机制

torch_tb_profiler 的 distributed view 让多 rank trace 可视化对齐。关键问题：不同 rank 的 CPU 时钟不一致（NTP 同步总有误差），怎么把它们的时间线对齐？

对齐的通用思路是"找共同事件当锚点"：多个 rank 上必然同时发生、且语义唯一的事件（典型就是同一次 collective 的两端），可以用来估算各 rank 时钟之间的偏移，再把所有时间戳平移到同一基准。具体用哪个锚点、误差多大，取决于 torch_tb_profiler 的版本实现 —— 它是独立于 pytorch/pytorch 的包（pip install torch_tb_profiler），细节要去它自己的仓库看，别想当然。

需要知道的实用结论是：跨 rank 的绝对时间不可尽信，跨 rank 的"谁先谁后"和"等了多久"才是可用信号。distributed view 用对齐后的 trace 做几件事：

  • 计算每个 collective 的 straggler：每 rank 进入 AllReduce 的时间差。差异大 → 有 rank 慢
  • overlap 比例：通信时间占总 step 时间多少
  • wait time on rank：每 rank 在 collective 上等了多久（暴露慢 rank）

典型排查形态：某个 rank 持续比其他晚进入 collective，其他 rank 都在 AllReduce 上干等它。这时问题不在通信而在那个 rank 自身 —— 可能是硬件降频 / ECC 纠错、可能是它分到的数据更重（变长序列没做 balanced batching）、也可能是它上面还跑着别的进程。collective 上的等待时间是症状，慢 rank 才是病因，别去调 NCCL 参数。

理解这点让你看到 distributed profile 的价值不在"绝对时间有多准"，而在于它把"谁在等谁"摆到了同一张图上。

21.27 自定义 profiler callback：Hook 进 RecordFunction

§21.12 提了 RecordFunction 是 profiler 实现机制。用户可以注册自家 callback 看每次 op 调用：

// C++ 端
#include <ATen/record_function.h>

at::addThreadLocalCallback(
    at::RecordFunctionCallback(
        [](const at::RecordFunction& fn) -> std::unique_ptr<at::ObserverContext> {
            // 算子开始时调 (StartCallback 只收一个参数, 返回 ObserverContext)
            std::cout << "op: " << fn.name() << " inputs: " << fn.inputs() << "\n";
            return nullptr;
        },
        [](const at::RecordFunction& fn, at::ObserverContext* ctx) {
            // 算子结束时调
            std::cout << "op done: " << fn.name() << "\n";
        }
    ).needsInputs(true)
);

aten/src/ATen/record_function.h:689 的 addThreadLocalCallback 与它旁边的 addGlobalCallback 是这套机制的公开入口；needsInputs(true) 会让 dispatcher 走带 inputs 的那个 runRecordFunction 重载（§21.12）。这是 C++ 侧的扩展点，Python 端没有对应的公开 API —— 想用得写 C++ 扩展（第 22 章）。

应用场景：

  • 统计每个 op 调用次数：自家训练 monitor，不依赖完整 profile
  • kernel-level 数据校验：每个 GEMM input 的 hash 上报，验证训练数据正确性
  • 自定义 trace 后端：把 trace 推到自家观测系统（如 Prometheus / 自家 timeline 平台）
  • op-level alert：某 op 输出 NaN 立刻报警（在 callback 里检测）

PyTorch 自己的 profiler 就是这套 callback 的一个使用者，它没有走任何私有通道。这意味着当内置 profiler 的输出形态不满足需求（比如要把 op 级事件实时推给自家观测系统）时，你可以在同一层挂自己的 callback，而不必去 fork profiler。

理解 RecordFunction callback 让你看到 PyTorch 的 profile 不是封闭系统，是开放的扩展点。

21.28 训练 vs 推理 profile 的不同关注点

profile 训练与推理用的指标 + 分析方法完全不同：

维度 训练 profile 推理 profile
关注 step time、GPU util、显存、通信 overlap latency、throughput、kv cache 利用率
主要瓶颈 通信 / DataLoader kernel launch / kv cache miss
优化目标 总训练时间最短 单 query 延迟 / 集群 QPS
关键 metric step/sec、samples/sec tokens/sec、TTFT、ITL
调优工具 torch.profiler + chrome trace + py-spy nsys + vLLM 内置 metric + Nsight Compute（ncu）

训练 profile 重点：

  • step time 分布（mean / p99 / outlier）
  • forward / backward / optim 各阶段占比
  • 通信 / 计算 overlap 比例
  • 显存峰值

推理 profile 重点：

  • TTFT (Time To First Token)：prefill 阶段时长
  • ITL (Inter-Token Latency)：decode 阶段每 token 时长
  • KV cache hit rate：prompt 复用情况
  • batched throughput：单 GPU tokens/sec
  • prefill / decode 调度均衡（continuous batching）

实战：调推理服务时，vLLM 暴露的服务级 metric 往往比单跑 torch.profiler 更直接 —— 它专门为推理服务设计、暴露 TTFT / ITL、调度队列状态等。

理解两条路径让你不会"用训练经验调推理"或反过来。两套 mindset、两套工具。

21.29 profiler 的分工格局

比起背版本号，更有用的是记住 v2.11 这一刻的分工格局：

组件 位置 归属
torch.profiler Python API torch/profiler/ PyTorch 主仓
RecordFunction（op 级采集） aten/src/ATen/record_function.h + Dispatcher.h PyTorch 主仓
Kineto（跨后端采集与 trace 格式） third_party/kineto pytorch/kineto，独立仓
CUPTI（GPU 硬件事件） CUDA Toolkit NVIDIA
TensorBoard 插件 torch_tb_profiler 独立 pip 包，不随 PyTorch 发布
Memory Snapshot viewer pytorch.org/memory_viz 网页端，独立

这张表解释了很多困惑：为什么 profiler 的 API 稳定但 TB 插件常年滞后（不同仓库、不同发布节奏）；为什么 _ExperimentalConfig 的文档里明写着不保证向后兼容（它直接透传给 Kineto）；为什么换个后端（ROCm / XPU）profiler 还能用（Kineto 这层做了抽象）。

老的 torch.autograd.profiler.profile 仍在 torch/autograd/profiler.py（key_averages 在 :544、export_chrome_trace 在 :515），torch.profiler 是包在它外面的新一层。新代码用 torch.profiler。

21.30 大 trace 文件的处理：filter + split

trace 的体积与事件数成正比，而事件数约等于"步数 × 每步算子数"。大模型训练随便录几步就能到几百 MB 甚至 GB 级，浏览器端的 viewer 加载会明显吃力。生产实战需要切割与过滤：

Filter 工具：

import json

def filter_events(input_path, output_path,
                  min_duration=10,        # 微秒，过滤短事件
                  exclude_names=None):    # 过滤特定 op
    exclude_names = set(exclude_names or [])
    with open(input_path) as f:
        data = json.load(f)

    filtered = []
    for e in data['traceEvents']:
        # 只过滤 ph == 'X' 的时长事件; metadata(M) / instant(i) / flow(s,f) 没有 dur, 要保留
        if e.get('ph') == 'X' and e.get('dur', 0) < min_duration:
            continue
        if e.get('name', '') in exclude_names:
            continue
        filtered.append(e)

    data['traceEvents'] = filtered
    with open(output_path, 'w') as f:
        json.dump(data, f)

filter_events('big_trace.json', 'small_trace.json',
              min_duration=100,
              exclude_names=['aten::contiguous', 'aten::view'])

实战常见 filter：

  • min_duration=100us：去掉 dispatcher 内部小事件
  • 过滤 aten::view 这类只改元数据的算子；aten::contiguous 只在输入已连续时不拷贝，输入不连续时会真的复制，过滤前先确认
  • 过滤 [Scheduler] 等运行时内部事件

Split 工具：

def split_by_step(input_path, output_dir):
    """按 ProfilerStep 切分 trace, 每步独立文件"""
    with open(input_path) as f:
        data = json.load(f)

    meta = [e for e in data['traceEvents'] if e.get('ph') == 'M']   # 泳道命名, 每个文件都带上
    steps = []
    current_step = []
    for e in sorted(data['traceEvents'], key=lambda x: x.get('ts', 0)):
        if e.get('ph') == 'M':
            continue
        if e.get('name', '').startswith('ProfilerStep#'):
            if current_step:
                steps.append(current_step)
            current_step = [e]
        else:
            current_step.append(e)
    if current_step:
        steps.append(current_step)

    for i, step_events in enumerate(steps):
        out = {'traceEvents': meta + step_events, 'displayTimeUnit': 'ns'}
        with open(f"{output_dir}/step_{i}.json", 'w') as f:
            json.dump(out, f)

按步切开后每个文件只剩原来的几分之一，viewer 打开就顺畅了。注意这里按事件开始时间切：某步 launch、下一步才执行完的 GPU kernel 会落进下一步的文件。这套 filter / split 脚本是生产 profile 的必备工具。

Perfetto 比 chrome tracing 处理大文件能力强得多（用 sqlite 后端），生产推荐用 Perfetto（ui.perfetto.dev）。

21.31 profiler 的内部架构图

把全章实现细节合起来看完整架构：

graph TB
    User["用户代码<br/>with profile():"]
    User --> PyAPI["Python: torch.profiler"]
    PyAPI --> CppPyBind["C++ pybind: profiler bindings"]
    CppPyBind --> RF[RecordFunction<br/>每 op 触发]
    CppPyBind --> Kineto[Kineto 库]
    RF --> CB[Profiler callback<br/>begin / end 时间]
    CB --> Buffer[Event Buffer<br/>thread-local]
    Kineto --> CUPTI[NVIDIA CUPTI]
    CUPTI --> Kernel[CUDA kernel events]
    CUPTI --> Mem[memcpy events]
    CUPTI --> HW[硬件 metric]
    Buffer --> Merge[合并 / 时间对齐]
    Kernel --> Merge
    Mem --> Merge
    HW --> Merge
    Merge --> Out{输出选择}
    Out --> Table[key_averages]
    Out --> Chrome[chrome trace JSON]
    Out --> TB[TensorBoard]
    Out --> Custom[用户自定义后处理]
    style RF fill:#fef3c7
    style Kineto fill:#dcfce7
    style Merge fill:#dbeafe

每层职责：

  • Python API：用户接口，控制 profile 生命周期
  • pybind：C++ 与 Python 桥接
  • RecordFunction：在 dispatcher 主路径捕获每个 op 调用
  • Kineto + CUPTI：拿 GPU 端原始数据
  • Merge 层：把 CPU events + GPU events 按时间对齐合并
  • Output：多种 view 共用同一份 merged data

理解架构让你看到 profile 不是单一组件，是 PyTorch 主仓 + Kineto 独立仓 + NVIDIA CUPTI + 用户代码四方协作的产物（分工见 §21.29）。每层独立可扩展，让 profiler 能持续演进而不破坏上层接口。

21.32 第三方工具综合对比

PyTorch profiler 不是唯一选择，根据场景选不同工具：

工具 强项 弱项 用法
torch.profiler 与 PyTorch 集成最好，看 ATen + GPU 大 trace 卡顿 训练调优 first choice
Nsight Systems (nsys) 硬件级 timeline，含 PCIe/DRAM 学习曲线陡 深入 GPU 性能
Nsight Compute (ncu) kernel 级硬件分析（occupancy / cache） 慢，每个 kernel 要重放多遍 微调单个 kernel
py-spy Python 火焰图，零代码改动 看不到 GPU Python 端瓶颈
scalene CPU + GPU + 内存综合分析 PyTorch 集成不完整 通用 Python profile
line_profiler 行级 CPU 时间 只 CPU 找慢 Python 函数
memory_profiler 行级内存 只 RAM 不看 GPU 找内存泄漏
viztracer 时间线 + 火焰图，开箱即用 不看 GPU Python 应用整体分析
vLLM profile 推理 LLM 专用，看 KV cache 仅 vLLM LLM 服务调优

实战工具组合：

  • 训练慢：torch.profiler → 看不出再用 nsys
  • 训练 OOM：memory_viz → 找不到再用 memory_profiler 看 CPU
  • 推理慢：vLLM 内置 metric → 看不出再用 nsys
  • kernel 慢：torch.profiler 找哪个 kernel → ncu 分析硬件 metric

这种"分层使用"是性能工程师的标配。每个工具学精都需要时间，但理解定位让你能为问题选对武器。

21.33 训练监控 dashboard 集成

profile 是"事件驱动"调优工具，dashboard 是"持续观测"工具。生产 LLM 训练通常两者结合：

graph LR
    Train[训练进程]
    Train --> Metrics[训练 metric<br/>loss, lr, step time, throughput]
    Train --> System[系统 metric<br/>GPU util, mem, network]
    Train --> Profile[采样 profile<br/>每 N step 一次]
    Metrics --> Wandb[Weights & Biases]
    Metrics --> TB[TensorBoard]
    System --> Prom[Prometheus + Grafana]
    Profile --> Storage[S3 / 公司存储]
    Wandb --> Alert[异常检测<br/>step time > 2x mean → alert]
    Prom --> Alert
    Storage --> OnDemand[需要时下载 trace 分析]
    style Alert fill:#fee2e2

实战集成：

import wandb
from torch.profiler import profile, schedule

wandb.init(project="llm-training")

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=500, warmup=2, active=3, repeat=1000),
    on_trace_ready=lambda p: upload_trace(p, wandb.run.id),
) as prof:
    for step, batch in enumerate(loader):
        t0 = time.time()
        loss, throughput = train_step(batch)
        step_time = time.time() - t0

        # 持续 metric → wandb
        wandb.log({
            'loss': loss,
            'lr': lr,
            'step_time': step_time,
            'throughput': throughput,
            'gpu_mem': torch.cuda.memory_allocated() / 1e9,
        }, step=step)

        # 偶发 profile → 上传 S3
        prof.step()

dashboard 上看：

  • 持续 metric：实时 loss 曲线、step time 曲线、显存曲线
  • 采样 trace：每 500 step 一次完整 profile，按需下载分析
  • 异常告警：step time / loss / 显存超阈值触发 PagerDuty / Slack

生产 LLM 训练 7×24 跑，把 profile 嵌进 monitoring 体系让团队能在故障扩大之前介入，而不是事后翻日志。

理解这套集成让你看到 profile 不只是"调试时用的工具"，是生产观测的核心数据源。

21.34 典型案例：定位训练偶发 hang

最难的性能问题是"训练偶尔卡死"——10 小时一次、每次卡一分钟左右然后自己恢复。临时开 profiler 抓不到、py-spy 也等不到现场。怎么办？（下面的数字都是示意。）

步骤 1：开 rolling profile + 异常检测

§21.23 的方案。每 1000 步采样、step time > 5x mean 触发 dump。

步骤 2：异常 trace 拿到后定位

正常 step: 5s
异常 step: 60s (12x)

打开异常 trace，对比正常 trace 看差异：

正常: forward 2s + backward 2.5s + optim 0.5s = 5s
异常: forward 2s + backward 2.5s + optim 55.5s   <- optim 暴涨!

optim 内部 profile 看：

正常 optim:  Adam.step (0.4s) + zero_grad (0.05s) = 0.45s
异常 optim:  cudaMemcpy DtoH (54.9s) + Adam.step (0.4s)

cudaMemcpy DtoH 卡 54.9 秒？ 这是 GPU → CPU 拷贝。检查代码：

# 偶尔触发的代码:
if step % 100 == 0 and loss.item() > 1.0:    # ← .item() 触发同步!
    log_extra_metrics(loss)

注意这里有个陷阱：loss.item() > 1.0 这个条件本身就已经同步了 —— 无论 loss 大小，每 100 步都会同步一次。真正偶发的是进入 if 之后的 log_extra_metrics，它内部又调了若干次 .item() / .cpu()，每次都是一轮 D2H + 等待。平时不进这个分支所以看不出来，偶尔进去就卡住。

（顺带：trace 里这类事件是 DtoH（device → host），不是 HtoD —— 把数值从 GPU 读回 CPU。看错方向会把排查引到 DataLoader 上去。）

步骤 3：修复

# 修复: 把 metric 累积改成 batched
log_buffer.append(loss.detach())
if step % 100 == 0:
    log_extra_metrics(torch.stack(log_buffer))    # 一次同步, batched
    log_buffer.clear()

修复后偶发 hang 消失。能这样定位，靠的是 profile 数据 + 异常检测让"偶发"变成"可观测"。

理解这种 case 让你看到 profile 不是"碰运气抓 bug"，是"长期监控 + 异常时拿到精确数据"的工程闭环。这是大模型训练 SRE 的核心方法论。

21.35 不同硬件后端的 profile 支持

PyTorch profiler 不只 CUDA，其他后端也支持但能力差异：

后端 profile 能力 限制
CUDA (NVIDIA GPU) 完整：CUPTI + Kineto + Memory Snapshot 最成熟
ROCm (AMD GPU) 较完整：rocprof + Kineto 部分 metric 缺失
XPU (Intel GPU) 基础：oneapi-pti + Kineto 还在迭代
MPS (Apple Silicon) 基础：torch.mps.profiler（OS Signpost，在 Xcode Instruments 里看） 不走 torch.profiler 的 activities（ProfilerActivity 里没有 MPS），缺硬件 metric
CPU 基础：函数级 timing 没有硬件 metric（用 perf 替代）
PrivateUse1 (国产芯片) 厂商自家 callback 各家实现不一

跨后端代码：

# 自适应不同后端
activities = [ProfilerActivity.CPU]
if torch.cuda.is_available():
    activities.append(ProfilerActivity.CUDA)
elif torch.xpu.is_available():
    activities.append(ProfilerActivity.XPU)

with profile(activities=activities) as prof:
    train_step()

实战：在 ROCm + AMD GPU 上 PyTorch profiler 能正常工作但 metric 比 CUDA 少；在 Apple M 系列上用 torch.mps.profiler 能在 Instruments 里看到 kernel 时长，但看不到 SM 占用率这类硬件指标。生产部署如果跨硬件，profile 流程要为每个后端验证。

理解后端差异让你看到 PyTorch 的 profile 不是"NVIDIA 专属"，是抽象层 + 各后端实现。但 NVIDIA 仍是事实标准（CUPTI 最成熟），其他后端在追赶。

21.36 一段总结：从 profile 到性能工程

把全章合起来看：profile 不是单一技能，是一套性能工程方法论：

graph TB
    Goal[确定优化目标<br/>throughput / latency / memory]
    Goal --> Baseline[Baseline 测量<br/>profile 当前性能]
    Baseline --> Hyp[假设瓶颈<br/>看 trace 形成 hypothesis]
    Hyp --> Fix[实施修复]
    Fix --> Verify[再次 profile<br/>验证 hypothesis]
    Verify --> Q{有改善?}
    Q -->|是| Next[找下一个瓶颈]
    Q -->|否| Hyp
    Next --> Goal2{达到目标?}
    Goal2 -->|是| Done[结束]
    Goal2 -->|否| Hyp
    style Baseline fill:#fef3c7
    style Verify fill:#dcfce7
    style Q fill:#dbeafe

每一步的关键技能：

  • 目标量化：不要"模糊地追求快"，定具体 metric（"step time 从 5s 降到 3s"）
  • Baseline 测量：修改前先 profile，知道改善多少
  • Hypothesis 驱动：看 trace 后形成"我认为是 X 慢"，不是凭直觉
  • 小步验证：每改一个修复就测，不要批量改后再 profile（不知道哪条改进有效）
  • 成本约束：性能优化是工程成本权衡，不是无限投入

这套方法论与软件工程的"测量驱动开发"一脉相承。理解 profile 工具是基础、形成工程闭环才是更高一层的能力。

实战：调好一个 LLM 训练任务的性能往往要持续迭代一段时间，每天 profile 几次、改几行、再 profile。最后总结成 SOP 给团队复用。这种"测量-优化-再测量"的工作模式是性能工程师的日常。

21.37 一个完整调优案例的复盘

最后给一个构造出来的示意案例（数字是为了讲清方法论编的，不是某次真实测量）：

目标：Llama-7B 训练 step time 从 1.2s 降到 0.8s（提升 33%）。

Day 1: Baseline + 第一轮分析

profile 录 5 步, step time 1.2s
breakdown:
  - DataLoader: 200ms (17%)
  - forward: 350ms (29%)
  - backward: 500ms (42%)
  - optim: 150ms (12%)

最大头是 backward。trace 显示 backward 含大量小 kernel、launch 密集 → 怀疑没用 compile。

修改：加 @torch.compile。

Day 2: 第二轮 profile

step time 0.95s (改善 21%)
  - DataLoader: 200ms (21%)
  - forward: 250ms (26%, fused)
  - backward: 350ms (37%, fused)
  - optim: 150ms (16%)

DataLoader 现在占比上升。trace 看 GPU 偶尔 idle 200ms → DataLoader 跟不上。

修改：num_workers=4 → 16、prefetch_factor=2 → 4、加 persistent_workers=True。

Day 3: 第三轮 profile

step time 0.85s (改善 11%)
  - DataLoader: 50ms (6%)
  - forward: 250ms
  - backward: 350ms
  - optim: 200ms (24%)  <- optim 占比上升

optim 不降反升（150ms → 200ms），占比涨到 24%。trace 看 optim 阶段是一长串读写 fp32 optimizer state（exp_avg / exp_avg_sq）的 memory-bound kernel。

修改：换 bnb.optim.AdamW8bit（8-bit Adam，optimizer state 以 8 bit 存储，读写字节数约为 fp32 的四分之一）。

Day 4: 第四轮 profile

step time 0.78s (达成目标!)
  - DataLoader: 50ms
  - forward: 250ms
  - backward: 350ms
  - optim: 130ms

总改善 35%（1.2s → 0.78s）。每天一次 profile-修改-验证循环，4 天达标。

复盘：

  • 3 个独立优化叠加：compile + DataLoader + 8-bit Adam
  • 每个优化都有 profile 数据驱动，不是凭直觉
  • 每次只改一处：知道哪条改善多少
  • 每次改动都可回滚：知道哪条有效、哪条没用

如果不 profile 直接猜着改，很容易：

  • 改了不该改的（增加复杂度）
  • 漏了关键优化
  • 没法量化每条的贡献
  • 总时间花更久

这是 profile 工具的真正价值 —— 让性能工程从猜测变成科学。这套 mindset + 工具应用比任何具体技术都更值得带走。

21.38 ProfilerStep 与 Schedule 状态机的源码细节

prof.step() 的内部 state machine 由 ProfilerAction 这个枚举驱动（torch/profiler/profiler.py:490）：

class ProfilerAction(Enum):
    NONE = 0
    WARMUP = 1
    RECORD = 2
    RECORD_AND_SAVE = 3

每次 prof.step()（torch/profiler/profiler.py:884）执行的逻辑，写成 if-else 大致是（真实实现在 :890-:894，用 (prev_action, current_action) 这个二元组去查 __init__ 里建好的动作表 action_map，由 _transit_action 执行，不是下面这串 if；下面只列常见转移）：

def step(self):
    prev = self.current_action
    self.step_num += 1
    self.current_action = self.schedule(self.step_num)
    cur = self.current_action

    if prev == ProfilerAction.RECORD_AND_SAVE:
        self.stop_trace()                # 上一个 cycle 录完了
        self.on_trace_ready(self)        # 在"切出" RECORD_AND_SAVE 时才触发用户回调
        if cur != ProfilerAction.NONE:
            self.prepare_trace()         # 紧接着下一个 cycle
    if prev == ProfilerAction.NONE and cur != ProfilerAction.NONE:
        self.prepare_trace()             # 进入 WARMUP: 开采集但数据不要
    if prev in (ProfilerAction.NONE, ProfilerAction.WARMUP, ProfilerAction.RECORD_AND_SAVE) \
            and cur in (ProfilerAction.RECORD, ProfilerAction.RECORD_AND_SAVE):
        self.start_trace()               # 真正开始录

schedule(step_num) 返回当前应该的 action。schedule(wait=1, warmup=2, active=3, repeat=2) 的状态序列：

step 0: NONE
step 1: WARMUP
step 2: WARMUP
step 3: RECORD
step 4: RECORD
step 5: RECORD_AND_SAVE     <- repeat 第 1 次结束
step 6: NONE                 <- 进入下一 cycle 的 wait
step 7: WARMUP
...

理解状态机让你能写自定义 schedule：

def my_schedule(step_num):
    # 1005 步一个周期: 1 步 warmup + 录 5 步, 其余 999 步不录
    cycle_pos = step_num % 1005
    if cycle_pos == 0:
        return ProfilerAction.WARMUP
    elif cycle_pos < 5:
        return ProfilerAction.RECORD
    elif cycle_pos == 5:
        return ProfilerAction.RECORD_AND_SAVE
    else:
        return ProfilerAction.NONE

with profile(activities=[...], schedule=my_schedule) as prof:
    for batch in loader:
        train_step(batch)
        prof.step()

这种自定义 schedule 让 long-running 训练能精确控制 profile 频率，不浪费空间也不漏掉关键步。生产 LLM 训练常用 1000+ 步一次的稀疏 schedule。

21.39 profile 与 CI/CD：性能回归检测

成熟工程团队把 profile 加进 CI 流水线 —— 每次 commit 自动跑 benchmark + profile，发现性能回归直接 block。

# CI workflow 片段
- name: Run perf benchmark
  run: |
    python benchmark.py --output results.json
    python compare_with_baseline.py results.json baseline.json

- name: Detect regression        # 伪代码: 实际判断写在 compare_with_baseline.py 里, 超阈值 exit 1
  run: |
    if perf_regression > 5%:
      exit 1

benchmark.py 内部用 profile：

import time
import torch
from torch.profiler import profile, ProfilerActivity

def benchmark():
    model = build_model()
    input = make_input()

    # warmup
    for _ in range(5):
        _ = model(input)
    torch.cuda.synchronize()

    # 测时间: 总时长用墙钟, profile 只用来拆解
    t0 = time.perf_counter()
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
        for _ in range(20):
            out = model(input)
        torch.cuda.synchronize()
    total_time = time.perf_counter() - t0

    # 提取关键 metric (key_averages() 本身不排序, 要自己按 device 时间排)
    top = max(prof.key_averages(), key=lambda e: e.self_device_time_total)
    return {
        'total_time': total_time,                       # 含 profiler 自身开销
        'top_op': top.key,
        'top_op_device_time_us': top.self_device_time_total,
        'memory_peak': torch.cuda.max_memory_allocated(),
    }

CI 上每次 commit 跑这个、与 baseline 对比、回归 > 5% 报警。

实战收益：

  • 早期发现：性能回归在 commit 当时就暴露，而不是几周后在生产上被发现 —— 那时候要从一堆改动里倒查根因
  • bisect 简单：知道哪个 commit 引起，git bisect 直接定位
  • 持续优化：team 看到 perf 回归压力大、自然写更优代码

PyTorch 自家也有这套（pytorch/benchmark repo，即 TorchBench），用一批真实模型持续跟踪性能回归。把 profile 工具化用到 CI 是工程成熟度的标志。

理解这种用法让你看到 profile 不只是"调试工具"，是质量保障基础设施的一部分。生产 ML team 必备能力。

21.40 几条实战调优经验

1. 训练慢先开 profiler：不要凭直觉猜瓶颈。3 行代码加 schedule 录 5 步、看 chrome trace 通常能立刻发现问题

2. GPU 利用率波动 (sawtooth) 多半是 DataLoader：第 11 章 §11.38 完整诊断流程

3. trace 里 cudaStreamSynchronize 多 = 隐式同步：找代码里的 .item() / .cpu() / print(tensor)，挪到训练 hot loop 外

4. autograd profiler vs profiler：老 torch.autograd.profiler.profile 仍在（torch/autograd/profiler.py），torch.profiler 是它外面的新一层，新代码用后者

5. profiler overhead：record_shapes=True 让 dispatcher 走带 inputs 的 runRecordFunction 重载、with_stack=True 要抓 Python 栈，两者都比默认更贵。只在真需要那份信息时开

6. flame graph for Python：补充 py-spy record -- python train.py 看 Python 端火焰图，与 profiler 互补

7. nsys（Nsight Systems）是更专业的工具：NVIDIA 官方的系统级 profiler，能把 CUDA API、kernel、NVTX、OS 线程放在一条时间线上，还能采 GPU 硬件指标。生产调优需要 nsys 时再上

8. 长跑训练 profiler 滚动落盘：每 1000 步触发一次 on_trace_ready 落盘 + 上传，监控异常步骤

9. trace 上传到对象存储 + URL 分享：team 协作 profile 数据。把 trace.json 上 S3 / GCS、配合 perfetto.dev 的 URL 加载，让团队成员零本地 setup 就能查看

10. profile 配合 unit test：核心算子写 assert kernel_time < threshold 单测，让性能回归在 PR review 阶段就暴露。注意阈值要留足余量，否则 CI 机器抖动会造成大量假阳性

11. profile 数据进 wandb / mlflow：与训练 metric 一起跟踪，可以在事后回看"哪个 commit 引起 step time 涨"

12. 多 GPU 调优先单卡再扩：单卡 profile 找 forward / backward 瓶颈，扩到多卡再看通信瓶颈。混在一起调很难

13. 训练 vs 推理：训练用 torch.profiler、推理用 vLLM / SGLang 内置 metric。两套 mindset 不同（§21.28）

21.41 完整调优流程

flowchart TB
    Start[训练慢 / OOM]
    Start --> P1[1 加 profiler 录 5 步]
    P1 --> Look{看 chrome trace 哪类问题}
    Look -->|GPU sawtooth idle| DL[DataLoader 优化<br/>第 11 章 §11.38]
    Look -->|kernel gap 多| Sync[找隐式同步<br/>.item / .cpu 挪走]
    Look -->|kernel 调用密集| Comp[torch.compile<br/>fuse 小算子]
    Look -->|GPU bound 但慢| Algo[换 dtype<br/>第 20 章 AMP]
    Look -->|OOM| Mem[Memory Snapshot<br/>第 4 章 §4.11 / §4.13]
    Mem --> Frag{碎片化?}
    Frag -->|是| EXP["PYTORCH_CUDA_ALLOC_CONF=<br/>expandable_segments:True"]
    Frag -->|否| Save[activation_checkpoint<br/>FSDP / cpu_offload]
    style P1 fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

这是一条常用的排查路径。每一步对应本专栏前面具体章节，profiler 是把它们串起来的诊断工具。OOM 分支这里只画到"是不是碎片化"这一层；按 cudaMalloc retries、reserved 与 allocated 之差、snapshot 里的长寿小块逐步排查的完整步骤，在第 4 章 §4.13。

21.42 相关章节

  • 第 4 章 §4.11 Memory Snapshot：本章的显存分析底层
  • 第 11 章 §11.38 DataLoader 瓶颈诊断：实战流程
  • 第 17 章 §17.28 TORCH_NCCL_BLOCKING_WAIT vs ASYNC_ERROR_HANDLING：分布式 profile 与 NCCL 调试开关配合
  • 第 14 章 Inductor：profile 看 fused kernel 与 output_code 联合调试
  • 第 15 章 §15.11.2 TORCH_COMPILE_DEBUG：profile 与 compile debug 双工具配合
  • 第 20 章 §20.18 量化精度调试：量化模型 profile 看 INT8 GEMM kernel 是否真的生效

21.43 设计启示

profiler 设计的核心思想：

第一：多源数据合并到统一 timeline：CPU 函数 + CUDA kernel + 显存事件 + 用户标注同步在一条时间轴上，让"为什么这里慢"能一眼看清

第二：Schedule API 让 profiler 可在生产用：录全程不可行，schedule 录采样让它适合长跑训练

第三：多层可视化：text 表（快速看 top）、chrome trace（看时间线）、TB 插件（自动分析）。每个层级解决不同诊断问题，组合起来覆盖完整诊断空间

第四：用户级标注 + 系统级采集合并：record_function 让 profiler 能体现用户的代码结构，与底层 ATen 算子事件融合

第五：profile 与编译栈共生：torch.compile 让 fused kernel 出现在 trace 里，profile 反过来验证 compile 收益。两套机制互相 feedback —— 改 compile 后 profile 看效果、profile 发现的瓶颈再回到 compile 调。这种循环是 v2.x 时代性能优化的核心 workflow

第六：测量是优化的前提：本章每个 case 都先 profile 测、再优化、再 profile 验证。这是与"凭直觉调"的根本差异。"没测量就不要优化"是性能工程的第一原则，profile 是这条原则的具体实现工具 —— 每次性能调优前先跑 5 步 profile、看 trace 形成 hypothesis，比直接改代码更快达到目标

下一章拆自定义算子 —— 当 PyTorch 内置算子不够用时，怎么写自家的 kernel 接进 dispatcher / autograd / torch.compile，把每一层都注册得完整、让 profile / autograd / compile 都能正确处理你的算子。