预训练工程:从单机到千卡
预训练是 LLM 最昂贵的一步:数千张 GPU 连续跑数周数月,一次"训崩"就是数百万美元。本篇讲工业级预训练的关键工程决策与故障诊断。
一、规模规划:先算账再开训
算力公式
训练总 FLOPs ≈ 6 · N · D
N = 参数量,D = 训练 token 数
例:7B 模型训 1.4T token(Chinchilla 最优)→ 6 × 7e9 × 1.4e12 ≈ 5.9e22 FLOPs。A100 实测 MFU(算力利用率)约 40-50%,单卡有效算力 ~150 TFLOPs → 约 4.5 万卡时,即 512 卡跑约 4 天。
超参参考表(经验起点)
| 模型规模 | d_model | n_layer | n_head | 峰值 lr | 全局 batch(token) | 上下文 |
|---|---|---|---|---|---|---|
| 125M | 768 | 12 | 12 | 6e-4 | ~50 万 | 2048 |
| 1.3B | 2048 | 24 | 16 | 2e-4 | ~100 万 | 2048 |
| 7B | 4096 | 32 | 32 | 3e-4 | ~200-400 万 | 4096 |
| 70B | 8192 | 80 | 64 | 1.5e-4 | ~400 万 | 8192 |
规律:模型越大,lr 越小,batch 越大(大 batch 平滑梯度噪声,允许更大步长)。µP(Maximal Update Parametrization)理论允许你在小模型上调好 lr 后直接迁移到大模型。
二、分布式并行:四种维度组合
单卡放不下 7B 以上的模型,必须切分。工业界用 3D/4D 并行:
数据并行 DP:每张卡一份完整模型,切 batch → 通信:梯度 all-reduce
张量并行 TP:把单层的矩阵乘切到多卡 → 通信最频繁,限单机 NVLink 内
流水线并行 PP:按层切分到不同卡 → 气泡问题,用 1F1B 调度缓解
序列并行 SP/CP:把超长序列切到多卡 → 长上下文(128k+)必需
ZeRO / FSDP:数据并行的显存革命
朴素 DP 每张卡都存完整的参数+梯度+优化器状态(Adam 状态下 = 16×N 字节 fp32 混合精度)。ZeRO-3 / PyTorch FSDP 把这些也切分,用时再通信聚合:
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD,
mixed_precision=MixedPrecision(param_dtype=torch.bfloat16))
经验搭配:单机内 TP=8(NVLink),跨机 PP 或 DP,万卡级再加 PP。
三、训练稳定性:loss 曲线诊断手册
健康曲线
平滑下降,初期快后期慢,无尖峰。train/val gap 随数据量增大而缩小(大数据下基本无过拟合)。
常见故障
| 现象 | 可能原因 | 对策 |
|---|---|---|
| loss spike(突刺) | 坏 batch(脏数据)、lr 过大、bf16 下溢 | 梯度裁剪;跳过该 batch;spike 后回滚到上一个 checkpoint 并跳过出问题数据段 |
| loss 平台期不降 | lr 过小;数据混洗不充分;模型容量不足 | 升温 lr / 检查 shuffle / 扩模型 |
| loss 发散 NaN | lr 过大;fp16 溢出;归一化缺失 | 降 lr 加 warmup;换 bf16;检查 Pre-Norm |
| grad norm 缓慢爬升 | 发散前兆 | 提前降 lr,或启用 z-loss |
| val loss 上升 | 数据重复 epoch 过多 → 过拟合 | 增数据 / 加 dropout / 早停 |
大模型特有的坑
- 注意力 logits 爆炸(Grok-1/PaLM 曾遇到):加 QK-Norm(对 q、k 做归一化)或 z-loss
λ·log²(Z)抑制 softmax 分母过大。 - Checkpoint 策略:每 N 步存一份 + 保留最近 K 份 + 每天永久存档;spike 后回滚是常态操作,存储系统必须扛得住并发写。
- 容错:千卡集群平均每天坏几张卡,训练框架要支持故障节点热替换与自动重启续训。
四、吞吐优化
单步耗时 = 计算 + 通信 + 显存换入换出
- FlashAttention:必开,O(T) 显存、2-4x 加速。
- 算子融合 / torch.compile:RMSNorm、RoPE、SwiGLU 融合后省 kernel launch 开销。
- 通信重叠:梯度 all-reduce 与反向计算重叠(bucket 化)。
- 梯度累积:显存不够时用小 micro-batch 多步累积成一个大 batch:
accum = 8
for i, (x, y) in enumerate(loader):
loss = model(x, y) / accum
loss.backward()
if (i + 1) % accum == 0:
optimizer.step(); optimizer.zero_grad(set_to_none=True)
- MFU 监控:实际 FLOPs / 理论峰值。>45% 为优秀,低于 30% 说明通信或数据加载在拖后腿。
五、训练中的评估
不能等训完才知道模型好不好。每 K 步评估:
- PPL:在留出验证集上(代码、网页、论文分领域看);
- 下游基准:MMLU、HellaSwag、GSM8K、HumanEval 的小样本子集;
- 定性探测:固定一批 prompt 看生成质量的"感觉"变化。
能力通常在 loss 的特定阈值后涌现式出现(如 GSM8K 准确率在 loss < ~1.8 后陡升),这也是 scaling 的迷人之处。
六、完整预训练 checklist
- 小模型 + 小数据冒烟跑通全流程(含 checkpoint 恢复);
- 用 µP 或文献值确定 lr/batch;
- 数据打包、混洗、去污染复核;
- 开训前 1 天观察 loss、grad norm、吞吐、显存;
- 设置 spike 告警与自动回滚;
- 定期评估与 checkpoint 归档;
- 末期退火:混入高质量数据 + lr 衰减到峰值 10%。
基座模型到手后,它只会续写不会对话——进入 微调与对齐。