第45篇｜混合精度、显存分析与 profiler

本篇定位：专栏四「PyTorch 与 AI Infra 入门」的第九篇。承接第 44 篇可运行的 mini-GPT，把“能训练”推进到“能测量、能解释、能优化”：先理解 dtype 与混合精度，再建立显存账本，最后用 profiler 找到时间和内存真正花在哪里。

读完你将能够：区分 float32、float16 与 bfloat16 的取舍；用 autocast 和 GradScaler 写出更稳健的混合精度训练循环；读取 PyTorch 的显存计数器；用 torch.profiler 做一次带 warmup 的 CPU/CUDA trace；从训练和推理两种 AI Infra 场景解释常见 OOM。

预计用时：阅读 60～80 分钟，动手 60～90 分钟。

第 44 篇的 mini-GPT 已经能够完成前向、反向和生成，但当我们把 batch_size、上下文长度或层数调大时，问题会迅速从“代码对不对”变成“显存够不够、哪一步最慢、换半精度是否稳定”。只看总耗时，无法知道时间是在 Python、数据拷贝、矩阵乘法还是同步上；只看 nvidia-smi，也无法知道是哪一批激活把峰值推高。

混合精度和 profiler 不是两个互不相干的技巧。前者改变张量的存储与计算路径，后者帮助我们验证这种改变是否真的减少了显存、提升了吞吐，以及有没有把瓶颈转移到别处。本篇继续使用第 44 篇的 [B,T,D] mini-GPT 心智模型，但示例尽量拆成可以独立运行的小实验。最后我们会把训练的参数、梯度、优化器状态与推理的 KV Cache 放进同一份资源账本。

---

一、先建立 dtype 与显存的直觉

1.1 dtype 不只是“数字更小”

dtype（数据类型）同时决定一个元素占多少字节、能表示多大的数、有效精度是多少，以及 GPU 是否能调用更快的 Tensor Core 路径。常见浮点类型可以先记成下面这张表：

类型	每元素字节	特点	常见用途
float32	4	范围和精度都比较均衡	默认训练、基准与排错
float16	2	节省显存、吞吐高，但动态范围较小	CUDA 混合精度训练与推理
bfloat16	2	指数范围接近 float32，尾数精度较低	新一代 GPU 训练、推理
float64	8	精度高但慢且占用大	科学计算，不是深度学习默认

例如，形状为 [1024, 4096] 的激活有 4,194,304 个元素。用 float32 仅计算这一块就约 16 MiB，换成 float16 或 bfloat16 约 8 MiB。这里的“约”是指张量本身，不包括 allocator 的缓存、对齐、梯度和其他临时张量。

import torch

shape = (1024, 4096)
for dtype in (torch.float32, torch.float16, torch.bfloat16):
    x = torch.empty(shape, dtype=dtype)
    mib = x.numel() * x.element_size() / 1024**2
    print(dtype, x.element_size(), f"{mib:.1f} MiB")

torch.float32 4 16.0 MiB
torch.float16 2 8.0 MiB
torch.bfloat16 2 8.0 MiB

字节数变少，不代表所有算子都应该永久改成半精度。归一化、损失计算、指数运算和某些归约对数值范围更敏感；权重更新也需要足够精细的梯度。工程上的关键不是“全模型强制 .half()”，而是让安全的算子使用低精度，让敏感的算子保留 float32，并用实验确认 loss 与指标没有异常。

1.2 float16 与 bfloat16 的直观区别

浮点数可以粗略看作符号位、指数位和尾数位的组合。float16 的指数位少，能表示的最大值和最小正规数范围较窄；梯度很小时容易下溢成 0，数值很大时容易变成 inf。bfloat16 与 float32 使用相近的指数位宽，因此动态范围大得多，但尾数位少，表示相邻数的精细程度不如 float16。

这解释了一个常见经验：在支持良好的训练硬件上，bfloat16 往往更容易直接训练；float16 通常有更高的兼容性和成熟的 Tensor Core 路径，但更需要梯度缩放。两者都不是“精度更高”的替代品。模型质量还取决于初始化、学习率、归一化、损失实现和数据规模。

import torch

for dtype in (torch.float16, torch.bfloat16, torch.float32):
    info = torch.finfo(dtype)
    print(dtype, "tiny=", info.tiny, "max=", info.max)

