从零手搓AI工程一个“不调包”主义者的踩坑与重构实录这两年“AI工程师”这个头衔被叫得越来越响但如果你真去翻招聘要求会发现一个尴尬的现实大部分岗位要的是“会调LangChain、会写Prompt、会接OpenAI API”的人而不是真正理解一个AI系统从数据到推理再到服务化整条链路的人。我见过太多简历上写着“熟悉大模型应用开发”的候选人问一句“Transformer里Q和K做点积之后为什么要除以根号d_k”就卡住了再问“推理时KV Cache到底缓存了什么”就直接沉默。这不是人的问题是学习路径的问题——大家都是从“调包”开始接触AI的框架把一切都封装好了反而让人失去了对底层机制的手感。ai-engineering-from-scratch这个标题之所以让我有共鸣就是因为它戳中了这个痛点从零开始不依赖高层框架把AI工程里那些被封装掉的核心环节一个个拆出来用最朴素的方式重新实现一遍。这不是要你抛弃PyTorch去手写CUDA而是说当你理解了每一层在干什么之后再用框架时你才知道自己到底在用什么、什么时候会出问题、出了问题该往哪个方向排查。这篇文章我会围绕这个思路把从零构建AI工程能力的关键环节拆开讲包括数学基础的补法、推理引擎的最小实现、训练循环的手写逻辑、服务化部署的坑以及我自己在实践中踩过的那些“看起来能跑但一上量就崩”的典型问题。1. 为什么“从零实现”比“直接调包”更值得投入时间1.1 调包能让你跑通Demo但跑不通生产我先说一个我自己的真实经历。早期做文本分类项目时我用HuggingFace的pipeline接口三行代码就跑通了情感分析当时觉得“AI工程不过如此”。后来业务方提了一个需求要在推理时对特定类别的输出做干预比如某些敏感词必须被压制。我第一反应是去翻pipeline的参数发现根本没有这个口子。再往下翻源码发现后处理逻辑被埋在postprocess函数里要改就得继承整个类重写。那一刻我才意识到我根本不知道这个模型从输入到输出中间经过了哪些步骤自然也就不知道该在哪里插入自己的逻辑。这就是调包最大的问题它给你一个黑盒你只能用它暴露出来的旋钮。一旦需求超出旋钮范围你就束手无策。而从零实现过一遍的人知道tokenizer怎么把文本变成id、模型怎么把id变成logits、logits怎么经过softmax变成概率、概率怎么经过解码策略变成最终输出。每一个环节都是你可以介入的点。1.2 从零实现建立的是“排查直觉”生产环境里AI系统出问题往往不是模型本身的问题而是工程链路的问题。比如推理延迟突然从50ms涨到500ms你第一反应是什么如果只会调包你可能只会去重启服务、加机器。但如果你手写过推理循环你会立刻想到几个可能是不是batch size变了导致padding变多是不是KV Cache没命中导致重复计算是不是某个请求的序列长度触发了attention的O(n²)爆炸这种排查直觉不是看书能看出来的是你自己写代码时踩过坑才会有的。我手写第一个attention实现时忘了做mask结果padding token也参与了attention计算模型输出全是乱的。排查了整整一个下午才找到原因。从那以后我看到任何attention相关的代码第一反应就是检查mask逻辑。这种“肌肉记忆”式的直觉是调包永远给不了你的。1.3 从零不等于拒绝框架而是理解框架这里要澄清一个误区从零实现不是让你在生产环境里手写所有东西。生产环境该用PyTorch还是用PyTorch该用vLLM还是用vLLM。但从零实现的过程让你具备了“框架选型”的能力——你知道vLLM的PagedAttention解决了什么问题所以你知道什么场景下它比朴素实现快什么场景下它反而有额外开销。我自己的做法是每学一个新概念先用最朴素的方式实现一个能跑的最小版本然后再去看主流框架是怎么优化的。比如学KV Cache我先写了一个不带cache的推理循环测出延迟然后加上cache再测延迟最后去看vLLM的PagedAttention论文理解它怎么解决显存碎片问题。这个顺序走下来框架对你来说就不再是黑盒而是一个“我知道它在干什么”的工具。2. 数学基础别从头推公式但要能看懂公式在说什么2.1 你需要补的不是数学课是“数学翻译能力”很多人一听到“从零实现AI”就想到要去补线性代数、概率论、微积分。这个方向对但方法不对。你不需要像数学系学生那样从头推公式你需要的是“看到公式能翻译成代码”的能力。举个例子attention的公式是Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V你不需要会证明这个公式为什么长这样但你需要能把它翻译成代码def attention(Q, K, V, maskNone): d_k Q.shape[-1] scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) weights torch.softmax(scores, dim-1) return torch.matmul(weights, V)这个翻译过程里你会自然理解几个关键点为什么要除以sqrt(d_k)防止点积结果过大导致softmax梯度消失、mask是怎么起作用的把不该看的位置设成负无穷、softmax是在哪个维度上做的最后一维即每个query对所有key的权重。这些理解比会推公式重要得多。2.2 几个必须搞清楚的数学概念根据我的经验从零做AI工程以下几个数学概念是绕不开的我按重要性排序概念为什么重要最小理解标准矩阵乘法几乎所有神经网络层都是矩阵乘法知道(m,n) (n,p) (m,p)能心算维度变化Softmax把分数变成概率分布知道它会让大的更大、小的更小且和为1交叉熵分类任务的标配损失函数知道它衡量的是预测分布和真实分布的差异梯度下降训练的本质知道参数沿着梯度反方向更新能降低损失链式法则反向传播的基础知道复合函数求导是逐层相乘这些概念你不需要会手推但需要能用代码实现。比如softmax你应该能不看资料写出数值稳定的版本def stable_softmax(x): x_max x.max(dim-1, keepdimTrue).values exp_x torch.exp(x - x_max) return exp_x / exp_x.sum(dim-1, keepdimTrue)为什么要减x_max因为exp对大数会溢出。这个细节在教科书里可能一笔带过但在实际工程里是必须处理的。这就是“从零实现”的价值——它逼你面对这些工程细节。2.3 我的数学补法按需补不囤积我不建议花三个月系统学完线性代数再开始做项目。我的做法是“按需补”遇到不懂的概念花半天到一天专门搞懂它然后立刻回到项目里用起来。比如我学LayerNorm的时候先看公式y (x - mean) / sqrt(var eps) * gamma beta然后自己实现一遍def layer_norm(x, gamma, beta, eps1e-5): mean x.mean(dim-1, keepdimTrue) var x.var(dim-1, keepdimTrue, unbiasedFalse) normalized (x - mean) / torch.sqrt(var eps) return gamma * normalized beta实现完就理解了gamma和beta是可学习参数让网络可以“撤销”归一化eps是防止除零unbiasedFalse是因为这里用的是有偏方差。这些细节看公式是看不出来的写一遍代码就全清楚了。3. 手写一个最小推理引擎从tokenizer到采样3.1 为什么值得手写推理引擎你可能会问现在vLLM、TGI、TensorRT-LLM这么多推理框架为什么还要手写我的答案是手写一遍你才知道这些框架在优化什么。我手写推理引擎的动机很实际当时线上服务遇到一个诡异问题同一个模型同样的输入有时候输出会突然变短。排查了半天发现是采样策略的问题——框架默认的top-p采样在某些概率分布下会截断得太狠。如果我不懂采样逻辑就只能去翻框架文档碰运气。但我手写过采样我知道top-p是怎么工作的就能快速定位问题。3.2 最小推理循环的骨架一个最小的自回归推理循环核心逻辑其实很简单def generate(model, input_ids, max_new_tokens, temperature1.0, top_kNone): model.eval() generated input_ids.clone() with torch.no_grad(): for _ in range(max_new_tokens): logits model(generated).logits[:, -1, :] logits logits / temperature if top_k is not None: top_k_values, _ torch.topk(logits, top_k) threshold top_k_values[:, -1].unsqueeze(-1) logits torch.where(logits threshold, torch.full_like(logits, -float(inf)), logits) probs torch.softmax(logits, dim-1) next_token torch.multinomial(probs, num_samples1) generated torch.cat([generated, next_token], dim-1) if next_token.item() eos_token_id: break return generated这段代码不到30行但它包含了推理引擎的所有核心要素前向传播、温度缩放、top-k截断、概率采样、终止条件。你把这30行搞懂了再看vLLM的源码就不会觉得是天书。3.3 温度、top-k、top-p到底在干什么这三个参数是推理时最常调的但很多人调的时候是凭感觉。我手写实现之后对它们的理解清晰了很多温度temperature是在softmax之前对logits做除法。温度小于1logits之间的差距被放大分布更尖锐输出更确定温度大于1差距被缩小分布更平坦输出更随机。我实测下来温度调到0.1以下基本就是贪心解码调到1.5以上输出就开始胡言乱语。top-k是只保留概率最高的k个token其余设为负无穷。这个操作在softmax之前做所以被截断的token概率直接变成0。k1就是贪心解码k越大越随机。top-pnucleus sampling是按概率从高到低累加直到累计概率超过p只保留这些token。它比top-k更灵活因为保留的token数量是动态的。我手写的时候犯过一个错先做softmax再做top-p截断结果截断后概率和不为1采样就出问题了。正确做法是在softmax之前截断或者截断后重新归一化。提示如果你在线上服务里发现输出长度不稳定优先检查top-p参数。top-p太小会导致候选集过窄模型容易提前输出终止符。3.4 KV Cache手写一遍才知道省了多少计算不带KV Cache的推理每生成一个token都要把整个序列重新过一遍模型。序列长度是n生成m个token计算量是O(m * n²)。加上KV Cache之后每个token的Key和Value被缓存下来新token只需要计算自己的Query然后和缓存的K、V做attention计算量降到O(m * n)。我实测过一个7B模型序列长度512生成128个token不带cache耗时约8.2秒带cache耗时约1.3秒差了6倍多。这个差距在长序列上会更夸张。手写KV Cache的关键是理解缓存的形状。以多头注意力为例缓存的K和V形状是(batch_size, num_heads, seq_len, head_dim)。每次新token进来计算完自己的K、V之后要和缓存的K、V在seq_len维度上拼接class KVCache: def __init__(self): self.k_cache None self.v_cache None def update(self, new_k, new_v): if self.k_cache is None: self.k_cache new_k self.v_cache new_v else: self.k_cache torch.cat([self.k_cache, new_k], dim2) self.v_cache torch.cat([self.v_cache, new_v], dim2) return self.k_cache, self.v_cache这个实现很朴素每次拼接都会重新分配内存效率不高。但理解了这个逻辑之后你再看vLLM的PagedAttention就知道它在解决什么问题——预分配显存块避免频繁拼接导致的内存碎片和拷贝开销。4. 手写训练循环损失函数、反向传播和优化器4.1 训练循环的骨架比你想的简单一个最小的训练循环核心就四步前向传播算损失、反向传播算梯度、优化器更新参数、梯度清零。def train_step(model, optimizer, input_ids, labels): model.train() outputs model(input_ids) logits outputs.logits loss F.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index-100 ) loss.backward() optimizer.step() optimizer.zero_grad() return loss.item()这十几行代码就是所有训练框架的核心。你把这个搞懂了再看PyTorch Lightning、HuggingFace Trainer的源码就会发现它们无非是在这个骨架上加了分布式、混合精度、梯度累积、日志记录等功能。4.2 交叉熵损失的实现细节交叉熵是分类任务最常用的损失函数但它的实现有几个坑。最朴素的实现是def naive_cross_entropy(logits, labels): probs softmax(logits) log_probs torch.log(probs) return -log_probs[range(len(labels)), labels].mean()这个实现有两个问题一是softmax之后取log数值上不稳定概率很小时log会变成负无穷二是没有处理ignore_indexpadding位置也会参与损失计算。PyTorch的F.cross_entropy内部用的是log-softmax把softmax和log合并成一个操作数值上更稳定。我手写的时候踩过这个坑用朴素实现训练loss偶尔会变成nan排查半天才发现是数值溢出。后来改用log-softmax就再也没出现过。4.3 优化器SGD、Adam、AdamW的区别优化器的选择直接影响训练效果。我手写实现过SGD和Adam对它们的区别有了更直观的理解SGD就是参数沿着梯度反方向更新param param - lr * grad。简单但对学习率很敏感学习率设大了会震荡设小了收敛慢。Adam在SGD基础上加了两个改进一是动量momentum用梯度的指数移动平均来平滑更新方向二是自适应学习率根据梯度平方的指数移动平均来调整每个参数的学习率。这让Adam对学习率的鲁棒性更好但代价是多了两个状态变量显存占用增加。AdamW是Adam的改进版主要区别是权重衰减的实现方式。Adam是在梯度里加L2正则AdamW是直接在参数上做权重衰减。看起来差别很小但实测下来AdamW的泛化性能更好现在基本是训练Transformer的标配。我手写Adam的时候最容易被忽略的是偏差修正bias correction。Adam的动量估计在训练初期是有偏的因为初始化为0前几步的估计会偏向0。偏差修正就是除以(1 - beta^t)来纠正这个偏差。忘了这一步训练初期的更新会偏小收敛变慢。4.4 梯度累积小显存跑大batch显存不够是训练时的常见问题。梯度累积的思路是把一个大batch拆成几个小batch分别前向反向梯度累加最后一起更新。def train_with_accumulation(model, optimizer, dataloader, accum_steps): optimizer.zero_grad() for i, batch in enumerate(dataloader): loss compute_loss(model, batch) loss loss / accum_steps loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()这里的关键是loss / accum_steps因为梯度是累加的不除以累积步数的话等效学习率会变大。我一开始忘了除结果训练直接发散排查了好久才找到原因。5. 服务化部署从单机脚本到线上服务5.1 为什么你的模型在本地跑得好上线就崩本地跑推理脚本和线上服务是两回事。本地你一次处理一个请求显存够、延迟无所谓。线上你要处理并发请求显存有限、延迟有SLA。我见过太多“本地跑通就上线”的案例结果一上量就出问题。最常见的问题是显存OOM。本地跑的时候batch size是1线上并发上来了如果每个请求都单独跑一次前向显存很快就不够。解决方案是动态批处理dynamic batching把多个请求攒成一个batch一起前向。但这里有个权衡——攒batch会增加延迟因为要等请求凑够。我一般会设一个最大等待时间比如10ms超过就立刻处理当前攒到的请求。5.2 动态批处理的实现要点动态批处理的核心是一个请求队列和一个调度循环class DynamicBatcher: def __init__(self, model, max_batch_size, max_wait_ms): self.model model self.max_batch_size max_batch_size self.max_wait_ms max_wait_ms self.queue [] def add_request(self, input_ids): self.queue.append(input_ids) def process_batch(self): if not self.queue: return batch self.queue[:self.max_batch_size] self.queue self.queue[self.max_batch_size:] # padding到相同长度 max_len max(len(ids) for ids in batch) padded [ids [pad_token_id] * (max_len - len(ids)) for ids in batch] input_tensor torch.tensor(padded) with torch.no_grad(): outputs self.model(input_tensor) return outputs这里有个细节padding会增加计算量。如果batch里有一个超长序列其他短序列都要padding到同样长度浪费显存和算力。优化方法是按长度分桶bucketing把长度相近的请求放在一起。我实测下来分桶能减少30%左右的padding浪费。5.3 流式输出用户体验的关键大模型生成是逐token的如果等全部生成完再返回用户要等好几秒。流式输出是每生成一个token就返回一个用户能看到文字逐渐出现体验好很多。流式输出的实现关键是异步生成器async def stream_generate(model, input_ids, max_new_tokens): generated input_ids.clone() for _ in range(max_new_tokens): with torch.no_grad(): logits model(generated).logits[:, -1, :] next_token sample(logits) generated torch.cat([generated, next_token], dim-1) yield next_token.item() if next_token.item() eos_token_id: break服务端用SSEServer-Sent Events或WebSocket把token推给前端。这里要注意的是流式输出和动态批处理有冲突——批处理要攒请求流式要尽快返回。我的做法是流式请求单独走一条路径不参与批处理牺牲一点吞吐换延迟。5.4 监控上线只是开始服务上线之后监控比开发更重要。我一般会监控这几个指标指标含义告警阈值P99延迟99%请求的响应时间超过SLA的80%显存使用率GPU显存占用超过90%队列长度等待处理的请求数持续超过batch size的3倍错误率请求失败比例超过1%输出长度分布生成token数的分布突然变短或变长输出长度分布这个指标很多人不监控但它很有用。如果输出长度突然变短可能是采样参数出了问题如果突然变长可能是模型陷入了重复循环。我遇到过几次线上输出异常都是靠这个指标先发现的。6. 那些只有踩过才知道的坑6.1 数值精度float16不是万能的为了省显存很多人会把模型转成float16。但float16的动态范围很窄遇到大数会溢出遇到小数会下溢。我在做attention的时候遇到过一个问题QK^T的结果在float16下溢出变成infsoftmax之后全是nan。解决方案是混合精度attention的score计算用float32其他部分用float16。PyTorch的torch.cuda.amp可以自动处理这个但你需要知道哪些操作是数值敏感的。我一般会把softmax、layer norm、loss计算放在float32下做。6.2 显存碎片为什么重启能解决80%的问题显存碎片是推理服务的隐形杀手。PyTorch的显存分配器会缓存已分配的显存块如果请求的显存大小和缓存块不匹配就会产生碎片。碎片多了之后即使总显存够也分配不出连续的大块显存导致OOM。我遇到过好几次“显存明明够但就是OOM”的情况重启服务就好了。后来才知道是碎片问题。解决方案有几个一是设置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True让分配器支持可扩展段二是定期重启服务三是用vLLM这类自带显存管理的框架它的PagedAttention就是专门解决碎片问题的。6.3 tokenizer的坑为什么你的模型输出乱码tokenizer看起来简单但坑不少。我遇到过几个典型问题一是特殊token处理。有些tokenizer会把[PAD]、[CLS]这些特殊token编码成id如果你不处理模型会把它们当成正常文本。我一般会在tokenize之后检查一下确保特殊token被正确mask。二是中文分词。有些tokenizer对中文支持不好一个汉字被拆成多个byte导致序列长度暴涨。我实测过同一个中文句子好的tokenizer编码后长度是20差的能到60。选tokenizer的时候一定要用中文语料测一下。三是padding方向。有些模型是左padding有些是右padding。用错了会导致生成结果完全不对。我一般会看模型的config里padding_side字段或者直接看tokenizer的默认设置。6.4 batch size和延迟的权衡batch size越大吞吐越高但延迟也越高。因为batch里的请求要等最慢的那个完成。我实测过一个7B模型batch size1时延迟50msbatch size8时延迟120msbatch size32时延迟350ms。吞吐从20 QPS涨到90 QPS但延迟涨了7倍。怎么选看业务需求。如果是离线批处理batch size拉满如果是在线交互batch size要控制一般不超过8。我一般会设一个动态策略队列短的时候用小batch保延迟队列长的时候用大batch保吞吐。6.5 模型量化省显存但别省质量量化是省显存的有效手段。4-bit量化能把7B模型的显存从14GB降到4GB左右。但量化会损失精度我实测下来4-bit量化在生成任务上质量下降明显尤其是长文本生成容易出现重复和逻辑断裂。8-bit量化质量损失小很多显存也能降到7GB左右是我比较推荐的折中方案。量化的实现方式也有讲究。GPTQ、AWQ、GGUF这些方法各有优劣。我一般用AWQ因为它在推理速度和质量之间平衡得比较好。但要注意量化后的模型不能直接用于训练只能推理。7. 从零实现之后我如何看待框架手写了一遍推理引擎和训练循环之后我对框架的态度发生了变化。以前是“框架能跑就行”现在是“我知道框架在干什么所以我知道什么时候该用它、什么时候该绕过它”。比如做推理服务小规模场景我直接用PyTorch手写循环因为可控性高出问题好排查。规模上来了我会切到vLLM因为它的PagedAttention和连续批处理确实能带来数倍的吞吐提升。但切过去之后我依然知道它在底层做了什么所以遇到问题不会慌。再比如训练小模型我手写训练循环因为灵活想改什么改什么。大模型我用HuggingFace Trainer或DeepSpeed因为它们处理了分布式、混合精度、梯度检查点这些复杂逻辑。但我知道这些逻辑的本质是什么所以配置的时候不会瞎填。这种“知其然也知其所以然”的状态是我认为AI工程师最应该追求的能力。框架是工具工具会变但底层原理不会变。你把底层搞懂了换什么框架都能快速上手。最后分享一个我自己的学习习惯每学一个新概念我都会问自己三个问题——如果不用框架我怎么实现它它的计算复杂度是多少它在什么情况下会出问题这三个问题逼着我把每个概念都落到代码和工程层面而不是停留在“知道有这么个东西”的层面。这个习惯让我在遇到线上问题时总能比别人快一步定位到根因。