GPU/CUDA 与混合精度

上一课我们把数据用数据加载器（DataLoader）稳定地喂进训练循环——现在数据流通了，下一个问题是：怎么让它更快、更省显存。这一课不改训练逻辑，而是给你的训练循环装上两个加速器：把计算交给 GPU 的成千上万个线程并行做，以及用混合精度（mixed precision，AMP）让前向更快、激活省一半显存。这两件事看起来高级，但底层只是两个朴素的工程直觉：搬运数据有成本、低精度数有表示范围。把这两点吃透，你才真正理解后面 LoRA/QLoRA「在小显存上微调大模型」为什么成立。

GPU 并行直觉：为什么大矩阵乘在 GPU 上飞快

一次矩阵乘 \( (N, D) \times (D, M) \) 要算 \( N\times M \) 个输出，每个输出是一条长度为 \( D \) 的点积。关键在于：这些输出之间互不依赖，可以同时算。CPU 有几个到几十个强力核心，擅长复杂的串行逻辑；GPU 有成千上万个弱小核心，擅长「同一种简单运算、海量数据、各自独立」的场景——矩阵乘恰好是教科书式的例子。神经网络的绝大部分计算就是矩阵乘和逐元素运算，所以天然适配 GPU。

但有一个隐藏成本经常被忽略：CPU 和 GPU 是两块独立的内存，数据来回搬运要走 PCIe 总线，很慢。张量在 GPU 上算得再快，每次你把它拉回 CPU，都要等总线传输 + 等 GPU 把手头排队的活算完。这就引出本课第一个高频误区。

误区：循环里频繁 .item() / .cpu() / print(tensor)

GPU 的计算是异步的：你写 loss = model(x) 时，Python 只是把任务塞进 GPU 的队列就立刻返回，GPU 在后台慢慢算。这种「下了单不等菜」的流水线让 CPU 能持续派活、GPU 满负荷运转。但下面这些操作会强制同步（synchronize）——CPU 必须停下来，等 GPU 把队列里所有活都算完，才能拿到那个具体数值：

• loss.item()：把标量从 GPU 取回 Python float。
• tensor.cpu() / tensor.numpy()：把张量搬回主机内存。
• print(some_gpu_tensor)：打印要先把值读回来。

每个 step 里只要有一次这类调用，流水线就被打断一次。把它们从「每个 step」挪到「每 N 个 step 打一次日志」，往往是零成本的提速。

# 慢：每个 step 都同步一次
for x, y in loader:
    loss = step(x, y)
    print(loss.item())          # 每步都强制 CPU-GPU 同步

# 快：累加在 GPU 上，少数步才同步一次
running = 0.0
for i, (x, y) in enumerate(loader):
    loss = step(x, y)
    running += loss.detach()    # 仍在 GPU 上，不同步
    if i % 50 == 0:
        print(running.item() / 50)   # 每 50 步才同步一次
        running = 0.0

显存四大块：你的显卡装得下吗

「显存够不够」几乎是训练能不能跑起来的第一道门槛。训练时的显存主要由四块构成，记住它们的相对大小，你就能在不跑代码的情况下估出一个模型要多大的卡。设参数量为 \( P \)，fp32 下每个数占 4 字节：

这里最反直觉的一点：很多人只数参数，却忘了后面三块。光是「参数 + 梯度 + Adam 状态」就已经是参数本身的 4 倍（\(4P+4P+8P=16P\)）。注意区分两个倍数基准：Adam 状态本身是参数的 2 倍（\(8P\)），而「参数+梯度+Adam」三块合计才是参数的 4 倍（\(16P\)）。激活则是另一条独立的轴——它不随参数量、而随 batch size 和序列长度涨，这正是「显存爆了，把 batch 调小就好了」的原因。下面这个纯 Python 估算器把前三块算给你看。

def estimate_memory(P, bytes_per=4):
    GB = 1024 ** 3
    params = P * bytes_per          # 参数
    grads  = P * bytes_per          # 梯度 ≈ 参数
    adam   = P * bytes_per * 2      # Adam 的 m, v ≈ 2×参数
    total  = params + grads + adam
    return params/GB, grads/GB, adam/GB, total/GB

P = 1.3e8   # 1.3 亿参数（约等于 BERT-base 量级）
p, g, a, t = estimate_memory(P)
print(f"参数        : {p:.2f} GB")
print(f"梯度        : {g:.2f} GB")
print(f"Adam 状态   : {a:.2f} GB  <- 占大头")
print(f"合计(不含激活): {t:.2f} GB")
print(f"= 参数的 {t/p:.0f} 倍")

