第 11 讲 · 在 GPU 上加速 Transformer（二）：训练加速 · Accelerating Transformer on GPU (II): Training with LightSeq2

对应 Maps to：CMU 11-868 / 11-968 Large Language Model Systems（Spring 2026，Lei Li），Lecture 11 · Wed 2/18 · Accelerating Transformer on GPU Part 2。 阅读材料：Wang et al. (2022) LightSeq2: Accelerated Training for Transformer-based Models on GPUs（SC22）；原课 HW4 说明（Softmax / LayerNorm 的前向与反向内核）。 依据说明：原课第 10、11 讲共用一份 slides（llmsys-10-transformer-acc.pdf）。本讲覆盖其中训练相关的部分：LayerNorm 反向的归约改写（第 22 页）、softmax 的模板参数（第 24–26 页，本讲用在反向上）、技术 3：混合精度计算与加速的混合精度更新（第 27–29 页）、技术 4：显存复用（第 30 页，自注意力反向的显存管理）、训练实验（第 31–40 页：机器翻译、GPT-2、BERT、ViT 的加速，加速分解，可扩展性，显存）和训练加速的总结（第 41 页）。论文中关于训练器”去掉 FP32 副本”、“训练前扫描数据集按最大形状分配”的细节取自 LightSeq2 论文 §IV-C、§IV-D 和 §VI。模拟器里的反向内核与训练器内核、slides 写法与 HW4 写法的数值对比、γ/β 梯度的三种并行方式与 bank 冲突、softmax 反向要存的显存核算、FP16 存储吞掉更新的分析与小实验、⏱ CPU 上的优化器计时、静态显存规划器与缓存分配器的模拟、一整步训练的纸面账，都是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture11/（5 个脚本 + train_kernels.py；都在 10 秒以内） 预计阅读时间：3 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 说出训练比推理多出的三笔开销：反向传播（每个融合内核都要有反向，反向里同样有归约）、参数更新（几百个参数张量、FP16 与 FP32 之间的拷贝）、显存（为反向保存的激活、反向中一个接一个的临时张量、长度不同的批次）。
2. 推导 LayerNorm 和 softmax 的反向公式，并解释怎样把”有依赖的两次归约”重排成”可以一起做的两次归约”（一次同步）；说出 slides 写法和 HW4 写法的数值差别。
3. 为按列求和的 γ、β 梯度设计并行方式，说出”一个线程一列""原子加""32 × 32 分块 + 转置归约”各自的代价，以及共享内存为什么要多留一列。
4. 解释 LightSeq2 的混合精度训练器：参数和梯度”链接”成一整块 FP16 workspace、一个内核在寄存器里转换精度；算出它省的内核数、字节和显存，并说出去掉 FP32 主参数的前提。
5. 对一段反向计算做生命周期分析，算出临时显存的下界（例如自注意力反向的 3BLH + max{BL²N, 3BLH}），并解释为什么要训练前一次性按最大形状分配。
6. 把一整步训练拆成阶段做纸面账，解释为什么批越大、加速比越小。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  背景：训练比推理多了什么；LightSeq2 的四类技术在训练里各管哪一段
§2  LayerNorm 的反向：公式、两次同步变一次、数值代价；γ、β 梯度的并行      ← 代码 01
§3  softmax 的反向：只要 y、一次归约；与 dropout 反向融合；反向要存什么      ← 代码 02
§4  混合精度训练器：几百个张量 → 一整块 workspace；去掉 FP32 副本的前提     ← 代码 03
§5  显存管理：生命周期、静态规划、预先分配                                 ← 代码 04
§6  一整步训练的账；LightSeq2 的实验结果                                    ← 代码 05
§7 小结   §8 常见误区   §9 练习   §10 延伸阅读   §11 下一讲预告

1. 背景：训练比推理多了什么 · Why training is harder

1.1 一次训练迭代

第 10 讲 §1.2 说过，一次训练迭代是：前向 → 反向 → （多卡时）梯度同步 → 参数更新。推理只有第一步。LightSeq2 论文 §II 总结了训练的三点不同：

1. 精度要求更高：每一步的误差会在成千上万次更新里累积放大，所以参数更新要用 FP32 计算。
2. 显存管理更难：反向要用前向存下的激活；反向过程中还会产生一串临时的梯度张量。
3. 需要快的训练器（trainer，也就是优化器）：它要碰模型的每一个参数，一个参数也不能少。

