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

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

- 作者：David（道雾轩）
- 专栏：PyTorch 源码漫游（https://daiw.org/manual/pytorch-internals.md）
- 最后更新：2026-07-13
- 原文：https://daiw.org/manual/pytorch-internals/nn-module
- 转载与引用：请注明出处并附原文链接（https://daiw.org/about/copyright）

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

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

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

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

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

真正的计算，交给张量运算；真正的求导，交给 [autograd](https://daiw.org/manual/pytorch-internals/autograd-graph)。Module 只是**把这些零件收拾得井井有条**。

## 小魔法：`__setattr__` 自动登记

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

```python
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`。

<Callout type="info">
  **这解释了「为什么你从不手动注册参数」。** `self.weight = nn.Parameter(...)` 一句，参数就自动进了名册——全靠这个 `__setattr__` 拦截。而 **`nn.Parameter` 本身，不过是个 `requires_grad=True`、且「会被 Module 自动登记」的 [Tensor](https://daiw.org/manual/pytorch-internals/tensor-and-storage)**，没有别的。所以[nanoGPT 里 `self.transformer.wte.weight = self.lm_head.weight` 那个权重绑定](https://daiw.org/manual/nanogpt/gpt-assembly)，本质就是让两个名册项指向同一个 Parameter。
</Callout>

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

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

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

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

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

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

- **Parameter**：`requires_grad=True`，是 [autograd 图的叶子](https://daiw.org/manual/pytorch-internals/autograd-graph)，会被优化器更新。
- **buffer**：模型状态的一部分，但**不求梯度、不被优化**。比如 BatchNorm 的 `running_mean`，或者[nanoGPT 注意力里那个因果掩码 `bias`](https://daiw.org/manual/nanogpt/attention)——它是固定的掩码，`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` 源码**的前提下，插进去看/改某层的输入、输出、甚至梯度——抓中间激活、做特征可视化、梯度裁剪、调试，都靠它。

## Module 不实现 backward——因为不需要

最后点破一件事：**`nn.Module` 从不实现反向。** 你只写 `forward`，为什么就能 `backward`？

因为 `forward` 里全是**被 autograd 追踪的张量运算**。[前向一跑，计算图就自动建好了](https://daiw.org/manual/pytorch-internals/autograd-graph)；`loss.backward()` 一调，[引擎自动回放](https://daiw.org/manual/pytorch-internals/autograd-backward)。**Module 只负责「持有参数（叶子）+ 定义 forward」，求导整个甩给了 autograd。** 这就是为什么[nanoGPT 的 `GPT` 类只有 `forward`、没有一行 `backward`](https://daiw.org/manual/nanogpt/gpt-assembly)，却能训练——上层的糖(Module)和底层的引擎(autograd)，各司其职、严丝合缝。

## 向上看完，该向下看金属了

`nn.Module` 这层糖拆完，你对「从模型定义到求导」的上半场就通透了。接下来向下，看张量在 **GPU** 上到底怎么活：CUDA 流、异步执行，以及那个让无数人困惑的「PyTorch 一上来就吃光显存」的缓存分配器。👉 [下到 GPU：CUDA 流与缓存分配器](https://daiw.org/manual/pytorch-internals/cuda-and-memory)

## 参考文献

- [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)。