ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

从零构建AI工程化:手写Transformer到推理部署的完整链路

从零构建AI工程化:手写Transformer到推理部署的完整链路 看到这个标题我的第一反应是这不是又一个“从入门到放弃”的学习计划就是一套把AI领域所有知识堆在一起的大纲合集。但真正点进去、动手把链路走通之后我的看法变了。这个项目想解决的问题非常明确不依赖现成的大模型API也不用东拼西凑地看教程而是从数学基础、模型结构、训练推理到工程部署把AI工程化这条链路完整地“手工打磨”一遍。适合谁适合那些已经会写Python、用过PyTorch但总觉得对AI的了解“隔着一层”的开发者。也适合想从算法岗转工程岗、或者纯工程背景想系统性补AI知识的朋友。这篇文章我会围绕这个项目标题展开讲讲我从中拆解出来的核心思路、技术主线和实操经验。1. 整体设计思路为什么“从零构建”才是最快的学习路径1.1 盖房子 vs 精装修搞清楚你缺的是哪一层很多人学AI的路径是反的。他们一开始就跑去调 Hugging Face 上的预训练模型用现成的transformers库跑通一个情感分析或者用diffusers生成几张图就觉得自己“会了”。等到真正需要自己训练一个模型、优化推理速度、或者把一个demo部署成高并发的服务时瞬间就懵了。这个项目标题里的“from scratch”之所以值钱就是因为它选择了另一条路先不急着用库而是把房子的地基、承重墙、管线先摸清楚。就好比装修一栋精装房你确实可以直接拎包入住但如果你想改造户型、调整水电你就必须知道墙里面埋的是什么、哪面墙是承重墙、管线是怎么走的。AI工程也是一样当你掌握了从零构建一个Transformer的能力再看那些封装好的框架你看到的就不再是黑盒而是一层薄薄的包装纸随时可以剥开。1.2 链路拆解从张量到服务的四个阶段把整个项目解剖开你会发现合理的路径不是按“知识领域”划分的而是按“工程链路”划分的这也是我认为这个项目最聪明的地方。完整链路可以划分为四个递进阶段。阶段一理论基建——线性代数、微积分、概率论里跟AI强相关的部分以及反向传播的数学推导。这个阶段的目标不是成为数学家而是做到“看到公式不心虚”。阶段二模型实现——用NumPy从零写全连接网络、CNN、RNN最后手写Transformer。不调PyTorch的nn.Module纯用矩阵运算把前向和反向撸出来。阶段三框架迁移——把NumPy实现“翻译”成PyTorch代码理解张量、自动求导、DataLoader等框架机制开始利用GPU加速。阶段四工程部署——训练脚本工程化、模型量化、推理优化、服务化部署FastAPI Docker、性能压测。这四个阶段环环相扣没有第一阶段第二阶段的代码写出来就是玄学没有第二阶段第三阶段你用PyTorch写模型就只是“照葫芦画瓢”没有第四阶段前面学的一切在工作场景中都落不了地。1.3 为什么这条路更适合“工程师”而非“研究员”这里要特别说明一下为什么这套路径更贴合“AI工程”而非“AI研究”。研究者的核心任务是探索未知他们需要的是快速验证想法所以大量使用现成框架是合理的效率优先。但工程师的核心任务是稳定、可控地交付系统这就意味着你必须理解底层机制。比如当你需要把一个模型的显存占用从12GB压到6GB时你就必须知道哪些层是显存大头、你的推理框架在计算图层面做了什么优化。这些信息只调API是永远学不到的。所以“从零构建”这种“慢功夫”恰恰是工程师最需要的“快路径”它帮你一次性建立完整的系统认知后面再遇到新模型、新框架学起来都非常快。2. AI工程的本质它“工程”在哪里2.1 工程思维的核心可复现、可测试、可观测聊到“AI工程”很多人脑子里的画面还是“调参”和“炼丹”。但实际上AI工程和传统软件工程在核心追求上没有本质区别都是三件事可复现、可测试、可观测。只不过AI系统引入了一个传统系统没有的变量——数据漂移这让工程化难度上了一个台阶。具体来说可复现意味着你记录的不只是代码版本还有数据版本、模型超参数、随机种子。我见过太多“昨天还能跑到95%今天变成93%”的惨案最后查下来不是代码问题而是数据文件被悄悄覆盖了。可测试在AI场景下更加复杂除了单测你的工具函数还需要对模型做“行为测试”比如对于文本分类模型你需要准备一些“对抗样本”——把“这家餐厅太难吃了服务也差”中的“难吃”换成“不好吃”看模型是不是还能正确识别。可观测则是把训练过程的loss曲线、显存占用、吞吐量、推理延迟全部用监控面板暴露出来一旦出问题能立刻定位是哪一环。2.2 数据、模型、算力AI工程的三驾马车在AI工程这个语境下数据和算力往往比模型本身更值得投入。这三个要素大约可以对应到三层问题。数据层数据从哪来、怎么清洗、怎么标注、怎么做数据增强、怎么保证训练集和验证集不泄漏。这部分能占到整个项目时间的60%以上而且是决定模型上限的关键。优秀的开源模型用相同的数据量能比你多跑5个点迭代的关键常常是“数据提纯”而非“模型改动”。模型层选什么架构、用多少参数、怎么初始化、用什么损失函数、怎么设置学习率。这层是大家最关注的但实际技术含量恰恰不在“跑通”而在于理解每个组件为什么存在。算力层单机多卡怎么并行、数据并行和模型并行怎么选、混合精度怎么开、显存不够怎么用梯度累积。这层是工程化的重头戏也是很多人从“算法能跑”到“系统能扛”的鸿沟所在。2.3 批评一些“重模型轻工程”的普遍误区现在行业里有一个普遍的误区觉得“AI工程师”每天的工作就是研究新模型架构。真实情况完全不是这样。在绝大多数落地场景中你用的模型大概率还是几年前就被提出的Transformer甚至BERT都还在大量生产环境中服役。真正决定一个AI系统能不能活下去的往往是那些不起眼的工程细节数据管道的稳定性、推理服务的延迟、模型的持续监控和更新机制。我见过一个团队花了一个月时间把模型准确率从88%提到了89%非常高兴。结果上线后发现因为推理服务没有做超时控制高峰期一个慢请求拖垮了整个服务链路用户端整体报错率上升了3个百分点那一个点的准确率提升瞬间显得毫无意义。这就是工程的价值它不显山不露水但它是模型价值兑现的唯一通道。3. 技术主线详解从主流的工程栈说起3.1 核心技术栈Python/C/CUDA 的层次划分这个项目涉及的编程语言和技术栈非常典型可以分为三层。层级语言/技术定位典型场景顶层Python快速原型、数据科学、训练逻辑PyTorch 训练脚本、数据预处理中间C高性能推理、算子融合、系统调度TensorRT、vLLM 等推理引擎底层CUDAGPU 并行计算极致性能榨取手写 Kernel、FlashAttention 实现对于绝大多数从事AI工程的朋友来说Python是绝对的主战场日常90%的代码都跑在这一层。但当你开始做推理优化时C和CUDA的知识就变得不可或缺。比如你用PyTorch导出一个模型部署到生产环境时通常会用TensorRT或者ONNX Runtime来做加速而这些引擎的核心算子都是用C和CUDA写的。你不需要成为CUDA专家但你需要理解“算子融合”“半精度计算”这些概念它们是你调优推理延迟的理论基础。3.2 核心框架选型PyTorch 为什么成了主流在框架层面PyTorch 几乎已经成为AI工程领域的“默认选项”。原因不复杂它的动态图机制让你可以在写代码的同时调试张量的形状和取值这种调试体验对研究极其友好而它的生产部署生态TorchScript、TorchServe、ONNX导出也日趋成熟。当然工程上也有其他选择。JAX在纯研究场景和部分大规模并行场景下越来越流行它的函数式纯计算模型在编译器优化上有天然优势TensorFlow在传统工业界仍有大量存量系统。但如果你是新手我强烈建议你从PyTorch入手社区的教程、开源项目、工作岗位需求都是最多的遇到问题一搜就有解决方案。当你把PyTorch用熟了再迁移到JAX或者TensorFlow成本远比你想象的低因为核心的深度学习概念——张量、自动微分、优化器——是相通的。3.3 混合精度与分布式工程实战的关键利器进阶到工程实战有两大块能力是硬门槛混合精度训练和分布式训练。混合精度训练的核心逻辑是用FP16半精度浮点数来做前向和反向计算用FP32来存储主权重master weights。为什么要这么做因为FP16占用的显存只有FP32的一半计算速度在支持FP16加速的硬件上如V100之后的NVIDIA GPU可以提升2到4倍。但FP16的数值范围很小直接训练很容易梯度下溢变成0所以实际操作中会做一个“损失缩放”loss scaling在反向传播时先放大损失值计算完梯度再缩小回去。用PyTorch实现非常简单只需要在训练循环里加一行from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for batch in dataloader: with autocast(): loss model(batch) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()分布式训练则复杂一些。数据并行DDP是最常用的方案它的思路是每张GPU卡放一份完整的模型副本但数据切分成多份分给不同的卡前向计算时每张卡处理自己的那批数据反向传播算完梯度后各卡之间做一次梯度同步再统一更新参数。PyTorch的DistributedDataParallel模块把这件事包装得非常干净你只需要用启动脚本唤起多进程即可。但当模型大到单卡放不下时比如训练千亿参数模型你就得深入研究模型并行、流水线并行、张量并行这些更复杂的策略了。4. 核心细节拆解从零实现一个Transformer有多难4.1 注意力机制的“为什么”从向量点乘到缩放点积要说这个项目里最核心、也最值得下功夫啃的部分绝对是手写Transformer。而Transformer的心脏就是自注意力机制Self-Attention。自注意力的原始公式是Attention(Q, K, V) softmax(Q K.T / sqrt(d_k)) V这个公式看着简单但每一个设计都是有讲究的。QQuery、KKey、VValue三个矩阵是从输入经过三个不同的线性变换得到的。理解它们最直观的方式是搜索引擎的类比你有一次搜索请求Query搜索引擎去库里面匹配相关的文档Key然后把匹配到的内容返回给你Value。在自注意力中序列里的每个词都会跟序列里的其他词做这样的检索从而建模词与词之间的依赖关系。为什么要除以sqrt(d_k)而不是其他数这是为了防止点积结果过大导致Softmax进入饱和区。假设Q和K的每个元素是均值为0、方差为1的独立随机变量那么它们的点积的方差就是d_k标准差就是sqrt(d_k)。点积结果的标准差越大数值就越分散Softmax的梯度就越容易消失。除以sqrt(d_k)就像给数据归一化一样把方差拉回1让梯度保持稳定。4.2 位置编码的“为什么”为什么RNN不需要而Transformer需要RNN之所以不需要显式的位置编码是因为它天生按顺序处理序列位置信息天然编码在“时间步”里。但Transformer的结构是并行计算的它一次看到整个序列如果把词序搞乱注意力机制的计算结果是完全一样的。这也是Transformer最初被批评的原因它好像不“在乎”词序。最初的Transformer论文使用了一种固定频率的正弦余弦位置编码这种方式的好处是无需学习、能外推到训练时没见过的序列长度。但后来实践发现这种固定编码在长序列上的表现不够好所以后来大家更常用可学习位置编码把位置信息当作一组可训练的参数随模型一起优化。而到了更现代的LLM时代比如GPT系列用的则是旋转位置编码RoPE。RoPE的思想是把位置信息通过旋转矩阵的方式注入到Q和K向量上让注意力的分数天然依赖于相对位置。我在实际工程中踩过一个坑用绝对位置编码的旧模型拼接超过训练长度的文本时效果会明显下降后来换用RoPE的底座模型外推能力好了很多。所以做文本类AI应用时选底层模型一定要优先考虑支持RoPE或类似相对位置编码的架构。4.3 手写实现的关键前向容易反向“炸裂”在从零实现阶段前向传播的逻辑其实不难难的是反向传播。用PyTorch不会觉得反向有多可怕因为有自动求导兜底。但当你想用NumPy自己撸一个完整的Transformer来验证理解时你会发现反向传播的推演过程“炸裂”Softmax的雅可比矩阵、LayerNorm的梯度回传、多头注意力的reshape和转置操作每一个都要小心处理。我给一个实操建议不要直接试图一口气写完整个Transformer的反向传播。正确的方式是先写一个只有一层注意力的简化版本手动验证梯度接着加上LayerNorm和Feed-Forward网络最后再把多头、Residual Connection这些组件逐个加回来。每一步都用PyTorch的自动求导结果作为“标准答案”跟自己的手推梯度对比误差控制在1e-6以内就说明理解到位了。这个过程确实是耗时但它的回报是你对反向传播的理解能深入到条件反射的层面日后调模型时哪些层容易出现梯度消失或爆炸你基本能靠直觉判断。5. 推理优化与部署实战从模型到服务5.1 张量布局与内存访问为什么“形状”会影响性能模型训练完真正的战斗才刚刚开始也就是部署和推理优化。很多AI工程师在训练阶段如鱼得水一到部署就水土不服模型跑起来了但速度慢得要命。这时候你需要关注的第一个指标是张量在内存中的布局。以经典的图像分类模型为例PyTorch里NCHW批次数、通道数、高度、宽度格式是默认的。但当模型转到推理引擎如TensorRT时操作往往变成NHWC格式也就是把通道维放最后。为什么因为GPU内存访问的局部性极高而NHWC布局在很多常见算子如卷积、像素级操作中能更好地利用GPU的缓存和向量化指令从而显著减少内存搬移的时间。这不是什么高深数学问题纯粹是一个“数据摆放优化”问题。5.2 算子融合与KV Cache两个必须理解的核心优化推理优化的两大核心技巧一个是算子融合另一个是KV Cache。展开讲讲这两点它们你在任何框架的调优文档里都会反复遇到。算子融合的意思是把多个连续的操作合并成单个操作减少内存读写的次数。一个典型例子是Conv BN ReLU三合一。在训练时BN层的参数是动态的但在推理时BN的参数是固定的所以它的数学运算可以提前折算到卷积层的权重上。这样推理时你只需要跑一个卷积算子少了两轮内存读写性能立刻就能上来。TensorRT这类引擎做的最核心的工作就是这种“计算图优化”。KV Cache则是生成式模型如GPT推理的核心优化。自回归生成时每次只生成下一个token但注意力计算需要用到所有历史token的K和V向量。如果不缓存每次生成都要把历史重新算一遍时间复杂度是O(n²)序列一长直接崩溃如果做过缓存每次只需要算新token的K、V再拼接到缓存的KV矩阵上时间复杂度变成O(n)。这里面涉及一个内存大小的计算问题以7B模型为例层数为32、注意力头数为32、头维度为128那么每层KV缓存的大小是2 * batch_size * 序列长度 * 32 * 128 * 2字节FP16。算一下就知道当并发请求多、序列长时KV Cache的显存占用相当可观。这也是为什么很多高性能推理框架如vLLM、TensorRT-LLM都在KV Cache的内存管理上做文章比如PagedAttention就是为了把显存利用率提到极致。5.3 从PyTorch到TensorRT一个简单的部署链路部署链路本身不复杂核心是一个“导出”和“优化”的过程。以TensorRT为例标准的落地流程如下。PyTorch训练得到权重先导出为ONNX通用格式。用TensorRT解析ONNX做计算图优化和算子融合。选定推理精度一般用FP16追求极致性能可以尝试INT8。生成TensorRT引擎文件一个针对特定GPU架构优化的二进制文件。用TensorRT的Python/C API加载引擎做一次“预热”推理然后正式对外提供服务。这里面有几个实操中容易踩的坑。第一ONNX导出时如果你的模型里有动态shape比如输入长度不固定需要在导出时指定动态轴否则推理时序列长度一变就会报错。第二TensorRT引擎是绑定显卡架构的你在A100上生成的engine文件换到4090上可能无法加载必须重新构建。第三INT8量化需要校准数据选有代表性的数据集做校准否则精度会大幅度掉点。我自己就吃过亏图省事随便找了一百张验证集图片做校准结果上线后模型的输出明显变差后来换了和线上分布一致的校准集才恢复正常。5.4 服务化FastAPI Docker 就够了模型推理本质上是计算密集型的“函数调用”所以服务化部署最自然的方式就是HTTP服务。对于大多数场景FastAPI加Docker的搭配完全够用不需要一上来就上Kubernetes那套重量级方案。一个基本的推理服务包括模型加载只加载一次到显存、请求预处理文本转为token、图像转为Tensor、模型推理、结果后处理logits转为概率或文本、并发控制。并发控制是容易被忽略的点GPU是共享资源如果同时来100个请求全部塞进模型显存可能瞬间爆掉。所以要么在应用层做排队要么用消息队列削峰要么用支持continuous batching的推理框架vLLM来动态调度。日志和监控也要从一开始就做好每个请求的延迟、输入token数、输出token数、GPU显存占用都要记录下来这是事后排查问题的重要依据。6. 生产环境中的模型选型与硬件考量6.1 开源模型选型不是越新越好而是越合适越好AI工程落地的第一件事往往是选模型而选模型绝不是看排行榜找最强者。最适合你的模型需要考虑几个维度的权衡效果、体积、推理速度、生态成熟度、许可证合规性。举个例子如果做的是中文场景的文本分类、信息抽取这类任务早期大家爱用BERT系列因为体积小110M参数、推理快、社区资料多。后来大家发现LLM能力更强开始用7B、13B的底座模型做微调。但7B模型在CPU上做推理几乎不可用必须依赖GPU如果你的线上环境只有T4这种入门级显卡内存也只有16GB那7B模型用FP16加载勉强装下但推理延迟会很高。这时候可能更合适的选择是量化到4bit的版本或者干脆选一个更小的模型如1.5B级别通过对特定任务的微调把效果追上来。6.2 训练硬件与服务器配置把钱花在刀口上硬件选型方面我给不出“万能答案”但可以给一套决策思路。训练阶段和推理阶段的需求是不同的。训练阶段追求的是高吞吐量和显存容量主流选择是NVIDIA的A100、H100或者性价比更高的L40S。推理阶段则更关注延迟和成本4090这类消费级显卡虽然在大规模并发下不稳定但做中小规模的私有化部署完全够用性价比极高。还要特别注意显卡互联NVLink、InfiniBand和CPU、内存的配套。很多人只盯着GPU结果CPU太弱数据加载跟不上或者内存太小数据预处理直接卡死。我曾在一台“GPU很强但CPU极弱”的机器上做数据加载一张大图加载预处理要半秒训练的时候GPU大部分时间在空转。后来换了个思路把CPU数据预处理改成异步多进程才把GPU利用率从30%拉回90%。工程问题往往就是这样你以为瓶颈在GPU其实在别的角落。7. 踩坑记录与问题排查速查表7.1 训练阶段最常见的六个问题训练是个试错的过程我把自己和数据、模型、训练相关的高频问题整理成了一张排查表遇到问题先对表自查能省下大把时间。现象可能原因排查方法loss不下降学习率过大/过小、数据没归一化输出loss数值检查是否存在NaN尝试用学习率扫描LR Finder梯度爆炸网络过深、学习率过大加梯度裁剪clip_grad_norm_设置阈值1.0训练acc高但验证acc低过拟合加正则化、数据增强、早停检查是否数据泄漏验证acc“跳变”不稳定Batch size过小增大batch size或使用梯度累积模拟大batch显存OOM单batch显存占用过高减小batch size、开启梯度累积、用混合精度训练多GPU训练结果不一致未同步BN层、种子未固定DDP模式下设置torch.manual_seed并固定所有随机源7.2 推理部署阶段最常见的五个问题推理部署的问题往往比训练更隐蔽因为环境复杂、依赖多样。下面这五个是高频出现的。CUDA版本与PyTorch版本不匹配报错信息可能很奇怪最常见的解决办法是用PyTorch官网的pip install命令重新安装匹配版本注意选择cu118、cu121这样的对应索引。模型加载慢大模型从磁盘加载到GPU需要时间可以考虑用torch.compile或者模型并行加载也可以把权重文件放在本地高速盘上如果是生产服务加载一次后常驻内存是常规操作。首次推理格外慢这通常是因为GPU没有预热CUDA kernel是懒加载的。解决办法是在服务启动后跑一次小输入“预热”让CUDA上下文建立并加载相关kernel。并发一高延迟就飙升大概率是GPU算力、显存带宽饱和或者服务端排队队列设计不合理。先用压测工具如wrk、locust找到瓶颈再决定是横向加GPU还是优化推理引擎。量化后精度掉得离谱校准数据集与线上数据分布不一致是首要嫌疑重新选择更接近线上分布的校准数据或者对敏感层保留FP16精度。7.3 定位问题的底层方法论二分定位法排查算法和工程问题最有效的方法论其实特别朴素——二分定位。训练loss异常了你先把模型看成一条“数据流”从头到尾二分排查组件。比如先检查数据加载是否正确看一眼预处理后的张量值是不是合理范围再检查前向输出是否正常最后检查反向传播的梯度是否有效。部署服务出问题也是同理。延迟高了先定位是网络层慢了还是推理慢了推理慢了再定位是算子慢还是显存搬运慢了。用这种“二分切割”的方式配合日志和监控数据绝大多数问题都能在十几分钟内定位到具体模块。切记不要东一榔头西一棒子去乱试。8. 学习路线建议与最终心得8.1 用“项目驱动”而非“知识驱动”来规划学习关于学习路径我最大的体会是不要按教科书顺序线性学习要用项目反向驱动。如果你想掌握Transformer不要先啃完Attention Is All You Need论文再动手而是直接给自己定一个目标用NumPy实现一个能跑通文本生成的小Transformer。在这个过程中你会自然遇到位置编码、掩码、LayerNorm、多头注意力、梯度消失等等一系列问题然后带着问题去查论文、看博客效率高出十倍。这个项目标题的“from scratch”真正想传达的并不是让你完全不用任何库、一切从零造轮子而是让你在关键路径上亲手走一遍理解核心机制。等你把核心链路走通后续再学大模型、多模态、Agent这些前沿方向你会发现它们都是“旧知识的新组合”。8.2 每日修炼的“工程师心法”平时保持技术敏感度的一个习惯值得分享每天固定花半小时看一次Hacker News、arXiv论文列表、以及GitHub上热门项目。特别要关注那些把论文复现成开源项目的仓库读它们的代码比读十篇教程有用。比如想学分布式训练直接把DeepSpeed或者Megatron-LM的源码挑一个关键模块精读比什么都学得快。8.3 最后的一点点感想跑通训练、部署上线、压测达标把这些流程完整走一遍之后我对AI工程的感受是最难的从来不是模型结构而是对全链路的掌控力。AI工程是一门平衡的艺术在效果、速度、成本之间找最优解。这个“从零构建”的过程恰恰是培养这种掌控感的最佳途径。它不轻松但每一步踩下去都踏实。如果你也正在这条路上摸索希望这些经验能帮你少踩几个坑。以上就是我从这个项目里拆解出的全部分享。最后再补充一个实战小技巧当你面临一个全新的AI工程项目时先花两小时把“最小可行性链路”跑通——哪怕只是加载一个最小的模型、跑一个batch的数据、输出一个粗糙的demo——然后再去迭代优化。这个习惯能帮你避免大量无意义的“过早优化”把精力花在真正影响交付的核心环节上。
返回列表