ARTICLE DETAIL

资讯详情

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

从零构建AI工程链路:手写Transformer与推理模型实战指南

从零构建AI工程链路:手写Transformer与推理模型实战指南 去年我给自己挖了一个不小的坑在一个中小规模项目里不用任何现成的模型库从零动手搭建一条完整的AI工程链路。项目标题就叫“ai-engineering-from-scratch”。当时好几个朋友都劝我说现成的框架一把一把抓费这个劲图什么但当我真正把最后一个模块跑通、看到自己训练的模型在推理任务上输出有逻辑的答案时我觉得这个坑挖得值。这篇博文不讲大道理就把我从零搞AI工程的全过程拆开揉碎讲清楚覆盖数据、模型、训练、推理、调参和排错。如果你想搞清大语言模型和推理模型内部到底发生了什么或者你正打算做一个类似的从零构建项目这篇文章应该能帮你省下几个月的摸索时间。1. 为什么非要从零开始搞AI工程1.1 打开黑盒从调用者变成建造者我们用HuggingFace加载一个GPT模型时几行代码就能完成推理。表面上看效率很高但背后隐藏了一个大问题模型内部对你而言是一个黑盒。你只知道它接收文本、输出文本中间发生了什么无能为力。可一旦生产环境出了诡异的问题比如生成结果突然变得重复、注意力权重错乱、生成序列越界不懂底层就只能干瞪眼。从零开始构建本质上强迫你把黑盒打开。你会被迫处理每个关键模块词表怎么建立、embedding查表是什么、多头注意力里的张量维度怎么变换、因果掩码为什么要用上三角、LayerNorm归一化放在哪个位置、残差连接怎么加、梯度在反向传播中怎么流动。这些细节光靠读书掌握不了必须亲手实现一遍甚至写一遍、改一遍、debug一遍才能变成自己的东西。我一直认为“build a large language model from scratch”这个方向之所以流行不单单是因为大家想省几个API的钱而是因为这种重造轮子的过程能把理论真正变成自己的。训练完成那一刻你不仅能回答“模型为什么这样输出”还能回答“如果我想改变输出该改哪里”。这种掌控感用现成库的人体会不到。1.2 从零开始的范围到什么程度才算from scratch“from scratch”在不同人嘴里含义可能完全不同。我见过用NumPy纯手工实现反向传播的朋友连自动微分都不用也见过从PyTorch的nn.Module开始自己写模型模块的人。这两种做法都有意义但适用场景截然不同。纯NumPy从头实现深度学习框架适合对学习深度要求极高的人。你需要手写梯度推导、实现矩阵运算、考虑数值稳定性这些工作耗时巨大但能把深度学习的底层基础打得很扎实。另一方面如果你的目标是快速验证一个产品想法那就没有必要连框架也重写直接用现成的自动微分工具自己实现模型结构和训练流程已经足以建立对工程链路的完整掌控。我在这个项目里选择了中间路线基于PyTorch的自动微分引擎自己实现Transformer的全部结构、Tokenizer、数据管线、训练循环、评估和推理逻辑。这个选择背后有一个朴素的权衡我的核心目标不是复刻一个深度学习框架而是完整掌握语言模型从数据到推理的整条链路。如果你也准备做一个“ai-engineering-from-scratch”项目建议先给自己准确定义from scratch的边界不然后面很容易被各种实现细节带偏。1.3 与推理模型的交集从语言模型到reasoning模型这个项目里我还特意做了点扩展不满足于做一个普通语言模型而是尝试构建一个具备基础推理能力的模型。从技术角度看“build a reasoning model from scratch”比单纯做语言模型多了一层难度模型不仅要学会文本的统计规律还要学会在输出答案前先组织逻辑步骤。这也改变了我的数据构造方式。普通的预训练语料只需要“上文→下文”的预测关系但做推理训练时我需要给模型构建大量“问题→思考过程→答案”的三段式样本。模型在训练时不仅学习答案还要同时学习怎么一步步推导。这种形式是当前开源社区探索推理模型最主流的方向之一。个人感受是语言模型和推理模型之间的界限没有以前想象得那么绝对。只要能提供合适的训练数据和推理时的生成策略一个中小规模的模型也能表现出不错的推理能力。所以这篇文章里讲的技术不只是为了复现一个玩具而是指向当前AI工程里最热门的真实需求。2. 核心工程链路设计与方案选型2.1 端到端设计数据、模型、训练、推理怎么串起来我习惯把一个AI工程拆成四个阶段数据、模型、训练、推理。很多新手容易犯的错误是过分关注模型结构下载一堆Transformer结构图反复看结果数据清洗和评估环节草草了事。但真实工程里数据质量决定了模型天花板模型结构只是逼近天花板的手段。我甚至见过另一个项目组在模型结构上翻了新花样但数据里全是重复内容和格式错乱最后训练出来的模型效果依然惨不忍睹。我的链路设计顺序是这样的先定任务本项目目标是做一个能做基础问答和逻辑推理的中文小模型。再定数据收集约5亿token的高质量语料并额外构造10万条推理样本。三定模型使用120M参数的Transformer适合单卡训练。四定训练策略AdamW加余弦退火配合梯度累积与混合精度。最后定义评估不只看loss还要看生成质量和推理准确率。这里有经验的成分也有反复试错后的总结。把顺序反过来做比如先选一个大模型再找数据往往会在后期遇到资源瓶颈。我最初就是先定了模型结构再去攒数据结果发现语料规模远远不够只能回头调低模型容量白白浪费了两天时间。2.2 模型结构选型为什么是Transformer而不是别的现在讨论语言模型结构几乎绕不开Transformer。RNN和LSTM也能做序列建模但在长文本依赖和并行训练上吃亏。CNN在文本上有一定应用但很难天然建模长距离依赖。Transformer的核心是注意力机制它让任何两个位置都可以直接建立联系正好匹配文本建模的核心需求。我在模型里采用了12层Transformer decoder-only结构embedding维度768注意力头数12前馈隐藏层维度3072。这样参数规模约1.2亿比动辄几十亿的大模型小很多但足以支撑训练和推理实验。选这个规模的主要依据是显存和时间成本。在RTX 4090 24GB显存上120M模型搭配batch size 64、序列长度512单卡完整训练约需2到3天。如果再大一倍显存和训练时长都会翻倍个人项目很可能就直接放弃了。做工程和做实验不一样工程需要在一个可接受的资源约束下把链路端到端打通。结构并行训练长距离依赖实现复杂度适合场景RNN/LSTM较差较弱低短序列时序任务CNN较好受限中局部特征提取Transformer好强中高语言建模、推理任务2.3 关键参数是怎么算出来的参数选择不是靠猜而是可以计算推导的。先说Transformer参数量估算。以我的结构为例每层有四个主要部分多头注意力、前馈网络、两个LayerNorm。多头注意力部分的参数量约为4乘以d_model的平方对应Q、K、V、输出投影四个矩阵前馈网络部分约为8乘以d_model的平方因为先扩大到4倍隐藏维度再缩回来LayerNorm等小项忽略不计。所以单层参数量约等于12乘以768的平方约707万。12层就是约8484万。再加上embedding层词表50000乘768约3840万以及输出层总参数量约1.2亿。训练数据量的选择也有一个经验公式对中小规模模型训练token数最好是参数量的20倍以上这样模型能学到的语言模式才足够。120M参数理想对应24亿token但我手上的语料只有约5亿token所以一开始就明确这是资源受限下的可行方案效果不能对标大模型。后面验证时发现5亿token已经足以让模型学会基本句法和简单推理说明这个规模在实践中是可行的。2.4 Tokenizer与技术栈选型Tokenizer我选择了BPE算法自己实现。原因很简单既然标题写了from scratchTokenizer如果直接调现成包链路就不是完整的。自己实现BPE需要处理字节配对、词表大小、特殊token、序列长度控制。这套代码写下来对NLP预处理的理解会清晰很多。实现时有几个小细节值得注意。第一文本要预处理成字节序列避免出现未登录词第二词表大小我定为50000太小会导致编码碎片化太大会增加embedding矩阵的参数量第三要统一处理起始符、结束符、填充符、未知符四个特殊token否则后续batch处理和推理时会很混乱。技术栈方面用PyTorch 2.1作为自动微分和计算后端Python 3.10管理代码CUDA负责GPU计算。我没有选用HuggingFace Trainer所有训练循环和评估逻辑都自己写。这样做的直接收益是当训练出现问题或想改动一个细节时不需要去追第三方库内部的封装逻辑直接在自有代码里定位。刚开始会有点繁琐但项目后期越改越顺手。3. 实操从零搭建一个推理模型训练项目3.1 项目目录结构设计直接给出一份我最后落地的目录结构照着抄作业基本不会出错ai-engineering-from-scratch/ ├── configs/ │ └── config.yaml # 全局训练配置 ├── data/ │ ├── build_tokenizer.py # 训练BPE tokenizer │ ├── preprocess.py # 数据清洗与格式转换 │ ├── dataset.py # Dataset与DataLoader │ └── reasoning_data.py # 推理样本构造 ├── models/ │ ├── transformer.py # Transformer主体 │ ├── attention.py # 多头注意力模块 │ └── tokenizer.py # BPE tokenizer推理用实现 ├── train.py # 训练循环 ├── eval.py # 评估脚本 ├── generate.py # 推理生成脚本 └── utils/ ├── logger.py # 日志记录 ├── checkpoint.py # 模型保存与恢复 └── metrics.py # 评估指标这个目录的设计原则是模块边界清晰。数据、模型、训练、推理、工具各自独立互不纠缠。实际开发时如果多个文件之间互相依赖改一个地方就要牵连一片会非常痛苦。保持边界清晰对后期调试价值巨大尤其是后面想替换数据源或者改模型结构边界清晰的项目能少走很多弯路。3.2 数据准备从原始语料到训练样本数据准备是整个项目中最耗时但最容易被低估的部分。我使用的语料来自多个公开的清洁文本源拉下来之后第一件事就是清洗。清洗流程分为四步去重用MinHash方案识别重复段落并去掉。过滤删除所有非文本内容、乱码、超短片段、低质量重复文本。规范化统一标点、处理空白、压缩连续换行。编码处理统一转为UTF-8处理不可见字符。清洗之后的数据还需要tokenize。先训练BPE词表然后把每个文档切分成token序列按固定长度512组织成训练样本。这里要注意直接按512硬切会把一句话切成两半影响模型的学习。我采用段落边界感知的拼接策略优先在句号、换行等处切断实在不行再硬切。推理样本的构造稍微特殊。我给每条训练样本设计了统一的模板格式如下问题 问题内容 /问题 思考 逐步推理过程 /思考 答案 最终答案 /答案训练时把这三段拼接在一起作为模型的输入输出序列。模型从问题标签开始预测后面的token通过这种方式学会先思考、后回答的格式。对于已有数据我需要先通过规则或人工标注生成中间推理过程。这个环节的工作量比想象中要大但对最终推理能力的影响非常关键。3.3 模型实现手写Transformer的关键代码我把核心模型拆成了几个小文件。注意力模块是理解Transformer的钥匙下面给一个简化但可跑通的实现思路class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() self.n_heads n_heads self.d_head d_model // n_heads self.w_q nn.Linear(d_model, d_model) self.w_k nn.Linear(d_model, d_model) self.w_v nn.Linear(d_model, d_model) self.out nn.Linear(d_model, d_model) def forward(self, x, maskNone): B, T, C x.shape q self.w_q(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2) k self.w_k(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2) v self.w_v(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2) attn q k.transpose(-2, -1) / (self.d_head ** 0.5) if mask is not None: attn attn.masked_fill(mask 0, float(-inf)) attn F.softmax(attn, dim-1) out attn v out out.transpose(1, 2).reshape(B, T, C) return self.out(out)多头注意力里最容易出错的地方是维度变换。输入是(B, T, C)先线性投影成(B, T, C)再拆成多头的形状(B, n_heads, T, d_head)把每个头的计算并行化最后拼接回原来的维度。每一步debug最好都打印一下张量形状很多新人在这里栽跟头。我调试时会在forward里临时加print确认形状后再删掉。因果掩码是decoder-only模型的关键位置i只能看到它之前的位置。实现方式是构造一个上三角矩阵掩住未来的位置。注意生成时虽然是一个一个token地输出但训练阶段是一口气预测整个序列所以掩码必不可少。前馈网络部分相对直接就是两个线性层加激活函数中间扩到3072维再缩回768维。LayerNorm我用pre-norm也就是先归一化再进入子层而不是post-norm。pre-norm在深层网络里训练更稳定对学习率也相对宽容。这两个选择在实际调优中体现出了明显的稳定性差异。3.4 训练循环与策略自己写一遍才知道坑在哪训练循环用最朴素的思路实现加载batch前向计算计算loss反向传播梯度裁剪优化器更新。听起来简单但细节很多。我按照下面的顺序实现了训练主逻辑for step, batch in enumerate(train_loader): x batch[input_ids].to(device) labels batch[labels].to(device) logits model(x) loss F.cross_entropy(logits.view(-1, vocab_size), labels.view(-1)) loss / accum_steps loss.backward() if (step 1) % accum_steps 0: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step() optimizer.zero_grad() if step % log_interval 0: logger.log(loss.item(), learning_ratecurr_lr)其中accum_steps是梯度累积步数。因为单卡每次只能吃下很小的batch所以我设micro_batch_size为8通过累积8步达到64的等效batch size。这直接解决了显存不足的问题几乎没有额外代价。梯度裁剪很关键。训练早期注意力logits偶尔会出现非常大的值导致梯度爆炸。把梯度范数限制在1.0之后loss就平稳多了。不设裁剪时模型在某个随机step突然出现NaN整个训练都要从checkpoint恢复这是新手最容易踩的坑。学习率调度我用的是warmup加余弦衰减前2000步线性从0升到3e-4之后余弦衰减到1e-5。之所以要warmup是因为Adam在训练初期容易因为方差估计不稳定而走得太猛。如果不预热前几百步loss很可能剧烈震荡。训练过程中我持续打印loss、当前学习率、token吞吐量等信息同时每个epoch存一次checkpoint。模型崩溃时checkpoint是我们最后的安全网。个人建议是训练日志至少要记录到本地文件不要只输出到控制台因为有些崩溃发生在半夜起床后只能靠日志回放定位原因。3.5 评估指标与模型保存很多人把loss当作唯一指标但loss低不代表模型真的会推理。我的评估策略分三层第一层看训练集loss和验证集loss是否同步下降判断是否过拟合第二层看perplexity值它衡量模型对数据的预测能力第三层看实际生成的答案质量我会准备30条带标准答案的推理题计算回答准确率。当时我观察到验证集loss在训练后期不再下降但推理准确率却在缓慢提升。这说明模型在loss维度已经接近饱和但在推理策略层面仍有进步空间。如果只看loss我可能会提前停止训练错过这个缓慢提升的阶段。所以多维度评估很重要。模型保存我采用了两种策略每个epoch保存完整checkpoint包括优化器状态训练结束后单独导出一份推理用权重。推理时不需要优化器直接加载模型权重就行。这个分离让生成脚本轻量很多。3.6 推理生成与端到端验证训练完成后编写推理脚本。生成时采用上下文生成的方式输入问题模型先输出思考部分思考结束后继续输出答案部分。这里的生成策略也可以调节。我实现了一个带温度的采样生成def generate(model, tokenizer, prompt, max_new_tokens128, temperature0.7, top_p0.9): model.eval() input_ids tokenizer.encode(prompt) with torch.no_grad(): for _ in range(max_new_tokens): logits model(input_ids)[0, -1, :] / temperature probs F.softmax(logits, dim-1) sorted_probs, sorted_indices torch.sort(probs, descendingTrue) cumulative torch.cumsum(sorted_probs, dim-1) mask cumulative top_p sorted_probs[mask] 0 sorted_probs / sorted_probs.sum() next_token torch.multinomial(sorted_probs, 1) input_ids torch.cat([input_ids, next_token], dim-1) return tokenizer.decode(input_ids)推理时temperature的调节比较重要。太高会让输出杂乱无章太低会让模型重复同一句话。在我这个小模型上0.7附近效果最稳。top_p设为0.9把那些概率极低的离谱token裁掉整体生成质量提升明显。端到端验证时我会拿几个经典的推理题去测比如逻辑判断题、简单数学应用题。刚开始效果一般后来我把采样参数调低并固定了思维链格式效果明显改善。说明推理模型的性能不只取决于模型大小推理时的prompt格式和采样策略同样重要。4. 常见问题与排查技巧实录4.1 训练不收敛从NaN到loss不降这是每个从零搭建训练系统的人都会遇到的问题。我遇到的第一个典型症状是loss变成NaN。排查一圈后发现两个线索一是在混合精度训练下注意力矩阵的softmax输入有时会出现无穷大值二是某个样本里出现了超长或编码异常的内容导致数值溢出。解决方式分三层在数据管线里把token长度超过上限或包含异常token的样本过滤掉。在模型前向计算里给注意力score加一个极小值保护避免logits溢出。在训练循环里加上梯度裁剪。还有一个现象是loss不降。这种情况下我先检查学习率是否太大或太小然后检查数据是否一致也就是标签和输入有没有对齐最后再看模型初始化是否合理。这里的排查顺序很重要不要一上来就怀疑模型结构有问题。模型结构即使有bug通常loss也会有反应不会一点不降。排查方法小结先用很小的学习率跑100步看loss是否单调下降。如果还不行取一个极小的batch单步调试检查logits的形状和数值。加上断言检查张量维度尽早定位bug。4.2 推理效果差生成重复、答非所问训练loss很低但生成的文本却乱来这是典型的问题。原因可能有三个过拟合、采样参数不对、推理格式不匹配。过拟合在小语料上很常见。我训练的模型在训练集上loss很低但测试集上表现不佳。对策是增加数据多样性、降低训练轮数或者干脆缩小模型规模。采样参数不匹配也很常见。把temperature拉到1.0以上小模型很容易输出一堆毫无关联的词。相比之下0.6到0.8之间对多数问答场景更合适。还有一个容易被忽视的问题推理时prompt格式必须和训练时完全一致。训练时用的是问题、思考、答案模板推理时如果只喂一句干巴巴的问题模型就不知道该怎么组织输出。后来我把生成提示词也套进模板里输出结构立刻稳定了。4.3 资源不足时的应对清单个人项目最常遇到的就是GPU资源紧张。我总结了一份优先级排序的应对清单降低序列长度从512降到256显存需求几乎减半对短文本任务影响不大。减小batch size并增加梯度累积步数在显存允许的范围内尽可能提高等效batch。使用混合精度训练能省约一半显存同时加速训练。降低模型层数或隐藏维度比如从12层降到8层参数减少三分之一。尽早保存checkpoint一旦发现资源不稳定随时中断也不亏。这些措施不影响对工程链路的理解但能把资源门槛降得足够低让更多人在日常电脑上完成实验。如果你连本地GPU都没有也可以考虑用小模型加CPU训练只是时间成本会高不少。4.4 从零构建项目避坑速查表常见坑具体现象解决方案数据质量差loss下降异常模型生成乱码清洗、去重、过滤异常样本位置编码缺失模型对顺序不敏感效果差实现位置编码并确认维度正确因果掩码错误训练时偷看未来信息推理错误百出画矩阵图核对掩码形状学习率过大loss震荡或出现NaNwarmup加降低峰值学习率未做梯度裁剪训练中途突然崩溃设置max_norm为1.0推理模板不一致训练loss低但生成效果差保证prompt格式与训练一致采样温度过高输出发散、答非所问调低temperature并加top_p显存不足OOM报错减小batch加梯度累积加混合精度这张表是我在两个项目里反复踩坑后整理出来的。如果你的项目也面临类似问题建议先照着表里的方案排查一遍绝大多数情况都能解决。真要说这个项目给我带来了什么最大的收获不是那个120M的模型文件而是对AI工程这四个字的完整感知。写代码的时候我意识到很多平时调包时完全无感的细节背后全是工程和算法的权衡。踩坑踩到怀疑人生的那些晚上反而让我对从零构建有了更强的信心因为我知道每个环节都能自己改、自己修不再依赖别人的封装。最后再分享一个小建议如果你也打算启动一个类似的ai-engineering-from-scratch项目别一上来就追求完美的分布式训练或者超大规模模型。先把一个100M左右的小模型从数据到推理完整跑通再逐步加复杂度。这个先搭骨架、再填肌肉的顺序是我走过的最值得的一条路。这个项目后续还能往很多方向扩展比如给模型加上LoRA微调、把训练流程迁到多机多卡、或者接入产品里做真实业务的推理服务这些都会是在从零构建基础上长出来的新能力。
返回列表