从零实现 nanoGPT（五）：保存模型，写一个独立的文本生成程序

上一篇，我们实现了一个小型 GPT，将多头注意力、前馈神经网络、残差连接和 LayerNorm 组合成 Transformer Block，并通过堆叠多个 Block，完成了 GPT 的主体结构。

到目前为止，我们的模型已经能够完成两个任务：

训练：根据输入字符预测下一个字符，并通过反向传播更新参数。
生成：给定一段起始文本，让模型不断预测下一个字符，最终生成一段新文本。

但是，目前还有一个实际问题：

训练结束后，怎样把模型保存下来，并在以后重新使用？

如果每次生成文本都需要重新训练模型，显然不现实。尤其是当模型规模变大、训练时间从几分钟增长到几小时甚至更长时，保存模型就成为必不可少的一步。

因此，这一篇我们暂时不改变 GPT 的主体结构，而是完善它的使用流程。

我们将完成以下几个任务：

使用 torch.save() 保存模型检查点。
理解模型参数与模型结构之间的关系。
在原有生成方法中加入 temperature 和 top-k 采样。
编写独立的 sample.py，加载模型并生成文本。
验证保存前后的模型是否一致。

完成之后，我们就能做到：

训练一次模型，保存到磁盘，以后直接加载并生成文本，而不需要重新训练。

本文继续使用前四篇构建的莎士比亚字符级语言模型。

一、为什么需要保存模型？

回顾一下我们的训练过程。

在 train.py 中，我们先创建模型：

scss 复制代码

model = GPTLanguageModel(**model_args).to(device)

然后不断执行：

scss 复制代码

logits, loss = model(x, y)

optimizer.zero_grad()
loss.backward()
optimizer.step()

随着训练进行，模型内部的参数会不断更新。

例如，我们之前实现的字符嵌入层：

ini 复制代码

self.token_embedding_table = nn.Embedding(
    vocab_size,
    n_embd
)

它内部维护着一个可学习的权重矩阵：
$Wembedding∈RV×d W_{\text{embedding}} \in \mathbb{R}^{V\times d}$ Wembedding∈RV×d

其中：

$VV$ V 表示字符词表大小。
$dd$ d 表示嵌入维度。

训练过程中，优化器不断调整这些参数，让模型预测的字符分布更加接近训练数据。

但这些参数默认保存在内存中。

当 Python 程序结束后，模型对象和对应的参数也就不再存在于程序内存中。

因此，我们需要将训练得到的参数保存到磁盘。

1.1 state_dict 是什么？

PyTorch 为我们提供了一个非常方便的方法：

scss 复制代码

model.state_dict()

它会返回一个类似字典的对象，保存模型中需要持久化的参数和缓冲区。

例如：

scss 复制代码

for name, param in model.state_dict().items():
    print(name, param.shape)

输出中可能包含：

erlang 复制代码

token_embedding_table.weight
position_embedding_table.weight
blocks.0.sa.heads.0.key.weight
blocks.0.sa.heads.0.query.weight
blocks.0.sa.heads.0.value.weight
...
ln_f.weight
ln_f.bias
lm_head.weight
lm_head.bias

具体名称取决于我们在 model.py 中定义的模块名称。

这里的每一项通常对应一个张量。

例如：

复制代码

token_embedding_table.weight

表示字符嵌入层的权重。

而：

erlang 复制代码

blocks.0.sa.heads.0.query.weight

表示第一个 Transformer Block 中，第一个注意力头的 Query 投影矩阵。

此外，state_dict() 还会保存通过 register_buffer() 注册且 persistent=True 的缓冲区。

例如，如果我们在因果注意力中注册了遮罩：

less 复制代码

self.register_buffer(
    "tril",
    torch.tril(torch.ones(block_size, block_size))
)

那么这个遮罩也会出现在 state_dict() 中。

需要注意，buffer 不一定是可训练参数，但可能是模型运行所需要的状态。

1.2 只有参数就够了吗？

假设我们直接保存：

less 复制代码

torch.save(model.state_dict(), "model.pt")

当然可以。

但以后加载时，我们仍然需要知道模型的结构。

例如：

ini 复制代码

GPTLanguageModel(
    vocab_size=65,
    n_embd=64,
    n_head=4,
    n_layer=2,
    block_size=64,
    dropout=0.1,
)

如果重新创建的模型结构和保存参数时的模型不一致，加载过程就可能失败。

此外，我们还需要保存字符与 ID 之间的映射。

回顾第一篇，我们通过字符词表构造：

css 复制代码

stoi = {
    ch: i
    for i, ch in enumerate(chars)
}

itos = {
    i: ch
    for i, ch in enumerate(chars)
}

其中：

stoi：将字符转换为整数 ID。
itos：将整数 ID 转换为字符。

例如：

arduino 复制代码

"R" → 24
"O" → 21
"M" → 19

这些数字仅用于说明，并不代表实际训练词表中的编号。

模型真正学习的是字符 ID 之间的关系。

如果训练时字符 "R" 对应 ID 24，而生成时错误地把它映射为 ID 18，模型收到的输入就不是原来那个字符。

因此，即使模型参数成功加载，也可能产生错误的结果。

所以，我们至少需要保存三类信息：

内容	用途
模型参数 `model`	恢复训练得到的权重
模型配置 `model_args`	重新创建相同结构的模型
字符映射 `meta`	保证编码和解码规则一致

