"""
=============================================================================
第 4 讲 · 代码 4：IsoFLOP 曲线与计算最优的模型大小；Kaplan 和 Chinchilla 为什么得出不同的结论
=============================================================================

对应讲义：lectures/lecture04_scaling_laws.md 第 5、6 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
给定一笔算力 C，该训练多大的模型、读多少数据？Kaplan et al. (2020) 和 Chinchilla (2022) 给出了很不一样的答案。
这个文件在合成语言上走一遍 Chinchilla 的三种方法，再检查两个被认为造成分歧的原因：

    演示 1：模型的"形状"重要吗？参数量差不多、层数和宽度不同的 4 个模型（Kaplan：在很宽的范围内影响很小）。
    演示 2：IsoFLOP 曲线——固定算力，横轴是模型大小，损失是一个"山谷"；谷底随算力移动，
            得到 N_opt ∝ C^a。和方法 1（每个算力下哪个宽度最好）、方法 3（参数化拟合）对照。
    演示 3：学习率调度没有和训练长度匹配会怎样？"一个长训练的中途检查点" vs "一个就训练这么长的模型"。
            用中途检查点拟合出的规模定律变成了什么样子（Chinchilla 对 Kaplan 的主要批评）。
    演示 4：参数怎么数：算不算嵌入参数，会改变拟合出的指数（Pearce & Song 2024 对分歧的解释）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 04_isoflop_and_schedules.py     # 需要代码 02 的缓存；第一次再训练 8 个模型（约 1.5 分钟），之后约 3 秒

需要 numpy、PyTorch；依赖本目录的 scaling_tools.py 和第 1 讲的 console.py（只读）。
如果代码 02 还没运行过，这里会先训练主网格（再加约 2 分钟）。
"""

import math
import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from console import pad, print_header                                          # noqa: E402
from scaling_tools import (DATA_SIZES, TRUNC_EVALS, WIDTHS, build_language, compute_optimal, ensure_runs,  # noqa: E402
                           fit_joint, grid_specs, grid_table, interp_loglog, parabola_min)

# 为了让 IsoFLOP 曲线的"小模型一侧"也有点，几个窄模型要训练得更久（4M、8M token）
EXTRA_SPECS = [dict(d_model=8, n_tokens=4_000_000), dict(d_model=8, n_tokens=8_000_000),
               dict(d_model=16, n_tokens=4_000_000), dict(d_model=24, n_tokens=4_000_000),
               dict(d_model=32, n_tokens=4_000_000)]
# 形状实验：非嵌入参数都在 2.5 万–2.9 万之间，层数 1、2、4、8
SHAPE_SPECS = [dict(d_model=48, n_tokens=1_000_000, n_layer=1), dict(d_model=24, n_tokens=1_000_000, n_layer=4),
               dict(d_model=16, n_tokens=1_000_000, n_layer=8)]
D_MIN = 500_000          # 只用 D ≥ 0.5M 的点（D = 0.25M 属于"小数据区"，见代码 02）
BUDGETS = [1e10 * 2 ** (i / 2) for i in range(14)]     # 1e10 到约 9e11，每档 × √2


def curves_from(runs):
    """所有"调度和训练长度匹配"的 2 层模型：{宽度: {D: 损失}}，以及 {宽度: (总参数, 非嵌入参数)}。"""
    curves, sizes = {}, {}
    for r in runs.values():
        if r["n_layer"] == 2 and r["seed"] == 0 and r["schedule"] == r["D"]:
            curves.setdefault(r["d_model"], {})[r["D"]] = r["loss"]
            sizes[r["d_model"]] = (r["n_params"], r["n_nonemb"])
    return curves, sizes


def isoflop(curves, sizes, count="total", verbose=False):
    """
    对每个算力预算 C，在每个宽度的"损失 vs D"曲线上读出 D = C / (6N) 处的损失（对数坐标插值，不外推），
    得到一条 IsoFLOP 曲线；用抛物线找谷底（Chinchilla 方法 2），同时记下网格里最好的宽度（方法 1 的粗糙版本）。

    【参数】count："total" 用总参数量记 N 和 C；"nonemb" 用非嵌入参数量（Kaplan 的数法）
    【返回】list[dict(C, points, n_opt, inside, best_d)]，只保留至少有 3 个点的预算
    """
    k = 0 if count == "total" else 1
    out = []
    for C in BUDGETS:
        pts = []
        for d in sorted(curves):
            Ds = sorted(x for x in curves[d] if x >= D_MIN)
            N = sizes[d][k]
            loss = interp_loglog(C / (6 * N), Ds, [curves[d][x] for x in Ds])
            if loss is not None:
                pts.append((d, N, C / (6 * N), loss))
        if len(pts) < 3:
            continue
        n_opt, l_min, inside = parabola_min([math.log(p[1]) for p in pts], [p[3] for p in pts])
        best = min(pts, key=lambda p: p[3])
        out.append(dict(C=C, points=pts, n_opt=n_opt, l_min=l_min, inside=inside, best_d=best[0], best_N=best[1]))
    return out


