"""
=============================================================================
第 10 讲 · 代码 2：训练不稳定——注意力 logit 失控、损失尖峰，以及 QK 归一化和 z-loss
=============================================================================

对应讲义：lectures/lecture10_training_in_practice.md 第 3.3、4.2 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
PaLM 540B 训练中出现了约 20 次损失尖峰；OPT-175B 的日志里反复记录"梯度范数和激活范数开始飙升、损失发散"。
这些问题通常只在最大的模型上出现，研究起来极其昂贵。Wortsman et al. (2023) 发现：**用很大的学习率训练小模型，
能复现同样的不稳定**——这就给了我们一个在笔记本上研究它的办法。

    演示 1：学习率扫描（4 层、d = 64 的小模型，400 步）：学习率越大，注意力 logit（q·k/√d 的最大值）越大；
            加上 QK 归一化之后呢？
    演示 2：学习率 0.1、不加 QK 归一化的那次训练，每 10 步的"体征"：损失、梯度范数、注意力 logit、log Z——
            看 logit 怎样一路涨上去、损失怎样停滞和跳起来。
    演示 3：学习率 0.1，3 个种子 × {什么都不加, QK 归一化, z-loss}：哪个办法稳定地有效。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_training_instability.py      # 约 100 秒（训练 13 个小模型，各 400 步）；需要 PyTorch

依赖本目录的 train_lab.py，以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py
（从随机初始化开始训练，不读也不写第 1 讲的缓存）。

【Python 小知识】
    - statistics.mean / min / max：对一组数求平均、最小、最大。
    - dict 的 setdefault(key, []).append(x)：键不存在时先放一个空列表，再往里加东西。
"""

import math
import os
import statistics
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  # noqa: E402
from train_lab import FlexConfig, eval_bpc, make_flex, train_flex  # noqa: E402

STEPS, BATCH, LAYERS = 400, 16, 4
VARIANTS = {"什么都不加": (dict(), 0.0), "QK 归一化": (dict(qk_norm=True), 0.0), "z-loss 1e-4": (dict(), 1e-4)}


def run(data, val, lr, variant, seed, log_every=20):
    """
    训练一个 4 层的小模型（不做梯度裁剪，好让不稳定显露出来），返回 (验证集 bit/字符, history)。

    【参数】lr：峰值学习率；variant：VARIANTS 的键；seed：模型初始化的种子（数据的种子 = seed + 1）
    【边界】训练发散（损失变成 nan）时 bit/字符记为 nan。
    """
    kw, zl = VARIANTS[variant]
    model = make_flex(FlexConfig(n_layer=LAYERS, **kw), seed=seed)
    hist, _ = train_flex(model, data, STEPS, batch_size=BATCH, lr=lr, seed=seed + 1, clip=None, z_loss=zl,
                         log_every=log_every)
    bpc = eval_bpc(model, val) if math.isfinite(hist[-1]["loss"]) else float("nan")
    return bpc, hist


def demo1_lr_sweep(data, val, cache):
    print_header("演示 1：学习率扫描（4 层，d = 64，batch 16，400 步，不裁剪梯度；种子 0）")
    print(f"{pad('学习率', 8)}{pad('', 14)}{pad('验证集 bit/字符', 18, '>')}{pad('训练中最大的注意力 logit', 26, '>')}{pad('最后的 log Z', 14, '>')}")
    for lr in [3e-3, 3e-2, 1e-1]:
        for variant in ("什么都不加", "QK 归一化"):
            bpc, hist = run(data, val, lr, variant, 0, log_every=10)
            cache[(lr, variant, 0)] = (bpc, hist)
            ml = max(r["max_logit"] for r in hist)
            log_z = f"{hist[-1]['log_z']:.2f}"          # 先算成字符串，避免 f-string 里嵌套同种引号
            print(f"{pad(f'{lr:g}', 8)}{pad(variant, 14)}{pad(f'{bpc:.3f}', 18, '>')}{pad(f'{ml:.1f}', 26, '>')}{pad(log_z, 14, '>')}")
    print("\n观察 Observations：")
    print("  1. 不加 QK 归一化时，注意力 logit 随学习率迅速变大：学习率 3e-3 时十几，0.1 时近千。")
    print("     logit 上千意味着 softmax 几乎变成了 one-hot：注意力只盯着一个位置，梯度几乎为 0，这一层'卡死'了。")
    print("     学习率 0.1 时模型明显变差。")
    print("  2. 加上 QK 归一化（先对每个头的 q、k 做 LayerNorm 再算点积），logit 被限制在十几以内，")
    print("     三个学习率下模型都正常——对学习率的敏感度小得多。这是 Wortsman et al. (2023) 的核心发现之一。")
    print("  3. log Z（softmax 分母的对数）在训练中一路涨到 10 左右：输出层的 logit 整体在漂移。")
    print("     这本身不影响损失（softmax 对整体平移不敏感），但 PaLM 认为它是另一种不稳定的来源，于是加了 z-loss（演示 3）。")


