先问一个问题你平时调用 GPT 接口、用 LangChain 做应用、或者浏览各种大模型的论文解读时有没有想过 GPT 内部的 transformer 层到底是怎么写的注意力机制里的 Q、K、V 矩阵是如何把一句话逐步编码成上下文向量的当我们在说“大语言模型”时它究竟是一堆什么样的张量运算和参数更新市面上关于大语言模型的文章很多但大多是“怎么调用 API”“怎么写提示词”真正从零开始用 PyTorch 搭建 GPT 结构的内容反而比较零散。本文将围绕“用 PyTorch 从零构建 GPT”这一主题把 GPT 的核心原理、模型结构、训练流程、文本生成推理完整走一遍。文章会以实战代码为主适合这几类读者已经会用 PyTorch 搭建简单神经网络想进一步理解 Transformer 和 GPT 源码的开发者。准备阅读 GPT、LLaMA 等大模型源码但被 attention、layer norm、位置编码绕晕的新手。需要在一个小规模数据集上自己做“类 GPT”实验的学生或算法工程师。通过本文你会掌握 GPT 的关键技术点并拿到一份可以在本地训练和推理解的完整代码。1. 背景与核心概念1.1 从语言模型说起“大语言模型”这个词听起来很庞大但把它拆到最小单元核心任务仍然是给定一段已知文本预测下一个词。比如输入“今天天气真”模型要计算下一个字是“好”“差”“热”还是“冷”的概率分布。传统语言模型基于 n-gram 统计只能考虑附近几个词无法处理长距离依赖。后来循环神经网络RNN、LSTM能建模序列但存在训练速度慢、长期记忆衰减的问题而且难以并行处理序列中的所有位置。2017 年《Attention Is All You Need》提出 Transformer彻底改变了序列建模的方式。Transformer 不再像 RNN 那样按时间步逐个读取输入而是把所有 token 一次性送入网络通过自注意力机制计算 token 之间的相关度。这种并行计算方式大幅提升了训练效率也让长距离依赖建模成为可能。GPT 全称是 Generative Pre-trained Transformer属于 Transformer 的 decoder 分支只保留自回归结构。它的训练思路分两阶段第一阶段在海量语料上做无监督预训练目标是预测下一个 token第二阶段可以用带标签数据做有监督微调SFT让模型学会用户期望的回答方式。1.2 GPT 与 Transformer 的关系Transformer 原论文包含 Encoder 和 Decoder 两个部分。Encoder 负责把输入序列编码成语义向量适合理解类任务Decoder 负责逐 token 生成输出适合生成类任务。GPT 只使用 Decoder 的 masked self-attention 部分。它不需要 Encoder因为语言建模的目标就是从历史 token 预测下一个 token当前时刻只能看到过去的信息不能看到未来的 token。这种限制由 attention mask 实现。这里要重点区分三个概念Encoder双向注意力能同时看到上下文适合 BERT 这类理解模型。Decoder单向注意力只能看左侧历史 token适合 GPT 这类生成模型。Encoder-Decoder完整 Transformer适合机器翻译、文本摘要等“序列到序列”任务。GPT 系列模型架构的核心是 Transformer decoder block一个 block 内包含masked multi-head self-attentionfeed-forward networkFFN每个子层都做残差连接和层归一化。堆叠多个这样的 decoder block就得到 GPT 的主干结构。模型规模增大时主要的参数来自 embedding 层、注意力层和 FFN 层。1.3 为什么要用 PyTorch 从零构建 GPT现在有很多现成的大模型框架比如 HuggingFace Transformers 里直接调用GPT2LMHeadModel就能加载预训练模型。那为什么还要自己手写一遍原因主要有三个。第一理解源码。GPT 的 huggingface 实现包含大量工程化逻辑新手很难一眼看出核心张量流。自己写一遍能把nn.Embedding、nn.Linear、masked attention这些组件的输入输出形状理清楚。第二方便实验。在做毕业设计、论文复现或者公司内部小规模预训练实验时我们经常需要改动 attention 结构、换激活函数、改位置编码方式。如果只会调用现成接口改动成本很高自己写核心模型改一行就能测试新想法。第三更深入理解大模型训练中的细节。比如为什么 GPT 需要独立的位置编码表为什么推理时要缓存 KV为什么采样时 temperature 会影响文本多样性这些问题的答案都藏在模型结构的细节里。2. 环境准备与版本说明2.1 操作系统与 Python 环境本文示例代码可以在 Windows、macOS 或 Linux 上运行。因为涉及矩阵运算和自动求导建议先安装 Anaconda 或 Miniconda 来管理 Python 环境。创建独立环境conda create -n gpt_from_scratch python3.10 conda activate gpt_from_scratchPython 版本建议 3.10 或更高。需要说明的是PyTorch 对不同 Python 版本的支持范围不同3.10 是一个非常稳妥的选择。2.2 安装 PyTorchPyTorch 的安装命令取决于你本机的 CUDA 版本。在开始之前先用下面的命令确认显卡驱动和 CUDA 能力nvidia-smi如果输出正常说明有 NVIDIA GPU 可用如果你的机器没有独立显卡也不用担心本文示例模型很小用 CPU 也可以完成训练只是速度会慢一些。以 CUDA 12.x 环境为例安装命令pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121如果你只需要 CPU 版本pip install torch版本说明本文示例代码基于 PyTorch 2.x 编写核心 API 在 1.13 及以上版本都能运行。如果你使用的是更早版本建议升级到 2.x因为 PyTorch 2.x 对编译优化和显存管理做得更好。2.3 安装 tiktoken训练 GPT 时我们需要一个分词器tokenizer。tiktoken 是 OpenAI 开源的 BPE 分词库可以让我们像 GPT 模型一样把文本切分成 token。pip install tiktoken安装完成后用 Python 简单验证一下import tiktoken enc tiktoken.get_encoding(gpt2) tokens enc.encode(Hello, world!) print(tokens) print(enc.decode(tokens))预期输出类似[15496, 11, 995, 0] Hello, world!能看到文本被切分成 token ID 序列并且可以解码回原文本说明环境没问题。3. GPT 核心原理拆解在写代码之前先把 GPT 的几个核心组件逐一拆开。不要小看这一部分后面代码能不能看懂关键就在这里。3.1 Tokenizer文本如何变成数字计算机无法直接处理文本神经网络也一样。我们需要把一段话变成数字序列这个过程叫 tokenization。一种最简单的方法是字符级切分每个字符对应一个 ID。这种方法实现简单但序列很长且无法体现词和子词的语义信息。GPT 使用的是 BPEByte Pair Encoding字节对编码它先把文本按 UTF-8 编码成字节序列然后迭代合并出现频率最高的字节对生成一个子词词表。使用 BPE 分词器后一段文本会被切分成类似这样的 tokenGPT is great - [GPT, is, great]tiktoken的get_encoding(gpt2)使用的就是 GPT-2 时代的 BPE 词表词汇量 50257。在实际训练中如果训练数据比较大一般会重新训练一个 tokenizer为了简化我们这里直接使用现成的 gpt2 分词器。3.2 Token Embedding 与 Position Embedding分词后每个 token 对应一个 ID例如15496。这个整数需要转换成向量才能进入神经网络这一步由nn.Embedding完成。nn.Embedding本质上是一个查找表输入形状为[batch_size, seq_len]的 token ID 矩阵输出形状为[batch_size, seq_len, n_embd]的向量序列。但是 Transformer 模型本身不包含任何序列顺序信息。对于注意力机制来说把句子“我打你”和“你打我”输入进去如果不加位置信息模型无法区分词的先后顺序。所以需要额外加入位置编码。GPT 使用可学习的位置编码表learned positional embedding每个位置都有一个独立向量维度与 token embedding 相同。位置编码表的形状是[block_size, n_embd]其中block_size是模型支持的最大序列长度。在代码中最终输入模型的特征是x token_embedding(tokens) position_embedding(positions)3.3 多头自注意力机制注意力机制是 GPT 的核心。它的作用是让每个 token 在计算特征表示时关注序列中其他 token 的信息。3.3.1 Self-Attention 的 Q、K、V自注意力Self-Attention的输入是同一个序列经过线性变换得到的三组矩阵Query查询向量表示当前 token “想查找什么”Key键向量表示当前 token “能提供什么”Value值向量表示当前 token “实际携带的信息”。对于序列中每个 token我们用它的 Query 向量与所有 token 的 Key 向量做点积得到注意力分数注意力分数经过 softmax 归一化后对 Value 向量加权求和。公式Attention(Q, K, V) softmax(Q * K^T / sqrt(d_k)) * V其中d_k是 Key 向量的维度除以sqrt(d_k)是为了防止点积值过大导致 softmax 梯度消失。Python 伪代码如下import torch import torch.nn as nn import torch.nn.functional as F def scaled_dot_product_attention(Q, K, V, maskNone): d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / d_k ** 0.5 if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights F.softmax(scores, dim-1) return torch.matmul(attn_weights, V), attn_weights3.3.2 Masked Attention为什么 GPT 只能看到过去GPT 是自回归模型生成第t个 token 时只能看到前t-1个 token。如果在训练时让当前位置看到了未来的 token模型就相当于照抄答案学不到真正的预测能力。实现方式很直接构造一个上三角矩阵把未来位置的注意力分数设置为负无穷softmax 之后这些位置的权重就变成 0。例如序列长度为 4合法的 mask 矩阵为1 0 0 0 1 1 0 0 1 1 1 0 1 1 1 1第i行第j列为 1 表示位置i可以关注位置j为 0 表示不允许关注。3.3.3 多头注意力多头注意力就是把注意力计算分成多个头并行做。假设特征维度是 768使用 12 个头那么每个头的维度是 64。每个头独立计算注意力最后拼接起来再过一层线性变换。多头注意力的好处是让模型从不同表示子空间捕捉信息。有的头可能关注语法关系有的头可能关注远距离指代关系。代码中可以用一个nn.Linear同时生成全部头的 Q、K、V再通过 reshape 拆分到不同头这样实现高效很多。3.4 前馈神经网络与残差、LayerNorm每个 Transformer block 在 attention 之后还会接一个前馈神经网络FFN通常包含两层全连接和一个激活函数FFN(x) GELU(x W1 b1) W2 b2GPT 中 FFN 的中间维度一般是模型维度的 4 倍也就是说 768 维的输入会先升到 3072 维再投影回 768 维。同时每个子层都使用“残差连接 层归一化”的结构残差连接Residual Connection让梯度能直接跨层传播避免深层网络梯度消失层归一化LayerNorm对每个样本的特征维做标准化加速训练收敛。注意 GPT 使用的是 pre-LayerNorm 结构即先做 LayerNorm 再做 Attention/FFN。这与原始 Transformer 论文中的 post-LayerNorm 略有不同但现代大模型几乎都采用 pre-LayerNorm训练更稳定。3.5 语言模型头与损失函数经过多层 Transformer block 后我们得到形状为[batch_size, seq_len, n_embd]的向量序列。要预测下一个 token需要把向量映射回词表大小这一层称为语言模型头LM Headlogits nn.Linear(n_embd, vocab_size)(x)得到每个位置上所有 token 的概率分布后训练目标就是交叉熵损失计算预测分布与真实下一个 token 的差距。假设输入序列是[1, 2, 3, 4]模型会同时预测给定[1]预测2给定[1, 2]预测3给定[1, 2, 3]预测4。这就是 GPT 的“并行 teacher forcing”训练方式一次前向计算得到整个序列所有位置的预测结果。4. 完整实战用 PyTorch 从零构建 GPT现在进入本文的主体部分。我们将搭建一个 mini 版 GPT 模型并在一个小型数据集上完成训练和文本生成。模型规模很小完全可以在 CPU 上运行。4.1 项目结构建议创建如下目录gpt_from_scratch/ ├── data.py # 数据加载与 tokenizer 封装 ├── model.py # GPT 模型结构 ├── train.py # 训练脚本 └── generate.py # 文本生成脚本4.2 数据准备为了演示我们使用一个英文小语料莎士比亚的《哈姆雷特》部分片段或者任何一本公开领域的英文书。这里为了代码简单直接内置一段文本作为训练语料你也可以用自己准备的.txt文件替换。数据加载代码# 文件路径data.py import tiktoken import torch from torch.utils.data import Dataset class TextDataset(Dataset): def __init__(self, text, block_size128): self.enc tiktoken.get_encoding(gpt2) self.tokens self.enc.encode(text) self.block_size block_size def __len__(self): return len(self.tokens) - self.block_size def __getitem__(self, idx): chunk self.tokens[idx: idx self.block_size 1] x torch.tensor(chunk[:-1], dtypetorch.long) y torch.tensor(chunk[1:], dtypetorch.long) return x, y def load_text(): sample_text To be, or not to be, that is the question: Whether tis nobler in the mind to suffer The slings and arrows of outrageous fortune, Or to take arms against a sea of troubles And by opposing end them. To die, to sleep; No more; and by a sleep to say we end The heart-ache and the thousand natural shocks That flesh is heir to: tis a consummation Devoutly to be wishd. To die, to sleep; To sleep, perchance to dream; ay, theres the rub. return sample_text这里的关键点每个样本的输入x和标签y是错位一个 token 的序列对应“预测下一个 token”的任务。tiktoken编码后文本中的空格和换行也会被保留为 token不需要额外清理。4.3 定义 GPT 模型下面编写完整的 GPT 模型代码。代码尽量简洁但保留 GPT 所有核心组件。# 文件路径model.py import torch import torch.nn as nn import torch.nn.functional as F class LayerNorm(nn.Module): def __init__(self, n_embd, eps1e-5): super().__init__() self.eps eps self.gamma nn.Parameter(torch.ones(n_embd)) self.beta nn.Parameter(torch.zeros(n_embd)) def forward(self, x): mean x.mean(-1, keepdimTrue) var x.var(-1, keepdimTrue, unbiasedFalse) x (x - mean) / torch.sqrt(var self.eps) return x * self.gamma self.beta class CausalSelfAttention(nn.Module): def __init__(self, n_embd, n_head, block_size, dropout0.0): super().__init__() assert n_embd % n_head 0 self.n_head n_head self.n_embd n_embd self.head_dim n_embd // n_head self.c_attn nn.Linear(n_embd, 3 * n_embd, biasFalse) self.c_proj nn.Linear(n_embd, n_embd, biasFalse) self.dropout nn.Dropout(dropout) # 上三角掩码矩阵用于屏蔽未来位置 self.register_buffer( mask, torch.tril(torch.ones(block_size, block_size)).view( 1, 1, block_size, block_size ), ) def forward(self, x): B, T, C x.shape # 一次性生成 Q、K、V qkv self.c_attn(x) # [B, T, 3*C] q, k, v torch.split(qkv, self.n_embd, dim-1) # 拆分成多头: [B, n_head, T, head_dim] q q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) v v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) # 缩放点积注意力 att (q k.transpose(-2, -1)) * (self.head_dim ** -0.5) att att.masked_fill(self.mask[:, :, :T, :T] 0, float(-inf)) att F.softmax(att, dim-1) att self.dropout(att) y att v # [B, n_head, T, head_dim] y y.transpose(1, 2).contiguous().view(B, T, C) y self.c_proj(y) return y class MLP(nn.Module): def __init__(self, n_embd, dropout0.0): super().__init__() self.fc1 nn.Linear(n_embd, 4 * n_embd) self.fc2 nn.Linear(4 * n_embd, n_embd) self.dropout nn.Dropout(dropout) def forward(self, x): x self.fc1(x) x F.gelu(x) x self.fc2(x) x self.dropout(x) return x class TransformerBlock(nn.Module): def __init__(self, n_embd, n_head, block_size, dropout0.0): super().__init__() self.ln1 LayerNorm(n_embd) self.attn CausalSelfAttention(n_embd, n_head, block_size, dropout) self.ln2 LayerNorm(n_embd) self.mlp MLP(n_embd, dropout) def forward(self, x): x x self.attn(self.ln1(x)) x x self.mlp(self.ln2(x)) return x class GPT(nn.Module): def __init__( self, vocab_size50257, n_embd128, n_head4, n_layer4, block_size128, dropout0.0, ): super().__init__() self.token_embedding nn.Embedding(vocab_size, n_embd) self.position_embedding nn.Embedding(block_size, n_embd) self.blocks nn.Sequential( *[ TransformerBlock(n_embd, n_head, block_size, dropout) for _ in range(n_layer) ] ) self.ln_f LayerNorm(n_embd) self.lm_head nn.Linear(n_embd, vocab_size, biasFalse) # 权重初始化 self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean0.0, std0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean0.0, std0.02) def forward(self, idx, targetsNone): B, T idx.shape assert T self.position_embedding.weight.shape[0], 序列长度超出位置编码范围 pos torch.arange(0, T, dtypetorch.long, deviceidx.device) pos pos.unsqueeze(0) # [1, T] x self.token_embedding(idx) self.position_embedding(pos) x self.blocks(x) x self.ln_f(x) logits self.lm_head(x) loss None if targets is not None: B, T, C logits.shape logits logits.view(B * T, C) targets targets.view(B * T) loss F.cross_entropy(logits, targets) return logits, loss def generate(self, idx, max_new_tokens100, temperature1.0): 自回归生成新 token self.eval() for _ in range(max_new_tokens): idx_cond idx[:, -self.position_embedding.weight.shape[0]:] logits, _ self(idx_cond) logits logits[:, -1, :] / temperature probs F.softmax(logits, dim-1) next_token torch.multinomial(probs, num_samples1) idx torch.cat((idx, next_token), dim1) return idx这个模型需要注意几个地方register_buffer注册的 mask 会随模型一起移动到 GPU/CPU不需要手动管理设备。c_attn一次性输出 3 倍维度的 Q、K、V比分别定义三个线性层效率更高。generate方法里只取最后一个位置的 logits 做采样这是因为自回归模型当前位置的预测只依赖之前所有 token。采样时除以 temperature让概率分布更尖锐或更平滑从而控制生成文本的随机性。4.4 编写训练脚本训练脚本中我们使用 AdamW 优化器。AdamW 是 GPT 训练的标配优化器因为它对权重衰减的处理更合理能提升泛化性能。# 文件路径train.py import torch from torch.utils.data import DataLoader from data import TextDataset, load_text from model import GPT def main(): # 超参数 batch_size 8 block_size 128 n_embd 128 n_head 4 n_layer 4 learning_rate 3e-4 epochs 50 device cuda if torch.cuda.is_available() else cpu text load_text() dataset TextDataset(text, block_sizeblock_size) loader DataLoader(dataset, batch_sizebatch_size, shuffleTrue) model GPT( vocab_size50257, n_embdn_embd, n_headn_head, n_layern_layer, block_sizeblock_size, dropout0.1, ).to(device) optimizer torch.optim.AdamW(model.parameters(), lrlearning_rate, weight_decay0.01) print(f数据集 token 数: {len(dataset.tokens)}) print(f使用设备: {device}) for epoch in range(epochs): model.train() total_loss 0 for x, y in loader: x, y x.to(device), y.to(device) optimizer.zero_grad() logits, loss model(x, y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_loss loss.item() avg_loss total_loss / len(loader) print(fEpoch {epoch1:3d}/{epochs}, loss {avg_loss:.4f}) torch.save(model.state_dict(), gpt_model.pt) print(训练完成模型已保存到 gpt_model.pt) if __name__ __main__: main()训练过程中值得关注的点clip_grad_norm_是梯度裁剪防止梯度爆炸。对语言模型来说这是稳定训练的重要技巧。数据量很小时loss 下降会比较快但生成的文本可能仍然不通顺。实际项目中需要更大规模的语料和更多的训练步数。这里模型参数量大约 800 万远小于真实的 GPT 模型所以 CPU 完全可以跑。4.5 编写文本生成脚本训练完成后加载模型权重并生成文本。# 文件路径generate.py import torch import tiktoken from model import GPT def main(): device cuda if torch.cuda.is_available() else cpu model GPT( vocab_size50257, n_embd128, n_head4, n_layer4, block_size128, dropout0.0, ).to(device) model.load_state_dict(torch.load(gpt_model.pt, map_locationdevice)) model.eval() enc tiktoken.get_encoding(gpt2) prompt To be, or not to be input_ids enc.encode(prompt) input_ids torch.tensor([input_ids], dtypetorch.long, devicedevice) output_ids model.generate(input_ids, max_new_tokens50, temperature0.8) output_text enc.decode(output_ids[0].tolist()) print(生成结果) print(output_text) if __name__ __main__: main()4.6 运行与预期结果依次运行python train.py python generate.py训练时如果数据量较小loss 可能在 0.1 以内如果你替换成了自己的语料loss 会在几十个 epoch 后逐渐下降。生成的部分文本大概率不完全通顺因为我们的训练数据只有一小段莎士比亚原文。但你能明显看出模型学到了一些英文单词拼写和常见词搭配这说明模型结构本身是有效的。换用更大的语料、更长的训练时间后生成质量会显著改善。5. 常见问题与排查思路5.1 常见报错汇总下表整理了手写 GPT 过程中最常见的几个问题问题现象常见原因解决思路训练时 loss 没有下降学习率过大/过小或 tokenizer 词表与模型 vocab_size 不匹配检查学习率确认tiktoken.get_encoding(gpt2)的词汇量与模型vocab_size一致运行时提示“序列长度超出位置编码范围”输入 token 长度超过block_size对输入做截断或增大block_size生成的文本全是乱码或重复字符采样温度过低或模型训练不充分调高 temperature增加训练数据量CUDA out of memory模型过大或 batch_size 过大降低 batch_size减少 block_size或使用梯度累积注意力分数出现 NaN数据中存在过大的数值网络未做梯度裁剪增加 LayerNorm 稳定性添加梯度裁剪5.2 关键排查步骤如果你训练出来的模型效果很差建议按下面顺序排查第一步确认数据预处理。打印一个样本的x和y确认y确实是x的下一个 token。如果错位逻辑写错模型永远学不对。第二步确认 mask 是否正确。把 mask 打印出来看是否上三角全为 0。如果 mask 没生效模型在训练时会偷看未来 token导致训练 loss 很低但生成效果很差。第三步确认 embedding 层维度。GPT 中 token embedding 和 position embedding 矩阵都是随机初始化的如果初始化不合理训练可能非常慢。建议使用均值为 0、标准差 0.02 的正态分布初始化这也是 GPT-2 原始代码的做法。第四步确认优化器和学习率。语言模型常用的学习率区间是1e-4到3e-4过大会导致 loss 震荡。同时不要忘记梯度裁剪。6. 最佳实践与工程建议6.1 从小规模实验出发不要一上来就训练一个 7B 参数的大模型。工程实践中先用一个几十 MB 的小语料、小模型几百万参数把代码流程跑通确认 loss 能正常下降、文本能正常生成然后逐步扩大规模。这样做能节省大量调试时间。6.2 使用梯度累积扩展 batch size真实的大语言模型训练往往需要很大的 batch size但单张显卡的显存有限。梯度累积的思路是先计算多个 mini-batch 的梯度然后将梯度累加最后做一次参数更新效果等价于更大的 batch size。PyTorch 中的实现思路如下accumulation_steps 4 optimizer.zero_grad() for i, (x, y) in enumerate(loader): x, y x.to(device), y.to(device) loss model(x, y)[1] / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意 loss 除以accumulation_steps是为了让梯度累加后的量级与正常 batch size 训练保持一致。6.3 混合精度训练PyTorch 2.x 中可以使用torch.autocast和GradScaler做混合精度训练。在支持 CUDA 的环境下混合精度能显著减少显存占用并提速训练。scaler torch.amp.GradScaler(cuda) for x, y in loader: x, y x.to(device), y.to(device) optimizer.zero_grad() with torch.amp.autocast(cuda): _, loss model(x, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度并不是必须的但对大规模训练非常有帮助。CPU 环境下不需要使用。6.4 推理阶段的 KV Cache在文本生成时模型需要逐个 token 生成。如果不做任何缓存每生成一个新 token都要重新计算之前所有 token 的 Key 和 Value计算量很大。真实 GPT 推理例如 HuggingFace 的generate会使用 KV Cache把之前 step 计算出来的 K、V 矩阵缓存下来下一轮只计算新 token 的 Q、K、V与缓存拼接后再做注意力计算。KV Cache 是理解大模型推理性能的核心概念。本文为了保持代码简洁没有实现但在深入研究大模型部署、vLLM、TensorRT-LLM 等推理框架时这个机制必须掌握。6.5 权重保存与加载训练完成后保存权重时最好同时保存模型配置和优化器状态方便断点续训。# 保存 torch.save({ model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), epoch: epoch, config: { n_embd: n_embd, n_head: n_head, n_layer: n_layer, block_size: block_size, vocab_size: 50257, } }, checkpoint.pt) # 恢复 checkpoint torch.load(checkpoint.pt) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict])6.6 数据质量决定模型上限大语言模型的训练遵循一个规律模型结构决定能力的下限数据质量决定能力的上限。如果你在自己的项目里微调 GPT不要只关心模型参数量更要投入时间清洗训练数据过滤重复文本、去掉格式错乱内容、尽量保证语料的多样性。7. 总结与下一步学习建议这篇文章从零实现了一个 mini 版 GPT核心代码只有不到 300 行。你已经掌握了 Token Embedding、Position Embedding、多头因果自注意力、LayerNorm、残差连接、FFN 和语言模型头的完整实现也能用自己的脚本训练一个小模型并生成文本。下一步可以从这几个方向继续深入把字符级或者 BPE 分词器替换成 SentencePiece训练支持中文的 tokenizer在模型中加入 KV Cache对比推理速度的提升参考 GPT-2/LLaMA 的源码把 RoPE旋转位置编码、SwiGLU 激活函数等现代改进加入模型尝试用更大的自己的语料例如维基百科导出数据训练一个 1 亿参数左右的模型感受训练 loss 和生成质量的真实变化。如果本文对你有帮助建议收藏备用如果你在复现代码时遇到问题可以把报错信息和模型配置发在评论区我们一起排查。动手跑一遍代码比看十篇文章的印象都深刻。 用PyTorch从零构建GPT：Transformer核心原理与实战
先问一个问题你平时调用 GPT 接口、用 LangChain 做应用、或者浏览各种大模型的论文解读时有没有想过 GPT 内部的 transformer 层到底是怎么写的注意力机制里的 Q、K、V 矩阵是如何把一句话逐步编码成上下文向量的当我们在说“大语言模型”时它究竟是一堆什么样的张量运算和参数更新市面上关于大语言模型的文章很多但大多是“怎么调用 API”“怎么写提示词”真正从零开始用 PyTorch 搭建 GPT 结构的内容反而比较零散。本文将围绕“用 PyTorch 从零构建 GPT”这一主题把 GPT 的核心原理、模型结构、训练流程、文本生成推理完整走一遍。文章会以实战代码为主适合这几类读者已经会用 PyTorch 搭建简单神经网络想进一步理解 Transformer 和 GPT 源码的开发者。准备阅读 GPT、LLaMA 等大模型源码但被 attention、layer norm、位置编码绕晕的新手。需要在一个小规模数据集上自己做“类 GPT”实验的学生或算法工程师。通过本文你会掌握 GPT 的关键技术点并拿到一份可以在本地训练和推理解的完整代码。1. 背景与核心概念1.1 从语言模型说起“大语言模型”这个词听起来很庞大但把它拆到最小单元核心任务仍然是给定一段已知文本预测下一个词。比如输入“今天天气真”模型要计算下一个字是“好”“差”“热”还是“冷”的概率分布。传统语言模型基于 n-gram 统计只能考虑附近几个词无法处理长距离依赖。后来循环神经网络RNN、LSTM能建模序列但存在训练速度慢、长期记忆衰减的问题而且难以并行处理序列中的所有位置。2017 年《Attention Is All You Need》提出 Transformer彻底改变了序列建模的方式。Transformer 不再像 RNN 那样按时间步逐个读取输入而是把所有 token 一次性送入网络通过自注意力机制计算 token 之间的相关度。这种并行计算方式大幅提升了训练效率也让长距离依赖建模成为可能。GPT 全称是 Generative Pre-trained Transformer属于 Transformer 的 decoder 分支只保留自回归结构。它的训练思路分两阶段第一阶段在海量语料上做无监督预训练目标是预测下一个 token第二阶段可以用带标签数据做有监督微调SFT让模型学会用户期望的回答方式。1.2 GPT 与 Transformer 的关系Transformer 原论文包含 Encoder 和 Decoder 两个部分。Encoder 负责把输入序列编码成语义向量适合理解类任务Decoder 负责逐 token 生成输出适合生成类任务。GPT 只使用 Decoder 的 masked self-attention 部分。它不需要 Encoder因为语言建模的目标就是从历史 token 预测下一个 token当前时刻只能看到过去的信息不能看到未来的 token。这种限制由 attention mask 实现。这里要重点区分三个概念Encoder双向注意力能同时看到上下文适合 BERT 这类理解模型。Decoder单向注意力只能看左侧历史 token适合 GPT 这类生成模型。Encoder-Decoder完整 Transformer适合机器翻译、文本摘要等“序列到序列”任务。GPT 系列模型架构的核心是 Transformer decoder block一个 block 内包含masked multi-head self-attentionfeed-forward networkFFN每个子层都做残差连接和层归一化。堆叠多个这样的 decoder block就得到 GPT 的主干结构。模型规模增大时主要的参数来自 embedding 层、注意力层和 FFN 层。1.3 为什么要用 PyTorch 从零构建 GPT现在有很多现成的大模型框架比如 HuggingFace Transformers 里直接调用GPT2LMHeadModel就能加载预训练模型。那为什么还要自己手写一遍原因主要有三个。第一理解源码。GPT 的 huggingface 实现包含大量工程化逻辑新手很难一眼看出核心张量流。自己写一遍能把nn.Embedding、nn.Linear、masked attention这些组件的输入输出形状理清楚。第二方便实验。在做毕业设计、论文复现或者公司内部小规模预训练实验时我们经常需要改动 attention 结构、换激活函数、改位置编码方式。如果只会调用现成接口改动成本很高自己写核心模型改一行就能测试新想法。第三更深入理解大模型训练中的细节。比如为什么 GPT 需要独立的位置编码表为什么推理时要缓存 KV为什么采样时 temperature 会影响文本多样性这些问题的答案都藏在模型结构的细节里。2. 环境准备与版本说明2.1 操作系统与 Python 环境本文示例代码可以在 Windows、macOS 或 Linux 上运行。因为涉及矩阵运算和自动求导建议先安装 Anaconda 或 Miniconda 来管理 Python 环境。创建独立环境conda create -n gpt_from_scratch python3.10 conda activate gpt_from_scratchPython 版本建议 3.10 或更高。需要说明的是PyTorch 对不同 Python 版本的支持范围不同3.10 是一个非常稳妥的选择。2.2 安装 PyTorchPyTorch 的安装命令取决于你本机的 CUDA 版本。在开始之前先用下面的命令确认显卡驱动和 CUDA 能力nvidia-smi如果输出正常说明有 NVIDIA GPU 可用如果你的机器没有独立显卡也不用担心本文示例模型很小用 CPU 也可以完成训练只是速度会慢一些。以 CUDA 12.x 环境为例安装命令pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121如果你只需要 CPU 版本pip install torch版本说明本文示例代码基于 PyTorch 2.x 编写核心 API 在 1.13 及以上版本都能运行。如果你使用的是更早版本建议升级到 2.x因为 PyTorch 2.x 对编译优化和显存管理做得更好。2.3 安装 tiktoken训练 GPT 时我们需要一个分词器tokenizer。tiktoken 是 OpenAI 开源的 BPE 分词库可以让我们像 GPT 模型一样把文本切分成 token。pip install tiktoken安装完成后用 Python 简单验证一下import tiktoken enc tiktoken.get_encoding(gpt2) tokens enc.encode(Hello, world!) print(tokens) print(enc.decode(tokens))预期输出类似[15496, 11, 995, 0] Hello, world!能看到文本被切分成 token ID 序列并且可以解码回原文本说明环境没问题。3. GPT 核心原理拆解在写代码之前先把 GPT 的几个核心组件逐一拆开。不要小看这一部分后面代码能不能看懂关键就在这里。3.1 Tokenizer文本如何变成数字计算机无法直接处理文本神经网络也一样。我们需要把一段话变成数字序列这个过程叫 tokenization。一种最简单的方法是字符级切分每个字符对应一个 ID。这种方法实现简单但序列很长且无法体现词和子词的语义信息。GPT 使用的是 BPEByte Pair Encoding字节对编码它先把文本按 UTF-8 编码成字节序列然后迭代合并出现频率最高的字节对生成一个子词词表。使用 BPE 分词器后一段文本会被切分成类似这样的 tokenGPT is great - [GPT, is, great]tiktoken的get_encoding(gpt2)使用的就是 GPT-2 时代的 BPE 词表词汇量 50257。在实际训练中如果训练数据比较大一般会重新训练一个 tokenizer为了简化我们这里直接使用现成的 gpt2 分词器。3.2 Token Embedding 与 Position Embedding分词后每个 token 对应一个 ID例如15496。这个整数需要转换成向量才能进入神经网络这一步由nn.Embedding完成。nn.Embedding本质上是一个查找表输入形状为[batch_size, seq_len]的 token ID 矩阵输出形状为[batch_size, seq_len, n_embd]的向量序列。但是 Transformer 模型本身不包含任何序列顺序信息。对于注意力机制来说把句子“我打你”和“你打我”输入进去如果不加位置信息模型无法区分词的先后顺序。所以需要额外加入位置编码。GPT 使用可学习的位置编码表learned positional embedding每个位置都有一个独立向量维度与 token embedding 相同。位置编码表的形状是[block_size, n_embd]其中block_size是模型支持的最大序列长度。在代码中最终输入模型的特征是x token_embedding(tokens) position_embedding(positions)3.3 多头自注意力机制注意力机制是 GPT 的核心。它的作用是让每个 token 在计算特征表示时关注序列中其他 token 的信息。3.3.1 Self-Attention 的 Q、K、V自注意力Self-Attention的输入是同一个序列经过线性变换得到的三组矩阵Query查询向量表示当前 token “想查找什么”Key键向量表示当前 token “能提供什么”Value值向量表示当前 token “实际携带的信息”。对于序列中每个 token我们用它的 Query 向量与所有 token 的 Key 向量做点积得到注意力分数注意力分数经过 softmax 归一化后对 Value 向量加权求和。公式Attention(Q, K, V) softmax(Q * K^T / sqrt(d_k)) * V其中d_k是 Key 向量的维度除以sqrt(d_k)是为了防止点积值过大导致 softmax 梯度消失。Python 伪代码如下import torch import torch.nn as nn import torch.nn.functional as F def scaled_dot_product_attention(Q, K, V, maskNone): d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / d_k ** 0.5 if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights F.softmax(scores, dim-1) return torch.matmul(attn_weights, V), attn_weights3.3.2 Masked Attention为什么 GPT 只能看到过去GPT 是自回归模型生成第t个 token 时只能看到前t-1个 token。如果在训练时让当前位置看到了未来的 token模型就相当于照抄答案学不到真正的预测能力。实现方式很直接构造一个上三角矩阵把未来位置的注意力分数设置为负无穷softmax 之后这些位置的权重就变成 0。例如序列长度为 4合法的 mask 矩阵为1 0 0 0 1 1 0 0 1 1 1 0 1 1 1 1第i行第j列为 1 表示位置i可以关注位置j为 0 表示不允许关注。3.3.3 多头注意力多头注意力就是把注意力计算分成多个头并行做。假设特征维度是 768使用 12 个头那么每个头的维度是 64。每个头独立计算注意力最后拼接起来再过一层线性变换。多头注意力的好处是让模型从不同表示子空间捕捉信息。有的头可能关注语法关系有的头可能关注远距离指代关系。代码中可以用一个nn.Linear同时生成全部头的 Q、K、V再通过 reshape 拆分到不同头这样实现高效很多。3.4 前馈神经网络与残差、LayerNorm每个 Transformer block 在 attention 之后还会接一个前馈神经网络FFN通常包含两层全连接和一个激活函数FFN(x) GELU(x W1 b1) W2 b2GPT 中 FFN 的中间维度一般是模型维度的 4 倍也就是说 768 维的输入会先升到 3072 维再投影回 768 维。同时每个子层都使用“残差连接 层归一化”的结构残差连接Residual Connection让梯度能直接跨层传播避免深层网络梯度消失层归一化LayerNorm对每个样本的特征维做标准化加速训练收敛。注意 GPT 使用的是 pre-LayerNorm 结构即先做 LayerNorm 再做 Attention/FFN。这与原始 Transformer 论文中的 post-LayerNorm 略有不同但现代大模型几乎都采用 pre-LayerNorm训练更稳定。3.5 语言模型头与损失函数经过多层 Transformer block 后我们得到形状为[batch_size, seq_len, n_embd]的向量序列。要预测下一个 token需要把向量映射回词表大小这一层称为语言模型头LM Headlogits nn.Linear(n_embd, vocab_size)(x)得到每个位置上所有 token 的概率分布后训练目标就是交叉熵损失计算预测分布与真实下一个 token 的差距。假设输入序列是[1, 2, 3, 4]模型会同时预测给定[1]预测2给定[1, 2]预测3给定[1, 2, 3]预测4。这就是 GPT 的“并行 teacher forcing”训练方式一次前向计算得到整个序列所有位置的预测结果。4. 完整实战用 PyTorch 从零构建 GPT现在进入本文的主体部分。我们将搭建一个 mini 版 GPT 模型并在一个小型数据集上完成训练和文本生成。模型规模很小完全可以在 CPU 上运行。4.1 项目结构建议创建如下目录gpt_from_scratch/ ├── data.py # 数据加载与 tokenizer 封装 ├── model.py # GPT 模型结构 ├── train.py # 训练脚本 └── generate.py # 文本生成脚本4.2 数据准备为了演示我们使用一个英文小语料莎士比亚的《哈姆雷特》部分片段或者任何一本公开领域的英文书。这里为了代码简单直接内置一段文本作为训练语料你也可以用自己准备的.txt文件替换。数据加载代码# 文件路径data.py import tiktoken import torch from torch.utils.data import Dataset class TextDataset(Dataset): def __init__(self, text, block_size128): self.enc tiktoken.get_encoding(gpt2) self.tokens self.enc.encode(text) self.block_size block_size def __len__(self): return len(self.tokens) - self.block_size def __getitem__(self, idx): chunk self.tokens[idx: idx self.block_size 1] x torch.tensor(chunk[:-1], dtypetorch.long) y torch.tensor(chunk[1:], dtypetorch.long) return x, y def load_text(): sample_text To be, or not to be, that is the question: Whether tis nobler in the mind to suffer The slings and arrows of outrageous fortune, Or to take arms against a sea of troubles And by opposing end them. To die, to sleep; No more; and by a sleep to say we end The heart-ache and the thousand natural shocks That flesh is heir to: tis a consummation Devoutly to be wishd. To die, to sleep; To sleep, perchance to dream; ay, theres the rub. return sample_text这里的关键点每个样本的输入x和标签y是错位一个 token 的序列对应“预测下一个 token”的任务。tiktoken编码后文本中的空格和换行也会被保留为 token不需要额外清理。4.3 定义 GPT 模型下面编写完整的 GPT 模型代码。代码尽量简洁但保留 GPT 所有核心组件。# 文件路径model.py import torch import torch.nn as nn import torch.nn.functional as F class LayerNorm(nn.Module): def __init__(self, n_embd, eps1e-5): super().__init__() self.eps eps self.gamma nn.Parameter(torch.ones(n_embd)) self.beta nn.Parameter(torch.zeros(n_embd)) def forward(self, x): mean x.mean(-1, keepdimTrue) var x.var(-1, keepdimTrue, unbiasedFalse) x (x - mean) / torch.sqrt(var self.eps) return x * self.gamma self.beta class CausalSelfAttention(nn.Module): def __init__(self, n_embd, n_head, block_size, dropout0.0): super().__init__() assert n_embd % n_head 0 self.n_head n_head self.n_embd n_embd self.head_dim n_embd // n_head self.c_attn nn.Linear(n_embd, 3 * n_embd, biasFalse) self.c_proj nn.Linear(n_embd, n_embd, biasFalse) self.dropout nn.Dropout(dropout) # 上三角掩码矩阵用于屏蔽未来位置 self.register_buffer( mask, torch.tril(torch.ones(block_size, block_size)).view( 1, 1, block_size, block_size ), ) def forward(self, x): B, T, C x.shape # 一次性生成 Q、K、V qkv self.c_attn(x) # [B, T, 3*C] q, k, v torch.split(qkv, self.n_embd, dim-1) # 拆分成多头: [B, n_head, T, head_dim] q q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) v v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) # 缩放点积注意力 att (q k.transpose(-2, -1)) * (self.head_dim ** -0.5) att att.masked_fill(self.mask[:, :, :T, :T] 0, float(-inf)) att F.softmax(att, dim-1) att self.dropout(att) y att v # [B, n_head, T, head_dim] y y.transpose(1, 2).contiguous().view(B, T, C) y self.c_proj(y) return y class MLP(nn.Module): def __init__(self, n_embd, dropout0.0): super().__init__() self.fc1 nn.Linear(n_embd, 4 * n_embd) self.fc2 nn.Linear(4 * n_embd, n_embd) self.dropout nn.Dropout(dropout) def forward(self, x): x self.fc1(x) x F.gelu(x) x self.fc2(x) x self.dropout(x) return x class TransformerBlock(nn.Module): def __init__(self, n_embd, n_head, block_size, dropout0.0): super().__init__() self.ln1 LayerNorm(n_embd) self.attn CausalSelfAttention(n_embd, n_head, block_size, dropout) self.ln2 LayerNorm(n_embd) self.mlp MLP(n_embd, dropout) def forward(self, x): x x self.attn(self.ln1(x)) x x self.mlp(self.ln2(x)) return x class GPT(nn.Module): def __init__( self, vocab_size50257, n_embd128, n_head4, n_layer4, block_size128, dropout0.0, ): super().__init__() self.token_embedding nn.Embedding(vocab_size, n_embd) self.position_embedding nn.Embedding(block_size, n_embd) self.blocks nn.Sequential( *[ TransformerBlock(n_embd, n_head, block_size, dropout) for _ in range(n_layer) ] ) self.ln_f LayerNorm(n_embd) self.lm_head nn.Linear(n_embd, vocab_size, biasFalse) # 权重初始化 self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean0.0, std0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean0.0, std0.02) def forward(self, idx, targetsNone): B, T idx.shape assert T self.position_embedding.weight.shape[0], 序列长度超出位置编码范围 pos torch.arange(0, T, dtypetorch.long, deviceidx.device) pos pos.unsqueeze(0) # [1, T] x self.token_embedding(idx) self.position_embedding(pos) x self.blocks(x) x self.ln_f(x) logits self.lm_head(x) loss None if targets is not None: B, T, C logits.shape logits logits.view(B * T, C) targets targets.view(B * T) loss F.cross_entropy(logits, targets) return logits, loss def generate(self, idx, max_new_tokens100, temperature1.0): 自回归生成新 token self.eval() for _ in range(max_new_tokens): idx_cond idx[:, -self.position_embedding.weight.shape[0]:] logits, _ self(idx_cond) logits logits[:, -1, :] / temperature probs F.softmax(logits, dim-1) next_token torch.multinomial(probs, num_samples1) idx torch.cat((idx, next_token), dim1) return idx这个模型需要注意几个地方register_buffer注册的 mask 会随模型一起移动到 GPU/CPU不需要手动管理设备。c_attn一次性输出 3 倍维度的 Q、K、V比分别定义三个线性层效率更高。generate方法里只取最后一个位置的 logits 做采样这是因为自回归模型当前位置的预测只依赖之前所有 token。采样时除以 temperature让概率分布更尖锐或更平滑从而控制生成文本的随机性。4.4 编写训练脚本训练脚本中我们使用 AdamW 优化器。AdamW 是 GPT 训练的标配优化器因为它对权重衰减的处理更合理能提升泛化性能。# 文件路径train.py import torch from torch.utils.data import DataLoader from data import TextDataset, load_text from model import GPT def main(): # 超参数 batch_size 8 block_size 128 n_embd 128 n_head 4 n_layer 4 learning_rate 3e-4 epochs 50 device cuda if torch.cuda.is_available() else cpu text load_text() dataset TextDataset(text, block_sizeblock_size) loader DataLoader(dataset, batch_sizebatch_size, shuffleTrue) model GPT( vocab_size50257, n_embdn_embd, n_headn_head, n_layern_layer, block_sizeblock_size, dropout0.1, ).to(device) optimizer torch.optim.AdamW(model.parameters(), lrlearning_rate, weight_decay0.01) print(f数据集 token 数: {len(dataset.tokens)}) print(f使用设备: {device}) for epoch in range(epochs): model.train() total_loss 0 for x, y in loader: x, y x.to(device), y.to(device) optimizer.zero_grad() logits, loss model(x, y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_loss loss.item() avg_loss total_loss / len(loader) print(fEpoch {epoch1:3d}/{epochs}, loss {avg_loss:.4f}) torch.save(model.state_dict(), gpt_model.pt) print(训练完成模型已保存到 gpt_model.pt) if __name__ __main__: main()训练过程中值得关注的点clip_grad_norm_是梯度裁剪防止梯度爆炸。对语言模型来说这是稳定训练的重要技巧。数据量很小时loss 下降会比较快但生成的文本可能仍然不通顺。实际项目中需要更大规模的语料和更多的训练步数。这里模型参数量大约 800 万远小于真实的 GPT 模型所以 CPU 完全可以跑。4.5 编写文本生成脚本训练完成后加载模型权重并生成文本。# 文件路径generate.py import torch import tiktoken from model import GPT def main(): device cuda if torch.cuda.is_available() else cpu model GPT( vocab_size50257, n_embd128, n_head4, n_layer4, block_size128, dropout0.0, ).to(device) model.load_state_dict(torch.load(gpt_model.pt, map_locationdevice)) model.eval() enc tiktoken.get_encoding(gpt2) prompt To be, or not to be input_ids enc.encode(prompt) input_ids torch.tensor([input_ids], dtypetorch.long, devicedevice) output_ids model.generate(input_ids, max_new_tokens50, temperature0.8) output_text enc.decode(output_ids[0].tolist()) print(生成结果) print(output_text) if __name__ __main__: main()4.6 运行与预期结果依次运行python train.py python generate.py训练时如果数据量较小loss 可能在 0.1 以内如果你替换成了自己的语料loss 会在几十个 epoch 后逐渐下降。生成的部分文本大概率不完全通顺因为我们的训练数据只有一小段莎士比亚原文。但你能明显看出模型学到了一些英文单词拼写和常见词搭配这说明模型结构本身是有效的。换用更大的语料、更长的训练时间后生成质量会显著改善。5. 常见问题与排查思路5.1 常见报错汇总下表整理了手写 GPT 过程中最常见的几个问题问题现象常见原因解决思路训练时 loss 没有下降学习率过大/过小或 tokenizer 词表与模型 vocab_size 不匹配检查学习率确认tiktoken.get_encoding(gpt2)的词汇量与模型vocab_size一致运行时提示“序列长度超出位置编码范围”输入 token 长度超过block_size对输入做截断或增大block_size生成的文本全是乱码或重复字符采样温度过低或模型训练不充分调高 temperature增加训练数据量CUDA out of memory模型过大或 batch_size 过大降低 batch_size减少 block_size或使用梯度累积注意力分数出现 NaN数据中存在过大的数值网络未做梯度裁剪增加 LayerNorm 稳定性添加梯度裁剪5.2 关键排查步骤如果你训练出来的模型效果很差建议按下面顺序排查第一步确认数据预处理。打印一个样本的x和y确认y确实是x的下一个 token。如果错位逻辑写错模型永远学不对。第二步确认 mask 是否正确。把 mask 打印出来看是否上三角全为 0。如果 mask 没生效模型在训练时会偷看未来 token导致训练 loss 很低但生成效果很差。第三步确认 embedding 层维度。GPT 中 token embedding 和 position embedding 矩阵都是随机初始化的如果初始化不合理训练可能非常慢。建议使用均值为 0、标准差 0.02 的正态分布初始化这也是 GPT-2 原始代码的做法。第四步确认优化器和学习率。语言模型常用的学习率区间是1e-4到3e-4过大会导致 loss 震荡。同时不要忘记梯度裁剪。6. 最佳实践与工程建议6.1 从小规模实验出发不要一上来就训练一个 7B 参数的大模型。工程实践中先用一个几十 MB 的小语料、小模型几百万参数把代码流程跑通确认 loss 能正常下降、文本能正常生成然后逐步扩大规模。这样做能节省大量调试时间。6.2 使用梯度累积扩展 batch size真实的大语言模型训练往往需要很大的 batch size但单张显卡的显存有限。梯度累积的思路是先计算多个 mini-batch 的梯度然后将梯度累加最后做一次参数更新效果等价于更大的 batch size。PyTorch 中的实现思路如下accumulation_steps 4 optimizer.zero_grad() for i, (x, y) in enumerate(loader): x, y x.to(device), y.to(device) loss model(x, y)[1] / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意 loss 除以accumulation_steps是为了让梯度累加后的量级与正常 batch size 训练保持一致。6.3 混合精度训练PyTorch 2.x 中可以使用torch.autocast和GradScaler做混合精度训练。在支持 CUDA 的环境下混合精度能显著减少显存占用并提速训练。scaler torch.amp.GradScaler(cuda) for x, y in loader: x, y x.to(device), y.to(device) optimizer.zero_grad() with torch.amp.autocast(cuda): _, loss model(x, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度并不是必须的但对大规模训练非常有帮助。CPU 环境下不需要使用。6.4 推理阶段的 KV Cache在文本生成时模型需要逐个 token 生成。如果不做任何缓存每生成一个新 token都要重新计算之前所有 token 的 Key 和 Value计算量很大。真实 GPT 推理例如 HuggingFace 的generate会使用 KV Cache把之前 step 计算出来的 K、V 矩阵缓存下来下一轮只计算新 token 的 Q、K、V与缓存拼接后再做注意力计算。KV Cache 是理解大模型推理性能的核心概念。本文为了保持代码简洁没有实现但在深入研究大模型部署、vLLM、TensorRT-LLM 等推理框架时这个机制必须掌握。6.5 权重保存与加载训练完成后保存权重时最好同时保存模型配置和优化器状态方便断点续训。# 保存 torch.save({ model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), epoch: epoch, config: { n_embd: n_embd, n_head: n_head, n_layer: n_layer, block_size: block_size, vocab_size: 50257, } }, checkpoint.pt) # 恢复 checkpoint torch.load(checkpoint.pt) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict])6.6 数据质量决定模型上限大语言模型的训练遵循一个规律模型结构决定能力的下限数据质量决定能力的上限。如果你在自己的项目里微调 GPT不要只关心模型参数量更要投入时间清洗训练数据过滤重复文本、去掉格式错乱内容、尽量保证语料的多样性。7. 总结与下一步学习建议这篇文章从零实现了一个 mini 版 GPT核心代码只有不到 300 行。你已经掌握了 Token Embedding、Position Embedding、多头因果自注意力、LayerNorm、残差连接、FFN 和语言模型头的完整实现也能用自己的脚本训练一个小模型并生成文本。下一步可以从这几个方向继续深入把字符级或者 BPE 分词器替换成 SentencePiece训练支持中文的 tokenizer在模型中加入 KV Cache对比推理速度的提升参考 GPT-2/LLaMA 的源码把 RoPE旋转位置编码、SwiGLU 激活函数等现代改进加入模型尝试用更大的自己的语料例如维基百科导出数据训练一个 1 亿参数左右的模型感受训练 loss 和生成质量的真实变化。如果本文对你有帮助建议收藏备用如果你在复现代码时遇到问题可以把报错信息和模型配置发在评论区我们一起排查。动手跑一遍代码比看十篇文章的印象都深刻。