输出中的具体科学计数法会因版本格式略有差异，但可以观察到 bfloat16 的 max 更接近 float32 的数量级。选择 dtype 前还要检查设备：并非所有 GPU 都对两种半精度同样快，CPU 的 bfloat16 支持也依赖 PyTorch、指令集和具体算子。

1.3 参数、梯度、状态和激活是一笔账

一个训练任务的峰值显存不能只用“参数量乘 2”估计。以 AdamW 为例，至少要考虑：模型参数、参数梯度、优化器的一阶动量、二阶动量、前向保存的激活、临时 workspace，以及 CUDA allocator 保留但暂时未使用的缓存。

若参数个数为 P，每个元素占 s 字节，那么参数本身约为 P×s。混合精度训练中，常见实现仍会保留 float32 的优化器状态；某些方案还保留 float32 master weights。因此低精度主要立刻影响参数和激活，不一定让优化器账本按同样比例缩小。反向图保存的激活通常与 B×T×D×layers 同量级；显式 Attention 的权重还可能出现 B×heads×T×T，上下文长度翻倍时这一项会变成四倍。

AI Infra 视角：训练集群上的“显存利用率”不是单一指标。参数、梯度、optimizer state、激活、通信 buffer 和 allocator cache 由不同生命周期管理。排查 OOM 时要问“哪个生命周期的哪一块增长”，而不是只问“模型有多少参数”。

图片展示了训练显存账本，说明一次step时模型参数、优化器状态、激活与临时量谁在占空间。模型参数包括权重/embedding（FP32、FP16、BF16）和梯度（反向传播后生成，生命周期短于状态）；优化器状态有AdamW（一阶+二阶）、momentum/variance（常保留FP32）和master weights（取决于训练方案）；激活与临时量为B×T×D×layers（前向保存激活）和Attention：T²。该图与上下文紧密相关，直观呈现了训练显存的占用情况。

图1 训练显存账本：参数、梯度、状态与激活

---

二、AMP：让安全的算子自动选精度

2.1 autocast 做了什么

AMP（Automatic Mixed Precision，自动混合精度）通常由两部分组成：autocast 负责在一个代码区域内为不同算子选择合适 dtype，GradScaler 负责在需要时放大 loss，减少 float16 梯度下溢。autocast 不等于把整个模型的参数永久改成半精度，它更像一个按算子分类的调度器。

在 CUDA 上，矩阵乘法和卷积通常适合低精度；而某些归约、归一化和损失会留在 float32。不同 PyTorch 版本、设备和算子实现的白名单会变化，所以不要靠猜测某一层的 dtype；必要时在关键位置打印 x.dtype，或用 profiler 看实际 kernel。

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.bfloat16
x = torch.randn(4, 8, device=device)
weight = torch.randn(8, 8, device=device)

with torch.autocast(device_type=device, dtype=dtype, enabled=True):
    y = x @ weight
    z = torch.softmax(y, dim=-1)

print("device=", device, "matmul=", y.dtype, "softmax=", z.dtype)

CPU 上不能简单照搬 CUDA 的 float16 配置；示例选择 bfloat16 只是为了展示接口，某些老 CPU 或算子可能不支持。训练代码最好把 device_type、dtype 和 enabled 做成配置，并在启动时打印，而不是把 torch.cuda.amp 写死在各个函数里。

2.2 一个最小的 AMP 训练循环

从第 44 篇的训练循环出发，混合精度版本最重要的顺序是：进入 autocast，计算 loss；如果使用 scaler，调用 scaler.scale(loss).backward()；更新前用 unscale_ 让梯度回到可检查的尺度；可选地裁剪梯度；最后 scaler.step 和 scaler.update。

import torch
from torch import nn

use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
use_amp = use_cuda
amp_dtype = torch.float16

model = nn.Sequential(nn.Linear(128, 256), nn.GELU(), nn.Linear(256, 10)).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scaler = torch.amp.GradScaler("cuda", enabled=use_amp)

for step in range(3):
    x = torch.randn(32, 128, device=device)
    target = torch.randint(0, 10, (32,), device=device)
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast("cuda", dtype=amp_dtype, enabled=use_amp):
        logits = model(x)
        loss = nn.functional.cross_entropy(logits, target)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()
    print(step, round(loss.item(), 4), scaler.get_scale())

