""" ============================================================================= 第 10 讲 · 代码 1：16 位浮点数——溢出、下溢、被"吞掉"的更新，以及混合精度的三个补丁 ============================================================================= 对应讲义：lectures/lecture10_training_in_practice.md 第 2 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 大模型都用 16 位浮点数训练（快一倍、省一半显存，第 8 讲 §3.4）。但 16 位数能表示的范围和精度都很有限。 OPT-175B 的日志里反复出现"loss scale 掉到 0"；BLOOM 的 104B 试验用 fp16 一直发散，换成 bf16 才稳定。 演示 1：fp32、fp16、bf16 的范围和精度：最大值、最小正规数、最小次正规数、机器精度 eps。 演示 2：两种"吞掉"：(a) 很多小数加起来，加到一定程度就加不动了；(b) 1 + 很小的更新 = 1。 演示 3：第 1 讲 base model 的真实梯度，在 fp16 里反向传播：loss scale 太小 → 小梯度变成 0； 太大 → 溢出成 inf。合适的范围在哪里。 演示 4：动态 loss scaling：遇到 inf 就跳过这一步、把 scale 减半；连续若干步没事就加倍。 演示 5：权重直接存成 16 位（没有 fp32 主参数）时，有多少比例的 Adam 更新被舍入掉； 从零训练 300 步：fp32 / 权重存 bf16 / 权重存 fp16 / bf16 + fp32 主参数。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 01_number_formats.py # 约 30 秒；需要 PyTorch 依赖本目录的 train_lab.py（不用它的模型，只为统一路径），以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py （演示 3、4 只读加载第 1 讲的 base model，先 copy.deepcopy 再转成 16 位）。 【关于"模拟"】 在这台 CPU 上，真正的 16 位矩阵乘法比 32 位慢好几倍（CPU 没有专门的 16 位单元）。所以： - 演示 3、4 真的用 16 位做前向和反向（只做几十步，慢一点没关系）； - 演示 5 的训练用 32 位计算，但**每次更新之后把权重舍入成 16 位**（p.to(bfloat16).float()）—— 这精确地模拟了"权重以 16 位存储"，但没有模拟 16 位计算本身的舍入。讲义会说明这个区别。 【PyTorch 小知识】 - torch.finfo(torch.float16)：查一个浮点类型的 max、tiny（最小正规数）、eps（1 和下一个数的距离）。 - t.half() / t.bfloat16() / t.float()：转换类型；model.half() 把模型的所有参数转成 fp16。 - torch.isfinite(t)：哪些元素不是 inf / nan。 """ import copy import math import os import sys import torch import torch.nn.functional as F 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 BASE_CONFIG, CharTokenizer, bits_per_char, get_batch, load_or_train_base_model, lr_at, make_model # noqa: E402 # ============================================================================= # 演示 1：三种格式 # ============================================================================= def demo1_formats(): print_header("演示 1：fp32、fp16、bf16 的范围与精度") rows = [("fp32", torch.float32, "1 + 8 + 23"), ("fp16", torch.float16, "1 + 5 + 10"), ("bf16", torch.bfloat16, "1 + 8 + 7")] print(f"{pad('格式', 6)}{pad('符号+指数+尾数', 16)}{pad('最大值', 14, '>')}{pad('最小正规数', 14, '>')}" f"{pad('最小次正规数', 14, '>')}{pad('eps', 12, '>')}") for name, dt, bits in rows: fi = torch.finfo(dt) sub = torch.tensor(fi.tiny, dtype=dt) # 最小的次正规数：从最小正规数开始反复减半，直到再减半就变成 0 while (sub / 2).item() > 0 and (sub / 2).to(dt).item() > 0: sub = (sub / 2).to(dt) print(f"{pad(name, 6)}{pad(bits, 16)}{pad(f'{fi.max:.4g}', 14, '>')}{pad(f'{fi.tiny:.4g}', 14, '>')}" f"{pad(f'{sub.item():.4g}', 14, '>')}{pad(f'{fi.eps:.4g}', 12, '>')}") print("\n观察 Observations：") print(" 1. fp16 的最大值只有 65504，比它大的数变成 inf；比 2^-24 ≈ 6e-8 小的数变成 0（下溢）。") print(" 在 6e-8 到 6.1e-5 之间是'次正规数'，精度越来越差。范围窄，是 fp16 训练的主要麻烦。") print(" 2. bf16 用了和 fp32 一样多的指数位（8 位），范围和 fp32 几乎相同（最大 3.4e38）——基本不会溢出或下溢；") print(" 代价是尾数只有 7 位：eps = 0.0078，只有大约 2–3 位十进制有效数字。") print(" 3. 两种 16 位格式各缺一样：fp16 缺范围，bf16 缺精度。混合精度训练的'补丁'就是分别针对它们的。") # ============================================================================= # 演示 2：两种"吞掉" # ============================================================================= def demo2_swamping(): print_header("演示 2：两种'吞掉'——累加停滞与小更新消失") print("(a) 把 10000 个 1e-4 加起来（正确答案 1.0）：") for name, dt in [("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)]: s = torch.zeros((), dtype=dt) inc = torch.tensor(1e-4, dtype=dt) for _ in range(10000): s = s + inc print(f" {pad(name, 6)}累加结果 {s.item():.6g}") print("(b) 1 + δ 是否等于 1：") for delta in [1e-2, 1e-3, 1e-4]: cells = [] for name, dt in [("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)]: one = torch.ones((), dtype=dt) cells.append(f"{name}: {'被吞掉' if (one + torch.tensor(delta, dtype=dt)).item() == 1.0 else '保留'}") print(f" δ = {delta:g}：" + "，".join(cells)) print("\n观察 Observations：") print(" 1. fp16 的累加停在约 0.25：和变大以后，相邻两个可表示数的间距超过了 2 × 1e-4，再加 1e-4 就被舍入掉了。") print(" bf16 停得更早（约 0.03）。所以混合精度的第三个补丁是：**矩阵乘法、softmax 的求和用 fp32 累加**") print(" （tensor core 的硬件就是这样做的；BLOOM 也把梯度累加和 softmax 放在 fp32 里）。") print(" 2. 1 + 0.001 在 bf16 里还是 1：一个大小为 1 的权重，加上 0.001 的更新，等于没加。") print(" 这就是第二个补丁——fp32 主参数——要解决的问题（演示 5）。") # ============================================================================= # 演示 3、4：真实梯度在 fp16 里 # ============================================================================= def fp16_grads(base, x, y, scale): """ 用 fp16 模型做一次前向 + 反向（损失乘以 scale），返回除以 scale 之后的梯度（fp32 张量拼成的向量）。 【思路】这就是 loss scaling：反向传播是线性的，损失乘以 S，所有梯度都乘以 S—— 小梯度被"抬"进 fp16 能表示的范围；用之前再除以 S（在 fp32 里除，不会再下溢）。 """ m16 = copy.deepcopy(base).half() loss = F.cross_entropy(m16(x).float().view(-1, 44), y.view(-1)) (loss * scale).backward() return torch.cat([p.grad.float().flatten() for p in m16.parameters()]) / scale def demo3_loss_scale(base, data): print_header("演示 3：第 1 讲 base model 的真实梯度，在 fp16 里反向传播（batch 32）") x, y = get_batch(data, 32, 96, torch.Generator().manual_seed(0)) m32 = copy.deepcopy(base) F.cross_entropy(m32(x).view(-1, 44), y.view(-1)).backward() ref = torch.cat([p.grad.flatten() for p in m32.parameters()]) a = ref.abs() print(f"fp32 梯度：中位数 {a.median().item():.2e}；小于 fp16 最小正规数 6.1e-5 的占 {(a < 6.1e-5).float().mean().item():.1%}，" f"小于 2^-24 的占 {(a < 2 ** -24).float().mean().item():.2%}") print(f"\n{pad('loss scale', 12)}{pad('被冲成 0 的梯度', 16, '>')}{pad('变成 inf 的', 12, '>')}{pad('与 fp32 梯度的相对误差', 24, '>')}") for p in [0, 4, 8, 12, 16, 20, 24]: g = fp16_grads(base, x, y, 2.0 ** p) zero = ((g == 0) & (ref != 0)).float().mean().item() bad = (~torch.isfinite(g)).float().mean().item() rel = f"{((g - ref).norm() / ref.norm()).item():.4f}" if bad == 0 else "—（有 inf）" print(f"{pad(f'2^{p}', 12)}{pad(f'{zero:.4%}', 16, '>')}{pad(f'{bad:.2%}', 12, '>')}{pad(rel, 24, '>')}") print("\n观察 Observations：") print(" 1. 这个小模型的梯度不算太小：不缩放时只有 0.007% 的梯度被冲成 0，但整体的相对误差 1.07%——") print(" 12.6% 的梯度落在'次正规数'区间，精度很差。乘以 2^4 以上，误差就减半到 0.55%，再往上不再改善") print(" （剩下的误差来自 fp16 的前向本身）。大模型、大 batch 的梯度更小，下溢的问题严重得多。") print(" 2. scale 太大（2^24）时，近 20% 的梯度溢出成 inf——这一步的更新必须整个跳过。") print(" Micikevicius et al. (2018) 报告的 scale 从 8 到 32K 不等，取决于网络：没有一个固定值对所有模型都合适。") def demo4_dynamic_scaling(base, data): print_header("演示 4：动态 loss scaling（fp16 前向反向，fp32 主参数 + AdamW，40 步，batch 8）") model16 = copy.deepcopy(base).half() master = [p.detach().float().clone().requires_grad_(False) for p in model16.parameters()] opt = torch.optim.AdamW(master, lr=1e-4, weight_decay=0.1, betas=(0.9, 0.95)) scale, good, growth_interval = 2.0 ** 24, 0, 10 # 故意从一个太大的 scale 开始 gen = torch.Generator().manual_seed(2) events = [] for step in range(40): x, y = get_batch(data, 8, 96, gen) model16.zero_grad() loss = F.cross_entropy(model16(x).float().view(-1, 44), y.view(-1)) (loss * scale).backward() grads = [p.grad.float() / scale for p in model16.parameters()] if not all(torch.isfinite(g).all() for g in grads): events.append((step, "溢出：跳过这一步，scale 减半", scale)) scale /= 2 good = 0 continue for mp, g in zip(master, grads): mp.grad = g opt.step() with torch.no_grad(): for p, mp in zip(model16.parameters(), master): p.copy_(mp.half()) # 更新在 fp32 主参数上做，再转回 fp16 good += 1 if good == growth_interval: events.append((step, f"连续 {growth_interval} 步没有溢出：scale 加倍", scale)) scale *= 2 good = 0 for step, what, s in events: print(f" 第 {step:2d} 步 scale = 2^{int(math.log2(s))} {what}") print(f" 最终 scale = 2^{int(math.log2(scale))}；40 步里跳过了 {sum('溢出' in e[1] for e in events)} 步") print("\n观察 Observations：") print(" 1. 从一个太大的 scale 开始，前几步连续溢出，每次减半，几步之后就落进了合适的范围；") print(" 之后每过一段时间试着加倍一次，太大了再减半——scale 在'刚好不溢出'的边缘来回。") print(" 2. 跳过的步骤什么也不更新。偶尔跳一步无所谓；但 OPT-175B 的日志里，训练发散之前 loss scale 会一路掉到 0：") print(" 梯度大到怎么缩都溢出，每一步都被跳过——这是训练已经坏掉的信号，不是原因（讲义 §4.3）。") print(" 3. bf16 的范围和 fp32 一样，基本不需要 loss scaling——这是 BLOOM 选 bf16 的主要原因。") # ============================================================================= # 演示 5：被吞掉的更新与 fp32 主参数 # ============================================================================= def train_with_storage(data, val, storage, master, steps=300, lr=3e-3): """ 从零训练 base model 结构 steps 步，模拟"权重以 storage 类型存储"。 【思路】storage = None：普通 fp32 训练。 storage = bf16/fp16、master = False：每次更新之后把权重舍入成 16 位——小于半个间距的更新被吞掉。 master = True：AdamW 在一份 fp32 的主参数上更新，前向用的是它舍入到 16 位的副本。 计算本身都是 fp32（这台 CPU 上真正的 16 位计算太慢），所以这里只模拟了"存储"的精度。 【返回】验证集 bit/字符 """ model = make_model(BASE_CONFIG, seed=0) params = list(model.parameters()) master_p = [p.detach().clone() for p in params] opt = torch.optim.AdamW(master_p if master else params, lr=lr, weight_decay=0.1, betas=(0.9, 0.95)) gen = torch.Generator().manual_seed(1) def rnd(t): return t if storage is None else t.to(storage).float() with torch.no_grad(): for p, mp in zip(params, master_p): p.copy_(rnd(mp)) for step in range(steps): for group in opt.param_groups: group["lr"] = lr_at(step, steps, lr) x, y = get_batch(data, 32, 96, gen) loss = F.cross_entropy(model(x).view(-1, 44), y.view(-1)) model.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(params, 1.0) if master: for p, mp in zip(params, master_p): mp.grad = p.grad.clone() opt.step() with torch.no_grad(): for p, mp in zip(params, master_p): p.copy_(rnd(mp)) else: opt.step() with torch.no_grad(): for p in params: p.copy_(rnd(p)) return bits_per_char(model, val) def demo5_master_weights(base, data, val): print_header("演示 5：权重直接存成 16 位，Adam 的更新有多少被吞掉；fp32 主参数的作用") w = torch.cat([p.detach().flatten() for p in base.parameters()]) print(f"base model 的权重：绝对值中位数 {w.abs().median().item():.3f}；LayerNorm 的增益参数初始化为 1.0") print(f"{pad('每一步的更新大小（≈ 学习率）', 30)}{pad('存成 bf16：被吞掉的比例', 24, '>')}{pad('存成 fp16：被吞掉的比例', 24, '>')}") for lr in [3e-3, 3e-4, 3e-5]: cells = [] for dt in (torch.bfloat16, torch.float16): wd = w.to(dt) lost = ((wd + torch.full_like(w, lr).to(dt)) == wd).float().mean().item() cells.append(pad(f"{lost:.1%}", 24, ">")) print(f"{pad(f'{lr:g}', 30)}" + "".join(cells)) print("（Adam 的每一步更新大约是 学习率 × 符号，所以'更新大小 ≈ 学习率'。）\n") print("从零训练 300 步（batch 32，学习率 3e-3 → 3e-4 余弦），验证集 bit/字符：") rows = [("fp32（参照）", None, False), ("权重存 bf16，没有主参数", torch.bfloat16, False), ("权重存 fp16，没有主参数", torch.float16, False), ("bf16 + fp32 主参数", torch.bfloat16, True)] for name, st, ms in rows: print(f" {pad(name, 30)}{train_with_storage(data, val, st, ms):.4f}") print("\n观察 Observations：") print(" 1. 学习率 3e-4 时，权重存成 bf16 会吞掉 12% 的更新；3e-5（微调常用的学习率）时吞掉 90%。") print(" fp16 的尾数多 3 位，吞得少得多，但在 3e-5 时也有约 40%。训练后期学习率衰减到很小，这个问题更严重。") print(" 2. 从零训练 300 步：权重存 bf16、没有主参数，比 fp32 差 0.12 bit/字符；加上 fp32 主参数就和 fp32 一样。") print(" 这就是混合精度的第二个补丁（Micikevicius et al. 2018）：前向、反向用 16 位，更新在 32 位的主参数上做——") print(" 代价是每个参数多存 4 字节（第 9 讲的 16 字节里就有这 4 字节）。") if __name__ == "__main__": world, tok = build_world(), CharTokenizer(ALPHABET) base, _ = load_or_train_base_model(world, tok) data = torch.tensor(tok.encode(world.train_text()), dtype=torch.long) val = torch.tensor(tok.encode(world.val_text()), dtype=torch.long) demo1_formats() demo2_swamping() demo3_loss_scale(base, data) demo4_dynamic_scaling(base, data) demo5_master_weights(base, data, val)