首页 / 资讯中心 / 文章详情

从零构建AI工程能力:手写自动微分与Transformer推理引擎实战

从零构建AI工程能力:手写自动微分与Transformer推理引擎实战 ★ FEATURED ARTICLE
1. 从零搭建AI工程能力为什么我劝你别再当“调包侠”这两年AI岗位的招聘需求翻了何止三倍但真正能扛住面试官追问的人少得可怜。我面过不下五十个候选人简历上清一色写着“熟悉PyTorch、TensorFlow”结果一问“手写一个带反向传播的两层神经网络”能当场在白板上推明白的不到两成。这就是“ai-engineering-from-scratch”这个项目标题背后最扎心的现实——绝大多数人所谓的AI工程能力其实只是熟练使用几个高层API一旦底层出问题连日志都看不懂。所谓“from scratch”不是让你从晶体管开始造计算机而是从数学原理、数据结构、训练循环、工程部署这四个层面把AI系统的每一块积木都亲手搭一遍。这个项目适合三类人第一类是刚入行的算法工程师想摆脱“只会调参”的标签第二类是有后端或数据背景的开发者想转AI方向但被各种框架的抽象层绕晕第三类是在校学生课程只教了理论没告诉你工业界真正跑一个模型要踩多少坑。我花了大概三个月时间把这条路径完整走了一遍从用NumPy实现第一个线性回归到用纯Python写一个能跑通的Transformer推理引擎再到把模型封装成可观测的服务。整个过程踩的坑比我想象的多得多但也正是这些坑让我真正理解了什么叫“工程能力”。下面我把这套方法拆开揉碎讲清楚你照着做至少能少走半年弯路。2. 整体设计思路为什么我选择“自底向上”而不是“自顶向下”2.1 主流学习路径的致命缺陷市面上大部分AI教程走的是“自顶向下”路线先教你调sklearn的fit和predict再教你用Keras搭几层网络最后告诉你“底层原理很复杂先会用就行”。这种路径的问题在于你永远在黑盒外面打转。模型不收敛你不知道是学习率问题还是梯度消失推理延迟高你不知道是算子实现问题还是内存布局问题。我试过带一个新人用这种方式入门三个月后他能跑通图像分类的demo但当我让他把模型部署到边缘设备上、把延迟从200ms压到50ms时他完全无从下手。因为他不知道卷积在底层是怎么算的不知道内存复用和算子融合是什么更不知道量化到底损失了什么。“from scratch”的思路完全反过来先理解最小可运行单元再逐层向上抽象。就像学开车你得先知道离合器、油门、变速箱各自干什么再去学怎么在高速上超车。这个顺序不能反。2.2 四层能力模型的设计逻辑我把整个学习路径拆成四个层次每一层都有明确的产出物和验收标准层次核心内容产出物验收标准数学与数据结构层矩阵运算、自动微分、张量存储纯NumPy实现的自动微分引擎能对任意标量函数求导并验证模型与训练层前向传播、反向传播、优化器手写两层神经网络完成分类在MNIST上准确率95%架构与推理层Attention、Transformer、KV Cache纯Python的GPT推理引擎能加载预训练权重并生成文本工程与部署层服务封装、性能优化、监控带指标暴露的推理服务QPS100P99延迟100ms这个分层的关键在于每一层都必须能独立运行。你写自动微分引擎时不需要任何深度学习框架写推理引擎时不需要PyTorch的nn.Module。这种“断奶式”训练才能真正逼你理解每个环节的本质。2.3 为什么不用现成框架“抄近路”有人会问既然PyTorch已经这么成熟了为什么还要自己造轮子我的回答是造轮子不是为了用而是为了懂。你不需要在生产环境用自己写的自动微分但你必须知道PyTorch的autograd在背后做了什么否则你连loss.backward()为什么报“梯度为None”都搞不清楚。我自己的经验是手写一遍反向传播之后再看PyTorch的retain_graph参数、detach()方法、torch.no_grad()上下文瞬间就通透了。这些概念在文档里看十遍不如自己踩一次“计算图被释放”的坑来得深刻。3. 核心细节解析从零实现自动微分引擎的完整思路3.1 计算图到底该怎么建自动微分的核心是计算图。每个张量是图上的一个节点每个运算是一条边。前向传播时记录运算关系反向传播时沿着图反向求导。听起来简单但实现起来有几个关键决策。第一个决策是用静态图还是动态图。静态图先定义再运行适合优化但调试困难动态图边运行边建图调试友好但性能略差。我选择动态图因为学习阶段调试便利性远比性能重要。PyTorch早期能打败TensorFlow靠的就是动态图的易用性。第二个决策是节点存储什么。每个节点需要存数据data、梯度grad、产生该节点的运算_op、父节点列表_parents、以及反向传播函数_backward。这里有个容易忽略的点梯度必须累加而不是覆盖。因为一个张量可能被多个下游节点使用反向传播时它的梯度是各分支梯度之和。我一开始写成覆盖结果多分支网络训练完全不收敛排查了一整天才发现这个问题。class Tensor: def __init__(self, data, _parents(), _op): self.data np.array(data, dtypenp.float64) self.grad np.zeros_like(self.data) self._backward lambda: None self._parents _parents self._op _op def __add__(self, other): other other if isinstance(other, Tensor) else Tensor(other) out Tensor(self.data other.data, (self, other), ) def _backward(): self.grad out.grad other.grad out.grad out._backward _backward return out上面这段代码里self.grad out.grad就是累加的关键。如果你写成self.grad out.grad多分支场景下梯度就错了。3.2 反向传播的拓扑排序为什么必不可少反向传播必须从输出往输入方向、按拓扑序依次调用每个节点的_backward函数。如果你随便遍历节点可能出现某个节点的梯度还没从所有下游收集完就提前传播给上游了结果就是梯度不完整。拓扑排序的实现有两种一种是DFS后序遍历一种是Kahn算法。我推荐DFS后序代码更简洁def backward(self): topo [] visited set() def build_topo(v): if v not in visited: visited.add(v) for parent in v._parents: build_topo(parent) topo.append(v) build_topo(self) self.grad np.ones_like(self.data) for node in reversed(topo): node._backward()这里有个细节输出节点的初始梯度设为1。因为∂L/∂L 1这是链式法则的起点。很多人会忘记这一步导致所有梯度都是0。3.3 广播机制带来的梯度求和陷阱NumPy的广播机制让不同形状的张量能直接运算但反向传播时必须把广播后的梯度还原回原始形状。比如形状(3,4)的张量和形状(4,)的偏置相加前向传播时偏置被广播成(3,4)反向传播时偏置的梯度应该是(3,4)梯度在第一个维度上求和得到(4,)。这个坑我踩过两次。第一次是偏置梯度形状不对第二次是批量矩阵乘法时梯度维度对不上。解决方法是在每个运算的_backward里显式处理形状还原def __mul__(self, other): other other if isinstance(other, Tensor) else Tensor(other) out Tensor(self.data * other.data, (self, other), *) def _backward(): self.grad _unbroadcast(out.grad * other.data, self.data.shape) other.grad _unbroadcast(out.grad * self.data, other.data.shape) out._backward _backward return out def _unbroadcast(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_unbroadcast这个函数是通用解法任何涉及广播的运算都可以复用它。写一次后面所有运算都省心。注意梯度累加前必须确保形状一致否则NumPy会静默广播导致梯度悄悄出错。建议在开发阶段加断言检查形状。4. 实操过程手写两层神经网络完成MNIST分类4.1 数据加载与预处理的关键细节MNIST数据集虽然简单但预处理有几个容易翻车的点。第一是归一化像素值从0-255缩放到0-1否则梯度爆炸。第二是标签格式如果用交叉熵损失标签应该是one-hot编码如果用稀疏交叉熵标签是整数索引。我选择one-hot因为手写交叉熵更直观。第三是数据打乱每个epoch开始前必须打乱训练集顺序否则模型会学到样本顺序的虚假相关性。我一开始忘了打乱训练准确率卡在60%上不去排查了半天才发现是数据顺序问题。def load_mnist(): from sklearn.datasets import fetch_openml X, y fetch_openml(mnist_784, version1, return_X_yTrue, as_frameFalse) X X.astype(np.float64) / 255.0 y y.astype(np.int64) X_train, X_test X[:60000], X[60000:] y_train, y_test y[:60000], y[60000:] y_train_onehot np.zeros((len(y_train), 10)) y_train_onehot[np.arange(len(y_train)), y_train] 1 return X_train, y_train_onehot, X_test, y_test4.2 网络结构设计与参数初始化两层网络的结构是输入784维 → 隐藏层128维ReLU激活→ 输出10维Softmax。参数量计算784×128 128 128×10 10 101770个参数。这个规模用纯NumPy训练完全可行一个epoch大概30秒。参数初始化用He初始化因为ReLU激活函数会把负半轴梯度清零如果用标准正态分布初始化深层网络的方差会逐层衰减。He初始化把权重标准差设为√(2/fan_in)保证前向传播时每层输出方差稳定。def init_params(): W1 np.random.randn(784, 128) * np.sqrt(2.0 / 784) b1 np.zeros(128) W2 np.random.randn(128, 10) * np.sqrt(2.0 / 128) b2 np.zeros(10) return W1, b1, W2, b2这里b1和b2初始化为0是安全的因为ReLU和Softmax对偏置的梯度不会消失。但如果你用Sigmoid激活偏置也建议小随机初始化。4.3 前向传播与损失计算的数值稳定性Softmax 交叉熵的组合有个经典陷阱指数运算溢出。如果logits里有大于88的值np.exp就会返回inf。解决方法是在Softmax之前减去每行的最大值def softmax(x): x_shifted x - np.max(x, axis1, keepdimsTrue) exp_x np.exp(x_shifted) return exp_x / np.sum(exp_x, axis1, keepdimsTrue)这个操作数学上等价因为Softmax对输入平移不变但数值上稳定得多。交叉熵损失也要用log_softmax的形式计算避免先算Softmax再取log导致的精度损失def cross_entropy_loss(logits, y_onehot): log_probs logits - np.max(logits, axis1, keepdimsTrue) log_probs log_probs - np.log(np.sum(np.exp(log_probs), axis1, keepdimsTrue)) return -np.mean(np.sum(y_onehot * log_probs, axis1))4.4 反向传播的逐层推导与实现反向传播的推导必须自己动手推一遍否则永远记不住。以两层网络为例输出层梯度∂L/∂z2 softmax(z2) - y这是Softmax交叉熵的经典结论推导过程用链式法则展开即可。隐藏层梯度∂L/∂z1 (∂L/∂z2 · W2ᵀ) ⊙ ReLU(z1)其中ReLU在z10时为1否则为0。参数梯度∂L/∂W2 a1ᵀ · ∂L/∂z2∂L/∂b2 sum(∂L/∂z2, axis0)。def backward(X, y_onehot, cache): a1, z1, a2, z2 cache m X.shape[0] dz2 (softmax(z2) - y_onehot) / m dW2 a1.T dz2 db2 np.sum(dz2, axis0) da1 dz2 W2.T dz1 da1 * (z1 0) dW1 X.T dz1 db1 np.sum(dz1, axis0) return dW1, db1, dW2, db2注意dz2除以了m因为损失是batch平均。如果你忘了这个除法梯度会大m倍学习率需要相应调整。4.5 训练循环与超参数调优实录训练循环的核心是前向传播 → 计算损失 → 反向传播 → 更新参数。超参数方面我实测下来比较稳的配置是batch size 64学习率0.1配合SGD训练20个epoch。学习率太大会震荡太小收敛慢。你可以用学习率预热前5个epoch用0.01后面逐步升到0.1。W1, b1, W2, b2 init_params() lr 0.1 batch_size 64 for epoch in range(20): indices np.random.permutation(len(X_train)) for i in range(0, len(X_train), batch_size): batch_idx indices[i:ibatch_size] X_batch X_train[batch_idx] y_batch y_train_onehot[batch_idx] cache forward(X_batch, W1, b1, W2, b2) loss cross_entropy_loss(cache[3], y_batch) grads backward(X_batch, y_batch, cache) W1 - lr * grads[0] b1 - lr * grads[1] W2 - lr * grads[2] b2 - lr * grads[3]我实测这个配置在MNIST测试集上能到97.5%的准确率。如果你想更高可以加动量或者用Adam但那就偏离“from scratch”的初衷了。实操心得训练初期loss不下降先检查数据预处理和梯度形状训练中期loss震荡调小学习率训练后期loss不降可能是过拟合加Dropout或L2正则。5. 从训练到推理手写Transformer推理引擎的核心要点5.1 为什么推理和训练是两码事训练时你关心的是梯度能不能算对、loss能不能降下去推理时你关心的是延迟、吞吐、内存。这两个阶段的优化目标完全不同。训练可以用大batch、混合精度、梯度累积推理必须考虑KV Cache、算子融合、量化。我见过太多人训练模型一把梭部署时发现推理速度慢十倍。原因就是训练和推理的计算图不一样。比如Dropout只在训练时启用BatchNorm在推理时用滑动平均Attention在推理时可以用KV Cache避免重复计算。5.2 KV Cache的原理与内存计算自回归生成时每生成一个token都要重新计算所有历史token的Key和Value这是巨大的浪费。KV Cache的思路是把历史token的K和V缓存起来新token只需要计算自己的Q、K、V然后和缓存的K、V做Attention。内存占用计算假设模型有L层每层有H个注意力头每个头维度为D序列长度为Sbatch size为B精度为FP162字节。KV Cache总内存 2 × L × H × D × S × B × 2字节。以GPT-2为例L12H12D64S1024B1算下来约37MB。序列越长内存越大这就是为什么长上下文推理这么吃显存。class KVCache: def __init__(self, num_layers, num_heads, head_dim, max_seq_len): self.cache [ { k: np.zeros((num_heads, max_seq_len, head_dim)), v: np.zeros((num_heads, max_seq_len, head_dim)) } for _ in range(num_layers) ] self.seq_len 0 def update(self, layer_idx, new_k, new_v): self.cache[layer_idx][k][:, self.seq_len] new_k self.cache[layer_idx][v][:, self.seq_len] new_v5.3 Attention的纯NumPy实现Attention的核心公式是softmax(QKᵀ/√d)V。实现时有几个细节第一是缩放因子√d不能忘否则点积结果太大会让Softmax饱和第二是mask处理生成时不能看到未来的token第三是数值稳定性Softmax前要减最大值。def attention(Q, K, V, maskNone): d_k Q.shape[-1] scores Q K.transpose(0, 2, 1) / np.sqrt(d_k) if mask is not None: scores scores mask * -1e9 scores scores - np.max(scores, axis-1, keepdimsTrue) weights np.exp(scores) / np.sum(np.exp(scores), axis-1, keepdimsTrue) return weights Vmask通常是一个上三角矩阵未来位置为负无穷Softmax后权重为0。注意-1e9而不是-inf因为-inf参与运算会产生NaN。5.4 文本生成策略与采样技巧生成文本时最简单的策略是贪心搜索每次选概率最大的token。但贪心搜索容易生成重复、无聊的文本。更好的策略是温度采样把logits除以温度TT1让分布更尖锐T1让分布更平坦。T0.7通常是比较好的平衡点。还有Top-K采样只从概率最高的K个token中采样避免长尾噪声。Top-P采样核采样从累积概率达到P的最小token集合中采样动态调整候选集大小。我实测Top-P0.9配合温度0.8生成质量比较稳定。def sample(logits, temperature0.8, top_p0.9): logits logits / temperature probs softmax(logits) sorted_idx np.argsort(probs)[::-1] sorted_probs probs[sorted_idx] cumsum np.cumsum(sorted_probs) cutoff np.searchsorted(cumsum, top_p) sorted_probs[cutoff1:] 0 sorted_probs / np.sum(sorted_probs) chosen np.random.choice(sorted_idx, psorted_probs) return chosen注意采样前必须把logits转成概率分布且要处理数值稳定性。温度不能设得太低否则退化成贪心搜索也不能太高否则生成乱码。6. 工程化落地把模型封装成可观测的推理服务6.1 服务框架选型与性能对比推理服务的框架选择很多Flask简单但性能差FastAPI异步支持好Triton Inference Server功能全但学习曲线陡。我推荐从FastAPI起步因为它异步支持好、自动生成API文档、生态成熟。实测下来FastAPI在单机上的QPS能到500以上足够中小规模场景。如果你追求极致性能可以考虑用ONNX Runtime或TensorRT做推理后端FastAPI只做请求转发。但那是另一个话题了from scratch阶段先把服务跑通再说。6.2 请求批处理与动态合并推理服务的吞吐瓶颈往往不在计算而在请求调度。单个请求进来就推理一次GPU利用率极低。解决方案是动态批处理把短时间内到达的多个请求合并成一个batch一起推理再拆分结果返回。实现上可以用一个队列加一个后台worker请求进来先入队worker每隔几毫秒取一批凑够batch size或超时就推理。这样能把吞吐提升5-10倍。import asyncio from collections import deque class BatchScheduler: def __init__(self, max_batch_size32, max_wait_ms10): self.queue deque() self.max_batch_size max_batch_size self.max_wait_ms max_wait_ms async def add_request(self, input_data): future asyncio.Future() self.queue.append((input_data, future)) return await future async def process_loop(self): while True: if not self.queue: await asyncio.sleep(0.001) continue batch [] start asyncio.get_event_loop().time() while self.queue and len(batch) self.max_batch_size: batch.append(self.queue.popleft()) if (asyncio.get_event_loop().time() - start) * 1000 self.max_wait_ms: break inputs [item[0] for item in batch] results self.model_inference(inputs) for (_, future), result in zip(batch, results): future.set_result(result)6.3 监控指标暴露与告警配置服务上线后必须能观测。核心指标有四个QPS、P99延迟、错误率、GPU利用率。QPS和延迟用Prometheus的Histogram记录错误率用CounterGPU利用率用pynvml采集。from prometheus_client import Histogram, Counter, start_http_server REQUEST_LATENCY Histogram(inference_latency_seconds, Inference latency) REQUEST_COUNT Counter(inference_requests_total, Total requests, [status]) app.post(/predict) async def predict(request: Request): start time.time() try: result await model.predict(request) REQUEST_COUNT.labels(statussuccess).inc() return result except Exception as e: REQUEST_COUNT.labels(statuserror).inc() raise finally: REQUEST_LATENCY.observe(time.time() - start)告警规则建议设三条P99延迟超过200ms告警、错误率超过1%告警、GPU利用率持续低于20%告警说明资源浪费。6.4 常见线上问题与排查手册线上问题排查讲究快准狠。我整理了一份速查表现象可能原因排查方法解决方案延迟突然飙升请求量突增或内存泄漏看QPS曲线和内存曲线扩容或重启服务错误率上升输入格式异常或模型加载失败看错误日志和输入样本加输入校验和重试吞吐上不去batch size太小或GPU未打满看GPU利用率和batch分布调大batch或加动态批处理生成结果重复温度太低或重复惩罚未开检查采样参数调高温度或加重复惩罚内存持续增长KV Cache未释放或请求泄漏看内存增长曲线加LRU淘汰或定期重启实操心得线上服务一定要加输入长度限制否则一个超长请求就能把内存打满。我吃过这个亏一个用户传了10万字的输入直接把服务搞崩了。7. 踩坑记录与能力进阶路线7.1 那些让我熬夜的典型Bug第一个坑是梯度爆炸。手写网络训练时loss突然变成NaN排查发现是学习率太大加上没有梯度裁剪。解决方法是在反向传播后加np.clip(grad, -1, 1)或者用更小的学习率。第二个坑是内存泄漏。推理服务跑几个小时后内存爆满排查发现是KV Cache没有按请求释放。每个请求的cache用完必须显式清理否则会一直累积。第三个坑是数值精度。用FP16推理时某些层的输出会溢出。解决方法是关键层用FP32或者用动态缩放loss scaling技术。第四个坑是并发竞争。多个请求同时访问同一个模型实例导致状态混乱。解决方法是每个请求用独立的session或者加锁串行化。7.2 从“能跑”到“跑得好”的进阶方向当你把上面这些全部跑通后可以往这几个方向深入算子融合把多个小算子合并成一个大算子减少内存访问、量化INT8量化能把模型压缩4倍速度提升2-3倍、蒸馏用大模型教小模型保持效果的同时减小体积、分布式推理张量并行、流水线并行支撑超大模型。每一个方向都够写一本书但核心思路是一样的理解底层原理才能做出正确的工程取舍。这也是“from scratch”最大的价值——它给你的不是某个具体技能而是一种看透本质的能力。我个人在实际操作中的体会是手写一遍之后再看任何AI框架的源码都不再发怵。你知道每个模块在干什么知道哪里可能出问题知道怎么优化。这种底气是调包永远给不了的。最后再分享一个小技巧把你手写的每个模块都加上单元测试尤其是梯度检查用数值梯度验证解析梯度能帮你省下大量调试时间。
阅读完成 · 觉得有帮助?
咨询建站