第44篇｜实战：从零写一个 mini-GPT
本篇定位：专栏四「PyTorch 与 AI Infra 入门」的第八篇。承接第 43 篇手写 Attention 与 Transformer Block，把已经验证过的 block 接到 token、位置、训练和自回归生成，完成一个教学用途的 mini-GPT。
读完你将能够：解释 decoder-only 语言模型的完整数据流；用 PyTorch 写出 embedding、位置 embedding、多个 Transformer Block 和 LM head；用 next-token prediction 训练字符级小语料；实现 greedy 与 temperature/top-k 采样；从 prefill、decode、KV Cache 和 batching 的角度理解它如何走向服务系统。
预计用时：阅读 70～90 分钟，动手 90～120 分钟。
第 43 篇我们把 Attention 的公式拆开，并用 pre-norm、残差和 FFN 组成了一个 Transformer Block。可是一个 block 还不是语言模型：它不知道整数 7 对应什么 token，也没有词表上的输出，更没有训练目标。真实工程中，最容易让人迷路的地方正是“把几个正确的零件接起来”：输入应该错位几位？位置向量加在哪里？最后一个时间步的 logits 怎么变成下一个 token？
本篇沿着一条能运行的最小链路回答这些问题。我们不追求训练出有用的聊天模型，而是用一小段重复、有规律的文本，训练一个字符级 decoder-only 模型。模型足够小，CPU 也能完成 smoke test；代码足够透明，可以把每个 shape 打印出来。最后再把训练阶段和推理阶段放回 AI Infra 的地图：训练是整段并行的矩阵计算，生成是逐 token 的服务循环，系统优化会从模型结构延伸到缓存、调度和带宽。
一、先画出 mini-GPT 的全链路
1.1 GPT 在这里指什么
GPT（Generative Pre-trained Transformer）在本文中只表示一种结构方向：decoder-only Transformer 用左到右的 causal mask 预测下一个 token。它不是某个商业模型的复刻，也不意味着参数规模、词表、数据和训练配方与大模型相同。我们写的是一个可读的教学模型，目的是把第 43 篇的 block 变成完整的语言建模任务。
一条最小数据流如下：
文本 → tokenizer → token ids [B,T]
→ token embedding + position embedding [B,T,D]
→ N 个 causal Transformer Block [B,T,D]
→ final LayerNorm [B,T,D]
→ lm_head [B,T,V]
→ 每个位置预测下一个 token
B 是 batch size，T 是上下文长度，D 是隐藏维度，V 是词表大小。lm_head 的每一个位置都输出 V 个 logits，但训练时位置 t 的目标是原序列位置 t+1 的 token。因此输入和目标看起来像下面这样：
原序列： 我 喜 欢 P y t h o n
input： 我 喜 欢 P y t h o
target： 喜 欢 P y t h o n
模型并不是一次只预测最后一个字符。训练时所有位置可以并行计算，causal mask 保证位置不能偷看右侧；这是 GPU 训练吞吐高的原因。生成时才会反复取最后一个位置的 logits，一次追加一个 token。
1.2 为什么先用字符级 tokenizer
大模型通常使用 BPE、Unigram 或其他子词 tokenizer，把常见片段压成一个 token。为了把注意力放在模型本身，本文用字符级 tokenizer：先收集语料中出现的字符，按排序建立 char_to_id，未知字符映射到 <unk>。字符级词表小、编码过程能用几行 Python 写完，但序列会更长；这正好让我们看见 T 对计算和显存的影响。
text = "Python让模型学习下一个字符。Python让模型学习下一个字符。"
chars = sorted(set(text))
itos = ["<unk>"] + chars
stoi = {ch: i for i, ch in enumerate(itos)}
def encode(s: str) -> list[int]:
return [stoi.get(ch, 0) for ch in s]
def decode(ids: list[int]) -> str:
return "".join(itos[i] for i in ids if i != 0)
ids = encode(text)
print(len(itos), ids[:8], decode(ids[:8]))
15 [5, 3, 8, 10, 1, 7, 2, 6] Python
实际输出中的数字会随着语料字符集合改变，不能把 token id 当成有语义的编号。embedding 只认识整数索引，id=5 的含义由训练出来的那一行向量决定。保存 checkpoint 时，tokenizer 的词表必须和模型一起保存；只保存权重而忘记词表，模型就无法把输入还原成相同的 id。
图1 mini-GPT 的端到端数据流
AI Infra 视角：tokenizer 是模型服务的第一道边界。线上吞吐常用 tokens/s 衡量，而不是字符数/s；同一个请求换一种 tokenizer，token 数和上下文占用都可能改变。预处理、padding 和 tokenizer 版本也应纳入基准条件。
1.3 位置不是 token 的属性
nn.Embedding(vocab_size, dim) 只表示“这个 token 是什么”，不表示它位于第几个位置。decoder-only 模型常把 token embedding 和 position embedding 相加：同一个字符出现在不同位置，会加上不同的位置向量。这里采用绝对位置 embedding，最大长度固定为 max_tokens，直观但不能自动外推到更长上下文。
若 ids 是 [B,T]，位置索引应是 [1,T]，查表后得到 [1,T,D]。利用广播与 token embedding 相加，结果仍为 [B,T,D]。写模型时要显式检查 T <= max_tokens，否则位置表越界的报错常常出现在深层 embedding 内部，不容易读懂。
二、组装模型：embedding、Block 与 LM head
2.1 把第 43 篇的 block 当成黑盒接口
在第 43 篇，TransformerBlock 接收 [B,T,D] 并返回相同 shape。它内部完成 pre-norm causal self-attention、残差、FFN 和残差。今天不重新发明公式，只依赖这个接口：只要每个 block 保持形状契约，nn.ModuleList 就可以把 N 个 block 串起来。
下面给出一个可独立运行的简化 block。为了让本文代码从零可复制，它保留显式的分数和 mask；这不是生产实现。register_buffer 让 causal mask 随模型移动到 CPU 或 GPU，但它不是参数，不会被优化器更新。
import math
import torch
from torch import nn
class CausalSelfAttention(nn.Module):
def __init__(self, dim: int, heads: int, max_tokens: int):
super().__init__()
if dim % heads != 0:
raise ValueError("dim must be divisible by heads")
self.heads = heads
self.head_dim = dim // heads
self.qkv = nn.Linear(dim, 3 * dim)
self.proj = nn.Linear(dim, dim)
future = torch.triu(torch.ones(max_tokens, max_tokens), diagonal=1).bool()
self.register_buffer("future", future[None, None], persistent=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
b, t, d = x.shape
if t > self.future.size(-1):
raise ValueError("sequence is longer than max_tokens")
q, k, v = self.qkv(x).chunk(3, dim=-1)
shape = (b, t, self.heads, self.head_dim)
q = q.view(shape).transpose(1, 2)
k = k.view(shape).transpose(1, 2)
v = v.view(shape).transpose(1, 2)
scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
scores = scores.masked_fill(self.future[:, :, :t, :t], float("-inf"))
weight = scores.softmax(dim=-1)
y = weight @ v
y = y.transpose(1, 2).contiguous().view(b, t, d)
return self.proj(y)
class TransformerBlock(nn.Module):
def __init__(self, dim: int, heads: int, max_tokens: int):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = CausalSelfAttention(dim, heads, max_tokens)
self.norm2 = nn.LayerNorm(dim)
self.ffn = nn.Sequential(
nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x + self.attn(self.norm1(x))
return x + self.ffn(self.norm2(x))
注意 shape 中 d 必须等于构造时的 dim，而 heads 必须整除 dim。transpose 后调用 contiguous() 再 view，是为了让内存布局和新 shape 一致。教学代码把注意力权重保留在局部变量里，便于理解；训练长序列时，它可能产生 [B,H,T,T] 的大张量，后面会解释为什么工程版本要换 fused kernel。
2.2 最小 GPT 模块
模型外壳只做四件事：查 token 表、查位置表、循环通过 blocks、投影到词表。最后的 LayerNorm 是常见的 final norm。lm_head 可以有 bias，也可以像很多 GPT 实现一样不带 bias；这里选择 bias=False，并保留一个可选的权重绑定说明。
class MiniGPT(nn.Module):
def __init__(self, vocab_size: int, dim=64, heads=4,
layers=2, max_tokens=64):
super().__init__()
self.max_tokens = max_tokens
self.token_emb = nn.Embedding(vocab_size, dim)
self.pos_emb = nn.Embedding(max_tokens, dim)
self.blocks = nn.ModuleList([
TransformerBlock(dim, heads, max_tokens)
for _ in range(layers)
])
self.norm = nn.LayerNorm(dim)
self.lm_head = nn.Linear(dim, vocab_size, bias=False)
def forward(self, ids: torch.Tensor) -> torch.Tensor:
b, t = ids.shape
if t > self.max_tokens:
raise ValueError(f"T={t} exceeds max_tokens={self.max_tokens}")
positions = torch.arange(t, device=ids.device).unsqueeze(0)
x = self.token_emb(ids) + self.pos_emb(positions)
for block in self.blocks:
x = block(x)
return self.lm_head(self.norm(x))
model = MiniGPT(vocab_size=32, dim=32, heads=4, layers=2, max_tokens=16)
ids = torch.randint(0, 32, (2, 8))
print(ids.shape, model(ids).shape)
torch.Size([2, 8]) torch.Size([2, 8, 32])
这就是最重要的形状检查：每个输入位置都拿到一个长度为 V 的 logits 向量。logits 还不是概率，不要在训练前手动 softmax；cross_entropy 内部会用数值稳定的 log-softmax。推理时才根据 logits 计算概率并选择 token。
2.3 参数量和配置的关系
可以用 sum(p.numel() for p in model.parameters()) 统计参数量。增大 layers 会近似线性增加 block 参数和计算；增大 dim 会同时影响 embedding、QKV、输出投影、FFN，参数量大致呈平方增长；增大词表会让 embedding 和 LM head 线性增加。max_tokens 主要影响位置表和注意力激活，并不直接让每个线性层变大。
n_params = sum(p.numel() for p in model.parameters())
print(f"parameters={n_params:,}")
for name, p in model.named_parameters():
print(name, tuple(p.shape))
parameters=31,360
# 后面还会打印 token_emb、blocks.0.attn.qkv 等参数形状
如果把输入 embedding 和输出 head 绑定权重，可以写 self.lm_head.weight = self.token_emb.weight，但要明确两个层的词表维度一致，且保存、加载和优化器参数列表都要按绑定后的语义验证。为了教学清晰，本篇不默认绑定；先理解 shape，再做参数共享实验。
图2 mini-GPT 的模块堆叠与 shape
三、训练：让模型学会下一个 token
3.1 从一维语料切 batch
字符序列太短会导致模型只背下几个片段，所以准备一段重复但包含标点和换行的文本。我们把全部 id 放进一维 Tensor，然后随机抽取若干起点，每个样本取连续的 block_size + 1 个 id：前 block_size 个作为输入，后 block_size 个作为 target。这样每个 batch 的 x、y shape 都是 [B,T]。
def get_batch(data: torch.Tensor, batch_size: int,
block_size: int, device: torch.device):
starts = torch.randint(0, len(data) - block_size - 1, (batch_size,))
x = torch.stack([data[i:i + block_size] for i in starts])
y = torch.stack([data[i + 1:i + block_size + 1] for i in starts])
return x.to(device), y.to(device)
y 只比 x 向右移动一格，不能把整段目标再错一次。现在模型在位置 t 看到 x[:, :t+1]，应该预测 y[:, t]。causal mask 负责限制可见范围，数据切片负责定义监督信号；两者缺一不可。
3.2 logits 展平后算交叉熵
model(x) 产生 [B,T,V]，交叉熵需要类别维在最后或使用常见的 [N,V] 形式。我们把前两维合并成 B*T，target 变为 [B*T]：
import torch.nn.functional as F
x, y = get_batch(train_ids, batch_size=8, block_size=32, device=device)
logits = model(x)
loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
print(logits.shape, y.shape, float(loss))
torch.Size([8, 32, 32]) torch.Size([8, 32]) 3.4...
随机模型的初始 loss 通常接近 log(V)，但具体值受词表、初始化和 batch 影响。训练后 loss 下降，表示平均而言正确 token 的概率上升；它不保证生成文本语法正确，更不等于真实世界知识。验证集 loss 需要用没有参与更新的文本片段，否则只是在测记忆。
3.3 一个短而完整的训练循环
下面是教学版训练骨架。它沿用了第 41 篇的 DataLoader/batch 思想和第 42 篇的 AdamW、梯度裁剪、scheduler 思想，但为了突出 GPT 数据流，直接从一维 Tensor 取 batch。真实项目还应加入验证、日志、checkpoint 和恢复训练。
import torch
import torch.nn.functional as F
text = ("Python让模型预测下一个字符。" * 200)
chars = sorted(set(text))
stoi = {ch: i for i, ch in enumerate(chars)}
itos = chars
all_ids = torch.tensor([stoi[ch] for ch in text], dtype=torch.long)
split = int(0.9 * len(all_ids))
train_ids, val_ids = all_ids[:split], all_ids[split:]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MiniGPT(len(chars), dim=64, heads=4, layers=2, max_tokens=64).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
for step in range(300):
model.train()
x, y = get_batch(train_ids, 16, 48, device)
optimizer.zero_grad(set_to_none=True)
logits = model(x)
loss = F.cross_entropy(logits.reshape(-1, len(chars)), y.reshape(-1))
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
if step % 50 == 0:
print(step, round(loss.item(), 3), round(float(grad_norm), 3))
0 2.1 1.8
50 0.4 0.5
100 0.2 0.3
...
输出数字只作格式示意；换 CPU、PyTorch 版本或随机种子，loss 不会逐行相同。这里设置的学习率只适合这个极小任务，不能照搬到大模型。若 loss 变成 nan，先检查学习率、输入 id 是否越界、mask 是否造成全 -inf，再检查 dtype 和梯度范数。
3.4 训练阶段要观察什么
只看 loss 不够。至少记录训练 loss、验证 loss、step time、每秒处理 token 数、峰值显存和当前学习率。训练 token 数可近似为 batch_size × block_size × step；同样的 step 数，如果 T 不同，工作量完全不同。model.train() 要在训练阶段打开，若没有 Dropout，它暂时看不出区别，但养成正确习惯能避免后续模型行为不一致。
可以每隔若干步用 torch.no_grad() 或 torch.inference_mode() 做验证，避免保存反向图。验证结束后记得恢复 model.train()。checkpoint 至少保存 model.state_dict()、optimizer.state_dict()、随机种子、tokenizer 词表和当前 step；否则恢复后可能 loss 曲线断裂，生成结果也无法复现。
图3 训练循环：输入错位到参数更新
AI Infra 视角：训练吞吐要以有效 token 计量。数据加载、tokenizer、CPU 到 GPU 的拷贝、前向、反向和 optimizer step 都可能成为瓶颈。小模型在 GPU 上反而可能被 Python 调度和 kernel launch 影响，不能因为“用了 CUDA”就默认更快。
四、生成：从 logits 走到文本
4.1 greedy generation
训练完成后，给模型一个 prompt。每一轮只取最后一个位置的 logits，选择最大值对应的 token，拼回上下文，再进行下一轮。这是 greedy decoding：每步都选当前概率最大的 token，结果可复现，但容易陷入重复，未必是全局最优序列。
@torch.inference_mode()
def generate_greedy(model, prompt_ids, max_new_tokens):
model.eval()
ids = prompt_ids
for _ in range(max_new_tokens):
context = ids[:, -model.max_tokens:]
logits = model(context)
next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
ids = torch.cat([ids, next_id], dim=1)
return ids
prompt = torch.tensor([[stoi["P"]]], device=device)
out = generate_greedy(model, prompt, 30)
print("".join(itos[i] for i in out[0].tolist()))
Python让模型预测下一个字符。Python让模型预测下一个字符。
示例输出只表示一个可能结果。ids[:, -model.max_tokens:] 是教学版的上下文截断：当生成长度超过位置表时保留最后窗口，避免越界；它同时意味着模型看不到更早的内容。服务系统会用更明确的最大上下文策略，并在请求超长时拒绝、截断或采用模型支持的长上下文方案。
4.2 temperature 与随机采样
贪心选择会把所有不确定性丢掉。把 logits 除以 temperature 后再 softmax：温度小于 1，分布更尖锐；大于 1，分布更平坦。然后按概率 multinomial 抽样。temperature 只能作用于 logits，不能直接作用于已经 argmax 的 id。
@torch.inference_mode()
def generate_sample(model, ids, max_new_tokens,
temperature=0.8, top_k=None):
model.eval()
for _ in range(max_new_tokens):
context = ids[:, -model.max_tokens:]
logits = model(context)[:, -1, :] / temperature
if top_k is not None:
values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
cutoff = values[:, [-1]]
logits = logits.masked_fill(logits < cutoff, float("-inf"))
probs = torch.softmax(logits, dim=-1)
next_id = torch.multinomial(probs, num_samples=1)
ids = torch.cat([ids, next_id], dim=1)
return ids
top_k 先保留分数最高的 k 个 token，其余置为 -inf，再归一化抽样。它减少长尾噪声，但 k 不是越大越好。生产系统还常见 top-p（nucleus）等策略；本篇只实现足够理解接口的版本。temperature 必须为正数，top_k 不应超过词表大小，抽样前要检查概率是否有限且总和接近 1。
greedy 和 sample 的差别不在 Transformer 主干，而在最后的 logits-to-token 后处理。把后处理独立出来，未来可以加入停止 token、重复惩罚、词表屏蔽和约束解码。服务端也需要给每个请求保存生成配置，否则同一个模型在不同调用间的行为很难解释。
4.3 一次前缀重算为什么慢
上面的教学循环每次把完整 context 重新送入模型。生成第 1 个 token 时计算长度为 T 的序列，第 2 个 token 又计算长度为 T+1 的序列；历史 token 的 Q、K、V 被反复投影。短 prompt 看不出问题，长上下文和高并发下，重复计算会直接变成延迟。
真正的 decoder 推理通常分为 prefill 和 decode：prefill 一次处理完整 prompt，并把每一层的 K/V 保存下来；decode 每次只输入一个新 token，计算它的 Q/K/V，把新 K/V 追加到 cache，然后用新 Q 读取历史 cache。这样用更多显存和读带宽，换取不重复处理历史 token。第 43 篇已经埋下这个概念，第 47 篇会专门实现 KV Cache、batching 和采样。
图4 训练、prefill 与逐 token decode
五、从小实验连接到 AI Infra
5.1 训练和服务的瓶颈不同
训练主要追求吞吐：同一批序列的所有位置可以并行，GPU 擅长大矩阵乘法；反向还要保存或重算激活，显存压力大。生成主要追求单请求延迟和并发吞吐：每个 decode step 的新 token 很少，计算呈增量形式，却要频繁读取不同请求的 KV Cache。
同一个 T 在两种阶段的含义也不同。训练时 B×T 是一次处理的 token 数，causal mask 允许并行计算；生成时每个请求的 T 逐步增长，batch 内请求可能处于不同长度。服务框架需要调度“谁现在需要一个 token”，并把不同长度请求放进合适的 batch，而不是简单地把所有请求补齐到最长长度。
5.2 一个粗略的资源账本
参数显存可以用“参数个数 × 每个元素字节数”估算；训练还要加梯度、AdamW 的一阶二阶状态、master weight、激活和临时 workspace。激活至少和 B×T×D×layers 同量级，显式 Attention 的权重还和 B×heads×T×T×layers 有关。推理没有反向激活，但每层 KV Cache 会随历史 token 数和并发请求数增长。
例如把 T 从 256 增到 1024，注意力分数矩阵的元素数量不是四倍，而是十六倍；把 batch 从 1 增到 4，通常则近似线性增加。实际 fused kernel 可以避免把完整权重写回显存，但不会改变“长序列交互更贵”的基本事实。估算后再做 profiler，比盲目增大 batch 更安全。
def bytes_to_mib(n: int) -> float:
return n / (1024 ** 2)
def attention_elements(batch, heads, tokens, layers):
return batch * heads * tokens * tokens * layers
for t in (128, 256, 512):
n = attention_elements(4, 4, t, 2)
print(t, n, f"fp16={bytes_to_mib(n * 2):.2f} MiB")
128 524288 fp16=1.00 MiB
256 2097152 fp16=4.00 MiB
512 8388608 fp16=16.00 MiB
这是只看一个中间张量的估算，不是完整峰值显存。allocator 缓存、参数、梯度和其他激活都没有算入。报告资源时要写明 dtype、batch、序列长度、是否训练和是否保存 attention weights。
5.3 手写版和生产版的边界
本篇显式写 scores、softmax、weight @ v，因为它让 mask 和 shape 可见，适合教学、单元测试和定位错误。线上模型一般使用 torch.nn.functional.scaled_dot_product_attention、更高层 Transformer 实现或推理框架提供的 kernel，让系统根据 device、dtype、mask 选择更高效的实现。不要把“我能写出来”误解为“适合生产”。
迁移到服务的检查清单可以是：
AI Infra 视角：模型只是服务的一部分。一次请求还经过 tokenizer、队列、batch 调度、GPU kernel、缓存管理和输出后处理。把本篇的
[B,T,D]、[B,H,T,T]和[B,T,V]形状写进监控或压测报告，往往比只写“模型很小”更能解释性能变化。
5.4 为第 45 篇留接口
下一步我们不会立刻换更大的模型，而是先测清楚它。混合精度会改变 Tensor Core 使用、显存占用和数值稳定性；profiler 会把 Python、CPU、CUDA kernel 和内存事件串起来。今天的 mini-GPT 已经提供了足够小的实验对象：固定 batch 和 T，比较 float32、自动混合精度、不同层数与不同 Attention 实现，才能知道优化到底改善了什么。
六、做一次端到端 smoke test
6.1 最小验收清单
把模型交给训练循环前，建议先完成五项检查：第一，随机 ids 的范围在 [0,V)；第二，输出 shape 是 [B,T,V]；第三，T=1 和 T=max_tokens 都能前向；第四，改变未来 token 时，当前位置的 logits 不应改变；第五，loss 能反向，所有梯度都是有限数。
# 假设 model 已构造，ids 是 [B,T]
ids = torch.randint(0, len(chars), (2, 8), device=device)
logits = model(ids)
assert logits.shape == (2, 8, len(chars))
assert torch.isfinite(logits).all()
loss = F.cross_entropy(logits.reshape(-1, len(chars)), ids.reshape(-1))
loss.backward()
assert all(p.grad is None or torch.isfinite(p.grad).all()
for p in model.parameters())
print("smoke test passed")
smoke test passed
验证 causal 性质时，可以构造两份序列，只修改位置 t 之后的 token，比较位置 t 的 logits。若前缀相同，allclose 应为真；这比只检查 mask 矩阵长什么样更接近模型行为。注意 Dropout 要关闭，且比较时使用 eval() 和固定 dtype。
6.2 用一个规律任务判断是否真的在学习
训练语料越随机，越难区分模型没学会还是任务本来不可预测。重复句子和固定模式让我们快速判断：训练若正常，loss 应明显下降，greedy 生成会出现语料中的片段。可以把训练集缩到一小段，故意让模型过拟合；如果连一个 batch 都记不住，优先排查 target 错位、学习率、模型是否进入 train 模式和 optimizer 是否真的拿到了参数。
过拟合小数据只是正确性测试，不是泛化能力证明。验证集 loss 下降、换 prompt 仍有合理延续，才说明模型捕捉到了一点规律。字符模型的输出可能重复标点或缺少语义，这属于任务规模的限制，不应通过堆参数掩盖数据和评测设计问题。
6.3 保存能够复现的产物
checkpoint = {
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"itos": itos,
"config": {"dim": 64, "heads": 4, "layers": 2, "max_tokens": 64},
"step": 300,
}
torch.save(checkpoint, "mini-gpt.pt")
print("saved mini-gpt.pt")
saved mini-gpt.pt
保存 itos 和配置不是附属信息，而是推理契约：缺少词表，整数没有可读含义；缺少 max_tokens、head 数或层数，无法正确构造模型。生产系统通常还保存代码版本、依赖版本、数据版本、随机种子和评测结果。一个能加载但不能解释的 checkpoint，运维成本会很高。
图5 从 mini-GPT 到模型服务的边界
七、踩坑大全
坑 1：把输入和 target 设成同一份序列。 现象是 loss 似乎下降，却没有学到“下一个 token”。原因是语言模型需要 x = ids[:, :-1]、y = ids[:, 1:] 的错位监督；生成时也应根据最后一个位置预测新 token。
坑 2：在训练前手动 softmax。 cross_entropy 接受 logits，内部会做稳定的 log-softmax。提前 softmax 可能损失数值精度，也容易把维度弄错；只有采样时才把最后一步 logits 变成概率。
坑 3：位置索引放在 CPU。 模型和 ids 在 GPU 时，torch.arange(t) 默认在 CPU，和 GPU 上的 embedding 相加会报 device 错误。应使用 device=ids.device，并把长度校验放在模型入口。
坑 4：dim 不能被 heads 整除。 reshape 会失败或 head 的宽度不一致。在构造函数里提前检查，不要等到第一次前向才发现配置错误。
坑 5：causal mask 的方向反了。 语言模型应该屏蔽严格未来，即上三角、diagonal=1。若把下三角屏蔽，模型会看见未来或只能看到错误的上下文；用 T=4 打印右上角并写测试。
坑 6：把完整前缀重算当成高性能生成。 它适合教学，但长 prompt 会重复计算历史 Q/K/V。生产 decode 需要 KV Cache；下一篇之后会进一步讨论缓存布局和 batch 调度。
坑 7：把随机采样的不同输出当成模型坏了。 temperature 和 multinomial 本来就引入随机性。调试时固定随机种子、使用 greedy 或保存生成配置；比较质量时要用多次采样和明确评测集。
坑 8：训练和推理模式混用。 生成前调用 model.eval()，并使用 torch.inference_mode()；eval() 关闭 Dropout，但不自动关闭梯度，两个开关职责不同。
坑 9：只报告参数量，不报告 token 和上下文。 参数量不能直接推断延迟或显存。至少同时记录 B、T、D、heads、layers、dtype、是否训练和是否使用 fused attention。
坑 10：直接把教学实现上线。 显式 [B,H,T,T] 的权重、Python 循环和单请求重算只为易读服务。生产要经过正确性对照、profile、压测和故障边界验证，再替换为成熟 kernel 与调度组件。
八、小结
- mini-GPT 的完整链路是 token embedding 加 position embedding，经过多个 causal Transformer Block、final norm 和 LM head，输出
[B,T,V]的 logits。 - 字符级 tokenizer 适合教学，词表简单但序列较长；tokenizer 词表、模型配置和 checkpoint 必须作为同一推理契约保存。
- 训练用
x与右移一位的y做 next-token prediction，训练阶段整段并行；cross_entropy直接接 logits，target 展平为类别索引。 - greedy 每次取最大 logits；temperature、top-k 和 multinomial 把 logits 变成可控的随机生成，生成配置应独立于模型主干。
- 教学代码显式保存 Attention 中间张量并重算前缀，便于理解但不适合生产；服务要关心 fused kernel、KV Cache、batching、tokens/s 和尾延迟。
- 一个可靠的 smoke test 要覆盖 shape、device、有限值、causal 不泄漏、loss 反向和 checkpoint 可恢复，而不只是看一段“像文本”的输出。
九、课后练习
- （基础） 把字符级 tokenizer 改成带
<unk>和<eos>的版本，验证未知字符能编码、生成遇到<eos>会停止。 - （基础） 给
MiniGPT加一个参数统计函数，分别改变layers、dim、max_tokens，记录参数量和一次前向的输出 shape。 - （理解） 写测试比较
T=1、T=8和T=max_tokens的前向；再传入T=max_tokens+1，确认错误信息清楚。 - （理解） 用同一个 prompt 比较 greedy、temperature 为
0.5和1.2的输出，固定 seed 后解释为什么结果可复现。 - （进阶） 实现 top-p 采样：按概率降序累加，只保留累积概率达到阈值的最小 token 集合，再抽样并恢复原词表索引。
- （AI Infra） 对比
B=1/4、T=128/256/512的前向时间和峰值显存，报告 CPU/GPU、dtype、warmup、是否保存 attention weights。 - （进阶） 为单层 Attention 实现一个不带梯度的 KV Cache，对比“每次重算前缀”和“只算新 token”的 logits，检查数值误差和缓存 shape。
- （综合） 保存 checkpoint 后重新加载，在相同 prompt 上比较生成结果；再记录训练 token 数、验证 loss、tokens/s 和峰值显存，写一份小型实验报告。
下一篇预告
第45篇｜混合精度、显存分析与 profiler。有了本篇可运行的 mini-GPT，下一篇我们会把 float32、float16、bfloat16 和自动混合精度放在同一个实验里，观察速度、显存和数值稳定性的平衡；再用 torch.profiler 追踪 CPU、CUDA kernel、内存分配和数据拷贝，建立一份可解释的显存账本。后续第 46 篇进入 torch.compile 与算子融合，第 47 篇把 KV Cache、batching 和采样做成推理基础，第 48 篇再把这些能力映射到 vLLM、SGLang 与 Agent 服务。