在 CPU 上运行时，torch.autocast("cuda", ...) 这一行会因为 enabled=False 而不执行 CUDA autocast，但为了让代码更清楚，生产代码应把 device_type=device.type 传入，而不是依赖短路行为：

device_type = device.type
amp_dtype = torch.float16 if device_type == "cuda" else torch.bfloat16
with torch.autocast(device_type=device_type, dtype=amp_dtype, enabled=use_amp):
    logits = model(x)

不同 PyTorch 版本的 scaler 写法可能是 torch.amp.GradScaler("cuda") 或兼容旧代码的 torch.cuda.amp.GradScaler()。新代码优先使用 torch.amp 命名空间，并在项目的依赖版本中固定行为。enabled=False 时 scaler 会退化成不缩放的路径，便于同一份脚本做 CPU smoke test。

2.3 为什么需要 GradScaler

假设半精度能表示的最小非零梯度有限，反向传播中的某些梯度可能先下溢为 0。GradScaler 把 loss 乘以一个 scale，让反向得到的梯度暂时变大；在真正更新前再除回这个 scale。如果发现梯度包含 inf 或 nan，它会跳过本次 optimizer step，并降低 scale；如果连续稳定，还会逐步增大 scale。

因此 scaler 不是“让模型更准确”的魔法，而是给窄动态范围 dtype 留出安全余量。它的自动调整也意味着：训练日志中最好记录 scale；如果 scale 长期快速下降，通常要检查学习率、数据、loss 实现和半精度算子，而不是只把 scale 手动写成一个巨大常数。

梯度裁剪必须放在 scaler.unscale_(optimizer) 之后，否则裁剪面对的是放大过的梯度，阈值没有原本的含义。梯度累积时，每个 micro-batch 的 loss 是否除以累积步数、何时 unscale 和 step，也必须固定规则，不能一半按放大值、一半按原始值操作。

2.4 训练、验证和推理的三个开关

model.train()/model.eval() 控制 Dropout、BatchNorm 等模块行为；torch.no_grad() 或 torch.inference_mode() 控制是否建立 autograd 图；autocast 控制区域内的算子精度。它们是三组不同开关：eval() 不会自动关闭梯度，no_grad() 也不会把模型切到评估模式。

model.eval()
with torch.inference_mode(), torch.autocast(
    device_type=device.type, dtype=amp_dtype, enabled=use_amp
):
    prediction = model(x)

验证时忘记 inference_mode()，可能因为保留图而浪费显存；生成时忘记 eval()，Dropout 会让同一 prompt 的输出变化。反过来，在训练循环外包一个永久的 inference_mode()，则会让反向传播失败。把这三个开关写在各自函数的入口，边界会更清晰。

图片展示了AMP训练循环流程，分为输入batch、autocast、loss、GradScaler、unscale_、clip、step + update等步骤。autocast将matmul等敏感算子从FP16/BF16转为FP32，loss通常保留稳定计算，GradScaler避免梯度下溢，unscale_恢复原始梯度尺度，clip可选裁剪梯度，step + update更新参数与scale。顺序要点是backward后先unscale，再clip，最后scaler.step；发现inf/nan时跳过更新并调整scale。该图与上下文介绍的AMP训练、验证和推理的三个开关及AMP训练循环内容相关。

图2 autocast 与 GradScaler 的训练顺序

AI Infra 视角：AMP 的收益最终要用有效吞吐衡量，例如每秒训练 token、step time 和达到相同验证 loss 所需的总时间。若低精度让 kernel 变快，却因为数据搬运、同步或频繁转换而没有提升端到端吞吐，优化就没有完成闭环。

---

三、显存分析：从计数器到峰值

3.1 allocated、reserved 和峰值有什么不同

CUDA caching allocator 为了减少频繁向驱动申请显存，会保留已经申请过的 block。于是 PyTorch 常见的两个数含义不同：memory_allocated() 是当前活跃 Tensor 使用的字节数；memory_reserved() 是 allocator 从 CUDA 驱动保留的总字节数，通常大于等于 allocated。nvidia-smi 更接近进程向驱动持有的总量，不能直接等同于某一时刻的活跃 Tensor。

