GPU、混合精度与分布式训练入门

一句话定义

GPU 用数千核心的并行吞吐匹配张量运算的天然并行性；混合精度（AMP）用 fp16/bf16 存储与前向、fp32 主权重与累加换取约 2 倍提速与减半显存；分布式数据并行（DDP）把批次切给多卡、各卡算梯度再同步平均——三者是「模型训得动」的三级火箭。

直觉

CPU 像「几位教授解题」，每人又快又聪明但人数有限；GPU 像「几千名小学生同时做加法」——单题不快、人海无穷。混合精度像「草稿纸用便宜纸、正式答卷用贵纸」；DDP 像「把一沓卷子分给几个考场批改，最后取平均」。

为什么重要

kp-016/kp-021 的项目规模下 CPU 尚可，一旦模型或数据上一个量级，「显存 OOM」与「训练一周跑不完」就是家常便饭。理解显存构成（激活缓存是 kp-007 的直接后果）、混合精度的数值保障（loss scaling）与数据并行的梯度同步（全归约），是从「会训模型」到「训得起模型」的工程分水岭；更大规模（流水线/张量并行、ZeRO 分片）属于大模型站点的领地，此处只立路标。

前置知识

• kp-007（激活缓存——显存大头）、kp-011（有效批大小与学习率联动）。

核心概念

• 显存构成：参数 + 梯度 + 优化器状态 + 激活缓存（反向要用，kp-007），训练时大头常是后者。
• fp32/fp16/bf16：单精度 / 半精度（动态范围窄）/ 脑浮点（指数位与 fp32 同宽、范围同、精度低）。
• AMP 自动混合精度：autocast 前向用半精度，梯度缩放（loss scaling）防 fp16 下溢。
• 梯度累积（gradient accumulation）：多个小批梯度累加后再更新，单卡凑等效大批（kp-011 的呼应）。
• DDP（DistributedDataParallel）：每卡持有完整模型副本，各卡处理不同批次，反向时梯度全归约（All-Reduce）取平均。
• 有效批大小 = 单卡批 × 卡数 × 累积步数——学习率与缩放规则必须跟着它走（kp-011）。

原理与机制

GPU 的机制基础：张量运算（矩阵乘、卷积）可分解为海量独立的乘加，恰合 GPU 的 SIMT 架构；batch 维天然并行，因此 kp-002 的形状设计直接决定吞吐——好形状让核函数满载，坏形状（碎片化小张量）让 GPU 空转。显存账本：以 13 层网络为例，参数与梯度各占 P 字节（fp32），Adam 额外存 m/v 两份（2P），激活缓存随批大小与序列/图像尺寸线性增长且常超参数本体——这解释了「加大批次最先爆显存」与梯度检查点（重算换显存）技术的存在理由。混合精度的数值机制：fp16 只有 5 位指数，微小梯度（<6e−8）会整体下溢为 0——AMP 的对策是损失先乘一个大因子 S（如 65536）再反向，把梯度整体抬进可表示区间，优化器更新前再除回来；bf16 的指数位与 fp32 同宽，无需 loss scaling，是新一代 GPU 的默认。主权重保持 fp32 的原因：参数更新量 η·g 常小于 fp16 的最小精度，若权重本身是 fp16 会「更新丢失」。DDP 的机制：每卡一份完整模型，前向各算各的批次；反向时梯度分桶广播做 All-Reduce 平均——各卡最终拿到相同的平均梯度，数学上严格等效于用「卡数×批」的大批做一次更新（这正是 kp-011 线性缩放规则的用武之地）。启动姿势的抽象：进程组初始化 → DistributedSampler 切数据 → DDP 包装模型 → 只在 rank 0 记日志存盘（kp-028 的档案纪律在多卡下同样成立）。

公式或模型

混合精度与数据并行的核心关系式：

显存 ≈ 参数 P + 梯度 P + 优化器状态 cP（Adam c=2） + 激活缓存（∝ batch×序列长×宽度）
AMP：  loss' = S·loss → backward → 梯度/S 后交给优化器（S≈2^16，fp16 专用）
DDP：  g_i = ∇L(batch_i) ，All-Reduce 后每卡 θ ← θ − η·(1/N)·Σ_{i=1..N} g_i
有效批： B_eff = B_card × N_card × A_accum