def slope(xs, ys):
    """log-log 直线拟合的斜率。"""
    return float(np.polyfit(np.log(xs), np.log(ys), 1)[0])


# =============================================================================
# 演示 1：形状
# =============================================================================
def demo_shape(runs, table):
    print_header("演示 1：参数量差不多时，模型的'形状'（层数 vs 宽度）重要吗？（D = 1M）")
    rows = []
    for key in ("L1_d48_D1000000_S1000000_seed0", "L2_d32_D1000000_S1000000_seed0",
                "L4_d24_D1000000_S1000000_seed0", "L8_d16_D1000000_S1000000_seed0"):
        r = runs[key]
        rows.append((r["n_layer"], r["d_model"], r["n_nonemb"], r["n_params"], r["loss"]))
    print("  " + pad("层数", 6) + pad("宽度", 6) + pad("宽度/层数", 10) + pad("非嵌入参数", 12) + pad("总参数", 10) + "损失")
    for L, d, ne, tot, loss in rows:
        print("  " + pad(L, 6) + pad(d, 6) + pad(f"{d / L:g}", 10) + pad(f"{ne:,}", 12) + pad(f"{tot:,}", 10)
              + f"{loss:.4f}")
    spread = max(r[4] for r in rows) - min(r[4] for r in rows)
    print(f"""
观察 Observations：
  · 在我们的玩具里，形状的影响**不小**：同样约 2.7 万个非嵌入参数，最好和最差差了 {spread:.2f} bit/字符，
    越深越窄越差。Kaplan et al. 的结论是"在很宽的范围内，形状只影响几个百分点"——
    但他们比较的是宽度/层数在几十到几百之间的模型；我们这里的比值只有 2–48，学习率只按宽度调过，
    而且每个模型都只训练了 977 步（越深的网络往往越需要更多步数和更小心的学习率）。
    极端的形状、没调好的超参数、太短的训练，都会让"形状不重要"失效。
  · 这也说明了为什么代码 02 只改宽度、层数固定为 2：扫描规模时要一次只变一个东西，
    否则"大模型更好"里混进了"形状更好"或"更难优化"。""")


