v0 单卡:一个 training step 花多少 compute 和 memory
v0 是一张 H100,直接训你在 TikTok 训过的 Qwen3-VL-8B。这一课把一个 training step 要多少 compute 和 memory 算清楚。结论是 memory 先爆,光 model state 就要 128 GB,超过 H100 的 80 GB。
一个 step 做四件事
训练就是把这几行重复几万次:
loss = loss_fn(model(x), y) # 1. forward 2. loss
loss.backward() # 3. backward
optimizer.step() # 4. 更新 parameter
optimizer.zero_grad()
forward 得到预测,loss 是预测和答案差多少。backward 从 loss 往回走,算出每个 parameter 的 gradient,也就是该往哪边改、改多少。optimizer.step 按 gradient 改 parameter,Adam 还会参考过去几步的 gradient。
Compute:为什么是 6ND
- forward:每个 parameter、每个 token 2 FLOPs。在 \(Y = XW\) 里,每个 weight 和输入乘一次,再加进输出一次
- backward 是 4 FLOPs。每个 \(Y = XW\) 要做两个一样大的 matmul:\(dX = dY\,W^\top\)(把 gradient 传给前一层)和 \(dW = X^\top dY\)(这个 weight 自己的 gradient)
- 合计:2 + 4 = 6。N 个 parameter(本课程记作 \(\Psi\))、D 个 token,训练要 \(C \approx 6ND\) FLOPs。8B 模型每个 token 要 48 GFLOPs
- 漏了什么:attention 的 \(QK^\top\) 和乘 V 不对应 parameter,6ND 没算。它约是 6ND 的 \(s/(6h)\),s = 8192、h = 4096 时约 1/3,长序列不能忽略
训练要多久
\[ \text{时长} \approx \frac{6ND}{\text{卡数} \times \text{峰值} \times \text{MFU}} \]例子:8B 模型训 1B tokens(假设的数)。
- C = 6 × 8e9 × 1e9 = 4.8e19 FLOPs
- 一张 H100,MFU 40%:989 × 0.4 ≈ 400 TFLOPS,4.8e19 ÷ 4e14 = 1.2e5 秒 ≈ 33 小时
- 32 张卡:理想情况约 1 小时(通信全藏在计算后面)
这里用的是 BF16 的峰值,v0 用 fp32 还要慢约 15 倍(第 2 课)。
MFU:厨师有多少时间真在炒菜
- 定义:MFU(model FLOPs utilization)= 每秒实际完成的模型 FLOPs ÷ 峰值 = tokens/s × 6Ψ ÷(卡数 × 峰值)
- 直觉:厨师(Tensor Core)没在炒的时间,是在等菜(data loading、memory-bound 的算子)、等别的厨房(communication),或者每次只炒一小盘(kernel 太小)
- 正常值:大模型训练一般在 35–50%,Llama 3 405B 在 16K 张 H100 上是 38–43%
- 分子只算 6Ψ:重算(第 2 课)多花的 FLOPs 不算,算上的叫 HFU,所以 MFU 能跨系统直接比
Parameter 从哪来:每层约 12h²
- 标准 Transformer 一层:attention 的 Q、K、V、O 共 4h²,MLP(h → 4h → h)共 8h²,合计 12h²。h = 4096 时每层约 201M,36 层约 7.2B
- Qwen3 的 GQA 只有 8 个 KV head,K、V 只有 h/4 宽,attention 是 2.5h²。SwiGLU 的 MLP 有三个 h × 3h 的矩阵(12288 = 3h),是 9h²。合计 11.5h² ≈ 193M
- 再加 embedding(151936 × 4096 ≈ 0.62B)、LM head 和 vision encoder,总数约 8B
Memory 分四块
打个比方:模型是一本菜谱,每个 parameter 是一条配方。
- gradient:这一轮试吃后写在每条配方旁边的修改意见
- optimizer state:Adam 的笔记本,每条配方记两笔:最近的意见平均往哪边(m),摇摆得多厉害(v)
- activation:做菜时留在台面上的半成品,backward 还要用,不能马上倒掉
parameter、gradient、optimizer state 合称 model state,只跟 parameter 个数有关。每个 parameter 占多少 bytes:
| 部分 | fp32 训练 | mixed precision | 8B,mixed precision |
|---|---|---|---|
| parameter | 4 | 2(bf16) | 16 GB |
| gradient | 4 | 2(bf16) | 16 GB |
| master weight(fp32 正本) | — | 4 | 32 GB |
| Adam m | 4 | 4 | 32 GB |
| Adam v | 4 | 4 | 32 GB |
| 合计 | 16 | 16 | 128 GB |
- fp32 的 Adam 要存 w、g、m、v 四份,这就是“Adam 占 4 倍 parameter 的 memory”
- mixed precision 用 bf16 省下的 4 bytes 被 fp32 master weight 吃回去,还是 16(原因见第 2 课)。其中 optimizer 那三行占 12,是大头
Activation:为什么要存,存多少
- 为什么要存:\(Y = XW\) 的 \(dW = X^\top dY\),没有 X 就算不出 W 的 gradient,所以每个 matmul 和 SwiGLU 的输入都得留到 backward
- 存多少(Korthikanti 2022 的估法,每个数按 16-bit 存):每层 \(sbh(34 + 5as/h)\) bytes,s 是 sequence length,b 是 batch size,a 是 head 数。34 那项是每个 token 的中间向量,\(5as/h\) 那项是 attention 每个 head 一张 s × s 的 score 表
- 每层输出多大(NVIDIA 电面问过):[b, s, h] = 1 × 8192 × 4096 个数,bf16 下 64 MiB
- 例子(s = 8192 是假设的长度,b = 1):sbh = 8192 × 4096 ≈ 3.36e7;5as/h = 5 × 32 × 8192 ÷ 4096 = 320;每层 3.36e7 × 354 ≈ 11.9 GB,36 层 ≈ 430 GB。其中 90% 是 s × s 那项
- 峰值在 backward 开头:forward 一路攒,backward 边用边释放,所以峰值是 model state + 全部 activation
- 公式按 GPT 的结构推导,对 Qwen3 只是数量级的估计
总账:v0 在一张 80 GB 的卡上 OOM
| 部分 | 8B,s = 8192,b = 1 |
|---|---|
| parameter + gradient(bf16) | 32 GB |
| optimizer state(master weight、m、v) | 96 GB |
| activation | ≈ 430 GB |
| 合计 | ≈ 560 GB,是 80 GB 的 7 倍 |
model state 跟着 Ψ 走,只能切到多张卡上(第 4 课)。activation 跟着 b × s 走,单卡上就能压(第 2 课)。
你的经验8B × 16 bytes = 128 GB,这就是你在 TikTok 没法单卡训 Qwen3-VL-8B、要用 DeepSpeed 把它摊到 32 张 H100 上的原因(怎么摊见第 4 课)。
去查一下翻出当时的训练 log,找到每张卡的 tokens/s(或者 step time 和每个 step 的 token 数),算一下 MFU:tokens/s × 48 GFLOPs ÷ 989 TFLOPS(vision token 也算进去,粗估即可)。面试官问 “What MFU did you get?” 时,要能报出一个数,并说出它为什么高或低。
30 秒英文回答
One step is forward, loss, backward and an optimizer update. Forward costs 2 FLOPs per parameter per token, a multiply and an add. Backward needs two matmuls of the same size, one for the input gradient and one for the weight gradient, so training costs about 6ND. An 8B model on a billion tokens is 4.8e19 FLOPs, about 33 hours on one H100 at 40% MFU. For memory, mixed-precision Adam needs 16 bytes per parameter for weights, gradients and optimizer state: 128 GB, already over 80 GB. Activations come on top and grow with batch times sequence length, about 430 GB for one 8K sequence.
自测
14B 模型的 bf16 weight 多大?用 Adam 训练要多少 memory?
weight 14B × 2 = 28 GB。训练的 model state 是 14B × 16 = 224 GB,再加 activation(看 b 和 s)。
8B 模型全用 fp32,SGD、momentum、Adam 各要多少 model state?
SGD 只有 w、g,8 bytes → 64 GB;momentum 多一份速度,12 bytes → 96 GB;Adam 多 m、v 两份,16 bytes → 128 GB。
8B 模型在一张 H100 上实测每秒 3,000 个 token,MFU 多少?说明什么?
3,000 × 48 GFLOPs = 144 TFLOPS,÷ 989 ≈ 15%,远低于正常的 35–50%,说明 GPU 大部分时间在等。下一步看 profiler trace,分清时间卡在数据、通信还是小 kernel 上(第 8 课)。
LoRA r = 8,加在 Qwen3-8B 每层的 W_q 和 W_v 上,有多少 trainable parameter?
每个矩阵 r × (d_in + d_out)。W_q 是 4096 → 4096:8 × 8192 = 65,536;W_v 因为 GQA 是 4096 → 1024:8 × 5120 = 40,960。每层 106,496,36 层 ≈ 3.8M,约占 8B 的 0.05%。初始化 A 随机、B = 0,所以训练从原模型出发。
s 从 4K 加到 8K,activation 涨几倍?
约 3.7 倍(117 → 430 GB),因为占大头的 s × s 随 s² 涨。去掉这一项就正好 2 倍,这正是下一课 FlashAttention 做的事。
想看细一点 6.10 training memory(四块怎么估、LoRA)· 6.11 FLOPs、6ND、MFU(反推 MFU、6ND 漏了什么)· 2.6 optimizer(SGD、Adam 的更新式和 memory)· 3.4 Transformer 复杂度 · 3.14 LoRA