"""
=============================================================================
第 4 讲 · 代码 1：幂律为什么在对数坐标下是直线、为什么会出现幂律，以及 C ≈ 6ND 从哪里来
=============================================================================

对应讲义：lectures/lecture04_scaling_laws.md 第 2、3 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
规模定律的三个基本零件：幂律、不可约损失、算力的记账方式。这个文件不训练任何模型，几秒钟跑完。

    演示 1：幂律 y = A·x^(−α) 在对数坐标下是一条直线；指数 α = 0.076 意味着什么（每翻一倍降多少）；
            加上不可约项 E 之后曲线为什么会"弯"，减掉 E 又为什么会变直（用 Chinchilla 发表的拟合式）。
    演示 2：为什么会出现幂律？三个最小的例子：
            (a) 估计均值：误差 ∝ 1/n（CS324 课程笔记里的例子）；
            (b) 最近邻回归：误差 ∝ n^(−2/d)，维数 d 越高指数越小（"内在维数"的解释）；
            (c) Zipf 分布的"知识块"：学会前 n 块，剩下的概率质量 ∝ n^(−α)（我们合成语言的词表就是 Zipf 的）。
    演示 3：C ≈ 6ND：一个线性层前向 2 次、反向 4 次浮点运算 / 参数 / token；
            用 PyTorch 的 FLOP 计数器实测 TinyGPT，看 6N 在什么时候准、什么时候不准；核对 GPT-3 的算力。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 01_power_laws_and_compute.py      # 约 5 秒

需要 numpy、PyTorch（只用来数 FLOPs）；依赖第 1 讲的 tiny_gpt.py、console.py（只读）。
"""

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"))

import torch                                                        # noqa: E402
import torch.nn as nn                                               # noqa: E402
import torch.nn.functional as F                                     # noqa: E402
from torch.utils.flop_counter import FlopCounterMode                # noqa: E402

from console import pad, print_header                               # noqa: E402
from tiny_gpt import GPTConfig, count_parameters, make_model        # noqa: E402

# Chinchilla（Hoffmann et al., 2022）发表的参数化拟合：L(N, D) = E + A/N^α + B/D^β（单位 nats/token）
CHIN = dict(E=1.69, A=406.4, alpha=0.34, B=410.7, beta=0.28)
KAPLAN_ALPHA_N = 0.076          # Kaplan et al. (2020) 式 (1.1)：L(N) = (N_c / N)^0.076


# =============================================================================
# 演示 1：幂律与对数坐标
# =============================================================================
def demo_power_law():
    print_header("演示 1：幂律在对数坐标下是直线；加上不可约项 E 就会弯")
    a = KAPLAN_ALPHA_N
    print(f"(a) 纯幂律 L = A·N^(−α)。取 log：log L = log A − α·log N，是一条斜率为 −α 的直线。")
    print(f"    Kaplan et al. 测得的参数指数 α_N ≈ {a}。这么小的指数意味着：")
    print(f"      N 翻一倍，损失乘以 2^(−{a}) = {2 ** -a:.4f}（降 {1 - 2 ** -a:.1%}）")
    print(f"      N 乘以 10，损失乘以 10^(−{a}) = {10 ** -a:.4f}（降 {1 - 10 ** -a:.1%}）")
    print(f"      要让损失减半，N 要乘以 2^(1/{a}) = {2 ** (1 / a):.2e}")
    print("    对照：α = 0.5 时，N 乘以 10 损失乘以 {:.3f}；α = 1（估计均值）时乘以 0.1。".format(10 ** -0.5))

    E, A, al = CHIN["E"], CHIN["A"], CHIN["alpha"]
    print(f"\n(b) 加上不可约项：L(N) = E + A·N^(−α)，用 Chinchilla 拟合式里的参数项（E = {E}, A = {A}, α = {al}）")
    print("    " + pad("N", 10) + pad("L(N)", 9) + pad("L − E", 9) + pad("L 的局部斜率", 15) + "L − E 的局部斜率")
    Ns = [1e6, 1e7, 1e8, 1e9, 1e10, 1e11, 1e12]
    Ls = [E + A * n ** (-al) for n in Ns]
    for i, n in enumerate(Ns):
        if i + 1 < len(Ns):
            s1 = -(math.log(Ls[i + 1]) - math.log(Ls[i])) / math.log(10)
            s2 = -(math.log(Ls[i + 1] - E) - math.log(Ls[i] - E)) / math.log(10)
            tail = pad(f"{s1:.4f}", 15) + f"{s2:.4f}"
        else:
            tail = ""
        print("    " + pad(f"{n:.0e}", 10) + pad(f"{Ls[i]:.4f}", 9) + pad(f"{Ls[i] - E:.4f}", 9) + tail)
    print("""
观察 Observations：
  · 纯幂律的"每翻一倍降几个百分点"是**固定**的——不管从 1 亿还是 1000 亿参数出发。这就是"可预测"的来源：
    在小模型上量出斜率，就能推出大模型。
  · 指数很小（0.076）意味着规模的回报很"贵"：只靠加参数，损失减半要把参数乘以约 9000 倍。真实的进步来自
    同时增加 N、D、C，以及不断改进数据和架构（它们改变的是曲线的位置，而不是斜率——第 8 节）。
  · 一旦有不可约项 E，直接看 L 的局部斜率会越来越小（从 0.2 多降到 0.01 以下），看上去像"规模失效了"；
    其实"还能减掉的部分" L − E 一直在按同一个指数 0.34 下降。画图前先想清楚：纵轴是 L 还是 L − E？""")


