训练基础：一个 Step 的执行图、前向、反向与优化器

推理一个 step 只做一次前向；训练一个 step 要做「前向 + 反向 + 参数更新」三段，并且显存里同时住着四份数据：权重、梯度、优化器状态和激活。理解训练系统——显存为什么爆、分布式为什么要切片、checkpointing 为什么值得——都从这一个 step 的执行图开始。

1. 一个 Step 总览

训练循环的最小形态：

for batch, labels in dataloader:
    optimizer.zero_grad()        # 清空上一轮梯度
    logits = model(batch)        # ① 前向：逐层算子，顺带保存反向要用的激活
    loss = criterion(logits, labels)
    loss.backward()              # ② 反向：autograd 沿执行图逆序求 ∇W
    optimizer.step()             # ③ 更新：AdamW 用 m/v 状态改写 Θ

把这三段摊开，就是训练的执行图：

三条流向对应三种数据：

• 蓝色（前向数据流）：batch tokens 与权重 Θ 进入 forward，逐层产出激活，最后得到 logits 和标量 loss。前向顺手把反向要用的激活「记入账本」。
• 琥珀色（反向梯度流）：从 loss 出发逆序遍历算子，每一步做两件事——算出对输入的梯度 dL/dX 继续向上游传，算出对权重的梯度 dL/dW 留下来。反向要回头读前向保存的激活（虚线）。
• 绿色（参数更新回路）：optimizer 把 ∇W 和自己的状态（Adam 的 m、v）合成更新量，写回 Θ；下一个 step 用新权重重复全过程。

2. 执行图是什么样子的

前向代码里的每个算子（GEMM、LayerNorm、Softmax、GELU……）是图的一个节点，张量是节点之间的边。以一个两层 MLP 为例：

forward:   x ─▶ [Linear1] ─▶ a ─▶ [GELU] ─▶ h ─▶ [Linear2] ─▶ y ─▶ [CE] ─▶ loss
backward:  dx ◂─ [Linear1'] ◂─ da ◂─ [GELU'] ◂─ dh ◂─ [Linear2'] ◂─ dy ◂─ [CE'] ◂─ 1
                 └─▶ dW1                                            └─▶ dW2

要点：

• 反向图是前向图的逆序镜像。PyTorch eager 模式下，每执行一个前向算子就在输出张量上挂一个 grad_fn，反向时 autograd 从 loss 沿 grad_fn 链逆序回调，每个算子的 backward 实现产出 dL/dX（继续上传）和 dL/dW（挂到对应参数上）。反向不需要也不应该手写。
• 反向依赖前向的中间值。比如 Y = X·W 的 wgrad 是 dW = Xᵀ·dY——必须留着前向的输入 X 才能算。这就是 save_for_backward 的来源，也是激活显存的来源：反向不是免费重放，前向必须替它「记账」。
• 动态图 vs 静态捕获。eager 每个 step 重新建图，灵活但有 Python 调度开销；torch.compile、XLA 这类路径先捕获整张图（含反向）再执行，可以做算子融合、省 kernel launch。图是逻辑视图，物理执行仍然是一串 GPU kernel，在 trace 里能直接看到 forward/backward/step 三段的分界（见 GPU Trace 分析）。

3. 四份数据：权重、梯度、优化器状态、激活

一个 step 期间显存里的常驻与流动数据：

数据	符号	大小	生命周期	混合精度下的精度
权重	Θ	P	常驻全程	bf16 前向用 + fp32 master
梯度	∇W	P	step 内产生，step 末消费	bf16
优化器状态	m, v	2P	常驻全程	fp32
master 权重	Θ_fp32	P	常驻全程	fp32
激活	activations	∝ B·S·L·H	forward 累积，backward 逐层释放	bf16

前三份半（Θ、∇W、m、v、master）只随参数量 P 变化，与 batch、序列长度无关，合称静态显存。bf16 混合精度 + AdamW 的标准账本：

$$M_{\text{static}} = \underbrace{2P}_{\text{bf16 } \Theta} + \underbrace{2P}_{\text{bf16 } \nabla W} + \underbrace{4P}_{\text{fp32 master}} + \underbrace{4P}_{m} + \underbrace{4P}_{v} = 16 \text{ bytes/param}$$

7B 模型仅静态部分就是 112 GB——单卡 80 GB 的 H100 连「站着不动」都做不到，这是 ZeRO/FSDP 把这些状态切片分置的根本动机（切分后的组合见 Megatron 并行）。精度的具体分配见 混合精度训练。

第四份激活是训练特有、且唯一随 B 和 S 增长的部分，它是 batch size 的实际约束。作为对照：推理侧没有反向、没有优化器，显存只有权重（2P）加一份随上下文增长的 KV Cache；训练侧的状态账比推理重一个量级。

4. 前向：算 loss 的同时记账

前向逐层执行 Y = f(X; W)，每个算子在算出输出的同时，把反向要用的量留下来：