本篇使用的检查点结构为：

css 复制代码

checkpoint = {
    "model": model.state_dict(),
    "model_args": model_args,
    "meta": meta,
}

这样，一个文件就包含了独立生成文本所需要的主要信息。

这里暂时不保存优化器状态和训练步数，因为本篇的目标是模型推理。

下一篇实现断点续训时，再完善训练状态的保存。

二、在 train.py 中保存模型

首先，我们需要在之前的训练程序中加入保存逻辑。

2.1 整理模型配置

上一篇创建 GPT 时，我们使用了多个超参数。

为了方便保存和重新加载，可以把它们统一放在一个字典中：

makefile 复制代码

model_args = {
    "vocab_size": meta["vocab_size"],
    "n_embd": 64,
    "n_head": 4,
    "n_layer": 2,
    "block_size": 64,
    "dropout": 0.1,
}

然后通过字典解包创建模型：

scss 复制代码

model = GPTLanguageModel(**model_args).to(device)

这里的：

复制代码

**model_args

相当于：

ini 复制代码

model = GPTLanguageModel(
    vocab_size=meta["vocab_size"],
    n_embd=64,
    n_head=4,
    n_layer=2,
    block_size=64,
    dropout=0.1,
)

这是一种 Python 字典解包写法。

使用字典保存参数还有一个好处：以后修改模型结构时，只需要维护这一份配置。

注意，这里的 model_args 应当与实际训练时使用的参数完全一致，不能在保存时随意更改。

2.2 保存检查点

在 train.py 顶部增加：

javascript 复制代码

from pathlib import Path

然后在训练循环结束后加入：

ini 复制代码

# 创建输出目录
out_dir = Path(__file__).resolve().parent / "out"
out_dir.mkdir(parents=True, exist_ok=True)

# 构造检查点
checkpoint = {
    "model": model.state_dict(),
    "model_args": model_args,
    "meta": meta,
}

# 保存模型
checkpoint_path = out_dir / "model.pt"

torch.save(checkpoint, checkpoint_path)

print(f"模型已保存到：{checkpoint_path}")

下面解释几个关键位置。

第一步：创建输出目录

ini 复制代码

out_dir = Path(__file__).resolve().parent / "out"

这里：

__file__：当前 Python 脚本的位置。
resolve()：获取对应的绝对路径。
parent：当前脚本所在目录。
/ "out"：在该目录下拼接一个 out 文件夹。

假设项目目录是：

kotlin 复制代码

my-nanogpt/
├── train.py
├── model.py
└── data/

那么最终输出目录就是：

bash 复制代码

my-nanogpt/out/

这种写法比直接使用：

ini 复制代码

out_dir = "out"

更加稳定。

因为相对路径会受到当前工作目录影响，而我们希望模型始终保存在项目目录中。

第二步：确保目录存在

ini 复制代码

out_dir.mkdir(
    parents=True,
    exist_ok=True
)

其中：

parents=True：必要时创建缺失的父目录。
exist_ok=True：目录已经存在时不报错。

第三步：保存模型

scss 复制代码

torch.save(checkpoint, checkpoint_path)

torch.save() 会将检查点序列化到磁盘。

生成文件：

bash 复制代码

out/model.pt

保存完成后，我们就不需要重新训练才能使用这些参数了。

需要注意，重复运行训练程序会覆盖同一个 model.pt。

如果希望保留不同实验结果，可以通过不同的输出目录区分。

三、怎样重新加载模型？

保存模型之后，下一步就是加载。

整个过程可以理解为：

markdown 复制代码

读取检查点
    ↓
恢复模型配置
    ↓
创建 GPT 模型
    ↓
加载训练好的参数
    ↓
切换到评估模式
    ↓
输入文本并生成

其中最重要的一点是：

模型配置负责创建网络结构，state_dict 负责恢复网络参数。

3.1 读取检查点

首先，通过 torch.load() 读取模型文件：

ini 复制代码

checkpoint = torch.load(
    "out/model.pt",
    map_location="cpu",
    weights_only=True,
)

这里需要解释两个参数。

map_location="cpu" 表示先把检查点中的张量加载到 CPU 上。

例如，我们之前在 Mac 上使用 MPS 训练：

ini 复制代码

device = "mps"

但以后可能希望在 CPU 上运行模型。

先加载到 CPU，再根据需要移动到其他设备，可以减少设备相关问题。

而：

ini 复制代码

weights_only=True

会使用受限制的加载方式，适合我们当前主要由张量、数字、字符串和普通字典组成的检查点。

不过，仍然应该只加载自己生成或来源可信的检查点文件。

3.2 根据配置创建模型

读取检查点后，可以取得模型配置：

ini 复制代码

model_args = checkpoint["model_args"]

随后创建模型：

ini 复制代码

model = GPTLanguageModel(**model_args)

这里的模型刚刚初始化，参数仍然是新初始化的值。

因此，还需要恢复之前训练得到的参数。

3.3 加载模型参数

使用：

css 复制代码

model.load_state_dict(
    checkpoint["model"]
)

load_state_dict() 会将保存的参数复制到对应模块中。

例如：

markdown 复制代码

保存文件中的：
token_embedding_table.weight
          ↓
加载到新模型中的：
token_embedding_table.weight

注意，默认情况下 load_state_dict() 使用严格匹配。

如果模块名称不一致，或者参数形状不兼容，就会报错。