对应地，max_memory_allocated() 和 max_memory_reserved() 是从统计重置以来的峰值。一次实验应在 warmup 后、正式测量前调用 torch.cuda.reset_peak_memory_stats()，否则前面加载模型或旧 batch 的峰值会污染结果。

import torch

def report_memory(tag: str):
    if not torch.cuda.is_available():
        print(tag, "CUDA unavailable")
        return
    mib = 1024 ** 2
    print(
        tag,
        f"allocated={torch.cuda.memory_allocated()/mib:.1f} MiB",
        f"reserved={torch.cuda.memory_reserved()/mib:.1f} MiB",
        f"peak={torch.cuda.max_memory_allocated()/mib:.1f} MiB",
    )

if torch.cuda.is_available():
    torch.cuda.reset_peak_memory_stats()
    report_memory("after reset")
    x = torch.randn(2048, 2048, device="cuda")
    report_memory("after tensor")
    del x
    report_memory("after delete")
    torch.cuda.empty_cache()
    report_memory("after empty_cache")

del x 只删除 Python 引用；当没有其他引用时，Tensor 的 block 可以回到 allocator，但未必立刻归还给驱动。empty_cache() 释放未使用的缓存给驱动，可能帮助别的进程看到空间，却不会释放仍被活跃 Tensor 使用的显存，也不是修复 OOM 的常规按钮。频繁调用它还可能损害性能。

3.2 一次可复现的前向峰值实验

测量必须固定输入、dtype、batch、序列长度和是否带梯度。下面的骨架使用第 44 篇的模型接口；如果当前文件中没有 MiniGPT，可以换成任意 nn.Module，重点是测量边界。

import torch

@torch.inference_mode()
def measure_inference(model, batch, tokens, vocab, device):
    ids = torch.randint(0, vocab, (batch, tokens), device=device)
    model.eval()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats(device)
    _ = model(ids)                    # warmup / allocation
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats(device)
    _ = model(ids)
    torch.cuda.synchronize()
    return torch.cuda.max_memory_allocated(device) / 1024**2

真正做比较时，应该分别为 float32、AMP 和不同的 B/T 创建干净实验，避免把上一次运行留下的模型引用带进来。GPU 是异步设备：如果在 kernel 完成前立刻读取时间或峰值，结果可能不完整；torch.cuda.synchronize() 是测量边界的一部分，而不是平时每一步都要插入的“优化”。

可以把实验结果写成表格：

模式	batch	tokens	dtype	峰值 allocated	step time	tokens/s
FP32 train	8	256	float32	记录	记录	记录
AMP train	8	256	float16	记录	记录	记录
AMP train	8	512	float16	记录	记录	记录
inference	1	512	bfloat16	记录	记录	记录

这比一句“半精度省了一半显存”更可靠，因为模型加载、激活保存和 allocator 碎片会改变实际比例。

3.3 显存碎片与生命周期

OOM 不一定表示总活跃 Tensor 已经超过卡的容量。动态 shape、不同大小的临时 Tensor、长短请求交错，会让 allocator 出现无法满足某个大 block 的空洞。日志里的 “allocated” 可能不算特别高，但 reserved 很高；此时要用 torch.cuda.memory_summary() 观察 allocator 状态，并优先减少 shape 抖动、统一 batch 或缩短上下文。

更重要的是追踪引用生命周期。把每一步的输出 append 到 Python list，会让 autograd 图和相关激活一直活着；记录 loss 应使用 loss.item()，保存 Tensor 时通常需要 tensor.detach().cpu()。在循环里重复写 total_loss += loss，也可能把每一轮的计算图串起来；更安全的是累加 Python 浮点数或 loss.detach()。

losses = []
for x, y in loader:
    logits = model(x)
    loss = criterion(logits, y)
    loss.backward()
    optimizer.step()
    losses.append(float(loss.detach()))  # 不保存计算图

这类问题在小模型上不明显，到了长上下文语言模型就会表现为显存逐步增长。观察“每个 step 结束后的 allocated 是否回到相近水平”，往往比只看最终 OOM 更快找到泄漏。

