优化器 · 哪些参数该做权重衰减 · nanoGPT 源码解读

逐行读 configure_optimizers——一个值得停下来看的小设计：把参数按维度分两组，二维以上的权重矩阵做权重衰减、一维的偏置和 LayerNorm 增益不做。外加 fused AdamW 的小心机。

优化器 · 哪些参数该做权重衰减

优化器是「怎么根据梯度更新参数」的那套规则。nanoGPT 用的是 AdamW——这本身没啥稀奇，业界默认。但 configure_optimizers 里藏着一个很值得停下来看的小设计：不是所有参数都一视同仁地做权重衰减。 这一章就拆这个。

先扫一眼代码

model.py 的 configure_optimizers（约 263–287 行）：

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}
    ]
    # 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)

    return optimizer

先补：权重衰减是什么

权重衰减(weight decay) 是一种正则化：每次更新，除了顺着梯度走，还额外把参数往 0 拉一点点。作用是防止某些权重变得过大、过拟合训练数据，让模型更「泛化」。它约等于给损失加了个「参数越大罚越多」的 L2 惩罚（AdamW 的 W 就是把这个衰减和 Adam 解耦、做得更干净）。

关键问题来了：是不是所有参数都该被往 0 拉？

核心设计：按维度分两组

nanoGPT 的答案是不。它把参数按维度分成两组，区别对待：

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]   # 一维及标量

dim >= 2 的参数 → 做权重衰减。 这些是矩阵乘法里的权重矩阵和词嵌入表（wte、wpe、各 Linear.weight、lm_head）。它们是模型真正「记东西」的地方，容易过拟合，该正则化。
dim < 2 的参数 → 不做权重衰减(weight_decay=0)。 这些是**偏置(bias)**和 LayerNorm 的增益/偏移——它们是一维向量。

为什么一维的不衰减？ 想想 LayerNorm 那个 weight（增益）：它的职责是把归一化后的数据缩放回合适的幅度。你要是把它也往 0 拉，等于在跟归一化「对着干」，破坏它的功能。偏置同理——它管的是「平移」，往 0 拉没有正则化意义，只会添乱。所以业界共识是：只给「矩阵类权重」做权重衰减，给「偏置 / 归一化参数」放行。 nanoGPT 用一个漂亮的 p.dim() >= 2 就干净利落地把这两类分开了。

这是那种「不看源码就不会注意、看了才恍然大悟」的细节：权重衰减该不该加，取决于这个参数是干嘛的。 一个 p.dim() >= 2 的判断，背后是「矩阵权重会过拟合、归一化增益不能乱动」这条实打实的经验。这也正是本专栏想让你体会的——魔鬼在细节，而细节只在代码里。

fused AdamW：能用快的就用快的

最后几行是个性能小心机：

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)

用 inspect.signature 动态检查当前 PyTorch 的 AdamW 支不支持 fused 参数；
支持、且在 CUDA 上，就开 fused=True。fused AdamW 把优化器更新的一堆小操作融合成一个 GPU 内核，少了很多次内核启动和显存往返，更新更快。
betas（默认 (0.9, 0.95)）是 Adam 的两个动量系数，控制「一阶/二阶动量」的记忆长度。

这又是一个「有更快的就用、没有就退回普通版」的稳健写法——和注意力那章 Flash Attention 的检测一个套路。

优化器就绪，该转动循环了

配置好了（上一章）、优化器建好了（本章），万事俱备。可训练循环每一步都要新鲜的数据——一批批 token 从哪来、怎么高效地从那个几百 MB 的 train.bin 里随机取？下一章看 get_batch。👉 数据加载：get_batch 怎么喂饭

参考文献

[1] Kingma & Ba. “Adam: A Method for Stochastic Optimization.” ICLR 2015. arXiv:1412.6980. —— Adam。
[2] Loshchilov & Hutter. “Decoupled Weight Decay Regularization (AdamW).” ICLR 2019. arXiv:1711.05101. —— AdamW，以及「解耦权重衰减」。
源码：Andrej Karpathy, nanoGPT， model.py 的 configure_optimizers. https://github.com/karpathy/nanoGPT

优化器 · 哪些参数该做权重衰减