文章详情

互动

CS336（二）：用 PyTorch 理解 LLM 的计算与训练

257413分钟2026-09-0540江苏苏州0

文章摘要

本文讨论训练语言模型时落到 PyTorch 里的具体计算问题：Tensor 怎么存、矩阵乘法要多少算力、反向传播为什么更贵、训练时显存里除了参数还放了什么。从浮点格式与内存布局讲到 FLOP 估算和硬件利用率，再沿着初始化、数据加载、优化器演进和 Checkpoint 把一次完整训练串起来，为后续的分布式训练与模型结构讨论做铺垫。

训练一个语言模型，真正落到 PyTorch 里以后，很多问题都会变成很具体的计算问题：一个 Tensor 怎么存，矩阵乘法到底有多少 FLOP，反向传播为什么比前向更贵，训练时显存里除了参数还存了什么。

这篇就从这些底层问题开始，把一次完整训练过程拆开来看。

1. 内存计算

1.1 数据类型

训练中常见的浮点格式：

FP32：1 + 8 + 23，每个元素 4 Byte；
FP16：1 + 5 + 10，内存减半，但大型模型训练时表示范围可能不够；
BF16：1 + 8 + 7，保留 FP32 的指数范围，同时维持 FP16 的内存占用，代价是有效精度更低；
FP8：常见格式为 1 + 4 + 3 或 1 + 5 + 2。

实际训练通常使用混合精度：

长期保存、对数值稳定性要求高的状态保留 FP32；
前向传播等临时计算使用 BF16；
推理阶段通常可以比训练阶段更加激进地降低精度。

1.2 Tensor 的存储、Stride 与 View

PyTorch Tensor 并不是“多维数组本体”，它更像是对一块底层 Storage 的描述。底层内存连续存放元素，Tensor 再记录 Shape 和各维度的 Stride。例如：

PYTHON

assert x.stride(0) == 4
assert x.stride(1) == 1

含义是：

沿第 0 维移动一次，需要在 Storage 中跨过 4 个元素；
沿第 1 维移动一次，只需要跨过 1 个元素。

索引一个 Tensor 元素时，会根据各维度的 Stride 计算它在底层 Storage 中的偏移。因此，像取某一行或列、transpose 这类操作通常只需要修改 Tensor 的 Shape 和 Stride，并不会复制底层数据；只要新的 Shape 仍然能够用现有 Stride 描述，就可以直接创建 View。

如果当前内存布局无法直接 View，可以先调用 x.contiguous() 生成一份连续存储的数据，这一步会发生实际 Copy。reshape 则会优先尝试复用原来的 Storage，只有无法复用时才会复制数据。

1.3 Tensor 运算与 einops

高维 Tensor 的矩阵乘法只作用在最后两个维度。例如：

PYTHON

x = torch.ones(4, 8, 16, 32)
w = torch.ones(32, 2)
y = x @ w

结果 Shape 为：(4, 8, 16, 2)。

einsum

einops 可以把很多维度变换直接写在表达式里。

传统写法：

PYTHON

z = x @ y.transpose(-2, -1)

可以写成：

PYTHON

z = einsum(
    x, y,
    "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2"
)

也可以使用省略维度：

PYTHON

z = einsum(
    x, y,
    "... seq1 hidden, ... seq2 hidden -> ... seq1 seq2"
)

reduce

单维度聚合：

PYTHON

y = reduce(x, "... hidden -> ...", "sum")

rearrange

拆开一个维度：

PYTHON

x = rearrange(
    x,
    "... (heads hidden) -> ... heads hidden",
    heads=2
)

再重新展平：

PYTHON

x = rearrange(
    x,
    "... heads hidden -> ... (heads hidden)"
)

1.4 FLOP 与计算成本

FLOP 指一次浮点运算。

稀疏计算和稠密计算的 FLOP/s 不能直接按同一口径比较；在部分硬件规格中，稀疏峰值可以达到稠密峰值的约 2 倍。以 8 张 H100 连续运行一周为例，理论计算量可以达到约 $4.7\times10^{21}$ FLOP。

对于矩阵乘法：

[A,B]@[B,C]=[A,C]

每个输出元素需要 $B$ 次乘法和累加，因此：

\text{FLOPs}=2ABC