因此，本篇的 sample.py 应当与保存检查点时使用的 model.py 配套。

检查点本身不会自动保存整个 GPTLanguageModel 类的 Python 实现。

3.4 切换到评估模式

完成参数加载后：

scss 复制代码

model.to(device)
model.eval()

其中：

javascript 复制代码

model.eval()

表示将模型切换到评估模式。

回顾上一篇，我们的 Transformer Block 包含 Dropout：

ini 复制代码

self.dropout = nn.Dropout(0.1)

训练时，Dropout 会随机将部分元素置零。

但生成文本时，我们通常希望关闭这一随机操作，让输出的不确定性主要由采样过程控制。

因此，需要调用：

javascript 复制代码

model.eval()

还有一个容易混淆的问题：

model.eval() 并不等于关闭梯度计算。

它只负责切换 Dropout 等模块的行为。

真正关闭梯度记录，可以使用：

csharp 复制代码

with torch.no_grad():
    ...

我们之前在 generate() 方法中已经使用了：

less 复制代码

@torch.no_grad()

所以不需要在每次调用生成方法时重复创建梯度上下文。

四、回顾 GPT 是怎样生成文本的

在加入新的采样方法之前，先回顾第四篇的文本生成过程。

模型的输入是一组字符 ID：

yaml 复制代码

R O M E O :
↓ ↓ ↓ ↓ ↓ ↓
整数 ID 序列

经过字符嵌入、位置嵌入以及多个 Transformer Block 后，模型输出：

ini 复制代码

logits, loss = model(idx)

其中：

ini 复制代码

logits.shape = (B, T, vocab_size)

这里：

B：批次大小。
T：输入序列长度。
vocab_size：字符词表大小。

例如：

ini 复制代码

logits.shape = (1, 6, 65)

表示模型对 6 个输入位置分别给出了整个字符词表的预测分数。

但我们现在想预测的是下一个字符。

因此，只需要取最后一个位置：

ini 复制代码

logits = logits[:, -1, :]

此时：

ini 复制代码

logits.shape = (1, 65)

接下来，将 logits 转换为概率：

ini 复制代码

probs = F.softmax(logits, dim=-1)

然后通过：

ini 复制代码

idx_next = torch.multinomial(
    probs,
    num_samples=1
)

按照概率分布随机采样一个字符 ID。

最后，将新字符拼接到输入后面：

ini 复制代码

idx = torch.cat(
    (idx, idx_next),
    dim=1
)

这样就完成了一次生成。

不断重复这个过程，就能逐个字符生成文本。

但目前的采样方法还比较简单。

接下来，我们加入两个常见参数：

temperature
top-k

它们并不改变 GPT 的模型结构，而是改变我们从模型预测结果中选择下一个字符的方式。

五、Temperature：调整生成的随机程度

5.1 为什么需要 Temperature？

假设模型在某一步预测时，对四个候选字符给出分数：

css 复制代码

候选字符     logits
A             3.0
B             2.0
C             1.0
D             0.0

经过 Softmax：
$P(xi)= ezi ∑jezj P(x_i)= \frac{e^{z_i}} {\sum_j e^{z_j}}$ P(xi)=∑jezjezi

得到的概率约为：

字符	概率
A	0.644
B	0.237
C	0.087
D	0.032

可以看到，A 的概率最高，但其他字符也有被采样到的机会。

如果我们希望模型更加倾向于选择高概率字符，应该怎么办？

一个简单的方法就是调整 Softmax 之前的 logits。

引入温度参数 $TT$ T：
$P(xi)= e zi/T ∑je zj/T P(x_i)= \frac{e^{z_i/T}} {\sum_j e^{z_j/T}}$ P(xi)=∑jezj/Tezi/T

其中：
$T>0T>0$ T>0

这个参数就叫作 Temperature。

5.2 Temperature 如何影响概率？

当：
$T<1T<1$ T<1

logits 之间的差距会被放大，最终概率分布更加集中。

当：
$T>1T>1$ T>1

logits 之间的差距会缩小，概率分布更加平缓。

例如，我们分别设置：

ini 复制代码

temperature = 0.5
temperature = 1.0
temperature = 2.0

可以运行下面的代码观察：

ini 复制代码

import torch
import torch.nn.functional as F

logits = torch.tensor([
    3.0, 2.0, 1.0, 0.0
])

for temperature in [0.5, 1.0, 2.0]:
    probs = F.softmax(
        logits / temperature,
        dim=-1
    )

    print(
        f"T={temperature}:",
        probs.round(decimals=3)
    )

得到的概率近似为：

字符	T=0.5	T=1.0	T=2.0
A	0.865	0.644	0.455
B	0.117	0.237	0.276
C	0.016	0.087	0.167
D	0.002	0.032	0.102

可以发现：

当 temperature=0.5 时，A 的概率变得非常高。

当 temperature=2.0 时，其他字符被选择的机会明显增加。

因此，我们可以总结为：

Temperature	分布特点	常见现象
小于 1	更集中	更倾向高概率字符，可能更加重复
等于 1	原始分布	按模型原始概率采样
大于 1	更平缓	随机性更强，更容易选到低概率字符

5.3 在代码中实现 Temperature

实现非常简单：

原来是：

ini 复制代码

probs = F.softmax(logits, dim=-1)

现在改成：

ini 复制代码

logits = logits / temperature