另一方面，训练也有比推理简单的地方：不用处理自回归解码那种”每步多一个 token”的增长形状。

1.2 动机：反向和更新各占多少？

先看 05 号脚本的纸面账（GPT-2 large，36 层，d = 1280，批 8 × 1024 个 token，A100，基线是”PyTorch 不跨运算融合、逐张量更新”）：

反向比前向还长（GEMM 翻倍，非 GEMM 运算也各有自己的反向）；参数更新虽然只占 6%，却要发射 1,308 个内核。第 10 讲只解决了第一行的一部分。这一讲按 slides 的顺序，把剩下的几段逐个拆开。

1.3 LightSeq2 的四类技术在训练里管哪一段

slides 第 8 页的四类技术（第 10 讲 §1.3）在训练里的分工：

2. LayerNorm 的反向 · LayerNorm backward

2.1 公式

前向：x̂ᵢ = (xᵢ − μ)/σ，σ = √(var + ε)，yᵢ = γᵢ x̂ᵢ + βᵢ，m 是行长。前向把每行的 μ 和 r = 1/σ 存下来。记 gᵢ = γᵢ∇yᵢ。对 x̂ⱼ 求 xᵢ 的偏导（σ 里含 ε 也成立）：

$\frac{\partial \hat x_j}{\partial x_i} = \frac{\delta_{ij} - 1/m}{\sigma} - \frac{\hat x_i \hat x_j}{m\sigma}$

于是

$\nabla x_i = \sum_j g_j \frac{\partial \hat x_j}{\partial x_i} = \frac{1}{\sigma}\Big[g_i - \frac{1}{m}\Big(\sum_j g_j + \hat x_i \sum_j g_j \hat x_j\Big)\Big]$

这是 HW4 的写法。把 x̂ = (x − μ)/σ 代进去，按 Σg 和 Σg·x 整理，就是 slides 第 22 页（和论文）的写法：

$\nabla x_i = \frac{g_i}{\sigma} + \alpha_i \sum_j g_j + \beta_i \sum_j g_j x_j,\quad \alpha_i = \frac{(x_i - \mu)\mu - \sigma^2}{m\sigma^3},\quad \beta_i = \frac{\mu - x_i}{m\sigma^3}$

01 号脚本第 1 部分：两种写法都和 PyTorch 的自动求导差 4.4e-16（float64）。γ、β 的梯度简单得多：∇γⱼ = Σ_行 ∇y·x̂，∇βⱼ = Σ_行 ∇y（§2.4）。

2.2 为什么要”重排”：两次同步变一次

如果照计算图让自动求导去做（前向是第 10 讲的 7 个小运算），反向是一串有依赖的运算：先求 ∇x̂ = γ⊙∇y，再归约 Σ∇x̂·(x − μ) 得到对方差的梯度，用它算出每个元素对 x − μ 的梯度 dxc，再归约 Σdxc 才能减去均值。第二次归约要等第一次的结果——两次块内同步，一次接一次。

重排后的两个和 Σg 和 Σg·x̂（或 Σg·x）互不依赖，可以放进同一次树形归约里同时累加，只要一次同步。这和第 10 讲 LayerNorm 前向把方差改写成 E[x²] − E[x]² 是同一个思路。01 号脚本第 2 部分（8 行 × 64 列，块大小 32）：

放大到 8192 个 token × d = 1024（fp16）：自动求导式 302.8 MB、230.7 μs；融合 71.3 MB、49.9 μs。融合省的是字节和发射，重排省的是同步——两者是不同的技术，LightSeq2 两个都用。

要说明的是，这里的”自动求导式”是把 LayerNorm 拆成基本运算的写法（HW3 的 miniTorch 就是这样）。PyTorch 自带的 layer_norm 已经有专门的反向内核；LightSeq2 报告的 LayerNorm 加速约 4 倍（slides 第 34 页、论文图 16）是和它比的，来源是一次同步、float4 向量化读写，以及按形状挑选的模板参数。HW4 要求的 LayerNorm 反向平均快约 3.7 倍。

2.3 重排的数值代价

第 10 讲的教训是：E[x²] − E[x]² 在输入均值远大于标准差时彻底出错。反向的重排呢？01 号脚本第 3 部分（每行 ~ N(均值, 1)，m = 256，∇x 的最大相对误差）：