LLM 里矩阵乘法占据绝大部分计算，所以很多估算可以忽略其他小算子。如果 $A$ 对应 Token 数， $B$ 、 $C$ 主要由权重维度决定，那么一次前向传播可以粗略估算为：

\text{Forward FLOPs} \approx 2\times \text{Token 数}\times \text{参数量}

1.5 MFU

MFU（Model FLOPs Utilization）衡量模型真正得到的计算吞吐相对于硬件理论峰值的比例：

\text{MFU} = \frac{\text{实际 FLOP/s}}{\text{理论 FLOP/s}}

通信、访存、Kernel Launch 等都会让实际值低于峰值。训练 MFU 能超过 50% 通常是不错的利用率。

1.6 反向传播为什么更贵

考虑两层线性模型：

flowchart LR
    X[x] -->|W1| H1[h1]
    H1 -->|W2| H2[h2]
    H2 --> L[loss]

设 $x:[B,D]$ 、 $W_1:[D,D]$ 、 $W_2:[D,K]$ ，因此 $h_1:[B,D]$ 、 $h_2:[B,K]$ 。

对于第二层：

h_2=h_1W_2

反向传播时，已知上游梯度 $h_2.grad=\partial loss/\partial h_2$ 。权重梯度为：

W_2.grad=h_1^T h_2.grad

展开到单个元素：

W_2.grad[j,k] = \sum_i h_1[i,j]\cdot h_2.grad[i,k]

因此计算 $W_2.grad$ 需要一次 $[D,B]@[B,K]$ 的矩阵乘，FLOP 约为 $2BDK$ 。

同时还需要继续把梯度传回上一层：

h_1.grad=h_2.grad W_2^T

这同样是一轮 $2BDK$ 的矩阵乘。

第一层也是一样的结构：

W_1.grad=x^T h_1.grad

x.grad=h_1.grad W_1^T

如果 $W_1:[D,D]$ ，这两次矩阵乘分别需要约 $2BD^2$ FLOP。

所以对于一个线性层，前向传播只需要一次矩阵乘，而反向传播通常要计算两次：一次得到权重梯度，一次把梯度继续传给上一层。

因此可以近似记成：反向传播的矩阵计算量约为前向传播的 2 倍，完整训练一步约为纯前向的 3 倍。

2. 模型

2.1 参数初始化

设 $x:[A,B]$ 、 $W:[B,C]$ ，线性层输出为 $h=xW$ 。

假设 $x_{mean}$ 、 $w_{mean}$ 的标准差均为 $\sqrt{5}$ 。对 $h$ 中任意一个元素来说，本质上都是 $B$ 个浮点数相乘后再求和，因此输出的方差会随着输入维度 $B$ 增长：

\operatorname{Var}(h)\propto B\times x_{mean}\times w_{mean}

对应的输出尺度可以近似写成 $h_{mean}=\sqrt{5B}$ 。层数不断增加以后，这种尺度变化会逐层累积，很容易导致参数和激活值爆炸。

Xavier 初始化

Xavier 初始化的思路，就是让权重的初始尺度随着输入维度增大而减小。可以先对 $W$ 做正态分布采样，再按 $\sqrt{B}$ 缩放：

W \leftarrow \frac{W}{\sqrt{B}}

这样可以抵消输入维度增加带来的方差放大。

初始化时还可以进一步做区间截断，避免正态分布采样到极端大的权重。

2.2 设计结构

完成参数初始化以后，就可以按照前面定义好的模型结构组织各层计算。这里不再展开具体网络结构，后面的基础架构部分会单独讨论 Norm、Attention、FFN、位置编码等组件。

2.3 训练

训练前通常需要固定随机种子torch.manual_seed(seed)，保证参数初始化、数据采样等随机过程尽可能可复现。

如果涉及 CUDA，还需要同时处理对应设备上的随机状态。

2.4 数据加载

大规模训练数据通常无法一次性完整加载到内存中，可以使用 np.memmap() 把磁盘文件映射成数组，并在访问时按需读取：

PYTHON

data = np.memmap(
    path,
    dtype=np.uint16,
    mode="r"
)

这样能够直接对大文件进行切片，而不需要先把完整数据复制到内存。

2.5 优化器

几种常见优化器之间的演进关系：

