""" ============================================================================= 第 1 讲 · 工具模块：toy world —— "多步算术题" + 一个可以精确控制的 simulated LLM ============================================================================= 对应讲义：lectures/lecture01_overview.md 第 0.4 节（本讲的玩具世界） ----------------------------------------------------------------------------- 【为什么要造一个 toy world？】 ----------------------------------------------------------------------------- 这门课的主角是 "self-improving AI agents"：模型通过多次尝试（sampling）、检查（verification）、 反馈（feedback）和训练（training）让自己越来越好。要真正看懂这些方法，我们需要回答一些很细的问题： · 模型为什么会错？错误是"随机手滑"，还是"每次都犯的同一个错"（systematic error）？ · 多采样 100 次能救回多少题？为什么 majority voting 会卡住（plateau）？ · verifier 有 10% 的 false positive，会把最终准确率限制在多少？ 用真实的 LLM 做实验，这些问题很难说清楚：我们不知道模型"心里"每一步犯错的概率，也很贵、很慢。 所以这里造一个小世界，**每个错误从哪里来都是已知的**。它刻画的是**机制**，不代表任何真实模型的数值。 ----------------------------------------------------------------------------- 【这个世界里有什么？】 ----------------------------------------------------------------------------- 1. 题目（Problem）：一个"迷你 GSM8K"——多步算术链。例如 从 47 开始：+ 38，× 12，- 95。最后得多少？ 每道题有 k 个步骤（step），每步是加（add）、减（sub）、乘（mul）之一。 答案可以用 Python 精确算出来 → 这是一个 **verifiable** 的任务（像数学题和带 unit test 的代码）。 2. 模型（ToyLLM）：模拟一个"会做算术、但会犯错"的语言模型。它一步一步地算（像 chain-of-thought）， 每一步从三种 "procedure"（做法）里**按概率**选一种： correct ：做对。 bug ：一个固定的错误套路（misconception），例如乘法"少加了一份"：a × (b − 1)。 它是确定性的 —— 同一步犯同一个 bug，得到的错误答案每次都一样。 这就是 systematic error，第 02 号文件会看到它让 majority voting 卡住。 careless ：随机手滑：在正确结果上随机偏一点（±1、±2、±10）。每次错得不一样。 选择概率由 softmax(logits / temperature) 决定（第 1 讲第 5.3 节讲 temperature）： logit(correct) = ALPHA[op] + BETA × log10(size / 1e9) − hardness(这一步) logit(bug) = LOGIT_BUG （固定为 0） logit(careless) = LOGIT_CARELESS （固定为 −0.5） · size（参数量）越大，correct 的 logit 越高 → 每步更可靠。这是我们对 "scaling" 的最简单的刻画： 每扩大 10 倍，logit 加 BETA。 · hardness：有进位的加法、有借位或结果为负的减法、乘两位数、乘大数，都更难。 于是不同的题难度差别很大 —— 这对理解 coverage 曲线（第 02 号文件）至关重要。 3. 两种回答模式： · cot=True（chain-of-thought）：把每一步写出来，每一步只需要"算一步"。 · cot=False（direct answer）：不写中间步骤，直接报答案 → 所有步骤必须"在一次前向传播里"完成。 模型一次前向传播能串行做的步数有限（depth_capacity，随 size 缓慢增长）；超过这个步数的步骤 会被重罚（DIRECT_PENALTY），几乎只能靠猜。这是对"CoT 提供了额外的串行计算"这一理论结论 的卡通化（讲义第 3.3 节）。 4. 工具（tool use）：generate() 可以让某些运算交给计算器（calculator）做。计算器不会算错， 但调用本身有一定概率失败（格式写错等），失败时模型只能自己算。第 03 号文件和第 4 讲会用到。 5. 成本：一次采样的计算量 ≈ 2 × 参数量 × token 数（FLOPs）。这是 Transformer 推理成本的常用估算 （每个参数每个 token 做一次乘法和一次加法）。第 02 号文件用它做 "compute-matched" 的比较。 ----------------------------------------------------------------------------- 【上帝视角 vs 模型视角】 ----------------------------------------------------------------------------- Problem.answer() 是标准答案，只有 verifier（例如"用 Python 重算一遍"）和评估代码可以用。 模型生成答案时**不看**标准答案。这和 CS329H、CS224V 课程代码里的约定一样：上帝视角只用于评估。 ----------------------------------------------------------------------------- 【用法】 ----------------------------------------------------------------------------- import random from toy_llm import ToyLLM, make_problem rng = random.Random(0) problem = make_problem(rng, n_steps=4) model = ToyLLM(size=1e10) # 10B 参数的模拟模型，默认 CoT sample = model.generate(problem, rng) # 采样一次 print(problem.text(), sample.answer, sample.correct) 直接运行本文件会打印一道题、两个大小不同的模型每一步的概率，以及几次采样的结果： python3 toy_llm.py """ # 【Python 小知识】 # import math / random / unicodedata：Python 自带的"标准库"模块，不用另外安装。 # math ：数学函数（exp、log10、comb……） # random ：随机数。本课所有随机数都来自 random.Random(种子)，种子相同 → 结果完全相同（可复现）。 # unicodedata ：查询字符属性。这里只用来判断一个字符是不是宽字符（汉字），让终端里的表格对齐。 # from dataclasses import dataclass： # @dataclass 是一个"装饰器"（写在类定义上方、以 @ 开头）。被它装饰的类会自动生成 __init__ 等方法， # 适合定义"主要用来装数据"的类。@dataclass(frozen=True) 表示对象创建后不能修改（像常量一样）。 import math import random import unicodedata from dataclasses import dataclass # ============================================================================= # 全局常量（全大写的变量名是 Python 社区的约定：表示"常量，不要修改"） # ============================================================================= OPS = ("add", "sub", "mul") # 三种运算 OP_SYMBOL = {"add": "+", "sub": "-", "mul": "×"} # 打印题目时用的符号 OP_WEIGHTS = (0.40, 0.35, 0.25) # 出题时三种运算的比例 PROCEDURES = ("correct", "bug", "careless") # 每一步可能的三种做法 REF_SIZE = 1e9 # 参照规模：1B（10 亿）参数。log10(size / REF_SIZE) 在 1B 时为 0。 ALPHA = {"add": 2.0, "sub": 1.5, "mul": 0.5} # 1B 模型做"简单"一步时 correct 的 logit：加法最容易，乘法最难 BETA = 1.2 # 参数量每扩大 10 倍，correct 的 logit 增加多少（"scaling" 的强度） LOGIT_BUG = 0.0 # 固定的错误套路的 logit LOGIT_CARELESS = -0.5 # 随机手滑的 logit # hardness：让某些步骤更难（从 correct 的 logit 里减掉） HARD_CARRY = 0.6 # 加法需要进位（个位相加 ≥ 10） HARD_BORROW = 0.6 # 减法需要借位（被减数个位 < 减数个位） HARD_NEGATIVE = 0.8 # 减法结果是负数 HARD_MUL_2DIGIT = 1.5 # 乘数是两位数（10~15） HARD_MUL_BIG = 0.7 # 被乘数的绝对值 ≥ 100 DIRECT_PENALTY = 3.0 # direct 模式下，超出"一次前向传播能做的步数"的步骤要减掉的 logit CARELESS_OFFSETS = (-10, -2, -1, 1, 2, 10) # 随机手滑时偏离正确结果的量 # token 数的粗略估计（只用于估算成本，不影响对错） PROMPT_TOKENS_BASE = 12 # 题目开头的固定部分 PROMPT_TOKENS_PER_STEP = 4 # 题目里每一步大约 4 个 token（"× 12，"） COT_TOKENS_PER_STEP = 10 # CoT 模式下，写出每一步推理大约 10 个 token ANSWER_TOKENS = 6 # 最后报答案的部分 # ============================================================================= # 第一部分：题目（Problem） # ============================================================================= @dataclass(frozen=True) class Step: """ 题目里的一步运算，例如 "× 12"。 【字段】 op ：运算类型，"add" / "sub" / "mul" 之一 operand ：运算数（加减是 10~99，乘法是 2~15） """ op: str operand: int def apply_op(op, a, b): """ 正确地做一步运算：a op b。 【参数】op："add"/"sub"/"mul"；a：当前值（int，可以是负数）；b：运算数（int） 【返回】int 【边界】op 不是这三种之一时抛出 ValueError（防止拼写错误悄悄地产生错误结果）。 """ if op == "add": return a + b if op == "sub": return a - b if op == "mul": return a * b raise ValueError(f"未知的运算：{op!r}") def apply_bug(op, a, b): """ 按"固定的错误套路"（misconception）做一步运算。 【思路】每种运算一个典型错误，**确定性**的：同样的 a、b，bug 永远给出同一个错误结果。 add：a + b − 10 （十位上少算了 1，像"忘了进位"） sub：a − b + 10 （借位之后忘了减，结果多了 10） mul：a × (b − 1) （"少加了一份"：12 个 47 算成了 11 个 47） 【参数】同 apply_op 【返回】int 【边界】mul 的 bug 在 a = 0 时和正确结果相同（0 × 任何数都是 0）。这种情况下这一步"碰巧对了"， 评估时我们按**数值**判断对错，所以会被算作正确 —— 真实世界里也常有"方法错、答案对"的情况， 这正是第 3 讲 process verification（检查过程而不只是答案）的动机之一。 """ if op == "add": return a + b - 10 if op == "sub": return a - b + 10 if op == "mul": return a * (b - 1) raise ValueError(f"未知的运算：{op!r}") @dataclass(frozen=True) class Problem: """ 一道多步算术题。 【字段】 pid ：题号（int），只用来打印和区分 start ：起始数（10~99） steps ：Step 组成的 tuple。用 tuple 而不是 list，因为 frozen 的 dataclass 里放可修改的 list 会让"不可修改"名不副实。 【方法】 n_steps ：步数 k correct_values ：每一步之后的正确值（list） answer ：标准答案（上帝视角，只给 verifier 和评估用） text ：题目的文字形式 """ pid: int start: int steps: tuple # 【Python 小知识】@property 让一个方法可以像属性一样访问：写 problem.n_steps 而不是 problem.n_steps()。 @property def n_steps(self): return len(self.steps) def correct_values(self): values = [] x = self.start for step in self.steps: x = apply_op(step.op, x, step.operand) values.append(x) return values def answer(self): return self.correct_values()[-1] def text(self): # 【Python 小知识】"，".join(列表)：用"，"把一串字符串连起来。 parts = [f"{OP_SYMBOL[s.op]} {s.operand}" for s in self.steps] return f"从 {self.start} 开始：" + "，".join(parts) + "。最后得多少？" def sample_index(probs, rng): """ 按概率 probs 随机抽一个下标（0, 1, 2, ...）。 【思路】"轮盘赌"：把 [0, 1) 按概率切成几段，扔一个均匀随机数 u，看它落在哪一段。 例如 probs = [0.7, 0.2, 0.1]：u < 0.7 → 0；0.7 ≤ u < 0.9 → 1；否则 → 2。 【参数】probs：非负数的列表，和应为 1（允许有微小的浮点误差）；rng：random.Random 对象 【返回】int 下标 【边界】 · 浮点误差可能让累加和略小于 1，u 恰好落在最后的缝隙里 → 循环结束仍未返回， 这时返回最后一个概率 > 0 的下标，保证不会返回概率为 0 的选项。 · 为什么不用 random.choices？它在不同 Python 版本里的内部实现可能不同；自己写可以保证 Python 3.9 和 3.13 得到逐字相同的输出。 """ u = rng.random() cumulative = 0.0 for i, p in enumerate(probs): cumulative += p if u < cumulative: return i for i in range(len(probs) - 1, -1, -1): if probs[i] > 0: return i raise ValueError("probs 里没有正数") def make_problem(rng, n_steps, pid=0): """ 随机生成一道 n_steps 步的算术题。 【参数】rng：random.Random 对象；n_steps：步数（≥ 1）；pid：题号 【返回】Problem 【边界】n_steps < 1 时抛出 ValueError（0 步的题没有意义）。 乘法会让数字迅速变大（几步之后可能上万），Python 的 int 没有上限，不会溢出。 """ if n_steps < 1: raise ValueError("n_steps 必须 ≥ 1") start = rng.randint(10, 99) # 【Python 小知识】randint(a, b) 包含两端：10~99 都可能 steps = [] for _ in range(n_steps): # 【Python 小知识】用 _ 当循环变量，表示"这个值用不到" op = OPS[sample_index(OP_WEIGHTS, rng)] operand = rng.randint(2, 15) if op == "mul" else rng.randint(10, 99) steps.append(Step(op, operand)) return Problem(pid, start, tuple(steps)) def make_problem_set(rng, n_problems, min_steps, max_steps): """ 生成一组题目，每道题的步数在 [min_steps, max_steps] 之间均匀随机。 【参数】rng；n_problems：题数；min_steps ≤ max_steps：步数范围（都 ≥ 1） 【返回】Problem 的列表，题号从 0 开始 【边界】min_steps > max_steps 时 randint 会抛出 ValueError。 """ return [make_problem(rng, rng.randint(min_steps, max_steps), pid=i) for i in range(n_problems)] def step_hardness(op, a, b): """ 这一步有多"难"：返回要从 correct 的 logit 里减掉的量（≥ 0）。 【思路】人做算术时，进位、借位、负数、两位数乘法都更容易出错。我们假设模型也一样。 难度只和"这一步的输入"有关 —— 所以同一道题里，前面算错了，后面的难度也会变（输入变了）。 【参数】op；a：这一步的输入值；b：运算数 【返回】float 【边界】用 abs(a) % 10 取个位，负数也能处理（Python 里 -47 % 10 == 3，不是 7，所以要先取绝对值）。 """ hard = 0.0 if op == "add": if abs(a) % 10 + b % 10 >= 10: hard += HARD_CARRY elif op == "sub": if abs(a) % 10 < b % 10: hard += HARD_BORROW if a - b < 0: hard += HARD_NEGATIVE elif op == "mul": if b >= 10: hard += HARD_MUL_2DIGIT if abs(a) >= 100: hard += HARD_MUL_BIG return hard # ============================================================================= # 第二部分：数学小工具 # ============================================================================= def softmax(logits, temperature=1.0): """ 把一组 logit 变成概率：p_i ∝ exp(logit_i / temperature)。 【思路】这正是语言模型选下一个 token 的方式。temperature（温度）控制"多样性"： · temperature = 1：按模型本来的概率采样。 · temperature → 0：越来越集中到 logit 最大的选项上（greedy decoding）。 · temperature 很大：越来越接近均匀分布（什么都可能出现，包括胡言乱语）。 【参数】logits：数字列表；temperature：≥ 0 的 float 【返回】概率列表，和为 1 【边界】 · temperature == 0：直接返回 one-hot（logit 最大的为 1，其余为 0）；并列时取下标最小的 （在本模块里下标 0 是 correct，所以并列时算"做对"）。 · 数值稳定：先减去最大的 logit 再取 exp，避免 exp(很大的数) 溢出（OverflowError）。 减掉同一个常数不改变结果，因为分子分母同乘了 exp(−常数)。 """ if temperature < 0: raise ValueError("temperature 不能为负") if temperature == 0: best = max(range(len(logits)), key=lambda i: (logits[i], -i)) return [1.0 if i == best else 0.0 for i in range(len(logits))] m = max(logits) exps = [math.exp((z - m) / temperature) for z in logits] total = math.fsum(exps) return [e / total for e in exps] def mean(xs): """ 平均值。用 math.fsum 精确求和，保证不同 Python 版本输出相同。 【边界】空列表抛出 ValueError（平均值没有定义），而不是悄悄返回 0。 """ xs = list(xs) if not xs: raise ValueError("不能对空列表求平均") return math.fsum(xs) / len(xs) def pass_at_k(n, c, k): """ pass@k 的无偏估计（Chen et al. 2021, Codex 论文）：从 n 个样本中有 c 个正确， 估计"随机抽 k 个，至少有一个正确"的概率。 【思路】直接的做法是"只看前 k 个样本有没有对的"，但这样浪费了其余 n − k 个样本的信息，方差大。 更好的做法：考虑从 n 个里不放回地抽 k 个的**所有**组合。全都抽到错误样本的组合数是 C(n−c, k)， 总组合数是 C(n, k)，所以 pass@k = 1 − C(n−c, k) / C(n, k) 这对每道题是"真实 pass@k"的无偏估计（期望恰好相等）。 【参数】n：样本总数（≥ 1）；c：其中正确的个数（0 ≤ c ≤ n）；k：预算（1 ≤ k ≤ n） 【返回】0~1 的 float 【边界】 · n − c < k：错误样本不够 k 个，任何 k 个里一定有正确的 → 返回 1.0。 · c == 0：返回 0.0（C(n, k) / C(n, k) = 1）。 · k > n：无法估计，抛出 ValueError。想估计 pass@1000，就至少要采 1000 个样本。 【自我改进联系】Large Language Monkeys 等论文的 "coverage" 就是对所有题目的 pass@k 取平均。 【Python 小知识】math.comb(n, k) 是组合数 C(n, k)，结果是精确的整数（可以非常大）。 """ if not (1 <= k <= n) or not (0 <= c <= n): raise ValueError(f"参数不合法：n={n}, c={c}, k={k}") if n - c < k: return 1.0 return 1.0 - math.comb(n - c, k) / math.comb(n, k) def majority_vote(answers): """ 多数投票（majority voting，也叫 self-consistency）：返回出现次数最多的答案。 【参数】answers：答案列表（int 或任何可以当字典键的值） 【返回】(答案, 票数) 【边界】 · 空列表抛出 ValueError。 · 并列时返回**最先出现**的那个答案。这让结果可复现；真实系统里并列时的规则各不相同。 【Python 小知识】dict 会记住键的插入顺序（Python 3.7+），所以"最先出现"可以这样实现： 遍历计数字典时，第一个达到最大票数的就是最先出现的。 """ if not answers: raise ValueError("answers 不能为空") counts = {} for a in answers: counts[a] = counts.get(a, 0) + 1 # 【Python 小知识】dict.get(键, 默认值)：键不存在时返回默认值 best = max(counts.values()) for a, cnt in counts.items(): if cnt == best: return a, cnt def fmt_size(size): """ 把参数量写成人类习惯的形式：1e8 → "100M"，3.16e9 → "3.2B"，1e12 → "1T"。 【边界】小于 1e6 的直接写成整数。只用于打印。 """ for unit, value in (("T", 1e12), ("B", 1e9), ("M", 1e6)): if size >= value * 0.9999: x = size / value return f"{x:.0f}{unit}" if abs(x - round(x)) < 0.05 else f"{x:.1f}{unit}" return f"{size:.0f}" def display_width(s): """字符串在终端里占几格：汉字等宽字符算 2 格，其余算 1 格。""" return sum(2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1 for ch in s) def pad(s, width, align="left"): """ 把字符串补空格到指定的显示宽度，让中英文混排的表格对齐。 【参数】s：任何值（会先转成字符串）；width：目标宽度；align："left" 或 "right" 【边界】字符串比 width 还宽时原样返回（不截断）。 """ s = str(s) gap = max(0, width - display_width(s)) return s + " " * gap if align == "left" else " " * gap + s def bar(x, width=30): """ 用字符画一个横条：x ∈ [0, 1] → 0~width 个 "█"。在终端里"画图"用。 【边界】x 超出 [0, 1] 时截断到边界。 """ x = min(1.0, max(0.0, x)) return "█" * int(round(x * width)) # ============================================================================= # 第三部分：模拟的语言模型 # ============================================================================= @dataclass(frozen=True) class Sample: """ 模型的一次生成结果（一个 "sample"）。 【字段】 answer ：最终答案（int） correct ：最终答案是否等于标准答案（bool）。由 generate() 用上帝视角算好，只供评估使用。 values ：每一步算出来的值（tuple）。CoT 模式下这些就是"写出来的推理过程"。 procs ：每一步用的做法（"correct"/"bug"/"careless"/"tool"） local_ok ：每一步"就它自己的输入而言"是否算对（tuple of bool）。 注意：前面错了，后面的步骤仍然可能"局部正确"（在错误的输入上正确地运算）。 第 3 讲的 process reward model 检查的就是这种局部正确性。 n_tool_calls ：调用计算器的次数（包括失败的调用） """ answer: int correct: bool values: tuple procs: tuple local_ok: tuple n_tool_calls: int = 0 class ToyLLM: """ 模拟的语言模型。 【思路】见文件开头。核心只有一个公式：每一步从 (correct, bug, careless) 三种做法里按 softmax(logits / temperature) 抽一种；correct 的 logit 随参数量的对数线性增长。 【参数】 size ：参数量（> 0），例如 1e9 表示 1B。 cot ：True = chain-of-thought（写出每一步）；False = direct answer（直接报答案）。 name ：打印用的名字，默认由 size 生成（例如 "10B"）。 【边界】 · 这是一个**卡通**：真实模型的错误远比"三种做法"复杂，参数量也不是唯一的因素（数据、训练方法……）。 它只用来演示机制：scaling 让每一步更可靠、systematic error 和 random error 的区别、 temperature 与多样性、CoT 与串行计算、采样 + 验证的威力与局限。 · 模型是"无状态"的：每次 generate 互相独立（除了共用同一个 rng）。 第 6 讲会给它加上"可以训练的参数"，让它从自己的正确样本中学习（STaR）。 """ def __init__(self, size, cot=True, name=None): if size <= 0: raise ValueError("size 必须为正数") self.size = float(size) self.cot = cot self.name = name if name is not None else fmt_size(size) def depth_capacity(self): """ direct 模式下，一次前向传播里能"在脑子里"可靠完成的步数。 【思路】Transformer 的层数是固定的，一次前向传播能做的**串行**计算有限。不写中间步骤， 模型就必须把 k 步运算都压缩进这有限的层数里；写出中间步骤（CoT）相当于把计算"摊"到很多个 token 上，每个 token 只需要做一步。这里用一个随规模缓慢增长的上限来卡通化这一点： 100M → 1 步，1B → 2 步，10B → 3 步，100B → 4 步，1T → 5 步。 【返回】int（≥ 1） 【边界】size < 100M 时也返回 1（至少能做一步）。加 1e-9 是为了防止 log10(1e9/1e8) 这类计算 因为浮点误差得到 0.9999999 而被向下取整成 0。 """ return max(1, 1 + int(math.floor(math.log10(self.size / 1e8) + 1e-9))) def step_logits(self, op, a, b, step_index): """ 某一步三种做法的 logits：[correct, bug, careless]。 【参数】op；a：这一步的输入值；b：运算数；step_index：这是第几步（从 0 开始） 【返回】长度为 3 的 float 列表 【边界】只有 direct 模式、且 step_index ≥ depth_capacity() 时才加 DIRECT_PENALTY。 """ z = ALPHA[op] + BETA * math.log10(self.size / REF_SIZE) - step_hardness(op, a, b) if not self.cot and step_index >= self.depth_capacity(): z -= DIRECT_PENALTY return [z, LOGIT_BUG, LOGIT_CARELESS] def step_probs(self, op, a, b, step_index, temperature=1.0): """某一步三种做法的概率（softmax 之后）。参数同 step_logits，外加 temperature。""" return softmax(self.step_logits(op, a, b, step_index), temperature) def generate(self, problem, rng, temperature=1.0, tool_ops=(), tool_fail_prob=0.0): """ 对一道题生成一次回答（采样一次）。 【思路】从起始数出发，一步一步地算：每一步先按概率选做法，再按做法算出这一步的值； 下一步的输入是**上一步自己算出来的值**（不是正确值）—— 一步错，后面全跟着错， 这就是多步任务里 "error compounding"（错误累积）的来源。 【参数】 problem ：Problem rng ：random.Random 对象（决定这次采样的随机性） temperature ：采样温度（≥ 0）。0 = greedy，每次结果完全相同。 tool_ops ：哪些运算交给计算器，例如 ("mul",)。默认不用工具。 tool_fail_prob ：每次调用计算器失败的概率（0~1），失败时模型只好自己算。 【返回】Sample 【边界】 · 计算器成功时这一步一定正确，procs 里记为 "tool"。 · temperature = 0 时每一步都选 logit 最大的做法；如果某一步 bug 的 logit 比 correct 高， greedy 解码会**每次都**犯同一个错 —— 多采样也救不了，必须提高 temperature 才有机会做对。 【Python 小知识】tool_ops=() 用空 tuple 做默认值。不要用空 list 做默认值（def f(x=[])）， 因为 list 是可修改的，默认值只创建一次，修改会"串"到下一次调用里；tuple 不可修改，安全。 """ x = problem.start values, procs, local_ok = [], [], [] n_tool_calls = 0 for i, step in enumerate(problem.steps): right = apply_op(step.op, x, step.operand) # 在"当前输入"上的正确结果 if step.op in tool_ops: n_tool_calls += 1 if rng.random() >= tool_fail_prob: # 工具调用成功 x = right values.append(x) procs.append("tool") local_ok.append(True) continue # 【Python 小知识】continue：跳过本轮剩下的代码，进入下一轮循环 probs = self.step_probs(step.op, x, step.operand, i, temperature) proc = PROCEDURES[sample_index(probs, rng)] if proc == "correct": new_x = right elif proc == "bug": new_x = apply_bug(step.op, x, step.operand) else: # careless：在正确结果上随机偏一点 new_x = right + CARELESS_OFFSETS[rng.randrange(len(CARELESS_OFFSETS))] values.append(new_x) procs.append(proc) local_ok.append(new_x == right) x = new_x return Sample(answer=x, correct=(x == problem.answer()), values=tuple(values), procs=tuple(procs), local_ok=tuple(local_ok), n_tool_calls=n_tool_calls) def n_tokens(self, problem): """ 一次生成大约要处理多少个 token（题目 + 输出）。只用于估算成本。 【返回】int。CoT 模式的输出长得多（每步约 10 个 token），direct 模式只输出答案。 """ prompt = PROMPT_TOKENS_BASE + PROMPT_TOKENS_PER_STEP * problem.n_steps output = ANSWER_TOKENS + (COT_TOKENS_PER_STEP * problem.n_steps if self.cot else 0) return prompt + output def flops_per_sample(self, problem): """ 一次采样的计算量（FLOPs）≈ 2 × 参数量 × token 数。 【思路】Transformer 推理时，每个 token 要和每个参数做一次乘法、一次加法，约 2N 次浮点运算。 （忽略了 attention 里和上下文长度有关的部分；对短文本这是很好的近似。） 【自我改进联系】test-time compute 的"预算"通常就用 FLOPs 或"采样次数 × 模型大小"来衡量。 比较"小模型采很多次"和"大模型采一次"时，必须在**相同的 FLOPs**下比较才公平（第 02 号文件）。 """ return 2.0 * self.size * self.n_tokens(problem) # ============================================================================= # 直接运行本文件时的小演示 # ============================================================================= def _demo(): rng = random.Random(0) problem = make_problem(rng, n_steps=4, pid=0) print("一道题：", problem.text()) print("标准答案（上帝视角）：", problem.answer(), "；每一步的正确值：", problem.correct_values()) print() for size in (1e9, 1e11): model = ToyLLM(size) print(f"模型 {model.name}（CoT）每一步选三种做法的概率（temperature = 1，输入取正确的前一步值）：") x = problem.start for i, step in enumerate(problem.steps): probs = model.step_probs(step.op, x, step.operand, i) print(f" 第 {i + 1} 步 {x} {OP_SYMBOL[step.op]} {step.operand}： " f"correct {probs[0]:.2f} bug {probs[1]:.2f} careless {probs[2]:.2f}") x = apply_op(step.op, x, step.operand) print(f" 5 次采样：") for _ in range(5): s = model.generate(problem, rng) mark = "✓" if s.correct else "✗" print(f" {mark} 答案 {s.answer:>6d} 做法 {', '.join(s.procs)}") print() # 【Python 小知识】if __name__ == "__main__": # 只有"直接运行这个文件"时，下面的代码才会执行；被别的文件 import 时不会执行。 if __name__ == "__main__": _demo()