"""
=============================================================================
第 10 讲 · 代码 4：PaLM 的架构选择，一个一个地在小模型上试
=============================================================================

对应讲义：lectures/lecture10_training_in_practice.md 第 3.1 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
PaLM 在标准的 decoder-only Transformer 上做了几处改动：SwiGLU、并行层、多查询注意力（MQA）、RoPE、
去掉所有偏置（以及共享输入输出嵌入，第 1 讲的 base model 本来就共享）。论文给出的理由各不相同：
有的为了质量（SwiGLU），有的为了速度（并行层快约 15%），有的为了推理（MQA），有的为了稳定（去掉偏置）。

    演示 1：在第 1 讲的 base model 结构上（2 层，d = 64），每次只打开一个开关，再把它们全部打开，
            同样训练 400 步、2 个种子，比较参数个数、验证集 bit/字符、每步用时。
    演示 2：多查询注意力为什么对推理重要：PaLM 540B 生成时要缓存多少 K、V（KV cache）。

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

依赖本目录的 train_lab.py，以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py
（从随机初始化开始训练，不读也不写第 1 讲的缓存）。以 ⏱ 开头的行含计时，每次运行略有不同。
"""

import os
import statistics
import sys
import time

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, count_params, eval_bpc, make_flex, train_flex  # noqa: E402

STEPS, SEEDS = 400, (0, 1)
VARIANTS = [("基线（GELU、学出来的位置、带偏置）", {}),
            ("+ SwiGLU（隐藏维度 8d/3）", dict(mlp="swiglu")),
            ("+ 并行层", dict(parallel=True)),
            ("+ 多查询注意力 MQA", dict(mqa=True)),
            ("+ RoPE（去掉位置嵌入）", dict(pos="rope")),
            ("− 去掉所有偏置", dict(bias=False)),
            ("PaLM 式：以上全部", dict(mlp="swiglu", parallel=True, mqa=True, pos="rope", bias=False))]


def demo1_ablations(data, val):
    print_header(f"演示 1：一次只改一处（2 层，d = 64，batch 32，学习率 3e-3，{STEPS} 步，2 个种子）")
    print(f"{pad('配置', 36)}{pad('参数个数', 10, '>')}{pad('种子 0', 9, '>')}{pad('种子 1', 9, '>')}{pad('平均', 9, '>')}{pad('与基线之差', 12, '>')}")
    base_mean, timings = None, []
    for name, kw in VARIANTS:
        vals, secs = [], 0.0
        for seed in SEEDS:
            model = make_flex(FlexConfig(**kw), seed=seed)
            t = time.perf_counter()
            train_flex(model, data, STEPS, batch_size=32, lr=3e-3, seed=seed + 1, log_every=100)
            secs += time.perf_counter() - t
            vals.append(eval_bpc(model, val))
        mean = statistics.mean(vals)
        base_mean = mean if base_mean is None else base_mean
        delta = "—" if mean == base_mean and name.startswith("基线") else f"{mean - base_mean:+.3f}"
        print(f"{pad(name, 36)}{pad(f'{count_params(model):,}', 10, '>')}" + "".join(pad(f"{v:.3f}", 9, ">") for v in vals)
              + f"{pad(f'{mean:.3f}', 9, '>')}{pad(delta, 12, '>')}")
        timings.append((name, secs / (len(SEEDS) * STEPS)))
    print()
    for name, sec in timings:
        print(f"⏱ {pad(name, 36)}每步用时 {sec * 1e3:.1f} ms")
    print("\n观察 Observations：")
    print("  1. 两处改动在这个小模型上效果很大：RoPE（−0.405）和 SwiGLU（−0.319）。两个种子之间的差别在 0.01–0.12，")
    print("     它们带来的改进是这个噪声的好几倍。")
    print("     RoPE 的优势在这么短的训练里尤其明显：学出来的位置嵌入要从零学起，RoPE 一开始就把'相对位置'写进了注意力。")
    print("     训练更久，差距可能会缩小——这是一个短预算下的结果。")
    print("  2. MQA 的差别（−0.050）在种子噪声以内，参数少了约 11%——和 PaLM 的报告一致：对质量和训练速度'中性'，")
    print("     好处在推理（演示 2）。")
    print("  3. 并行层几乎没有差别（+0.014，在噪声以内），去掉偏置让模型变差（+0.110）。PaLM 的说法是：并行层在 8B 上有轻微的质量下降，")
    print("     到 62B 就看不出了，所以他们推断 540B 上是中性的，换来约 15% 的训练速度；去掉偏置是为了大模型的训练稳定性，")
    print("     不是为了质量。小模型上看到的'代价'，正是论文愿意付的价钱。")
    print("  4. 全部打开的 PaLM 式配置比基线好（参数还少了约 19%），但不如只加 RoPE：")
    print("     几处改动的好处和代价叠在一起，不是简单相加。每一个架构选择都要在目标规模上验证——这也是 PaLM 做消融的原因。")
    print("  5. 在这台 CPU 上，每步用时的差别很小（并行层的速度优势来自大模型里把两个矩阵乘法合并、减少通信，")
    print("     在 2 层 64 宽的模型上体现不出来）。")


def demo2_kv_cache():
    print_header("演示 2：生成时的 KV cache——PaLM 540B，多头 vs 多查询注意力（16 位）")
    L, H, dh = 118, 48, 256                    # PaLM 540B：118 层，48 个头，头维度 256
    per_token_mha = 2 * L * H * dh * 2         # K、V 各一份，每层每个头 dh 个数，2 字节
    per_token_mqa = 2 * L * 1 * dh * 2         # 所有头共享一组 K、V
    print(f"每个 token 要缓存的 K、V：多头注意力 {per_token_mha / 1e6:.2f} MB，多查询注意力 {per_token_mqa / 1e3:.0f} KB（{per_token_mha // per_token_mqa} 倍）")
    print(f"{pad('上下文 × batch', 16)}{pad('多头注意力', 14, '>')}{pad('多查询注意力', 14, '>')}")
    for ctx, b in [(2048, 1), (2048, 64), (8192, 64)]:
        print(f"{pad(f'{ctx} × {b}', 16)}{pad(f'{per_token_mha * ctx * b / 1e9:.1f} GB', 14, '>')}{pad(f'{per_token_mqa * ctx * b / 1e9:.2f} GB', 14, '>')}")
    print("（对照：PaLM 540B 的 16 位参数约 1.08 TB。）")
    print("\n观察 Observations：")
    print("  1. 生成时每一步都要读一遍整个 KV cache（第 8 讲 §3.3：生成受带宽限制）。多头注意力下，")
    print("     64 个 2048 token 的请求就要 760 GB 的 KV cache，比读参数还贵；MQA 把它缩小 48 倍。")
    print("  2. 所以 MQA 对训练几乎没影响（演示 1），却决定了能同时服务多少请求、每个 token 多快。")
    print("     后来的 GQA（分组查询注意力，第 3 讲 §9）在两者之间折中，Llama 2 70B 起被广泛采用。")


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)
    demo1_ablations(data, val)
    demo2_kv_cache()