flowchart LR
    BGD[BGD] --> SGD[SGD]
    SGD --> Adagrad[Adagrad]
    SGD --> Momentum[Momentum]

    Adagrad --> RMSGrad[RMSGrad]
    Momentum --> Adam[Adam]
    RMSGrad --> Adam

    Momentum --> Muon[Muon]

BGD

最基础的 Batch Gradient Descent 使用完整数据计算梯度，再更新参数：

w=w-lr\cdot w.grad

SGD

SGD 在 BGD 的基础上，不再每次使用完整训练集，而是随机取一个 Mini-Batch 估计梯度。

这样单步计算成本更低，但梯度也会引入更明显的随机噪声。

Momentum

Momentum 在 SGD 的基础上加入历史梯度信息：

v=\beta v+(1-\beta)grad

再使用累计后的方向更新参数：

w=w-lr\cdot v

$\beta$ 一般取 $0.9$ 左右。

这样可以让梯度更新具有一定惯性，减少单个 Mini-Batch 对更新方向的剧烈影响。

Adagrad

Adagrad 会累计历史梯度平方，并使用：

\frac{\text{当前 grad}} {\sqrt{\text{历史 grad 平方和}}}

调整每个参数的更新幅度。

历史梯度较大的参数，后续有效学习率会逐渐降低；但因为平方和会持续累积，训练后期学习率可能下降到接近 0。

RMSGrad

RMSGrad 在 Adagrad 的基础上，把历史梯度平方和的简单累加改成动量形式，只保留近期统计，从而避免分母持续增大。

Adam

Adam 可以看成 RMSGrad 和 Momentum 的组合：

Momentum 负责保留历史梯度方向的惯性；
梯度平方统计负责缩放不同参数的更新幅度，避免同一个梯度矩阵中不同参数尺度差异过大。

Muon

Muon 主要保留 Momentum，并进一步对动量矩阵做正交化，使不同方向上的更新尺度更加均衡。

2.6 参数量统计

假设模型包含 num_layers 个 $D\times D$ 全连接层，再加一个大小为 $D$ 的投影层。

状态	规模
Weight	$D^2\times \text{num\_layers} + D$
Activation	$B\times D\times \text{num\_layers}$
Gradient	与 Weight 同规模
Optimizer State	与 Weight 同量级；Adam 需要两份状态

其中 Activation 是前向传播过程中产生的中间结果。

Gradient 的规模与可训练参数量一致；中间激活对应的梯度只在反向传播过程中临时使用，完成对应计算后就可以释放。

Adam 需要额外保存一阶动量和二阶统计量，因此 Optimizer State 的内存占用会明显高于只保存权重本身。这也是后面 ZeRO / FSDP 会重点切分 Parameter、Gradient 和 Optimizer State 的原因。

2.7 Checkpoint

一个能够真正恢复训练状态的 Checkpoint 至少需要：

模型参数；
优化器参数。

只保存模型权重，只能恢复模型本身；Momentum、Adam 的历史统计都会丢失，训练轨迹也就不是原来的了。

xxxxavierrr

兴趣是最好的老师。 —— 爱因斯坦

本文是原创文章，采用 CC BY-NC-ND 4.0 协议，完整转载请注明来自 xxxxavierrr

CS3363 LLM3

你可能还想看

CS336（三）：现代 LLM 的基础架构

CS336（一）：Tokenizer 与 BPE

隐私政策

评论内容

0/500

昵称

邮箱

网址

滚动到此处加载评论...

雲自說

来自雲自說最新设计与科技的文章

查看全部

文章详情

互动

CS336（二）：用 PyTorch 理解 LLM 的计算与训练

257413分钟2026-09-0540江苏苏州0

文章摘要

这篇就从这些底层问题开始，把一次完整训练过程拆开来看。

1. 内存计算

1.1 数据类型

训练中常见的浮点格式：

FP32：1 + 8 + 23，每个元素 4 Byte；
FP16：1 + 5 + 10，内存减半，但大型模型训练时表示范围可能不够；
BF16：1 + 8 + 7，保留 FP32 的指数范围，同时维持 FP16 的内存占用，代价是有效精度更低；
FP8：常见格式为 1 + 4 + 3 或 1 + 5 + 2。

实际训练通常使用混合精度：

长期保存、对数值稳定性要求高的状态保留 FP32；
前向传播等临时计算使用 BF16；
推理阶段通常可以比训练阶段更加激进地降低精度。

