1. 这个项目到底在解决什么问题第一次看到 ai-engineering-from-scratch 这个标题我脑子里蹦出来的第一个念头是又是一个教人调包的教程但仔细琢磨了一下 from scratch 这四个字我意识到它想做的事情可能完全不一样。市面上讲 AI 工程的内容绝大多数都是从pip install transformers或者import openai开始的你跟着敲一遍模型跑起来了loss 也降下去了但真要问你一句这个 attention 里的 Q、K、V 到底是怎么算出来的很多人就卡壳了。这个项目要解决的恰恰就是这个断层——它试图让你从最底层的矩阵乘法开始一步步把 AI 工程里那些被封装得严严实实的核心组件亲手搭出来。说白了这是一个面向想真正搞懂 AI 系统内部运转机制的工程师的学习路径项目。它适合的人很明确你已经会用 PyTorch 或者 TensorFlow 跑通一些 demo 了但总觉得自己像个调包侠遇到模型不收敛、显存爆炸、推理延迟高这些问题时只能靠搜索引擎和玄学调参。它不适合完全零基础的人——如果你连 Python 的列表推导式都写不利索那还是先去补编程基础。但只要你有一点点深度学习的使用经验想往底层钻一钻这个项目的价值就非常大。我之所以对这个方向特别有感触是因为我自己就经历过那个阶段。早些年做推荐系统的时候模型跑不通我就换优化器换学习率换 batch size折腾一整天可能只是碰巧好了但根本不知道问题出在哪。后来逼着自己手写了一遍反向传播手写了一遍 attention那种原来如此的感觉是看一百篇博客都给不了的。所以当我看到 from scratch 这个定位时我第一反应是这个方向对了它补的是工程能力和原理理解之间那道最容易被忽略的鸿沟。这个项目的核心价值可以拆成三层。第一层是认知层让你知道每个组件为什么存在、解决什么问题第二层是实现层让你能用最基础的算子把它搭出来不依赖高级封装第三层是工程层让你理解真实系统里这些组件是怎么组合、怎么优化、怎么排查问题的。三层缺一不可只讲原理不讲实现是纸上谈兵只讲实现不讲原理是另一种调包。2. 从零构建的核心思路拆解2.1 为什么非要从零而不是从框架很多人会问现在框架这么成熟为什么还要费劲从零写这不是重复造轮子吗我的理解是从零构建的目的从来不是让你在生产环境里真的手写一个 Transformer而是让你获得一种透视能力。当你亲手实现过一遍 softmax 的数值稳定性处理你再看到框架里那个F.softmax的时候你就知道它在背后帮你做了什么遇到 NaN 的时候你也能立刻反应过来可能是哪里溢出了。从工程角度看框架帮你屏蔽了细节但同时也屏蔽了调试的抓手。举个我踩过的坑有一次模型训练到一半 loss 突然变成 NaN我查了半天数据、学习率、梯度裁剪都没问题最后发现是某一层的 attention score 在 fp16 下溢出了。如果我没有手写过 attention我根本不会想到去检查那个位置。这就是从零带来的直觉——你对每个数值的流动范围有了体感。2.2 自底向上的分层设计一个合理的 from scratch 路径我倾向于按这样的层次来组织这也是我在带新人的时候反复验证过的顺序数值计算基础层向量、矩阵、广播机制、数值稳定性。这一层看起来最枯燥但它是后面所有内容的地基。很多人跳过这层直接写模型结果遇到维度对不上、广播出错就懵了。自动微分层计算图、前向传播、反向传播、链式法则的工程实现。这是整个深度学习的心脏理解了它你就理解了为什么 PyTorch 要设计成动态图。神经网络基础组件层全连接层、激活函数、损失函数、优化器。这一层开始有模型的感觉了。序列建模层从 RNN 到 attention 再到 Transformer。这是现代 AI 工程的核心也是最能体现从零价值的地方。训练与推理工程层数据加载、批处理、混合精度、梯度累积、KV cache、量化。这一层直接对应真实生产环境的问题。这个分层的关键在于每一层都只依赖它下面的层你可以在任何一层停下来都获得一个完整可用的能力。比如你只做到自动微分层你就能自己实现一个简单的线性回归训练你做到 Transformer 层你就能手搓一个迷你 GPT。2.3 选型背后的取舍逻辑在实现语言和工具的选择上我的建议是用 Python NumPy 起步逐步过渡到 PyTorch 的底层 API。为什么不直接用 PyTorch 的高层 API因为那样你就又回到调包的老路了。NumPy 的好处是它足够裸矩阵乘法就是矩阵乘法没有 autograd 帮你兜底你必须自己想清楚每一步的梯度怎么传。但纯 NumPy 也有问题——它没有 GPU 加速跑到 Transformer 这种规模就慢得没法忍。所以合理的过渡点是基础层和自动微分层用 NumPy 手写把原理吃透到了神经网络组件层之后开始用 PyTorch 的torch.Tensor和torch.autograd但依然自己定义 forward 和 backward不用nn.Module的高级封装。这样既保证了原理的透明度又获得了工程上的可行性。提示不要一上来就追求性能。from scratch 阶段代码的可读性和正确性远比速度重要。我见过太多人为了优雅把代码写得极其抽象结果自己都看不懂了调试的时候痛苦不堪。3. 核心组件的实操实现要点3.1 自动微分引擎的手写细节自动微分是整个项目里最烧脑但也最值得的部分。核心思路其实不复杂把每个运算都表示成一个节点节点记录自己的输入、输出和局部梯度前向传播时构建计算图反向传播时沿着图反向应用链式法则。听起来简单但魔鬼在细节里。第一个坑是计算图的构建时机。动态图和静态图的区别在你手写一遍之后就彻底明白了。动态图是边算边建图每次前向传播都重新构建灵活但开销大静态图是先建图再算效率高但不灵活。PyTorch 用的是动态图这也是为什么它调试起来比早期的 TensorFlow 舒服。第二个坑是梯度的累加。同一个节点可能被多条路径用到反向传播时它的梯度是所有路径梯度的和。这个累加逻辑如果写错了梯度就会偏而且偏得很隐蔽——模型还能训就是效果差一点你很难发现。我当时的做法是写一个数值梯度检验函数用有限差分法算出来的梯度和反向传播算出来的梯度对比误差在 1e-6 以内才算通过。def numerical_gradient(f, x, eps1e-6): grad np.zeros_like(x) it np.nditer(x, flags[multi_index]) while not it.finished: idx it.multi_index old x[idx] x[idx] old eps fx1 f(x) x[idx] old - eps fx2 f(x) x[idx] old grad[idx] (fx1 - fx2) / (2 * eps) it.iternext() return grad这个函数看起来笨但它是你验证自己 autograd 实现是否正确的唯一可靠手段。我强烈建议每实现一个新的算子都用它检验一遍。3.2 Attention 机制的数值稳定性处理Attention 是 Transformer 的灵魂也是数值问题的高发区。标准公式是softmax(QK^T / sqrt(d_k)) V但直接这么写在 fp16 下几乎必炸。问题出在QK^T的值可能很大指数运算exp之后直接溢出成 inf。解决办法是减去最大值softmax(x) softmax(x - max(x))。数学上等价但数值上稳定得多因为减完之后最大的指数是exp(0) 1不会溢出。这个技巧看起来简单但它是无数人踩坑之后总结出来的你在框架源码里能看到但如果不手写一遍你永远不会真正记住它。def stable_softmax(x, axis-1): x_max np.max(x, axisaxis, keepdimsTrue) exp_x np.exp(x - x_max) return exp_x / np.sum(exp_x, axisaxis, keepdimsTrue)还有一个容易被忽略的点是mask 的处理。在自回归生成里我们需要屏蔽掉未来的 token通常是把对应位置的 score 设成一个极大的负数这样 softmax 之后概率就趋近于 0。但如果你设成-inf在某些实现里会出 NaN。稳妥的做法是设成一个足够大的负数比如-1e9而不是真正的负无穷。3.3 反向传播中的梯度检查清单手写反向传播的时候我总结了一个检查清单每次实现新层都过一遍检查项常见错误验证方法维度匹配梯度形状和参数形状不一致打印 shape 对比链式法则顺序乘法顺序反了数值梯度检验梯度累加多路径梯度没求和构造共享节点的测试用例原地操作修改了前向的中间变量检查是否有 in-place 赋值广播回传广播后的梯度没求和还原对广播维度做 sum这个表格里的每一条都是我真金白银踩出来的。尤其是广播回传这一条特别隐蔽。比如你把一个(batch, 1)的 bias 加到(batch, hidden)的激活上前向传播时广播自动扩展了但反向传播时 bias 的梯度必须沿着 hidden 维度求和还原成(batch, 1)。如果你忘了这一步梯度形状就不对但可能因为后续运算的广播机制而看起来能跑实际上已经错了。4. 完整训练流程的搭建实录4.1 数据管道的从零实现真实工程里数据管道往往是性能瓶颈但在 from scratch 阶段我们更关注它的正确性和可理解性。一个最小的数据管道需要解决三件事读取、批处理、打乱。读取这块我建议先用内存里的 NumPy 数组模拟不要一上来就搞复杂的文件格式。批处理的核心是切片但要注意最后一个不完整的 batch 怎么处理——是丢弃还是保留训练时通常丢弃推理时通常保留。打乱这块关键是每个 epoch 重新打乱而且打乱的是索引而不是数据本身这样避免大数组的拷贝。class SimpleDataLoader: def __init__(self, data, labels, batch_size, shuffleTrue): self.data data self.labels labels self.batch_size batch_size self.shuffle shuffle def __iter__(self): n len(self.data) indices np.arange(n) if self.shuffle: np.random.shuffle(indices) for start in range(0, n, self.batch_size): end min(start self.batch_size, n) batch_idx indices[start:end] yield self.data[batch_idx], self.labels[batch_idx]这个实现很朴素但它把数据管道的本质暴露得很清楚。等你理解了它再去看 PyTorch 的DataLoader那些collate_fn、num_workers、sampler的概念就一目了然了。4.2 训练循环的关键控制点训练循环看起来就是前向、算 loss、反向、更新四步但真正写好需要处理很多细节。我列几个最关键的学习率调度。固定学习率往往不是最优的warmup cosine decay 是目前比较通用的方案。warmup 的作用是让模型在初期不要因为随机初始化的大梯度而震荡cosine decay 则是让后期学习率平滑下降帮助收敛。梯度裁剪。尤其是 RNN 和 Transformer梯度爆炸是家常便饭。通常用全局范数裁剪把所有参数的梯度拼成一个向量算它的 L2 范数超过阈值就整体缩放。def clip_gradients(params, max_norm): total_norm 0.0 for p in params: total_norm np.sum(p.grad ** 2) total_norm np.sqrt(total_norm) if total_norm max_norm: scale max_norm / (total_norm 1e-6) for p in params: p.grad * scale return total_norm参数更新。SGD 太慢Adam 是默认选择。但 Adam 的eps参数很关键默认 1e-8 在 fp16 下可能太小导致除零通常要调到 1e-6 甚至 1e-5。4.3 推理阶段的性能优化训练完了要推理推理阶段的核心矛盾是延迟和吞吐。最基础的优化是KV cache自回归生成时每次只生成一个新 token但 attention 需要看所有历史 token。如果不缓存每步都要重算所有历史的 K 和 V复杂度是 O(n²)缓存之后每步只需要算新 token 的 K、V复杂度降到 O(n)。KV cache 的实现要点是维护一个不断增长的 K 和 V 的缓存每次新 token 进来把它的 K、V 追加进去然后用完整的缓存做 attention。这里有个细节是位置编码的处理——缓存之后新 token 的位置索引要正确否则位置编码就错了。另一个优化是批处理推理。把多个请求拼成一个 batch能显著提升 GPU 利用率。但要注意不同请求的生成长度可能不同需要 padding 和 mask而且 padding 的位置不能参与 attention。注意KV cache 会占用大量显存尤其是长序列场景。一个实用的技巧是设置最大缓存长度超过之后用滑动窗口或者丢弃最早的 token。这个取舍要根据具体任务来定。5. 常见问题与排查技巧实录5.1 训练不收敛的排查路径训练不收敛是最高频的问题我一般按这个顺序排查先看数据。把 batch 里的样本打印出来看看标签对不对有没有全零、全一的异常样本。我遇到过好几次是数据预处理写错了归一化用了错误的均值方差。再看 loss。loss 完全不降可能是学习率太大或者太小。太大表现为 loss 震荡甚至爆炸太小表现为 loss 几乎不动。可以先跑几个 step手动调学习率试试。然后看梯度。打印每层梯度的范数如果某层梯度一直是 0可能是激活函数饱和了比如 sigmoid 在两端如果梯度特别大可能是初始化有问题。最后看初始化。权重初始化对训练影响巨大Xavier 和 Kaiming 初始化是有理论依据的不要随便用np.random.randn了事。5.2 显存爆炸的定位与解决显存爆炸通常有几个来源我整理成一个速查表现象可能原因解决方案训练一开始就 OOMbatch size 太大减小 batch用梯度累积训练中途 OOM中间激活没释放检查是否有变量被意外持有长序列 OOMattention 矩阵太大用 flash attention 或分块计算推理 OOMKV cache 太大限制缓存长度或用量化梯度累积是个特别实用的技巧用小 batch 跑多次前向反向累积梯度再统一更新。这样等效于大 batch但显存占用只有小 batch 的水平。实现上就是在更新前不清零梯度累积够步数再清零。5.3 数值不稳定的典型场景数值问题最爱出现在三个地方softmax、log、除法。softmax 前面说过了减最大值。log 的问题是log(0)会变成负无穷通常加一个极小值log(x 1e-8)。除法的问题是分母为 0同样加极小值。还有一个隐蔽的场景是混合精度训练。fp16 的动态范围很窄小梯度会下溢成 0大梯度会溢出成 inf。解决办法是用 loss scaling先放大 loss反向传播得到放大的梯度更新前再缩回去。这样小梯度在放大后就不会下溢了。# 简化的 loss scaling 逻辑 scale 1024.0 scaled_loss loss * scale scaled_loss.backward() for p in params: p.grad / scale5.4 我踩过的几个印象深刻的坑第一个坑是忘记zero_grad。PyTorch 的梯度是累加的如果你不在每个 step 开始时清零梯度会越积越大模型直接飞掉。这个错误新手特别容易犯而且现象很迷惑——前几个 step 看起来正常突然就炸了。第二个坑是在 eval 模式忘了no_grad。推理时如果不禁用梯度会白白占用大量显存而且速度慢很多。更严重的是如果你在推理时不小心更新了 BatchNorm 的统计量模型的行为就变了。第三个坑是随机种子没固定。做实验对比的时候如果种子不固定你根本分不清效果差异是来自你的改动还是随机性。我的习惯是每个实验都固定种子而且要固定所有相关的随机源——Python 的、NumPy 的、框架的。6. 从手写实现到工程落地的距离6.1 手写代码和生产代码的差距必须承认from scratch 写出来的代码离生产可用还有很大距离。手写版本追求的是看得懂生产版本追求的是跑得快、稳得住。这个差距主要体现在几个方面性能。手写的矩阵运算没有经过任何优化没有向量化没有利用 SIMD更没有 GPU 加速。生产代码会用 cuBLAS、cuDNN 这些高度优化的库性能差距可能是几十上百倍。鲁棒性。手写代码基本没有错误处理输入形状不对就直接崩。生产代码需要处理各种边界情况——空输入、超长输入、异常值。可维护性。手写代码往往是一个大文件从头到尾生产代码需要模块化、可测试、可扩展。但这不意味着 from scratch 没意义。恰恰相反只有理解了手写版本的每个细节你才能在生产代码出问题时快速定位。框架帮你做了优化但优化本身也可能引入 bug这时候你的底层知识就是排查的钥匙。6.2 性能优化的渐进路径如果你想让手写代码逐步接近生产性能我建议按这个顺序优化第一步是向量化。把 Python 的 for 循环换成 NumPy 的批量运算这一步通常能带来几十倍的提升。比如计算 attention不要一个 token 一个 token 算要整个矩阵一起算。第二步是减少内存拷贝。NumPy 的很多操作会创建新数组频繁的拷贝很浪费。可以用out参数指定输出数组或者用视图view代替拷贝。第三步是利用 GPU。把 NumPy 换成 PyTorch 的 Tensor.cuda()一下性能立刻上一个台阶。但要注意 CPU 和 GPU 之间的数据传输是瓶颈尽量让数据待在 GPU 上。第四步是算子融合。把多个小算子合并成一个大算子减少 kernel launch 的开销。这一步通常需要写 CUDA门槛较高但收益也大。6.3 这套知识在实际工作中的价值最后说说这套 from scratch 的知识在实际工作中到底有什么用。我的体会是它最大的价值不是让你去手写模型而是让你在遇到问题时知道往哪看。模型不收敛你知道去检查梯度、初始化、学习率推理慢你知道瓶颈可能在 attention 的 O(n²) 或者 KV cache 的显存显存爆炸你知道是激活、梯度还是缓存的问题。这些判断力是调包调不出来的。而且当你需要做一些框架不支持的自定义操作时这套知识就是你的底气。比如你想实现一个新的 attention 变体或者一个新的归一化层你能自己写 forward 和 backward而不是干等框架更新。我个人在实际操作中的体会是from scratch 的学习曲线很陡前期会很痛苦尤其是手写反向传播那一段但一旦跨过去你对整个 AI 系统的理解会发生质变。那种我能看透它的感觉是任何高级教程都给不了的。如果你正在这个阶段我的建议是别急着求快一个组件一个组件地啃每个都用手写代码验证一遍慢就是快。
从零手写AI核心组件：突破调包侠瓶颈的底层实践指南
1. 这个项目到底在解决什么问题第一次看到 ai-engineering-from-scratch 这个标题我脑子里蹦出来的第一个念头是又是一个教人调包的教程但仔细琢磨了一下 from scratch 这四个字我意识到它想做的事情可能完全不一样。市面上讲 AI 工程的内容绝大多数都是从pip install transformers或者import openai开始的你跟着敲一遍模型跑起来了loss 也降下去了但真要问你一句这个 attention 里的 Q、K、V 到底是怎么算出来的很多人就卡壳了。这个项目要解决的恰恰就是这个断层——它试图让你从最底层的矩阵乘法开始一步步把 AI 工程里那些被封装得严严实实的核心组件亲手搭出来。说白了这是一个面向想真正搞懂 AI 系统内部运转机制的工程师的学习路径项目。它适合的人很明确你已经会用 PyTorch 或者 TensorFlow 跑通一些 demo 了但总觉得自己像个调包侠遇到模型不收敛、显存爆炸、推理延迟高这些问题时只能靠搜索引擎和玄学调参。它不适合完全零基础的人——如果你连 Python 的列表推导式都写不利索那还是先去补编程基础。但只要你有一点点深度学习的使用经验想往底层钻一钻这个项目的价值就非常大。我之所以对这个方向特别有感触是因为我自己就经历过那个阶段。早些年做推荐系统的时候模型跑不通我就换优化器换学习率换 batch size折腾一整天可能只是碰巧好了但根本不知道问题出在哪。后来逼着自己手写了一遍反向传播手写了一遍 attention那种原来如此的感觉是看一百篇博客都给不了的。所以当我看到 from scratch 这个定位时我第一反应是这个方向对了它补的是工程能力和原理理解之间那道最容易被忽略的鸿沟。这个项目的核心价值可以拆成三层。第一层是认知层让你知道每个组件为什么存在、解决什么问题第二层是实现层让你能用最基础的算子把它搭出来不依赖高级封装第三层是工程层让你理解真实系统里这些组件是怎么组合、怎么优化、怎么排查问题的。三层缺一不可只讲原理不讲实现是纸上谈兵只讲实现不讲原理是另一种调包。2. 从零构建的核心思路拆解2.1 为什么非要从零而不是从框架很多人会问现在框架这么成熟为什么还要费劲从零写这不是重复造轮子吗我的理解是从零构建的目的从来不是让你在生产环境里真的手写一个 Transformer而是让你获得一种透视能力。当你亲手实现过一遍 softmax 的数值稳定性处理你再看到框架里那个F.softmax的时候你就知道它在背后帮你做了什么遇到 NaN 的时候你也能立刻反应过来可能是哪里溢出了。从工程角度看框架帮你屏蔽了细节但同时也屏蔽了调试的抓手。举个我踩过的坑有一次模型训练到一半 loss 突然变成 NaN我查了半天数据、学习率、梯度裁剪都没问题最后发现是某一层的 attention score 在 fp16 下溢出了。如果我没有手写过 attention我根本不会想到去检查那个位置。这就是从零带来的直觉——你对每个数值的流动范围有了体感。2.2 自底向上的分层设计一个合理的 from scratch 路径我倾向于按这样的层次来组织这也是我在带新人的时候反复验证过的顺序数值计算基础层向量、矩阵、广播机制、数值稳定性。这一层看起来最枯燥但它是后面所有内容的地基。很多人跳过这层直接写模型结果遇到维度对不上、广播出错就懵了。自动微分层计算图、前向传播、反向传播、链式法则的工程实现。这是整个深度学习的心脏理解了它你就理解了为什么 PyTorch 要设计成动态图。神经网络基础组件层全连接层、激活函数、损失函数、优化器。这一层开始有模型的感觉了。序列建模层从 RNN 到 attention 再到 Transformer。这是现代 AI 工程的核心也是最能体现从零价值的地方。训练与推理工程层数据加载、批处理、混合精度、梯度累积、KV cache、量化。这一层直接对应真实生产环境的问题。这个分层的关键在于每一层都只依赖它下面的层你可以在任何一层停下来都获得一个完整可用的能力。比如你只做到自动微分层你就能自己实现一个简单的线性回归训练你做到 Transformer 层你就能手搓一个迷你 GPT。2.3 选型背后的取舍逻辑在实现语言和工具的选择上我的建议是用 Python NumPy 起步逐步过渡到 PyTorch 的底层 API。为什么不直接用 PyTorch 的高层 API因为那样你就又回到调包的老路了。NumPy 的好处是它足够裸矩阵乘法就是矩阵乘法没有 autograd 帮你兜底你必须自己想清楚每一步的梯度怎么传。但纯 NumPy 也有问题——它没有 GPU 加速跑到 Transformer 这种规模就慢得没法忍。所以合理的过渡点是基础层和自动微分层用 NumPy 手写把原理吃透到了神经网络组件层之后开始用 PyTorch 的torch.Tensor和torch.autograd但依然自己定义 forward 和 backward不用nn.Module的高级封装。这样既保证了原理的透明度又获得了工程上的可行性。提示不要一上来就追求性能。from scratch 阶段代码的可读性和正确性远比速度重要。我见过太多人为了优雅把代码写得极其抽象结果自己都看不懂了调试的时候痛苦不堪。3. 核心组件的实操实现要点3.1 自动微分引擎的手写细节自动微分是整个项目里最烧脑但也最值得的部分。核心思路其实不复杂把每个运算都表示成一个节点节点记录自己的输入、输出和局部梯度前向传播时构建计算图反向传播时沿着图反向应用链式法则。听起来简单但魔鬼在细节里。第一个坑是计算图的构建时机。动态图和静态图的区别在你手写一遍之后就彻底明白了。动态图是边算边建图每次前向传播都重新构建灵活但开销大静态图是先建图再算效率高但不灵活。PyTorch 用的是动态图这也是为什么它调试起来比早期的 TensorFlow 舒服。第二个坑是梯度的累加。同一个节点可能被多条路径用到反向传播时它的梯度是所有路径梯度的和。这个累加逻辑如果写错了梯度就会偏而且偏得很隐蔽——模型还能训就是效果差一点你很难发现。我当时的做法是写一个数值梯度检验函数用有限差分法算出来的梯度和反向传播算出来的梯度对比误差在 1e-6 以内才算通过。def numerical_gradient(f, x, eps1e-6): grad np.zeros_like(x) it np.nditer(x, flags[multi_index]) while not it.finished: idx it.multi_index old x[idx] x[idx] old eps fx1 f(x) x[idx] old - eps fx2 f(x) x[idx] old grad[idx] (fx1 - fx2) / (2 * eps) it.iternext() return grad这个函数看起来笨但它是你验证自己 autograd 实现是否正确的唯一可靠手段。我强烈建议每实现一个新的算子都用它检验一遍。3.2 Attention 机制的数值稳定性处理Attention 是 Transformer 的灵魂也是数值问题的高发区。标准公式是softmax(QK^T / sqrt(d_k)) V但直接这么写在 fp16 下几乎必炸。问题出在QK^T的值可能很大指数运算exp之后直接溢出成 inf。解决办法是减去最大值softmax(x) softmax(x - max(x))。数学上等价但数值上稳定得多因为减完之后最大的指数是exp(0) 1不会溢出。这个技巧看起来简单但它是无数人踩坑之后总结出来的你在框架源码里能看到但如果不手写一遍你永远不会真正记住它。def stable_softmax(x, axis-1): x_max np.max(x, axisaxis, keepdimsTrue) exp_x np.exp(x - x_max) return exp_x / np.sum(exp_x, axisaxis, keepdimsTrue)还有一个容易被忽略的点是mask 的处理。在自回归生成里我们需要屏蔽掉未来的 token通常是把对应位置的 score 设成一个极大的负数这样 softmax 之后概率就趋近于 0。但如果你设成-inf在某些实现里会出 NaN。稳妥的做法是设成一个足够大的负数比如-1e9而不是真正的负无穷。3.3 反向传播中的梯度检查清单手写反向传播的时候我总结了一个检查清单每次实现新层都过一遍检查项常见错误验证方法维度匹配梯度形状和参数形状不一致打印 shape 对比链式法则顺序乘法顺序反了数值梯度检验梯度累加多路径梯度没求和构造共享节点的测试用例原地操作修改了前向的中间变量检查是否有 in-place 赋值广播回传广播后的梯度没求和还原对广播维度做 sum这个表格里的每一条都是我真金白银踩出来的。尤其是广播回传这一条特别隐蔽。比如你把一个(batch, 1)的 bias 加到(batch, hidden)的激活上前向传播时广播自动扩展了但反向传播时 bias 的梯度必须沿着 hidden 维度求和还原成(batch, 1)。如果你忘了这一步梯度形状就不对但可能因为后续运算的广播机制而看起来能跑实际上已经错了。4. 完整训练流程的搭建实录4.1 数据管道的从零实现真实工程里数据管道往往是性能瓶颈但在 from scratch 阶段我们更关注它的正确性和可理解性。一个最小的数据管道需要解决三件事读取、批处理、打乱。读取这块我建议先用内存里的 NumPy 数组模拟不要一上来就搞复杂的文件格式。批处理的核心是切片但要注意最后一个不完整的 batch 怎么处理——是丢弃还是保留训练时通常丢弃推理时通常保留。打乱这块关键是每个 epoch 重新打乱而且打乱的是索引而不是数据本身这样避免大数组的拷贝。class SimpleDataLoader: def __init__(self, data, labels, batch_size, shuffleTrue): self.data data self.labels labels self.batch_size batch_size self.shuffle shuffle def __iter__(self): n len(self.data) indices np.arange(n) if self.shuffle: np.random.shuffle(indices) for start in range(0, n, self.batch_size): end min(start self.batch_size, n) batch_idx indices[start:end] yield self.data[batch_idx], self.labels[batch_idx]这个实现很朴素但它把数据管道的本质暴露得很清楚。等你理解了它再去看 PyTorch 的DataLoader那些collate_fn、num_workers、sampler的概念就一目了然了。4.2 训练循环的关键控制点训练循环看起来就是前向、算 loss、反向、更新四步但真正写好需要处理很多细节。我列几个最关键的学习率调度。固定学习率往往不是最优的warmup cosine decay 是目前比较通用的方案。warmup 的作用是让模型在初期不要因为随机初始化的大梯度而震荡cosine decay 则是让后期学习率平滑下降帮助收敛。梯度裁剪。尤其是 RNN 和 Transformer梯度爆炸是家常便饭。通常用全局范数裁剪把所有参数的梯度拼成一个向量算它的 L2 范数超过阈值就整体缩放。def clip_gradients(params, max_norm): total_norm 0.0 for p in params: total_norm np.sum(p.grad ** 2) total_norm np.sqrt(total_norm) if total_norm max_norm: scale max_norm / (total_norm 1e-6) for p in params: p.grad * scale return total_norm参数更新。SGD 太慢Adam 是默认选择。但 Adam 的eps参数很关键默认 1e-8 在 fp16 下可能太小导致除零通常要调到 1e-6 甚至 1e-5。4.3 推理阶段的性能优化训练完了要推理推理阶段的核心矛盾是延迟和吞吐。最基础的优化是KV cache自回归生成时每次只生成一个新 token但 attention 需要看所有历史 token。如果不缓存每步都要重算所有历史的 K 和 V复杂度是 O(n²)缓存之后每步只需要算新 token 的 K、V复杂度降到 O(n)。KV cache 的实现要点是维护一个不断增长的 K 和 V 的缓存每次新 token 进来把它的 K、V 追加进去然后用完整的缓存做 attention。这里有个细节是位置编码的处理——缓存之后新 token 的位置索引要正确否则位置编码就错了。另一个优化是批处理推理。把多个请求拼成一个 batch能显著提升 GPU 利用率。但要注意不同请求的生成长度可能不同需要 padding 和 mask而且 padding 的位置不能参与 attention。注意KV cache 会占用大量显存尤其是长序列场景。一个实用的技巧是设置最大缓存长度超过之后用滑动窗口或者丢弃最早的 token。这个取舍要根据具体任务来定。5. 常见问题与排查技巧实录5.1 训练不收敛的排查路径训练不收敛是最高频的问题我一般按这个顺序排查先看数据。把 batch 里的样本打印出来看看标签对不对有没有全零、全一的异常样本。我遇到过好几次是数据预处理写错了归一化用了错误的均值方差。再看 loss。loss 完全不降可能是学习率太大或者太小。太大表现为 loss 震荡甚至爆炸太小表现为 loss 几乎不动。可以先跑几个 step手动调学习率试试。然后看梯度。打印每层梯度的范数如果某层梯度一直是 0可能是激活函数饱和了比如 sigmoid 在两端如果梯度特别大可能是初始化有问题。最后看初始化。权重初始化对训练影响巨大Xavier 和 Kaiming 初始化是有理论依据的不要随便用np.random.randn了事。5.2 显存爆炸的定位与解决显存爆炸通常有几个来源我整理成一个速查表现象可能原因解决方案训练一开始就 OOMbatch size 太大减小 batch用梯度累积训练中途 OOM中间激活没释放检查是否有变量被意外持有长序列 OOMattention 矩阵太大用 flash attention 或分块计算推理 OOMKV cache 太大限制缓存长度或用量化梯度累积是个特别实用的技巧用小 batch 跑多次前向反向累积梯度再统一更新。这样等效于大 batch但显存占用只有小 batch 的水平。实现上就是在更新前不清零梯度累积够步数再清零。5.3 数值不稳定的典型场景数值问题最爱出现在三个地方softmax、log、除法。softmax 前面说过了减最大值。log 的问题是log(0)会变成负无穷通常加一个极小值log(x 1e-8)。除法的问题是分母为 0同样加极小值。还有一个隐蔽的场景是混合精度训练。fp16 的动态范围很窄小梯度会下溢成 0大梯度会溢出成 inf。解决办法是用 loss scaling先放大 loss反向传播得到放大的梯度更新前再缩回去。这样小梯度在放大后就不会下溢了。# 简化的 loss scaling 逻辑 scale 1024.0 scaled_loss loss * scale scaled_loss.backward() for p in params: p.grad / scale5.4 我踩过的几个印象深刻的坑第一个坑是忘记zero_grad。PyTorch 的梯度是累加的如果你不在每个 step 开始时清零梯度会越积越大模型直接飞掉。这个错误新手特别容易犯而且现象很迷惑——前几个 step 看起来正常突然就炸了。第二个坑是在 eval 模式忘了no_grad。推理时如果不禁用梯度会白白占用大量显存而且速度慢很多。更严重的是如果你在推理时不小心更新了 BatchNorm 的统计量模型的行为就变了。第三个坑是随机种子没固定。做实验对比的时候如果种子不固定你根本分不清效果差异是来自你的改动还是随机性。我的习惯是每个实验都固定种子而且要固定所有相关的随机源——Python 的、NumPy 的、框架的。6. 从手写实现到工程落地的距离6.1 手写代码和生产代码的差距必须承认from scratch 写出来的代码离生产可用还有很大距离。手写版本追求的是看得懂生产版本追求的是跑得快、稳得住。这个差距主要体现在几个方面性能。手写的矩阵运算没有经过任何优化没有向量化没有利用 SIMD更没有 GPU 加速。生产代码会用 cuBLAS、cuDNN 这些高度优化的库性能差距可能是几十上百倍。鲁棒性。手写代码基本没有错误处理输入形状不对就直接崩。生产代码需要处理各种边界情况——空输入、超长输入、异常值。可维护性。手写代码往往是一个大文件从头到尾生产代码需要模块化、可测试、可扩展。但这不意味着 from scratch 没意义。恰恰相反只有理解了手写版本的每个细节你才能在生产代码出问题时快速定位。框架帮你做了优化但优化本身也可能引入 bug这时候你的底层知识就是排查的钥匙。6.2 性能优化的渐进路径如果你想让手写代码逐步接近生产性能我建议按这个顺序优化第一步是向量化。把 Python 的 for 循环换成 NumPy 的批量运算这一步通常能带来几十倍的提升。比如计算 attention不要一个 token 一个 token 算要整个矩阵一起算。第二步是减少内存拷贝。NumPy 的很多操作会创建新数组频繁的拷贝很浪费。可以用out参数指定输出数组或者用视图view代替拷贝。第三步是利用 GPU。把 NumPy 换成 PyTorch 的 Tensor.cuda()一下性能立刻上一个台阶。但要注意 CPU 和 GPU 之间的数据传输是瓶颈尽量让数据待在 GPU 上。第四步是算子融合。把多个小算子合并成一个大算子减少 kernel launch 的开销。这一步通常需要写 CUDA门槛较高但收益也大。6.3 这套知识在实际工作中的价值最后说说这套 from scratch 的知识在实际工作中到底有什么用。我的体会是它最大的价值不是让你去手写模型而是让你在遇到问题时知道往哪看。模型不收敛你知道去检查梯度、初始化、学习率推理慢你知道瓶颈可能在 attention 的 O(n²) 或者 KV cache 的显存显存爆炸你知道是激活、梯度还是缓存的问题。这些判断力是调包调不出来的。而且当你需要做一些框架不支持的自定义操作时这套知识就是你的底气。比如你想实现一个新的 attention 变体或者一个新的归一化层你能自己写 forward 和 backward而不是干等框架更新。我个人在实际操作中的体会是from scratch 的学习曲线很陡前期会很痛苦尤其是手写反向传播那一段但一旦跨过去你对整个 AI 系统的理解会发生质变。那种我能看透它的感觉是任何高级教程都给不了的。如果你正在这个阶段我的建议是别急着求快一个组件一个组件地啃每个都用手写代码验证一遍慢就是快。