PyTorch 与大模型训练基础

一、语言模型训练概览

大语言模型（Large Language Model，LLM）本质上是一个包含大量可训练参数的函数。文本通过分词器变成数字，数字经神经网络计算得到预测，模型再根据预测误差修改参数。

flowchart TD
    A[原始文本] --> B[Tokenizer：编码为 Token ID]
    B --> C[Embedding：把 ID 变成向量]
    C --> D[神经网络：计算预测结果]
    D --> E[Loss：衡量预测误差]
    E --> F[Backward：求参数梯度]
    F --> G[Optimizer：更新参数]
    G -.下一训练步.-> D

这个过程可以分成三个问题：

模型在学什么？调整参数，让预测尽可能接近训练目标。
模型用什么计算？用张量（Tensor）存数据，用矩阵乘法等算子处理数据。
训练为什么贵？既要做大量浮点运算，还要保存参数、梯度、优化器状态和中间激活值。

1.1 基本术语

机器学习（Machine Learning，ML）是从数据中学习规律，而不是由程序员逐条写规则。深度学习（Deep Learning，DL）是机器学习的一类方法，主要利用多层神经网络（Neural Network，NN）学习特征和预测结果。大语言模型属于深度学习模型。数据集（Dataset）是样本的集合；标签或目标（Label / Target）是训练时希望模型得到的正确结果。

术语	含义	例子
模型 Model	从输入计算输出的函数	`y = w * x`
参数 Parameter	训练时会更新的数值	权重 `w`、偏置 `b`
超参数 Hyperparameter	人为设置的训练配置	学习率、批大小
样本 Sample	一条训练数据	一句话及其预测目标
特征 Feature	模型处理的数值表示	Token 的 4 维向量
前向传播 Forward	用当前参数算出预测和损失	`pred = model(x)`
损失函数 Loss	用数值表示预测误差	平方误差、交叉熵
反向传播 Backward	计算损失对参数的梯度	`loss.backward()`
优化器 Optimizer	利用梯度更新参数	SGD、AdamW
推理 Inference	用训练好的模型做预测	语言模型生成下一 Token
Epoch	完整遍历一遍训练数据	10 万条数据全部用过一次
Training Step	通常指完成一次参数更新	一次 `optimizer.step()`

训练和推理最重要的区别是：训练通常需要求梯度、更新参数；推理通常只进行前向计算，不更新参数。

二、神经网络训练原理

2.1 单参数模型

假设我们希望模型学会“输入乘以 2”，使用：

\hat y=wx

$x$ ：输入。
$w$ ：模型唯一的可训练参数。
$\hat y$ ：模型预测值，帽子表示“估计”，与正确答案 $y$ 区分。

现在给出样本 $x=3$ ，正确答案 $y=6$ ，但初始 $w=1$ 。

第一步：前向传播。

\hat y=1\times3=3

模型本应预测 6，却只预测了 3。

第二步：计算损失。先用平方误差（Squared Error）：

L=(\hat y-y)^2=(3-6)^2=9

把 $\hat y=wx$ 代入：

L(w)=(3w-6)^2

损失 $L$ 现在是关于参数 $w$ 的函数。训练的目标就是寻找让 $L$ 更小的 $w$ 。

第三步：求梯度。梯度（Gradient）描述参数发生微小变化时，损失如何变化。此处只有一个参数，因此梯度就是导数。

用链式法则：

\begin{aligned} \frac{\partial L}{\partial w} &=\frac{\partial(3w-6)^2}{\partial(3w-6)} \cdot\frac{\partial(3w-6)}{\partial w}\\ &=2(3w-6)\times3 \end{aligned}

在 $w=1$ 时：

\frac{\partial L}{\partial w}=2(3-6)\times3=-18

这里 $-18$ 表示：在当前参数附近，增大 $w$ 会让损失下降。它不是“损失本身是负数”；损失仍然是 9。

第四步：根据梯度更新参数。使用梯度下降（Gradient Descent）：

\boxed{w_{\mathrm{new}}=w-\eta\frac{\partial L}{\partial w}}

其中 $\eta$ （eta）是学习率（Learning Rate），决定每次沿梯度方向走多远。设 $\eta=0.05$ ：

\begin{aligned} w_{\mathrm{new}}&=1-0.05\times(-18)=1.9\\ \hat y_{\mathrm{new}}&=1.9\times3=5.7\\ L_{\mathrm{new}}&=(5.7-6)^2=0.09 \end{aligned}

一次更新就让损失从 9 下降到 0.09。不断重复“预测 → 计算损失 → 求梯度 → 更新参数”，就是模型训练的基本过程。

2.2 线性层与激活函数

真实神经网络会使用大量线性层（Linear Layer）和激活函数（Activation Function）。

