"""
=============================================================================
第 9 讲 · 代码 5：规模与算力 —— 参数量、训练成本、Chinchilla 规模定律
=============================================================================

对应讲义：lectures/lecture09_pretraining.md 第 9 节（详细内容见 CS336 第 2、9、11 讲）

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
预训练最显著的特点是"大"：大模型、大数据、大算力。这个文件不训练任何东西，只做"算术"：

    演示 1：参数量 —— 用第 7 讲的公式算出 GPT-2 四个版本的参数量
    演示 2：训练要多少计算？—— C ≈ 6ND，以及在 GPU 上要算多久
    演示 3：给定算力，模型该多大、数据该多少？—— Chinchilla 规模定律
    演示 4：训练只算一次，推理要算无数次 —— 为什么今天的小模型都"训练过度"

只需要 numpy，约 1 秒。
"""

import numpy as np


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


def gpt_params(V, ctx, d, L):
    """
    GPT 式解码器的参数量（可学习位置嵌入、输出层与词嵌入共享、每块 12d² + 13d，见第 7 讲第 7.5 节）。

    【参数】V：词表大小；ctx：最大长度；d：模型维度；L：层数
    """
    return V * d + ctx * d + L * (12 * d * d + 13 * d) + 2 * d


def chinchilla_loss(N, D, E, A, B, alpha, beta):
    """
    Chinchilla 规模定律（Hoffmann 等 2022 的第 3 种拟合方法）：

        L(N, D) = E + A / N^α + B / D^β

    【参数】
        N : 参数量；D : 训练的词元数
        E : "不可约"损失（文本本身的不确定性，再大的模型也降不下去）
        A, α : 模型太小带来的损失；B, β : 数据太少带来的损失

    【返回】预测的交叉熵损失（每个词元，自然对数）
    """
    return E + A / N ** alpha + B / D ** beta


def compute_optimal(C, consts):
    """
    给定算力 C（浮点运算次数），在 C = 6ND 的约束下找让损失最低的 N（D = C / 6N）。

    【思路】在对数尺度上网格搜索 N（从 1e6 到 1e14），取损失最小的那个。
    【返回】(N_opt, D_opt, 最低损失)
    """
    Ns = np.logspace(6, 14, 4001)
    Ds = C / (6 * Ns)
    losses = chinchilla_loss(Ns, Ds, *consts)
    i = int(np.argmin(losses))
    return Ns[i], Ds[i], losses[i]


def fmt(x):
    """把大数字写成 1.2e9 这样的形式。"""
    return f"{x:.2e}".replace("e+0", "e").replace("e+", "e")


