第 6 章推理与采样

文档摘要

第 6 章推理与采样本章目标：搞懂自回归生成 + temperature + top-k 采样的实现细节。 6.1 章节地图推理层干三件事： 6.2 生成封装类两个关键设置：禁用自动求导。推理时不需要反向传播，禁用后不存中间激活，显存占用大幅降低（可省 5-10 倍），速度提升（少了构建计算图的开销）。关闭 dropout，保证推理结果确定性（除了采样本身）。 💡 这两个加在一起是推理的标配。忘任何一个都有问题：忘 nograd 显存爆；忘 eval 推理结果每次不同。为什么封装成类状态封装：模型、编码器、设备、配置都存在 self 上，不用每次都传一遍。复用友好：Web UI 能持有一个实例反复调用。接口清晰：外部只调，不关心内部细节。 6.

第 6 章推理与采样

本章目标：搞懂自回归生成 + temperature + top-k 采样的实现细节。

6.1 章节地图

推理层干三件事：


┌─────────────────────────────────────────────────────────┐
│  1. 生成封装类    封装模型 + 编码器 + generate()         │
│  2. generate() 实现    自回归循环 + 三种采样策略         │
│  3. 智能加载    自动从 checkpoint 恢复配置               │
└─────────────────────────────────────────────────────────┘

6.2 生成封装类


class GPTGenerator:
    def __init__(self, model, gpt_config, device=None):
        self.model = model
        self.config = gpt_config
        self.encoder = get_encoder()
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.model.eval()   # ★ 推理模式：关 dropout

    @torch.no_grad()        # ★ 推理不建计算图，省显存、加速
    def generate(self, prompt, max_new_tokens=100, temperature=0.8, top_k=40):
        ...

两个关键设置：

`@torch.no_grad()`

禁用自动求导。推理时不需要反向传播，禁用后不存中间激活，显存占用大幅降低（可省 5-10 倍），速度提升（少了构建计算图的开销）。

`model.eval()`

关闭 dropout，保证推理结果确定性（除了采样本身）。

💡 这两个加在一起是推理的标配。忘任何一个都有问题：忘 no_grad 显存爆；忘 eval 推理结果每次不同。

为什么封装成类

状态封装：模型、编码器、设备、配置都存在 self 上，不用每次都传一遍。
复用友好：Web UI 能持有一个实例反复调用。
接口清晰：外部只调 generate(prompt, ...)，不关心内部细节。

6.3 自回归生成主循环


# 编码 prompt
if not prompt:
    input_ids = torch.tensor([[self.config.vocab_size - 1]], ...)  # <|endoftext|>
else:
    tokens = self.encoder.encode(prompt)
    input_ids = torch.tensor([tokens], ...)

for _ in range(max_new_tokens):
    # ① 截断到 block_size
    cond = input_ids if input_ids.size(1) <= block_size else input_ids[:, -block_size:]
    # ② 前向，取最后一位的 logits
    logits = self.model(input_ids=cond).logits
    next_logits = logits[:, -1, :]
    # ③ 采样（见 6.4）
    next_token = ...
    # ④ 拼接
    input_ids = torch.cat([input_ids, next_token], dim=1)
    # ⑤ 遇到结束符提前终止
    if next_token.item() == self.config.vocab_size - 1 and not prompt:
        break

return self.encoder.decode(input_ids[0].tolist())

5 步逐个讲

步骤 ①：截断到 block_size。模型最大上下文是 block_size。生成过程中序列越来越长，一旦超过，必须截掉前面的，只保留最后 block_size 个 token。input_ids[:, -block_size:] 是切片，负索引表示从末尾往前数。

⚠️ 这种朴素实现的性能问题：每生成一个 token 都要把整段 cond（最长 block_size）重新前向一遍，复杂度 O(block_size²)。生产中用 KV Cache（缓存每层 K/V 避免重算），框架的 model.generate() 默认就用。本项目手写循环是为了教学清晰。

步骤 ②：取最后一位的 logits。自回归生成时，我们只关心「给定当前全部 token，下一个 token 是什么」。位置 -1（最后一个 token）的 logits 就是这个预测。前面位置的 logits 在生成场景下是「废物」（它们预测的是「已经知道的下一个」），丢弃即可。

