nn.Module · 参数、buffer 与 hook 的那层糖 · PyTorch 源码漫游

nn.Module 自己不做任何数学，它只「组织」——用一个 __setattr__ 的小魔法，把你赋值的 Parameter / 子模块 / buffer 自动登记起来，于是 .parameters()、state_dict、.to(device) 全都水到渠成。看清它如何只靠「持有参数 + 定义 forward」，把求导整个甩给 autograd。

nn.Module · 参数、buffer 与 hook 的那层糖

从引擎深处爬回上层。你写 nanoGPT 时，class GPT(nn.Module)、self.wte = nn.Embedding(...)、self.parameters()、state_dict() 用得飞起。这一章揭穿 nn.Module：它其实什么数学都不做，只干一件事——组织。

一句话：Module 是个「参数容器 + forward + hook」

nn.Module 不实现加法、不实现矩阵乘，更不实现反向。它是张量和 autograd 之上一层薄薄的 Python 糖，职责只有三样：

把散落的参数、子模块、buffer 登记、管理起来；
定义 forward（数据怎么流）；
提供 hook（在前后插自己的逻辑）。

真正的计算，交给张量运算；真正的求导，交给 autograd。Module 只是把这些零件收拾得井井有条。

小魔法：`setattr` 自动登记

Module 最聪明的一手，是重写了 __setattr__。当你写：

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)   # 子模块
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size)  # 子模块

每一个 self.xxx = ... 赋值，都被 Module.__setattr__ 拦下，它一看右边是什么，就自动归档：

是 nn.Parameter → 登记进 self._parameters；
是 nn.Module（子模块）→ 登记进 self._modules；
是 buffer（通过 register_buffer）→ 登记进 self._buffers。

这解释了「为什么你从不手动注册参数」。 self.weight = nn.Parameter(...) 一句，参数就自动进了名册——全靠这个 __setattr__ 拦截。而 nn.Parameter 本身，不过是个 requires_grad=True、且「会被 Module 自动登记」的 Tensor，没有别的。所以nanoGPT 里 self.transformer.wte.weight = self.lm_head.weight 那个权重绑定，本质就是让两个名册项指向同一个 Parameter。

登记好之后，一切水到渠成

有了这三本名册，你平时用的那些方法，全是在名册上递归遍历而已：

.parameters() / .named_parameters()：遍历 _parameters + 递归所有子模块的。nanoGPT 的 configure_optimizers 靠它拿到所有待优化参数，按维度分组——现在你知道这些参数是怎么被「找齐」的了。
.state_dict() / load_state_dict()：把所有 parameter + buffer 收进一个字典 / 从字典灌回去。nanoGPT 存 / 读 checkpoint、from_pretrained 搬 GPT-2 权重，用的都是它。
.to(device)、.cuda()：递归把每个参数和 buffer 搬到目标设备。
.train() / .eval()：递归设置一个 training 标志（dropout、BatchNorm 据此切换行为）。nanoGPT 评估时 model.eval()、之后 model.train()，拨的就是这个标志。
.zero_grad()：递归把每个参数的 .grad 清零。

全是「在名册上走一遍」。 Module 的「强大」，不过是把「登记 + 递归遍历」做扎实了。

Parameter 和 buffer：都在名册，区别在「要不要梯度」

两者都被 Module 管理、都进 state_dict（所以都随 checkpoint 存取），区别只有一个：

Parameter：requires_grad=True，是 autograd 图的叶子，会被优化器更新。
buffer：模型状态的一部分，但不求梯度、不被优化。比如 BatchNorm 的 running_mean，或者nanoGPT 注意力里那个因果掩码 bias——它是固定的掩码，register_buffer 登记进来，随模型走（存 checkpoint、跟着 .to('cuda') 搬），但绝不参与训练。「该不该被优化器动」，就看它是 Parameter 还是 buffer。

`call` 与 hook：为什么是 `model(x)` 不是 `model.forward(x)`

你调 model(x)，触发的是 Module.__call__，它不只是调 forward，而是：

model(x)  →  Module.__call__:
                跑 forward pre-hooks
                out = model.forward(x)
                跑 forward hooks(可改 out)
                (反向时)跑 backward hooks(可改梯度)

所以建议用 model(x) 而不是直接 model.forward(x)——后者会跳过所有 hook。hook 是个强大的口子：你能在不改 forward 源码的前提下，插进去看/改某层的输入、输出、甚至梯度——抓中间激活、做特征可视化、梯度裁剪、调试，都靠它。

[1] PyTorch 文档：“torch.nn.Module.” https://pytorch.org/docs/stable/generated/torch.nn.Module.html 与 “Forward and backward hooks.”
[2] Paszke et al. “PyTorch: An Imperative Style, High-Performance Deep Learning Library.” NeurIPS 2019. arXiv:1912.01703.
源码：torch/nn/modules/module.py（__setattr__、__call__、_parameters/_buffers/_modules）、torch/nn/parameter.py(pytorch/pytorch)。

nn.Module · 参数、buffer 与 hook 的那层糖