probs = F.softmax(logits, dim=-1)

还需要限制：

css 复制代码

if temperature <= 0:
    raise ValueError("temperature 必须大于 0")

需要强调的是：

Temperature 只能改变采样概率，不能提高模型已经学到的知识和语言能力。

如果模型本身训练不足，即使把 temperature 调低，也不意味着它能生成逻辑连贯的莎士比亚文本。

六、Top-k：只从概率较高的字符中采样

Temperature 调整的是整个概率分布。

但即使某个字符的概率非常低，它仍然可能被采样到。

例如，一个字符的预测概率只有：
$P(x)=0.0001P(x)=0.0001$ P(x)=0.0001

虽然概率很小，但并不等于零。

而在连续生成几百个字符的过程中，这些低概率候选也可能被选中。

我们可以进一步限制采样范围。

这就是 Top-k Sampling。

6.1 Top-k 的基本思想

假设模型预测的 logits 为：

ini 复制代码

logits = torch.tensor([
    [1.2, 3.5, 0.8, 2.7, 4.1]
])

对应 5 个候选字符。

现在设置：

ini 复制代码

top_k = 3

意味着只保留分数最高的 3 个候选。

也就是：

复制代码

4.1
3.5
2.7

剩下两个候选不参与采样。

我们可以通过：

ini 复制代码

values, indices = torch.topk(
    logits,
    k=3
)

得到：

lua 复制代码

values  = [[4.1, 3.5, 2.7]]
indices = [[4, 1, 3]]

其中：

values：最大的 3 个分数。
indices：这些分数在原始张量中的位置。

接下来，需要把低于第 3 高分数的候选排除。

6.2 为什么使用负无穷？

可以将不保留的 logits 设置为：

bash 复制代码

float("-inf")

例如：

ini 复制代码

原始 logits:

[1.2, 3.5, 0.8, 2.7, 4.1]

Top-k 处理后:

[-inf, 3.5, -inf, 2.7, 4.1]

因为：
$e−∞=0e^{-\infty}=0$ e−∞=0

所以经过 Softmax 后，被排除候选的概率为零。

这就保证模型不会从被排除的字符中采样。

6.3 使用 masked_fill 实现

我们可以写成：

ini 复制代码

k = 3

values, _ = torch.topk(
    logits,
    k
)

threshold = values[:, [-1]]

logits = logits.masked_fill(
    logits < threshold,
    float("-inf")
)

其中：

css 复制代码

values[:, [-1]]

取的是 top-k 结果中最后一个，也就是第 k 大的分数。

为什么写成：

css 复制代码

[:, [-1]]

而不是：

css 复制代码

[:, -1]

因为前者保留二维形状：

css 复制代码

(B, 1)

便于与：

css 复制代码

(B, vocab_size)

进行广播比较。

这里使用的是阈值筛选方式。如果第 k 大的分数与其他候选并列，实际保留下来的候选可能超过 k 个。

对于我们当前的小型字符级 GPT，这种实现已经足够用于实验。

6.4 Temperature 和 Top-k 可以同时使用吗？

可以。

在生成过程中，我们先通过 Temperature 缩放 logits，再进行 Top-k 筛选：

ini 复制代码

logits = logits / temperature

if top_k is not None:
    k = min(top_k, logits.shape[-1])

    values, _ = torch.topk(logits, k)

    logits = logits.masked_fill(
        logits < values[:, [-1]],
        float("-inf")
    )

probs = F.softmax(logits, dim=-1)

最后再进行随机采样。

两者作用不同：

Temperature：控制保留下来的候选之间的概率分布。
Top-k：控制允许哪些候选参与采样。

在 Temperature 大于零的情况下，缩放不会改变 logits 的大小顺序，因此先缩放还是先选择 Top-k，不会改变按分数确定的候选集合。

七、更新 GPTLanguageModel 的 generate 方法

现在，我们将前面的内容整合起来。

打开之前的 model.py。

在 GPTLanguageModel 类中，找到第四篇实现的：

scss 复制代码

generate()

将其替换为下面的版本。

ini 复制代码

@torch.no_grad()
def generate(
    self,
    idx,
    max_new_tokens,
    temperature=1.0,
    top_k=None,
):
    # 参数检查
    if temperature <= 0:
        raise ValueError(
            "temperature 必须大于 0"
        )

    if max_new_tokens < 0:
        raise ValueError(
            "max_new_tokens 不能小于 0"
        )

    if top_k is not None and top_k <= 0:
        raise ValueError(
            "top_k 必须大于 0"
        )

    for _ in range(max_new_tokens):

        # 只使用最近 block_size 个字符
        idx_cond = idx[:, -self.block_size:]

        # 前向传播
        logits, _ = self(idx_cond)

        # 只取最后一个位置的预测
        logits = logits[:, -1, :]

        # Temperature
        logits = logits / temperature

        # Top-k
        if top_k is not None:
            k = min(
                top_k,
                logits.shape[-1]
            )

            values, _ = torch.topk(
                logits,
                k
            )

            logits = logits.masked_fill(
                logits < values[:, [-1]],
                float("-inf")
            )

        # 转换为概率分布
        probs = F.softmax(
            logits,
            dim=-1
        )

        # 采样下一个字符
        idx_next = torch.multinomial(
            probs,
            num_samples=1
        )

        # 拼接到当前序列
        idx = torch.cat(
            (idx, idx_next),
            dim=1
        )

    return idx

