这两周可以压成一个问题: 当一个模型训练不动、放不下、跑不快、跑不稳时,你从哪一层动手? 下面这张表就是阶段 2 的排查入口。

GPT-2 small 量级的模型是很好的阶段项目: num_layers=12, hidden_size=768, num_attention_heads=12, seq_length=1024, 参数量约 124M 到 125M,足够暴露训练系统问题,又不会让调试成本过高。 数据可以先用小规模文本数据集或 synthetic token dataset 跑通系统, 再换真实 WebDataset/Streaming pipeline。

模型配置

# GPT-125M training shape
num_layers = 12
hidden_size = 768
num_attention_heads = 12
ffn_hidden_size = 3072
seq_length = 1024
vocab_size = 50257

# tokens_per_step = global_batch_size * seq_length
global_batch_size = 128
tokens_per_step = 128 * 1024

Megatron-LM 风格启动参考

# 2 GPU baseline: TP=1, PP=1, DP=2
torchrun --nproc_per_node=2 pretrain_gpt.py \
  --num-layers 12 \
  --hidden-size 768 \
  --num-attention-heads 12 \
  --seq-length 1024 \
  --max-position-embeddings 1024 \
  --micro-batch-size 4 \
  --global-batch-size 128 \
  --tensor-model-parallel-size 1 \
  --pipeline-model-parallel-size 1 \
  --bf16 \
  --log-interval 10 \
  --save-interval 500

# 8 GPU TP+PP 对照: TP=2, PP=2, DP=2
torchrun --nproc_per_node=8 pretrain_gpt.py \
  --num-layers 12 \
  --hidden-size 768 \
  --num-attention-heads 12 \
  --seq-length 1024 \
  --micro-batch-size 2 \
  --global-batch-size 128 \
  --tensor-model-parallel-size 2 \
  --pipeline-model-parallel-size 2 \
  --bf16 \
  --log-interval 10

MFU(Model FLOPs Utilization) 的目标是回答: 模型理论需要的 FLOPs,有多少比例真正转化成了硬件峰值算力? 今天使用近似公式即可: 每 token 训练 FLOPs 约为 6 * num_params, 再用实际 tokens/sec 得到 achieved FLOPs/sec。 如果要更精确,再把 attention 的二次项和 embedding/head 等细节加进去。

def estimate_mfu(num_params, global_batch, seq_len, step_time_s, peak_flops_per_gpu, num_gpus):
    tokens_per_step = global_batch * seq_len
    tokens_per_second = tokens_per_step / step_time_s
    flops_per_token = 6 * num_params
    achieved = tokens_per_second * flops_per_token
    peak = peak_flops_per_gpu * num_gpus
    return achieved / peak

# 示例:125M 参数,global batch 128,seq 1024,step 0.8s
mfu = estimate_mfu(
    num_params=125_000_000,
    global_batch=128,
    seq_len=1024,
    step_time_s=0.8,
    peak_flops_per_gpu=312e12,  # 示例:按自己的 GPU BF16/FP16 峰值替换
    num_gpus=2,
)
print(f"MFU={mfu:.2%}")

记录表

这篇笔记的重点不是宣布谁更强,而是说明: 在我的硬件、我的模型、我的数据、我的 batch 约束下,哪个方案更合适。 一篇好的 Infra 笔记要能让别人复现实验,也能看懂你的判断边界。

# 取舍笔记结构
1. 实验环境:GPU 型号、数量、互联、CPU、内存、存储、软件版本
2. 模型与数据:125M GPT 配置、seq_len、global batch、数据格式
3. 三组配置:DDP baseline、ZeRO-3、TP+PP
4. 指标表:step time、tokens/s、MFU、peak memory、data_time、checkpoint time
5. nsys 证据:compute、comm、bubble、data stall 截图或摘要
6. 结论:在这台机器上我会优先选择什么,什么条件下会改选

1. 如果 step time 变慢,你如何区分 data stall、communication、pipeline bubble 和 kernel compute?
2. ZeRO-3 和 TP 解决的显存问题分别是什么?
3. 为什么 TP 通常优先放在单机高速互联内?
4. MFU 的分子和分母分别是什么?为什么不同 GPU 不能只看 MFU 横向比较?
5. 125M GPT 的训练实验里,哪些变量必须固定才能公平比较并行策略?
6. 如果 ZeRO-3 MFU 低于 TP+PP,你需要哪些证据才能说明原因?
7. 阶段 3 进入推理 Infra 前,你最想补哪一个训练系统短板?