""" ============================================================================= 第 4 讲 · 代码 5：真实模型的算力账本 —— Kaplan vs Chinchilla、推理成本、数据不够时重复几遍 ============================================================================= 对应讲义：lectures/lecture04_scaling_laws.md 第 3、5、6、9、10、11 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 代码 02–04 在玩具上"量"规模定律；这个文件只用论文里**发表的数字和拟合式**做计算，不训练任何东西。 演示 1：用 C ≈ 6ND 给 GPT-3、Gopher、Chinchilla、PaLM、LLaMA、Llama 3 记账：算力、每个参数看了多少 token。 演示 2：同样的算力，Kaplan 的分配规则和 Chinchilla 的分配规则各会怎样分给参数和数据？ 演示 3：用 Chinchilla 的参数化拟合 L(N, D) = E + A/N^α + B/D^β 比较 Gopher 和 Chinchilla； 以及"指数只保留两位小数"会让外推出的最优分配差多少。 演示 4：把推理成本算进去：一个模型要回答很多请求时，最省的方案是"更小的模型、更多的数据"。 演示 5：数据不够时重复几遍——Muennighoff et al. (2023) 的"有效数据量"公式。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 05_compute_budget_calculator.py # < 1 秒 只用 Python 标准库和第 1 讲的 console.py。所有常数都标明了出处（论文与表号）。 """ import math import os import sys 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 PF_DAY = 1e15 * 86400 # 1 PF-day = 每秒 10^15 次浮点运算 × 一天 = 8.64e19 FLOPs # 模型：(名字, 参数量, 训练 token 数, 出处) MODELS = [ ("GPT-3", 175e9, 300e9, "Brown et al. 2020；Chinchilla 表 1"), ("Gopher", 280e9, 300e9, "Rae et al. 2021；Chinchilla 表 1"), ("Chinchilla", 70e9, 1.4e12, "Hoffmann et al. 2022"), ("PaLM 540B", 540e9, 780e9, "Chowdhery et al. 2022"), ("LLaMA-65B", 65.2e9, 1.4e12, "Touvron et al. 2023 表 2"), ("LLaMA-7B", 6.7e9, 1.0e12, "Touvron et al. 2023 表 2"), ("Llama 3 405B", 405e9, 15.6e12, "Llama 3 论文（报告 3.8e25 FLOPs）"), ("Llama 3 8B", 8e9, 15e12, "Meta 博客：超过 15T token"), ] # Kaplan et al. (2020) 表 6：计算最优的模型大小 N_opt = 1.3e9 · C^0.73（C 以 PF-day 为单位，N 不含嵌入参数） KAPLAN_N_COEF, KAPLAN_N_EXP = 1.3e9, 0.73 # Chinchilla（Hoffmann et al. 2022）方法 3 的拟合（正文里四舍五入后的数字；单位 nats/token） CHIN_ROUNDED = dict(E=1.69, A=406.4, B=410.7, alpha=0.34, beta=0.28) # 同一个拟合的未四舍五入版本（Besiroglu et al. 2024 从论文的 TeX 源文件里读出的 α、β） CHIN_UNROUNDED = dict(E=1.6934, A=406.4, B=410.7, alpha=0.3392, beta=0.2849) # Besiroglu et al. (2024) 重新拟合 Chinchilla 数据得到的参数（他们的主结果） EPOCH_REFIT = dict(E=1.8172, A=482.01, B=2085.43, alpha=0.3478, beta=0.3658) # Chinchilla 表 3（方法 1 的外推）：参数量 → (FLOPs, token 数) CHIN_TABLE3 = [(400e6, 1.92e19, 8.0e9), (1e9, 1.21e20, 20.2e9), (10e9, 1.23e22, 205.1e9), (67e9, 5.76e23, 1.5e12), (175e9, 3.85e24, 3.7e12), (280e9, 9.90e24, 5.9e12), (520e9, 3.43e25, 11.0e12), (1e12, 1.27e26, 21.2e12), (10e12, 1.30e28, 216.2e12)] # Muennighoff et al. (2023) 附录 A：重复数据的"半衰期"常数 R_D* MUENNIGHOFF_RD_STAR = 15.387756 def loss(p, N, D): """参数化损失 L(N, D) = E + A/N^α + B/D^β。""" return p["E"] + p["A"] / N ** p["alpha"] + p["B"] / D ** p["beta"] def chin_optimal(p, C): """ 在 C = 6ND 约束下最小化 L(N, D) 的 (N, D)。推导见 scaling_tools.compute_optimal： N_opt = G · (C/6)^(β/(α+β))，G = (αA / (βB))^(1/(α+β))。 """ a, b = p["alpha"], p["beta"] G = (a * p["A"] / (b * p["B"])) ** (1 / (a + b)) n = G * (C / 6) ** (b / (a + b)) return n, C / 6 / n def fmt(x): """把大数写成 175B、1.4T 这样的形式。""" for unit, v in (("T", 1e12), ("B", 1e9), ("M", 1e6)): if x >= v: return f"{x / v:.3g}{unit}" return f"{x:.3g}" # ============================================================================= # 演示 1：算力账本 # ============================================================================= def demo_ledger(): print_header("演示 1：用 C ≈ 6ND 给真实模型记账") print(" " + pad("模型", 14) + pad("参数 N", 9) + pad("token D", 9) + pad("C = 6ND", 11) + pad("PF-days", 11) + pad("token/参数", 12) + "出处") for name, N, D, src in MODELS: C = 6 * N * D print(" " + pad(name, 14) + pad(fmt(N), 9) + pad(fmt(D), 9) + pad(f"{C:.2e}", 11) + pad(f"{C / PF_DAY:,.0f}", 11) + pad(f"{D / N:,.1f}", 12) + src) print(""" 观察 Observations： · 6ND 和论文自己报告的数字吻合：GPT-3 约 3.1e23；PaLM 论文表格里 540B 是 2527.2 ZFLOPs（= 2.53e24）、 Gopher 504.0、Chinchilla 588.0 ZFLOPs；Llama 3 405B 报告 3.8e25。记账方式简单，但对大模型很准（代码 01）。 · 最后一列是这一讲的主角。2020–2021 年的大模型（GPT-3、Gopher、PaLM）每个参数只看了 1–2 个 token； Chinchilla（2022）是 20；LLaMA（2023）把小模型推到 150 左右；Llama 3 8B 接近 2000——远远超过 "计算最优"。为什么会这样变？演示 2–4 一步步回答。""") # ============================================================================= # 演示 2：Kaplan vs Chinchilla 的分配 # ============================================================================= def demo_allocation(): print_header("演示 2：同样的算力，Kaplan 和 Chinchilla 各会怎样分配？") print("(a) Chinchilla 表 3（方法 1）：计算最优的参数量与 token 数") print(" " + pad("参数", 8) + pad("FLOPs", 11) + pad("token", 9) + pad("token/参数", 12) + "6ND ÷ 表中 FLOPs") for N, C, D in CHIN_TABLE3: print(" " + pad(fmt(N), 8) + pad(f"{C:.2e}", 11) + pad(fmt(D), 9) + pad(f"{D / N:.1f}", 12) + f"{6 * N * D / C:.3f}") print("\n(b) 把算力乘以 10，参数和数据各该乘以多少？") print(f" Kaplan（N ∝ C^0.73，D ∝ C^0.27）：参数 × {10 ** 0.73:.2f}，数据 × {10 ** 0.27:.2f}") print(f" Chinchilla（N ∝ C^0.5，D ∝ C^0.5）：参数 × {10 ** 0.5:.2f}，数据 × {10 ** 0.5:.2f}") print("\n(c) 几个真实的算力预算，按不同规则分配（N 为参数量，D 为 token 数）：") print(" " + pad("算力（来自）", 26) + pad("实际 N / D", 16) + pad("Kaplan 表 6", 16) + pad("每参数 20 token", 18) + "Chinchilla 方法 3") for name in ("GPT-3", "Gopher", "PaLM 540B", "Llama 3 405B"): _, N, D, _ = [m for m in MODELS if m[0] == name][0] C = 6 * N * D nk = KAPLAN_N_COEF * (C / PF_DAY) ** KAPLAN_N_EXP n20 = math.sqrt(C / 120) # 6 · N · 20N = C → N = sqrt(C / 120) nc, dc = chin_optimal(CHIN_ROUNDED, C) print(" " + pad(f"{C:.2e}（{name}）", 26) + pad(f"{fmt(N)} / {fmt(D)}", 16) + pad(f"{fmt(nk)} / {fmt(C / 6 / nk)}", 16) + pad(f"{fmt(n20)} / {fmt(20 * n20)}", 18) + f"{fmt(nc)} / {fmt(dc)}") print(""" 观察 Observations： · Chinchilla 表 3 每一行的"token/参数"都在 20 上下：这就是"每个参数约 20 个 token"这条经验法则的来源 （论文本身没有用这句话，是读者从表 3 和 Chinchilla 70B / 1.4T 总结出来的）。 · 算力 × 10 时，Kaplan 把大部分给参数（× 5.37），数据只 × 1.86；Chinchilla 两边各 × 3.16。 规模越大，两种规则的差距越大：在 GPT-3 的预算上，Kaplan 的公式会建议比 GPT-3 还大的模型、更少的数据； "每参数 20 token"则建议约 50B 参数、1T token——GPT-3 被认为"训练不足"（undertrained）就是这个意思。 · Chinchilla 方法 3（发表的参数化拟合）给出的模型比"20 token/参数"还小、数据还多： 同一篇论文的三种方法并不完全一致（讲义 §6.2；Besiroglu et al. 2024 发现方法 3 的拟合本身有问题，见演示 3）。""") # ============================================================================= # 演示 3：Chinchilla 的参数化拟合 # ============================================================================= def demo_parametric(): print_header("演示 3：Chinchilla 的参数化拟合：Gopher vs Chinchilla，以及四舍五入的代价") p = CHIN_ROUNDED print(f"L(N, D) = {p['E']} + {p['A']}/N^{p['alpha']} + {p['B']}/D^{p['beta']} （nats/token，MassiveText 数据）") print(" " + pad("模型", 14) + pad("E", 7) + pad("A/N^α（参数不够）", 20) + pad("B/D^β（数据不够）", 20) + "预测 L") for name in ("Gopher", "Chinchilla", "GPT-3", "PaLM 540B", "LLaMA-65B"): _, N, D, _ = [m for m in MODELS if m[0] == name][0] tn, td = p["A"] / N ** p["alpha"], p["B"] / D ** p["beta"] print(" " + pad(name, 14) + pad(f"{p['E']:.3f}", 7) + pad(f"{tn:.4f}", 20) + pad(f"{td:.4f}", 20) + f"{p['E'] + tn + td:.4f}") C = 5.76e23 # Chinchilla 论文给 Gopher / Chinchilla 预算用的数字 print(f"\n在 C = {C:.2e} FLOPs（Gopher 的预算）上，三组参数各自推出的计算最优分配：") print(" " + pad("参数来源", 40) + pad("α", 8) + pad("β", 8) + pad("N_opt", 9) + pad("D_opt", 9) + "token/参数") for name, q in (("Chinchilla 正文（两位小数）", CHIN_ROUNDED), ("Chinchilla 未四舍五入（TeX 源文件）", CHIN_UNROUNDED), ("Besiroglu et al. 2024 重新拟合", EPOCH_REFIT)): n, d = chin_optimal(q, C) print(" " + pad(name, 40) + pad(q["alpha"], 8) + pad(q["beta"], 8) + pad(fmt(n), 9) + pad(fmt(d), 9) + f"{d / n:.1f}") print(""" 观察 Observations： · 同样的算力，Gopher（280B / 300B token）的预测损失比 Chinchilla（70B / 1.4T）高：差别几乎全在 "数据不够"那一项。Gopher 的参数项已经很小了，再加参数几乎没用；数据项却还很大。 Chinchilla 论文的实验结果和这个预测一致：Chinchilla 在几乎所有评测上超过了 Gopher。 · 这些"预测损失"只能在**同一个数据集、同一个分词器**下比较。GPT-3、PaLM、LLaMA 的训练数据不同， 把它们代进 Chinchilla 的式子，只是在问"如果它们在 MassiveText 上用这样的 N、D 训练会怎样"。 · 外推对指数极其敏感：α、β 只差千分之几（0.34 → 0.3392、0.28 → 0.2849），最优模型大小就变了约四分之一， token/参数从 90 多变成 60 左右。Besiroglu et al. 重新拟合后得到约 20 token/参数，和方法 1、2 一致。 教训：报告拟合结果要给足位数和置信区间；用别人的拟合外推前，先检查它能不能复现原论文的结论。""") # ============================================================================= # 演示 4：把推理成本算进去 # ============================================================================= def cheapest_for_target(p, target, inference_tokens): """ 在 L(N, D) = target 的约束下，让"训练 + 推理"总 FLOPs 最小的 N 和 D。 【思路】训练 6ND，推理每个 token 约 2N（只有前向），推理 T 个 token 共 2NT。 对每个候选 N，解出达到 target 需要的 D：B/D^β = target − E − A/N^α → D = (B / 剩余)^(1/β)； 在 N 的对数网格上找总成本最小的点（Sardana & Frankle 2023 的思路，这里用 Chinchilla 式的损失）。 【返回】(N, D, 训练 FLOPs, 总 FLOPs) """ best = None for i in range(2001): N = 10 ** (9 + 3 * i / 2000) # 1B 到 1T rest = target - p["E"] - p["A"] / N ** p["alpha"] if rest <= 0: continue D = (p["B"] / rest) ** (1 / p["beta"]) train = 6 * N * D total = train + 2 * N * inference_tokens if best is None or total < best[3]: best = (N, D, train, total) return best def demo_inference(): print_header("演示 4：把推理成本算进去 —— 用得越多，越该训练一个更小、读更多数据的模型") p = EPOCH_REFIT C = 5.76e23 n0, d0 = chin_optimal(p, C) target = loss(p, n0, d0) print(f"目标：达到 Chinchilla 预算（{C:.2e} FLOPs）下计算最优模型的损失 L = {target:.4f}" f"（用 Besiroglu et al. 的重新拟合；计算最优点 N = {fmt(n0)}，D = {fmt(d0)}）") print(" " + pad("推理 token 总数", 16) + pad("最省的 N", 10) + pad("D", 9) + pad("token/参数", 12) + pad("训练 FLOPs", 12) + pad("总 FLOPs", 12) + "比'计算最优 + 同样推理'省") for T in (0, 1e11, 1e12, 1e13, 1e14): N, D, train, total = cheapest_for_target(p, target, T) baseline = 6 * n0 * d0 + 2 * n0 * T print(" " + pad(f"{T:.0e}" if T else "0", 16) + pad(fmt(N), 10) + pad(fmt(D), 9) + pad(f"{D / N:.0f}", 12) + pad(f"{train:.2e}", 12) + pad(f"{total:.2e}", 12) + f"{max(0.0, 1 - total / baseline):.1%}") print(""" 观察 Observations： · 推理量为 0 时，最省的就是计算最优点（约 20 token/参数）。推理量越大，最优点越往"小模型 + 多数据"移： 训练多花一些 FLOPs，换来每次推理都便宜。推理量到 1e13–1e14 token 时（一个被大量使用的产品）， 最优的 token/参数从 20 涨到 100–600。 · 这就是 LLaMA、Llama 3 远远"过度训练"（over-train）小模型的经济学：LLaMA 论文说 Hoffmann et al. 建议 在 200B token 上训练 10B 模型，但他们发现 7B 模型在 1T token 之后还在变好；Llama 3 的 8B 模型 读了超过 15T token（Meta 说它的 Chinchilla 最优量约 200B），仍然按对数线性继续变好。 · Sardana & Frankle (2023) 用同样的思路算过：预期有大量推理需求（约 10 亿次请求）时，应该训练比 Chinchilla 最优更小、更久的模型；他们训练的模型在 token/参数高达 10,000 时仍在变好。""") # ============================================================================= # 演示 5：数据不够时重复几遍 # ============================================================================= def demo_repetition(): print_header("演示 5：数据不够时重复几遍 —— Muennighoff et al. (2023) 的有效数据量") rs = MUENNIGHOFF_RD_STAR print(f"有效数据 D' = U + U · R* · (1 − e^(−R/R*))，U = 不重复的 token，R = 重复次数（epoch − 1），R* = {rs:.2f}") print(" " + pad("epoch 数", 10) + pad("实际读的 token（× U）", 24) + pad("有效 token（× U）", 20) + "平均每个 token 值多少") for ep in (1, 2, 4, 8, 16, 40, 100): R = ep - 1 eff = 1 + rs * (1 - math.exp(-R / rs)) print(" " + pad(ep, 10) + pad(f"{ep}", 24) + pad(f"{eff:.2f}", 20) + f"{eff / ep:.1%}") print(f" " + pad("∞", 10) + pad("∞", 24) + pad(f"{1 + rs:.2f}", 20) + "→ 0") print(""" 观察 Observations： · 重复 4 遍时，平均每个 token 还值 93%：Muennighoff et al. 的实验里，在固定算力下重复到 4 个 epoch， 损失和"全部用新数据"几乎没有差别（8.7B 的模型 4 个 epoch 只比单遍高 0.5%）。 · 之后价值衰减得很快：16 遍时平均每个 token 只值三分之二，40 遍不到 40%；无论重复多少遍， 有效数据最多只有 U 的约 16 倍。他们的结论是：数据受限时，多出来的算力应该更多地花在"多读几遍"上， 而不是继续加参数；但不要指望无限重复。 · 这把"规模定律"从两个资源（N、D）变成了三个（N、不重复的数据 U、重复次数）。 当高质量的公开文本快被用完时（讲义 §11），这个问题变得非常现实。""") def main(): demo_ledger() demo_allocation() demo_parametric() demo_inference() demo_repetition() if __name__ == "__main__": main()