• GEMM：保存输入 X（wgrad 需要）；
• 非线性（GELU/SiLU）：保存输入或输出（求导需要函数值）；
• Softmax：保存输出（反向用输出重算）；
• Dropout：保存 mask（反向按同一 mask 置零）；
• LayerNorm：保存均值、方差和归一化后的值。

最终 logits 对照 labels 算出标量 loss。loss 是整张图的终点，也是反向的起点——反向传播的种子就是 dL/dL = 1。

前向的算力账很整齐：每个参数在每层 GEMM 里对经过的每个 token 做一次乘加（2 FLOPs），即

$$C_{\text{forward}} \approx 2P \quad \text{FLOPs/token}$$

5. 反向：链式法则的工程化

反向就是链式法则的机械执行：

$$\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} \cdot \frac{\partial Y}{\partial X}, \qquad \frac{\partial L}{\partial W} = X^{\mathsf{T}} \cdot \frac{\partial L}{\partial Y}$$

对每个 GEMM 算子，backward 要算两个矩阵乘：dgrad（∂L/∂X，继续向上游传）和 wgrad（∂L/∂W，留下给 optimizer）。各需 2P FLOPs/token：

$$C_{\text{total}} = \underbrace{2P}_{\text{forward}} + \underbrace{2P}_{\text{dgrad}} + \underbrace{2P}_{\text{wgrad}} = 6P \quad \text{FLOPs/token}$$

这就是「训练算力 ≈ 6ND」（N = 参数量，D = 训练 token 数）的来源，估算训练时长和算力成本的第一公式。

反向同时是激活显存的「释放过程」：每算完一层的反向，该层保存的激活就可以释放。所以一个 step 的显存曲线是——前向爬坡（激活逐层堆积）→ loss 处到峰 → 反向下坡（激活逐层释放、梯度累积）→ step 原地更新。

Activation checkpointing 打破「前向必须全量记账」的默认：只保存少数层边界激活，反向到某层时先用边界激活就地重跑一遍该层前向，把中间激活补出来。显存从随 L 线性压到只存边界，代价是多约一次前向的算力（6P → 8P per token）。它改执行图不改语义，是「时间换显存」的元操作，也是长序列大 batch 训练的标配。机制细节与 selective 重算见 Activation Checkpointing。

6. 优化器：从 SGD 到 AdamW

优化器回答一个问题：拿到 ∇W 之后，这一步走多大、朝哪走。

• SGD：Θ ← Θ − η·g。零额外状态，但对 loss 曲面尺度敏感，大模型训练基本不用裸 SGD。
• SGD + Momentum：m ← β·m + g; Θ ← Θ − η·m。多存一份动量（+4P fp32），方向更稳。
• AdamW（当前主流）：

$$m \leftarrow \beta_1 m + (1-\beta_1)\, g, \qquad v \leftarrow \beta_2 v + (1-\beta_2)\, g^2$$ $$\hat{m} = \frac{m}{1-\beta_1^t}, \qquad \hat{v} = \frac{v}{1-\beta_2^t}, \qquad \Theta \leftarrow \Theta - \eta \left( \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon} + \lambda \Theta \right)$$

m 是梯度的一阶滑动平均（方向），v 是二阶平均（逐参数的步长归一），bias correction 修正冷启动，λ 是解耦的 weight decay。代价是第 3 节里那 12P 的 fp32 状态。

优化器还有三个常配套出现的部件：

• Gradient clipping：按全局范数截断梯度，防单步爆炸，先于 step 执行；
• LR schedule：warmup + cosine/linear decay，η 不是常数；
• 省状态变体：8-bit Adam、Adafactor、Lion 等压缩 m/v 的存储，改写但不推翻 16P 账本的结论。

7. 从单卡到分布式

本文的执行图是单卡语义视图。分布式并行改的是「数据住在哪里、通信插在哪条边上」，不改图本身：

• DP：每卡一份完整执行图和静态状态，反向末尾对 ∇W 做 all-reduce；
• TP/PP：把图本身切开，层内切张量或层间切段，通信插进前向反向的算子之间（切分见 Megatron 并行、调度气泡见 流水线并行、推理侧对照见 推理并行）；
• ZeRO/FSDP：把静态三份（Θ、∇W、optimizer states）按卡切片，用通信换回显存（机制见 ZeRO 与 FSDP）；
• Activation checkpointing：改图的记账方式（第 5 节），与上面任意组合叠加。

排查训练 OOM 的顺序也按这个账本来：先动激活（调小 B/S、开 checkpointing），再动静态（上分片或换省状态优化器）。

相关页面

• 混合精度训练16 B/param 账本的精度分配细节
• Activation Checkpointing激活显存的省法
• ZeRO 与 FSDP静态三份的分片机制
• Megatron 并行TP/PP/EP 如何切这张执行图
• Distributed Training 索引并行策略与训练系统导航
• GPU 架构前向反向落到 SM 与 HBM 上的物理基础
• GPU Trace 分析在 trace 里辨认 forward/backward/step 三段
• Compute-bound vs Memory-bound算术强度与 Roofline，训练的 GEMM 形状分析同样适用
• KV Cache推理侧的状态账，与训练侧四份数据对照