资讯详情

从零手搓AI工程:深入理解Transformer、Tokenizer与推理优化

📅 2026/11/3 13:03:00 | 华诺云谱 👁 阅读
从零手搓AI工程:深入理解Transformer、Tokenizer与推理优化
1. 这个项目到底在解决什么问题第一次看到 ai-engineering-from-scratch 这个标题我脑子里蹦出来的第一个念头是又是一个教你怎么调库的教程仓库但翻了一圈之后发现它想做的事情比调包要硬核得多——它试图回答一个被大多数人跳过的问题当你把 transformers、langchain、pytorch 这些库全部拿掉之后你还剩多少真本事这个项目本质上是一套从零手搓 AI 工程的学习路径集合。它不满足于让你会model.fit()而是要求你理解 attention 是怎么算出来的、tokenizer 是怎么把一句话切碎的、反向传播的梯度是怎么一层层传回去的、推理服务是怎么把显存榨干的。适合谁看适合那些已经会用框架、但心里总觉得我好像只是个大号 API 调用员的工程师也适合刚入门但不想被框架惯坏的新人。我自己的判断是这类内容的价值不在于让你真的在生产环境里手写一个 transformer而在于当框架出问题、当性能不达标、当你要做别人没做过的定制时你能知道该往哪个方向挖。框架是别人抽象好的东西抽象就意味着有信息丢失而from scratch就是把这些丢失的信息补回来。下面我按自己的理解把这个项目涉及的核心内容拆成几个部分来讲包括它为什么这么设计、每个环节的关键点在哪、实操时容易踩什么坑以及我自己在类似路径上的一些经验。2. 为什么从零实现比直接调库更值得投入2.1 框架抽象带来的认知盲区先说一个我观察到的现象。现在很多做 AI 应用的人工作流是这样的找一个开源模型套一个推理框架写个 prompt接个 API上线。这套流程能跑通但一旦遇到下面这些情况就抓瞎了模型输出不符合预期不知道是数据问题、tokenizer 问题还是采样参数问题推理速度慢不知道该优化哪一层想改模型结构做定制发现连 forward 函数里张量的形状变化都理不清显存爆了不知道是 KV cache 还是中间激活占的这些问题的共同点是它们都发生在框架抽象的缝隙里。你调model.generate()的时候框架帮你做了几十件事但你不知道是哪一件出了问题。从零实现一遍就是把这些黑盒一个个拆开看。2.2 从零实现的三层价值我把这个项目的价值分成三层从浅到深第一层是祛魅。很多人觉得 transformer 很神秘其实把 attention 的矩阵运算手推一遍你会发现它就是几个矩阵乘法加一个 softmax没什么魔法。这种祛魅能极大降低你的心理门槛。第二层是可控。当你自己写过一遍 tokenizer你就知道 BPE 的合并规则是怎么定的遇到中文、emoji、特殊符号时该怎么处理。当你自己写过一遍采样你就知道 temperature、top-k、top-p 到底在做什么而不是盲目调参。第三层是可改。这是最高价值。当你要做一个别人没做过的模型结构改动时如果你只有调库经验你连从哪下手都不知道但如果你手搓过你知道改哪一层、影响哪些张量、需要重新训练哪些参数。2.3 一个反直觉的观点这里我要说一个可能有点反直觉的观点从零实现的目的不是让你以后都手写而是让你以后调库时调得更明白。我见过一些人学完从零实现之后就走火入魔了什么都要自己写结果项目进度慢得要死。这是走偏了。正确的姿势是你知道框架内部在干什么所以你调库的时候知道哪些参数重要、哪些坑要避、出问题往哪查。框架依然是你的主力工具只是你不再是个盲人。3. 核心模块拆解一个 AI 工程从零实现要过哪几关3.1 第一关Tokenizer——把文字变成数字这是整个流程的入口也是最容易被忽视的一环。很多人觉得 tokenizer 就是个工具调一下就行。但实际上tokenizer 的设计直接决定了模型的很多行为。从零实现一个 tokenizer核心是理解BPEByte Pair Encoding或者WordPiece的合并逻辑。BPE 的思路很朴素从字符级别开始统计相邻字符对的出现频率把最高频的对合并成一个新 token重复这个过程直到达到目标词表大小。我自己手写过一遍 BPE踩过的坑包括中文和英文的处理差异。英文按字符切分很自然但中文如果按字符切词表会爆炸如果按词切又需要分词工具。实践中常见做法是对中文按字符处理因为中文字符本身信息密度高。特殊 token 的处理。[PAD]、[UNK]、[CLS]这些特殊 token 必须在训练 BPE 之前就预留好位置否则后面会乱。词表大小的权衡。词表太小序列会很长推理慢词表太大embedding 层参数多小数据集上容易过拟合。常见的选择是 32k 到 50k 之间。提示手写 tokenizer 时建议先用一个小语料比如几千行跑通流程确认合并逻辑正确后再上大规模语料。直接上大语料调试一次训练几十分钟效率极低。3.2 第二关Attention——transformer 的心脏Attention 是 transformer 的核心也是从零实现里最值得花时间的地方。它的数学形式很简单Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V但真正手写的时候有几个细节必须搞清楚为什么要除以 sqrt(d_k)因为当 d_k 很大时QK^T 的点积结果方差会变大导致 softmax 的输入进入饱和区梯度接近零。除以 sqrt(d_k) 是把方差拉回到 1 附近保证梯度健康。这个细节如果你只调库永远不会注意到。Multi-Head 是怎么切的不是把 Q、K、V 各切成几份分别算而是把 d_model 维度切成 num_heads 份每份独立做 attention最后 concat 起来再过一个线性层。切的时候要注意 reshape 和 transpose 的顺序这里张量形状最容易搞错。Mask 怎么加Decoder 里的 causal mask 是一个上三角为负无穷的矩阵加到 attention score 上softmax 之后这些位置就变成 0。手写的时候要注意 mask 的形状要能 broadcast 到 (batch, heads, seq_len, seq_len)。我自己的经验是手写 attention 时每一步都打印张量形状。这是最笨但最有效的方法。形状对了逻辑基本就对了形状错了后面全是错的。3.3 第三关训练循环——梯度是怎么流的从零实现训练循环核心是理解这几件事前向传播输入经过每一层得到输出和 loss反向传播loss 对每个参数的梯度通过链式法则一层层传回去参数更新用优化器SGD、Adam 等根据梯度更新参数如果你用 PyTorchloss.backward()和optimizer.step()两行就搞定了。但如果你想真正理解可以尝试用 numpy 手写一个简单的反向传播。我试过用 numpy 实现一个两层 MLP 的训练写完之后对梯度的理解上了一个台阶。关键点在于梯度消失和梯度爆炸不是抽象概念而是你能亲眼看到的现象。当你手写反向传播打印每一层的梯度范数你会看到深层网络的梯度确实在变小这时候你就理解为什么需要残差连接、为什么需要 LayerNorm、为什么需要好的初始化。3.4 第四关推理优化——把显存和速度榨干训练和推理是两回事。训练时你关心的是收敛推理时你关心的是延迟和吞吐。从零实现推理核心是理解KV Cache。KV Cache 的思路是自回归生成时每次生成一个新 token只需要计算这个新 token 的 Q而 K 和 V 可以复用之前所有 token 的。这样就把每次生成的计算量从 O(n^2) 降到 O(n)。但 KV Cache 也带来显存问题。序列越长KV Cache 占的显存越多。一个 7B 模型序列长度 4096batch size 1KV Cache 大概占 1-2GB。这就是为什么长上下文推理这么吃显存。我实测下来推理优化里性价比最高的几个手段是优化手段效果实现难度KV Cache显著降低计算量中量化INT8/INT4降低显存占用中Batch 推理提高吞吐低Flash Attention降低显存加速高通常用现成实现注意量化不是万能的。INT4 量化在小模型上可能掉点明显在大模型上相对安全。做量化前一定要在自己的任务上评估效果不要盲目上。4. 实操路径我是怎么安排学习顺序的4.1 不要一上来就啃 transformer很多人一上来就想手写一个完整的 GPT结果卡在 attention 那里就放弃了。我的建议是分阶段推进每个阶段都有可运行的产出。阶段一基础工具链。先手写一个简单的 tokenizerBPE再手写一个线性回归的训练循环用 numpy。这两个任务都很小但能让你把数据→模型→loss→梯度→更新这条链路走通。阶段二核心组件。手写 attention、LayerNorm、FeedForward、残差连接。每个组件单独测试确认输入输出形状和数值正确。阶段三组装模型。把组件拼成一个完整的 transformer block再堆成多层。这时候可以跑一个小规模的语言模型训练比如在莎士比亚文本上训练一个字符级模型。阶段四推理优化。实现 KV Cache、采样策略temperature、top-k、top-p、batch 推理。阶段五工程化。模型保存加载、服务化部署、性能监控。这个顺序的好处是每一步都有正反馈不会因为目标太大而中途放弃。4.2 每个阶段的关键检查点我在每个阶段都会设一个检查点通过了才往下走阶段一检查点tokenizer 能把一句话正确编码解码训练循环能让 loss 下降阶段二检查点attention 的输出和 PyTorch 官方实现数值接近误差 1e-5 以内阶段三检查点小模型能过拟合一个小数据集loss 降到接近 0阶段四检查点KV Cache 版本的输出和朴素版本一致速度有提升阶段五检查点模型能保存后重新加载推理结果一致这些检查点看起来简单但能帮你快速定位问题。比如阶段三过拟合不了说明模型或训练循环有问题阶段四输出不一致说明 KV Cache 实现有 bug。4.3 一个具体的实操记录我拿阶段三举例说一下我实际跑的过程。数据集用的是莎士比亚全集大概 1MB 文本。字符级 tokenizer词表大概 65 个字符。模型配置4 层 transformerd_model128num_heads4context_length128。训练参数batch_size32learning_rate3e-4Adam 优化器训练 5000 步。第一次跑loss 从 4.2 降到 2.1 就卡住了。排查发现是学习率太大改成 1e-4 后 loss 降到 1.5。继续排查发现是 LayerNorm 的位置放错了应该用 Pre-LN 而不是 Post-LN改过来之后 loss 降到 1.2。最后生成的文本大概是这样的ROMEO: What sayest thou, my lord? JULIET: I know not what to say.虽然语法还有点问题但已经能看出莎士比亚的风格了。这个过程让我对 transformer 的训练动态有了直观感受。5. 常见问题与排查技巧实录5.1 训练不收敛怎么办这是最常见的问题。我的排查顺序是先检查数据。打印几条样本确认输入输出对应正确。我遇到过一次数据加载时 shuffle 和 label 没对齐导致模型学了个寂寞。再检查 loss。loss 一开始应该是 ln(vocab_size) 左右。如果一开始就很小说明模型可能作弊了比如 label 泄漏。然后检查学习率。太大 loss 震荡太小 loss 下降慢。可以先用一个很小的学习率跑几百步确认 loss 在下降再逐步调大。最后检查模型结构。用一个小数据集比如 100 条让模型过拟合如果过拟合不了说明模型结构有问题。5.2 显存不够怎么办显存问题分训练和推理两种情况。训练时显存不够可以减小 batch size最直接用梯度累积模拟大 batch用混合精度训练fp16/bf16用 gradient checkpointing用时间换显存推理时显存不够可以量化模型减小 batch size限制序列长度用 PagedAttention 之类的技术如果框架支持我自己的经验是先定位显存被什么占了。用torch.cuda.memory_summary()能看到显存分配情况。有时候是中间激活占的有时候是优化器状态占的有时候是 KV Cache 占的不同原因对应不同解法。5.3 推理结果和训练时不一致这个问题通常有几个原因Dropout 没关。推理时要model.eval()否则 dropout 还在随机丢弃。LayerNorm 的统计量问题。训练时用 batch 统计量推理时用滑动平均如果实现不对会不一致。位置编码问题。如果位置编码是动态生成的推理时序列长度和训练时不同可能出问题。KV Cache 实现 bug。这是最常见的建议先用小序列对比 KV Cache 版本和朴素版本的输出。5.4 常见问题速查表问题现象可能原因排查方法loss 不下降学习率、数据、模型结构小数据集过拟合测试loss 震荡学习率太大、batch 太小降低学习率、增大 batch显存爆了batch 太大、序列太长、KV Cachememory_summary 定位推理慢没用 KV Cache、没量化对比不同配置的延迟输出乱码tokenizer 问题、解码问题检查 encode/decode 一致性梯度爆炸学习率太大、没做梯度裁剪加 gradient clipping6. 我踩过的几个坑和一点个人体会6.1 不要追求一次写对我一开始手写 attention 的时候总想一次写对结果调了两天才跑通。后来我改变策略先写一个最笨的版本跑通了再优化。比如 attention 先写一个 for 循环版本确认逻辑对了再改成矩阵运算。这样虽然一开始慢但总体效率高得多。6.2 数值稳定性是个大坑手写 softmax 的时候如果不做数值稳定处理减去最大值遇到大数会溢出。手写 LayerNorm 的时候如果方差算错会出 NaN。这些细节在调库时都被框架处理好了但手写时必须自己注意。我的做法是每个涉及指数、除法、开方的操作都检查一遍数值范围。softmax 减最大值、LayerNorm 加 epsilon、除法检查分母这些都是基本操作。6.3 测试驱动开发很有用手写每个组件时我都会先写一个测试给定输入期望输出是什么。比如 attention 的测试是输入全零输出应该也是全零输入相同输出应该相同。这些测试能帮你快速定位问题。6.4 最后分享一个小技巧如果你觉得从零实现太慢可以用一个折中方案用 PyTorch 的底层 API但不用高层封装。比如用torch.nn.functional而不是torch.nn用torch.autograd手动算梯度而不是loss.backward()。这样既能理解底层又不用完全从 numpy 开始。这个项目后续还可以这样扩展加入分布式训练的从零实现、加入 RLHF 的从零实现、加入多模态模型的从零实现。每一个方向都是一个大坑但每填一个坑你对 AI 工程的理解就深一层。我个人在实际操作中的体会是从零实现最大的收获不是那些代码而是一种我知道底下发生了什么的底气。这种底气让你在面对新框架、新模型、新问题时不会慌。
📝

华诺云谱内容团队

资深建站顾问 · 行业研究员

10年+企业数字化服务经验,专注智能建站、SEO优化与品牌营销,持续输出建站技巧、行业洞察与营销干货,已帮助5000+企业实现数字化增长。

你可能需要的服务

订阅华诺云谱资讯周报

每周一封,精选建站技巧、SEO与营销干货,直达邮箱。已有 8,000+ 企业主订阅,助你少走弯路。

↑