def demo2_trace(cache):
    print_header("演示 2：学习率 0.1、什么都不加——训练过程中的'体征'（每 10 步）")
    _, hist = cache[(0.1, "什么都不加", 0)]
    print(f"{pad('步', 6, '>')}{pad('训练损失', 10, '>')}{pad('梯度范数', 10, '>')}{pad('注意力 logit', 14, '>')}{pad('log Z', 8, '>')}")
    for r in hist:
        if r["step"] <= 200 or r["step"] % 50 == 0 or r["step"] == STEPS - 1:
            flag = "  ← 尖峰" if r["step"] > 0 and r["loss"] > min(h["loss"] for h in hist if h["step"] < r["step"]) + 0.3 else ""
            cells = [f"{r['loss']:.3f}", f"{r['grad_norm']:.2f}", f"{r['max_logit']:.1f}", f"{r['log_z']:.1f}"]
            print(f"{pad(r['step'], 6, '>')}{pad(cells[0], 10, '>')}{pad(cells[1], 10, '>')}{pad(cells[2], 14, '>')}{pad(cells[3], 8, '>')}{flag}")
    _, qk = cache[(0.1, "QK 归一化", 0)]
    print(f"\n对照：同样学习率、加 QK 归一化的那次训练，第 200 步的训练损失 {[r['loss'] for r in qk if r['step'] == 200][0]:.3f}，"
          f"第 399 步 {qk[-1]['loss']:.3f}")
    print("\n观察 Observations：")
    print("  1. 前 100 步：损失正常下降，但注意力 logit 一路从 0 涨到几百。")
    print("  2. 第 100–260 步：logit 到了近千，损失停在 2.0 左右下不去——注意力卡死了，模型学不动。")
    print("     其间出现损失尖峰（第 130–140 步，2.0 → 2.5），梯度范数同时跳高。")
    print("  3. 之后 logit 慢慢回落，损失才重新开始下降，但已经浪费了一大段训练，最后比加 QK 归一化的模型差得多。")
    print("     这和大模型训练日志里的描述很像：先是某个'体征'（激活范数、梯度范数）慢慢变坏，然后损失才出问题——")
    print("     所以 OPT 的团队盯着激活范数和梯度范数，而不只是损失。")


def demo3_seeds(data, val, cache):
    print_header("演示 3：学习率 0.1，3 个种子 × 三种做法（验证集 bit/字符，越低越好）")
    results = {}
    for variant in VARIANTS:
        for seed in range(3):
            if (0.1, variant, seed) not in cache:
                cache[(0.1, variant, seed)] = run(data, val, 0.1, variant, seed)
            results.setdefault(variant, []).append(cache[(0.1, variant, seed)][0])
    print(f"{pad('做法', 14)}{pad('种子 0', 10, '>')}{pad('种子 1', 10, '>')}{pad('种子 2', 10, '>')}{pad('平均', 10, '>')}{pad('最差', 10, '>')}")
    for variant, vals in results.items():
        print(f"{pad(variant, 14)}" + "".join(pad(f"{v:.3f}", 10, ">") for v in vals) +
              f"{pad(f'{statistics.mean(vals):.3f}', 10, '>')}{pad(f'{max(vals):.3f}', 10, '>')}")
    ref = cache[(3e-3, "什么都不加", 0)][0]
    print(f"（对照：学习率 3e-3、什么都不加，种子 0：{ref:.3f}）")
    print("\n观察 Observations：")
    print("  1. 什么都不加时，三个种子的结果差别很大——不稳定的训练像抽签。")
    print("  2. QK 归一化在每个种子上都是三者里最好的（平均 0.999），它对准了这里的病因（注意力 logit 失控）；")
    print("     但它也没有完全消除种子之间的差别（0.744–1.189）：学习率 0.1 对这个模型还是太大了。")
    print("  3. z-loss 在每个种子上都比什么都不加好一些（平均 1.761 → 1.434），但远不如 QK 归一化：")
    print("     它管的是**输出层**的 logit 漂移（log Z），不是注意力 logit。只看种子 0（0.808 vs QK 的 0.744）")
    print("     会以为两者差不多——这正是要多跑几个种子的原因。")
    print("  4. 小模型 + 大学习率只是大模型不稳定的一个'代理'：它能复现现象、比较对策，但大模型里尖峰的具体原因")
    print("     （某些数据、某种参数状态、数值精度）未必和这里相同。")


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)
    cache = {}
    demo1_lr_sweep(data, val, cache)
    demo2_trace(cache)
    demo3_seeds(data, val, cache)