一句话总结：推理就是用 checkpoint 重建同一个模型和 tokenizer，把 prompt 编码成 token IDs，再一个 token 一个 token 地生成。代码不长，但模型契约、随机性和长度边界一个都不能丢。

从受信任的 checkpoint 重建模型与 tokenizer，将 prompt 编码后自回归生成，最后解码成文本

📦 2024 年原始演示仓库：github.com/waylandzhang/Transformer-from-scratch。历史 demo 把数据、模型、训练和生成放在一个 model.py 里；本章的 inference.py 与第 18、19 章审校后的独立文件配套。

20.1 推理不是“不训练”这么简单

	训练	推理
输入	连续序列与右移一位的 targets	prompt，以及已生成的 tokens
每次 forward	同时对多个位置计算 loss	只取最后一个位置来选下一 token
参数更新	有	没有
Dropout	`train()` 下随机丢弃	`eval()` 下关闭
Autograd	需要	用 `inference_mode()` 关闭

model.eval() 和 torch.inference_mode() 不是一件事。前者切换模块的训练/评估行为；后者不记录 autograd 状态。本书模型有 Dropout，没有 BatchNorm，所以这里 eval() 的直接作用是关闭 Dropout，不是“固定 BatchNorm”。

eval() 也不会把采样变成确定性。temperature > 0 时仍会从概率分布采样；temperature = 0 才是第 18 章定义的 greedy decoding。

20.2 先重建模型契约，再加载权重

第 19 章的 checkpoint 已经包含 model_config、tokenizer_name 和 model_state_dict。推理不应该在脚本里再猜一遍层数、宽度或词表大小。

def load_model(checkpoint_path, device):
    checkpoint = torch.load(
        checkpoint_path,
        map_location="cpu",
        weights_only=True,
    )
    required = {"model_state_dict", "model_config", "tokenizer_name"}
    if not isinstance(checkpoint, dict) or not required <= checkpoint.keys():
        raise ValueError("checkpoint is missing inference metadata")

    model_config = ModelConfig(**checkpoint["model_config"])
    tokenizer = tiktoken.get_encoding(checkpoint["tokenizer_name"])
    if tokenizer.n_vocab != model_config.vocab_size:
        raise ValueError("tokenizer vocabulary does not match the model")

    model = Model(model_config)
    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
    model.to(device)
    model.eval()
    return model, tokenizer

先 map_location="cpu" 再把重建好的模型搬到当前设备，不会因为 checkpoint 当年在某张 GPU 上保存，就强迫今天的机器找同一个设备。strict=True 要求权重名与 shape 完整对上。

weights_only=True 会限制 unpickler 能构造的对象，但它不是“任意下载文件都安全”的承诺。仍然只加载自己创建或来源可信的 checkpoint。

20.3 Prompt 也要遵守 tokenizer 契约

继续用原来的商品名称语料：

prompt = "农夫山泉 "
prompt_ids = tokenizer.encode(prompt, disallowed_special=())
if not prompt_ids:
    raise ValueError("prompt must contain at least one token")
x = torch.tensor(prompt_ids, dtype=torch.long, device=device)[None, :]

这里不写死 token IDs。Tokenizer 版本、前后空格、甚至一个标点都可能改变结果，应该让当前实际 tokenizer 打印它。

如果 prompt 超过 context_length，第 18 章的教学版 generate() 会在每一步只取最后 context_length 个 token。较早部分仍保留在最终输出 tensor 里，但已经不再影响新 token。

本书用的是固定正弦位置表。裁成新的可见窗口后，窗口内的位置编号也会重新从 0 开始。这是教学实现的明确约定，并非所有 serving system 的默认行为。

20.4 生成参数有边界，没有万能推荐值

20.4.1 Temperature

第 18 章的定义是：

temperature = 0：直接取最大 logit，是 greedy decoding；
0 < temperature < 1：分布更尖，高 logit token 更容易被选中；
temperature = 1：不改变 logits 的相对尺度；
temperature > 1：分布更平。

它不是“事实问答用 0.2、创意写作用 0.9”这种固定处方。小模型的 logits 校准、语料与任务都会改变合适值。

20.4.2 Top-K

top_k=K 先保留 logit 最高的 K 个 token，再在它们之间 Softmax 与采样。本书代码会把大于词表的 K 截成词表大小，None 表示不过滤。Top-K 会去掉长尾，却不保证文本正确、不重复或连贯。

20.4.3 Max New Tokens 与停止条件

这个教学模型没有定义 EOS token，也没有 stop string。因此 max_new_tokens=80 就会生成恰好 80 个新 token；它是硬上限，不是“最多生成，模型会自己停”。

20.5 完整 `inference.py`

import argparse
from pathlib import Path

import tiktoken
import torch

from model import Model, ModelConfig