1.2 Tensor 的存储、Stride 与 View

PyTorch Tensor 并不是“多维数组本体”，它更像是对一块底层 Storage 的描述。底层内存连续存放元素，Tensor 再记录 Shape 和各维度的 Stride。例如：

PYTHON

assert x.stride(0) == 4
assert x.stride(1) == 1

含义是：

沿第 0 维移动一次，需要在 Storage 中跨过 4 个元素；
沿第 1 维移动一次，只需要跨过 1 个元素。

1.3 Tensor 运算与 einops

高维 Tensor 的矩阵乘法只作用在最后两个维度。例如：

PYTHON

x = torch.ones(4, 8, 16, 32)
w = torch.ones(32, 2)
y = x @ w

结果 Shape 为：(4, 8, 16, 2)。

einsum

einops 可以把很多维度变换直接写在表达式里。

传统写法：

PYTHON

z = x @ y.transpose(-2, -1)

可以写成：

PYTHON

z = einsum(
    x, y,
    "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2"
)

也可以使用省略维度：

PYTHON

z = einsum(
    x, y,
    "... seq1 hidden, ... seq2 hidden -> ... seq1 seq2"
)

reduce

单维度聚合：

PYTHON

y = reduce(x, "... hidden -> ...", "sum")

rearrange

拆开一个维度：

PYTHON

x = rearrange(
    x,
    "... (heads hidden) -> ... heads hidden",
    heads=2
)

再重新展平：

PYTHON

x = rearrange(
    x,
    "... heads hidden -> ... (heads hidden)"
)

1.4 FLOP 与计算成本

FLOP 指一次浮点运算。

对于矩阵乘法：

[A,B]@[B,C]=[A,C]

每个输出元素需要 $B$ 次乘法和累加，因此：

\text{FLOPs}=2ABC

\text{Forward FLOPs} \approx 2\times \text{Token 数}\times \text{参数量}

1.5 MFU

MFU（Model FLOPs Utilization）衡量模型真正得到的计算吞吐相对于硬件理论峰值的比例：

\text{MFU} = \frac{\text{实际 FLOP/s}}{\text{理论 FLOP/s}}

通信、访存、Kernel Launch 等都会让实际值低于峰值。训练 MFU 能超过 50% 通常是不错的利用率。

1.6 反向传播为什么更贵

考虑两层线性模型：

flowchart LR
    X[x] -->|W1| H1[h1]
    H1 -->|W2| H2[h2]
    H2 --> L[loss]

设 $x:[B,D]$ 、 $W_1:[D,D]$ 、 $W_2:[D,K]$ ，因此 $h_1:[B,D]$ 、 $h_2:[B,K]$ 。

对于第二层：

h_2=h_1W_2

反向传播时，已知上游梯度 $h_2.grad=\partial loss/\partial h_2$ 。权重梯度为：

W_2.grad=h_1^T h_2.grad

展开到单个元素：

W_2.grad[j,k] = \sum_i h_1[i,j]\cdot h_2.grad[i,k]

因此计算 $W_2.grad$ 需要一次 $[D,B]@[B,K]$ 的矩阵乘，FLOP 约为 $2BDK$ 。

同时还需要继续把梯度传回上一层：

h_1.grad=h_2.grad W_2^T

这同样是一轮 $2BDK$ 的矩阵乘。

第一层也是一样的结构：

W_1.grad=x^T h_1.grad

x.grad=h_1.grad W_1^T

如果 $W_1:[D,D]$ ，这两次矩阵乘分别需要约 $2BD^2$ FLOP。

所以对于一个线性层，前向传播只需要一次矩阵乘，而反向传播通常要计算两次：一次得到权重梯度，一次把梯度继续传给上一层。

因此可以近似记成：反向传播的矩阵计算量约为前向传播的 2 倍，完整训练一步约为纯前向的 3 倍。

2. 模型

2.1 参数初始化

设 $x:[A,B]$ 、 $W:[B,C]$ ，线性层输出为 $h=xW$ 。

\operatorname{Var}(h)\propto B\times x_{mean}\times w_{mean}

对应的输出尺度可以近似写成 $h_{mean}=\sqrt{5B}$ 。层数不断增加以后，这种尺度变化会逐层累积，很容易导致参数和激活值爆炸。

Xavier 初始化