slides 的写法用原始的 x 求第二个和，αᵢΣg 和 βᵢΣg·x 里都有一项正比于 μ(xᵢ − μ)Σg/(mσ³)，相加时互相抵消，误差是 HW4 写法的 1.9–3.2 倍。但远没有前向那么糟：前向从 μ² 量级的两个数里减出 σ² 量级的差，误差 ∝ ε·μ²/σ²；这里抵消掉的项只比结果大约 μ/(σ√m) 倍，误差 ∝ ε·μ/(σ√m)。HW4 的写法先用存下的 μ、r 把 x 变成 x̂（量级 1）再求和，同样一次同步，没有这层抵消。同样是”一次同步”，选哪两个量一起归约，数值上不一样。

2.4 γ、β 的梯度：沿着行求和

∇γ、∇β 是”每一列一个和”：N 个 token（几千）加成 d 个数（上千）。一行一个块的写法用不上了。01 号脚本第 4 部分比较了三种并行方式：

• 一个线程一列：相邻线程读相邻的列，访存合并；但只有 d 个线程，每个顺序加 N 次——几乎是串行的。
• 原子加：并行度最高，可同一列的 8192 个原子加落在同一个地址上要排队；顺序还不确定，结果每次可能不同（第 10 讲 §5.1）。
• HW4 的 ker_ln_bw_dgamma_dbetta：一个 32 × 32 的线程块负责 32 列。线程 (x, y) 负责第 x 列、从第 y 行开始每隔 32 行累加一次（读仍是合并的）；部分和写进共享内存 buf[y][x]；再”转置着读” buf[x][y]，于是同一个线程束（同一个 y）拿到同一列的 32 个部分和，用 g.shfl_down 在寄存器之间加起来。

转置着读时，一个线程束的 32 个线程读共享数组的同一列：数组是 [32][32] 时地址相差 32 个字，全落在同一个 bank 上（32 路冲突）；声明成 [32][33]，地址相差 33 个字，32 个 bank 各一个。脚本里整个内核的平均轮数从 2.63 降到 1.00——这就是第 4 讲矩阵转置里的 padding 技巧。

分块版的并行度仍只有 32·d。要占满 GPU，可以再分两段：先把 N 行切成 P 段、每段算一份部分和（P·d 个线程），第二个小内核再把 P 份加起来（练习 7）。

3. softmax 的反向 · Softmax backward

3.1 只要 y、只要一次归约

y = softmax(x) 的雅可比矩阵 ∂yᵢ/∂xⱼ = yᵢ(δᵢⱼ − yⱼ) 是 L × L 的，但它是”对角阵减秩一矩阵”，乘一个向量只要一次点积：

$\nabla x_j = \sum_i \nabla y_i\, y_i(\delta_{ij} - y_j) = y_j\Big(\nabla y_j - \sum_i \nabla y_i y_i\Big)$

02 号脚本第 1 部分：这个式子和先构造雅可比再乘都与 PyTorch 差 2.8e-17。两点要紧的推论：

1. 只要前向的输出 y，不要输入 x：前向存 y 就够了。
2. 每行只要一次归约（点积 Σ∇y·y），前向要两次（最大值、和）。所以 HW4 的 ker_attn_softmax_bw 可以一行交给一个线程束，用一次 __shfl_xor_sync 归约搞定，不需要块级栅栏。

3.2 融合，以及和注意力 dropout 一起融合

未融合：P = ∇y ⊙ y、按行求和 s、D = ∇y − s、∇x = y ⊙ D，4 个内核、两个和分数矩阵一样大的中间张量。注意力里 softmax 后面紧跟着 dropout，slides 第 30 页的反向里 ∇S ← ∇Dropout(∇S)、∇S ← ∇Softmax(∇S) 是相邻的两步，可以合进一个内核：先在寄存器里算 g′ = ∇P ⊙ M · keep，再做 softmax 反向，g′ 不写回显存。02 号脚本第 2 部分（64 × 64 的分数矩阵）：

放大到一层注意力（批 8 × 16 头 × 512 × 512，fp16），纸面快 3.1–3.2 倍。HW4 要求的融合 softmax 反向平均快约 5.5 倍。

3.3 为了反向，前向要存什么

公式决定了要存什么：softmax 反向要 y，dropout 反向要掩码。两样都是 B × h × L × L 的。02 号脚本第 3 部分（y 按 fp16、掩码按 1 字节）：

省法有三层：掩码按位存（1/8）；只存随机数种子、反向时重新生成掩码；连 y 都不存、反向时重算——第 21 讲的 FlashAttention 走的是最后一条路。