用行向量表示一个样本时，线性层可写成：

\mathbf y=\mathbf xW+\mathbf b

$\mathbf x$ ：输入特征向量。
$W$ ：权重矩阵，决定怎样组合输入特征。
$\mathbf b$ ：偏置（Bias），相当于额外的可学习平移量。
$\mathbf y$ ：线性层输出的新特征。

只堆叠线性层不够，因为连续的线性变换仍可合并成一次线性变换。例如：

(xW_1)W_2=x(W_1W_2)

因此通常加入非线性激活函数。最简单的是 ReLU（Rectified Linear Unit，修正线性单元）：

\operatorname{ReLU}(x)=\max(0,x)

即负数变成 0，正数保持不变。例如 [-2, 0, 3] 经过 ReLU 后是 [0, 0, 3]。现代大模型还常用 GELU、SiLU、SwiGLU 等机制。

2.3 交叉熵损失

平方误差适用于“预测具体数字”的回归（Regression）任务。下一 Token 预测则属于“在词表中选择正确 Token”的分类（Classification）任务。

如果词表中有 $V$ 个 Token，模型会对每个候选 Token 给出一个分数 $z_i$ （Logit，未归一化分数）。Softmax 把分数转成概率：

p_i=\frac{e^{z_i}}{\sum_{j=1}^{V}e^{z_j}}

$V$ ：词表大小。
$z_i$ ：第 $i$ 个 Token 的分数。
$p_i$ ：第 $i$ 个 Token 的预测概率，所有 $p_i$ 相加为 1。

假设正确 Token 的编号是 $c$ ，交叉熵损失为：

\boxed{L=-\ln p_c}

交叉熵可以看作对模型预测正确的一个惩罚正确 Token 概率越高，损失越低。例如预测正确 Token 的概率从 $0.1$ 增至 $0.9$ ，损失就会明显下降。

三、分词与词嵌入

3.1 Token 与 Tokenizer

计算机不能把字符串直接送进矩阵乘法。首先需要分词器（Tokenizer），把字符串切分成 Token（词元），再编码成整数 Token ID。

Token 不一定等于一个汉字或一个英文单词，也可能是子词、标点或其他文本单位。例如：

文本            → Tokenizer → Token ID
"hello world"   →（若干词元）→ [整数, 整数, ...]

上面只说明处理过程；不同分词器的分词结果及整数编号可能不同。

BPE（Byte Pair Encoding，字节对编码）是常见的子词切分思路：从较小的单位开始，反复合并训练语料中频繁共现的相邻单位。Byte-level BPE 使用字节作为基本单位，与按字符举例的简化版本不完全相同。

需要记住：Token ID 只是查表编号，不是有数学大小关系的语义数值。例如 ID 为 10 的 Token 不代表在含义上比 ID 为 2 的 Token“大五倍”。

3.2 Embedding 词嵌入

Embedding（嵌入层）可以理解为一张可训练的查找表。假设词表有 100 个 Token，每个 Token 用 4 个数表示，那么表的形状就是 $(100,4)$ 。

import torch
from torch import nn

# B=2：两条序列；T=3：每条序列三个 Token。
# Token ID 必须是整数类型；long 即 int64。
ids = torch.tensor([
    [1, 5, 2],
    [7, 1, 0]
], dtype=torch.long)

# 词表大小 V=100，隐藏维度 H=4。
embedding = nn.Embedding(num_embeddings=100, embedding_dim=4)

# 对每个 ID 查表，得到相应的 4 维浮点向量。
x = embedding(ids)

print(ids.shape)  # torch.Size([2, 3])
print(x.shape)    # torch.Size([2, 3, 4])

形状变化是：

\boxed{(B,T)\;\xrightarrow{\text{Embedding}}\;(B,T,H)}

$B$ （Batch Size）：每次处理多少条样本。
$T$ （Sequence Length）：每条样本有多少个 Token。
$H$ （Hidden Size）：每个 Token 用多少个数描述其特征。

这里 (2, 3, 4) 读成“2 条序列，每条 3 个 Token，每个 Token 有 4 维特征”，而不是“24 条样本”。这三个字母是阅读语言模型代码最常见的维度记号。

四、张量基础

4.1 张量形状与维度

Tensor（张量）可以先理解为多维数组。

数据	示例	Shape（各轴长度）	Rank / ndim（轴数）
标量 Scalar	`5`	`()`	0
向量 Vector	`[1,2,3]`	`(3,)`	1
矩阵 Matrix	2 行 3 列	`(2,3)`	2
三维张量	2 个 3×4 矩阵	`(2,3,4)`	3

Shape 告诉你“每个轴有多长”，Rank 告诉你“一共有几个轴”。例如 (2,3,4) 的 Rank 是 3，元素总数是 $2\times3\times4=24$ 。

