从零实现 GPT 类 LLM：分步工程路径与 PyTorch 实现细节

在构建大型语言模型时，理解每个模块内部机制与掌握端到端的工程路径同样重要。本文以 Sebastian Raschka 的《Build a Large Language Model (From Scratch)》配套代码仓库为蓝本，系统梳理从原始文本到预训练 GPT 模型的全链路实现细节，重点给出可直接落地的参数配置与显存估算方法。

1. 数据预处理与滑动窗口采样

数据质量直接决定预训练效果上限。工程实践中通常经历三步：首先使用正则表达式对原始文本做初步清洗与分句；然后通过 BPE（Byte Pair Encoding）分词器将文本转换为整数 token 序列；最后构造滑动窗口数据集用于自回归训练。

在具体实现中，GPTDatasetV1 类的核心逻辑如下：输入完整语料后按 stride 步长滑动，每次截取长度为 max_length 的 token 块作为输入，其向右偏移一位的序列作为目标。步长设置决定相邻样本的重叠程度 ——stride 等于 max_length 时样本完全不重叠，有助于降低过拟合；stride 较小（如 max_length/2）则可增加样本数量但引入信息冗余。训练常用的默认参数为 batch_size=4, max_length=256, stride=128, drop_last=True。

class GPTDatasetV1(Dataset):
    def __init__(self, txt, tokenizer, max_length, stride):
        token_ids = tokenizer.encode(txt, allowed_special={"<|endoftext|>"})
        for i in range(0, len(token_ids) - max_length, stride):
            input_chunk = token_ids[i:i + max_length]
            target_chunk = token_ids[i + 1: i + max_length + 1]
            self.input_ids.append(torch.tensor(input_chunk))
            self.target_ids.append(torch.tensor(target_chunk))

2. BPE 分词器的工程选型

GPT-2 采用 BPE 分词而非传统 word-level 方案，其优势在于将 Out-of-Vocabulary（OOV）词拆解为已知子词片段，从而实现对任意文本的无损编码。工程实践中推荐直接使用 OpenAI 开源的 tiktoken 库 —— 其核心算法以 Rust 实现，吞吐量比纯 Python 实现高约 5 倍。初始化时仅需一行代码：tokenizer = tiktoken.get_encoding("gpt2")，即可获得包含 50257 个 token 的 GPT-2 词汇表。

若需从零实现 BPE 训练流程，基本步骤包括：初始化字节级字符词汇、统计所有相邻 token 对频率、迭代合并最高频 pair 直到达到目标词汇量。工程实现时需注意合并操作的可逆性 —— 编码时记录的 merge 顺序必须在解码时严格还原，否则会产生无法对齐的乱码。

3. 位置编码：从查找表到旋转编码

GPT 类模型使用可学习的位置嵌入（Position Embedding）而非原始 Transformer 中的正弦 / 余弦编码。具体实现中，位置嵌入层与 token 嵌入层结构相同：nn.Embedding(context_length, emb_dim)。在模型前向传播时，将 token 嵌入与位置嵌入逐元素相加后送入 Transformer 块。

对于希望扩展到更长上下文（超过预训练长度）的场景，旋转位置编码（RoPE）是一个值得关注的替代方案。RoPE 将相对位置信息编码到 Q/K 的旋转矩阵中，允许模型处理任意长度的位置序列而无需额外训练。实现时需对注意力计算的 query 和 key 向量施加复数旋转，但工程复杂度相对可控。

4. 多头自注意力的 PyTorch 实现

多头注意力是 GPT 的核心组件。实现时需注意三个关键细节：因果掩码、头维度分割与dropout 正则化。

因果掩码通过 register_buffer 注册一个上三角矩阵（对角线含 0），在前向传播中将未来位置的注意力分数置为负无穷：

self.register_buffer("mask", torch.triu(torch.ones(context_length, context_length), diagonal=1))
# 前向传播中
attn_scores.masked_fill_(mask_bool[:num_tokens, :num_tokens], -torch.inf)

头维度分割使用 .view() 和 .transpose() 将 (batch, seq_len, d_out) 的投影结果重塑为 (batch, num_heads, seq_len, head_dim)，确保多头的计算可并行化。GPT_CONFIG_124M 的推荐配置为 num_heads=12, head_dim=64（即 d_out=768）。

5. Transformer 块的残差与归一化顺序

每个 Transformer 块包含两个子层：多头自注意力与前馈网络（FFN），每个子层外围包裹残差连接与层归一化。标准顺序为 LayerNorm → SubLayer → Dropout → + Shortcut。代码实现中，LayerNorm 使用自定义类而非 nn.LayerNorm，以确保 eps=1e-5 并支持可学习的 scale/shift 参数：

class TransformerBlock(nn.Module):
    def forward(self, x):
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_shortcut(x)
        x = x + shortcut
        # FFN 残差块同理

6. 前馈网络与 GELU 激活

GPT-2 使用 GELU（Gaussian Error Linear Unit）而非 ReLU，其数学形式为 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))。GELU 在负半轴有非零输出，能够更好地保留梯度信息。FFN 的中间维度通常扩展为 4 * emb_dim，这也是 transformer 中参数量最大的子层。124M 参数模型的 FFN 约占总参数量的三分之一。

7. 端到端 GPT 模型的组装

整合以上所有组件，完整的 GPTModel 由以下模块顺序堆叠而成：Token Embedding → Position Embedding → Dropout → N × TransformerBlock → Final LayerNorm → 线性输出头。输出头 nn.Linear 不带偏置（bias=False），这是 GPT 系列的标准做法，有助于减少参数量并与权重初始化策略更好地配合。

128M 参数模型的默认配置：

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": False
}

8. 预训练的显存估算与关键工程参数

显存占用主要由模型参数、梯度和激活值三部分构成。以 124M 模型为例：vocab_size=50257, emb_dim=768, n_layers=12 的模型约占用 500MB；混合精度训练（FP16）可将梯度与优化器状态压缩至约 600MB；激活值的显存消耗与 batch_size 和序列长度成正比 ——batch_size=4、序列长度 = 256 时，单次前向传播的激活值约消耗 1.2GB。

实用估算公式：总显存（GB）≈ 参数量（B）×（2× 精度字节 + 优化器状态系数 + 激活倍数）。使用 AdamW 优化器时，参数与动量状态各占用一份；使用 torch.compile 或梯度检查点（gradient checkpointing）可将激活显存降低约 60%。

推荐训练关键参数：

• 学习率初始值：3e-4，配合线性 warmup 约 100 步
• 梯度裁剪阈值：1.0
• 权重衰减系数：0.1
• 批大小：8-16（消费级 GPU），64-128（高端训练集群）

9. 推理时 KV Cache 与生成策略

推理阶段可使用 KV Cache 避免重复计算历史 token 的注意力：每生成一个新 token 时，将当前的 K/V 向量缓存起来，仅计算新增 token 与缓存的拼接。实现时需注意 register_buffer 用于存储静态掩码，确保因果注意力不泄露未来位置信息。

生成策略上，Greedy Decoding（贪心选择概率最高的 token）计算最简单但易陷入重复循环；Temperature Sampling（top-p/nucleus sampling）通过调整 softmax 温度系数可获得更多样化输出。实践中推荐 top_p=0.9, temperature=0.8 作为初始调试点。

资料来源

• GitHub: rasbt/LLMs-from-scratch
• 书籍: Build a Large Language Model (From Scratch), Sebastian Raschka, Manning, 2024, ISBN 978-1633437166

ai-systems