3.4 模板参数：ITERATIONS

HW4 的反向内核一行一个线程束，每个线程负责 ITERATIONS = ⌈L/32⌉ 个元素，把它们的 ∇y、y 放在寄存器数组里。CUDA 里寄存器数组的大小必须编译时确定、循环要能完全展开，否则数组会被放到慢得多的本地内存里。所以 ITERATIONS 只能是模板参数：为每一档最大长度（HW4 要求 {32, 64, 128, 256, 384, 512, 768, 1024, 2048}）编译一个实例，运行时按实际长度选。L = 2048 时每个线程要 128 个数，占掉每线程寄存器上限（255 个）的一半，再长就该换成一行多个线程束的写法。LightSeq2 在训练开始前把各种模板参数跑一遍、按每种形状挑最快的（论文 §IV-B；slides 第 24–26 页）。

4. 混合精度训练器 · The mixed-precision trainer

4.1 问题：几百个张量，每个都要拷来拷去

混合精度训练（P2 §4）：前向、反向用 FP16，更新用 FP32——更新量 lr × 梯度往往很小，直接加到 FP16 参数上会被舍入掉。slides 第 29 页左图是”原来”的做法，对每一个参数张量：

1. 把 FP16 梯度 g 拷成 FP32 副本 g′；
2. 在 FP32 主参数 p′ 上做 Adam；
3. 把 p′ 拷回 FP16 的 p，给下一步的前向用。

论文指出它的两个缺点：许多”很快就返回”的小内核拉低 GPU 利用率；FP32 的 g′、p′ 占用额外显存。张量到底有多少？03 号脚本第 1 部分：

六成的张量是一千来个元素的偏置和 LayerNorm 的 γ、β——每个只搬几 KB，时间几乎全是发射开销。

4.2 LightSeq2 的做法：链接成一整块 workspace

slides 第 29 页右图、论文 §IV-C 和图 7：

• 符号链接（symbolic tensor linking）：训练器初始化时把所有参数按顺序拷进一块连续的显存（workspace），梯度同样一块；然后让每个参数变成这块显存的一个片段。模型照常用各个参数，训练器只看见一整块。
• 一个内核：每一步只对整块 workspace 发射一次更新内核。
• 即时转换（on-the-fly conversion）：内核从 workspace 读 FP16 的参数和梯度，在寄存器里转成 FP32，更新，再转回 FP16 写回。没有 FP32 的 g′、p′（图里的虚线框表示”不占实际显存”）；Adam 的 m、v 仍是 FP32。

在 NumPy 里，“片段”就是视图：flat[a:b].reshape(shape) 不拷贝数据。03 号脚本第 3 部分验证了每个参数都与整块共享内存，而且逐张量调用 Adam 和对整块调用一次的结果逐位相同——Adam 是逐元素的，切成多少块算都一样。

每个参数每一步的读写（Adam，m、v 为 FP32）：

03 号脚本第 2 部分在模拟器里数出来正好是每个参数 40.0 和 22.0 字节。第 3 部分的纸面账：

论文在 Transformer-big 上报告：训练器的时间减少 54.9%、显存少 2 GB（和 Apex 高度融合的 Fairseq 训练器比）；Adam 内核比 Apex 快 2.3 倍（slides 第 34 页）。

⏱ 本机 CPU 上也看得到”发射开销”（03 号脚本第 3 部分，Transformer-big 的 257 个张量缩小 8 倍）：PyTorch Adam 逐张量约 4.5 ms、foreach=True 约 3.9 ms、fused=True 约 1.7 ms；NumPy 逐张量约 5.4 ms、整块约 3.8 ms。三种 PyTorch 实现的结果逐位相同。PyTorch 的 foreach / fused 是多张量批处理（multi-tensor apply，最早见于 NVIDIA Apex）：不要求参数在内存里连续，而是把许多张量的指针打包进少数几个内核——和 workspace 殊途同归。

4.3 去掉 FP32 主参数的前提

workspace 每一步都把更新后的参数舍入回 FP16。不够半个 FP16 间距的更新量就被”吞掉”了；原来的做法把这些零头留在 FP32 副本里慢慢攒。03 号脚本第 2 部分：第 1 步两种做法逐位相同，之后开始出现最后一位的差别（20 步后最大差 3.7e-04）。第 4 部分算了哪些参数会受影响（Adam 每步的更新量约为 ±lr）：

