""" tiny_autograd.py —— 一个只有一百多行的标量自动微分引擎（反向模式 + 前向模式） ===================================================================== 【这个文件要回答的问题】 PyTorch 里写一句 loss.backward()，所有参数的梯度就都算好了。它是怎么做到的？ 这里从零实现一个"玩具版"： - Value：一个标量（一个数）。每次用 Value 做运算（+、×、exp、tanh……）， 结果都会记住"我是由谁、经过什么运算得到的"，于是整个计算自动连成一张图 （计算图，computational graph）。 - Value.backward()：从最终结果出发，按"拓扑顺序"倒着走一遍这张图， 每一步用链式法则把梯度传给输入。这就是**反向模式自动微分**（reverse-mode autodiff）， 也叫反向传播（backpropagation）。 - Dual：前向模式自动微分（forward-mode），用"对偶数" a + b·ε（ε² = 0） 在算函数值的同时顺便算出一个方向导数。用来对比两种模式的代价。 【为什么值得自己写一遍】 PyTorch 的 autograd 原理和这里完全一样，只是： (1) 它的节点是张量而不是标量（一次处理成千上万个数）； (2) 每种运算的"局部求导规则"用 C++/CUDA 写，非常快； (3) 它还处理了 in-place 修改、显存释放等工程细节。 看懂这一百多行，你就知道 backward() 不是魔法。 【怎样使用】 from tiny_autograd import Value x = Value(2.0); y = Value(3.0) z = x * y + x.tanh() z.backward() x.grad, y.grad # dz/dx, dz/dy 【LLM 联系】 训练一个 LLM，每一步就是：前向算出损失（一张巨大的计算图）→ backward() 把损失对几十亿个参数的梯度全部算出来 → 优化器按梯度更新参数。 反向模式的关键性质是：**不管有多少个参数，算出全部梯度只需要一次反向传播**， 代价和一次前向计算同一个量级。这正是深度学习能训练大模型的前提。 """ import math class Value: """ 一个参与自动微分的标量。 【属性】 data：这个数的值（float）。 grad：最终输出对这个数的导数 ∂out/∂self（调用 backward() 之后才有意义；初始为 0）。 _parents：产生它的输入（元组）。叶子节点（直接创建的数）没有父节点。 _backward：一个函数，负责把 self.grad 按局部导数"分发"给父节点。 op：产生它的运算名字（只用于打印和调试）。 【边界】 - grad 是**累加**的（+=）：同一个数在图里被用了两次，两条路径的梯度要加起来。 这也意味着第二次调用 backward() 前要先把 grad 清零（和 PyTorch 的 zero_grad 一样）。 """ def __init__(self, data, _parents=(), op=""): self.data = float(data) self.grad = 0.0 self._parents = _parents self._backward = lambda: None # 叶子节点：什么都不用做 self.op = op def __repr__(self): # 【Python 小知识】__repr__ 决定 print(对象) 时显示什么。 return f"Value(data={self.data:.4f}, grad={self.grad:.4f})" # ---------------- 基本运算 ---------------- # 每个运算做两件事：(1) 算出结果；(2) 定义结果的 _backward： # "如果已知 out.grad = ∂L/∂out，那么 ∂L/∂输入 += ∂L/∂out × ∂out/∂输入" # 这就是链式法则在一个节点上的样子。 def __add__(self, other): other = other if isinstance(other, Value) else Value(other) # 允许 Value + 3 out = Value(self.data + other.data, (self, other), "+") def _backward(): # ∂(a+b)/∂a = 1，∂(a+b)/∂b = 1：加法把梯度原样"复制"给两个输入 self.grad += out.grad other.grad += out.grad out._backward = _backward return out def __mul__(self, other): other = other if isinstance(other, Value) else Value(other) out = Value(self.data * other.data, (self, other), "*") def _backward(): # ∂(a·b)/∂a = b，∂(a·b)/∂b = a：乘法把梯度"交叉"乘上另一个输入的值 self.grad += other.data * out.grad other.grad += self.data * out.grad out._backward = _backward return out def __pow__(self, k): """只支持指数是普通数字（常数）的情况：x ** 2、x ** -1。""" assert isinstance(k, (int, float)), "这个玩具引擎只支持常数指数" out = Value(self.data ** k, (self,), f"**{k}") def _backward(): self.grad += k * self.data ** (k - 1) * out.grad out._backward = _backward return out def exp(self): out = Value(math.exp(self.data), (self,), "exp") def _backward(): self.grad += out.data * out.grad # (e^x)' = e^x，正好是 out.data out._backward = _backward return out def log(self): """自然对数。【边界】self.data 必须 > 0，否则 math.log 会报错。""" out = Value(math.log(self.data), (self,), "log") def _backward(): self.grad += (1.0 / self.data) * out.grad out._backward = _backward return out def tanh(self): t = math.tanh(self.data) out = Value(t, (self,), "tanh") def _backward(): self.grad += (1 - t * t) * out.grad # tanh'(x) = 1 - tanh(x)² out._backward = _backward return out def relu(self): out = Value(max(0.0, self.data), (self,), "relu") def _backward(): self.grad += (1.0 if self.data > 0 else 0.0) * out.grad out._backward = _backward return out # ---------------- 由基本运算组合出来的运算 ---------------- # 【Python 小知识】实现了 __add__、__mul__ 之后，再定义下面这些"魔法方法"， # 就能写 -x、x - y、x / y、3 + x、3 * x 这样的式子。它们都只是基本运算的组合， # 所以不需要单独写求导规则——这正是自动微分的威力。 def __neg__(self): return self * -1 def __sub__(self, other): return self + (-other if isinstance(other, Value) else Value(-other)) def __truediv__(self, other): other = other if isinstance(other, Value) else Value(other) return self * other ** -1 def __radd__(self, other): # 3 + x 会调用 x.__radd__(3) return self + other def __rmul__(self, other): # 3 * x return self * other def __rsub__(self, other): # 3 - x return Value(other) - self def __rtruediv__(self, other): # 3 / x return Value(other) / self # ---------------- 反向传播 ---------------- def backward(self): """ 从 self（通常是损失）出发，计算图中每个节点的 grad = ∂self/∂节点。 【思路】 1. 拓扑排序：把图里的所有节点排成一个列表，保证"每个节点都排在它的输入之后"。 2. 设 self.grad = 1（∂self/∂self = 1）。 3. 按拓扑序**倒着**依次调用每个节点的 _backward()：轮到某个节点时， 所有用到它的下游节点都已经把梯度传给了它，它的 grad 已经完整，可以继续往上传。 【边界】 - 叶子节点（直接创建的数，例如参数）的 grad 不会自动清零：重复调用会把梯度再累加一遍。 - 中间节点的 grad 每次反向传播前都会先清零（它们只是"传递梯度的中转站"）。 这和 PyTorch 的行为一致：PyTorch 只把梯度累加到叶子张量的 .grad 上。 - build() 是递归函数，图特别深（上千层）时会超过 Python 的递归深度限制； 本讲的例子都很小，不会遇到。 """ topo, visited = [], set() def build(v): # 【Python 小知识】id(v) 是对象在内存里的身份编号；用它判断"是不是同一个对象"。 if id(v) not in visited: visited.add(id(v)) for p in v._parents: build(p) topo.append(v) build(self) for v in topo: if v._parents: # 中间节点：清零 v.grad = 0.0 self.grad = 1.0 for v in reversed(topo): v._backward() return topo # 返回拓扑序，方便调用者数一数图里有多少个节点 def zero_grad(values): """把一组 Value 的 grad 清零。""" for v in values: v.grad = 0.0 # ===================================================================== # 前向模式：对偶数 # ===================================================================== class Dual: """ 对偶数 a + b·ε，其中 ε² = 0。 【思路】把 x 换成 x + 1·ε 代入任何由 +、×、exp、tanh…组成的函数 f， 由泰勒展开 f(x + ε) = f(x) + f'(x)·ε（ε² 及更高次都是 0）， 所以结果的 ε 系数恰好就是导数。多变量时，给每个输入一个"方向" v， 结果的 ε 系数就是方向导数 ∇f · v。 【代价】一次前向计算只能得到**一个**方向导数。想要 n 个参数的完整梯度， 要做 n 次（每次把方向设为一个坐标轴）。这就是 LLM 训练不用前向模式的原因。 【属性】val：函数值部分 a；der：导数部分 b。 """ def __init__(self, val, der=0.0): self.val = float(val) self.der = float(der) def _wrap(self, o): return o if isinstance(o, Dual) else Dual(o, 0.0) def __add__(self, o): o = self._wrap(o) return Dual(self.val + o.val, self.der + o.der) def __mul__(self, o): o = self._wrap(o) # (a + bε)(c + dε) = ac + (ad + bc)ε + bd·ε² = ac + (ad + bc)ε return Dual(self.val * o.val, self.val * o.der + self.der * o.val) __radd__ = __add__ __rmul__ = __mul__ def __sub__(self, o): o = self._wrap(o) return Dual(self.val - o.val, self.der - o.der) def exp(self): e = math.exp(self.val) return Dual(e, e * self.der) def tanh(self): t = math.tanh(self.val) return Dual(t, (1 - t * t) * self.der)