import torch

x = torch.randn(2, 3, 4)  # 创建 2×3×4 的随机浮点张量
print(x.shape)           # torch.Size([2, 3, 4])
print(x.ndim)            # 3：一共有三个轴
print(x.numel())         # 24：元素数量
print(x.dtype)           # torch.float32：每个元素的数据类型
print(x.element_size())  # 4：每个元素占 4 字节
print(x.device)          # cpu（默认情况下）

4.2 dtype、device 与 stride

dtype（Data Type）：每个元素用什么编码，例如 float32、bfloat16、int64。
device（设备）：张量位于 CPU 还是 GPU，例如 cpu、cuda:0。
stride（步长）：某个轴的索引加 1，底层存储地址要跨过多少个元素。
contiguous（连续）：张量是否符合默认的行优先连续内存布局；不连续张量也可以被正常计算。

看一个具体例子：

import torch

m = torch.arange(6).reshape(2, 3)
# m = [[0, 1, 2],
#      [3, 4, 5]]

print(m.stride())          # (3, 1)
print(m.T.stride())        # (1, 3)
print(m.T.is_contiguous()) # False

m.stride()=(3,1) 的含义：往下一行，要跨过 3 个元素；往右一列，只跨过 1 个元素。

m.T 是转置（Transpose）：把行与列交换。通常它只改变张量的访问视图，不会立即重新复制底层数据，因此步长会变成 (1,3)。

容易混淆：transpose 是交换轴；reshape 是改变形状，通常不等价于转置。reshape 是否复制数据取决于存储布局。

4.3 张量内存计算

设张量的形状为 $(s_1,s_2,\ldots,s_k)$ ，每个元素占 $q$ 字节。则元素数据本身约占：

\boxed{M=\left(\prod_{i=1}^{k}s_i\right)q}

$s_i$ ：第 $i$ 个轴的长度。
$\prod$ ：连乘符号，即把所有轴的长度相乘。
$q$ ：单个元素的字节数（Bytes）。

例如张量形状 (8, 1024, 4096)，类型为 BF16（每元素 2 字节）：

\begin{aligned} M&=8\times1024\times4096\times2\\ &=67{,}108{,}864\ \text{Bytes}\\ &=64\ \text{MiB} \end{aligned}

注意 MB 和 MiB 不相同：

1\ \text{MB}=10^6\ \text{Bytes},\qquad 1\ \text{MiB}=2^{20}\ \text{Bytes}

这个公式只统计张量元素存储，不是完整训练显存：训练还需要参数、梯度、优化器状态、激活值和临时缓冲区。多个视图若共享同一存储，也不能把底层内存重复加总。

五、矩阵计算

5.1 线性层的矩阵乘法

设一个 Token 的隐藏特征有 $H=4$ 维，希望把它变成 $K=5$ 维。用权重矩阵 $W\in\mathbb R^{4\times5}$ ：

\mathbf y=\mathbf xW,\qquad (1,4)\times(4,5)\rightarrow(1,5)

如果一次处理 $B=2$ 条文本，每条 $T=3$ 个 Token，则：

\boxed{(B,T,H)\times(H,K)\rightarrow(B,T,K)}

也就是：

(2,3,4)\times(4,5)\rightarrow(2,3,5)

可以理解为：保留前面的样本数和 Token 数，只把每个 Token 的特征从 4 维映射到 5 维。

import torch

x = torch.randn(2, 3, 4)  # B=2，T=3，H=4
w = torch.randn(4, 5)     # H=4，K=5

y = torch.matmul(x, w)   # 沿最后一维执行矩阵乘法
print(y.shape)            # torch.Size([2, 3, 5])

这里手写权重 w 的形状是 (H,K)。如果换成 nn.Linear(4,5)，PyTorch 内部存放权重的形状却是 (5,4)，因为其计算约定为 $Y=XW^\top+b$ 。两种写法在数学上可对应，不能把内部存储形状与数学公式直接混为一谈。

5.2 einsum 张量求和

einsum（Einstein Summation，爱因斯坦求和）把矩阵计算写成带字母的规则。

import torch

# 创建输入 (B,T,H) 与权重 (H,K)。
x = torch.randn(2, 3, 4)
w = torch.randn(4, 5)
y = torch.einsum("bth,hk->btk", x, w)
print(y.shape)  # (2, 3, 5)

bth,hk->btk 可以逐个字母读：

b：批次维，在输入、输出都保留。
t：Token 维，在输入、输出都保留。
h：输入特征维，出现在输入中却不出现在输出中，因此沿这个维度求和。
k：输出特征维，在权重和输出中出现。

对应数学公式：

\boxed{Y_{b,t,k}=\sum_{h=1}^{H}X_{b,t,h}W_{h,k}}

