""" ============================================================================= CS324 · 第 4 讲共享模块：scaling_tools.py —— "知道熵"的合成语言、扫描训练与幂律拟合 ============================================================================= 【这个模块是干什么的】 规模定律（scaling laws）说：损失 L 随参数量 N、数据量 D、算力 C 按幂律下降，并趋向一个"不可约"的下限 E。 在真实语料上，E 是多少谁也不知道——它是"英语的熵"之类的量，只能靠拟合猜出来。 这里造一个**熵可以精确算出来**的合成语言，于是： · 拟合出来的 E 对不对，可以直接和真值比； · "损失 → E"不再是一个神秘常数，而是"模型把能学的都学会了"。 本模块提供四样东西： 1. SyntheticLanguage 合成语言：Zipf 分布的词表 + 词与词之间的一阶马尔可夫链；精确的熵率（bit/字符） 2. train_run(...) 在合成语言上训练一个小 Transformer（复用第 1 讲的 tiny_gpt），可以在中途评测 3. ensure_runs(...) 扫描训练 + 缓存：一组 (宽度, 数据量, 调度长度, 种子) 的训练结果存进 cache/sweep.json 4. 拟合工具 纯幂律、带下限的幂律、Chinchilla 式的 L(N, D) = E + A/N^α + B/D^β 【运行方式】 python3 scaling_tools.py # 自检：语言的熵、拟合工具能否找回已知参数（约 3 秒） 依赖 numpy、PyTorch，以及第 1 讲的 tiny_gpt.py、console.py（只读导入）。 【和第 1 讲的关系】 模型就是第 1 讲的 TinyGPT（GPTConfig / make_model / count_parameters），只是更窄、上下文更短（32 个字符）。 训练循环自己写（train_run），原因写在它的文档字符串里。 **本模块从不读写第 1 讲的 checkpoints/**；所有缓存都在 code/lecture04/cache/ 下。 """ import json import math import os import sys import time 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.functional as F # noqa: E402 from tiny_gpt import GPTConfig, count_parameters, make_model # noqa: E402 （导入时会 torch.set_num_threads(1)） CACHE_PATH = os.path.join(HERE, "cache", "sweep.json") # ============================================================================= # 1. 合成语言：Zipf 词表 + 词的一阶马尔可夫链 # ============================================================================= # 字母表：空格 + 15 个辅音 + 5 个元音，一共 21 个字符。词由"辅音 + 元音"音节拼成（例如 "mozimi"）。 CONSONANTS = "bdfghklmnprstvz" VOWELS = "aeiou" ALPHABET_L4 = " abdefghiklmnoprstuvz" # 语言的全部设定。改任何一个值，熵率和所有训练结果都会变（缓存会自动失效）。 LANG_CONFIG = dict( n_words=300, # 词表大小 W zipf_s=1.2, # "背景"词频服从 Zipf：第 k 常见的词概率 ∝ 1 / k^1.2 gamma=0.5, # 下一个词有 50% 的概率来自"搭配"（前一个词偏爱的两个词），50% 来自背景 Zipf 分布 fav_weights=(0.6, 0.4), # 两个偏爱的后继词各占搭配概率的 60% / 40% max_syllables=3, # 每个词 1–3 个音节（2–6 个字母） seed=0, ) class SyntheticLanguage: """ 一个"熵可以精确计算"的合成语言。 【思路】 生成过程：先按一阶马尔可夫链生成一串**词**，再把每个词拼写出来、后面加一个空格。 P(下一个词 = j | 前一个词 = i) = (1 − γ) · zipf(j) + γ · fav_i(j) 其中 zipf 是所有词共用的背景分布（少数词很常见、大量词很罕见），fav_i 是"词 i 偏爱的两个后继词" （像真实语言里的搭配："peanut" 后面常跟 "butter"）。 为什么熵能精确算？字符串和词序列一一对应（空格把词分开），所以 每个字符的熵率 = 每个词的熵率 / 每个词平均占几个字符 = Σ_i π_i · H(P(·|i)) / Σ_i π_i · (len_i + 1) π 是马尔可夫链的平稳分布（长期来看每个词出现的比例）。 更具体地：给定前面足够长的上下文（至少看到上一个完整的词），下一个字符的条件分布和"给定全部历史" 完全一样，所以一个上下文 ≥ 14 个字符的模型，最好也只能做到这个熵率——它就是不可约损失 E。 【FM 联系】真实语料的熵不知道，规模定律里的 E 只能拟合出来（Chinchilla 拟合出 E = 1.69 nats/token）。 这里 E 已知，于是可以检查"拟合出来的 E"靠不靠谱（代码 03）。 【边界】词表里的词互不相同，但可以互为前缀（"mo" 和 "mozimi"）；空格保证了切分唯一。 """ def __init__(self, n_words, zipf_s, gamma, fav_weights, max_syllables, seed): rng = np.random.default_rng(seed) # ---- 词表：W 个互不相同的"音节词" ---- words, seen = [], set() while len(words) < n_words: n_syl = int(rng.integers(1, max_syllables + 1)) w = "".join(CONSONANTS[rng.integers(len(CONSONANTS))] + VOWELS[rng.integers(len(VOWELS))] for _ in range(n_syl)) if w not in seen: seen.add(w) words.append(w) self.words = words W = n_words # ---- 背景 Zipf 分布 ---- ranks = np.arange(1, W + 1, dtype=np.float64) zipf = ranks ** (-zipf_s) self.zipf = zipf / zipf.sum() # ---- 每个词偏爱的后继词（随机挑，不能重复） ---- self.fav = [rng.choice(W, size=len(fav_weights), replace=False).tolist() for _ in range(W)] self.gamma = gamma self.fav_weights = np.array(fav_weights, dtype=np.float64) # ---- 转移矩阵 P[i, j] = P(下一个词 j | 当前词 i) ---- P = np.tile((1 - gamma) * self.zipf, (W, 1)) for i in range(W): for w_k, j in zip(self.fav_weights, self.fav[i]): P[i, j] += gamma * w_k self.P = P # ---- 平稳分布 π：反复做 π ← π P，直到不再变化（幂迭代） ---- pi = np.full(W, 1.0 / W) for _ in range(10000): new = pi @ P if np.abs(new - pi).sum() < 1e-15: pi = new break pi = new self.pi = pi / pi.sum() # ---- 熵率 ---- # 【numpy 小知识】np.log2 对整个数组逐元素取对数；(P * np.log2(P)).sum(axis=1) 是每一行的 Σ p log p row_entropy = -(P * np.log2(P)).sum(axis=1) # 每个"前一个词"条件下，下一个词的熵 self.bits_per_word = float((self.pi * row_entropy).sum()) lengths = np.array([len(w) + 1 for w in words], dtype=np.float64) # +1 是空格 self.chars_per_word = float((self.pi * lengths).sum()) self.entropy_rate = self.bits_per_word / self.chars_per_word # bit / 字符：不可约损失 E # --------------------------------------------------------------------- def sample_word_ids(self, n_words, seed, chunk=100_000): """ 按马尔可夫链采样 n_words 个词的编号。 【思路】一个词一个词地走马尔可夫链。为了快，先把三类随机数（"用搭配还是背景"、"选第几个搭配"、 "背景里抽哪个词"）成块地预先抽好，循环里只做查表。 每 chunk 个词用一个独立的随机数生成器（种子 = [seed, 块编号]），所以**短序列永远是长序列的前缀**： 训练 2M 字符的模型看到的数据，正好是训练 4M 字符的模型的前一半——不同数据量的实验可以公平比较。 【参数】n_words：要几个词；seed：整数种子 【返回】list[int] 【Python 小知识】.tolist() 把 numpy 数组变成 Python 列表；在纯 Python 循环里，列表下标比 numpy 下标快得多。 """ ids = [] cur = None for c in range(0, n_words, chunk): m = min(chunk, n_words - c) rng = np.random.default_rng([seed, c // chunk]) # 总是抽满一整块的随机数（即使只用其中前 m 个），这样块内的随机数和 n_words 无关 use_fav = (rng.random(chunk) < self.gamma).tolist() which_fav = rng.choice(len(self.fav_weights), size=chunk, p=self.fav_weights).tolist() background = rng.choice(len(self.words), size=chunk, p=self.zipf).tolist() fav = self.fav for t in range(m): if cur is not None and use_fav[t]: cur = fav[cur][which_fav[t]] else: cur = background[t] # 第一个词（cur 为 None）总是从背景分布里抽 ids.append(cur) return ids def text(self, ids): """词编号 → 文本：每个词后面加一个空格。""" return "".join(self.words[i] + " " for i in ids) def oracle_bits(self, ids): """ "真实模型"（知道生成过程的全部参数）在这段词序列上的平均编码长度（bit / 字符）。 【思路】链式法则：一个词的所有字符（含空格）的条件概率连乘起来，恰好是 P(这个词 | 前一个词)。 所以不必逐字符计算：总 bit 数 = Σ_t −log2 P(w_t | w_{t−1})，再除以字符数。 第一个词没有前一个词，不计入（字符数也不计入）。 【返回】浮点数。样本足够长时应该非常接近 entropy_rate（只差抽样噪声）。 """ a = np.array(ids[:-1]) b = np.array(ids[1:]) bits = float(-np.log2(self.P[a, b]).sum()) n_chars = sum(len(self.words[i]) + 1 for i in ids[1:]) return bits / n_chars def build_language(cfg=None): """用 LANG_CONFIG（或传入的设定）造一个合成语言。""" return SyntheticLanguage(**(cfg or LANG_CONFIG)) _LUT = np.zeros(256, dtype=np.uint8) for _i, _c in enumerate(ALPHABET_L4): _LUT[ord(_c)] = _i def encode(text): """ 文本 → 字符编号的 uint8 张量（和 tiny_gpt.CharTokenizer 的编号规则相同：字母表里的第 i 个字符 ↔ i）。 【思路】几百万个字符时，逐字符查字典太慢；先把字符串变成字节数组，再用一张 256 格的查找表一次性换成编号。 【边界】只接受 ALPHABET_L4 里的字符（合成语言只会产生这些字符）。 【numpy 小知识】np.frombuffer(bytes, dtype=np.uint8) 把字节串"看成"一个数组，不复制数据； _LUT[数组] 是"花式索引"：对数组里每个值查表。 """ arr = _LUT[np.frombuffer(text.encode("ascii"), dtype=np.uint8)] return torch.from_numpy(arr.copy()) def make_stream(lang, n_chars, seed): """ 采样一段至少 n_chars 个字符的文本，返回 uint8 张量（长度恰好 n_chars）。 【边界】同一个 seed 下，n_chars 小的结果是 n_chars 大的结果的前缀（见 sample_word_ids）。 """ n_words = int(n_chars / lang.chars_per_word * 1.02) + 100 # 多采一点，保证够长 ids = lang.sample_word_ids(n_words, seed) data = encode(lang.text(ids)) if len(data) < n_chars: raise RuntimeError("采样的文本不够长") return data[:n_chars] # ============================================================================= # 2. 模型与训练 # ============================================================================= TRAIN_CONFIG = dict( block_size=32, # 上下文 32 个字符：足够看到"上一个完整的词"（最多 14 个字符） batch_size=32, # 每步 32 段 × 32 字符 = 1024 个 token warmup_frac=0.05, # 前 5% 的步数线性热身 min_lr_frac=0.1, # 余弦衰减到峰值的 10% weight_decay=0.1, betas=(0.9, 0.95), grad_clip=1.0, eval_chars=100_000, # 验证集大小（字符） train_seed=1, # 训练文本的种子 val_seed=2, # 验证文本的种子（和训练文本互不相干） head_dim=8, # 每个注意力头 8 维：n_head = d_model / 8 ) TOKENS_PER_STEP = TRAIN_CONFIG["batch_size"] * TRAIN_CONFIG["block_size"] def lr_for_width(d_model): """ 每个宽度用的峰值学习率：宽度 ≤ 32 时 0.04，更宽时按 (32/d)² 缩小（d = 48 → 0.018，d = 64 → 0.01）。 【思路】写代码前做过一个小的学习率扫描（见 README"关于实验结果"）：越宽的模型能承受的学习率越小， 用同一个学习率会让某一端的模型吃亏，拟合出的指数就会被歪曲。这里用一个简单的规则让每个宽度都 接近各自的最优学习率。 【FM 联系】Kaplan et al. 也发现最优学习率随模型变大而变小（他们给了一个随 log N 线性下降的经验式）； 2024 年的复现工作（Porian et al.）把"没有为每个规模重新调优化器"列为 Kaplan 和 Chinchilla 结论不一致的原因之一。muP 等方法试图让超参数在不同宽度之间直接"搬过去"。 """ return 0.04 * min(1.0, (32.0 / d_model) ** 2) def model_config(d_model, n_layer=2): """ 宽度 d_model、层数 n_layer 的 GPTConfig（词表 21，上下文 32，每个头 8 维）。 【参数】d_model：宽度（必须是 8 的倍数，才能整除成 8 维的头）；n_layer：层数 【返回】tiny_gpt.GPTConfig 【思路】扫描规模时只改宽度、固定层数和头的维度——一次只变一个东西（代码 04 演示 1 说明了为什么）。 """ return GPTConfig(vocab_size=len(ALPHABET_L4), block_size=TRAIN_CONFIG["block_size"], n_layer=n_layer, n_head=max(1, d_model // TRAIN_CONFIG["head_dim"]), d_model=d_model) def param_counts(cfg): """ (总参数量, 非嵌入参数量)。 【参数】cfg：GPTConfig 【返回】(int, int) 【思路】"嵌入参数"= 字符嵌入表（vocab × d）+ 位置嵌入表（block × d）。输出层和字符嵌入共享权重，只算一次。 Kaplan et al. 的 N 不含嵌入参数；Chinchilla 的 N 含全部参数。小模型里嵌入占的比例大得多， 所以两种数法会让拟合出的指数不同（代码 04 演示 4）。 """ model = make_model(cfg, seed=0) total = count_parameters(model) emb = (cfg.vocab_size + cfg.block_size) * cfg.d_model return total, total - emb class EvalSet: """ 验证集：把验证文本切成互不重叠的、长 32 的窗口，**只在每个窗口的后 16 个位置上**算损失。 【思路】窗口开头的几个字符看不到上一个完整的词，任何模型（包括"真实模型"）在那里都做不到熵率。 只用后半段，保证每个被评测的位置都有 ≥ 16 个字符的上下文，于是"最好可能的损失"恰好等于 语言的熵率 E——这样"损失 − E"才是干净的"还没学会的部分"。 """ def __init__(self, data): T = TRAIN_CONFIG["block_size"] n = (len(data) - 1) // T self.x = data[:n * T].long().view(n, T) self.y = data[1:n * T + 1].long().view(n, T) @torch.no_grad() def loss(self, model, chunk=512): """平均交叉熵（bit / 字符），只算每个窗口的后半段。""" model.eval() half = TRAIN_CONFIG["block_size"] // 2 total, count = 0.0, 0 for i in range(0, len(self.x), chunk): logits = model(self.x[i:i + chunk])[:, half:] target = self.y[i:i + chunk, half:] total += F.cross_entropy(logits.reshape(-1, logits.size(-1)), target.reshape(-1), reduction="sum").item() count += target.numel() model.train() return total / count / math.log(2) def lr_schedule(step, schedule_steps, peak_lr): """ 第 step 步的学习率：前 5% 的步数线性热身，然后余弦衰减到峰值的 10%（在 schedule_steps 步时到底）。 【参数】step：当前步（从 0 开始）；schedule_steps：调度按多少步来排；peak_lr：峰值学习率 【返回】浮点数 【思路】和第 1 讲 tiny_gpt.lr_at 的形状一样，但热身长度按比例而不是固定 100 步 （最短的训练只有 244 步，固定 100 步热身就占了四成）。 【边界】step > schedule_steps 时保持在最低值。 """ warm = max(1, int(TRAIN_CONFIG["warmup_frac"] * schedule_steps)) if step < warm: return peak_lr * (step + 1) / warm progress = min(1.0, (step - warm) / max(1, schedule_steps - warm)) lo = TRAIN_CONFIG["min_lr_frac"] return peak_lr * (lo + (1 - lo) * 0.5 * (1 + math.cos(math.pi * progress))) def train_run(cfg, train_data, eval_set, n_tokens, schedule_tokens=None, eval_at=(), seed=0, lr=None): """ 训练一个模型：从训练文本的开头起，**按顺序**读 n_tokens 个字符（每个字符只读一遍），然后在验证集上评测。 【为什么不用第 1 讲的 train_lm】 1. 规模定律的实验要求"学习率调度的长度"可以和"实际训练长度"不同（schedule_tokens）： Chinchilla 指出 Kaplan 的一个问题就出在这里（代码 04 演示 3）。train_lm 的调度长度总是等于训练步数。 2. train_lm 随机切 batch（可能重复读同一段文本）；这里按顺序读，D 就是"见过的不同 token 数"， 和真实大模型"单遍（one epoch）训练"的设定一致。 3. 热身长度按比例设置（见 lr_schedule）。 其余部分（AdamW、梯度裁剪、交叉熵）和 train_lm 完全相同。 【参数】 cfg : GPTConfig train_data : uint8 张量（训练文本），长度至少 n_tokens + 1 eval_set : EvalSet n_tokens : 训练读多少个 token（D）；步数 = n_tokens / 1024 schedule_tokens : 学习率调度按多少个 token 来排（默认 = n_tokens，即"调度和训练长度匹配"）。 大于 n_tokens 时，相当于"一个更长的训练跑到一半就停下"。 eval_at : 训练途中在这些 token 数处评测（用来得到"中途检查点"的损失） seed : 模型初始化的种子 lr : 峰值学习率（默认 lr_for_width） 【返回】dict：loss（训练结束时的验证损失，bit/字符）、evals（{token 数: 损失}）、steps、seconds 【边界】n_tokens 必须是 1024 的整数倍（否则多出来的零头不训练）。 【PyTorch 小知识】data[idx] 里 idx 是二维整数张量时，结果是同样形状的张量（"花式索引"）： 一次就切出整个 batch，比用 Python 循环一段一段切再 torch.stack 快得多。 """ B, T = TRAIN_CONFIG["batch_size"], TRAIN_CONFIG["block_size"] steps = n_tokens // (B * T) sched_steps = (schedule_tokens or n_tokens) // (B * T) if len(train_data) < steps * B * T + 1: raise ValueError("训练文本不够长") peak = lr if lr is not None else lr_for_width(cfg.d_model) model = make_model(cfg, seed=seed) opt = torch.optim.AdamW(model.parameters(), lr=peak, weight_decay=TRAIN_CONFIG["weight_decay"], betas=TRAIN_CONFIG["betas"]) eval_steps = {e // (B * T): e for e in eval_at} offsets = torch.arange(B)[:, None] * T + torch.arange(T)[None, :] # (B, T)：一个 batch 里每个字符的相对位置 evals = {} t0 = time.time() model.train() for step in range(steps): if step in eval_steps: evals[eval_steps[step]] = eval_set.loss(model) for group in opt.param_groups: group["lr"] = lr_schedule(step, sched_steps, peak) idx = offsets + step * B * T # 第 step 步读第 step 块（每块 B·T 个字符），不重复 x = train_data[idx].long() y = train_data[idx + 1].long() logits = model(x) loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1)) opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), TRAIN_CONFIG["grad_clip"]) opt.step() final = eval_set.loss(model) return dict(loss=final, evals={str(k): v for k, v in sorted(evals.items())}, steps=steps, seconds=time.time() - t0) # ============================================================================= # 3. 扫描训练 + 缓存 # ============================================================================= CACHE_VERSION = 1 def run_key(d_model, n_tokens, n_layer=2, schedule_tokens=None, seed=0): """一次训练的"身份证"：宽度、层数、数据量、调度长度、种子。缓存按它查找。""" sched = schedule_tokens or n_tokens return f"L{n_layer}_d{d_model}_D{n_tokens}_S{sched}_seed{seed}" def _cache_meta(): return dict(version=CACHE_VERSION, lang=LANG_CONFIG, train=TRAIN_CONFIG, torch=torch.__version__.split("+")[0]) def load_cache(path=CACHE_PATH): """ 读缓存；配置（语言、训练设定、版本）对不上时返回空缓存——避免读到用旧设定训练出的结果。 【Python 小知识】json 会把元组存成列表，所以比较前先把当前配置也"过一遍 json"（dumps 再 loads）。 """ meta = json.loads(json.dumps(_cache_meta())) if os.path.exists(path): with open(path) as f: saved = json.load(f) if saved.get("meta") == meta: return saved print("（缓存的配置和当前配置不同，忽略旧缓存）") return dict(meta=meta, runs={}) def ensure_runs(specs, path=CACHE_PATH, verbose=True): """ 保证 specs 里的每一次训练都在缓存里：缺的就训练并立刻写回缓存。 【参数】specs：list[dict]，每个 dict 有 d_model、n_tokens，可选 n_layer、schedule_tokens、seed、eval_at 【返回】dict：run_key → 结果（loss、evals、n_params、n_nonemb、D、d_model、n_layer……） 【思路】训练是确定性的（固定种子、单线程、按顺序读数据），所以"读缓存"和"重新训练"得到逐字相同的数字。 每训练完一个就保存一次：中途按 Ctrl-C 打断，已经训练好的不会丢。 【边界】训练文本只生成一次，长度取 specs 里最大的 n_tokens。 """ cache = load_cache(path) runs = cache["runs"] def key_of(s): return run_key(s["d_model"], s["n_tokens"], s.get("n_layer", 2), s.get("schedule_tokens"), s.get("seed", 0)) missing = [s for s in specs if key_of(s) not in runs] if missing: lang = build_language() need = max(s["n_tokens"] for s in missing) + 1 train_data = make_stream(lang, need, TRAIN_CONFIG["train_seed"]) eval_set = EvalSet(make_stream(lang, TRAIN_CONFIG["eval_chars"], TRAIN_CONFIG["val_seed"])) if verbose: total_tok = sum(s["n_tokens"] for s in missing) print(f"（缓存里缺 {len(missing)} 次训练，共 {total_tok / 1e6:.1f}M 个 token，开始训练；单线程 CPU……）") t_all = time.time() for s in missing: n_layer = s.get("n_layer", 2) cfg = model_config(s["d_model"], n_layer) res = train_run(cfg, train_data, eval_set, s["n_tokens"], s.get("schedule_tokens"), s.get("eval_at", ()), s.get("seed", 0)) total, nonemb = param_counts(cfg) res.update(d_model=s["d_model"], n_layer=n_layer, D=s["n_tokens"], schedule=s.get("schedule_tokens") or s["n_tokens"], seed=s.get("seed", 0), n_params=total, n_nonemb=nonemb) key = key_of(s) runs[key] = res os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) with open(path, "w") as f: json.dump(cache, f, indent=1, sort_keys=True) if verbose: print(f" 训练完成 {key:<34} 损失 {res['loss']:.4f} 用时 {res['seconds']:.1f} 秒") if verbose: print(f"（训练用时合计 {time.time() - t_all:.0f} 秒；结果已存入 cache/sweep.json）") return {key_of(s): runs[key_of(s)] for s in specs} # 主扫描网格（代码 02 训练；代码 03、04 读取）：6 个宽度 × 4 个数据量，调度长度和训练长度匹配。 # 每个宽度 2M token 的那一次训练，途中在 0.125M / 0.25M / 0.5M / 1M 处各评测一次—— # 这就是"一个长调度的训练跑到一半"的损失（Kaplan 式的数据，代码 04 演示 3 用）。 WIDTHS = (8, 16, 24, 32, 48, 64) DATA_SIZES = (250_000, 500_000, 1_000_000, 2_000_000) TRUNC_EVALS = (125_000, 250_000, 500_000, 1_000_000) def grid_specs(): """主网格的训练清单（24 次 + 2 次换种子的重复，用来估计噪声）。""" specs = [] for d in WIDTHS: for D in DATA_SIZES: spec = dict(d_model=d, n_tokens=D) if D == DATA_SIZES[-1]: spec["eval_at"] = TRUNC_EVALS specs.append(spec) specs.append(dict(d_model=16, n_tokens=500_000, seed=1)) specs.append(dict(d_model=48, n_tokens=500_000, seed=1)) return specs def grid_table(runs): """ 把主网格的结果整理成表。 【参数】runs：ensure_runs 的返回值 【返回】(table, sizes)：table = {宽度: {D: 损失}}，sizes = {宽度: (总参数, 非嵌入参数)} 【边界】只收 2 层、种子 0、调度长度 = 训练长度的训练；其他训练（换种子、改形状）被忽略。 """ table, sizes = {}, {} for r in runs.values(): if r["n_layer"] != 2 or r["seed"] != 0 or r["schedule"] != r["D"]: continue table.setdefault(r["d_model"], {})[r["D"]] = r["loss"] sizes[r["d_model"]] = (r["n_params"], r["n_nonemb"]) return table, sizes # ============================================================================= # 4. 拟合工具 # ============================================================================= def fit_power_law(x, y): """ 纯幂律 y = A · x^(−α)：在对数坐标下就是一条直线 log y = log A − α log x，用最小二乘拟合直线。 【返回】(A, α) 【边界】要求 x、y 都 > 0。没有下限项：x → ∞ 时预测值 → 0，这在有"不可约损失"的问题上一定会出错（代码 03）。 """ lx, ly = np.log(np.asarray(x, float)), np.log(np.asarray(y, float)) slope, intercept = np.polyfit(lx, ly, 1) return float(math.exp(intercept)), float(-slope) def _nonneg_wls(columns, y): """ 给定几列"基函数"，求非负系数 c 使 Σ c_k · column_k 在**相对误差**意义下最接近 y。 【思路】我们关心的是对数空间的误差 log y − log ŷ ≈ (y − ŷ) / y，所以每一行除以 y 再做普通最小二乘。 系数必须非负（E、A、B 都是非负的）：先解不加约束的最小二乘，系数都非负就直接用； 否则把一部分系数固定为 0 再解，在所有组合里挑对数误差最小的（只有 2–3 个系数，枚举很便宜）。 【返回】(对数空间的平方误差和, 系数数组)；没有可行解时返回 None """ X = np.stack(columns, axis=1) Xw = X / y[:, None] ones = np.ones_like(y) k = X.shape[1] # 先解不加约束的最小二乘：如果系数都非负（绝大多数情况），它就是答案，不必再枚举 coef, *_ = np.linalg.lstsq(Xw, ones, rcond=None) pred = X @ coef if (coef >= 0).all() and (pred > 0).all(): return float(((np.log(y) - np.log(pred)) ** 2).sum()), coef best = None for mask in range(1, 2 ** k): cols = [i for i in range(k) if mask >> i & 1] coef, *_ = np.linalg.lstsq(Xw[:, cols], ones, rcond=None) if (coef < 0).any(): continue full = np.zeros(k) full[cols] = coef pred = X @ full if (pred <= 0).any(): continue err = float(((np.log(y) - np.log(pred)) ** 2).sum()) if best is None or err < best[0]: best = (err, full) return best EXPONENT_GRID = np.exp(np.linspace(math.log(0.02), math.log(3.0), 241)) def fit_power_law_with_floor(x, y): """ 带下限的幂律 y = E + A · x^(−α)：E 是不可约部分，A · x^(−α) 是"还能学会的部分"。 【思路】"变量投影"（variable projection）：指数 α 一旦固定，模型对 E、A 就是线性的，可以直接解； 所以只需要在 α 的一个网格上逐个试（每个 α 解一次线性最小二乘），挑对数误差最小的那个， 再在它附近用更细的网格精修。确定性、不需要初始值，也不会"卡在局部最优" （Chinchilla 用 L-BFGS + 一个初始值网格解决同样的问题）。 【参数】x、y：一维数组，至少 3 个点 【返回】dict(E, A, alpha, rmse)，rmse 是对数空间的均方根误差 【边界】只有 3 个点时自由度为 0（三个参数正好穿过三个点）；点少、范围窄时 E 和 α 会互相"抵消"， 同样好的拟合可以对应很不同的 (E, α)——代码 03 演示 4 专门看这个问题。 """ x = np.asarray(x, float) y = np.asarray(y, float) best = None for grid in (EXPONENT_GRID, None): if grid is None: # 第二轮：在第一轮最优值附近细化 a0 = best[3] grid = a0 * np.exp(np.linspace(-0.03, 0.03, 121)) for a in grid: r = _nonneg_wls([np.ones_like(x), x ** (-a)], y) if r is not None and (best is None or r[0] < best[0]): best = (r[0], r[1][0], r[1][1], float(a)) err, E, A, a = best return dict(E=float(E), A=float(A), alpha=a, rmse=math.sqrt(err / len(x))) def fit_joint(N, D, y): """ Chinchilla 的"方法 3"：L(N, D) = E + A / N^α + B / D^β，一次用上所有 (N, D) 点。 【思路】同样是变量投影：固定 (α, β) 时对 (E, A, B) 线性。先在 (α, β) 的粗网格上搜索，再在最优点附近细化。 【返回】dict(E, A, alpha, B, beta, rmse) 【FM 联系】Chinchilla 拟合出 E = 1.69、A = 406.4、B = 410.7、α = 0.34、β = 0.28（nats/token）； 用它可以推出"给定算力时最优的 N 和 D"（compute_optimal）。 """ N = np.asarray(N, float) D = np.asarray(D, float) y = np.asarray(y, float) coarse = np.exp(np.linspace(math.log(0.05), math.log(3.0), 61)) best = None # 【Python 小知识】内层函数要给外层函数的变量重新赋值，必须先声明 nonlocal； # 否则 Python 会把 best 当成内层函数自己的新变量。 def try_pair(a, b): nonlocal best r = _nonneg_wls([np.ones_like(N), N ** (-a), D ** (-b)], y) if r is not None and (best is None or r[0] < best[0]): best = (r[0], r[1], float(a), float(b)) for a in coarse: for b in coarse: try_pair(a, b) for _ in range(2): # 两轮局部细化 a0, b0 = best[2], best[3] for a in a0 * np.exp(np.linspace(-0.07, 0.07, 29)): for b in b0 * np.exp(np.linspace(-0.07, 0.07, 29)): try_pair(a, b) err, coef, a, b = best return dict(E=float(coef[0]), A=float(coef[1]), alpha=a, B=float(coef[2]), beta=b, rmse=math.sqrt(err / len(y))) def joint_predict(p, N, D): """用 fit_joint 的结果预测损失。""" return p["E"] + p["A"] * np.asarray(N, float) ** (-p["alpha"]) + p["B"] * np.asarray(D, float) ** (-p["beta"]) def compute_optimal(p, C): """ 在算力约束 C = 6ND 下，让 L(N, D) = E + A/N^α + B/D^β 最小的 N 和 D。 【思路】把 D = C/(6N) 代进去，对 N 求导令其为 0，得到 N_opt = G · (C/6)^(β/(α+β)), D_opt = (C/6) / N_opt, G = (αA / (βB))^(1/(α+β)) 所以"算力增加时参数该增长多快"的指数 a = β/(α+β)，数据的指数 b = α/(α+β)，a + b = 1。 哪个资源的指数（α 或 β）更小、回报衰减得更慢，就该把更多算力分给另一个……反过来说： β 大（数据的回报衰减快）→ a 大 → 多给参数。 【返回】(N_opt, D_opt) """ a, b = p["alpha"], p["beta"] G = (a * p["A"] / (b * p["B"])) ** (1.0 / (a + b)) n_opt = G * (C / 6.0) ** (b / (a + b)) return n_opt, (C / 6.0) / n_opt def parabola_min(log_n, losses): """ IsoFLOP 曲线（固定算力，横轴 log N）上拟合一条抛物线，返回最低点的 N 和那里的损失。 【思路】Chinchilla 的"方法 2"：同一算力下，模型太小（数据多但容量不够）和太大（容量够但数据太少）都不好， 损失对 log N 是一个"山谷"；在谷底附近用抛物线近似，顶点就是这个算力下的最优模型大小。 【返回】(N_opt, L_min, 是否在数据范围之内)。抛物线开口向下（c2 ≤ 0）时返回 (nan, nan, False)。 【边界】顶点落在采样范围之外时，说明山谷没有被"包住"，结果是外推，要打折扣。 """ c2, c1, c0 = np.polyfit(np.asarray(log_n, float), np.asarray(losses, float), 2) if c2 <= 0: return float("nan"), float("nan"), False x0 = -c1 / (2 * c2) inside = min(log_n) <= x0 <= max(log_n) return float(math.exp(x0)), float(c0 - c1 * c1 / (4 * c2)), bool(inside) def interp_loglog(x_query, xs, ys): """ 在 (log x, log y) 上做分段线性插值；查询点超出 [min xs, max xs] 时返回 None（不外推）。 【FM 联系】Chinchilla 的方法 1 也是这样：每个模型大小有一条"损失 vs token 数"的曲线， 在给定算力处读出（插值）它的损失。 """ xs = np.asarray(xs, float) ys = np.asarray(ys, float) if x_query < xs.min() * (1 - 1e-9) or x_query > xs.max() * (1 + 1e-9): return None return float(math.exp(np.interp(math.log(x_query), np.log(xs), np.log(ys)))) # ============================================================================= # 5. 自检 # ============================================================================= if __name__ == "__main__": lang = build_language() print(f"合成语言：{len(lang.words)} 个词，熵率 E = {lang.entropy_rate:.4f} bit/字符") ids = lang.sample_word_ids(50_000, seed=7) print(f" '真实模型'在 5 万个词的样本上：{lang.oracle_bits(ids):.4f} bit/字符（应接近 E）") a = make_stream(lang, 5000, 1) b = make_stream(lang, 20000, 1) assert torch.equal(a, b[:5000]), "短文本应该是长文本的前缀" print(" 前缀一致性 ✓") # 拟合工具：造一组已知参数的数据，看能否找回来 xs = np.array([1e3, 3e3, 1e4, 3e4, 1e5]) ys = 0.7 + 50 * xs ** (-0.5) f = fit_power_law_with_floor(xs, ys) print(f" 带下限拟合：E = {f['E']:.4f}（真值 0.7），α = {f['alpha']:.4f}（真值 0.5）") assert abs(f["E"] - 0.7) < 1e-3 and abs(f["alpha"] - 0.5) < 1e-2 Ns = np.array([1e3, 1e4, 1e5] * 3) Ds = np.repeat([1e5, 1e6, 1e7], 3) ys = 0.7 + 30 * Ns ** (-0.4) + 200 * Ds ** (-0.3) p = fit_joint(Ns, Ds, ys) print(f" 联合拟合：E = {p['E']:.3f} α = {p['alpha']:.3f} β = {p['beta']:.3f}（真值 0.7、0.4、0.3）") assert abs(p["alpha"] - 0.4) < 0.02 and abs(p["beta"] - 0.3) < 0.02 print("自检通过。")