从零手写一个GPT：预训练循环的完整实现

大多数教程在注意力公式处止步，而从公式到"能在单卡上跑完且 loss 真的下降"的完整训练循环之间，还隔着一大段没人细讲的工程。这篇文章把这段路走完：一个约六百行的最小实现，涵盖数据管线、模型、混合精度训练循环与评估，每一处都是真实的坑。

一、数据管线：预训练的第一性问题其实是 IO

预训练中 GPU 常常在等数据，而不是在算。文本语料必须被切成固定长度的 token 序列并预打包，否则每个 batch 都要做动态拼接与移位，数据加载器会成为瓶颈。正确的做法是离线把语料写成一个巨大的二进制 token 文件，训练时按 batch 随机偏移读取：

import numpy as np

class TokenDataset:
"""预打包的 token 二进制文件读取器。文件内容是 uint16/uint32 的 token 流。"""
def __init__(self, path, block_size):
self.data = np.memmap(path, dtype=np.uint16, mode="r")
self.block = block_size

def sample_batch(self, batch_size):
# 随机起点, 取 block+1 个 token: 输入与目标错开一位
ix = np.random.randint(0, len(self.data) - self.block - 1,
size=batch_size)
x = np.stack([self.data[i:i + self.block] for i in ix]).astype(np.int64)
y = np.stack([self.data[i+1:i+1+self.block] for i in ix]).astype(np.int64)
return torch.from_numpy(x), torch.from_numpy(y)

两个细节决定成败。其一，目标序列是输入序列右移一位，这是自回归训练的全部秘密------模型在位置 i 看到的上下文是 x[:i]，要预测的是 x[i]，实现上只需错位打包，模型内部再用因果掩码保证"看不见未来"。其二，memmap 是关键：几十 GB 的 token 文件不需要载入内存，操作系统按页缓存，随机读的局部性也足够好。若要更进一步，可按文档边界分 shard、训练时按 epoch 内 shuffle shard，避免语料中的文档级顺序记忆。

二、模型：从多头注意力到完整块

模型部分的核心是因果自注意力的实现。现代实现与 2017 年论文的差异集中在三处：预归一化（LayerNorm 放在注意力与 FFN 之前而不是之后，训练稳定性天差地别）、旋转位置编码或可学习位置嵌入、以及门控线性单元式的 FFN（SwiGLU）。下面是带 KV 缓存接口的注意力层，它同时服务训练与推理两个场景：

class CausalSelfAttention(nn.Module):
def __init__(self, d_model, n_head):
super().__init__()
self.n_head = n_head
self.qkv = nn.Linear(d_model, 3 * d_model)
self.proj = nn.Linear(d_model, d_model)

def forward(self, x, kv_cache=None):
B, T, C = x.shape
q, k, v = self.qkv(x).split(C, dim=2)
# (B, n_head, T, head_dim)
q, k, v = map(lambda t: t.view(B, T, self.n_head, -1).transpose(1, 2),
(q, k, v))
if kv_cache is not None:
# 推理: 把新 token 的 k,v 追加到缓存, 重新取完整序列
k = torch.cat([kv_cache[0], k], dim=2)
v = torch.cat([kv_cache[1], v], dim=2)
kv_cache = (k, v)
att = (q @ k.transpose(-2, -1)) / math.sqrt(k.size(-1))
# 因果掩码: 只允许关注自己及之前
T_full = k.size(2)
mask = torch.triu(torch.ones(T, T_full, device=x.device), 1).bool()
att.masked_fill_(mask[:, :, -T:, :], float("-inf"))
y = (att.softmax(-1) @ v).transpose(1, 2).reshape(B, T, C)
return self.proj(y), kv_cache

注意掩码维度处理：训练时 T_full == T，推理生成时缓存里的键比查询长，掩码要对齐到查询的最后 T 个位置。这个"训练与推理共用一个 forward、靠 kv_cache 分流"的设计，能避免两套实现各自腐化。

三、训练循环：混合精度、梯度累积与稳定性