这里 $\sum$ 表示“把很多项加起来”。例如两个向量点积：

[1,2,3]\cdot[4,5,6]=1\times4+2\times5+3\times6=32

einsum 的省略号 ... 表示一组没有逐一命名的轴，例如：

import torch

x = torch.randn(2, 3, 4)
w = torch.randn(4, 5)
z = torch.einsum("...h,hk->...k", x, w)
print(z.shape)  # (2,3,5)

这里 ... 对应 b,t；它不是永远指 Batch 或 Token，而是取决于输入张量的实际轴数。

5.3 einops 维度变换

einops（Einstein Operations）是一套处理张量维度的工具，特别常用的两个操作是：

rearrange：改变轴的排列、合并或拆分；通常不会对数值求和。
reduce：沿指定轴求和、求均值等，减少数据维度。

# 首次使用需安装：pip install einops
import torch
from einops import rearrange, reduce

x = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)

# 把 b=2 和 t=3 合并成一个长度为 6 的轴。
flat = rearrange(x, "b t h -> (b t) h")
print(flat.shape)  # (6, 4)

# 对每条样本的三个 Token，在 t 维取平均。
avg = reduce(x, "b t h -> b h", "mean")
print(avg.shape)   # (2, 4)
print(avg[0])      # tensor([4., 5., 6., 7.])

第一行新形状 (6,4) 仍有 24 个元素；第二行 (2,4) 只有 8 个元素，因为每条样本的 3 个 Token 被平均汇总。

选择方法：做矩阵乘法用 matmul；想清楚表达“按哪个轴乘、按哪个轴求和”，用 einsum；想整理张量形状，用 einops。

六、浮点精度

6.1 浮点数表示

浮点数（Floating-point Number）类似计算机里的科学计数法。对于正常规格化的有限二进制浮点数：

\boxed{x=(-1)^s\times(1.f)_2\times 2^{e-\mathrm{bias}}}

$s$ （Sign）：符号位，决定正负。
$e$ （Exponent）：指数编码，主要决定数值的动态范围，即能表示多大或多小的数。
$f$ （Fraction / Significand）：尾数字段，主要决定在相近数值大小下有多精细。
$\mathrm{bias}$ ：指数偏置，把有符号的指数映射到无符号编码。

上式说明正常规格化数的原理；零、非正规数、无穷大与 NaN 有特殊编码。

6.2 FP32、FP16 与 BF16

数据格式	总位数	符号位	指数位	尾数字段位	每元素大小
FP32	32	1	8	23	4 Bytes
FP16	16	1	5	10	2 Bytes
BF16	16	1	8	7	2 Bytes

FP16 和 BF16 都占 2 字节，但侧重点不同：

FP16 尾数字段更长，表示相近大小的数时通常更精细；但指数位少，最大有限值是 65,504。
BF16 保留 8 位指数，动态范围接近 FP32，最大有限值约为 $3.39\times10^{38}$ ；但尾数更短，同一数量级的精度不如 FP16。
FP32 精度更高，但存储与数据传输成本更大。

比如一个含 1 亿个元素的张量：FP32 元素数据约 400 MB；BF16 或 FP16 约 200 MB。

Overflow（上溢）是结果大到超过格式可表示的范围；Underflow（下溢）是值过小，以致精度丢失甚至变成 0。BF16 更大的指数范围有利于避免某些上溢问题，但不等于不存在舍入误差。

6.3 自动混合精度 AMP

Mixed Precision（混合精度）是不同运算或状态使用不同数据格式。例如大矩阵乘法可以使用 BF16，而某些累积、归一化、优化器状态保留 FP32。

AMP（Automatic Mixed Precision，自动混合精度）可由 PyTorch 的 autocast 实现：

import torch
from torch import nn

# 只有具备 CUDA、且硬件支持 BF16 时才执行。
if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
    model = nn.Linear(128, 64).cuda()      # 默认权重通常仍是 FP32
    x = torch.randn(8, 128, device="cuda")

    # 自动为合适的计算选择 BF16 等精度。
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        y = model(x)

    print(y.shape)  # (8, 64)：8 条样本，每条输出 64 维

必须区分：参数的存储精度（内存中用什么格式保存）与算子的计算精度（执行计算时采用什么格式）。autocast 不代表所有参数永久变成 BF16，也不意味着整个训练显存必定减半。

七、计算量 FLOPs

7.1 FLOP、FLOPs 与 FLOP/s

FLOP（Floating-point Operation）：一次浮点运算。
FLOPs：一项计算需要多少次浮点运算，表示工作量。
FLOP/s：每秒完成多少次浮点运算，表示计算吞吐； $1\ \text{TFLOP/s}=10^{12}\ \text{FLOP/s}$ 。

