ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

从零构建AI工程:手写反向传播、Attention与推理优化实战

2026/10/4 13:49:46 拓冰建站 浏览量
从零构建AI工程:手写反向传播、Attention与推理优化实战 1. 这个项目到底在解决什么问题第一次看到 ai-engineering-from-scratch 这个标题我脑子里蹦出来的第一个念头是又是一个教人调包的教程但仔细琢磨了一下 from scratch 这四个字我意识到它想做的事情可能完全不一样。市面上讲 AI 工程的内容绝大多数都是从pip install transformers开始然后from_pretrained加载一个模型跑个 demo 就结束了。这种内容看的时候很爽但看完之后你脑子里其实什么都没留下——你不知道 attention 到底怎么算的不知道反向传播的梯度是怎么流动的更不知道一个推理服务在生产环境里为什么会突然延迟飙升。ai-engineering-from-scratch这个项目要解决的核心问题就是把 AI 工程从调包拉回到理解的层面。它面向的不是已经在大厂做模型训练的老手而是那些会用 PyTorch 但说不清楚 autograd 原理、能跑通 LoRA 微调但不知道秩分解为什么有效、部署过推理服务但遇到 OOM 只会重启的工程师。说白了它想让你从AI 使用者变成AI 建造者。我自己带过不少刚入行的同学最常见的困境就是面试的时候被问到手写一个 self-attention写不出来线上服务挂了看日志一脸懵想优化推理速度除了换更小的模型之外毫无头绪。这些问题的根源都是一样的——你跳过了从零构建的那一步。而这个项目标题里的 from scratch恰恰就是冲着这个痛点来的。这篇文章我会从项目设计的整体思路讲起然后拆解它涉及的核心技术点接着给出可以照着做的实操路径最后分享一些我在实际动手过程中踩过的坑和总结出来的排查技巧。不管你是刚转行做 AI 的新人还是已经工作几年但想补基础的工程师应该都能从中找到对自己有用的东西。2. 项目整体设计与思路拆解2.1 为什么从零实现比调包更有价值先说说为什么我认同 from scratch 这个方向。很多人觉得现在框架这么成熟何必重复造轮子这个想法在工程效率上没错但在能力建设上是大错特错的。我给你打个比方你会开车不代表你懂发动机。平时通勤当然没问题但车在高速上抛锚了你连引擎盖都不敢打开那就只能等救援。AI 工程也是一样调包能让你完成 80% 的日常工作但剩下 20% 的关键时刻——性能调优、bug 排查、架构设计——靠的全是你对底层原理的理解。从零实现的价值体现在三个层面。第一是调试能力当你自己手写过一遍反向传播再遇到梯度消失或者梯度爆炸的时候你就知道该去看哪一层的激活值而不是盲目地调学习率。第二是优化能力你手写过 attention 的矩阵运算就会明白为什么 flash attention 要做分块计算为什么 KV cache 能省显存这些直觉是调包调不出来的。第三是创新能力所有的架构创新本质上都是对底层组件的重新组合你只有把组件拆开看过才有可能组装出新的东西。提示从零实现不等于拒绝框架。正确的姿势是先用框架跑通再自己手写一遍对照最后回到框架但带着理解去用。这个循环走一遍效果比看十篇论文都好。2.2 项目的知识体系是怎么组织的一个合格的 ai-engineering-from-scratch 项目它的知识体系应该是有层次的不能一上来就让你手写 Transformer那样只会劝退。我理解它的组织逻辑应该是这样的从最基础的数学运算开始到单个神经元到多层网络到序列模型最后到完整的 Transformer 和大模型推理。每一层都建立在前一层的基础上每一层都有可以运行的代码。具体来说我把它分成四个阶段。第一阶段是数值基础矩阵乘法、softmax、交叉熵、梯度计算这些是后面所有内容的砖块。第二阶段是神经网络核心前向传播、反向传播、优化器、正则化这里要手写一个能训练 MNIST 的完整网络。第三阶段是序列与注意力RNN、LSTM、attention 机制、Transformer block这是理解现代大模型的必经之路。第四阶段是工程化模型量化、推理优化、服务部署、监控把前面学的东西真正跑在生产环境里。这个分层的妙处在于每一层你都可以独立验证。比如第一阶段学完你可以手写一个矩阵乘法然后跟 NumPy 的结果对比误差在 1e-6 以内就说明对了。这种即时反馈对学习来说太重要了比看视频看半天不知道自己对不对强太多。2.3 技术选型背后的考量既然是 from scratch语言和工具的选择就很关键。我的建议是Python NumPy 起步然后过渡到 PyTorch。为什么不用纯 Python 列表因为列表做矩阵运算太慢了你训练一个小网络要等半天学习热情直接被磨没了。NumPy 的好处是它足够底层你能看到每一个数组的形状变化同时它的向量化运算又足够快能在几秒内完成一次前向传播。为什么不直接上 PyTorch因为 PyTorch 的 autograd 太方便了你loss.backward()一调梯度就自动算好了你根本不会去想链式法则到底怎么展开的。所以正确的路径是先用 NumPy 手写前向和反向把梯度算清楚然后再用 PyTorch 的 autograd 去验证你的实现。这样你既理解了原理又掌握了工具。到了推理优化阶段工具链会扩展到ONNX、TensorRT、vLLM这些。但注意这些工具的使用前提是你已经理解了模型的计算图否则你连哪里是瓶颈都找不到。我见过太多人上来就装 TensorRT结果模型转换失败报了一堆看不懂的错最后只能放弃。根因就是跳过了理解计算图这一步。3. 核心细节解析与实操要点3.1 手写反向传播最容易出错的地方反向传播是整个 from scratch 过程中最容易翻车的环节。我见过太多人前向传播写得漂漂亮亮一到反向就各种维度对不上、梯度爆炸。这里我把关键点拆开讲。首先是计算图的构建。你在做前向传播的时候其实是在隐式地构建一张计算图。比如z W x b这个操作在计算图里就是一个节点它的输入是 W、x、b输出是 z。反向传播就是沿着这张图从后往前用链式法则把梯度一层层传回去。手写的时候你需要为每个操作记录它的输入和输出这样才能在反向的时候用上。其次是梯度的形状。这是新手最容易搞混的地方。记住一个原则梯度的形状永远和对应变量的形状一致。W 是 (m, n)那 dW 也必须是 (m, n)。如果你算出来的 dW 是 (n, m)那大概率是转置搞反了。我自己的习惯是每算完一个梯度就打印一下形状跟变量对比不一致立刻停下来查。再就是数值稳定性。softmax 和交叉熵放在一起的时候如果你先算 softmax 再算 log很容易出现 log(0) 的情况导致 inf 或者 nan。正确的做法是把 softmax 和交叉熵合并成一个操作用 log-sum-exp 技巧来算。这个技巧在数学上等价但数值上稳定得多。我实测过不合并的话训练到一半 loss 突然变成 nan 是家常便饭。注意手写反向传播的时候一定要做梯度检查gradient check。方法很简单用数值微分算一个近似梯度跟你反向传播算出来的解析梯度对比相对误差在 1e-5 以内就算通过。这个检查能帮你抓出 90% 的实现 bug。3.2 Attention 机制的实现细节Attention 是理解现代大模型的钥匙手写一遍绝对值得。但这里面有几个细节不看清楚很容易写错。第一个是缩放因子。标准的 scaled dot-product attention 里Q 和 K 点积之后要除以 sqrt(d_k)其中 d_k 是 key 的维度。为什么要除因为当 d_k 很大的时候点积的结果会变得很大softmax 之后会变得非常尖锐梯度就消失了。除以 sqrt(d_k) 能把方差拉回到 1 左右保证 softmax 的输出比较平滑。这个细节很多人写的时候会漏掉结果训练不收敛还找不到原因。第二个是mask 的处理。在 decoder 里你需要一个 causal mask让每个位置只能看到它前面的位置。实现的时候是在 softmax 之前把不该看的位置加上一个很大的负数比如 -1e9这样 softmax 之后这些位置的权重就接近 0 了。注意是加负数不是直接置零因为 softmax 对 0 也会分配权重。第三个是多头注意力的拆分。多头不是简单地把多个 attention 拼起来而是把 d_model 维度的向量拆成 h 个头每个头维度是 d_model/h各自做 attention最后再拼回去。拆的时候要注意 reshape 和 transpose 的顺序顺序错了结果就完全不对。我建议你写完之后用一个小的随机输入手动算一遍单头的结果跟多头拆出来的第一个头对比能对上就说明拆分逻辑对了。3.3 训练循环里的那些坑训练循环看起来简单无非就是前向、算 loss、反向、更新参数。但真正写起来坑多得很。学习率的选择是最关键的。太大不收敛太小收敛慢。我的经验是从 1e-3 开始试如果 loss 震荡就降到 1e-4如果 loss 下降太慢就升到 3e-3。但更靠谱的做法是用学习率预热warmup加余弦退火前期用小学习率慢慢升上去后期再降下来。这个策略在 Transformer 训练里几乎是标配。批次大小的选择也有讲究。批次太小梯度噪声大训练不稳定批次太大显存吃不下而且泛化性能可能变差。一般从 32 或 64 开始试。如果显存不够可以用梯度累积攒几个小批次再更新一次参数效果等价于大批次。参数初始化经常被忽略但它对训练的影响巨大。全零初始化会让所有神经元学到一样的东西对称性破缺不了。全用大的随机数又会导致梯度爆炸。正确的做法是用 Xavier 或者 Kaiming 初始化根据层的类型和激活函数来调整方差。我自己的习惯是如果训练一开始 loss 就是 nan第一件事就是检查初始化。过拟合的判断也很重要。训练 loss 一直降但验证 loss 开始升那就是过拟合了。这时候可以加 dropout、weight decay或者做数据增强。但要注意dropout 在推理的时候要关掉否则结果会不稳定。3.4 推理优化的核心手段模型训练好了怎么让它跑得快、跑得省这是工程化的重头戏。我总结下来有四个方向。量化是最直接的。把 FP32 的权重和激活值降到 FP16 或者 INT8显存占用直接减半甚至减到四分之一速度也能提升。但量化会带来精度损失需要做校准calibration来找到合适的缩放因子。INT8 量化在视觉模型上通常没问题但在大语言模型上要小心某些层对精度很敏感。算子融合是另一个大头。比如把 LayerNorm 和后面的线性层融合成一个算子减少内存读写次数。这个在推理框架里通常是自动做的但你得知道它在做什么才能判断性能瓶颈在哪。KV Cache是大语言模型推理的必备优化。自回归生成的时候每次都要重新计算前面所有 token 的 attention太浪费了。KV Cache 把前面算过的 key 和 value 缓存起来每次只算新 token 的速度能提升好几倍。代价是显存占用增加需要根据序列长度和批次大小来权衡。批处理能显著提升吞吐量。把多个请求攒在一起做推理GPU 利用率能上去。但批处理会增加延迟因为要等请求攒够。所以在线服务通常用连续批处理continuous batching动态地把新请求加进来兼顾吞吐和延迟。4. 实操过程与核心环节实现4.1 环境搭建与依赖管理动手之前环境得先弄利索。我的建议是用 conda 建一个独立环境Python 版本选 3.10 或 3.11太新太旧都容易出兼容问题。conda create -n ai-scratch python3.11 conda activate ai-scratch pip install numpy matplotlib jupyter pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu注意这里我先装了 CPU 版的 PyTorch因为前期手写 NumPy 实现用不到 GPU等到了训练大一点的模型再换 GPU 版。这样能避免一开始就被 CUDA 版本问题卡住。依赖管理我强烈建议用requirements.txt或者pyproject.toml固定版本。AI 领域的库更新太快今天能跑的代码明天可能就报错了。把版本锁死能省掉很多莫名其妙的调试时间。提示如果你用的是 Apple Silicon 的 MacPyTorch 的 MPS 后端可以用但有些算子支持不全。遇到报错就切回 CPU别在这上面浪费时间。4.2 从矩阵运算到单层网络第一步手写矩阵乘法。别小看这个它是后面所有内容的基础。import numpy as np def matmul(A, B): assert A.shape[1] B.shape[0], 维度不匹配 M, K A.shape K2, N B.shape C np.zeros((M, N)) for i in range(M): for j in range(N): for k in range(K): C[i, j] A[i, k] * B[k, j] return C这个三重循环版本慢得离谱但它的价值在于让你看清楚矩阵乘法到底在算什么。写完之后用np.allclose(matmul(A, B), A B)验证一下通过了就说明逻辑对了。然后你可以用向量化的方式重写把内层循环去掉速度能快几十倍。接下来是单个神经元。一个神经元就是y activation(W x b)。手写的时候要注意x 是列向量W 是行数等于输出维度、列数等于输入维度的矩阵。前向算完反向要算三个梯度dW、db、dx。dW 等于上游梯度乘 x 的转置db 等于上游梯度dx 等于 W 的转置乘上游梯度。这三个公式背下来后面写多层网络就是套娃。4.3 搭建并训练一个完整网络有了单层的基础搭一个多层网络就是堆叠。我建议从两层开始输入层到隐藏层用 ReLU隐藏层到输出层用 softmax。损失函数用交叉熵。class TwoLayerNet: def __init__(self, input_dim, hidden_dim, output_dim): self.W1 np.random.randn(input_dim, hidden_dim) * np.sqrt(2.0 / input_dim) self.b1 np.zeros(hidden_dim) self.W2 np.random.randn(hidden_dim, output_dim) * np.sqrt(2.0 / hidden_dim) self.b2 np.zeros(output_dim) def forward(self, X): self.X X self.h np.maximum(0, X self.W1 self.b1) scores self.h self.W2 self.b2 scores - scores.max(axis1, keepdimsTrue) exp_scores np.exp(scores) self.probs exp_scores / exp_scores.sum(axis1, keepdimsTrue) return self.probs def backward(self, y): N y.shape[0] dscores self.probs.copy() dscores[range(N), y] - 1 dscores / N self.dW2 self.h.T dscores self.db2 dscores.sum(axis0) dh dscores self.W2.T dh[self.h 0] 0 self.dW1 self.X.T dh self.db1 dh.sum(axis0)这段代码里有个细节值得说scores - scores.max(axis1, keepdimsTrue)这一步是数值稳定性的处理防止 exp 溢出。很多人写 softmax 不加这一步训练小数据没问题数据一大就 nan 了。训练循环就是反复调用 forward 和 backward然后用梯度下降更新参数。学习率从 1e-3 开始跑个几百轮看 loss 是不是稳定下降。如果 loss 不降先检查梯度对不对如果 loss 降了但准确率不涨检查数据标签有没有搞错。4.4 实现一个完整的 Transformer Block到了这一步你已经具备了实现 Transformer 的所有前置知识。一个 Transformer block 包含多头注意力、前馈网络、残差连接和 LayerNorm。class TransformerBlock: def __init__(self, d_model, n_heads, d_ff): self.d_model d_model self.n_heads n_heads self.d_k d_model // n_heads self.W_q np.random.randn(d_model, d_model) * 0.02 self.W_k np.random.randn(d_model, d_model) * 0.02 self.W_v np.random.randn(d_model, d_model) * 0.02 self.W_o np.random.randn(d_model, d_model) * 0.02 self.W_ff1 np.random.randn(d_model, d_ff) * 0.02 self.W_ff2 np.random.randn(d_ff, d_model) * 0.02 def attention(self, X): Q X self.W_q K X self.W_k V X self.W_v batch, seq_len, _ Q.shape Q Q.reshape(batch, seq_len, self.n_heads, self.d_k).transpose(0, 2, 1, 3) K K.reshape(batch, seq_len, self.n_heads, self.d_k).transpose(0, 2, 1, 3) V V.reshape(batch, seq_len, self.n_heads, self.d_k).transpose(0, 2, 1, 3) scores Q K.transpose(0, 1, 3, 2) / np.sqrt(self.d_k) mask np.triu(np.ones((seq_len, seq_len)), k1).astype(bool) scores np.where(mask, -1e9, scores) scores - scores.max(axis-1, keepdimsTrue) weights np.exp(scores) / np.exp(scores).sum(axis-1, keepdimsTrue) out weights V out out.transpose(0, 2, 1, 3).reshape(batch, seq_len, self.d_model) return out self.W_o这里面的 reshape 和 transpose 顺序是最容易错的。我的经验是写完之后用一个 batch1、seq_len2 的小输入手动算一遍单头的结果跟多头拆出来的第一个头对比。能对上说明拆分和合并的逻辑就对了。前馈网络就是两个线性层夹一个激活函数通常中间维度是 d_model 的 4 倍。残差连接是把输入加到输出上LayerNorm 是对最后一维做归一化。这两个组件看起来简单但它们是深层网络能训练起来的关键。没有残差梯度传不了几层就消失了没有 LayerNorm每层的输入分布一直在变训练很难稳定。4.5 推理服务的部署与监控模型训练好了最后一步是让它能对外提供服务。最简单的做法是用 FastAPI 包一层。from fastapi import FastAPI import numpy as np app FastAPI() model load_model() app.post(/predict) def predict(input_data: dict): x np.array(input_data[input]) probs model.forward(x) return {probs: probs.tolist()}但生产环境要考虑的东西多得多。并发处理是第一位的单个请求一个线程太浪费要用异步或者线程池。批处理能提升吞吐但要控制等待时间。超时和重试机制必须有否则一个慢请求能把整个服务拖垮。监控更是不能少要记录每个请求的延迟、显存占用、错误率出问题的时候才能快速定位。我自己的经验是推理服务的瓶颈通常不在计算而在内存拷贝和数据预处理。GPU 算得再快数据在 CPU 和 GPU 之间来回搬时间全耗在传输上了。所以优化的时候先看数据管道再看计算本身。5. 常见问题与排查技巧实录5.1 训练不收敛的排查清单训练不收敛是最高频的问题我整理了一个排查顺序按这个顺序走基本能定位到原因。排查项检查方法常见原因梯度是否正确梯度检查对比数值微分反向传播公式写错初始化是否合理打印初始 loss全零初始化或方差过大学习率是否合适尝试 1e-4 到 1e-2太大震荡太小不降数据是否有问题检查标签和输入对应标签错位或归一化遗漏数值是否稳定检查是否有 nan/infsoftmax 溢出或除零梯度检查是我最推荐的手段。具体做法是对某个参数用(loss(weps) - loss(w-eps)) / (2*eps)算数值梯度跟你反向传播算的解析梯度对比。相对误差小于 1e-5 就说明反向实现是对的。这个检查能抓出绝大多数实现 bug。5.2 显存不够用的应对策略显存不够是另一个高频问题尤其是想在自己机器上跑大模型的时候。我按性价比排个序。第一招是减小批次大小。这是最简单的批次减半显存基本减半。代价是训练速度可能变慢但至少能跑起来。第二招是梯度累积。小批次跑几次把梯度攒起来再更新一次效果等价于大批次。这样显存占用小但训练效果不打折。第三招是混合精度训练。用 FP16 存激活值和梯度FP32 存模型参数。显存能省一半左右速度还能提升。PyTorch 里用torch.cuda.amp就能开几行代码的事。第四招是梯度检查点。前向传播的时候不存中间激活值反向的时候重新算一遍。显存能省很多代价是计算时间增加。这个在训练超大模型的时候是标配。第五招是模型并行。把模型拆到多张卡上每张卡存一部分。这个复杂度最高一般不到万不得已不用。注意显存不够的时候先看是不是有内存泄漏。比如在训练循环里不断往列表里 append tensor那显存肯定爆。用torch.cuda.memory_summary()能看到显存分配情况排查起来很方便。5.3 推理速度慢的优化路径推理速度慢先别急着换模型按这个顺序排查。先看是不是在做无用的计算。比如每次推理都重新加载模型或者重复做数据预处理。这些用缓存就能解决。再看是不是没有批处理。单个请求推理GPU 利用率可能只有 10%。攒一批一起算利用率能到 80% 以上吞吐量提升好几倍。然后看是不是没有用量化。FP32 转 FP16速度通常能提升 1.5 到 2 倍精度损失很小。INT8 更快但需要校准。最后看是不是没有用 KV Cache。自回归生成的时候不用 KV Cache 的话每次都要重算前面所有 token复杂度是 O(n^2)。用了之后降到 O(n)长序列上差距巨大。我实测过一个 7B 的模型在单张卡上不做任何优化的话生成速度大概是每秒几个 token。加上批处理、FP16、KV Cache 之后能到每秒几十个 token提升非常明显。5.4 我踩过的那些坑最后分享几个我自己踩过的坑都是文档里不会写的。第一个坑是 NumPy 的广播机制。有次我写(A - B)A 是 (N, 1)B 是 (N,)结果 NumPy 广播成了 (N, N)我还纳闷怎么显存突然爆了。后来养成习惯做减法之前先assert A.shape B.shape能省很多事。第二个坑是 PyTorch 的 in-place 操作。有次我用x 1更新张量结果反向传播报错说版本不对。原因是 in-place 操作会破坏计算图。改成x x 1就好了。这个坑在写自定义层的时候特别容易遇到。第三个坑是随机种子没固定。有次调参同样的配置跑两次结果差很多查了半天才发现是随机种子的问题。后来我在代码开头固定np.random.seed(42)和torch.manual_seed(42)结果就可复现了。第四个坑是数据加载成了瓶颈。GPU 利用率一直上不去查了半天发现是数据加载太慢GPU 一直在等数据。后来把num_workers调大开了预取利用率立马上去了。这个坑在训练大模型的时候特别常见因为数据预处理往往比模型计算还慢。第五个坑是忘了关 dropout。训练的时候开了 dropout推理的时候忘了关结果每次推理结果都不一样还以为是模型有问题。后来在推理前加model.eval()问题解决。这个坑新手几乎都会踩一次。这些坑看起来都是小问题但每一个都能让你卡半天。我把它们记下来就是希望后来的人能少走点弯路。从零构建 AI 工程能力这件事最难的不是理解原理而是把原理变成能跑的代码再把能跑的代码变成跑得好的服务。这个过程没有捷径但有了正确的路径和避坑指南能快很多。