注意，这个方法仍然位于 GPTLanguageModel 类内部，并且沿用第四篇的 forward() 接口：

ini 复制代码

logits, loss = self(idx)

7.1 为什么要截取最后 block_size 个字符？

来看这一行：

ini 复制代码

idx_cond = idx[:, -self.block_size:]

假设：

ini 复制代码

block_size = 64

那么模型每次最多处理 64 个字符。

如果目前已经生成了 100 个字符，就不能直接把全部 100 个字符送进模型。

因为我们之前创建的位置嵌入表大小为：

markdown 复制代码

nn.Embedding(
    block_size,
    n_embd
)

它只支持对应范围内的位置索引。

因此，我们只截取最近的 64 个字符作为当前上下文。

例如：

css 复制代码

完整生成序列：

[字符1, 字符2, ..., 字符100]

当前送入 GPT：

[字符37, 字符38, ..., 字符100]

这样可以保证输入长度不超过模型支持的上下文长度。

需要注意，这种做法并不是让模型真正拥有无限长的记忆。

超过窗口范围的字符不再参与当前这一步的预测。

7.2 为什么只取最后一个位置？

ini 复制代码

logits = logits[:, -1, :]

假设输入序列长度为 10：

复制代码

x1 x2 x3 ... x10

由于我们的 GPT 使用因果注意力，每个位置只能依赖当前位置及之前的字符。

所以第 10 个位置的输出表示：
$P(x11∣x1,x2,...,x10) P(x_{11}\mid x_1,x_2,\ldots,x_{10})$ P(x11∣x1,x2,...,x10)

也就是根据当前 10 个字符预测第 11 个字符。

而前面 9 个位置对应的是其他位置的下一字符预测。

生成时我们只需要最后一个位置的结果。

7.3 为什么使用 torch.multinomial？

ini 复制代码

idx_next = torch.multinomial(
    probs,
    num_samples=1
)

因为我们希望按照模型给出的概率分布进行随机采样。

如果始终使用：

ini 复制代码

torch.argmax(probs, dim=-1)

每一步都会选择概率最高的字符。

这样属于贪心解码，容易出现重复。

而 torch.multinomial() 可以在高概率字符和其他可能字符之间进行随机选择。

不同的随机采样过程，可能生成不同的文本。

7.4 最终返回什么？

最后：

kotlin 复制代码

return idx

返回的是完整的 token 序列。

包括：

复制代码

原始提示词 + 新生成字符

假设输入提示词长度为 6：

makefile 复制代码

ROMEO:

并且设置：

ini 复制代码

max_new_tokens = 100

那么最终输出长度就是：
$6+100=1066+100=106$ 6+100=106

对于当前字符级模型，一个字符对应一个 token。

但对于以后使用 BPE 等分词方法的语言模型，一个 token 就不一定是一个字符了。

八、编写独立的 sample.py

模型已经可以保存，也拥有了新的生成方法。

现在，我们要实现本篇最重要的目标：

将文本生成从训练程序中独立出来。

这样，即使关闭 train.py，仍然可以通过 sample.py 加载模型并生成文本。

首先，在项目根目录新建：

复制代码

sample.py

为了方便使用，我们通过命令行参数控制生成过程。

例如：

arduino 复制代码

python sample.py --prompt "ROMEO:" --tokens 300

表示：

使用 "ROMEO:" 作为提示词。
生成 300 个新字符。

8.1 完整代码

下面是 sample.py 的完整实现：

ini 复制代码

from pathlib import Path
import argparse

import torch

from model import GPTLanguageModel


BASE_DIR = Path(__file__).resolve().parent


def main(argv=None):
    # 命令行参数
    parser = argparse.ArgumentParser(
        description="加载字符级 GPT 并生成文本"
    )

    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=BASE_DIR / "out" / "model.pt",
    )

    parser.add_argument(
        "--prompt",
        default="ROMEO:\n"
    )

    parser.add_argument(
        "--tokens",
        type=int,
        default=300
    )

    parser.add_argument(
        "--temperature",
        type=float,
        default=0.8
    )

    parser.add_argument(
        "--top-k",
        type=int,
        default=20
    )

    parser.add_argument(
        "--seed",
        type=int,
        default=1337
    )

    parser.add_argument(
        "--device",
        choices=["auto", "cpu", "mps", "cuda"],
        default="auto"
    )

    args = parser.parse_args(argv)

    # 参数检查
    if not args.prompt:
        parser.error("提示词不能为空")

    if args.tokens < 0:
        parser.error("--tokens 不能小于 0")

    if args.temperature <= 0:
        parser.error(
            "--temperature 必须大于 0"
        )

    if args.top_k < 0:
        parser.error(
            "--top-k 不能小于 0，0 表示不限制"
        )

    if not args.checkpoint.is_file():
        parser.error(
            f"检查点不存在：{args.checkpoint}"
        )

    # 选择设备
    device = args.device

    if device == "auto":
        if torch.cuda.is_available():
            device = "cuda"
        elif torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"

    # 加载检查点
    checkpoint = torch.load(
        args.checkpoint,
        map_location="cpu",
        weights_only=True
    )

    # 恢复字符映射
    meta = checkpoint["meta"]

    stoi = meta["stoi"]
    itos = meta["itos"]

    # 检查提示词是否存在未知字符
    unknown = sorted(
        set(args.prompt) - set(stoi)
    )

    if unknown:
        parser.error(
            f"提示词包含词表外字符：{unknown!r}"
        )

    # 恢复模型
    model = GPTLanguageModel(
        **checkpoint["model_args"]
    )

    model.load_state_dict(
        checkpoint["model"]
    )

    model.to(device).eval()

    # 编码提示词
    start = torch.tensor(
        [[stoi[ch] for ch in args.prompt]],
        dtype=torch.long,
        device=device
    )

    # 设置随机种子
    torch.manual_seed(args.seed)

    # 生成文本
    generated = model.generate(
        start,
        max_new_tokens=args.tokens,
        temperature=args.temperature,
        top_k=(
            None
            if args.top_k == 0
            else args.top_k
        )
    )

    # 解码
    text = "".join(
        itos[i]
        for i in generated[0].tolist()
    )

    print(text)

    return text


