论小规模 LM 预训练调参魔法

这一篇主要是用于记录关于 LM 预训练中关于参数设置的一些经验之谈。这些经验的基础来自于 nanochat/nanoGPT 一系列公开的仓库（论 karpathy 的伟大之处），以及其他一些相关的论文，并结合了自己经常观察到的一些实验现象。之所以加上“小规模”，是因为笔者常用的训练规模差不多就是 0.4-0.6B，这个规模允许在 4x4090 或者 2xA100/H100 上训练，并且能够在 48-72h 内进行验证，对于算力并不充足的情况是一个比较友好的区间。以后有条件的话，再来尝试一些 scaling。

nanoGPT/nanochat 题解

nanoGPT 和 nanochat 均是由 karpathy 公开的仓库，它们都是在 GPT-2 规模上的标准预训练仓库。前者比后者要更老一些，karpathy 本人的描述是“very old and deprecated”。nanochat 在 nanoGPT 的基础上使用了显著更现代的优化器和调度设计，还包括了 SFT 阶段的设计。在这里，我们只关注预训练阶段。

nanoGPT 使用的是基于 AdamW 的优化器配置，只有两个参数组：

def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
    # start with all of the candidate parameters
    param_dict = {pn: p for pn, p in self.named_parameters()}
    # filter out those that do not require grad
    param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
    # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
    # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
    decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
    nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
    optim_groups = [
        {'params': decay_params, 'weight_decay': weight_decay},
        {'params': nodecay_params, 'weight_decay': 0.0}
    ]
    num_decay_params = sum(p.numel() for p in decay_params)
    num_nodecay_params = sum(p.numel() for p in nodecay_params)
    print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
    print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
    # Create AdamW optimizer and use the fused version if it is available
    fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
    use_fused = fused_available and device_type == 'cuda'
    extra_args = dict(fused=True) if use_fused else dict()
    optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
    print(f"using fused AdamW: {use_fused}")

    return optimizer

这确实是相当简易的配置了，只有矩阵参数和非矩阵参数两种，后者不用 weight_decay。值得注意的是，这里面的 p.dim() >= 2 更准确的应该是 p.dim() == 2。类似 causalconv 的权重应该也和 bias 一样进入到非矩阵参数中，尽管它的形状是三维的。

在训练 GPT-2 124M 时，nanoGPT 使用的 batchsize 为一次 0.5M token（计梯度累加），weight_decay 为 0.1，lr 范围是 [6e-5, 6e-4]。调度策略是经典的线性 warmup + 余弦退火，warmup 占用的比例是 0.3%。不过实际占比应该会大于这个数，它的总训练 token 数量达到 300B，对于一个 ~100M 的模型来说肯定是过于夸张了。betas 的设置是 (0.9, 0.95)。注意这个设置其实不是 torch 的默认值，默认值应该是 (0.9, 0.999)。nanochat 中对于这个值有更加激进的配置。

到了 nanochat 之后，整体设置就显著复杂了许多。彼时 Muon 已经开始得到验证，nanochat 的配置也是基于 Muon 的：

# Scale the LR for the AdamW parameters by ∝1/√dmodel (tuned for 768 dim model)
dmodel_lr_scale = (model_dim / 768) ** -0.5
print0(f"Scaling the LR for the AdamW parameters ∝1/√({model_dim}/768) = {dmodel_lr_scale:.6f}")

# Build param_groups with all required fields explicit
param_groups = [
    # AdamW groups (embeddings, lm_head, scalars)
    dict(kind='adamw', params=lm_head_params, lr=unembedding_lr * dmodel_lr_scale, betas=(0.8, 0.96), eps=1e-10, weight_decay=0.01),
    dict(kind='adamw', params=embedding_params, lr=embedding_lr * dmodel_lr_scale, betas=(0.8, 0.995), eps=1e-10, weight_decay=0.001),
    dict(kind='adamw', params=value_embeds_params, lr=embedding_lr * dmodel_lr_scale * 0.5, betas=(0.8, 0.995), eps=1e-10, weight_decay=0.01),
    dict(kind='adamw', params=resid_params, lr=scalar_lr * 0.01, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.05),
    dict(kind='adamw', params=x0_params, lr=scalar_lr, betas=(0.96, 0.95), eps=1e-10, weight_decay=0.0),  # higher beta1 for x0
    dict(kind='adamw', params=smear_params, lr=0.2, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.0),
]
# Muon groups (matrix params, grouped by shape for stacking)
for shape in sorted({p.shape for p in matrix_params}):
    group_params = [p for p in matrix_params if p.shape == shape]
    param_groups.append(dict(
        kind='muon', params=group_params, lr=matrix_lr,
        momentum=0.95, ns_steps=5, beta2=0.9, weight_decay=weight_decay,
    ))

optimizer = MuonAdamW(param_groups)

简单来看，属于 Muon 的只有一组，也就是 Linear 的矩阵参数。属于 AdamW 的看上去有很多组，这是因为 nanochat 使用了改良的 GPT-2 结构；每一层的输出会以一定的比例与原始 embedding 混合，双方的比例由逐层标量权重 resid_params 和 x0_params 控制；每一层会额外学习一个 embedding，这个 embedding 会经过门控后混合到 attention 的 value 上；此外还有控制前一个 token 和后一个 token 混合的廉价 bigram，也就是里面的 smear 参数。这些东西或者是逐层标量参数，或者是 embedding。除去这些乱七八糟的东西，通用的部件只有两个，即 embedding 和 unembedding。AdamW 的 lr 部分相对于 dim=768（GPT-2 124M 尺寸）进行了倒数平方根缩放，模型越宽，学习率越低。betas 的设置也是来到了惊人的 (0.8, 0.95) 附近。