运行会得到：参数 0.48 GB，梯度 0.48 GB，Adam 状态 0.97 GB，合计约 1.94 GB——是参数的整整 4 倍，而且还没算激活。把 P 改成 7e9（7B 模型）你立刻明白：为什么 70 亿参数的模型，全量 + Adam 微调要 100 GB 量级的显存，而 LoRA 这类只更新极小一部分参数、从而砍掉绝大部分「梯度 + 优化器状态」的方法，才能让大模型在一张消费级卡上微调。

混合精度的数值原理：fp16 的范围有多窄

提速和省显存的另一把钥匙是降低精度：把计算从 32 位浮点（fp32）切到 16 位。16 位的数搬得快、算得快，激活只占一半显存。但天下没有免费的午餐——16 位能表示的数范围和精度都缩水了。要理解后面的 GradScaler，必须先亲眼看看 fp16 窄到什么程度。

下面这段纯 NumPy 代码演示三件事：fp16 的最大值、超过它会溢出（overflow）成 inf、以及很小的「梯度」会下溢（underflow）成 0——这正是训练 fp16 时梯度凭空消失的元凶。

import numpy as np

# 1) fp16 能表示的最大值
fmax = np.finfo(np.float16).max
print("fp16 最大值 =", fmax)          # 约 6.55e4（numpy 打印 65500.0；fp16 真正最大值是 65504）

# 2) 溢出：超过最大值 -> inf
#    这一步会打印一条 "RuntimeWarning: overflow encountered in cast"，
#    那不是报错，正是 numpy 在告诉你溢出了；下一行的结果仍是 inf。
big = np.float16(70000.0)
print("70000 转 fp16 =", big)         # inf

# 3) 下溢：很小的梯度 -> 0
grad = 1e-8
print("1e-8 转 fp16 =", np.float16(grad))   # 0.0  <- 梯度没了！

第三行是混合精度训练的核心痛点：深度网络里很多梯度天然就很小（\(10^{-7}\) 量级很常见），一转 fp16 就被压成 0，参数得不到更新，训练悄悄崩掉。注意 bf16 不太有这个问题（下一节解释），这个下溢主要是 fp16 的麻烦。

多说一句关于「多小才会下溢成 0」：fp16 除了正规数（normal），还有次正规数（subnormal）把可表示范围一直往下延伸到约 \(6\times10^{-8}\)（精确是 \(2^{-24}\approx5.96\times10^{-8}\)）。所以判 0 的真正门槛不是「最小正规数」（约 \(6.1\times10^{-5}\)），而是这个更低的最小可表示数约 \(6\times10^{-8}\)——比它还小（严格说小到约 \(3\times10^{-8}\) 以下）才会被四舍五入到 0。后文凡是讲「下溢成 0」，都以这个约 \(6\times10^{-8}\) 的下限为准。

GradScaler：用 loss scaling 把小梯度抬回可表示区

救法巧妙又朴素：既然小梯度落在了 fp16 表示不了的「太小」区域（小于约 \(6\times10^{-8}\) 就会被压到 0），那就在反向传播前把 loss 乘一个大数（如 1024）。由链式法则，loss 放大多少倍，所有梯度就同比放大多少倍——原本下溢成 0 的小梯度被抬回 fp16 能表示的区间。等梯度算完、要更新参数前，再除回这个倍数，数值就还原了。这个「乘大数 → backward → 除回来」的把戏就叫 loss scaling，PyTorch 用 GradScaler 自动帮你做（还会自动挑倍数、遇到 inf 就跳过这一步）。下面演示里的 1024 只是个固定的演示值，真实的 GradScaler 初始倍数约为 65536，并会在训练中动态调整。

import numpy as np

grad = 1e-8
scale = 1024

# 不缩放：直接转 fp16 -> 下溢成 0
print("不缩放:", np.float16(grad))                 # 0.0

# 缩放：先乘 1024 再转 fp16 -> 救回来了
scaled = np.float16(grad * scale)
print("乘 1024 后转 fp16:", scaled)                # 约 1.02e-05，非 0

# 更新参数前，把梯度除回原倍数（用 fp32 还原）
recovered = np.float32(scaled) / scale
print("除回 1024 还原:", recovered)                # 约 1.0e-8，回来了

