第 11 讲:Scaling Laws(下)
上一讲我们学习了 Scaling Laws 的基础——Power Law 关系和 Chinchilla-optimal 分配。但将这些理论应用到实际训练中还面临许多挑战:如何在不同规模之间转移超参数?如何处理训练过程中的不稳定性?Chinchilla-optimal 是否总是最佳策略?本讲将深入探讨这些实际问题。
超参数转移:μP (Maximal Update Parameterization)
Scaling Laws 告诉我们最优的 N 和 D,但没有告诉我们如何设置模型的超参数(learning rate、initialization scale 等)。在小模型上调好的超参数,直接用到大模型上通常效果很差。
μP(muParameterization)由 Yang et al. (2022) 提出,目标是实现超参数的零成本转移——在小模型上调好超参数,可以直接用到任意大小的模型上。
标准参数化的问题
在 Standard Parameterization(SP)中,所有模型使用相同的 learning rate。但随着模型宽度 d_model 增大:
- 每层的输入
x的大小(L2 norm)保持 O(1) - 权重矩阵
W的初始化通常是O(1/√d) - 输出
Wx的大小是 O(1) - 但权重的更新量
η × gradient的大小也是 O(1) - 当 d 很大时,更新量相对于权重的影响变得不一致
μP 的核心思想
μP 通过调整初始化、learning rate 和某些层的参数化方式,使得在任意宽度下,每层的更新-权重比(update-to-weight ratio)保持一致。具体规则:
| 组件 | Standard Param. | μP |
|---|---|---|
| Input embedding init | O(1) | O(1) |
| Hidden layer init | O(1/√d) | O(1/√d) |
| Output layer init | O(1/√d) | O(1/d) |
| Hidden layer LR | η | η / d (相对于 base width) |
| Input/Output LR | η | η |
| Attention logits | QK^T / √d_head | QK^T / d_head |
μP 的实际价值
μP 最大的实际价值是:你可以在一个很小的 proxy model(如 40M 参数)上做 hyperparameter sweep(尝试几十组 learning rate),然后直接将最优 learning rate 应用到目标大模型(如 7B)。这节省了大量实验成本。在 Cerebras-GPT 等模型的训练中,μP 已被成功应用。
μP 的实验验证
验证方法:在不同宽度的模型上,用相同的(μP 调整后的)learning rate 训练,观察是否都能达到接近最优的 loss。
# muP transfer: tune LR on small model, apply to large model
# Step 1: Define base model (small proxy)
base_width = 256 # d_model of proxy model
target_width = 4096 # d_model of target model
# Step 2: Sweep LR on proxy model
best_lr = hyperparameter_sweep(
model_width=base_width,
lr_candidates=[1e-4, 3e-4, 1e-3, 3e-3, 1e-2],
metric="validation_loss"
) # e.g., best_lr = 3e-3
# Step 3: Scale LR for target model using muP rules
# Hidden layer LR scales as base_width / target_width
target_lr = best_lr # Base LR stays the same in muP
# The parameterization handles the scaling internally
# Step 4: Train target model with transferred LR
train_model(model_width=target_width, lr=target_lr)
Learning Rate 与模型大小的关系
即使不使用 μP,理解 learning rate 与模型大小的经验关系也很重要。
经验观察
在 Standard Parameterization 下,最优 learning rate 随模型大小增大而减小:
一些具体的例子:
| 模型 | 参数量 | Peak Learning Rate |
|---|---|---|
| GPT-3 Small | 125M | 6 × 10⁻⁴ |
| GPT-3 Medium | 350M | 3 × 10⁻⁴ |
| GPT-3 Large | 760M | 2.5 × 10⁻⁴ |
| GPT-3 XL | 1.3B | 2 × 10⁻⁴ |
| GPT-3 | 175B | 6 × 10⁻⁵ |
| LLaMA-7B | 7B | 3 × 10⁻⁴ |
| LLaMA-65B | 65B | 1.5 × 10⁻⁴ |
Batch Size 的 Scaling
Batch size 也是一个需要随规模调整的超参数。有两个维度需要考虑:
Critical Batch Size
McCandlish et al. (2018) 定义了 Critical Batch Size (B_crit):当 batch size 小于 B_crit 时,增大 batch size 可以等比提高训练效率(训练步数减半,总 token 数不变);当 batch size 大于 B_crit 时,进一步增大 batch size 只减少步数但增加总 token 消耗。
直觉上:loss 越低(训练越深入),梯度的 signal-to-noise ratio 越低,需要更大的 batch 来平均化噪声。
实际经验
常见做法是在训练过程中逐渐增大 batch size:
- 训练初期:使用较小的 batch size(如 256K tokens),学习率 warmup
- 训练中后期:逐步增大到更大的 batch size(如 4M tokens)
- GPT-4 据报道在训练过程中将 batch size 从小逐步增大
训练稳定性问题
大模型训练中最棘手的实际问题之一是训练不稳定——loss 在训练过程中突然飙升(loss spike),有时甚至无法恢复。
Loss Spikes
Loss spike 的特征是 loss 突然大幅上升,然后可能缓慢恢复,也可能导致训练发散。常见原因包括:
- 数据质量问题:batch 中出现异常数据(如重复文本、乱码、极长序列)
- 数值不稳定:某些层的 activation 或 gradient 出现 NaN 或 Inf
- Learning rate 过大:大模型对 learning rate 更敏感
- BF16 精度问题:BF16 的动态范围足够,但精度有限(只有 7 位尾数)
Loss Spike 的处理方法
- 跳过坏 batch:检测到 loss 异常高时,丢弃当前 batch 的梯度更新(gradient skipping)
- 回滚到 checkpoint:回到最近的稳定 checkpoint,跳过导致问题的数据
- Gradient Clipping:限制梯度的最大范数,防止爆炸式更新
# Gradient clipping example torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) - 降低 Learning Rate:如果频繁出现 spike,考虑降低 learning rate
- Z-Loss:PaLM 引入了一个辅助 loss,鼓励 logits 的对数范数保持较小:
# Z-loss: penalize large logits to improve stability z_loss = 1e-4 * torch.logsumexp(logits, dim=-1).pow(2).mean()
PaLM 的经验
Google 在训练 PaLM-540B 时遇到了约 20 次 loss spike。他们的解决方案是:从 spike 前约 100 步的 checkpoint 重启,跳过约 200-500 个 batch 的数据。他们发现是特定的数据 batch 触发了不稳定,跳过这些数据后训练可以正常进行。
Warmup 策略
Learning Rate Warmup 是大模型训练的标准做法——在训练开始时,将 learning rate 从 0(或很小的值)线性增加到 peak learning rate。
为什么需要 Warmup?
- 初始化时的梯度不准确:模型刚初始化时,参数接近随机,early batches 的梯度方向和大小都不可靠
- AdamW 的 second moment 需要预热:AdamW 用梯度的指数移动平均来自适应 learning rate,初始阶段这些统计量还不准确
- 避免 early training 的不稳定:大 learning rate 在初始阶段可能导致 loss 发散
Warmup 的典型配置
常见做法是线性 warmup,持续训练 token 总量的 0.1% - 1%:
| 模型 | Total Tokens | Warmup Tokens | Warmup 比例 |
|---|---|---|---|
| GPT-3 | 300B | ~375M | ~0.1% |
| LLaMA | 1.4T | ~2B | ~0.14% |
| Chinchilla | 1.4T | ~3.5B | ~0.25% |
Cooldown 策略
Learning Rate Cooldown(或 Decay)是训练末期逐步降低 learning rate 的过程。主流的两种 schedule:
Cosine Decay
其中 t 是当前步数,T 是总步数。Cosine schedule 从 peak LR 平滑下降到 η_min(通常是 peak 的 1/10 或 1/100)。
Linear Decay
简单地从 peak LR 线性衰减到 0 或 η_min。虽然简单,但在一些实验中效果与 cosine 接近。
Cooldown 的重要性
多个研究表明,最后的 cooldown 阶段对最终 loss 有显著影响:
- MiniCPM:在训练接近结束时加入一个较短的 cooldown 阶段(如最后 10% 的 tokens),loss 会有一个显著的下降
- WSD Schedule(Warmup-Stable-Decay):先 warmup,然后保持稳定的 learning rate,最后快速 decay。这种 schedule 的好处是可以随时决定"何时结束训练"——只需要在想停的时候开始 decay
# WSD (Warmup-Stable-Decay) schedule
def wsd_schedule(step, warmup_steps, stable_steps, decay_steps, peak_lr, min_lr):
if step < warmup_steps:
# Linear warmup
return peak_lr * step / warmup_steps
elif step < warmup_steps + stable_steps:
# Stable phase
return peak_lr
else:
# Cosine decay
decay_progress = (step - warmup_steps - stable_steps) / decay_steps
decay_progress = min(decay_progress, 1.0)
return min_lr + 0.5 * (peak_lr - min_lr) * (1 + math.cos(math.pi * decay_progress))
Downstream Task 的 Scaling 行为
Pre-training loss 遵循平滑的 power law,但 downstream task 的表现并不总是如此平滑。
两种 Scaling 模式
- 平滑 scaling:多数 downstream task 的 accuracy 随模型大小增大而平滑提升(在 log-log 坐标下近似线性)
- 非线性/突变 scaling:某些任务在小模型上表现接近随机,当模型达到某个规模时突然"涌现"出能力
Emergent Abilities 讨论
Wei et al. (2022) 提出了 Emergent Abilities(涌现能力)的概念:某些能力(如 chain-of-thought reasoning、arithmetic)只在模型超过某个规模阈值后才出现。
但后续研究(Schaeffer et al., 2023)提出了不同观点:
"涌现" 可能是评估指标的假象
Schaeffer 等人发现,如果使用非线性的评估指标(如 exact match accuracy,答对了才得分),小模型会显示为 0 分,看起来像"突然涌现"。但如果改用连续的指标(如 per-token accuracy 或 log-likelihood),能力的增长其实是平滑的。这意味着"涌现"可能更多是评估方式的产物,而非模型行为的本质变化。
Beyond Chinchilla:过度训练的价值
Chinchilla Scaling Laws 告诉我们compute-optimal 的训练配置。但在实际部署中,compute-optimal 不一定是最佳策略。
LLaMA 的选择
Meta 的 LLaMA 系列模型有意选择了远超 Chinchilla-optimal 的训练数据量:
| 模型 | 参数量 N | 训练 Tokens D | D/N 比例 | Chinchilla D/N |
|---|---|---|---|---|
| Chinchilla | 70B | 1.4T | 20 | 20 (optimal) |
| LLaMA-1 7B | 7B | 1T | 143 | 20 |
| LLaMA-1 65B | 65B | 1.4T | 22 | 20 |
| LLaMA-2 7B | 7B | 2T | 286 | 20 |
| LLaMA-3 8B | 8B | 15T | 1875 | 20 |
| LLaMA-3 70B | 70B | 15T | 214 | 20 |
为什么要"过度训练"?
核心原因是 Chinchilla-optimal 优化的是训练成本,而非推理成本。
- 推理成本与模型大小正相关:一个 7B 模型的推理成本约是 70B 模型的 1/10。如果 7B 模型通过更多数据训练能接近 70B 的效果,推理性价比大幅提升
- 总拥有成本(TCO):训练是一次性成本,推理是持续成本。对于大规模部署的模型,推理成本远超训练成本
- 数据质量在提升:高质量数据(教科书、代码等)的效用可能高于一般的 web text,D/N = 20 可能不再是最优
Inference-Optimal vs Compute-Optimal
如果我们把优化目标从"最小化训练 FLOPs"改为"最小化达到目标性能的推理 FLOPs",最优策略就变了。一个训练了 15T tokens 的 8B 模型(LLaMA-3 8B),虽然训练成本高于 Chinchilla-optimal,但它的推理速度远快于性能相当的更大模型。这就是 over-training(过度训练)的价值。
Over-training 的 Scaling Laws
Sardana and Frankle (2023) 提出了考虑推理成本的 scaling laws。他们发现:
也就是说,预期推理量越大,就越应该用更小的模型训练更多数据。这完美解释了 LLaMA-3 8B 训练 15T tokens 的决策。
Scaling Laws 的局限性
- 只预测 pre-training loss:不直接预测 downstream task 性能或 chat 质量
- 数据质量未纳入:Scaling Laws 假设数据分布固定,但实际中数据质量差异巨大
- 架构依赖:不同架构(如 MoE vs dense)的 scaling behavior 不同
- 训练效率未考虑:没有考虑硬件利用率、通信开销等实际因素
- 可能在极端规模失效:Power Law 是经验拟合,可能在未探索的规模上偏离
小结
- μP 通过特殊的参数化,实现超参数在不同模型宽度之间的零成本转移
- Learning Rate 通常随模型增大而减小;Batch Size 的 critical batch size 随 loss 降低而增大
- 训练稳定性是大模型的关键挑战:loss spike 可通过 gradient clipping、z-loss、checkpoint 回滚等方法处理
- Warmup + Cosine/WSD Decay 是标准的 learning rate schedule
- Emergent Abilities 可能是评估指标的假象,连续指标下能力增长通常是平滑的
- Over-training(如 LLaMA-3 的 D/N >> 20)在推理密集场景下是合理的策略
- 实际训练配置需要综合考虑训练成本、推理成本、数据质量、硬件约束等多个因素