# =============================================================================
# 演示 2：为什么会出现幂律
# =============================================================================
def fitted_slope(xs, ys):
    """log-log 坐标下最小二乘直线的斜率。"""
    return float(np.polyfit(np.log(xs), np.log(ys), 1)[0])


def mean_estimation(ns, reps, rng):
    """
    用 n 个样本的平均值估计一个分布的均值，均方误差随 n 怎样变化。

    【思路】样本来自 N(μ, 1)。统计学告诉我们 E[(x̄ − μ)²] = σ²/n——这本身就是一条斜率 −1 的规模定律。
    【返回】list[float]，每个 n 的均方误差（重复 reps 次取平均）
    """
    out = []
    for n in ns:
        samples = rng.normal(loc=3.0, scale=1.0, size=(reps, n))
        out.append(float(((samples.mean(axis=1) - 3.0) ** 2).mean()))
    return out


def nearest_neighbor_mse(d, ns, rng, n_test=500):
    """
    d 维单位立方体里的 1-最近邻回归：用离测试点最近的训练点的函数值当预测，均方误差随训练集大小 n 怎样变化。

    【思路】n 个点均匀撒在 d 维空间里，最近邻的距离大约 ∝ n^(−1/d)；函数是光滑的，误差 ∝ 距离，
            平方误差 ∝ n^(−2/d)。维数越高，同样多的数据把空间"填"得越稀，指数越小——
            CS324 课程笔记用这个例子说明：神经网络的数据指数小，可能是因为数据的"内在维数"高。
    【参数】d：维数；ns：训练集大小的列表
    【返回】list[float]
    【numpy 小知识】‖a − b‖² = ‖a‖² + ‖b‖² − 2 a·b：用矩阵乘法一次算出所有测试点到所有训练点的距离。
    """
    def f(x):
        return np.sin(2 * np.pi * x).sum(axis=1)          # 一个光滑的目标函数

    test = rng.random((n_test, d))
    out = []
    for n in ns:
        train = rng.random((n, d))
        d2 = (test ** 2).sum(1)[:, None] + (train ** 2).sum(1)[None, :] - 2 * test @ train.T
        pred = f(train)[d2.argmin(axis=1)]
        out.append(float(((pred - f(test)) ** 2).mean()))
    return out


def zipf_pieces(alpha, n_items=10 ** 6):
    """
    "知识块"模型：世界上有很多块知识，第 k 块被用到的概率 p_k ∝ k^(−(1+α))（Zipf 分布）。

    【思路】
        容量受限：一个"能记住 n 块"的模型记住最常用的 n 块，其余的都会答错——损失的"可减部分"∝ 剩下的
                  概率质量 Σ_{k>n} p_k ≈ n^(−α)。
        数据受限：看了 D 个样本，一块知识至少出现一次才学得会；第 k 块一次都没出现的概率是 (1 − p_k)^D。
                  可减部分 ∝ Σ_k p_k (1 − p_k)^D ≈ D^(−α/(1+α))。
        这是 Michaud et al. (2023)"量子化模型"和 Hutter (2021) 的思路。我们的合成语言里，词表就是 Zipf 分布的：
        罕见的词、罕见的搭配，就是这里的"尾部知识块"。
    【返回】(容量曲线 [(n, 剩余质量)], 数据曲线 [(D, 期望剩余质量)])
    """
    k = np.arange(1, n_items + 1, dtype=np.float64)
    p = k ** (-(1 + alpha))
    p /= p.sum()
    tail = 1.0 - np.cumsum(p)                       # tail[n−1] = Σ_{k>n} p_k
    cap = [(n, float(tail[n - 1])) for n in (10, 100, 1000, 10000)]
    dat = [(D, float((p * np.exp(D * np.log1p(-p))).sum())) for D in (1e3, 1e4, 1e5, 1e6)]
    return cap, dat


