资讯详情 从零构建AI工程:基于PyTorch手写Transformer到推理模型全流程
📅 2026/10/3 21:33:21
“ai-engineering-from-scratch”这个标题我盯着看了很久。它既不是一个具体项目的名字也不是某篇论文的标题更像是一整条学习路线的代号从零开始自己动手构建一套AI工程能力——不做调包侠不靠黑盒API把数据、模型、训练、评估、部署这条路亲手走一遍。这几年大模型火起来之后很多人直接站在“应用层”往上看会用别人的接口就算会AI了但真正遇到问题的时候无从下手。从一个从业者的视角看从零构建这件事不只是补基础更是建立“我能掌控它”的底气。这篇文章适合两类人一类是有一定编程基础、想从应用开发转向模型训练和AI工程方向的人另一类是已经在用大模型API做产品但总觉得底层是个黑盒、想拆开看一看到底发生了什么的人。我会按自己实际走过一遍的路线来拆解把从零训练模型、微调、构建推理模型这条路上绕不开的环节、踩过的坑、值得借鉴的思路都讲清楚。文章篇幅会比较长建议收藏后分段阅读。1. 从零构建AI工程先看懂整张路线图1.1 这里的“从零”到底指什么很多人一听到“from scratch”第一反应是“我要自己写一个Transformer”或者“我要复现GPT-4”。如果抱着这个目标大概率一个月后就放弃了。从我自己的经验来看“从零”正确的理解方式应该是不依赖现成的预训练模型和AutoML平台亲手完成数据准备、模型搭建、训练、调优、评估这一整套流程每一步都知道它在做什么、为什么这么做。这就好比学做饭。有人说“从零开始学做菜”你不是要从种水稻开始而是要学会自己买菜、洗菜、切菜、掌握火候最终能做出一桌像样的菜。你能分清生抽和老抽的区别知道什么时候该爆锅什么时候该小火慢炖——这才是“从零”真正的意义。至于能不能复现米其林餐厅的招牌菜那是后话。所以我在规划这条学习路线的时候给自己定了一个很务实的目标不用任何第三方封装好的大模型库用PyTorch从张量运算层面写出一个可训练的语言模型在合理的硬件条件下让它跑通训练流程产出能用的结果。然后再往推理方向延伸让模型不仅会“续写”还会“思考”。1.2 技术栈选型Python、PyTorch、CUDA与硬件技术栈的选择直接决定你在这条路上能走多远。我推荐的组合很固定Python PyTorch HuggingFace生态工具 一张至少8GB显存的NVIDIA显卡。为什么是PyTorch而不是TensorFlow或JAX因为PyTorch的调试体验是目前几个主流框架里最好的动态计算图意味着你可以在训练循环里随意打印张量、打断流程做检查这对新手期特别重要。Python这边没什么好说的AI工程的主战场语言。但我要强调一点不要在“学Python语法”上花太多时间直接在项目里学。你只需要掌握张量操作、类、循环、条件判断、文件读写这些基础然后马上下水。我见过太多人把《Python从入门到精通》翻了一遍又一遍结果连一个Dataloader都没写过这属于本末倒置。硬件方面如果你打算从零训练一个小型语言模型比如千万到亿级参数8GB显存勉强够用但会非常痛苦。16GB以上的显存会让你舒服很多24GB比如RTX 3090/4090基本可以覆盖大多数实验场景。没有N卡的话可以考虑云GPU实例按小时租用前期实验完全够用。需要注意苹果的M系列芯片虽然能跑PyTorch的MPS后端但很多CUDA生态的优化特性用不上踩坑成本会高不少。1.3 两条典型的“从零”路线经典机器学习与大语言模型规划路线的时候你会发现自己面前有两条路。一条是经典机器学习路线线回归、逻辑回归、决策树、SVM、梯度提升树然后延伸到基础神经网络。另一条是大语言模型路线直接上手Transformer、注意力机制、预训练、微调、对齐。我的建议是两条路都要走但顺序有讲究。先从经典机器学习入手不需要学得很深重点理解梯度下降、过拟合、交叉验证这几个核心概念因为它们在不同模型间是通用的。然后直接跳到语言模型路线把Transformer的每一个细节吃透。很多人想跳过第一步直接学大模型结果连“学习率”为什么要调整都搞不明白更别提理解训练曲线了。这两条路线我用一张脑图来整理过简单说就是经典路线负责建立“直觉”LLM路线负责建立“体系”。直觉让你面对一个陌生问题时知道往哪个方向找答案体系让你知道答案的每一层是怎么拼起来的。只有直觉没体系你是民间高手只有体系没直觉你是理论复读机。从零构建的核心就是两条腿走路缺一不可。2. 核心原理拆解训练一个模型绕不开的四个环节2.1 数据工程的起点文本清洗与分词器很多人一想到训练语言模型脑子里浮现的全是“注意力机制”“Transformer”这些高深名词。但真正动手做的时候你会发现第一个拦住你的不是模型架构而是数据。模型吃进去的是数字不是文字。把文本变成数字的过程就是分词器Tokenizer的工作。分词器的作用是把一段文本切分成子词单元再映射成ID序列。“我爱AI”这样一句话可能会被切成“我”、“爱”、“AI”三个token也可能被切成“我爱”、“AI”两个token取决于分词器的词表设计。常用的BPEByte Pair Encoding算法会从字符级开始反复合并出现频率最高的相邻字符对最终得到一份词表。这个方案的好处是它既不会像整词分词那样遇到生词就抓瞎也不会像纯字符分词那样让序列长到无法训练。从零训练分词器的时候需要重点考虑词表大小。词表太小句子会被切得很碎序列变长、计算量变大词表太大Embedding矩阵参数量暴涨小数据集根本学不好。我个人的经验是中英文混合场景下词表大小设在16000到32000之间比较合适。另外分词器的训练语料要和训练语料保持一致否则会出现大量token被切成碎片的情况。数据清洗也值得多说两句。网上爬来的原始文本里HTML标签、URL、乱码符号、重复字符、敏感信息全都要处理。我踩过最大的坑是语料里去重没做干净模型在训练时会“背诵”大段重复文本生成的内容就像复读机。清洗流程至少要包含统一编码、去HTML、去URL、去重、按质量过滤、长度过滤这几步。每一步都会直接影响最终效果别偷懒。2.2 模型架构从零手写Transformer核心组件如果只学一个模型架构那就是Transformer。从零构建AI工程核心中的核心就是亲手实现一个Transformer编码器或解码器。不要先用现成的nn.Transformer先自己写一遍多头注意力层、前馈网络、残差连接和层归一化再回头看现成实现会轻松很多。多头注意力机制是Transformer的灵魂。它做的事情用一句话概括让序列中的每个位置都能“看到”其他位置并根据相关性加权整合信息。拆开来看输入向量通过三个不同的权重矩阵映射成Q查询、K键、V值然后计算Q与K的点积获得注意力分数经过Softmax归一化后加权求和V。多头的意思是做多次这种映射每次学习不同的关系子空间最后拼在一起再过一层线性层。从工程角度看需要注意的点包括注意力分数通常要除以sqrt(d_k)做缩放否则点积结果过大会把Softmax推到梯度饱和区因果语言模型Causal LM还需要做mask让每个位置只能看到当前位置之前的信息预测下一个token。这些细节不亲手写一遍代码很容易在使用现成库的时候忽视掉。训练一个语言模型的时候输入和输出之间的关系很微妙。输入是一串token ID比如“我 爱 学 习”输出是这串token往右平移一位的“爱 学 习 ”。模型要做的事情是给定“我爱学”预测下一个位置最可能的token是“习”。这个设计就是自回归语言模型的基本范式ChatGPT系列的大模型都是这样训练的。2.3 训练循环前向传播、反向传播与梯度更新核心训练循环其实并不复杂多年来基本没变过。循环里包含四步取一个batch的数据、计算损失、反向传播得到梯度、更新参数。但是真要从零实现里面的细节多到你头皮发麻。前向传播阶段数据经过Embedding层变成向量然后层层通过Transformer Block最后一层线性层把隐状态映射到词表大小的logits。接下来计算损失函数常用的是交叉熵损失但需要做一点调整因为每个序列都有seq_len个位置每个位置都对应一次预测所以要把[batch, seq_len, vocab_size]的logits和[batch, seq_len]的labels对齐只对真实存在的token位置计算损失填充位padding要mask掉不参与梯度计算。反向传播阶段PyTorch的自动求导机制会替你完成大部分脏活累活但你需要知道哪些张量需要梯度、哪些不需要。比如Embedding层的参数需要梯度而padding mask生成的布尔矩阵就只是普通张量。loss.backward()执行完梯度就累积在每个requires_gradTrue的叶子张量上了。参数更新阶段优化器的选择会对训练产生直接影响。AdamW是目前训练大模型的主流选择它把权重衰减从L2正则化中解耦出来对Transformer这类架构更友好。学习率通常配合warmup策略一起使用前几千步从0线性增长到峰值然后按余弦曲线衰减。至于峰值的取值小模型可以用1e-3到3e-4之间大模型要降到1e-4甚至更低具体要靠实验观察loss曲线来定。2.4 损失函数与优化器交叉熵、AdamW与学习率策略交叉熵损失在语言模型里通用得几乎不需要思考但我还是建议你亲手实现一次并对比一下PyTorch内置版本的结果。实现的细节在于logits要先经过LogSoftmax再和真实标签计算负对数似然。你可能会想“这有什么好对比的”实际做一次之后你会彻底理解为什么PyTorch的CrossEntropyLoss底层把LogSoftmax和NLLLoss合并在一起了——为了数值稳定性。优化器方面AdamW和Adam的区别值得认真理解。Adam自带L2正则化的方式会让带有大梯度的参数受到更强的正则化这在某些情况下会让泛化变差。权重衰减的初衷是对所有参数施加相同的衰减比例让它只与参数本身的数值相关而AdamW就是修正这个问题的版本。实践中用AdamW替换Adam通常能带来更稳定的训练。学习率策略对新手来说是最难调的部分之一。我见过不少朋友刚开始训练loss居高不下第一反应是“模型有问题”结果只是学习率太高导致震荡或者没有warmup导致早期梯度爆炸。从零构建的早期阶段建议固定住一个标准配置AdamW学习率3e-4warmup步数占总步数的5%到10%然后观察loss曲线一次只改一个变量。这个习惯能帮你省掉大量无效实验。3. 实操记录从零训练一个ChatGPT风格小模型3.1 环境搭建与项目结构规划动手之前先把环境理清楚。我用的是Python 3.10 PyTorch 2.x CUDA 11.8的组合项目根目录下划分成data/、tokenizer/、model/、train.py、config.py这五块。很多初学者喜欢把所有代码塞进一个文件里实验一多就完全失控。从零训练模型本身已经够难了别让工程结构拖后腿。配置管理的优先级容易被低估。训练一个模型涉及到的参数非常多词表大小、嵌入维度、层数、头数、学习率、batch大小、梯度累积步数、最大序列长度、训练步数、warmup步数、随机种子……如果这些参数散落在代码各处复现实验会变成一场灾难。我用一个config.py把全部参数集中起来每个实验复制一份配置并命名好版本号。别嫌麻烦你一定会为一个“昨天还正常今天就不收敛”的模型找半天原因这时候配置归档能救你的命。环境层面如果你用的是云GPU记得提前测一下CUDA版本和PyTorch的编译版本是否匹配torch.cuda.is_available()返回False是新手最容易撞上的墙。本地RTX系列的话直接用官方推荐的pip安装命令即可。另外固定随机种子torch.manual_seed很重要否则你每次训练结果都不一样无法判断改动到底是好是坏。3.2 数据准备从原始文本到Tensor数据集数据准备这一步我以训练一个中文对话小模型为例。原始数据来源可以是开源的中文语料比如维基百科的中文dump、社区问答数据等。从零构建的阶段数据量不需要太大几千万token就够了重点是流程要完整。拿到原始文本后先做清洗和过滤然后按段落切分成训练样本再用自己训练好的分词器把每段文本编码成token序列。编码完成之后我需要把样本组装成固定长度的序列。我采用的方式是把全部token拼成一个超长的一维数组然后按max_seq_len 1切块每块的前max_seq_len个token作为输入后移一位的max_seq_len个token作为标签。这样做的好处是避免逐条padding带来的算力浪费也让模型在训练时看到更连贯的长文本上下文。数据装载方面直接使用PyTorch的Dataset和DataLoader设置好batch_size和shuffle。有一个容易忽略的细节DataLoader的num_workers可以根据CPU核数调大默认值0意味着主进程处理数据会成为训练瓶颈。GPU训练的时候数据加载往往才是那个隐性瓶颈。# 数据切块示例把长文本切成固定长度的训练样本 import torch from torch.utils.data import Dataset class TextDataset(Dataset): def __init__(self, tokens, seq_len): # tokens: 完整的编码后token ID列表 self.seq_len seq_len total (len(tokens) - 1) // seq_len self.samples [] for i in range(total): start i * seq_len x tokens[start:start seq_len] y tokens[start 1:start seq_len 1] self.samples.append((torch.tensor(x), torch.tensor(y))) def __len__(self): return len(self.samples) def __getitem__(self, idx): return self.samples[idx]3.3 训练脚本一个迷你但完整的PyTorch训练循环有了数据之后就到了最核心的部分训练循环。我直接给出一个精简版本这个结构你以后做任何模型训练都能复用。需要说明的是为了可读性我简化了部分细节实际工程中你还需要加入梯度累积、EMA、早停、checkpoint保存等逻辑。模型定义方面我建议直接用PyTorch的nn.TransformerEncoder或者从HuggingFace借用GPT2LMHeadModel的小配置先跑通流程然后再自己从零写一遍核心层。为什么这么建议因为如果一开始就完全手写出bug的概率太大了你不知道到底是“数据错了”还是“模型写错了”还是“训练逻辑错了”。先跑通再重构这是工程思维不是走捷径。import torch import torch.nn as nn from torch.utils.data import DataLoader from torch.optim import AdamW def train_step(model, batch_x, batch_y, optimizer, criterion): model.train() optimizer.zero_grad() logits model(batch_x) # [batch, seq_len, vocab_size] loss criterion(logits.view(-1, logits.size(-1)), batch_y.view(-1)) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() return loss.item() model GPTConfig.build_model() optimizer AdamW(model.parameters(), lr3e-4) criterion nn.CrossEntropyLoss(ignore_index-100) for step, (batch_x, batch_y) in enumerate(train_loader): batch_x, batch_y batch_x.to(device), batch_y.to(device) loss train_step(model, batch_x, batch_y, optimizer, criterion) if step % 500 0: print(fstep {step}, loss {loss:.4f})梯度裁剪这行代码经常被忽略但对语言模型训练来说它是稳定训练的关键。Transformer的训练早期梯度范数可能突然飙升到几百甚至上千不做裁剪的话一个step就能把参数打飞到loss变成NaN。我把max_norm设为1.0这是一个很保守但稳妥的起点。另外CrossEntropyLoss里我用ignore_index-100来屏蔽padding位的损失这是处理变长序列的通用技巧。3.4 评估与生成训练完不等于能用训练循环跑完之后模型loss降到了一个相对平稳的值但这只能说明“模型记住了训练数据里的统计规律”。要判断模型好不好最直接的办法是让它生成文本亲眼看一看。生成文本的流程很简单给模型一个起始token序列让它预测下一个token然后把这个token拼到输入末尾继续预测下一个循环往复。这个自回归生成过程就是ChatGPT回复你的基本机制。需要注意的是每步直接选概率最大的token贪心搜索会让生成结果陷入重复你可以试一下在采样时引入随机性比如temperature和top-p。def generate(model, tokenizer, prompt, max_new_tokens100, temperature0.8, top_p0.9): model.eval() ids tokenizer.encode(prompt) for _ in range(max_new_tokens): input_ids torch.tensor(ids).unsqueeze(0).to(device) with torch.no_grad(): logits model(input_ids)[0, -1, :] / temperature probs torch.softmax(logits, dim-1) sorted_probs, sorted_idx torch.sort(probs, descendingTrue) cumsum torch.cumsum(sorted_probs, dim-1) cutoff (cumsum top_p).nonzero()[0].item() sorted_probs[cutoff:] 0 sorted_probs / sorted_probs.sum() next_id sorted_idx[torch.multinomial(sorted_probs, 1)].item() ids.append(next_id) if next_id tokenizer.eos_id(): break return tokenizer.decode(ids)temperature的作用是调节概率分布的“锐利程度”值越低越倾向于选高概率token输出更保守值越高随机性越强输出更多样但可能胡言乱语。top-p则是只在累积概率达到p的候选集中采样避免从概率极低的垃圾token中抽样。这是工程上最常用的两个采样参数值得反复体会。4. 从语言模型到推理模型让模型学会“思考”4.1 为什么预训练后的模型还不会“思考”当你把一个小型语言模型训练到loss不错、能生成通顺文本之后很快会发现一个尴尬的事实它能造句但不会算数、不会推理、不会按照指令要求回答问题。比如你问它“小明有3个苹果吃了1个还剩几个”它可能会一本正经地回答“还剩3个”。原因在于预训练的目标只是预测下一个token它学到的是文本中的统计相关性而不是因果逻辑。所谓推理能力reasoning是怎么来的如果你读过大模型技术发展的脉络会发现思路已经从“让模型记住更多知识”转向了“让模型在推理时多想几步”。OpenAI的o1系列模型展示了一个事实在回答复杂问题之前让模型生成一长串内部思维过程可以显著提升正确率。这就像解数学题直接写答案容易错但如果先在草稿纸上一步步推演正确率就会高很多。从零构建一个推理模型核心思路是先用普通预训练模型做基础再通过微调和强化学习让模型学会“长思考”。这不再只是“续写文本”的问题而是“如何在答案前生成有效的思考过程”的问题。4.2 从零构建推理模型的几种主流路线目前从零构建一个reasoning model工程上比较成熟的方案有几种。最直接的路线是做监督微调SFT准备一批包含“思维链”Chain of Thought的训练数据让模型学会模仿这种思考格式。训练数据和普通SFT没有本质差别只是label部分变成了“思考过程最终答案”。再往前走一步就是用强化学习来优化推理过程。这条路线因为DeepSeek-R1的公开而被广泛讨论核心是GRPOGroup Relative Policy Optimization一种比PPO更节省显存的强化学习算法。做法是针对同一道数学题或逻辑题让当前模型生成多个回答每个回答算出一个奖励分数然后以“这组回答内的相对优劣”作为信号更新模型参数让平均分更高的回答被强化。4.3 一个简化的RL训练流程示例这里我必须先说清楚完整的GRPO实现涉及策略模型、参考模型、奖励模型、KL散度约束等多个组件不是一个小工程。但为了让你理解核心逻辑我给你一个精简的思路把一次训练迭代拆成三步。第一步采样阶段给定一组prompt让策略模型当前正在训练的模型生成N个回答。第二步评估阶段用一个规则型奖励函数给每个回答打分。推理任务的最大好处是可以用规则的确定性来评估比如数学题可以比对最终答案是否正确、格式是否符合要求。第三步更新阶段计算这组回答中每个回答的奖励相对值用策略梯度方法更新模型注意需要同时约束模型参数不要偏离原始模型太远防止它为了拿奖励而输出胡乱内容。# GRPO中计算相对优势值的简化示意 def compute_advantages(rewards, group_size): import torch rewards torch.tensor(rewards, dtypetorch.float) mean rewards.mean() std rewards.std() 1e-4 advantages (rewards - mean) / std return advantages你可能已经注意到这整条路线和“预测下一个token”的语言模型训练很不一样。它不再只是模仿数据而是一种“在环境里试错、根据奖励调整行为”的学习方式。从零构建的后期如果能把这条路线跑通你基本上就掌握了当前大模型领域最前沿的一整套训练方法。5. 常见问题与避坑技巧实录5.1 训练不收敛先把loss曲线拆开看训练不收敛是最常见的问题但它可能是好几种原因叠加造成的。我的排查顺序向来很固定先看loss是不是完全静止再看是不是在震荡然后看是不是NaN。完全静止通常是学习率太小或者模型本身有问题震荡通常是学习率太大或者batch太小NaN则复杂得多可能是数据里有NaN、学习率过高、梯度爆炸、或者层归一化数值不稳定。一个特别容易被忽略的坑是数据问题。有些文本清洗不干净里面夹杂着异常字符经分词器处理后会产生巨大数值。我自己遇到过一种情况文本编码后长度超过预设max_seq_len截断后出现了没有任何有效信息的长段padding导致模型学习效率极低。排查方法是先单独写脚本做数据检查统计token分布、样本长度分布、空样本数量确认数据本身没问题再怀疑模型。5.2 显存不够梯度累积、混合精度与序列长度取舍显存不够几乎是每个人都会撞上的墙。最直接的缓解方案是减小batch_size但batch太小会导致训练不稳定。替代方案是梯度累积每过几个小batch才做一次参数更新等效于增大了batch_size同时控制住显存峰值。我经常用gradient_accumulation_steps4让“每4个batch更新一次参数”。混合精度训练AMP是另一个立竿见影的优化手段。PyTorch的torch.cuda.amp可以自动把部分计算变成FP16精度显存占用直接砍半同时通过梯度scaler保证训练精度不受影响。在RTX 30系及以上显卡上AMP几乎是必开的选项。另外序列长度也不是越长越好如果显存紧张可以先从256的序列长度开始训练验证思路后再逐步拉长。5.3 数据质量陷阱模型的“上限”是数据决定的训练模型的时间长了你会越来越深刻地体会到一句话垃圾进垃圾出。模型的参数只是从数据里提取规律的工具数据里没有的信息模型不可能无中生有。我在数据集上栽过最大的跟头是“不知道样本重复率有多高”。网上爬的语料经过简单去重后我以为干净了结果训练出来的模型总是复读某些固定段落。后来用MinHash做了一遍去重才发现数据里存在大量同源改写重复度远超预期。数据质量还体现在多样性和分布上。如果你的语料全是技术文档那模型生成任何内容都会带着技术腔如果语料里网络用语扎堆生成内容就会轻浮且缺乏事实性。从零构建实验时建议先人工浏览几万个样本建立对数据分布的直觉。这个工作虽然枯燥但对后续所有训练实验都有指导意义。5.4 书籍与学习资源如何系统补充知识文章最后补一个学习资源板块。很多人会问“从零开始到底该读什么书”我的建议是如果只读一本选《Build a Large Language Model (From Scratch)》——注意要去官网或正规书店购买正版电子书不要贪图网盘里的盗版资源。这本书几乎就是冲着“ai-engineering-from-scratch”这个标题写的从数据准备、注意力机制、预训练、微调一直讲到部署代码全部配套可跑适合放在手边边看边敲。再推荐几篇值得反复读的论文《Attention Is All You Need》是Transformer的原始论文必读DeepSeek-R1的technical report是理解推理模型强化学习路线的最佳材料之一PPO和GRPO的原始论文用于理解RL训练细节。论文阅读有门槛但配合实操读一遍收获会非常大。我个人在实际操作中的体会是从零构建AI工程真正的分水岭不在你是否能跑通代码而在你是否能回答出每一个“为什么”。代码跑通只是“知其然”能把训练曲线的每一个波动都解释清楚、能在出问题时快速定位到数据、模型或训练逻辑中的具体环节这才是“知其所以然”。这种能力没有任何捷径只能靠一次一次的实验、一次一次的踩坑、一次一次的复盘堆出来。希望这篇内容能帮你少走一些弯路也欢迎你在评论区聊聊自己从零构建过程中踩过的坑我们互相学习。