★ FEATURED ARTICLE 1. 从业务代码到算力底座一个一线工程师的转型思考干了六年业务后端从CRUD写到微服务治理再到后来带小团队做高并发系统我一直觉得自己在技术这条路上走得还算扎实。直到去年公司启动了一个模型推理平台的项目我被抽调过去做服务化那一层才真正意识到AI-Infra这个领域跟传统后端完全是两套知识体系。不是说后端那套东西没用而是你面对的问题域变了——以前关心的是接口响应时间、数据库连接池、缓存命中率现在关心的是显存碎片、算子融合、通信拓扑、吞吐与延迟的权衡曲线。这些词你光看文档是理解不透的必须亲手踩过坑才能建立起直觉。这篇文章是我个人从传统后端转向AI-Infra方向的第一阶段总结主要面向那些有一定工程基础、想往这个方向靠但不知道从哪里下手的同行。我不会讲太玄乎的东西更多是从一个“刚入门的工程师”视角把我在这个过程中理解到的核心概念、动手实操的关键路径、以及踩过的典型坑尽量完整地分享出来。整个系列我会分成几个阶段来写这是第一章重点放在整体认知框架的搭建和最基础的环境实操上。先说一个我自己的判断AI-Infra不是一个单一的技术栈它更像是一个交叉地带——底层是GPU体系结构和高速互联中间是推理/训练框架和算子库上层是服务编排和资源调度。你不需要每一层都精通但必须知道每一层在干什么、边界在哪里、出问题的时候该往哪个方向排查。这个认知框架的建立比学会某个具体工具重要得多。2. AI-Infra到底在解决什么问题2.1 从一次推理服务的性能异常说起我刚到项目组的时候接到的第一个任务是优化一个已经上线的模型推理服务。这个服务在测试环境跑得好好的单次推理延迟大概80毫秒但上了生产之后P99延迟直接飙到800毫秒以上而且吞吐量死活上不去。当时我第一反应是查日志、看监控、排查网络——完全是后端的惯性思维。结果折腾了两天发现问题根本不在应用层GPU显存碎片化导致频繁的显存分配和释放加上请求批次batch的动态拼接策略不合理导致GPU利用率忽高忽低。这件事给我上了一课在AI-Infra领域性能瓶颈往往不在你熟悉的那个层面。传统后端优化你关注的是CPU、内存、磁盘IO、网络带宽而在AI场景下GPU的计算利用率、显存带宽、卡间通信效率才是决定性的因素。你得学会用nvidia-smi、nsight这些工具去看GPU的真实状态而不是只盯着应用层的指标。2.2 AI-Infra的核心问题域拆解我把AI-Infra要解决的问题归纳为四个层面这个分类不一定严谨但对我自己理解这个领域帮助很大第一层是单卡效率。一块GPU怎么跑满这里面涉及算子优化、内存管理、计算图优化、精度选择FP32/FP16/INT8。你要知道你的模型在GPU上到底是怎么执行的哪些算子是瓶颈显存是怎么分配的。第二层是多卡扩展。单卡跑不动了怎么用多卡这里有数据并行、模型并行、流水线并行等不同策略每种策略对应的通信模式和通信量完全不同。你得理解AllReduce、AllGather、ReduceScatter这些集合通信操作知道NVLink和PCIe的带宽差异对性能的影响。第三层是服务化与调度。模型怎么部署成服务怎么处理并发请求怎么做动态批处理怎么在多个模型之间共享GPU资源这一层跟传统后端最接近但多了GPU资源管理这个维度。第四层是集群与运维。多机多卡怎么组网故障怎么恢复怎么做弹性扩缩容这一层涉及Kubernetes、容器运行时、RDMA网络等基础设施。我个人的经验是新手最容易犯的错误是跳过第一层直接去看第三层的东西。觉得服务化、调度这些跟后端经验接近上手快。但实际上如果你不理解单卡上模型是怎么执行的后面所有的优化都是盲人摸象。2.3 为什么传统后端工程师有优势也有劣势先说优势。传统后端工程师在服务设计、接口定义、并发处理、可观测性建设这些方面是有积累的。AI-Infra最终也要落地成服务也要考虑高可用、可扩展、可监控这些能力是通用的。我在做推理服务的时候很多工程规范、代码组织方式、异常处理逻辑直接复用了之前做微服务的经验。劣势也很明显。对GPU体系结构不熟悉对数值计算没有直觉对通信模式理解不到位。举个例子传统后端做负载均衡你考虑的是请求怎么分发到不同的CPU核心或不同的服务实例但在AI场景下一个请求可能需要在多张GPU上协同完成数据要在卡之间来回传输这个通信开销可能比计算本身还大。你不理解这一点就设计不出合理的并行策略。我的建议是不要试图一下子把所有底层知识都补完而是带着问题去学。比如你在优化推理延迟的时候发现瓶颈在显存拷贝那就去查显存带宽、PCIe带宽、数据搬运的优化方法。这种问题驱动的方式比从头啃CUDA编程高效得多。3. 动手之前环境搭建与工具链准备3.1 硬件环境的选择与考量如果你是自己学习不可能上来就搞多机多卡。我的建议是分阶段来第一阶段单卡环境。一张消费级显卡就够了比如RTX 3060 12GB或者RTX 4090 24GB。显存尽量大一点因为很多模型即使做推理显存占用也不小。如果预算有限也可以考虑云上的按量付费实例用的时候开不用的时候关成本可控。第二阶段多卡单机。如果有条件搞两张同型号的卡可以做一些数据并行和模型并行的实验。注意做多卡实验最好用同型号的卡否则会出现木桶效应慢的那张卡拖累整体性能。第三阶段多机多卡。这个一般只有公司环境才具备个人学习可以通过云服务商的GPU集群来体验但成本较高建议先把前两个阶段吃透。这里有一个坑很多人买卡的时候只看算力TFLOPS忽略了显存带宽和显存容量。实际上对于大模型推理来说显存容量决定了你能不能跑起来显存带宽决定了你跑得快不快。算力反而没那么关键因为很多场景是memory-bound而不是compute-bound。3.2 软件栈的版本管理AI-Infra的软件栈版本兼容性是一个让人头疼的问题。CUDA版本、显卡驱动版本、深度学习框架版本、算子库版本这四者之间有严格的对应关系。我踩过的坑是兴冲冲装了一个最新版的框架结果发现它依赖的CUDA版本跟我驱动支持的版本不匹配折腾了一整天才搞定。我的做法是用容器来管理环境。具体来说用NVIDIA提供的CUDA基础镜像然后在上面装框架和依赖。这样做的好处是环境隔离不同项目可以用不同的镜像互不干扰。下面是一个典型的Dockerfile结构FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04 # 安装Python和基础工具 RUN apt-get update apt-get install -y \ python3.10 \ python3-pip \ git \ wget \ rm -rf /var/lib/apt/lists/* # 安装PyTorch注意版本对应关系 RUN pip3 install torch2.1.0 torchvision0.16.0 torchaudio2.1.0 \ --index-url https://download.pytorch.org/whl/cu121 # 安装其他依赖 RUN pip3 install transformers4.35.0 \ accelerate0.24.0 \ triton2.1.0这个Dockerfile的关键点是基础镜像的CUDA版本、PyTorch的CUDA版本、宿主机的驱动版本三者必须兼容。一般来说驱动版本要大于等于CUDA版本要求的最低驱动版本。你可以通过nvidia-smi查看当前驱动版本然后去NVIDIA的官方文档查对应关系。3.3 必备工具清单以下是我日常工作中用得最多的工具建议都装上工具名称用途安装方式nvidia-smi查看GPU状态、显存占用、温度、功耗随驱动安装nvtop类似htop的GPU监控工具更直观apt install nvtopnsight-systems系统级性能分析看CPU-GPU交互官网下载nsight-compute内核级性能分析看算子执行细节官网下载PyTorch Profiler框架级性能分析看算子耗时pip install torchtriton写自定义GPU内核的Python库pip install triton其中PyTorch Profiler是我用得最多的。它可以在不离开Python环境的情况下给你一个相对完整的性能视图哪些算子耗时最长、显存分配在哪里、CPU和GPU的交互是否有瓶颈。对于刚入门的人来说这个工具的上手门槛最低信息量也足够大。4. 第一个实操让模型在GPU上跑起来并观测性能4.1 模型加载与推理的完整流程我们以一个文本分类模型为例走一遍从加载到推理的完整流程并在每一步观测GPU的状态。这里我用的是一个虚构的轻量级模型结构类似BERT但层数更少方便在单卡上快速实验。import torch from transformers import AutoTokenizer, AutoModelForSequenceClassification import time # 检查GPU是否可用 device torch.device(cuda if torch.cuda.is_available() else cpu) print(fUsing device: {device}) print(fGPU name: {torch.cuda.get_device_name(0)}) print(fGPU memory total: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB) # 加载模型和分词器 model_name fake-bert-base # 虚构的模型名称 tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForSequenceClassification.from_pretrained(model_name) model model.to(device) model.eval() # 准备输入 text This is a test sentence for performance profiling. inputs tokenizer(text, return_tensorspt, paddingTrue, truncationTrue, max_length128) inputs {k: v.to(device) for k, v in inputs.items()} # 预热非常重要 with torch.no_grad(): for _ in range(10): _ model(**inputs) # 正式计时 torch.cuda.synchronize() start time.perf_counter() with torch.no_grad(): for _ in range(100): outputs model(**inputs) torch.cuda.synchronize() end time.perf_counter() print(fAverage inference time: {(end - start) / 100 * 1000:.2f} ms) print(fGPU memory allocated: {torch.cuda.memory_allocated() / 1024**2:.1f} MB) print(fGPU memory reserved: {torch.cuda.memory_reserved() / 1024**2:.1f} MB)这段代码有几个关键点需要解释预热为什么重要GPU的时钟频率是动态调整的刚启动的时候可能处于低频状态。另外CUDA的上下文初始化、cuDNN的算法选择、显存的首次分配这些都有开销。如果不预热你测出来的第一次推理时间会明显偏长影响平均值。torch.cuda.synchronize()的作用CUDA操作是异步的你调用model(**inputs)之后CPU会立即返回但GPU可能还在计算。如果不加同步你测出来的时间只是CPU下发指令的时间不是真正的GPU计算时间。synchronize()会阻塞CPU直到所有GPU操作完成。memory_allocated和memory_reserved的区别allocated是当前实际被张量占用的显存reserved是PyTorch向GPU申请的总显存包括缓存池中未使用的部分。PyTorch的显存分配器会预留一部分显存作为缓存避免频繁向GPU申请和释放。这也是为什么有时候你看到nvidia-smi显示的显存占用比allocated大很多。4.2 用PyTorch Profiler定位性能瓶颈上面的代码只能告诉你总的推理时间但不知道时间花在哪里了。这时候就需要Profiler出场from torch.profiler import profile, ProfilerActivity, schedule with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], scheduleschedule(wait1, warmup1, active3, repeat1), on_trace_readytorch.profiler.tensorboard_trace_handler(./log), record_shapesTrue, profile_memoryTrue, ) as prof: with torch.no_grad(): for _ in range(10): _ model(**inputs) prof.step() # 打印按CUDA时间排序的算子耗时 print(prof.key_averages().table(sort_bycuda_time_total, row_limit20))运行之后你会看到一张表列出了每个算子的CPU时间、CUDA时间、调用次数、显存占用等信息。我第一次看到这张表的时候发现大部分时间花在了矩阵乘法GEMM和LayerNorm上这跟模型结构是吻合的。但有一个意外发现显存拷贝Memcpy占了不少时间原因是输入数据从CPU传到GPU、以及中间某些操作触发了隐式的设备间拷贝。这个发现直接引导我做了两个优化一是把输入数据的预处理也放到GPU上做减少Host-to-Device的拷贝二是检查模型中是否有不必要的.cpu()调用避免GPU到CPU的回传。这两个改动加起来把端到端延迟降低了大概15%。实操心得Profiler的输出信息量很大新手容易看花眼。我的建议是先看cuda_time_total排序找到Top 5的算子然后逐个分析。如果Top算子都是计算密集型的如GEMM、Conv那说明模型本身计算量大优化空间在算法层面如果Top算子中有很多Memcpy或ElementWise操作那说明有工程层面的优化空间。4.3 显存管理的常见陷阱显存管理是AI-Infra中最容易出问题的地方之一。我总结了几个典型的陷阱陷阱一显存泄漏。PyTorch使用引用计数来管理显存如果一个张量还被引用着它的显存就不会被释放。常见的情况是在循环中不断创建新的张量但没有及时释放旧的。比如你把每次推理的中间结果都存到一个列表里那显存就会一直涨。陷阱二显存碎片。频繁地申请和释放不同大小的显存块会导致显存碎片化。即使总空闲显存足够也可能因为找不到连续的大块而OOM。PyTorch的缓存分配器在一定程度上缓解了这个问题但在动态形状的场景下仍然会出现。陷阱三多进程显存冲突。如果你用DataLoader的num_workers0每个worker进程都会独立地申请显存。如果配置不当可能导致显存被多个进程占满。解决办法是设置pin_memoryTrue并使用较小的prefetch_factor或者改用GPU上的数据加载方案。# 查看显存碎片情况的实用代码 print(torch.cuda.memory_summary(deviceNone, abbreviatedFalse))这个命令会输出详细的显存分配情况包括每个大小区间的已分配块数、空闲块数、碎片率等。当你遇到OOM但nvidia-smi显示还有空闲显存的时候大概率就是碎片问题。5. 常见问题与排查技巧实录5.1 典型问题速查表问题现象可能原因排查方向解决方案CUDA out of memory显存不足或碎片用memory_summary查看碎片率减小batch size、清理缓存、重启进程推理延迟远高于预期未预热、同步问题、数据拷贝用Profiler看时间分布预热、加synchronize、减少H2D拷贝GPU利用率低数据加载瓶颈、CPU预处理慢nvtop看GPU利用率波动异步数据加载、GPU预处理多卡训练速度不升反降通信开销大于计算收益看NCCL通信时间占比调整并行策略、增大batch size模型加载后显存暴涨精度问题、冗余副本检查dtype和device使用FP16、删除冗余变量首次推理特别慢未预热、cuDNN算法选择对比首次和后续推理时间增加预热轮次5.2 几个让我印象深刻的排查经历经历一诡异的显存增长。有一次我写了一个流式推理的服务每来一个请求就处理一次。跑了一段时间后发现显存持续增长最终OOM。查了半天代码没发现明显的泄漏点。后来用torch.cuda.memory_summary才发现是tokenizer的输出没有及时释放。因为tokenizer返回的tensor默认在CPU上但我把它转到了GPU然后在后续处理中又创建了新的tensor旧的tensor因为还被某个变量引用着一直没有释放。解决办法是在每次推理结束后显式地del掉不再需要的变量并调用torch.cuda.empty_cache()。经历二多卡通信的隐形瓶颈。在做数据并行实验的时候我发现用两张卡比用一张卡还慢。一开始以为是通信库的问题换了几个版本都没改善。后来用Nsight Systems抓了一次trace发现NCCL的AllReduce操作占用了大量时间而且通信和计算没有重叠。原因是我的模型太小计算量不足以掩盖通信开销。解决办法是增大batch size让每次通信传输的数据量更大摊薄通信的固定开销。这个经历让我深刻理解了计算通信比这个概念——不是卡越多越好而是要找到计算和通信的平衡点。经历三版本不匹配导致的玄学问题。有一次升级了PyTorch版本之后模型推理结果跟之前不一致了差异很小但确实存在。排查了很久最后发现是新版本PyTorch默认启用了TF32精度在Ampere及以上的GPU上导致矩阵乘法的精度略有下降。解决办法是设置torch.backends.cuda.matmul.allow_tf32 False。这个问题在文档里其实有写但很容易被忽略。我的教训是升级框架版本后一定要做数值一致性的回归测试。5.3 独家避坑技巧技巧一始终保留一个最小可复现脚本。当你遇到性能问题或正确性问题的时候不要在大项目里改来改去而是抽出一个最小的可复现脚本。这个脚本只包含必要的代码能稳定复现问题。这样做的好处是排查范围小、变量少而且方便跟别人交流。技巧二用环境变量控制调试开关。我习惯在代码里加一些环境变量控制的调试开关比如DEBUG_PROFILE1时启用ProfilerDEBUG_MEMORY1时打印显存摘要。这样在生产环境关闭调试在排查问题时打开不需要改代码。import os if os.environ.get(DEBUG_PROFILE) 1: # 启用Profiler prof torch.profiler.profile(...) prof.start() if os.environ.get(DEBUG_MEMORY) 1: print(torch.cuda.memory_summary())技巧三记录每次实验的配置和结果。AI-Infra的优化是一个实验驱动的过程你改一个参数性能可能变好也可能变差。如果不记录很快就会忘记哪个配置对应哪个结果。我用一个简单的CSV文件来记录日期、改动内容、batch size、精度、延迟、吞吐、显存占用。这个习惯帮我避免了很多重复劳动。技巧四不要迷信默认配置。框架的默认配置是为了通用性不是为了你的特定场景。比如PyTorch的默认显存分配策略、cuDNN的默认算法选择、NCCL的默认通信参数这些在特定场景下可能都不是最优的。花时间读一读相关文档了解这些配置的含义往往能带来意想不到的收益。6. 从单卡到多卡下一步的学习路径单卡跑通之后下一步自然是多卡。但我要提醒的是多卡不是简单的“卡越多越快”。在进入多卡之前你需要确保单卡上的性能已经调优到位否则多卡只会放大问题。多卡学习的第一个关键概念是并行策略。数据并行是最简单的每张卡持有完整的模型副本处理不同的数据批次然后通过AllReduce同步梯度。模型并行则是把模型切分到不同的卡上每张卡只持有部分参数。流水线并行是模型并行的一种特殊形式按层切分不同卡处理不同层的计算。实际生产中往往是多种策略混合使用。第二个关键概念是通信原语。AllReduce、AllGather、ReduceScatter、Broadcast、Send/Recv这些是分布式训练的基础操作。你需要理解它们的语义、通信量、以及在不同拓扑下的性能特征。NCCL是NVIDIA的集合通信库提供了这些原语的高效实现。学会用NCCL的调试工具如NCCL_DEBUGINFO来看通信过程是排查多卡问题的基本功。第三个关键概念是计算通信重叠。理想情况下通信应该在计算的同时进行而不是串行执行。这需要框架层面的支持如PyTorch的DDP在反向传播过程中就开始同步梯度也需要你在模型设计时考虑通信的粒度。如果通信太频繁、每次传输的数据量太小重叠的效果就很差。我目前还在多卡这个阶段摸索踩了不少坑也积累了一些经验。这部分内容我打算放在第二章详细展开包括数据并行的实操配置、NCCL调优参数、以及多卡场景下的性能分析方法。7. 一些个人体会从业务后端转到AI-Infra最大的感受是知识体系需要重建。以前积累的那些经验有些能用上有些完全用不上。但这个过程并不痛苦反而很有意思——你每天都在学到新东西每天都能感受到自己对系统的理解在加深。另一个感受是工具的重要性。在传统后端领域你可能靠日志和监控就能解决大部分问题。但在AI-Infra领域没有好的性能分析工具你基本上就是瞎子摸象。所以我建议刚入门的朋友先把Profiler、Nsight这些工具用熟它们能帮你节省大量时间。最后一点是保持耐心。AI-Infra的优化往往不是一蹴而就的你需要反复实验、反复调整。有时候你花了一整天只为了降低5%的延迟但这5%在生产环境中可能意味着每年节省几万块的算力成本。这种投入产出比在规模化之后是非常可观的。这个系列我会继续写下去下一章重点讲多卡并行的实操和调优。如果你也在往这个方向走欢迎一起交流踩坑经验。
阅读完成 · 觉得有帮助？
