斯坦福 CS336 “从零开始构建语言模型”第二讲（Lecture 2）：PyTorch与资源核算（PyTorch, Resource Accounting）

斯坦福 CS336 “从零开始构建语言模型”的第二讲（Lecture 2）的主题是PyTorch与资源核算（PyTorch, Resource Accounting）。讲师从底层机制拆解了PyTorch的运行原理，并教导学生如何通过简单的“纸笔计算（Napkin math）”来精确评估大模型训练时的显存占用量和算力需求。

以下是第二讲中涵盖的所有核心知识点：

1. 数据类型与混合精度（Data Types & Mixed Precision）

在语言模型中，张量（Tensor）是构建参数、梯度、优化器状态和激活值的基础单元。讲师对比了不同的浮点数表示方法：

• FP32（单精度）：占据4个字节（1位符号，8位指数，23位小数）。它是深度学习的默认标准，但占用显存极大。
• FP16（半精度）：占据2个字节。虽然省显存，但动态范围较小（指数位仅5位），在遇到极小值或极大值时容易发生下溢（Underflow）或上溢，导致训练不稳定。
• BF16（Bfloat16）：占据2个字节。专为深度学习设计，虽然精度（小数位）不如FP16，但动态范围与FP32完全相同（8位指数），这使它成为目前大模型计算的标配。
• 混合精度训练的最佳实践：为了在速度、显存和稳定性之间取得平衡，通常在前向和反向传播的矩阵计算中使用 BF16 或 FP8，但模型参数和优化器状态必须保留在 FP32 以确保累加更新时不会丢失精度。

2. PyTorch底层机制与 einops 库

• Tensor的本质：PyTorch中的Tensor实际上是一个指向一维内存数组的指针，并通过元数据（Metadata，如 stride 和 shape）来决定如何读取这些数据。
• 视图（View）与内存拷贝：大部分操作（如切片 x、重塑 .view()、转置 .transpose()）都不会复制内存，只是改变了元数据的读取方式。但像 transpose 这样的操作会使得张量在内存中变得不连续（non-contiguous），后续某些操作可能需要通过 .contiguous() 强制进行内存拷贝。
• 推崇使用 einops：讲师强烈建议使用基于爱因斯坦求和约定的 einops 库（如 einsum, rearrange, reduce）来替代难懂的 .transpose(-1, -2) 等操作。通过用字符串直接命名维度（例如 ‘batch seq_1 hidden, batch seq_2 hidden -> batch seq_1 seq_2’），代码的可读性极大增强，极难写出隐藏的bug。

3. 算力核算（Compute Accounting & FLOPs）

对于大模型，绝大多数计算成本都花在了矩阵乘法（MatMul）上。我们要学会计算浮点运算次数（FLOPs）：

• 矩阵乘法的法则：两个矩阵相乘，所需操作数约为 2 × 参与运算的三个维度的乘积（因为每一个元素都需要一次乘法和一次加法）。
• 前向传播（Forward Pass）算力：对于常见的线性层，前向传播消耗的算力等于 2 × Token数量 × 模型参数量。
• 反向传播（Backward Pass）算力：反向传播不仅需要计算相对于权重的梯度，还需要计算相对于输入的梯度（以便继续往前传）。因此反向传播的计算量大约是前向的两倍，即 4 × Token数量 × 模型参数量。（550行）
• 总训练算力：一次完整的训练迭代，所需的总算力近似为 6 × Token数量 × 模型参数量。
• MFU（模型算力利用率，Model FLOPs Utilization）：这是衡量你写的代码对硬件榨取程度的指标。计算方法是：实际发生的有效FLOPs ÷ 硬件标称的理论峰值FLOPs。如果MFU能大于50%（0.5），就说明利用率非常优秀了。

4. 显存核算（Memory Accounting）

许多人以为显存只要能装下“模型权重”就行了，这是大错特错的。在训练时，显存被以下四大组件瓜分：

