26

自动微分——PyTorch 怎么帮你算梯度

阅读约 5 分钟第 26/60 课

上一章节手动算了两个参数的梯度。但 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 处理。

下一章节来看一个训练中的实际问题:梯度有时候会变得极大,需要人为限制它的大小——这叫梯度裁剪。