本讲主线

这节课围绕一个核心问题展开：如何精确估算训练成本。主要分为三块：

  • 内存占用
  • FLOPs 估算
  • 模型与训练配置

这和 Lecture 1 提到的“效率优先”直接相关。

例题：70B 模型训练时长

问题：使用 1024 张 H100，在 15T tokens 上训练一个 70B 模型，需要多久？

根据 Transformer 常用训练量估算公式（$6 \times$ 参数量 $\times$ token 数）：

$$ \text{总 FLOPs} = 6 \times 70 \times 10^9 \times 15 \times 10^{12} = 6.3 \times 10^{24}\,\text{FLOPs} $$

单卡 H100（BF16）理论性能约为 $2 \times 10^{15}$ FLOPs/s。若考虑大模型训练实际效率约为 30%，则单卡有效 FLOPs/s 为：

$$ 0.3 \times 2 \times 10^{15} = 6 \times 10^{14}\,\text{FLOPs/s} $$

总时长：

$$ \text{总时间} = \frac{\text{总 FLOPs}}{\text{每秒有效 FLOPs} \times \text{GPU 数量}} = \frac{6.3 \times 10^{24}}{6 \times 10^{14} \times 1024} $$

约等于 119 天。

内存与数值精度

PyTorch 使用 tensor 存储参数、梯度、优化器状态、输入数据和激活值。

  • float32：单精度，32 bit（4 byte）。
  • float16：半精度，省内存但 dynamic range 较小，容易下溢，训练稳定性受影响。
  • bfloat16：内存占用与 float16 相同，但 dynamic range 接近 float32，在训练中更稳。
  • H100 支持两种 FP8：E4M3 与 E5M2。

当前主流是混合精度训练：例如前向用 {bfloat16, fp8}，而梯度和主参数以 float32 维护，从而平衡稳定性与效率。NVIDIA Transformer Engine 也在推动 FP8 在线性层中的应用。

FLOPs 估算

FLOPs 有两种常见语义：

  1. Floating Point Operations Per Second（性能指标）
  2. Floating Point Operations（计算总量）

假设线性层将 $d$ 维向量映射到 $k$ 维输出，batch 中有 $n$ 个样本。前向总计算量近似：

$$ \text{flops} = 2 \times n \times d \times k $$

其中乘以2是因为每个乘法操作通常伴随着一个加法操作。

反向传播中，参数梯度与激活梯度都会引入额外矩阵运算，因此总量通常近似为前向的 3 倍，即：

$$ \text{flops} = 6 \times n \times d \times k $$

模型部分

这一部分偏基础实践，主要围绕两层线性模型展开：参数量计算、模型搭建、数据加载、优化器设置与训练循环。原课内容较基础，这里不展开细记。

参考资料