变量说明：N 为卡数、A 为累积步数。第三式的 (1/N)·Σ 即梯度平均——与单卡大批梯度在期望上等价（样本均匀切分时）。

图示

单卡：  GPU0: 模型副本 ← batch0 → 梯度 g0 → 更新

DDP：   GPU0: batch0 → g0 ─┐
        GPU1: batch1 → g1 ─┼─► All-Reduce 平均 ḡ ─► 各卡同步更新（参数保持一致）
        GPU2: batch2 → g2 ─┤
        GPU3: batch3 → g3 ─┘
        等效批 = 4×B_card → 学习率按 kp-011 线性缩放规则重调

直观类比

loss scaling 像「体重秤量不准 50 克以内的变化，就把人连货一起抬上秤再去皮」——整体放大让变化落到量程内，读数后再换算回来。

实例或案例

kp-021 项目的三级升级对照（同一模型同一数据）：

# ① 单卡混合精度（AMP 三件套）
scaler = torch.cuda.amp.GradScaler()
for x, y in loader:
    with torch.cuda.amp.autocast():          # 前向半精度
        loss = lossf(model(x), y)
    opt.zero_grad()
    scaler.scale(loss).backward()            # 梯度放大
    scaler.step(opt); scaler.update()        # 除回并自适应调整 S
# ② 梯度累积：累积 4 步再 step（等效批 ×4）
# ③ DDP 启动：torchrun --nproc_per_node=4 train.py
#    model = DDP(model.to(dev), device_ids=[rank])
#    sampler = DistributedSampler(dataset)   # 每卡切不同数据

实测参考（ResNet 变体 @CIFAR-10）：fp32 约 55 分钟 → AMP 约 28 分钟、显存约降 35%；四卡 DDP 再除以约 3.2（非 4.0，因同步开销）。踩坑实录三则：AMP 下 BN 统计量自动保持 fp32（框架处理，无需手管）；DDP 下日志与存盘只在 rank 0 做，否则文件互相覆盖；多卡后有效批变大，学习率不重调则收敛明显变慢（kp-011 规则的现实教育）。

常见误区

• 「GPU 一切都比 CPU 快」：小张量、碎片化操作、数据预处理常是 CPU 瓶颈；GPU 利用率低时先查数据管道（kp-015）。
• fp16 训练不加 loss scaling：微小梯度下溢归零，模型「静悄悄地学不动」（kp-027 的 silent failure 之一种）。
• 「DDP 就是把 batch 传给多卡」：核心是梯度同步（All-Reduce）而非数据分发；同步缺失会让各卡参数悄悄分叉。
• 忽略有效批大小变化：多卡/累积后不重调学习率与调度器 T_max，收敛行为全变（kp-011）。

与其他知识点的关系

显存账本来自 kp-007 的激活缓存；批大小与学习率联动是 kp-011 的规则放大版；BN 统计量在 AMP 下的处理呼应 kp-014；多卡下的日志/检查点纪律是 kp-028 的多进程版；更大的并行维度（流水线/张量并行、ZeRO）在大模型站点展开，此处一句话指路。

自测题

1. 训练显存的四大构成中，哪个随批大小线性增长且常是最大项？

答案要点：激活缓存（反向传播需要的中间量，kp-007）；参数/梯度/优化器状态与批大小无关。

1. AMP 为什么需要 loss scaling？bf16 为什么不需要？

答案要点：fp16 指数位少、小梯度下溢为 0，需放大保精度；bf16 指数位与 fp32 同宽，动态范围足够。

1. 四卡 DDP、单卡批 64、不累积时有效批多大？学习率应如何调整？

答案要点：256；按线性缩放规则约放大 4 倍并配合预热（kp-011）。

延伸阅读

• Micikevicius 等《Mixed-Precision Training» ICLR 2018（AMP 数值方案的原始出处）。
• 李沐《动手学深度学习》第 13 章「计算机视觉」前置的 GPU 与多设备计算章节。