if __name__ == "__main__":
    main()

这里假设第一篇构建的 meta 至少包含：

makefile 复制代码

meta = {
    "vocab_size": vocab_size,
    "stoi": stoi,
    "itos": itos,
}

并且 itos 是以整数 ID 为键的字典。

下面我们逐步理解这段程序。

8.2 使用 argparse 接收命令行参数

首先：

arduino 复制代码

import argparse

argparse 是 Python 标准库，用来处理命令行参数。

例如：

ini 复制代码

parser = argparse.ArgumentParser()

创建一个参数解析器。

然后：

go 复制代码

parser.add_argument(
    "--tokens",
    type=int,
    default=300
)

表示允许用户使用：

css 复制代码

python sample.py --tokens 500

传入生成长度。

如果用户没有指定，则使用默认值：

复制代码

这样就不需要每次修改 Python 源代码。

本篇主要提供以下参数：

参数	含义	默认值
`--checkpoint`	检查点文件路径	`out/model.pt`
`--prompt`	起始提示词	`ROMEO:\n`
`--tokens`	生成的新字符数	300
`--temperature`	采样温度	0.8
`--top-k`	候选字符数量	20
`--seed`	随机种子	1337
`--device`	推理设备	auto

其中：

ini 复制代码

args = parser.parse_args(argv)

会解析命令行参数，并将结果保存到 args 对象中。

因此，我们可以通过：

复制代码

args.prompt
args.tokens
args.temperature

取得对应的值。

8.3 自动选择运行设备

我们之前在 Mac 上训练时，使用的是：

ini 复制代码

device = "mps"

但并不是所有电脑都支持 MPS。

因此，在独立生成程序中，可以自动选择可用设备：

ini 复制代码

if device == "auto":
    if torch.cuda.is_available():
        device = "cuda"
    elif torch.backends.mps.is_available():
        device = "mps"
    else:
        device = "cpu"

这里：

CUDA：适用于支持 CUDA 的 NVIDIA GPU 环境。
MPS：适用于支持 Metal 的 Apple Silicon Mac。
CPU：通用计算设备。

如果当前设备同时不支持 CUDA 和 MPS，就会自动使用 CPU。

也可以手动指定：

css 复制代码

python sample.py --device cpu

需要注意，自动选择的是程序可用的设备，并不代表一定拥有最佳运行性能。

8.4 从检查点恢复字符词表

接下来：

ini 复制代码

meta = checkpoint["meta"]

stoi = meta["stoi"]
itos = meta["itos"]

这里我们没有重新读取第一篇的数据文件。

而是直接使用检查点中保存的字符映射。

这样可以确保：

生成文本时使用的字符编号与训练时完全一致。

接下来，还需要检查提示词中有没有未知字符：

scss 复制代码

unknown = sorted(
    set(args.prompt) - set(stoi)
)

假设词表只包含英文、数字、换行和部分标点。

如果我们输入：

复制代码

你好，世界

其中的中文字符大概率不在当前莎士比亚词表中。

那么就会得到明确的错误提示。

为什么不直接给未知字符分配一个新 ID？

因为我们的字符嵌入表只学习过固定范围内的 ID。

随意增加新的字符编号，可能导致嵌入索引越界，也无法让模型自动理解新字符。

因此，目前最简单可靠的方式就是提前检查。

8.5 将提示词转换为 Tensor

我们已经拿到了：

复制代码

stoi

接下来，对用户输入的提示词进行编码：

ini 复制代码

start = torch.tensor(
    [[stoi[ch] for ch in args.prompt]],
    dtype=torch.long,
    device=device
)

例如：

ini 复制代码

prompt = "ROMEO:"

经过编码后可能得到：

csharp 复制代码

[24, 21, 19, 17, 21, 10]

随后转换成张量：

ini 复制代码

shape = (1, 6)

为什么外面要使用两层中括号？

因为我们的 GPT 输入格式是：
$(B,T)(B,T)$ (B,T)

而不是：
$(T,)(T,)$ (T,)

这里的：

lua 复制代码

[[...]]

让输入拥有一个 batch 维度。

对于单条提示词：

ini 复制代码

B = 1

如果不增加这一维，模型的 forward() 可能无法按照之前约定的输入形状正常工作。

8.6 设置随机种子

scss 复制代码

torch.manual_seed(args.seed)

用于固定 PyTorch 随机数生成器的起始状态。

因为：

scss 复制代码

torch.multinomial()

包含随机采样过程。