Xavier 初始化的思路，就是让权重的初始尺度随着输入维度增大而减小。可以先对 $W$ 做正态分布采样，再按 $\sqrt{B}$ 缩放：

W \leftarrow \frac{W}{\sqrt{B}}

这样可以抵消输入维度增加带来的方差放大。

初始化时还可以进一步做区间截断，避免正态分布采样到极端大的权重。

2.2 设计结构

2.3 训练

训练前通常需要固定随机种子torch.manual_seed(seed)，保证参数初始化、数据采样等随机过程尽可能可复现。

如果涉及 CUDA，还需要同时处理对应设备上的随机状态。

2.4 数据加载

大规模训练数据通常无法一次性完整加载到内存中，可以使用 np.memmap() 把磁盘文件映射成数组，并在访问时按需读取：

PYTHON

data = np.memmap(
    path,
    dtype=np.uint16,
    mode="r"
)

这样能够直接对大文件进行切片，而不需要先把完整数据复制到内存。

2.5 优化器

几种常见优化器之间的演进关系：

flowchart LR
    BGD[BGD] --> SGD[SGD]
    SGD --> Adagrad[Adagrad]
    SGD --> Momentum[Momentum]

    Adagrad --> RMSGrad[RMSGrad]
    Momentum --> Adam[Adam]
    RMSGrad --> Adam

    Momentum --> Muon[Muon]

BGD

最基础的 Batch Gradient Descent 使用完整数据计算梯度，再更新参数：

w=w-lr\cdot w.grad

SGD

SGD 在 BGD 的基础上，不再每次使用完整训练集，而是随机取一个 Mini-Batch 估计梯度。

这样单步计算成本更低，但梯度也会引入更明显的随机噪声。

Momentum

Momentum 在 SGD 的基础上加入历史梯度信息：

v=\beta v+(1-\beta)grad

再使用累计后的方向更新参数：

w=w-lr\cdot v

$\beta$ 一般取 $0.9$ 左右。

这样可以让梯度更新具有一定惯性，减少单个 Mini-Batch 对更新方向的剧烈影响。

Adagrad

Adagrad 会累计历史梯度平方，并使用：

\frac{\text{当前 grad}} {\sqrt{\text{历史 grad 平方和}}}

调整每个参数的更新幅度。

历史梯度较大的参数，后续有效学习率会逐渐降低；但因为平方和会持续累积，训练后期学习率可能下降到接近 0。

RMSGrad

RMSGrad 在 Adagrad 的基础上，把历史梯度平方和的简单累加改成动量形式，只保留近期统计，从而避免分母持续增大。

Adam

Adam 可以看成 RMSGrad 和 Momentum 的组合：

Momentum 负责保留历史梯度方向的惯性；
梯度平方统计负责缩放不同参数的更新幅度，避免同一个梯度矩阵中不同参数尺度差异过大。

Muon

Muon 主要保留 Momentum，并进一步对动量矩阵做正交化，使不同方向上的更新尺度更加均衡。

2.6 参数量统计

假设模型包含 num_layers 个 $D\times D$ 全连接层，再加一个大小为 $D$ 的投影层。

状态	规模
Weight	$D^2\times \text{num\_layers} + D$
Activation	$B\times D\times \text{num\_layers}$
Gradient	与 Weight 同规模
Optimizer State	与 Weight 同量级；Adam 需要两份状态

其中 Activation 是前向传播过程中产生的中间结果。

Gradient 的规模与可训练参数量一致；中间激活对应的梯度只在反向传播过程中临时使用，完成对应计算后就可以释放。

2.7 Checkpoint

一个能够真正恢复训练状态的 Checkpoint 至少需要：

模型参数；
优化器参数。

只保存模型权重，只能恢复模型本身；Momentum、Adam 的历史统计都会丢失，训练轨迹也就不是原来的了。

xxxxavierrr

兴趣是最好的老师。 —— 爱因斯坦

本文是原创文章，采用 CC BY-NC-ND 4.0 协议，完整转载请注明来自 xxxxavierrr

CS3363 LLM3

你可能还想看

CS336（三）：现代 LLM 的基础架构

CS336（一）：Tokenizer 与 BPE

隐私政策

评论内容

0/500

昵称

邮箱

网址

滚动到此处加载评论...

雲自說

来自雲自說最新设计与科技的文章

查看全部