FP16 在 1 附近的间距是 2⁻¹⁰ ≈ 9.8e-4。小实验：2000 个 γ 从 1 出发、目标 1 + U(−0.05, 0.05)，Adam 1000 步后的平均误差：

学习率 1e-4 时只存 FP16 的 γ 一步也走不动；3e-4 时目标小于 1 的那一半走到了（1 以下间距减半），大于 1 的一半停在原地。

所以”去掉 FP32 副本”是一个有前提的取舍：论文在它测的任务上报告精度不降——常见的学习率和参数量级下，更新量大于半个间距；但学习率衰减得很小、又有量级接近 1 的参数时，FP16 存储会让它们停住。主流做法（Micikevicius et al. 2018、Apex、Megatron、DeepSpeed）都保留 FP32 主参数；第 18 讲的 ZeRO 则把这份副本切到多张卡上来省显存。

5. 显存管理 · Memory management

5.1 自注意力反向的临时张量

slides 第 30 页（论文图 8）把自注意力的反向写成十步，每一步产生一个或两个梯度张量，它们大多只活一两步。用 B、L、H、N 表示批大小、长度、隐藏维度、头数，BLH 是一个 B × L × H 张量的大小，BL²N 是一个分数矩阵的大小。04 号脚本第 1 部分对每一步列出”此刻活着的临时张量”（B = 16、L = 512、H = 1024、N = 16，BL²N = 8 BLH）：

最挤的两步是第 6 步（3BLH + BL²N）和第 7 步（6BLH）。取大的那个：

$\text{临时显存下界} = 3BLH + \max\{BL^2N,\ 3BLH\}$

正是论文的式子（∇out 是上一层传下来的输入，不算临时张量）。

5.2 静态规划：训练前算一次

论文图 8 右边把张量排成”列”：同一列的张量生命周期互不重叠，共用同一块预先分配的显存。04 号脚本的 plan_blocks 用一个贪心算法做这件事（按产生的先后，把每个张量放进一个此刻空闲、放进去后增长最少的块；块的大小是放进去的最大张量），结果是 4 块：

块 0（1 BLH）：∇Y → ∇Y′ → ∇Q → ∇Y″
块 1（8 BLH）：∇Z → ∇S →（dropout、softmax 原地）→ ∇Q̃∇K̃∇Ṽ
块 2（1 BLH）：∇V → ∇in
块 3（1 BLH）：∇K

贪心规划正好达到下界。序列越长，BL²N 越占主导——注意力的反向显存又被 L × L 的分数矩阵支配。

5.3 预先分配：为什么不边用边申请

cudaMalloc 很慢，而且会让 GPU 同步，所以 PyTorch 用缓存分配器：释放的块不还给 GPU，留着给下次用。可是长度不同的批次轮流到来时，空闲块被切得零零碎碎，后来的大请求找不到连续空间，只好再要新段。04 号脚本第 2 部分模拟了一个简化的缓存分配器（6 层编码器，每批约 4096 个 token，长度随机取自 32–256，共 300 批）：

300 批后已占用 1,315 MiB，是任何一批真正需要的最大值（904 MiB）的 1.45 倍，多出来的都是碎片。

LightSeq2 的做法（论文 §IV-D）：把显存分成永久的（参数、梯度，大小固定）和临时的（中间张量）；训练前扫描训练集，估出临时显存的上界，一次性分配好，之后每一批都复用，训练结束才释放。论文观察到 Fairseq 的显存占用随训练慢慢上涨（遇到长句子就要再申请），Transformer-Base 的 GPU 利用率在 80%–93% 之间波动；LightSeq2 的显存不变、利用率约 99%，少用约 6 GB（slides 第 36 页），在 32 GB 的 V100 上能用更大的批训练更深的模型。代价是要提前知道最大形状，短批次时预留的显存闲着。

6. 一整步训练的账，与实验结果 · The whole step, and results

6.1 纸面账

05 号脚本把 GPT-2 large 的一步训练按阶段拆开，基线 vs LightSeq2 式（前向融合 + 反向融合 + workspace 训练器）：

几点观察：

• 反向的加速比前向小：反向的 GEMM 是前向的两倍，非 GEMM 的比例更低（论文图 19 也是”前向的加速比反向大”）。
• 交叉熵：第 10 讲的融合把每个 logit 的读写从 37.9 字节降到 4.2 字节（fp16），词表 5 万时这一段快 1.79 倍。
• 整步纸面 1.39 倍，slides 第 37 页 GPT-2 large 在 A100 上实测 1.6–1.9 倍。纸面账假设基线的每个小内核都跑满带宽，真实的 PyTorch 小内核做不到，所以实测更大。