使用相同随机种子，可以让同一环境中的生成实验更容易复现。

这里把随机种子设置放在模型创建之后，是为了避免模型初始化消耗随机数，从而改变后续采样的随机起点。

不过，不同设备、不同 PyTorch 版本或不同计算实现之间，仍然可能出现差异。

因此，我们不要求 CPU 和 MPS 生成的内容逐字相同。

8.7 调用 generate 生成文本

现在可以直接调用：

ini 复制代码

generated = model.generate(
    start,
    max_new_tokens=args.tokens,
    temperature=args.temperature,
    top_k=(
        None
        if args.top_k == 0
        else args.top_k
    )
)

其中：

ini 复制代码

top_k=None

表示不限制候选数量。

由于命令行参数不方便直接传入 Python 的 None，我们约定：

css 复制代码

--top-k 0

表示关闭 Top-k 限制。

最后，将生成的字符 ID 解码为文本：

ini 复制代码

text = "".join(
    itos[i]
    for i in generated[0].tolist()
)

这里：

css 复制代码

generated[0]

取出 batch 中的第一条序列。

scss 复制代码

.tolist()

将张量中的 ID 转换为 Python 列表。

再通过：

css 复制代码

itos[i]

转换回对应字符。

最后使用：

bash 复制代码

"".join(...)

把字符拼接成完整文本。

到这里，我们的独立生成程序就完成了。

九、运行 sample.py，观察生成结果

现在，我们可以关闭训练程序。

只需要确保以下两个文件存在：

bash 复制代码

model.py
out/model.pt

然后在项目根目录执行：

复制代码

python sample.py

程序会自动读取：

bash 复制代码

out/model.pt

并使用默认提示词生成 300 个新字符。

9.1 指定提示词

例如：

lua 复制代码

python sample.py \
  --prompt "ROMEO:" \
  --tokens 300

程序会使用：

makefile 复制代码

ROMEO:

作为起始上下文。

最终生成结果会以这段提示词开头。

但模型是否能够继续生成合理的戏剧对话，取决于它的训练效果。

9.2 尝试较低的 Temperature

lua 复制代码

python sample.py \
  --prompt "ROMEO:" \
  --temperature 0.6 \
  --top-k 10

这里：

ini 复制代码

temperature = 0.6
top_k = 10

意味着：

更倾向选择高分字符。
每一步只在分数最高的一组候选中采样。

这种配置通常会让输出更集中，但也可能产生重复内容。

9.3 尝试较高的 Temperature

lua 复制代码

python sample.py \
  --prompt "ROMEO:" \
  --temperature 1.0 \
  --top-k 0

这里：

ini 复制代码

top-k = 0

表示不限制候选范围。

模型会根据原始 Softmax 概率，从整个字符词表中采样。

相比上一种配置，更多低概率字符可能被选中。

9.4 怎样比较两种采样方法？

为了让实验更有意义，我们可以固定：

ini 复制代码

prompt = "ROMEO:"
tokens = 300
seed = 1337

只改变：

css 复制代码

temperature
top-k

例如：

实验	Temperature	Top-k
A	0.6	10
B	0.8	20
C	1.0	不限制
D	1.2	不限制

然后观察生成文本的差异。

可以重点关注：

是否出现大量重复字符或词语。
是否能模仿莎士比亚戏剧中的人物名称与换行格式。
是否出现明显不符合上下文的字符组合。
不同采样参数是否改变输出的多样性。

这里不需要急着追求某一组最好的参数。

我们更希望理解：

同一个模型，即使参数完全不变，也可以因为采样策略不同而产生不同的文本。

这也是语言模型生成过程中的一个重要特点。

十、怎样验证模型保存和加载正确？

到这里，我们已经完成了模型保存与独立生成。

但还有一个问题：

我们怎样确认加载后的模型，真的恢复了训练结束时的参数？

只观察生成文本并不可靠。

因为文本生成包含随机采样，即使模型参数相同，也可能产生不同的文字。

因此，我们采用一个更加直接的验证方法：

比较模型保存前后，对同一个输入产生的 logits。

10.1 构造一个固定输入

例如，我们取换行符作为输入：

ini 复制代码

probe = torch.tensor(
    [[meta["stoi"]["\n"]]],
    dtype=torch.long,
    device=device
)

这里：

ini 复制代码

probe.shape = (1, 1)

表示只有一个字符。

10.2 保存前计算 logits

首先，把训练好的模型切换到评估模式：

javascript 复制代码

model.eval()

然后：

ini 复制代码

with torch.no_grad():
    before, _ = model(probe)

此时：

复制代码

before

保存着训练模型对输入字符的预测结果。

10.3 重新加载并计算 logits

接下来，从磁盘重新加载检查点：

ini 复制代码

saved = torch.load(
    checkpoint_path,
    map_location="cpu",
    weights_only=True
)

重新创建一个 GPT：

ini 复制代码

restored = GPTLanguageModel(
    **saved["model_args"]
)

加载参数：

css 复制代码

restored.load_state_dict(
    saved["model"]
)

最后：

scss 复制代码

restored.to(device).eval()

计算恢复后模型的输出：

ini 复制代码

with torch.no_grad():
    after, _ = restored(probe)

10.4 比较输出

PyTorch 提供了：

scss 复制代码

torch.testing.assert_close(
    before,
    after
)

它会检查两个张量是否在允许的数值误差范围内接近。