# =============================================================================
# 演示 2：IsoFLOP
# =============================================================================
def demo_isoflop(curves, sizes):
    print_header("演示 2：IsoFLOP 曲线 —— 固定算力，模型太小和太大都不好")
    res = isoflop(curves, sizes)
    for r in res:
        pts = "  ".join(f"d{p[0]}:{p[3]:.3f}" for p in r["points"])
        mark = "谷底在采样范围内" if r["inside"] else "谷底不在采样范围内，不用"
        n_txt = f"{r['n_opt']:,.0f}" if r["n_opt"] == r["n_opt"] else "—"
        print(f"  C = {r['C']:.2e}   {pts}")
        print(f"      抛物线谷底 N_opt = {n_txt}（{mark}）；网格里最好的宽度 d = {r['best_d']}")
    ok = [r for r in res if r["inside"]]
    a2 = slope([r["C"] for r in ok], [r["n_opt"] for r in ok])
    a1 = slope([r["C"] for r in res], [r["best_N"] for r in res])
    print("\n  计算最优点的 token/参数（方法 2）：" + "  ".join(
        f"C={r['C']:.1e}: {r['C'] / 6 / r['n_opt'] ** 2:.0f}" for r in ok))

    # 方法 3：用所有匹配调度的点（包括多训练的窄模型）联合拟合
    N, D, L = [], [], []
    for d, cv in curves.items():
        for x, loss in cv.items():
            if x >= D_MIN:
                N.append(sizes[d][0])
                D.append(x)
                L.append(loss)
    p = fit_joint(N, D, L)
    a3 = p["beta"] / (p["alpha"] + p["beta"])
    print(f"\n  方法 3（{len(L)} 个点联合拟合）：E = {p['E']:.3f}，α = {p['alpha']:.3f}，β = {p['beta']:.3f}；"
          f"在 C = 1e11 时 N_opt = {compute_optimal(p, 1e11)[0]:,.0f}")
    print("\n  N_opt ∝ C^a 的指数 a：")
    print("  " + pad("", 34) + pad("我们的玩具", 12) + "Chinchilla 论文（表 2）")
    print("  " + pad("方法 1（每个算力下最好的模型）", 34) + pad(f"{a1:.2f}", 12) + "0.50")
    print("  " + pad("方法 2（IsoFLOP 抛物线谷底）", 34) + pad(f"{a2:.2f}", 12) + "0.49")
    print("  " + pad("方法 3（参数化拟合）", 34) + pad(f"{a3:.2f}", 12) + "0.46")
    print("  " + pad("（Kaplan et al. 2020）", 34) + pad("", 12) + "0.73")
    print(f"""
观察 Observations：
  · 每条 IsoFLOP 曲线都是一个山谷：固定算力，太小的模型（d = 8、16）读了很多数据也学不会——容量不够；
    太大的模型（d = 48、64）只读了很少的数据——步数不够。谷底就是这个算力下"计算最优"的模型。
  · 谷底随算力右移：最好的宽度从 d = 16 移到 d = 32、再到 d = 48。但移动是"一步一步跳"的：前三个预算的
    谷底都在约 1 万参数，之后跳到 2.5 万、3.2 万、4.6 万、5.5 万——每条曲线只有 3–4 个点、相邻宽度的参数量
    差 2 倍左右，谷底的位置很粗糙。三种方法给出的指数在 0.64–0.99 之间，彼此差得比 Chinchilla 的三种方法
    （0.46–0.50）大得多；方法 1 只能在 6 个宽度里挑，最粗糙，抛物线谷底（方法 2）平滑一些。
  · 计算最优点的 token/参数在 35–177 之间，最大的三个预算上是 36–61，比真实 LLM 的约 20 大；
    方法 2、3 都说参数分到的份额 a > 0.5（代码 03 演示 2 的解释：这个玩具里数据的回报衰减得快）。
    数字和 Chinchilla 不同，是两个世界的指数不同，不是方法的问题。""")
    return res