图片展示了显存指标中的活跃Tensor、缓存与峰值情况。横轴为时间，纵轴为容量。图中有两条曲线，蓝色代表活跃Tensor，橙色代表缓存。曲线在不同时间点有波动，且在某些点有明显上升。图中还标注了“max_reserved / 峰值边界”和“reserved: allocator保留”“allocated: 活跃Tensor”等信息。右下角说明删除Tensor后allocated可下降，reserved可能继续存在，nvidia-smi不等于活跃Tensor。该图与上下文关于显存分析的内容相关，用于直观呈现显存指标变化。

图3 显存指标与实验测量边界

---

四、torch.profiler：把“慢”拆成事件

4.1 profiler 能回答什么

torch.profiler 可以记录 CPU 活动、CUDA kernel、算子形状、内存分配和调用栈，并把结果导出给 TensorBoard。它适合回答“哪几个算子占总时间最多”“CPU 是否在等 GPU”“数据拷贝是否成为瓶颈”“某个算子是否产生了异常大的 allocation”等问题。

它不应被当成免费日志。收集 trace 会增加运行时间和内存，尤其是记录 shape、stack 或 profile_memory 时。正确做法是先 warmup，让 CUDA kernel、缓存和编译路径稳定，再只采集少量 active steps；不要把整个长训练都开启 profiler。

4.2 一个最小的 profiling sketch

import torch
from torch import nn
from torch.profiler import profile, ProfilerActivity, schedule, tensorboard_trace_handler

use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
model = nn.Sequential(nn.Linear(1024, 2048), nn.GELU(), nn.Linear(2048, 10)).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
activities = [ProfilerActivity.CPU]
if use_cuda:
    activities.append(ProfilerActivity.CUDA)

with profile(
    activities=activities,
    schedule=schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=tensorboard_trace_handler("./profiler_logs"),
    record_shapes=True,
    profile_memory=True,
) as prof:
    for step in range(5):
        x = torch.randn(64, 1024, device=device)
        y = torch.randint(0, 10, (64,), device=device)
        optimizer.zero_grad(set_to_none=True)
        loss = nn.functional.cross_entropy(model(x), y)
        loss.backward()
        optimizer.step()
        prof.step()

print("trace written to profiler_logs")

schedule 中 wait 不记录，warmup 让 profiler 准备，active 才写有效 trace。每个训练 step 结束必须调用 prof.step()，否则 schedule 不会前进。若没有 GPU，代码只记录 CPU；不要为了生成 trace 把 CPU 机器强行配置成 CUDA。

终端汇总可以这样打印：

# profile 上下文退出后
print(prof.key_averages().table(
    sort_by="cuda_time_total" if use_cuda else "cpu_time_total",
    row_limit=12,
))

CUDA 版本中，排序字段常用 cuda_time_total；CPU-only 时使用 cpu_time_total。不同 PyTorch 版本的表格列略有差异，重点观察 self time、总 time、调用次数和输入 shape。若某个 kernel 看起来很快但调用几十万次，Python 循环和 launch overhead 仍可能是整体瓶颈。

4.3 从 trace 到优化假设

看到 aten::mm 占比高，不要马上认为它“有问题”：训练 Transformer 时矩阵乘法本来就应占较大比例。更有价值的问题是矩阵 shape 是否足够大、dtype 是否走到预期 kernel、CPU 是否在每个小算子间隙等待。如果 Memcpy HtoD 很突出，检查 DataLoader、pin_memory 和 .to(device, non_blocking=True) 的边界；如果 aten::copy_ 很多，检查是否在循环中频繁 .float()、.cpu() 或隐式转换。

record_shapes=True 能显示 [B,T,D] 等 shape，但会增加开销；with_stack=True 能定位 Python 调用点，代价更高。内存分析则关注 allocation 的大小、调用次数和生命周期。一次 profile 只能说明一组配置：batch、序列长度、warmup、硬件、驱动和是否打开 AMP 都必须写在报告里。

AI Infra 视角：profile 的价值是把优化从“感觉”变成可证伪的假设。比如“AMP 会更快”应拆成：矩阵乘法 kernel 使用了 Tensor Core、GPU busy time 下降、tokens/s 上升、验证 loss 没恶化、峰值 allocated 下降。任何一项不成立，都值得继续定位。

