☰
从零手搓AI工程:自底向上构建自动微分与Transformer推理引擎
2026/11/3 15:23:03 网站建设 项目流程

1. 从零手搓AI工程:为什么我不建议你直接调包

很多人一提到AI工程,脑子里第一反应就是pip install transformers,然后写个pipeline就觉得自己入门了。我刚开始也这么干过,直到有一次线上推理服务在高峰期直接雪崩,日志里全是显存溢出的报错,我才意识到——只会调包的人,永远不知道模型在底层到底发生了什么。

ai-engineering-from-scratch这个方向,说白了就是把AI工程里那些被封装好的黑盒,一层一层拆开,用最原始的方式重新实现一遍。它不是让你去造一个比PyTorch还快的框架,而是让你在拆解的过程中真正理解:一个张量是怎么在内存里排布的、反向传播的梯度到底怎么流动、注意力机制为什么能work、推理服务的瓶颈究竟卡在哪里。

这套东西适合谁?我认为有三类人值得花时间啃:第一类是刚入行做AI应用开发,只会调API但遇到性能问题就抓瞎的工程师;第二类是想转行做AI基础设施,但缺乏底层认知的后端开发者;第三类是在校学生,课本上学的理论和工业界的实践之间隔着一道鸿沟,这个项目方向正好能帮你把桥搭起来。

我自己走这条路大概花了三个月,从手写一个标量级别的自动微分开始,到后来能独立实现一个简化版的Transformer推理引擎。过程很痛苦,但收获是——现在看到任何一个AI框架的报错,我基本能猜到问题出在哪一层。下面我把整个学习路径和实操细节完整拆解一遍,你照着走能少踩不少坑。

2. 整体学习路径设计与技术选型思路

2.1 为什么选择"自底向上"而不是"自顶向下"

大部分人学AI工程的路径是自顶向下的:先学怎么用HuggingFace加载模型,再学怎么用vLLM做推理加速,最后才可能接触到CUDA kernel。这条路看起来快,但有个致命问题——你永远在别人的抽象层上打转。一旦遇到抽象层没覆盖的场景,比如自定义一个特殊的注意力mask,或者想优化某个算子的内存占用,你就完全无从下手。

自底向上的路径则相反:先搞明白一个标量怎么求导,再扩展到向量、矩阵,然后搭建神经网络,最后才是Transformer和大模型。这条路前期慢,但每一步都是你自己推导出来的,知识是长在身上的,不是借来的。

我选择的技术栈是这样的:Python作为主语言(生态好、上手快),NumPy做第一阶段的张量运算(理解底层原理),PyTorch做第二阶段的对照验证(确认自己的实现是否正确),CUDA C++作为可选的进阶(理解GPU并行)。这个选型逻辑是——先用纯Python和NumPy把算法逻辑跑通,不引入任何框架的魔法,等逻辑完全正确了,再逐步引入性能优化。

2.2 分阶段目标拆解

整个学习路径我分成了五个阶段,每个阶段都有明确的产出物,避免学到最后变成"什么都懂一点但什么都做不出来"。

阶段核心目标产出物预计耗时
第一阶段理解自动微分手写标量版autograd引擎1周
第二阶段张量运算与神经网络基于NumPy的MLP实现2周
第三阶段注意力与Transformer手写单头/多头注意力2周
第四阶段训练工程化完整训练循环+梯度检查2周
第五阶段推理优化KV Cache+量化推理3周

这个时间安排是基于每天投入2-3小时来估算的。如果你是全职学习,可以压缩到一个月左右。但我要提醒一句:不要为了赶进度跳过任何一个阶段,尤其是第一阶段。我见过太多人觉得自动微分"不就是链式法则嘛"然后直接跳到Transformer,结果在反向传播调试的时候完全不知道梯度是从哪来的。

2.3 环境准备与工具链

环境这块我不建议搞太复杂,初期一个干净的Python虚拟环境就够了。但有几个细节要注意:

python -m venv ai-scratch source ai-scratch/bin/activate pip install numpy matplotlib pytest

