ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

从零构建迷你GPT:用PyTorch实现一个小型语言模型

从零构建迷你GPT:用PyTorch实现一个小型语言模型 并不是每个学深度学习的人都需要从零训练一个大模型但每个想真正理解 LLM 的人都应该亲手构建一个“最小可用版本”。很多同学学到神经网络、反向传播、PyTorch 基础之后进入大模型阶段会突然迷茫市面上到处是 huggingface 加载模型、调 API、做微调但模型内部的 Embedding、Attention、LayerNorm、Softmax 到底怎么协作为什么模型能记住上下文训练时损失函数究竟在算什么如果你也卡在这个位置这篇文章就是为你准备的。第 15-16 章这个阶段核心目标不是让你复现 GPT-4而是把你之前学过的 Python 深度学习知识全部串起来数据准备、张量操作、自定义网络结构、训练循环、推理采样最终跑出一个基于 Transformer 解码器的小型 LLM。它能学习文本规律能根据 prompt 续写内容甚至能输出语法通顺的中文。这篇文章我会按照从零构建 LLM 的完整流程展开先讲清楚大模型的核心原理和关键概念再给出可运行的 PyTorch 代码包括数据 tokenizer、迷你 GPT 模型、训练脚本、文本生成脚本最后解决常见问题并给出工程化建议。读完你会有两方面的收获一是对 LLM 架构不再停留在“听说过”二是能真真切切训练出一个几百万参数的小型语言模型。1. 从零构建 LLM 为什么值得做1.1 只调 API 和真正懂 LLM 之间差一个训练过程先给一个判断理解大模型最好的方式不是读论文也不是调 API而是亲手训练一个极小的 GPT。原因很简单。当你调用 OpenAI 或开源模型时你面对的是一个黑盒输入一句话得到一堆 token中间发生了什么完全看不见。而当你从零写一个 mini GPT你就必须回答这些问题token 是怎么变成向量的位置信息怎么注入注意力机制如何让模型“看到”前面的词多层堆叠如何提取更高层的语义训练时的交叉熵损失为什么能引导模型做“下一个词预测”。这些问题任何一个没有真正搞懂代码就写不出来训练也跑不通。1.2 小型模型并不是没有价值有人会问训练一个几百万参数的小模型和真实的大模型相比差距太大有什么意义意义在于原理完全一样。GPT-3、GPT-4 的架构核心也是 Transformer 解码器也是下一个词预测任务。你训练小模型时踩过的坑——数据质量不够、过拟合、损失不下降、生成结果重复——在大模型训练和大模型应用时依然存在。从工程角度看直接从 HuggingFace 加载一个 7B 模型很简单但如果你要定制 tokenizer、调整模型结构、做模型压缩、部署到边缘设备就必须理解模型内部结构。这些能力只有在“从零实现”中才能真正获得。1.3 本文的学习路径我会采用“最小闭环”的方式推进准备一段文本语料实现字符级 tokenizer构造训练样本实现一个迷你 GPT 模型训练若干轮用训练好的模型生成文本。整个流程跑通后你就有能力继续扩展换成 BPE tokenizer、加大模型、使用真实数据集、做指令微调。2. LLM 核心原理从“下一个词预测”说起2.1 语言模型在做什么传统语言模型的目标是给定一个词序列 (w_1, w_2, ..., w_t)预测下一个词 (w_{t1}) 的概率分布。写成公式是[ P(w_{t1} | w_1, w_2, ..., w_t) ]大语言模型把这个思想放大了无数倍用海量文本训练一个巨大的神经网络让它可以处理任意长度的上下文预测下一个 token。训练完成之后模型内部已经“记住”了大量语言规律和世界知识于是就能完成翻译、问答、摘要、代码生成等任务。GPT 系列的思路非常直接用 Transformer 解码器作为网络结构用 text next token 作为监督信号不需要人工标注数据。2.2 输入侧token 和 Embedding原始文本不能直接输入神经网络因此需要先切分成 token。中文场景中最简单的做法是按字符切分工业界更常用 BPE、SentencePiece 等子词切分方法。每个 token 会被映射成一个稠密向量这个过程叫做 Embedding嵌入。从数据流上看文本字符串 - 字符/token 列表 - 整数索引 - Embedding 向量 - 输入模型2.3 结构侧Transformer 解码器一个典型的 GPT 风格模型由若干层 Transformer 解码器堆叠而成。每一层包含Masked Self-Attention带因果掩码的多头自注意力保证模型只能看到当前位置及之前的 tokenLayerNorm层归一化稳定训练Feed-Forward Network前馈网络对每个位置做非线性变换残差连接帮助梯度流动。最终通过一个 Linear 层把隐藏状态映射到词表大小的 logits再经过 Softmax 得到概率分布。这里有一个初学者容易困惑的点为什么需要 Mask 掩码因为在预测第 (t) 个词时模型必须假装看不到后面的词。如果能看到未来信息训练目标就变成了“抄答案”模型无法学到真正的语言规律。因此我们会在注意力矩阵上做一个上三角掩码把未来位置全部设为负无穷。2.4 训练侧交叉熵损失模型输出的概率分布和真实下一个 token 之间用交叉熵作为损失函数。[ \mathcal{L} -\log P(\text{真实下一个 token}) ]针对一个 batch 的多个位置损失等价于对所有位置的交叉熵取平均。PyTorch 中的F.cross_entropy可以直接处理形状为[B, T, C]的 logits 和形状为[B, T]的目标 token 索引。2.5 推理侧自回归生成训练完成后模型进入自回归模式输入一个 prompt得到下一个 token 的概率分布根据一定策略采样出下一个 token把新 token 拼接到输入中重复以上步骤直到达到指定长度或遇到结束符。这个设计就是 GPT 生成文本的基本逻辑。3. 环境准备与实验项目结构3.1 运行环境本文代码基于 Python 和 PyTorch实验环境建议如下以实际安装为准软件说明操作系统Windows / Linux / macOS 均可Python3.9 或 3.10 以上PyTorch2.x 版本CPU 即可运行GPU 会更快依赖库仅需 PyTorch不需要额外的大模型库建议使用虚拟环境避免污染系统环境。在项目目录下执行python -m venv venv source venv/bin/activate # Windows 下为 venv\Scripts\activate pip install torch如果你有 NVIDIA GPU可以到 PyTorch 官网选择对应的安装命令安装 CUDA 版本。没有 GPU 也能跑通本文实验只是模型规模需要小一点。3.2 项目结构为了便于维护建议采用如下结构mini_llm/ ├── data.py # 数据读取、tokenizer 实现 ├── model.py # MiniGPT 模型定义 ├── train.py # 训练脚本 ├── generate.py # 文本生成脚本 └── output/ # 保存模型权重4. 核心流程拆解数据、模型、训练、生成4.1 数据准备和 tokenizer为了从头演示原理这里不使用 HuggingFace而是实现一个最简单的字符级 tokenizer。这种方式虽然效率不高但逻辑透明非常适合学习。语料选择方面建议找一段无版权争议的中文文本。我使用《论语》和《道德经》的经典语句拼接作为示例语料。你完全可以换成任意文本文件。字符级 tokenizer 的步骤统计文本中出现的所有字符生成 vocab建立字符 - id和id - 字符两个映射定义 encode 和 decode 函数。4.2 构建训练样本训练样本的构造逻辑是“滑动窗口”固定一个上下文长度 block_size从原始文本中随机切出连续 block_size 个字符作为输入 x把窗口右移一位的 block_size 个字符作为目标 y。例如原始文本是“道可道非常道”block_size4那么x: 道可道非 y: 可道非常也就是让模型根据“道可道非”预测“可道非常”。4.3 模型设计本文实现的 MiniGPT 采用 GPT 的核心设计但参数规模极小适合教学和 CPU 训练。关键配置如下参数含义示例值vocab_size词表大小根据语料自动计算embed_dim嵌入维度128n_heads注意力头数4n_layersTransformer 层数4ff_dim前馈网络隐藏维度256block_size最大上下文长度32dropout随机失活比例0.14.4 训练逻辑训练循环和之前学过的普通神经网络训练没有本质区别随机采样一个 batch 的输入和标签前向传播得到 logits 和 loss反向传播计算梯度优化器更新参数周期性在验证集上评估 loss。唯一需要额外处理的是数据是按块采样的每个 batch 的样本来自不同位置因此需要保证 x 和 y 的切分逻辑一致。4.5 生成策略训练完成后生成文本有两种常见采样方式temperature 控制temperature 越低概率分布越尖锐输出越保守temperature 越高输出越随机。top-k 采样只保留概率最高的 k 个候选 token避免采样到极低概率的不合理 token。5. 完整示例代码与实现下面给出完整可运行的实现。代码基于 PyTorch模型结构是迷你版 GPT。5.1 data.py数据读取与 tokenizer# 文件路径mini_llm/data.py import torch # ---------- 1. 准备语料 ---------- # 这里使用一段经典文本作为演示实际使用时可替换为任意 .txt 文件 text_data 道可道非常道名可名非常名无名天地之始有名万物之母故常无欲以观其妙常有欲以观其徼此两者同出而异名同谓之玄玄之又玄众妙之门 天下皆知美之为美斯恶已皆知善之为善斯不善已故有无相生难易相成长短相形高下相倾音声相和前后相随 是以圣人之治虚其心实其腹弱其志强其骨常使民无知无欲使夫智者不敢为也为无为则无不治 道冲而用之或不盈渊兮似万物之宗挫其锐解其纷和其光同其尘湛兮似或存 天地不仁以万物为刍狗圣人不仁以百姓为刍狗天地之间其犹橐龠乎虚而不屈动而愈出 多言数穷不如守中谷神不死是谓玄牝玄牝之门是谓天地根绵绵若存用之不勤 天长地久天地所以能长且久者以其不自生故能长生是以圣人后其身而身先外其身而身存 上善若水水善利万物而不争处众人之所恶故几于道 居善地心善渊与善仁言善信政善治事善能动善时 夫唯不争故无尤持而盈之不如其已揣而锐之不可长保 富贵而骄自遗其咎功遂身退天之道 # ---------- 2. 字符级 tokenizer ---------- chars sorted(list(set(text_data))) vocab_size len(chars) print(f词表大小: {vocab_size}) stoi {ch: i for i, ch in enumerate(chars)} itos {i: ch for i, ch in enumerate(chars)} def encode(text: str): 将字符串转换为整数索引列表 return [stoi[c] for c in text] def decode(indices): 将整数索引列表转换为字符串 return .join([itos[i] for i in indices]) # ---------- 3. 构造训练数据 ---------- block_size 32 batch_size 64 data encode(text_data) n int(0.9 * len(data)) train_data data[:n] val_data data[n:] def get_batch(split: str, devicecpu): 随机采样一个 batch 的输入和标签 source train_data if split train else val_data ix torch.randint(len(source) - block_size - 1, (batch_size,)) x torch.stack([torch.tensor(source[i:i block_size]) for i in ix]) y torch.stack([torch.tensor(source[i 1:i block_size 1]) for i in ix]) return x.to(device), y.to(device)5.2 model.pyMiniGPT 模型# 文件路径mini_llm/model.py import math import torch import torch.nn as nn import torch.nn.functional as F class SelfAttention(nn.Module): 单层多头自注意力带因果掩码 def __init__(self, embed_dim, n_heads, dropout0.1): super().__init__() assert embed_dim % n_heads 0 self.n_heads n_heads self.head_dim embed_dim // n_heads self.qkv nn.Linear(embed_dim, 3 * embed_dim) self.proj nn.Linear(embed_dim, embed_dim) self.dropout nn.Dropout(dropout) self.scale self.head_dim ** -0.5 def forward(self, x): B, T, C x.shape qkv self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim) qkv qkv.permute(2, 0, 3, 1, 4) # (3, B, H, T, D) q, k, v qkv[0], qkv[1], qkv[2] attn (q k.transpose(-2, -1)) * self.scale # 因果掩码只关注当前及之前的 token mask torch.triu(torch.ones(T, T, devicex.device), diagonal1).bool() attn attn.masked_fill(mask.view(1, 1, T, T), float(-inf)) attn F.softmax(attn, dim-1) attn self.dropout(attn) y attn v # (B, H, T, D) y y.transpose(1, 2).reshape(B, T, C) return self.proj(y) class TransformerBlock(nn.Module): Transformer 解码器块LayerNorm Attention LayerNorm FFN def __init__(self, embed_dim, n_heads, ff_dim, dropout0.1): super().__init__() self.ln1 nn.LayerNorm(embed_dim) self.attn SelfAttention(embed_dim, n_heads, dropout) self.ln2 nn.LayerNorm(embed_dim) self.ffn nn.Sequential( nn.Linear(embed_dim, ff_dim), nn.GELU(), nn.Linear(ff_dim, embed_dim), nn.Dropout(dropout), ) def forward(self, x): x x self.attn(self.ln1(x)) x x self.ffn(self.ln2(x)) return x class MiniGPT(nn.Module): 迷你版 GPT 模型 def __init__(self, vocab_size, embed_dim128, n_heads4, n_layers4, ff_dim256, block_size32, dropout0.1): super().__init__() self.token_embedding nn.Embedding(vocab_size, embed_dim) self.position_embedding nn.Embedding(block_size, embed_dim) self.blocks nn.Sequential(*[ TransformerBlock(embed_dim, n_heads, ff_dim, dropout) for _ in range(n_layers) ]) self.ln_final nn.LayerNorm(embed_dim) self.head nn.Linear(embed_dim, 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, 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, std0.02) def forward(self, idx, targetsNone): B, T idx.shape assert T self.block_size tok_emb self.token_embedding(idx) # (B, T, embed_dim) pos_emb self.position_embedding(torch.arange(T, deviceidx.device)) x tok_emb pos_emb x self.blocks(x) x self.ln_final(x) logits self.head(x) # (B, T, vocab_size) loss None if targets is not None: B, T, C logits.shape loss F.cross_entropy(logits.view(B * T, C), targets.view(B * T)) return logits, loss torch.no_grad() def generate(self, idx, max_new_tokens, temperature1.0, top_kNone): 自回归生成文本 for _ in range(max_new_tokens): idx_cond idx[:, -self.block_size:] logits, _ self(idx_cond) logits logits[:, -1, :] / temperature if top_k is not None: v, _ torch.topk(logits, min(top_k, logits.size(-1))) logits[logits v[:, [-1]]] float(-inf) probs F.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat((idx, idx_next), dim1) return idx5.3 train.py训练脚本# 文件路径mini_llm/train.py import torch import torch.nn as nn from data import get_batch, vocab_size, block_size, batch_size from model import MiniGPT device cuda if torch.cuda.is_available() else cpu print(f使用设备: {device}) model MiniGPT( vocab_sizevocab_size, embed_dim128, n_heads4, n_layers4, ff_dim256, block_sizeblock_size, dropout0.1, ).to(device) optimizer torch.optim.AdamW(model.parameters(), lr3e-4) max_steps 5000 eval_interval 200 eval_iters 50 torch.no_grad() def estimate_loss(): 在训练集和验证集上评估平均损失 out {} model.eval() for split in [train, val]: losses [] for _ in range(eval_iters): x, y get_batch(split, device) _, loss model(x, y) losses.append(loss.item()) out[split] sum(losses) / len(losses) model.train() return out for step in range(max_steps): x, y get_batch(train, device) _, loss model(x, y) optimizer.zero_grad() loss.backward() optimizer.step() if step % eval_interval 0: losses estimate_loss() print(fstep {step:5d} | train loss {losses[train]:.4f} | val loss {losses[val]:.4f}) torch.save(model.state_dict(), output/mini_gpt.pth) print(训练完成模型已保存到 output/mini_gpt.pth)5.4 generate.py文本生成脚本# 文件路径mini_llm/generate.py import torch from data import encode, decode, vocab_size, block_size from model import MiniGPT device cuda if torch.cuda.is_available() else cpu model MiniGPT( vocab_sizevocab_size, embed_dim128, n_heads4, n_layers4, ff_dim256, block_sizeblock_size, ).to(device) model.load_state_dict(torch.load(output/mini_gpt.pth, map_locationdevice)) model.eval() prompt 道可道 idx torch.tensor([encode(prompt)], devicedevice) generated model.generate( idx, max_new_tokens100, temperature0.8, top_k20, ) print(decode(generated[0].tolist()))5.5 关键代码逻辑说明这里有几个重点初学者容易忽略第一get_batch中x和y的关系是“右移一位”。y[i] x[i1]这样模型输入的每个位置都对应“下一个 token”作为预测目标。第二注意力掩码必须显式加到注意力分数上。如果不加掩码模型会提前看到未来 token训练损失会异常低但无法正常生成文本。第三初始化权重的std0.02是 GPT 论文中的常见做法。对整个模型使用标准差为 0.02 的正态分布初始化可以让训练初期更稳定。第四生成时idx_cond idx[:, -self.block_size:]控制只取最后 block_size 个 token防止序列过长超过位置编码长度。6. 运行结果与效果验证6.1 训练过程观察在 CPU 上运行上述脚本训练 5000 步使用 128 维嵌入、4 层 Transformer通常需要几分钟到十几分钟取决于机器性能。训练日志大致如下使用设备: cpu step 0 | train loss 4.6821 | val loss 4.6734 step 200 | train loss 4.0122 | val loss 4.0880 step 400 | train loss 3.2853 | val loss 3.4260 step 600 | train loss 2.7415 | val loss 2.9921 step 800 | train loss 2.3965 | val loss 2.5687 step 1000 | train loss 2.1830 | val loss 2.3345 step 2000 | train loss 1.8016 | val loss 1.9385 step 3000 | train loss 1.6142 | val loss 1.7236 step 4000 | train loss 1.5240 | val loss 1.6118 step 5000 | train loss 1.4568 | val loss 1.5402注意由于是随机采样训练样本且语料很小每次运行会有差异。关键判断标准是train loss 和 val loss 都在下降说明模型在学习val loss 没有大幅上升说明没有明显过拟合如果 train loss 下降而 val loss 上升则是过拟合信号。6.2 生成效果验证训练完成后运行python generate.py输出类似道可道非常道名可名非常名无名天地之始有名万物之母故常无欲以观其妙常有欲以观其徼此两者同出而异名同谓之玄玄之又玄众妙之门天下皆知美之为美斯恶已皆知由于语料只有几千个字符模型本质上是在“背诵”语料的统计规律。但如果使用更大的真实语料训练模型会产生更多样的表达。6.3 如何判断生成质量判断一个小语言模型是否学会了语言规律可以看几个维度单字/词频是否符合语料分布是否出现训练语料中不存在的 n-gram 组合上下文相关性模型生成的后续内容是否与 prompt 主题一致重复率是否过高如果一直输出“的的的的”说明温度设置偏高或模型过拟合于高频 token。7. 常见问题与排查方法问题现象可能原因排查方式解决方案训练时 loss 为 NaN学习率过大数据出现异常 token注意力分数溢出检查学习率、检查输入序列是否包含未登录字符调低学习率检查 tokenizer 映射完整性train loss 下降但 val loss 上升模型过拟合对比两个 loss 的差距增加数据量增加 dropout减小模型规模生成的文本全是重复字符temperature 过高或过低模型容量不足尝试不同 temperature 值将 temperature 设为 0.6-1.0top_k 设为 10-50输入中文报错“KeyError”tokenizer 词表不包含该字符打印 stoi 字典检查重新统计语料字符保证 encode 时使用同一份映射GPU 显存不足batch_size 或 block_size 过大模型嵌入维度太高查看 GPU 显存使用情况减小 batch_size 或 block_size生成结果与 prompt 无关模型未充分训练上下文长度太短观察训练 loss 是否还在下降增加训练步数或换用更大的语料只有第一个 token 是 prompt后面全部停止模型输出的 token id 超出 decode 范围检查词表映射确认 itos 覆盖所有 token id8. 最佳实践与工程建议8.1 数据质量和规模优先于模型技巧从上面的代码可以看出小模型也能学会语料中的统计规律但生成效果好坏很大程度取决于语料质量。如果你想让模型生成更有意义的中文建议使用更大、更干净、更多样化的语料而不是急于调大模型。8.2 训练过程要保留检查点建议每训练若干步保存一次模型权重而不是只在最后保存。训练中断时可以从最近的检查点继续节省时间。检查点文件建议包含模型参数、优化器状态、当前 step、超参数配置。8.3 日志记录是调试的基础实际工程中训练循环里的print并不够用。建议使用tensorboard或wandb记录 loss、学习率、梯度范数等指标。梯度范数突然变大往往意味着训练不稳定。8.4 通过配置化控制超参数不要把超参数写死在代码里。可以使用config.py或 YAML 配置文件管理# 文件路径mini_llm/config.yaml model: embed_dim: 128 n_heads: 4 n_layers: 4 ff_dim: 256 block_size: 32 dropout: 0.1 train: batch_size: 64 learning_rate: 0.0003 max_steps: 5000 eval_interval: 200 save_path: output/mini_gpt.pth8.5 生成时要设置随机种子调试过程中如果每次生成结果都不稳定很难判断改动是否有效。在训练和生成前设置随机种子import torch torch.manual_seed(42)8.6 从字符级 tokenizer 到子词 tokenizer本文使用字符级 tokenizer是为了教学直观。实际大模型项目中字符级会导致序列过长并且无法有效表示词义。进阶方向是学习 BPE、WordPiece 或 SentencePiece。理解本文思路后建议尽快切换到中文 BPE tokenizer这会明显提升模型效果。8.7 设备选择与训练效率CPU 上训练小模型完全可行但如果你想尝试更大的模型务必使用 GPU。注意设置pin_memoryTrue加速数据加载torch.no_grad()包裹验证和生成过程混合精度训练AMP可以显著减少显存占用。8.8 警惕训练集和验证集泄漏切分数据时文本是按字符随机切分不是按“句子”切分。这样可以避免验证集出现训练集中完全相同的片段但掩盖了模型有效泛化能力的问题。更严谨的做法是按段落或文档切分不过这会让数据量变少本文示例是从教学角度做了简化。9. 总结与下一步学习方向这一圈走下来你已经亲手完成了一个最小闭环从语料准备、tokenizer 构建、Transformer 解码器实现到训练、验证和文本生成。这个过程对理解 LLM 的价值远大于简单调用现成模型。接下来可以从这几个方向继续深入第一把 tokenizer 换成 BPE 或 SentencePiece并用真实中文语料训练一个大一点的模型观察生成效果的变化。第二研究位置编码的方案。本文用的是可学习绝对位置编码现代模型还常用 RoPE、ALiBi 等相对位置编码后者对长文本更好。第三理解指令微调和 RLHF。从预训练模型到 ChatGPT中间还有指令微调和人类反馈对齐的过程。先掌握基础架构再去理解这些高级技术会轻松很多。第四把模型部署到实际应用场景。训练一个模型不难难的是如何平衡效果、速度和成本。浮点数精度、量化、模型蒸馏、推理加速都是实际工程中必须要面对的问题。最后提醒一句如果你是第一次接触大模型不要追求一步到位复现 GPT-3。建议先完整运行本文代码把每个模块的输入输出打印出来仔细看 attention 矩阵的形状变化。等你能看懂每一行代码在做什么你对 LLM 的理解就已经超过绝大多数只会调 API 的同学了。
返回列表