def select_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def load_model(checkpoint_path, device):
    checkpoint = torch.load(
        checkpoint_path,
        map_location="cpu",
        weights_only=True,
    )
    required = {"model_state_dict", "model_config", "tokenizer_name"}
    if not isinstance(checkpoint, dict) or not required <= checkpoint.keys():
        raise ValueError("checkpoint is missing inference metadata")

    model_config = ModelConfig(**checkpoint["model_config"])
    tokenizer = tiktoken.get_encoding(checkpoint["tokenizer_name"])
    if tokenizer.n_vocab != model_config.vocab_size:
        raise ValueError("tokenizer vocabulary does not match the model")

    model = Model(model_config)
    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
    model.to(device)
    model.eval()
    return model, tokenizer


def encode_prompt(tokenizer, prompt, device):
    prompt_ids = tokenizer.encode(prompt, disallowed_special=())
    if not prompt_ids:
        raise ValueError("prompt must contain at least one token")
    return torch.tensor(
        prompt_ids, dtype=torch.long, device=device
    ).unsqueeze(0)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=Path("model/checkpoint.pt"),
    )
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--max-new-tokens", type=int, default=80)
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--top-k", type=int)
    parser.add_argument("--seed", type=int, default=1337)
    return parser.parse_args()


def main():
    args = parse_args()
    device = select_device()
    model, tokenizer = load_model(args.checkpoint, device)
    x = encode_prompt(tokenizer, args.prompt, device)

    torch.manual_seed(args.seed)

    with torch.inference_mode():
        y = model.generate(
            x,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            top_k=args.top_k,
        )

    parameter_count = sum(p.numel() for p in model.parameters())
    print(f"device={device} prompt_tokens={x.size(1)}")
    print(f"parameters={parameter_count:,}")
    print("---")
    print(tokenizer.decode(y[0].tolist()))


if __name__ == "__main__":
    main()

运行：

python inference.py --checkpoint model/checkpoint.pt --prompt "农夫山泉 " --temperature 0.7 --top-k 50

输出应该由当前 checkpoint 真实生成。旧版的四行商品名称没有对应可复现 checkpoint，所以本版不再把它写成实验结果。

20.6 想看每一步，就打印真实候选

不要手写一组看起来很合理的概率。可以直接从当前模型取候选：

@torch.inference_mode()
def show_next_candidates(model, tokenizer, x, k=5):
    x_crop = x[:, -model.config.context_length:]
    logits, _ = model(x_crop)
    probs = torch.softmax(logits[0, -1], dim=-1)
    top_probs, top_ids = torch.topk(probs, min(k, probs.numel()))

    for probability, token_id in zip(top_probs, top_ids):
        piece = tokenizer.decode_single_token_bytes(
            token_id.item()
        ).decode("utf-8", errors="backslashreplace")
        print(repr(piece), f"{probability.item():.4f}")

单个 token 可能只是一段 UTF-8 bytes，不一定能独立解码成完整字符。因此这里先取 decode_single_token_bytes()，而不是假设每个 token 就是一个“词”。

20.7 常见问题要从边界开始查

权重名或 shape 不匹配：不要改成 strict=False 把问题藏起来，先确认 model.py 与 checkpoint 是否来自同一契约。
输出重复：可能来自语料重复、过拟合、校准差或采样设置。调大 temperature 只是试验，不是保证。
输出不连贯：先看真实 train/valid loss、prompt 是否接近训练分布、tokenizer 是否一致，再调 temperature 和 Top-K。
生成慢：本章的 generate() 每次都重算当前窗口。GPU 可能有帮助，但核心冗余要到第 22 章用 KV Cache 才消掉；Flash Attention 改善的是每次 Attention 的内存 I/O 与计算实现，两者不是同一件事。

20.8 本章总结

从 checkpoint 的 model_config 和 tokenizer_name 重建契约，不在推理脚本里猜；
权重先安全地加载到 CPU，再把模型搬到当前设备；
eval() 处理模块行为，inference_mode() 处理 autograd，采样随机性由 decoding 规则决定；
Temperature、Top-K 和长度都有精确边界，不应写成普遍适用的经验处方；
本章无 EOS，会生成恰好 max_new_tokens 个 token；
不编造 token IDs、概率、参数量或生成结果，全部让当前代码与 checkpoint 给出。

本章交付物

能用 checkpoint 重建完全一致的 Model 和 tokenizer
能区分 eval()、inference_mode() 与 greedy/sampling
能说清 prompt 截断、Temperature、Top-K 和生成长度的边界
能用当前 checkpoint 运行 inference.py，而不依赖预写结果

Part 5 总结

第 18 章定义模型，第 19 章训练并保存完整状态，第 20 章只取推理需要的部分来生成文本。三个文件加起来已经超过旧版声称的“不到 400 行”，因为本版补上了配置校验、数据边界和恢复状态。行数不是目标：它仍是教学实现，但每一个简化都应该说清。

理解这条链路后，你能看懂 GPT 式 decoder 的主干，但不等于已经复制了 ChatGPT 或生产级 serving system。真实系统还要处理 KV Cache、batching、定制 kernel、并行、内存管理与服务调度。

下一章预告

现在的生成循环每走一步，都会重新计算当前窗口里的 Attention。第 21 章先看 Flash Attention：它不改变 Attention 的数学结果，而是重排计算与内存访问，让同一次 Attention 做得更高效。第 22 章再用 KV Cache 消掉 decoding 步之间的 K/V 重算。