拆开 LLM（一）：纯 numpy 手写 mini-autograd——反向传播不是魔法

2026-07-17

llm-anatomy 是一个"拆开 LLM"的教学项目：不调库，把大语言模型的每个零件——自动微分、注意力、训练循环——亲手造一遍（前两章纯 numpy，第 3 章起 PyTorch 登场）。上一篇（第零章）：知识地图画了整张路线图；这一篇造第一个零件：自动微分引擎，外加一整条优化器推导链。

先把结论放在最前面：loss.backward() 这一行里没有魔法，只有链式法则加记账。 PyTorch 替你做的事，本质上是在你算 loss 的时候偷偷记了一本账——谁由谁算出来、局部导数是多少——然后从 loss 出发把账倒着结一遍。这本账，不到两百行 numpy 就能记明白。

不信？那先看效果。

先看效果

本章代码就四个文件，加起来 499 行：

ch01/
├── autograd.py    自动微分引擎（193 行）
├── gradcheck.py   数值梯度检验：给引擎判卷
├── optimizers.py  SGD / Momentum / RMSProp / Adam 从零实现
└── train.py       双螺旋二分类，四个优化器同台竞技

跑起来：

$ uv run python ch01/train.py
双螺旋 300 点 | MLP 2-64-1 | lr=0.1 | steps=2000 | batch=32 | 同一初始化

optimizer   final loss  accuracy
SGD             0.6267     64.7%
Momentum        0.6286     62.0%
RMSProp         0.0767     97.0%
Adam            0.0743     96.7%

双螺旋是二分类里的经典刁钻数据：两条螺旋线缠在一起，线性模型完全无解。一个两层 MLP，靠我们自己写的 autograd 算梯度、自己写的 Adam 更新参数，2000 步之后 96.7% 的准确率。没有 import torch，从张量到优化器每一行都是我们自己的。

同一张图里还有一个刺眼的事实：同样的初始化、同样的学习率，SGD 在 64% 上躺平了，Adam 已经收工。这个差距从哪来的，是本章后半场的主线。先拆前半场：梯度到底是怎么算出来的。

反向传播：一本账，两个字段

每个节点只记两件事

引擎的核心是一个 Tensor 类，初始化就这几行：

1
2
3
4
5
6
7
class Tensor:
    def __init__(self, data, _prev=(), _op=""):
        self.data = np.asarray(data, dtype=np.float64)
        self.grad = np.zeros_like(self.data)
        self._prev = tuple(_prev)          # 上游节点：图的边
        self._backward = lambda: None      # 梯度分发闭包
        self._op = _op                     # 仅用于调试打印

data 和 grad 不用解释，真正的机关是后面两个：

_prev：我是由哪几个张量算出来的。所有张量的 _prev 连起来，就是计算图——它不是什么需要单独构建的数据结构，就是你做正向计算时顺手留下的脚印。
_backward：一个闭包，内容是"等我拿到自己的梯度后，怎么把它分给上游"。

每个运算符负责在产出结果的同时写好这两笔账。看乘法：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
    def __mul__(self, other):
        other = _ensure(other)
        out = Tensor(self.data * other.data, (self, other), "*")

        def _backward():
            # 乘法：交换律——各自的梯度是"对方的值"
            self.grad += _unbroadcast(other.data * out.grad, self.data.shape)
            other.grad += _unbroadcast(self.data * out.grad, other.data.shape)

        out._backward = _backward
        return out

链式法则在这里的样子朴素到让人失望：z = x * y，那么 dL/dx = dL/dz * y。下游传来的梯度 out.grad，乘上局部导数（对方的值），加到自己头上。每个算子只需要懂自己那一小段导数，链式法则负责把它们串成任意深的网络——这就是"自动"微分的全部含义。matmul、relu、sigmoid、log 全是同一个模子，只是局部导数换了张脸。

拓扑排序不是可选项

账记好了，怎么结？从 loss 出发，把每个节点的 _backward() 调一遍。但顺序不能乱：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
    def backward(self):
        # 拓扑排序：保证一个节点开始分发梯度时，
        # 它的所有下游已经把梯度交齐（否则分发的是半成品）
        topo, visited = [], set()

        def build(t):
            if id(t) not in visited:
                visited.add(id(t))
                for p in t._prev:
                    build(p)
                topo.append(t)

        build(self)
        self.grad = np.ones_like(self.data)  # dL/dL = 1
        for t in reversed(topo):
            t._backward()

