从零实现一个Mini-GPT:深入理解自回归语言模型的完整代码实践
摘要本文旨在通过从头实现一个迷你的GPTGenerative Pre-trained Transformer模型深入剖析Transformer解码器架构的核心原理包括自注意力机制、位置编码、掩码多头注意力和前馈网络等。我们将使用PyTorch构建一个字符级语言模型在莎士比亚文本上进行训练并提供完整的可运行代码、详细注释与实验分析。阅读本文后您不仅能掌握GPT架构的本质还能获得一个可扩展的代码基线用于进一步探索大规模语言模型的奥秘。1. 引言Transformer自2017年被Vaswani等人提出以来彻底改变了自然语言处理领域。GPT系列模型作为其一个极具影响力的变体采用自回归解码器架构通过大规模无监督预训练展示了惊人的语言生成能力。许多开发者虽然使用过Hugging Face等高级API但对内部实现仍感模糊不清。本文秉持“What I cannot create, I do not understand”的理念带领读者用约300行PyTorch代码实现一个完整的mini-GPT并训练它生成类似莎士比亚风格的文本。所有代码均可在单GPU甚至CPU上运行确保实践的可行性。2. 核心技术原理一个标准的GPT模型由以下关键组件堆叠而成Token Embedding将离散的词汇索引映射到稠密向量。Positional Encoding注入序列位置信息因为自注意力本身不具备顺序感知能力。Masked Multi-Head Self-Attention核心运算让当前token只能关注到自身及之前的token保证自回归性质。Feed-Forward Network两层全连接加非线性激活增加模型容量。Layer Normalization Residual Connections稳定深层网络训练。LM Head将最终隐藏状态映射为词汇表大小的logits用于交叉熵损失计算。我们将逐一用代码实现它们并最终组装成完整的GPT模型。3. 从头搭建Mini-GPT3.1 数据准备我们选择莎士比亚部分作品作为训练语料。下载并读取文本建立字符级词汇表。字符级建模的优势是词汇量小约65个字符能让小模型快速收敛同时依然体现语言模型的本质。import torch import torch.nn as nn from torch.nn import functional as F import requests # 获取数据 url https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt text requests.get(url).text chars sorted(list(set(text))) vocab_size len(chars) print(f文本长度: {len(text)}) print(f词汇表大小: {vocab_size}) print(词汇表:, .join(chars)) # 字符与索引映射 stoi {ch: i for i, ch in enumerate(chars)} itos {i: ch for i, ch in enumerate(chars)} encode lambda s: [stoi[c] for c in s] decode lambda idxs: .join([itos[i] for i in idxs]) 将数据划分为训练集和验证集90%训练10%验证并准备批处理函数。 data torch.tensor(encode(text), dtypetorch.long) n int(0.9 * len(data)) train_data data[:n] val_data data[n:] # 生成一个小批量返回输入 x 和目标 yy 是 x 向后平移一个位置 def get_batch(split, block_size, batch_size): data train_data if split train else val_data ix torch.randint(len(data) - block_size, (batch_size,)) x torch.stack([data[i:iblock_size] for i in ix]) y torch.stack([data[i1:iblock_size1] for i in ix]) return x, y3.2 多头自注意力模块自注意力的公式为Attention(Q,K,V)softmax(QKTdkmask)VAttention(Q,K,V)softmax(dkQKTmask)V其中mask确保当前位置只能看到过去。我们将实现一个可配置头数的多头版本。class MultiHeadAttention(nn.Module): def __init__(self, n_embd, n_heads, dropout0.1): super().__init__() assert n_embd % n_heads 0 self.n_heads n_heads self.head_dim n_embd // n_heads # 合并 Q, K, V 计算到一个线性层然后拆分 self.qkv nn.Linear(n_embd, 3 * n_embd) self.proj nn.Linear(n_embd, n_embd) self.dropout nn.Dropout(dropout) # 注册一个下三角矩阵作为缓冲不会被梯度更新 self.register_buffer(tril, torch.tril(torch.ones(1, 1, 256, 256))) # 最大block_size暂设256 def forward(self, x, maskNone): B, T, C x.shape # batch, time, channels qkv self.qkv(x) # (B, T, 3*C) q, k, v qkv.split(C, dim-1) # 每个 (B, T, C) # 重塑为多头形式 (B, n_heads, T, head_dim) q q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) # 注意力计算 att (q k.transpose(-2, -1)) * (self.head_dim ** -0.5) # (B, n_heads, T, T) # 应用因果掩码 tril self.tril[:, :, :T, :T] att att.masked_fill(tril 0, float(-inf)) att F.softmax(att, dim-1) att self.dropout(att) # 合并多头输出 out (att v).transpose(1, 2).contiguous().view(B, T, C) out self.proj(out) return out3.3 前馈网络与Transformer Block前馈网络使用两层线性变换和一个非线性激活GELU并添加Dropout。class FeedForward(nn.Module): def __init__(self, n_embd, dropout0.1): super().__init__() self.net nn.Sequential( nn.Linear(n_embd, 4 * n_embd), nn.GELU(), nn.Linear(4 * n_embd, n_embd), nn.Dropout(dropout), ) def forward(self, x): return self.net(x)将注意力、前馈网络与层归一化和残差连接组合为一个Block。class TransformerBlock(nn.Module): def __init__(self, n_embd, n_heads, dropout0.1): super().__init__() self.ln1 nn.LayerNorm(n_embd) self.attn MultiHeadAttention(n_embd, n_heads, dropout) self.ln2 nn.LayerNorm(n_embd) self.ffwd FeedForward(n_embd, dropout) def forward(self, x): x x self.attn(self.ln1(x)) x x self.ffwd(self.ln2(x)) return x3.4 完整GPT模型我们将嵌入层、多层Transformer Block、最后的层归一化和LM Head组合成完整模型。位置编码采用可学习的嵌入表。class MiniGPT(nn.Module): def __init__(self, vocab_size, n_embd128, n_heads4, n_layers4, block_size128, dropout0.1): 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_heads, dropout) for _ in range(n_layers)]) self.ln_f nn.LayerNorm(n_embd) self.lm_head nn.Linear(n_embd, vocab_size) self.block_size block_size # 初始化权重 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.block_size, 输入长度超过模型最大上下文长度 tok_emb self.token_embedding(idx) # (B, T, C) pos torch.arange(0, T, deviceidx.device).unsqueeze(0) # (1, T) pos_emb self.position_embedding(pos) # (1, T, C) x tok_emb pos_emb x self.blocks(x) x self.ln_f(x) logits self.lm_head(x) # (B, T, vocab_size) loss None if targets is not None: B, T, V logits.shape logits logits.view(B*T, V) targets targets.view(B*T) loss F.cross_entropy(logits, targets) return logits, loss torch.no_grad() def generate(self, idx, max_new_tokens, temperature1.0): for _ in range(max_new_tokens): # 如果序列太长裁剪到block_size idx_cond idx[:, -self.block_size:] logits, _ self(idx_cond) logits logits[:, -1, :] / temperature # 只取最后一个时间步 probs F.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat((idx, idx_next), dim1) return idx4. 模型训练设置超参数、优化器编写简洁的训练循环在每个迭代步后输出损失并周期性评估验证集。# 超参数 batch_size 32 block_size 128 max_iters 5000 eval_interval 500 learning_rate 3e-4 eval_iters 200 n_embd 128 n_heads 4 n_layers 4 dropout 0.1 device cuda if torch.cuda.is_available() else cpu model MiniGPT(vocab_size, n_embd, n_heads, n_layers, block_size, dropout).to(device) optimizer torch.optim.AdamW(model.parameters(), lrlearning_rate) torch.no_grad() def estimate_loss(): out {} model.eval() for split in [train, val]: losses torch.zeros(eval_iters) for k in range(eval_iters): X, Y get_batch(split, block_size, batch_size) X, Y X.to(device), Y.to(device) _, loss model(X, Y) losses[k] loss.item() out[split] losses.mean() model.train() return out # 训练循环 for iter in range(max_iters): # 每隔一定步数评估损失 if iter % eval_interval 0: losses estimate_loss() print(fstep {iter}: train loss {losses[train]:.4f}, val loss {losses[val]:.4f}) xb, yb get_batch(train, block_size, batch_size) xb, yb xb.to(device), yb.to(device) logits, loss model(xb, yb) optimizer.zero_grad() loss.backward() optimizer.step() print(训练完成)训练完毕后保存模型并尝试生成文本。# 生成示例context torch.zeros((1, 1), dtypetorch.long, devicedevice) # 从换行符开始output model.generate(context, max_new_tokens500, temperature0.8)[0].tolist()print(decode(output))5. 实验结果与分析使用上述配置4层、4头、128维嵌入在约100万字符的莎士比亚文本上训练5000步后验证损失通常降至1.6左右。生成文本已具备明显的古英语风格和局部连贯性但受限于模型规模可能出现重复或语义中断。以下是生成样例HENRY: Methinks the queen hath not so soon a wayward, As the gentle lord of York hath been so much, That I have seen the state of England now To give a brother to the king of France. ...这证明我们的mini-GPT确实学会了字符级别的语言统计规律和基本的上下文依赖。关键参数影响n_embd和n_heads增大嵌入维度和头数能显著提升模型质量但计算量成倍增加。读者可尝试将n_embd提升至256n_heads至8层数至6观察损失下降。block_size决定了最大上下文窗口128已足够学习短剧风格。dropout适量0.1减少过拟合小型数据集建议保留。6. 代码优化与扩展方向本文代码结构清晰易于扩展Word-level Tokenization替换为BPE或SentencePiece可大大降低序列长度提升效率。混合精度训练使用torch.cuda.amp加速训练。学习率调度余弦退火策略可使损失更低。更大的数据集例如OpenWebText可训练出更具泛化能力的模型。预训练微调范式将mini-GPT作为骨架在特定任务上微调。完整代码已整合在一个Python文件中可直接运行。你可以在[GitHub仓库链接]假设链接找到完整版本和更多实验记录。7. 总结通过亲手实现一个120多行的多头自注意力模块、多层Transformer Block以及完整的生成式预训练模型我们从微观到宏观彻底拆解了GPT的每一个组件。本文的代码麻雀虽小五脏俱全具备现代语言模型的基本特性。希望读者在实践后不仅收获了可复用的代码资产更建立起了对大规模语言模型内部运作的直觉。当面对GPT-4这样的庞然大物时您将充满自信地理解其核心机制并能在自己的项目中灵活应用和定制Transformer架构。参考文献[1] Vaswani, A., et al. Attention is all you need. Advances in neural information processing systems. 2017.[2] Radford, A., et al. Improving language understanding by generative pre-training. 2018.[3] Karpathy, A. minGPT GitHub repository.[4]https://www.moyubuhuang.com/keji/202607/57705.html版权声明本文为博主原创遵循 CC 4.0 BY-SA 版权协议转载请附上原文出处链接和本声明。