为什么装pytest?因为你在手写这些底层组件的时候,单元测试是你唯一的救命稻草。我当时的做法是,每实现一个算子,就写一个测试用例,用PyTorch的对应算子做数值对照。比如我实现了一个矩阵乘法,就用np.allclose(my_matmul(a,b), torch.matmul(a,b))来验证。这个习惯帮我省了无数调试时间。

注意:不要一开始就装CUDA和cuDNN,那会让你陷入环境配置的泥潭。前四个阶段纯CPU就够了,等到第五阶段做推理优化的时候再考虑GPU。

3. 核心模块拆解与实操要点

3.1 自动微分引擎:一切的起点

自动微分是整个深度学习的地基。你调loss.backward()的时候,PyTorch在背后做的事情,就是沿着计算图反向传播梯度。手写一个简化版的autograd引擎,能让你彻底理解这件事。

核心思路是这样的:每个张量(或者标量)都记录两样东西——它的值和它的梯度,以及它是怎么被计算出来的(也就是它的父节点和对应的反向传播函数)。当你调用backward()的时候,从最终的loss节点开始,沿着计算图反向遍历,每个节点调用自己的反向函数,把梯度传给父节点。

我当时的实现大概长这样:

class Value: def __init__(self, data, children=(), op=''): self.data = data self.grad = 0.0 self._backward = lambda: None self._prev = set(children) self._op = op def __add__(self, other): other = other if isinstance(other, Value) else Value(other) out = Value(self.data + other.data, (self, other), '+') def _backward(): self.grad += out.grad other.grad += out.grad out._backward = _backward return out def __mul__(self, other): other = other if isinstance(other, Value) else Value(other) out = Value(self.data * other.data, (self, other), '*') def _backward(): self.grad += other.data * out.grad other.grad += self.data * out.grad out._backward = _backward return out def backward(self): topo = [] visited = set() def build_topo(v): if v not in visited: visited.add(v) for child in v._prev: build_topo(child) topo.append(v) build_topo(self) self.grad = 1.0 for node in reversed(topo): node._backward()

这段代码看起来简单,但里面有几个关键点值得展开说。第一,为什么用+=而不是=?因为一个节点可能被多条路径使用,梯度需要累加。这是很多人手写反向传播时最容易犯的错误,写成=会导致梯度被覆盖,训练完全不收敛。第二,为什么要做拓扑排序?因为反向传播必须保证一个节点的所有下游都处理完了才能处理它自己,拓扑排序保证了正确的计算顺序。

实操心得:写完这个引擎后,一定要用数值梯度做验证。具体做法是取一个很小的epsilon(比如1e-6),计算(f(x+eps) - f(x-eps)) / (2*eps),和你反向传播算出来的梯度对比。如果误差在1e-6以内,说明实现正确。这个技巧叫梯度检查,是训练工程里非常重要的调试手段。

3.2 张量运算与广播机制

从标量扩展到张量,最大的坑就是广播机制。NumPy的广播规则看起来简单,但当你需要手写反向传播的时候,广播的反向操作(也就是梯度的求和降维)很容易写错。

广播的核心规则是:从右往左对齐维度,维度大小为1的可以扩展成任意大小,缺失的维度视为1。比如形状(3, 1)和(1, 4)相加,结果是(3, 4)。反向传播的时候,(3, 4)的梯度要分别降维回(3, 1)和(1, 4),也就是沿着被广播的维度求和。

def broadcast_backward(grad, original_shape): # 把被广播掉的维度求和回去 while grad.ndim > len(original_shape): grad = grad.sum(axis=0) for i, dim in enumerate(original_shape): if dim == 1: grad = grad.sum(axis=i, keepdims=True) return grad

这个函数我调了整整一个下午才写对。关键点在于先处理维度数量不一致的情况,再处理维度为1的情况,顺序反了就会出错。

3.3 注意力机制的手写实现

注意力机制是Transformer的核心,也是面试里被问得最多的部分。手写一遍你会发现,公式Attention(Q,K,V) = softmax(QK^T/√d)V看起来简单,但实现的时候有一堆细节。

