v0 单卡:一个 training step 花多少 compute 和 memory

v0 是一张 H100,直接训你在 TikTok 训过的 Qwen3-VL-8B。这一课把一个 training step 要多少 compute 和 memory 算清楚。结论是 memory 先爆,光 model state 就要 128 GB,超过 H100 的 80 GB。

面试会这样问How much GPU memory does it take to train an 8B model with Adam? · Where does 6ND come from, and how long will training take? · What is MFU, and what does a low MFU tell you?

考过NVIDIA Senior MLE 电面 2025-10(全是计算题:SGD、Adam 各需要多少 parameter;每层输出多大、trainable parameter 多少;LoRA/QLoRA 的 parameter count 和 initialization,Fail,1p3a 1152213)· NVIDIA DevTech NG 2026(MLP 的维度变化,1p3a 1161658)· 字节("14B:FP16 权重大概多大;训练还要加哪些")· 滴滴("AdamW 优化器为什么会占用4倍模型参数量显存?")· 拼多多("训练时占用显存的数据有哪些")· 牛客("mfu计算公式,6Nd公式是什么"、"给定训练所需的Tokens,怎么估计模型训练所需的完整时间?")

一个 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

训练要多久

\[ \text{时长} \approx \frac{6ND}{\text{卡数} \times \text{峰值} \times \text{MFU}} \]

例子:8B 模型训 1B tokens(假设的数)。

这里用的是 BF16 的峰值,v0 用 fp32 还要慢约 15 倍(第 2 课)。

MFU:厨师有多少时间真在炒菜

Parameter 从哪来:每层约 12h²

Memory 分四块

打个比方:模型是一本菜谱,每个 parameter 是一条配方。

parameter、gradient、optimizer state 合称 model state,只跟 parameter 个数有关。每个 parameter 占多少 bytes:

部分fp32 训练mixed precision8B,mixed precision
parameter42(bf16)16 GB
gradient42(bf16)16 GB
master weight(fp32 正本)—432 GB
Adam m4432 GB
Adam v4432 GB
合计1616128 GB

Activation:为什么要存,存多少

总账: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 做的事。