""" ============================================================================= 第 2 讲 · 代码 6：从零实现优化器 —— SGD、动量、AdaGrad、RMSProp、Adam、AdamW ============================================================================= 对应讲义：lectures/lecture02_pytorch_resource_accounting.md 第 9 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 优化器决定"拿到梯度之后，参数怎么更新"。它有两方面值得关心： - 效果：收敛得快不快、稳不稳； - 成本：要为每个参数额外存几个数（直接决定显存，见 05 号文件）。 这个文件沿着历史的脉络，一步一步从 SGD 走到 AdamW，每一步都说明"它解决了上一个的什么问题"， 并且： - 用 PyTorch 官方的 AdamW 验证我们的实现； - 实现 CS336 作业 1 要求的"预热 + 余弦退火"学习率调度和梯度裁剪。 ----------------------------------------------------------------------------- 【整体思路】 ----------------------------------------------------------------------------- 第 1 部分：PyTorch 优化器的写法（继承 torch.optim.Optimizer） 第 2 部分：六个优化器的实现 第 3 部分：与 torch.optim.AdamW 对比 第 4 部分：在一个"病态"问题上比较它们（为什么需要自适应学习率） 第 5 部分：学习率调度：线性预热 + 余弦退火 第 6 部分：梯度裁剪 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 06_optimizers.py """ import math import torch def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1、2 部分：优化器实现 # ============================================================================= # 【PyTorch 小知识】自定义优化器的写法 # class MyOpt(torch.optim.Optimizer): # def __init__(self, params, lr): # super().__init__(params, dict(lr=lr)) # dict 里是"默认超参数" # @torch.no_grad() # 更新参数时不需要记录计算图 # def step(self): # for group in self.param_groups: # 参数可以分组，每组有自己的超参数 # for p in group["params"]: # if p.grad is None: continue # 没有梯度的参数（比如被冻结的）跳过 # state = self.state[p] # 每个参数一个字典，存它的"状态"（动量等） # ... 用 p.grad 更新 p ... # # 【为什么要 @torch.no_grad()】参数更新本身不是模型计算的一部分，不应该被 autograd 记录。 # 【为什么要"原地"修改 p】p -= ... 直接修改参数张量的数据；如果写成 p = p - ...， # 只是让局部变量 p 指向了一个新张量，模型里的参数根本没变。 class SGD(torch.optim.Optimizer): """ 随机梯度下降：θ ← θ - lr × g 【状态】无（0 个额外的数 / 参数）。 【问题】所有参数用同一个学习率。如果不同方向的"陡峭程度"差别很大， 学习率大了会在陡峭方向震荡发散，小了在平缓方向走不动（见第 4 部分）。 """ def __init__(self, params, lr=1e-3): if lr < 0: raise ValueError("学习率不能为负") super().__init__(params, dict(lr=lr)) @torch.no_grad() def step(self): for group in self.param_groups: for p in group["params"]: if p.grad is None: continue p -= group["lr"] * p.grad class Momentum(torch.optim.Optimizer): """ 带动量的 SGD：v ← β v + g；θ ← θ - lr × v 【直觉】像一个在山坡上滚动的球：持续同一方向的梯度会让速度越积越大； 来回震荡的梯度会相互抵消。 【状态】v（1 个数 / 参数）。 """ def __init__(self, params, lr=1e-3, beta=0.9): super().__init__(params, dict(lr=lr, beta=beta)) @torch.no_grad() def step(self): for group in self.param_groups: for p in group["params"]: if p.grad is None: continue state = self.state[p] if "v" not in state: state["v"] = torch.zeros_like(p) v = state["v"] v.mul_(group["beta"]).add_(p.grad) # 【PyTorch 小知识】带下划线的方法是"原地"操作 p -= group["lr"] * v class AdaGrad(torch.optim.Optimizer): """ AdaGrad：G ← G + g²；θ ← θ - lr × g / (√G + ε) 【直觉】每个参数有"自己的"学习率：历史梯度大的方向（陡峭）步子小，历史梯度小的方向（平缓）步子大。 【状态】G（1 个数 / 参数）。 【问题】G 只增不减，学习率会越来越小，训练后期几乎停滞。 """ def __init__(self, params, lr=1e-2, eps=1e-10): super().__init__(params, dict(lr=lr, eps=eps)) @torch.no_grad() def step(self): for group in self.param_groups: for p in group["params"]: if p.grad is None: continue state = self.state[p] if "G" not in state: state["G"] = torch.zeros_like(p) G = state["G"] G += p.grad ** 2 p -= group["lr"] * p.grad / (G.sqrt() + group["eps"]) class RMSProp(torch.optim.Optimizer): """ RMSProp：G ← ρ G + (1-ρ) g²；θ ← θ - lr × g / (√G + ε) 【直觉】把 AdaGrad 的"累加"换成"指数移动平均"：只记住最近一段时间梯度的大小， 学习率不会无限衰减。 【状态】G（1 个数 / 参数）。 """ def __init__(self, params, lr=1e-3, rho=0.99, eps=1e-8): super().__init__(params, dict(lr=lr, rho=rho, eps=eps)) @torch.no_grad() def step(self): for group in self.param_groups: for p in group["params"]: if p.grad is None: continue state = self.state[p] if "G" not in state: state["G"] = torch.zeros_like(p) G = state["G"] G.mul_(group["rho"]).add_((1 - group["rho"]) * p.grad ** 2) p -= group["lr"] * p.grad / (G.sqrt() + group["eps"]) class AdamW(torch.optim.Optimizer): """ AdamW（CS336 作业 1 的写法，来自 Loshchilov & Hutter 2019，Algorithm 2）： m ← β1 m + (1-β1) g 一阶矩：梯度的移动平均（= 动量） v ← β2 v + (1-β2) g² 二阶矩：梯度平方的移动平均（= RMSProp） α_t ← lr × √(1-β2^t) / (1-β1^t) 偏差修正后的学习率 θ ← θ - α_t × m / (√v + ε) Adam 更新 θ ← θ - lr × λ × θ 解耦的权重衰减（W 的由来） 【为什么要"偏差修正"】 m、v 都从 0 开始。前几步时它们被"拉向 0"：比如第 1 步 m = (1-β1) g = 0.1 g。 除以 (1-β1^t) 正好把这个偏差修正回来（t=1 时 m / 0.1 = g）。 随着 t 增大，β^t → 0，修正因子 → 1，不再起作用。 【为什么叫"解耦的"权重衰减】 早期做法是把 L2 正则 λθ 加进梯度（g ← g + λθ），但这样它也会被 √v 缩放， 对梯度大的参数衰减得少、梯度小的参数衰减得多，效果不可控。 AdamW 把权重衰减从梯度里拿出来，直接作用在参数上。 【状态】m、v（2 个数 / 参数）→ 显存是参数的 2 倍（FP32 时每参数 8 字节）。 【参数】 torch_compatible : bool，默认 False（按作业 1 的公式）。 True 时按 PyTorch 官方实现的顺序计算：先做权重衰减、再做 Adam 更新， 并且 ε 加在"√v 除以 √(1-β2^t) 之后"，即 θ ← θ - lr/(1-β1^t) × m / (√v/√(1-β2^t) + ε)。 两种写法在数学上只差在 ε 的位置（和权重衰减的先后），见演示 2。 【边界】 - betas 必须在 [0, 1) 之间，eps > 0。 - 【PyTorch 小知识】m.lerp_(g, w) 是"线性插值"：m ← m + w × (g - m) = (1-w) m + w g。 取 w = 1-β1，正好就是 m ← β1 m + (1-β1) g。 v.addcmul_(g, g, value=c) 是 v ← v + c × g × g。 我们用和 PyTorch 官方实现相同的原地操作，是为了在 torch_compatible=True 时做到逐比特一致。 """ def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01, torch_compatible=False): if not (0 <= betas[0] < 1 and 0 <= betas[1] < 1): raise ValueError("betas 必须在 [0, 1) 之间") if eps <= 0: raise ValueError("eps 必须大于 0") super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)) self.torch_compatible = torch_compatible @torch.no_grad() def step(self): for group in self.param_groups: lr, (b1, b2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"] for p in group["params"]: if p.grad is None: continue state = self.state[p] if not state: # 第一次见到这个参数：初始化状态 state["t"] = 0 state["m"] = torch.zeros_like(p) state["v"] = torch.zeros_like(p) state["t"] += 1 t, m, v, g = state["t"], state["m"], state["v"], p.grad bc1, bc2 = 1 - b1 ** t, 1 - b2 ** t # 偏差修正因子 if self.torch_compatible: p.mul_(1 - lr * wd) # 先做权重衰减 m.lerp_(g, 1 - b1) # m ← β1 m + (1-β1) g v.mul_(b2).addcmul_(g, g, value=1 - b2) # v ← β2 v + (1-β2) g² if self.torch_compatible: denom = (v.sqrt() / math.sqrt(bc2)).add_(eps) p.addcdiv_(m, denom, value=-lr / bc1) else: alpha_t = lr * math.sqrt(bc2) / bc1 p -= alpha_t * m / (v.sqrt() + eps) p -= lr * wd * p # 解耦的权重衰减（作业 1 放在最后） # 不带权重衰减的 Adam 就是 weight_decay=0 的 AdamW def Adam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8): """Adam = weight_decay 为 0 的 AdamW。""" return AdamW(params, lr=lr, betas=betas, eps=eps, weight_decay=0.0) def optimizer_state_numbers_per_param(opt_class, **kwargs): """ 实测一个优化器每个参数存了几个数。 【思路】建一个小参数，做一次 step，统计状态里所有张量的元素个数 / 参数元素个数。 """ p = torch.nn.Parameter(torch.randn(100)) opt = opt_class([p], **kwargs) p.grad = torch.randn(100) opt.step() n = sum(v.numel() for v in opt.state[p].values() if torch.is_tensor(v)) return n / p.numel() def demo_state_sizes(): """演示 1：各优化器的状态大小。""" print_header("演示 1：每个优化器为每个参数额外存几个数") for name, cls in [("SGD", SGD), ("Momentum", Momentum), ("AdaGrad", AdaGrad), ("RMSProp", RMSProp), ("Adam", Adam), ("AdamW", AdamW)]: print(f" {name:<10}{optimizer_state_numbers_per_param(cls):.0f}") print(" → Adam/AdamW 效果好，代价是显存：每个参数额外 2 个 FP32 数（8 字节）。") # ============================================================================= # 第 3 部分：与官方实现对比 # ============================================================================= def train_pair(our_opt_kwargs, steps): """ 用同样的初始参数和数据，分别用我们的 AdamW 和 torch.optim.AdamW 训练， 返回每个记录点上两组参数的最大差异。 【参数】 our_opt_kwargs : 传给我们的 AdamW 的额外参数（例如 torch_compatible=True） steps : 训练步数 【返回】 dict：{步数: 最大参数差异} """ torch.manual_seed(0) model_a = torch.nn.Sequential(torch.nn.Linear(32, 64), torch.nn.ReLU(), torch.nn.Linear(64, 1)) model_b = torch.nn.Sequential(torch.nn.Linear(32, 64), torch.nn.ReLU(), torch.nn.Linear(64, 1)) model_b.load_state_dict(model_a.state_dict()) # 两个模型的初始参数完全相同 hp = dict(lr=1e-2, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1) opt_a = AdamW(model_a.parameters(), **hp, **our_opt_kwargs) opt_b = torch.optim.AdamW(model_b.parameters(), **hp) x, y = torch.randn(128, 32), torch.randn(128, 1) diffs = {} for step in range(1, steps + 1): for model, opt in [(model_a, opt_a), (model_b, opt_b)]: opt.zero_grad() loss = ((model(x) - y) ** 2).mean() loss.backward() opt.step() if step in (1, 10, 50): diffs[step] = max((pa - pb).abs().max().item() for pa, pb in zip(model_a.parameters(), model_b.parameters())) return diffs def demo_compare_with_torch(): """演示 2：我们的 AdamW 与 torch.optim.AdamW 对比。""" print_header("演示 2：我们的 AdamW vs torch.optim.AdamW") exact = train_pair(dict(torch_compatible=True), 50) print(f" (a) torch_compatible=True：第 1/10/50 步参数最大差异 = " f"{exact[1]:.1e} / {exact[10]:.1e} / {exact[50]:.1e}") print(" → 按 PyTorch 的计算顺序实现，结果逐比特一致，说明我们完全理解了官方实现。") assert exact[50] < 1e-6 cs336 = train_pair(dict(torch_compatible=False), 50) print(f" (b) 作业 1 的公式：第 1/10/50 步参数最大差异 = " f"{cs336[1]:.1e} / {cs336[10]:.1e} / {cs336[50]:.1e}") print(""" 为什么 (b) 会有差异，而且越来越大？ - 两个公式只在 ε 的位置上不同：作业 1 是 m/(√v + ε)，PyTorch 相当于 m/(√v + ε×√(1-β2^t))。 对绝大多数参数（|梯度| 远大于 ε=1e-8），两者几乎相同； 但有些参数的梯度本身就接近 0（比如 ReLU 几乎不激活的神经元），此时 ε 的大小会显著改变更新量。 - Adam 的更新量约等于 lr × sign(梯度)，与梯度大小无关，所以哪怕一个极小的差异， 也会让某些参数朝不同方向走一步；之后的梯度随之不同，差异就会越滚越大。 这不是 bug，而是两种合理写法的差别。实践中两者的训练效果没有区别。 【边界】这也提醒我们：同样叫"AdamW"，不同框架的实现细节可能不同， 复现别人的实验时，连 ε 的位置都可能影响"逐比特一致"。""") # ============================================================================= # 第 4 部分：病态问题上的比较 # ============================================================================= def run_on_quadratic(opt_class, lr, steps=200, **kwargs): """ 在二维"病态"二次函数上运行优化器，返回每一步的损失。 【问题】 L(w) = 0.5 × (100 × w1² + 1 × w2²)，起点 w = (1, 1)，最优点 (0, 0)。 w1 方向的曲率是 w2 方向的 100 倍（"条件数"= 100）。 想象一个又窄又长的山谷：横向很陡，纵向很平。 【返回】list[float]，损失曲线；如果发散（出现 inf/nan）就提前停止。 """ w = torch.nn.Parameter(torch.tensor([1.0, 1.0])) curv = torch.tensor([100.0, 1.0]) opt = opt_class([w], lr=lr, **kwargs) losses = [] for _ in range(steps): opt.zero_grad() loss = 0.5 * (curv * w ** 2).sum() loss.backward() opt.step() losses.append(loss.item()) if not math.isfinite(losses[-1]) or losses[-1] > 1e10: break return losses def demo_ill_conditioned(): """演示 3：为什么需要"每个参数自己的学习率"。""" print_header("演示 3：在又窄又长的山谷里下山（条件数 = 100）") configs = [ ("SGD lr=0.019", SGD, 0.019, {}), ("SGD lr=0.021", SGD, 0.021, {}), ("Momentum lr=0.005", Momentum, 0.005, {}), ("AdaGrad lr=0.5", AdaGrad, 0.5, {}), ("RMSProp lr=0.05", RMSProp, 0.05, {}), ("Adam lr=0.05", Adam, 0.05, {}), ] print(f" {'优化器':<22}{'第 10 步':>12}{'第 50 步':>12}{'第 200 步':>12}") for name, cls, lr, kw in configs: losses = run_on_quadratic(cls, lr, **kw) def at(i): return f"{losses[i]:.2e}" if i < len(losses) else "发散" print(f" {name:<22}{at(9):>12}{at(49):>12}{at(199):>12}") print(""" （起点的损失是 0.5 × (100 + 1) = 50.5） - SGD：学习率受陡峭方向限制，必须 < 2/100 = 0.02，否则在 w1 方向来回震荡、越震越大（lr=0.021 发散）。 但这么小的学习率在平缓的 w2 方向每步只缩小约 2%，200 步后损失还有 2e-4。 - 动量：平缓方向的速度会累积起来，陡峭方向的来回震荡会相互抵消，200 步后比 SGD 低好几个数量级。 - AdaGrad / RMSProp：用梯度大小给每个方向"归一化"，两个方向的步长差不多， 可以用大得多的学习率，几十步就收敛了。代价是每个参数要多存状态。 - Adam = 动量 + RMSProp。这里它前 10 步反而冲过了头（动量带来的"惯性"），但很快也收敛了。 【边界】这个例子里曲率恰好沿坐标轴方向，"每个坐标一个学习率"正好完美对症； 真实网络的曲面方向是斜的，自适应方法的优势没这么夸张，但仍然很可观。 【大模型联系】神经网络在不同参数方向上的曲率可以相差好几个数量级， 这是 Adam 家族在大模型训练中几乎一统天下的主要原因。""") # ============================================================================= # 第 5 部分：学习率调度 # ============================================================================= def lr_cosine_schedule(t, lr_max, lr_min, warmup_iters, cosine_cycle_iters): """ 带线性预热的余弦退火学习率（CS336 作业 1 的定义）。 【公式】 t < T_w ：lr = t / T_w × lr_max （线性预热） T_w ≤ t ≤ T_c ：lr = lr_min + ½(1 + cos(π (t - T_w)/(T_c - T_w))) × (lr_max - lr_min) t > T_c ：lr = lr_min （退火结束后保持最小值） 【为什么要预热】 训练刚开始时参数是随机的，梯度方向很不稳定，Adam 的二阶矩 v 也还没估计准， 一上来就用大学习率容易"一步跳飞"。先用很小的学习率走几百到几千步，再逐渐升到最大值。 【为什么要退火】 训练后期需要小步长才能"落进"更低的谷底，而不是在谷底附近来回跳（演示 3 里 Adam 的现象）。 余弦曲线开始下降慢、中间快、结尾又慢，实践效果很好。 【参数】 t : 当前步数（>= 0 的整数） lr_max, lr_min : 最大 / 最小学习率 warmup_iters : 预热步数 T_w（>= 0） cosine_cycle_iters : 余弦退火结束的步数 T_c（>= T_w） 【返回】float 【边界】warmup_iters = 0 时跳过预热；cosine_cycle_iters == warmup_iters 时直接从 lr_max 跳到 lr_min。 """ if t < warmup_iters: return t / warmup_iters * lr_max if t <= cosine_cycle_iters: if cosine_cycle_iters == warmup_iters: return lr_min progress = (t - warmup_iters) / (cosine_cycle_iters - warmup_iters) return lr_min + 0.5 * (1 + math.cos(math.pi * progress)) * (lr_max - lr_min) return lr_min def demo_schedule(): """演示 4：打印学习率曲线。""" print_header("演示 4：线性预热 + 余弦退火") lr_max, lr_min, Tw, Tc = 1e-3, 1e-4, 100, 1000 for t in [0, 50, 100, 300, 550, 800, 1000, 1200]: lr = lr_cosine_schedule(t, lr_max, lr_min, Tw, Tc) bar = "#" * int(lr / lr_max * 40) print(f" t = {t:>5}: lr = {lr:.2e} {bar}") print(" （第 11 讲会介绍另一种流行的调度：WSD，预热-稳定-衰减。）") # ============================================================================= # 第 6 部分：梯度裁剪 # ============================================================================= def clip_gradients(params, max_norm, eps=1e-6): """ 按全局 L2 范数裁剪梯度（CS336 作业 1 的定义）。 【思路】 把所有参数的梯度拼成一个大向量，算它的 L2 范数 ||g||。 如果 ||g|| > max_norm，就把所有梯度乘以 max_norm / (||g|| + eps)， 让整体范数降到约 max_norm；方向保持不变，只缩短长度。 【为什么需要】 训练中偶尔会遇到"梯度爆炸"（某个批次的梯度异常大），一步大更新就可能毁掉之前的训练。 裁剪给每一步的更新幅度加了一个上限。 【参数】 params : 可迭代的参数（通常是 model.parameters()） max_norm : 最大允许的全局范数（> 0） eps : 防止除以 0 的小数 【返回】 float，裁剪**前**的全局范数（方便记录日志） 【边界】 - 没有梯度的参数（p.grad is None）跳过； - 必须是"全局"范数，而不是对每个参数分别裁剪，否则会改变梯度的方向。 """ grads = [p.grad for p in params if p.grad is not None] if not grads: return 0.0 total_norm = math.sqrt(sum((g.float() ** 2).sum().item() for g in grads)) if total_norm > max_norm: scale = max_norm / (total_norm + eps) for g in grads: g.mul_(scale) return total_norm def demo_clipping(): """演示 5：梯度裁剪，并与 torch.nn.utils.clip_grad_norm_ 对比。""" print_header("演示 5：梯度裁剪") torch.manual_seed(0) model_a = torch.nn.Linear(10, 10) model_b = torch.nn.Linear(10, 10) model_b.load_state_dict(model_a.state_dict()) x = torch.randn(4, 10) * 100 # 故意放大输入，制造很大的梯度 for m in [model_a, model_b]: m(x).pow(2).sum().backward() before = clip_gradients(model_a.parameters(), max_norm=1.0) torch.nn.utils.clip_grad_norm_(model_b.parameters(), max_norm=1.0) after = math.sqrt(sum((p.grad ** 2).sum().item() for p in model_a.parameters())) diff = max((pa.grad - pb.grad).abs().max().item() for pa, pb in zip(model_a.parameters(), model_b.parameters())) print(f" 裁剪前全局范数 {before:,.1f} → 裁剪后 {after:.4f}") print(f" 与 torch.nn.utils.clip_grad_norm_ 的最大差异：{diff:.2e}") def main(): demo_state_sizes() demo_compare_with_torch() demo_ill_conditioned() demo_schedule() demo_clipping() print("\n全部演示运行完毕。") if __name__ == "__main__": main()