6.2 批越大，加速越小

训练器的时间只和参数个数有关，与批大小无关；GEMM 的时间随 token 数线性增长。所以批越大，GEMM 越占主导，加速越小（论文图 15）。反过来，用梯度累积、小批次训练时训练器的占比更大，一个内核的训练器最划算。GEMM 只跑到峰值 60% 时，整步加速从 1.39 倍降到 1.31 倍：GEMM 越快，非 GEMM 的优化越重要。这也解释了 slides 第 32 页的观察：“A100 的 GEMM 更高效”，所以 A100 上的加速（1.5–3.5 倍）比 V100（1.4–2.8 倍）大。

6.3 LightSeq2 的实验（slides 第 31–40 页）

• 时间线（slides 第 33 页）：一次翻译训练迭代从 Fairseq 的 457 ms 降到 214 ms，编码器 / 解码器的前向与反向、输出投影、优化各段都变短，优化那一段缩得最多。
• 算子加速（slides 第 34 页）：LayerNorm 4 倍、softmax 2.5–3.4 倍、dropout 1.1–2.5 倍、训练器 2.3 倍。
• 可扩展性（slides 第 35 页）：1–5 台机器（每台 8 × A100），加速 1.12–1.41 倍。机器越多，梯度同步占的时间越多，单卡计算的加速被摊薄（第 14–15 讲）。
• 显存（slides 第 36 页）：Transformer-Base / Big 都少用约 6 GB。
• 与 DeepSpeed 的区别（slides 第 38 页）：DeepSpeed 当时只融合了 Transformer 层，LightSeq2 还覆盖了嵌入、交叉熵和训练器。

slides 第 41 页的总结：高性能的前向 / 反向内核（融合 + 代数改写）、高效的参数更新（混合精度、整块更新）、显存管理（复用，依赖于具体结构）。

7. 小结 · Summary

1. 训练比推理多出三笔：反向（每个融合内核都要有反向，里面同样有归约）、参数更新（几百个张量、精度转换）、显存（存下的激活、反向中的临时张量、变化的形状）。
2. LayerNorm 反向：∇x = r[g − (Σg + x̂Σgx̂)/m]。照依赖做要两次同步；重排后两个和一起归约，一次同步。slides 的写法用 x、HW4 的写法用 x̂，后者误差小 2–3 倍。融合把读写从 18 倍降到约 4 倍。
3. γ、β 梯度是按列求和：一个线程一列太串行，原子加有竞争且不确定，32 × 32 分块 + 转置 + shuffle 是折中；共享数组 [32][33] 消除 bank 冲突。
4. softmax 反向：∇x = y ⊙ (∇y − Σ∇y·y)，只要 y、一次归约；与 dropout 反向融合后读写 12.38 → 4 倍；为反向存下的 y 和掩码随 L² 增长。
5. 训练器：把参数链接成一整块 FP16 workspace，一个内核、寄存器里转换精度；每个参数 40 → 22 字节，几百次发射 → 1 次，省下每个参数 8 字节的 FP32 副本。前提是更新量不能小于半个 FP16 间距。
6. 显存：按生命周期复用，自注意力反向的临时显存下界是 3BLH + max{BL²N, 3BLH}，静态规划能达到；训练前按最大形状一次性分配，避免缓存分配器的碎片和反复 cudaMalloc。
7. 整步训练纸面快 1.39 倍（实测 1.6–1.9 倍）；批越大、GEMM 越快，加速越小。

8. 常见误区 · Common pitfalls