按常用计数方式，一次浮点乘法约算 1 FLOP、一次加法约算 1 FLOP；一次融合乘加 FMA（Fused Multiply-Add）通常按 2 FLOPs 计算。

如果任务计算量为 $F$ ，实际吞吐为 $P$ ，时间近似是：

\boxed{t\approx\frac{F}{P}}

7.2 矩阵乘法计算量

假设：

A:(M,K),\qquad B:(K,N),\qquad C=AB:(M,N)

输出有 $M\times N$ 个元素。每个元素是两个长度为 $K$ 的向量点积：

C_{ij}=\sum_{k=1}^{K}A_{ik}B_{kj}

一个点积要做 $K$ 次乘法和 $K-1$ 次加法，合计 $2K-1$ 次标量加乘运算。因此：

F=MN(2K-1)\approx\boxed{2MNK}

为什么可以用近似号？因为当 $K$ 较大， $2K-1$ 与 $2K$ 的差异很小。硬件性能估算通常采用这一近似计数约定。

例子： $(2,3)$ 乘 $(3,4)$ ，即 $M=2,K=3,N=4$ ：

F\approx2\times2\times4\times3=48\ \text{FLOPs}

若输入是 (B,T,H)、权重是 (H,K)，一共有 $B\times T$ 个 Token 各做一次映射，所以：

\boxed{F_{\mathrm{linear}}\approx2BTHK}

符号提醒：本节矩阵乘法公式 $2MNK$ 中的 $N$ 是矩阵的列数；下一节 $6ND$ 中的 $N$ 则是模型参数量，不要混淆。

7.3 反向传播计算量

把 Token 维和 Batch 维合并，记输入为 $X:(M,H)$ 、权重为 $W:(H,K)$ ：

Y=XW

前向传播做一次矩阵乘法：

F_{\mathrm{forward}}\approx2MHK

反向传播已知输出梯度：

G=\frac{\partial L}{\partial Y},\quad G:(M,K)

根据矩阵求导，需要获得两个主要梯度：

\frac{\partial L}{\partial X}=GW^\top, \qquad \frac{\partial L}{\partial W}=X^\top G

$\partial L/\partial X$ ：损失对输入的梯度，使梯度能继续传回更前面的网络层。
$\partial L/\partial W$ ：损失对权重的梯度，用来更新当前层参数。
$\top$ ：矩阵转置，即行列交换。

这两项各是一次主要矩阵乘法，因此：

\begin{aligned} F_{\mathrm{backward}}&\approx 2MHK+2MHK=4MHK\\ F_{\mathrm{train}}&\approx 2MHK+4MHK=\boxed{6MHK} \end{aligned}

结论：对于矩阵乘法主导、且需要完整梯度的线性层，前向与反向的主计算量大约是 $1:2$ 。这不是所有算子都严格成立的定律，也不保证反向运行时间恰好是前向两倍。

7.4 6ND 训练量估算

对以稠密矩阵乘法为主要计算的 Transformer，做一个粗略近似：

每处理一个 Token，前向计算约需要 $2N$ FLOPs。
反向计算约是前向的 2 倍，约需要 $4N$ FLOPs。
每个 Token 的一次前向和反向合计约 $6N$ FLOPs。
训练过程中共处理 $D$ 个 Token。

所以：

\boxed{F_{\mathrm{train}}\approx6ND}

$F_{\mathrm{train}}$ ：总训练工作量，单位 FLOPs。
$N$ ：模型参数量（此处针对稠密模型的粗略估算）。
$D$ ：整个训练过程中处理的 Token 总数，不是隐藏维度。

例如 $N=10^9$ （10 亿参数）、 $D=10^{10}$ （100 亿 Token）：

F_{\mathrm{train}}\approx6\times10^9\times10^{10} =6\times10^{19}\ \text{FLOPs}

适用范围：这是规划时的近似值，忽略了部分 Attention、归一化、Embedding 等计算。长序列 Attention、MoE、激活重计算等都会让实际 FLOPs 偏离此式；不能把它当成精确计数。

八、训练显存

8.1 显存组成

训练要保存的不止模型权重，还包括：

类型	作用	是否一定只在训练时需要
Parameters 参数	保存当前模型权重	训练、推理都需要
Gradients 梯度	保存损失对参数的导数	主要在训练中需要
Optimizer States 优化器状态	保存更新参数需要的历史统计量	训练中需要
Activations 激活值	保存反向求导需要的中间结果	训练期间通常更多
Temporary Buffers 临时缓冲	算子工作区、临时计算数据	训练和推理都可能需要

因此：

M_{\mathrm{total}} \approx M_{\mathrm{params}}+M_{\mathrm{grads}} +M_{\mathrm{optimizer}}+M_{\mathrm{activations}}+M_{\mathrm{temporary}}