首先是缩放因子√d。为什么要除以√d?因为当维度d很大的时候,Q和K的点积结果会很大,经过softmax之后会变得非常尖锐(接近one-hot),导致梯度消失。除以√d可以把点积的方差拉回到1附近,保持softmax的输出分布比较平滑。

其次是mask的处理。在decoder的自注意力里,需要屏蔽掉未来位置的信息,做法是在softmax之前把对应位置的分数设成负无穷。这里有个坑:如果你用-1e9而不是-inf,在某些极端情况下softmax之后仍然会有微小的概率泄漏。我建议直接用float('-inf'),或者用torch.finfo(dtype).min。

def scaled_dot_product_attention(Q, K, V, mask=None): d_k = Q.shape[-1] scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) attn = softmax(scores, axis=-1) return attn @ V

注意:手写softmax的时候一定要做数值稳定处理,也就是先减去每行的最大值再取指数。不然遇到大数值会直接溢出成inf。这个细节在框架里是自动处理的,但你自己写的时候必须手动加上。

3.4 训练循环的工程化细节

很多人觉得训练循环就是forward -> loss -> backward -> step四步走,但真正工程化的时候,魔鬼全在细节里。

梯度累积是第一个要处理的。当显存不够放不下大batch的时候,可以把一个大batch拆成几个小batch,分别计算梯度然后累加,最后再统一更新。这里的关键是在累积过程中不能清零梯度,只在更新之后清零。

学习率调度是第二个。我常用的是warmup + cosine decay:前10%的步数线性升温,后面余弦衰减。warmup的作用是避免训练初期梯度不稳定导致模型发散,cosine decay则是让模型在后期精细收敛。

梯度裁剪是第三个。当梯度范数超过阈值的时候,按比例缩放梯度。这个操作在训练RNN和Transformer的时候几乎是必须的,不然一个异常batch就能把模型带偏。

def train_step(model, optimizer, batch, grad_accum_steps=4, max_grad_norm=1.0): for i, micro_batch in enumerate(split(batch, grad_accum_steps)): loss = compute_loss(model, micro_batch) / grad_accum_steps loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm) optimizer.step() optimizer.zero_grad()

4. 推理优化:从能跑到跑得快

4.1 KV Cache的原理与实现

训练完了模型,下一步就是推理。如果你直接每次生成一个token就把整个序列重新跑一遍forward,那速度会慢到无法接受。KV Cache的核心思想是:自回归生成的时候,前面token的Key和Value是不变的,可以缓存起来复用,只需要计算新token的Query。

这个优化能把推理复杂度从O(n²)降到O(n),效果非常显著。但实现的时候要注意内存管理——KV Cache会随着序列长度线性增长,长序列场景下可能比模型本身还占内存。我当时的做法是预分配一个固定大小的cache,用滑动窗口的方式管理。

4.2 量化推理的基本思路

量化是把模型的权重从FP16/FP32降到INT8甚至INT4,用精度换速度和内存。最粗暴的方式是直接round到最近的整数,但这样精度损失很大。实践中常用的是per-channel量化,也就是每个通道单独计算scale和zero_point。

def quantize_int8(tensor): scale = tensor.abs().max() / 127 zero_point = 0 quantized = (tensor / scale).round().clamp(-128, 127).to(torch.int8) return quantized, scale, zero_point def dequantize_int8(quantized, scale, zero_point): return (quantized.to(torch.float32) - zero_point) * scale

实操心得:量化之后一定要做精度回归测试。我遇到过量化后模型在某个特定任务上精度掉了10个点的情况,后来发现是那一层的权重分布特别不均匀,per-tensor量化完全hold不住。改成per-channel之后就正常了。

4.3 批处理与连续批处理

推理服务的吞吐量很大程度上取决于批处理策略。静态批处理是攒够一个batch再一起推理,但问题是不同请求的生成长度不一样,短的请求要等长的请求,浪费算力。连续批处理(continuous batching)的思路是:不等整个batch完成,只要有请求结束就立刻把新的请求塞进去,保持GPU始终满载。

