前向传播与反向传播概述

一、前向传播的产物

一次前向传播结束后，除了标量 loss，还会留下两项反向传播必须依赖的东西。

为什么 loss 必须是标量

这不是实现约定，而是 VJP 机制的硬性要求。反向传播的起点需要一个种子梯度；标量 loss 天然提供 ∂L/∂L = 1，一次逆序遍历即可得到全部参数的梯度。若 loss 是 n 维向量，就必须为每个分量提供种子并各跑一次扫描，等价于把反向计算量乘以 n。这也是 CrossEntropyLoss 默认 reduction='mean' 的原因。

缓存的内容

被缓存的并非前向全部中间结果，而是「局部导数足以算出」的最低限度：

激活缓存的逐项拆解

以 7B 模型（32 层 / hidden 4096 / 32 head / 词表 151936）、batch 1、序列长度 512 为例：

该表定位了长上下文场景的瓶颈。序列长度从 512 增至 4096 时，Q/K/V、mask、FFN 中间激活按长度线性增长 8 倍；Attention 分数矩阵按长度平方增长 64 倍（FlashAttention 因此选择不缓存、反向重算）。同一模型的激活缓存会由 1.5 GB 增至 10 GB 以上。

对应的估算公式：激活显存 ≈ batch × seq × hidden × layers × 常数。这是长上下文的瓶颈落在激活而非权重上的原因。

二、反向传播的机制

反向传播本身不引入新的数学，它是链式法则在计算图上的逆序编排。每个算子被调用时的动作固定：

1. 确定种子。loss.backward() 等价于 loss.backward(torch.ones_like(loss))。对非标量张量直接调用 backward 会报 grad can be implicitly created only for scalar outputs。
2. 逆拓扑序遍历。按计算图依赖的逆序处理节点，保证节点被访问时，其所有下游分支的梯度均已到达并可合并。顺序错误会导致重复计算或结果不完整。
3. 只做 VJP，不构造完整 Jacobian。对 y = Wx，完整 Jacobian 是 out × in 个元素；实际只需针对当前的上游梯度 g 计算 gᵀ · x，一次矩阵乘法。计算量因此从参数量级降到激活量级——这是 autograd 可行的前提。
4. 梯度累加而非覆盖。同一张量被多个算子消费（残差分支）或权重共享（tied embedding、共享专家）时，梯度按加法合并。.grad 是累加缓冲，因此每个 step 必须显式清零。

三、微调策略的差异

LoRA 冻结底模权重，仅训练旁路矩阵 A、B。反向传播的覆盖范围因此收缩到旁路支路——但这不改变激活缓存的需求。

LoRA 的前向为 h = Wx + (α/r)·B·A·x，求导后两个待训练矩阵的梯度分别为：

• dL/dB = (α/r)·g·(Ax)ᵀ：需要前向缓存的 Ax（A 的输出）。
• dL/dA = (α/r)·Bᵀg·xᵀ：需要前向缓存的 x（该层输入）。

两个梯度都依赖前向缓存。LoRA 缩小的是求导范围，并未减少激活的保留量。

把这一机制折算到显存（7B 量级估算）：

三种策略对反向传播的覆盖范围与代价：

误解一：LoRA 的显存收益来自激活

并非如此。激活缓存量不变，前向计算量也不变。LoRA 的收益集中在优化器状态：Adam 为每个参与训练的参数维护两份 fp32 状态（一阶矩 m、二阶矩 v）。7B 模型全量微调时该项为 参数量 × 4 字节 × 2 = 56 GB，加上 fp32 主权重 28 GB，合计约 84 GB，即图 4 中占比最大的橙色段。LoRA 仅对 A、B 维护优化器状态，合计不足 0.5 GB。

结论：LoRA 使训练能够在小显存设备上启动，但不解决上下文长度增长带来的激活压力。前者由优化器状态决定，后者只能依靠 gradient checkpointing 或长上下文并行策略。

误解二：gradient checkpointing 是零成本优化

不是。选择不保留中间激活，就必须在反向传播到达该段时重新执行一次前向以重建它们。反向传播本身约为 2 倍前向 FLOPs，重算再叠加 1 倍前向，总开销由 2 倍升至 3 倍，即约 +50% 训练时间。这是用计算换显存，实践中通常只对部分层启用（use_reentrant=False + 选择性 checkpoint）。

四、实现要点

dL/dx = g · W
dL/dW = gᵀ · x              # 依赖前向缓存的输入 x
dL/db = Σ_batch g           # 在 batch 维上求和

dL/dz_i = p_i - y_i         # p 为前向输出的概率，y 为 onehot
                            # 因此只需缓存 p，无需保留 logits

dL/dx = (1/σ) · ( g − mean(g) − ŷ · mean(g · ŷ) )
                            # 依赖前向缓存的 μ、σ

# 1. 冻结底模：前向正常执行，反向不产生梯度
for p in model.base_model.parameters():
    p.requires_grad = False

# 2. loss 必须标量化，反向传播才具备种子
optimizer.zero_grad(set_to_none=True)   # 比 zero_grad() 少一次 memset
loss = criterion(model(**batch).logits, batch["labels"])
assert loss.dim() == 0

# 3. AMP：先还原梯度尺度，再裁剪
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(trainable_params, 1.0)
scaler.step(optimizer)
scaler.update()

# 4. 校验：冻结参数不应出现 .grad
for n, p in model.base_model.named_parameters():
    assert p.grad is None, n