拓十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

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

从零搭建AI工程全链路:手写Transformer与部署复盘

从零搭建AI工程全链路:手写Transformer与部署复盘 ai-engineering-from-scratch这个名字乍一看像某个开源仓库的标题但它其实是我花了小半年时间维护的一套个人项目记录不依赖任何现成的AI应用框架从零开始搭建一条完整的AI工程链路。这里的从零不是指从空文件夹开始写Python而是指不直接调用封装好的大模型API、不直接用现成的训练框架从数据处理、模型结构、训练循环到推理服务每一层都自己动手实现一遍。这个项目做完之后我对AI工程这四个字的理解发生了很大的变化——它不只是调接口、写提示词而是一套涉及数据、模型、算力、服务化的系统工程。这篇文章就是这套项目从0到1的完整复盘内容包括整体架构怎么拆、数据管线怎么做、Transformer核心组件怎么手写、训练和推理阶段有哪些容易踩的坑以及我对整个学习路径的思考。如果你已经会用Python也跑过一些现成的AI模型但总觉得这些东西内部像个黑盒想自己动手把盒子拆开看看那这篇文章应该能给你一条清晰的路线参考。1. 项目整体设计为什么要把AI工程从零做一遍1.1 核心需求解析这个项目到底在解决什么问题先说清楚我最初的想法。那时候我已经用别人的训练脚本微调过一些模型也用过封装好的推理服务效果都还行但有一个问题始终让我不安出问题的时候我完全不知道问题出在哪一层。Loss不降的时候我不知道是数据的问题还是学习率的问题推理变慢的时候我分不清是显存不够、算子效率低还是框架本身有瓶颈。这种只会用不会修的状态对一个做工程的人来说是很危险的。所以我给自己定了一个目标不借助任何一键式框架徒手搭建一条AI应用链路。具体来说包括四个核心环节——数据工程分词器、数据清洗、采样策略、模型工程Transformer的完整实现、训练循环、推理优化KV Cache、批处理、量化、服务化部署模型加载、并发处理、监控。这不是为了造轮子而是为了在拆轮子的过程中理解每一处设计背后的权衡。这个项目的关键取舍是工程优先而非研究优先。我不去追求刷榜指标不搞全新的模型结构而是把已经成熟的技术一条条落地。就像学开车一样研究的任务是设计更好的发动机而工程的任务是把它装进车里、调好变速箱、处理好路况让整车能稳定跑起来。AI工程做的就是这个装车的活儿。1.2 方案选型在全手工和“调接口”之间找一条折中路线一开始我也想得很激进想着连GPU都用不着纯CPU手写一个完整的Transformer训练一个模型。后来实践证明这条路太极端了。训练一个能用的小模型就算参数不多在纯CPU上也要跑到天荒地老。所以我在实际操作中把路线调整为核心逻辑全部自己实现比如分词器、Dataset构建、模型结构、训练循环、推理服务这些代码一行不省但算力调度、混合精度、底层算子这些基础设施我选择站在巨人的肩膀上。这个折中非常关键它让我把精力花在了工程链路上而不是去重写GPU算子。你可以理解为自己动手建房子但钢筋水泥用现成的建材不自己去烧砖。技术上我选了PyTorch作为底层因为它生态成熟但所有模型代码我都用nn.Module从零搭不用现成的TransformerEncoder。分词器参考了BPE的原始论文自己实现。服务化部分用的是FastAPI加自写的推理封装不用企业级的大模型部署框架。这套选型让我既看到了每一层的真实面貌又保证了项目能在合理的时间内有产出。1.3 适合谁的路线图三种人建议这么学做完这套项目之后我对AI工程的能力地图有了清楚的认知也大概知道了什么样的人适合走这条路。第一种是刚入门想深入理解大模型原理的学习者第二种是被框架困住、想提升排障能力的工程师第三种是想从算法转工程、补足工程化能力的从业者。如果你只是调用现成API做应用开发那这个项目的很多内容对你来说属于进阶题可以挑着看。我自己的体会是这个项目不适合作为AI学习的第一步。它要求你至少有基础的Python功力最好还训练过一两个模型知道loss、梯度、过拟合这些基本概念。有了这些基础再从头造轮子才有意义否则容易在细节里绕不出来还打击信心。2. 数据工程部分从零手写BPE分词器与数据集管线2.1 分词器不是小事为什么选择自己实现BPE很多做AI应用的同学不太关注分词器觉得它就是调用一下tokenizer.encode。但如果你要自己训练模型分词器是第一个必须亲自动手的地方因为它直接决定了模型的语言单位。词表大小、分词质量、未登录词的处理方式都会影响下游训练效果和推理时的泛化能力。我选择实现BPEByte Pair Encoding而不是更复杂的WordPiece或者Unigram因为BPE在效果和实现复杂度之间是最平衡的。BPE的核心逻辑说白了很有趣从字符级开始不断统计相邻字节对的出现频率把最高频的字节对合并成一个新token反复执行直到词表达到目标大小。这里有一个工程细节词表大小设多少合适。我基于中文语料实验了一段时间发现8000到16000这个区间对中小型任务性价比最高。词表太小句子会被切得很碎序列太长影响训练效率词表太大模型embedding层的参数会暴涨而收益却不明显。BPE有个问题是在多语言场景下容易出现奇怪的token边界比如中文里一个词被切得七零八落。我的处理方式是在数据预处理阶段先用jieba做一次粗粒度的分词然后再在词级别做BPE合并这样中文场景下的语义完整性会好很多。这个方法不算标准做法但实测下来对中文下游任务的稳定性提升很明显。2.2 数据清洗与采样策略高质量的喂饭比模型结构更影响下限动手训练过模型的人都有一个体会模型效果的上限取决于数据和算法但很多时候数据比算法更决定成败。我在项目中专门搭了一条数据清洗管线规则不复杂但很实用——去重MD5指纹加局部敏感哈希、去除低质量文本长度过滤、标点符号密度过滤、乱码检测、去掉与任务无关的噪声比如网页里的导航文本、无意义的重复字符。这里我多说一句规则细节很多新手清洗数据就只去个重这是远远不够的。我遇到过语料里有大量啊啊啊啊啊这类重复字符的情况BPE会把它们合并成高频token直接压缩了词表里真正有语义token的空间。后来我在清洗阶段加了一条规则如果一行文本里同一个字符连续出现超过10次就做压缩处理或者直接丢弃。一个小规则训练出来的模型生成质量明显更稳定。采样策略上我一开始犯过一个大错按原始文件大小来采样。结果某个超大文件的主导内容完全覆盖了其他小文件的信息模型学到的知识严重偏科。后来我改成按类别分层采样每个类别设定一个目标占比在超大数据集里先random shuffle再按比例抽取。这个细节直接影响了模型在验证集上的表现可以说是数据工程阶段成本最低、收益最高的一次改动。2.3 上下文长度与token化效率训练前的最后一公里数据准备的最后一步是构造训练样本。这里面有一个容易忽略的坑上下文长度的选择。一开始我参考论文直接用512后来发现我的任务大多数句子长度集中在30到80个token之间把上下文设成512意味着大量padding计算白占显存也拖慢训练速度。我把长度调整到128之后训练速度提升了近40%而效果几乎不变。更关键的一个细节是token化效率。我犯过蠢事在Dataset类的__getitem__里实时调用分词器结果每个epoch都要重新分词浪费时间。后来我把token化提到数据预处理阶段一次性把语料全部转成token id训练时直接加载整数数组读取速度大幅提升。这个优化看起来很简单但在数据量有几GB的时候就非常明显。还有一个实用的工程技巧是做一个token比例图。我会对每一个batch统计真实token占比去掉padding之后的有效token数除以总token数这个数字如果低于0.6说明上下文尺寸设置不合理或者数据padding太多需要回炉检查采样策略。这个指标成了我衡量数据管线是否健康的一个核心依据。3. 模型核心组件手写Transformer的关键模块与值得注意的细节3.1 整体架构设计与模块拆解不直接用框架的Encoder模型结构我没有用一行现成的Transformer实现。拆开来看主要包括Token Embedding层、可学习的位置编码、多层Decoder Block每层包含自注意力模块和前馈网络、LayerNorm和残差连接。我选择做Decoder-only结构因为我的目标场景是文本生成而且这个结构在自回归任务上更自然。每个Transformer Block都包含一个因果自注意力层带mask防止看到未来token和两层前馈网络。前馈网络的隐藏层维度我设置成模型维度的4倍这是一个工程上广泛使用的经验值。模型总参数量大致可以通过公式粗略估算参数量约等于嵌入维度乘以词表大小加上层数乘以注意力参数加前馈参数。以我的配置为例模型维度2568层词表大小1万参数量在2000万左右这个量级在单张消费级显卡上可以轻松训练。我强烈建议在实现时先写一个参数统计函数每加一个模块就打印一次总参数量。这能帮你从直觉上建立每个设计选择都会带来参数成本的概念。我自己在加位置编码的时候发现可学习位置编码虽然能用但参数量微小而收益不明显后来直接换成了基于频率的三角函数式位置编码效果差不多、参数量更少生成时还能处理更长的序列。3.2 缩放点积注意力的为什么除以根号d_k不是玄学自注意力机制里有一个看似不起眼但至关重要的细节在计算Q和K的点积之后要除以根号d_k。很多教程会说这是为了梯度稳定但真正理解它需要算一下。假设Q和K的每个维度都是均值为0、标准差为1的随机变量那么点积的均值是0标准差是根号d_k。当d_k变大时点积的方差也会变大导致softmax输入值分布过于极端——某些值非常大softmax输出接近one-hot梯度消失。除以根号d_k正好把方差还原为1让softmax区域保持在梯度敏感的范围内。这个保持方差一致的思想在transformer里还出现在初始化策略上是我认为整个结构里最优雅的一处设计。我的注意力实现里还处理了两个工程细节。第一个是因果mask在计算时加了一个极小负值而不是零因为加零会让softmax里那些位置的exp值为1而不是0注意力权重不会被完全屏蔽。第二个是注意力权重矩阵的维度是batch, head, seq_len, seq_len这个矩阵在中长序列下会非常占显存所以后面我在推理阶段用KV Cache来规避这个问题。3.3 训练稳定性设计LayerNorm位置、初始化与梯度裁剪Transformer训练的稳定性很大程度上取决于两个设计一是LayerNorm的位置二是初始化策略。关于LayerNorm现代transformer普遍采用Pre-LN结构也就是先做归一化再做注意力/前馈计算。我在实践中的体会是Post-LN在深层网络下很容易出现训练初期loss震荡而Pre-LN收敛更稳。代价是Pre-LN在部分任务上的最终效果略逊但工程上我更看重稳定收敛因为一个epoch几百万token的训练成本摆在那里训练崩了重来的代价太大了。初始化方面我采用了标准差为0.02的正态分布初始化所有参数这是很多大模型普遍使用的做法。特别注意两点一是embedding层的初始化方差不宜过大否则初始输出分布会偏移二是残差分支的初始化我采用了GPT-2的缩放方法把每个残差分支的初始化标准差按层数开根号缩小确保深层网络的输出方差保持在合理范围。梯度裁剪也建议开着。我在训练时默认把全局梯度范数裁剪到1.0别问为什么问就是有过惨痛的教训。某一个epoch我用3e-4的学习率训练忘记开裁剪结果loss曲线像过山车一样上下乱窜后来开了裁剪才恢复正常。对比实验后我认为裁剪不会损害模型质量但对训练稳定性有显著帮助。4. 训练工程实践从损失设计到训练循环的调优过程4.1 优化器选型与学习率策略为什么AdamW是默认选项训练一个小型模型的时候你同样会遇到大模型训练时的所有基本问题选择优化器、设置学习率、防止过拟合。我在项目中把优化器从SGD一路实验到AdamW最后固定使用AdamW。关于优化器我核心关注的不是收敛速度而是稳定性。SGD在小型任务上也能收敛但对学习率极其敏感需要精心调整。Adam系列对学习率的宽容度高很多。而AdamW和Adam的差别在于weight decay的实现位置解耦权重衰减这确实是改善泛化能力的有效手段。在我的任务上AdamW的验证集loss始终比传统Adam低一点点这说明正则化是真实有效的。学习率策略上我采用了预热加余弦退火。具体来说前1000步学习率从0线性升到峰值3e-4之后按余弦曲线衰减到峰值的1/10。前期预热的意义是让模型参数在一个较小的步长下先稳定下来避免一开始就走太快扯到蛋后期余弦退火则能让模型在loss盆地底部进一步精调。这个策略不需要每轮手动调参是工程上非常省心的选择。4.2 混合精度与批量大小在单卡上把算力吃满训练效率优化是我在这个项目里花时间最多的部分之一。我的硬件只是一张消费级显卡能不能把算力吃满直接决定了项目周期。最先做的是混合精度训练用FP16计算梯度用FP32保存主权重再用损失缩放保证梯度不下溢。实测下来训练速度提升基本能稳定在1.5倍左右显存占用降到原来的一半以下。这里有一个细节值得说一下混合精度训练里如果你的模型里有一些对精度特别敏感的操作比如softmax或者LayerNorm最好把它们放在FP32下计算PyTorch的autocast会自动处理这些但如果你用自定义算子就要自己注意。我因为没注意曾经在自定义的LayerNorm里强制在FP16下计算结果训练到一半loss变成了NaN后来把所有中间计算改成FP32才好。批量大小方面我尝试过能塞进显存的最大值。量化地说当批量大小从16提高到64时训练吞吐量提升了接近两倍同时稳定性也更好。但这里有个误差太大的批量会让模型收敛到sharp minima泛化能力可能下降所以我的做法是在显存允许的边缘试探设定一个上限后就不再加了。平衡点在64左右单卡刚好能跑。4.3 训练过程中的实时监控与中断恢复不做好会想哭训练过程不是配置好就万事大吉的。我每次开启一轮训练都会同时启动一个后台监控线程记录loss、学习率、梯度范数、token吞吐量、显存占用这几个指标。梯度范数是我用来判断训练健康度的关键信号如果出现异常波动大概率是数据里有问题或者学习率偏高了。有一个常见的教训训练中断恢复。如果你训练到一半断点了没有checkpoint恢复机制就得从头再来。我在项目里定义了每500步保存一次checkpoint加一个对应的optimizer状态字典的策略这样即使中途崩溃最多丢失500步的计算量。这个问题的工程解法极其枯燥但谁经历过凌晨一点训练到95%崩掉然后又要重来的痛苦谁就会感谢自己当初写过这样的代码。我还加了两个观察维度的习惯。第一个是样本可视化每500步我会从验证集抽几个真实样本让模型做一次生成看看输出是通顺的还是胡言乱语这比盯着loss曲线更直观。第二个是loss曲线绘制我不用tensorboard直接用一个简单的日志库收集数据再用matplotlib画图轻量且够用。这个组合拳让我能在一轮训练进行到一半的时候就预判结果而不必傻等整个训练结束。5. 推理与部署从模型checkpoint到稳定对外服务5.1 KV Cache的原理与实现让推理速度翻倍的技巧训练阶段结束之后真正让模型产生商业价值的是推理。很多人训练完模型就把推理当个附属品处理但其实推理优化直接影响用户体验和部署成本。我在项目中实现的第一个推理优化就是KV Cache。在自回归生成中模型逐token生成输出每一步都要重新计算之前所有token的Key和Value向量。没有KV Cache的时候生成第100个token要重复计算前99个token的attention状态。KV Cache的思路就是把之前算好的Key和Value缓存下来每一步只需要计算当前token的Key和Value然后拼接上之前缓存的K和V再进行attention计算。这个优化在代码层面听着简单但实际实现有几个细节要处理好。第一个是显存分配K和V的缓存需要按最大生成长度预分配否则每步都做显存动态分配会引入不必要的开销。第二个是缓存的生命周期管理一个batch里每个序列的长度都可能不同短的序列已经结束了它的缓存要尽快释放不然显存会像泄漏一样被慢慢耗尽。我在实现时给每个序列维护了一个是否活跃的标志位序列结束时立刻清理对应的KV缓存块。用上KV Cache之后我实测的生成吞吐量大约提升了3到4倍显存占用虽然变高了因为缓存需要额外的空间但换来的速度提升是非常值得的。这里再补充一个很小的组件我写了一个贪心解码函数每次选择概率最高的token作为输出后来又加上了温度采样、Top-k采样和Top-p采样三种策略让推理时有更多控制手段。5.2 批处理与并发设计为什么需要动态Batching而不是批大小为1推理阶段最常见的问题不是模型算不过来而是并发请求来了的时候每个请求自己占一个batchGPU利用率惨不忍睹。批大小为1跑推理的时候单卡利用率经常只有个位数百分比。解决这个问题的方法叫动态Batch核心思想是把多个请求攒在一起凑够一定数量或者到了最大等待时间后一起送进模型推理。我实现动态Batching的思路不复杂一个请求队列加一个调度线程。调度线程每隔10毫秒查看队列如果队列里的请求数量达到了预设的batch大小我常用的是8或者距离上次推理已经过去了50毫秒就触发一次批量推理。这样做的效果非常直接在并发量低的时候延迟控制得不错在并发量高的时候吞吐量接近理论上限。这里面有一个关于显存的上限计算需要提前做好。比如模型是2000万参数FP16推理权重本身其实只占约40MB但是KV Cache和注意力矩阵加起来可能再吃几百MB。我在部署前会先写一个前向函数用最大的batch size和最大生成长度试一遍算出峰值显存再留出30%的缓冲这样不会上线之后才OOM还找不到原因。5.3 服务化架构从本地脚本到可调用的AI服务最后一步就是把模型封装成稳定可调用的服务。技术选型上我用了FastAPI一方面因为用起来确实简单另一方面异步支持挺适合处理并发请求。在这个环节最有价值的工作是设计并实现了推理服务的三个分层结构。第一层是接口层负责接收HTTP请求做基本的参数校验第二层是模型管理层负责加载模型、管理KV Cache的分配和释放第三层是推理执行层负责运行模型前向、把输出转为文本、做日志统计。这个分层让我在处理问题时能快速定位是接口问题、模型问题还是资源分配问题而不用在一个大函数里翻来翻去。部署时还加了一些实用保障措施比如请求超时处理如果一个生成长度很长的请求迟迟不结束就主动中断返回已生成的部分再比如生成长度上限控制避免恶意或异常请求导致显存不可控膨胀。服务上线之后我测试了持续打几个小时请求的稳定性确认没有内存泄漏、没有线程泄漏日志里token吞吐量也保持稳定。这个稳定的服务最后部署到了线上跑了一段时间让我真正体会到模型训练只是AI工程的一半另一半模型上线后的工程保障。6. 常见问题排查那些在实操中反复折磨我的Bug与解法6.1 训练阶段的经典故障Loss不降、梯度爆炸、NaN训练过程中我遇到的最常见故障有三个Loss完全不下降、梯度爆炸导致Loss骤升、以及训练到一半突然出现NaN。这三个问题各有各的原因但排查路径是有套路可循的。Loss不降的时候我第一步看的不是模型而是数据把一条样本打印出来看看输入和标签是否对齐了。我曾经犯过一个低级错误就是没有shift标签导致模型要预测的token永远滞后一位。这个小问题让我排查了整整一个晚上。如果你的数据和标签没问题第二步检查的是学习率——如果学习率太小模型更新太慢Loss看着就像不动一样这时要用小数形式打印出实际学习率确认一下。梯度爆炸和NaN通常相伴出现。排查顺序是这样先把梯度裁剪打开看是否解决如果还不行就把混合精度关闭测试看是不是FP16中间计算溢出了如果还不行就去检查数据里有没有包含巨大数值的特殊token或者损坏样本。我最后一次遇到NaN就是数据里混进了一段带有特殊无意义字符的文本导致embedding层计算出极大值。清洗数据之后NaN问题彻底消失。这里提供一个排查技巧写一个训练健康检查函数在每个batch结束之后输出梯度范数、权重范数和激活值均值。如果你能看到这些指标的变化趋势大多数训练故障都能在几分钟内定位出来而不是靠重启训练碰运气。6.2 推理阶段的性能瓶颈显存OOM与生成速度慢推理阶段的问题类型和训练阶段完全不一样。我踩过的最大的坑是生成速度莫名变慢排查下来发现是没有复用KV Cache导致的。加上KV Cache之后速度立刻提上来了这个在前面已经说过不再重复。显存OOM则是另一个大坑。第一次部署上线时我的服务在低并发下一切正常但一到高峰期就报OOM。查下来发现有两个原因第一是KV Cache为每个请求动态分配显存请求多且生成长度长的时候缓存叠加起来非常可观第二是我没有限制最大并发数导致多个长请求同时挤在一起。解决办法是预先分配固定大小的缓存池超出容量就直接返回繁忙提示而不是让显存无限增长。还有一个不太起眼但经常拖慢速度的因素CPU和GPU之间的数据传输。每次只把单个token传回CPU再传下一个请求会产生巨大的传输张量开销。正确做法是一次把整段需要处理的数据打包传过去减少设备间的通信次数。这个优化在长序列推理时效果尤其明显。6.3 数据质量导致的隐藏问题为什么模型会生成重复内容最后说一下我通过这个项目深刻认识到的一个问题很多看似模型能力不行的表现根源其实在数据质量上。我训练出来的模型有一段时间生成内容非常爱重复总是循环输出同样的句子。我一开始以为是模型结构问题后来对比实验发现训练数据里存在大量重复的文本片段模型直接把复读这个模式学进去了。解决方法是给数据清洗管线加了一条规则算一下每条文本里相邻n-gram的重叠度如果一个句子中连续的n-gram大量重复就判定为低质量文本并过滤掉。这个规则加入之后模型生成内容的多样性有了肉眼可见的提升。从那以后我养成了一个习惯分析模型输出之前先分析它的训练数据。如果你也是在训练自己的模型我建议你把数据质量监控当作一个持续迭代的工程来做而不是一锤子买卖。数据分布会随着来源变化而变化只清洗一次远远不够定期抽样检查数据能看到很多你没预料到的问题。7. 全文总结与后续拓展方向从动手实现BPE分词器开始到把一个可调用的AI推理服务部署上线这个名为ai-engineering-from-scratch的项目让我把AI工程从抽象词汇真正变成了肌肉记忆。我最大的体会是训练一个模型只是AI工程的一小部分一个稳定、高效、可维护的AI系统需要你把数据、模型、训练、推理、服务化每个环节都理解到位。这个项目做完之后我后续的计划是往两个方向拓展。第一个方向是做模型评估体系给这个模型搭建一套自动化评测框架包括生成质量指标、鲁棒性测试和性能基准测试让模型迭代有数据驱动依据。第二个方向是继续优化训练效率研究更高效的注意力变体、稀疏化方法和分布式训练方案让更大规模的模型在资源有限的情况下也能训练起来。如果你也想尝试这种从零构建AI工程的路线我给的建议是先选一个极小的任务比如做一个中文笑话生成器或短文本续写模型让整条链路先完整跑通再逐步加复杂度。这条路最迷人的地方在于你踩过的每一个坑都会让你比那些只调接口的开发者多一层对这个系统的理解。愿你在拆轮子的过程中也能享受到那种把黑盒一点点打开的快感。
返回列表