上一章节手动算了两个参数的梯度。但 Sophia 有十亿个参数,每个参数的梯度都要通过链式法则算——如果要手写这些梯度公式,工程量是不可能完成的。
PyTorch 的自动微分(autograd)解决了这个问题:你只需要写前向传播的代码,PyTorch 自动帮你算所有参数的梯度。
计算图
PyTorch 在执行前向传播时,不只是算出结果——它同时悄悄记录了每一步运算,形成一个计算图。
import torch
x = torch.tensor(2.0)
w1 = torch.tensor(3.0, requires_grad=True) # requires_grad=True 表示"这是需要算梯度的参数"
w2 = torch.tensor(4.0, requires_grad=True)
y = w1 * x # PyTorch 记录:y 是 w1 和 x 的乘积
z = w2 * y # PyTorch 记录:z 是 w2 和 y 的乘积
loss = (z - 20) ** 2 # PyTorch 记录:loss 是 (z-20) 的平方
执行完这段代码后,PyTorch 内部有了一张图:
w1, x → 乘法 → y, w2 → 乘法 → z, 20 → 减法 → 差 → 平方 → loss
一行代码算所有梯度
loss.backward() # 沿着计算图反向传播,自动算出所有梯度
调用 backward() 之后,每个标记了 requires_grad=True 的变量都有了一个 .grad 属性:
print(w1.grad) # tensor(64.) —— 跟上一章节手算的一样
print(w2.grad) # tensor(48.)
你不需要写任何链式法则的公式。PyTorch 知道每一步运算的局部导数(乘法、加法、平方、softmax、ReLU……每一种运算都预先定义好了导数规则),沿着计算图自动串联起来。
这对训练意味着什么
训练 Sophia 的代码本质上就是这个模式的放大版:
# 伪代码
for batch in training_data:
# 前向传播
logits = model(batch.input_tokens) # 前向传播,PyTorch 自动记录计算图
loss = cross_entropy(logits, batch.targets) # 算 loss
# 反向传播
loss.backward() # 自动算出所有十亿个参数的梯度
# 更新参数
optimizer.step() # 用梯度更新参数(下一个单元详讲)
optimizer.zero_grad() # 清空梯度,准备下一步
四行核心代码。前向传播、反向传播、更新、清空——这个循环重复几十亿次,就是 Sophia 315 个 GPU 小时的训练。
PyTorch 不仅自动算梯度,还把计算放在 GPU 上并行执行。矩阵乘法、激活函数、softmax——所有的前向和反向运算都有针对 GPU 优化的实现。研究者只需要关注模型的设计(架构选择、超参数调整),底层的数学和工程由 PyTorch 处理。
下一章节来看一个训练中的实际问题:梯度有时候会变得极大,需要人为限制它的大小——这叫梯度裁剪。