• “融合了就等于减少了同步。” 照依赖关系融合的 LayerNorm 反向，栅栏一个都没少；减少同步要靠代数重排。
• “代数上等价的重排在数值上也等价。” slides 的 α、β 写法和 HW4 写法同样一次同步，误差差 2–3 倍；前向的 E[x²] − E[x]² 差得更多。
• “softmax 的反向要存 softmax 的输入。” 只要输出 y；而且只要一次归约。
• “γ、β 的梯度和 ∇x 一样，一行一个块就行。” 它们是沿着行（token）的方向求和，要另一种并行方式。
• “训练器只占几个百分点，不值得优化。” 它要碰每一个参数、发射上千个内核；小批次、梯度累积时占比更大。
• “一个内核更新全部参数，必须把参数拷到一起再拷回去。” workspace 是一次性的”链接”：参数本身就是整块显存的片段，之后不再拷贝。
• “去掉 FP32 主参数没有任何代价。” 更新量小于半个 FP16 间距时会被吞掉；量级接近 1 的参数、很小的学习率最危险。
• “缓存分配器不还显存，所以不会有问题。” 不还显存恰恰让碎片一直占着；形状变化时已占用会越涨越高。
• “LightSeq2 的整体加速和算子加速一样大。” GEMM 不变，整体加速被 GEMM 的比例稀释；批越大，GEMM 占的比例越高，稀释得越多。

9. 练习 · Exercises

1. 一行 softmax 的输出 y = (0.5, 0.3, 0.2)，∇y = (1, 0, 0)。求 ∇x，并验证 Σ∇x = 0。为什么 softmax 反向的结果总是和为 0？
2. 用 HW4 的公式证明 LayerNorm 反向满足 Σᵢ∇xᵢ = 0。这在直觉上是什么意思？
3. 按 §4.2 的表，核对”原来”每个参数 40 字节、LightSeq2 22 字节的读写（Adam 的 m、v 为 FP32）。
4. 一个 7B 参数的模型，去掉 FP32 的参数和梯度副本能省多少显存？如果还保留 FP32 主参数、只去掉 FP32 梯度副本呢？
5. B = 4、L = 2048、H = 4096、N = 32、fp16：自注意力反向的临时显存下界是多少 MiB？
6. 什么时候 BL²N ≥ 3BLH？用每个头的维度 dₕ = H/N 表示。dₕ = 64 时 L 至少多长？
7. N = 8192、d = 1024 的 γ、β 梯度：设计一个”两段”的归约，让线程数接近 A100 的 221,184 个常驻线程。第一段多少线程、每线程几步？第二段做什么？
8. GPT-2 large 有 436 个参数张量、7.74 亿参数。按”原来”的训练器（每个张量 3 个内核、每个参数 40 字节），发射开销和读写时间各多少？LightSeq2 呢？
9. 一个 LayerNorm 的 γ = 1.5，学习率 3e-4，用 FP16 存储：更新会被吞掉吗？γ = 0.6 呢？在 γ = 0.6 时，实际加上去的量是多少？
10. 缓存分配器里有两段各 100 MB 的空闲块（不相邻），来了一个 150 MB 的请求。会发生什么？LightSeq2 的预先分配怎样避免它？
11. 05 号脚本里，1024 个 token 时训练器占基线的 32%，16384 个 token 时只占 3%。为什么？这对”梯度累积”有什么启示？
12. 注意力 dropout 反向和 softmax 反向可以融合成一个内核。下一步 ∇Q = ∇S·K 是一个 GEMM，能不能把 softmax 反向也融合进这个 GEMM？代价是什么？

参考答案