为什么必须拓扑排序？因为一个张量可能被下游用了不止一次——比如隐层输出既进了下一层、又进了残差连接。它的完整梯度是所有下游贡献之和。如果它在只收到一半贡献时就开始给上游分发，分出去的就是半成品，而且这个错误会沿着链一路放大。拓扑序保证每个节点开口说话时，它听到的信息已经是完整的。 DFS 后序遍历天然产出这个序，倒着走就行。

对了，self.grad = 1 这行也值得看一眼：反向传播的起点不是什么神秘初始化，就是"loss 对自己的导数是 1"。

坑：广播是最大的那一个

原理三段讲完了，但真写起来，让我调最久的不是链式法则，是下面这几个坑。逐条编号，每一条都踩过：

坑 1：广播的梯度必须 sum 回去——最大的坑，没有之一。 x @ W1 + b1 里，b1 的 shape 是 (64,)，却和 (300, 64) 的矩阵相加——numpy 广播替你把 b1 复制了 300 份。正向复制，反向呢？300 份的梯度必须求和收回来，否则 b1.grad 的 shape 直接对不上，或者更阴险：shape 恰好能再广播一次，梯度悄悄错掉，网络照样训练，只是永远差口气。处理它的是全引擎最值钱的十几行：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
def _unbroadcast(grad, shape):
    """广播梯度归约：把被广播撑大的梯度 sum 回参数原本的 shape。

    正向广播 = 复制数据；反向就必须把复制份的梯度加回来。
    """
    # 多出来的前导维度：整维求和压掉
    while grad.ndim > len(shape):
        grad = grad.sum(axis=0)
    # 原本是 1 的维度：被复制过 n 份，求和收回（保持维度）
    for i, n in enumerate(shape):
        if n == 1 and grad.shape[i] != 1:
            grad = grad.sum(axis=i, keepdims=True)
    return grad

一句话记住它：正向怎么复制，反向就得怎么求和。

坑 2：梯度要 +=，不能 =。 上面说了，一个张量可能有多个下游，每个下游的 _backward 都会给它交一份梯度。用赋值就是后来者覆盖前者，只有最后一个下游的贡献活下来。所有 _backward 里清一色 +=，对应地，每步训练前必须 zero_grad()——这两件事是同一枚硬币的两面。

坑 3：log(0) 会炸。 交叉熵里有 log(p)，sigmoid 输出贴近 0 或 1 时直接 -inf。工程解法很不优雅但很有效：(prob + 1e-8).log()。你在 PyTorch 里不用操心这个，是因为它的 binary_cross_entropy 内部替你 clamp 过了——福利背后都有人写过丑代码。

坑 4：relu 在 0 点没有导数。 数学上 0 点导数未定义，代码里 (self.data > 0) 悄悄裁决了：0 点算 0。所有框架都是这么拍板的，没人告诉你而已。

梯度检验：不对拍的 autograd 都是许愿

引擎写完了，怎么知道梯度算对了？肉眼 review 链式法则？我不信任何人手推的导数，包括我自己的。好在有个几乎不会错的裁判：导数的定义本身。

把某个参数挪动一丁点，看 loss 变多少——中心差分：

1
2
3
4
5
6
7
            old = p.data[i]
            p.data[i] = old + h
            lp = forward(params).data
            p.data[i] = old - h
            lm = forward(params).data
            p.data[i] = old                    # 恢复现场
            num[i] = (lp - lm) / (2 * h)       # 中心差分

gradcheck.py 搭了一个随机小网络（matmul、广播加法、relu、sigmoid、log、mean 全都过一遍），对每个参数的每一个标量元素都这么算一遍数值梯度，再和 autograd 给的解析梯度对比相对误差：

$ uv run python ch01/gradcheck.py
loss = 0.952288
param shape        max_rel_err
W1    (3, 5)         7.624e-09
b1    (5,)           2.117e-10
W2    (5, 2)         3.663e-09
b2    (2,)           2.116e-10

所有参数最大相对误差: 7.624e-09
gradcheck 通过（阈值 1e-4）

最大相对误差 7.6e-9。float64 下中心差分的理论精度就在这个量级，说明解析梯度和数值梯度在机器精度意义上是同一个东西。这一步没有任何技术含量，但它是整章的地基：没过 gradcheck 的 autograd 引擎，跑起训练来不是在优化，是在许愿。 广播归约那种"网络照常收敛但永远差口气"的暗伤，只有它能查出来。

一条公式，三个优化器

梯度有了，该更新参数了。先上最笨的方案：

1
2
3
4
5
6
class SGD(Optimizer):
    """θ = θ - lr·g，就一行。"""

    def step(self):
        for p in self.params:
            p.data -= self.lr * p.grad