跑出来：直接转是 0.0，乘 1024 后变成约 1.02e-05（活了），除回去约 1.0e-8（还原）。这就是 GradScaler 在底层做的全部事情——理解了这 8 行，你就不会再对「为什么 fp16 训练要套一个 scaler」感到神秘。

bf16 vs fp16：别把它俩当一回事

这是第二个高频误区：以为 bf16 = fp16。它们都是 16 位，但 16 位的「预算」分配方式完全不同。浮点数由符号位、指数位（exponent，决定范围）、尾数位（mantissa，决定精度）三部分组成：

关键洞察：bf16 把预算砸在指数位上，指数位和 fp32 一样多（8 位），所以动态范围和 fp32 几乎一致——小梯度不会下溢成 0，因此一般不需要 GradScaler。代价是尾数只有 7 位，精度低（相邻可表示的数之间间隔更大），但对深度学习这通常无所谓。而 fp16 指数位只有 5 位，范围窄，才会有上一节的溢出/下溢问题，必须配 GradScaler。

• 新硬件（Ampere 及以后的 NVIDIA 卡、TPU）→ 优先 bf16，省心，无需 loss scaling。
• 只有 fp16 可用（较老的卡）→ 用 fp16 + GradScaler。
• 记忆口诀：bf16 拼范围、fp16 拼精度；范围不够才会下溢，所以是 fp16 需要救、bf16 不太需要。

把它们拼进训练循环：标准 AMP 三件套

第 6 课你已经手写过完整训练循环（forward → loss → backward → step → zero_grad）。混合精度只是在这个骨架上加两样东西，逻辑一行没改：

1. with autocast(...) 包住前向 + 算 loss：PyTorch 自动把里面的运算切到 bf16/fp16（矩阵乘等），算得快、激活省一半显存。
2. 用 GradScaler 包住反向与更新（fp16 时需要；bf16 可省）：scaler.scale(loss).backward() 先放大再反向，scaler.step(opt) 内部除回缩放再更新，scaler.update() 调整下次的缩放倍数。

注意 AMP 主要省的是激活（前向中间结果用低精度存），它不会把参数、梯度、Adam 状态那 \(16P\) 也减半——为了数值稳定，这些通常仍保留 fp32 主副本。所以 AMP 省显存的口径是「激活减半」，而不是「整块显存减半」。

import torch
from torch.amp import autocast, GradScaler

device = "cuda"
model.to(device)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler("cuda")                 # fp16 才需要；bf16 可不用（推荐显式传设备）

for x, y in loader:
    x, y = x.to(device), y.to(device)
    opt.zero_grad()

    with autocast(device_type="cuda", dtype=torch.float16):  # ① 前向降精度
        pred = model(x)
        loss = loss_fn(pred, y)

    scaler.scale(loss).backward()           # ② 放大 loss 再反向（救小梯度）
    scaler.step(opt)                         #    内部把梯度除回再 step
    scaler.update()                          #    动态调整缩放倍数
    # 注意：不要在这里 loss.item()，攒几步再打日志

# 用 bf16 时更简单：不需要 GradScaler
with autocast(device_type="cuda", dtype=torch.bfloat16):
    pred = model(x)
    loss = loss_fn(pred, y)
loss.backward()
opt.step()
opt.zero_grad()

另外两个你会反复听到、但这一课只需知道名字的工具：torch.compile(model) 把模型编译成更快的融合算子（一行包一下，常有几成提速）；分布式数据并行（DDP，DistributedDataParallel）把同一个模型复制到多张卡、各算一部分 batch 再同步梯度，是多卡训练的标准做法。两者的细节留到规模化的模块再展开。

调一调，观察现象

微任务 1：把缩放倍数调小，看小梯度还能不能救回来。把 scale 从 1024 改成 1。预期现象：1e-8 * 1 = 1e-8 仍然小于 fp16 能表示的最小可表示数（约 \(6\times10^{-8}\)，含次正规数），转 fp16 后还是 0，救不回来；而 scale=8 时 8e-8 已经超过这个下限，能被表示成约 6e-08，就救回来了。为什么：loss scaling 要把梯度抬到 fp16 的最小可表示数（约 \(6\times10^{-8}\)）之上，倍数太小就抬不够——这也是为什么 GradScaler 会从一个很大的初始倍数（如 65536）开始动态调。下面代码把这条「倍数从抬不够到抬够」的分界线一次跑给你看。