图片展示了torch.profiler从step到事件的时间线。上方为schedule，包含wait 1、warmup 1、active 3 → trace。下方从上至下依次为CPU、CUDA、memory，分别标注了dataloader、Python / aten dispatch、copy、next step；matmul / Tensor Core、GELU、backward kernels；allocation、activation / workspace、free。底部说明先用trace找到热点，再用普通benchmark验证优化后的端到端收益，每个step末尾调用prof.step()；profile不是免费的benchmark。

图4 profiler 从 step 到 CPU、CUDA 与 memory 事件

4.4 profile 与 benchmark 的区别

Profiler 是诊断工具，benchmark 是比较工具。开启 profiler 后的 step time 不能直接与普通运行的 step time 对比；trace、同步和 shape 记录会改变行为。基准测试应先做若干 warmup，再用固定次数测量，并只在需要时同步：

import time

for _ in range(10):
    _ = model(x)
if use_cuda:
    torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(50):
    _ = model(x)
if use_cuda:
    torch.cuda.synchronize()
elapsed = time.perf_counter() - start
print(f"step_time={elapsed / 50 * 1000:.2f} ms")

比较 AMP 时要确保输入、模型副本、warmup 次数和随机数据形状一致。只跑一次往往会被首次 kernel 加载、频率变化或系统后台任务干扰；至少报告平均值，并在服务场景补充 p50、p95 或 p99。

---

五、训练和推理的显存账本不同

5.1 训练：激活常常是峰值来源

训练一次 mini-GPT，前向需要参数和输入；为了反向，autograd 会保存若干中间激活；反向还会产生梯度，optimizer step 读取并更新状态。即使参数本身能放进 GPU，也可能因为 B、T 或层数太大而在反向阶段 OOM。

降低训练峰值的顺序通常是：减小 micro-batch；用梯度累积保持有效 batch；缩短序列或使用更省内存的 attention；启用 AMP；使用 activation checkpoint 以计算换显存；最后再考虑参数、优化器或模型并行。梯度累积不是“零成本”：它减少单次激活峰值，却增加 step 次数和同步设计的复杂度。分布式训练还要把通信 buffer 和每张卡的 shard 纳入账本。

5.2 推理：参数之外还有 KV Cache

推理时使用 inference_mode() 不保存反向图，通常比训练省很多显存；但 decoder-only 模型在生成时会保存每层的 K/V。粗略估计，KV Cache 元素量与 batch × layers × tokens × 2 × kv_heads × head_dim 成正比，具体还取决于 GQA/MQA 和 cache dtype。并发数或上下文长度增加时，KV Cache 会持续增长。

因此推理服务常见的 OOM 并不来自一次模型 forward，而是来自多个请求的 cache 叠加。continuous batching、分页 KV Cache、最大上下文限制和请求排队，都是显存管理的一部分。把参数从 float32 换成 bfloat16 可能节省一半参数存储，却不能解决无限增长的 cache；服务端仍需要 admission control。

5.3 训练与推理的对照

维度	训练	推理
主要目标	吞吐、有效 token/s	延迟、并发吞吐、尾延迟
反向图	需要保存或重算	通常关闭
额外显存	梯度、optimizer state、激活	KV Cache、batch workspace
典型峰值	backward 前后	长上下文高并发 decode
常见手段	AMP、checkpoint、累积、并行	cache 管理、batching、量化、限流

第 44 篇提到的 prefill 和 decode 在这里有不同资源画像：prefill 一次处理较长 prompt，矩阵并行度高，激活和临时 attention 更明显；decode 每次只处理新 token，但反复读取越来越大的 KV Cache，更受显存带宽和调度影响。profiling 这两个阶段时应分开打标签，不要只给一条混合平均线。

图片展示了训练与推理两种不同显存增长路径。训练部分，参数和梯度经optimizer state处理，激活量为B×T×D×layers，峰值常在backward附近；推理部分，参数与KV Cache相关，并发和上下文让cache增长，高并发decode可能OOM。两者都需明确配额与测量，图片与上下文紧密相关，直观呈现了训练和推理在显存使用上的差异。

图5 训练与推理的显存增长路径

---

六、一次从 baseline 到优化的实验流程

6.1 先定义可比较的 baseline

拿第 44 篇 mini-GPT 做实验时，先固定随机种子、模型配置、数据 batch、B、T、设备和训练步数。baseline 可以是 FP32 的 20 个 step；记录 loss、平均 step time、有效 token/s、max_memory_allocated 和 max_memory_reserved。接着只改变一个变量，例如打开 AMP，再重复同样实验。