这是显存预算的分项思路，实际实现中还会有框架、缓存、碎片等额外开销，不能把各项简单视为完全固定。

8.2 AdamW 显存开销

SGD（Stochastic Gradient Descent，随机梯度下降）的简单更新是：

w_{t+1}=w_t-\eta g_t

其中 $g_t$ 是本次梯度。

AdamW（Adaptive Moment Estimation with Decoupled Weight Decay）除了权重和梯度，还维护两类历史信息：

一阶矩 $m$ ：对历史梯度做指数移动平均，用于估计更新方向。
二阶矩 $v$ ：对历史梯度平方做指数移动平均，用于估计梯度尺度。

假设参数、梯度、一阶矩和二阶矩都用 FP32：

项目	每参数大小
参数 $w$	4 Bytes
梯度 $g$	4 Bytes
一阶矩 $m$	4 Bytes
二阶矩 $v$	4 Bytes
合计	16 Bytes

模型有 $N$ 个参数时：

\boxed{M_{\mathrm{states}}\approx16N\ \text{Bytes}}

例如 10 亿参数：

16\times10^9\ \text{Bytes}=16\ \text{GB}

注意：16 GB 还不含激活值和临时缓冲区。它也不是任何训练方式都适用的固定数字：混合精度实现可能使用 BF16 参数、FP32 优化器状态，甚至额外的 FP32 主权重（Master Weights）。

8.3 梯度累积

Batch Size（批大小）是一次计算送进模型的样本数量；Micro Batch（微批）则是把一个更大的有效 Batch 拆成若干小批。

在单张 GPU 上：

\boxed{B_{\mathrm{effective}}=B_{\mathrm{micro}}\times K}

$B_{\mathrm{micro}}$ ：每次前向/反向处理的样本数。
$K$ ：累计多少个微批后才更新参数。

例如每次处理 2 条、连续累积 4 次，最后只更新一次，那么有效 Batch 是 $2\times4=8$ 。

为什么可以累积？因为梯度对求和满足线性：

\nabla_\theta\left(\frac1K\sum_{i=1}^{K}L_i\right) =\frac1K\sum_{i=1}^{K}\nabla_\theta L_i

这里 $L_i$ 是第 $i$ 个等大微批的平均损失， $\nabla_\theta$ 表示损失对所有模型参数求梯度。将每次损失除以 $K$ 再反向，就能累积出整体平均损失的梯度。

import torch
from torch import nn

torch.manual_seed(0)                       # 固定随机种子，便于复现
model = nn.Linear(4, 1)                 # 输入 4 维，输出 1 维
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

xs = torch.randn(8, 4)                  # 共 8 条样本
targets = torch.randn(8, 1)             # 每条样本一个目标值

micro_batch_size = 2
accum_steps = 4                        # 8 / 2 = 4 个微批

optimizer.zero_grad()                  # 一个有效 Batch 开始前清梯度

for x, target in zip(xs.split(micro_batch_size),
                     targets.split(micro_batch_size)):
    pred = model(x)                    # 本次只输入 2 条样本
    loss = ((pred - target) ** 2).mean()  # 当前微批的平均损失
    (loss / accum_steps).backward()    # 梯度相加，暂时不更新权重

optimizer.step()                       # 4 个微批结束后只更新 1 次

理解代码的关键：循环内部的 backward() 不会自动清除已有梯度，而是累加到 .grad；循环外的 step() 才真正更新权重。

这一方法能减少单次微批的激活值显存，但不会按比例减少参数、参数梯度或优化器状态。若有 BatchNorm、不同大小微批、随机层或数值舍入，不一定与完整大 Batch 在数值上完全一致；不同微批大小还需要按样本数量加权损失。

8.4 激活检查点

通常，反向传播需要用到前向传播中的中间激活值。如果全部保存，显存压力较大。

Activation Checkpointing（激活检查点）只保存一部分必要信息，其他中间结果在反向传播时重新做一次部分前向计算。

\boxed{\text{少存一部分激活值}\quad\Longleftrightarrow\quad \text{增加一部分重计算}}

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

# 例子：一个线性层接 ReLU。
block = nn.Sequential(nn.Linear(64, 64), nn.ReLU())

# 8 条样本，每条 64 维；允许对输入求导。
x = torch.randn(8, 64, requires_grad=True)

# block 的部分中间结果不必像普通前向那样全部长期保留。
out = checkpoint(block, x, use_reentrant=False)

# 正常定义损失并反向传播；需要时重新计算 block 的部分前向。
loss = out.square().mean()
loss.backward()

两种优化方式不要混淆：