1. s = Σ∇y·y = 0.5；∇x = y ⊙ (∇y − s) = (0.5 × 0.5, 0.3 × (−0.5), 0.2 × (−0.5)) = (0.25, −0.15, −0.10)，和为 0。原因：softmax 对所有输入同时加一个常数不变，所以沿 (1, 1, …, 1) 方向的导数为 0；代数上 Σⱼ yⱼ(∇yⱼ − s) = s − s·Σy = 0。
2. Σᵢ∇xᵢ = r[Σg − (m·Σg + Σᵢx̂ᵢ·Σgx̂)/m] = r[Σg − Σg − 0] = 0，因为 Σx̂ᵢ = 0。直觉：LayerNorm 对每行同时加一个常数不变（减均值抵消了），所以梯度没有”整体平移”的分量。
3. 原来：拷梯度读 FP16 2 + 写 FP32 4 = 6；Adam 读 p′、g′、m、v 16，写 p′、m、v 12，共 28；拷参数读 4 + 写 2 = 6；合计 40。LightSeq2：读 p、g（FP16）4，读 m、v 8，写 p（FP16）2，写 m、v 8，合计 22。
4. FP32 参数 + 梯度副本每个参数 8 字节：7e9 × 8 = 56 GB。只去掉梯度副本：7e9 × 4 = 28 GB。
5. BLH = 4 × 2048 × 4096 × 2 字节 = 64 MiB；BL²N = 4 × 2048² × 32 × 2 字节 = 1024 MiB = 16 BLH。下界 3 × 64 + max(1024, 192) = 1216 MiB。
6. BL²N ≥ 3BLH ⇔ LN ≥ 3H ⇔ L ≥ 3dₕ。dₕ = 64 时 L ≥ 192：几乎所有实际的训练长度下，分数矩阵那一块都更大。
7. 把 8192 行切成 P 段：第一段 P·d 个线程，每个线程累加 8192/P 行，写出 P × d 个部分和；第二段 d 个线程（或 d/32 个线程束）把每列的 P 个部分和加起来。P = 216 时第一段 221,184 个线程，每个约 38 步；第二段每列只加 216 个数，代价很小。
8. 原来：发射 3 × 436 × 4 μs = 5.2 ms，读写 40 × 7.74e8 ÷ 1.555e12 = 19.9 ms（合计 25.1 ms，与 05 号脚本一致）。LightSeq2：发射 4 μs，读写 22 × 7.74e8 ÷ 1.555e12 = 10.95 ms。
9. γ = 1.5 在 [1, 2) 里，FP16 间距 2⁻¹⁰ ≈ 9.77e-4，一半 4.88e-4 > 3e-4：被吞掉。γ = 0.6 在 [0.5, 1) 里，间距 4.88e-4，一半 2.44e-4 < 3e-4：不会被吞掉，但舍入后实际加上的是一个整间距 4.88e-4（比 3e-4 多 63%）——FP16 存储不只吞掉小更新，还会放大”刚过半个间距”的更新。
10. 两个空闲块总量 200 MB 够，但都不够 150 MB 连续：分配器只好 cudaMalloc 一个新段，已占用涨到至少 350 MB。预先分配时，临时显存的布局在训练前就按最大形状规划好（§5.2），每批都用同样几块，不会切碎。
11. 训练器的时间只取决于参数个数（每步都要更新全部参数），GEMM 和非 GEMM 的时间随 token 数增长。梯度累积时，多个小批次只做一次更新，训练器被摊薄；但每个小批次本身的前向、反向更小、更受发射开销影响，融合内核更重要。
12. 可以，但 GEMM 需要 ∇S 的整块数据：要把 softmax 反向放进 GEMM 的分块循环里，在每个 tile 读进来时现算 y ⊙ (∇P − s)，而行和 s 需要整行的信息（可以提前单独算好，每行一个数）。这正是 FlashAttention 反向的做法（第 21 讲）：不把 ∇S 写回显存，代价是更复杂的内核，并且要重算 y。

10. 延伸阅读 · Further reading

• Wang et al. (2022) LightSeq2: Accelerated Training for Transformer-based Models on GPUs（SC22）：本讲的主要来源，§IV-B（归约改写与模板）、§IV-C（训练器）、§IV-D（显存管理）、§VI（实验）。
• Micikevicius et al. (2018) Mixed Precision Training（ICLR）：FP32 主权重、损失缩放、FP32 累加。
• NVIDIA Apex 的 multi_tensor_apply 与 FusedAdam；PyTorch 的 torch.optim.Adam(foreach=True / fused=True)：多张量批处理。
• PyTorch CUDA 缓存分配器的文档（“CUDA semantics → Memory management”）：缓存、碎片、expandable_segments。
• Chen et al. (2016) Training Deep Nets with Sublinear Memory Cost：用重算换显存（激活检查点）。
• 原课 HW4 说明：ker_attn_softmax_bw、ker_ln_bw_dinp、ker_ln_bw_dgamma_dbetta。
• Dao et al. (2022) FlashAttention：反向时重算注意力，不存 L × L 的矩阵（第 21 讲）。

11. 下一讲预告 · Next lecture

第 10–11 讲的路线是”手写融合内核”：每一个融合都要人去发现、去写 CUDA、去为每种形状调模板参数。能不能让编译器来做？第 12 讲换一条完全不同的技术路线：Google 的 TPU 与 JAX。JAX 把 Python 函数追踪（trace）成一种中间表示（jaxpr），在它上面做函数变换——grad 求导、vmap 自动批处理、jit 交给 XLA 编译器，由编译器自动做算子融合、显存规划；TPU 则用脉动阵列（MXU）做矩阵乘法，配上 HBM / VMEM / SMEM 的存储层次。我们会写一个迷你的追踪器和 jit，逐拍模拟脉动阵列，并看 XLA 式的融合能自动得到本讲手写的多少好处。