完整的训练循环要处理四个正交机制：混合精度（fp16/bf16 计算加速）、梯度累积（用小显存模拟大 batch）、梯度裁剪（防止 loss spike）、以及学习率预热与余弦衰减。下面是可以直接跑的核心循环：

def train(model, data, steps, batch_size, accum_steps=8,
lr=6e-4, warmup=200, max_grad_norm=1.0):
opt = torch.optim.AdamW(model.parameters(), lr=lr,
betas=(0.9, 0.95), weight_decay=0.1)
scaler = torch.cuda.amp.GradScaler(enabled=False)  # bf16 不需要 scaler
model.train()

def lr_at(step):  # 线性预热 + 余弦衰减
if step < warmup:
return lr * step / warmup
p = (step - warmup) / (steps - warmup)
return lr * 0.5 * (1 + math.cos(math.pi * p))

for step in range(steps):
for g in opt.param_groups:
g["lr"] = lr_at(step)

total_loss = 0.0
opt.zero_grad(set_to_none=True)
for _ in range(accum_steps):        # 梯度累积
x, y = data.sample_batch(batch_size)
x, y = x.cuda(), y.cuda()
with torch.autocast("cuda", dtype=torch.bfloat16):
_, loss = model(x, y)
loss.backward()
total_loss += loss.item() / accum_steps

torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
opt.step()

if step % 100 == 0:
print(f"step {step}  loss {total_loss:.4f}  "
f"lr {lr_at(step):.2e}")

几个稳定性经验。betas=(0.9, 0.95) 而非默认的 0.999：训练大模型时二阶动量的长记忆会放大早期 loss 尖峰。bf16 优先于 fp16------它的指数位与 fp32 相同，不需要 loss scaling，数值下溢风险大幅降低。梯度裁剪几乎总是需要的，一次 loss spike 不裁剪就可能让训练永久性跑偏。另外，loss 的解读要有基准：一个随机初始化的模型初始 loss 约等于 ln(vocab_size)（5 万词表约 10.8），训练的第一个小时主要就是在把 loss 从这里压下来；若开局就远低于此，多半是数据泄漏或目标错位。

四、评估与采样：训练的另一半

预训练的评估必须是 held-out 数据上的 per-token loss，且要用与训练相同的分块方式；更直观的方式是定期采样观察文本连贯性：

@torch.no_grad()
def generate(model, prompt_ids, max_new_tokens, temperature=1.0, top_k=50):
model.eval()
ids = list(prompt_ids)
kv = None
for _ in range(max_new_tokens):
x = torch.tensor([ids[-1]], device="cuda").unsqueeze(0)
logits, kv = model(x, kv_cache=kv)     # 增量解码: 每步只算 1 个 token
logits = logits[0, -1] / temperature
kth = torch.topk(logits, top_k).values[-1]
logits[logits < kth] = -float("inf")   # top-k 截断
ids.append(torch.multinomial(logits.softmax(-1), 1).item())
return ids

增量解码正是 KV 缓存的意义：每步只对最新的一个 token 做前向，历史 token 的键值全部复用。温度与 top-k 的调节有一个实用规律------temperature 调节的是分布的"锐度"，top-k 调节的是"尾巴长度"；采样文本出现复读时先降低温度，出现无意义词串时检查 top-k 是否过大。

五、剩下的坑

把上面拼起来，单卡就能训练一个字符级或 BPE 级的小模型并得到流畅输出。真正会咬人的坑还有三个：权重初始化（残差投影层要缩小初始化方差，否则深层网络开局就梯度爆炸）、dropout 的位置（注意力权重与残差分支各一处，小数据才有意义，大数据一般关掉）、以及 checkpoint 的保存粒度（优化器状态是参数的两倍体积，断点续训时记得保存 opt.state_dict 与步数，否则学习率计划会从头再来）。预训练从来不是黑魔法，它是一套可以被完整理解、逐行验证的工程------而亲手把 loss 曲线压下去一次，胜过读十篇解读文章。