它能用，但前面的表格你也看到了：2000 步，64.7%。往下的三个优化器看起来一个比一个花哨，但我要说的重点是：Momentum、RMSProp、Adam 的全部差别，就是把同一条指数加权平均公式喂了不同的东西。

这条公式是：

v = β·v + (1-β)·g

它算的是 g 的滑动平均——β 越大，记忆越长，曲线越平滑。整个现代优化器体系就悬在这一根线上。

Momentum：对梯度做平均

minibatch 的梯度是带噪声的——这一批说往东，下一批说往西。把 EWA 直接套在梯度上，噪声相互抵消，剩下的是一致的方向：

1
2
3
4
5
    def step(self):
        for p, v in zip(self.params, self.v):
            v *= self.beta
            v += (1 - self.beta) * p.grad      # v = β·v + (1-β)·g
            p.data -= self.lr * v

RMSProp：对梯度的平方做平均

换个问题：有的参数梯度常年很大，有的常年很小，一个全局 lr 伺候不了所有人。RMSProp 的答案是给每个参数记一本"梯度通常有多大"的账——同一条 EWA，喂的是 g²：

1
2
3
4
5
    def step(self):
        for p, s in zip(self.params, self.s):
            s *= self.beta
            s += (1 - self.beta) * p.grad ** 2  # s = β·s + (1-β)·g²
            p.data -= self.lr * p.grad / (np.sqrt(s) + self.eps)

更新量除以 √s：梯度常年大的参数被压小步子，常年小的被放大步子。每个参数拿到了自己的学习率——这就是"自适应"三个字的全部内容。

Adam：两本账一起记

Adam 没有新想法。它就是把上面两本账同时记了：m 是梯度的 EWA（Momentum 那本），v 是梯度平方的 EWA（RMSProp 那本）：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
    def step(self):
        self.t += 1
        for p, m, v in zip(self.params, self.m, self.v):
            m *= self.beta1
            m += (1 - self.beta1) * p.grad          # 一阶矩 EWA
            v *= self.beta2
            v += (1 - self.beta2) * p.grad ** 2     # 二阶矩 EWA
            if self.bias_correction:
                m_hat = m / (1 - self.beta1 ** self.t)
                v_hat = v / (1 - self.beta2 ** self.t)
            else:
                m_hat, v_hat = m, v
            p.data -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

唯一的新面孔是 bias_correction 那四行。它值得单独一节。

bias correction：为什么非有不可

m 和 v 都从 0 起步。第一步之后 m = 0.1·g（因为 β1=0.9），v = 0.001·g²（因为 β2=0.999）——两个都被严重低估，且低估的程度不一样。做个除法：第一步的更新量是 0.1·g / √(0.001·g²) ≈ 3.16·g/|g|。前期步长不是缩水，是膨胀了 3.16 倍——m 低估 10 倍，v 低估 1000 倍，开方后 31.6 倍，一除反而超了。

修正的办法直白：既然第 t 步的 EWA 被系统性缩小了 (1 - β^t) 倍，那就除回去。t 很大时这个因子趋近 1，修正自动退场——它只管开头那几百步。

空口无凭，实测（train.py 的第二段输出，同一初始化、lr=0.1）：

Adam 偏差修正实测（同一初始化，lr=0.1）:
  bias_correction=on  首步更新范数 1.6030 | loss step  1: 1.2690  step 10: 0.7451  step 50: 0.5936  step200: 0.2399
  bias_correction=off 首步更新范数 5.0606 | loss step  1: 3.2829  step 10: 0.9961  step 50: 0.4975  step200: 0.2550

对一下账：5.0606 / 1.6030 = 3.16，和理论值 (1-β1)/√(1-β2) = 3.162 严丝合缝。关掉修正，第一步就把 loss 从 0.73 干到 3.28——起手先摔一跤，再花几十步爬起来。在这个玩具问题上它 200 步内还能追回来；换成一个几十亿参数、warmup 都要精心设计的大模型，开局步长凭空放大三倍就不是摔一跤的事了。所以别误会这个机制的定位：bias correction 不是让训练更快，是让第 1 步就服从你设的学习率。

顺带一提，RMSProp（β=0.99）没有这个修正，首步更新量是 lr/√(1-β) = 10 倍 lr。回头看 loss 曲线图里 RMSProp 的原始曲线开局直接蹿到 7 以上（初始 loss 只有 0.73），就是这 10 倍步长干的——Adam 论文补上 bias correction，就是在给这类前科擦屁股。

四个优化器同台

回到开头那张表，现在配上曲线（粗线是 EWA 平滑后的，淡线是原始值——没错，画图的平滑用的还是同一条公式）：