• 参数（Parameters）：模型自身的权重。
• 梯度（Gradients）：大小与参数完全一致。
• 优化器状态（Optimizer States）：不同的优化器占用不同。例如AdamW需要存储动量（一阶矩）和梯度平方的运行平均值（二阶矩），这会额外占用2倍的模型参数空间。
• 激活值（Activations）：前向传播过程中产生的所有中间特征矩阵。由于反向传播求导时必须用到它们，因此必须把它们保存在显存中（这也是为什么Batch Size开大会爆显存的原因）。

将上述各项相加，再乘以数据类型的字节数（例如FP32乘4），就能精准算出你到底需要多少显存。

5. 初始化与优化器（Initialization & Optimizers）

• 模型初始化：如果在初始化模型时简单使用标准正态分布（randn），经过矩阵乘法后，数值的方差会随着隐藏维度的变大而爆炸式增长。必须使用类似 Xavier初始化 的方法（即除以输入维度的平方根）来保持方差稳定，否则训练很容易崩溃。
• 优化器回顾：从基础的SGD，到引入动量（Momentum），再到自适应学习率的AdaGrad、RMSprop，最后演化为目前大模型最常用的 Adam/AdamW。

6. 工程最佳实践

• 断点保存（Checkpointing）：由于大模型训练极其漫长且不可避免会崩溃，必须定期保存。切记不仅要保存模型权重，还要原封不动地保存“优化器状态”以及当前的迭代步数，否则训练恢复时会破坏优化器的动量积淀。
• 控制随机种子（Determinism）：深度学习中有太多随机性（初始化、Dropout、数据乱序）。为了能够复现和调试bug，必须养成严格设置全局随机种子的习惯。

---

第二讲复习题 (Lecture 2: PyTorch, Resource Accounting)

一、 数据类型与混合精度

1. 相比于传统的半精度浮点数（FP16），BF16 (Brain Float 16) 在大模型训练中有什么核心优势？为什么它更适合深度学习？
2. 在混合精度训练的最佳实践中，通常会将哪些数据保留在 FP32（单精度）格式，又在哪些具体计算步骤中使用 BF16 或 FP8？

二、 PyTorch底层机制与代码工程

3. PyTorch 中的张量（Tensor）在底层到底是什么？当我们进行切片（Slicing）或转置（Transpose）等“视图（View）”操作时，PyTorch 会复制底层内存数据吗？
4. 讲师为什么强烈建议使用 einops 库（如 einsum、rearrange）来代替 PyTorch 原生的诸如 .transpose(-2, -1) 这样的张量维度操作？
5. 在大模型训练漫长的过程中，为了防止崩溃需要定期保存检查点（Checkpoint）。除了模型自身的权重和当前的迭代步数，还有什么是必须保存的？

三、 算力核算（FLOPs Accounting）

6. 根据“纸笔计算”的经验法则，对于一个线性层，其前向传播的总浮点运算次数（FLOPs）如何通过“Token数量（数据量）”和“模型参数量”来估算？
7. 为什么**反向传播（Backward pass）**的算力消耗是前向传播的 2 倍（即 4 × Token数量 × 模型参数量）？反向传播中具体需要计算哪两部分梯度？
8. 什么是 MFU（Model FLOPs Utilization，模型算力利用率）？如何计算它？在工程实践中，MFU 大于多少会被认为是一个比较好的利用率？

四、 显存核算与初始化（Memory & Initialization）

9. 很多人误以为只要显存能装下“模型参数”就可以进行训练，但实际上在训练语言模型时，显存是被哪四大组件瓜分的？
10. 为什么在初始化模型权重时不能单纯使用标准正态分布，而通常需要根据输入维度进行缩放（例如除以输入维度的平方根，即类似 Xavier 初始化）？

---

💡 参考答案与知识点解析

1. BF16 与 FP16 的对比