步骤 ③：采样。见 6.4 节，三种策略。

步骤 ④：拼接。把新 token 拼到序列末尾，下一轮循环它就成了已知上下文。

6.4 三种采样策略（重难点）

分支 A：贪心解码（`temperature <= 0`）


if temperature <= 0:
    next_token = torch.argmax(next_logits, dim=-1, keepdim=True)

每步取概率最大的 token。确定性强（同输入永远同输出），但容易陷入重复循环：


prompt: "To be, or"
输出:  "To be, or not to be not to be not to be not to be..."

分支 B：温度采样（`temperature > 0`）


next_logits = next_logits / temperature

温度的数学本质：softmax(logits/T)。

T → 0⁺：分布越来越尖，趋近 argmax（贪心）。
T = 1：原分布不变。
T → ∞：分布越来越平，趋近均匀随机（乱码）。

直觉：温度「平滑」或「锐化」概率分布。低温保守确定，高温奔放随机。

温度的数学演示：


logits = [2.0, 1.0, 0.5]

T=0.1:  softmax([20, 10, 5])    → [0.9999, 0.0001, 0.0000]  ← 几乎必选第一个
T=1.0:  softmax([2, 1, 0.5])    → [0.59,  0.22,  0.13]      ← 偶尔选其他
T=2.0:  softmax([1, 0.5, 0.25]) → [0.39, 0.28, 0.21]        ← 经常选其他

分支 C：Top-K 过滤（`top_k > 0`）


if top_k > 0:
    top_k = min(top_k, next_logits.size(-1))
    top_values, top_indices = torch.topk(next_logits, top_k, dim=-1)
    mask = torch.full_like(next_logits, float("-inf"))
    mask.scatter_(1, top_indices, top_values)
    next_logits = mask

Top-K 的作用

只保留概率最高的 K 个 token，其余 logits 设为 -inf（softmax 后概率为 0），再从这 K 个里按概率采样。

为什么要 top-k

纯温度采样的痛点：长尾。即使温度合适，词表里仍有大量「明显不对」的 token，它们概率很小但偶尔被抽中会让生成崩坏。Top-K 把这些「无希望」的候选直接踢掉。


词表 50257 个 token 的概率分布（示意）

排序后：[0.30, 0.20, 0.15, 0.10, 0.05, ... 一堆 0.001 ...]

top_k=40：只在前 40 个里采样，其余强制为 0
                ▲
                │ 这些是「合理候选」
   ────────────┼──────────────────────
                │
                │ 这些是「噪声」，丢弃

代码逐行

top_k = min(top_k, next_logits.size(-1))：钳制，万一用户传 99999，但词表只有 50257，取 min 防越界。torch.topk 返回最大的 K 个值及其位置。构造全 -inf 的 mask，用 scatter_ 把 top K 个位置填回原值。probs = F.softmax(next_logits, dim=-1)：-inf 位置的概率是 0（因为 exp(-inf) = 0），top K 位置正常分布。

`torch.multinomial` 采样


probs = F.softmax(next_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)

multinomial 按 probs 给的概率分布有放回抽样 1 次。这正是「按概率采一个 token」的标准实现。

三者组合的实际效果

组合	行为
`temperature=0`	贪心，确定但易重复
`temperature=1, top_k=0`	纯随机采样，混乱
`temperature=0.8, top_k=40`（默认）	推荐：有变化但不离谱
`temperature=0.3, top_k=10`	保守、聚焦、更连贯
`temperature=1.2, top_k=80`	多样、发散、易跑题

调参经验：

写代码补全：低温度（0.2-0.4），需要确定性和正确性。
写诗/创意：中温度（0.7-1.0），需要一点惊喜。
brainstorming：高温度（1.0-1.3），最大化多样性。
几乎所有场景都开 top_k=40，防止采到离谱 token。

6.5 智能加载


