1. 为什么从零手搓AI工程比调包更值得投入第一次看到ai-engineering-from-scratch这个项目名的时候我脑子里冒出来的第一个念头是又来了一个从零实现的教程仓库。这类项目我见过太多大部分是拿 NumPy 写个两层神经网络跑通 MNIST 就收工README 里写一句帮助你理解深度学习原理然后就没有然后了。但真正翻进去看之后我发现这个项目的定位和那些教学玩具完全不是一回事。它要解决的是一个更实际、也更痛的问题当你离开那些高度封装的框架之后还能不能把一个AI系统从数据到推理完整地搭起来这个问题听起来有点自虐毕竟现在pip install transformers之后三行代码就能跑一个对话模型。但我在实际工作中反复遇到这样的场景模型在本地跑得好好的一上生产就出问题显存莫名其妙爆掉推理延迟比预期高一个数量级换个硬件平台整个 pipeline 就崩了。这些问题的根因几乎都藏在那些被封装掉的细节里——张量的内存布局、算子的融合方式、KV Cache 的管理策略、量化误差的累积路径。ai-engineering-from-scratch的价值就在这里。它不是教你AI是什么而是教你AI系统是怎么被工程化地组装起来的。适合的读者有三类一是已经会用框架、但想搞清楚底层到底发生了什么的工程师二是准备做推理优化、模型部署、边缘端落地的开发者三是面试时被问到手写一个 attention就卡壳的人。我打算按这个项目的思路把从零搭建AI工程这件事拆成几个真正能落地的模块来讲。不讲空泛的概念只讲我实际动手时踩过的坑、验证过的方案、以及那些文档里不会写的经验。全程用最朴素的工具链能不用框架就不用框架目的是让你看清楚每一层到底在干什么。2. 环境与工具链的取舍为什么我坚持用最笨的配置2.1 从零开始不等于从汇编开始很多人对from scratch有个误解以为要连矩阵乘法都自己写。这是走极端了。我的原则是凡是能帮我省时间、又不掩盖核心逻辑的工具就用凡是把关键过程封装成黑盒的就自己写。具体到这套工程我的工具链是这样的层级选择理由数值计算NumPy矩阵运算的语义清晰能直接看到 shape 变化自动微分自己实现约200行这是理解反向传播的唯一途径张量加速可选接 PyTorch 做对照用来验证自己实现的正确性数据处理纯 Python NumPy避免 DataLoader 隐藏的采样逻辑推理服务FastAPI轻量能看清请求到响应的完整链路性能分析cProfile 手写计时不依赖任何 profiler 黑盒这套配置看起来土但它的好处是任何一个环节出问题你都能顺着代码一路追到底不会撞到某个框架内部的 C 墙上。2.2 自动微分那200行是整个工程的灵魂我见过太多人跳过这一步直接上 PyTorch 的autograd。结果就是模型不收敛的时候完全不知道从哪查起。自己写一遍自动微分你会被迫回答几个关键问题计算图是动态构建还是静态构建每个算子需要保存哪些中间变量用于反向梯度累加是在节点上做还是在边上做广播broadcasting在反向时怎么处理维度对齐我实现的时候用的是基于 Tape 的反向模式自动微分。核心思路是前向计算时每个操作都往一个磁带上记录一条(输出, 输入, 反向函数)的记录反向传播时从损失开始倒着遍历磁带把梯度按链式法则传回去。class Tensor: def __init__(self, data, requires_gradFalse): self.data np.asarray(data, dtypenp.float32) self.grad None self.requires_grad requires_grad self._backward lambda: None self._prev set() def __add__(self, other): other other if isinstance(other, Tensor) else Tensor(other) out Tensor(self.data other.data, requires_gradself.requires_grad or other.requires_grad) def _backward(): if self.requires_grad: self.grad (self.grad or 0) out.grad if other.requires_grad: other.grad (other.grad or 0) out.grad out._backward _backward out._prev {self, other} return out这段代码看起来简单但里面有个新手最容易踩的坑梯度的累加必须用而不是。因为一个张量可能被多个下游节点使用比如残差连接里的那个恒等分支它的梯度是所有这些路径贡献的总和。我第一版写成结果残差网络怎么训都不收敛查了整整一个下午。2.3 广播的反向传播一个被严重低估的难点(3, 1)和(1, 4)相加得到(3, 4)前向很简单。但反向的时候(3, 4)的梯度要分别缩回(3, 1)和(1, 4)这中间涉及沿被广播的维度求和。def _sum_to_shape(grad, shape): # 把 grad 沿着被广播的维度求和还原到原始 shape while grad.ndim len(shape): grad grad.sum(axis0) for i, dim in enumerate(shape): if dim 1: grad grad.sum(axisi, keepdimsTrue) return grad这个函数我建议你亲手写一遍。因为一旦广播处理错了梯度形状对不上报错信息往往指向一个完全无关的地方排查起来极其痛苦。我的经验是在自动微分实现完成后立刻用数值梯度有限差分做一次校验把每个算子的解析梯度和数值梯度对比误差超过1e-4就说明反向写错了。这一步能帮你省下后面几天的调试时间。3. 手写注意力机制从公式到能跑通的代码3.1 先把公式翻译成人话注意力机制的公式网上一搜一大把Attention(Q,K,V) softmax(QK^T / sqrt(d_k)) V。但公式和代码之间隔着一层工程直觉这层直觉才是真正值钱的东西。我的理解方式是这样的注意力本质上是一个加权平均操作权重由查询和键的相似度决定。你可以把它想象成一个查字典的过程——Q 是你要查的词K 是字典里所有词条V 是每个词条的解释。你查一个词会得到所有词条解释的加权组合权重取决于哪个词条和你要查的词最像。sqrt(d_k)这个缩放因子经常被忽略但它的作用很关键。当维度d_k很大时QK^T的点积结果方差会随维度线性增长导致 softmax 进入饱和区梯度几乎为零。除以sqrt(d_k)就是把方差拉回到 1 附近让 softmax 工作在梯度健康的区间。3.2 一个能跑的最小实现def softmax(x, axis-1): x_max x.max(axisaxis, keepdimsTrue) exp_x np.exp(x - x_max) # 减去最大值防止溢出 return exp_x / exp_x.sum(axisaxis, keepdimsTrue) def attention(Q, K, V, maskNone): d_k Q.shape[-1] scores Q K.transpose(0, 1, 3, 2) / np.sqrt(d_k) if mask is not None: scores np.where(mask, scores, -1e9) weights softmax(scores, axis-1) return weights V, weights这段代码里有三个实战中必须注意的点第一softmax里减最大值这一步不是可选的。我见过有人为了简洁省掉它结果在长序列上直接inf溢出。exp(1000)在任何浮点格式下都是inf而减去最大值之后最大的指数项变成exp(0)1其余都是小数数值稳定性瞬间拉满。第二mask 的实现用-1e9而不是-inf。用-inf的话如果整行都被 mask 掉比如 padding 位置softmax 会出现0/0的 NaN。用-1e9虽然理论上不完美但实践中足够安全而且不会产生 NaN。第三transpose(0,1,3,2)这个操作在 NumPy 里返回的是视图不复制数据。但如果你后面接的是自己写的算子一定要确认它能不能正确处理非连续内存。我就遇到过自己写的矩阵乘法在非连续输入上结果错误的情况排查了半天才发现是 stride 的问题。3.3 多头注意力拆分的维度选择有讲究多头注意力的核心是把d_model维的输入拆成h个头每个头处理d_model/h维。但拆分方式有两种一种是按维度切分reshape后transpose另一种是交错切分。我推荐用reshape transpose的方式因为它在内存上是连续的后续算子处理起来效率更高def split_heads(x, num_heads): batch, seq_len, d_model x.shape d_head d_model // num_heads x x.reshape(batch, seq_len, num_heads, d_head) return x.transpose(0, 2, 1, 3) # (batch, heads, seq, d_head)这里有个容易被忽略的细节reshape之后的transpose会让张量变成非连续的。如果你后面要把它喂给一个对内存布局敏感的算子比如自己写的 CUDA kernel记得先ascontiguousarray一下。在纯 NumPy 环境下这个问题不明显但一旦涉及性能优化它就是隐藏的杀手。4. 训练循环里那些看起来对但实际错的写法4.1 损失函数交叉熵不只是取对数再求平均交叉熵损失看起来简单但手写的时候有几个坑。最典型的是数值稳定性如果你先算 softmax 再取 log中间结果可能下溢成 0log(0)就是-inf。正确的做法是把 softmax 和 log 合并成一个操作数学上等价于log_softmaxdef cross_entropy(logits, targets): # logits: (N, C), targets: (N,) max_logits logits.max(axis1, keepdimsTrue) shifted logits - max_logits log_sum_exp np.log(np.exp(shifted).sum(axis1, keepdimsTrue)) log_probs shifted - log_sum_exp N logits.shape[0] return -log_probs[np.arange(N), targets].mean()这个实现里shifted的最大值是 0exp不会溢出log_sum_exp也不会出现log(0)。整个计算过程在数值上是安全的。4.2 参数初始化不是随便乘个 0.01 就行我早期训练的时候参数初始化就是np.random.randn(...) * 0.01结果深层网络根本训不起来。后来才明白初始化的方差需要根据层的输入输出维度来调整。对于tanh或sigmoid激活函数用Xavier 初始化方差为2 / (fan_in fan_out)。对于ReLU用He 初始化方差为2 / fan_in。def xavier_init(fan_in, fan_out): limit np.sqrt(6.0 / (fan_in fan_out)) return np.random.uniform(-limit, limit, (fan_in, fan_out)) def he_init(fan_in, fan_out): std np.sqrt(2.0 / fan_in) return np.random.randn(fan_in, fan_out) * std为什么这个重要因为如果初始化方差太小信号每过一层就衰减一次到深层几乎为零梯度也传不回去方差太大则相反信号逐层放大最后爆炸。初始化的本质是让信号在前向和反向传播中保持方差稳定这是训练能收敛的前提。4.3 学习率调度warmup 不是玄学Transformer 类模型训练时学习率通常需要一个 warmup 阶段——从很小的值线性增加到峰值然后再衰减。很多人觉得这是调参玄学其实背后有明确的道理。训练初期参数是随机初始化的梯度方向噪声很大。如果这时候用大学习率参数会被推向一个糟糕的区域后面很难拉回来。Warmup 的作用就是在梯度方向还不稳定的时候用小步走等方向稳定了再加速。def lr_schedule(step, d_model, warmup_steps): # 经典的 Transformer 学习率公式 return d_model ** (-0.5) * min(step ** (-0.5), step * warmup_steps ** (-1.5))这个公式里step * warmup_steps ** (-1.5)是线性增长阶段step ** (-0.5)是衰减阶段两者取最小值就实现了先增后减。d_model ** (-0.5)是缩放因子让不同模型规模下的学习率有一个统一的基准。5. 推理优化从能跑到跑得快之间隔着什么5.1 KV Cache自回归生成的必备优化自回归生成的时候每生成一个新 token都要重新计算整个序列的注意力。但前面 token 的 Key 和 Value 其实没变重复计算纯属浪费。KV Cache 的思路就是把已经算过的 K 和 V 缓存起来新 token 只需要算自己的 Q然后和缓存的 K、V 做注意力。class KVCache: def __init__(self, max_len, num_heads, d_head): self.k_cache np.zeros((max_len, num_heads, d_head)) self.v_cache np.zeros((max_len, num_heads, d_head)) self.pos 0 def update(self, k, v): seq_len k.shape[0] self.k_cache[self.pos:self.posseq_len] k self.v_cache[self.pos:self.posseq_len] v self.pos seq_len return self.k_cache[:self.pos], self.v_cache[:self.pos]这个优化能把生成阶段的复杂度从O(n^2)降到O(n)每步只算一个新 token 的注意力。但要注意显存占用KV Cache 的大小是2 * num_layers * max_len * num_heads * d_head * dtype_size。对于长序列这个数字可能比模型本身还大。我实测过一个 7B 模型在 4096 序列长度下KV Cache 占了将近 2GB 显存。5.2 量化省显存但别省掉精度量化是把 FP32 或 FP16 的权重压缩成 INT8 或 INT4能显著降低显存占用和带宽需求。但量化不是免费的午餐误差会累积。我的经验是权重量化对精度影响较小激活量化影响较大。因为权重是静态的可以提前校准激活是动态的范围随输入变化量化误差更难控制。一个实用的做法是per-channel 量化即每个输出通道用独立的缩放因子而不是整个张量共用一个def quantize_per_channel(weight, bits8): # weight: (out_features, in_features) qmax 2 ** (bits - 1) - 1 scales np.abs(weight).max(axis1, keepdimsTrue) / qmax scales np.where(scales 0, 1e-8, scales) # 防止除零 quantized np.round(weight / scales).clip(-qmax, qmax).astype(np.int8) return quantized, scalesscales里那个1e-8的兜底很重要。如果某个通道的权重全是 0训练中可能出现scales就是 0除法直接产生 NaN。这种边界情况在测试集上不一定触发但生产环境里迟早会遇到。5.3 算子融合减少内存往返才是关键现代 GPU 的算力早就过剩了真正的瓶颈是内存带宽。每次算子切换都要把数据从显存读进来、算完再写回去这个往返开销远大于计算本身。算子融合就是把多个连续的操作合并成一个 kernel中间结果留在寄存器或共享内存里不落显存。最典型的例子是LayerNorm Linear的融合或者GELU的近似实现。在纯 NumPy 环境下我们没法做真正的 kernel 融合但可以用**原地操作in-place**来减少内存分配# 不推荐每次操作都分配新内存 x x bias x x * scale x np.maximum(x, 0) # 推荐原地操作 x bias x * scale np.maximum(x, 0, outx)这个优化在 NumPy 里效果有限但一旦你迁移到 GPU 或者用 C 重写这个思路就是性能提升的核心。养成能原地就原地的习惯是从零做AI工程必须有的意识。6. 踩坑实录那些让我熬夜到凌晨的bug6.1 梯度爆炸loss 突然变成 NaN这是我最常遇到的问题。训练到一半loss 突然从 2.3 跳到 NaN。排查思路是这样的第一步确认是前向还是反向出的问题。在前向的每个算子后面加断言检查有没有inf或nan。如果前向正常那就是反向的梯度出了问题。第二步检查梯度范数。在backward之后打印每一层的梯度 L2 范数如果某一层突然变成1e10这种量级就是梯度爆炸。第三步定位到具体算子。梯度爆炸通常发生在exp、log、除法这些操作上。比如 softmax 如果没有减最大值exp就会溢出。解决方案有两个梯度裁剪和更稳定的算子实现。梯度裁剪是兜底手段def clip_gradients(params, max_norm1.0): total_norm np.sqrt(sum(np.sum(p.grad ** 2) for p in params if p.grad is not None)) if total_norm max_norm: scale max_norm / (total_norm 1e-6) for p in params: if p.grad is not None: p.grad * scale但裁剪只是治标治本还是要找到数值不稳定的根源。我的经验是90% 的 NaN 问题都出在 softmax、log、除法这三个操作上把它们的数值稳定性处理好问题基本就消失了。6.2 形状对不上报错信息指向错误的位置NumPy 的广播机制很强大但也很容易掩盖 bug。比如你期望(batch, seq, dim)和(batch, dim)相加结果写成了(batch, seq, dim)和(seq, dim)NumPy 不会报错而是默默地广播成(batch, seq, dim)但语义完全错了。我的应对方法是在关键算子入口处显式检查 shape。def linear(x, weight, bias): assert x.shape[-1] weight.shape[1], \ fdim mismatch: x last dim {x.shape[-1]} vs weight in {weight.shape[1]} out x weight.T if bias is not None: assert bias.shape[0] weight.shape[0] out bias return out这些断言在训练时几乎不花时间但能在 bug 刚出现的时候就抓住它而不是等到 loss 不收敛才回头查。6.3 内存泄漏训练越跑越慢有段时间我发现训练循环跑了几百步之后越来越慢一开始以为是正常的计算波动后来发现是内存泄漏。原因是计算图没有释放。在基于 Tape 的自动微分里每次前向都会往磁带上追加记录。如果不在每个训练步结束后清空磁带记录会无限增长内存越占越多GC 压力越来越大。for step in range(num_steps): tape.clear() # 关键每步开始前清空 loss forward_and_loss(batch) loss.backward() optimizer.step()这个坑的隐蔽之处在于小规模测试时看不出来因为几十步的内存增长不明显。只有跑到几千步以上问题才暴露。所以我的建议是任何训练循环都要在早期用tracemalloc或者简单的 RSS 监控跑一次长程测试确认内存是稳定的。7. 从这套工程里能带走什么把ai-engineering-from-scratch这条路走一遍最大的收获不是我会手写 Transformer 了而是建立起了一套排查问题的思维框架。以前遇到模型不收敛我的反应是换个学习率试试换个初始化试试本质上是碰运气。现在我会按顺序问自己前向的数值稳定吗梯度的量级正常吗形状有没有隐式广播内存有没有泄漏每一步都有明确的检查手段而不是盲目试错。另一个收获是对性能的理解变了。以前觉得性能优化就是用更快的框架现在明白真正的瓶颈在内存带宽、在算子融合、在数据布局。这些认知在调包的时候是永远学不到的因为框架把这些细节全藏起来了。如果你也想走这条路我的建议是别贪多先把自动微分和注意力这两块吃透。这两个是AI工程的地基地基打牢了上面盖什么楼都稳。至于那些花哨的优化技巧等你把基础版本跑通、跑对之后自然就知道该往哪个方向优化了。最后分享一个我自己的习惯每实现一个新算子我都会写一个数值梯度校验的测试。虽然多花十分钟但它能帮我在后续几天的调试里省下无数个小时。这个投入产出比是我做AI工程这些年里见过最高的。