大模型预训练入门与 nanoGPT 源码详解

这份教程面向会写基本 Python、希望逐步掌握 PyTorch 与大模型训练的工程师。起步只需要能运行 Python 和 PyTorch；第 1 章可在 CPU 上完成。先把一个小模型训练、验证、保存和加载的闭环走通，再进入 GPT、GPU 优化和多卡训练。

使用说明：单机单卡 → 单机多卡 → 多机多卡

主线按能力和设备规模递进：第 1 章解释一次参数更新，第 2 章构建最小语言模型，第 3 章把完整流程映射到 nanoGPT，第 4 章用两卡 DDP 协同处理数据，第 6 章增加第二台机器。第 5 章是显存分片与可靠恢复分支；第 7 章在理解 nanoGPT 与 DDP 之后，用 Picotron 深入显式通信与调度。每次只增加一个因素，按“问题与预期 → 最小改动 → 运行观察 → 正确性 → 性能”完成闭环。

路径	最终能做什么	交付证据
第 1 章：基础	解释 scalar/vector/matrix/tensor 与一次参数更新；定位形状、类型、设备错误	修改线性关系，核对预测、梯度与参数
第 2—3 章：单机单卡	从 token 追到 loss、更新与 checkpoint；按运行过程定位 nanoGPT	训练、验证、恢复、生成；一组单卡基线
第 4 章：单机两卡 DDP	解释 rank、数据分工与梯度平均；改变有效 batch	固定权重/输入的一次更新对照，独立吞吐记录
第 6 章：多机多卡	将同一 nanoGPT 扩到 $2\,\text{台机器}\times 2\,\text{GPU/台}=4\,\text{GPU}$	全 rank 通信、训练与恢复通过；固定工作量比较
第 5 / 7 章：按问题分支	FSDP2：分片状态与恢复；Picotron：DP/TP/PP/CP 的通信与调度	状态生命周期、数值对照和修改后的可运行实验
第 8—9 章及附录	用问题选择论文、实验与工程方案	注明版本、假设、测量条件与适用边界

学习路线图：主线 1→2→3→4→6。FSDP2 是第 4 章后的状态分片分支；Picotron 在 nanoGPT 与 DDP 基础上展开。多机起步规模为 $2\,\text{台机器}\times 2\,\text{GPU/台}=4\,\text{GPU}$，再扩到 $2\,\text{台机器}\times 8\,\text{GPU/台}=16\,\text{GPU}$。

打开 SVG 原图

主线环境基线沿用 PyTorch 2.9.x、Python 3.11；CUDA 示例以兼容 CUDA 12.9 的环境为参考。第 7 章 Picotron 固定源码版本使用较旧的依赖组合，需要单独环境，详见 7.2。GPU 驱动、PyTorch 构建版本和设备可见性需要匹配，本机不必为了使用预编译 PyTorch 另外安装同版本 CUDA Toolkit。安装命令按 PyTorch 官方安装选择器和内部环境规范确定。性能数字与训练效果应以自己的设备、数据和版本实测为准。

1. 解释一次参数更新：从 Tensor 到 PyTorch 训练

完成本章后，你应能从输入和目标解释预测、scalar loss、梯度与参数更新，修改线性关系并验证结果，定位 shape、dtype、device 错误。先用 $y=2x+1$ 跑通一次训练，再用一个两层 MLP 理解非线性与跨层反向；随后把各张量的 shape、保存时间和字节数对应起来。先在 CPU 验证，再理解迁移 GPU 增加的数据搬运、显存和计时问题。

1.1 起步：确认环境，先让 CPU 例子跑起来

PyTorch 提供两项核心能力：用张量表示和计算数据，以及自动计算梯度。先确认它能导入；没有 CUDA GPU 也可以完成本章。下面的设备查询先判断 CUDA 是否可用，避免在 CPU 环境直接访问第 0 张 GPU 而报错。

import torch
print("PyTorch:", torch.__version__)
print("构建使用的 CUDA:", torch.version.cuda)
print("CUDA 可用:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("设备:", torch.cuda.get_device_name(0))
device = "cuda" if torch.cuda.is_available() else "cpu"
print("本次使用:", device)

环境自检

如果 CUDA 不可用，可能是安装了 CPU 构建、驱动不兼容、容器未暴露 GPU，或 CUDA_VISIBLE_DEVICES 设置不正确。先用 CPU 继续学习，再结合 nvidia-smi 与环境配置排查；不要仅凭 is_available() 为 False 就判定某一种原因。

1.2 张量基础：scalar、vector、matrix 与 tensor

在本教程的数值计算语境中，Tensor 是统一的数组对象，可以有 0、1、2 或更多个轴。标量 scalar 是一个数，以 0 个轴表示；向量 vector 是按一个轴排列的一组分量；矩阵 matrix 按行、列两个轴排列；三阶及以上 Tensor 再增加轴。它们都能用 torch.Tensor 保存，不能把 Tensor 仅理解为“三维以上的数组”。

打开 SVG 原图

1.2.1 从一个数到一批矩阵

在 PyTorch 中，这些对象的类型都可以是 torch.Tensor。scalar、vector、matrix 描述形状特征；tensor 是统一的数据对象名称。这里采用数值计算中的理解，不展开数学中张量的坐标变换定义。

名称	具体数据 / shape	ndim / numel
scalar 标量	$$3.0\ \longrightarrow\ ()$$	0 个轴 / 1 个元素
vector 向量	$$[1,2,3]\ \longrightarrow\ (3,)$$	1 个轴 / 3 个元素
行矩阵 / 列矩阵	$\begin{bmatrix}1&2&3\end{bmatrix}\ \longrightarrow\ (1,3)$；$\begin{bmatrix}1\\2\\3\end{bmatrix}\ \longrightarrow\ (3,1)$	2 个轴 / 都是 3 个元素
matrix 矩阵	$$\begin{bmatrix}1&2&3\\4&5&6\end{bmatrix}\ \longrightarrow\ (2,3)$$	2 个轴 / 6 个元素
三阶张量	两个 $(2,3)$ 矩阵沿新轴叠起 → $(2,2,3)$	3 个轴 / 12 个元素

“维度”必须说明口径：向量空间维数表示分量数，例如三维向量有 3 个分量；PyTorch 的 ndim 表示轴数，所以 $\operatorname{shape}=(3,)$ 的向量 $\mathrm{ndim}=1$。$\operatorname{shape}=(2,3)$ 表示 2 行、3 列，$\mathrm{ndim}=2$、$\mathrm{numel}=6$；轴索引从 0 开始，dim=-1 是最后一轴。$\operatorname{shape}=()$ 是标量，$\operatorname{shape}=(1,)$ 是长度 1 的向量，$\operatorname{shape}=(1,1)$ 是 $1\times1$ 矩阵；三者元素数都为 1，但不是同一形状。矩阵的秩表示线性独立行/列的数量，与 ndim、分布式 rank 均不同，速查见 A.11。

import torch
s = torch.tensor(3.0)
v = torch.tensor([1., 2., 3.])
m = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
t = torch.stack([m, m], dim=0)
for name, x in [("scalar", s), ("vector", v), ("matrix", m), ("tensor", t)]:
    print(name, tuple(x.shape), x.ndim, x.numel())
# scalar () 0 1；vector (3,) 1 3；matrix (2,3) 2 6；tensor (2,2,3) 3 12
print(s.item())                         # 3.0：取成 Python 数值
print(v.unsqueeze(0).shape)             # (1,3)：增加行轴
print(v.unsqueeze(1).shape)             # (3,1)：增加列轴
print(m[1, 2].item())                   # 6.0：第 1 行、第 2 列，索引从 0 开始
print(m[0].shape, m[0:1].shape)          # (3,) 与 (1,3)：整数索引会去掉该轴

运行并观察：同一种 Tensor，不同的轴数

1.2.2 给每个轴赋予含义，再选择类型和设备

shape 本身不包含业务含义。例如 $(32,128)$ 可以是 32 条样本、每条 128 个特征，也可以表示 32 个时间位置、每处 128 个通道。常用约定：图像批次 $(B,C,H,W)$ 表示批量、通道、高、宽；文本 token id 为 $(B,T)$，嵌入后为 $(B,T,C)$。B、T、C 都是轴的名字，不是 PyTorch 自动保存的标签；同一个字母在不同上下文中也可能有不同含义。

dtype 决定元素的表示方式：模型计算通常用浮点数，类别索引和 token id 常用 torch.long，掩码常用 torch.bool。torch.tensor([1,2]) 默认推断为整数；[1.,2.] 通常是默认浮点类型。整数索引不能直接求梯度；常规 autograd 针对浮点或复数张量。float32 每元素 4 字节，float16/bfloat16 为 2 字节；低精度细节见 B.3。

device 指明存储与执行设备；本节先用 CPU。参与同一常规运算的输入与参数应在兼容设备上。x = x.to(device) 要接住返回值；设备与类型匹配时，to 可能直接返回原张量。CPU 张量转换到 CUDA 的过程涉及数据搬运，1.7 会解释这一步。

1.3 第一次训练

先用单个样本理解预测 $\hat y=xw+b$：x 和目标 y 是输入数据，w、b 是需要学习的参数，loss 用一个数衡量误差。下面把 $N=100$ 个样本一起放进矩阵 $X\in\mathbb{R}^{N\times1}$，用矩阵乘 $\mathrm{pred}=Xw+b$ 同时预测。训练先求 scalar loss，再计算它对参数的梯度，最后由更新规则改变参数；梯度与优化器详见 1.5—1.6 和 A.3/A.12。

词语	在 $y\approx w\times x+b$ 中是什么意思
sample / feature / target	样本是一组输入与目标；本例特征 x 是一个数，目标 y 是想让模型预测的值。真实任务的目标来自数据，本例由公式合成。
model / prediction	模型规定计算形式 $w\times x+b$；prediction 是用当前参数算出的 pred。
parameter / hyperparameter	w、b 是训练得到的参数；学习率、训练步数等是事先选择的超参数。
loss / gradient	loss 衡量预测偏差，本例用平均平方误差；gradient 给出当前位置 loss 对参数变化的敏感方向和大小。
learning rate / update	学习率控制一次更新的幅度。梯度下降使用 $w\leftarrow w-\mathrm{lr}\times\mathrm{grad}$；backward 只算梯度，更新另做。

对象	本例 shape / ndim	含义与生命周期
x、y	$(100,1)$ / 2	100 条样本，每条一个特征/目标；输入数据不在本例中求梯度
w、b	$(1,1)$ / 2；$(1,)$ / 1	可训练参数；更新后保留到下一步
pred、$\mathrm{pred}-y$	$(100,1)$ / 2	本次前向预测与误差；用于生成 loss
loss	$()$ / 0	所有 100 个误差平方的均值；$\mathrm{numel}=1$
w.grad、b.grad	$(1,1)$ / 2；$(1,)$ / 1	分别与对应参数同 shape；backward 累积，清梯度时移除或归零

"""CPU 即可：python linear.py"""
import torch

torch.manual_seed(42)
x = torch.linspace(-1, 1, 100).reshape(-1, 1)
y = 2 * x + 1
w = torch.zeros(1, 1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
for step in range(200):
    pred = x @ w + b
    loss = ((pred - y) ** 2).mean()
    loss.backward()
    with torch.no_grad():
        w -= 0.1 * w.grad
        b -= 0.1 * b.grad
    w.grad = None
    b.grad = None
    if step % 50 == 0:
        print(step, round(loss.item(), 6))
print("w, b =", w.item(), b.item())
with torch.no_grad():
    print("x=3 的预测：", (torch.tensor([[3.]]) @ w + b).item())

第一次训练：保存为 linear.py，运行 python linear.py

预期观察：loss 逐步接近 0，w 接近 2，b 接近 1；给出未参与训练的 $x=3$，预测接近 7。这个无噪声例子只用于解释训练机制，不能据此判断真实任务的泛化能力。先跑出这些现象，再读后面的解释。

训练结构与状态流图：输入和目标生成 scalar loss，backward 形成与参数同 shape 的梯度，手工 SGD 更新参数后清梯度。后续换成 optimizer.step 只是接管更新规则；Adam 还会保留跨 step 的历史状态，见 A.3。

打开 SVG 原图

代码	作用
requires_grad=True	追踪参数参与的运算，以便计算梯度
loss.backward()	把梯度累积到参数的 .grad；这一步不会更新参数
with torch.no_grad()	更新参数时不把更新动作继续计入计算图
w -= lr * w.grad	以学习率 lr 调整参数；学习率过大可能让 loss 发散
w.grad = None	丢弃本轮梯度，下一轮重新累积

动手练习：把真实关系换成 $y=-3x+2$，再运行一次。解释学到的参数应接近什么；把学习率调大，观察是否仍能收敛。能解释结果，比背下循环更重要。

1.4 回看张量运算：形状为什么决定训练含义

回看 1.3 的 pred=x@w+b：代码中的 @ 是矩阵乘，$(N,1)\mathbin{@}(1,1)$ 得到 $(N,1)$；加 $b\,(1,)$ 时沿样本轴广播。误差平方使用逐元素运算，mean 把所有误差聚合成 scalar。单特征例子容易掩盖 * 与 @ 的区别，下面用非方阵继续核对。

1.4.1 归约、变形与拼接：哪些轴保留

import torch
m = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
print(m.sum(dim=0))                     # [5,7,9]：两行相加，消去第 0 轴
print(m.mean(dim=1))                    # [2,5]：每行求平均，消去第 1 轴
print(m.mean(dim=1, keepdim=True))       # [[2],[5]]，shape (2,1)
print(m.mean().shape)                   # ()：所有元素求平均，得到标量
print(m.reshape(-1))                    # [1,2,3,4,5,6]，由元素总数推断长度
print(m.T)                             # [[1,4],[2,5],[3,6]]，交换两个轴
v = torch.tensor([1., 2., 3.])
print(torch.cat([v, v], dim=0).shape)    # (6,)：沿已有轴拼接
print(torch.stack([v, v], dim=0).shape)  # (2,3)：增加新轴再堆叠

先预测输出，再运行核对

dim=0 指第 0 轴，dim=-1 指最后一轴。sum/mean 是把该轴上的数聚合起来，默认消去该轴；keepdim=True 保留长度为 1 的轴，便于广播。reshape 重新组织形状，元素总数必须不变；它不能代替 transpose/permute 交换轴。unsqueeze 增加长度为 1 的轴，squeeze 去掉这样的轴；小批量可能 $B=1$ 时，优先指定 squeeze(dim)，避免误删 batch 轴。存储是否复制见 B.1。

1.4.2 广播、点积与矩阵乘：计算的含义不同

广播用于逐元素运算。把两个 shape 从右侧对齐：每个位置的长度要相同，或其中一个是 1；缺少的前导维按 1 理解。例如 $(3,1)+(1,4)$ 得到 $(3,4)$，而 $(2,3)+(2,4)$ 不兼容。

广播示意图：行、列按坐标取数相加；虚拟扩展不复制输入，结果矩阵需要存储。

打开 SVG 原图

col = torch.tensor([[10], [20], [30]])
row = torch.tensor([[0, 1, 2, 3]])
result = col + row
print(result)
print(result[1, 2])  # 20 + 2 = 22
# 常见陷阱：(N,1) 与 (N,) 相减，会广播成 (N,N)。
pred = torch.zeros(4, 1)
target = torch.arange(4.)
print((pred - target).shape)              # (4,4)
print((pred - target.unsqueeze(1)).shape) # (4,1)

广播与形状陷阱

矩阵乘 $XW$ 的规则是 $(m,k)\mathbin{@}(k,n)\longrightarrow(m,n)$：左矩阵第 i 行与右矩阵第 j 列做点积，沿共同长度 k 求和，得到 scalar 输出 $Z_{i,j}$。逐元素乘 $X\odot X$ 则保留 X 的 shape，不做该求和。图中 $B=2$ 是样本数，$C_{\mathrm{in}}=3$ 是输入特征数，$C_{\mathrm{out}}=2$ 是输出特征数；shape 的每项单位都是元素个数。

矩阵乘与线性层图：$X\,(2,3)\mathbin{@}W\,(3,2)\longrightarrow Z\,(2,2)$，再加 $b\,(2,)$ 得到 $Y\,(2,2)$。高亮的一行与一列算出一个 scalar；下方对照 nn.Linear 的实际权重布局。

打开 SVG 原图

import torch
X = torch.tensor([[1., 2., 3.], [4., 5., 6.]])  # (B,C_in)=(2,3)
W = torch.tensor([[1., 0.], [0., 1.], [1., 1.]]) # (C_in,C_out)=(3,2)
b = torch.tensor([0.5, -0.5])                    # (C_out,)=(2,)
Z = X @ W                                      # (2,2)
Y = Z + b                                      # (2,2)，b 沿样本轴广播
torch.testing.assert_close(Z, torch.tensor([[4.,5.],[10.,11.]]))
torch.testing.assert_close(Y, torch.tensor([[4.5,4.5],[10.5,10.5]]))
layer = torch.nn.Linear(3, 2)
with torch.no_grad():
    layer.weight.copy_(W.T)                     # weight.shape=(2,3)
    layer.bias.copy_(b)                         # bias.shape=(2,)
torch.testing.assert_close(layer(X), Y)
print('X/W/Z/Y:', X.shape, W.shape, Z.shape, Y.shape)
print('逐元素平方:', (X * X).shape)              # (2,3)
v = torch.tensor([1., 2., 3.])                  # (3,)
u = torch.tensor([4., 5., 6.])                  # (3,)
print('点积:', (v @ u).item(), (v @ u).shape)    # 32, ()
print('外积:', torch.outer(v, u).shape)          # (3,3)

非方阵计算与 nn.Linear 对照：CPU 可运行

nn.Linear(C_in,C_out) 的 weight shape 是 $(C_{\mathrm{out}},C_{\mathrm{in}})$，bias 是 $(C_{\mathrm{out}},)$，所以计算为 $X\,\mathrm{weight}^{\mathsf T}+b$。图中数学矩阵 W 的 $\operatorname{shape}=(C_{\mathrm{in}},C_{\mathrm{out}})$，$\mathrm{weight}=W^{\mathsf T}$；两种写法表示同一映射。输入若是 $(B,T,C_{\mathrm{in}})$，线性层在最后一个轴上独立计算，输出 $(B,T,C_{\mathrm{out}})$，前面的 batch 与位置轴保留。不要根据“行/列并行”名称猜切分轴，应核对实际 weight 布局；完整形状链见 A.14。

两个向量 v、u 的 shape 都为 $(3,)$ 时，$v\cdot u$ 得到 scalar $()$，$v\odot u$ 得到逐元素乘积 $(3,)$，outer(v,u) 得到矩阵 $(3,3)$。其中 v[0] 是 scalar $()$，v[0:1] 保留一个轴，$\operatorname{shape}=(1,)$；矩阵 X[0] 取一行得到向量 $(3,)$，X[0:1] 则保持行矩阵 $(1,3)$。

本节自测：能否解释 $()$、$(3,)$、$(1,3)$ 的区别；写出 $(B,T,C)$ 沿 dim=1 求平均后的 shape；说明 cat 与 stack、reshape 与 transpose 的区别；在运行前判断一次广播或矩阵乘是否合法。

1.5 自动求导：梯度从哪里来，何时累积和释放

梯度描述 scalar loss 对参数的局部敏感程度。参数 θ 有多个分量时，每个分量对应一个偏导数；把这些偏导按原形状排列，就是 $\frac{\partial L}{\partial\theta}$，因此常规稠密训练中梯度与参数同 shape。正梯度表示该分量略增会使 loss 局部增大，梯度下降沿相反方向更新；学习率过大仍可能越过合适区域。下面 $f=a\times b+a$ 只是求导例子，a、b、f 及它们的标量导数 shape 都为 $()$；与训练损失相连的手算见 A.12。

自动求导示意图：上层计算 $f=-4$，下层分别回传两条路径对 a 的贡献，得到 $\frac{\partial f}{\partial a}=-2$、$\frac{\partial f}{\partial b}=2$。

打开 SVG 原图

import torch
a = torch.tensor(2., requires_grad=True)
b = torch.tensor(-3., requires_grad=True)
f = a * b + a
f.backward()
print(f.item(), a.grad.item(), b.grad.item()) # -4, -2, 2
# a 被两条路径使用，所以 df/da = b + 1 = -2；df/db = a = 2。
# 下一次建立新图并 backward，会继续累加到叶子张量的 .grad。
a.grad = None
b.grad = None

对照图中的数值

同一张图中，多条路径对同一参数的梯度贡献需要相加；多次 backward 则会继续累积到叶子参数的 .grad。这是两个层次。清梯度可放在一组 micro-batch 的反向之前，或上一次 optimizer.step 之后；本教程示例两种位置都有，但都保证相邻参数更新之间只累积计划中的数据。梯度累积到 A 批才更新时，中间不能 zero_grad；batch 术语和归一化见 A.13。

计算图由这一次实际执行的运算构成，每轮 forward 会建立新图。直接创建且 requires_grad=True 的参数通常是叶子张量；默认把梯度累积到这些叶子的 .grad，中间结果需要 retain_grad() 才保留自己的 .grad。loss 常先聚合成标量再 backward；对非标量直接反传，需要指定上游梯度。

例如线性层 $\mathrm{weight}\in\mathbb{R}^{2\times3}$、$\mathrm{bias}\in\mathbb{R}^{2}$，scalar loss 反向后 weight.grad 为 $(2,3)$、bias.grad 为 $(2,)$；输入 $X\,(4,3)$ 只有在启用梯度追踪且需要保留结果时才观察对应梯度。对非标量输出 $Y\in\mathbb{R}^{4\times2}$，backward 必须接收同 shape 的上游梯度 $G\,(4,2)$，求的是它与 Y 的导数的收缩，即向量—雅可比积；不会默认给出完整雅可比矩阵。PP 的阶段边界梯度正需要这种含义。

操作	作用和边界
zero_grad / grad=None	清掉之前累积的梯度；不重建或保留计算图。多轮使用新图反传也会累积 .grad。
backward()	默认反传后释放为反向保存的中间量；不要把 retain_graph=True 当作梯度累加的开关。需要再次反传同一张图时才考虑保留它。
torch.no_grad()	让上下文内的普通运算不被 autograd 记录；适合手动更新、验证等。不会改变原参数的 requires_grad 标志。
x.detach()	返回与原图断开的张量，通常共享存储；原地修改可能影响原张量。
x.clone()	复制数据；在梯度开启且输入需梯度时，仍保留求导关系。要独立且断图的副本可用 x.detach().clone()。

记录日志时，可存 loss.item() 或仅保存必要的 detach 后数据；不要无意把每轮带图的 loss、输出一直放进列表。no_grad 与 model.eval() 用途不同，2.2 会用完整验证流程展示。

1.6 用 Module 和 optimizer 管理同一训练过程

优化器 optimizer 管理“怎样根据当前梯度和历史状态更新参数”。它持有 nn.Module 中 Parameter 的引用；loss.backward() 计算并累积 .grad，optimizer.step() 执行更新，zero_grad() 清理梯度。下面无动量、无权重衰减的 SGD 对应 $\theta\leftarrow\theta-\mathrm{lr}\times\mathrm{grad}$；Adam 额外维护梯度与梯度平方的指数移动平均，AdamW 还使用解耦权重衰减。定义、公式、状态 shape 和一个手算例子统一见 A.3。

import torch
from torch import nn
torch.manual_seed(42)
x = torch.linspace(-1, 1, 100).reshape(-1, 1)
y = 2 * x + 1
model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.MSELoss()
for step in range(200):
    optimizer.zero_grad(set_to_none=True)
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()
    optimizer.step()
print(model.weight.item(), model.bias.item())

同一个线性回归，换成标准 API

在 nn.Linear(1,1) 中，weight 与 weight.grad 都是 $(1,1)$，bias 与 bias.grad 都是 $(1,)$；loss 是 $()$。这段 SGD 无额外动量张量。若改用 Adam，通常在参数首次获得梯度并更新时建立 exp_avg 和 exp_avg_sq，它们与对应参数同 shape，跨更新保留；zero_grad 不清空这两份历史，恢复训练也必须恢复 optimizer.state_dict。

1.6.1 从现成层到自己的 Module

import torch
from torch import nn

class TinyRegressor(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)  # 子模块及其参数自动注册
        self.offset = nn.Parameter(torch.zeros(1))
        self.register_buffer("input_mean", torch.zeros(1))
        self.description = "回归示例"  # 普通属性

    def forward(self, x):
        return self.linear(x - self.input_mean) + self.offset

model = TinyRegressor()
print([name for name, _ in model.named_parameters()])
# offset、linear.weight、linear.bias（顺序以打印为准）
print(list(model.state_dict()))        # 还包含 buffer：input_mean
print(model(torch.ones(4, 1)).shape)   # (4,1)

自定义模型：forward 定义计算，Module 管理状态

nn.Parameter 是登记为模型参数的张量；普通 Tensor 属性不会因此自动成为参数。buffer 用于应随模型搬设备、通常应保存，但不交给优化器更新的张量，例如统计量；默认持久 buffer 会进入 state_dict，persistent=False 的 buffer 不会。state_dict 主要保存参数与持久 buffer，不是完整的训练现场；继续训练还需保存优化器等状态。

把多个层放进 nn.Sequential 会注册它们并定义顺序前向；nn.ModuleList 负责注册层列表，但需要自己写 forward。普通 Python list 中的层不会自动注册到父模块，可能被 parameters()、to() 和保存过程漏掉。训练时调用 model(x)，让 Module 的调用机制处理 hooks；先把模型移到目标设备，再构造优化器。

1.6.2 从线性回归到两层 MLP：跟完一次非线性变换

线性回归已经说明一次更新怎样发生。现在只增加一个因素：在两层线性变换之间加入非线性激活。MLP（多层感知机）把一批输入特征变成一批预测；每一层的参数由同一个 loss 训练。若去掉中间激活，$(XW_1+b_1)W_2+b_2$ 可以合并为 $X(W_1W_2)+(b_1W_2+b_2)$，仍是一次仿射变换，单纯多堆一层不会得到新的非线性表达能力。

沿用 1.4 的 $X=\begin{bmatrix}1&2&3\\4&5&6\end{bmatrix}$ 和 $W_1=\begin{bmatrix}1&0\\0&1\\1&1\end{bmatrix}$，只将第一层偏置改为 $b_1=[-5,-6]$，再接 ReLU 与第二层 $W_2=\begin{bmatrix}1\\-0.5\end{bmatrix}$、$b_2=[0.5]$。这里 $B=2$ 个样本、$C_{\mathrm{in}}=3$ 个输入特征、$C_{\mathrm{hidden}}=2$ 个隐藏特征、$C_{\mathrm{out}}=1$ 个输出特征。$\operatorname{ReLU}(z)=\max(z,0)$ 逐元素作用，shape 保持不变；本例没有恰为 0 的输入，避开不可导点的约定。

数据或操作	shape / 参数布局	本例结果
输入 X 与目标 Y	$X\,(2,3)$，$Y\,(2,1)$	$$Y=\begin{bmatrix}1\\4\end{bmatrix}$$
$$Z_1=XW_1+b_1$$	$$(2,3)\mathbin{@}(3,2)+(2,)\longrightarrow(2,2)$$	$$Z_1=\begin{bmatrix}-1&-1\\5&5\end{bmatrix}$$
$$H=\operatorname{ReLU}(Z_1)$$	$$(2,2)\longrightarrow(2,2)$$	$$H=\begin{bmatrix}0&0\\5&5\end{bmatrix}$$
$$\mathrm{pred}=HW_2+b_2$$	$$(2,2)\mathbin{@}(2,1)+(1,)\longrightarrow(2,1)$$	$$\mathrm{pred}=\begin{bmatrix}0.5\\3\end{bmatrix}$$
$$L=\operatorname{mean}\!\left((\mathrm{pred}-Y)^2\right)$$	$(2,1)$ 归约为 scalar $()$	$$L=0.625$$
对应 nn.Linear(3,2) → ReLU → nn.Linear(2,1)	$\mathrm{weight}_1\,(2,3)$，$\mathrm{bias}_1\,(2,)$；$\mathrm{weight}_2\,(1,2)$，$\mathrm{bias}_2\,(1,)$	11 个参数元素；PyTorch weight 为数学 W 的转置

反向也沿着这条链走。设 $dP=\frac{\partial L}{\partial\mathrm{pred}}=\frac{2}{B}(\mathrm{pred}-Y)$，$\operatorname{shape}=(2,1)$；第二层得到 $dW_2=H^{\mathsf T}dP$，$\operatorname{shape}=(2,1)$，$db_2=\operatorname{sum}_{0}(dP)$，$\operatorname{shape}=(1,)$。再将 $dH=dP\,W_2^{\mathsf T}$，$\operatorname{shape}=(2,2)$，传回 ReLU；$dZ_1=dH\odot\mathbf{1}_{Z_1>0}$ 仍为 $(2,2)$，第一层得到 $dW_1=X^{\mathsf T}dZ_1$，$\operatorname{shape}=(3,2)$，$db_1=\operatorname{sum}_{0}(dZ_1)$，$\operatorname{shape}=(2,)$。这里 $\odot$ 是逐元素乘；转换成 nn.Linear 的 weight.grad 时，采用转置后的参数布局。

第一条样本的隐藏值被 ReLU 置零，所以它在这一点不给第一层参数贡献梯度；这不等于整个网络没有梯度，第二层 bias 仍收到它的贡献。下面的图同时标出前向结果与反向需要的依赖；参数要等 backward 完成后再由 optimizer.step 更新。

打开 SVG 原图

运行 C.7 的 mlp_bridge_check.py，先核对 $\mathrm{loss}=0.625$，再核对解析梯度与 autograd 一致。按数学布局，$dW_1=\begin{bmatrix}-4&2\\-5&2.5\\-6&3\end{bmatrix}$、$db_1=[-1,0.5]$、$dW_2=\begin{bmatrix}-5\\-5\end{bmatrix}$、$db_2=[-1.5]$。脚本用 float64 减少核对误差，并检查学习率 0.001 的一次 SGD 更新使本例 loss 降低；这只验证一次更新，不代表模型已训练好或对新数据有效。

接下来沿用同一 MLP：1.6.3 先解释 PyTorch 各部分怎样协作，1.7 再看 GPU 架构、显存与时间。第 2.3 节把同样的矩阵推导用于 GPT，第 5—7 章继续追踪并行化后的数据、状态和通信。

1.6.3 PyTorch 运行架构：从训练代码到设备执行

前面已经算清两层 MLP 的结果，现在沿同一次更新看框架的分工。读完本节，应能解释：模型对象怎样管理参数，一次前向怎样留下求导依赖，一项张量运算怎样在 CPU 或 GPU 上执行。后面遇到“参数没更新”“显存持续上涨”“GPU 等待 CPU”等现象，就能沿对应的状态和执行环节定位。

仍用 Linear(3,2) → ReLU → Linear(2,1)：输入 $X$ 的 shape 为 $(2,3)$，第一层输出与隐藏激活为 $(2,2)$，预测和目标为 $(2,1)$，均方损失 $L$ 为标量，shape 为 $()$。第一层 PyTorch 权重的 shape 是 $(2,3)$，前向计算 $Z_1=XW_{1,\mathrm{pt}}^{\mathsf T}+b_1$；这里 $W_{1,\mathrm{pt}}$ 表示 PyTorch 存放的权重，与 1.6.2 数学推导中的 $W_1$ 互为转置。

下图上半部说明计算经过哪些执行设施，下半部说明一次训练怎样改变状态。Autograd 在前向进行时记录依赖；反向计算和优化器更新也会调用张量算子。它们不是三套彼此隔离的 GPU 引擎。

打开 SVG 原图

模型层：Module 组织计算，Tensor 承载数据。model(x) 通过 Module 的调用机制执行 forward；nn.Linear 管理注册的 weight 和 bias，因此 model.parameters() 能把它们交给优化器。Module 的子模块树描述模型的组织，不等于当次自动求导图。普通稠密 Tensor 还记录 shape、stride、dtype 和 device，用这些信息解释底层 Storage 中的数据；转置视图可共享原数据，未必发生复制。布局的具体含义见 B.1。

算子层：选择适用实现，再完成数值计算。nn.Linear.forward 调用 F.linear。PyTorch 的算子分派机制（Dispatcher）根据设备、布局及当前求导、自动混合精度等状态选择实现；ATen 是其中的核心张量算子库。后端再根据 dtype、shape 和布局选择具体计算路径。以核对过的 PyTorch v2.9.0 实现为例，二维输入且有偏置的 Linear 会走 addmm，另一些条件会走其他路径；不能把某次 profiler 中的名字当成所有版本的固定调用链。

设备层：CPU 与 CUDA 后端执行相同的运算语义。CPU Tensor 的数值计算由 CPU 实现完成；输入与参数在 CUDA 上时，适用的 GPU 实现可以调用 cuBLAS 等数学库，或运行原生、编译生成的 kernel。CUDA 运行时与驱动负责相关设备操作和工作提交，GPU 执行这些工作。cuDNN 常用于其支持的神经网络算子，并非每次张量运算都要经过它。Python API、框架算子和 GPU kernel 不一一对应：一次调用可展开成多项工作，多项运算可以融合，某些视图操作只改变元信息。

求导层：边做前向，边记录反向需要的依赖。默认 eager 模式下，开启梯度记录且涉及需梯度输入的可微运算会形成当次 Autograd 图，并保存导数需要的输入或输出。loss.backward() 从标量损失出发，按依赖执行反向计算，把参数梯度累积到 p.grad。第一层权重梯度的 shape 仍为 $(2,3)$，与该参数相同。训练参数通常是叶子 Tensor，因此 p.grad_fn is None 很正常；中间结果的 grad_fn 指向反向节点，它不是梯度数值。

更新层：优化器读取梯度，改变参数与持久状态。本例的前向和 backward() 都不改变参数值；optimizer.step() 才依据梯度和优化器状态更新参数。AdamW 的矩估计跨更新保留；下一轮需要清理的是累计梯度，使用 zero_grad(set_to_none=True)，而不是重新初始化模型。默认反向结束后，用完的求导保存值可以释放；外部仍持有的 Tensor 引用和缓存分配器会影响实际内存表现。数学定义见 A.3、A.12，显存账本在 1.7.2 继续展开。

编译是可选的执行优化。torch.compile(model) 尝试捕获并优化模型调用中的运算片段，可能减少 Python 开销、融合算子和中间读写；训练仍保留前向、反向与更新的职责，反向仍由 Autograd 引擎调度。只编译 model 不代表外层 optimizer.step() 自动纳入编译。先理解和验证 eager 路径，再在 2.4 比较编译后的正确性与预热后的性能；编译不保证每个场景都更快。

动手观察。C.11 的独立脚本复用本节 MLP，依次记录前向、反向、更新和清梯度四个区段。先在 CPU 上核对参数是否变化、梯度与 AdamW 状态的 shape，再在具备 CUDA 的环境中增加 GPU trace。这里观察执行机制；1.7.4 与 2.4 再讨论怎样测量耗时，避免把一次带检查和 profiler 的运行当成性能结果。

本节核对依据：PyTorch 2.9 Module、Autograd mechanics、CUDA semantics、torch.compiler，以及 v2.9.0 ATen Linear 实现。核心术语集中在 A.10；这些执行职责会在后续 DDP、FSDP、TP 和 PP 中继续出现。

1.7 GPU 架构与训练：计算在哪里，数据怎样搬运

本节以 NVIDIA GPU 与 CUDA 为背景，接着 1.6.3 的 PyTorch 执行链路解释硬件。先分清执行计算的 SM、组织工作的线程层次、保存数据的存储层次，再用同一个 MLP 建立显存账本，最后测量设备时间。目标是能定位“放不下、算得慢、搬得慢、等得久”分别与什么有关；第 4～5 章再研究 GPU 之间的协作，第 6 章扩展到多机。

1.7.1 GPU 的计算与存储：SM 怎样执行 kernel

CPU 擅长控制流程、通用程序和数据准备；GPU 通过大量并行线程提高吞吐，适合足够规模的矩阵运算等负载。常见独立 GPU 有自己的设备内存；CPU 内存里的 Tensor 不会因为机器装了 GPU 就自动在 GPU 上计算，需要把输入和模型放到相应设备。

1.6.3 已说明 PyTorch 如何选择 CPU 或 CUDA 实现。进入 GPU 后，kernel 是设备执行的函数；它的实现规定每个线程处理哪些数据，一次启动用 grid 组织全部 block，block 内包含多个 thread。硬件按 block 内每组 32 个线程组成的 warp 来组织指令执行。这里说的是计算的组织方式，Tensor 的 shape 仍描述数据；两者由 kernel 算法建立映射。

图中先沿 CPU 提交工作到 GPU 执行看主线，再看放大的 SM 和线程层次。SM（Streaming Multiprocessor，流式多处理器）是 GPU 的硬件执行单元，包含 warp 调度器、算术单元、寄存器、共享内存等资源。普通 CUDA block 在一个 SM 上执行，一个 SM 可以同时驻留多个 block；grid 的 block 可以分批调度到多个 SM，因而 block 数不必等于 SM 数。

打开 SVG 原图

把架构接回矩阵乘。MLP 的权重、输入和激活通常保存在设备显存中；矩阵乘实现把工作划为 tile（数据块），让线程协作计算，并通过寄存器、shared memory 与缓存复用数据。CUDA Core 与 Tensor Core 是 SM 内不同类型的计算单元；Tensor Core 加速其支持的矩阵运算，但创建 torch.Tensor 并不保证使用它。每个线程处理多少元素、采用多大的 tile，由实现决定，不能从输入 shape 直接推断。

因此，增大 batch 可能增加可并行工作和数据复用，也会提高激活显存；对这里极小的 MLP，GPU 启动和调度开销还可能超过计算收益。优化要分别观察显存容量、数据带宽和计算吞吐，不能只数 CUDA Core 或线程。SM、warp、block、grid 的完整定义见 A.9；下面先用训练中实际保留的状态说明显存消耗。

1.7.2 参数、梯度和激活为什么占显存

训练持续保存模型参数和优化器状态；前向为反向保留必要激活，backward 产生梯度并释放不再需要的中间值。GPU 上的这些张量通常占用设备显存。扩大 batch 或序列长度主要增加激活等工作量，扩大模型还会增加参数、梯度及优化器状态；这解释了为什么单看参数文件大小不足以预测训练能否放下。缓存、shared memory 和寄存器的作用见 A.9。

显存容量决定“放不放得下”，显存带宽决定“单位时间能搬多少数据”，计算吞吐决定“单位时间能执行多少运算”。三者要分开看。一个连续稠密 float32 张量有 $1024\times 1024$ 个元素，逻辑数据量为 4 MiB；训练总显存还包含梯度、激活、优化器状态、工作区与分配器保留空间，不能只拿参数字节数估算。

矩阵乘通过分块和复用数据提高每次读取带来的计算量；逐元素加法只做少量计算，较容易受内存带宽影响。FlashAttention 的出发点之一是减少中间注意力矩阵在 HBM 的读写，但具体瓶颈仍取决于形状、精度和实现，见 8.3。

沿 1.6.2 的两层 MLP 建立账本：先假设所有训练张量都为 FP32，每元素 4 字节；11 个参数占 44 字节，完整梯度占 44 字节，Adam 的 m、v 各占 44 字节。参数、梯度、两份状态合计 176 字节，只是这四类逻辑数据的合计；$X\,(2,3)$ 另占 24 字节、$Y\,(2,1)$ 另占 8 字节，还没计入激活、临时缓冲和分配器。C.7 为数值核对使用 float64，上述每项字节数在那段代码中会翻倍。

一次更新中的时刻	哪些状态同时存在	保留或释放的原因
稳态起点：上一轮已清梯度	参数 44 B；Adam m/v 共 88 B；grad 为 None	共 132 B 状态；首次 step 前 m/v 通常尚未建立
前向完成	上述状态，加 X、隐藏激活 $H\,(2,2)$、预测等；H 的逻辑大小为 16 B	dW₂ 需要 H，dW₁ 需要 X；激活反向还需要输出或掩码等信息
反向进行中	逐步形成总计 44 B 的参数梯度；同时还有尚未用完的保存值	每层用完依赖后可释放其保存引用，但 Python 变量等外部引用可能继续持有张量
optimizer.step 之后	参数和 m/v 已改变，参数梯度仍可能存在	step 不负责清梯度；zero_grad(set_to_none=True) 后回到稳态起点

峰值是“某一时刻仍存活的独立分配总量”的最大值，不能把整个训练曾出现过的张量全部相加。视图或多个保存引用可能共用同一 storage；autograd 具体保存哪些输入、输出或掩码取决于算子实现。上表说明依赖与生命周期，不声称是 PyTorch 实际分配峰值；对如此小的模型，分配粒度和运行时开销尤其明显。释放张量也不保证 memory_reserved 立刻下降。

混合精度时逐项计算 $P\times\left(\text{参数字节数}+\text{梯度字节数}+m\,\text{字节数}+v\,\text{字节数}\right)$，P 是参数元素数；仅当实现确实维护独立的 FP32 主权重时再加这一份，不能把同一份 FP32 参数重复记账。激活重计算减少部分前向保存值，以额外计算换空间；FSDP 主要分担训练状态。两者解决账本中的不同项目，分别在 2.4 和第 5 章验证。

1.7.3 异步执行：为什么计时要等待

GPU 工作通常对 CPU 异步：Python 调用返回时，GPU 可能仍在计算。stream 可以理解为设备操作队列；同一 stream 内按顺序执行，不同 stream 只有在依赖与硬件资源允许时才可能并发。不能假定“多开 stream 一定更快”，也不能忽略跨 stream 的数据依赖。

直接用 CPU 时钟包住一次矩阵乘，可能主要测到提交工作的时间。教学计时可在开始和结束处同步，或用同一 stream 的 CUDA Event；同步应放在明确的测量边界，日常训练中频繁同步会破坏重叠。CUDA 张量的 .item()、取回 CPU 结果等操作也可能需要等待，因此逐步打印 loss 会干扰吞吐测量。

pin_memory 配合适当的异步传输有助于数据搬运，但 non_blocking=True 本身不保证传输与计算重叠；还需要合适的 stream、依赖与硬件支持。先建立单卡基线，再按 profiling 结果优化。

1.7.4 小实验：核对字节与测量范围

import torch

if not torch.cuda.is_available():
    print("当前没有可用 CUDA GPU：可先在 CPU 完成前面的基础练习。")
else:
    device = torch.device("cuda:0")
    props = torch.cuda.get_device_properties(device)
    print("设备：", props.name)
    print("SM 数量：", props.multi_processor_count)
    print("显存 GiB：", round(props.total_memory / 1024**3, 2))
    print("Compute capability：", torch.cuda.get_device_capability(device))
    print("PyTorch 构建所用 CUDA 版本：", torch.version.cuda)

    x = torch.randn(1024, 1024, device=device, dtype=torch.float32)
    w = torch.randn_like(x)
    print("x 逻辑大小 MiB：", x.numel() * x.element_size() / 1024**2)  # 4.0
    with torch.inference_mode():
        for _ in range(10):
            y = x @ w                   # 预热；分配和库初始化可能影响首次耗时
        torch.cuda.synchronize(device)
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        repeats = 50
        with torch.cuda.device(device):
            start.record()
            for _ in range(repeats):
                y = x @ w
            end.record()
        end.synchronize()
    print("每次矩阵乘平均毫秒：", start.elapsed_time(end) / repeats)
    print("已分配 / 已保留 MiB：",
          torch.cuda.memory_allocated(device) / 1024**2,
          torch.cuda.memory_reserved(device) / 1024**2)

有 CUDA 时运行；无 CUDA 时安全跳过

这段代码测量当前 stream 上重复矩阵乘的区间，不包含前面的数据生成和传入 GPU；区间可能包含提交间隙，并非严格隔离某个内核。memory_allocated 是 PyTorch 张量等当前占用的分配量，memory_reserved 是缓存分配器管理的保留量，两者都不等于整卡所有进程的总占用。Compute capability 表示硬件能力，不是 CUDA Toolkit 版本，也不是 PyTorch 版本。

不要用本实验的某个固定毫秒数作为通关标准。先固定形状、dtype、设备和版本，再改变一项条件比较；更完整的端到端吞吐、数据加载与优化实验见 2.4。occupancy 是活跃驻留 warp 数与硬件上限之比，既不是监控面板上的 GPU 利用率，也不是性能越高越好的保证。

1.7.5 带着设备、显存和时间问题进入单卡训练

单卡阶段先回答：数据是否已在 GPU 上；显存主要被什么占用；时间花在加载、搬运、计算还是等待。单机多卡阶段再加入 PCIe/NVLink 等互联和梯度通信：DDP 通常让每卡保留一份模型，FSDP 分片部分训练状态。多机多卡阶段再增加节点之间的网络与进程协作。多放几张卡不代表程序会自动把显存拼成一块，必须由并行策略决定如何分工。

本节自测：沿一个 Linear 解释 Module、算子、CUDA 实现与 GPU 执行的关系；说明 backward 与 optimizer.step 分别改变什么；区分 Tensor、Tensor Core、CUDA thread 与 DDP rank；解释 grid、block、warp 与 SM 的对应关系；说明显存容量和带宽的差别，以及 GPU 计时为什么必须考虑异步执行。完成 1.8 的形状和状态检查后进入第 2 章。

参考：用户提供的 CUDA C++ Programming Guide，Release 13.0，2025-09-02。按需阅读 §3.1（GPU 的收益，印刷页 5）、§5.1–5.4（kernel、线程层级、内存与异构编程，页 11–18）、§5.6（compute capability，页 20–21）、§6.2.8.5（streams，页 50）、§7.1–7.2（SIMT 与硬件多线程，页 141–143）、§8.2–8.3（并行与访存，页 146 起）。该 PDF 的主文档物理页码比印刷页码大 16。本节引用原理，不要求将现有 PyTorch 环境升级到 CUDA 13.0。

官方在线版本：CUDA C++ Programming Guide 13.0。希望继续写 CUDA 时，读附录 A.9；PyTorch 设备搬运的补充见 B.2。

1.8 本章验收：解释一次更新，定位形状与状态错误

现象	先检查什么	常见处理
Expected all tensors to be on the same device	输入和参数的 .device	将模型和数据放到同一设备
mat1 and mat2 shapes cannot be multiplied	矩阵乘的内维是否相等	先解释各轴含义，再修正布局，勿盲目 reshape
交叉熵标签类型/形状不对	分类索引标签应是 long，shape 通常为 $(B,)$	本教程用整数类别；概率标签是另一种受支持用法
CUDA out of memory	batch、序列长度、模型、保留的计算图	先减小工作量；排查是否保存带图的 loss/tensor
loss 不下降或 NaN	数据标签、学习率、梯度、数值范围	先在小数据上验证；打印 shape 和有限值

验收练习：将真实关系改为 $y=-3x+2$，记录更新前后 w、b 与 loss；说明为何期望 w 接近 −3、b 接近 2。故意把 target 从 $(N,1)$ 改成 $(N,)$，在计算 loss 前用 shape 断言定位错误。再临时去掉清梯度操作，解释多次 backward 为何改变更新量。完成后恢复正确版本，分别指出前向临时值、累积梯度与跨 step 保留的参数。

再用两层 MLP 复习：写出 X、H、pred、loss 和四个参数梯度的 shape；解释 ReLU 为什么不改变 shape 却改变梯度路径。运行 C.7 验算，再按 FP32 计算 11 个参数及 Adam 状态的逻辑字节数；指出这些字节为什么不是一次训练的实际峰值。能够把“计算需要谁”与“谁必须保留到何时”对应起来，再进入 GPT。

2. 完成单卡语言模型：训练、验证、恢复与生成

本章的产物是一个能重新加载并生成文本的小 GPT。你应能解释 token 输入如何变成 logits 和 loss，区分训练、验证、测试与生成，修改一个形状相关配置并验证结果。沿用 tiny Shakespeare 字符文本；CPU 可检查小规模流程，GPU 用于记录单卡基线。图像分类作为同一训练骨架的旁证放在 C.1。

本章沿用同一个训练循环，逐步增加模型和数据的复杂度。CPU 适合验证短例子和小规模 smoke test；完整 GPT 训练可放到单卡 H20/H800 或其他合适 GPU。学习所需设备由模型规模决定，不要求某个特定卡型。

2.1 Dataset 与 DataLoader：从单条样本到一个 batch

Dataset 说明“第 i 条样本是什么”；DataLoader 负责抽样、打乱、组 batch。下面是索引式 Dataset 的最小实现；另一类流式 IterableDataset 使用迭代接口，不要求同样的方法。

import torch
from torch.utils.data import Dataset, DataLoader
class Pairs(Dataset):
    def __init__(self):
        self.x = torch.linspace(-1, 1, 100).reshape(-1, 1)
        self.y = 2 * self.x + 1
    def __len__(self):
        return len(self.x)
    def __getitem__(self, i):
        return self.x[i], self.y[i]
loader = DataLoader(Pairs(), batch_size=16, shuffle=True, num_workers=0)
xb, yb = next(iter(loader))
print(xb.shape, yb.shape)  # (16,1), (16,1)

独立可运行的数据加载例子

batch 是一批样本，batch size 是样本条数。在 mini-batch SGD 的算法描述中，mini-batch 指一次参数更新所依据的样本批次；global batch 强调合计整个训练系统参与这次更新的全部样本。本例单卡每取一批就更新一次，DataLoader 返回的一批恰好就是这次更新的完整批次。分布式资料也可能把每个副本的本地部分称为 mini-batch；本文需要这个含义时明确写作 local mini-batch（每副本、每次更新），详见 A.13。

本例 Dataset 的单条 x[i] 为 $(1,)$，DataLoader 默认把 16 条堆成 $(16,1)$，新增的第 0 轴就是 batch。100 条数据、batch_size=16 时，drop_last=False 得到 7 个 batch（最后一个只有 4 条）；drop_last=True 得到 6 个完整 batch，并舍弃本轮余下 4 条。不要把 batch 大小硬编码进 reshape，也不要默默丢掉验证样本。

计算梯度和更新参数是两个动作：backward 计算并累积梯度，optimizer.step 才修改参数。若连续处理多个 DataLoader batch、每批执行一次前向和反向，最后才更新一次，那么每次实际执行的一批是 micro-batch；这些批次合起来构成该副本一次更新的 local mini-batch。在本节的单卡场景中，它也是这次更新的 global batch。因此单卡也可以累积多个微批后更新。本文 step 指一次参数更新，epoch 指训练数据被遍历一轮。

shuffle=True 改变每轮抽样顺序，同时保留输入与标签的配对关系；通常只打乱训练集。collate_fn 决定如何把样本列表合成 batch，默认逻辑会堆叠同形状张量。变长文本不能直接堆叠，需在组 batch 时补齐长度（padding）并保留 mask，或使用其他明确的变长表示。

数据变换分两类：确定性预处理（如转 Tensor、归一化）与训练用的随机增强。验证/测试应使用确定的评估流程；从数据估计归一化统计量时只使用训练集，避免信息泄漏。num_workers 是 CPU 数据加载进程数，不是 GPU 数或 CUDA 线程数。

2.2 先走完整流程：从文本到 checkpoint，再恢复和生成

本章先统一符号：B 是一个 micro-batch 中的序列条数，T 是每条序列的预测位置数，C 是每个 token 的隐藏特征数，V 是词表中候选 token 数；shape 的各项是元素个数。输入和目标是整数 ID 张量 $(B,T)$，隐藏激活是浮点张量 $(B,T,C)$，logits 是未归一化分数 $(B,T,V)$，loss 是 scalar。第 1 章的 vector/matrix/tensor 在这里分别承担 token 表示、参数表和批量计算的职责。

先把系统看成三条路径：训练读取文本窗口并更新模型；验证读取留出的文本，只前向并帮助选择 checkpoint；生成加载 checkpoint 和字符表，从提示词逐个追加字符。先按 C.2 保存三个完整脚本，再运行本节命令。理解输入、输出和状态后，2.3 再拆开模型内部；不要求先读完每层实现才能看到完整流程。

训练和验证的区别是数据用途与是否更新参数。model.eval() 改变 Dropout 等层行为，inference_mode/no_grad 控制梯度记录，二者分别设置；测试数据留到模型选择完成后评估。checkpoint 要区分仅供推理的模型权重和供续训的模型、优化器、进度等状态。

准备一个 UTF-8 文本文件 input.txt，与脚本放在同一目录。可以使用自己的文本，或从 Karpathy tinyshakespeare 示例数据下载；字符级玩具模型只用于学习。先保证训练段和验证段都长于 block_size，短文本可减小窗口。

下载 tinyshakespeare 示例文本

训练脚本默认 2 层、128 维、64 字符窗口。每隔若干更新步，在固定抽样窗口上分别评估训练段和验证段；best.pt 保存验证 loss 最优状态，last.pt 保存最近状态。checkpoint 同时包含模型配置、字符表、参数、优化器和步数，避免只有权重却无法正确恢复。恢复时应使用相同数据文件与 batch 设置；随机状态不保证跨硬件、跨版本逐位复现。

生成使用独立脚本，重新构建模型并加载 best.pt。每步只取最后位置的 logits，转成概率后采样一个字符。教学版本每次重算上下文，没有 KV cache；缓存与部署引擎放到 9.2。

python train_gpt.py --steps 500
python generate.py --checkpoint best.pt --prompt "To be"
python train_gpt.py --steps 1000 --resume last.pt

运行顺序；steps 是累计更新次数

检查是否学会：能解释 x 与 y 的错位、未来位置为什么必须被 mask、$(B,T,V)$ 如何与标签计算交叉熵；能展示 train/val loss，重新启动进程后仍能加载模型生成。恢复训练还要确认 step 接续、优化器状态被加载。生成文本未必流畅，训练数据、容量与步数都会影响效果。

修改练习：复制项目并使用独立输出目录，在 train_gpt.py 新建模型的配置中将 block_size 从 64 改为 32，重新训练 500 次更新，再运行相同生成命令。先预测每个窗口的目标数减半、位置嵌入表变短；运行后打印 input/target/logits 的 shape，确认生成仍只追加一个新字符。此实验同时改变上下文和每次更新目标数，只用于检查形状与流程，不能据此认定吞吐或模型质量提升。

2.3 从预测下一个字符，逐步搭建最小 GPT

本节模型是教学用 decoder-only Transformer：字符级词表、可学习位置嵌入、LayerNorm、GELU。它与 LLaMA 共享因果自注意力和残差堆叠的思路，但位置编码、归一化、MLP、tokenizer 等设计不同。先逐层理解这个小模型，再到 8.2 学习 LLaMA 的组件替换。

2.3.1 训练目标：每个位置预测它后面的一个字符

假设一段文本是“我爱学习”，输入 $x=[\text{我},\text{爱},\text{学}]$，标签 $y=[\text{爱},\text{学},\text{习}]$。标签窗口比输入的起点向后移动一个字符。训练时所有位置的预测可以一次计算；生成时没有未来答案，只能预测一个、接到序列后面，再继续预测。

位置	输入	允许看到的输入	预测目标
0	我	我	爱
1	爱	我、爱	学
2	学	我、爱、学	习

训练段和验证段按原始文本的连续区间划分，各自取窗口，不让窗口跨越边界。先查看训练 loss 是否下降，再看验证 loss 是否改善；只有训练 loss 下降而验证 loss 上升，可能已经过拟合。

2.3.2 Embedding：把整数编号变成可学习的向量

token id 是查表下标。Embedding 保存一张 $(V,C)$ 的参数表，V 是词表大小，C 是向量长度；输入 ids 为 $(B,T)$ 时，查出的向量为 $(B,T,C)$。同一 id 取同一行，表中的向量会通过训练更新。位置嵌入提供序列位置的信息。

import torch
from torch import nn
ids = torch.tensor([[0, 1, 2], [2, 1, 0]]) # B=2,T=3
embedding = nn.Embedding(4, 8)             # V=4,C=8
x = embedding(ids)
print(x.shape)                            # (2,3,8)

查表后的 shape

2.3.3 单头 Attention：评分、遮住未来、归一化、加权

先看一个头：输入 $X\,(B,T,C)$ 经三个投影得到 Q、K、V，均为 $(B,T,D)$，D 是每头特征数。$QK^{\mathsf T}$ 的形状是 $(B,T,T)$，表示每个查询位置对各键位置的分数；除以 $\sqrt D$、遮住未来、沿最后一个 T 轴 softmax 后，权重仍为 $(B,T,T)$。$\mathrm{weights}\,V$ 得到 $(B,T,D)$。下方程序把 $H=4$ 个头放在一起独立计算，$B=2$、$T=5$、$D=8$，所以 $\operatorname{shape}(q)=\operatorname{shape}(k)=\operatorname{shape}(v)=(2,4,5,8)$、$\operatorname{shape}(\mathrm{scores})=\operatorname{shape}(\mathrm{weights})=(2,4,5,5)$、$\operatorname{shape}(\text{输出})=(2,4,5,8)$；H 是头数，不是高度或隐藏维。

因果注意力示意图：第 i 行只能使用第 0 到 i 列，未来列的权重为 0；权重用于对 V 求加权和。

打开 SVG 原图

"""手写因果注意力，与 SDPA 对照；CPU 即可运行。"""
import torch
from torch.nn import functional as F

torch.manual_seed(42)
q, k, v = [torch.randn(2, 4, 5, 8) for _ in range(3)]  # B,H,T,D
scores = q @ k.transpose(-2, -1) / (q.size(-1) ** 0.5)
allowed = torch.ones(5, 5, dtype=torch.bool).tril()
scores = scores.masked_fill(~allowed, float("-inf"))
weights = scores.softmax(dim=-1)
manual = weights @ v
sdpa = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)
print("输出 shape：", manual.shape)
print("每行权重和：", weights.sum(dim=-1)[0, 0])
print("未来位置最大权重：", weights[..., ~allowed].abs().max().item())
torch.testing.assert_close(manual, sdpa, atol=1e-5, rtol=1e-5)
print("手写 attention 与 SDPA 对照通过")

保存为 attention.py；先跑通手写计算与 SDPA 对照

注意这里有两种 softmax：Attention 内部沿上下文位置归一化；语言模型输出沿词表候选字符归一化。它们作用的轴和对象不同。减去每行最大值再取指数是稳定实现的一部分；直接使用 softmax/log_softmax 等库函数。SDPA 根据设备、dtype、shape 等选择后端，不保证每次都使用 FlashAttention。

2.3.4 多头、残差与 Pre-Norm：一个 Block 内部发生什么

多头把 C 维拆成 H 个头，每个头的维度 $D=\frac{C}{H}$。将 $(B,T,C)$ 重排成 $(B,H,T,D)$ 后，矩阵乘只作用于最后两维；前面的 B、H 是批维。分数 shape 为 $(B,H,T,T)$，加权 V 后回到 $(B,H,T,D)$，再合头成为 $(B,T,C)$。

步骤	shape
输入/每组 Q、K、V	$$(B,T,C)$$
拆头、交换轴	$$(B,H,T,D)$$
$$QK^{\mathsf T}$$	$$(B,H,T,T)$$
$$\mathrm{weights}\,V$$	$$(B,H,T,D)$$
合头并输出投影	$$(B,T,C)$$

残差相加保留原输入，并提供直接的梯度通路；相加的两个张量形状必须匹配。Pre-Norm 表示先归一化，再进入注意力或 MLP 分支。本教程 Block 的两步是 $x\leftarrow x+\operatorname{Attention}(\operatorname{LayerNorm}(x))$，随后 $x\leftarrow x+\operatorname{MLP}(\operatorname{LayerNorm}(x))$。MLP 对每个 token 独立做同一组非线性变换；Attention 负责位置之间的信息交互。

结构位置	符号 shape	小例子 $B=2,\ T=3,\ C=8,\ H=2,\ D=4$
Block 输入 / LayerNorm 输出	$$(B,T,C)$$	$(2,3,8)$；归一化最后的特征轴
Q/K/V 拆头后	$$(B,H,T,D)$$	$(2,2,3,4)$；$C=H\times D$
注意力分数 / 权重	$$(B,H,T,T)$$	$(2,2,3,3)$；最后两轴为查询/键位置
加权 V、合头、输出投影	$$(B,H,T,D)\longrightarrow(B,T,C)$$	$$(2,2,3,4)\longrightarrow(2,3,8)$$
MLP 两层（本教学模型中间宽度 $4C$）	$$(B,T,C)\longrightarrow(B,T,4C)\longrightarrow(B,T,C)$$	$$(2,3,8)\longrightarrow(2,3,32)\longrightarrow(2,3,8)$$
两条残差相加	$$(B,T,C)+(B,T,C)$$	每次都得到 $(2,3,8)$，不会拼接成 16 维

GPT Block 内部图：两次 Pre-Norm、两条残差旁路和两次相加，输出 shape 保持 $(B,T,C)$。

打开 SVG 原图

2.3.5 回到完整模型：把局部机制接回训练循环

GPT 架构图的 shape 读法：token IDs 为整数张量 $(B,T)$，词嵌入表 $(V,C)$ 查表得到 $(B,T,C)$，位置嵌入 $(T,C)$ 沿 B 广播相加；各 Block 保持 $(B,T,C)$，最终投影得到 $\mathrm{logits}\,(B,T,V)$。训练时展平为 $(B\times T,V)$，目标为 $(B\times T,)$，交叉熵归约成 $\mathrm{loss}\,()$；生成时只取 logits[:,-1,:] 得到 $(B,V)$。完整符号、小尺寸例子与权重布局见 A.14。

打开 SVG 原图

完整 model.py、train_gpt.py 和 generate.py 已集中在 C.2。这里回看整条前向：输入编号 → Embedding → Block 堆叠 → 最后归一化 → 词表 logits；训练分支把 logits 与向后错一位的目标计算 loss，生成分支只取最后位置。按此链路去找函数，不需要按文件顺序记住全部代码。

2.3.6 从 shape 推导计算量：这个 Block 做了多少工作

先对矩阵乘 $X\,(M,K)\mathbin{@}W\,(K,N)\longrightarrow Y\,(M,N)$ 计数：$MN$ 个输出元素，每个累加 K 项乘积，常用估算为 $2MKN$ FLOPs，一次乘加按两次浮点运算计。FLOPs 在这里表示运算次数，FLOP/s 才是吞吐速率。1.6.2 的两层 MLP 仅计矩阵乘为 $2\times2\times3\times2+2\times2\times2\times1=32$ FLOPs；偏置和 ReLU 的运算另外计，省略它们是估算口径，不是说它们没有耗时。

沿用前面的 B、T、C、H 与 $D=\frac{C}{H}$：线性层可把 B、T 展平为 $M=B\times T$，Attention 则是 $B\times H$ 组矩阵乘。下面只计算本教学模型的标准多头 Attention 与中间宽度 $4C$ 的两层 MLP；列出的权重数不含 bias、LayerNorm，也不含整个模型的 Embedding、词表投影。

主要运算	输出 shape / 权重元素数	前向矩阵乘 FLOPs
Q/K/V 投影合计	$(B,T,3C)$；$3C^2$	$$6BTC^2$$
$$QK^{\mathsf T}$$	$(B,H,T,T)$；无可训练权重	$$2BHT^2D=2BT^2C$$
$$\mathrm{Attention\ weights}\,\mathrm{Value}$$	$(B,H,T,D)$；无可训练权重	$$2BT^2C$$
Attention 输出投影	$(B,T,C)$；$C^2$	$$2BTC^2$$
MLP 第一层 $C\to4C$	$(B,T,4C)$；$4C^2$	$$8BTC^2$$
MLP 第二层 $4C\to C$	$(B,T,C)$；$4C^2$	$$8BTC^2$$
上述矩阵乘合计	$12C^2$ 个权重元素	$$24BTC^2+4BT^2C$$

代入已有小例子 $B=2$、$T=3$、$C=8$、$H=2$，得到 9,792 FLOPs；显式 Attention 分数有 $BHT^2=36$ 个元素，FP32 为 144 B。若仅将 T 翻倍，线性层矩阵乘量翻倍，$QK^{\mathsf T}$ 与 $AV$ 的量及完整注意力矩阵元素数变为四倍。这给 CP、重计算和 FlashAttention 的实验提供了可验证的预期。

这是稠密矩阵乘的前向估算；softmax、归一化、激活、残差与数据搬运仍有成本。因果 mask 在朴素实现中不会自动省去半次 GEMM；支持因果分块的内核可能跳过部分工作。FlashAttention 仍计算必要的分数块，但避免把完整 $T\times T$ 中间矩阵写回 HBM。训练还要计算参数/输入梯度，重计算会增加前向工作；GQA、SwiGLU、不同词表投影等结构必须按实际 shape 重算，不能直接套本式。

2.4 单卡优化：一次改变一项，再看收益

2.4.1 先形成预期：计算、搬运与实测分别是什么

对 $X\,(M,K)\mathbin{@}W\,(K,N)$，若每元素占 s 字节，假设两份输入各从 HBM 读一次、输出写一次，则最低搬运量估算为 $s(MK+KN+MN)$ 字节。它是理想模型：实际缓存复用、重复读取、融合、输出累加与临时缓冲会改变流量。定义算术强度 $I=\frac{\mathrm{FLOPs}}{\text{搬运字节数}}$，含义是每搬一个字节能做多少计算。

设计算吞吐上限 P 的单位为 FLOP/s，带宽上限 BW 的单位为 byte/s，则 $t_{\mathrm{compute}}=\frac{2MKN}{P}$，$t_{\mathrm{memory}}=\frac{s(MK+KN+MN)}{\mathrm{BW}}$，理想时间下界约为 $\max(t_{\mathrm{compute}},t_{\mathrm{memory}})$。若用实测有效吞吐和带宽代入，则是在构造经验估计。启动延迟、同步、无法重叠的阶段等可能使真实时间更长；不要把两个理想时间机械相加，也不要把下界填进实测结果栏。

同一个线性层：$K=1024$，$N=4096$，FP32	$$M=1$$	$$M=512$$
矩阵乘运算量	8,388,608 FLOPs	4,294,967,296 FLOPs
理想搬运量	16,797,696 B（约 16.02 MiB）	27,262,976 B（26 MiB）
理想算术强度	约 0.50 FLOP/B	约 157.54 FLOP/B
应验证的预期	权重复用少；启动与搬运可能占较大比重	更多输入行复用同一权重，可能更充分利用计算单元

表格只改变 M，不改变权重形状。它解释了小 batch 或逐 token 计算为什么可能难以发挥峰值算力；不能据此保证增加 M 就一定更快，更大的激活可能先触及容量限制。C.8 的 resource_probe.py 在 CPU 上核对公式和正确性，在 CUDA 上才测完整训练更新；完整更新包含前向、反向和 AdamW，不能直接与表中单次 GEMM 的时间相比。

参考阅读：从小 MLP 推导显存、计算量与访存。本教程沿用这种分析顺序，并将理论估算、具体实现与实测结果分开记录。

2.4.2 单项改造：保持工作量与验证口径一致

先保留能稳定收敛的 FP32 基线，再比较优化。每次使用相同模型、数据、序列长度和有效 batch；记录 warmup 后的 tokens/s、峰值已分配显存、验证 loss，并单独记录编译或初始化耗时。固定若干 step 的显存/耗时改善不等于收敛效果改善。

顺序	只改变这一项	观察与边界
0	FP32 基线	先确认训练/验证/保存/加载闭环
1	bf16 autocast	硬件支持时尝试；参数和优化器状态不一定减半，测总显存与验证 loss
2	梯度累加	保持有效 batch 比较；仅在一组 micro-batch 完成后更新参数
3	梯度裁剪	在累加完成后、optimizer.step 之前裁剪，记录梯度范数
4	warmup 与 cosine	按参数更新步数推进，保存和恢复 scheduler 状态
5	fused AdamW	确认当前设备、dtype、版本支持，再测优化器耗时
6	torch.compile	先保持其他配置不变；首次编译单独计时，留意 graph break

单项实验 A：只启用 bf16。将训练脚本 forward/loss 的两行放入下面的 autocast 上下文，其他行为先保持不变。bf16 通常不需要 GradScaler；fp16 的动态范围与缩放策略不同，见 B.3。

if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
    raise RuntimeError("本实验需要支持 bf16 的 CUDA GPU")
with torch.amp.autocast("cuda", dtype=torch.bfloat16):
    _, loss = model(x, y)
loss.backward()

替换 train_gpt.py 中对应的 forward/backward 片段

单项实验 B：只使用梯度累加。下列代码替换一次参数更新中的取 batch、反向和 step。每个 micro-batch 大小相同时，loss 除以累加次数；若样本/有效 token 数不同，应按实际数量加权。与大 batch 比较时，还要考虑 Dropout、BatchNorm 等可能带来的差异。

accum_steps, micro_batch = 4, 4  # 有效 batch 仍是 16
optimizer.zero_grad(set_to_none=True)
for _ in range(accum_steps):
    x, y = get_batch(train, config["block_size"], micro_batch, device)
    _, loss = model(x, y)
    (loss / accum_steps).backward()
optimizer.step()

梯度累加：先保持有效 batch 相同

测 CUDA 耗时要在计时边界同步，先完成预热；峰值显存统计在同一范围开始前 reset_peak_memory_stats。tokens/s 按$\frac{\text{更新步数}\times\text{每步有效 batch}\times\text{序列长度}}{\text{秒数}}$计算。训练循环频繁 .item() 会引入同步，但日志应以可观测性和实际测量为准。更完整的计时解释见 B.2。

2.4.3 重计算实验：少保留激活，需要付出什么

1.7 的账本说明，前向保存激活是为了反向使用。activation checkpointing 保留选定边界，在反向需要内部数据时重新执行必要的前向计算，从而减少部分保存值。预计是峰值显存可能降低、计算量增加；边界输入、参数和优化器状态仍要保留，不能把它解释为“训练不需要激活”。

运行 C.8：先执行 python resource_probe.py --device cpu，确认普通前向与 checkpoint(use_reentrant=False) 在同一权重、输入和目标上的输出、loss 与参数梯度一致；再在 CUDA 机器执行 python resource_probe.py --device cuda。后者固定 FP32、$M=512$、$1024\to 4096\to 1024$ 的 MLP、AdamW 和种子，只切换是否 checkpoint，并分别预热、清理上一分支后测量。

观察每步耗时、测量前 allocated、区间 peak allocated 及二者差值；不要把 allocated 与 reserved 混为一谈。Adam 状态在预热中建立，避免把首次状态分配与稳态更新混在一起。若峰值没有下降，先判断峰值是否发生在重计算以外的阶段、该边界是否确实省下足够大的保存值；若更慢，则量化时间换空间的代价。这些都是实验结论，不能预设优化必然有效。

小 MLP 验证通过后，再对最小 GPT 的一个或一组 Block 采用相同方法，其他配置保持不变；带 Dropout 等随机操作时还要保持重计算的随机状态语义，并再次对齐梯度。本教程新增脚本的 CPU 正确性已验证；CUDA 性能结果需要在目标设备上实际填写。

完成单项实验之后，可用附录 B.6 整合优化。先确认单卡训练、验证、保存、加载与生成都正常，再进入第 3 章，沿真实 nanoGPT 仓库完成同一个闭环。第 4—5 章随后在同一台机器上增加 GPU，第 6 章再增加机器；不要用改变硬件规模掩盖尚未解决的单卡问题。

3. 单机单卡实战：沿 nanoGPT 跑通训练与源码

nanogpt-learning-labs.zip

本章沿一次单机单卡训练展开：文本变成样本，样本经过模型得到 loss，梯度形成参数更新，checkpoint 用于验证、恢复和生成。先用 CPU 小配置检查数据与代码，再在一张 GPU 上跑通相同流程并记录基线。进入条件是能理解第 2 章的 GPT 前向和训练循环；多进程和多卡分工在第 4—5 章学习。

第 3—6 章固定阅读 Karpathy nanoGPT 的提交 3adf61e154c3fe3fca428ad6bc3818b27a3b8291（2025-11-12），源码于 2026-09-13 核对。上游 README 已将 nanoGPT 标为历史项目并推荐 nanochat；这里选择它，是因为模型和训练入口紧凑，适合学习 GPT-2 风格实现。nanoGPT、build-nanogpt 和 nanochat 是不同项目，代码与启动命令不能混用。

第 3—6 章的运行与检查入口已整理到上方 nanogpt-learning-labs.zip。下载并解压后，把 nanogpt-learning-labs 目录中的 Python 文件复制到固定版 nanoGPT 仓库根目录；所有命令都从该根目录运行。包内 README 标明每个脚本的用途与验证范围；原版 nanogpt-guide.zip 仍保留在第 5 章末尾。

阶段	本次只增加什么	留下什么结果
第 3 章：单卡	先运行原版，再观察一次更新并使用官方示例配置	数据校验、checkpoint、续训和生成；一份可解释的单卡记录
第 4 章：单机两卡 DDP	同一模型分配给两个数据副本；再修改梯度同步范围	固定输入更新对照、同步改造记录、同工作量吞吐
第 5 章：FSDP2 分支	在同一台机器上分片训练状态，并改变保存方式	分片 shape、一次更新对照、连续与恢复运行对照、导出
第 6 章：多机	保留已通过的训练入口，增加节点和跨机网络	通信检查、跨机训练与恢复、相同总卡数的对照

3.1 先走完整条链路：训练结束后，我们得到了什么

以 tiny Shakespeare 字符级模型为例：输入是一份剧本文本，训练目标是在每个位置预测下一个字符，最终产物是模型 checkpoint。训练反复取样、预测、计算误差，再用误差的梯度调整参数；生成阶段加载这些参数，从提示词开始逐个采样新字符。

准备数据 → 建立训练状态。 prepare.py 将 input.txt 编成 train.bin / val.bin，meta.pkl 保存编号与字符映射。train.py 读取配置、构造 GPT、创建 AdamW；恢复训练时载入 checkpoint 中的模型与优化器。

取样 → 前向 → loss → 反向 → 更新。 每批取输入 X 和向后错一位的目标 Y；GPT 输出每个位置对词表的分数 logits，交叉熵衡量预测误差。backward 累积梯度，达到设定的累加次数后，先裁剪梯度，再由优化器更新参数并清空梯度。重复这一循环，模型逐步学会预测下一个字符。

验证与保存 → 生成与恢复。 验证段用于估计对未训练文本的预测误差；原版在下一次更新前触发验证与条件保存。sample.py 用 checkpoint、词表和提示词生成新文本；train.py 的 resume 路径则恢复训练状态，继续更新参数。

先记住三个时间尺度：一个 token 是一个字符编号；一个 micro-batch 是 B 段长度为 T 的窗口；一次参数更新由 A 个 micro-batch 累积而成。下面的小配置 $B=4$、$T=32$、$A=2$，因此一次更新使用 8 段窗口、256 个 token 预测目标。验证和生成只做前向，不更新参数。

后续按这条主线阅读：先用项目关系图定位三个入口，跑通短训练与生成，再沿 get_batch → GPT.forward → Block / Attention / MLP → 训练循环 → generate 下钻源码。每一步都先确认输入、输出和参数是否发生变化。

3.2 从运行过程定位源码：三条入口与共享模型

把 nanoGPT 看成三条入口：prepare.py 把文本变成训练数据，train.py 更新参数，sample.py 用已有参数生成文本。model.py 提供共同使用的模型定义；配置决定模型大小、数据位置和训练方式。阅读时先走完这三条路径，再深入 Attention 内部。

文件 / 对象	职责与阅读重点
data/shakespeare_char/prepare.py	下载或读取 input.txt，建立字符词表；按前 90% / 后 10% 切成训练与验证段，输出 train.bin、val.bin、meta.pkl。
config/train_shakespeare_char.py + configurator.py	配置文件是 Python 代码；默认值先定义在入口脚本中，然后按命令行顺序覆盖。把配置文件放前面，--key=value 放后面。
train.py	初始化进程和模型；get_batch 取数据；训练、验证、记录指标和保存。它是理解工程组织的主线。
model.py	GPTConfig、GPT、Block、CausalSelfAttention、MLP；还包括 optimizer 配置、GPT-2 权重导入与 generate。
sample.py	从 checkpoint 重建 GPT，按 dataset 找到 meta.pkl，编码提示词、循环采样并解码。
out_dir/ckpt.pt	原版保存模型、优化器、模型配置、迭代号、验证记录和运行配置；词表映射仍在数据目录。

项目关系图：配置、数据文件、训练入口与生成入口各自承担什么职责。

打开 SVG 原图

建议按“配置 → prepare.py → train.py 的 get_batch → GPT.forward → 训练循环 → sample.py”阅读。先回答数据从哪里来、shape 是什么、参数什么时候变化，再读融合 QKV、权重共享和性能优化。源码入口：train.py、model.py、sample.py、data/shakespeare_char/prepare.py。

这里的模型采用可学习绝对位置 Embedding、LayerNorm、GELU MLP 和多头因果自注意力。它没有直接采用 LLaMA 的 RoPE、RMSNorm、SwiGLU/GQA。transformer.wte.weight 与 lm_head.weight 是同一个参数，这个共享关系在保存、FSDP 分组和 PP 拆阶段时都需要保留或明确处理。

3.3 跑通一次：数据准备、短训练、恢复与生成

先准备 Python 3.12，并建立独立环境。下面给出 CPU 起步命令；已有符合要求的环境可直接激活。CUDA 机器改用与驱动兼容的 PyTorch 官方安装命令，随后核对 torch.version.cuda 与 torch.cuda.is_available()。本教程已验证的 CPU 环境是 Python 3.12.14、PyTorch 2.9.1。字符路径需要 numpy、requests；sample.py 即使读取字符词表也会在启动时导入 tiktoken。

git clone https://github.com/karpathy/nanoGPT.git
cd nanoGPT
git checkout 3adf61e154c3fe3fca428ad6bc3818b27a3b8291
python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install torch==2.9.1 numpy requests tiktoken
python -c "import torch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available())"
python data/shakespeare_char/prepare.py
# 把下载的 nanogpt-learning-labs 目录放到当前仓库根目录，再执行：
cp nanogpt-learning-labs/*.py .

官方 tiny Shakespeare 数据准备后，词表有 65 个字符，训练段 1,003,854 个 token，验证段 111,540 个 token。train.bin/val.bin 中保存 uint16 编号，meta.pkl 保存 vocab_size、stoi 和 itos。分割是文本的连续前后段，随机训练窗口不会跨越这条分界；窗口之间可能重叠。若替换成自己的文本，重新建立词表并重新准备数据；词表超过 uint16 容量时要同时修改写入和读取的数据类型。

将下列内容保存为仓库根目录的 config_smoke.py。它把实验缩到 2 层、2 头、64 维、32 个字符上下文，关闭编译与混合精度，方便先检查流程。梯度累加次数设为 2，后面可以直接用于 2 卡 DDP。

# Run from the pinned nanoGPT repository root:
# python train.py /path/to/config_smoke.py
out_dir = 'out-guide'
dataset = 'shakespeare_char'
device = 'cpu'
dtype = 'float32'
compile = False
wandb_log = False
n_layer = 2
n_head = 2
n_embd = 64
block_size = 32
batch_size = 4
gradient_accumulation_steps = 2
dropout = 0.0
learning_rate = 1e-3
decay_lr = False
max_iters = 20
eval_interval = 10
eval_iters = 2
log_interval = 5
always_save_checkpoint = True

config_smoke.py：本章 CPU 起步配置

python check_inputs.py
python train.py config_smoke.py
python inspect_checkpoint.py out-guide/ckpt.pt
python sample.py --out_dir=out-guide --device=cpu --dtype=float32 --num_samples=1 --max_new_tokens=40 --start=ROMEO:
python train.py config_smoke.py --init_from=resume --max_iters=30
python inspect_checkpoint.py out-guide/ckpt.pt

从 CPU 检查转到单机单卡。上面的 CPU 命令用于确认代码和数据；有一张 CUDA GPU 后，用同一模型完成下面的训练、生成和恢复，建立后续多卡实验的基线。先保留 FP32、dropout=0 和 compile=False；CUDA_VISIBLE_DEVICES=0 只向进程暴露一张卡，可按环境调整为目标设备。

CUDA_VISIBLE_DEVICES=0 python train.py config_smoke.py --device=cuda --dtype=float32 --out_dir=out-guide-1gpu
CUDA_VISIBLE_DEVICES=0 python sample.py --out_dir=out-guide-1gpu --device=cuda --dtype=float32 --num_samples=1 --max_new_tokens=40 --start=ROMEO:
CUDA_VISIBLE_DEVICES=0 python train.py config_smoke.py --device=cuda --dtype=float32 --out_dir=out-guide-1gpu --init_from=resume --max_iters=30

单机单卡 nanoGPT：训练、生成、恢复

先确认单卡训练、生成与恢复均成功：check_inputs.py 应报告数据有效，并核对当前提交、词表和文件校验值；inspect_checkpoint.py 应显示保存的 iter_num、模型配置、非空优化器状态以及有限参数。吞吐和峰值显存使用 3.8 的观测入口记录，原版日志的 time/MFU 不能直接代替这一测量。GPU 结果由目标机器实测。

先观察日志是否识别 vocab_size=65、是否打印 train/val loss，以及 out-guide/ckpt.pt 是否生成。几十步训练通常只会生成杂乱字符；本节用它验证完整流程，语言质量需要更充分的训练。提示词字符必须存在于 stoi 中，空提示词也不适合直接喂给原版 generate。生成时保留原来的 data/shakespeare_char/meta.pkl；若丢失，sample.py 会尝试回退到 GPT-2 编码，不能用这个回退解释字符级模型的输出。

理解两个时间点：原版在循环开头做验证与条件保存，且 iter_num>0 才真正落盘；停止条件在更新后检查 iter_num>max_iters。因此 max_iters=20 会执行编号 0—20 的 21 次更新，最后保存的 step 20 文件对应第 20 次已完成更新，内存随后还会再更新一次。恢复训练从保存的 iter_num 继续，不代表从退出前最后一次更新继续。要做严格实验，可以改成 while completed_updates<target，并在完成更新后递增计数和保存。

日志或文件	它对应的训练状态	读者如何判断
step 0: train loss …, val loss …	首次参数更新前的抽样验证	基线误差；不是已经训练了一轮
iter 0: loss …, time …	编号 0 的第一次更新已经执行；训练 loss 来自最后一个微批	不要把它当作所有微批的精确均值，也不要要求逐行下降
step 20 与 ckpt.pt 的 iter_num=20	循环开头验证、保存，已有 20 次更新完成	原循环随后还会执行编号 20 的更新；磁盘状态不是退出前最后一次内存状态

原版 resume 恢复模型和优化器，但没有保存 CPU/CUDA RNG、数据采样进度和 FP16 GradScaler 状态；所以“能继续训练”不等于“与不中断训练逐步一致”。配置文件仍会重新执行，未强制恢复的配置可能改变运行行为。保留同一数据和词表，确认学习率、batch、精度、dropout 等设置，再解释恢复后的曲线。

再用官方示例配置跑一次。config_smoke.py 用于检查链路；下面改读仓库的 config/train_shakespeare_char.py，保留官方的 6 层、6 头、384 维、256 字符上下文和 dropout=0.2。只把每次计算批量缩到 4、训练缩到约 200 次更新，并相应缩短学习率调度和验证周期。命令行参数放在配置文件后，覆盖文件中的同名值。使用新的输出目录，避免与小配置 checkpoint 混用。

CUDA_VISIBLE_DEVICES=0 python train.py config/train_shakespeare_char.py --device=cuda --dtype=float32 --compile=False --batch_size=4 --max_iters=200 --lr_decay_iters=200 --warmup_iters=20 --eval_interval=100 --eval_iters=10 --always_save_checkpoint=True --out_dir=out-official-char
python inspect_checkpoint.py out-official-char/ckpt.pt
python sample.py --out_dir=out-official-char --device=cpu --dtype=float32 --num_samples=1 --max_new_tokens=100 --start=ROMEO:

这轮验收的是“会使用官方配置”：核对 checkpoint 中的模型配置、保存步数与数据集，解释哪些值来自官方文件、哪些被命令行覆盖。它改变了模型和训练条件，不与 smoke 的速度直接比较；短跑也不保证生成可读文本。CPU 学习者可先完成 smoke 和源码练习，再到 GPU 环境运行这一轮。

现象	对应检查
没有 ckpt.pt	检查是否走到非零 eval_interval、保存条件是否满足；短跑使用本章 always_save_checkpoint=True。
配置覆盖失败	原版采用 --key=value；bool 写 True/False；数值类型要匹配默认值，例如浮点参数写 0.0。
采样 KeyError 或文本异常	核对提示词字符、meta.pkl、模型 vocab_size 和 checkpoint.config.dataset；不要混用字符级与 BPE 数据。
显存不足 / CPU 太慢	先减小 batch_size、block_size、n_layer、n_embd；保证 n_embd 能整除 n_head。
旧 AMP 接口出现 FutureWarning	固定版本仍使用 torch.cuda.amp.GradScaler；本章 CPU FP32 路径已运行通过，扩展代码采用独立入口。

3.4 沿一批数据：串起输入、前向与参数更新

源码入口：prepare.py：字符词表与数据切分，以及 train.py：get_batch。prepare 先用排序后的字符集合建立 stoi / itos，再编码文本；get_batch 每次随机取窗口，不是按 epoch 顺序遍历整个训练集。X 和 Y 共用同一段文本，只错开一个位置，因此不需要人工标注答案。

ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])

磁盘上使用 uint16 节省空间，取样时转成 int64，对应 PyTorch 的 torch.long。起步配置得到 $\operatorname{shape}(X)=\operatorname{shape}(Y)=(4,32)$。训练、验证使用各自的 bin 文件；验证也随机抽窗口。代码没有 DataLoader / DistributedSampler，也没有保证窗口互不重叠，这些行为将在第 4 章用于理解多卡数据分配。

train.py 的 get_batch 不使用第 2 章示例中的 DataLoader。它读取 memmap，在 token 流上随机选 B 个起点，每个位置取长度 T 的输入及向后错一位的目标。若 token 为 $[10,11,12,13,14]$、$T=4$，则 $X=[10,11,12,13]$，$Y=[11,12,13,14]$。这是同一前向同时学习多个“下一个 token”预测的原因。

下表中 $B$ 是一次实际前向的窗口数，$T$ 是每个窗口的 token 数，$C$ 是隐藏维度，$H$ 是注意力头数，$V$ 是词表大小。本例 $B=4, T=32, C=64, H=2, V=65$；后续增加 GPU 时先保持模型维度和这些统计口径明确。

阶段	形状与含义
X / Y	$(B,T)$，torch.long；每个位置是词表中的整数编号。
wte(X) + wpe(position)	$(B,T,C)$；词向量与位置向量相加。
c_attn → split → reshape	先得到 $(B,T,3C)$，按 Q/K/V 各 C 分开，再变为 $(B,H,T,C/H)$。
因果 Attention → c_proj	每个位置只使用允许访问的上下文；多头合回 $(B,T,C)$，再做输出投影。
Block 与 ln_f	每个 Block 两条 Pre-Norm 残差，维持 $(B,T,C)$；最后再做 LayerNorm。
lm_head 与 loss	训练输出 $(B,T,V)$，展平成 $(B\times T,V)$ 与目标 $(B\times T,)$ 计算交叉熵。未传 targets 时原版只计算最后位置的 logits。

训练流程图：梯度累加形成一次更新；验证、保存与生成具有各自的入口和时机。

打开 SVG 原图

一次 optimizer 更新包含 A 个 micro-step：各次前向的 loss 除以 A 后 backward，梯度累积到参数上；最后才做 unscale、梯度裁剪、optimizer 更新和清零。原版通过 scaler.step/update 兼容 FP16；FP32/BF16 时 scaler 不启用。前向之后立即取下一批，再对当前 loss 反向，这是一种预取安排，并不会让当前 loss 改用下一批标签。

有两个指标容易读错。训练日志里的 loss 是最后一个 micro-step 的 loss 乘回 A，并非所有 micro-step 的精确平均；estimate_loss 则在 eval 模式和 no_grad 下对若干随机批次求平均。另一个是 MFU：estimate_mfu 硬编码 A100 BF16 峰值 312 TFLOPS，不能把 CPU、H800 或其他设备打印的百分比直接当作硬件利用率。

model.py 中名为 flash 的开关只检查 scaled_dot_product_attention 这个 API 是否存在；实际采用哪种 SDPA 后端由设备、dtype、shape 等决定，不能仅凭这个变量断言启用了 FlashAttention。没有该 API 时，源码使用显式 $QK^{\mathsf T}$、因果 mask、softmax 和与 V 相乘的路径。

3.5 前向下钻：GPT → Block → Attention / MLP → loss

先看外层 GPT.forward。以 $B=4$、$T=32$、$C=64$、$H=2$、$V=65$ 为例，wte 把编号变成 $(4,32,64)$，wpe 生成 $(32,64)$ 并沿 batch 广播相加。两个 Block 和末尾 LayerNorm 保持 hidden 的形状，lm_head 将每个位置投影到 65 个字符的分数。这里的 logits 尚未经过 softmax。

GPT.forward 核心片段：先沿模型外层读一遍，再进入其中的 Block。

tok_emb = self.transformer.wte(idx)
pos_emb = self.transformer.wpe(pos)
x = self.transformer.drop(tok_emb + pos_emb)
for block in self.transformer.h:
    x = block(x)
x = self.transformer.ln_f(x)
if targets is not None:
    logits = self.lm_head(x)
    loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

loss 将 $(4,32,65)$ 展平成 $(128,65)$，与 128 个目标编号计算交叉熵，默认对有效目标求平均。F.cross_entropy 内部处理 log-softmax，不要先把 logits 做 softmax 再传进去。因果 Attention 阻止位置 t 读取未来输入，而 Y[t] 正是下一字符；训练一次并行算出 32 个位置的预测，这与生成时逐个追加 token 是两个阶段。

继续进入 Block.forward：先对 x 做 LayerNorm，再做 Attention，与原来的 x 相加；随后对新的 x 做 LayerNorm、MLP，再残差相加。归一化放在子层之前，所以称 Pre-Norm。残差要求两个子层都输出相同的 C 维。

def forward(self, x):
    x = x + self.attn(self.ln_1(x))
    x = x + self.mlp(self.ln_2(x))
    return x

再进入 CausalSelfAttention.forward。c_attn 是一次 $C\to3C$ 的融合线性变换，输出排列是完整 Q、完整 K、完整 V。split 之后各自变成 $(B,H,T,C/H)$，小配置中就是 $(4,2,32,32)$。每个头计算带因果限制的注意力，多头再合并成 $(4,32,64)$，经过 c_proj 回到残差路径。

以下摘出有 SDPA API 时的主路径；完整源码还保留了显式 $QK^{\mathsf T}$、mask、softmax 与乘 V 的后备实现。

q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
y = torch.nn.functional.scaled_dot_product_attention(
    q, k, v, attn_mask=None,
    dropout_p=self.dropout if self.training else 0,
    is_causal=True)
y = y.transpose(1, 2).contiguous().view(B, T, C)
y = self.resid_dropout(self.c_proj(y))

最后进入 MLP：c_fc 将每个 token 的 64 维扩展为 256 维，GELU 提供非线性，再由 c_proj 投回 64 维。MLP 对各位置分别计算；跨位置的信息交流发生在 Attention。这也是附录 B.8 和第 7.4 节从线性层理解张量并行的原因。

def forward(self, x):
    x = self.c_fc(x)
    x = self.gelu(x)
    x = self.c_proj(x)
    x = self.dropout(x)
    return x

3.6 训练下钻：梯度如何变成更新，更新如何被记录

回到 train.py：scaler 与 optimizer 初始化。模型构造后，configure_optimizers 按参数维数分成两组：矩阵和 Embedding 参与 weight decay，bias 与 LayerNorm 参数不衰减，然后创建 AdamW。优化器保存自己的状态；前向计算出 loss 时，参数尚未改变，backward 只向参数的 .grad 累积梯度。

以下是原训练循环的单设备阅读版，省略 DDP 分支与日志，保留原有操作顺序。FP32 起步配置中 scaler 是空操作；通常在 scaler.step 内部调用 optimizer.step 时参数才改变，FP16 检测到非有限梯度时可能跳过这次更新。

for micro_step in range(gradient_accumulation_steps):
    with ctx:
        logits, loss = model(X, Y)
        loss = loss / gradient_accumulation_steps
    X, Y = get_batch('train')  # 为下一次前向取数
    scaler.scale(loss).backward()
if grad_clip != 0.0:
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)

固定大小、没有 padding 的 A 个 micro-batch，各自 loss 除以 A 后累加，等价于对这 A 批 token 的平均 loss 求梯度；dropout 与浮点归约顺序会影响逐值对照。zero_grad 不会清空模型参数，它释放或清空梯度，为下一次更新准备。前向后的 get_batch 提前准备下一批数据，不会让当前 loss 改用下一批标签。

训练还要管理时间：get_lr 实现 warmup、cosine decay 与最小学习率，循环开头将结果写入 optimizer.param_groups；本章 decay_lr=False，始终使用配置的学习率。estimate_loss 在 eval 与 no_grad 下抽取 train/val 批次，结束后切回 train。eval 改变 Dropout 等层的行为，no_grad 关闭梯度记录，两者作用不同。

保存发生在验证与 checkpoint 分支；恢复路径见模型恢复和优化器恢复。ckpt.pt 包含模型权重、AdamW 状态、模型结构参数、迭代号和配置；meta.pkl 仍在数据目录。结合 3.3 的保存时机与恢复限制，区分“能生成”“能续训”和“能精确复现中断前轨迹”。

3.7 生成下钻：checkpoint 怎样续写 ROMEO:

源码从 sample.py：恢复模型 开始：用 checkpoint.model_args 构造相同 GPT，加载权重，切到 eval 并放到指定设备；根据 checkpoint.config.dataset 找到 meta.pkl。提示词 ROMEO: 被编码为形状 $(1,6)$ 的 long tensor，再传给 GPT.generate。这条路径没有目标 Y，不计算训练 loss，也不调用 backward 或 optimizer。

generate 核心片段：裁取上下文、采样下一个 token，再追加到输入末尾。

for _ in range(max_new_tokens):
    idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
    logits, _ = self(idx_cond)
    logits = logits[:, -1, :] / temperature
    if top_k is not None:
        v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
        logits[logits < v[:, [-1]]] = -float('Inf')
    probs = F.softmax(logits, dim=-1)
    idx_next = torch.multinomial(probs, num_samples=1)
    idx = torch.cat((idx, idx_next), dim=1)
return idx

每轮只取最后位置的分数，除以正的 temperature；top_k 可将候选限制在高分字符中，然后 softmax 得到概率，multinomial 随机抽一个编号，cat 追加到输入末尾。要求生成 40 个新字符时，最终序列长度为 $6+40=46$；原版没有自动按结束标记停下的逻辑。相同提示词在不同随机状态下可能生成不同续写。

当输入超过 block_size=32，generate 只保留最近 32 个 token 作为下一次模型输入，返回值仍保留原提示词与完整生成结果。原版每次重新计算整个上下文，没有 KV cache。训练时一次并行预测窗口内所有位置，生成时必须先选出下一 token 才能继续，这解释了两条路径的速度与计算方式为何不同。

作为对照，带 KV cache 的推理通常分为 prefill 与 decode：prefill 一次处理提示词的 $(B,T,C)$，后续 decode 在每层只处理新增位置的 $(B,1,C)$，并读取之前保存的 K、V。标准 MHA 中每层缓存可写为 K/V 各 $(B,H,T,D)$，新增一个位置后 T 增一；通常不需要缓存历史 Q。缓存增加了跨生成步骤保留的状态，与训练中为 backward 暂存激活的目的不同。

prefill 也属于推理，却可能具有较大的矩阵计算；小 batch decode 则往往权重复用较少，并持续读取 KV。因此要用 2.4 的形状、计算量和字节数判断瓶颈，不能只按“训练/推理”二分。这里介绍通用优化思路；本章固定版 nanoGPT 的 generate 没有实现 KV cache，不能直接按缓存 decode 的成本解释该代码。

3.8 本章通关：能运行，也能解释一次训练

第一项改造：给原训练循环增加观察点。运行补充包的生成器，在同一仓库生成 train_observe.py、training_observer.py 和 train_observe.patch。它校验固定版 train.py 的内容；打开 patch，可以逐处看到配置、前向、反向、更新和清梯度位置新增了什么。默认模式保留原训练与保存边界。

python make_observed_train.py --repo .
python train_observe.py config_smoke.py --out_dir=out-observe-cpu
# 有一张 GPU 时再运行同一实验：
CUDA_VISIBLE_DEVICES=0 python train_observe.py config_smoke.py --device=cuda --out_dir=out-observe-1gpu

观察对象是首个 Block 的 MLP 权重 transformer.h.0.mlp.c_fc.weight，其形状为 $(4C,C)=(256,64)$。输出 GUIDE 记录将这个参数与第一轮更新对应起来：

阶段	应看到什么
before_forward	该参数已有数值，.grad 为 None
after_forward	$X,Y:(4,32)$，$logits:(4,32,65)$，loss 是标量；参数未改变
after_backward	两次 backward 累积出有限梯度，参数仍未改变；global_mean_loss 是本次全部等大微批的平均
after_step → after_clear	一次 step 改变参数；清理后 .grad=None，并输出 state_check=PASS

然后仅把命令行 gradient_accumulation_steps 改为 1，使用另一个 out_dir 重跑，先预测 backward_calls 从 2 变为 1、每次更新的目标数从 256 变为 128，再核对日志。这一练习改变了更新批量，用来认识训练语义；固定输入下不同累积方式的数值等价另由 C.6 验证。观察单个参数通过，不代表整个模型已被全面验证。

python train_observe.py config_smoke.py --gradient_accumulation_steps=1 --out_dir=out-observe-a1

再建立单卡测量基线。下面先完成 5 次预热，随后测 10 次完整更新；单卡保持配置的每次更新 256 个目标 token，总计测量 2560 个。GUIDE benchmark_result 给出 measured_updates、global_tokens、slowest_rank_seconds 和 global_tokens_per_second；benchmark_rank 给出本 rank 的 peak_allocated_bytes。

CUDA_VISIBLE_DEVICES=0 python train_observe.py config_smoke.py --device=cuda --out_dir=out-bench-1gpu-g2 --observe_first_update=False --benchmark=True --benchmark_warmup=5 --benchmark_updates=10 --max_iters=20

计时窗前后同步 CUDA，窗口内跳过验证、保存和原逐步日志，统计取样/预取、前向、反向、更新与清梯度。关闭窗口内验证会改变后续随机采样轨迹，因此这类运行用于测量，不用它的 checkpoint 证明与普通训练逐值一致。CPU 可用 --device=cpu 检查工具，显存字段为 null；不将 CPU 时间当作 GPU 性能。

用第 3.3 节命令得到 checkpoint、生成文本并恢复训练后，尝试不看文件清单，按顺序解释：X/Y 为什么错开一位；logits 为什么是 $(B,T,V)$；A 次 backward 为什么只更新一次；验证为什么不改变参数；恢复训练为什么还需要优化器与随机状态；生成为什么只使用最后位置。能把这六个问题定位到源码，并能完成观察与批量改动实验，就具备了进入第 4 章的单设备基线。

原小配置的已有 CPU 实跑记录中，验证 loss 从 step 0 的约 4.1699 降到 step 20 的约 3.2548，保存后能生成 40 个字符，也能恢复训练；模型因果性、权重共享及最后位置 logits 检查通过。这是链路跑通的证据，短训练与少量验证批次不能据此评价语言能力。接下来第 4 章只增加数据并行，第 5 章单独研究状态分片，第 6 章再增加节点；TP、PP、CP 的进一步分工在第 7 章及对应附录展开。

4. 两卡 DDP：让同一个 nanoGPT 协同处理数据

进入条件：完成第 3 章 nanoGPT 的单卡训练、恢复与生成。本章继续用同一个固定版本、config_smoke.py 和字符数据，在一台机器的两张 GPU 上只增加数据并行。完成后应能解释进程与样本分工、定位梯度同步、修改有效 batch，并验证一次更新与单卡参考是否一致。

DDP 要解决的是完整模型能放入单卡时，如何让多张卡协作处理数据。输入仍是 $(B,T)$ 的字符窗口，输出仍是 logits/loss；每卡保存完整参数，各自前向反向，通过梯度平均让副本保持同样的更新。先验证这个语义，再看吞吐；显存分片属于第 5 章，跨机运行属于第 6 章。

本章统一 $B=4$、$T=32$，nanoGPT 命令行累加值 $G=2$。总进程数 $W=2$ 时每卡累加 $A=\frac{G}{W}=1$，一次更新含 256 个预测目标；单卡使用 $A=2$ 处理同样规模。初始化权重、实际样本、损失权重也要对齐，才能做逐值正确性比较。

4.1 分布式的运行方式：torchrun 与进程模型

分布式训练的核心运行模型是 SPMD（单程序多数据）：同一份脚本被 torchrun 在每张卡上各起一个进程同时运行，每个进程通过环境变量知道「我是第几个」，从而分工。这和第二部分的单进程完全不同，是首先要建立的认知。

打开 SVG 原图

把下方完整代码保存为 setup_dist.py，或使用第 3 章补充包中的同名文件，然后在两张可用 GPU 上运行：

torchrun --standalone --nproc_per_node=2 setup_dist.py

import os, torch, torch.distributed as dist

def setup():
    dist.init_process_group("nccl")            # GPU 间通信用 NCCL 后端
    local_rank = int(os.environ["LOCAL_RANK"]) # 本节点内的第几张卡
    torch.cuda.set_device(local_rank)          # 绑定当前进程到这张卡
    return local_rank

local_rank = setup()
rank = dist.get_rank()          # 全局第几个进程（跨所有节点）
world_size = dist.get_world_size()  # 总进程数 = 总卡数
is_main = rank == 0             # 只在主进程打印日志/存 checkpoint
print(f"rank={rank}, local_rank={local_rank}, world_size={world_size}", flush=True)
dist.destroy_process_group()

每个进程开头的标准初始化

常见一进程一卡配置中，local_rank 是本机进程编号（用于选择 GPU），rank 是全局进程编号，world_size 是总进程数。普通日志和 DDP 全量模型权重通常只由 rank 0 写；DCP 等分布式 checkpoint 的集合操作需要所有 rank 参与，不能照搬 rank 0 条件。

4.2 数据分工与梯度平均：一次更新保持什么不变

DDP 包装模型并在 backward 中归约梯度，数据怎样分配由训练器决定。nanoGPT 让各 rank 用不同采样随机数，从同一 token 流独立取窗口，不使用 DistributedSampler，也不保证窗口不重叠。下图关注完整模型副本与梯度同步；若要单独研究按索引分配样本，见 C.3。

打开 SVG 原图

这一节使用仓库已有 train.py。torchrun 给每个进程设置 RANK、LOCAL_RANK 和 WORLD_SIZE；代码据此初始化通信组、绑定本地 GPU，再把模型放进 DDP。每卡持有完整参数、梯度和优化器状态，各自完成前向/反向，通过梯度同步保持更新一致。原版 DDP 路径会强制选择 CUDA 设备，不能仅把 backend 改为 gloo 就视为 CPU 多进程示例。

设 $\texttt{world\_size}=W$，命令行中的梯度累加值为 G。原版要求 G 能被 W 整除，并把每卡累加次数改为 $A=\frac{G}{W}$。全局每次更新的窗口数是 $B\times A\times W=B\times G$，token 数再乘 T。这里 $B=4$、$T=32$、$G=2$：单卡 $A=2$，两卡各 $A=1$，两者每次更新都是 256 个 token。保持 G 不变增加卡数时，累加次数会减小；若想保持每卡 A 不变，则要随 W 一起调大 G。

原配置 train_shakespeare_char.py 的 $G=1$，直接 torchrun 两卡会触发整除断言，这就是示例显式覆盖 G 的原因。数据方面，各 rank 用 $1337+\mathrm{rank}$ 的随机种子，从同一 token 流独立采窗口；它不是 DistributedSampler，也不保证不同卡的窗口互不重叠。固定有效 batch 只说明梯度平均的规模可比较，不保证两种运行抽到了同样数据。

原版只在最后一个 micro-step 开启 require_backward_grad_sync，前面的反向先在本卡累积。自行改写时优先使用公开的 DDP.no_sync()，并把对应前向和反向都放在上下文中。rank 0 负责日志、验证和保存；原版验证结果是 rank 0 抽样批次的估计。迁移到分片模型后，必须重新设计这些职责，不能照搬“验证只让 rank 0 前向”。

先使用已有 DDP，再改一处同步行为。原版 train.py 已有 DDP 路径，使用 torchrun 不等于亲手完成全部改造。先沿下表在源码定位接入点，再在 train_observe.py 副本中将私有同步标志改为公开 no_sync 上下文；生成器已给出可读 patch，读者可以对照复现。

接入点	相对于单卡增加什么	如何核对
进程与设备	读取 RANK/LOCAL_RANK/WORLD_SIZE，初始化进程组，绑定当前 GPU	两条 rank 日志应分别为 0/1，本地设备不同
参数和数据	DDP 初始化同步参数；各 rank 独立取随机窗口	区分相同初始参数与不同采样；随机窗口允许重叠
模型包装	将同一个 GPT 放进 DDP，forward/loss 定义保持一致	模型 shape 不变，梯度同步由 DDP 反向触发
累积与同步	前几个微批只在本地累积，最后一个微批同步	no_sync 必须包住前向和反向；核对实际累积次数
验证和保存	普通日志与完整 checkpoint 由 rank 0 负责	这是本章原版 DDP 的职责，不能照搬给分片模型

sync_now = ddp_sync_each_micro or micro_step == gradient_accumulation_steps - 1
sync_context = model.no_sync() if ddp and not sync_now else nullcontext()
with sync_context:
    with ctx:
        logits, loss = model(X, Y)
        loss = loss / gradient_accumulation_steps
    X, Y = get_batch('train')
    scaler.scale(loss).backward()

这里展示同步改造的核心片段，完整可执行入口由 make_observed_train.py 生成；首次运行必须先完成第 3 章的生成步骤。参数更新仍在微批循环结束后执行。

4.3 实验：两卡短训练、保存、恢复与生成

# 在有两张可用 CUDA GPU 的机器运行；本章未做硬件实测。
torchrun --standalone --nproc_per_node=2 train.py config_smoke.py --device=cuda --dtype=float32 --out_dir=out-guide-ddp --gradient_accumulation_steps=2

2 卡 DDP：复用起步配置

# 仍在固定版 nanoGPT 根目录，使用同一份 config_smoke.py。
torchrun --standalone --nproc_per_node=2 train.py config_smoke.py --device=cuda --dtype=float32 --out_dir=out-guide-ddp --gradient_accumulation_steps=2 --init_from=resume --max_iters=30
python sample.py --out_dir=out-guide-ddp --device=cpu --dtype=float32 --num_samples=1 --max_new_tokens=40 --start=ROMEO:

短训练完成并出现 checkpoint 后，全部 rank 一起恢复

预期两个进程完成更新、rank 0 输出日志并保存 ckpt.pt；恢复后核对迭代号与优化器状态，重新启动 sample.py 应能读取同一词表生成。沿用第 3 章解释的保存时点与 max_iters 边界，不能把最近 checkpoint 当成退出前内存里的最终参数。随机窗口和 RNG 未完整恢复，因此这只证明能够续训，不证明与连续训练逐步一致。

改造实验：只改变梯度同步范围。两卡设置命令行 $G=4$，实际每卡 $A=2$，每次全局更新 512 个目标 token。ddp_sync_each_micro=False 只在最后一个微批同步，True 让每个微批都同步。两者都累积两批后更新一次；预期更新语义相同，但通信机会不同，速度需实测。若仍用 $G=2$，每卡只有一个微批，这个对比没有意义。

torchrun --standalone --nproc_per_node=2 train_observe.py config_smoke.py --device=cuda --gradient_accumulation_steps=4 --out_dir=out-sync-last --ddp_sync_each_micro=False --max_iters=1 --eval_interval=1
torchrun --standalone --nproc_per_node=2 train_observe.py config_smoke.py --device=cuda --gradient_accumulation_steps=4 --out_dir=out-sync-every --ddp_sync_each_micro=True --max_iters=1 --eval_interval=1
python inspect_checkpoint.py out-sync-last/ckpt.pt --compare out-sync-every/ckpt.pt

原版保存边界下，这两个 checkpoint 都应为 iter_num=1，即第一次已完成更新；比较器应报告 model_comparison.passed=true。它检查保存步数、结构和参数，不是完整状态恢复验收。再运行 4.4 的固定输入参考对照，排除仅比较两份相似错误实现的可能。发生差异时先检查实际输入和归一化，不能先扩大容差。

4.4 正确性实验：与同一批数据的单卡更新比较

保存 C.6 的 verify_ddp_step.py 到固定版 nanoGPT 根目录。脚本使用同一份 GPT 权重和固定 token 窗口：单卡参考一次处理全部窗口，两卡分别处理各自窗口并平均梯度；逐项比较 logits、全局平均 loss、参数梯度和一次 AdamW 更新。它验证 DDP 与累积的数学语义，不代替原版随机采样、保存和恢复的集成验收。

torchrun --standalone --nproc_per_node=2 verify_ddp_step.py --device cuda --accum 1
torchrun --standalone --nproc_per_node=2 verify_ddp_step.py --device cuda --accum 2

分别验证每卡累积 1 次和 2 次；FP32、dropout=0

第一轮每次更新是 $2\times4\times32=256$ 个 token；第二轮是 512 个。期望 passed=true，并记录每类最大绝对误差；初始容差为 atol=5e-6、rtol=1e-4。若失败，先查权重/样本切片、loss 是否多除或少除、no_sync 是否包住前向和反向，再查精度与内核差异。不要先放宽容差。CPU/Gloo 可用 --device cpu 检查机制，CUDA/NCCL 仍需在目标机器验证。

4.5 性能实验：固定工作量，再判断两卡是否有收益

只在正确性检查通过后测性能。沿用 3.8 的观察入口，但关闭首步状态记录；下面两条命令都固定 $B=4, T=32, G=4$，单卡 $A=4$，两卡各 $A=2$。先预热 5 次，再测 10 次更新，每次全局 512 个目标 token。

CUDA_VISIBLE_DEVICES=0 python train_observe.py config_smoke.py --device=cuda --gradient_accumulation_steps=4 --out_dir=out-bench-1gpu-g4 --observe_first_update=False --benchmark=True --benchmark_warmup=5 --benchmark_updates=10 --max_iters=20
torchrun --standalone --nproc_per_node=2 train_observe.py config_smoke.py --device=cuda --gradient_accumulation_steps=4 --out_dir=out-bench-2gpu-g4 --observe_first_update=False --benchmark=True --benchmark_warmup=5 --benchmark_updates=10 --max_iters=20

benchmark_result 的 global_tokens 应为 $10\times512=5120$；耗时取各 rank 最大值，$\text{吞吐}=5120/\text{最大秒数}$。两卡加速比为两卡吞吐除以单卡吞吐，效率再除以 2；峰值读取各 rank 的 peak_allocated_bytes 后取最大值。这个值是测量区间的绝对 allocated 峰值，不是显存增量或设备总占用。计时范围沿用 3.8；比较同步策略的速度时再只改变 ddp_sync_each_micro，不能同时改变批量。

先算同步的规模，再看 trace。设完整参数梯度缓冲大小为 S 字节，D 为数据并行副本数；均匀环式 AllReduce 的理想每 rank 发送量为 $\frac{2(D-1)S}{D}$，接收量与之相同但这里不重复相加。S 是梯度张量的字节数，不是 global batch；只统计这次参与同步的梯度，精度与未使用参数都会影响它。实际 NCCL 可能选择其他算法、协议与分桶，这个公式是核对量级的模型。

两卡例子：$D=2$	单卡	两卡 DDP
同一次全局更新的数据	处理全部样本	每卡处理一半独立样本
参数、梯度和普通 Adam 状态	完整一份	每卡仍完整一份，状态不因 $D=2$ 减半
若梯度共 $S=4\,\mathrm{MB}$（十进制）	无跨 rank 梯度同步	理想环式 AllReduce 每卡发送约 4 MB；接收另计
固定全局 batch 后的时间预期	计算量较大，无该项同步	每卡计算减少，同步与启动开销不一定同比减少

进一步可写 $t_{\mathrm{comm}}\approx2(D-1)\alpha+\frac{2(D-1)S}{D\beta}$：$\alpha$ 是环中每步的启动延迟（秒），$\beta$ 是有效传输速率（byte/s）。分桶会改变启动次数、消息大小和与 backward 的重叠程度。观察 trace 中未被计算遮住的通信尾部，而不是把全部 NCCL 时长直接加到计算时间。本文固定 global batch 的比较属于强扩展；若改成每卡 batch 不变，应另记全局工作量。

小模型可能因通信与调度开销而变慢，这仍是有效结论。检查 trace 中梯度通信与计算重叠、数据读取等待和不同 rank 的耗时；一次只改一个因素。完成标志是能解释数据与状态分工、通过一次更新与恢复验收，并用同口径记录说明收益或瓶颈。之后可直接进入多机 DDP；需要分片状态时进入第 5 章。

5. FSDP2：分片 nanoGPT 训练状态并可靠恢复

进入条件：理解第 4 章的数据并行与有效 batch，并已有可对照的单卡小模型。完成本章后，应能解释参数从分片到临时聚合再到重分片的生命周期，选择 Block 边界，接入全 rank 验证与 DCP 保存/恢复，并分开检查数值正确性和显存收益。仍使用第 3 章的 nanoGPT 和字符数据；不新增另一套主线模型。

5.1 为什么要分片：参数、梯度与优化器状态的生命周期

DDP 每卡持有完整训练状态，增加卡数并不会自动减小这部分占用。FSDP2 让同一数据并行组分担状态，并在计算当前模块时临时聚合所需参数。先用单卡也放得下的小模型观察这一变化，便于与完整模型作对照；之后才考虑必须靠分片才能容纳的规模。

FSDP2 将参数、梯度和优化器状态分布到多个 rank；计算一个分片单元时，通过 all-gather 临时恢复所需参数。是否在前向后立即重新分片受 reshard_after_forward 等配置影响；反向通常还需要参数聚合，并通过 reduce-scatter 汇总、分片梯度。峰值显存还包括激活、临时完整参数与通信缓冲，不能把整个训练显存简单除以卡数。

把 1.7 的状态账本扩大到 P 个参数：暂定参数、梯度、Adam 的 m/v 全为 FP32，无额外主权重，D 个 rank 均匀分片。下表只比较“参数梯度已形成时”的四类逻辑数据量，使用十进制 MB；激活、临时聚合、通信工作区和分配器未计入，不是实际峰值显存。

状态分配方式	每卡四类状态的理想合计	$P=1{,}000{,}000$，$D=2$
DDP	$$16P\,\mathrm{B}$$	16 MB
ZeRO-1：只分片优化器状态	$$\left(8P+\frac{8P}{D}\right)\mathrm{B}$$	12 MB
ZeRO-2：再分片梯度	$$\left(4P+\frac{12P}{D}\right)\mathrm{B}$$	10 MB
ZeRO-3 / 全分片思想	$$\frac{16P}{D}\,\mathrm{B}$$	8 MB

FSDP2 在本章采用全分片方式，可用最后一行理解状态量级，但它不是调用 DeepSpeed ZeRO 的实现。当前 Block 计算时还需要临时完整参数：若该单元有 U 个 FP32 参数，完整参数大小为 $4U\,\mathrm{B}$；不能直接在 $\frac{16P}{D}$ 之外无条件再加 $4U$，因为是否复用本地存储、保留哪些缓冲取决于实现，预取还可能让多个单元同时在场。用第 5.5 节的峰值统计检验分片粒度的影响。

分片也改变通信账本。对一个完整参数大小为 S 字节的单元，在均匀环式模型中一次 AllGather 每 rank 发送约 $\frac{(D-1)S}{D}$；同字节量梯度的 ReduceScatter 也约为 $\frac{(D-1)S}{D}$。若前向后重新分片，并在反向再聚合，则本单元一次前后向可有两次参数 AllGather 和一次梯度 ReduceScatter。混合精度、预取、重计算、重新分片策略或通信算法变化时要按实际调用重新统计。

本教程以 PyTorch 2.9.x 的 FSDP2 fully_shard API 为基线。FSDP2 使用 DTensor 表示分片参数；历史工程仍可能使用 FSDP1，迁移时按版本文档核对初始化、state_dict 与混合精度行为。

打开 SVG 原图

小模型通过后，再逐项引入 bf16、激活重计算和更大的模型。若完整模型初始化就超出单卡显存，需要另行设计 meta device 建模、分片、物化或分布式加载，以及共享权重和 checkpoint 语义；附件入口尚未实现这一能力。论文背景见 8.4，八进程 DeviceMesh 示意见 B.7。

以一个 Transformer Block 作为分片单元，前向前聚合本 Block 参数；配置 reshard_after_forward=True 时，前向结束后释放对应完整参数，反向需要时再聚合。反向得到的梯度归约并分片，优化器更新本地状态。分片粒度决定临时参数峰值、调用次数和通信大小；激活重计算另行解决激活保存压力，应在基本正确性通过后单独评估。

5.2 改造与运行：定义分片边界，再训练和验证

FSDP2 适合完整训练状态放不下单卡的情形。先把 nanoGPT 中哪些参数由哪个分片组管理讲清楚：对 transformer.h 的每个 Block 调用 fully_shard，再对整个 GPT 调用 fully_shard。前者提供逐 Block 的聚合/重分片边界，后者管理尚未分组的 Embedding、末尾归一化和输出层。wte 与 lm_head 共用参数，应留在同一个根组里。

from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard

# 所有 rank 已初始化 NCCL，绑定自己的 CUDA 设备。
# 使用相同模型初始化种子；另设按 rank 区分的数据采样 Generator。
model = GPT(gpt_config).to(device)
mesh = init_device_mesh("cuda", (world_size,))
for block in model.transformer.h:
    fully_shard(block, mesh=mesh, reshard_after_forward=True)
fully_shard(model, mesh=mesh, reshard_after_forward=True)
# 优化器必须在 fully_shard 之后创建；此路径不再套 DDP。
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

FSDP2 核心接入点：结构示意，完整入口见附件 train_fsdp2.py

第 3 章补充包包含完整 train_fsdp2.py，内容与原附件入口一致；复制到固定版 nanoGPT 根目录后继续导入 model.py。它使用 FP32、dropout=0、每卡每次更新一个微批，提供全 rank 验证、DCP 保存/恢复和可选生成导出。当前未接入编译、混合精度、梯度裁剪或累积，且先在每卡构造完整小模型；因此先用小配置验收，性能比较前再对齐训练设置。

改造点	在 nanoGPT 中如何落实
模型初始化与优化器	先绑定设备，以相同种子构建相同参数；逐 Block 分片再分片 root，最后创建 AdamW。FSDP2 不能靠 DDP 的初始广播替它同步随机初始化。
训练与数据	各 rank 各自采 batch，所有 rank 执行前向/反向；区分模型初始化随机数与数据采样随机数。新增累加时要用 FSDP2 的同步控制，不复用 DDP 标志。
验证	所有 rank 都进入模型前向，以相同调用次数参与参数通信；汇总 loss 后只让 rank 0 打印。变长或带 ignore 标签的数据按有效 token 数加权。
恢复训练	全部 rank 参与 DCP 保存/加载模型和优化器；额外保留完成更新数、配置以及各 rank 的 RNG/采样器状态。
生成导出	DCP 分片目录不能直接交给原版 sample.py；全 rank 参与聚合，rank 0 导出普通 ckpt.pt，保留 model_args、dataset 和对应 meta.pkl。

# 将附件 train_fsdp2.py 放到 nanoGPT 根目录；需要 CUDA/NCCL。
torchrun --standalone --nproc_per_node=2 train_fsdp2.py --steps 20 --out-dir out-fsdp2 --export
torchrun --standalone --nproc_per_node=2 train_fsdp2.py --steps 30 --out-dir out-fsdp2 --resume out-fsdp2/step-20 --export
python sample.py --out_dir=out-fsdp2/step-30/sample --device=cpu --dtype=float32 --num_samples=1 --max_new_tokens=40 --start=ROMEO:

FSDP2 目标机器验证命令：本章提供代码，尚未在多卡硬件执行

这个入口使用 argparse，参数写法为 --steps 20；原版 configurator.py 使用 --max_iters=20，两者不要混淆。steps 在新增入口里表示已完成的参数更新次数，和原版 max_iters 的边界语义不同。每次保存在新的 step-N 目录，completed.json 标识该次保存已完成。恢复示例要求相同 world size、模型和关键运行配置；虽然 DCP 可以支持更广的重分片加载，本例没有把弹性恢复与可复现性混在一起。多机时这些路径应在所有 rank 可访问的共享存储上。

导出完整模型仍需要足够的 rank 0 CPU 内存。新增检查器已完成语法、PyTorch 2.9.1 API 核对及 CPU DCP 比较器正负例测试；CUDA/NCCL 一次更新、真实分布式保存恢复与导出尚需按 5.4 在目标机器实测。先完成这些检查，再扩大模型。官方接口参考：fully_shard 与 Distributed Checkpoint。

5.3 保存与恢复：所有 rank 协作完成状态转换

分布式 checkpoint 应保存各 rank 的分片状态，避免为保存模型而把完整状态集中到一张 GPU。DCP 的 save/load 是所有参与 rank 协作的操作；还应保存模型配置、词表、训练步数，以及使用中的 scheduler/缩放器状态。下面展示模型和优化器状态的核心调用，需要在进程组销毁之前执行。

# 前置：model、optimizer 来自 5.2 的 FSDP2 训练脚本（分布式进程内）
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

# 保存：先取得适配分布式模型与优化器的 state dict，各 rank 并行写入分片
model_sd, optim_sd = get_state_dict(model, optimizer)
dcp.save(
    {"model": model_sd, "optim": optim_sd},
    checkpoint_id="ckpt/step1000",
)

# 加载：先用相同结构创建 model 与 optimizer，再准备占位 state dict
model_sd, optim_sd = get_state_dict(model, optimizer)
state = {"model": model_sd, "optim": optim_sd}
dcp.load(state, checkpoint_id="ckpt/step1000")
set_state_dict(
    model,
    optimizer,
    model_state_dict=state["model"],
    optim_state_dict=state["optim"],
)

FSDP2 + DCP：保存和加载分片状态

5.4 正确性验收：一次更新、恢复与导出分别验证

第一项：固定输入的一次更新。补充包的 verify_fsdp2_step.py 在每卡保留一份未分片参考模型，比较相同初始权重、固定窗口下的 logits、全局平均 loss、完整梯度和一次 AdamW 更新。所有 rank 按同一参数顺序参与 full_tensor 聚合；这份小模型检查器用于正确性，不能作为显存节省或吞吐基准。

torchrun --standalone --nproc_per_node=2 verify_fsdp2_step.py

预期 passed=true，输出四类 max_abs_error 和参数 global/local shape。默认 $b=4, T=32, A=1, D=2$，一次全局更新使用 8 个窗口、256 个目标 token；初始容差为 $\mathrm{atol}=5\times10^{-6},\ \mathrm{rtol}=10^{-4}$。若失败，先核对分片组、权重、样本和 loss 归一化，再排查数值实现。

对齐条件不能只写“相同 seed”：附件固定 bias=False、FP32、dropout=0；AdamW 的 lr=1e-3、betas=(0.9,0.95)，二维及以上参数 weight_decay=0.1，其余为 0，关闭 fused/foreach、裁剪、编译和累加。原版 nanoGPT 默认 bias 与训练设置可能不同。先导入同一份完整初始权重，再固定每个 rank 实际使用的窗口；用全 rank 协作取得的完整参数/梯度与未分片参考按名称比较。4.4 的 DDP 通过不等于 FSDP2 已通过。

第二项：连续训练与中断恢复。在两个新的输出目录中分别运行“连续到 30”和“到 20 保存、再恢复到 30”；相同 world size、模型、batch 和评估频率，两个最终运行都开启 --export。下面的比较器在单进程 CPU 上重建 DCP，逐项比较模型、优化器、各 rank RNG、导出权重和固定输入 logits。

torchrun --standalone --nproc_per_node=2 train_fsdp2.py --steps 30 --out-dir out-fsdp2-continuous --export
torchrun --standalone --nproc_per_node=2 train_fsdp2.py --steps 20 --out-dir out-fsdp2-split --export
torchrun --standalone --nproc_per_node=2 train_fsdp2.py --steps 30 --out-dir out-fsdp2-split --resume out-fsdp2-split/step-20 --export
python compare_fsdp2_runs.py --continuous out-fsdp2-continuous/step-30 --resumed out-fsdp2-split/step-30 --repo .

第三项：导出与生成。比较器同时核对 sample/ckpt.pt 与对应 DCP 的权重和输出，预期 passed=true、failures 为空，再执行下方采样命令。生成文本相似不能替代数值检查。脚本为检查完整状态需要足够的 CPU 内存；原训练器拒绝覆盖已有 step-N，重复实验应换新的输出目录。

python sample.py --out_dir=out-fsdp2-split/step-30/sample --device=cpu --dtype=float32 --num_samples=1 --max_new_tokens=40 --start=ROMEO:

补充的 CPU 模型机制检查可运行下面命令；它只检查未分片模型的因果性、绑权及数学边界，不承担上述 FSDP2 更新或恢复验收。

python verify_model.py --repo .

检查项	本次结果 / 下一步要求
原版数据准备	实际 tiny Shakespeare：65 字符，训练 1,003,854 token，验证 111,540 token。
CPU 短训练与生成	PyTorch 2.9.1、Python 3.12.14；验证 loss 约 $4.1699\rightarrow3.2548$；完成 checkpoint 读取及字符采样。
原版恢复训练	从保存的 step 20 继续到 step 30 并再次保存；未宣称与连续训练逐步一致。
模型语义	因果性、绑权、最后位置 logits 与生成形状检查通过。
TP / PP 数学边界	CPU 上 MLP 两分片输出/梯度对照通过；两逻辑 stage 输出/梯度对照通过。
FSDP2 扩展代码	训练入口和新增检查器完成语法、导入与 2.9 API 核对；CPU DCP 比较器正负例通过。CUDA/NCCL 一次更新、实际 DCP 恢复和导出仍需目标机器验证。
真实 DDP / TP / PP 与性能	本次未做多 GPU / 多机实测，不填写加速比和显存节省结论。

本阶段记录提交版本、PyTorch/CUDA/驱动、GPU 数与互联、B/T/G 或每卡累加次数、精度、编译开关、loss/梯度误差、恢复结果、峰值显存和 tokens/s。单机两卡通过后，再按第 6 章扩到两台机器；如果 FSDP2 尚未通过单机验收，先沿已验证的 DDP 路线完成跨机实验。

5.5 性能实验：状态分片后，峰值显存为何仍可能较高

预期变化是参数、梯度和优化器状态的常驻分片变小；训练峰值还包含激活、当前聚合参数与通信缓冲，所以它不一定按卡数等比例下降。先测相同小模型和有效 token 数，再分别改变分片粒度或激活重计算。若没有收益，用 profiler 区分计算、参数通信、数据读取和同步等待，独立示例见 C.5。

测吞吐前先预热，排除编译和验证/保存耗时，在明确的计时范围前后同步设备；用全局有效 token 数除以各 rank 中最大的耗时。记录每卡 max_memory_allocated 的峰值并取最大值，同时注明是否包含验证、checkpoint 和导出。只改一个因素，保持模型、序列长度、有效 batch、数据与精度一致；参数存储减少不等于训练峰值按卡数等比例下降。

容量实验：每次只改变本地 batch。补充包的 make_observed_fsdp2.py 校验原入口后生成 train_fsdp2_observe.py，在模型分片后打印参数 shape，在训练更新外增加同步、计时和峰值记录。在生成文件中搜索 fsdp2_parameter_shape 和 fsdp2_training_observation，定位参数形状与测量代码，再对照原入口的训练、RNG 和 DCP 逻辑。

python make_observed_fsdp2.py
torchrun --standalone --nproc_per_node=2 train_fsdp2_observe.py --steps 20 --batch-size 4 --out-dir out-fsdp2-observe-b4 --observe --warmup-steps 5
torchrun --standalone --nproc_per_node=2 train_fsdp2_observe.py --steps 20 --batch-size 2 --out-dir out-fsdp2-observe-b2 --observe --warmup-steps 5

查看各 rank 的 fsdp2_parameter_shape：同一参数的 global_shape 与 local_shape 应反映分片；减小 batch 不应改变该参数的分片布局。fsdp2_training_observation 报告 measured_updates_per_rank、global_tokens、seconds_max_rank_sum、tokens_per_second 和 peak_allocated_bytes_max_rank。前 5 次原有更新预热，其余 15 次测量，没有额外训练。

每个测量更新分别同步 CUDA、重置 allocated 峰值、执行取数与训练更新、同步并记录，再进入验证。耗时包含更新结束同步开销；排除验证、保存、导出操作本身，峰值包含更新开始时仍存活的参数、优化器和运行时缓冲。这是有观测开销的教学测量，与 4.5 连续窗口的计时方法不同，不能直接横比二者的 tokens/s。

预期 batch 减小后激活需求降低，但实际峰值可能被状态或缓冲主导，不保证明显下降。两次运行每更新目标数分别为 256 和 128，因而用于容量观察；要比较 DDP/FSDP2 吞吐，还应按 5.4 对齐所有训练设置并采用同一计时方法。

手段	效果	代价
torch.compile	可能通过图优化与 kernel 融合提速	首次编译慢；注意 graph break
SDPA / FlashAttention 后端	合适的融合后端可降低 attention 中间显存并提速	收益取决于设备、dtype、形状与可用后端
Activation checkpointing	激活显存大降	反向时重算部分前向操作，增加计算量
bf16 混合精度	转为 bf16 的张量每元素由 4 字节变为 2 字节，并可能提速；训练总显存不一定减半	精度较低；需验证 loss、收敛与算子支持
提高 DataLoader 吞吐	喂数据不卡 GPU	调 num_workers / 预取 / 预分词

调 OOM 用 torch.cuda.memory_summary() 和 memory snapshot（torch.cuda.memory._record_memory_history() 导出后可视化，能精确看到显存被什么占了）。判断是不是数据加载拖后腿：看 trace 里 GPU 是否频繁空闲等待 CPU。

本章完成标志：能指出一份参数何时是 DTensor 分片、何时为计算而聚合；固定输入下的一次更新、DCP 恢复和导出分别通过后，再提交显存与吞吐记录。若只完成静态 API 检查，应保留“待 CUDA/NCCL 验证”状态。多机 FSDP2 需先通过本章；多机 DDP 可从第 4 章直接进入第 6 章。

---

5.6 根据瓶颈选择下一步：更多机器，或另一种分工

理解分布式训练，可以先问三个问题：每张 GPU 保存什么、负责哪部分计算、需要和其他 GPU 交换什么。下面统一用两张 GPU 举例，对照数据、训练状态、单层矩阵、模型层和上下文位置这几种分工方式。除了看“切什么”，还要看计算时是否临时恢复完整参数，以及通信发生在哪里。

打开 SVG 原图

打开 SVG 原图

方式	各卡如何分工	主要通信	适用目的与边界
DDP：数据并行	各卡处理不同样本；每卡保存完整模型、梯度和普通优化器状态。	常规同步反向中，按梯度桶做 all-reduce；可与反向计算重叠。	完整训练配置能放入单卡，希望提高吞吐；收益取决于计算和通信比例。
FSDP2：全分片数据并行	各卡处理不同样本；参数、梯度和优化器状态按配置分片，计算当前单元时临时聚合所需参数。	参数 all-gather；反向可能再次 all-gather，梯度以 reduce-scatter 汇总并分片。重新分片与预取受配置影响。	降低训练状态占用；峰值仍含激活、临时参数和通信缓冲。详见 5.1—5.2。
TP：张量并行	同一层的权重与计算分给不同卡，例如按线性层输出维度分块。	依切法交换或归约中间结果，常见 all-reduce、all-gather 或 reduce-scatter。	单层计算或存储压力大时考虑；优先高速互联域，跨节点是否划算需实测。
PP：流水线并行	将不同层交给不同 stage；多个 micro-batch 按调度穿过各阶段。	前向传激活，反向传对应的梯度；需协调阶段负载与流水线气泡。	分散模型层的存储与计算，可跨设备或节点；效率依赖调度和负载均衡。
CP：上下文并行	同一序列的不同位置分给不同卡，分担长上下文计算与激活。	Attention 需要跨段信息；具体实现可交换 K/V、使用环式通信等，反向也需相应通信。	长上下文带来的激活与 Attention 压力较大时考虑。
SP：序列并行（此处指 Megatron 风格）	通常配合 TP，将部分激活及 LayerNorm、Dropout 等计算沿序列维分片。	常见实现通过 all-gather / reduce-scatter 衔接 TP 算子；与 CP 的分工范围不同。	减少部分激活的重复保存；不能直接与 CP 当作同一种机制。
EP：专家并行	MoE 层的不同专家放在不同卡上；token 按路由结果送到目标专家，同一 token 可选择多个专家。	典型实现先按目的 rank 做 AllToAll 分发，专家算完后再反向重分布并合并。	分担专家参数和计算；需要处理路由、负载不均与通信缓冲。它要求 MoE 架构，不能直接给 dense nanoGPT 打开一个开关。见 7.9。

这些方式可以组合，例如在不同并行维度上使用 FSDP2、TP 与 PP。建议先用 DDP 理解数据分配和梯度同步，再用 FSDP2 理解状态分片与参数聚合；这是学习顺序，实际训练不必逐级采用。随后依据峰值显存、吞吐和通信测量选择其他方式。图中 CP 展示上下文分工，SP 的含义则随框架有所差异，表中已注明本节采用的定义；更复杂的组合可结合 torchtitan 等成熟框架学习。

5.6.1 选择配置：先放得下，再对齐 batch，最后比较速度

Ultrascale Playbook 给出的有效方法，是把配置搜索分为三个有先后依赖的问题。先确定模型、完整序列长度、精度、目标 global batch 和可用硬件；这些是本轮实验条件。增加 GPU 或改变并行方式之前，先说明当前卡在容量、计算还是通信上遇到了什么限制。

当前问题	怎样行动	进入下一步需要的证据
完整训练步能否放入显存？	至少运行到 optimizer.step，使延迟创建的 Adam 状态也出现，再测后续稳定训练步。激活过大时先减小 micro-batch、考虑重计算；状态过大时考虑 FSDP/ZeRO；单层、长上下文或专家压力再分别考虑 TP、CP、EP。	成功完成前向、反向和更新，记录各 rank 峰值及余量。仅加载模型或跑通 forward 不算通过。
是否达到目标 global batch？	用真正处理不同样本的 DP 副本数、micro-batch 和累积次数重新计数；保持完整序列长度一致。TP/PP/CP/EP 度数不能直接当作额外 batch 倍数。	每次更新的样本数与有效 token 数对齐；padding、mask 或变长样本需要按有效目标位置计数。
在相同训练任务下能否更快？	从可行配置中一次只改一项，先调整 micro-batch 与累积的分配，再比较并行组合和拓扑。复用对应章节的实验入口；跨 DDP/FSDP 入口比较前，统一计时窗口和同步方式。	正确性先通过，再比较预热后的端到端 tokens/s、峰值显存和最慢 rank；最后回到 trace 解释等待来自哪里。

先做一组不增加 GPU 的搜索。沿用第 4 章两卡 nanoGPT，固定 $D=W=2$、完整序列长度 $T=32$、每次更新 16 条序列，即 512 个有效目标 token。记每卡 micro-batch 为 $b$、实际累积次数为 $A$；约束为 $DAb=16$。本固定版本 nanoGPT 的命令行 gradient_accumulation_steps 记作 $G$，入口会除以 $W$，所以这里 $G=WA$。下表中的两个累积数字不能混用。

每卡 micro-batch $b$	每卡实际累积 $A$	nanoGPT CLI $G$	每次更新 token 数
1	8	16	512
2	4	8	512
4	2	4	512
8	1	2	512

在已生成 train_observe.py 的 nanoGPT 仓库中运行下列命令。它保持模型与 global batch 不变，只改变一次实际计算有多大。预计较大的 micro-batch 可能提高算子利用率，但也会增加激活显存；哪个组合更快必须由结果决定。

for b in 1 2 4 8; do
  g=$((16 / b))
  torchrun --standalone --nproc_per_node=2 train_observe.py config_smoke.py \
    --device=cuda --block_size=32 --batch_size="$b" \
    --gradient_accumulation_steps="$g" --out_dir="out-scan-b$b-g$g" \
    --observe_first_update=False --benchmark=True \
    --benchmark_warmup=5 --benchmark_updates=10 --max_iters=20
done

每项记录 benchmark_result 中的 global_tokens（应为 5120）与吞吐；显存读取各 rank 的 benchmark_rank.peak_allocated_bytes 后取最大值，OOM 也要记为不可行配置。先前 4.4 的固定输入检查负责数值正确性；本扫描使用各次训练的实际取样与计时窗口，不能用不同配置的 checkpoint 逐值相同来判定通过。它也只测系统效率，不能证明改变 global batch 后的收敛质量。

这套搜索顺序借鉴 Ultrascale Playbook 的 Finding the Best Training Configuration。原文特定 GPU 数量下的推荐组合是其集群上的经验，不是本教程的通用阈值。CP 能让更长序列的训练成为可能；在完整序列长度不变时，它不会直接增加 global batch。

第 3—6 章共用下方原版附件：起步配置、CPU 验证脚本、FSDP2 教学入口及配套图。包内旧正文编号仅作历史记录，当前章节与运行步骤以在线正文为准；新的一次 DDP 更新对照代码见 C.6。

nanogpt-guide.zip

6. 多机多卡：从 2 台机器、每台 2 GPU 开始

进入本章前，先在一台机器上完成第 4 章的两卡 DDP；如果要走 FSDP2 路线，还应完成单机 FSDP2 的一次更新、DCP 保存、恢复与导出。多机阶段继续使用同一份 nanoGPT、数据和训练目标，先以两台机器、每台两张 GPU 验证跨机训练，再扩大到每台八张 GPU。这里提供待目标环境验证的运行步骤，不把命令模板当作已完成的多机实验。

6.1 进入多机：先固定环境、数据和进程编号

一台机器上，进程主要通过本机互联通信；增加第二台机器后，还要处理节点会合、跨机网络、文件可见性和节点故障。训练算法可以沿用单机版本，但每个节点必须运行同一份代码与配置，所有参与通信的 rank 必须按一致的顺序执行集合操作。

准备项	两台机器应满足的条件
代码与运行环境	nanoGPT 固定提交一致；Python、PyTorch、CUDA 构建及依赖一致，驱动与构建兼容；可见 GPU 数和所选精度符合配置。
数据与词表	train.bin、val.bin、meta.pkl 内容相同。相同路径名不代表相同数据，应对照文件大小与校验值；各节点从仓库根目录运行。
会合与训练网络	使用所有节点都能访问的节点 0 地址，会合端口空闲且可达。会合成功只代表控制连接建立；NCCL 使用的网卡、RDMA 和网络策略仍需单独验证。
输出与恢复	DDP 普通 checkpoint 由 rank 0 写入；恢复前所有 rank 都应能读到同一文件。DCP 的所有 rank 共同保存与加载，示例统一使用可读写的共享目录。

在每台机器仓库根目录运行补充包的数据与环境检查，比较输出中的源码、config_smoke.py、train.bin、val.bin 和 meta.pkl 校验值。提交与数据应一致；CUDA 可用性和可见卡数符合本次运行。生成 train_observe.py 后再做测量实验。

python check_inputs.py
python make_observed_train.py --repo .

先约定 $N=\text{节点数}$、$K=\text{每节点进程数}$，常见一进程一卡配置下 $W=N\times K$。2 台机器 × 每台 2 GPU时 $W=4$；各节点 LOCAL_RANK 都从 0 到 1，RANK 是全局编号。本章固定规模示例中，节点 0 的全局 rank 为 0—1，节点 1 为 2—3；代码始终读取 torchrun 设置的环境变量，不自行硬编码编号。

6.2 先验证跨机通信，再启动模型

这里先分清通信层次：torch.distributed 用进程组表达“哪些 rank 协作、做什么操作”；NCCL 在 NVIDIA GPU 上实现集合通信及点对点通信，并选择算法、协议与传输路径。单机可能使用 NVLink、PCIe 等互联，跨机通常还要经过网卡和网络；NCCL 与 RDMA 不是同一层的概念。

RDMA（远程直接内存访问）使网卡参与主机间内存数据传输，减少常规数据路径中的 CPU 拷贝与介入；它不代表完全没有 CPU 控制开销。InfiniBand 与 RoCE 是常见的 RDMA 网络环境。GPUDirect RDMA 则进一步允许符合条件的网卡访问 GPU 内存，依赖硬件、驱动和拓扑支持，不能仅凭“使用 NCCL”就断言已启用。

先辨认通信改变了什么数据。以下 $W$ 是当前通信组内的 rank 数，$r$ 为组内编号；$n$ 是每份分片的行数，$H$ 是每行的元素数。所有形状均指每个 rank 持有的张量。表中按相同形状、均匀切分说明，实际 API 还需满足对应后端的形状、类型与调用顺序要求。

操作	输入 → 输出与含义	在训练中补全什么
AllReduce（SUM）	各 rank 输入 $X_r\in\mathbb{R}^{n\times H}$；每个 rank 得到同形状的 $\sum_r X_r$。	合并梯度或同一运算的部分和。SUM 本身不除以组大小；DDP 的平均语义来自训练框架的缩放规则。
AllGather	各 rank 输入不同分片 $X_r\in\mathbb{R}^{n\times H}$；按 rank 顺序拼接后，每个 rank 得到 $(Wn,H)$ 的完整结果。	例如 FSDP 收集参数分片。它只收集，不求和，也不取平均。
ReduceScatter（SUM）	各 rank 提供 $(Wn,H)$ 输入；先逐元素求和，再沿首维切成 $W$ 份，rank $r$ 只得到自己的 $(n,H)$ 分片。	例如汇总梯度后只保留本 rank 负责的分片。它不是把某一个 rank 的输入简单切开。
AllToAll	每个 rank 把输入分成给各目的 rank 的块 $C_{r,j}\in\mathbb{R}^{n\times H}$；rank $j$ 收到来自所有源 rank 的 $C_{r,j}$，拼成 $(Wn,H)$。	按目的地重新分布数据，例如 EP 把不同 token 送到不同专家。没有求和；非均匀路由时，每对 rank 的块长可以不同。

AllReduce 的 SUM 结果可以用 ReduceScatter 后接 AllGather 来得到；这是理解环式算法与 FSDP 的桥梁，不代表每次 NCCL AllReduce 都必须采用同一种实现。PP 常用点对点 send/recv 传边界激活；CP 的环式 K/V 收发也不能仅因数据跨过多个 rank 就称为一次 AllToAll。参与同一 collective 的 rank 必须按相容顺序进入；某个 rank 跳过一次调用，其他 rank 就可能一直等待。

因此下面的 sum=6 只证明这次集合通信结果正确。后续还应从 NCCL 日志核对所选传输、网卡与拓扑，再在相同节点和消息尺寸下测有效带宽；小标量成功不证明大梯度缓冲传得快。将第 4.5 节的梯度字节数、7.4 的 TP 激活字节数带入第 6.6 节分析，才知道应该测多大的消息。

将下面脚本保存为两个节点仓库根目录中的 check_dist.py。它先打印主机、全局 rank、本地 GPU 和 world size，再做一次 NCCL all-reduce。2 台机器 × 每台 2 GPU时各 rank 输入自己的编号，求和结果应都为 $0+1+2+3=6$；这个小测试用于验证连接和基本集合通信，不代表网络已达到目标带宽。

import os, socket
from datetime import timedelta
import torch
import torch.distributed as dist

local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group("nccl", timeout=timedelta(seconds=120))
try:
    rank, world = dist.get_rank(), dist.get_world_size()
    value = torch.tensor(float(rank), device=f"cuda:{local_rank}")
    dist.all_reduce(value)
    expected = world * (world - 1) / 2
    print(socket.gethostname(), "rank=", rank, "local_rank=", local_rank,
          "world=", world, "sum=", value.item(), flush=True)
    assert value.item() == expected
finally:
    dist.destroy_process_group()

check_dist.py：四个 rank 都应输出 sum=6

以下命令分别在两台机器启动，两个终端都要运行。NODE0_ADDR 必须替换为实际可达地址，不能使用 localhost；本例采用固定节点数和明确的 node_rank。首次检查可设置 NCCL_DEBUG=INFO。不要预先随意设置 NCCL_IB_HCA 等环境变量，应根据机器实际网卡、RDMA 配置和日志确定。

# 节点 0；替换地址后执行
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
NCCL_DEBUG=INFO torchrun --nnodes=2 --nproc_per_node=2 --node_rank=0 --master_addr="$NODE0_ADDR" --master_port=29500 check_dist.py

# 节点 1；填写同一个地址
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
NCCL_DEBUG=INFO torchrun --nnodes=2 --nproc_per_node=2 --node_rank=1 --master_addr="$NODE0_ADDR" --master_port=29500 check_dist.py

两机 × 每机两卡：先运行通信检查

从“连得通”走到“传得怎样”。完成上面的标量求和后，运行附录 C.10 的 collective_probe.py：先检查四种操作的输出，再扫描 4 KiB、1 MiB、16 MiB、64 MiB 的 AllReduce。先单机两卡，再在两机上用同样的消息尺寸；比较时固定总 rank 数，避免把多了 GPU 和跨了网络混为一谈。

每条记录包含消息字节数、预热与测量次数、由最大 rank 窗口耗时换算的每次平均毫秒数。小消息常受启动延迟影响，大消息更容易体现传输能力；不应拿一次标量 all-reduce 的耗时预测完整模型同步。用 4.5 的实际梯度桶大小、7.4 的 TP 激活大小选择补充测点，再回到训练 trace 看未被计算遮住的通信尾部。

附录输出的“消息大小 ÷ 平均耗时”只是本次操作的有效载荷速率，不是 NIC 线速，也没有按 Ring 算法换算 bus bandwidth。隔离测得的通信时间也不能直接相加为训练时间：通信可能与计算重叠，两者还可能争用 SM、显存带宽和链路。是否改善性能，最终看同口径的完整训练步。

6.3 第一轮跨机 DDP：2 台机器 × 每台 2 GPU

通信测试通过后，继续使用第 3 章的 config_smoke.py 和第 4 章已跑通的原生 DDP 入口。先保持 FP32、dropout=0、compile=False；本例选 $B=4$、$T=32$、$G=4$。nanoGPT 会将每卡累加次数设为 $A=\frac{G}{W}$，因此四卡时 $A=1$，每次更新共处理 $B\times T\times G=512$ 个有效 token。

为了对比同一个有效 batch，先在一台机器上分别保存单卡 $G=4$、两卡 $G=4$ 的基线，再运行下面的四卡实验。第 3 章原来的 $G=2$ 对应 256 token；若直接沿用到四卡，不满足 G 能被 W 整除的要求。独立随机采样不会自动得到完全相同的数据；逐值正确性对照仍需使用固定输入和初始权重。

python train.py config_smoke.py --device=cuda --dtype=float32 --gradient_accumulation_steps=4 --out_dir=out-1gpu-g4
torchrun --standalone --nproc_per_node=2 train.py config_smoke.py --device=cuda --dtype=float32 --gradient_accumulation_steps=4 --out_dir=out-2gpu-g4

在节点 0 建立相同有效 batch 的单机基线

# 节点 0
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
RUN_DIR='/REPLACE_WITH_SHARED_DIR/nanogpt-2n2g-g4'
torchrun --nnodes=2 --nproc_per_node=2 --node_rank=0 --master_addr="$NODE0_ADDR" --master_port=29500 train.py config_smoke.py --device=cuda --dtype=float32 --gradient_accumulation_steps=4 --out_dir="$RUN_DIR"

# 节点 1
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
RUN_DIR='/REPLACE_WITH_SHARED_DIR/nanogpt-2n2g-g4'
torchrun --nnodes=2 --nproc_per_node=2 --node_rank=1 --master_addr="$NODE0_ADDR" --master_port=29500 train.py config_smoke.py --device=cuda --dtype=float32 --gradient_accumulation_steps=4 --out_dir="$RUN_DIR"

在两台机器分别启动 DDP；地址与共享目录须替换

预期看到四个训练进程持续运行、rank 0 输出训练与验证 loss，并在共享输出目录保存 ckpt.pt。沿用第 3 章的保存时机解释：短跑结束时的内存参数不一定等于最近 checkpoint。验证恢复时，所有节点使用相同启动参数，追加 --init_from=resume --max_iters=30；不要只重启 rank 0。原版未保存完整随机状态，因此这里只验证能接续训练，不声称与连续运行逐步一致。

跨机短训练通过后，可将下面的训练入口放到 6.2 两个节点各自的 torchrun 命令末尾，保留对应 node_rank 与相同会合地址。四个 rank 用固定输入比较单设备参考与 DDP，一次全局更新为 512 个目标 token；预期 passed=true。这一步单独验证跨机梯度更新，不替代随机数据训练与 checkpoint 集成验收。

verify_ddp_step.py --device cuda --accum 1

6.4 再扩展 FSDP2：跨机参数通信与共享 checkpoint

只在第 5.4 节全部单机 FSDP2 验收通过后扩到多机。仍使用同一 train_fsdp2.py，保持模型、精度与每卡批量；下面两条完整命令分别在两个节点运行，地址相同、node_rank 不同。输出必须位于所有 rank 可读写的新共享目录。

# 节点 0
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
FSDP_RUN_DIR='/REPLACE_WITH_SHARED_DIR/fsdp2-2n2g'
torchrun --nnodes=2 --nproc_per_node=2 --node_rank=0 --master_addr="$NODE0_ADDR" --master_port=29500 train_fsdp2.py --steps 20 --out-dir "$FSDP_RUN_DIR" --export
# 节点 1：使用相同地址与共享目录
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
FSDP_RUN_DIR='/REPLACE_WITH_SHARED_DIR/fsdp2-2n2g'
torchrun --nnodes=2 --nproc_per_node=2 --node_rank=1 --master_addr="$NODE0_ADDR" --master_port=29500 train_fsdp2.py --steps 20 --out-dir "$FSDP_RUN_DIR" --export
# 恢复时，两个节点都把末尾训练参数改为：
# --steps 30 --out-dir "$FSDP_RUN_DIR" --resume "$FSDP_RUN_DIR/step-20" --export

此入口每卡每次更新一个 batch，增加 world size 会增大全局有效 batch；它没有复用原版的 $\frac{G}{W}$ 换算。比较 DDP 与 FSDP2 前，按各自实际 B、T、A、W 重新对齐 token 数和优化器设置。跨机 FSDP2 的参数 all-gather 与梯度 reduce-scatter 经过节点间网络，显存收益和通信代价都需要实测。

全部 rank 参与训练、验证和 DCP 保存/加载，不能只让 rank 0 进入分片模型前向。先检查所有进程完成保存且 completed.json 存在，再用相同 world size 恢复，核对完成更新数、模型和优化器状态；按需加 --export，并在全 rank 聚合完成后由 rank 0 导出供 sample.py 使用。跨节点异常时，应先终止并清理本次相关训练进程，再从最后一个完整 checkpoint 重启固定规模任务。

附件 FSDP2 入口仍属于待多 GPU 验证的教学扩展，本次结构调整没有增加硬件实测结果。节点内分片、节点间复制的混合分片可在网络成为瓶颈后研究；它需要修改 DeviceMesh 与分片配置，不能通过增加 --nnodes 自动获得。

6.5 扩大规模：从总计 4 GPU 到 16 GPU

2 台机器 × 每台 2 GPU的训练和恢复通过后，再把每节点进程数改为 8，总 world size 变为 16。下面的图展示2 台机器 × 每台 8 GPU拓扑；启动命令继续复用固定版 nanoGPT。原版要求 G 能被 16 整除，本例使用 $G=16$，每卡累加一次。若要与较小规模公平比较，应先在那些规模上也使用 $G=16$ 重跑基线，不能把增加有效 batch 带来的变化当成单纯的扩卡收益。

打开 SVG 原图

# 节点 0
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
RUN_DIR='/REPLACE_WITH_SHARED_DIR/nanogpt-2n8g-g16'
torchrun --nnodes=2 --nproc_per_node=8 --node_rank=0 --master_addr="$NODE0_ADDR" --master_port=29500 train.py config_smoke.py --device=cuda --dtype=float32 --gradient_accumulation_steps=16 --out_dir="$RUN_DIR"

# 节点 1
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
RUN_DIR='/REPLACE_WITH_SHARED_DIR/nanogpt-2n8g-g16'
torchrun --nnodes=2 --nproc_per_node=8 --node_rank=1 --master_addr="$NODE0_ADDR" --master_port=29500 train.py config_smoke.py --device=cuda --dtype=float32 --gradient_accumulation_steps=16 --out_dir="$RUN_DIR"

两机 × 每机八卡：使用 G=16，对应 16 个 rank

多机常见坑	排查方向
卡在 init_process_group 不动	节点间端口不通 / MASTER_ADDR 填错 / 防火墙
通信极慢	核对 NCCL 是否使用预期 RDMA 路径或回退到 Socket/TCP，并检查网卡、拓扑与拥塞。RoCE 本身运行在以太网上；按实际配置核对 NCCL_SOCKET_IFNAME、NCCL_IB_HCA
NCCL 报错难懂	设 NCCL_DEBUG=INFO 看详细日志
各节点结果不一致	代码/数据/依赖版本不一致；确保各节点环境镜像相同

6.6 比较扩展效果：固定工作量，再看增加节点的收益

先隔离“增加节点”这个因素。单机有四张 GPU 时，用下面的一机四卡对照两机各两卡：两者都 $W=4, G=4, A=1$，每次全局更新 512 个目标 token。首次生成观察入口后，先完成同配置正确性检查，再运行测量。

# 一台机器的四张 GPU
torchrun --standalone --nproc_per_node=4 train_observe.py config_smoke.py --device=cuda --gradient_accumulation_steps=4 --out_dir=out-bench-1n4g --observe_first_update=False --benchmark=True --benchmark_warmup=5 --benchmark_updates=10 --max_iters=20
# 两机各两卡：以下分别在对应节点运行
# 节点 0
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
BENCH_RUN_DIR='/REPLACE_WITH_SHARED_DIR/bench-2n2g'
torchrun --nnodes=2 --nproc_per_node=2 --node_rank=0 --master_addr="$NODE0_ADDR" --master_port=29500 train_observe.py config_smoke.py --device=cuda --gradient_accumulation_steps=4 --out_dir="$BENCH_RUN_DIR" --observe_first_update=False --benchmark=True --benchmark_warmup=5 --benchmark_updates=10 --max_iters=20
# 节点 1
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
BENCH_RUN_DIR='/REPLACE_WITH_SHARED_DIR/bench-2n2g'
torchrun --nnodes=2 --nproc_per_node=2 --node_rank=1 --master_addr="$NODE0_ADDR" --master_port=29500 train_observe.py config_smoke.py --device=cuda --gradient_accumulation_steps=4 --out_dir="$BENCH_RUN_DIR" --observe_first_update=False --benchmark=True --benchmark_warmup=5 --benchmark_updates=10 --max_iters=20

两次 benchmark_result 的 measured_updates 都应为 10、global_tokens 都应为 5120；比较 slowest_rank_seconds 与各 rank 峰值，计时口径和 4.5 完全相同。设单卡吞吐为 $Q_1$、$W$ 卡吞吐为 $Q_W$，加速比为 $Q_W/Q_1$，效率为 $Q_W/(WQ_1)$。只有同型号设备、相同配置与测量范围时，才能把同卡数对比主要用于分析跨节点成本；没有单机四卡资源时，明确保留这个对照未完成。

每个实验使用相同提交、数据、模型、序列长度、有效 batch、精度和计时范围。预热后同步设备，统计全局有效 token 数，并用各 rank 中最大的耗时计算 tokens/s；记录每卡峰值显存并取最大值。验证、保存与首次编译是否计入耗时要写清楚。第 3 章的小模型适合检查流程，计算量过小时多卡反而可能更慢，不能据此推断大模型扩展能力。

实验	固定什么	重点观察
单机单卡 → 单机两卡	相同有效 token 数、初始权重、数据口径和精度	梯度与更新是否一致，本机通信开销和吞吐变化
单机四卡 → 2 台机器 × 每台 2 GPU	总卡数和训练工作量相同；尽量使用同型号 GPU	节点间通信、慢 rank、网络与数据读取等待
2 台机器 × 每台 2 GPU → 2 台机器 × 每台 8 GPU	先设统一 G，例如 16，再重跑各规模基线	每卡累加次数变化、全局吞吐和扩展效率

用同一张时间账本解释扩展结果：计算由每 rank 的实际张量 shape 决定，HBM 读写由算子与缓存复用决定，网络开销由消息字节数、调用频次、有效带宽和延迟决定。跨机先比较相同总卡数的一机四卡与两机各两卡；若 FLOPs 基本相同而时间增加，检查未被隐藏的网络通信、拓扑、慢 rank 与数据读取，不能直接归因于“GPU 算力不够”。

观察到的现象	先形成的假设	验证时保留的证据
小消息很多，每条耗时下降不明显	启动与同步延迟占比大	collective 类型、尺寸、次数及时间线
大消息耗时随字节数近似增长	传输带宽成为主要限制	同一组 rank、同尺寸的有效带宽，单位统一为 byte/s
NCCL 总时长很长，但 step 增幅较小	部分通信与计算重叠	关键路径上剩余的通信尾部，而非各事件简单求和
同组某个 rank 明显更慢	等待可能来自上游计算、数据或不均衡	各 rank 时间线、数据耗时及实际通信路径

网络标称速率常用 bit/s，账本通常用 byte/s，换算时先除以 8，再考虑协议、拓扑及实测有效利用率。RDMA 路径选择、NCCL 算法选择与应用层重叠是三个分别需要证据的问题。先验证单项，再讨论 TP 留在高速互联域、PP/DP 跨节点等放置；这些是依据消息规模作出的选择，不是适用于所有模型的固定拓扑。

6.7 本章通关：完成跨机训练、恢复与一次有效对比

两台机器上的所有 rank 都能完成通信检查，随后完成短训练、验证、完整 checkpoint 保存与恢复。能够解释 N、K、W、rank、local_rank 以及每次更新的有效 token 数；遇到卡住时能区分会合、NCCL、数据和集合调用不一致几类问题。最后记录至少一组固定工作量的单机与多机结果，说明瓶颈在哪里，再决定是否扩到更多节点或引入新的并行维度。

到这里，主线已完成从单卡到单机多卡、再到多机多卡的同一训练任务。希望深入通信与调度时，可读第 7 章 Picotron；它也允许在第 4 章后进入。第 8 章论文、第 9 章工程资源围绕实际问题选读。

7. 读懂 Picotron：从一次更新到组合并行

完成 nanoGPT 单卡闭环与 DDP 基础后，即可进入本章；不要求先完成 FSDP2 或多机训练。你应能解释一次更新中各 rank 持有什么数据和状态，定位通信与调度，修改一个并行维度并验证数学结果。跨机部分再结合第 6 章；Picotron 的依赖与运行配置始终放在独立环境。

7.1 先看全貌：训练器的输入、输出与职责

Picotron 是 Hugging Face 的简洁分布式训练教学项目，支持 Llama 风格模型的 Data、Tensor、Pipeline、Context 四种并行。项目明确强调教育和实验用途；主训练、模型和四种并行实现各约数百行，适合逐步阅读和修改。

输入是模型/并行配置与 token 数据；启动阶段建立模型参数、通信组、优化器和数据迭代器。每次更新依次取 micro-batch、算 logits 和 loss、反向形成梯度、按需要通信，再更新本 rank 的参数。输出是训练指标与 checkpoint；参数和优化器跨 step 保留，激活按反向需要保留/释放，梯度累积到本轮结束后清空。先在一张卡看清这条链，再引入协作。

前面的学习对象	主要回答的问题	在这里怎样衔接
nanoGPT	一个 GPT 怎样完成数据→前向→loss→反向→更新→保存	保留整条训练链路的理解，再观察并行后每个 rank 分到哪部分
PyTorch DDP / FSDP2	如何用框架 API 同步梯度、分片训练状态和恢复	继续作为工程 API 主线；用 Picotron 看更显式的通信与调度
Picotron	模型、数据、通信组与调度怎样组合成分布式训练	重点读实现，验证梯度和一次更新，再扩设备规模

这里的 4D 是 $\mathrm{DP}\times\mathrm{TP}\times\mathrm{PP}\times\mathrm{CP}$，不包含 FSDP/ZeRO；也不是 GPU、block、warp、thread 四层硬件结构。Picotron 的 DP 路径主要通过复制模型和梯度 all-reduce 工作，不能把它直接当成前面 FSDP2 参数分片实现的替代品。

本章源码基于 2026-09-24 检查到的 main 提交 59714b1bf25541dbeb6f5a721ed15f6d0768e786，最近一次 main 提交日期为 2025-08-26。README 的“仍在积极开发”和性能数字应结合此版本背景理解，不能据此推断当前维护频率或在其他硬件上必然达到相同吞吐。

建议把配套 picotron_tutorial 当作分步教材：step1_modeling → step2_process_group_manager → step3_dataloader → step4_tensor_parallel → step5_data_parallel_naive → step6_data_parallel_bucket → step7_pipeline_parallel_afab → step8_pipeline_parallel_1f1b。教程仓库的默认分支是 master；主仓库是 main，两者的命令行参数不同，不要混用。完整组合与 CP 再回主仓库读。

7.2 单卡闭环：从一组 token 到一次参数更新

为本章准备独立的 Linux/CUDA 学习环境。该提交 requirements.txt 固定 torch==2.1.0、triton==2.1.0、transformers==4.47.0、flash-attn==2.5.0 等依赖；它与本文前面 PyTorch 2.9.x 的 FSDP2 环境不同。先验证驱动、PyTorch 构建、CUDA 扩展和 FlashAttention 的兼容性，再记录环境；不要把一份依赖清单直接覆盖另一章环境。以下命令是根据源码整理的实验步骤，未在本次工作中完成 GPU 实跑。

git clone https://github.com/huggingface/picotron.git picotron-study
cd picotron-study
git checkout 59714b1bf25541dbeb6f5a721ed15f6d0768e786
# 在为本章建立的独立环境中安装；CUDA 扩展的构建条件需先满足。
python -m pip install -e .

固定学习版本；环境安装成功后再做训练实验

先选择较小的公开模型配置，保留其注意力头数与 KV 头数，只减少层数、序列长度和 micro-batch。需要 Hugging Face 访问凭据时，通过已有 HF_TOKEN 环境变量提供，不把真实 token 写进示例、配置或仓库。当前代码会检查这一环境变量，并下载模型配置、tokenizer、safetensors 与数据。

本版本的初始化并不等于“下载权重后继续训练”：checkpoint.py 先用 safetensors 物化参数形状，随后 init_model_parameters() 调用 reset_parameters()，重新初始化模型。需要明确区分“借用现有架构/形状”“随机初始化训练”和“从训练 checkpoint 恢复”；只看到 from_pretrained 或下载文件不能判断训练起点。

# 保存为 make_study_config.py，在 picotron-study 根目录运行。
from pathlib import Path
import json

cfg = json.loads(Path("template/base_config.json").read_text())
cfg["distributed"].update(tp_size=1, cp_size=1, pp_size=1, dp_size=1)
cfg["model"].update(name="HuggingFaceTB/SmolLM-360M-Instruct",
                    num_hidden_layers=4, use_fused_adam=False)
# 未显式指定头数时，train.py 会保留下载模型配置的头数。
cfg["model"].pop("num_attention_heads", None)
cfg["model"].pop("num_key_value_heads", None)
cfg["training"].update(seq_length=128, micro_batch_size=1,
                       gradient_accumulation_steps=1, total_train_steps=20,
                       num_samples=512, max_tokens=None)
cfg["environment"]["HF_TOKEN"] = None       # 使用进程已有的 HF_TOKEN
cfg["environment"]["FLASH_ATTEN"] = "0"    # 先走常规算子；扩展仍需能导入
cfg["model"]["use_flash_attention"] = False
cfg["logging"]["use_wandb"] = False
cfg["checkpoint"].update(save_frequency=10, save_dir="study_runs/single",
                         load_path="")
out = Path("study_runs/single")
out.mkdir(parents=True, exist_ok=False)     # 已存在就停止，避免覆盖以前的实验
(out / "config.json").write_text(json.dumps(cfg, indent=2))

用基础配置生成 20 次更新的单卡实验；不调用会删除同名目录的生成器

python make_study_config.py
torchrun --standalone --nproc_per_node=1 train.py --config study_runs/single/config.json

单机单卡基线

观察 20 次更新是否完成、loss 是否有限、每十步 checkpoint 是否出现，并记录输入 shape、首步 loss、一次更新后的选定参数、峰值 allocated/reserved 显存及吞吐。训练日志里的 Memory usage 是 reserved 值，不能直接当作峰值张量显存；短训练也不能证明模型已经收敛。之后每个配置使用独立目录，保持模型、数据和度量口径一致。

此处将 FLASH_ATTEN 设为 0 是选择代码中的普通计算分支，不代表能卸载 flash-attn；model.py 顶部仍有无条件导入。README 的 CPU 示例也不等于当前版本已支持纯 CPU 安装：utils.py 的参数统计仍创建 CUDA tensor。CPU 上可先运行 7.7.1 的分组练习；完整无 GPU 训练需要另做适配与验证。

手算沿用一组小例子：每个 micro-batch 有 $B=2$ 条序列，完整长度 $T=8$，隐藏维度 $C=16$，词表 $V=64$；4 个 Query 头、2 个 KV 头，单头维度 $D=4$，MLP 中间维度 $I=32$，共 4 个 Decoder 层。A 表示每 rank 每次更新累积的 micro-batch 数，默认 $A=1$。源码中 seq_length 对应 T，hidden_size 对应 C，grad_acc_steps 对应 A；源码片段与原图若使用 S/H/M，分别对应这里的 T/C/A；这些形状例子用于解释计算，实际运行继续使用上面的 SmolLM 配置。

7.2.1 沿输入、激活和 loss 看前向与反向

MicroBatchDataLoader.tokenizer_group_text 先把文本 token 拼接，再取 $T+1$ 长度的窗口；多出的那个 token 是最后一个位置的预测目标。例如一行 $[10,11,12,13,14,15,16,17,18]$，input_ids 是前 8 个，target_ids 是后 8 个。位置 0 的输出预测 token 11，位置 7 的输出预测 token 18。这把“自回归预测下一个 token”落实成了两个错开一位的整数张量。

位置	单卡手算 shape	这一步在做什么
input_ids / target_ids	$$\operatorname{shape}(\mathrm{input\_ids})=\operatorname{shape}(\mathrm{target\_ids})=(2,8)$$	整数 token ID；target 比 input 向后偏移一位
Embedding.forward	$$(2,8)\rightarrow(2,8,16)$$	每个 token 查到长度 C 的向量
Attention 的 Q、K、V	$\operatorname{shape}(Q)=(2,4,8,4)$；$\operatorname{shape}(K)=\operatorname{shape}(V)=(2,2,8,4)$	投影后拆头，Q/K 加 RoPE；K/V 再 repeat_interleave 到 4 头
Attention 输出与残差	$$(2,8,16)\rightarrow(2,8,16)$$	各头结果合并、out_proj，再与层输入相加
MLP 与第二条残差	$$(2,8,16)\rightarrow(2,8,32)\rightarrow(2,8,16)$$	$\operatorname{silu}\!\left(\operatorname{gate\_proj}(x)\right)\odot\operatorname{up\_proj}(x)$，再 down_proj 并加残差
final_norm / final_proj	$$(2,8,16)\rightarrow(2,8,64)$$	每个位置对 V 个候选 token 产生 logits
cross_entropy	$\operatorname{shape}(\mathrm{logits})=(16,64)$；$\operatorname{shape}(\mathrm{target})=(16,),\quad\operatorname{shape}(\mathrm{loss})=()$	对 $B\times T$ 个位置取均值，得到 0 维 Tensor，即 scalar loss

DecoderLayer.forward 的两条主语句是 $x\leftarrow x+\operatorname{Attention}\!\left(\operatorname{RMSNorm}(x)\right)$，随后 $x\leftarrow x+\operatorname{MLP}\!\left(\operatorname{RMSNorm}(x)\right)$。RMSNorm 在最后一个维度归一化；MLP 的 gate 分支与 up 分支做逐元素乘法，不是矩阵乘法。这里正好对应前面学过的 axis、broadcast、elementwise multiply 与 matmul。Attention 的普通路径调用 PyTorch SDPA；可把内部注意力分数理解为 $(B,\text{头数},T,T)$，但不能据此断言 SDPA 后端一定显式存下这整个矩阵。

# train_step 的核心计算；变量名沿用源码。
outputs = model(input_ids=input_ids)  # [B, S, V]
batch_size, seq_len = input_ids.shape
target_ids = target_ids.reshape(-1)  # [B*S]
outputs = outputs.view(seq_len * batch_size, -1)  # [B*S, V]
loss = F.cross_entropy(outputs, target_ids, reduction="mean") / data_loader.grad_acc_steps
loss.backward()

单卡/非 PP 路径：从三维 logits 到标量，再交给 autograd

cross_entropy 接收 logits 并内部处理归一化，不需要先手工 softmax。loss.backward() 计算当前 micro-batch 对参数的梯度；A 次反向完成后，外层 while 才调用一次 optimizer.step。$A>1$ 时每个 micro-batch 的均值 loss 再除以 A，在各 micro-batch 大小相同的条件下，这等价于对本次更新的所有位置取均值。遇到不同有效长度、padding 或末批变小，须按有效 token 数重新核对权重。

对照源码：$T+1$ 窗口；输入、目标与位置切片；Attention.forward；两条残差；Llama.forward；train_step。

单卡本轮的 loss 依次对 A 个 micro-batch 反向，梯度先累积在参数 .grad；外层 AdamW.step 只调用一次，然后清梯度、计步、按频率保存。此时没有跨 rank 通信，也没有 DP 梯度桶。先记录一个参数的初值、梯度与更新量，后续 DP/TP/PP/CP 都与这条参考计算对照。

7.3 DP：从不同样本的梯度到一致的参数更新

当完整模型能放在单卡、希望协作处理更多样本时，先只打开 $\mathrm{DP}=2$。模型副本相同，样本按 dp_rank 分配；需要平均的是同一参数来自不同样本的梯度。一次只改这个维度，TP/PP/CP 均为 1；有效 batch、初始权重与实际样本的对照方式必须同时确定。

从 7.2 的单卡配置复制一份 JSON 到新的文件，修改 distributed.dp_size=2，其他并行维为 1，并为 checkpoint 配置独立目录；沿用 7.2 的启动方式，把 --nproc_per_node 改为 2、配置路径指向新文件。TP 和 PP 单项实验同样每次只修改对应 size。所有性能对照都另行对齐有效 batch；下面的 $A=2$ 例子专门解释累积与桶同步。

DP 阅读顺序：先看逐参数同步的教学类 DataParallelNaive，再看实际训练入口使用的 DataParallelBucket、bucket.py 和 autograd hook。不要只依据 Naive 名称就把它当作正确性基准：应以 PyTorch DDP 或单卡参考计算核对梯度。代码里的实现注释也需要结合当前 PyTorch 行为理解，例如混合精度训练本身并不意味着 PyTorch 无法做梯度累积。

先回到 $\mathrm{TP}=\mathrm{PP}=\mathrm{CP}=1$，再单独设置 $\mathrm{DP}=2$、$A=2$。两个 DP rank 各处理 2 个 micro-batch，每个 micro-batch 有 $B=2$ 条、每条 $T=8$ 个预测位置。因此一次更新共处理 $2\times2\times2=8$ 条序列、64 个 token 位置。两张卡保存相同模型参数，处理不同样本；各自得到梯度后，再在对应参数之间取均值。

DataParallelBucket 默认用 FP32 梯度桶。BucketManager 为一组参数创建连续的 grad_data，param.main_grad 是其中一段按参数 shape 看待的 view；它和桶共享存储，不是第二份独立的同尺寸副本。注册在参数梯度累积节点上的 hook 会把当前 param.grad 加进 main_grad，随后把 param.grad 设为 None。因此在 micro-batch 中途看到 .grad 为空，不能直接判断“反向没有发生”，还要看 main_grad。

源码接力图：以 $\mathrm{DP}=2$、$\mathrm{TP}=\mathrm{PP}=\mathrm{CP}=1$ 为例，观察一次更新内的 micro-batch 循环、FP32 梯度累积、桶就绪通信，以及 optimizer.step 前的等待。

打开 SVG 原图

时点	源码动作	此时应该怎样理解
第 1 个 micro-batch 反向	main_grad += param.grad；param.grad=None	累积含 $\frac{1}{A}$ 缩放的局部梯度；require_backward_grad_sync=False，暂不发 DP 桶通信
第 2 个 micro-batch 反向	继续累积；mark_param_as_ready	本轮最后一次反向中，每个桶等待自己所有参数就绪；无需等其他桶全部完成
一个桶就绪	grad_data /= group_size；all_reduce(async_op=True)	先除组大小，再求和；$\mathrm{CP}=1$ 时就是两个 DP 副本的均值
最后一次 backward 结束前的回调	_post_backward → bucket_manager.wait → p.grad=p.main_grad.to(p.dtype)	等通信完成，再把梯度交回标准优化器接口；BF16 参数会接收 BF16 的 .grad
optimizer.step 后	model.reset → bucket_manager.reset	清空 FP32 桶、就绪集合与通信句柄，为下一次更新重新累积

数值上，若每个 micro-batch 的平均损失记为 $L_{d,m}$，最终目标梯度是 $\frac{1}{\mathrm{DP}}\sum_d\left[\frac{1}{A}\sum_m\nabla L_{d,m}\right]$。前一个 $\frac{1}{A}$ 来自 train_step，后一个 $\frac{1}{\mathrm{DP}}$ 来自桶归约。average_loss_across_dp_cp_ranks 处理的是用于日志的 Python 数值，它发生在 backward 之后；平均日志不能替代梯度同步。

DP 桶同步开关只控制这条复制参数的梯度归约路径。若同时启用 TP、PP、CP，它们为完成数学计算所需的层内归约、激活收发、K/V 收发仍要执行。“前 $A-1$ 个 micro-batch 不同步 DP 梯度”不能理解为“前 $A-1$ 个 micro-batch 完全没有通信”。另外，optimizer.zero_grad 不会替你清空自定义 main_grad，所以外层的 model.reset 不能漏掉。

对照源码：桶存储与 main_grad view；参数 hook；桶归约与等待；回填 .grad 与 reset；日志 loss 的单独归约。

7.4 TP：把一层矩阵运算分给两张卡

TP 解决单层权重或计算需要分担的问题。同一批输入由两张卡共同计算，因而不能沿用 DP 的“各卡取不同样本”。先把同一份完整权重按特征维切开，再看局部结果在哪里合成完整输出；这里的通信是补全同一次矩阵运算。

保持同一批数据，单独将 TP 设为 2。apply_tensor_parallel 的替换清单最值得先读：q/k/v、up/gate 都是 ColumnParallelLinear；out_proj、down_proj 是 RowParallelLinear；embedding 按词表切；final_proj 也按输出切，但额外指定 gather_output=True。Column/Row 来自论文里的矩阵记法；在 PyTorch 中 weight 的存储 shape 是 $(\mathrm{out\_features},\mathrm{in\_features})$，因此要以实际切分维度为准。

模块	$\mathrm{TP}=2$ 时每卡拥有/产生什么	通信发生在哪里
q_proj；k_proj/v_proj	Q 权重 $(8,16)$，$\operatorname{shape}(Q)=(2,2,8,4)$；K/V 权重 $(4,16)$，重复前各为 $(2,1,8,4)$	前向保留本地头；反向把各分片对共同输入 x 的梯度在 TP 组求和
attention.out_proj	权重 $(16,8)$；本地 2 头结果先合并成 $(2,8,8)$	各卡产生 $(2,8,16)$ 的部分和；前向 all-reduce 相加后恢复完整 C
MLP up/gate → down	up/gate 各产生 $(2,8,16)$；down 权重 $(16,16)$	中间 $I=32$ 被切为两半；down 的部分和在前向 all-reduce
embedding	各持有 32 行词表权重 $(32,16)$	本卡词表外的 token 查表结果置零，再 all-reduce 得完整向量
final_proj	各持有权重 $(32,16)$，先产生 $(2,8,32)$	all-gather 得 $(2,8,64)$，每个 TP rank 再用普通 cross_entropy

把表中的 shape 换成通信字节：本例 $B=2$、$T=8$、$C=16$，若激活为 FP32，down_proj 的局部部分和 $(2,8,16)$ 为 1,024 B。$\mathrm{TP}=2$ 的理想环式 AllReduce 对这一次调用每 rank 发送约 1,024 B，接收另计；它恢复的是同一输出的各部分之和。up/gate 的反向还各有输入梯度归约，final_proj 有自己的 AllGather，因此不能把这 1 KiB 当成整个 Block 或整个更新的总通信量。

这一点也解释了为什么列切分后未必立刻 AllGather：如果下一步逐元素激活和行切分线性层能直接消费本地隐藏特征，就把中间结果继续保留为分片。只有上下游的数据布局要求完整结果时才恢复；先看消费者需要的 shape，再判断通信是否必要。第 8.5 节的两层 MLP 正好可以与 1.6.2 的完整计算对照。

为什么 Column 的反向要通信？对某个分片 $Y_i=XW_i^{\mathsf{T}}$，本地算出的输入梯度是 $\mathrm{d}X_i=\mathrm{d}Y_i W_i$；完整 $\mathrm{d}X$ 必须把各个输出分片的贡献相加。CopyToModelParallelRegion 把这一规则放进 autograd：前向直接返回 X，反向 all-reduce。Row 的各卡前向本来就是同一输出的部分和，所以 ReduceFromModelParallelRegion 前向求和、反向直接传递。这里的 TP 求和是在补全同一数学运算，不是 DP 对不同样本梯度取均值。

final_proj 的 gather_output=True 还有一个直接影响：虽然词表权重分片了，每个 TP rank 仍会拿到完整词表 logits。$(B,T,V)$ 的这部分显存和收集通信仍然存在；此实现没有改成词表并行 cross-entropy。GatherFromModelParallelRegion 的反向只把完整 logits 梯度按词表切回本 rank。

ColumnParallelLinear 还提供 LinearWithAsyncAllReduce：先发起 $\mathrm{d}X$ 的异步归约，再计算 $\mathrm{d}W$，最后 wait，目的是尝试重叠通信与计算。但本版本 apply_tensor_parallel 没有给构造器传 async_all_reduce=True，默认走同步分支。读到一个优化实现，不等于训练入口已经启用了它。

对照源码：模块替换清单；Column 前向与默认开关；Copy / Reduce 的反向规则；Gather 的反向切片；可选异步线性层。

TP 配置必须满足模型的整除条件：注意力头数、KV 头数及对应的特征/词表切分都要合法。不要只把 --tp 改大，然后用任意 reshape 消除报错。比较单卡与 TP 时，先把同一份完整初始权重按规则切给各 rank；相同随机种子不自动保证不同拓扑构造出的权重逐项相同。

源码入口：picotron/data_parallel/data_parallel.py、picotron/data_parallel/bucket.py、picotron/tensor_parallel/tensor_parallel.py、picotron/tensor_parallel/tp_communications.py。

7.4.1 延伸：SP 怎样减少 TP 区域之间的激活副本

本节 SP 专指 Megatron 风格的 sequence parallelism。普通 TP 中，行并行层求和后，各 rank 往往都保留完整序列的激活，接着重复做 Norm、Dropout 等逐 token 操作。SP 在这些区域沿序列维分片，但为每个 token 保留完整隐藏维，所以 LayerNorm/RMSNorm 仍可对该 token 的隐藏特征正确归约。它减少的是这些区域的重复激活和计算。

沿用本章 $B=2,T=8,C=16,I=32$ 的 MLP，令 $\mathrm{TP}=2$。下面是典型 TP+SP 的前向布局，并非本固定版本 Picotron 已启用的路径。

当前位置	每 rank 的 shape	为什么转换
SP 区域：Norm	$$(B,T/\mathrm{TP},C)=(2,4,16)$$	只处理本地序列片段；每个 token 的完整隐藏维仍在本卡。
进入列并行线性层前	AllGather 后 $(2,8,16)$	让每份权重分片都为完整序列计算自己负责的输出特征。
列并行 → 激活 → 行并行	中间特征 $(B,T,I/\mathrm{TP})=(2,8,16)$；行并行产生 $(2,8,16)$ 的部分和。	中间按特征分片；行并行输出虽然同形状，各卡内容仍只是部分贡献。
返回 SP 区域	ReduceScatter 后 $(2,4,16)$	把行并行的部分和相加，同时沿序列切分输出，不再让每卡保存完整输出副本。

反向必须配套正确的布局变换；Norm 等复制参数也需要汇总各序列片段对它们的梯度。只把前向 AllReduce 换成 ReduceScatter，还不能得到正确的 SP 训练。临时 AllGather、Attention、logits 等仍会占显存，因此不能把整个模型的激活峰值一律除以 TP。

SP 与 7.6 的 CP 区别在于：这里在 Norm 等区域切序列，进入相应 TP 计算区域时收集所需输入；CP 则让不同 rank 的 Query 持续负责不同上下文位置，并通过通信访问其他位置的 K/V，以完成长序列 Attention。它们可以组合，但要分别交代布局与通信组。参考 Ultrascale Playbook 的 Sequence Parallelism；本章 Picotron 的 apply_tensor_parallel 没有提供这套 SP 接入，不能当作现成配置运行。

7.5 再打开 PP：把 forward/backward 排成流水线

PipelineParallel 将 decoder 层按 stage 分配；第一阶段保留 embedding，最后阶段保留 final_norm 与 final_proj。中间阶段接收 hidden_states，反向时将输入激活的梯度传回上一阶段。PP 传的是阶段边界上的激活和梯度，不是把每个 micro-batch 当成一次独立参数更新。

将 4 层 Decoder 分给 $\mathrm{PP}=2$：stage 0 保留 embedding 和第 0、1 层，stage 1 保留第 2、3 层、final_norm 与 final_proj。distribute_layers 按层数尽量均分，余数给前面的 stage；这不保证各阶段实际耗时相同。先看 $A=1$ 的 AFAB：stage 0 输出 $(2,8,16)$ 激活，send_forward 发给 stage 1；stage 1 接收一个 requires_grad=True 的张量，继续前向并计算 scalar loss。

跨进程 send/recv 不会自动把两个 PyTorch 计算图连接成一个全局图。stage 1 先对自己的图反向，得到边界输入的梯度 dH（与边界激活同 shape，这里的 H 表示激活张量），再 send_backward。stage 0 用收到的 dH 作为 torch.autograd.backward(output_tensor, grad_tensors=dH) 的上游梯度，对自己保留的局部图反向。最后阶段以 scalar loss 的梯度 1 为起点；前面阶段用下游传回的张量梯度为起点。PipelineParallel.backward 正是在处理这两种情况。

AFAB 的 input_tensors/output_tensors 列表持有每个 micro-batch 的输入、输出及其关联的局部 autograd 图，随后按 FIFO 取出反向，并不是在 backward 时重跑 forward。1F1B 使用相同的模型阶段与边界梯度规则，主要改变列表中同时等待反向的 micro-batch 数量以及收发顺序。

本版本需要先修正或规避一个数值差异：普通 train_step 将反传的 loss 除以 grad_acc_steps；AFAB/1F1B 只在 logging_loss 中做这个除法，参与 backward 的 output_tensor 没有同样缩放。因此 accumulation>1 时，“打印 loss 相近”不能证明梯度与单卡一致。第一轮 $\mathrm{PP}=2$ 可先用 accumulation=1 对齐；学习多 micro-batch 调度前，将反传目标与日志统计分别理清，避免对日志再多除一次，再比较参数梯度和一次更新。该问题来自静态源码检查，本章没有声称已替上游修复或完成 GPU 回归。

$\mathrm{PP}=2$，$A=4$	每个 stage 的本地计算顺序	何时打开该 stage 的 DP 同步
AFAB：两个 stage	$$F_0\rightarrow F_1\rightarrow F_2\rightarrow F_3\rightarrow B_0\rightarrow B_1\rightarrow B_2\rightarrow B_3$$	最后的 $B_3$；期间通信和等待另计
1F1B：stage 0	$$F_0\rightarrow F_1\rightarrow B_0\rightarrow F_2\rightarrow B_1\rightarrow F_3\rightarrow B_2\rightarrow B_3$$	cooldown 的 $B_3$
1F1B：stage 1	$$F_0\rightarrow B_0\rightarrow F_1\rightarrow B_1\rightarrow F_2\rightarrow B_2\rightarrow F_3\rightarrow B_3$$	steady state 的 $B_3$

这张表列的是各阶段本地次序，不是等宽时间轴。stage 0 的 $\mathrm{warmup}=\min(\mathrm{PP}-\mathrm{pp\_rank}-1,A)=1$，stage 1 为 0；实际时间还包括依赖等待、通信与不同阶段耗时。$A=4$ 只用于理解调度，本版本 PP 反传 loss 缺少 $\frac{1}{A}$ 缩放的问题仍须按 7.5 先修正，再做数值对比。optimizer.step 始终在外层，等该轮所有 micro-batch 反向完成后才执行。

组合 TP 时，PP 边界张量仍是 $(B,T/\mathrm{CP},C)$，本例是 $(2,8,16)$，因为 Attention/MLP 的 RowParallel 输出已还原完整 C。TP rank t 与下一 stage 的同一 t 收发，而不是把 hidden_size 再除以 TP。通信文件虽然使用 isend/irecv 和 batch_isend_irecv，随后会 wait 并调用 cuda.synchronize；仅看到异步 API 名字，不能据此断言已经充分隐藏通信开销。

PP 的账本沿阶段边界计数：本例一条边界激活 $(2,8,16)$ 若为 FP32，占 1,024 B，其反向梯度也是 1,024 B。两 stage、$A=4$ 时，在每一条匹配的 TP 通道上，每次参数更新前向发送 4 KiB、反向发送 4 KiB，两个方向合计 8 KiB；未计协议开销、DP/TP 内部通信，也未把张量重复除以 TP。AFAB 与 1F1B 可有相同的边界数据总量，却因调度不同而有不同的等待时间和激活峰值。

先区分两种收益：激活驻留与流水线气泡。AFAB 要为这一轮多个前向保留待反向的图；经典 1F1B 更早完成一部分反向并释放相应激活，主要收益是降低同时驻留的激活数量。固定 stage 数和 micro-batch 数时，它仍有填充和排空阶段，不能把“换成 1F1B”理解成自动消除气泡。

在各 stage 均衡、前向/反向耗时固定、忽略通信且不交错虚拟 stage 的理想模型下，设 stage 数为 $P$，一次更新含 $A$ 个 micro-batch，气泡时间占比约为 $\frac{P-1}{A+P-1}$。本节 $P=2,A=4$ 对应 20% 的理论空闲比例，不是对本项目的实测。增加 $A$ 可摊薄填充/排空开销；若要固定 global batch，就要同时减小每个 micro-batch，单次矩阵运算效率也可能下降。实验应先修正前述 loss 缩放，再对同一输入比较 AFAB 与 1F1B 的梯度、峰值和端到端时间；zero-bubble、interleaving 属于后续调度，不是当前 pp_engine=1f1b 的承诺。

对照源码：阶段持有的模块；显式连接边界梯度；AFAB 的保存与 FIFO 反向；1F1B 的 warmup / steady / cooldown；阶段收发与等待。

建议练习：两 stage、四个 micro-batch，手工列出每阶段的 $F_0\ldots F_3$、$B_0\ldots B_3$；在源码中找到 send_forward、recv_forward、send_backward、recv_backward，以及 send_fwd_recv_bwd 双向通信。解释哪一刻触发最后一次 DP 梯度同步，哪些激活可以释放。

源码入口：picotron/pipeline_parallel/pipeline_parallel.py、picotron/pipeline_parallel/pp_communications.py；对照单卡 train_step 的 loss 缩放。

7.6 最后单独理解 CP：同一序列怎样跨卡计算 Attention

CP 沿序列长度切分同一批样本。以完整长度 $T=8$、$\mathrm{CP}=2$ 为例，一个 rank 持有位置 0…3，另一个持有 4…7；它们不是两条独立训练样本。后半段的查询仍需要前半段的 K/V，所以不能把两个半段各自独立做普通 attention 就结束。输入、target 的一位偏移，以及 RoPE 的全局位置都必须对应。

需要 CP 的背景是长序列带来的 Attention 计算与激活压力。它切的是同一条序列的位置，不能把两段当成互不相关的样本。下面先解释两段如何共同计算完整因果注意力，再讨论模型参数梯度的同步。当前局部 attention 用 matmul/exp 等实现，不把教学 RingAttention 等同于优化后的 Ring FlashAttention 内核。

先固定 $\mathrm{TP}=\mathrm{PP}=1$，只研究 $\mathrm{CP}=2$ 的局部 attention 计算。仍以 9 个 token 的样本为例：cp_rank=0 的 $\mathrm{input}=[10,11,12,13]$、$\mathrm{target}=[11,12,13,14]$；cp_rank=1 的 $\mathrm{input}=[14,15,16,17]$、$\mathrm{target}=[15,16,17,18]$。每卡张量为 $(2,4)$，Query 长度是 4，但后半段 Query 仍需要看前半段的 Key/Value。

RingAttentionFunc.forward 让本地 Q 留在当前 rank，沿 CP 环发送和接收 K/V。step=0 计算本地块并应用因果 mask；后续只把历史块加入当前 Query 的计算。两卡例子中，rank 0 收到的是未来块，需要跳过计算；rank 1 收到的是历史块，4 个本地 Query 都可看这 4 个历史 Key。即使某块不参与当前 Query 的计算，环上的配对收发仍然要按一致顺序执行。

同样长的 Query 片段，计算负担也可能不同。因果 mask 下，后面的 Query 能看见更多历史位置。沿用一条 $T=8$ 序列：位置 0～3 的有效 Query–Key 对共 $1+2+3+4=10$ 个，位置 4～7 共 $5+6+7+8=26$ 个。Ultrascale Playbook 的 Zig-Zag Ring Attention 将早、晚位置搭配分配，例如两卡分别持有位置 $\{0,1,6,7\}$ 与 $\{2,3,4,5\}$，两边各有 18 个有效位置对。

这是负载分配的概念推导，不是本版本的现成开关，也不是实际 kernel 耗时比。真正改造需要同时调整输入/目标切片、全局位置、因果 mask、环中块次序以及输出拼回；仅重新排序输入会破坏语义。先完成本节连续切片 CP 的前向、激活梯度和参数更新验收，再考虑此优化；一般稠密 kernel 也未必能跳过所有被 mask 的计算。

每个块的 softmax 分母不同，不能简单把 block_out 平均。若一个块的 log-sum-exp 是 $\ell_1$，输出是 $O_1$，另一个是 $\ell_2$、$O_2$，则合并输出为 $\exp(\ell_1-\ell)O_1+\exp(\ell_2-\ell)O_2$，其中 $\ell=\log\!\left(\exp(\ell_1)+\exp(\ell_2)\right)$。例如两块的未归一化概率总量为 1 和 3，合并权重应是 $\frac14$ 和 $\frac34$。update_out_and_lse 用 sigmoid/logsigmoid 写成数值稳定的等价形式。

反向时，本地 Q 的梯度累加在本 rank，K/V 的梯度贡献则沿环累积并回到它们所属的 rank。这解决的是“同一个 Attention 运算中，各位置彼此依赖”的导数；模型的 q_proj、k_proj、MLP 等参数在 CP rank 间仍是副本，还需要把各段序列对同一参数的贡献在 cp_dp_group 中归约。两条通信链承担不同工作，所以 7.6 所指出的纯 CP 包装条件缺口，不能靠 ring backward 自动补齐。

位置编码也要沿真实代码追下去。collate_batch 返回全局 position_ids，但普通 train_step 没有传它，DecoderLayer.forward 当前也不靠这个参数索引 RoPE。实际路径是在构造 DecoderLayer 时先生成完整 cos/sin，再由 update_rope_for_context_parallel 按 cp_rank 切片；后半段拿到的是位置 4～7，而不是重新从 0 开始。修改模型或支持变长序列时，不能只改 DataLoader 的 position_ids 就以为位置语义已经生效。

对照源码：Ring 前向与 causal 条件；K/V 梯度环；在线合并输出；RoPE 分片；RoPE 实际接入位置。

纯 CP 训练还有一个明确的入口检查点：train.py 仅在 dp_world_size>1 时包装 DataParallelBucket，但 bucket 同步使用的是 cp_dp_group；apply_context_parallel 只设置环境变量。因此 $\mathrm{DP}=1$、$\mathrm{CP}>1$ 时，当前入口不会为复制参数建立这条梯度同步路径。要学习两卡纯 CP，需将是否包装的条件按 $\mathrm{CP}\times\mathrm{DP}$ 的同步需求修正，并验证同一参数的梯度/更新一致；环上传递 $\mathrm{d}K,\mathrm{d}V$ 不能替代复制参数的梯度归约。不要将未核对的纯 CP 配置直接作为训练基线。

验证顺序：先在小形状下比较普通完整 attention 与 CP 拼回后的前向输出，再比较 $\mathrm{d}Q,\mathrm{d}K,\mathrm{d}V$，最后比较模型参数梯度和一次更新。seq_length 必须能被 CP 整除；global_batch 仍只随 DP、micro_batch 和累积步数变化。短序列可能因通信开销变慢，这是性能观察，不必自动归因于实现错误。

源码入口：picotron/context_parallel/context_parallel.py、picotron/context_parallel/cp_communications.py、CP 数据切片、CP 与 DP 包装入口。

7.7 组合与跨机：先算清 rank，再启动作业

7.7.1 建立通信坐标：本 rank 与谁协作

读代码的顺序	观察对象	读完应回答
train.py	配置→进程组→数据→模型→TP/PP→参数物化→CP→DP→训练	本 rank 的 device、并行坐标、模型分片和数据分片在哪里确定？
process_group_manager.py	grid 的 $\operatorname{shape}=(\mathrm{DP},\mathrm{PP},\mathrm{CP},\mathrm{TP})$；沿某一维建通信组	谁和谁做 TP、PP、DP、CP 通信？为什么全局 rank 不等于某个组内 rank？
data.py / model.py	batch、token、position 与 Llama 的 forward	哪些 rank 用相同样本？哪些 rank 只持有同一序列的一段？
四类 parallel 模块与 communications 文件	参数替换、autograd、collective、点对点收发	计算量、存储量和通信量各自发生了什么变化？
checkpoint.py / utils.py	分片保存、恢复、初始化、loss 与吞吐统计	恢复了哪些状态？日志中的数值代表什么？

第一条恒等式：$\mathrm{world\_size}=\mathrm{DP}\times\mathrm{PP}\times\mathrm{CP}\times\mathrm{TP}$。GPU 模式下一般一进程对应一张卡。第二条是每次更新的样本数：$\mathrm{global\_batch}=\mathrm{DP}\times\mathrm{micro\_batch\_size}\times\mathrm{gradient\_accumulation\_steps}$；有效 token 数再乘完整 seq_length。TP/PP/CP 让同一工作协作完成，不应再把它们乘进 global_batch。CP 将完整序列长度分成每卡 $\mathrm{seq\_length}/\mathrm{CP}$，不能把“局部长度”当成完整长度后再重复计数。

Ultrascale Playbook 用“5D”概括 DP、TP、PP、CP、EP 五类分工；它不是一个与框架无关的进程数公式。这里仍使用本固定版本 Picotron 的四维 grid。MoE 框架经常在已有数据并行 rank 内组织专家组，EP 会改变专家参数的复制与同步范围；应先画清各组成员，再计数。不能向这里的四维乘积再随意乘一个 EP，也不能虚构 Picotron 的 ep_size 选项。EP 的数据流见 7.9。

Picotron 进程分组图：用 $\mathrm{DP}=2$、$\mathrm{PP}=2$、$\mathrm{CP}=1$、$\mathrm{TP}=2$ 的八进程例子，分别看副本、模型阶段、层内切分与通信对象。图中的 rank 编排来自本版本的 ProcessGroupManager。

打开 SVG 原图

from itertools import product

DP, PP, CP, TP = 2, 2, 1, 2
def rank(d, p, c, t):
    return ((d * PP + p) * CP + c) * TP + t

tp_groups = [[rank(d,p,c,t) for t in range(TP)]
             for d,p,c in product(range(DP),range(PP),range(CP))]
pp_groups = [[rank(d,p,c,t) for p in range(PP)]
             for d,c,t in product(range(DP),range(CP),range(TP))]
dp_groups = [[rank(d,p,c,t) for d in range(DP)]
             for p,c,t in product(range(PP),range(CP),range(TP))]
print("TP:", tp_groups)  # [0,1] [2,3] [4,5] [6,7]
print("PP:", pp_groups)  # [0,2] [1,3] [4,6] [5,7]
print("DP:", dp_groups)  # [0,4] [1,5] [2,6] [3,7]

不用 GPU，先把八个 rank 的通信组算出来

源码入口：train.py、picotron/process_group_manager.py、picotron/data.py。

7.7.2 启动时怎样让数据、模型和优化器采用同一分工

torchrun 启动的每个进程都会执行 train.py。LOCAL_RANK 决定本机设备，RANK/WORLD_SIZE 决定全局通信身份；ProcessGroupManager 再把全局 rank 映射到 $(\mathrm{dp},\mathrm{pp},\mathrm{cp},\mathrm{tp})$ 坐标。这个对象被各模块通过 pgm.process_group_manager 读取，相当于共享的分工说明：DataLoader 用 dp_rank 选样本，模型用 tp_world_size 算本地头数，PP 用前后邻居收发激活，CP 用环上邻居收发 K/V。

DistributedSampler 只按 DP 维度分样本。相同 dp_rank 的 TP/PP/CP 进程会在相同迭代位置取到对应的完整样本，再由 collate_batch 按 cp_rank 切出序列片段。PP 的各阶段也各自推进 DataLoader：第一阶段用 input_ids，最后阶段用 target_ids，中间阶段主要计算收到的 hidden_states。因此它们必须保持 micro-batch 次序一致。

模型构建的顺序有实际意义。先把注册参数放在 meta 设备上建立结构，进行 TP 模块替换和 PP 层选择，再从 safetensors 物化当前 rank 所需的参数形状并 reset_parameters。meta 参数可以描述 shape，却不能用于实际数值计算；这里主要减少完整模型参数常驻的开销，也不能据此宣称初始化期间完全没有其他内存分配。随后转换 dtype/device、建立 DP 梯度桶，最后把最终的 model.parameters() 交给 AdamW。若先创建优化器、再替换 Parameter 对象，优化器就可能仍握着旧参数。

四种并行的接入方式也不同：apply_tensor_parallel 用 setattr 真正替换 Linear/Embedding；PipelineParallel 创建只包含本阶段层的模块；apply_context_parallel 设置环境变量，让 Attention.forward 选择 ring_attention 分支；DataParallelBucket 包住模型并给参数的梯度累积节点挂 hook。学习时要找到“行为在哪里改变”，不能只看函数名中是否带 parallel。

对照源码：设备与进程组初始化；模型构建与包装顺序；AdamW 接收最终参数；meta 参数注册；参数物化后重新初始化。

7.7.3 跟踪 rank 5：从取样到更新和保存

回到 7.7.1 的八进程布局 $\mathrm{DP}=2$、$\mathrm{PP}=2$、$\mathrm{CP}=1$、$\mathrm{TP}=2$，并使用本节的 4 层小尺寸例子、$A=1$。rank 5 的坐标是 $(\mathrm{dp}=1,\mathrm{pp}=0,\mathrm{cp}=0,\mathrm{tp}=1)$。它与 rank 4 做 TP，与 rank 7 做 PP，与 rank 1 做 DP。若两机各运行 4 个连续 rank，它在第二台机器的 local_rank=1，DP 梯度通信跨机，TP 和 PP 通信在本机。

一次更新中，rank 5 取 DP 分片 1 的 $B=2$ 条完整序列；它持有词表后半段 $[32,64)$ 的 embedding，以及第 0、1 层中第二份 TP 参数分片。每层负责 2 个 Query 头和 1 个 KV 头。它与 rank 4 共同完成词表查表和每层的 TP 运算，然后把 $(2,8,16)$ 激活发给 rank 7。rank 7 完成本地后两层和输出投影，与 rank 6 收集完整词表 logits、计算 loss，再把边界梯度传回 rank 5。

rank 5 对自己的前两层反向时，TP 通信补全输入梯度；参数 hook 把本地分片梯度放入桶，与持有对应分片的 rank 1 在 DP 组归约。桶 wait 完成，AdamW 更新 rank 5 持有的参数以及本地优化器状态。它不保存全模型，也不等待 rank 0 替它做优化器更新；每个 rank 都执行自己那份 optimizer.step。这里一次更新的有效工作量仍是 $\mathrm{DP}\times B\times A\times T=32$ 个 token 位置。

外层随后累计 trained_tokens、递增 step、reset 梯度桶，按保存频率调用 CheckpointManager。此布局中由 $\mathrm{dp\_rank}=0$、$\mathrm{cp\_rank}=0$ 的 rank 0、1、2、3 分别写出四份 TP/PP 模型与优化器状态，rank 5 不重复写副本。日志则选 $\mathrm{dp}=0$、$\mathrm{tp}=0$、$\mathrm{cp}=0$ 且处于最后阶段的 rank 2 输出，并非固定 global rank 0。保存时所有 rank 进入函数，也不代表所有 rank 都写文件。

从 checkpoint 继续时，仍先构造同一拓扑下的模型与优化器，再加载各自对应的文件、恢复 step/token 计数。源码只保存 model、optimizer、trained_steps、trained_tokens；这里的“继续训练”不能推导出数据迭代位置和随机数轨迹也精确恢复。

对照源码：四维 rank 布局与通信组；日志 rank 的选择；一次更新的外层控制流；分片命名、写入资格与恢复。

7.7.4 从单项配置进入组合，再跨机

实验顺序	DP / TP / PP / CP	进程数	进入条件
单卡基线	1 / 1 / 1 / 1	1	短训练、checkpoint 和度量口径已确认
分别尝试 DP、TP、PP	仅对应的一维为 2，其他为 1	各 2	模型与输入/梯度基线对齐；PP 首轮累积为 1
CP 专项	1 / 1 / 1 / 2	2	完成 7.6 同步入口修正及前向/反向比较
$\mathrm{DP}\times\mathrm{TP}$ 或 $\mathrm{DP}\times\mathrm{PP}$	2 / 2 / 1 / 1；或 2 / 1 / 2 / 1	4	相关单项已验证，PP 累积缩放已对齐
$$\mathrm{DP}\times\mathrm{TP}\times\mathrm{PP}$$	2 / 2 / 2 / 1	8	可先单机八卡；也可两机、每机四卡
完整四维组合	2 / 2 / 2 / 2	16	选做：先验证 CP，再核对通信组和完整 token 口径

配置编辑的是 JSON 中 distributed 的 dp_size、tp_size、pp_size、cp_size，流水线调度是 pp_engine（afab 或 1f1b）。torchrun 的进程总数必须等于四维乘积。每次改变 DP 后重新核算有效 batch；保持一次更新的 token 数相同时，可以调整 micro_batch/gradient_accumulation_steps，但比较实验还要固定初始权重和实际样本。

对 $\mathrm{DP}=2$、$\mathrm{PP}=2$、$\mathrm{CP}=1$、$\mathrm{TP}=2$，若两机各运行四个进程，默认连续 global rank 让 rank 0…3 位于节点 0、4…7 位于节点 1：TP 与 PP 组在节点内，DP 组跨节点。这是本例由 rank 布局产生的结果，不能推断任何配置都自动选择最优通信拓扑。先用第 6 章的方法验证节点间通信，再比较布局与真实互联是否匹配。

# 两台机器使用同一份已核对的 eight/config.json：DP=2 TP=2 PP=2 CP=1。
# MASTER_ADDR 在两台机器上都设为节点0可达地址；在各自环境中设置 NODE_RANK=0 或 1。
torchrun --nnodes=2 --nproc_per_node=4 --node_rank "$NODE_RANK" \
  --master_addr "$MASTER_ADDR" --master_port 29570 \
  train.py --config study_runs/eight/config.json

跨机启动模板：先完成单项正确性与环境/网络检查

源码只让 global rank 0 下载模型文件，随后假设所有 rank 都能从同一路径读取 hf_model_safetensors。多节点必须有共享路径，或在启动前让每个节点具备一致文件与工作目录；barrier 不会替你复制模型文件。数据集、tokenizer 缓存和 checkpoint 可见性也需分别确认。单机能启动，不能直接证明跨机准备已完成。

7.8 恢复、性能与通关：用证据判断实现是否可靠

CheckpointManager 按 TP/PP 坐标保存模型与 optimizer state，并保存 trained_steps/trained_tokens；DP/CP 的第 0 个副本负责写文件。当前加载代码假设拓扑不变，没有实现任意 TP/PP 度数之间的自动重分片。文件中也没有完整的数据迭代位置与 RNG 状态，因此“能加载继续训练”不等于“恢复后逐步复现原训练轨迹”。先在相同拓扑验证状态恢复，再讨论更完整的 checkpoint 设计；生产方向与 PyTorch DCP 对照。

7.8.1 留下可对照的记录：样本、梯度、更新与恢复

第一次跟踪建议选择 decoder_layers.0.attention.q_proj.weight；DP 包装后名字会带 module. 前缀，PP 下层号保留全局编号。先在单卡记录，再在两卡的单项实验里追踪对应分片。用同一份参考权重和实际样本对齐，不能只对比“相同 seed”或“loss 大约差不多”。

观察位置	记录什么	应该回答的问题
collate_batch 返回时	dp/cp 坐标、input/target/position 的首行与 shape	样本分给谁？输入与目标是否错位一位？CP 的边界是否连续？
模型物化后、优化器创建前	选定参数的名称、shape、dtype、device	这个 Parameter 属于哪层、哪份 TP 分片？优化器会持有谁？
Attention / final_proj 输出	Q/K/V、hidden_states、logits 的 shape	哪些维度被切了？哪些地方又被归约或收集成完整形状？
参数 hook 与 _post_backward	选定参数的 main_grad/.grad 范数、桶就绪时点、通信组	当前看到的是局部梯度、累积梯度，还是已经同步的梯度？
optimizer.step 前后、reset 后	选定参数的更新量；桶是否归零；step/token 计数	一次更新到底消费了几个 micro-batch？梯度是否泄漏到下一步？

实验顺序继续保持单卡→单机两卡逐项→单机组合→多机组合。TP 比较时拼回对应分片；PP 比较时按全局层号对应；CP 比较时拼回序列并分别检查激活梯度与参数梯度。记录绝对/相对误差并结合 dtype 解释，再讨论吞吐。多进程调试可按 rank 输出有限记录，避免单独长时间暂停一个参与 collective 的 rank，导致其余进程等待超时。

本节已对照固定版本的调用、shape、切片、分组和调度顺序；小例子的数值口径另做了本地推导核对。真实训练的损失、梯度一致性、通信重叠和性能仍需在 CUDA 环境按前述实验验证。读完后应能从一条 token 样本追到一个参数的更新，并说明中途每次通信究竟补齐了哪一部分计算。

性能记录至少包含：提交号、环境、GPU/互联、四维度数、模型与完整序列长度、micro-batch/累积步数、每步有效 token、预热后的端到端吞吐、各 rank 峰值显存及阶段耗时。PP 要留意最慢阶段和气泡，TP/CP 要看通信与计算占比，DP 要看梯度 bucket。

README 中的 H100 MFU 数据是项目作者报告；本次没有复现。utils.get_mfu 使用固定 theoretical_flops 默认值，并按模型参数和序列长度估算工作量；换成 H20/H800 等卡时不能直接沿用其百分比比较。先看统一口径下的 tokens/s 和耗时，再说明 MFU 的精度、峰值与计算量假设。

通关任务	应提交的证据
解释一次训练与四种分工	画出一组实际 rank 的 TP/PP/DP/CP 组，标出每 rank 的参数、样本和序列范围
数值正确性	相同权重/样本下，对齐前向、梯度和一次更新；单独验证 PP 累积缩放与 CP 复制参数同步
故障恢复	相同拓扑保存、加载并检查模型/优化器/进度；说明数据位置和 RNG 尚有哪些缺口
一次有效性能实验	固定总工作量，对比单项/组合并行的吞吐和峰值显存，解释通信与阶段不均的影响

资源顺序：先用 picotron_tutorial 的分阶段代码建立理解，再读本章固定版本的主仓库，最后按问题阅读第 8 章论文，并在第 9 章比较更完整的工程框架。Picotron 的教学价值是让复杂机制变得可检查；更高吞吐、容错和生产运维能力需要额外验证与工程工作。

核对来源：项目定位与作者性能报告；依赖版本；初始化与 checkpoint；MFU 与跨 rank 指标；提交历史；配套分步教程快照。

7.9 延伸到 EP：MoE 的 token 怎样找到专家并返回

MoE（Mixture of Experts）把某些 dense MLP 替换为多个专家 MLP，再由 router 为每个 token 选择少数专家。EP（expert parallelism）把这些专家放到不同 rank：token 先去专家所在的卡计算，再把结果送回原位置。它能分担专家参数和计算，但引入路由、重分布通信与负载均衡问题；并不是把同一个 MLP 的矩阵按特征切开，那是 TP 的分工。

先固定模型与实现边界。下面说明通用的 top-k MoE/EP 机制，采用不丢 token、对所选专家权重重新归一化的教学约定；不同 MoE 的共享专家、容量限制和权重归一化规则可能不同。本文固定的 nanoGPT 与 Picotron 都使用 dense MLP，没有可直接运行的 EP 接入。所参考 Playbook 的该版本也将更完整的 EP 示例列为计划新增；以下不能当作本章源码已有能力。

本节先不叠加 TP/PP/CP。记一个源 rank 输入的隐藏状态为 $X\in\mathbb{R}^{N\times C}$，其中 $N=BT$ 是本地这次 micro-batch 的 token 位置数，$B$ 是本地序列条数，$T$ 是每条序列的完整长度，$C$ 是隐藏宽度。设总专家数 $E=4$、每个 token 选择 $K=2$ 个专家，EP 组有两个 rank，各保存两个专家。所有专家均把一个 $C$ 维向量映射回 $C$ 维向量；专家的内部参数彼此独立。

步骤	数据及 shape 怎样变化	必须保留什么
路由：选择专家	router 权重 $W_R\in\mathbb{R}^{E\times C}$，分数 $R=XW_R^{\mathsf T}\in\mathbb{R}^{N\times E}$。按分数选出 top-k 专家编号，再对选中的分数做 softmax，得到合并权重；编号与权重的 shape 均为 $(N,K)$。	编号是整数；权重是浮点数。编号用于分发，权重用于最终合并，二者不是同一个张量。
Dispatch：按专家目的地分发	把每个 token 对每个选中专家展开为一项任务，形成 $NK$ 行 $C$ 维隐藏向量；按目的 rank 排列并交换。典型实现使用 AllToAll，可为非均匀长度。	源 rank、原 token 编号、所选专家及排列映射。各方向的发送/接收数量必须匹配。
专家计算	专家 $e$ 收到来自 EP 组各源 rank 的任务，输入 $X_e\in\mathbb{R}^{n_e\times C}$，输出同形状。	$n_e$ 是本轮实际路由到该专家的 token 任务数，会随数据改变；不同专家不一定相等，也可能为零。
送回并 Combine	结果按源 rank 送回，典型实现再做一次 AllToAll；逆置换后按原 token 加权求和，恢复 $Y\in\mathbb{R}^{N\times C}$。	原 token 次序与路由权重。收到几个专家结果，不意味着产生几条新的训练样本。

令 $\mathcal{K}_i$ 为 token $i$ 选中的专家集合，$f_e$ 为专家 $e$ 的函数，$\alpha_{i,e}$ 为本节归一化后的路由权重，则 $y_i=\sum_{e\in\mathcal{K}_i}\alpha_{i,e}f_e(x_i)$，且 $\sum_{e\in\mathcal{K}_i}\alpha_{i,e}=1$。因此一次 forward 中的基本链路是“路由 → 分发 → 专家 → 返回 → 合并”。若多个源 rank 参与，记 $N_r$ 为源 rank $r$ 的 token 数，所有专家的任务总数为 $\sum_e n_e=K\sum_r N_r$；这是专家任务数，不能计成 global batch 中额外的样本或 token。

反向也要沿这些映射传递梯度：合并处把输出梯度按权重送回对应专家结果，专家产生参数梯度与输入梯度，再把输入梯度送回原 token 位置累加。输入 $X$ 还进入了 router，因此最终输入梯度必须累加专家路径和路由连续权重路径的贡献。路由中的连续权重可以收到梯度，离散 top-k 编号本身不按普通可微运算求导。前向能收发成功，并不代表这条反向链已经正确；只调用普通通信 API 也不会自动连接跨进程 autograd 图。

瓶颈来自最忙的专家和通信。热门专家的 $n_e$ 大，会增加该 rank 的计算与缓冲需求，其他 rank 可能等待。负载均衡约束、router 设计、capacity 与丢弃/不丢弃 token 的策略必须结合模型定义解释；不能静默丢掉溢出 token 后仍宣称与原计算等价。性能记录要包含每个专家的任务数、最大/平均负载、dispatch/combine 耗时、峰值和完整训练步耗时。

专家参数只需在保存同一专家副本的 rank 间同步；Attention、Norm、router 等非专家参数可能有不同的复制组。例如不使用 TP/PP/CP 时，四个 rank 各取不同数据，可在两个二卡 EP 组内分布专家，再让两个组中的同名专家副本同步。仍只有四份独立数据，不能因为 EP 度数为 2 就把 batch 再乘 2。具体组的定义必须以采用的 MoE 框架为准。

阅读与改造的验收顺序。先在单进程固定权重和路由，比较“直接逐 token 调用所选专家”与“按专家重排再逆置换”的输出、输入梯度和专家梯度，确认每项任务恰好处理一次；再用 C.10 理解 AllToAll 的目的地语义，最后才做两卡分发与返回，继续与同一个参考计算比较。MoE 路由还要专门覆盖某专家为空、负载不均和非均匀交换长度。C.10 只验证通信原语，不是完整 EP 训练器。

参考：Ultrascale Playbook：Expert Parallelism 与 5D Parallelism；MoE 模型背景可选读其链接的 Hugging Face MoE Explained。本节增加的是机制与验收方法，真实跨卡 MoE 实现需要独立的模型改造和实验。

8. 论文与现代大模型组件：按实践中的问题选读

阅读方式：这一章不是新的必修主线，而是把前面已经运行过的代码映射回论文。建议按问题查读：想理解注意力看 8.1/8.3，想理解现代模型结构看 8.2，想理解多卡显存与通信看 8.4/8.5，想做低成本微调看 8.6。

这一章的能力目标是把一个观测到的问题，映射为论文机制、具体代码改动和可证伪的实验。阅读时分别记下原论文结论、当前库的实现与自己的测量；论文中的设备和模型规模不是本教程的实验结果。先复用前面的小形状或 nanoGPT，再决定是否需要新的模型和数据。

📖

读论文的实用姿势（省时间）：先读摘要 + 引言 + 图表 + 结论建立全局；再定位它的核心贡献那 1-2 个公式/算法框；然后去找官方或第三方的最小实现对照读（代码永远比公式好懂）；最后自己跑通一个玩具版。不要一上来逐字啃，90% 的细节第一遍不需要。

8.1 问题：因果 Attention 怎样把上下文变成预测

原始 Transformer（2017）采用 encoder-decoder 架构。Encoder 的自注意力可以访问整个输入序列；Decoder 使用因果自注意力，并通过 cross-attention 读取 Encoder 输出。现代 GPT 常采用 decoder-only 架构，不保留 Encoder 与 cross-attention。2.3 已经学习因果注意力；本节在此基础上对照论文结构。

打开 SVG 原图

对应到 2.3：本教程 Block 复用了多头注意力、残差和前馈网络等思想，但采用 Pre-Norm、GELU 和 decoder-only 结构，不能直接等同于原论文的完整 Decoder。$\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^{\mathsf{T}}}{\sqrt{D}}\right)V$ 是共同的核心计算；因果 mask 决定能读取哪些位置。

打开 SVG 原图

延伸实验：2.3.3 已用手写因果注意力与 SDPA 做数值对照。这里改变序列长度、头数和 head_dim，再比较输出、梯度与耗时；确保两个实现使用相同 mask、dropout 和精度。

import torch, torch.nn.functional as F, math

B, H, T, d = 2, 4, 16, 32
q, k, v = torch.randn(B, H, T, d), torch.randn(B, H, T, d), torch.randn(B, H, T, d)

# 手写版：完全照抄论文公式 + 因果 mask
scores = q @ k.transpose(-2, -1) / math.sqrt(d)          # (B,H,T,T) 相关性
mask = torch.triu(torch.ones(T, T), diagonal=1).bool()   # 上三角=未来，屏蔽
scores = scores.masked_fill(mask, float("-inf"))
attn = F.softmax(scores, dim=-1)
out_manual = attn @ v

# API 版：一行
out_sdpa = F.scaled_dot_product_attention(q, k, v, is_causal=True)

print(torch.allclose(out_manual, out_sdpa, atol=1e-5))   # True —— 你写的和 FlashAttn 数学等价

手写 attention，对照 SDPA

8.2 问题：从小 GPT 到 LLaMA，哪些组件改变了

解决什么：LLaMA 系列把多项成熟设计组合成现代 decoder-only LLM 的代表架构。LLaMA 1 采用 MHA、RoPE、RMSNorm 与 SwiGLU；GQA 则主要出现在 LLaMA 2 70B、Llama 3 等后续版本。它们共同展示了现代 LLM 架构如何逐步演进。

组件	原始 Transformer	LLaMA 系列（现代）	为什么换
位置编码	绝对/正弦位置嵌入	RoPE 旋转位置编码	以旋转方式编码相对位置信息；超出训练长度通常还需 RoPE scaling 等适配，不能保证直接外推
归一化	LayerNorm + Post-Norm	RMSNorm + Pre-Norm	更省算、深层训练更稳
激活/MLP	ReLU MLP	SwiGLU	同参数量下效果更好
注意力	多头 MHA	LLaMA 2 70B、Llama 3 等采用 GQA	推理时 KV cache 显存大降

打开 SVG 原图

对应你的代码：把玩具 GPT 的 LayerNorm→RMSNorm、可学习位置嵌入 →RoPE、GELU-MLP→SwiGLU，可以得到 LLaMA 1 风格的主干；再把 MHA→GQA，则更接近 LLaMA 2 70B 与 Llama 3 等后续版本。这里讨论的是架构家族的演进，不应把所有组件都归到 2023 年的原始 LLaMA。

动手复现：先实现最有代表性的 RMSNorm 和 GQA 的 KV 复制，这两个最能体现「现代」和「旧」的差别。

import torch, torch.nn as nn

class RMSNorm(nn.Module):
    """比 LayerNorm 少算一个均值：只按均方根缩放，更省更稳。"""
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
    def forward(self, x):
        rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x * rms * self.weight

def repeat_kv(x, n_rep):
    """GQA：KV 头数少于 Q 头数，把每个 KV 头复制 n_rep 次去匹配 Q。
    推理时 KV cache 只存少量 KV 头 -> 显存大降。"""
    B, T, n_kv, d = x.shape
    if n_rep == 1:
        return x
    return (x[:, :, :, None, :]
            .expand(B, T, n_kv, n_rep, d)
            .reshape(B, T, n_kv * n_rep, d))

RMSNorm 与 GQA 的 repeat_kv（LLaMA 核心片段）

🔗

推荐对照读：Meta 官方 llama 仓库的 model.py（RoPE/GQA/RMSNorm 标准实现，仅几百行）和 Karpathy nanoGPT。把论文里的图和这些代码并排看，理解会非常快。

8.3 问题：Attention 的中间显存与读写怎样减少

解决什么：朴素注意力要显式生成 $T\times T$ 的分数矩阵，序列一长，这个中间矩阵的显存和读写就爆炸（$\mathcal{O}(T^2)$ 显存）。FlashAttention 的洞见是「瓶颈不在算力，在显存读写（IO）」：它把 Q/K/V 分块（tiling）载入片上高速 SRAM，用在线 softmax 逐块累加，全程不把完整的 $T\times T$ 矩阵写回显存，从而省显存又更快。这是 IO-aware 算法设计的典范。

打开 SVG 原图

对应你的代码：F.scaled_dot_product_attention 会根据硬件、dtype、张量形状和运行环境，在 Flash、memory-efficient 或数学实现之间调度。H800 通常能从 Flash 类后端中明显受益，但具体实现及版本取决于当前 PyTorch、CUDA 与算子条件，不能仅凭调用 SDPA 就断言一定是 FlashAttention-3。

import torch, torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend

q, k, v = (torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.bfloat16) for _ in range(3))

# 强制只允许 PyTorch 的 Flash SDPA 后端；能跑通说明该输入满足 Flash 后端条件
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(out.shape)   # 具体底层实现与版本取决于 PyTorch/CUDA 构建和输入条件

8.4 问题：训练状态为何冗余，ZeRO / FSDP 怎样分担

ZeRO 提出按阶段分担优化器、梯度和参数状态；《PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel》（VLDB 2023）讨论相应思想在 PyTorch FSDP 中的工程实现。先以第 5 章的参数生命周期为坐标阅读，再区分论文中的 FSDP1 与本教程采用的 FSDP2 fully_shard API。

先按明确的 dtype 估算常驻状态：纯 FP32 参数、FP32 梯度与 Adam 两份 FP32 动量，合计约 16 字节/参数，未计激活、临时聚合参数、通信缓冲和分配器缓存。混合精度是否另有 FP32 主权重由实现决定，不能重复计算。分片降低其中可分担状态的冗余，训练峰值仍需测量。

打开 SVG 原图

方式	跨 rank 分担什么	估算范围
DDP	通常保留完整参数、梯度与优化器状态副本	本例 FP32 Adam 常驻状态约 16 字节/参数
ZeRO-1	优化器状态	其他状态仍复制
ZeRO-2	优化器状态和梯度	参数仍复制
ZeRO-3 / FSDP 全分片思想	再分担参数；计算时按需聚合	可分片部分均匀时约为 $\frac{1}{W}$，整体峰值不按 $\frac{1}{W}$ 保证

这篇论文真正的干货（读它主要就读这四点，都是工程经验）：

论文机制	解决什么	对应到你的实践
FlatParameter（把一个分片单元里所有参数摊平拼成一个大张量再切）	小张量各自通信太碎、效率低；合成一个大 tensor 通信更高效	FSDP2 同样可把一个 Transformer block 的参数组织成集合通信组，但参数仍保持各自的 DTensor 身份，不再合并为一个 FlatParameter
Deferred Initialization（延迟初始化：先在 meta 设备上「录制」建模操作，再逐单元真正初始化并分片）	模型大到单卡连「完整建出来」都会 OOM	为什么能训远超单卡显存的模型；对应 torch.device("meta") 建模
通信-计算重叠（backward prefetch：算这一层时提前 all-gather 下一层参数，靠算子重排把通信气泡藏起来）	分片必然引入通信，会阻塞计算产生气泡	FSDP 高效的关键，框架自动做；你无需手写
Rate Limiter（限制在途未分片参数占的显存块数，必要时暂停 CPU 下发）	CUDA caching allocator + 多 stream 下显存会被 prefetch 撑爆	解释了 FSDP 为什么能稳定压住峰值显存

论文还区分了 Full Sharding（全切，每卡 $\frac{1}{N}$）与 Hybrid Sharding（混合：节点内分片、节点间复制），后者在多机时能大幅减少跨机通信——这正是 6.5 多机场景该关心的调优点。

关于 FSDP2 有没有论文

本教程对照的 VLDB 2023 论文讲 FSDP1：以 FlatParameter 管理一组参数。FSDP2 的行为以 PyTorch 2.9.x 官方 API 与设计说明为准，采用 per-parameter DTensor 表达参数分片；TorchTitan 的论文与源码可用于理解多维并行的工程组合。不要把旧论文中的限制、优化器处理或预取行为直接套在另一个版本。

	FSDP1（VLDB 2023 论文）	FSDP2（第 5 章 用的）
API	FullyShardedDataParallel 包一层（类）	fully_shard(module)（函数式）
分片方式	FlatParameter：整单元摊平成一个大张量再切	per-parameter：每个参数各自按 DTensor 切
和 TP/PP 组合	较别扭	原生顺畅（都在 DeviceMesh 上）
文档出处	VLDB 2023 论文	PyTorch RFC/文档 + torchtitan 论文

复现实验回到 5.4—5.5：先验证同一初始权重与样本的一次更新，再比较同一训练状态和计时区间的显存。下面保留一个接口串联片段，GPT/get_batch 需自行接入；它只显示 rank 0 的单次峰值，不能作为完整性能报告。超单卡模型的 meta 初始化属于后续能力，附件尚未实现。

import torch, torch.distributed as dist, os
# torchrun --nproc_per_node=8 compare.py
dist.init_process_group("nccl")
lr = int(os.environ["LOCAL_RANK"]); torch.cuda.set_device(lr)

model = GPT(vocab=32000, dim=2048, n_layers=24).to(lr)  # 对比实验需选择能被 DDP 单卡容纳的规模

USE_FSDP = True
if USE_FSDP:
    from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
    mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
    for blk in model.blocks:
        fully_shard(blk, mp_policy=mp)
    fully_shard(model, mp_policy=mp)
else:
    from torch.nn.parallel import DistributedDataParallel as DDP
    model = DDP(model, device_ids=[lr])

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, fused=True)
x, y = get_batch()
_, loss = model(x, y); loss.backward(); opt.step()

peak = torch.cuda.max_memory_allocated() / 1e9
if dist.get_rank() == 0:
    print(f"{'FSDP2' if USE_FSDP else 'DDP'} 单卡峰值显存: {peak:.2f} GB")
# 切换 USE_FSDP 跑两次对比：模型越大，FSDP 省得越多

对比 DDP 与 FSDP2 的单卡峰值显存

🔗

论文：ZeRO《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》(2019)；PyTorch FSDP《PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel》(VLDB 2023，即 vldb.org/pvldb/vol16/p3848-huang.pdf)；FSDP2/多维并行看 torchtitan 论文 (arXiv 2410.06511)。想看生产级实现，直接读 PyTorch 官方 torchtitan 仓库——它把 FSDP2 + TP + PP 的论文思想落成了干净代码。

FSDP2 的原理：为什么改成 per-parameter 分片

FSDP1 的 FlatParameter 把一个单元里所有参数摊平拼成一个大一维张量再切——通信高效，但副作用很烦：切出来的分片和「原始的某个权重」对不上，导致 state_dict 要做复杂映射、和 TP（张量并行本身也要切同一个权重）组合时两套切法打架。FSDP2 的做法是不摊平，每个参数各自按第 0 维切成分片，用 DTensor 记录「这块是完整参数的哪一片」。DTensor = 普通张量 + 一份「分布信息（在哪个 DeviceMesh 上、按哪维切/复制）」的元数据。

打开 SVG 原图

FSDP2 保留具名参数的 DTensor 身份，便于 state_dict 表达与分布式布局组合；实际保存仍需使用正确的状态转换接口，并非无需处理。与 TP 组合要显式设计 DeviceMesh、分片布局和边界。混合精度策略决定计算与归约 dtype，其粒度和限制应按当前 API 检查，不能归因于“有 DTensor 就自动完成”。

动手对照：亲眼看到 per-parameter 分片

下面这段能直接跑，验证两件事：分片后每个参数变成了 DTensor、且它记得自己按哪一维切；用 full_tensor() 又能聚合回完整权重。跑一次胜过读十遍论文。

import torch, torch.nn as nn, torch.distributed as dist, os
from torch.distributed.fsdp import fully_shard
from torch.distributed.tensor import DTensor
# torchrun --nproc_per_node=2 inspect_fsdp2.py
dist.init_process_group("nccl")
lr = int(os.environ["LOCAL_RANK"]); torch.cuda.set_device(lr)

model = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 8)).to(lr)
fully_shard(model)                       # FSDP2：每个参数各自切成 DTensor

w = model[0].weight                      # 原本 (16, 8) 的完整权重
print(type(w))                           # <class 'torch.distributed.tensor.DTensor'> —— 不再是普通 Tensor
print("完整逻辑形状:", w.shape)          # torch.Size([16, 8]) 对用户仍是完整的
print("本卡实际持有分片:", w.to_local().shape)  # 2 卡时约 (8, 8) —— 只存了一半
print("分布信息 placements:", w.placements)     # (Shard(dim=0),) —— 明确按第0维切

full = w.full_tensor()                   # all-gather 回完整权重
print("聚合回完整:", full.shape)         # (16, 8)，且 rank 间一致
dist.destroy_process_group()

打印 FSDP2 分片后的 DTensor placement

以本节的 Linear(8, 16) 为例：DDP 会在每个 rank 保留完整的 $(16,8)$ 权重副本；FSDP2 的参数是带有全局形状与分片布局信息的 DTensor。在示例的按第 0 维分片且均匀可分条件下，2 卡的本地权重为 $(8,8)$，4 卡为 $(4,8)$。用 to_local().shape 观察本地分片，并结合计算期间的 all-gather 理解其生命周期；这不意味着峰值训练显存等于单卡的 $\frac{1}{N}$。

进阶实验：把 to_local().shape 与同口径峰值显存记录放在一起。前者证明参数如何分片，后者反映整段训练的资源需求；若本地分片变小但峰值未下降，继续检查激活、临时聚合参数和通信缓冲。

8.5 问题：怎样把一层矩阵计算分给多个 rank

解决什么：当单层权重矩阵本身就大到放不下、或想进一步降低单卡显存和延迟时，把单层内部的矩阵按行/列切到多张卡上并行算，这就是张量并行（TP）。Megatron-LM 给出了 attention 和 MLP 的标准张量并行切法。每个“列并行 → 行并行”子块通常在前向和反向各需要一次集合通信；一个完整 Transformer 层同时包含 attention 与 MLP，不能简单概括为整层总共只有两次通信。

对应 5.6 的方式对比、7.4 的显式 TP 通信和 B.7 的 DeviceMesh。TP 通常优先放在高速互联域，是否跨节点应依据网络、模型形状与实际通信时间判断。下面用 PyTorch parallelize_module 观察切分，并与完整权重参考输出比较；它不是 Picotron 自定义 TP 类的直接替代。

Megatron 的原理：列并行与行并行如何减少通信

Megatron 的精髓是把一层拆成两个 Linear、一列切一行切，让中间激活天然对齐，从而整段只在末尾做一次通信。以 MLP（Linear1 → GELU → Linear2）为例：第一层按「输出维（列）」切（Column Parallel），每张卡各算出一部分隐藏神经元；GELU 是逐元素的，各卡对自己那部分直接算、无需通信；第二层按「输入维（行）」切（Row Parallel），每张卡拿自己那半隐藏层乘出一个局部部分和，最后一次 all-reduce 把各卡的部分和相加，就得到完整输出。对这一个 MLP 子块而言，前向与反向各发生一次关键集合通信。完整 Transformer 层还包含 attention 子块，因此总通信次数需要按具体并行计划计算。

打开 SVG 原图

Attention 也可以按 head 分工：各 rank 计算部分 Q/K/V 头，输出投影采用行并行后再汇总结果。经典 Megatron 方案中的 f 在前向为恒等、反向做 all-reduce；g 在前向做 all-reduce、反向为恒等。列并行接行并行可以减少子层间的通信，但每层仍有频繁的集合通信。TP 通常优先放在 NVLink 等高速互联域内；跨节点是否合适取决于网络带宽、延迟、拓扑、模型形状和并行实现，需要实测，不能概括为只能用于单节点。

动手对照：用 parallelize_module 在 2 卡切一层 MLP

下面的两卡实验将 fc1 按输出维切、fc2 按输入维切。所有 rank 使用相同完整权重和输入，先保存未分片参考，再切分并核对输出；print shape 只验证布局，assert_close 才检查这组输入下的数值。完整训练正确性还需比较梯度和一次更新。

import copy
import torch, torch.nn as nn, torch.distributed as dist, os
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    parallelize_module, ColwiseParallel, RowwiseParallel)
# torchrun --nproc_per_node=2 tp_mlp.py
dist.init_process_group("nccl")
lr = int(os.environ["LOCAL_RANK"]); torch.cuda.set_device(lr)
mesh = init_device_mesh("cuda", (2,), mesh_dim_names=("tp",))

class MLP(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.fc1 = nn.Linear(dim, hidden)      # 列并行：按输出维切
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden, dim)      # 行并行：按输入维切
    def forward(self, x):
        return self.fc2(self.act(self.fc1(x)))

torch.manual_seed(1337)
model = MLP(512, 2048).to(lr)
reference = copy.deepcopy(model)
# 关键：fc1 列并行、fc2 行并行 —— 中间激活天然对齐，整段只在 fc2 后做一次 all-reduce
parallelize_module(model, mesh, {
    "fc1": ColwiseParallel(),
    "fc2": RowwiseParallel(),
})

r = dist.get_rank()
print(f"rank{r} fc1.weight 本卡分片:", model.fc1.weight.to_local().shape)  # (1024, 512) 各持一半输出
print(f"rank{r} fc2.weight 本卡分片:", model.fc2.weight.to_local().shape)  # (512, 1024) 各持一半输入

x = torch.randn(4, 512, device=lr)
y = model(x)                       # fc1 各切各算 -> fc2 出局部部分和 -> 一次 all-reduce 求和
print(f"rank{r} 输出形状:", y.shape)   # (4, 512)，与单卡数值等价
torch.testing.assert_close(y, reference(x), atol=1e-5, rtol=1e-4)
print(f"rank{r} forward comparison passed")
dist.destroy_process_group()

tp_mlp.py：布局与前向数值对照，待两卡验证

对照 8.4 的 FSDP2：FSDP2 切的是不同参数在卡间的存放（用到时 all-gather 出完整层再算），本质是省显存；Megatron TP 切的是同一层内部的矩阵，让一层的计算真正拆到多卡并行，本质是拆开单层的计算与显存。两者正交、可叠加——真实千亿模型就是把它们挂在同一个 DeviceMesh 的不同维度上组合使用（B.7 里 init_device_mesh(..., mesh_dim_names=("dp","tp")) 就是为此）。

前向比较通过后，再对同一 scalar loss 比较输入梯度、拼回后的参数梯度和一次优化器更新。只有相同权重、输入、精度与合理容差下的对照，才能作为该配置的数值证据。

8.6 问题：如何只训练少量参数完成微调

LoRA 冻结原权重 W，只训练低秩增量 $BA$。若 W 的形状是 $(\mathrm{out},\mathrm{in})$，$\operatorname{shape}(A)=(r,\mathrm{in})$、$\operatorname{shape}(B)=(\mathrm{out},r)$，新增可训练参数是 $r\times(\mathrm{in}+\mathrm{out})$；节省比例取决于 r 与挂载模块。激活和冻结权重仍占内存。QLoRA 进一步量化基础权重，具体能否在某张卡上运行取决于模型、上下文、batch、量化格式和内核。

打开 SVG 原图

对应你的实践：9.2 提到的 LoRA/QLoRA。实战直接用 HuggingFace peft，几行就能给任意模型挂上 LoRA。

import torch, torch.nn as nn

class LoRALinear(nn.Module):
    """冻结原 Linear，只训练低秩的 A、B。"""
    def __init__(self, base: nn.Linear, r=8, alpha=16):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False              # 冻结原权重
        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)
        self.B = nn.Parameter(torch.zeros(base.out_features, r))  # 初始为0 -> 起点等于原模型
        self.scale = alpha / r
    def forward(self, x):
        return self.base(x) + self.scale * (x @ self.A.t() @ self.B.t())

# 实战里不用手写，一行挂上：
# from peft import LoraConfig, get_peft_model
# model = get_peft_model(model, LoraConfig(r=8, target_modules=["q_proj","v_proj"]))

LoRA 的本质：一个 nn.Module 就能说清

8.7 还该知道的几篇（按需精读）

论文	一句话贡献	关联
GPT-2/3	只靠 decoder + next-token 预测 + 规模 = 通用能力	2.3 的语言建模目标
Chinchilla (scaling law)	给定算力，模型大小和数据量要按比例配，别只堆参数	决定你该训多大模型、喂多少数据
Mixed Precision Training	fp16/bf16 训练的理论与 loss scaling	2.4 / 第 5 章后的 bf16 autocast
RoPE	用旋转方式注入相对位置信息；长上下文外推通常还需 RoPE scaling 等策略	8.2 LLaMA 的位置编码
PagedAttention (vLLM)	像 OS 分页管理 KV cache，提升推理吞吐	9.2 推理部署

本章验收按所选问题完成一项即可：写出问题、预期、唯一改动与运行条件；提交数值正确性证据，再记录资源或质量变化。Attention 对齐 mask/梯度，FSDP 对齐训练状态，TP 对齐完整权重，LoRA 核对可训练参数与基座冻结。没有改善时，也要能用证据说明瓶颈与适用边界。

---

9. 学习资源与后续工程方向

9.1 学习路线与核心资源

主线为 1→2→3→4→6；第 5 章按状态分片需求学习，第 7 章可在 nanoGPT 与 DDP 基础后学习。第 8 章与附录按当前问题检索。选择下一份材料之前，先指出自己缺少的是模型机制、状态恢复、容量、速度还是部署能力，再确定要交付的实验记录。GPU 运行概念在 1.7，SM/warp/内存与 CUDA kernel 细节在 A.9。

类别	资源
官方	pytorch.org 文档/教程；torchtitan（3D 并行必读）；gpt-fast（推理优化最小实现）；torchao（量化）；PyTorch 官方博客
从零理解	Karpathy nanoGPT / build-nanogpt / 「Zero to Hero」系列；Meta llama 模型代码
生态	HuggingFace transformers / datasets / accelerate / peft / trl；vLLM / SGLang
论文	Attention Is All You Need；LLaMA / LLaMA 2；RoPE、FlashAttention 1/2/3、GQA、SwiGLU、Chinchilla（scaling law）

结合本教程阅读 Ultrascale Playbook。它适合在 PyTorch 与 nanoGPT 单卡基础之后阅读：2.4 对照 First Steps 的内存与重计算；4～6 章对照 Data Parallelism、ZeRO、NCCL 和 profiling；7.4～7.9 对照 TP/SP、PP、CP、EP；5.6.1 对照 Finding the Best Training Configuration。先完成本教程的小实验，再用原文的更大规模案例解释为什么配置会随模型与互联改变。

参考资料为 The Ultra-Scale Playbook: Training LLMs on GPU Clusters，本次核对日期为 2026-10-13，并保留 Space 提交 c93890dd 的正文快照。原文配置搜索的热图基于其 H100 集群、完整序列长度 4096、约 100 万 token 的 global batch 等条件；曲线、MFU 和最优配置均为作者报告，不是本教程复现实测。本文新增脚本的验证范围单独写在 C.10。

9.2 后续方向：推理与部署

训练结束后，先确定推理工作负载：模型与上下文长度、并发数、首 token 延迟、持续生成速度、显存和成本约束。用同一组请求比较候选设备和引擎，再作选型；不能只凭 H20/H800 型号或训练配置判断部署方案。KV cache、量化和调度分别改变不同资源开销。

方向	改变什么	怎样验证
KV cache	保存历史 K/V，减少逐 token 生成中的重复计算；占用随上下文增长	固定提示与生成长度，比较输出、首 token 延迟、后续 token 延迟和显存
PagedAttention	分页组织 KV cache，改善动态请求下的内存管理	固定请求长度分布、并发和调度策略，比较容量、延迟与吞吐
量化	改变权重/激活的表示与相应计算内核	核对设备与算子支持；同时测任务质量、延迟和显存
推理引擎	采用 vLLM、SGLang 或 TensorRT-LLM 等实现调度和执行	先核对模型架构支持，再做相同请求集测试
LoRA / QLoRA	减少需要训练的参数；量化基座进一步改变存储	统计可训练参数、检查冻结状态并报告验证集效果

# 前置：model 是你训练好或加载的模型（如 2.3 的 GPT，或下面 vLLM 直接拉 HF 权重）
# 量化：torchao int4 权重量化
from torchao.quantization import quantize_, int4_weight_only
quantize_(model, int4_weight_only())

# 生产推理直接用 vLLM（命令行起 OpenAI 兼容服务）
# pip install vllm
# vllm serve meta-llama/Llama-3.1-8B-Instruct --tensor-parallel-size 2

9.3 成熟框架：从教学实验走向工程实践

从教学实现进入工程时，用同一能力清单比较框架：支持的模型与并行布局、checkpoint/恢复边界、可观测性、依赖维护和部署环境。先用已经掌握的小模型完成一次训练与恢复，再迁移规模；框架名称本身不能替代验收。

框架	定位
torchtitan（PyTorch 官方）	FSDP2 + TP + PP 的参考实现，代码干净，强烈推荐通读
Megatron-LM / DeepSpeed	工业级大规模训练；DeepSpeed ZeRO 与 FSDP 属于相近的全分片思路，但实现与接口不同
HuggingFace accelerate / Trainer	快速起步，少写样板
vLLM / SGLang	推理部署事实标准

最终交付应是一份可复查的实验记录：固定提交和环境、输入/状态/输出说明、改动位置、运行命令、正确性结果、独立性能结果及仍未验证的边界。读完后能解释、定位、修改并验证同一个训练系统，比记住多少框架或并行缩写更能说明学习是否完成。

附录 A：关键概念速查

使用方式：本附录用于查漏补缺，不要求顺序通读。正文已经讲清的概念会在这里以更短的形式重新归纳，方便忘记时快速回来查询。

概念入口：优化器与 Adam/AdamW 见 A.3；标量/向量/矩阵/Tensor 与维度见 A.11；梯度与一次更新见 A.12；mini-batch、micro-batch、global batch 见 A.13；矩阵计算与架构形状链见 A.14；DP、TP、PP、CP、SP 及 DDP/FSDP 的关系见 A.15；MoE、router 与 EP 见 A.16。本附录以概念定义、必要公式和张量形状为主；完整机制、改造步骤与运行实验请沿各节链接阅读正文和附录 C。

A.1 激活函数：给网络加上非线性

没有激活函数，多层线性层叠起来还是一个线性变换，学不了复杂模式。激活函数就是夹在层与层之间的那个非线性「掰弯」操作。

函数	公式 / 输出范围	特点与用途
Sigmoid	$\sigma(x)=1/(1+e^{-x})$，输出 $(0,1)$	把任意数压到 0~1，适合表示概率。缺点是两端「饱和」会导致梯度消失，现在隐藏层基本不用，只保留在二分类输出或门控里
Tanh	输出 $(-1,1)$	相当于零中心版的 Sigmoid，比它好训练，但同样会饱和
ReLU	$$\max(0,x)$$	正区间不饱和、计算极快，是深层网络最常用的默认选择。缺点是负区间恒为 0，可能出现「死神经元」
GELU / SiLU	ReLU 的平滑版	在 0 附近更平滑，Transformer 常用：GPT 用 GELU，LLaMA 用 SiLU（配合 SwiGLU，见 8.2）

A.2 Softmax 与 CrossEntropyLoss：查漏补缺

Softmax 是归一化函数：把一组实数分数转成非负、总和为 1 的分布；它本身不是损失函数。交叉熵是损失函数，用来衡量预测分布与目标分布的差异。

设有 $K$ 个类别，logits 为 $z\in\mathbb{R}^{K}$，类别索引为 $0,\ldots,K-1$。Softmax 输出 $p\in\mathbb{R}^{K}$，其中 $p_i=\frac{e^{z_i}}{\sum_{j=0}^{K-1}e^{z_j}}$。本教程使用整数目标类别 $y$；单样本交叉熵为 $L=-\log p_y$。目标类别的预测概率越高，这个损失越小。

PyTorch 的 CrossEntropyLoss 直接接收 logits 和目标类别，内部稳定地计算 log-softmax 与负对数似然，不需要先调用 Softmax。对 $B$ 个分类样本，logits 的 shape 为 $(B,K)$，整数目标为 $(B,)$；默认 mean 归约得到 shape 为 $()$ 的标量损失。

归一化轴由含义决定：分类和语言模型的输出沿类别或词表轴归一化；Attention 权重沿键位置轴归一化。稳定实现先减去该轴最大值再取指数；有限精度下极小概率可能下溢为零。模型中的使用位置见 2.2—2.3。

A.3 优化器：从梯度到 SGD、Adam 与 AdamW

优化器接收参数的当前梯度，并按规则改变参数和自身状态。梯度 $g_t=\frac{\partial L_t}{\partial\theta}$ 描述本次目标对当前参数的导数，优化器不是生成梯度的算法。这里 $t=1,2,\ldots$ 表示参数更新次数，$\theta$ 是某个参数张量；下列运算对它的各元素分别执行。入门对应 1.5—1.6，矩阵梯度见 A.12。

SGD 的基本更新是 $\theta_t=\theta_{t-1}-\eta g_t$，$\eta$ 是学习率。随机性来自抽样：可以每次用一个样本，也可以用 mini-batch；torch.optim.SGD 这个名称不表示 batch 必须为 1。加入 momentum 后需额外保存历史方向；“加动量一定能越过局部极小值”或“一定更快”都不是保证。

Adam（Adaptive Moment Estimation）同时跟踪一阶矩 m：梯度的指数移动平均，以及二阶原始矩 v：梯度平方的指数移动平均。v 不是减去均值后的方差，也不是 Hessian/二阶导数。它按分量调整步长尺度；仍要选择全局学习率 $\eta$、衰减系数 $\beta_1,\beta_2$ 和数值稳定项 $\varepsilon$，收敛速度与泛化效果依赖任务和配置。

AdamW 在 Adam 的自适应更新之外，直接对参数做权重衰减。下列讲解采用标准 Adam（无权重衰减、无 AMSGrad）和解耦的 AdamW；PyTorch 某些版本也支持通过 Adam 的选项启用解耦，阅读代码时应核对版本与参数。权重衰减不等于梯度裁剪，Adam 与 AdamW 的区别也不是有无学习率调度。

A.3.1 Adam 怎样使用两份历史状态

初始化 $m_0=0$、$v_0=0$。对本次梯度 $g_t$，先按下列公式更新状态，再作偏差修正，最后改变参数。$\beta_1$、$\beta_2$、$\eta$、$\varepsilon$、更新计数 t 是标量；$\theta$、g、m、v 的 shape 相同。$g_t^2$、平方根、乘除都逐元素执行，不是矩阵乘。

$m_t=\beta_1m_{t-1}+(1-\beta_1)g_t$；$v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2$。

$\hat m_t=\frac{m_t}{1-\beta_1^t}$；$\hat v_t=\frac{v_t}{1-\beta_2^t}$；$\theta_t=\theta_{t-1}-\eta\frac{\hat m_t}{\sqrt{\hat v_t}+\varepsilon}$。

两份历史状态从零开始时，早期估计会偏小；偏差修正补偿这一初始化影响。二阶原始矩决定逐分量的缩放尺度，不是 Hessian，也不代表某个参数的重要性。具体数值与代码验证见 C.9.2。

A.3.2 参数、梯度与优化器状态的 shape

张量或操作	weight$(2,3)$ 示例	何时改变 / 是否跨 step 保留
参数 $\theta$	$(2,3)$，6 个权重	step 更新，跨 step 保留
梯度 $g=\texttt{.grad}$	$$(2,3)$$	backward 累积；zero_grad 清理
一阶矩 $m=\texttt{exp\_avg}$	$$(2,3)$$	Adam step 更新，zero_grad 不清理
二阶原始矩 $v=\texttt{exp\_avg\_sq}$	$$(2,3)$$	Adam step 更新，zero_grad 不清理
bias 及其 g/m/v	每项 $(2,)$	和 weight 分开管理；不是共用一个标量状态
计数 t 与超参数	逻辑上 scalar	计数/状态表示细节随实现而定；随 optimizer checkpoint 恢复

因此训练恢复需要模型参数、优化器状态和更新进度；只加载模型权重会重新开始 m/v 的历史。纯 FP32 的参数、梯度和两份 Adam 状态，按每个元素 4 字节计算约为 16 字节/参数，未含激活、临时缓冲等。混合精度与分片会改变实际布局和 dtype，不应直接套用这个估算。

A.3.3 AdamW：把权重衰减与梯度更新分开

对启用衰减的参数，标准 AdamW 可写为 $\theta_t=(1-\eta\lambda)\theta_{t-1}-\eta\frac{\hat m_t}{\sqrt{\hat v_t}+\varepsilon}$，$\lambda$ 为权重衰减系数。m/v 基于损失梯度更新；若把 $\lambda\theta$ 直接加进 g 再交给自适应分母，是耦合 L2 正则的另一种行为。bias、归一化参数是否衰减取决于参数分组，本教程 nanoGPT 的具体分组见 3.6。

PyTorch 2.9 Adam；PyTorch 2.9 AdamW。

A.4 各种 Norm：稳定训练的归一化

不同归一化方法使用不同的统计量和轴。LayerNorm 按指定特征维做中心化与缩放，并可带可学习仿射参数；RMSNorm 按均方根缩放，不减均值，因此不能概括为所有 Norm 都把输出变成均值 0、方差 1。2.3.4 先学习 LayerNorm 在 Pre-Norm Block 中的位置。

方法	统计范围与核心计算	训练与评估时的区别
BatchNorm	训练时按通道汇总 batch 及适用的空间或序列轴，计算均值和方差，再归一化；通常带可学习缩放和平移。	常见默认配置在评估时使用累计统计量，因此同一个输入在 train/eval 下可能采用不同统计量。
LayerNorm	对每个样本的指定特征轴计算均值和方差。例如输入 $(B,T,C)$ 时，可对每个位置的 $C$ 个特征归一化；通常带可学习仿射参数。	通常在训练与评估时都使用当前输入在指定特征轴上的统计量。
RMSNorm	对指定特征轴按均方根缩放，不减去均值；通常带可学习缩放。	通常在训练与评估时都根据当前输入计算均方根。

这里 $B$ 是样本条数，$T$ 是序列位置数，$C$ 是特征数。三者的关键区别是统计哪些轴、是否减去均值，以及是否在评估时改用累计统计；具体参数和默认行为仍以所用模块为准。

A.5 QKV：注意力机制快速复习

QKV 是 Query（查询）、Key（键）、Value（值）的缩写，是自注意力（self-attention）的核心。每个 token 的向量分别乘上三个不同的线性投影矩阵，得到自己的 q、k、v。可以用检索来类比：q 是「我要找什么」，k 是每个 token「我能提供什么」的标签，v 是它实际携带的信息。

单头时，Q、K、Value 为 $(B,T,D)$，$QK^{\mathsf T}$ 得到分数 $(B,T,T)$，沿最后的键位置轴 softmax，权重再乘 Value 得到 $(B,T,D)$。多头只增加 H 轴为 $(B,H,T,D)$，并在最后两个轴进行上述矩阵乘；因果 mask 约束键位置，输出合头后回到 $(B,T,C)$。符号维度与具体例子见 2.3、A.14。

A.6 训练稳定与防过拟合的常用技巧

下面几个不是网络结构，而是训练过程中的「调节旋钮」，几乎每个训练脚本都会用到。

Dropout。训练时随机把一部分神经元的输出临时置 0（比如按 p=0.1 的比例），强迫网络不依赖某几个特定神经元，从而缓解过拟合。它只在训练时生效，推理时自动关闭——这正是 model.train() 和 model.eval() 要切换的原因之一。

学习率调度（learning rate schedule）。学习率不必全程固定。主流做法是 warmup + cosine：训练初期把学习率从 0 线性升上去（warmup，避免一开始就大步乱走），之后按余弦曲线缓慢降到接近 0（让模型后期精细收敛）。大模型训练几乎都用这个组合，见正文 2.4。

梯度裁剪（gradient clipping）。当某一步梯度突然变得很大时，参数会被猛地推飞导致训练发散（loss 变 NaN）。梯度裁剪用 clip_grad_norm_ 给梯度的整体范数设一个上限，超过就等比缩小，是大模型训练防炸的标准保险，见正文 2.4。

梯度累加：在一次参数更新之前处理多个 micro-batch，把按目标正确归一化的梯度累积到参数。等大小批次常把每批 mean loss 除以累积次数；变长或末批不足时按有效目标数加权。累积不增加模型容量，也不保证与大 batch 在 Dropout/BatchNorm 等条件下逐值一致；换算和验算见 A.13，OOM 定位见 1.8。

A.7 参数规模与超参数：两个容易混淆的「参数」

初学时最容易混的两个词，其实是完全不同的东西：一个是模型自己学出来的，一个是你训练前手动设定的。

模型参数规模是参数张量元素数量之和，例如 7B 表示约 70 亿个参数；这里 $\mathrm{B}=\text{Billion}$，与 batch 轴符号 B 无关。参数 $\theta$、梯度和 Adam 的 m/v 各自有明确 shape。若仅把 70 亿参数保存为每元素 2 字节，权重约 14 GB（十进制），训练还要加梯度、优化器状态、激活和临时缓冲；状态分片见第 5 章。模型质量还依赖架构、数据与训练，不能由参数量单独保证。

超参数（hyperparameter）。指训练开始前由你设定、训练过程中不通过梯度更新的配置项，用来控制「怎么训练」。常见的有：学习率（最关键）、batch size、训练轮数（epoch）、Dropout 比例、优化器选择、warmup 步数、层数与隐藏维度等。它们不会被模型自己学到，只能靠经验、实验或搜索来调——「调参」调的就是这些。一句话区分：参数是模型学出来的，超参数是你喂给训练过程的。

A.8 Logits：Softmax 之前那排「原始分数」

Logits 指模型最后一层还没经过 Softmax 的原始输出分数。它们是任意实数（可正可负、没有范围限制），数值本身没有直接的概率含义，只表示「模型对每个选项的偏好强弱」——越大代表越倾向。经过 Softmax（A.2）归一化后，这排 logits 才变成加起来为 1 的概率分布。

单样本三分类 $\mathrm{logits}=(2.0,1.0,-1.0)$ 是向量 shape=$(3,)$，softmax 后仍为 $(3,)$，数值约 $(0.7054,0.2595,0.0351)$。B 条样本的分类 logits 为 $(B,3)$；语言模型完整序列 logits 为 $(B,T,V)$，生成时取最后位置变成 $(B,V)$，再为每条序列选择下一个 token。

分类交叉熵接收 logits$(B,V)$ 与整数目标$(B,)$；语言模型可将 $(B,T,V)$ 展平为 $(B\times T,V)$，目标展平为 $(B\times T,)$，mean reduction 返回 $\text{scalar},\quad\operatorname{shape}=()$。不需要预先 softmax；只取 argmax 时也可直接比较 logits。图像分类完整例子见 C.1，语言模型见 2.2—2.3。

A.9 GPU 与 CUDA：SM、thread、warp、block、grid

这些术语回答两个不同问题：SM、计算单元和存储层次描述硬件有什么；grid、block、thread 描述一次 kernel 的逻辑工作怎样组织，warp 则是 block 内线程的执行分组。先分清这两层，再看它们怎样对应。正文 1.7 用这些概念解释训练；本节提供独立定义和边界，CUDA 代码练习留在 C.9.5。依据 CUDA C++ Programming Guide 13.0 的 §5.1–5.3、§7.1–7.2 与 §8.2–8.3，页码索引见 1.7。

A.9.1 SM、CUDA Core 与 Tensor Core：硬件执行资源

SM（Streaming Multiprocessor，流式多处理器）是 NVIDIA GPU 上承载线程块执行的硬件单元，一块 GPU 包含多个 SM。每个 SM 有 warp 调度器、算术执行单元、寄存器文件、共享内存等资源，能够驻留并交错推进多个 warp，用其他就绪 warp 的工作掩盖部分等待时间。一个 SM 能同时容纳多少 block/warp，受到寄存器、共享内存和硬件上限约束；数量不是越多就必然越快。

矩阵乘 kernel 的算法实现先决定如何把矩阵划成 tile、每个 block/thread 处理什么；硬件调度器再将 block 放到有可用资源的 SM 上。硬件调度器并不会自行把任意 Python 任务拆成矩阵小块。分块、访存方式与资源使用共同决定性能，不能只看启动线程够不够多。

CUDA Core通常指 SM 内执行一类标量算术指令的硬件单元，不等于功能完整的 CPU 核心，也不等于一个 CUDA thread。thread 是逻辑执行实例，会随着指令执行使用硬件资源，并非永久占据一个物理核心。SM 数量、每个 SM 的单元数量和支持的指令随 GPU 架构而变；设备查询方法见 1.7.4。

Tensor Core 是加速特定矩阵乘累加等操作的硬件单元，支持的精度与能力随架构变化。torch.Tensor 是软件的数据抽象，既可在 CPU 上，也可在 GPU 上；创建 Tensor 并不意味着一定使用 Tensor Core。是否走 Tensor Core 路径取决于硬件、算子、dtype、形状与库的实现，bf16/FP16 加速会在 2.4 与 B.3 继续讨论。

A.9.2 thread、warp、block 与 grid：定义与对应关系

一次 kernel 启动创建一个 grid；grid 由 block 组成，block 由 thread 组成。warp 不是需要额外指定的另一层启动维度：它是同一 block 内按线性线程编号划分的 32 个线程一组的执行分组。grid 和 block 的索引可以用一至三维表达，最终由 kernel 算法把线程索引映射到数据。

概念	是什么、为什么需要	关键边界
thread（线程）	kernel 的一个逻辑执行实例，有自己的索引、局部值和执行状态。它执行同一 kernel 代码中分配给自己的工作	一个线程可以处理多个元素；多个线程也可以协作计算一个结果。它不是操作系统线程或物理核心
warp（线程束）	同一 block 内按连续线性编号划分的每组 32 个线程，SM 以它组织 SIMT 指令执行与调度	尾组可能不满；分支或边界判断会使部分线程暂不参与指令。warp 不跨 block，也不保证线程隐式同步
block（线程块）	可以通过共享内存交换数据、进行块内同步的一组线程，提供明确的协作边界	普通 block 在一个 SM 上执行；一个 SM 可驻留多个 block。不同 block 不能假定执行顺序或普通的全 grid 同步
grid（网格）	一次 kernel 启动的全部 block，表达本次工作的总体规模	block 可以分批分配到多个 SM；grid 不等于 GPU 数量，也不要求全部 block 同时驻留

只用一个计数例子连接这些定义：长度为 $N=1000$ 的向量加法，若规定每个线程处理一个元素、每 block 有 $256$ 个线程，则需 $\lceil 1000/256\rceil=4$ 个 block，共启动 $1024$ 个线程。每 block 含 $256/32=8$ 个 warp，全 grid 共 $32$ 个 warp；用 i < N 保护多出的 24 个线程不访问越界元素。这是教学映射，PyTorch 实际 kernel 可以采用不同划分；4 个 block 也不要求 GPU 恰好有 4 个 SM。

张量的轴描述数据，例如 $(B,T,C)$；grid/block 的维度描述线程索引，两者由算法映射，不能按维数机械对应。batch_size=32 不表示启动 32 个 CUDA 线程，num_workers=4 也不表示 4 个 block。DDP 中常说的 rank 则是进程编号，与矩阵秩、张量轴数、warp 都是不同概念。

SIMT 表示单指令、多线程。一个 warp 内的分支使不同路径上的线程分批参与指令，可能降低有效并行度，不能简单理解成“所有 if 都慢”。Volta 起支持独立线程调度，不能假设同一 warp 的线程始终隐式同步；需要协作时必须采用正确的同步原语。指南中的 thread block cluster 是较新架构上的扩展，留到内核开发时选读。

打开 SVG 原图

block 与 SM 的关系是“逻辑工作单元被放到硬件上执行”。普通 block 的线程在同一个 SM 上使用相应资源；SM 可以先后处理许多 block，也可以在资源允许时同时驻留多个 block。块内共享内存用量和每线程寄存器用量都会限制驻留数量，所以仅把 block size 调大不保证更快。跨 block 的协作需要适用的专门机制，不能用块内同步代替。

A.9.3 内存作用域与位置：显存、缓存和片上资源

显存是 GPU 可访问的设备内存，H800 等设备使用 HBM，高带宽内存并不等于所有 GPU 显存的通用名称，其他设备也可能使用 GDDR。模型参数、激活和优化器状态通常存放在设备内存中；寄存器与共享内存更靠近计算单元，容量也更小。

层次	位置与用途	PyTorch 中的直觉
主机内存	CPU 一侧保存原始数据、加载后的 batch 等	DataLoader 通常先产出 CPU 张量
设备显存 / global memory	容量较大；数据中心 GPU 常用 HBM，其他设备也可用 GDDR	参数、梯度、激活、优化器状态等通常在这里占用空间
L2 / L1 缓存	L2 由多个 SM 共享，L1 靠近 SM；自动缓存部分访问	复用和访问局部性会影响访存成本
shared memory 共享内存	SM 上的片上资源，普通情形由同一 block 内线程显式协作使用	kernel 可暂存数据块，减少重复读取显存
register 寄存器	逻辑上属于线程；寄存器文件是 SM 的有限资源	保存正在使用的局部值；用量影响同时驻留的线程数

这是存储作用域与位置的示意，不是一条所有数据都必须顺序经过的流水线。shared memory 与 L1 的物理组织依架构而变；CUDA 的 local memory 虽然线程私有，却可能位于设备内存中（例如寄存器溢出），名称 local 不代表一定是快速片上存储。

打开 SVG 原图

memory-bound 指性能主要受数据访问限制；compute-bound 指主要受运算吞吐限制。同一算法在不同形状、精度、实现下也可能改变瓶颈。FlashAttention（8.3）利用分块与在线 softmax 等方法，减少大注意力中间矩阵在 HBM 的读写，并用片上资源复用数据；它并不是把整个 attention 永久装进 shared memory，也不能据此断言所有 attention 都受访存限制。

A.9.4 kernel、stream 与 occupancy：执行和测量术语

kernel是在 GPU 上执行的函数；一次普通主机发起的 kernel 启动用 grid/block 配置组织线程。stream（流）表达设备操作的顺序：同一 stream 内的操作按序执行，不同 stream 之间能否并发取决于依赖和硬件资源。stream、grid 与 SM 分别描述执行顺序、工作组织和硬件资源，三者不能混为一谈。向量加法代码及启动语法见 C.9.5，异步计时见 1.7.3。

普通 PyTorch 模型通常依靠框架、数学库或编译器生成的 kernel。等 profiling 显示瓶颈且现有实现不满足需求时，再考虑 Triton/CUDA 自定义算子。occupancy 的含义是某 SM 上活跃驻留 warp 数与硬件允许上限的比例；寄存器、shared memory、block/thread 上限都会约束它。较高 occupancy 可能帮助隐藏延迟，但不是利用率读数，也不保证性能更好。

💡

对照 1.7 复习：kernel 启动一个 grid；block 内线程划为 warp，并在 SM 上执行；数据访问涉及显存、缓存与片上资源。性能由计算、搬运、同步、启动成本与并行度共同决定，需要测量后再选优化方式。

A.10 PyTorch 核心概念：数据、模型、求导与执行

本节解释阅读 PyTorch 训练代码时必须分清的对象、状态和操作。沿 1.6.3 的运行架构理解它们，遇到术语时再回来查；完整的状态观测实验见 C.11。标量、向量与矩阵的数学概念和维度见 A.11，优化器算法见 A.3，梯度的数学含义见 A.12。

A.10.1 数据对象：Tensor、Storage、shape、dtype 与 device

Tensor（张量对象）是 PyTorch 表达数据及其运算的基本对象。对普通稠密张量，Storage 承载底层数据，Tensor 通过元信息解释这些数据。多个视图可以共享同一份 Storage，所以“有两个 Tensor 对象”不等于“存了两份数据”。这里讨论普通稠密 Tensor；稀疏等特殊布局有不同的存储表达。

属性	定义	应分清的边界
shape / ndim	各轴长度 / 轴数；shape 为 $(2,3)$ 的张量有 2 个轴、6 个元素	轴的业务含义由代码约定；shape 本身不说明哪一轴是 batch
dtype	元素的数据类型，如 float32、bfloat16、int64	影响数值范围、精度和每元素字节数；与 shape、device 是不同属性
device	张量的数据所在设备，如 cpu、cuda:0	在 Python 中持有 CUDA Tensor 对象，不表示其数值数据在 CPU 上
stride / storage_offset	沿各轴移动一步在 Storage 中跨过的元素数 / 起始元素偏移	shape 相同也可能布局不同；转置可以只改布局，reshape 则不保证零复制

view（视图）与原张量共享数据；clone 创建数据副本；detach 切断结果与原求导历史的关系，但仍可能共享数据。三者分别涉及布局、复制和求导关系，不能互相替代。连续布局与实际 stride 示例见 B.1。

A.10.2 模型对象：Module、Parameter 与 buffer

nn.Module 是组织模型的对象：它管理子模块、参数和缓冲区，并通过 forward 定义前向计算。通常调用 model(x)，让 Module 的调用机制执行 forward 及相关 hooks。nn.Parameter 是用于表示模型参数的 Tensor 子类；把它赋给 Module 的属性时会注册为参数，能由 parameters() 找到。普通 Tensor 即使 requires_grad=True，也不会因此自动注册为模型参数；已注册的参数也可以设置为不求梯度。

buffer（缓冲区）是随模型管理、但不作为注册参数交给优化器的 Tensor，例如 BatchNorm 的运行均值。用 register_buffer 注册后，它会随 model.to(device) 迁移；默认持久 buffer 也会进入模型的 state_dict，非持久 buffer 不进入。仅把普通 Tensor 放在属性里，不会自动获得这些管理行为。

state_dict 是按名字组织状态的映射。模型的 state_dict 主要保存参数和持久 buffer；优化器有自己的 state_dict，保存参数组配置与优化器状态。checkpoint 是训练程序实际保存的快照，若要恢复训练，通常还需优化器、训练步数及随机状态等。只保存模型权重，不等于保存了一次可完整恢复的训练现场；第 3、5 章分别验证单卡与分片恢复。

A.10.3 自动求导：计算图、requires_grad、叶子、grad 与 grad_fn

Autograd 是自动求导系统。默认 eager 模式下，它在前向执行时记录可微运算的依赖，形成当次计算图，供反向应用链式法则。这个图描述实际执行过的运算，Module 的子模块树描述模型对象如何组织，两者不是同一结构；同一个 Module 在一次前向中被调用两次，也会产生两次运算的依赖。

requires_grad 表示张量是否需要参与梯度计算；它本身不是梯度值，也不表示已经执行过 backward。对需要梯度的 Tensor，leaf（叶子）指它不是由另一项受记录的运算产生；常见可训练 Parameter 就是叶子。在普通训练中，反向会把所需的叶子梯度累积到 .grad；中间的非叶子结果默认不保留自己的 .grad，需要观察时可显式调用 retain_grad()。

.grad 存放计算得到的梯度 Tensor；对参数 $W$，其梯度 $\partial L/\partial W$ 与 $W$ 的 shape 相同。.grad_fn 指向产生非叶子结果的反向节点，描述后续怎样求导，不存放梯度值。因此，叶子参数 grad_fn=None 可以正常训练；非叶子结果有 grad_fn，也不表示它已经有了 .grad。

A.10.4 一次更新：loss、backward、optimizer 与 zero_grad

loss（损失）量化预测与训练目标之间的偏差。常见训练将各样本或各 token 的损失归约为一个标量 Tensor，shape 为 $()$，再从它开始反向。loss 是本次前向计算出来的结果，不会随参数更新自动重算；要观察更新后的损失，需要重新前向。样本和 token 的归约范围见 A.13。

下表按本文不传入 closure（用于重新计算损失的回调）的常规 SGD/AdamW 调用方式说明。

操作	做什么	不会自动做什么
loss.backward()	沿计算图求导，把梯度累积到叶子参数的 .grad	不会自动调用优化器更新权重，也不会先清除已有梯度
optimizer.step()	按更新规则读取梯度，更新所管理的参数及优化器状态	不会自动重新前向，也不会自动清理梯度
optimizer.zero_grad(set_to_none=True)	将所管理参数的累计梯度置为 None，为下一轮累积准备	不会重置参数，也不会清空 Adam 的历史矩估计

optimizer（优化器）实现“怎样使用梯度修改参数”的规则；它管理的参数由创建时或后续参数组配置显式指定，不会自动寻找程序里所有需要梯度的 Tensor。学习率决定更新尺度，Adam/AdamW 还保存随训练演进的状态，具体公式见 A.3。梯度累积是先做多次 backward、再做一次 step 的组织方式，micro-batch 与 global batch 的关系见 A.13。

A.10.5 数据与模式：Dataset、DataLoader、train、eval 与 no_grad

Dataset 定义数据如何取得，可以按索引访问，也可以按迭代方式提供；DataLoader 按配置组织取样、组 batch、合并样本和加载工作进程。它不会替训练循环执行 backward 或 step，num_workers 也不是 GPU 数量。每次取得的数据规模与一次更新包含的数据规模如何对应，取决于累积和并行策略，见 2.1 与 A.13。

model.train() 与 model.eval() 切换模型中对模式敏感的层的行为，例如 Dropout 和 BatchNorm；eval 不关闭求导，也不会保证推理结果自动正确。torch.no_grad() 在作用域内关闭常规反向求导的运算记录，减少不必要的图与保存值，但不改变层的训练/评估模式。评估时通常同时使用 eval 与 no_grad；inference_mode 进一步减少跟踪开销，但其中创建的 Tensor 用于后续求导计算时限制更严格。完整循环见 C.9.4。

A.10.6 执行方式：operator、Dispatcher、backend、kernel 与 compile

operator（算子）是具有明确输入、输出与语义的计算操作，如矩阵乘、求和；ATen 是 PyTorch 的核心张量算子库。Dispatcher（算子分派器）按张量与当前执行状态选择注册实现，backend（后端）提供相应设备或执行方式下的实现。分派还涉及自动求导等机制，不能理解成“只看 device 选一次 CPU/GPU”。shape、dtype 和布局会进一步影响实现内部采用的计算路径。

GPU kernel（设备内核）是在 GPU 上执行的函数；CUDA 如何把它组织为 grid、block、warp 和 thread，见 A.9。框架中的一个算子实现入口，不保证只启动一个 GPU kernel；视图操作还可能完全不启动计算 kernel。CUDA stream用于表达设备工作的执行顺序和依赖，CPU 提交结束不等于 GPU 已完成，计时与同步见 1.7.3。

eager 表示程序运行到张量操作时即执行相应调用，并动态记录求导依赖；它不表示 CUDA 工作一定同步完成。torch.compile 则尝试捕获、优化并编译可处理的运算片段。graph break 指编译捕获在某处中断；默认允许混合执行时可以回到 eager，fullgraph=True 遇到这类中断会报错。编译不会取消 backward 和 optimizer 的职责，也不承诺把整个训练变成一个 kernel；其使用与验证见 2.4。

术语依据 PyTorch 2.9 官方文档：Tensor Views、Module、Parameter、Autograd mechanics、Optimizers、Compiler。训练、验证、测试、推理的完整流程见 2.2；API 的输入、输出、默认值与常见报错检查见 C.9.4 与 1.8。

补充教学资料：李宏毅 Machine Learning 2023 Spring — PyTorch Tutorial 1 & 2（TA 蕭淇元），可配合 PyTorch 官方文档 与 NumPy 文档 阅读。训练循环与 API 阅读练习见 C.9.4。

A.11 标量、向量、矩阵与 Tensor：维度口径速查

正文初次学习见 1.2。这里的 Tensor 指 PyTorch 数值数组；数学张量的坐标变换性质不在本入门范围。看到“维度”时先判断作者说的是轴数、某轴长度，还是向量空间维数。

概念	示例 / shape	ndim / numel
scalar 标量	$$3.0;\quad\operatorname{shape}=()$$	0 / 1
vector 向量	$(1,2,3)$；$(3,)$	1 / 3；数学上有 3 个分量
行/列表示	$\begin{bmatrix}1&2&3\end{bmatrix}$ $(1,3)$；$\begin{bmatrix}1\\2\\3\end{bmatrix}$ $(3,1)$	均为 2 / 3；不是 shape=$(3,)$
matrix 矩阵	两行三列；$(2,3)$	2 / 6
三阶 Tensor	两份两行三列矩阵；$(2,2,3)$	3 / 12
一般 Tensor	$$\operatorname{shape}=(d_0,\ldots,d_{k-1})$$	k 个轴；numel 为各轴长度乘积（标量为 1）

axis/dim 是轴的索引；shape 是长度的元组；ndim 是轴数；numel 是元素数。matrix rank 是矩阵秩，distributed rank 是进程编号；二者都不应简写成张量“几维”。shape 本身不带语义，必须标明各轴代表样本、位置、特征或专家等。

A.11.1 一维向量、行矩阵、列矩阵：相同分量，不同 shape

PyTorch 的一维向量没有行、列方向。“行”与“列”需要两个轴才能表达。因此，含有相同三个数的一维向量、行矩阵、列矩阵，元素数相同，但不是同一种 shape。行矩阵也常叫行向量，列矩阵也常叫列向量；谈 PyTorch 张量时，应同时说清它们是二维表示。

表示	PyTorch 创建方式与 shape	轴的含义
一维向量：分量为 1、2、3	torch.tensor([1., 2., 3.])；$(3,)$	ndim=1，只有一个长度为 3 的轴，没有行轴或列轴。
行矩阵 $\begin{bmatrix}1&2&3\end{bmatrix}$	torch.tensor([[1., 2., 3.]])；$(1,3)$	ndim=2，第 0 轴为 1 行，第 1 轴为 3 列。
列矩阵 $\begin{bmatrix}1\\2\\3\end{bmatrix}$	torch.tensor([[1.], [2.], [3.]])；$(3,1)$	ndim=2，第 0 轴为 3 行，第 1 轴为 1 列。

三者的 numel() 都为 3。线性代数说“3 维向量”，通常指有 3 个分量；PyTorch 说“一维 Tensor”，指只有 1 个轴。这两种“维度”不是同一口径。shape 中 $(3,)$ 的逗号表示只有一个元素的长度元组，不代表还有一个省略的行轴或列轴。

import torch

v = torch.tensor([1., 2., 3.])  # shape (3,)
row = v.unsqueeze(0)           # shape (1, 3)：在前面加一个长度为 1 的轴
col = v.unsqueeze(1)           # shape (3, 1)：在后面加一个长度为 1 的轴

print(tuple(v.shape), v.ndim, v.numel())      # (3,)   1 3
print(tuple(row.shape), row.ndim, row.numel())  # (1, 3) 2 3
print(tuple(col.shape), col.ndim, col.numel())  # (3, 1) 2 3

print(tuple(row.T.shape))       # (3, 1)
print(tuple(col.T.shape))       # (1, 3)
print(tuple(row.squeeze(0).shape))  # (3,)
print(tuple(col.squeeze(1).shape))  # (3,)

同一组分量的三种表示

unsqueeze 添加长度为 1 的轴，squeeze(dim) 删除指定的长度为 1 的轴；这里得到的是共享原数据的视图，元素数没有变化。二维矩阵的 .T 交换行、列轴；v.T 不能把一维向量变成列矩阵，因为它没有第二个轴可交换。PyTorch 2.9.1 中 v.T 的 shape 仍为 $(3,)$，并会提示非二维张量使用 .T 已弃用；需要行、列表示时应显式添加轴。

A.11.2 矩阵乘法：数值可能相同，输出形状却不同

@ 对应 torch.matmul。对两个二维矩阵，它遵守 $(m,n)@(n,k)\longrightarrow(m,k)$：中间两个长度必须相等。对一维输入，PyTorch 有专门规则：左侧一维向量在计算时视作一行，右侧一维向量在计算时视作一列，计算结束后再去掉这些临时加入的轴。这个规则不会永久改变原向量的 shape。

继续使用上面的 v、row、col。下面所有结果均来自同一组分量；$14=1^2+2^2+3^2$。

表达式	结果	输出 shape 与原因
v @ v	$$14$$	$()$，两个一维向量做内积，得到 scalar。
row @ col	$$\begin{bmatrix}14\end{bmatrix}$$	$(1,1)$；$(1,3)@(3,1)$，输入本来就是二维，不删除任何轴。
col @ row	$$\begin{bmatrix}1&2&3\\2&4&6\\3&6&9\end{bmatrix}$$	$(3,3)$；$(3,1)@(1,3)$，得到外积矩阵。
row @ v	含一个元素 14 的一维向量	$(1,)$，去掉右侧一维输入临时补出的列轴。
v @ col	含一个元素 14 的一维向量	$(1,)$，去掉左侧一维输入临时补出的行轴。
row @ row	维度不匹配，报错	$(1,3)@(1,3)$ 的中间长度 3 与 1 不相等。

因此，$()$、$(1,)$、$(1,1)$ 虽然都可以只存一个数，却分别是标量、一维向量、二维矩阵。后续计算会按不同的 shape 处理它们，不能只看打印出来都包含 14 就认为三者等价。

for name, result in [
    ("v @ v", v @ v),
    ("row @ col", row @ col),
    ("col @ row", col @ row),
    ("row @ v", row @ v),
    ("v @ col", v @ col),
]:
    print(name, "shape =", tuple(result.shape), "value =", result.tolist())

# row @ row 会抛出 RuntimeError；中间维度 3 和 1 不匹配。

打印矩阵乘法的实际数值和形状

A.11.3 逐元素计算与广播：列矩阵加向量为什么会变成方阵

* 是逐元素乘法，不是矩阵乘法。v * v 得到分量 1、4、9，shape 为 $(3,)$；v @ v 则进一步按内积规则相乘并求和，得到 scalar 14。加减乘等逐元素运算会先应用广播规则，而不是应用矩阵乘法的中间维度规则。

广播从最后一个轴向前对齐。缺少的前导轴按长度 1 处理；对应轴长度相同，或者其中一个为 1，才可以广播。因此 row + v 按 $(1,3)+(1,3)$ 计算，结果还是 $(1,3)$；col + v 则按 $(3,1)+(1,3)$ 计算，两条轴分别扩展为 3，结果为 $(3,3)$。这不是 PyTorch 自动认定 v 是列向量。

print((v * v).tolist())         # [1.0, 4.0, 9.0]
print((row + v).tolist())       # [[2.0, 4.0, 6.0]]
print(tuple((row + v).shape))   # (1, 3)

print((col + v).tolist())
# [[2.0, 3.0, 4.0],
#  [3.0, 4.0, 5.0],
#  [4.0, 5.0, 6.0]]
print(tuple((col + v).shape))   # (3, 3)

仍使用同一组 1、2、3，观察广播

以元素下标表示，col + v 的输出为 $Z_{ij}=v_i+v_j$：每个列元素与每个向量元素各加一次。它不是只把对应位置相加。row + col 也会得到同一个 $(3,3)$ 矩阵，因为使用的是同样的广播形状。

训练中要先对齐预测和标签的 shape。设 $B$ 为样本数，预测 pred 为 $(B,1)$、标签 target 为 $(B,)$，直接相减会得到 $(B,B)$，相当于把每个预测与每个标签都比较。若任务要求逐样本比较，应把两者统一为 $(B,)$ 或 $(B,1)$，再计算损失。

pred = col  # (3, 1)
target = v  # (3,)

wrong = pred - target
print(wrong.tolist())
# [[0.0, -1.0, -2.0],
#  [1.0,  0.0, -1.0],
#  [2.0,  1.0,  0.0]]
print(tuple(wrong.shape))       # (3, 3)：逐对比较，不能当作逐样本误差

error_1d = pred.squeeze(-1) - target  # (3,)，全部为 0
error_2d = pred - target.unsqueeze(-1)  # (3, 1)，全部为 0

用相同数据说明一个常见的损失计算错误

这里使用 squeeze(-1) 只删除末尾的单元素轴。若直接写不带参数的 squeeze()，当 $B=1$ 时，样本轴也可能被删除，结果变成 scalar。判断一个表达式是否符合训练含义，先写清输入与输出 shape，再判断数值。

本节例子已在 PyTorch 2.9.1 的 CPU 环境核对数值、shape、广播结果及错误分支。这里只引入一组分量来展示三种表示的区别；矩阵计算如何接入模型，继续看 1.4 与 A.14。

A.12 梯度：从一个 scalar loss 回到每个参数

梯度是标量损失对参数各元素的偏导数组成的张量。它描述：在当前参数附近，参数的微小变化会怎样影响损失。设损失为 $L$、参数为 $\theta$，则 $g=\nabla_{\theta}L$，其中 $g_i=\frac{\partial L}{\partial\theta_i}$。

梯度分量为正，表示只把对应参数略微增大时，损失在局部倾向于增大；为负则倾向于减小。基本梯度下降沿负梯度方向调整参数，学习率控制调整幅度。具体怎样使用当前梯度和历史状态，由优化器决定，见 A.3。

残差、损失与梯度是不同对象：残差是预测值与目标值的差；损失把预测质量汇总为需要优化的标量；梯度是这个损失对参数的导数。梯度不是预测误差本身，也不是更新后的参数。

对标量损失，参数梯度与参数具有相同形状：若 $\theta\in\mathbb{R}^{m\times n}$，则 $\nabla_{\theta}L\in\mathbb{R}^{m\times n}$；$m,n$ 分别是参数矩阵的行数和列数。每个梯度元素对应一个参数元素的偏导。

操作	改变什么
backward	沿计算依赖反向应用链式法则，计算梯度，并累积到参数的 .grad；它本身不更新参数。
optimizer.step	根据梯度和优化器规则更新参数；有状态的优化器还会更新历史状态。
zero_grad	清除上一轮累积的参数梯度，不清除参数，也不清除 Adam 的历史状态。

求导与梯度累积的机制见 1.5，矩阵乘法的梯度形状见 A.14。需要手算和运行核对时，见 C.9.1。

A.13 Batch：mini-batch、global batch 与更新范围

batch 是一批样本；batch size 是这批样本的条数。下面的批量大小都按样本条数计，序列长度和 token 数另行统计。

A.13.1 mini-batch 和 global batch 分别是什么

mini-batch（小批次）：在 mini-batch SGD 的算法描述中，指从训练集中选出、用于计算一次参数更新所需梯度的一批样本。“mini”相对于整个训练集而言，强调每次更新通常只使用其中一部分。在分布式实现或数据加载器文档中，这个词也可能指本地批次，因此需要结合上下文确认范围。

global batch（全局批次）：整个训练系统共同用于一次参数更新的全部样本，包含所有数据并行副本的贡献，也包含更新前通过梯度累积处理的所有批次。global batch size 是这些样本的总条数。这里的“全局”以一次参数更新为范围。

local mini-batch（每个副本一次更新的本地批次）：本节用这个完整名称，表示一个数据并行副本在一次参数更新中累计处理的全部样本。它是 global batch 中由该副本负责的部分；一个副本可以由多张 GPU 协作实现。

mini-batch 与 global batch 的区别取决于前者的范围。若 mini-batch 指算法层面一次更新的完整批次，它与 global batch 相同；若指上述 local mini-batch，两者的区别就是一个副本的本地部分与所有副本的总和。后一种大小差异来自数据并行，不能泛化为使用多张 GPU 就一定有差异。

A.13.2 micro-batch 与梯度累积：一次计算不等于一次更新

micro-batch（微批次）：一个数据并行副本一次实际前向、反向计算所处理的样本。local mini-batch 统计这个副本一次更新累计处理的数据，micro-batch 统计一次实际计算处理的数据。

多个 micro-batch 并不要求多卡并行。单卡为了控制显存，也可以把一次更新的 16 条样本拆成四个 micro-batch，每个 4 条：执行四次前向和反向，期间参数不变、梯度累积并按目标正确归一化，最后更新一次参数。此时 micro-batch size 为 4，local mini-batch size 和 global batch size 都为 16。这就是梯度累积；它决定算几批才更新，数据并行则决定哪些副本共同参与更新，两者可以独立使用或组合。

A.13.3 加入梯度累积和数据并行后，怎样计数

计数时先合计一个副本在一次更新中处理的微批，再合计所有数据并行副本。下面的批量大小均以样本条数为单位：

符号	含义
$$b$$	micro-batch size：每个数据副本一次前向、反向处理的样本条数
$$A$$	每个副本在一次参数更新前累积的 micro-batch 数
$$D$$	独立处理不同数据、共同形成这次更新的数据并行副本数
$$B_{\mathrm{local}}$$	local mini-batch size：一个副本一次更新累计处理的样本条数
$$B_{\mathrm{global}}$$	global batch size：所有副本一次更新共同使用的样本总条数

当各副本每次都处理 $b$ 条样本，并且都累积 $A$ 个微批后更新时：$B_{\mathrm{local}}=A\times b$；$B_{\mathrm{global}}=D\times B_{\mathrm{local}}=D\times A\times b$。local mini-batch size 已包含本地梯度累积的样本，因此从它换算 global batch size 时只乘 $D$，不能再乘一次 $A$。单卡时 $D=1$；每算一个微批就更新时 $A=1$。

$D$ 统计数据并行副本数。TP、PP、CP 可以让多张 GPU 协作处理同一批样本，EP 会把 token 路由给不同专家；这些计算上的分工不会自动增加样本条数，所以不能直接把总 GPU 数代入 $D$。PP 也常用多个 micro-batch 让不同阶段交错执行，但微批的存在本身不意味着进行了数据并行。

A.13.4 不同资料中的 mini-batch 怎样对应

资料中 mini-batch 的范围	对应的批量大小	global batch 的换算
算法层面一次更新使用的完整批次	$$B_{\mathrm{global}}$$	二者大小相同
每个数据并行副本一次更新累计的本地批次	$$B_{\mathrm{local}}=A\times b$$	$$B_{\mathrm{global}}=D\times B_{\mathrm{local}}$$
每个副本一次实际前向、反向处理的批次；此时与 micro-batch 同义	$$b$$	$$B_{\mathrm{global}}=D\times A\times b$$

本文做执行计数时明确使用 global batch、local mini-batch（每副本、每次更新）和 micro-batch。其他资料若只写 mini-batch、local batch 或 batch_size，应确认它统计一次计算还是一次更新、一个副本还是全部副本，以及是否已经包含梯度累积。具体代码字段保留项目原名，再按实际取批与更新循环确定含义。

effective batch size（有效批量）在本文与 global batch size 同口径；step 指一次参数更新；epoch 指训练数据被遍历一轮。这些名称分别描述批量和训练进度。

执行过程图、完整训练循环、框架字段含义与 CPU 验证见 C.9.3；本节用于查询概念和计数关系。

A.14 矩阵计算与模型架构：沿形状读前向和反向

线性代数小例子见 1.4：$\underbrace{X}_{(B,C_{\mathrm{in}})}\underbrace{W}_{(C_{\mathrm{in}},C_{\mathrm{out}})}=\underbrace{Z}_{(B,C_{\mathrm{out}})}$，再加 b$(C_{\mathrm{out}},)$；nn.Linear 的实际 weight 是 $W^{\mathsf T}$，shape=$(C_{\mathrm{out}},C_{\mathrm{in}})$。若收到上游梯度 dZ$(B,C_{\mathrm{out}})$，则 $dW=X^{\mathsf T}dZ$ 为 $(C_{\mathrm{in}},C_{\mathrm{out}})$，$dX=dZW^{\mathsf T}$ 为 $(B,C_{\mathrm{in}})$，$db=\sum_{i=1}^{B}dZ_{i,:}$ 为 $(C_{\mathrm{out}},)$；对实际 weight 求导需采用转置后的布局。

下面使用教学 GPT 的一组合法小形状：$B=2$ 条序列、$T=3$ 个位置、$C=8$ 个隐藏特征、$H=2$ 个头、$D=C/H=4$ 个每头特征、$V=4$ 个词表项。这里只说明 shape，不替换前面实际 nanoGPT 的训练配置；V 表示词表大小时不要与 Attention 的 Value 张量 V 混淆。

对象 / 运算	符号 shape	本例 shape 与轴含义
token IDs / targets	$$(B,T)$$	$(2,3)$，整数；targets 向后错一位
Embedding 参数表 / 输出	$(V,C)$ / $(B,T,C)$	$(4,8)$ / $(2,3,8)$
Q/K/Value 三个投影权重（数学写法）	各 $(C,C)$	各 $(8,8)$；合并 qkv 时 API weight 可为 $(3C,C)$
拆头后 Q/K/Value	$$(B,H,T,D)$$	$$(2,2,3,4)$$
$QK^{\mathsf T}$，转置最后两轴	$$(B,H,T,D)\mathbin{@}(B,H,D,T)$$	$$(2,2,3,4)\mathbin{@}(2,2,4,3)\longrightarrow(2,2,3,3)$$
$$\text{Attention weights}\cdot\text{Value}$$	$$(B,H,T,T)\mathbin{@}(B,H,T,D)$$	→$(2,2,3,4)$，合头回 $(2,3,8)$
输出投影和残差	$$(B,T,C)\longrightarrow(B,T,C)$$	$(2,3,8)$；相加两边 shape 相同
MLP，采用中间宽度 $4C$	$$(B,T,C)\longrightarrow(B,T,4C)\longrightarrow(B,T,C)$$	$$(2,3,8)\longrightarrow(2,3,32)\longrightarrow(2,3,8)$$
最终词表投影 weight / logits	$(V,C)$ / $(B,T,V)$	$(4,8)$ / $(2,3,4)$
交叉熵输入 / 目标 / mean loss	$(B\times T,V)$ / $(B\times T,)$$\quad/\quad()$	$(6,4)$ / $(6,)$ / scalar
生成取最后位置	logits[:,-1,:]	$(B,V)=(2,4)$，每条序列选下一个 token

架构图读法：沿箭头先写输入输出 shape，再标出参数属于哪个模块、哪些激活为反向保留。TP 改变单层特征切分和局部结果汇总，PP 在层边界传 $(B,T,C)$ 激活及同形状上游梯度，CP 分担位置轴但必须保持完整上下文依赖。EP 将 token 激活按路由分组给专家后还原；这些是后续学习的切入点，不表示本章小 GPT 已实现全部并行。

本节 shape 表描述逻辑计算；高效 SDPA 不一定显式物化完整 $(B,H,T,T)$ 分数矩阵。切片/transpose 后的 stride、是否需要 contiguous、局部/全局形状也必须分别核对，详见 B.1 和第 7 章。

A.15 并行训练：DP、TP、PP、CP 与 SP

并行训练是让多个进程协作完成同一个训练任务。rank 是进程在某个通信范围内的编号；进程组 是参与一组通信的进程集合；并行度 是某种分工所用的组规模。本文 GPU 实验通常一个进程对应一张 GPU，但 rank 本身是进程编号。同一进程可以同时属于不同的 TP、DP 等组，global rank 与组内 rank 也不一定相同。

下面先单独解释每一种分工，不叠加其他并行方式。涉及隐藏状态时，使用 $X\in\mathbb{R}^{B\times T\times C}$：$B$ 是一次 micro-batch 的序列条数，$T$ 是完整序列长度，$C$ 是每个 token 的隐藏特征数。符号 $p$ 在各小节中表示该种并行的度数；写成均匀分片时，默认被切维度能被 $p$ 整除。

A.15.1 DP：不同数据分片，执行同一个逻辑模型

DP（Data Parallelism，数据并行）把不同样本交给不同的数据并行副本。各副本执行相同的逻辑模型，得到各自数据贡献的梯度，再汇总这些贡献，使更新对应于所有副本共同处理的数据。它主要用于分担数据计算、提高训练吞吐；样本、损失和梯度的汇总口径必须一致。等量样本且各副本取平均损失时，可对梯度取平均；有效样本或 token 数不等时需要相应加权。

DP 是策略名称，不等同于某一个 PyTorch 类。torch.nn.DataParallel 是一个具体 API；DDP（DistributedDataParallel）是另一种实现，本文采用多进程方式，每个数据副本通常保存完整模型，并通过梯度同步保持更新一致。普通 DDP 不会把每卡的参数和优化器状态自动缩小为原来的几分之一。

FSDP（Fully Sharded Data Parallel，全分片数据并行）仍属于数据并行，但会分片存储参数、梯度和优化器状态，在计算相应模块时收集所需参数。它主要解决训练状态的显存压力。因此不能把 DP 一概定义为“每张卡始终保存完整模型”。FSDP 与 TP 的区别是：前者分片训练状态、按数据分工；后者把同一次层内计算也拆给多个 rank。

global batch 统计的是参与一次更新的样本总条数，按样本出现次数计，不做去重，也不是简单数 GPU。mini-batch、micro-batch 与累积次数的完整口径见 A.13；DDP 与 FSDP 的机制、实验分别见第 4、5 章。

A.15.2 TP：多个 rank 共同完成一层运算

TP（Tensor Parallelism，张量并行）把一层中的权重张量及其计算分给多个 rank。它们共同处理同一批逻辑输入，各自计算部分结果，再根据下一层需要的布局进行求和或拼接。它主要分担单层参数和计算压力；层内的通信往往处于计算依赖路径上。

暂不考虑偏置，线性层的 PyTorch 权重 $W\in\mathbb{R}^{F\times C}$，其中 $C$ 是输入宽度、$F$ 是输出宽度。按输出特征分片时，每 rank 权重为 $(F/p,C)$，局部输出为 $(B,T,F/p)$，各卡负责不同输出特征。按输入特征分片时，每 rank 权重为 $(F,C/p)$，接收对应的 $(B,T,C/p)$ 输入，产生 $(B,T,F)$ 的部分和；这些部分和要相加才是完整输出。局部张量即使形状相同，也未必代表相同或完整的数据。

TP 与 DP 的本质区别是：TP 合作完成同一次运算，DP 汇总不同样本对同一个模型的训练贡献。不能把 TP 各卡处理的逻辑同一批样本重复计入 global batch。前向与反向的具体通信见 7.4；nanoGPT 改造边界见 B.8。

A.15.3 PP：把模型的不同层放到不同阶段

PP（Pipeline Parallelism，流水线并行）按层或连续模块划分模型，每个阶段（stage）只执行自己负责的部分。前向将边界激活传给下一个阶段，反向将相应的上游梯度传回前一个阶段。以隐藏宽度不变的 Transformer 为例，未叠加其他并行时，边界激活和它的梯度通常同为 $(B,T,C)$。

把一次更新的数据分成多个 micro-batch，可以让不同阶段同时处理不同 micro-batch，形成流水线；所有相关反向完成后才进行这轮参数更新。某阶段等待输入、梯度或其他阶段时出现的空闲称为流水线气泡。PP 主要分担模型各层的存储与计算，效果受阶段负载、通信和调度影响；它与 TP 的区别是“不同层分工”与“同一层内分工”。调度与实验见 7.5，nanoGPT 的共享权重问题见 B.9。

A.15.4 CP：共同计算同一条长序列

CP（Context Parallelism，上下文并行）把同一条序列的不同位置分给不同 rank，以分担长上下文的激活和 Attention 计算压力。均匀切分时，每 rank 的隐藏状态为 $(B,T/p,C)$；这里减少的是本地持有的位置数，完整序列仍有 $T$ 个位置，样本条数仍是 $B$。

不同位置的 Query 可能需要访问其他 rank 持有的 Key/Value，因此各片段不能各自当作独立序列做普通 Attention。实现必须通过通信保留完整上下文依赖，同时保持全局位置与因果 mask 正确。CP 组中复制的模型参数也需要汇总各位置贡献的梯度。CP 切的是同一序列的位置，DP 分的是不同样本；不能因 CP 度数增加就把 global batch 再乘一次。环式计算和验证见 7.6。

A.15.5 SP：在 TP 配套区域分片序列激活

SP（Sequence Parallelism，序列并行）在本文特指 Megatron 风格、配合 TP 使用的激活布局分片。它将 Norm、Dropout 等区域的激活沿序列维分给不同 rank，每 rank 保留 $(B,T/p,C)$，仍为每个 token 保留完整隐藏维；进入相应的 TP 计算区域时再收集所需输入，离开时通过归约并分片恢复该布局。

SP 的主要目的，是减少这些区域重复保存的激活和重复计算；CP 的主要目的，是让不同 rank 负责同一长序列的不同 Query 位置，并在 Attention 中交换所需上下文信息。二者都可能出现沿序列维的分片，但作用范围和通信过程不同。有些资料也把 CP 类方法称为 sequence parallelism，阅读时要先确认作者的定义。形状转换见 7.4.1。

组合时分别判断三件事：哪些 rank 处理不同数据，哪些 rank 共同完成同一次运算，哪些 rank 保存同一参数的副本并需要同步。TP、PP、CP、EP 的度数都不能自动当作额外的 batch 倍数；进程总数怎样计算，也取决于具体框架的组组织方式。NCCL 是完成这些通信的软件库，RDMA 是网络数据传输机制，它们不是新的模型分工方式；层次与通信操作见 6.2。

A.16 MoE 与 EP：模型结构和专家分工

A.16.1 MoE：每个 token 选择少数专家计算

MoE（Mixture of Experts，混合专家）是一类模型结构。本文讨论语言模型中常见的稀疏 top-k MoE：将某些普通 MLP 替换为多个专家 MLP，再由一个可学习的router（路由器）决定每个 token 使用哪些专家，以及如何组合它们的输出。专家是独立的可学习模块，不要求预先指定“数学专家”“写作专家”等人工领域。

普通 dense MLP 让所有 token 使用同一组 MLP 参数；稀疏 MoE 让每个 token 只调用所选专家，可以增加模型总参数量，而不让每个 token 都计算所有专家。总参数量与每 token 参与计算的参数量不同。“稀疏”指每个 token 的选择；一个 batch 中的不同 token 合起来，仍可能使用多数甚至全部专家。MoE 也可以运行在单张 GPU 上，不天然要求分布式训练。

用 $N$ 表示本次输入的 token 数、$C$ 表示隐藏宽度、$E$ 表示专家总数、$K$ 表示每个 token 选择的专家数。输入 $X$ 的 shape 为 $(N,C)$；router 对专家打分，分数 shape 为 $(N,E)$；top-k 专家编号及其组合权重均为 $(N,K)$。每个专家保持 token 的隐藏宽度，合并后的输出仍为 $(N,C)$，不会因为选了多个专家而产生更多训练样本。

常见的加权合并写作 $y_i=\sum_{e\in\mathcal{K}_i}\alpha_{i,e}f_e(x_i)$：$x_i,y_i\in\mathbb{R}^{C}$ 是 token $i$ 的输入与输出，$\mathcal{K}_i$ 是它选中的专家集合，$f_e$ 是专家函数，$\alpha_{i,e}$ 是路由组合权重。如何从路由分数得到权重、是否有共享专家等，属于具体 MoE 的模型定义；7.9 使用对所选分数做 softmax 的教学约定。

A.16.2 EP：把专家放到不同 rank，并分发 token

EP（Expert Parallelism，专家并行）是 MoE 的一种并行执行方式：不同专家由不同 rank 保存和计算。源 rank 先根据路由结果，把 token 的隐藏向量发送给目标专家所在的 rank，这一步称为 dispatch（分发）；专家计算完成后，把结果送回来源，恢复原 token 顺序并按路由权重合并，称为 combine（合并）。典型实现用 AllToAll 完成按目的地的重分布。

设专家 $e$ 本轮收到 $n_e$ 项 token 任务，它的输入和输出 shape 都为 $(n_e,C)$。$n_e$ 随路由结果变化，不同专家可能收到不同数量的任务；最忙的专家会影响整轮耗时，这就是专家负载不均。专家数均分给 GPU，并不保证计算量也均分。分发缓冲、往返通信与梯度传递同样是 EP 必须处理的开销。

MoE 回答模型如何计算，EP 回答专家在哪里计算。MoE 不等于 EP，EP 也不等于 TP：EP 把不同专家分给不同 rank；TP 则可继续拆分一个专家内部的矩阵运算。EP 的结果送回原 token 后继续原模型的计算，所以不能把专家任务数计成额外的 global batch。专家参数只在保存同名专家副本的组内同步；它与非专家参数的复制组可能不同。

本文固定版本的 nanoGPT、Picotron 使用 dense MLP，没有现成的 EP 开关。完整路由、通信、反向与验证边界见 7.9；集合通信练习见 C.10。

附录 B：底层原理与工程选读

本附录在主线跑通后按需阅读。它解释内存、GPU、数值和求导实现，不作为第一次训练的门槛。

B.1 张量布局：storage、stride 与视图

对常见的稠密 strided 张量，可以把 storage 看成底层存储，把 shape、stride 与 storage_offset 看成寻址信息。张量的逻辑元素不一定连续，不同视图可以共享同一 storage。二维情况下，元素位置为 $\text{storage\_offset}+i\times\text{stride}[0]+j\times\text{stride}[1]$；稀疏等布局需要另行理解。

共享存储示意图：x 与转置 y 使用同一条 storage，shape 和 stride 改变，数据本身没有复制。

打开 SVG 原图

理解了 stride，很多「玄学」会清楚很多：transpose、permute 和基础切片通常只修改 shape/stride，返回共享 storage 的视图；view 只有在当前布局兼容时才能零复制完成。reshape 可能返回视图，也可能悄悄复制，而高级索引通常会产生新张量。某些 kernel 要求连续布局时，需要调用 .contiguous()，若当前张量在目标内存格式下不连续，会分配并复制；已经连续时通常返回原张量。

import torch
x = torch.arange(6).reshape(2, 3)      # shape (2,3) stride (3,1)
y = x.T                                # 转置：shape (3,2) stride (1,3)，没有复制数据
print(x.data_ptr() == y.data_ptr())    # True，指向同一块内存
print(y.is_contiguous())               # False，stride 非标准顺序

转置不复制数据，只改 stride

📌

第一性原理结论：许多形状变换只是「换个视角看同一块 storage」，成本很低；但 reshape、高级索引、contiguous 等操作可能复制数据。判断是否共享存储，才能知道显存和时间真正花在哪里。

contiguous() 在当前布局已经满足要求时可能返回自身；需要转换布局时才复制。data_ptr() 相等可验证本例转置共享起始地址，但不同 offset 的共享视图不一定有相同 data_ptr()。reshape 可能返回视图也可能复制，高级索引通常产生新存储。

B.2 GPU：设备搬运、异步执行与计时

GPU 架构、设备查询与 CUDA Event 计时实验已放在主线 1.7。这里保留 CPU/GPU 数据搬运的示意：CPU 与 GPU Tensor 的逻辑接口相近，实际存储和执行设备不同；足够大的并行工作可能获益，小任务、频繁搬运和同步则可能抵消收益。

GPU 数据搬运图：尽量复用设备上的数据，减少不必要的 CPU/GPU 往返。

打开 SVG 原图

import time, torch
if torch.cuda.is_available():
    x = torch.randn(2048, 2048, device="cuda")
    for _ in range(10):
        y = x @ x
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter()
    for _ in range(50):
        y = x @ x
    torch.cuda.synchronize()
    print("ms/次:", (time.perf_counter()-start)*1000/50)
    print("峰值已分配 MiB:", torch.cuda.max_memory_allocated()/2**20)

先预热，再同步计时；本例不声称固定加速倍数

CUDA 调用通常异步提交，Python 返回不代表 GPU 已完成；计时边界 synchronize 用来等待已提交工作。.item() 或把 GPU 结果搬回 CPU 往往需要等待结果。pin_memory 与 non_blocking 为异步搬运提供条件，是否与计算重叠还取决于 stream 和调度。

B.3 数值精度：fp32、fp16 与 bf16

浮点数的指数位主要决定动态范围，尾数位主要决定精度。fp32 每元素 4 字节，fp16/bf16 为 2 字节。bf16 的指数位数与 fp32 相同、尾数更少，动态范围接近但精度较低。

浮点格式示意图：比较指数范围与尾数精度，区分单个张量大小和整个训练状态大小。

打开 SVG 原图

import torch

# 1) 占用字节：从 fp32 转为 16 位后，每元素存储字节减半；整个训练的总显存不一定减半
for dt in [torch.float32, torch.float16, torch.bfloat16]:
    print(dt, "每个元素字节数:", torch.empty(0, dtype=dt).element_size())
# float32 -> 4, float16 -> 2, bfloat16 -> 2

# 2) 动态范围：fp16 最大约 6.5万，大数直接溢出成 inf；bf16 与 fp32 的指数位数相同，动态范围接近（最大有限值不完全相同）
big = torch.tensor(70000.0)
print("fp16 存 70000:", big.half())              # inf —— 溢出了！
print("bf16 存 70000:", big.bfloat16())          # 70144，能表示（只是精度略降）

# 3) 精度：bf16 尾数少，小数会被明显舍入
x = torch.tensor(1.0 / 3.0)
print("fp32:", x.item())                          # 0.3333333...
print("bf16:", x.bfloat16().float().item())       # 0.333984375，精度更粗

# 4) 常见搭配：部分矩阵运算用 bf16，数值敏感的计算或累加保留 fp32
#    autocast 按算子策略选择 dtype；FSDP2 MixedPrecisionPolicy 控制参数计算/梯度通信等 dtype。
#    两者作用不同；仍需核对具体算子、版本与设备行为。

亲手看到三种精度的范围与精度差异

fp16 的前向运算可能溢出，反向小梯度也可能下溢。GradScaler 主要用于缓解梯度下溢，并在检测到非有限梯度时跳过更新、调整缩放；它不能保证消除前向溢出。bf16 通常不需要缩放，但仍要验证稳定性。autocast 按算子选择计算精度，不能推导出模型、优化器状态和总显存都减半。

B.4 选做：手写标量 autograd

先完成 1.5 的 PyTorch 数值例子，再研究这个微型引擎。Value 保存数值、梯度、父节点和局部反向函数；拓扑逆序保证在处理一个节点时，它从后续节点收到的贡献已经累积。它只是标量加法和乘法的教学模型，不涵盖 PyTorch 自动求导的全部机制。

class Value:
    """一个标量 + 它的梯度 + 怎么反向传播。"""
    def __init__(self, data, _children=()):
        self.data = data
        self.grad = 0.0                 # 该值对最终 loss 的偏导
        self._backward = lambda: None   # 局部反向：把梯度传给上游
        self._prev = set(_children)     # 记住是谁生成了我 → 这就是"图"

    def __add__(self, other):
        out = Value(self.data + other.data, (self, other))
        def _backward():                # 加法：梯度原样传给两个输入
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward
        return out

    def __mul__(self, other):
        out = Value(self.data * other.data, (self, other))
        def _backward():                # 乘法：∂(a*b)/∂a = b（链式法则）
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward
        return out

    def backward(self):
        topo, visited = [], set()       # 拓扑排序，保证反向时上游 grad 已就绪
        def build(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build(child)
                topo.append(v)
        build(self)
        self.grad = 1.0                 # d(loss)/d(loss)=1，反向起点
        for v in reversed(topo):
            v._backward()

a = Value(2.0); b = Value(-3.0)
f = a * b + a
f.backward()
print(a.grad, b.grad)   # a.grad = b+1 = -2；b.grad = a = 2

手写标量 autograd 引擎

本例 $a=2$、$b=-3$、$f=a\times b+a$，预期 $\frac{\partial f}{\partial a}=-2$、$\frac{\partial f}{\partial b}=2$。沿多条路径的贡献必须相加。这个简化引擎不提供完整的重复 backward、图释放、张量广播求导和原地操作检查；扩展时需要明确这些语义。

B.5 API 选择与旧代码迁移速查

在遇到相应场景时查询下表即可。框架 API 的适用范围取决于版本和任务，教学实现与生产优化各有用途。

旧写法或非主线方案	本教程推荐	适用边界
Variable(tensor)	直接使用 Tensor 与 requires_grad=True	Variable 已并入 Tensor
.data 绕过求导	.detach() 或 torch.no_grad()	避免破坏 autograd 的版本与梯度跟踪
nn.DataParallel	DistributedDataParallel（DDP）	新建多卡训练项目优先 DDP
torch.distributed.launch	torchrun	launch 已废弃
NVIDIA apex.amp	原生 torch.amp	除非维护依赖 Apex 的旧工程
手写生产级 attention	F.scaled_dot_product_attention	学习原理时仍应手写一次；生产代码让 SDPA 按条件选择后端
FSDP1（类式 API）	新项目优先 FSDP2 fully_shard	维护既有工程时 FSDP1 仍可能合理
nn.Transformer、RNN、LSTM	学习现代 LLM 时以 decoder-only Transformer 为主线	这些模块并未失效，在教学及特定序列任务中仍有价值
torch.save(model) 保存整个对象	单卡保存 state_dict；分布式使用 DCP	避免 checkpoint 绑定 Python 类定义
默认使用 fp16 与 loss scaling	H800/H20 训练优先评估 bf16	具体精度仍取决于硬件、算子、推理引擎和数值稳定性

B.6 组合优化：在单项验证之后整合

先完成 2.4 的单项实验，再组合 bf16、梯度累加、裁剪与学习率调度。每组累加只更新一次参数和 scheduler；checkpoint 需包含优化器、scheduler、缩放器（若使用）、训练进度以及模型配置。

组合训练流程图：多个 micro-batch 累积后，统一裁剪、更新参数、推进调度并清梯度。

打开 SVG 原图

B.7 选读：DeviceMesh 表达逻辑通信拓扑

只有组合多种并行或研究混合分片时，才需要多维网格。DeviceMesh 描述 rank 的逻辑组织，并不自动发现最优物理互联。下面是八进程形状示意，不是两卡入门的运行要求。

当要组合多种并行（如数据并行 + 张量并行），用 DeviceMesh 描述卡的拓扑网格，FSDP2 / TP 都在 mesh 上工作：

# 前置：在 torchrun 启动的进程内运行（已 init_process_group("nccl")），共 8 卡
from torch.distributed.device_mesh import init_device_mesh
# 8 卡组织成 2(数据并行) x 4(张量并行) 的网格
mesh = init_device_mesh("cuda", (2, 4), mesh_dim_names=("dp", "tp"))
# fully_shard(model, mesh=mesh["dp"]) 等在对应子网格上工作

B.8 选读：nanoGPT 的 TP 改造边界

TP 与 PP 改变计算分工。先分别验证，再讨论组合：明确哪些张量完整、哪些是分片、在哪里交换信息。下图将改造定位到 nanoGPT 模块，DDP/FSDP2 的机制与 4.2、5.1—5.2、5.6 对照。

并行改造落点图：模型副本、Block 分片、MLP 张量切分与跨层流水线。

打开 SVG 原图

先切 MLP。 nanoGPT 的 MLP 是 $C\longrightarrow4C\longrightarrow C$：c_fc 后接 GELU，再接 c_proj。两卡可以把 c_fc 的输出特征各取一半，各自计算 GELU，再用 c_proj 对应的输入列得到一份部分输出，最后求和。PyTorch Linear.weight 的布局为 $(\text{out\_features},\text{in\_features})$，所以 c_fc.weight 沿第 0 维切，c_proj.weight 沿第 1 维切。若有 c_proj.bias，只能在合并后加一次。

from torch.distributed.tensor.parallel import (
    ColwiseParallel, RowwiseParallel, parallelize_module,
)

# tp_mesh 是一维 TP 子网格；各 rank 持有一致的模型和输入。
# 先用 dropout=0 验证下列 MLP 模式；这不是完整训练入口。
for block in model.transformer.h:
    parallelize_module(block.mlp, tp_mesh, {
        "c_fc": ColwiseParallel(),
        "c_proj": RowwiseParallel(),
    })
# 布局改造完成后再创建优化器，并验证输出与梯度。

TP 的第一步：MLP 列并行接行并行

注意原版 Attention 的 c_attn 一次输出 [Q 全部通道, K 全部通道, V 全部通道]。直接把这个 $3C$ 维连续对半切，可能得到不完整的 Q/K/V 配对；同时 forward 中 split(self.n_embd) 和 view(...,self.n_head,...) 仍假设全局维度。扩展 Attention TP 时，应显式重排或拆开 Q/K/V 投影，让每卡拿到对应局部头，调整 local n_head 和局部 reshape，再设计输出投影归约。必须让 n_head 能按 TP 度数合理分配。

TP 组内的各卡共同处理同一个 batch；从 DDP 路径迁移时，数据随机种子应按数据并行组区分，并在 TP 组内保持输入一致。LayerNorm 等算子所见的布局、dropout 随机性、共享 Embedding/lm_head 的分片和 checkpoint 布局也要一起考虑。官方接口参考：Tensor Parallelism。

B.9 选读：nanoGPT 的 PP 与共享权重边界

再按 Block 拆 PP。 第一阶段接收 token，执行 wte+wpe、dropout 和前半段 Blocks，输出 hidden states $(B,T,C)$；第二阶段执行后半段 Blocks、ln_f、lm_head，并用目标计算 loss。前向传 hidden，反向传对应梯度；最后阶段还需要拿到同一 micro-batch 的 Y。完整训练循环交给 PipelineStage 与 schedule 组织，不能仅把两段模块 .to 到不同设备就认为完成了流水线。

先在单进程中按边界顺序执行两段，验证与整模型输出及梯度一致，再换成两进程通信，并从简单的填充—排空调度进入 1F1B。明确 micro-batch 数、每段输入输出 shape、loss 的 sum/mean 与缩放规则，避免 schedule 与外层循环重复除以 micro-batch 数。负载也不能只按层数平均：首尾的 Embedding、词表投影和 loss 可能造成额外开销。

nanoGPT 的 wte 与 lm_head 若位于不同 stage，跨进程不会自动保持 Parameter 别名。保持绑权时，需要一致初始化，并把两处梯度贡献正确合并，协调更新，使两份权重继续一致；使用支持共享权重的框架可以减少手写同步。若为了学习先解绑，必须明确这是修改模型结构，再建立对应的 untied 基线。官方接口参考：Pipeline Parallelism。

附件 verify_model.py 在 CPU 上验证 MLP 两分片求和的输出/梯度，以及保留共享参数时按两个逻辑 stage 执行的输出/梯度。这能检验切分的数学关系；真实 TP/PP 的通信、调度、随机数和性能还要在多卡上单独验证。

附录 C：独立实验与完整代码

本附录保留完整可运行例子和局部机制实验。第 2 章的最小 GPT 与第 3～6 章的 nanoGPT 使用不同 model.py，分别放在独立目录；沿主线进入 nanoGPT 后，无需再切回最小 GPT。代码标注为片段时仍需嵌入相应上下文；未完成 GPU 验证的脚本不视为实测结论。

C.1 选做：用图像分类复习训练、验证与测试

FashionMNIST 每张图片为 $(1,28,28)$，标签是 0 到 9 的类别编号。Flatten 把每张图展开为 784 个特征；Linear 与 ReLU 组成多层感知机。没有非线性激活，多层线性变换仍等价于一个线性变换。最后输出 10 个 logits，即未归一化的类别分数。

MLP 数据流图：从图片特征到 10 个类别分数，每一层都保持 batch 轴。

打开 SVG 原图

CrossEntropyLoss 接收 logits 与标签，内部以稳定形式完成 log_softmax 与负对数似然计算，不要在模型末尾重复添加 softmax。验证集用于比较模型和选择 checkpoint；测试集留到选择结束后评估，避免把测试数据用于调参。

"""FashionMNIST：训练 / 验证 / 测试 / 保存 / 加载。需安装 torchvision。"""
import torch
from torch import nn
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms


def make_model():
    return nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(),
                         nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 10))


@torch.inference_mode()
def accuracy(model, loader, device):
    model.eval()
    correct, total = 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        correct += (model(x).argmax(dim=-1) == y).sum().item()
        total += len(y)
    return correct / total


def main():
    torch.manual_seed(42)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    full = datasets.FashionMNIST("data", train=True, download=True, transform=transforms.ToTensor())
    train, val = random_split(full, [55000, 5000], generator=torch.Generator().manual_seed(42))
    test = datasets.FashionMNIST("data", train=False, download=True, transform=transforms.ToTensor())
    # 起步用 num_workers=0，先排除多进程环境差异；之后再测数据加载吞吐。
    train_dl = DataLoader(train, batch_size=256, shuffle=True, num_workers=0)
    val_dl = DataLoader(val, batch_size=512, num_workers=0)
    test_dl = DataLoader(test, batch_size=512, num_workers=0)
    model = make_model().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()
    best = -1.0
    for epoch in range(5):
        model.train()
        for x, y in train_dl:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad(set_to_none=True)
            loss = loss_fn(model(x), y)
            loss.backward()
            optimizer.step()
        score = accuracy(model, val_dl, device)
        if score > best:
            best = score
            torch.save(model.state_dict(), "mlp_best.pt")
        print(f"epoch {epoch + 1}: val_accuracy={score:.4f}")
    restored = make_model().to(device)
    restored.load_state_dict(torch.load("mlp_best.pt", map_location="cpu", weights_only=True))
    print("test_accuracy:", accuracy(restored, test_dl, device))
    x, y = test[0]
    with torch.inference_mode():
        predicted = restored(x.unsqueeze(0).to(device)).argmax(dim=-1).item()
    print("预测类别编号 / 真实编号：", predicted, y)


if __name__ == "__main__":
    main()

保存为 classify.py；运行 python classify.py

model.eval() 改变 Dropout、BatchNorm 等层的行为，不会自动关闭求导；inference_mode() 关闭推理过程的梯度记录，二者用途不同。示例每轮观察验证准确率，保存最优模型，最后用重新创建并加载参数的模型评估测试集与单张图片。准确率会随版本、随机性和训练设置变化，不把某个固定小数作为唯一通关条件。

C.2 最小 GPT 完整代码：模型、训练与生成

对应第 2 章。按下列顺序保存 model.py、train_gpt.py、generate.py，与 UTF-8 文本 input.txt 放在同一目录，再回到 2.2 执行命令。模型定义导入时不启动训练；这三个文件构成同一个最小完整项目。

"""最小 GPT；导入本文件不会下载数据或启动训练。"""
from pathlib import Path
import torch
from torch import nn
from torch.nn import functional as F


class Block(nn.Module):
    def __init__(self, dim, n_heads):
        super().__init__()
        if dim % n_heads:
            raise ValueError("dim 必须能被 n_heads 整除")
        self.n_heads = n_heads
        self.ln1 = nn.LayerNorm(dim)
        self.qkv = nn.Linear(dim, 3 * dim)
        self.proj = nn.Linear(dim, dim)
        self.ln2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(C, dim=-1)
        q, k, v = (t.reshape(B, T, self.n_heads, C // self.n_heads).transpose(1, 2)
                   for t in (q, k, v))
        # 根据设备、dtype、shape 等选择后端，不保证一定使用 FlashAttention。
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)
        a = a.transpose(1, 2).reshape(B, T, C)
        x = x + self.proj(a)
        return x + self.mlp(self.ln2(x))


class GPT(nn.Module):
    def __init__(self, vocab, dim=128, n_heads=4, n_layers=2, block_size=64):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab, dim)
        self.pos_emb = nn.Embedding(block_size, dim)
        self.blocks = nn.ModuleList([Block(dim, n_heads) for _ in range(n_layers)])
        self.ln_f = nn.LayerNorm(dim)
        self.head = nn.Linear(dim, vocab, bias=False)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        if not 0 < T <= self.block_size:
            raise ValueError("序列长度必须在 1 到 block_size 之间")
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        for block in self.blocks:
            x = block(x)
        logits = self.head(self.ln_f(x))
        loss = None if targets is None else F.cross_entropy(
            logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
        return logits, loss


def load_corpus(path="input.txt", block_size=64):
    text = Path(path).read_text(encoding="utf-8")
    split = int(len(text) * 0.9)
    # 字符表描述数据格式；训练样本和验证样本按连续区间分开，不跨边界取窗口。
    chars = sorted(set(text))
    stoi = {c: i for i, c in enumerate(chars)}
    data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
    train, val = data[:split], data[split:]
    if min(len(train), len(val)) <= block_size:
        raise ValueError("文本太短：训练段、验证段都要至少有 block_size + 1 个字符")
    return train, val, chars


def get_batch(data, block_size, batch_size, device, generator=None):
    # 最后一个合法起点为 len(data)-block_size-1；randint 的上界不包含在内。
    ix = torch.randint(len(data) - block_size, (batch_size,), generator=generator)
    x = torch.stack([data[i:i + block_size] for i in ix])
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
    return x.to(device), y.to(device)


def build_dataset(path="input.txt", block_size=64):
    """兼容第 3 章：返回训练段、词表大小和窗口长度。各 rank 先共享准备好的文件。"""
    train, _, chars = load_corpus(path, block_size)
    return train, len(chars), block_size

model.py

"""python train_gpt.py --steps 500；python train_gpt.py --steps 1000 --resume last.pt"""
import argparse
import torch
from model import GPT, get_batch, load_corpus


@torch.inference_mode()
def estimate_loss(model, data, config, batch_size, device, batches=20):
    was_training = model.training
    model.eval()
    # 固定评估窗口，减少不同检查点之间的数据抽样波动。
    generator = torch.Generator().manual_seed(2026)
    losses = []
    for _ in range(batches):
        x, y = get_batch(data, config["block_size"], batch_size, device, generator)
        _, loss = model(x, y)
        losses.append(loss.item())
    model.train(was_training)
    return sum(losses) / len(losses)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--text", default="input.txt")
    parser.add_argument("--steps", type=int, default=500, help="累计更新步数，含已恢复的步数")
    parser.add_argument("--resume", default="")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--eval-every", type=int, default=100)
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()
    if min(args.steps, args.batch_size, args.eval_every) < 1:
        parser.error("steps、batch-size、eval-every 必须为正数")
    device = args.device
    torch.manual_seed(42)
    saved = torch.load(args.resume, map_location="cpu", weights_only=True) if args.resume else None
    config = saved["config"] if saved else dict(dim=128, n_heads=4, n_layers=2, block_size=64)
    train, val, chars = load_corpus(args.text, config["block_size"])
    config = dict(config, vocab=len(chars))
    if saved and chars != saved["chars"]:
        raise ValueError("恢复训练时字符表必须与保存时一致")
    model = GPT(**config).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
    start, best = 0, float("inf")
    if saved:
        model.load_state_dict(saved["model"])
        optimizer.load_state_dict(saved["optimizer"])
        start, best = saved["step"], saved["best_val"]
        torch.set_rng_state(saved["cpu_rng"])
        if str(device).startswith("cuda") and saved.get("cuda_rng") is not None:
            torch.cuda.set_rng_state_all(saved["cuda_rng"])
    if start >= args.steps:
        parser.error("steps 必须大于 checkpoint 已完成的更新步数")
    print("device:", device, "parameters:", sum(p.numel() for p in model.parameters()))
    for step in range(start, args.steps):
        model.train()
        x, y = get_batch(train, config["block_size"], args.batch_size, device)
        optimizer.zero_grad(set_to_none=True)
        _, loss = model(x, y)
        if not torch.isfinite(loss):
            raise RuntimeError("loss 非有限值，请停止并检查数据、学习率和精度")
        loss.backward()
        optimizer.step()
        completed = step + 1
        if completed % args.eval_every == 0 or completed == args.steps:
            train_loss = estimate_loss(model, train, config, args.batch_size, device)
            val_loss = estimate_loss(model, val, config, args.batch_size, device)
            improved = val_loss < best
            best = min(best, val_loss)
            state = dict(config=config, chars=chars, model=model.state_dict(),
                         optimizer=optimizer.state_dict(), step=completed, best_val=best,
                         cpu_rng=torch.get_rng_state(),
                         cuda_rng=torch.cuda.get_rng_state_all() if str(device).startswith("cuda") else None)
            torch.save(state, "last.pt")
            if improved:
                torch.save(state, "best.pt")
            print(f"step {completed}: train={train_loss:.3f}, val={val_loss:.3f}, best={best:.3f}")


if __name__ == "__main__":
    main()

train_gpt.py

"""python generate.py --checkpoint best.pt --prompt 'To be'"""
import argparse
import torch
from model import GPT


@torch.inference_mode()
def generate(model, idx, max_new_tokens):
    model.eval()
    for _ in range(max_new_tokens):
        logits, _ = model(idx[:, -model.block_size:])
        probs = logits[:, -1, :].softmax(dim=-1)
        idx = torch.cat((idx, torch.multinomial(probs, 1)), dim=1)
    return idx


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", default="best.pt")
    parser.add_argument("--prompt", default="")
    parser.add_argument("--tokens", type=int, default=200)
    args = parser.parse_args()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    saved = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
    model = GPT(**saved["config"]).to(device)
    model.load_state_dict(saved["model"])
    chars = saved["chars"]
    stoi = {c: i for i, c in enumerate(chars)}
    prompt = args.prompt or chars[0]
    unknown = set(prompt) - set(stoi)
    if unknown:
        raise ValueError(f"提示词含训练字符表之外的字符：{unknown}")
    idx = torch.tensor([[stoi[c] for c in prompt]], dtype=torch.long, device=device)
    result = generate(model, idx, args.tokens)
    print("".join(chars[i] for i in result[0].tolist()))


if __name__ == "__main__":
    main()

generate.py

C.3 机制实验：DistributedSampler 与最小 GPT DDP

只有想单独观察按索引分配样本时才做本实验，复用 C.2 的 model.py。主线 nanoGPT 使用独立随机窗口，不采用这里的 DistributedSampler。此脚本只保存模型权重，不能当作完整续训入口。

import os, torch, torch.distributed as dist
from torch.utils.data import Dataset, DataLoader, DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from model import GPT, build_dataset          # 复用 2.3 的 model.py

dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
is_main = dist.get_rank() == 0

# 把字符数据包成 Dataset，DistributedSampler 才能按样本切分
data, vocab_size, block_size = build_dataset()
class CharDataset(Dataset):
    def __len__(self):  return len(data) - block_size
    def __getitem__(self, i):
        return data[i:i+block_size], data[i+1:i+1+block_size]
train_ds = CharDataset()
num_epochs = 2

torch.manual_seed(42)
model = GPT(vocab=vocab_size, block_size=block_size).to(local_rank)
model = DDP(model, device_ids=[local_rank])   # 关键：包一层 DDP

# DistributedSampler 分配样本索引；默认补齐可能重复少量索引
sampler = DistributedSampler(train_ds)
loader = DataLoader(train_ds, batch_size=32, sampler=sampler,
                    num_workers=0, pin_memory=True)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, fused=False)

max_steps, completed_updates = 20, 0
for epoch in range(num_epochs):
    sampler.set_epoch(epoch)          # 必须！否则每个 epoch 会重复相同的 shuffle 顺序
    model.train()
    for x, y in loader:
        x, y = x.to(local_rank), y.to(local_rank)
        _, loss = model(x, y)
        optimizer.zero_grad()
        loss.backward()               # backward 时自动 all-reduce 梯度
        optimizer.step()
        completed_updates += 1
        if completed_updates >= max_steps:
            break
    if is_main:                       # 只在主进程打印/存
        print(f"updates {completed_updates}: loss={loss.item():.4f}")
        torch.save(model.module.state_dict(), f"ckpt_{epoch}.pt")  # 注意 .module

    if completed_updates >= max_steps:
        break

dist.destroy_process_group()
# 启动：torchrun --standalone --nproc_per_node=2 train_ddp.py

DDP 的三个易错点：（1）DistributedSampler 每个 epoch 必须调用 set_epoch(epoch)，否则每一轮会重复相同的全局 shuffle 顺序；不同 rank 仍会拿到不同分片，但跨 epoch 的顺序不会变化。（2）保存 checkpoint 使用 model.module.state_dict()，并且只在主进程执行。（3）梯度 all-reduce 由 backward() 自动触发，不需要手写通信。启动命令：torchrun --standalone --nproc_per_node=2 train_ddp.py。

首次只运行 20 次更新，确认两个 rank 都能完成训练、rank 0 保存权重，再增加训练时长。此处保存的是模型权重，可重新构建同一 GPT 后加载验证；接续训练还需要优化器、训练进度等状态，第 5.2 节再沿 nanoGPT 的完整 checkpoint 路径学习。

C.4 机制实验：最小 GPT 的 FSDP2 参数生命周期

复用 C.2 的 model.py，在小模型上看 fully_shard 的行为。它与 C.3 的采样方式不同，两个脚本的 loss 曲线不能直接作逐值正确性对照；主线训练、恢复与导出使用第 5 章 nanoGPT 入口。

import os, torch, torch.distributed as dist
from torch.distributed.fsdp import fully_shard
from model import GPT, build_dataset, get_batch

local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group("nccl")
device = torch.device("cuda", local_rank)
data, vocab_size, block_size = build_dataset()

# 先创建相同的小模型；大模型的 meta 初始化在小模型验收后研究。
torch.manual_seed(42)
model = GPT(vocab=vocab_size, block_size=block_size).to(device)
for block in model.blocks:
    fully_shard(block, reshard_after_forward=True)
fully_shard(model, reshard_after_forward=True)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

# 初始化参数后，再按 rank 区分数据采样随机数。
torch.manual_seed(1337 + dist.get_rank())
model.train()
for step in range(20):
    x, y = get_batch(data, block_size, batch_size=32, device=device)
    optimizer.zero_grad(set_to_none=True)
    _, loss = model(x, y)
    loss.backward()
    optimizer.step()
    if dist.get_rank() == 0 and step % 5 == 0:
        print(f"step {step}: local loss={loss.item():.4f}")

# 4.4 的 DCP save/load 必须放在销毁进程组之前，由所有 rank 参与。
dist.destroy_process_group()
# 启动：torchrun --standalone --nproc_per_node=2 fsdp_train.py

fsdp_train.py：先用与 DDP 相同的小模型和 FP32 验证两卡分片

这个两卡示例先用 FP32、小模型和相同的每卡 batch 观察分片生命周期；初始化时每卡暂存完整模型，因此适合验证原理。DDP 示例按索引遍历窗口，FSDP2 示例随机采窗口，日志也只是局部 loss；两段脚本的曲线不能直接作为逐值正确性对照。严谨对比时应固定两者的初始权重、每 rank 输入、loss 归一化与更新次数，再使用第 5.4 节的验收方法。

C.5 机制实验：查看最小 GPT 的算子耗时

此例复用 C.2 的 model.py，帮助认识 profiler 输出。它测的是一个独立小模型的训练区间，不能替代第 3～6 章 nanoGPT 的端到端、多 rank 性能记录；增加第二张卡之前先确认要测的范围。

import torch
from torch.profiler import profile, ProfilerActivity
from model import GPT, build_dataset, get_batch   # 复用 2.3 的 model.py

data, vocab_size, block_size = build_dataset()
model = GPT(vocab=vocab_size, block_size=block_size).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
x, y = get_batch(data, block_size, batch_size=16, device="cuda")

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for _ in range(10):
        _, loss = model(x, y); loss.backward(); optimizer.step(); optimizer.zero_grad()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
prof.export_chrome_trace("trace.json")   # 用 Perfetto / chrome://tracing 可视化

用 profiler 找瓶颈

C.6 nanoGPT：固定输入的一次 DDP 更新对照

保存为 verify_ddp_step.py，放到 nanoGPT 根目录，按 4.4 运行。此处的 model.py 来自固定版 nanoGPT；它与 C.2 最小 GPT 的 model.py 不是同一个接口。参考模型也占用设备内存，因此该脚本只用于小模型正确性检查。本次新增脚本完成语法检查；分布式运行仍待目标机器验证。

"""One-update check for pinned nanoGPT, PyTorch 2.9.x.

Copy beside nanoGPT/model.py. CPU/Gloo is a logic check; CUDA/NCCL must
still be checked on the target GPUs. This is not a throughput benchmark.
"""
import argparse
import copy
import json
import os
from contextlib import nullcontext

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from model import GPT, GPTConfig


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', choices=['cpu', 'cuda'], default='cuda')
    parser.add_argument('--accum', type=int, default=1)
    args = parser.parse_args()
    if args.accum < 1:
        raise ValueError('accum must be positive')
    torch.set_num_threads(1)
    local_rank = int(os.environ['LOCAL_RANK'])
    if args.device == 'cuda':
        torch.cuda.set_device(local_rank)
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
    device = torch.device(args.device, local_rank) if args.device == 'cuda' else torch.device('cpu')
    dist.init_process_group('nccl' if args.device == 'cuda' else 'gloo')
    try:
        rank, world = dist.get_rank(), dist.get_world_size()
        B, T, A = 4, 32, args.accum
        torch.manual_seed(1337)
        model = GPT(GPTConfig(vocab_size=65, block_size=T, n_layer=2,
                              n_head=2, n_embd=64, dropout=0.0, bias=False)).to(device)
        reference = copy.deepcopy(model)
        wrapped = DDP(model, device_ids=[local_rank] if args.device == 'cuda' else None)
        generator = torch.Generator().manual_seed(2027)
        windows = torch.randint(65, (world * A * B, T + 1), generator=generator).to(device)
        x, y = windows[:, :-1].contiguous(), windows[:, 1:].contiguous()

        def optimizer(m):
            params = list(m.parameters())
            return torch.optim.AdamW([
                {'params': [p for p in params if p.dim() >= 2], 'weight_decay': 0.1},
                {'params': [p for p in params if p.dim() < 2], 'weight_decay': 0.0},
            ], lr=1e-3, betas=(0.9, 0.95), fused=False, foreach=False)

        opt, ref_opt = optimizer(model), optimizer(reference)
        ref_logits, ref_loss = reference(x, y)
        ref_loss.backward()
        local_loss = torch.zeros((), device=device)
        local_logits = []
        for micro in range(A):
            start = (rank * A + micro) * B
            context = wrapped.no_sync() if micro < A - 1 else nullcontext()
            with context:  # no_sync covers both forward and backward
                logits, loss = wrapped(x[start:start+B], y[start:start+B])
                (loss / A).backward()
            local_logits.append(logits.detach())
            local_loss += loss.detach() / A
        dist.all_reduce(local_loss)
        local_loss /= world
        errors = {'logits': 0.0, 'loss': 0.0, 'gradient': 0.0, 'parameter': 0.0}
        failures = []

        def compare(kind, name, actual, expected):
            errors[kind] = max(errors[kind], (actual - expected).abs().max().item())
            if not torch.allclose(actual, expected, atol=5e-6, rtol=1e-4):
                failures.append(f'{kind}:{name}')

        start = rank * A * B
        compare('logits', 'local', torch.cat(local_logits), ref_logits[start:start+A*B].detach())
        compare('loss', 'global_mean', local_loss, ref_loss.detach())
        ref_params = dict(reference.named_parameters())
        for name, parameter in model.named_parameters():
            compare('gradient', name, parameter.grad, ref_params[name].grad)
        opt.step()
        ref_opt.step()
        for name, parameter in model.named_parameters():
            compare('parameter', name, parameter.detach(), ref_params[name].detach())
        metrics = torch.tensor(list(errors.values()) + [float(bool(failures))], device=device)
        dist.all_reduce(metrics, op=dist.ReduceOp.MAX)
        if rank == 0:
            print(json.dumps({'device': args.device, 'world': world, 'A': A,
                              'tokens': world*A*B*T, 'max_abs_error': dict(zip(errors, metrics[:-1].tolist())),
                              'passed': metrics[-1].item() == 0}, indent=2), flush=True)
        if metrics[-1].item():
            raise AssertionError(f'rank {rank}: failed comparisons {failures[:8]}; inspect before changing tolerance')
    finally:
        dist.destroy_process_group()


if __name__ == '__main__':
    main()

verify_ddp_step.py：对应第 4.4 节

C.7 两层 MLP：前向、手算梯度与一次更新

将下列完整代码保存为 mlp_bridge_check.py，运行 python mlp_bridge_check.py。只需 PyTorch，使用 CPU 与 float64，对应 1.6.2；不依赖 nanoGPT 仓库。脚本先核对每一步形状与数值，再将解析梯度同 autograd 比较，最后验证一次 SGD 更新。预期初始 $L=0.625$，更新后约 0.486185699；去掉 ReLU 后，两层仿射变换可合并。

"""Two affine layers + ReLU: dimensions, manual gradients, and one SGD step."""
import json

import torch
from torch import nn


def main():
    torch.manual_seed(0)
    dtype = torch.float64
    x = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=dtype)  # (B=2, D=3)
    target = torch.tensor([[1], [4]], dtype=dtype)          # (B=2, O=1)
    w1 = torch.tensor([[1, 0], [0, 1], [1, 1]], dtype=dtype)  # (D=3, H=2)
    b1 = torch.tensor([-5, -6], dtype=dtype)                  # (H=2,)
    w2 = torch.tensor([[1], [-0.5]], dtype=dtype)             # (H=2, O=1)
    b2 = torch.tensor([0.5], dtype=dtype)                    # (O=1,)
    model = nn.Sequential(
        nn.Linear(3, 2, dtype=dtype), nn.ReLU(), nn.Linear(2, 1, dtype=dtype)
    )
    with torch.no_grad():
        # nn.Linear stores weight as (out_features, in_features).
        model[0].weight.copy_(w1.T)
        model[0].bias.copy_(b1)
        model[2].weight.copy_(w2.T)
        model[2].bias.copy_(b2)

    z1 = model[0](x)       # (2, 2)
    h = model[1](z1)       # (2, 2)
    pred = model[2](h)     # (2, 1)
    loss = (pred - target).square().mean()  # scalar: ()
    torch.testing.assert_close(z1, x.new_tensor([[-1, -1], [5, 5]]))
    torch.testing.assert_close(h, x.new_tensor([[0, 0], [5, 5]]))
    torch.testing.assert_close(pred, x.new_tensor([[0.5], [3]]))
    torch.testing.assert_close(loss, x.new_tensor(0.625))

    # Manual backward, using the original mathematical W shapes.
    with torch.no_grad():
        d_pred = 2 * (pred - target) / pred.numel()  # (2, 1)
        d_w2 = h.T @ d_pred                         # (2, 1)
        d_b2 = d_pred.sum(dim=0)                    # (1,)
        d_h = d_pred @ w2.T                         # (2, 2)
        d_z1 = d_h * (z1 > 0)                       # (2, 2)
        d_w1 = x.T @ d_z1                           # (3, 2)
        d_b1 = d_z1.sum(dim=0)                      # (2,)
    expected = [d_w1.T, d_b1, d_w2.T, d_b2]
    loss.backward()
    gradients = {}
    for (name, parameter), manual in zip(model.named_parameters(), expected):
        assert parameter.grad.shape == parameter.shape
        torch.testing.assert_close(parameter.grad, manual)
        gradients[name] = {
            "shape": list(parameter.grad.shape), "value": parameter.grad.tolist()
        }

    # Without ReLU, two affine layers collapse into a single affine map.
    without_relu = (x @ w1 + b1) @ w2 + b2
    collapsed = x @ (w1 @ w2) + (b1 @ w2 + b2)
    torch.testing.assert_close(without_relu, collapsed)
    assert not torch.allclose(pred.detach(), collapsed)

    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
    with torch.no_grad():
        loss_after = (model(x) - target).square().mean()
    assert loss_after < loss
    print(json.dumps({
        "torch": torch.__version__, "device": "cpu", "dtype": str(dtype),
        "z1": z1.tolist(), "h": h.tolist(), "pred": pred.tolist(),
        "loss_before": loss.item(), "loss_after_sgd": loss_after.item(),
        "gradients": gradients,
        "without_relu": without_relu.tolist(),
        "collapsed_weight_shape": list((w1 @ w2).shape),
        "collapsed_bias_shape": list((b1 @ w2 + b2).shape),
        "checks": "forward, manual gradients, gradient shapes, affine collapse, SGD: passed"
    }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()

mlp_bridge_check.py：可独立运行的两层 MLP 验算

C.8 计算与显存账本：验证重计算，再测完整训练步

将下列代码保存为 resource_probe.py。python resource_probe.py --device cpu 核对矩阵运算量、逻辑字节数以及普通执行与重计算的数值一致性；python resource_probe.py --device cuda 才在 GPU 上测量。同一个模型、初始化、输入和优化器分别运行普通分支与 checkpoint 分支，各自至少预热 3 次、测量 10 次；无 CUDA 时明确报错，不生成虚构的 GPU 数据。

本实验用 $\operatorname{mean}(\mathrm{output}^2)$ 作为目标，相当于向零目标拟合，只用于机制验证。GPU 规模为 512 个输入行、$1024\longrightarrow4096\longrightarrow1024$ 的 MLP，FP32、无 Dropout、关闭 TF32；前向两次矩阵乘共 8,589,934,592 FLOPs，参数连 bias 共 8,393,728 个。参数、完整梯度与 Adam m/v 四类逻辑数据合计 134,299,648 B，实际峰值另含激活、临时量等。CPU 只验证小模型，避免为了教学核对运行大负载。

"""CPU: verify estimates and checkpoint correctness; optional CUDA measurement.

Run: python resource_probe.py
     python resource_probe.py --device cuda --warmup 3 --steps 10
CUDA results are steady-state, full forward/backward/AdamW steps in FP32.
"""
import argparse
import gc
import json
import platform
import time

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

SEED = 20261010


def make_mlp(d, h, dtype=torch.float32):
    return nn.Sequential(nn.Linear(d, h, dtype=dtype), nn.GELU(),
                         nn.Linear(h, d, dtype=dtype))


def estimates():
    # One matrix multiply: each input read once, output written once.
    m, k, n = 2, 3, 2
    x, w, y = torch.empty(m, k), torch.empty(k, n), torch.empty(m, n)
    traffic = sum(t.numel() * t.element_size() for t in (x, w, y))
    assert traffic == 4 * (m * k + k * n + m * n) == 64
    assert 2 * m * k * n == 24
    # CPU does NOT construct the larger CUDA benchmark model.
    tokens, d, h, bytes_per_element = 512, 1024, 4096, 4
    parameters = d * h + h + h * d + d
    tiny = make_mlp(16, 64)
    assert sum(p.numel() for p in tiny.parameters()) == 16 * 64 + 64 + 64 * 16 + 16
    return {
        "matrix_example": {
            "X_shape": [m, k], "W_shape": [k, n], "Y_shape": [m, n],
            "dtype": "float32", "matmul_flops": 2 * m * k * n,
            "ideal_bytes": traffic, "flops_per_ideal_byte": 24 / traffic,
            "assumptions": "2 FLOPs per multiply-add; X/W read once, Y written once; no bias, cache reuse, or extra workspace",
        },
        "cuda_mlp_estimate_only": {
            "tokens": tokens, "input_dim": d, "hidden_dim": h,
            "parameters_including_bias": parameters,
            "forward_matmul_flops": 2 * tokens * (d * h + h * d),
            "parameter_bytes": parameters * bytes_per_element,
            "gradient_bytes": parameters * bytes_per_element,
            "adam_m_v_bytes": 2 * parameters * bytes_per_element,
            "parameter_gradient_m_v_total_bytes": 4 * parameters * bytes_per_element,
            "input_bytes": tokens * d * bytes_per_element,
            "hidden_tensor_bytes": tokens * h * bytes_per_element,
            "assumptions": "FP32 weights/gradients/moments, no master-weight copy; excludes activations, optimizer step scalars, allocator overhead and workspaces; FLOPs exclude bias/GELU/backward/optimizer",
        },
    }


def verify_checkpoint():
    # This small CPU check runs even when --device cuda is requested.
    torch.manual_seed(SEED)
    plain = make_mlp(16, 64, torch.float64)
    recomputed = make_mlp(16, 64, torch.float64)
    recomputed.load_state_dict(plain.state_dict())
    x_plain = torch.randn(8, 16, dtype=torch.float64, requires_grad=True)
    x_recomputed = x_plain.detach().clone().requires_grad_()
    y_plain = plain(x_plain)
    y_recomputed = checkpoint(recomputed, x_recomputed, use_reentrant=False)
    torch.testing.assert_close(y_plain, y_recomputed)
    loss_plain = y_plain.square().mean()
    loss_recomputed = y_recomputed.square().mean()
    loss_plain.backward()
    loss_recomputed.backward()
    torch.testing.assert_close(x_plain.grad, x_recomputed.grad)
    grad_errors = {}
    for (name, p), (other_name, q) in zip(
        plain.named_parameters(), recomputed.named_parameters()
    ):
        assert name == other_name and p.grad.shape == q.grad.shape == p.shape
        torch.testing.assert_close(p.grad, q.grad)
        grad_errors[name] = (p.grad - q.grad).abs().max().item()
    return {
        "device": "cpu", "dtype": "float64", "input_shape": [8, 16],
        "model": "Linear(16,64) -> GELU -> Linear(64,16); dropout=0",
        "checkpoint": "whole MLP; use_reentrant=False",
        "loss": loss_plain.item(),
        "max_output_abs_error": (y_plain - y_recomputed).abs().max().item(),
        "max_input_gradient_abs_error": (x_plain.grad - x_recomputed.grad).abs().max().item(),
        "max_parameter_gradient_abs_errors": grad_errors, "passed": True,
    }


def training_step(model, x, optimizer, use_checkpoint):
    optimizer.zero_grad(set_to_none=True)
    y = checkpoint(model, x, use_reentrant=False) if use_checkpoint else model(x)
    loss = y.square().mean()
    loss.backward()
    optimizer.step()
    # No output tensor is retained by the caller across steps.


def measure_branch(state_cpu, input_cpu, use_checkpoint, warmup, steps):
    # Each branch creates/releases its own model, optimizer, input and gradients.
    gc.collect()
    torch.cuda.empty_cache()
    model = make_mlp(1024, 4096).cuda()
    model.load_state_dict(state_cpu)
    x = input_cpu.cuda()
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=1e-4, weight_decay=0.01, foreach=False, fused=False
    )
    for _ in range(warmup):
        training_step(model, x, optimizer, use_checkpoint)
    optimizer.zero_grad(set_to_none=True)
    torch.cuda.synchronize()
    # Adam m/v exist after warmup; no parameter gradients remain at baseline.
    baseline = torch.cuda.memory_allocated()
    torch.cuda.reset_peak_memory_stats()
    start_event, end_event = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    wall_start = time.perf_counter()
    start_event.record()
    for _ in range(steps):
        training_step(model, x, optimizer, use_checkpoint)
    end_event.record()
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - wall_start
    peak = torch.cuda.max_memory_allocated()
    result = {
        "checkpoint": use_checkpoint, "warmup_steps": warmup, "measured_steps": steps,
        "baseline_allocated_bytes": baseline, "peak_allocated_bytes": peak,
        "peak_extra_allocated_bytes": peak - baseline,
        "mean_step_wall_ms": elapsed * 1000 / steps,
        "mean_step_cuda_event_ms": start_event.elapsed_time(end_event) / steps,
        "baseline_definition": "after warmup and zero_grad(set_to_none=True); model/input/Adam moments live, gradients absent",
        "peak_extra_definition": "peak minus baseline; includes gradients, activations, recomputation and optimizer temporaries, NOT just saved activations",
    }
    # Synchronize and release all branch tensors before starting the other branch.
    del model, x, optimizer, start_event, end_event
    gc.collect()
    torch.cuda.synchronize()
    torch.cuda.empty_cache()
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
    parser.add_argument("--warmup", type=int, default=3)
    parser.add_argument("--steps", type=int, default=10)
    args = parser.parse_args()
    if args.warmup < 3 or args.steps < 10:
        parser.error("use at least 3 warmup steps and 10 measured steps")
    if args.device == "cuda" and not torch.cuda.is_available():
        parser.error("--device cuda was requested, but CUDA is unavailable; no GPU measurement was made")
    torch.manual_seed(SEED)
    report = {
        "environment": {"python": platform.python_version(), "torch": torch.__version__,
                        "seed": SEED, "requested_device": args.device},
        "estimates": estimates(), "checkpoint_correctness": verify_checkpoint(),
        "cuda_measurements": None,
    }
    if args.device == "cuda":
        torch.manual_seed(SEED)
        torch.set_float32_matmul_precision("highest")
        torch.backends.cuda.matmul.allow_tf32 = False
        initial = make_mlp(1024, 4096)
        state_cpu = {k: v.detach().clone() for k, v in initial.state_dict().items()}
        input_cpu = torch.randn(512, 1024)
        del initial
        report["environment"].update({
            "gpu": torch.cuda.get_device_name(), "cuda_runtime": torch.version.cuda,
            "dtype": "float32", "tf32": False,
            "model": "Linear(1024,4096) -> GELU -> Linear(4096,1024); dropout=0",
            "optimizer": "AdamW(lr=1e-4,weight_decay=.01,foreach=False,fused=False)",
            "loss": "mean(output**2)", "input_requires_grad": False,
        })
        report["cuda_measurements"] = [
            measure_branch(state_cpu, input_cpu, enabled, args.warmup, args.steps)
            for enabled in (False, True)
        ]
        report["measurement_scope"] = (
            "Same initialization/input/optimizer/warmup in both branches; only checkpoint differs. "
            "Single GPU, one steady-state window per branch; timing is descriptive, not a speedup guarantee. "
            "allocated is PyTorch tensor memory, not total device usage or reserved memory. "
            "A lower saved-activation cost need not lower the full-step peak if another phase dominates."
        )
    else:
        report["measurement_scope"] = "CPU correctness and analytical estimates only; no GPU timing or memory measurements"
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()

resource_probe.py：CPU 正确性与可选 CUDA 重计算实验

阅读结果时区分 estimates（理论账本）、checkpoint_correctness（CPU 数值验证）和 cuda_measurements（仅 CUDA 分支实际运行后产生）。$\text{peak\_extra\_allocated\_bytes}=\text{peak}-\text{baseline}$，包含新增梯度、激活与临时量，不是单独的激活显存。CUDA Event 区间与主机计时的完整 step 区间均记录，任何一次结果都需连同设备、版本、形状和测量范围解释。本文新增示例在 Python 3.12.14 / PyTorch 2.9.1 的 CPU 环境通过；未给出 GPU 性能结论。

C.9 概念验证与接口练习

本节收纳附录 A 对应的手算、运行代码与实现细节。查概念时阅读附录 A；需要逐步计算或运行验证时，再按条目进入这里。

C.9.1 梯度：手算、自动求导与有限差分

梯度是导数信息，误差 residual 是预测与目标的差，两者不同。设 X$(N,1)$、w$(1,1)$、b$(1,)$，$\mathrm{pred}=Xw+b$ 得到 $(N,1)$，目标 Y$(N,1)$，残差 $R=\mathrm{pred}-Y$ 为 $(N,1)$。本例 $L=\frac{1}{N}\sum_{i=1}^{N}R_i^2$ 是 $\text{scalar},\quad\operatorname{shape}=()$。

由链式法则：$\frac{\partial L}{\partial\mathrm{pred}}=\frac{2}{N}R$，shape=$(N,1)$；$\frac{\partial L}{\partial w}=X^{\mathsf T}\left(\frac{2}{N}R\right)$，$(1,N)\mathbin{@}(N,1)\longrightarrow(1,1)$；$\frac{\partial L}{\partial b}=\sum_{i=1}^{N}\frac{2}{N}R_i$，shape=$(1,)$。bias 被广播给每个样本，反向时要把这些路径的贡献加回同一份 bias。若输出有多列且 loss 对全部元素取 mean，分母还要包含输出列数。

用同一关系 $y=2x+1$ 手算两个样本：$X=\begin{bmatrix}1\\2\end{bmatrix}$，$Y=\begin{bmatrix}3\\5\end{bmatrix}$；$w=b=0$ 时 $\mathrm{pred}=0$，$L=17$，$\frac{\partial L}{\partial w}=-13$、$\frac{\partial L}{\partial b}=-8$。SGD 学习率 0.1 更新为 $w=1.3$、$b=0.8$，新 $L=1.685$。下面同时验证解析梯度、autograd 与有限差分；有限差分只用于小例子检查，不替代训练中的反向传播。

import torch
X = torch.tensor([[1.], [2.]], dtype=torch.float64)  # (N,1)=(2,1)
Y = 2 * X + 1                                     # (2,1): [[3],[5]]
w = torch.zeros(1, 1, dtype=X.dtype, requires_grad=True)
b = torch.zeros(1, dtype=X.dtype, requires_grad=True)
pred = X @ w + b                                  # (2,1)
residual = pred - Y                               # (2,1)
loss = residual.square().mean()                   # ()
loss.backward()
expected_w = (2 / X.shape[0]) * X.T @ residual.detach()  # (1,1)
expected_b = (2 / X.shape[0]) * residual.detach().sum(0) # (1,)
torch.testing.assert_close(w.grad, expected_w)
torch.testing.assert_close(b.grad, expected_b)
assert loss.item() == 17 and w.grad.item() == -13 and b.grad.item() == -8
with torch.no_grad():
    delta = 1e-6
    plus = ((X @ (w + delta) + b - Y) ** 2).mean()
    minus = ((X @ (w - delta) + b - Y) ** 2).mean()
    torch.testing.assert_close(((plus-minus)/(2*delta)).reshape_as(w), w.grad)
    w -= 0.1 * w.grad
    b -= 0.1 * b.grad
    after = ((X @ w + b - Y) ** 2).mean()
torch.testing.assert_close(after, torch.tensor(1.685, dtype=X.dtype))
print('loss 17 ->', after.item(), 'w/b shapes:', w.shape, b.shape)
w.grad = None
b.grad = None

gradient_check.py：CPU 核对导数、shape 与一次更新

对标量 loss L，$\frac{\partial L}{\partial\theta}$ 与 $\theta$ 同 shape，并不意味着完整雅可比也同 shape。参数得到 .grad 后，optimizer.step 才更新参数；多次 backward 会累积，zero_grad 清除梯度，Adam 的 m/v 则保留。默认参数梯度在叶子上保留，中间激活的 .grad 需 retain_grad；detach 会切断当前依赖路径，见 1.5。

C.9.2 Adam：首步数值与状态验证

偏差修正补偿从零状态开始导致的早期偏小；v 较大的分量会被较大的分母缩放，但不能把它简单解释为“该参数不重要”。例如 $\beta_1=0.9$、$\beta_2=0.999$、$g_1=-13$，得到 $m_1=-1.3$、$v_1=0.169$，修正后 $\hat m_1=-13$、$\hat v_1=169$。$\eta=0.1$、初值 0 时，Adam 首次更新后约为 0.1；同梯度下基本 SGD 更新为 1.3。不同更新幅度不说明哪种优化器必然更好。

下面使用 C.9.1 中线性回归的一个参数梯度，验证 Adam 首步数值、状态形状以及 zero_grad 的作用范围；这里只隔离优化器机制，真实训练仍由 backward 生成梯度。

import torch
p = torch.nn.Parameter(torch.tensor([[0.]], dtype=torch.float64))  # (1,1)
opt = torch.optim.Adam([p], lr=0.1, betas=(0.9, 0.999), eps=1e-8,
                       weight_decay=0, amsgrad=False, foreach=False, fused=False)
p.grad = torch.tensor([[-13.]], dtype=torch.float64)  # A.12 首步的 w 梯度
opt.step()
state = opt.state[p]
torch.testing.assert_close(state['exp_avg'], torch.tensor([[-1.3]], dtype=p.dtype))
torch.testing.assert_close(state['exp_avg_sq'], torch.tensor([[0.169]], dtype=p.dtype))
expected = 0.1 * 13 / (13 + 1e-8)
torch.testing.assert_close(p, torch.full_like(p, expected))
print('parameter/m/v shapes:', p.shape, state['exp_avg'].shape, state['exp_avg_sq'].shape)
saved_m = state['exp_avg'].clone()
opt.zero_grad(set_to_none=True)
assert p.grad is None
torch.testing.assert_close(state['exp_avg'], saved_m) # 清梯度不清历史
print('Adam first update:', p.item(), 'step:', state['step'].item())

adam_state_check.py：CPU 检查 Adam 的首步状态

接口参考：PyTorch Adam 与 AdamW 文档。学习时先关闭 fused/foreach 等执行优化以便观察，之后再分别验证数值和性能；优化实现不是另一种 batch 或梯度定义。

C.9.3 Batch：执行过程与梯度累积验证

先只记住两件事：一次 forward/backward 处理了多少条样本，以及一次 optimizer.step 使用了多少条样本的梯度。两者在最简单的训练中相等；加入梯度累积或数据并行后，就可能不同。下面始终用同一组 16 条样本，先看过程，再给这些数量命名。

C.9.3.1 单卡一次算完：16 条样本，更新一次参数

假设训练集有 64 条样本。对本文的字符语言模型，一条样本是一对长度为 32 的输入/目标窗口：$\operatorname{shape}(x)=(32,)$，$\operatorname{shape}(y)=(32,)$。所以一条样本不是一个 token；本例每条样本提供 32 个有效预测位置。先取编号 1—16 的样本，堆起来后 x、y 的 shape 都是 $(16,32)$。

这 16 条样本共同用于计算一次参数更新的梯度，因此组成一个 mini-batch；按全局范围统计，global batch 也是这 16 条样本。$\text{batch size}=16$ 在这里数的是样本条数。接着清梯度，前向得到预测，将全部有效目标的损失取 mean 得到 loss，再用 backward 求这个平均 loss 对参数的梯度，最后 optimizer.step 更新一次参数。

这一轮发生了 1 次前向、1 次 backward、1 次 optimizer.step，使用了 16 条样本。再取 17—32、33—48、49—64，各更新一次；若所有样本本轮恰好使用一次、不丢末批，也没有重复采样，4 次参数更新就遍历了整个训练集，这才叫 1 个 epoch。一个 step 不等于一个 epoch。

C.9.3.2 单卡分四批算：每次 4 条，累计 16 条后才更新

现在假设显存一次只能处理 4 条，但希望每次参数更新仍使用相同的 16 条样本。就把它们依次分成 1—4、5—8、9—12、13—16 四批；每批的 x、y 都是 $(4,32)$。一次实际送进模型做前向与反向的小批称为 micro-batch，本例 $\text{micro-batch size}=4$。

处理哪批样本	backward 之后的梯度来自谁	是否更新参数
开始前	zero_grad 清除上一轮梯度	否
1—4	第 1 批的贡献	否
5—8	第 1、2 批的贡献之和	否
9—12	第 1、2、3 批的贡献之和	否
13—16	四批的贡献之和	随后执行 1 次 optimizer.step

这里的四次 backward 都使用同一份尚未更新的参数，梯度不断加到 .grad；中间既不能清梯度，也不能执行 step。这叫梯度累积，累积次数为 4。把 optimizer.step 放到循环内部，会变成四次先后改变参数的更新，已经不是这里要模拟的一次 16 条样本更新。

本例四个小批的有效目标数相同，且每批 loss 都是对有效目标取 mean 时，每次反向前应将 loss 除以 4。于是累积后的梯度是四批平均梯度的平均，也就是这 16 条样本的整体平均梯度；若不除以 4，梯度量级会变成预期的四倍。下面只展示循环结构，model、optimizer、四组 (xb,yb) 来自前文训练过程；完整可运行核对在 C.9.3.6。

optimizer.zero_grad(set_to_none=True)
for xb, yb in four_micro_batches:  # 共 4 批；每批 4 条样本
    loss = compute_mean_loss(model(xb), yb)
    (loss / 4).backward()          # 只累积梯度，参数保持不变
optimizer.step()                  # 四批都结束后，才更新 1 次

循环骨架：compute_mean_loss 表示本任务对有效目标取均值的损失

因此本例有 4 次前向、4 次 backward，但只有 1 次参数更新。四个 micro-batch 合起来构成这次更新的 mini-batch；global batch size 为 $B_{\mathrm{global}}=4\times4=16$，本文也称为 effective batch size（有效批量）。这 16 条样本可以分次计算，无需同时放入显存。

C.9.3.3 两卡 DDP：每卡累计 8 条，合起来更新同一个模型

再增加一张卡，先只用 DDP。两张卡各保存同样的模型参数与优化器，但分别处理这 16 条样本中的一半。为了仍保持 $B_{\mathrm{global}}=16$，每卡一次处理 4 条，只需要累积两批：

数据并行副本	第 1 次前向/反向	第 2 次前向/反向	本轮累计样本
GPU 0 / rank 0	样本 1—4	样本 5—8	8 条
GPU 1 / rank 1	样本 9—12	样本 13—16	8 条

每卡将各自两批的 mean loss 分别除以 2 再 backward。标准 DDP 默认再对两个副本的参数梯度取平均，最终每卡得到这 16 条样本的整体平均梯度。因此实际 DDP 代码中此处除的是本地累积次数 2，不要再额外除一次 GPU 数 2，否则会重复缩小梯度。这里假设没有自定义通信 hook，两个副本和各小批的有效目标数相等。

梯度同步完成后，两张卡各自执行一次 optimizer.step；相同参数、优化器状态与梯度产生相同更新。虽然程序里有两个 step 调用，它们是在各副本上完成同一轮全局更新，不能把 global batch 再乘 2，也不能算成模型先后更新了两轮。默认 DDP 可在每次 backward 同步；用 no_sync 将前几批同步延后，是减少通信的实现方式，不改变本例的样本计数。

现在，micro-batch size 是每卡每次计算的 4 条；local mini-batch size（每副本、每次更新）是每卡累计的 8 条，也就是 local samples per update；global batch size 是两卡共同用于这次更新的 16 条。本地与全局的差别来自两个数据并行副本共同参与更新，本地批次与微批的差别来自每个副本累积了两次计算。

打开 SVG 原图

C.9.3.4 再归纳名字与公式：每次算多少、累计几次、几份数据

名称	本节用它回答什么	两卡例子
batch / batch size	一批数据 / 这批数据有几条样本	说大小时必须说明是哪一级
mini-batch（算法语境）	标准 mini-batch SGD 中，共同形成一次参数更新的样本批次	本次更新的 16 条样本，与 global batch 相同
micro-batch size，记作 $b$	每个数据副本一次 forward/backward 处理几条	$$b=4$$
gradient accumulation steps，记作 $A$	每个副本累计几批才更新一次参数	$$A=2$$
data parallel size，记作 $D$	有几份分别取数、同步梯度的模型副本	$$D=2$$
local mini-batch size，记作 $B_{\mathrm{local}}$	每个数据副本一次更新累计的样本条数，即 local samples per update	$$B_{\mathrm{local}}=A\times b=8$$
global / effective batch size	同一轮参数更新在所有数据副本上共使用多少条	$$D\times b\times A=16$$

这三层计数分别对应：micro-batch 统计一个副本一次计算，local mini-batch 统计一个副本一次更新，global batch 统计所有副本一次更新。表中“mini-batch（算法语境）”指最后一层的完整批次；框架若把本地批次也叫 mini-batch，应先对应到具体范围。阅读代码时，按 $b$、$A$、$D$ 和 optimizer.step 的位置核对字段含义。

当各副本的小批大小一致，且每个副本都累积 A 批时：$B_{\mathrm{global}}=D\times b\times A$。本文其他章节常用 B 表示某次输入的 batch 轴长度，在这些场景下 $B=b$。下面每行都完成同样的 16 条样本更新，只改变执行方式；反向次数一栏按每个数据副本计算。

执行方式	D	b	A / 每副本 backward 次数	global batch
单卡一次算完	1	16	1	16
单卡分四批	1	4	4	16
单机两卡 DDP	2	4	2	16
两机、每机两卡，纯 DDP	4	4	1	16
四卡：$\mathrm{DP}=2,\quad\mathrm{TP}=2$	2	4	2	16

最后一行只有两组独立取数的数据副本：每组内部两张 TP 卡协作处理同一份输入，所以 $D=2$，不能把总 GPU 数 4 代入 D。PP 的不同阶段、CP 的不同序列片段也可能共同服务同一批样本；EP 把 token 路由给多个专家，也不因此增加原始样本条数。纯 DDP 才有 $D=\text{world\_size}$；机器数本身不另外乘进公式。

本例每条样本有 $T=32$ 个有效预测位置，所以 $\text{tokens per update}=16\times32=512$。16 的单位是样本，512 的单位是有效目标 token；序列长度不会改变 $\text{batch size}=16$。样本条数按参与计算的出现次数计，即使两条样本内容相同也不去重。

C.9.3.5 看真实训练配置时，再检查三个边界

第一，配置字段可能已经包含某个乘数。本教程固定版 nanoGPT 的 batch_size 对应 b，但命令行 gradient_accumulation_steps 记为 G：纯 DDP 启动后，程序将每个 rank 实际执行的次数设为 $A=G/D$，要求整除。比如 $b=4$、$G=4$，单卡实际 $A=4$，两卡 DDP 实际 $A=2$，两者 global batch 都是 $b\times G=16$。不能再用 $D\times b\times G$，否则重复乘了 D。这里的 A 始终指实际循环次数；Picotron 等项目应按固定版本的取批与更新循环核对，不直接照搬 nanoGPT 的字段含义。

第二，样本数与 loss 的分母不是一回事。前面每批都恰好有 $4\times32$ 个有效目标，所以简单除以 A 正确。若最后一批不足、存在 padding/ignore_index、变长或打包，按实际有效目标数加权；全局 token 数应是各批有效目标数之和。两批分别有 4 和 2 个有效目标、各自 mean loss 为 $L_1,L_2$ 时，合并后的 mean loss 是 $\frac{4L_1+2L_2}{6}$，而不是 $\frac{L_1+L_2}{2}$。在 DDP 中还需计入其梯度平均因子；不要对不同大小批次的均值直接平均。

第三，统计口径相同不代表训练结果逐值相同。要比较一次更新，还要使用相同初始参数、相同实际样本、相同损失归一化与优化器状态；Dropout、BatchNorm、随机状态与浮点归约顺序会影响结果。iter/iteration 在代码里可能指取批循环，也可能指参数更新，先看 optimizer.step 实际在哪里；随机窗口或流式数据常用更新次数、累计 token 数记录进度，不强行套固定 epoch。

C.9.3.6 跑一次核对：四次 backward 是否等于四次更新

答案是否定的，关键看 step 放在哪里。下面沿用第 1 章的线性回归 $y=2x+1$，用 16 条标量输入样本验证更新过程；这里 $\operatorname{shape}(X)=(16,1)$，不运行前面的字符模型。两份模型从同一初始参数出发，一份一次处理 16 条，另一份分四批各处理 4 条；两份都只 step 一次，并对照梯度和更新后的参数。

"""CPU check: one batch of 16 and four micro-batches of 4, one update each."""
import copy
import torch

torch.manual_seed(7)
x = torch.linspace(-1, 1, 16, dtype=torch.float64).reshape(16, 1)
y = 2 * x + 1
full = torch.nn.Linear(1, 1).double()
accumulated = copy.deepcopy(full)  # same initial parameters


def one_update(model, micro_batch_size):
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    accumulation_steps = len(x) // micro_batch_size
    assert len(x) % micro_batch_size == 0
    optimizer.zero_grad(set_to_none=True)
    backward_calls = 0
    for start in range(0, len(x), micro_batch_size):
        xb = x[start:start + micro_batch_size]
        yb = y[start:start + micro_batch_size]
        loss = (model(xb) - yb).square().mean()
        (loss / accumulation_steps).backward()
        backward_calls += 1
    gradients = [p.grad.detach().clone() for p in model.parameters()]
    optimizer.step()  # outside the loop: exactly one parameter update
    print(f"micro_batch={micro_batch_size}, backward={backward_calls}, "
          f"optimizer.step=1, global_batch={len(x)}")
    return gradients


full_gradients = one_update(full, micro_batch_size=16)
accumulated_gradients = one_update(accumulated, micro_batch_size=4)
for a, b in zip(full_gradients, accumulated_gradients):
    torch.testing.assert_close(a, b, rtol=1e-12, atol=1e-12)
for a, b in zip(full.parameters(), accumulated.parameters()):
    torch.testing.assert_close(a, b, rtol=1e-12, atol=1e-12)
print("PASS: same mean gradient and same parameters after one update")

batch_update_check.py：CPU 对照一大批与四小批的一次 SGD 更新

运行 python batch_update_check.py，应看到两行分别为 micro_batch=16 / backward=1 和 micro_batch=4 / backward=4；两行的 optimizer.step 都为 1、global_batch 都为 16，最后输出 PASS。这是 CPU 上对单卡梯度累积的验证，不是 NCCL 或多机测试；真实 DDP 的一次更新对照见 4.4 和 C.6。

最后自测：如果固定 $b=4$、$A=2$，只把纯 DDP 从两卡变成四卡，global batch 会从 16 变为 32；要仍处理 16 条，可将 A 改为 1。若把 T 从 32 增至 64，则 batch 条数不变，每次更新的有效 token 数翻倍。若把 step 放进四批循环，每处理 4 条就更新一次，此时是 4 次 $B_{\mathrm{global}}=4$ 的更新。

C.9.4 PyTorch：训练循环与接口阅读

下面展示训练、验证和无标签推理的循环骨架。测试集有标签时，评估过程与验证类似，但应留到模型选择结束后使用；无标签数据只能产生预测，不能据此计算准确率。完整例子见 2.2。

# ---- 训练：唯一会更新参数的循环 ----
model.train()                       # 开启 dropout / BatchNorm 的训练行为
for x, y in tr_set:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad()           # 清空上一步的梯度
    pred = model(x)
    loss = criterion(pred, y)
    loss.backward()                 # 反向传播，算出梯度
    optimizer.step()                # 沿梯度反方向更新参数

# ---- 验证：只看指标，不更新参数 ----
model.eval()                        # 关闭 dropout，BatchNorm 用统计均值
total_loss = 0
for x, y in dv_set:
    x, y = x.to(device), y.to(device)
    with torch.no_grad():           # 不建梯度图，省显存也更快
        pred = model(x)
        loss = criterion(pred, y)
    total_loss += loss.cpu().item() * len(x)
avg_loss = total_loss / len(dv_set.dataset)   # 按样本数加权平均

# ---- 无标签推理：只输出预测，不计算评估指标 ----
model.eval()
preds = []
for x in tt_set:
    x = x.to(device)
    with torch.no_grad():
        pred = model(x)
    preds.append(pred.cpu())

训练、验证与无标签推理的循环骨架

写 PyTorch 卡住，多半不是算法难，而是不会读文档签名、或撞上几类反复出现的报错。先说读文档：同一个函数换参数、行为就变，看懂签名比猜快得多。以 torch.max 为例，它有三种形态：

调用形式	行为
torch.max(input) → Tensor	返回整个张量的全局最大值（一个标量）
torch.max(input, dim) → (Tensor, LongTensor)	沿某个维度求最大，返回(最大值, 对应下标) 两个张量
torch.max(input, other) → Tensor	两个同形状张量逐元素取较大者

读签名还要看两件事：一是签名里的 *，它之前是位置参数（可以不写参数名按顺序传），之后是关键字参数（必须写参数名，如 out=）——像 torch.max(x, 0, False, p) 这样把 out 当第四个位置参数传就会报错；二是带默认值的参数（如 keepdim=False）可以不传。至于四类高频报错——device 不一致、维度不匹配、CUDA out of memory、Long/Float 类型错——已在正文 §1.6 用一张表逐条给了根因和一句话解法，这里不再重复，遇到报错直接去 §1.6 对号入座即可。

C.9.5 CUDA：向量加法 kernel

写一个 CUDA kernel，思路和写 Python 循环相反：不是「我遍历每个元素」，而是「我描述单个线程该干什么，然后启动几万个线程各干各的」。下面这个向量加法是最经典的入门例子，看懂它，前面的 grid/block/thread 就落地了：

// 教学代码片段：展示 kernel 与启动配置，不是独立可编译程序。
// 假设 a、b、out 已在同一 CUDA 设备上分配，至少各有 n 个 float，输入已准备好。
__global__ void add(const float* a, const float* b, float* out, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) out[i] = a[i] + b[i];
}

// 下列为 host 侧调用片段；实际工程还需分配、拷贝、错误检查和释放。
int n = 1000;
int threads = 256;
int blocks = (n + threads - 1) / threads;  // 4 个 block，共 1024 个线程
add<<<blocks, threads>>>(a, b, out, n);
cudaDeviceSynchronize();                  // 教学中等结果；实际需检查返回值

向量加法片段：把 1000 个元素映射到 1024 个线程

三个要点：（1）blockIdx.x * blockDim.x + threadIdx.x 是 CUDA 里最常见的一行——每个线程靠它算出自己的全局编号，从而知道该处理哪个数据；（2）<<<blocks, threads>>> 这个三尖括号语法就是在指定 grid 和 block 的大小，和 A.9.2 完全对应；（3）线程总数（blocks × threads）通常会略多于实际元素数，所以要有 if (i < n) 的越界保护。

C.10 集合通信实验：核对数据，再测消息尺寸

把本节完整代码保存为 collective_probe.py。使用第 3～6 章的 PyTorch 2.9.x 环境；跨机时每台机器放置同一份脚本。它先检查四种 collective 的具体输入输出，全部 rank 通过后才可进入 AllReduce 测量。Gloo/CPU 路径用于学习相同的数据语义，不能用它的耗时推断 NCCL 或 RDMA。

# 可选：两进程 CPU 语义检查
torchrun --standalone --nproc_per_node=2 collective_probe.py --backend gloo

# 单机两卡：先检查语义并记录 NCCL 选择的传输路径
NCCL_DEBUG=INFO torchrun --standalone --nproc_per_node=2 collective_probe.py --backend nccl

# 单机两卡：再测默认的 4 KiB、1 MiB、16 MiB、64 MiB
torchrun --standalone --nproc_per_node=2 collective_probe.py \
  --backend nccl --bench --sizes-mib 0.00390625,1,16,64 --warmup 20 --iters 50

两 rank 时的理论预期如下；这些数值是代码输入对应的数学结果，不是硬件运行日志。所有语义记录都应出现 all_ranks_passed=true。AllGather API 返回的是张量列表，脚本只是为了打印，用 stack 将其显示为 $(W,2)$；这与 6.2 沿首维拼接成完整张量是两种展示方式，收集到的数据相同。

操作	rank 0 输出	rank 1 输出
AllReduce SUM	$[3,3]$，shape $(2,)$	$[3,3]$，没有自动除以 2
AllGather（stack 后展示）	$\begin{bmatrix}0&10\\1&11\end{bmatrix}$，shape $(2,2)$	与 rank 0 相同
ReduceScatter SUM	$[100,102]$，shape $(2,)$	$[104,106]$，是求和结果的另一份分片
AllToAll	$[0,1,100,101]$，shape $(4,)$	$[10,11,110,111]$，收到发给自己的块，没有相加

随后做两机、每机一卡实验，使总 rank 数仍为 2；这才与上面的单机两卡形成相同组规模的互联对比。两个节点分别运行下面对应的命令，并将地址替换为实际可达的节点 0 地址。首次可加 NCCL_DEBUG=INFO 核对传输与网卡；只看到集合通信成功，不能断言使用了 RDMA。

# 节点 0
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
torchrun --nnodes=2 --nproc_per_node=1 --node_rank=0 \
  --master_addr="$NODE0_ADDR" --master_port=29610 \
  collective_probe.py --backend nccl --bench --warmup 20 --iters 50

# 节点 1：同一个节点 0 地址
NODE0_ADDR='REPLACE_WITH_NODE0_REACHABLE_ADDRESS'
torchrun --nnodes=2 --nproc_per_node=1 --node_rank=1 \
  --master_addr="$NODE0_ADDR" --master_port=29610 \
  collective_probe.py --backend nccl --bench --warmup 20 --iters 50

测量中的 FP32 缓冲在计时前分配，以零值反复 SUM，避免溢出，也不把每轮重新填充的代价混进来。每个尺寸先预热 20 次，barrier 与设备同步后测连续 50 次操作，窗口末等待设备完成，再取各 rank 窗口耗时的最大值。mean_ms_using_slowest_rank 是这个最大耗时除以次数后的毫秒数；payload_bytes_per_second 是单 rank 输入字节数除以该平均时间，不是网卡带宽。计时包含 Python 提交和完成等待；它也不等同于独立 kernel 的 CUDA Event 时长。

保存每条消息尺寸、耗时与传输路径，先判断小消息延迟和大消息速率如何变化，再用接近真实梯度桶或 TP 张量的尺寸补测。--sizes-mib 使用二进制 MiB；脚本要求正整数个字节且 FP32 元素数能被组大小整除，这是本实验的检查约束，AllReduce 本身并不要求后一个条件。若某项失败或超时，先定位对应的 collective 和 rank，不要继续解释后面的性能。

验证范围。本次在 Python 3.12.14、PyTorch 2.9.1 上完成语法/API 核对、离线数值模型及脚本控制流检查。本机执行环境禁止绑定本地通信端口，真实两进程 Gloo 检查未完成；没有 CUDA/NCCL、跨机或 RDMA 实测。上面的预期值不能替代这些运行验收，也没有可报告的通信性能数字。

"""Collective semantics and an optional FP32 AllReduce payload sweep.

Run with torchrun; Gloo/CPU checks do not validate CUDA/NCCL or RDMA.
The sweep reports payload/time, not bus bandwidth or network-interface bandwidth.
For equal-sized rank chunks, this probe requires numel % world_size == 0;
AllReduce itself does not require this divisibility restriction.
"""
import argparse
import json
import os
import socket
import time
from datetime import timedelta
from decimal import Decimal, InvalidOperation

import torch
import torch.distributed as dist


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--backend', choices=['gloo', 'nccl'], default='gloo')
    parser.add_argument('--bench', action='store_true')
    parser.add_argument('--sizes-mib', default='0.00390625,1,16,64')
    parser.add_argument('--warmup', type=int, default=20)
    parser.add_argument('--iters', type=int, default=50)
    args = parser.parse_args()
    if args.warmup < 1 or args.iters < 1:
        parser.error('--warmup and --iters must be positive.')
    local_rank = int(os.environ['LOCAL_RANK'])
    world = int(os.environ['WORLD_SIZE'])
    if world < 2:
        parser.error('Use at least two processes to observe communication.')
    sizes = []
    try:
        for item in (args.sizes_mib.split(',') if args.bench else []):
            byte_count = Decimal(item.strip()) * (1024 ** 2)
            if not byte_count.is_finite() or byte_count <= 0 or byte_count != int(byte_count):
                raise ValueError('Each size must specify a positive integer number of bytes.')
            if int(byte_count) % (4 * world):
                raise ValueError('FP32 element counts must be divisible by world_size in this probe.')
            sizes.append(int(byte_count))
    except (InvalidOperation, ValueError, OverflowError) as error:
        parser.error(str(error))
    if args.backend == 'nccl':
        if not torch.cuda.is_available() or local_rank >= torch.cuda.device_count():
            parser.error('NCCL needs a visible CUDA device for each local rank.')
        torch.cuda.set_device(local_rank)
    device = torch.device('cuda', local_rank) if args.backend == 'nccl' else torch.device('cpu')
    torch.set_num_threads(1)

    def sync():
        if device.type == 'cuda':
            torch.cuda.synchronize(device)

    dist.init_process_group(args.backend, timeout=timedelta(seconds=120))
    try:
        rank = dist.get_rank()
        print(json.dumps({'kind': 'environment', 'host': socket.gethostname(), 'rank': rank,
            'local_rank': local_rank, 'world_size': world, 'backend': args.backend,
            'device': str(device), 'torch': torch.__version__, 'cuda_build': torch.version.cuda,
            'gpu': torch.cuda.get_device_name(device) if device.type == 'cuda' else None}), flush=True)
        checks = []

        def check(name, source, output, expected):
            passed = torch.equal(output, expected)
            checks.append({'op': name, 'input_shape': list(source.shape),
                'output_shape': list(output.shape), 'input_values': source.cpu().tolist(),
                'output_values': output.cpu().tolist(), 'passed': passed})

        source = torch.full((2,), float(rank + 1), device=device)
        output = source.clone()
        dist.all_reduce(output, op=dist.ReduceOp.SUM)
        check('all_reduce_SUM_not_mean', source, output,
              torch.full_like(output, world * (world + 1) / 2))
        source = torch.tensor([rank, rank + 10], dtype=torch.float32, device=device)
        gathered = [torch.empty_like(source) for _ in range(world)]
        dist.all_gather(gathered, source)
        output = torch.stack(gathered)  # Display the returned list as a (world, 2) tensor.
        expected = torch.tensor([[r, r + 10] for r in range(world)], dtype=torch.float32, device=device)
        check('all_gather', source, output, expected)
        source = torch.arange(2 * world, dtype=torch.float32, device=device) + 100 * rank
        output = torch.empty(2, device=device)
        assert source.numel() == world * output.numel()
        dist.reduce_scatter_tensor(output, source, op=dist.ReduceOp.SUM)
        expected = torch.arange(2 * rank, 2 * rank + 2, device=device) * world + 100 * world * (world - 1) / 2
        check('reduce_scatter_tensor_SUM', source, output, expected)
        source = torch.tensor([100 * rank + 10 * d + i for d in range(world) for i in range(2)],
                              dtype=torch.float32, device=device)
        output = torch.empty_like(source)
        assert source.numel() % world == 0
        dist.all_to_all_single(output, source)  # Equal chunks, ordered by destination/source rank.
        expected = torch.tensor([100 * s + 10 * rank + i for s in range(world) for i in range(2)],
                                dtype=torch.float32, device=device)
        check('all_to_all_single', source, output, expected)
        failed = torch.tensor(int(not all(c['passed'] for c in checks)), device=device)
        dist.all_reduce(failed, op=dist.ReduceOp.MAX)
        print(json.dumps({'kind': 'semantics', 'rank': rank, 'checks': checks,
                          'all_ranks_passed': not bool(failed.item())}), flush=True)
        if failed.item():
            raise AssertionError('At least one rank failed its expected-value check.')
        if args.bench:
            for byte_count in sizes:
                buffer = torch.zeros(byte_count // 4, dtype=torch.float32, device=device)
                # Zeros stay bounded under SUM; no refill or allocation enters the timed window.
                for _ in range(args.warmup):
                    dist.all_reduce(buffer, op=dist.ReduceOp.SUM)
                sync()
                dist.barrier()
                sync()
                start = time.perf_counter()
                for _ in range(args.iters):
                    dist.all_reduce(buffer, op=dist.ReduceOp.SUM)
                sync()
                elapsed = torch.tensor(time.perf_counter() - start, dtype=torch.float64, device=device)
                dist.all_reduce(elapsed, op=dist.ReduceOp.MAX)  # Outside the measured window.
                seconds = elapsed.item() / args.iters
                if rank == 0:
                    print(json.dumps({'kind': 'all_reduce_benchmark', 'backend': args.backend,
                        'world_size': world, 'dtype': 'float32', 'payload_bytes': byte_count,
                        'payload_mib': byte_count / (1024 ** 2), 'warmup': args.warmup, 'iters': args.iters,
                        'mean_ms_using_slowest_rank': 1000 * seconds,
                        'payload_bytes_per_second': byte_count / seconds,
                        'rate_definition': 'one_rank_input_payload / collective_time; not bus or NIC bandwidth'}), flush=True)
                del buffer
    finally:
        dist.destroy_process_group()


if __name__ == '__main__':
    main()

collective_probe.py：完整独立脚本

C.11 PyTorch 执行观测：从前向依赖到一次参数更新

本实验复用 1.6.2 的两层 MLP，观察 1.6.3 所述的执行机制。把下面完整代码保存为 pytorch_execution_probe.py，使用前文 PyTorch 2.9.x 环境运行。先用 CPU 核对状态，再在有 CUDA 的机器上观察设备执行；无需 nanoGPT 数据文件。FP32 与 C.7 的 float64 数值核对采用相同输入和权重，本例理论预测为 $\begin{bmatrix}0.5\\3\end{bmatrix}$、shape 为 $(2,1)$，均方损失为 $0.625$、shape 为 $()$。

python pytorch_execution_probe.py --device cpu --trace trace-cpu.json

# 在有可用 CUDA 的机器上再运行
python pytorch_execution_probe.py --device cuda --trace trace-cuda.json

程序打印 JSON 状态记录并保存 Chrome trace 文件。按四个命名区段阅读：01_forward 产生预测、loss 与反向依赖；02_backward 形成参数梯度；03_optimizer_step 更新参数及 AdamW 状态；04_zero_grad 清理梯度。额外的 inspect_* 区段用于检查，包含读取或比较结果的开销，不属于训练计算本身。

观察项	本例应满足的判断
前向与反向后参数	forward_parameters_unchanged 与 backward_parameters_unchanged 都为 true
梯度 shape	四组参数梯度分别为 $(2,3)$、$(2,)$、$(1,2)$、$(1,)$，与对应参数相同
更新后的状态	本例四组参数均改变；AdamW 的 exp_avg、exp_avg_sq 与参数同 shape，各参数 step 为 1
清梯度后	grad_none_after_zero_grad=true，参数和 AdamW 状态继续保留
求导节点与执行记录	参数是 leaf，grad_fn 为 null；中间结果具有反向节点。CPU 运行可见 aten 算子，CUDA 运行才可能观察对应设备 kernel

本次实际验证。Python 3.12.14、PyTorch 2.9.1、CPU、eager、FP32 下已运行通过：预测为 $[0.5,3.0]^{\mathsf T}$，loss 为 $0.625$，上述状态检查全部通过。trace 中观察到两次 aten::linear 和两次 aten::addmm。AddmmBackward0 是本次自动求导节点名，aten::addmm 是框架算子名，两者都不能直接当作固定 CUDA kernel 名。CUDA 路径尚未实跑。

可以用兼容 Chrome trace 格式的本地查看器打开 trace，按四个区段寻找算子；有 CUDA 记录时再关联 CPU 提交与 GPU 执行。程序只执行一步且包含检查和 profiler 开销，没有预热，输出的微秒数仅供定位事件，不能用于比较 CPU/GPU 性能。若检查失败，先查看失败区段、参数名和 shape；不要以“看见 GPU 有活动”代替正确性判断。

"""Observe one tiny MLP update; CPU by default, optional --device cuda."""
import argparse
import json
import platform
from pathlib import Path
import torch
from torch import nn
from torch.profiler import ProfilerActivity, profile, record_function


def metadata(t):
    return {"shape": list(t.shape), "dtype": str(t.dtype), "device": str(t.device),
            "stride": list(t.stride()), "requires_grad": t.requires_grad,
            "is_leaf": t.is_leaf, "grad_fn": type(t.grad_fn).__name__ if t.grad_fn else None}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
    parser.add_argument("--trace", type=Path, help="Chrome trace JSON output path")
    args = parser.parse_args()
    if args.device == "cuda" and not torch.cuda.is_available():
        parser.error("CUDA is unavailable; no CUDA observation was made")
    torch.manual_seed(0)
    device = torch.device(args.device)
    model = nn.Sequential(nn.Linear(3, 2), nn.ReLU(), nn.Linear(2, 1)).to(device)
    with torch.no_grad():
        model[0].weight.copy_(torch.tensor([[1., 0., 1.], [0., 1., 1.]], device=device))
        model[0].bias.copy_(torch.tensor([-5., -6.], device=device))
        model[2].weight.copy_(torch.tensor([[1., -.5]], device=device)); model[2].bias.fill_(.5)
    x = torch.tensor([[1., 2., 3.], [4., 5., 6.]], device=device)
    target = torch.tensor([[1.], [4.]], device=device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=.001, foreach=False, fused=False)
    before = {name: p.detach().clone() for name, p in model.named_parameters()}
    report = {"environment": {"python": platform.python_version(), "torch": torch.__version__,
              "device": str(device), "cuda_runtime": torch.version.cuda},
              "parameters": {name: metadata(p) for name, p in model.named_parameters()},
              "inputs": {"x": metadata(x), "target": metadata(target)}}
    activities = [ProfilerActivity.CPU]
    if device.type == "cuda":
        activities.append(ProfilerActivity.CUDA)
        report["environment"]["gpu"] = torch.cuda.get_device_name()
    with profile(activities=activities, record_shapes=True, profile_memory=True) as prof:
        with record_function("01_forward"):
            z = model[0](x); h = model[1](z); pred = model[2](h)
            loss = (pred-target).square().mean()
        report["autograd"] = {name: metadata(t) for name, t in (("z", z), ("h", h), ("pred", pred), ("loss", loss))}
        with record_function("inspect_forward_no_update"):
            report["forward_parameters_unchanged"] = all(torch.equal(p, before[n]) for n, p in model.named_parameters())
            report["forward_values"] = {"pred": pred.detach().cpu().tolist(), "loss": loss.item()}
            torch.testing.assert_close(pred, pred.new_tensor([[.5], [3.]]))
            torch.testing.assert_close(loss, loss.new_tensor(.625))
        with record_function("02_backward"):
            loss.backward()
        report["gradient_shapes"] = {name: list(p.grad.shape) for name, p in model.named_parameters()}
        assert all(p.grad.shape == p.shape for p in model.parameters())
        with record_function("inspect_backward_no_update"):
            report["backward_parameters_unchanged"] = all(torch.equal(p, before[n]) for n, p in model.named_parameters())
        with record_function("03_optimizer_step"):
            optimizer.step()
        with record_function("04_zero_grad"):
            optimizer.zero_grad(set_to_none=True)
    report["parameters_changed_after_step"] = {name: not torch.equal(p, before[name]) for name, p in model.named_parameters()}
    report["adam_state"] = {name: {key: metadata(value) for key, value in optimizer.state[p].items()}
                            for name, p in model.named_parameters()}
    report["adam_step"] = {name: optimizer.state[p]["step"].item() for name, p in model.named_parameters()}
    assert all(optimizer.state[p]["exp_avg"].shape == p.shape == optimizer.state[p]["exp_avg_sq"].shape for p in model.parameters())
    report["grad_none_after_zero_grad"] = all(p.grad is None for p in model.parameters())
    assert report["forward_parameters_unchanged"] and report["backward_parameters_unchanged"] and all(report["parameters_changed_after_step"].values())
    assert report["grad_none_after_zero_grad"] and all(step == 1 for step in report["adam_step"].values())
    trace = (args.trace or Path(__file__).with_name(f"trace-{device.type}.json")).resolve()
    trace.parent.mkdir(parents=True, exist_ok=True)
    prof.export_chrome_trace(str(trace))
    ops = [event for event in prof.key_averages() if event.key.startswith("aten::")]
    report["operators"] = [{"name": op.key, "calls": op.count, "self_cpu_us": op.self_cpu_time_total,
                            "self_device_us": op.self_device_time_total}
                           for op in sorted(ops, key=lambda op: op.self_cpu_time_total, reverse=True)[:8]]
    report["trace"] = str(trace)
    report["scope"] = "One observed step, no warmup: timings are not a benchmark. Inspection is a separate trace region. CUDA kernels appear only in a CUDA run."
    report["passed"] = True
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()

pytorch_execution_probe.py：完整独立脚本