""" ============================================================================= 第 3 讲 · 代码 6：真实模型的超参数，以及三个稳定性技巧 ============================================================================= 对应讲义：lectures/lecture03_architectures_hyperparameters.md 第 7、8 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 第一部分：把一批真实模型的配置放在一起，算出几个关键比例， 看看大家的选择是不是真的"惊人地一致"： d_ff / d_model 前馈层放大倍数（ReLU/GELU 约 4，GLU 类约 8/3） d_model / n_layers "宽深比"（大多在 100~130 左右） n_heads × d_head / d_model （大多等于 1） 第二部分：训练大模型时最常见的不稳定来源是 **softmax**（指数函数 + 除法）。 三个常用技巧分别针对不同的 softmax： z-loss → 输出层的 softmax（词表上的概率） QK-norm → 注意力里的 softmax logit 软截断 → 两者都可以用 用数值实验看它们各自在防什么。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 06_hyperparams_and_stability.py """ import math import torch import torch.nn.functional as F def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：真实模型的配置 # ============================================================================= # (名字, 层数, d_model, 查询头数, KV 头数, 头维度, d_ff, 词表, 前馈类型) # 数据来自各模型公开的论文或配置文件。 MODELS = [ ("GPT-2 XL", 48, 1600, 25, 25, 64, 6400, 50257, "GELU"), ("GPT-3 175B", 96, 12288, 96, 96, 128, 49152, 50257, "GELU"), ("T5 11B", 24, 1024, 128, 128, 128, 65536, 32128, "ReLU"), ("Llama 2 7B", 32, 4096, 32, 32, 128, 11008, 32000, "SwiGLU"), ("Llama 2 70B", 80, 8192, 64, 8, 128, 28672, 32000, "SwiGLU"), ("Mistral 7B", 32, 4096, 32, 8, 128, 14336, 32000, "SwiGLU"), ("Llama 3 8B", 32, 4096, 32, 8, 128, 14336, 128256, "SwiGLU"), ("Qwen2.5 7B", 28, 3584, 28, 4, 128, 18944, 152064, "SwiGLU"), ("Gemma 2 9B", 42, 3584, 16, 8, 256, 14336, 256000, "GeGLU"), ("OLMo 2 7B", 32, 4096, 32, 32, 128, 11008, 100278, "SwiGLU"), ] def demo_configs(): """演示 1：真实模型的超参数比例。""" print_header("演示 1：真实模型的超参数 —— 惊人地一致（除了几个例外）") print(f" {'模型':<12}{'前馈':>7}{'d_ff/d':>8}{'d/层数':>8}{'头×头维/d':>10}{'KV头':>6}{'词表':>9}") for name, L, d, h, kv, dh, dff, V, act in MODELS: print(f" {name:<12}{act:>7}{dff / d:>8.2f}{d / L:>8.1f}{h * dh / d:>10.2f}{kv:>4}/{h:<3}{V:>8,}") print(""" 读表： - d_ff/d：GELU/ReLU 模型基本都是 4；SwiGLU 模型在 8/3 ≈ 2.67 附近（Llama 2 7B 是 2.69）， 也有不少略大一些（3.5）。T5 11B 的 64 倍是著名的"异类"（它的后续版本 T5 v1.1 改回了常规的 2.5）。 - d/层数（宽深比）：大多在 100~130 之间；GPT-2 XL（33）和 T5 11B（43）偏"深而窄"。 讲义第 7.3 节：太深的模型难以并行、延迟高；在很宽的范围内效果差别不大，常由系统因素决定。 - 头数 × 头维度 / d：几乎都是 1；T5 11B（16）和 Gemma 2（1.14）是例外。 - KV 头：新模型普遍使用 GQA（KV 头少于查询头），为推理省 KV 缓存。 - 词表：单语（英文为主）3~5 万；多语言 / 生产系统 10 万~25 万。""") # ============================================================================= # 第 2 部分：z-loss —— 输出 softmax 的稳定性 # ============================================================================= def z_loss(logits, coef=1e-4): """ z-loss（PaLM 使用；Baichuan 2、DCLM、OLMo 2 等也用）： z = log Σ_j exp(logit_j) （softmax 的"归一化常数"取对数） z_loss = coef × z² 【为什么需要】 softmax 对所有 logits 加同一个常数不变：softmax(o + c) = softmax(o)。 所以交叉熵损失对"logits 整体往上漂"毫无约束，训练中 logits 可能越漂越大。 数值很大时，低精度（BF16）的相对精度不变、**绝对**精度变差：1000 附近的 BF16 间隔是 4！ logits 之间原本 0.5 的差距可能被舍入成 0 或 4，概率就全乱了。 z-loss 惩罚 z 偏离 0，把 logits 的整体水平"拴"在 0 附近。 【参数】logits : (..., V) """ z = torch.logsumexp(logits.float(), dim=-1) return coef * (z ** 2).mean() def demo_z_loss(): """演示 2：logits 整体漂移时，BF16 下的交叉熵误差。""" print_header("演示 2：z-loss —— 为什么要把 logits 拴在 0 附近") torch.manual_seed(0) V = 1000 logits = torch.randn(64, V) * 2 targets = torch.randint(0, V, (64,)) ref = F.cross_entropy(logits, targets).item() print(f" {'整体偏移 c':>10}{'FP32 交叉熵':>14}{'BF16 logits 交叉熵':>20}{'z = logsumexp':>16}{'z-loss(1e-4)':>14}") for c in [0, 10, 100, 1000, 10000]: shifted = logits + c ce32 = F.cross_entropy(shifted, targets).item() ce16 = F.cross_entropy(shifted.to(torch.bfloat16).float(), targets).item() z = torch.logsumexp(shifted, -1).mean().item() print(f" {c:>10}{ce32:>14.4f}{ce16:>20.4f}{z:>16.1f}{z_loss(shifted).item():>14.2f}") print(f""" - FP32 下交叉熵完全不受偏移影响（都等于 {ref:.4f}）：数学上 softmax 对整体平移不变。 - 但如果 logits 用 BF16 存储，偏移越大误差越大：c = 1000 时相邻可表示数的间隔已经是 4， logits 之间的细微差别被抹掉，损失严重偏离。 - z-loss 随偏移迅速变大，在训练中会把 logits 的整体水平往 0 拉，从而避免这种情况。""") # ============================================================================= # 第 3 部分：QK-norm —— 注意力 softmax 的稳定性 # ============================================================================= def rms_normalize(x, eps=1e-6): """不带可学习参数的 RMSNorm（演示用）：x / RMS(x)。""" return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps) def attention_stats(q, k): """ 计算注意力权重的平均熵和平均最大权重。 【熵】H = -Σ p log p。均匀分布时熵最大（log n），全部集中在一个位置时熵为 0。 """ d = q.shape[-1] p = torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(d), dim=-1) entropy = -(p * torch.log(p + 1e-12)).sum(-1).mean().item() return entropy, p.max(-1).values.mean().item() def demo_qk_norm(): """演示 3：q、k 的范数增大时，注意力会"塌缩"；QK-norm 可以防止。""" print_header("演示 3：QK-norm —— 防止注意力 softmax 塌缩") torch.manual_seed(0) n, d = 64, 64 q, k = torch.randn(n, d), torch.randn(n, d) print(f" {n} 个位置（均匀分布的熵 = ln {n} = {math.log(n):.2f}）") print(f" {'q、k 放大倍数 s':>16}{'普通：熵':>10}{'普通：最大权重':>14}{'QK-norm：熵':>14}{'QK-norm：最大权重':>18}") for s in [1, 2, 4, 8]: e1, m1 = attention_stats(q * s, k * s) e2, m2 = attention_stats(rms_normalize(q * s), rms_normalize(k * s)) print(f" {s:>16}{e1:>10.2f}{m1:>14.2f}{e2:>14.2f}{m2:>18.2f}") print(""" 注意力分数 = q·k/√d，q 和 k 同时放大 s 倍，分数就放大 s² 倍。 训练中如果 q、k 的范数不断增长，softmax 就越来越"尖"（熵趋于 0）：每个位置只盯着一个位置看， 梯度也几乎消失 —— 这是大模型训练中损失突然飙升（loss spike）的常见原因之一。 QK-norm 在计算点积前先对 q、k 做归一化（真实模型用带可学习增益的 RMSNorm），分数的大小就被控制住了， 与 q、k 本身有多大无关。OLMo 2、Gemma 3、Qwen3 等模型都使用它 （Gemma 2 用的是下一个演示的软截断，Gemma 3 改用了 QK-norm）。""") # ============================================================================= # 第 4 部分：logit 软截断 # ============================================================================= def soft_cap(logits, cap): """ logit 软截断（Gemma 2 使用）：cap × tanh(logits / cap)。 【思路】 tanh 把任意实数压到 (-1, 1)。乘以 cap 之后，输出被限制在 (-cap, cap) 之内； 而在 |logits| 远小于 cap 时，tanh(x) ≈ x，几乎不改变原值。 所以它像一个"软"的天花板：小的 logits 原样通过，大的 logits 被平滑地压住。 Gemma 2 对注意力分数用 cap = 50，对最终输出 logits 用 cap = 30。 【边界】截断会损失一部分表达能力（真的需要很大的 logit 时也被压住了），可能影响效果。 """ return cap * torch.tanh(logits / cap) def demo_soft_cap(): """演示 4：软截断的效果。""" print_header("演示 4：logit 软截断") x = torch.tensor([-200.0, -50.0, -10.0, 0.0, 5.0, 10.0, 30.0, 50.0, 200.0]) y = soft_cap(x, 30.0) for a, b in zip(x.tolist(), y.tolist()): print(f" {a:>8.1f} → {b:>7.2f}") print(" 小的值几乎不变（10 → 9.6），大的值被压到 ±30 以内（200 → 30.0）。") def main(): demo_configs() demo_z_loss() demo_qk_norm() demo_soft_cap() print("\n全部演示运行完毕。") if __name__ == "__main__": main()