1. 从零手搓AI工程为什么我不建议你直接调包1.1 一个让我彻底改变主意的真实场景去年帮一个朋友排查线上推理服务的问题现象很典型模型在测试集上指标漂亮得不行一上生产环境延迟直接飙到800msGPU利用率却只有30%出头。他第一反应是“模型太大换个小模型”第二反应是“加机器”。我让他把推理链路完整跑一遍profiling结果发现瓶颈根本不在模型本身——数据预处理阶段有个Python循环在做逐样本的tokenize单条耗时占了整个链路的60%以上。这个问题换成任何调包方案都看不出来因为框架把细节全封装了你只能看到“输入进去、输出出来”中间发生了什么完全是黑盒。这件事让我重新审视一个被很多人忽略的事实AI工程不是调包工程。你可以用HuggingFace的pipeline三行代码跑通一个demo但一旦进入真实生产环境延迟、吞吐、显存、并发、容错、版本管理、数据漂移每一个环节都能让你怀疑人生。而解决这些问题的前提是你得知道每一层到底在干什么。ai-engineering-from-scratch这个方向说白了就是把AI系统拆开从最底层开始自己搭一遍。不是让你重复造轮子去替代PyTorch而是通过亲手实现关键组件建立起对整条链路的“肌肉记忆”。这种能力在面试、排查线上问题、做架构选型的时候价值是碾压性的。1.2 这个方向到底适合谁不适合谁先说适合的人。如果你已经会用PyTorch或TensorFlow训练模型但说不清楚DataLoader的num_workers到底怎么影响吞吐如果你能跑通BERT微调但不知道attention的显存占用是怎么算出来的如果你部署过模型但遇到OOM只会调小batch size——那这个方向就是为你准备的。从零实现的价值不在于代码本身而在于建立“每一行代码对应什么计算、什么内存、什么时间开销”的直觉。不适合的人也很明确。如果你现在连Python的装饰器和生成器都用不利索或者矩阵乘法的维度变换还需要想半天那建议先把基础打牢。从零实现AI组件需要你同时具备三样东西Python工程能力、线性代数直觉、对神经网络的基本理解。缺任何一样过程都会变成抄代码学不到东西。还有一个常见的误区有人觉得“从零实现”就是要用纯Python写矩阵乘法连NumPy都不用。这属于走极端了。合理的边界是底层数值计算用NumPy或手写CUDA kernel进阶但框架层面的调度、内存管理、并行策略要自己实现。比如你可以用NumPy实现一个完整的Transformer前向传播但没必要自己写BLAS库。1.3 我踩过的三个认知坑第一个坑是过度追求“纯手写”。刚开始我试图用纯Python列表实现所有矩阵运算结果一个简单的矩阵乘法在1000x1000维度上跑了十几秒完全没法调试。后来想明白了从零实现的目标是理解计算图和内存布局不是重新发明数值计算库。用NumPy做底层运算自己实现调度逻辑这才是正确的打开方式。第二个坑是忽略工程化细节。我第一版手写推理引擎只关注了计算正确性没考虑内存复用。结果每处理一个batch就重新分配一次中间张量GC压力巨大实际吞吐比PyTorch低了两个数量级。后来引入了内存池和原地操作性能才追上来。这个教训让我意识到AI工程的“工程”二字一半以上的功夫在计算之外。第三个坑是低估了数值稳定性的重要性。手写softmax的时候直接exp(x) / sum(exp(x))在x较大时直接溢出成NaN。后来才学会减最大值、用log-sum-exp技巧。这些在调包时永远不会遇到的问题恰恰是理解模型行为的关键。2. 核心组件拆解从张量到推理引擎的完整链路2.1 张量抽象一切计算的基石任何AI框架的核心都是一个张量类。听起来简单但要做好需要处理至少四件事数据存储、形状管理、设备管理、自动微分。数据存储层面关键决策是内存布局。行优先C order还是列优先Fortran order对于深度学习行优先是主流因为卷积和矩阵乘法的实现更自然。但如果你要做转置操作列优先可能更高效。我自己的实现选择了行优先因为NumPy默认就是行优先互操作方便。形状管理需要支持广播机制。广播的本质是在维度不匹配时自动扩展规则是从右往左对齐维度为1或缺失的可以广播。实现的时候要注意广播不实际复制数据只在计算时虚拟扩展。我见过有人用np.broadcast_to之后直接做原地操作结果写坏了原始数据这种坑踩一次就记住了。设备管理在从零实现阶段可以先只支持CPU但接口要预留。设计上用一个device属性标记计算时检查所有输入设备一致。这个设计后来在我加GPU支持时省了大量重构工作。自动微分是张量类最复杂的部分。核心思路是构建计算图反向拓扑排序求导。每个操作记录输入输出和梯度函数反向传播时从loss开始链式求导。实现时要注意梯度累加而不是覆盖因为一个张量可能被多个下游使用。我第一版就犯了这个错误导致梯度只保留了最后一条路径的贡献。class Tensor: def __init__(self, data, requires_gradFalse): self.data np.asarray(data, dtypenp.float32) self.requires_grad requires_grad self.grad None self._backward lambda: None self._prev set() def __matmul__(self, other): out Tensor(self.data other.data, self.requires_grad or other.requires_grad) def _backward(): if self.requires_grad: self.grad (out.grad other.data.T) if self.grad is None \ else self.grad out.grad other.data.T if other.requires_grad: other.grad (self.data.T out.grad) if other.grad is None \ else other.grad self.data.T out.grad out._backward _backward out._prev {self, other} return out这段代码虽然简化但包含了自动微分的核心逻辑前向计算 反向闭包 梯度累加。理解了这个再看PyTorch的autograd就不会觉得神秘了。2.2 计算图与调度决定性能的关键计算图有两种模式静态图和动态图。静态图先定义后执行可以做全局优化但调试困难动态图边定义边执行灵活但优化空间小。PyTorch 2.0的torch.compile本质上是在动态图基础上做JIT编译兼顾两者。从零实现时我建议先做动态图因为逻辑直观。核心是一个拓扑排序 顺序执行的调度器。每次前向传播时记录操作序列反向时逆序执行。听起来简单但有几个优化点值得注意算子融合是最有效的优化之一。比如matmul bias_add relu可以融合成一个kernel减少内存读写。我实测下来融合后延迟能降低30%到50%因为省掉了中间结果的写回和读取。实现上可以用一个简单的模式匹配扫描计算图找到连续的可融合算子替换成融合版本。内存复用是另一个大头。中间张量的生命周期分析可以让你提前分配好内存池避免频繁malloc/free。我的做法是前向传播时记录每个张量的首次和末次使用位置反向传播时按逆序释放。这样峰值内存能降低40%左右。并行调度在多核CPU上效果明显。把独立的算子分配到不同线程用线程池管理。但要注意GIL的限制纯Python层面并行收益有限计算密集部分最好用NumPy释放GIL。2.3 手写Transformer理解attention的最好方式Transformer是当前AI工程的核心架构手写一遍能让你对attention的理解从“公式”变成“直觉”。先说多头注意力的维度变换。输入(batch, seq_len, d_model)经过QKV投影变成三个(batch, seq_len, d_model)然后拆分成(batch, num_heads, seq_len, d_head)。这里的关键是d_model num_heads * d_head拆分时用reshapetranspose。我第一版忘了transpose导致attention计算时维度对不上debug了半天。缩放点积注意力的公式是softmax(QK^T / sqrt(d_head)) V。缩放因子sqrt(d_head)的作用是防止点积结果过大导致softmax梯度消失。我做过实验不加缩放时d_head64的情况下attention权重的熵值会急剧下降几乎变成one-hot梯度接近零。因果掩码的实现要注意用-inf填充上三角softmax后这些位置变成0。但直接用-1e9可能因为浮点精度问题出bug建议用np.finfo(dtype).min。def scaled_dot_product_attention(Q, K, V, maskNone): d_head Q.shape[-1] scores Q K.transpose(0, 1, 3, 2) / np.sqrt(d_head) if mask is not None: scores scores mask # mask用-inf填充 # 数值稳定的softmax scores_max scores.max(axis-1, keepdimsTrue) exp_scores np.exp(scores - scores_max) attn_weights exp_scores / exp_scores.sum(axis-1, keepdimsTrue) return attn_weights V位置编码我推荐先实现正弦编码因为它是确定性的不引入额外参数。实现时注意频率的计算1 / (10000 ** (2i / d_model))偶数维用sin奇数维用cos。后来可以换成可学习的位置嵌入对比两者在长序列上的表现差异。2.4 推理引擎从模型到服务的最后一公里模型训练完只是开始推理引擎才是真正面对用户的环节。核心要解决四个问题批处理、缓存、并发、容错。动态批处理是提升吞吐的关键。思路是维护一个请求队列攒够一定数量或等待超时后一起推理。我实现过一个简单的版本队列长度达到8或等待超过10ms就触发推理。实测下来相比单条推理吞吐提升了5到8倍。但要注意批处理会增加单条请求的延迟需要根据业务场景权衡。KV Cache是自回归生成的核心优化。每次生成新token时之前token的K和V可以复用避免重复计算。实现上用一个预分配的缓存张量按位置写入。显存占用是2 * num_layers * batch_size * num_heads * max_seq_len * d_head * dtype_size。以7B模型为例fp16精度下batch1、seq_len2048时大约占用1.5GB。这个计算过程建议自己推一遍对显存规划很有帮助。并发处理在Python里是个难题。GIL导致多线程无法真正并行计算。我的方案是计算密集部分用NumPy/C扩展释放GILIO部分用asyncio。或者直接用多进程每个进程独立加载模型用共享内存传递数据。后者显存占用翻倍但实现简单。容错机制包括超时控制、降级策略、健康检查。超时控制用signal.alarm或异步超时降级策略可以是返回缓存结果或简化模型健康检查定期跑一个dummy输入验证服务正常。3. 实操全流程从零搭建一个可用的推理服务3.1 环境准备与依赖选择环境配置这块我的建议是尽量精简。核心依赖只有三个NumPy做数值计算FastAPI做HTTP服务uvicorn做ASGI服务器。不需要PyTorch不需要Transformers所有东西自己实现。Python版本选3.10以上因为要用到一些新的类型注解语法。虚拟环境用venv就够了conda太重。安装命令python -m venv venv source venv/bin/activate pip install numpy fastapi uvicorn pydanticNumPy版本建议1.24以上因为要用到新的np.float32类型别名。FastAPI用0.100以上Pydantic用2.0以上因为v2的性能和类型检查更好。目录结构这样组织ai-from-scratch/ ├── core/ │ ├── tensor.py # 张量类 │ ├── autograd.py # 自动微分 │ ├── nn.py # 网络层 │ └── attention.py # 注意力机制 ├── model/ │ ├── transformer.py # Transformer实现 │ └── weights.py # 权重加载 ├── engine/ │ ├── scheduler.py # 调度器 │ ├── cache.py # KV Cache │ └── batch.py # 动态批处理 ├── server/ │ ├── main.py # FastAPI入口 │ └── schemas.py # 请求响应模型 └── tests/ └── test_*.py这个结构的好处是关注点分离core是纯计算model是网络结构engine是推理优化server是对外接口。每一层可以独立测试和替换。3.2 权重加载与格式转换从零实现意味着不能用from_pretrained得自己解析权重文件。常见格式有safetensors和PyTorch的pickle。safetensors更安全更快推荐优先支持。safetensors的格式很简单前8个字节是header长度小端序然后是JSON header描述每个张量的名称、dtype、shape、数据偏移最后是连续的二进制数据。解析代码如下import json import struct import numpy as np def load_safetensors(path): with open(path, rb) as f: header_len struct.unpack(Q, f.read(8))[0] header json.loads(f.read(header_len)) data_start 8 header_len tensors {} for name, info in header.items(): if name __metadata__: continue dtype_map {F32: np.float32, F16: np.float16, I64: np.int64} shape info[shape] dtype dtype_map[info[dtype]] start, end info[data_offsets] f.seek(data_start start) raw f.read(end - start) tensors[name] np.frombuffer(raw, dtypedtype).reshape(shape) return tensors加载后要做权重映射。HuggingFace的命名和自定义实现往往不一致比如model.layers.0.self_attn.q_proj.weight可能对应你的layers[0].attention.q.weight。写一个映射表逐层转换。这个过程容易出错建议加载后跑一个数值对比测试用同样的输入对比你的实现和参考实现的输出误差在1e-5以内才算通过。3.3 推理流程的完整实现推理流程分四步预处理、前向传播、后处理、返回。预处理包括tokenize和padding。tokenize如果不想引入tokenizers库可以自己实现一个简单的BPE。核心逻辑是从字符级别开始统计相邻pair的频率合并频率最高的pair重复直到达到目标词表大小。推理时按合并规则贪心匹配。我实现过一个200行左右的BPE在英文上效果和GPT-2的tokenizer接近。padding要注意attention mask的生成。padding位置在attention计算时要被mask掉否则会影响结果。mask的生成规则是真实token位置为0padding位置为-inf。前向传播就是逐层调用你的Transformer实现。这里有个优化点逐层释放中间张量。因为推理不需要反向传播每层的输入在输出计算完后就可以释放。用Python的del加gc.collect()或者更优雅地用上下文管理器。后处理包括采样策略。贪心解码最简单但生成质量一般。beam search效果好但计算量大。top-k和top-p采样是折中方案。我推荐先实现top-pnucleus sampling因为它在多样性和质量之间平衡得最好。实现时注意排序后计算累积概率找到超过p的截断位置然后在这之前重新归一化采样。def top_p_sampling(logits, p0.9, temperature1.0): logits logits / temperature sorted_indices np.argsort(logits)[::-1] sorted_logits logits[sorted_indices] probs np.exp(sorted_logits - sorted_logits.max()) probs probs / probs.sum() cumsum np.cumsum(probs) cutoff np.searchsorted(cumsum, p) 1 top_probs probs[:cutoff] top_probs top_probs / top_probs.sum() chosen np.random.choice(cutoff, ptop_probs) return sorted_indices[chosen]3.4 性能测试与调优实录实现完成后必须做性能测试。我用的指标是首token延迟TTFT和每token延迟TPOT。测试方法固定输入长度生成100个token记录时间。第一版实测结果TTFT 450msTPOT 120ms。这个成绩很差7B模型在GPU上TTFT应该在50ms以内。排查发现三个问题问题一没有用KV Cache。每次生成新token都重新计算所有位置的K和V计算量随序列长度平方增长。加上KV Cache后TPOT降到35ms。问题二NumPy的矩阵乘法没有用多线程。NumPy默认可能没链接到OpenBLAS的多线程版本。检查方法np.show_config()看blas信息。解决装numpy时确保有openblas或者用numpy-mkl。改完后TTFT降到180ms。问题三Python循环开销。逐层调用有大量Python层面的函数调用和属性访问。优化方法把热点路径用numba的njit编译或者用Cython重写。我用numba后TTFT进一步降到90ms。最终成绩TTFT 90msTPOT 28ms。虽然还比不上vLLM这类专业推理引擎但作为从零实现已经可用了。更重要的是每一个优化点你都知道为什么有效效果有多大这种掌控感是调包永远给不了的。4. 常见问题与排查技巧实录4.1 数值稳定性问题速查数值问题是手写实现最容易踩的坑因为框架帮你处理了大部分边界情况。下面这张表是我实际遇到过的典型问题和解决方案问题现象根本原因解决方案softmax输出NaNexp溢出减最大值后再exp梯度爆炸链式求导累积梯度裁剪或归一化梯度消失sigmoid/tanh饱和换ReLU或加残差连接输出全零mask用-inf导致全被mask检查mask逻辑保留至少一个有效位置精度损失float16累加累加用float32存储用float16结果不稳定随机种子未固定固定np.random.seed重点说softmax的数值稳定实现。标准公式exp(x_i) / sum(exp(x_j))在x_i较大时会溢出。正确做法是减去最大值exp(x_i - max) / sum(exp(x_j - max))。数学上等价但数值上安全。这个技巧在attention、分类头、采样等所有用到softmax的地方都适用。梯度裁剪的实现要注意裁剪的是梯度的范数不是逐元素裁剪。先算所有梯度的L2范数如果超过阈值就整体缩放。逐元素裁剪会改变梯度方向影响优化效果。4.2 内存与性能问题排查内存问题通常表现为OOM或频繁GC。排查思路是先定位峰值内存位置再分析是否有优化空间。我常用的工具是tracemalloc可以追踪Python层面的内存分配。用法import tracemalloc tracemalloc.start() # 跑推理 snapshot tracemalloc.take_snapshot() top_stats snapshot.statistics(lineno) for stat in top_stats[:10]: print(stat)输出会显示内存分配最多的代码行。我靠这个定位过一次问题发现np.concatenate在循环里被调用了上千次每次分配新数组。改成预分配切片写入后峰值内存降了60%。性能问题排查用cProfilepython -m cProfile -s cumtime infer.py按累积时间排序找到最耗时的函数。常见瓶颈有Python循环、不必要的数组复制、未释放GIL的并行、缓存未命中。一个容易被忽略的点是内存对齐。NumPy数组默认可能不是64字节对齐的影响SIMD指令效率。可以用np.empty预分配对齐内存或者用np.ascontiguousarray确保连续。实测下来对齐后矩阵乘法能快5%到10%。4.3 我踩过的五个真实坑坑一忘记设置requires_gradFalse。推理时如果张量默认需要梯度会白白构建计算图内存和速度都受影响。解决推理前统一设置torch.set_grad_enabled(False)或自定义实现里默认不追踪。坑二KV Cache的索引错误。缓存是按位置写入的如果batch内序列长度不一致索引会错位。解决用attention_mask的累积和作为写入位置或者padding到统一长度。坑三tokenizer的special token处理。BOS、EOS、PAD的处理方式不同模型不一样搞错了生成结果会莫名其妙。解决对照参考实现逐token对比。坑四多线程下的随机数。NumPy的随机数生成器不是线程安全的多线程采样会出问题。解决每个线程独立seed或者用np.random.Generator的spawn方法。坑五浮点精度累积误差。长序列推理时误差会逐层累积最后输出完全跑偏。解决关键路径用float32定期做数值校验。4.4 从零实现到生产可用的差距最后说句实话从零实现的版本距离生产可用还有距离。生产环境需要考虑的模型版本管理、灰度发布、监控告警、自动扩缩容、安全防护这些在从零实现里都不会涉及。但这个差距恰恰是你需要知道的。知道差距在哪比不知道差距存在要安全得多。我见过太多人以为调个API就是AI工程的全部结果线上出问题连排查方向都没有。从零实现一遍你会清楚地知道每个环节的边界和风险点这才是真正的工程能力。如果要把从零实现推进到生产级别下一步建议是接入Prometheus做指标监控用Docker封装环境用Kubernetes做编排加一层API网关做限流和鉴权。这些工程化的工作建立在你对底层计算有清晰理解的基础上做起来会顺畅很多。我个人在实际操作中的体会是从零实现最大的收获不是代码本身而是建立了一套完整的排查方法论。遇到问题知道从哪一层开始查知道每个指标的正常范围知道优化手段的优先级。这种能力在任何AI工程项目里都是硬通货。