import numpy as np
grad = 1e-8
for scale in [1, 8, 256, 1024, 65536]:
    s = np.float16(grad * scale)
    print(f"scale={scale:>6}: fp16={s}  {'救回' if s != 0 else '仍为0'}")
# 预期：scale=1 仍为0；scale>=8 起都救回（8 -> 约 6e-08）

微任务 2：换成 bf16，看下溢问题是否消失。把下溢的值用 np.float16 和 bfloat16 各转一次（NumPy 无原生 bfloat16，我们用「指数位和 fp32 一样、所以范围一样」这个事实手动验证：bf16 能表示 \(10^{-8}\)）。预期现象：fp16 下溢成 0，而 bf16 的范围足够、不为 0。为什么：bf16 有 8 位指数，动态范围≈fp32，所以一般不需要 loss scaling。

import numpy as np
grad = 1e-8
print("fp16:", np.float16(grad))            # 0.0，下溢
# bf16 指数位=8，最小正规数约 1.18e-38，远小于 1e-8
# 故 bf16 能表示 1e-8（这里用 fp32 近似其“范围”验证）
print("fp32(模拟 bf16 的范围):", np.float32(grad))  # 1e-8，不为 0
print("=> bf16 范围足够，无需 loss scaling")
# 提醒：这里只借 fp32 验证“范围”（指数位相同），bf16 的精度（尾数 7 位）
#       仍远低于 fp32，别把 bf16 和 fp32 等同起来。

微任务 3：把模型规模放大，看显存怎么涨。在显存估算器里把 P 从 1.3e8 改成 7e9。预期现象：合计从约 1.9 GB 跳到约 100 GB，且 Adam 状态始终是参数的 2 倍、合计始终是参数的 4 倍。为什么：四大块里前三块都正比于 \(P\)，所以「能不能装下」基本由参数量线性决定——这也是 LoRA 砍掉梯度+优化器状态后能省一大截的根因。

def estimate_memory(P, bytes_per=4):
    GB = 1024 ** 3
    total = (P + P + 2*P) * bytes_per       # 参数+梯度+Adam
    return total / GB

for P in [1.3e8, 1.3e9, 7e9]:
    print(f"P={P:.1e}: 合计(不含激活) ≈ {estimate_memory(P):.1f} GB")

动手练习

1. 写一个「精度损失」演示。用 NumPy 验证 fp16 尾数只有 10 位带来的精度问题：计算 np.float16(1.0) + np.float16(0.0005)，打印结果，并解释为什么它没有正确加上 0.0005（提示：1.0 附近 fp16 的最小间隔约 0.001，0.0005 约为半个间隔，按「四舍五入到最近可表示数」实际会显示 1.001，即只前进了一个最小间隔，而不是停在 1.0——总之 0.0005 这个量级在 fp16 里没法被精确表示）。再用 fp32 做同样的加法对比（会得到 1.0005）。
2. 给显存估算器加上「激活」和「混合精度」两个开关。扩展 estimate_memory：(a) 接收一个 act_gb 参数表示激活占用，加进合计；(b) 增加一个 amp=True 选项，让激活减半（混合精度的主要省显存来源在激活）。用 \(P=1.3\times10^8\)、激活 2 GB，对比开/关 AMP 的合计。注意：AMP 通常只减少激活，参数/梯度/Adam 那 \(16P\) 仍按 fp32 主副本算，别把它们也减半。
3. （纸面题，可选）解释下面这段代码慢在哪、怎么改：循环里写了 if loss.item() < best: best = loss.item()。指出它每步同步两次，并改写成「在 GPU 上比较、每 50 步才取一次最优值」。

掌握自检

• 我能说出 GPU 擅长矩阵乘的一句话理由（海量独立的简单运算并行做）。
• 我知道 .item()/.cpu()/print(tensor) 为什么慢，并能把日志从「每步」改成「每 N 步」。
• 我能不查表写出训练显存的四大块，并算出「参数+梯度+Adam ≈ 4×参数」（其中 Adam 状态本身是 2×参数）。
• 我能解释 fp16 为何会下溢成 0（小于最小可表示数约 \(6\times10^{-8}\)），以及 loss scaling 的「乘大数→backward→除回来」三步。
• 我能说清 bf16（多指数位=范围大，一般不用 scaler）与 fp16（范围小=需 scaler）的区别，不再混为一谈。
• 我能读懂 AMP 三件套：autocast 包前向、scaler.scale(loss).backward()、scaler.step+scaler.update。