# =============================================================================
# 演示 3：学习率调度
# =============================================================================
def demo_schedule(runs, table, sizes):
    print_header("演示 3：学习率调度 —— '长训练的中途检查点' ≠ '就训练这么长的模型'")
    trunc = {}
    for d in WIDTHS:
        r = runs[f"L2_d{d}_D2000000_S2000000_seed0"]
        trunc[d] = {int(k): v for k, v in r["evals"].items()}
        trunc[d][DATA_SIZES[-1]] = r["loss"]
    cols = [x for x in TRUNC_EVALS if x in DATA_SIZES]
    print("每个宽度都有一次按 2M token 排学习率调度的训练；在它读到 0.25M / 0.5M / 1M 时停下来评测，")
    print("和'调度长度就等于 0.25M / 0.5M / 1M'的训练比较。表中是 中途检查点 − 匹配调度（正数 = 中途检查点更差）：")
    print("  " + pad("宽度 d", 8) + "".join(pad(f"D={x / 1e6:g}M（走了 {x / 2e6:.0%}）", 22) for x in cols))
    diffs = {x: [] for x in cols}
    for d in WIDTHS:
        row = []
        for x in cols:
            diffs[x].append(trunc[d][x] - table[d][x])
            row.append(pad(f"{trunc[d][x] - table[d][x]:+.4f}", 22))
        print("  " + pad(d, 8) + "".join(row))
    print("  " + pad("平均", 8) + "".join(pad(f"{np.mean(diffs[x]):+.4f}", 22) for x in cols))

    def fit_on(source):
        N, D, L = [], [], []
        for d in WIDTHS:
            for x in DATA_SIZES:
                if x >= D_MIN:
                    N.append(sizes[d][0])
                    D.append(x)
                    L.append(source[d][x])
        return fit_joint(N, D, L)

    pm, pt = fit_on(table), fit_on(trunc)
    print("\n用同样的 (N, D) 格子（D ≥ 0.5M）拟合 L = E + A/N^α + B/D^β：")
    print("  " + pad("数据来源", 26) + pad("E", 8) + pad("α", 8) + pad("β", 8) + "a = β/(α+β)")
    for name, p in (("匹配调度（Chinchilla 式）", pm), ("中途检查点（Kaplan 式）", pt)):
        print("  " + pad(name, 26) + pad(f"{p['E']:.3f}", 8) + pad(f"{p['alpha']:.3f}", 8) + pad(f"{p['beta']:.3f}", 8)
              + f"{p['beta'] / (p['alpha'] + p['beta']):.3f}")
    print(f"  （真正的熵率 E = {build_language().entropy_rate:.4f}）")
    print("""
观察 Observations：
  · 走到一半（1M）和四分之一（0.5M）时，中途检查点平均差 0.1 bit 以上（唯一的例外 d = 16、0.5M，
    恰好是代码 02 里种子噪声最大的那一次训练）：学习率还停在高位，模型没有经过
    "退火"（学习率降下来时损失会再降一截）。Chinchilla 的附录图 A1 显示了同样的现象：余弦周期比训练长 25% 以上，
    最终损失就明显变差。
  · 走到八分之一（0.25M）时反而是中途检查点更好：匹配调度的训练只有 244 步，学习率很快就降下来了，
    还没学会多少就"减速"了。这是我们的玩具步数太少造成的，真实的大训练不会这么短。
  · 用中途检查点拟合，得到的是另一条"规模定律"：E 掉到 0，数据指数 β 小了很多，由它推出的分配指数 a 也完全不同。
    误差不是一个常数偏移，而是随"走了多少"变化，所以它**改变了曲线的形状**——这正是 Chinchilla 对 Kaplan 的批评：
    Kaplan 的所有模型用同一个长调度、在中途读损失。Chinchilla 认为这让 Kaplan 低估了"少读数据"的模型、
    高估了大模型；在我们的玩具里偏差的方向不同（a 变小了），但"不匹配调度会系统地歪曲拟合"这一点是一样的。
  · 2024 年的复现（Porian et al.）发现：在仔细控制其他因素后，学习率衰减本身并不是分歧的主要原因；
    更重要的是最后一层的计算量怎么算、热身（warmup）的长度、以及优化器超参数有没有随规模重新调。""")


# =============================================================================
# 演示 4：参数怎么数
# =============================================================================
def demo_counting(curves, sizes):
    print_header("演示 4：参数怎么数 —— 算不算嵌入参数，IsoFLOP 推出的指数不同")
    print("  " + pad("宽度 d", 8) + pad("总参数", 10) + pad("非嵌入参数", 12) + "嵌入所占比例")
    for d in WIDTHS:
        tot, ne = sizes[d]
        print("  " + pad(d, 8) + pad(f"{tot:,}", 10) + pad(f"{ne:,}", 12) + f"{1 - ne / tot:.1%}")
    out = {}
    for count, name in (("total", "总参数（Chinchilla 的数法）"), ("nonemb", "非嵌入参数（Kaplan 的数法）")):
        res = [r for r in isoflop(curves, sizes, count) if r["inside"]]
        out[count] = slope([r["C"] for r in res], [r["n_opt"] for r in res])
        print(f"  用{name}记 N 和 C = 6ND：IsoFLOP 谷底 N_opt ∝ C^{out[count]:.2f}（{len(res)} 个预算）")
    print(f"""
观察 Observations：
  · 嵌入参数（字符嵌入 + 位置嵌入）在最窄的模型里接近两成，在最宽的模型里只占 3%。
    不算它们，小模型就显得"更小"，参数量的跨度被拉大，谷底随算力移动得更快：指数从 {out['total']:.2f} 变成 {out['nonemb']:.2f}。
  · Pearce & Song (2024) 认为这正是 Kaplan 和 Chinchilla 分歧的一大来源：Kaplan 用非嵌入参数、
    而且模型偏小（多数在 1 亿参数以下），在小规模上嵌入占的比例大；他们模拟后发现，在这些条件下
    会得到接近 Kaplan 的 0.73。他们建议今后的规模定律研究一律用总参数和总算力。""")


def main():
    specs = grid_specs() + EXTRA_SPECS + SHAPE_SPECS
    runs = ensure_runs(specs)
    table, sizes = grid_table(runs)
    curves, sizes2 = curves_from(runs)
    demo_shape(runs, table)
    demo_isoflop(curves, sizes2)
    demo_schedule(runs, table, sizes)
    demo_counting(curves, sizes2)


if __name__ == "__main__":
    main()