从零构建小型语言模型：PyTorch实现GPT的完整工程实践-北京尧图网络

如果有人问我真正理解 AI 工程的最短路径是什么我的答案一定会是自己动手从零开始构建一个模型哪怕它很小。过去两周我给自己定了一个叫 ai-engineering-from-scratch 的目标——不依赖任何第三方预训练模型从一行 Python 代码开始硬是搭出了一个可训练、可推理、参数规模约 1 亿的小型语言模型跑通了数据清洗、分词器、模型架构、预训练、采样生成、评估的完整链路。这篇记录不是什么保姆级教程更像是我在“从零构建 AI 工程”这条路上真实踩坑的经验总结包含所有我验证过的方案和能直接落地的细节。它适合已经会用 PyTorch、但总觉得大模型像黑盒想亲手拆开看一遍的工程师和学习者。1. 为什么坚持“从零构建”而不是直接调库1.1 调库很快但天花板也来得很快用 Hugging Face 的 transformers 加载一个 GPT-2 只需要三行代码生成一段文本也只差一个 for 循环。可我第一次意识到自己并没有真正理解大模型是在修改 attention 的时候我想给模型加一个自定义的 bias 项结果对着几百行封装好的源码找了半天最后还是决定不看代码自己写一个简化的模型。从零构建最大的价值就在这里——它逼着你面对所有原本被库隐藏的细节张量怎么排布、mask 怎么加、梯度怎么流动、loss 怎么算。这些东西不是背下来就够的它们是 AI 工程的地基。如果有人抱着“从零构建全都自己手写”的想法我建议立刻放下这个执念。框架还是要用的而且要熟练地用。PyTorch 的自动求导帮你省掉了反向传播的体力活但模型结构、训练策略、数据管道这些真正的工程问题框架不会替你做。我这次从零构建本质上是用 PyTorch 的算子去实现 GPT 的每一个部件而不是调一个封装好的GPT2LMHeadModel完事。1.2 我的目标一个 1 亿参数的小型 GPT从零构建不等于要把目标定成复刻 LLaMA-70B那既不现实也没必要。我的目标是一个约 1 亿参数的 decoder-only 模型具体配置如下表配置项数值层数12隐藏维度768注意力头数12上下文长度512词表大小16384总参数量约 1.02 亿选这个规模有两个原因。第一单张 RTX 4090 就能跑一天内能看到 loss 明显下降同时保留真实的预训练体验。第二训练数据用了大约 1GB 纯文本数量和规模匹配不至于小到学不出语言规律。《Build a Large Language Model from Scratch》这本书里的路线也是类似思路用小型数据集、小模型走通大模型的全流程优先级是先有完整链路再谈规模。1.3 技术选型框架、硬件和数据框架我用的是 PyTorch没有用 DeepSpeed 之类的重型工具因为 1 亿参数单卡还不需要。我的环境是 Python 3.10、PyTorch 2.x、CUDA 12.1一张 24GB 的 RTX 4090另外用 CPU 做了两天的数据处理。这里想提醒一句如果自己抓网页一定要注意版权和合规尽量使用公开数据集。我最终用的是 WikiText-103 的约 1 亿 token 子集混上 TinyStories 的约 2 亿 token两者混合作为预训练语料。为什么混 TinyStories因为它的句子结构简单、故事性强能让小模型更快学到基础语义不会一上来就被 WikiText 里密集的长句折磨。2. 从零构建 LLM 的核心流程拆解2.1 数据工程语料清洗与格式统一别看模型很酷真正花时间最多的是数据。原始数据里的 HTML 标签、控制字符、重复片段都会污染训练。我的处理流程是先用 BeautifulSoup 去掉 HTML再用正则过滤特殊符号按句子切分后去掉长度小于 20 个字符或大于 500 个字符的片段最后做 minhash 去重。这里有个非常值得分享的细节去重时不要对整个文档做 hash而是对滑窗片段做 minhash。原因是网上爬来的文本经常是“几乎一样但略有不同”整体 hash 根本发现不了重复而滑动窗口采样多个局部片段再对这些片段做相似度估计才能有效去掉复制粘贴产生的垃圾数据。清洗完成后我把所有文本统一成一行一个段落并在段落间插入|endoftext|作为分隔符方便 tokenizer 和训练循环识别边界。2.2 Tokenizer从字符级到 BPE最早的版本我图省事直接用字符级 tokenizer结果效果很差——序列长度爆炸、词汇表容量低模型根本学不了语义。后来换成 BPEByte Pair Encoding才真正解决了问题。BPE 的核心思路是把 byte 序列逐步合并成高频子词词表大小我设为 16384覆盖常见的英文单词和子词。推荐用 OpenAI 开源的 tiktoken 库来训练和编码它是 GPT 系列同款算法速度也快。需要特别注意tokenizer 的词汇表必须在训练前冻结后面不能中途增删否则嵌入层维度变化之前的 checkpoint 全废。我就在这上面吃过亏第一次训练到一半觉得词表不够用重新训练了一次 tokenizer结果所有历史权重都不能加载只能从头重训。现在回想起来这就是标准的工程陷阱——以为小改动无所谓实际破坏性极大。2.3 模型架构一个标准的 decoder-only Transformer我的模型基本是 GPT-2 的简化实现token embedding position embedding 12 层 transformer block每个 block 包含 causal self-attention、feed-forward network、LayerNorm 和残差连接。这里最重要的设计决定是使用 pre-LayerNorm 还是 post-LayerNorm。我试下来pre-LN在 attention 和 FFN 之前做 LayerNorm训练更稳定尤其当模型加深时post-LN 训练初期容易震荡需要更小的学习率。另一个容易忽略的点是 attention 的因果掩码。自回归模型只能看到左侧 token所以 QK 矩阵计算后要加上一个上三角全 -inf 的 mask再 softmax。写成代码只有几行mask torch.triu(torch.ones(seq_len, seq_len), diagonal1).bool() attn_scores scores.masked_fill(mask, float(-inf))但漏掉 mask整个模型就变成双向的了生成时行为会非常奇怪。我在做消融测试时故意去掉过一次 maskloss 降得比有 mask 还要快因为模型“作弊”看到了未来 token但真正生成时什么都输出不对。所以遇到“训练得好、生成很差”的情况第一个该查的就是这里。2.4 训练循环自回归语言建模训练目标是用前一个 token 预测下一个 token。具体实现是把序列 x 作为输入标签就是 x 整体右移一位后的结果最后一个 token 的预测不参与 loss。损失函数用 CrossEntropyLoss需要设置ignore_index-100来忽略 padding 部分。优化器我用 AdamWbeta10.9beta20.95eps1e-8weight_decay0.1学习率峰值 3e-4采用 warmup 1000 步后余弦衰减到 1e-5。batch size 我设成 32序列长度 512所以每个 batch 有 16384 个 token。梯度裁剪设到 1.0防止 loss spike。完整训练配置如下超参数设置优化器AdamW峰值学习率3e-4Warmup 步数1000学习率调度Cosine decay 到 1e-5Batch size328 × 4 梯度累积序列长度512梯度裁剪max_norm1.0训练步数50000这套配置不是凭空拍脑袋出来的而是参考了 GPT-2 的训练策略再针对小模型做了缩放。如果你自己调整核心记住一点学习率是训练稳定性的第一杠杆warmup 不要省特别是从零开始训随机初始化的模型前几百步一崩后面全崩。3. 核心环节实现与踩坑实录3.1 BPE Tokenizer 实现中的两个坑第一个坑是编码 / 解码不一致。我一开始自己写 BPE训练时在某一端用了正则预处理另一端忘了导致词表 ID 对不上训练 loss 始终不降。后来统一用 tiktoken所有文本先按字节编码再走 BPE避免了很多问题。如果你要自己实现强烈建议先在几条短文本上做“编码后解码是否还原”的单元测试。第二个坑是序列截断逻辑。512 上下文长度意味着超过 512 的文档要被切成长度 512 的多个块但如果切分时不考虑段落边界很容易把所有数据切成“半句话”导致训练语料异常。我的做法是先用|endoftext|分隔文档再按滑动窗口切块并且保证每个块从完整段落开始。同时切分时保留一个小 overlap避免把完整语义切断。别小看这种细节它直接影响模型对句子边界的感知。3.2 注意力机制从朴素实现到 Flash Attention朴素点积注意力实现很简单QK.T / sqrt(d_k)加上因果 masksoftmax再与 V 相乘。单卡上 1 亿模型、序列 512 完全够用但训练速度大约每秒 1200 token太慢。后来我换成 Flash Attention通过 PyTorch 2.x 的scaled_dot_product_attentionAPI 一行切换训练速度直接提升到每秒 2200 token几乎翻倍。原因是 Flash Attention 融合了掩码、注意力计算和显存访问减少了 HBM 读写。对大部分注意力计算都有效建议尽早引入不要等到最后再改。如果你不想引入额外依赖其实 PyTorch 2.x 内置的F.scaled_dot_product_attention已经能自动选择融合实现我在最终版本里用的就是它。3.3 训练稳定性loss 不降和 loss spike训练中最常见的问题是 loss 在某个 step 突然跳高甚至变 NaN。我的经验是先看梯度范数如果异常大很可能是数据里混入了超长或异常 token 序列。后来我在训练循环里加了梯度裁剪max_norm1.0问题少了很多。另一个有效手段是检查 loss 的平均窗口单步 loss 波动很大是正常的不要看到一次震荡就调参。如果 loss 持续不降我建议先跑一个小规模的 overfit 测试——在几千条样本上训练如果模型能过拟合说明数据和模型没问题问题出在训练配置或数据管线如果不能过拟合那八成是代码有 bug。这个排查方法非常值钱比调参管用一百倍。我当时发现 loss 卡在 6.0 附近不动一测 overfit 发现 loss 也降不下去最后定位到 label 移位写错了白白浪费了一下午。3.4 生成与采样temperature、top-k、top-p 怎么配训练完成后我最喜欢做的就是写一个 generate 函数从|endoftext|开始逐 token 自回归采样。这里最核心的是采样策略。我试过纯贪心解码结果文本高度重复改用 temperature0.7 又偏保守最后发现 temperature0.8、top-k40、top-p0.9 的组合最自然。原因是纯贪心只取最大概率 token很容易陷入循环top-k 限制候选集合top-p 动态保留累计概率达到 0.9 的 token两者配合可以兼顾多样性和质量。另外加上repetition_penalty1.1能明显减少重复短句。我整理了几组参数的实际效果采样设置效果贪心解码重复严重几乎无法阅读temperature0.7流畅但偏保守temperature0.8 top-k40多样性提升偶有重复temperature0.8 top-k40 top-p0.9自然且稳定再加 repetition_penalty1.1重复问题基本消除4. 资源消耗与工程优化4.1 算力预算模型规模、显存和训练时长很多人在从零构建前最关心“我这机器能不能跑”。我帮你算一笔账1 亿参数模型fp32 权重占 400MBAdamW 优化器状态占权重的约 2 倍800MB梯度占 400MB再加上激活值和临时变量24GB 显存单卡绰绰有余。如果只有 8GB 显存可以把模型降到 3000 万参数即 6 层、384 隐藏维度照样能跑通流程。训练时长上我训练了 5 万步总 token 数约 8 亿RTX 4090 上实际用了大约 8 小时。如果只想验证链路训练 5000 步、loss 降到 3.5 左右就够看效果了。下面给一个粗略参考模型规模参数量单卡训练速度建议数据量适合硬件Mini30M4000 token/s0.2B token8GB 显卡Small100M2200 token/s1B token24GB 显卡Base350M800 token/s5B token多卡或更高显存4.2 显存优化混合精度、梯度累积、梯度检查点如果你显存不够我按优先级推荐的方案是混合精度、梯度累积、梯度检查点。混合精度用 PyTorch 的amp.autocast权重和优化器状态保留 bf16 或 fp16显存大概能省 30% 到 40%。梯度累积可以解决 batch size 不够的问题我把微批次 batch size 设为 8累积 4 次等价于 32 的 batch size但显存占用只有四分之一。梯度检查点gradient checkpointing是拿计算换显存transformer block 的前向激活值不全部保留反传时重新计算显存能再省一半代价是训练时间增加约 20%。三招叠加后普通消费级显卡也能练千万参数级别的模型。不过我要提醒一句不要一上来就把所有优化叠满先按默认配置跑通再逐步降低显存。否则你会在“代码正确性”和“性能优化”两个问题之间傻傻分不清楚。4.3 数据加载优化别让 CPU 拖慢 GPU早期我的训练速度总上不去后来发现瓶颈不在 GPU而在 dataloader每次都是重新读原始文本再 tokenizeGPU 有一半时间在空转。解决办法是提前用 tokenizer 把所有数据编码成 token ID 数组保存成二进制格式训练时直接用随机索引读取。之后把 DataLoader 的num_workers设为 4prefetch_factor设为 8并用 pinned memory训练吞吐又提升了 30%。这一步听起来很基础但很多从零开始项目都会忽略。如果你对速度敏感我建议在训练循环里打印一个“每秒 token 数”的指标如果这个数字低于你预期的一半先去看数据加载而不是立刻怀疑模型架构。5. 常见问题与排查技巧实录5.1 训练 Loss 不下降先怀疑数据管线如果你遇到 loss 一直不下降不要先调学习率。我的经验是先检查 tokenizer 是否正常工作把输入的 batch 解码成文本看看是不是乱码、ID 和文本是否对应。再检查标签自回归标签应该是输入右移一位很多新手在这里写错。最后可以做 overfit 测试用 500 条样本训练 20 步如果 loss 有明显下降说明模型有学习能力如果完全不动大概率是计算图断了比如 label 和 logits 对不上。5.2 模型生成内容全是重复生成重复是自回归模型最常见的问题。原因通常是训练不充分、模型太小、或者采样策略不当。如果你的模型只有 1 亿参数而训练数据不够多它很容易学会高频的“安全词”。我建议先用更大的 top-k 和 top-p 拉高多样性然后尝试 temperature 从 0.6 到 1.2 进行扫描找到最小困惑度但文本不重复的点。如果还不行可以在损失函数里加入重复惩罚或者解码时加repetition_penalty。另外降低学习率或增加训练步数也能缓解重复因为模型能学到更丰富的上下文分布。我在训练第 3 万步时生成的文本明显比第 1 万步重复少说明数据学习到位后解码侧的压力会小很多。5.3 训练到一半显存不足训练中显存不足的常见原因有三个输入序列长度太长、batch size 太大、或激活值累积过高。如果你不想降低效果可以先试梯度检查点再不行就减小 batch size并用梯度累积补回来也可以把序列长度从 512 降到 256但要注意这会影响模型长距离依赖能力。另一个容易忽略的点是 PyTorch 的内存碎片化内存碎片会导致明明总显存够用、但单个 tensor 分配失败。遇到这种情况先尝试用torch.cuda.empty_cache()清一下缓存如果还不行就重启训练进程让显存重新整理。别迷信动态显存分配长时间训练以后碎片化几乎不可避免。5.4 Loss 降了但生成的文本依旧没有意义这种情况通常说明模型还在学习词法、句法的早期阶段还没有学到完整语义。我的建议是不要只看 loss还要看模型的生成样本和测试集困惑度。我从 8 小时训练中记录了不同阶段的生成结果前 2 小时基本是乱码第 4 小时开始出现单词第 6 小时有了短句最后 2 小时才对上下文一致性有所改善。如果你时间有限至少训练到生成文本通顺再停。另外可以加一个简单评估任务比如让模型补全一个常识句子看结果是否合理。我经常用的自测句子是“The capital of France is”然后看模型能不能接出“Paris”。虽然这有点刻板但作为项目门禁很有效。最后再说一点个人体会从零构建一趟最大的收获不是得到了一个能生成文本的模型而是脑子里有了整个 AI 系统的地图。以前用 API 时我只关心 prompt 和温度参数现在再想起来会不自觉地去推它是怎么分词、怎么过 attention、怎么采样出来的。这种“底层感”一旦建立起来后续无论是调库还是做微调思考方式都会完全不同。如果你也想尝试我的建议是别贪大先选 3000 万参数左右的迷你模型跑通全链再去加数据、加层数。先让它会说话再让它说得聪明。

从零构建小型语言模型：PyTorch实现GPT的完整工程实践

相关推荐

后见之明偏差：从复盘陷阱到决策引擎的认知升级

ILI2511触控芯片规格书精读：从引脚定义到上电时序的避坑指南

新能源线全功能处理机选购指南与关键技术解析

更多资讯

C++求整数数位之和：从循环取余到边界处理全解析

SSM+Vue民宿管理系统设计与实现指南

Buzz概念解析：从网络热度到技术传播机制

现代HTML API：9个替代第三方库的原生方案

现代HTML API：浏览器原生接口的性能优势与实践

Linux TCP网络编程：从基础到实战

今日更新

淘宝商品详情API开发实战与性能优化指南

Unity Package Manager刷新错误排查与解决方案

Java record 规范构造器能否被反序列化或反射绕过

本周精选

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

本月热门

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

读完文章，想聊聊您的网站？