方法	主要节省什么	主要代价
Gradient Accumulation 梯度累积	每次微批的激活显存	更多小批次执行，可能增加调度开销
Activation Checkpointing 激活检查点	需要长期保存的部分激活值	反向传播中额外重计算

另外，这里的“激活检查点”不是保存模型文件以便断点续训的“训练检查点（Training Checkpoint）”。

九、GPU 性能与 Roofline

9.1 容量、带宽与算力

显存容量（Memory Capacity）：最多能容纳多少数据，回答“装不装得下”。
显存带宽（Memory Bandwidth）：每秒能搬运多少字节，回答“数据搬得快不快”。
计算吞吐（Compute Throughput）：每秒能完成多少浮点运算，回答“算得快不快”。

实际执行可能是：

Memory-bound（带宽受限）：计算单元等数据，主要瓶颈是数据搬运。
Compute-bound（计算受限）：数据供应得上，主要瓶颈是计算吞吐。

显存有足够空间，只说明能装下；不保证运行速度快。

9.2 算术强度

Arithmetic Intensity（算术强度）的定义为：

\boxed{I=\frac{F}{Q}}

$F$ ：执行的浮点运算量（FLOPs）。
$Q$ ：在所分析内存层级搬运的数据量（Bytes），例如 GPU 显存读写量。
$I$ ：每字节对应多少 FLOPs，单位 FLOPs/Byte。

例如：

搬运 1000 Bytes，只执行 1000 FLOPs，则 $I=1$ 。
搬运同样 1000 Bytes，执行 10000 FLOPs，则 $I=10$ 。

后者每次搬数据能完成更多计算，说明数据复用程度更高。但算术强度高也不代表整个程序一定更快，还受硬件、内核实现和并行程度影响。

9.3 Roofline 模型推导

Roofline（屋顶线模型）给出程序计算吞吐的理论上限。

设：

$P_{\max}$ ：GPU 理论计算峰值，单位 FLOP/s。
$BW$ ：GPU 在相应内存层级的带宽，单位 Bytes/s。
$F$ ：任务计算量，单位 FLOPs。
$Q$ ：任务数据搬运量，单位 Bytes。

只考虑计算单元时，任务至少需要：

t_{\mathrm{compute}}\ge\frac{F}{P_{\max}}

只考虑内存搬运时，任务至少需要：

t_{\mathrm{memory}}\ge\frac{Q}{BW}

二者都必须完成，因此实际时间满足：

t\ge\max\left(\frac{F}{P_{\max}},\frac{Q}{BW}\right)

实际吞吐是 $P=F/t$ 。对上式两边变形：

\begin{aligned} P&\le\min\left(P_{\max},\frac{F}{Q}BW\right)\\ &=\boxed{\min(P_{\max},I\times BW)} \end{aligned}

这就是 Roofline 的核心：程序吞吐不可能超过计算峰值，也不可能超过“数据搬运能力 × 每字节的计算量”。

两个上限相等时，得到临界算术强度（Ridge Point）：

\boxed{I^*=\frac{P_{\max}}{BW}}

举一个假设的 GPU：计算峰值 $100\ \text{TFLOP/s}$ 、带宽 $1\ \text{TB/s}$ ：

I^*=\frac{100\times10^{12}}{1\times10^{12}}=100\ \text{FLOPs/Byte}

$I=20$ ：带宽上限为 $20\ \text{TFLOP/s}$ ，低于计算峰值，属于带宽受限区域。
$I=150$ ：带宽对应 $150\ \text{TFLOP/s}$ ，但计算峰值只有 100，所以属于计算受限区域。

Roofline 是“最高可能达到多少”的模型，而不是实际速度预测器。Kernel（GPU 计算内核）的实现质量、同步、访问模式等可能使实际表现更低。

9.4 MFU 与 GPU 利用率

MFU（Model FLOPs Utilization，模型 FLOPs 利用率）：模型有效计算吞吐占相应 GPU 理论峰值吞吐的比例。

\boxed{\mathrm{MFU}\approx\frac{F_{\mathrm{model}}}{t\times G\times P_{\max}}}

$F_{\mathrm{model}}$ ：统计时间内的模型理论有效 FLOPs。
$t$ ：实际运行时间，单位秒。
$G$ ：GPU 数量。
$P_{\max}$ ：每张 GPU 在对应精度/计算模式下的理论峰值 FLOP/s。

假设一张卡的理论峰值为 $100\ \text{TFLOP/s}$ ，模型实际有效计算吞吐为 $40\ \text{TFLOP/s}$ ：

\mathrm{MFU}=\frac{40}{100}=40\%

GPU Utilization（GPU 使用率）通常是 GPU 有多少时间处于忙碌状态，不能直接代表计算吞吐已接近峰值。GPU 可能显示很忙，却一直在等待显存读写或执行低效 Kernel。

