先问一个问题:你平时调用 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-attention;
- feed-forward network(FFN);
- 每个子层都做残差连接和层归一化。
堆叠多个这样的 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 python=3.10 conda activate gpt_from_scratchPython 版本建议 3.10 或更高。需要说明的是,PyTorch 对不同 Python 版本的支持范围不同,3.10 是一个非常稳妥的选择。
2.2 安装 PyTorch
PyTorch 的安装命令取决于你本机的 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 使用的是 BPE(Byte Pair Encoding,字节对编码),它先把文本按 UTF-8 编码成字节序列,然后迭代合并出现频率最高的字节对,生成一个子词词表。
使用 BPE 分词器后,一段文本会被切分成类似这样的 token:
"GPT 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, mask=None): 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 Head):
logits = 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_size=128): 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], dtype=torch.long) y = torch.tensor(chunk[1:], dtype=torch.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 wish'd. To die, to sleep; To sleep, perchance to dream; ay, there's 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, eps=1e-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, keepdim=True) var = x.var(-1, keepdim=True, unbiased=False) 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, dropout=0.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, bias=False) self.c_proj = nn.Linear(n_embd, n_embd, bias=False) 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, dropout=0.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, dropout=0.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_size=50257, n_embd=128, n_head=4, n_layer=4, block_size=128, dropout=0.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, bias=False) # 权重初始化 self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, idx, targets=None): B, T = idx.shape assert T <= self.position_embedding.weight.shape[0], "序列长度超出位置编码范围" pos = torch.arange(0, T, dtype=torch.long, device=idx.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_tokens=100, temperature=1.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_samples=1) idx = torch.cat((idx, next_token), dim=1) 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_size=block_size) loader = DataLoader(dataset, batch_size=batch_size, shuffle=True) model = GPT( vocab_size=50257, n_embd=n_embd, n_head=n_head, n_layer=n_layer, block_size=block_size, dropout=0.1, ).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.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(f"Epoch {epoch+1: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_size=50257, n_embd=128, n_head=4, n_layer=4, block_size=128, dropout=0.0, ).to(device) model.load_state_dict(torch.load("gpt_model.pt", map_location=device)) 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], dtype=torch.long, device=device) output_ids = model.generate(input_ids, max_new_tokens=50, temperature=0.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 和生成质量的真实变化。
如果本文对你有帮助,建议收藏备用;如果你在复现代码时遇到问题,可以把报错信息和模型配置发在评论区,我们一起排查。动手跑一遍代码,比看十篇文章的印象都深刻。