答案：FP16 虽然省显存，但只有5位分配给指数部分，动态范围较小，在表示极小数值（如 10−8）时容易发生**下溢（Underflow）**变成0，导致训练不稳定。而 BF16 将8个比特位分配给指数部分，拥有与标准 FP32 完全相同的动态范围。深度学习对动态范围非常敏感而对小数精度要求不高，因此 BF16 完美解决了下溢问题并保持了训练稳定性。

2. 混合精度训练的分配策略

答案：为了防止在累加更新时丢失微小的精度，模型的核心参数（Parameters）和优化器状态（Optimizer states）通常必须保留在 FP32 格式。而对于前向传播和反向传播中极其耗费算力的矩阵乘法（MatMul）运算，则可以将张量临时转换为 BF16 或 FP8 进行计算，以提升速度并减少显存占用。

3. PyTorch 张量的底层逻辑与视图

答案：PyTorch 的张量本质上是指向一段分配好的一维内存数组的指针，并带有描述如何读取这些数据的元数据（Metadata，包含尺寸 shape 和步长 stride）。进行切片或转置时，PyTorch 绝大部分时候不会复制底层内存，而是仅仅修改元数据创建一个“视图（View）”。只有调用类似 .contiguous() 时才可能强制进行内存拷贝。

4. 为什么推崇 einops

答案：使用原生的 .transpose(-2, -1) 依赖于数字索引，代码可读性极差，时间久了很容易忘记维度代表的具体含义并写出隐藏 Bug。einops 引入了类似爱因斯坦求和约定的理念，通过明确给每个维度赋予字符串命名（例如 batch sequence hidden），让维度的变换、重塑和聚合操作变得一目了然，极大地提高了代码的自文档化能力和准确性。

5. 保存检查点（Checkpointing）的关键

答案：除了模型权重和迭代步数，必须保存优化器的状态（Optimizer State）。对于诸如 AdamW 这样的现代优化器，其内部保存了梯度的动量和平方移动平均值等重要历史信息，如果丢失这些状态直接恢复训练，会破坏优化轨迹。

6. 前向传播算力公式

答案：前向传播的算力大约等于 2×Token数量×模型参数量。因为矩阵乘法中计算每一个元素都需要进行一次乘法和一次加法（计为2次浮点运算），将所涉及的维度相乘后正好得出这个公式。

7. 反向传播的 2 倍算力消耗

答案：在反向传播中，根据链式法则，模型不仅需要计算**“损失相对于当前层权重（Weights）的梯度”（用于更新参数），还必须计算“损失相对于当前层输入激活值（Activations / Hiddens）的梯度”**（用于将误差继续向更浅层反向传递）。这两个操作各自都需要进行一次规模等同于前向传播的矩阵乘法，因此总算力是前向的两倍（即 4×Token数量×参数量）。

8. MFU 的定义与优秀标准

答案：MFU 指的是在实际运行模型时，真正发生的有效浮点运算数（Actual FLOPs）除以硬件厂商承诺的峰值理论算力（Promised FLOPs）。通常情况下，因为不可避免的通信和系统开销，MFU 大于 0.5（即 50%） 就被认为是相当优秀的资源利用率了。

9. 训练时的四大显存占用组件

答案：

• 模型参数（Parameters）
• 模型参数对应的梯度（Gradients）
• 优化器状态（Optimizer states）（例如 Adam 会占用额外2倍参数大小的空间）
• 前向传播中产生并为反向传播保留的激活值（Activations）。

10. 权重的恰当初始化

答案：如果仅仅使用标准正态分布进行初始化，在经过多次矩阵乘法后，输出数值会随着隐藏维度（Hidden dimension）的增大而呈平方根级别地爆炸式增长。这会导致模型训练极不稳定。因此需要除以输入维度的平方根来进行缩放调整，确保前向传播过程中各层激活值的方差（分布范围）能够保持稳定