因此，估计大模型训练时间时，可把总工作量除以有效计算吞吐：

\boxed{t\approx\frac{6ND}{G\times P_{\max}\times\mathrm{MFU}}}

这是估算公式，成立前提是 $6ND$ 适用于模型、所选 MFU 有依据，并且没有其他重大开销被忽略。

十、稠密模型与稀疏 MoE

Dense Model（稠密模型）：处理每个 Token 时，通常会执行模型固定的主要稠密层计算路径。

Sparse MoE（Sparse Mixture of Experts，稀疏混合专家模型）：模型包含多个 Expert（专家子网络），由 Router（路由器）为每个 Token 选择一部分专家参与计算。专家通常是 FFN（Feed-Forward Network，前馈网络），不等于人为指定的“数学专家”“语言专家”。

稠密层：
Token → 同一组主要计算层 → 输出

稀疏 MoE 层：
Token → Router → 从多个 Expert 中挑选少数几个 → 合并输出

关键是区分两种参数量：

Total Parameters（总参数量）：模型所有参数加在一起。
Active Parameters（每 Token 激活参数量）：处理一个 Token 时实际参与计算的参数规模。

MoE 可能总参数量很大，但每个 Token 只激活一小部分参数。所以不能直接把 MoE 总参数量代入稠密模型 $6ND$ 公式，否则可能高估计算量。即便使用激活参数量估算，也需要考虑共享层、Router 和专家结构。

但总权重仍需存储或分布式管理。“每次少算一些”不等于“所有参数只占相同比例的显存”。此外，MoE 的专家选择稀疏与Weight Sparsity（权重稀疏，即矩阵内部有大量零权重）不是同一概念。

十一、训练资源估算实例

假设要训练一个稠密模型：参数量 $N=10^9$ ，训练 Token 总数 $D=10^{10}$ 。为了便于推算，先采用 FP32 AdamW，并假设一张 GPU 的相应计算峰值为 $100\ \text{TFLOP/s}$ ，模型训练 MFU 为 $40\%$ 。这些 GPU 数值只是例子，不代表某张真实显卡的实测值。

第一问：模型训练状态需要多少显存？

M_{\mathrm{states}}\approx16N=16\times10^9\ \text{Bytes}=16\ \text{GB}

这只包括参数、梯度和 AdamW 两个 FP32 状态；尚未加激活值及临时开销。

第二问：训练一共需要多少 FLOPs？

F\approx6ND=6\times10^{19}\ \text{FLOPs}

第三问：理论规划时间是多少？

\begin{aligned} t&\approx\frac{6\times10^{19}}{1\times(100\times10^{12})\times0.4}\\ &=1.5\times10^6\ \text{s}\\ &\approx17.36\ \text{天} \end{aligned}

这个时间是计算预算的粗估：若显存不足、MFU 达不到假设值，或者存在通信、I/O 和额外重计算，真实运行就会偏离；它不能直接保证训练能在这张卡上完成。

怎样思考优化？先判断显存够不够；不够时考虑缩小 Micro Batch、激活检查点或改变训练/模型方案。能放进显存却很慢时，再看计算量、数据搬运、Kernel 效率和 MFU，而不是只看 GPU 使用率。

十二、公式与概念速查

公式	含义	适用范围/提醒
$\hat y=f(x;\theta)$	模型由输入和参数计算预测	模型的抽象表示
$\theta_{\rm new}=\theta-\eta\nabla_\theta L$	梯度下降更新参数	学习率决定更新步幅
$M=(\prod_i s_i)q$	张量元素数据量	不包括训练其他状态
$(B,T)\to(B,T,H)$	Embedding 输出形状	ID 变成特征向量
$F\approx2MNK$	矩阵乘法 FLOPs	$M,K,N$ 是矩阵尺寸
$F\approx2BTHK$	Token 级线性映射 FLOPs	对 $(B,T,H)$ 的最后一维计算
$F_{\rm train}\approx6MHK$	线性层前后向计算量	主要矩阵乘法完整求导的近似
$F_{\rm train}\approx6ND$	稠密大模型总训练 FLOPs	$N$ 是参数量， $D$ 是 Token 总数
$M_{\rm states}\approx16N$ Bytes	FP32 AdamW 训练状态	不含激活；不是固定 BF16 公式
$B_{\rm eff}=B_{\rm micro}K$	有效 Batch	单 GPU、等大微批
$I=F/Q$	算术强度	运算量与数据搬运量之比
$P\le\min(P_{\max},I\cdot BW)$	Roofline 上限	不等于实际吞吐
$\mathrm{MFU}\approx F_{\rm model}/(tGP_{\max})$	模型有效 FLOPs 利用率	不等于 GPU 使用率