ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

从零构建大模型:手写Transformer训练到推理的完整工程实践

从零构建大模型:手写Transformer训练到推理的完整工程实践 把“从零做AI工程”这件事亲手跑通一遍是我最近一年做得最值的一件事。这个项目名字叫 ai-engineering-from-scratch核心不是调现成接口也不是套开源权重而是从数据、分词、模型结构、训练、推理这条完整链路自己动手搭一遍最后跑出一个能真的吐出连贯文本的小模型。参考像《Build a Large Language Model (From Scratch)》这类路线的人很多但真正动手把每个环节踩一遍、把每个细节问一遍为什么的其实没那么多。这篇文章就把我自己的完整过程、踩过的坑、以及最后怎么延伸到“推理模型reasoning model”这条线的经验一次性写清楚。1. 为什么我坚持“从零构建”而不是直接调API1.1 这个项目到底要解决什么问题现在学AI工程的人分两类。一类是“调用派”会写两三行代码把大模型API接好写个Prompt模板就当做了个项目另一类是“理解派”希望能深入模型内部知道一次前向传播里数据长什么样、损失函数在优化什么、为什么显存会爆、为什么同一个种子下生成结果还是飘。这个项目就是给“理解派”准备的同时也适合那些正处在调用派和理解派分叉口上的人。我自己的体会是单纯调API带来的能力增长是线性的你学一百个Prompt技巧遇到模型内部出问题照样束手无策。但从零构建一次哪怕只是几十层的小Transformer你也会把“大模型”这三个字彻底祛魅。它不再是玄学而是一套可以用公式、配置、显存字节数精确描述的系统。1.2 从零构建的收益边界必须诚实地说从零构建并不是人人都必需的。如果你的目标是三个月内做出一个产品Demo那直接基于开源模型微调确实是效率最高的路径。但如果你是想建立不可迁移的工程直觉想在以后排查问题时知道该往哪个方向检查那从零训练一个模型是你无论如何都绕不开的一课。我做这个项目时给自己定的边界很明确不追求模型多大、多强只追求“完整可控”。模型选型是GPT风格的decoder-only架构参数量控制在1亿左右这个级别用一张RTX 4090就能跑训练时间按天算而不是按周算。项目目标是让模型能生成语法正确、话题连贯的短文并且我能准确说出每一个超参数为什么这么设。1.3 适合什么阶段的人来跟做这个项目的实践难度我评估为中等偏上。你需要的基本功是Python能熟练写类和装饰器至少看得懂torch.nn.Module的继承结构。对神经网络基本概念有了解知道反向传播、梯度下降、过拟合这些词。最好有GPU显存8GB以上的NVIDIA卡都行没有的话用云GPU或Kaggle免费环境也能跑但训练速度会慢。如果你只会调用接口没关系这个项目正好逼你把基础补上。如果你已经读过不少模型论文但没动手训练过那这个项目就是你的动手验真环节。2. 从零搭建大模型的架构选型与关键决策2.1 技术路线为什么选GPT风格的decoder-only在动手前要做一个最关键的决策选哪种架构现在主流的路线无非三类纯编码器BERT风格、编码解码器T5风格、纯解码器GPT风格。三天两头出现的“大语言模型”概念绝大多数指的就是GPT风格的decoder-only架构。它的特点是只保留Transformer的Decoder部分用因果注意力Causal Attention做自回归生成也就是模型在生成每个token时只能看到它前面的内容。我选decoder-only不单是因为它主流而是因为自回归训练和推理的接口非常统一。训练时输入是一整段文本预测目标是同样的文本往右平移一位推理时只需要拿着上一个token的输出继续接到输入后面。这种“一个接口搞定训练和推理”的特性对从零构建的人来说友好度极高调试路径短。另一个理由是数据使用效率。编码器模型的训练通常需要设计遮罩语言模型masked language modeling之类的目标而decoder-only直接用“文本下一位预测”就能训练数据准备简单到只需要把语料切片拼好。2.2 初始模型配置与参数量估算我选的基础配置参考了GPT-2 small的规模但为了适配更长上下文的训练做了微调。具体配置如下参数名称配置值说明n_layer12Transformer Block层数n_embd768隐藏层维度n_head12注意力头数量每头64维ctx_len512最大上下文长度单位tokenvocab_size50257词表大小与GPT-2的BPE一致MLP ratio4前馈网络从768升维到3072这个配置下的参数量可以直接手算出来。词嵌入层是vocab_size * n_embd约50257乘768等于3860万每个Transformer Block包含注意力四组线性层Q、K、V、投影共4个n_embd * n_embd加MLP两个线性层升维和降维共2 * 4 * n_embd * n_embd再加两个LayerNorm各2 * n_embd单层约998万参数12层就接近1.2亿加上最终的输出投影层和位置嵌入总参数量大概1.27亿。这个手算很重要。因为只有你亲手算过一遍才知道网上那些“1B模型需要多少显存”的估算值是怎么来的。以bf16混合精度为例1.27亿参数每个占2字节光权重就要254MBAdamW优化器额外保存一阶和二阶动量又各需要254MB前向激活值在训练时又要按批量占用几倍的显存。知道这些你选显存卡的时候心里就有谱了。2.3 任务拆解数据、分词、训练、评估四条线一个大模型项目看着复杂拆开其实就是四条并行线数据线收集、清洗、切分语料做成可重复读取的批量数据。分词线训练或加载一个tokenizer把文本变成整数id序列并处理特殊token。模型线实现模型结构、初始化、前向传播、损失函数。训练评估线准备优化器、学习率调度、训练循环、采样生成、评估指标。四条线里最容易低估的是第2条。很多人一开始就直接用现成的tiktoken当BPE分词器这确实省事。但如果你想让模型跑在中文语料上直接用GPT-2的词表会导致中文被拆成零碎的单字或字节序列训练效率很低。我最后采用了在这条路线里比较折中的方式用开源的中文BPE词表作为底座再做一点词表压缩和特殊token补齐。后面第3章我会详细展开这个决策。3. 实操全流程数据、训练、推理3.1 环境与硬件准备先讲环境配置这部分是最容易被忽略但最影响心情的。我的环境是操作系统Ubuntu 22.04Windows跑WSL2也可以。Python3.11PyTorch 2.4以上越新越好因为很多编译器优化需要新版本。GPU单张RTX 4090 24GB。如果你的卡显存小于16GB建议把ctx_len降到384或者减少batch_size。CUDA12.4以上配合cuDNN。安装依赖时建议直接用conda或venv建一个独立环境避免污染系统Python。核心依赖就几个torch、numpy、tiktoken或tokenizers、sentencepiece、datasets、tqdm、wandb本地快速看loss不需要也可以去掉。硬件不足是很多人从零开始的第一个坎。我的建议是哪怕只有一张8GB显存的卡也一定要开始做。把模型缩小到0.5亿参数把ctx_len降到256照样能把流程跑通。跑通一遍之后你积累的工程经验比任何参数刷榜都有价值。3.2 数据准备与tokenizer选择先说我用的语料。英文我用了一部分The Pile的子集中文用了清洗后的开源中文语料。如果你不想折腾多语言混训我建议新手起步只选一种语言效果更可控排查问题更容易。我一开始就是中英混训结果loss曲线出现周期性波动排查了很久才发现是两种语言数据分布不均匀导致的后来改成按比例打包才稳定。数据处理流程很简单但清洗很耗时。我按下面这几步走按段落切分原始文本过滤太短的碎片少于200字符的直接丢弃。去掉网页标签、无意义符号、重复段落。使用datasets库生成arrow格式的数据集比直接存JSON快得多。将文本按ctx_len切块每个样本是一段连续的512个token序列样本之间可以重叠滑动。分词器的选择是我这个项目里花时间最多的部分之一。如果你用英文tiktoken的cl100k_base够用如果你想做中文生成直接套用就显得笨拙。我的做法是下载一个开源中文BPE词表按词频保留前5万个词再补上|endoftext|、|user|、|assistant|等特殊token。这样模型在生成中文时更自然也不会出现单字崩裂。这里有个容易踩的坑数据切块和tokenizer的词表必须在训练之前做好对齐验证。你说模型用的是5万词表数据却被原来的GPT-2 tokenizer切成了另一种id那模型训练就是白费。我专门写了一个校验脚本每次开始训练前随机抽100个样本检查tokenizer解码后能否大致还原原文能还原才继续。3.3 核心模型代码实现模型结构我用的是经典的GPT风格代码这里不再贴几百行完整代码只说几个最关键的部分和设计理由。第一个关键是因果注意力掩码。所谓“因果”就是每个位置只能看到自己及之前的token不能看到未来。实现时用一个上三角布尔矩阵作为mask在注意力分数上加一个极大的负偏置把所有未来位置的softmax权重压到接近零。import torch import torch.nn as nn class CausalSelfAttention(nn.Module): def __init__(self, n_embd, n_head, ctx_len): super().__init__() assert n_embd % n_head 0 self.n_head n_head self.head_dim n_embd // n_head self.c_attn nn.Linear(n_embd, 3 * n_embd) self.c_proj nn.Linear(n_embd, n_embd) mask torch.triu(torch.ones(ctx_len, ctx_len, dtypetorch.bool), diagonal1) self.register_buffer(mask, mask) def forward(self, x): B, T, C x.size() qkv self.c_attn(x) q, k, v qkv.split(self.n_embd, dim2) 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].view(1, 1, T, T), float(-inf)) att torch.softmax(att, dim-1) y att v y y.transpose(1, 2).contiguous().view(B, T, C) return self.c_proj(y)第二个关键是Transformer Block。标准结构是“注意力 → 残差 → 前馈 → 残差”并在每组前加LayerNorm。最近研究更推荐Pre-LNLayerNorm放在子层之前训练更稳定。class MLP(nn.Module): def __init__(self, n_embd, ratio4): super().__init__() self.fc nn.Linear(n_embd, ratio * n_embd) self.gelu nn.GELU(approximatetanh) self.proj nn.Linear(ratio * n_embd, n_embd) def forward(self, x): return self.proj(self.gelu(self.fc(x))) class TransformerBlock(nn.Module): def __init__(self, n_embd, n_head, ctx_len): super().__init__() self.ln1 nn.LayerNorm(n_embd) self.attn CausalSelfAttention(n_embd, n_head, ctx_len) self.ln2 nn.LayerNorm(n_embd) self.mlp MLP(n_embd) def forward(self, x): x x self.attn(self.ln1(x)) x x self.mlp(self.ln2(x)) return x第三个关键是输出层。模型最后一层是nn.Linear(n_embd, vocab_size)配合交叉熵损失计算预测下一个token的概率。为了让损失计算更快PyTorch建议把损失计算和输出投影融合起来用torch.nn.CrossEntropyLoss作用于模型输出的logits和真实id序列即可不需要单独softmax。我强烈建议写完代码后先做一次“单batch过拟合”小测试拿一小段数据让模型拟合几十步看看loss是否逼近0。如果几十步后loss死活不降那你的模型结构大概率有bug而不是学习率问题。这一步能帮你把模型结构的bug和训练过程的bug分开排查节省大量时间。3.4 训练超参与损失曲线解读训练超参这一环直接决定你能不能在一个可接受的时间范围内看到模型“开口说话”。我用的是业界非常通用的AdamW配置加上cosine学习率衰减。具体数值如下超参数设置值为什么这么设optimizerAdamW不解释大模型事实标准learning_rate3e-4GPT-2 paper给出的经验值过大初期爆loss过小收敛慢betas(0.9, 0.95)第二动量偏大一点让更新更平滑weight_decay0.1对Embedding和LayerNorm不加其余加抑制过拟合warmup_steps1000预热防止初期更新太猛导致不收敛batch_size16个样本每个样本512 token单batch约8192 token混合精度bf16显存减半训练加速约30%max_steps100000全套跑完几天时间我实际在35000步时看到合格效果我训练时最关注两个指标训练loss和验证loss。两者同步下降说明模型在学习训练loss下降但验证loss不动说明过拟合两者都不动先怀疑学习率或数据有问题。实际的loss曲线大致是前500步从10左右快速降到6之后进入缓慢下降期到15000步左右降到3左右。这个阶段模型已经能生成“看起来像样”的文本。不要追求loss降到极小对1亿参数模型来说验证loss到3.2左右就已经是一个能做出可用Demo的水平继续训练性价比不高。还有一个经验loss曲线如果出现突然的尖峰然后恢复大概率是碰到了一小批脏数据比如极端重复段落或者语料里的乱码。这时候不要急着调学习率先把脏数据抠出来。3.5 推理生成与评估训练完成后最爽的一步是让模型生成文本。推理和训练的前向过程几乎一样区别是推理时要自己控制采样策略。最简单的是贪心解码也就是每一步选概率最高的token但这种输出容易重复。我常用的生成参数范围是temperature0.7到0.9温度越低越保守越高越发散。top_k50只保留概率最高的50个token参与采样。top_p0.9按累积概率截断。用PyTorch实现采样就几行代码。关键是理解temperature就是在softmax之前把logits除以温度值top_k和top_p都是过滤掉概率质量过低的候选。我用一个简单的循环每生成一个token就把它拼回输入序列再跑一次前向直到生成指定长度或遇到|endoftext|。评估模型质量最常用的指标是困惑度PerplexityPPL数学上就是交叉熵损失取指数。PPL越低越好我训练的模型验证PPL在25左右这个水平的生成结果已经能稳定输出结构完整的段落。但PPL只反映统计层面的拟合真正评估生成质量还是要靠人看比如生成100条短文人工判断语法正确率、话题相关度和重复率。别迷信数字一定要看真实输出。4. 从零训练常见的问题与排查实录4.1 OOM和训练速度的排查优化显存溢出OOM是所有人都会遇到的第一道坎。我一开始直接按GPT-2 small配置跑batch_size32结果刚跑两步就OOM。第一次遇到OOM别慌按下面的优先级排查降低batch_size用梯度累积补偿。比如batch_size8梯度累积步数为4等效batch就是32显存占用却少很多。开启混合精度bf16或fp16。bf16在支持它的卡上几乎没有精度损失我强烈推荐。打开gradient_checkpointing用时间换显存训练速度会慢约30%但显存占用能少一半。检查有没有忘记把不需要的中间变量释放。写代码时用torch.no_grad包住验证环节的推理避免梯度图意外保留。训练速度慢则要区分是数据瓶颈还是计算瓶颈。我的排查方法是先看GPU利用率如果nvidia-smi显示利用率一直在90%以上说明计算在全力跑瓶颈在模型本身可以试试torch.compile如果利用率低说明数据加载跟不上这时候要增加num_workers、打开pin_memory或者检查是不是在GPU上做了不必要的CPU数据操作。4.2 loss不下降与生成乱码的根源Loss不下降有两种常见情况。第一种是前500步就完全不降大概率是学习率太高导致loss直接飞到无穷或者模型初始化有问题。解决方式是调低学习率到1e-4同时在模型定义时对最后一层投影用较小的标准差初始化。第二种是训练到中途loss突然平台期不再下降这种往往是数据不够多样或者学习率衰减过快把模型锁死在局部最优。我的经验是遇到平台期先别急着加数据把warmup之后的峰值学习率调小一半再观察2000步。生成乱码是我在这个项目里最想吐槽的问题。我踩过的坑有两个tokenizer词表与模型词表不匹配。有时候你加载了别人的预训练embedding但自己的输出层还是随机初始化的模型生成的id解码出来全是凭空的字符。出现这个情况先检查tokenizer.encode(你好)和tokenizer.decode(ids)能否还原再检查模型vocab_size是否与tokenizer长度一致。特殊token没处理好。训练时|endoftext|混进了正常文本导致模型以为段落总在奇怪的地方结束。解决方法是训练时明确把|endoftext|只放在每个文档的末尾其他位置一律屏蔽这个token的loss。4.3 生成重复和质量不稳定的处理模型生成到一半开始车轱辘话来回说是最让初学者头疼的现象。小模型特别容易出现这个问题因为参数少学到的长期依赖模式就弱。我试过几种方法最有用的是降低top_k从100降到50减少尾部随机性。增加重复惩罚。在logits上对已经出现过的token做惩罚衰减简单实现就是logits[token_id] - 1.2。在训练数据里过滤掉高重复率的样本比如用SimHash算法把重复度超过0.8的段落直接删除。但我要说一句实话对1亿参数的模型重复问题很难完全消除这本质上是模型容量和训练数据多样性不足的双重结果。如果你想让生成质量上一个台阶最终还是要增加参数规模或数据量这才是治本之道。5. 下一步从“能说话”到“能推理”的延伸5.1 为什么推理模型成了“从零构建”的新话题在我做完基础的GPT训练后正好赶上“推理模型”这个方向突然火爆起来。相关热词里总能看到“build a reasoning model from scratch”很多人想复现那种能一步步推导、展示长思维链的模型。其实推理模型并不是另起炉灶的新架构它仍然建立在自回归语言模型之上关键在于训练数据的形态从“问题-答案”变成了“问题-详细推理过程-答案”再通过强化学习进一步强化“遇到难题愿意多想几步”的行为。从零构建推理模型的难度比单纯训练LLM高一个维度难点主要在数据侧。你可以自己拼接开源数据集构造“问题-中间步骤-最终答案”的思维链Chain of ThoughtCoT样本。我在实际尝试中发现哪怕只构造几千条高质量的长思维链样本让上一章的小模型去做SFT它也能表现出一定程度的“多步思考”比如计算数学应用题时会先写“设未知数x”再逐步推导。效果虽然远不如大模型但工程线路是走得通的。5.2 强化学习阶段的基本工程框架如果想把推理能力真正“调教”出来下一步就要引入强化学习。这里不是让你从头实现PPO或GRPO而是要把整个训练框架在工程上盘顺。最简单的路线是先SFT训练出一个会输出CoT格式的模型这一步只是让模型学会“格式”。准备一个奖励函数可以是规则评分比如数学题的最终答案字符串是否匹配也可以是一个小的奖励模型。用GRPO这类梯度策略算法微调让模型在给定问题上生成多个候选推理过程奖励高的轨迹被加强奖励低的轨迹被压制。工程上最需要注意的地方是“采样过程中的内存释放”。强化学习阶段会在同一个batch里生成大量候选文本如果不及时清理GPU上的中间logits显存占用会呈指数级别增长。我的建议是用torch.no_grad()做推理采样只保留生成后的token序列和奖励值把所有涉及梯度的操作严格限制在策略模型再次前向计算时。5.3 给你的学习路线建议如果你连基础的从零LLM都还没跑过千万别直接跳到推理模型。我见过太多人一上来就去看R1的技术报告结果被强化学习部分的公式劝退。合理的学习顺序应该是先跑通一个GPT风格的小模型理解训练和生成的闭环。再在生成阶段自己实现一遍温度采样、top_k、top_p体会策略空间的含义。接着用现成的小模型做SFT学会构造CoT数据。最后才是强化学习微调。每一步之间最好都隔出几天时间把前一阶段的代码吃透再往前走。从零构建的路非常考验耐心但一旦走通你对模型的掌控感是指数级提升的后面无论是换大语料、换更大模型还是换强化学习算法都有了可以依仗的知识骨架。关于这个项目最后想说的把 ai-engineering-from-scratch 从头到尾做下来我最大的体会不是“我会训练模型了”而是“我理解了我正在用的每一个工具”。以前用别人的推理接口总觉得生成结果像个黑盒子现在看到一条loss曲线、一次显存OOM、一个重复生成问题我都能直接定位到那个具体环节。这种手感没有任何捷径只能靠亲手构建一遍来获得。最后分享一个小技巧当你想快速验证一个新模型架构是否靠谱不要急着上大数据先用一个不超过1亿参数的小模型跑通单batch过拟合测试loss能掉到接近0再考虑放量训练。这会帮你过滤掉绝大部分结构性bug省下的时间远比那点测试时间多。如果你也正在准备走这条路希望这篇记录能让你少踩几个我踩过的坑。
返回列表