def main():
    # ---------------- 演示 1 ----------------
    print_header("演示 1：GPT-2 四个版本的参数量（第 7 讲的公式）")
    print(f"\n    {'模型':<14s}{'d':>6s}{'层数':>6s}{'公式算出的参数量':>18s}{'公布的大小':>10s}")
    for name, d, L, known in [("GPT-2 small", 768, 12, "124M"), ("GPT-2 medium", 1024, 24, "355M"),
                              ("GPT-2 large", 1280, 36, "774M"), ("GPT-2 XL", 1600, 48, "1.5B")]:
        print(f"    {name:<14s}{d:>6d}{L:>6d}{gpt_params(50257, 1024, d, L):>20,}{known:>10s}")
    mini = gpt_params(182, 96, 64, 3)
    print(f"    {'本讲的迷你 GPT':<12s}{64:>6d}{3:>6d}{mini:>20,}")
    print("""
  观察：公式和 OpenAI 公布的大小完全吻合。参数量主要由 L·12d² 决定：
    d 加倍、L 加倍，参数量大约变成 8 倍。我们的迷你 GPT 比 GPT-2 small 小约 700 倍。""")

    # ---------------- 演示 2 ----------------
    print_header("演示 2：训练要多少计算？C ≈ 6ND")
    print("""
  每个词元：前向传播约 2N 次浮点运算（每个参数一次乘法、一次加法），反向传播约 4N 次，合计 6N。
  训练 D 个词元：C ≈ 6ND（CS336 第 2 讲有详细推导）。""")
    gpu = 989e12 * 0.4          # H100 的 bf16 稠密峰值约 989 TFLOP/s；实际利用率按 40% 算
    print(f"\n    {'模型':<22s}{'参数 N':>10s}{'词元 D':>10s}{'C = 6ND':>11s}{'1000 张 H100 要几天':>18s}")
    rows = [("本讲的迷你 GPT", mini, 4 * 1.08e6), ("GPT-3（2020）", 175e9, 300e9),
            ("Chinchilla（2022）", 70e9, 1.4e12), ("LLaMA-2 7B（2023）", 7e9, 2e12), ("LLaMA-3 8B（2024）", 8e9, 15e12)]
    for name, N, D in rows:
        C = 6 * N * D
        print(f"    {name:<20s}{fmt(N):>10s}{fmt(D):>10s}{fmt(C):>11s}{C / (1000 * gpu) / 86400:>16.3g}")
    print("""
  观察：
    - GPT-3 的 3.15e23 次运算和论文报告的数字一致。即使用 1000 张今天的 H100（按 40% 的利用率），也要算好几天；
      当年用的是更慢的 V100，而且还要算上调试、失败的实验、数据处理…… 实际花费远不止这些。
    - 我们的迷你 GPT 的算力大约是 GPT-3 的 10^11 分之一 —— 在笔记本上几分钟就能训练完。
    - 注意 LLaMA-3 8B：参数比 GPT-3 少 20 多倍，训练的词元却多 50 倍（演示 4 解释为什么）。
    （GPU 的峰值和利用率都是粗略的假设，只用来感受数量级。）""")

    # ---------------- 演示 3 ----------------
    print_header("演示 3：Chinchilla 规模定律 —— 给定算力，模型多大、数据多少最好？")
    fits = {
        "Hoffmann 等 2022": (1.69, 406.4, 410.7, 0.34, 0.28),
        "Besiroglu 等 2024 的重新拟合": (1.82, 482.01, 2085.43, 0.3478, 0.3658),
    }
    for name, consts in fits.items():
        print(f"\n    拟合参数：{name}（E={consts[0]}, A={consts[1]}, B={consts[2]}, α={consts[3]}, β={consts[4]}）")
        print(f"    {'算力 C':>10s}{'最优参数 N':>12s}{'最优词元 D':>12s}{'D / N':>8s}{'预测损失':>10s}")
        for C in [1e19, 1e21, 1e23, 1e25]:
            N, D, Lmin = compute_optimal(C, consts)
            print(f"    {fmt(C):>10s}{fmt(N):>12s}{fmt(D):>12s}{D / N:>8.1f}{Lmin:>10.3f}")
    consts = fits["Besiroglu 等 2024 的重新拟合"]
    C3 = 6 * 175e9 * 300e9
    N, D, Lmin = compute_optimal(C3, consts)
    print(f"\n    用 GPT-3 的算力（{fmt(C3)}）：")
    print(f"      GPT-3 的分配（N = 1.75e11，D = 3.00e11）预测损失 {chinchilla_loss(175e9, 300e9, *consts):.3f}")
    print(f"      最优分配（N = {fmt(N)}，D = {fmt(D)}）预测损失 {Lmin:.3f}")
    print("""
  观察：
    - Kaplan 等（2020）最早发现：损失随参数量、数据量、算力按**幂律**下降（在对数坐标上是直线）。
      规模定律让人可以先在小模型上做实验，再**预测**大模型的表现，决定怎样花钱。
    - Chinchilla（Hoffmann 等 2022）的结论：算力增加时，参数量和数据量应该**大致同比例**增加，
      每个参数大约对应 20 个词元。按这个标准，GPT-3（每个参数不到 2 个词元）的模型太大、数据太少。
    - 但看上面的表：原论文第 3 种方法给出的拟合参数，算出来的最优比例和"20"并不一致；
      Besiroglu 等（2024）重新分析原论文的数据，发现原来的拟合有问题，重新拟合之后才和"约 20 个词元 / 参数"吻合。
      规模定律是经验规律，具体系数取决于数据、模型和拟合方法，要谨慎对待（CS336 第 9、11 讲）。""")

    # ---------------- 演示 4 ----------------
    print_header("演示 4：训练只算一次，推理要算无数次")
    print("""
  生成一个词元约需 2N 次运算。假设一个模型上线后总共要生成 10^13 个词元（一个热门产品很快就能达到），""")
    served = 1e13
    print(f"\n    {'方案':<36s}{'训练 6ND':>11s}{'推理 2N × 1e13':>16s}{'合计':>11s}")
    for name, N, D in [("70B 参数，Chinchilla 最优（1.4e12 词元）", 70e9, 1.4e12),
                       ("8B 参数，训练过度（1.5e13 词元）", 8e9, 15e12)]:
        tr, inf = 6 * N * D, 2 * N * served
        print(f"    {name:<30s}{fmt(tr):>11s}{fmt(inf):>16s}{fmt(tr + inf):>11s}")
    print("""
  观察：Chinchilla 回答的是"训练算力固定时，怎样得到最好的模型"。但模型上线后，推理的总成本可能远超训练。
    小模型每次推理都便宜，所以今天的做法常常是：选一个比 Chinchilla 最优更小的模型，用多得多的数据把它
    "训练过度"（LLaMA 系列每个参数上千个词元）—— 训练时多花一些算力，换来推理时省下大量算力。""")


if __name__ == "__main__":
    main()