对于每一组的 lr，nanochat 的设置也是相当激进的：

parser.add_argument("--embedding-lr", type=float, default=0.3, help="learning rate for embedding parameters (Adam)")
parser.add_argument("--unembedding-lr", type=float, default=0.008, help="learning rate for unembedding parameters (Adam)")
parser.add_argument("--weight-decay", type=float, default=0.28, help="cautious weight decay for the Muon optimizer (for weights)")
parser.add_argument("--matrix-lr", type=float, default=0.02, help="learning rate for matrix parameters (Muon)")
parser.add_argument("--scalar-lr", type=float, default=0.5, help="learning rate for scalars (resid_lambdas, x0_lambdas)")
parser.add_argument("--warmup-steps", type=int, default=40, help="number of steps for LR warmup")
parser.add_argument("--warmdown-ratio", type=float, default=0.65, help="ratio of iterations for LR warmdown")
parser.add_argument("--final-lr-frac", type=float, default=0.05, help="final LR as fraction of initial LR")

speedrun 默认的参数配置下，模型有 24 层，head_dim=128，n_heads=12，vocab_size=32768，所以整个模型的大小应该在 600-800M 左右，比 nanoGPT 中的配置是大不少的。在调度器方面，nanochat 使用的也是更加激进的配置：

# Learning rate schedule (linear warmup, constant, linear warmdown)
def get_lr_multiplier(it):
    warmup_iters = args.warmup_steps
    warmdown_iters = round(args.warmdown_ratio * num_iterations)
    if it < warmup_iters:
        return (it + 1) / warmup_iters
    elif it <= num_iterations - warmdown_iters:
        return 1.0
    else:
        progress = (num_iterations - it) / warmdown_iters
        return progress * 1.0 + (1 - progress) * args.final_lr_frac

其实大致形状和 cosine decay 也差不多，但是 warmup 所用的周期显著更少，同样会经历一个比较漫长的 decay。

nanochat 甚至对 Muon 的 momemtum 和 weight decay 也做了调度。动量调度上，momemtum 会在一定 warmup 周期内从 0.85 上升到 0.97，然后保持不变，最后随着学习率的 warmdown 逐渐下降到 0.9。weight decay 则是在整个训练周期内逐渐退火到零。

MuP 缩放与经验魔法

MuP 缩放指的是随着模型尺度，主要是宽度的变化，学习率应该做出的调整。此处我们暂且不管 MuP 缩放率是怎么来的，这篇文章 中有详细的介绍。就结论而言，在 AdamW 的设置下，embedding 的学习率随模型尺度不变，hidden layer 和 unembedding 的学习率均随着模型尺度增长倒数下降。由于 embedding 和 unembedding 总是使用 AdamW 优化的，因此可以看到 embedding 的学习率设置总是相对偏高，unembedding 的参数设置总是相对偏低。至于到底缩放到什么程度，实验上 AdamW 的学习率设置总是在 1e-4 量级上下，不会太高。作为一个大规模的参考，我们看看 DeepSeek V4 的报告里关于 DeepSeek V4 Flash 是怎么说的：

We employ the Muon optimizer (Jordan et al., 2024; Liu et al., 2025) for the majority of parameters, but use the AdamW optimizer (Loshchilov and Hutter, 2017) for the embedding module, the prediction head module, and the weights of all RMSNorm modules. For AdamW,wesetits hyper-parameters to $\beta_{1}$ = 0.9, $\beta_{2}$ = 0.95, $\epsilon$ = $10^{-20}$, and weight_decay = 0.1. For Muon,we set the momentum to 0.95 and the weight decay to 0.1, and rescale the RMS of each update matrix to 0.18 for reutilization of the AdamW learning rate. … The learning rate is linearly warmed up in the first 2000 steps, maintained at $2.7\times 10^{-4}$ for most of the training. Near the end of the training, we finally decay the learning rate to $2.7 \times 10^{-5}$ following a cosine schedule.

到了 DeepSeek V4 Pro，这个数字也只是从 2.7 减到了 2.0，并没有量级上的改变。虽然报告中并没有说这个学习率是哪一个优化器的，不过可以看到缩放的程度大致也就是 1e-4 量级内的常数改动。

由于 AdamW 的这种性质，即使在纯 AdamW 的环境下，也可以使用一定的缩放，即调大 embedding 的学习率，调小 unembedding 的学习率。进行适当的缩放之后，收敛速度往往会比 nanoGPT 中的配置会好一些，更接近 Muon 的收敛速度。

至于 Muon 优化器，它在 MuP 缩放中有一个相当不错的性质，即 hidden layer 中的学习率与模型尺寸无关。这意味着我们不太需要在不同尺寸的模型中调整 Muon 相关的学习率。而在实验上，Muon 所需要的学习率往往要比 AdamW 大上几个量级，比如 1e-2 量级左右，这导致 Muon 的学习率设置总是偏大的。减小 Muon 部分的学习率反而会导致收敛速度异常慢，没有什么好处。

另外，尽管 nanochat 中使用了比较激进的 betas 配置，个人仍然不推荐这样做。$\beta_{1}$ 减小到 0.8 附近是非常危险的，很容易引发梯度的剧烈震荡，甚至是完全的梯度爆炸。(0.9, 0.95) 是一个更安全的配置。虽然在一定程度的梯度震荡下，loss 仍然能够平稳下降，但调整 betas 和 weight decay 使其平稳是更稳妥的做法。