四个优化器的 loss 曲线：同一初始化、同一学习率

optimizer	final loss	accuracy
SGD	0.6267	64.7%
Momentum	0.6286	62.0%
RMSProp	0.0767	97.0%
Adam	0.0743	96.7%

两个值得诚实交代的观察：

SGD 和 Momentum 几乎重合（图里蓝线被橙线盖住了）。这不是 bug：我们的 Momentum 用的是 v = β·v + (1-β)·g 这个平均形式，全批意义上它就是"平滑版 SGD"，平均步长和 SGD 一样。PyTorch 的 momentum 用的是 v = μ·v + g（不带 1-β），稳态下等效于把 lr 放大 1/(1-μ) = 10 倍——你感受到的"momentum 更快"，一大半是这个隐藏的 lr 放大。同一个名字，两种公式，敢用同一个 lr 对比的时候才会露馅。

真正的分水岭是自适应步长。lr=0.1 对这个损失面来说太小，SGD 和 Momentum 老老实实按梯度大小挪步，2000 步还在平台上（跑到 8000 步它们其实能到 92%——不是不能，是慢）；RMSProp 和 Adam 把小梯度参数的步子自动放大，同一预算下直接翻过去了。而 Adam 的曲线比 RMSProp 平滑得多、几乎没有尖刺——那是 m 那本账（Momentum 成分）在起作用。两本账，各管一件事，谁也替代不了谁。

这和 LLM 有什么关系

关系是：今天训练绝大多数大模型，optimizer.step() 里跑的还是这套 2014 年就定型的数学。 Adam 是 2014 年的论文，主流 LLM 用的 AdamW 只改了 weight decay 的挂载位置，m、v、bias correction 三件套原封不动。你在本章 92 行的 optimizers.py 里看到的每一行，放大到几千亿参数后依然逐字成立。说"绝大多数"是因为 2025 年起例外来了：Muon 系优化器已经在万亿参数的生产模型上顶掉 AdamW（Kimi K2 用 MuonClip 端到端预训练了 15.5T token，PyTorch 2.9 起也内置了 torch.optim.Muon）——但它干的仍是"给梯度记账、再决定往哪走多远"这件事的变奏，这条推导链照旧是地基。

放大之后反而多出一个本章就能算清的账：Adam 给每个参数记 m、v 两本账，优化器状态是参数量的两倍。7B 参数的模型，fp32 的 m + v 就是 56GB 显存——为什么大模型训练里有 ZeRO、有 8-bit Adam、有各种优化器状态分片，根源就是你刚刚亲手写的那两行 np.zeros_like。

autograd 这边同理：PyTorch 的 requires_grad、grad_fn、retain_graph，对应的就是本章的 _prev、_backward、拓扑排序。名字更体面，账本是同一本。

试一试（包括把它弄坏）

代码在 llm-anatomy 的 ch01/，正常跑通只要两条命令：

uv run python ch01/gradcheck.py
uv run python ch01/train.py

跑通之后，故意打破它：

把 lr 调大十倍（0.1 → 1.0）再跑 train.py。 猜猜谁先炸？我赌你猜的是 SGD。实测恰恰相反：SGD 反而收敛得更好了（0.1 对它本来就太小），最先死的是 RMSProp——首步 10 倍放大变成等效 lr=10，一步把 sigmoid 推进饱和区，loss 钉死在 9.2，准确率 50%，纯抛硬币。自适应优化器不是不会炸，是炸法不一样。
把 Adam 的 bias_correction 传成 False，只看前 50 步的 loss。 亲眼确认那个 3.16 倍的开局踉跄。再把 β2 改成 0.9999 试试——低估更狠，膨胀更猛。
破坏拓扑排序：把 backward() 里的 reversed(topo) 改成随机序，再跑 gradcheck.py。 你会看到相对误差从 1e-9 飙到接近 1——梯度不是差一点，是彻底错了，因为大量节点在梯度交齐之前就分发了半成品。顺手把 _unbroadcast 删掉跑一次，感受一下 shape 报错和"静默错误"哪个更可怕。

弄坏再修好，比只看它工作学得多。

下一章预告

引擎有了，下一章造 LLM 真正的心脏：注意力。还是纯 numpy——Q、K、V 三个矩阵到底在算什么，softmax 为什么要除 √d_k（不除会发生什么，我们直接做实验看方差爆炸），causal mask 怎么让模型"不偷看未来"。写完你会发现，attention 的核心代码比本章的 _unbroadcast 长不了多少。

仓库在 github.com/KerroKapple/llm-anatomy，每章一个目录，所有输出可复现。有问题欢迎开 issue 讨论。