def token_throughput(batch_size, tokens, steps, elapsed):
    return batch_size * tokens * steps / elapsed

# 报告中至少保存这些字段，而不是只打印一个 loss：
report = {
    "dtype": "float16 + autocast",
    "batch": 8,
    "tokens": 256,
    "steps": 20,
    "loss_last": 0.0,       # 替换为真实值
    "step_ms": 0.0,         # 替换为真实值
    "peak_allocated_mib": 0.0,
}
print(report)

一次实验只改一个因素，才能把差异归因给 AMP、T 或 batch。若为了让 AMP 跑通同时换了模型、数据和 optimizer，最后即使结果变好，也无法知道是哪项改变起作用。

6.2 正确性先于速度

低精度实验至少做四个检查：loss 是有限数；梯度没有 nan/inf；验证指标没有明显退化；小数据 overfit smoke test 仍能下降。不要要求半精度与 FP32 每一步 loss 完全相同，浮点运算次序本来就可能不同；应比较一段训练后的趋势和最终指标。

def assert_finite(model, loss):
    assert torch.isfinite(loss).item(), "loss is not finite"
    for name, p in model.named_parameters():
        if p.grad is not None:
            assert torch.isfinite(p.grad).all(), f"bad grad: {name}"

# 在 scaler.unscale_(optimizer) 后调用，才检查原始梯度尺度

如果 FP16 的 loss 变成 nan，先确认 scaler 顺序和学习率，再尝试 BF16 或把敏感的 loss、norm 放到 float32。如果速度反而下降，检查 GPU 是否支持目标 dtype、矩阵 shape 是否太小、是否被 CPU 数据加载或同步限制。优化的结果可以是“更快”“更省显存”或“更稳定”，三者未必同时发生。

6.3 生成一份最小实验报告

每次 profile 或显存实验至少保存配置和结论。推荐字段如下：

hardware: GPU model / CPU
software: Python, PyTorch, CUDA
model: layers, dim, heads, vocab
input: batch, tokens, dtype
mode: train / prefill / decode
amp: off / fp16 / bf16
warmup: N
measurement: steps, synchronization rule
result: loss, tokens/s, p50/p95, peak allocated/reserved
finding: bottleneck and next hypothesis

报告不是形式主义。未来换驱动、换 kernel、接入 torch.compile 或迁移到 vLLM 时，只有这些上下文能帮助你解释“为什么今天的数字不一样”。下一篇会接着这个 baseline 讨论编译和算子融合：在 profile 发现 Python 调度或许多小 kernel 后，怎样让 PyTorch 把一段算子组合成更少的执行单元。

---

七、踩坑大全

坑 1：直接对整个模型调用 .half()。 现象是 loss 变成 nan 或某些算子报 dtype 错误。原因是所有参数和敏感计算都被强制成 FP16；优先使用 autocast，让算子按设备和实现选择精度。

坑 2：忘记 GradScaler 的 unscale 顺序。 现象是梯度裁剪阈值失效或 scaler 报错。原因是裁剪发生在放大后的梯度上；应先 scaler.unscale_(optimizer)，再裁剪，最后 scaler.step()。

坑 3：把 BF16 当成“FP16 的更高精度版”。 两者都占 2 字节，但指数和尾数分配不同。BF16 动态范围更大，尾数精度更低；是否更适合要结合硬件、算子和训练结果验证。

坑 4：用 nvidia-smi 的数字判断活跃 Tensor。 现象是明明删了 Tensor，显存仍显示很高。原因是 caching allocator 保留缓存；同时查看 allocated、reserved 和峰值，必要时用 memory_summary()。

坑 5：把 empty_cache() 当成显存修复按钮。 它只能释放未使用的缓存，不能释放仍被引用的 Tensor；频繁调用还会增加申请开销。先修复引用生命周期、batch 和 shape，再考虑缓存策略。

坑 6：在日志列表里保存带图的 loss。 现象是每个 step 显存逐渐上涨。原因是 list 持有 loss 及其计算图；记录 loss.item() 或 loss.detach()，保存预测则按需移到 CPU。

