""" ============================================================================= 第 10 讲 · 代码 3：检查点、逐位复现，以及遇到损失尖峰时怎么办 ============================================================================= 对应讲义：lectures/lecture10_training_in_practice.md 第 3.3、4.2、5.2 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 大模型训练会被打断几十上百次（硬件故障、手动干预）。每次都要从检查点恢复。恢复得对不对？ 遇到损失尖峰时，PaLM 的做法是"回滚到尖峰前约 100 步的检查点，跳过 200–500 个 batch"； OPT 的做法是"回滚到更早的检查点，降低学习率"。在小模型上比一比。 演示 1：检查点与恢复：在第 100 步保存（参数 + 优化器状态；数据位置由步数决定），恢复后训练到第 400 步—— 和不中断地训练逐位相同吗？如果忘了保存优化器状态，或者恢复后数据顺序变了，会怎样？ 演示 2：从第 100 步的检查点（尖峰在第 130–140 步）出发，比较几种处置：什么都不做、跳过一段数据（PaLM）、 降低学习率（OPT）、打开梯度裁剪。 演示 3：PaLM 的处置花了多少代价：粗略的账。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 03_checkpoints_and_spikes.py # 约 50 秒（训练 1 个 400 步的模型 + 8 次从检查点恢复）；需要 PyTorch 依赖本目录的 train_lab.py，以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py （从随机初始化开始训练，不读也不写第 1 讲的缓存；检查点只保存在内存里，不写文件）。 【PyTorch 小知识】 - model.state_dict() / opt.state_dict()：模型参数、优化器状态（Adam 的 m、v、步数）的字典； load_state_dict(…) 把它们装回去。保存检查点就是把这两个字典写到磁盘（这里用 copy.deepcopy 存在内存里）。 - 一个真实的坑：optimizer.load_state_dict 在类型、设备都相同时**直接引用**传进来的张量，不复制。 恢复后的训练会原地修改它们——如果检查点还在内存里、还要再用，就被改坏了。写这个文件的第一版时 我就踩了这个坑（同一个检查点恢复第二次时，结果悄悄变了），所以 resume() 里先 deepcopy。 写到磁盘再读回来的检查点没有这个问题。 """ import copy import math import os import sys import torch HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import pad, print_header # noqa: E402 from fm_world import ALPHABET, build_world # noqa: E402 from tiny_gpt import CharTokenizer, lr_at # noqa: E402 from train_lab import FlexConfig, eval_bpc, make_flex, train_flex # noqa: E402 CFG = FlexConfig(n_layer=4) # 和代码 02 相同：4 层、不加 QK 归一化 STEPS, BATCH, PEAK_LR, CKPT = 400, 16, 0.1, 100 def schedule(step, scale=1.0): """代码 02 用的学习率调度（warmup 100 步 + 余弦到峰值的 10%），可以整体乘一个系数 scale。""" return scale * lr_at(step, STEPS, PEAK_LR) def new_optimizer(model): """和 train_flex 默认的一样的 AdamW。""" return torch.optim.AdamW(model.parameters(), lr=PEAK_LR, weight_decay=0.1, betas=(0.9, 0.95)) def resume(ckpt, data, keep_optimizer=True, data_seed=1, **kw): """ 从检查点恢复，训练到第 STEPS 步。 【参数】ckpt：(模型 state_dict, 优化器 state_dict)；keep_optimizer：False 时用一个全新的 AdamW（m、v 都是 0）； data_seed：数据的种子（和原来的 1 不同，就相当于"恢复后数据顺序变了"）；kw：传给 train_flex 的其他参数 【返回】(模型, history) """ model = make_flex(CFG, seed=0) model.load_state_dict(ckpt[0]) # 模型：load_state_dict 会把数值复制进参数，检查点不受影响 opt = new_optimizer(model) if keep_optimizer: # 优化器：load_state_dict **不会**复制 m、v 这些张量（类型、设备相同时直接引用）， # 之后 Adam 原地更新它们，就会悄悄改坏内存里的检查点——所以先 deepcopy 一份再装进去。 opt.load_state_dict(copy.deepcopy(ckpt[1])) kw.setdefault("lr_schedule", schedule) hist, _ = train_flex(model, data, STEPS, batch_size=BATCH, seed=data_seed, clip=kw.pop("clip", None), log_every=5, opt=opt, start_step=CKPT, **kw) return model, hist def window(hist, lo, hi): """history 里 [lo, hi) 这一段：最高的训练损失、最大的注意力 logit。""" seg = [r for r in hist if lo <= r["step"] < hi] return max(r["loss"] for r in seg), max(r["max_logit"] for r in seg) def demo1_checkpoints(data, val): print_header("演示 1：在第 100 步保存检查点，恢复后训练到第 400 步（学习率 0.1，4 层，种子 0）") full = make_flex(CFG, seed=0) h_full, _ = train_flex(full, data, STEPS, batch_size=BATCH, clip=None, log_every=5, lr_schedule=schedule) model = make_flex(CFG, seed=0) h0, opt = train_flex(model, data, CKPT, batch_size=BATCH, clip=None, log_every=5, lr_schedule=schedule) ckpt = (copy.deepcopy(model.state_dict()), copy.deepcopy(opt.state_dict())) rows = [("不中断地训练（参照）", full, h_full)] m_ok, h_ok = resume(ckpt, data) rows.append(("恢复：参数 + 优化器状态 + 数据位置", m_ok, h_ok)) m_noopt, h_noopt = resume(ckpt, data, keep_optimizer=False) rows.append(("恢复：忘了优化器状态（新的 AdamW）", m_noopt, h_noopt)) m_data, h_data = resume(ckpt, data, data_seed=7) rows.append(("恢复：数据顺序变了", m_data, h_data)) ref_after = {r["step"]: r["loss"] for r in h_full if r["step"] >= CKPT} print(f"{pad('做法', 36)}{pad('与参照的参数最大差', 20, '>')}{pad('第 100–110 步最高损失', 22, '>')}{pad('验证集 bit/字符', 18, '>')}") for name, m, h in rows: diff = max((p - q).abs().max().item() for p, q in zip(m.parameters(), full.parameters())) early = max(r["loss"] for r in h if CKPT <= r["step"] <= CKPT + 10) print(f"{pad(name, 36)}{pad(f'{diff:.1e}', 20, '>')}{pad(f'{early:.3f}', 22, '>')}{pad(f'{eval_bpc(m, val):.3f}', 18, '>')}") same = all(r["loss"] == ref_after[r["step"]] for r in h_ok if r["step"] in ref_after) print(f"\n正确恢复的那次，每一条记录的训练损失都和参照**完全相等**：{same}") print("\n观察 Observations：") print(" 1. 保存了参数、优化器状态，并且第 k 步的数据只由 (种子, k) 决定——恢复后的训练和不中断地训练逐位相同（差 0.0）。") print(" 这就是 PaLM 说的'从任何检查点出发都逐位可复现'：出了问题可以精确地重放，分清是数据的问题还是状态的问题。") print(" 2. 忘了保存优化器状态：Adam 的 m、v 从 0 开始，恢复后的第一步相当于按梯度的符号、以整个学习率更新每一个参数——") print(" 损失立刻从 2.07 跳到 3.39，最终 2.889，远差于参照。数据顺序变了：训练走上另一条轨迹，") print(" 这一次反而更好（0.994）——在这个不稳定的设置里，换一条轨迹就像换一个种子（代码 02），好坏都有可能。") print(" 无论哪种，'恢复后的训练是不是原来那次训练'都无法再确认了。") print(" 3. 检查点要存的不只是模型：优化器状态（第 9 讲：每个参数 12 字节，比模型本身大得多）、学习率调度的位置、") print(" 数据管线的位置、随机数生成器的状态。BLOOM 的完整检查点 2.3TB，只存参数是 329GB。") return ckpt, h_full, full def demo2_remedies(ckpt, h_full, full, data, val): print_header("演示 2：从第 100 步的检查点出发，几种处置损失尖峰的办法（尖峰在第 130–140 步）") rows = [("什么都不做（参照）", h_full, full)] for name, kw in [("跳过第 100–129 步的数据（PaLM 式）", dict(skip_batches=set(range(100, 130)))), ("跳过第 110–139 步的数据（PaLM 式）", dict(skip_batches=set(range(110, 140)))), ("学习率减半（OPT 式）", dict(lr_schedule=lambda s: schedule(s, 0.5))), ("打开梯度裁剪（范数 1.0）", dict(clip=1.0))]: m, h = resume(ckpt, data, **kw) rows.append((name, h, m)) print(f"{pad('处置', 36)}{pad('100–200 步最高损失', 20, '>')}{pad('其间最大 logit', 16, '>')}" f"{pad('第 200 步', 10, '>')}{pad('第 300 步', 10, '>')}{pad('第 399 步', 10, '>')}{pad('验证集 bit/字符', 18, '>')}") for name, h, m in rows: peak, ml = window(h, CKPT, 200) at = {r["step"]: r["loss"] for r in h} cells = [f"{at[k]:.3f}" for k in (200, 300, STEPS - 1)] print(f"{pad(name, 36)}{pad(f'{peak:.3f}', 20, '>')}{pad(f'{ml:.1f}', 16, '>')}" + "".join(pad(c, 10, ">") for c in cells) + f"{pad(f'{eval_bpc(m, val):.3f}', 18, '>')}") print("\n观察 Observations：") print(" 1. 学习率减半（OPT 的做法）效果最好，但方式和我预想的不同：100–200 步之间注意力 logit 照样涨到近千，") print(" 降学习率并没有马上把它压下来；改善来自模型更早走出了停滞——第 200 步损失 1.477（参照 1.933），") print(" 第 300 步 0.635（参照 1.423），最终验证集 0.690（参照 1.156）。它对准的是病因本身：学习率太大（代码 02）。") print(" 2. 跳过一段数据（PaLM 的做法）在这里像抽签：跳过第 100–129 步，结果和参照几乎一样（1.122）；") print(" 跳过第 110–139 步，结果更差（1.670）。学习率和模型状态没变，logit 照样失控，跳过数据绕不过去。") print(" PaLM 的尖峰是'某些 batch 和某种参数状态的组合'造成的，换一段数据就能绕过；这里的病因是学习率本身。") print(" **处置要对准病因**——而要知道病因，就需要演示 1 的逐位复现：能精确重放，才能做对照实验。") print(" 3. 梯度裁剪也没有帮助（1.306）：注意力 logit 的增长不是靠一两个巨大的梯度，而是很多步稳定的小更新累积起来的。") print(" （PaLM 一直开着梯度裁剪，仍然出现了约 20 次尖峰。）") def demo3_palm_cost(): print_header("演示 3：PaLM 处置损失尖峰的代价（粗略的账，数字来自 PaLM 论文第 5.1 节和表 1）") spikes, rollback, skip_lo, skip_hi = 20, 100, 200, 500 total_steps, final_batch_tokens = 255_000, 2048 * 2048 redo = spikes * rollback print(f"约 {spikes} 次尖峰 × 每次回滚约 {rollback} 步 = 重算约 {redo:,} 步，占 {total_steps:,} 步的 {redo / total_steps:.1%}") lo, hi = spikes * skip_lo, spikes * skip_hi print(f"跳过的数据：{spikes} × {skip_lo}–{skip_hi} 个 batch = {lo:,}–{hi:,} 个 batch；" f"按最后阶段的 batch（{final_batch_tokens / 1e6:.1f}M token）算，最多约 {hi * final_batch_tokens / 1e9:.0f}B token，" f"占 780B 的 {hi * final_batch_tokens / 780e9:.0%}（上限；前期的 batch 更小）") print("\n观察 Observations：") print(" 1. 回滚重算的代价不大（约 1% 的步数）；真正的代价是**人**：每次尖峰都要有人发现、判断、决定回滚到哪里、跳过多少。") print(" 2. 这些数字是上限的粗估。它说明 PaLM 的'跳过数据'是一个很便宜的处置——前提是尖峰确实是数据和状态的偶然组合。") if __name__ == "__main__": world, tok = build_world(), CharTokenizer(ALPHABET) data = torch.tensor(tok.encode(world.train_text()), dtype=torch.long) val = torch.tensor(tok.encode(world.val_text()), dtype=torch.long) ckpt, h_full, full = demo1_checkpoints(data, val) demo2_remedies(ckpt, h_full, full, data, val) demo3_palm_cost()