def load_for_inference(checkpoint_path, device=None, gpt_config=None):
    target_device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))

    if os.path.isdir(checkpoint_path):
        # 格式 A：transformers 风格目录
        model = GPT2LMHeadModel.from_pretrained(checkpoint_path)
    else:
        # 格式 B：本项目 .pt 文件
        ckpt = torch.load(checkpoint_path, map_location=target_device)
        # 优先用 checkpoint 里保存的配置（重要！）
        if "gpt_config" in ckpt:
            cfg = GPTConfig()
            cfg.__dict__.update(ckpt["gpt_config"])
            gpt_config = cfg

        model = build_model(gpt_config)
        state = ckpt.get("model_state_dict", ckpt)
        model.load_state_dict(state, strict=False)

    return GPTGenerator(model=model, gpt_config=gpt_config, device=target_device)

智能配置恢复

最关键的几行：


if "gpt_config" in ckpt:
    cfg = GPTConfig()
    cfg.__dict__.update(ckpt["gpt_config"])
    gpt_config = cfg

亮点：从 checkpoint 自动恢复配置。训练时你改了 n_layer=12，checkpoint 里就存了 n_layer=12；推理时不需要再传 --n-layer 12，加载器自己读出来。这是训练入口把配置存进 checkpoint 的回报（见第 5 章）。

6.6 命令行入口


def main():
    args = parse_args()
    generator = load_for_inference(args.checkpoint, device=args.device)
    text = generator.generate(
        prompt=args.prompt, max_new_tokens=args.max_tokens,
        temperature=args.temperature, top_k=args.top_k,
    )
    print("\n===== 生成结果 =====")
    print(text)

常用命令


# 默认
python inference.py --prompt "To be, or"

# 贪心（确定）
python inference.py --prompt "ROMEO:" --temperature 0

# 多样化
python inference.py --prompt "ROMEO:" --temperature 1.0 --top-k 50 --max-tokens 200

# 指定 checkpoint
python inference.py --checkpoint checkpoints/gpt_step2000.pt --prompt "To be, or"

参数表

参数	类型	默认值	说明
`--checkpoint`	str	`checkpoints/gpt_final.pt`	权重路径或 transformers 目录
`--prompt`	str	`""`	输入提示文本
`--max-tokens`	int	100	最多生成的新 token 数
`--temperature`	float	0.8	采样温度
`--top-k`	int	40	top-k 采样的 k 值
`--device`	str	-	推理设备

6.7 Python 中调用

不通过命令行，直接导入使用：


from inference import load_for_inference

gen = load_for_inference("checkpoints/gpt_final.pt")

# 多次生成
for _ in range(3):
    text = gen.generate(prompt="To be, or", max_new_tokens=100,
                        temperature=0.8, top_k=40)
    print(text)
    print("---")

这是 Web UI 的用法（见第 7 章）。

6.8 动手实验

确定性 vs 随机：同一个 prompt 跑 5 次——--temperature 0 观察输出是否完全一致（贪心确定）；--temperature 1.0 --top-k 0 观察是否几乎每次不同且混乱。
温度扫描：固定 prompt，跑 0.2 / 0.5 / 0.8 / 1.2 几个温度，对比效果。
长 prompt 测试：给一个超过 block_size 的长 prompt，观察生成时是否报错（应该不会，因为有截断）。
思考题：如何实现 top-p（nucleus）采样？提示：不再固定 K，而是选概率累加 ≥ p 的最小集合。
进阶：用框架的 model.generate(use_cache=True) 对比本项目的手写循环，体会 KV Cache 的加速。

6.9 小结

生成封装类封装模型+编码器+设备，提供 generate() 接口。
@torch.no_grad() + model.eval() 是推理标配（省显存 + 关 dropout）。
自回归循环 5 步：截断 → 前向取末位 logits → 采样 → 拼接 → 检查结束符。
三种采样：贪心（temperature<=0，确定易重复）；温度（temperature>0，softmax(logits/T) 控制随机性）；Top-K（踢掉长尾噪声）。
默认组合 temperature=0.8, top_k=40 是安全选择。
智能加载从 checkpoint 自动恢复配置，无需手动传架构。

6.10 下一章

推理逻辑通了，去《第 7 章 Web UI》把它包成浏览器可访问的界面。

第 6 章 推理与采样

文档摘要