def demo_why_power_laws():
    print_header("演示 2：为什么会出现幂律？三个最小的例子")
    rng = np.random.default_rng(0)
    ns = [4, 16, 64, 256, 1024]
    mse = mean_estimation(ns, reps=4000, rng=rng)
    print("(a) 估计均值（CS324 课程笔记的例子）：均方误差 vs 样本数 n")
    print("    " + "  ".join(f"n={n}: {m:.5f}" for n, m in zip(ns, mse)))
    print(f"    对数坐标下的斜率 = {fitted_slope(ns, mse):.3f}（理论值 −1：误差 = σ²/n）")

    print("\n(b) d 维空间里的 1-最近邻回归：均方误差 vs 训练集大小 n")
    ns2 = [64, 256, 1024, 4096]
    print("    " + pad("维数 d", 8) + "".join(pad(f"n={n}", 11) for n in ns2) + pad("拟合斜率", 10) + "理论 −2/d")
    for d in (1, 2, 4, 8):
        m = nearest_neighbor_mse(d, ns2, np.random.default_rng(d))
        print("    " + pad(d, 8) + "".join(pad(f"{v:.2e}", 11) for v in m) + pad(f"{fitted_slope(ns2, m):.3f}", 10)
              + f"{-2 / d:.3f}")

    alpha = 0.5
    cap, dat = zipf_pieces(alpha)
    print(f"\n(c) Zipf 分布的知识块（p_k ∝ k^(−{1 + alpha:g})，100 万块）：")
    print("    容量受限，记住最常用的 n 块：剩余概率质量 " +
          "  ".join(f"n={n}: {v:.4f}" for n, v in cap) + f"   斜率 {fitted_slope(*zip(*cap)):.3f}"
          + f"（理论 −α = {-alpha}）")
    print("    数据受限，见过才会：期望剩余质量       " +
          "  ".join(f"D={D:.0e}: {v:.4f}" for D, v in dat) + f"   斜率 {fitted_slope(*zip(*dat)):.3f}"
          + f"（理论 −α/(1+α) = {-alpha / (1 + alpha):.3f}）")
    print("""
观察 Observations：
  · (a) 统计学里最简单的估计问题就已经是一条规模定律：误差 ∝ n^(−1)。任何"误差按多项式速度下降"的学习问题，
    在对数坐标下都是直线。
  · (b) 但指数不一定是 −1：同样是"多看数据、预测更准"，维数越高指数越小（d = 8 时只有约 −0.25）。
    Hestness et al. (2017) 测到的神经网络数据指数：词级语言模型 −0.066、字符级 −0.094、机器翻译 −0.128、
    语音 −0.30 左右——都远小于 1。一种解释是：自然数据的"内在维数"很高。
  · (c) 另一种解释：知识是一块一块的，有用程度服从 Zipf 分布。不管是"容量只够记 n 块"还是"数据只够见到一部分"，
    剩下的损失都是幂律，而且**同一个分布同时决定了参数指数和数据指数**。这个图景能解释"为什么是幂律"，
    但不能告诉你真实语言的指数是多少——那只能量。""")


# =============================================================================
# 演示 3：C ≈ 6ND
# =============================================================================
def flops_per_token(model, tokens, vocab, backward):
    """
    用 PyTorch 的 FlopCounterMode 数一次前向（以及可选的反向）一共做了多少次浮点运算，除以 token 数。

    【思路】FlopCounterMode 会拦截每一个矩阵乘法，按"乘加 = 2 次浮点运算"记账（逐元素运算、softmax 等不计）。
    【PyTorch 小知识】with FlopCounterMode(display=False) as fc: ... 之后，fc.get_total_flops() 返回总数。
    """
    B, T = tokens.shape
    with FlopCounterMode(display=False) as fc:
        logits = model(tokens)
        if backward:
            F.cross_entropy(logits.view(-1, vocab), tokens.view(-1)).backward()
    return fc.get_total_flops() / (B * T)