坑 7：profiler 没有调用 prof.step()。 现象是 trace 为空或 schedule 一直停在第一阶段。每个要观察的迭代末尾调用 prof.step()，并让 warmup 足够覆盖首次分配和 kernel 初始化。

坑 8：拿 profiler 的时间和普通 benchmark 直接比较。 profiler 会改变执行成本，还可能引入同步。性能比较应关闭 profiler，固定 warmup、输入和同步边界；profiler 只用来解释瓶颈。

坑 9：只 profile 一次混合阶段。 prefill、decode、训练 forward/backward 的瓶颈不同。用标签或分开的 profile 记录它们，并写明 batch、上下文和请求并发。

坑 10：只降低 dtype，不限制服务并发。 推理参数节省了显存，但 KV Cache 随请求和 token 增长，仍会 OOM。服务端要设置最大上下文、cache 配额、batch 策略和 admission control。

坑 11：CPU smoke test 强行使用 CUDA API。 现象是没有 GPU 的环境启动就报错。用 torch.cuda.is_available() 配置 device、activities 和 scaler；CPU 仍可验证 shape、loss 和有限值。

坑 12：只追求某一项峰值数字。 峰值降低但 tokens/s 下降，或者速度提高却验证 loss 恶化，都不能算完整优化。至少同时报告正确性、吞吐、延迟和显存。

---

八、小结

1. dtype 同时影响元素字节数、动态范围、有效精度和 GPU kernel 路径；FP16 与 BF16 都是 2 字节，但取舍不同。
2. AMP 通常由 autocast 和 GradScaler 配合完成；autocast 选择算子精度，scaler 处理 FP16 梯度下溢风险。
3. 参数之外还要计算梯度、optimizer state、激活、临时 workspace 和 allocator cache；训练与推理的峰值来源不同。
4. allocated 是活跃 Tensor，reserved 是 allocator 保留量；测峰值前 reset，GPU 计时需要明确 warmup 和同步边界。
5. torch.profiler 要用 wait、warmup、active schedule 控制成本，并在每步调用 prof.step()；trace 的目标是验证可操作的优化假设。
6. 训练侧重点是激活和反向，推理侧重点是 KV Cache、并发和上下文；AI Infra 优化必须把正确性、吞吐、延迟与显存一起报告。

---

九、课后练习

1. （基础） 分别创建 float32、float16、bfloat16 的同形状 Tensor，打印 element_size()，验证理论字节数与实际数量级。
2. （基础） 把第 44 篇 mini-GPT 的 FP32 训练循环改成 autocast + GradScaler，记录 loss、scale 和每步耗时。
3. （理解） 比较 FP16 与 BF16 的 torch.finfo，再在你的 GPU 或 CPU 上跑一个小 batch，说明哪种 dtype 被支持、结果是否稳定。
4. （理解） 写一个显存报告函数，分别打印 allocated、reserved、max_memory_allocated；构造一个带 Tensor list 的循环，再用 detach() 修复显存增长。
5. （进阶） 用 torch.profiler profile 训练 5～10 步，设置 wait、warmup、active，比较 record_shapes=False/True 对 trace 和运行时间的影响。
6. （AI Infra） 固定模型和 dtype，对比 B=1/4/8、T=128/256/512 的峰值显存与 tokens/s，画出 attention 相关内存随 T² 增长的曲线。
7. （进阶） 分别 profile prefill 与逐 token decode，观察矩阵乘法、数据拷贝、同步和 KV Cache 读写的差异，写出一个可证伪的优化假设。
8. （综合） 为实验报告补齐硬件、软件、配置、warmup、同步规则、峰值和 p50/p95；下一次改动前先保存 baseline，避免“优化”没有可比对象。

---

下一篇预告

第46篇｜torch.compile 与算子融合入门。本篇我们用 AMP 和 profiler 建立了 baseline：知道哪些算子耗时、显存峰值发生在哪里，以及训练和推理的账本为什么不同。下一篇会从 eager mode 的 Python 调度出发，理解 torch.compile 如何捕获计算图、做算子融合与生成更高效的执行代码，并讨论首次编译开销、动态 shape、graph break 和如何用本篇的 benchmark 验证收益。之后第 47 篇进入 KV Cache、batching 与采样，第 48 篇再把这些能力接到 vLLM、SGLang 和 Agent 服务的完整 AI Infra 路线。