★ FEATURED ARTICLE 1. 从零手搓AI工程为什么我不建议你直接调包第一次看到ai-engineering-from-scratch这个项目名的时候我正坐在工位上啃一个调了三天都没收敛的推荐模型。说实话那一瞬间我是有点不屑的——这年头谁还从零写AI工程HuggingFace 上pipeline()一行代码就能跑推理LangChain 拖几个节点就能搭个 RAG连微调都有 LoRA 这种“穷人版”方案何必自讨苦吃去手搓矩阵乘法但后来我改主意了。原因很简单我在生产环境里被“黑盒”坑过太多次。模型输出突然漂移我不知道是 tokenizer 的问题还是 KV cache 的锅显存 OOM我连 attention 的中间张量长什么样都说不清楚想做个自定义的采样策略翻遍文档发现底层接口根本不给你留口子。那种感觉就像开着一辆引擎盖焊死的车仪表盘只告诉你“车坏了”但坏在哪、怎么修全靠猜。ai-engineering-from-scratch这个项目本质上就是把引擎盖重新打开。它不是教你“怎么用 AI”而是教你“AI 到底是怎么跑起来的”——从张量怎么在内存里排布到反向传播的梯度怎么一层层流回去再到推理时 KV cache 怎么省显存、量化怎么压精度。这套东西学下来你不会立刻变成调包侠里的高手但你会变成那种出了问题能定位到具体一行代码的人。这篇文章适合谁看如果你是刚入门的算法工程师天天调model.fit()但说不清loss.backward()里发生了什么那这篇就是给你写的。如果你是有几年经验的工程老手但一直停留在“调 API”层面想往下钻一层看看地基那也合适。至于纯小白我建议你先补一下 Python 和基础线代不然第二章的矩阵乘法可能会让你想砸键盘。我接下来会按这个项目的典型学习路径来拆先讲整体设计思路再拆核心模块的实现细节然后是完整的实操流程和参数计算最后是我踩过的坑和排查技巧。全程不藏私能抄的代码直接给能算的数直接算给你看。2. 整体设计思路为什么是“从零”而不是“从框架”2.1 从零实现的真正价值在哪很多人对“from scratch”有误解以为就是不用任何库、纯手写。其实不是。这个项目的核心思路是用最少的依赖把 AI 工程里最关键的几个环节亲手实现一遍。你可以用 NumPy 做矩阵运算可以用 PyTorch 的Tensor做自动求导的对照但核心逻辑——比如 attention 的缩放点积、LayerNorm 的均值方差计算、softmax 的数值稳定性处理——必须你自己写。为什么这么设计因为框架帮你做的事太多了多到你把“会用”当成了“懂”。举个例子PyTorch 的nn.Linear里权重初始化用的是 Kaiming 均匀分布偏置初始化为零。你知道为什么吗因为如果权重全初始化为零每一层的输出都一样反向传播时梯度也一样网络永远学不到东西——这叫对称性破缺问题。但如果你一直用nn.Linear你根本不会去想这件事直到某天你自己手写了一个全连接层发现 loss 死活不降才回头去查初始化。这个项目的设计逻辑就是逼你面对这些“默认值”。它把每个模块拆到最小可运行单元让你自己决定初始化方式、自己处理数值溢出、自己管理内存布局。做完一遍你再回头看框架的源码会有一种“哦原来你在这儿帮我擦屁股”的顿悟感。2.2 技术选型NumPy PyTorch 双轨制项目在工具选型上走的是双轨制底层计算用 NumPy自动求导和 GPU 加速用 PyTorch。这个选择很讲究。纯 NumPy 的好处是透明。你能看到每一个ndarray的 shape、stride、dtype能手动控制np.dot还是np.einsum能亲眼看到float32和float16在累加时的精度差异。但 NumPy 没有自动求导你得自己推导反向传播公式自己实现梯度更新。这对理解原理极有帮助但效率太低跑个 MNIST 都要等半天。所以项目在关键路径上引入 PyTorch 作为“对照系”你用 NumPy 手写一个 attention再用 PyTorch 的nn.MultiheadAttention跑同样的输入对比输出是否一致。如果不一致说明你的实现有问题——可能是缩放因子忘了除sqrt(d_k)也可能是 mask 的位置搞错了。这种对照调试法是我认为这个项目最聪明的设计之一。提示不要一上来就追求完全不用框架。先用 PyTorch 跑通再用手写实现去对齐最后把 PyTorch 的部分替换掉。这个渐进路径比硬啃纯 NumPy 高效得多。2.3 模块划分从张量到推理的完整链路项目的模块划分遵循 AI 工程的实际数据流大致分五层张量层手写Tensor类支持reshape、transpose、广播机制理解 stride 和内存连续性。算子层实现matmul、softmax、layernorm、gelu等基础算子处理数值稳定性。网络层搭Linear、Embedding、Attention、TransformerBlock理解残差连接和归一化的位置。训练层写loss、optimizer、lr_scheduler理解梯度裁剪和权重衰减。推理层实现 KV cache、量化、采样策略理解显存和延迟的权衡。这个划分的好处是每一层都可以独立测试。你写完softmax就能用numpy.allclose验证它和 PyTorch 的输出是否一致不用等到整个 Transformer 搭完才发现问题。我在实际做的时候每写完一个算子就写一个单元测试前后大概攒了 40 多个测试用例后面调 bug 的时候省了至少一半时间。3. 核心细节解析手写算子的那些坑3.1 Softmax 的数值稳定性不只是减最大值Softmax 的公式看起来人畜无害softmax(x_i) exp(x_i) / sum(exp(x_j))。但如果你直接这么写输入稍微大一点就会溢出。比如x [1000, 1001, 1002]exp(1000)直接变成inf整个结果就废了。标准做法是减去最大值softmax(x_i) exp(x_i - max(x)) / sum(exp(x_j - max(x)))。这样指数最大是exp(0) 1不会溢出。但这里有个细节减最大值之后如果输入全是负无穷比如 mask 掉的 padding 位置max 也是负无穷减法会变成nan。所以实际实现里要处理这种情况通常是把负无穷替换成一个极小的数或者直接在 mask 位置输出零。我在手写的时候还遇到一个更隐蔽的问题float16 下的精度损失。即使减了最大值exp的结果在 float16 下只有 10 位有效数字累加时如果序列很长比如 4096 个 token误差会累积到不可忽略。解决办法是在 softmax 内部用 float32 计算最后再转回 float16。这个技巧在混合精度训练里很常见但如果你不手写一遍根本不会意识到。def softmax(x, axis-1): # 转 float32 计算避免 float16 精度损失 x x.astype(np.float32) x_max np.max(x, axisaxis, keepdimsTrue) # 处理全负无穷的情况 x_max np.where(np.isfinite(x_max), x_max, 0.0) exp_x np.exp(x - x_max) return exp_x / np.sum(exp_x, axisaxis, keepdimsTrue)3.2 LayerNorm 和 BatchNorm 的选择逻辑LayerNorm 和 BatchNorm 的区别网上文章一抓一大把但大多停留在“LN 对序列建模更好”这种结论层面。这个项目逼你从实现角度去理解BatchNorm 的均值和方差是在 batch 维度上算的所以推理时要用训练时累积的 running mean 和 running var而 LayerNorm 是在特征维度上算的训练和推理行为完全一致。这个差异在实现上体现得很明显。BatchNorm 需要维护running_mean和running_var两个 buffer训练时用动量更新推理时直接用。LayerNorm 什么都不用存每次前向现算就行。所以在 Transformer 里用 LayerNorm 而不是 BatchNorm不只是因为“序列长度可变”更是因为推理时不需要额外的状态管理部署起来简单得多。但 LayerNorm 也有坑。它的均值和方差是在最后一维上算的如果你的输入是[batch, seq_len, hidden_dim]那 normalization 是在hidden_dim上做的。这意味着每个 token 的归一化是独立的不同 token 之间不共享统计量。这跟 BatchNorm 在 batch 维度上做归一化完全不同。我在第一次实现的时候把 axis 搞错了结果模型训练 loss 震荡得厉害查了半天才发现是归一化维度不对。3.3 Attention 的缩放因子为什么是 sqrt(d_k)Scaled Dot-Product Attention 的公式是Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。那个sqrt(d_k)的缩放因子很多人知道要除但说不清为什么。原因在于点积的方差。假设 Q 和 K 的每个元素都是独立同分布的均值为 0方差为 1那么QK^T的每个元素是d_k个乘积之和方差就是d_k。当d_k很大时比如 64 或 128点积的结果会很大softmax 之后会变得非常尖锐——几乎所有的概率质量都集中在一个位置上梯度接近于零训练不动。除以sqrt(d_k)之后方差被拉回 1softmax 的输出分布更平滑梯度更健康。这个解释在论文里只有一句话但如果你手写一遍把d_k4和d_k128的 attention 权重打印出来对比会直观地看到不缩放的时候d_k128的 attention 矩阵几乎是一个 one-hot 向量缩放之后才呈现出合理的分布。这种亲眼看到差异的体验比读十遍论文都管用。注意缩放因子是sqrt(d_k)而不是d_k因为方差是d_k标准差是sqrt(d_k)。归一化要除的是标准差不是方差。这个细节我在面试里问过不下二十个人能说清楚的不到一半。3.4 KV Cache 的显存计算与实现KV Cache 是推理加速的核心技术但它的显存开销经常被低估。我来算一笔账假设模型有L层每层有h个注意力头每个头的维度是d_head序列长度是Sbatch size 是B数据类型是 float162 字节。KV Cache 的显存占用是2 * L * h * d_head * S * B * 2字节。那个开头的 2 是因为要存 Key 和 Value 两份。拿 LLaMA-7B 举例L32h32d_head128S2048B1。算一下2 * 32 * 32 * 128 * 2048 * 1 * 2 2^1 * 2^5 * 2^5 * 2^7 * 2^11 * 2^1 2^30字节也就是1GB。这还只是 batch size 为 1 的情况如果 batch 开到 16KV Cache 就要 16GB比模型本身的权重7B * 2 字节 14GB还大。所以实际部署时KV Cache 的管理是个大问题。项目里实现了一个简单的 KV Cache用预分配的numpy数组来存避免动态扩容的开销。关键点是预分配 指针移动一开始就分配好最大序列长度的空间每次生成新 token 时把新的 K、V 写到指针位置然后指针前移。这样避免了频繁的np.concatenate速度能快好几倍。class KVCache: def __init__(self, max_seq_len, n_layers, n_heads, d_head, dtypenp.float16): # 预分配避免动态扩容 self.k_cache np.zeros((n_layers, max_seq_len, n_heads, d_head), dtypedtype) self.v_cache np.zeros((n_layers, max_seq_len, n_heads, d_head), dtypedtype) self.ptr 0 def update(self, layer_idx, new_k, new_v): seq_len new_k.shape[0] self.k_cache[layer_idx, self.ptr:self.ptrseq_len] new_k self.v_cache[layer_idx, self.ptr:self.ptrseq_len] new_v # 注意ptr 的更新要在所有层都更新完之后4. 实操过程从零搭一个 Mini Transformer4.1 环境准备与依赖管理项目对环境的依赖很克制核心就三个numpy、torch只用来做对照验证、pytest跑单元测试。Python 版本建议 3.10 以上因为用到了match-case语法和一些类型注解的新特性。我建议用conda建一个独立环境避免和系统里的包冲突conda create -n ai-scratch python3.10 conda activate ai-scratch pip install numpy torch pytest如果你有 GPU可以装 CUDA 版的 PyTorch但项目本身的手写部分不依赖 GPUCPU 就能跑。我一开始在 MacBook 上跑的M1 芯片的 NumPy 性能还不错跑一个小型 Transformer 的前向传播大概几百毫秒完全能接受。提示不要一上来就装一堆可视化工具。项目前期不需要 TensorBoard打印 loss 曲线用matplotlib就够了。依赖越少出问题的概率越小。4.2 手写张量类的关键设计张量类的核心是数据 形状 步长。数据是扁平的ndarray形状是逻辑上的维度步长是每个维度在内存里跨多少元素。理解步长是理解transpose、reshape、broadcast的关键。比如一个[2, 3]的矩阵内存里是[1, 2, 3, 4, 5, 6]步长是[3, 1]——第 0 维跨 3 个元素第 1 维跨 1 个元素。transpose之后形状变成[3, 2]步长变成[1, 3]但内存里的数据没动。这就是为什么transpose是零拷贝的而reshape可能需要拷贝——如果新的形状和原来的步长不兼容就得重新排布内存。我在实现的时候偷了个懒直接用numpy的ndarray作为底层存储自己封装了shape和stride的访问接口。这样既保留了 NumPy 的性能又能手动控制步长逻辑。如果你要完全从零写那就得自己管理一块bytearray用struct或memoryview来读写工作量会大很多但理解也会更深。4.3 搭建 Transformer Block 的完整流程一个标准的 Transformer Block 包含多头注意力、前馈网络、两个 LayerNorm、两个残差连接。顺序通常是LayerNorm - Attention - 残差 - LayerNorm - FFN - 残差这叫Pre-Norm结构。原始论文用的是 Post-NormAttention - 残差 - LayerNorm但后来大家发现 Pre-Norm 训练更稳定不需要 warmup 也能收敛。实现的时候有几个关键点残差连接要求输入输出维度一致。所以Attention和FFN的输出维度必须等于hidden_dim。FFN 内部通常会放大 4 倍hidden_dim - 4*hidden_dim - hidden_dim但最终会投影回来。多头注意力的拆分和合并。输入是[batch, seq_len, hidden_dim]要先 reshape 成[batch, seq_len, n_heads, d_head]再 transpose 成[batch, n_heads, seq_len, d_head]然后做 attention。做完之后再 transpose 回来reshape 成[batch, seq_len, hidden_dim]。这一串操作很容易搞错维度顺序我建议每一步都打印 shape 确认。mask 的处理。训练时要用 causal mask 防止看到未来的 token推理时如果用了 KV Cachemask 的形状会不一样。训练时 mask 是[seq_len, seq_len]的下三角矩阵推理时只需要处理当前 token 和 cache 里所有 token 的关系。def forward(self, x, maskNone): # Pre-Norm: 先归一化再进 attention residual x x self.ln1(x) x self.attention(x, mask) x x residual # 残差连接 residual x x self.ln2(x) x self.ffn(x) x x residual return x4.4 训练循环与梯度裁剪训练循环看起来简单但细节很多。首先是loss 的计算语言模型用交叉熵但要注意 ignore_index把 padding 位置的 loss 屏蔽掉。其次是梯度累积如果显存不够可以把 batch 拆成几个小 batch分别前向反向累积梯度后再更新。最后是梯度裁剪Transformer 训练时梯度容易爆炸通常把梯度的 L2 范数裁剪到 1.0。梯度裁剪的实现有个小技巧先算所有参数梯度的全局范数如果超过阈值就按比例缩放。不要对每个参数单独裁剪那样会改变梯度的方向。def clip_gradients(params, max_norm1.0): total_norm 0.0 for p in params: total_norm np.sum(p.grad ** 2) total_norm np.sqrt(total_norm) clip_coef max_norm / (total_norm 1e-6) if clip_coef 1.0: for p in params: p.grad * clip_coef return total_norm我在实际训练的时候一开始没加梯度裁剪loss 在第 500 步左右突然变成nan查了半天发现是某个 batch 的梯度范数到了 1e4 量级。加上裁剪之后训练稳定多了。这个教训告诉我Transformer 的训练稳定性一半靠架构一半靠这些工程细节。5. 常见问题与排查技巧实录5.1 数值不稳定问题速查表现象可能原因排查方法解决方案loss 变成 nan梯度爆炸打印每层梯度范数加梯度裁剪检查学习率loss 震荡不收敛学习率太大画 loss 曲线看周期降学习率加 warmup输出全是同一个 tokensoftmax 溢出检查 logits 范围softmax 减最大值float16 下精度差累加溢出对比 float32 结果关键算子用 float32attention 权重太尖忘了缩放打印 attention 矩阵除以 sqrt(d_k)这张表是我踩坑踩出来的每一条都对应一个真实 bug。比如“输出全是同一个 token”那个我一开始以为是模型没训练好后来发现是 softmax 在 float16 下溢出了logits 稍微大一点就变成infsoftmax 之后变成nan采样的时候全落到同一个 token 上。改成 float32 计算之后立刻正常。5.2 维度不匹配的调试方法维度错误是手写 AI 工程里最常见的 bug没有之一。我的经验是在每个关键操作前后打印 shape不要嫌麻烦。比如 attention 里Q是[B, H, S, D]K^T是[B, H, D, S]QK^T是[B, H, S, S]V是[B, H, S, D]最后输出是[B, H, S, D]。这一串如果哪个 transpose 忘了shape 就对不上报错信息还特别隐晦。我后来养成了一个习惯写一个assert_shape函数在每个模块的入口和出口调用。这样一旦出错报错信息直接告诉你哪个模块的输入输出不对比 NumPy 的广播错误清晰多了。def assert_shape(tensor, expected, nametensor): actual tensor.shape assert actual expected, f{name} shape mismatch: expected {expected}, got {actual}5.3 性能优化的几个实用技巧手写实现通常比框架慢但有几个技巧可以缩小差距用einsum代替循环。np.einsum(bhqd,bhkd-bhqk, Q, K)比手写循环快几十倍而且代码更清晰。避免不必要的拷贝。transpose是零拷贝的reshape尽量用view而不是copy。批量操作。能一次算完的不要拆成循环NumPy 的向量化是性能的关键。用float32而不是float64。精度够用的情况下float32 快一倍内存省一半。我在实现 attention 的时候一开始用双重循环算QK^T跑一个[8, 12, 128, 64]的输入要好几秒。改成einsum之后降到几十毫秒差距非常明显。这个优化在框架里是自动做的但手写的时候你得自己意识到。5.4 从手写实现到框架源码的对照阅读做完手写实现之后我强烈建议你去读一遍 PyTorch 的源码。你会发现很多之前忽略的细节比如nn.Linear的权重初始化是kaiming_uniform_nn.LayerNorm的eps默认是1e-5nn.MultiheadAttention的dropout是在 attention 权重上做的而不是输出上。这些细节在手写的时候你可能没考虑但框架帮你做了默认选择。理解这些默认值背后的原因你才能在需要的时候正确地覆盖它们。比如你知道eps1e-5是为了防止方差为零时除零那当你的输入方差特别小的时候就知道要把eps调大一点。提示读源码的时候不要从头读到尾带着问题去读。比如“为什么 LayerNorm 的 eps 是 1e-5 而不是 1e-8”然后去搜 issue、看 commit比通读源码高效得多。6. 我个人的实操体会这个项目我断断续续做了大概三周中间踩的坑比我预想的多得多。最大的体会是AI 工程的难点不在算法在工程细节。算法论文里一句话带过的东西实现起来可能有一堆边界情况要处理。比如 attention 的 mask论文里就一句“we apply a mask”但实际实现要考虑 padding mask 和 causal mask 的叠加、mask 的形状广播、mask 值的类型是-inf还是-1e9每一个细节都可能让结果出错。另一个体会是手写一遍之后调包的效率反而更高了。因为你知道每个参数在干什么知道哪些默认值可以改、哪些不能动。以前调model.generate()的时候temperature、top_p、repetition_penalty这些参数我都是凭感觉设的现在我能从采样原理上解释每个参数的作用调起来有方向多了。最后分享一个小技巧如果你也想做类似的项目不要一上来就搭完整的 Transformer。先从单个算子开始写一个softmax用numpy.allclose验证它和 PyTorch 一致再写layernorm再写attention每步都验证。这样虽然看起来慢但后面调 bug 的时候会省大量时间。我一开始想一口气搭完结果卡在 attention 的维度变换上整整两天后来退回去一步步验证半天就搞定了。这个项目后续还可以往几个方向扩展比如加一个简单的 BPE tokenizer把文本预处理也手写一遍或者实现一个量化推理的版本把 float16 压到 int8看看精度掉多少再或者加一个 beam search 的采样策略对比 greedy 和 beam search 的生成质量。每一个方向都能让你对 AI 工程的理解再深一层。
阅读完成 · 觉得有帮助？