def demo_six_nd():
    print_header("演示 3：C ≈ 6ND —— 每个参数、每个 token，前向 2 次 + 反向 4 次浮点运算")
    torch.manual_seed(0)
    layer = nn.Linear(256, 512, bias=False)
    x = torch.randn(64, 256, requires_grad=True)
    with FlopCounterMode(display=False) as fc:
        y = layer(x)
    fwd = fc.get_total_flops()
    with FlopCounterMode(display=False) as fc:
        layer(x).sum().backward()
    total = fc.get_total_flops()
    n_param = 256 * 512
    print(f"(a) 一个 256 → 512 的线性层（{n_param:,} 个参数），64 个 token：")
    print(f"    前向 {fwd:,} FLOPs = {fwd / (64 * n_param):.0f} × 参数 × token     （每个权重：一次乘法 + 一次加法）")
    print(f"    前向 + 反向 {total:,} FLOPs = {total / (64 * n_param):.0f} × 参数 × token"
          f"（反向要算两个梯度：对输入的、对权重的，各 2 次）")

    print("\n(b) 实测 TinyGPT 训练一步的 FLOPs / token，和 6N 比较（N 含嵌入参数）：")
    print("    " + pad("模型", 42) + pad("N", 12) + pad("前向/token", 13) + pad("训练/token", 13) + pad("训练 ÷ 6N", 11)
          + "上下文 ÷ 宽度")
    configs = [("代码 02 最窄的模型（2 层，d=8）", 2, 8, 32, 21),
               ("代码 02 最宽的模型（2 层，d=64）", 2, 64, 32, 21),
               ("第 1 讲的 base model（d=64，上下文 96）", 2, 64, 96, 44),
               ("更宽的模型（2 层，d=256）", 2, 256, 32, 21),
               ("更宽更深（8 层，d=512）", 8, 512, 32, 21)]
    for name, L, d, T, V in configs:
        cfg = GPTConfig(vocab_size=V, block_size=T, n_layer=L, n_head=max(1, d // 8), d_model=d)
        model = make_model(cfg, seed=0)
        tok = torch.randint(0, V, (2, T), generator=torch.Generator().manual_seed(0))
        f_fwd = flops_per_token(model, tok, V, backward=False)
        f_all = flops_per_token(model, tok, V, backward=True)
        n = count_parameters(model)
        print("    " + pad(name, 42) + pad(f"{n:,}", 12) + pad(f"{f_fwd:,.0f}", 13) + pad(f"{f_all:,.0f}", 13)
              + pad(f"{f_all / (6 * n):.3f}", 11) + f"{T / d:.2f}")

    print("\n(c) 用 C ≈ 6ND 核对 GPT-3（175B 参数，300B 训练 token）：")
    c = 6 * 175e9 * 300e9
    print(f"    6 × 1.75e11 × 3.0e11 = {c:.3e} FLOPs = {c / 8.64e19:,.0f} PF-days（1 PF-day = 1e15 FLOP/s × 1 天 = 8.64e19）")
    n_layer, n_ctx, d_model = 96, 2048, 12288
    attn = 2 * n_layer * n_ctx * d_model
    print(f"    注意力额外的前向 FLOPs（Kaplan 的式子 2·n_layer·n_ctx·d）：{attn:.2e} / token，"
          f"只占 2N = {2 * 175e9:.2e} 的 {attn / (2 * 175e9):.1%}")
    print("""
观察 Observations：
  · 6 = 2（前向）+ 4（反向）。"一个参数、一个 token、6 次浮点运算"只数了矩阵乘法，
    但 Transformer 里几乎所有计算都是矩阵乘法，所以对大模型非常准：训练 ÷ 6N 随宽度增大趋向 1。
  · 6N 漏掉的主要是注意力里 QKᵀ 和"按注意力加权求和"那两步：它们的计算量 ∝ 上下文长度 × 宽度，
    而参数量 ∝ 宽度²。所以"上下文 ÷ 宽度"越大，6N 低估得越多：我们最窄的玩具模型（上下文是宽度的 4 倍）
    实际多出 20% 以上；GPT-3 的上下文只有宽度的六分之一，注意力只多出一个多百分点。
    （FLOP 计数器把因果掩码遮住的那一半也算进去了，所以比 Kaplan 的式子还多一倍。）
  · GPT-3 的 6ND 和常被引用的约 3.1e23 FLOPs（第 1 讲 §3.2）一致。下面几段（代码 02–05）都用 C = 6ND 记账；
    对玩具模型这会低估真实计算量，但同一个记账方式用到底，比较和拟合仍然是自洽的。""")


def main():
    demo_power_law()
    demo_why_power_laws()
    demo_six_nd()


if __name__ == "__main__":
    main()