ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

从零搭建AI系统:Transformer训练推理全流程实战指南

从零搭建AI系统:Transformer训练推理全流程实战指南 “ai-engineering-from-scratch”这个话题说白了就是一句话不靠现成的模型接口和框架封装自己动手把一个能用的AI系统从零搭出来。过去一年我一直在做这件事从最开始的“调一下超参就以为自己在做AI”到后来真正把数据管线、模型结构、训练策略、推理部署全链路亲手过了一遍中间踩过的坑比学到的知识还多。但我可以负责任地说这一步绕不过去——你只有亲手写过一遍attention、亲手盯着loss曲线从4.2降到0.1才能真正理解那些开源框架的每一行配置到底在干什么也才能在模型出问题时不至于一脸懵。这篇文章面向两种人一是刚入门、想系统掌握AI工程能力的同学二是已经在用现成框架但总觉得“差点意思”、想补上底层功底的工程师。我会从思路拆解、核心原理、实操实现、踩坑实录四个角度把我亲测有效的最小闭环路径完整分享出来每一步都带着为什么要这么干的解释不是那种“照着敲一遍就完事”的教程。1. 先搞清楚“从零开始”到底是从哪里开始1.1 三个层次的“零”定位完全不一样我见过太多人一听到“deep learning from scratch”就开始手推反向传播一听到“LLM from scratch”就想复现GPT-4这是两个极端。实际工程里“from scratch”至少有三种层次各自的工作量和对你的价值完全不同。第一层是“从数学开始”。手写矩阵求导、手动实现反向传播推导Attention公式的梯度流。这一层适合彻底搞原理但说实话日常工程里你不会天天写反向传播除非你要做框架开发或者学术界发论文。我当初跟着课程手推过几轮链式法则收获很大但它不是主线任务更像是一个必要的“心智建设”阶段。第二层是“从框架原语开始”。用PyTorch/TensorFlow的基础算子比如torch.nn.Linear、torch.matmul、torch.nn.functional.scaled_dot_product_attention去拼出一个完整的Transformer自己写训练循环、自己实现数据采样、自己管理checkpoint。这一层才是“AI工程sense”的核心来源。你开始明白一个batch的数据长什么样、tokenizer的padding到底怎么影响训练效率、梯度累积和batch size之间是什么关系。第三层是“从裸代码开始”。完全不依赖深度学习框架用纯Python和NumPy实现前向传播和反向传播。我建议有兴趣的人做一次但别把它当日常。因为生产环境里99%的场景你会用框架裸写是为了破除“框架是魔法”的心理障碍。这篇文章的“from scratch”我定位在第二层为主、辅以部分第三层原理理解。也就是说你跟着文章走完能自己写出一个小规模的、能训练的、能推理的reasoning模型雏形而不是只会调用现成的API。1.2 为什么一定要亲手造一遍轮子直接说结论用现成框架三个月你对模型的理解可能停留在“换个数据集、调一下学习率、看acc涨没涨”。但亲手造一遍轮子哪怕是最小的模型你会被迫面对一个接一个“原来如此”的问题。比如你会发现训练一个语言模型最重要的其实不是模型结构有多花哨而是数据怎么切。我第一版自己写的training loop里因为偷懒没有按document边界做截断导致一个batch里的序列跨越了两篇完全无关的文档模型训练出来的效果惨不忍睹。这种问题你读十篇论文都不会有体感但自己写一次数据加载器就全懂了。再比如你会在“位置编码”这件事上彻底想通。之前我用position_ids只是当成一个常规参数直到自己实现RoPE旋转位置编码时才明白为什么很多论文说“外推能力”很重要因为三角函数编码在相对位置上的表达是连续的而可学习的位置嵌入在序列变长时就会失效。这些认知不是靠背诵得到的是写代码写出来的。所以如果你真的想走AI工程这条路把“再造一个轮子”当成必修课而不是浪费时间。2. 核心知识地图把Transformer主线彻底吃透2.1 从整体到局部一个语言模型被拆开是什么样一个最简的GPT式decoder-only模型从输入到输出会经历完整的七站tokenizer把文本映射成整数idid进入embedding层变成向量向量加上位置信息进入多头自注意力层去“彼此交换信息”然后进入前馈网络做非线性变换中间靠残差连接和层归一化稳定训练最后一层把隐藏状态映射成词表大小的概率分布再用softmax选下一个token。很多新手卡在“注意力机制”这关我提供一个自己的理解方式你就把序列里每一个token想象成一个参加圆桌讨论的人attention就是让每个人看完所有人的发言后再基于“哪些人的话跟我当前的问题相关、相关到什么程度”更新自己的发言稿。Query是“我在找什么”Key是“我能提供什么”Value是“我真正要说的内容”。三个矩阵的乘积算出的就是“关注度权重”权重再对Value做加权求和就是一次信息融合。具体到代码层面自注意力就那么几步把隐藏向量过三个线性层得到Q、K、VQ和K做点积后除以“根号下头维度”做缩放用softmax归一化成权重权重乘以V再拼回头数。缩放那一步常被忽略但它很重要——不缩放的话当维度变大时点积结果的方差会跟着变大softmax会过早饱和梯度就消失了。多头注意力是什么呢就是不要把“一次圆桌讨论”只做一遍而是做8遍甚至12遍每个头用不同的Q、K、V投影让模型并行关注不同位置、不同维度的关系。比如一个头可能专门关注句法依赖另一个头关注共现关系最后拼接起来再过一个线性层融合。2.2 三大工程支柱归一化、残差、位置编码在模型结构里有几个部件常年被当成“默认存在”但你真的自己实现一遍时才知道它们是为了解决什么具体问题而存在的。层归一化LayerNorm解决的是训练稳定性问题。深层网络中每层输出的分布一直在变下一层要不停去适应上一层的“漂移”训练就会慢。LayerNorm对每个token的特征向量做均值和方差归一化再乘可学习的缩放参数和平移参数相当于把每一层的输入分布拉回一个相对稳定的区间。我用过PreNorm和PostNorm两种排列方式实践经验是小模型上PostNorm往往收敛更快但训练更深更大的模型时PreNorm更稳不容易炸loss。残差连接解决的是梯度穿越问题。没有残差时梯度要从最后一层一路传回第一层中间经过几十上百个矩阵乘法要么消失要么爆炸。有了残差恒等映射梯度有了一条“高速公路”可以直接从输出端回到输入端。这算是让深层网络真正能训练起来的核心设计之一。位置编码解决的是“顺序感知”问题。attention本身是集合运算你把一句话的token顺序打乱算出的结果完全一样不考虑mask。所以必须额外注入位置信息。简单的做法是学一个位置嵌入表更强的做法是RoPE把位置信息通过旋转矩阵编码进Q和K。RoPE有个非常漂亮的性质——两个位置的相对夹角只和它们的差值有关这让模型更容易学到“距离越远关系越弱”这种相对位置规律。3. 实操主线从数据管线到推理模型的最小闭环3.1 数据是唯一的野路子自己造一份干净的训练数据很多人动手写第一个模型时最容易忽视的就是数据预处理。我建议第一版不要用别人处理好的现成数据集而是自己选一个小范围语料亲手走一遍清理流程。我当时的做法是找一批公开的中文技术文档统一转成UTF-8剔除所有乱码行、空行、表格碎片然后按段落切分。这一步看起来简单实际藏了很多坑——比如编码问题、换行符不统一、半角全角符号混杂都会直接影响tokenizer的效果。tokenizer是第二个被我低估的模块。你直接用一个预训练好的BPE tokenizer的话能跳过很多痛苦但我还是强烈建议你至少自己训练一个极小的BPE词表设成8000左右。做这件事的意义在于你会亲眼看到“分词”如何影响训练。比如英文单词“learning”如果被拆成“learn”和“ing”模型学到的语义就比整词更稳定而中文如果只用单字切分很多复合词“人工智能”的语义就要靠模型自己拼装学起来更费劲。如果数据量不大建议控制在100MB到1GB这个范围。第一版我用了约300MB的文本训练一个130M左右的小模型两到三张卡跑了大半天。数据不用多关键是让整套流程顺畅走通而不是一上来就想复现大厂千亿参数。3.2 代码骨架一个能跑的最小GPT实现长什么样我直接把我第一版骨架的要点列出来你可以照着搭。完整的代码其实网上很多但关键是理解每块为什么存在。import torch import torch.nn as nn class MiniGPT(nn.Module): def __init__(self, vocab_size, d_model, n_heads, n_layers, max_seq_len): super().__init__() self.token_embedding nn.Embedding(vocab_size, d_model) self.pos_embedding nn.Parameter(torch.zeros(1, max_seq_len, d_model)) self.blocks nn.ModuleList([ TransformerBlock(d_model, n_heads) for _ in range(n_layers) ]) self.ln_f nn.LayerNorm(d_model) self.lm_head nn.Linear(d_model, vocab_size, biasFalse) def forward(self, idx): B, T idx.shape x self.token_embedding(idx) self.pos_embedding[:, :T, :] for block in self.blocks: x block(x) logits self.lm_head(self.ln_f(x)) return logitsTransformerBlock里核心是两个子层先做带因果掩码的多头自注意力再做前馈网络。前馈网络我用的是Linear(d_model, 4 * d_model) GELU Linear(4 * d_model, d_model)这已经是GPT系列的标准配方。为什么中间要扩到4倍因为FFN层是模型“记忆知识”的主要场所维度太窄的话容量不够训练完会发现loss下降很慢。因果掩码是decoder-only模型最关键的细节它的作用是在预测第t个token时禁止模型看到第t个之后的token。实现上就是一个上三角矩阵值为负无穷加到注意力权重矩阵上丢弃未来信息。你在写这个骨架时注意两个地方第一lm_head通常和token_embedding共享权重能显著减少参数量词表越大省得越多第二LayerNorm的参数在初始化时一般设eps1e-5太小可能导致数值不稳定太大又会让归一化失真。3.3 训练循环除了反向传播你还要会“看曲线”训练循环写起来不复杂但工程师和调包侠的区别就在这里调包侠只盯最终指标工程师盯整个训练的动态过程。我训练时的标准流程是这样的每N步打印一次step / loss / lr / tokens_per_sec我习惯设N100每500步在验证集上算一次交叉熵损失记录log用梯度范数torch.nn.utils.clip_grad_norm_做梯度裁剪max_norm我默认设1.0每个epoch结束时保存checkpoint除了保存state_dict还要保存optimizer状态、当前lr、step编号这样断点续训才完整。关于学习率我建议走warmupcosine decay策略。前几百步线性升到一个峰值再用余弦曲线缓慢下降。峰值lr的选取和batch size强相关我自己的经验是batch size为64时峰值lr在1e-3到3e-3之间比较安全如果用梯度累积模拟出更大的batch可以适当调高。注意不要拿一个很小的batch直接配一个很大的lrloss会在头部直接炸掉。还有一个经验loss刚开始的那几百步会下降得非常快从10降到3左右但之后会进入一个漫长的平台期。很多新手觉得模型坏了其实不是这是正常的。真正要警惕的是“loss降到一定值后突然反弹”那大概率是学习率峰值太高或者数据里有异常batch。3.4 从训练到推理把模型从“背题模式”切换到“答题模式”训练时模型看到的是完整序列预测时则是一个一个token往外交。推理的经典做法是自回归拿[BOS] prompt输入模型取最后一个位置的logitssoftmax后按某种策略选token然后把这个token拼到输入末尾重复直到遇到eos或者达到max_new_tokens。最常用的采样策略是top-k加temperature。temperature的作用是调节概率分布的尖锐程度大于1时分布更平缓、生成更大胆小于1时分布更尖锐、生成更保守。top-k则是在候选token里只保留概率最高的K个再做采样避免模型从超长尾里捡到垃圾词。我实测下来temperature0.8、top_k50是一套比较稳妥的默认参数但具体还得看场景——写代码风格的内容可以更窄闲聊场景可以更发散。如果你想让推理更快可以先做KV Cache把前面算过的K和V存下来避免每次生成都重新计算一遍之前所有token的注意力。这个优化对长序列推理的提升非常明显不夸张地说序列长度翻倍之后不加KV Cache的推理时间几乎指数上涨。4. 常见问题与排查技巧实录4.1 loss不降先按顺序排查这五个地方这是每个从零开始的人都会卡住的问题。我整理了一个排查顺序按成本从低到高排列建议你照着做第一检查数据的标签对不对。如果是语言模型目标就是输入序列右移一位最后一位通常被pad token补齐。新手最容易犯的错是label和input没对齐导致模型学的是“预测自己”。第二看tokenizer的输出是否有大量unk。如果你拿中文数据硬套一个为英文训练的BPE词表可能会发现几乎所有中文都被映射成unk模型根本无从学起。解决办法是自己重新训一个中文BPE或者至少把词表扩充到包含中文常用字。第三关注学习率。loss纹丝不动很多时候是lr太小loss先降后直接nan是lr太大。建议对比三个量级1e-4、1e-3、5e-3观察前200步的表现再定。第四怀疑mask写错了。我犯过一个低级错误因果mask的实现里多加了一个偏移量等于把未来信息泄露了一部分给模型。表面看训练loss很低但验证loss一直很高这就是典型的“作弊”信号。第五数据量太少导致的欠拟合。如果训练集只有几千条数据你不可能训出一个泛化能力很好的模型。这时候先别急着调模型结构先去多找一倍数据。4.2 显存OOM不是卡的问题是你设计的问题训练大模型最常见的报错是CUDA out of memory。我的经验是先别急着换更大显存的卡有几个优化空间先看看。第一个优化是梯度累积。显存不够可以把一个大的batch拆成多个小batch分别计算梯度后累加更新一次参数。比如目标batch size是64显存只够跑16那就累积4个step再更新。要注意的是PyTorch里默认的梯度是累加的所以在每个更新step之后要手动optimizer.zero_grad()否则梯度会跨累积步骤叠加方向就乱了。第二个优化是用混合精度训练。用AMPAutomatic Mixed Precision把一部分计算降到FP16同时保留FP32的master weight显存占用能降低近一半而且在A100这类卡上速度反而更快。但要注意FP16的表示范围有限loss出现nan时先检查是不是数值溢出了尤其是算attention时先做缩放再做softmax。第三个优化是减少中间变量的生存周期。比如你在debug时打开了梯度记录但没有实际使用梯度多出来的计算图会白白占用显存。用torch.no_grad()包住推理和评估代码能省下一大块内存。4.3 推理结果全是废话怎么办模型训练完生成出来的内容要么是重复的“的的了了”要么是一堆语法错乱的中文。这种情况十有八九是采样参数有问题或者模型还没充分收敛。重复性问题可以试试加repition penalty对已经生成的token的logits做惩罚或者提高top-k的K值给模型更多候选空间。如果是采样参数已经调整过还是很差那就要看训练loss是不是真的降到了合理水平。对于一个小型模型比如300MB数据、1亿参数训练完的交叉熵loss如果停在7、8以上基本就是没学好。这时候先回头查数据质量——我遇到过一版数据里混进了大量HTML标签残片tokenizer把div和/div当成了高频token生成的东西自然一塌糊涂。再有一种可能是你的prompt格式和训练数据格式不一致。训练时每段文本都带|endoftext|标记推理时你也得在prompt末尾加上这个特殊token否则模型根本没有“停止”的触发条件。很多人在这一步忽略了“BOS和EOS的一致性”踩坑之后才明白特殊token不仅仅是装饰它们参与了训练时的分布建模。5. 避坑指南与工程化路线建议5.1 从第一个模型到工程化真正的分水岭在哪如果你已经走通最小闭环第二个里程碑就是“工程化”。工程化不是说你会用Docker和K8s而是你的代码能经得起换数据、换环境、换规模的考验。我在做完第一版小模型后花了很大力气重构代码核心改动有三个一是把所有数据路径、超参数集中到一个config对象里不再散落在代码各处二是把数据处理、训练、评估、推理四个阶段彻底解耦每个阶段都能独立运行三是把随机种子固定逻辑做成一个公共模块保证实验可复现。这些改动看起来朴实无华但它们决定了你后续能不能快速做消融实验能不能从一次训练中得到可迁移的结论。再一个建议是永远保留一个“最小可用复现”分支。我在迭代过程中经常改坏代码如果每次都要从头找bug会非常浪费时间。我把第一版能跑通训练循环的代码单独打了个tag之后所有实验都从那个版本派生改坏了就回滚效率高很多。5.2 下一步路线图往CLM、SFT还是RLHF走你有了from scratch的基础之后下一步有三条路可以选取决于你想做什么。第一条路是继续做预训练模型CLM方向把模型规模从一亿参数推到十亿以上这时候你会遇到分布式训练的问题张量并行、数据并行、流水线并行。这个方向适合想进大厂做训练框架的人。第二条路是往SFT监督微调方向走把预训练模型用指令数据微调让它能“听懂人话”、按照指令回答问题。这条路你会接触到对话数据构建、指令格式设计、评估集怎么搭等是做成产品最直接的一步。第三条路是往RLHF/DPO方向走让模型学会对齐人类偏好。这条路的数学要求更高也需要更深的强化学习功底但如果你真的靠自己的双手完成了“预训练SFTRLHF”全链路那你对LLM的理解深度绝对超过大多数只会调API的开发者。我个人走的是第二条路为主因为它能最快形成产品闭环。但不管哪条路请记住一件核心的事你已经不再是“用AI”的人而是“造AI”的人。这个身份转换的关键就是当你面对一个从未见过的问题时不是去找现成的解决方案而是能自己拆开模型定位问题修好它再搭回原样。这也是“from scratch”这个项目带给我最大的改变——它不只是让我学会了怎么训练模型更让我建立起了一种对系统的掌控感。再遇到任何新框架、新模型我都知道它大概是怎么转的它的瓶颈可能在哪里哪一层是值得信任的哪一层需要亲自验证。
返回列表