如果检查通过，说明模型保存前后的计算结果一致。

我们可以把前面步骤整理成下面的代码，临时放在 train.py 保存检查点之后：

ini 复制代码

# 保存与加载一致性检查
model.eval()

probe = torch.tensor(
    [[meta["stoi"]["\n"]]],
    dtype=torch.long,
    device=device
)

with torch.no_grad():
    before, _ = model(probe)

# 从磁盘读取
saved = torch.load(
    checkpoint_path,
    map_location="cpu",
    weights_only=True
)

# 重新创建模型
restored = GPTLanguageModel(
    **saved["model_args"]
)

# 恢复参数
restored.load_state_dict(
    saved["model"]
)

restored.to(device).eval()

# 计算恢复后的输出
with torch.no_grad():
    after, _ = restored(probe)

# 检查数值一致性
torch.testing.assert_close(
    before,
    after
)

print("保存与加载一致性检查通过")

这里特别注意：

javascript 复制代码

model.eval()
restored.eval()

两个模型都应当处于评估模式。

否则，Dropout 可能使两次输出存在差异。

这项测试验证的是：

同一个模型配置和同一组参数，在相同输入下能否得到一致的预测结果。

它并不代表模型生成的文本已经足够连贯。

这两个问题需要分别判断。

10.5 再进行独立生成测试

最后，可以做几个简单的功能检查：

训练程序结束后，sample.py 是否仍然能够运行。
输出是否以输入提示词开头。
输出新增字符数量是否等于 --tokens。
提示词包含词表外字符时，是否能给出明确错误。
修改 Temperature 和 Top-k 后，程序是否仍然能够正常生成。

如果这些测试都通过，就说明我们已经完成了模型推理所需的基本保存和加载流程。

十一、整理项目结构

完成本篇代码之后，我们的项目目录可以整理为：

kotlin 复制代码

my-nanogpt/
├── data/
│   ├── train.bin
│   ├── val.bin
│   └── meta.pkl
│
├── out/
│   └── model.pt
│
├── prepare.py
├── dataset.py
├── model.py
├── train.py
└── sample.py

各个文件的职责也更加清晰：

文件	主要作用
`prepare.py`	将原始文本转换为字符 ID
`dataset.py`	构造训练批次
`model.py`	定义 GPT 结构与生成方法
`train.py`	训练 GPT 并保存检查点
`sample.py`	加载检查点并生成文本
`out/model.pt`	保存训练后的模型参数、配置和字符映射

其中，model.py 同时被 train.py 和 sample.py 使用。

这是一种很重要的代码组织方式。

我们不需要在训练程序和生成程序中分别实现两个 GPT。

相反，只需要维护一份模型定义。

训练时：

javascript 复制代码

from model import GPTLanguageModel

生成时：

javascript 复制代码

from model import GPTLanguageModel

两者使用同一个模型类，只是运行方式不同。

训练阶段：

复制代码

数据 → GPT → Loss → 反向传播 → 更新参数

生成阶段：

复制代码

提示词 → GPT → Logits → 采样 → 新字符

这样，我们就把模型结构、训练过程和推理过程分离开了。

随着项目继续扩展，这种划分也会让代码更容易维护。

十二、总结与下一篇预告

回顾整个 nanoGPT 系列：

第一篇，我们将莎士比亚文本转换成字符 ID，并构造了用于下一字符预测的训练样本。

第二篇，我们实现 Bigram 语言模型，理解了语言模型最基本的训练与生成过程。

第三篇，我们引入单头因果自注意力，让模型能够根据上下文预测字符。

第四篇，我们加入多头注意力、前馈神经网络、残差连接和 LayerNorm，构建了完整的小型 GPT。

而这一篇，我们没有继续增加模型层数，也没有引入新的注意力结构。

相反，我们完善了一个能够实际使用的语言模型所需要的基本流程。

本篇最重要的几个知识点：

state_dict() 用于取得模型需要保存的参数与持久化缓冲区。
模型检查点需要同时保存模型参数、模型配置和字符映射。
torch.save() 和 torch.load() 可以完成模型状态的保存与恢复。
model.eval() 用于切换评估模式，torch.no_grad() 用于关闭梯度记录。
Temperature 调整采样概率的集中程度。
Top-k 限制参与采样的候选字符范围。
独立的 sample.py 可以在不运行训练程序的情况下加载模型并生成文本。

至此，我们已经完成了一条完整的流程：

css 复制代码

莎士比亚原始文本
        ↓
字符编码与数据集构造
        ↓
GPTLanguageModel
        ↓
模型训练
        ↓
保存 model.pt
        ↓
独立加载检查点
        ↓
输入 Prompt
        ↓
Temperature + Top-k
        ↓
生成文本

但随着模型规模增大，我们还会遇到一些实际问题。

例如：

如果训练进行了很长时间，程序突然中断，是否只能从头开始？

如果模型训练了几千步，我们怎样判断哪一次保存的模型效果最好？

如果希望训练更加稳定，学习率是否应该始终保持不变？

这些问题已经超出了简单的模型结构实现，开始涉及完整的训练系统。

下一篇，我们将继续完善 nanoGPT 的训练流程，加入验证集评估、最佳与最新检查点、断点续训、学习率调度和梯度累积。

让我们的小型 GPT 不仅能够训练和生成，还能够更加稳定、可靠地进行较长时间的训练。