这个技术在vLLM里叫PagedAttention,核心是把KV Cache分成固定大小的block来管理,像操作系统的虚拟内存一样。手写一个简化版能让你彻底理解为什么vLLM比朴素实现快那么多。

5. 常见问题与排查技巧实录

5.1 梯度相关问题的排查

梯度问题是手写AI工程里最常见的坑,我整理了一个速查表:

现象可能原因排查方法
loss不下降学习率过大/过小打印梯度范数,观察是否爆炸或消失
loss变成NaN数值溢出检查softmax、log等算子是否做了稳定处理
梯度全为0计算图断裂检查是否误用了.detach()或.data
梯度爆炸缺少裁剪加梯度裁剪,检查初始化方式
训练正常但推理异常BN/Dropout模式确认调用了model.eval()

梯度检查是我最常用的手段。具体做法是:取一个参数,手动扰动一个极小量,计算数值梯度,和你反向传播算出来的解析梯度对比。如果相对误差大于1e-4,基本可以确定反向传播实现有问题。

5.2 内存问题的定位

显存溢出是另一个高频问题。我的排查顺序是这样的:先看是不是batch size太大,再看是不是中间激活值没释放,最后看是不是有内存泄漏(比如在循环里不断往list里append tensor)。

一个很隐蔽的坑是Python的引用计数导致显存不释放。比如你把loss存到一个list里做记录,那整个计算图都会被保留,显存直接爆炸。正确做法是存loss.item()而不是loss本身。

5.3 数值精度的陷阱

FP16训练的时候经常会遇到loss scale的问题。因为FP16的表示范围比FP32小很多,小梯度直接下溢成0。解决方案是用动态loss scaling:先放大loss,反向传播得到放大的梯度,更新前再缩回去。如果检测到inf或nan,就跳过这一步并降低scale。

这个机制在PyTorch的amp里是自动的,但你自己实现的时候必须手动处理。我当时的做法是维护一个scale因子,每次更新前检查梯度里有没有inf,有就跳过并减半,连续N步没出问题就翻倍。

6. 我在这条路上踩过的几个大坑

第一个坑是过早优化。我一开始就想用CUDA写kernel,结果卡在环境配置上整整一周,算法逻辑反而没推进。后来我调整策略,先用纯Python把逻辑跑通,确认正确之后再考虑性能。这个顺序非常重要,正确性永远优先于性能。

第二个坑是不做单元测试。我手写多头注意力的时候,觉得逻辑很清晰就直接跑训练了,结果loss一直不降。debug了两天才发现是reshape的时候维度顺序搞错了。如果一开始就写个测试用例和PyTorch对照,五分钟就能定位问题。

第三个坑是忽视数值稳定性。softmax不减最大值、log不加epsilon、除法不检查零——这些在小规模测试的时候都不会出问题,但一旦数据规模上去就全暴露了。现在我的习惯是,任何涉及指数、对数、除法的算子,第一版就加上数值稳定处理。

第四个坑是只看不写。我早期看了大量Transformer的讲解文章,觉得自己懂了,但真正动手写的时候发现处处是问题。看懂和写出来之间隔着一条巨大的鸿沟,只有亲手实现过,知识才真正属于你。

7. 后续可以继续深入的方向

把上面这些走完之后,你基本对AI工程的底层有了扎实的理解。接下来有几个方向可以继续深入:一是分布式训练,理解数据并行、模型并行、流水线并行的原理和通信开销;二是算子融合与图优化,看看TVM、TensorRT这些编译器是怎么把多个算子合并成一个kernel的;三是自定义CUDA kernel,针对特定算子做极致的性能优化。

我个人的建议是,不要贪多,选一个方向深挖。比如你对推理性能感兴趣,就把KV Cache、量化、连续批处理这条线吃透,做到能独立优化一个真实模型的服务延迟。这种深度比广度更有价值,也更容易在面试和实际工作中体现出来。

最后分享一个我自己的学习习惯:每学完一个模块,我都会写一篇技术笔记,用自己的话把原理和实现讲一遍。如果讲不清楚,说明还没真正理解。这个习惯逼着我不断回头补基础,也是我能把这套东西坚持下来的关键。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询