1. 为什么一线工程师都在往AI-Infra方向靠这两年跟同行聊天话题绕来绕去最后总会落到一个点上纯做业务开发的岗位越来越卷但懂模型训练和推理底层设施的人市场上是真的缺。我自己是从传统后端转过来的踩了不少坑也攒了一些还算拿得出手的经验打算分几篇把这条路上的关键节点梳理清楚。这一篇先聊最顶层的东西——AI-Infra到底在做什么、需要哪些能力、怎么规划学习路径以及我踩过的那些典型的坑。AI-Infra这个词拆开看就是“AI基础设施”说白了就是让模型能跑起来、跑得快、跑得稳的那一整套工程体系。它跟算法工程师的区别在于算法工程师关心模型结构、损失函数、精度指标而Infra工程师关心的是显存够不够、通信带宽有没有打满、任务调度会不会死锁、集群利用率能不能再提两个百分点。你不需要发明新的注意力机制但你需要知道为什么某个算子在特定硬件上会慢三倍以及怎么把它优化回来。适合看这个系列的人我大致分三类一是做了几年后端或运维想往AI方向转但不知道从哪下手的二是算法岗出身发现自己训练脚本写得还行但一上多机多卡就各种报错、效率极低的三是刚入行不久对GPU集群、分布式训练这些概念只有耳闻、没有实操的。这三类人的切入点不一样但最终要掌握的核心能力是重叠的。我自己的背景是做了四年多Java后端转过来大概花了大半年才把整个链路摸清楚。回头看最难的不是某个具体技术而是建立对整个系统的全局认知——你得知道数据从磁盘读到显存、经过计算、再写回去中间每一步可能卡在哪里。这个认知建立起来之后具体工具和框架反而是查文档就能解决的事。2. AI-Infra的核心能力模型拆解2.1 硬件层从GPU架构到显存层级很多人一上来就去学分布式训练框架结果连GPU的基本结构都没搞明白遇到OOM就只会调小batch size遇到慢就只会加卡。这是典型的跳步。硬件层的东西不要求你懂电路设计但几个关键概念必须清楚。首先是显存层级。GPU的显存带宽和容量直接决定了你能跑多大的模型、用多大的batch。以常见的80GB显存卡为例一个7B参数的模型用FP16精度加载光权重就要占14GB左右加上优化器状态Adam的话每个参数要存一阶矩和二阶矩又是两倍再加上梯度、激活值实际占用可能是权重的好几倍。这就是为什么很多人发现“模型明明只有14GB怎么80GB的卡还是OOM”。其次是计算单元和通信带宽的关系。单卡内部的计算速度再快如果多卡之间的通信带宽跟不上整体效率就会被拖死。NVLink和PCIe的带宽差距是数量级的这直接决定了你的并行策略该怎么设计。我见过一个案例同样的模型和卡数只是把通信拓扑从PCIe换成NVLink训练吞吐直接翻了一倍多。注意不要只看卡的算力参数TFLOPS显存带宽和卡间互联带宽往往才是实际训练中的瓶颈。选卡的时候这三个指标要一起看。2.2 框架层训练框架与推理引擎的选型逻辑框架层的选择没有绝对的好坏只有适不适合你的场景。训练侧主流的就是PyTorch生态配合DeepSpeed、Megatron-LM、FSDP这些做分布式加速的方案。推理侧选择更多TensorRT、vLLM、TGI各有各的适用场景。选型的核心考量因素我总结为三个维度模型规模、延迟要求、团队技术栈。模型规模决定了你需要什么样的并行策略——单卡能放下就用最简方案放不下就得考虑张量并行或流水线并行。延迟要求决定了推理引擎的选择——在线服务要求低延迟高吞吐离线批处理可以牺牲延迟换吞吐。团队技术栈决定了维护成本——用一个团队没人熟悉的框架后期出问题排查成本极高。我个人的经验是训练侧优先用PyTorch原生FSDP起步简单场景够用遇到瓶颈再引入DeepSpeed或Megatron。推理侧如果追求极致性能且模型结构固定TensorRT是首选如果需要灵活支持多种模型、快速迭代vLLM的PagedAttention机制在显存利用率和吞吐上表现很均衡。2.3 平台层调度、监控与资源管理平台层是最容易被忽视但实际工作中花时间最多的地方。你不可能每次训练都手动ssh到每台机器上起进程这就需要调度系统。Kubernetes是目前的主流选择但原生K8s对GPU的支持需要额外装设备插件对分布式训练的支持也需要通过Operator来增强。监控体系同样关键。训练任务跑崩了你得知道是显存爆了、通信超时了、还是数据加载卡住了。没有完善的监控排查问题就是盲人摸象。我一般会至少监控这几个指标每张卡的显存使用率、GPU利用率、卡间通信带宽、数据加载队列长度、loss曲线。这几个指标基本能覆盖80%的常见问题。资源管理方面核心是提高集群利用率。GPU卡很贵闲置就是浪费。通过合理的排队策略、优先级调度、抢占机制可以把整体利用率从30%提到60%以上。这里面涉及到的配额管理、公平调度、碎片整理都是实打实的工程问题。3. 从零搭建一个最小可用的训练环境3.1 环境准备与依赖安装假设你手里有一台带GPU的机器想从零开始搭一个能跑分布式训练的环境。我以最常见的Ubuntu加NVIDIA显卡为例把关键步骤和容易踩坑的地方过一遍。第一步是驱动和CUDA的安装。这里最大的坑是版本匹配问题。显卡驱动版本、CUDA版本、PyTorch版本三者之间有严格的对应关系装错了就是各种莫名其妙的报错。我的建议是先去PyTorch官网查它推荐的CUDA版本然后根据CUDA版本去NVIDIA官网查需要的最低驱动版本最后确认你的系统内核版本是否支持这个驱动。# 查看当前驱动版本 nvidia-smi # 查看CUDA版本 nvcc --version # 查看系统内核 uname -r安装驱动的时候如果系统之前装过旧版本一定要先彻底卸载干净。我遇到过一次旧驱动残留导致新驱动装上了但nvidia-smi报错的情况排查了大半天。卸载命令根据安装方式不同有差异apt装的用apt purgerunfile装的用自带的卸载脚本。CUDA安装完之后还要装cuDNN。这个东西是NVIDIA的深度神经网络加速库PyTorch底层会调用它。版本同样要和CUDA严格对应装错了不会报错但性能会差很多。装完之后可以用一个简单的矩阵乘法测试一下实际算力是否正常。3.2 单机多卡通信测试环境装好之后别急着跑模型先做一个通信测试。这一步能帮你确认卡间通信是否正常、带宽是否达到预期。NCCL是NVIDIA的集合通信库PyTorch的分布式训练底层就是用它。import torch import torch.distributed as dist # 初始化进程组 dist.init_process_group(backendnccl) rank dist.get_rank() world_size dist.get_world_size() # 创建一个张量测试all_reduce tensor torch.ones(1024, 1024).cuda(rank) dist.all_reduce(tensor) print(fRank {rank}: all_reduce result sum {tensor.sum().item()})跑这个测试的时候重点观察两个东西一是能否正常完成不报错二是通信耗时是否合理。如果卡间通信带宽明显低于硬件标称值可能是拓扑配置有问题需要检查NCCL的拓扑检测日志。提示设置环境变量NCCL_DEBUGINFO可以看到详细的通信初始化日志排查通信问题时非常有用。3.3 分布式训练的最小示例通信测试通过之后可以跑一个最小的分布式训练示例。我建议从PyTorch原生的DDPDistributedDataParallel开始它封装程度高、上手快适合理解分布式训练的基本流程。import torch import torch.nn as nn import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): dist.init_process_group(nccl, rankrank, world_sizeworld_size) def cleanup(): dist.destroy_process_group() class SimpleModel(nn.Module): def __init__(self): super().__init__() self.fc nn.Linear(1024, 1024) def forward(self, x): return self.fc(x) def train(rank, world_size): setup(rank, world_size) model SimpleModel().cuda(rank) ddp_model DDP(model, device_ids[rank]) optimizer torch.optim.SGD(ddp_model.parameters(), lr0.01) for step in range(100): data torch.randn(32, 1024).cuda(rank) output ddp_model(data) loss output.sum() optimizer.zero_grad() loss.backward() optimizer.step() if rank 0 and step % 10 0: print(fStep {step}, Loss: {loss.item()}) cleanup()这个示例虽然简单但包含了分布式训练的核心要素进程组初始化、模型包装、梯度同步。跑通之后你可以逐步替换成真实的模型和数据观察显存占用和训练速度的变化。4. 实操中那些文档不会告诉你的坑4.1 显存碎片化导致的神秘OOM这个坑我踩过不止一次。明明nvidia-smi看显存还有好几个GB空闲但程序就是报OOM。原因通常是显存碎片化——反复申请和释放不同大小的显存块之后虽然总空闲量够但没有一块连续的空间能满足当前需求。PyTorch有一个环境变量可以缓解这个问题PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True。开启之后PyTorch会使用可扩展的显存段减少碎片。实测下来在一些动态shape的场景下效果很明显。另一个相关的设置是max_split_size_mb控制显存分配器单次分割的最大块大小。设置一个合理的值可以减少碎片但设得太小会增加分配开销。我一般会从128MB开始试根据实际情况调整。4.2 数据加载成为瓶颈的隐蔽表现GPU利用率上不去很多人第一反应是模型太慢或者通信太慢但实际原因可能是数据加载跟不上。表现是GPU利用率呈锯齿状波动——计算的时候冲上去等数据的时候掉下来。排查方法很简单在训练循环里记录数据加载的耗时和计算耗时。如果数据加载耗时占比超过20%基本可以确定是瓶颈。解决方案包括增加DataLoader的num_workers、使用更快的存储、提前把数据预处理成更高效的格式比如把小文件合并成大的二进制文件、使用GPU直接加载数据的技术。我遇到过一个典型案例数据存在网络文件系统上每个batch要读几百个小文件IO延迟极高。后来把数据转成单个大文件加索引的方式加载速度提升了将近十倍。4.3 多机训练中的网络配置陷阱单机多卡跑通了上多机的时候又是一堆新问题。最常见的是网络接口选错——机器上有多张网卡NCCL默认可能选了带宽低的那张。表现是单机测试正常多机通信极慢。解决方法是通过环境变量指定NCCL使用的网络接口NCCL_SOCKET_IFNAMEeth0替换成实际的高速网卡名称。另外如果机器支持RDMA一定要确保NCCL走了RDMA通道而不是TCP两者带宽差距巨大。还有一个坑是防火墙规则。多机通信需要开放一些端口如果防火墙没配好表现是进程卡在初始化阶段不动也不报错。这种问题最难排查因为看起来像是死锁。我的经验是先在机器之间用简单的网络工具测试连通性确认基础网络没问题再排查NCCL层面的配置。4.4 常见问题速查表现象可能原因排查方法解决方案显存够但报OOM显存碎片化打印显存分配详情开启expandable_segmentsGPU利用率锯齿波动数据加载瓶颈记录加载与计算耗时增加workers或优化存储多机通信极慢网卡选择错误查看NCCL日志指定NCCL_SOCKET_IFNAME进程卡在初始化防火墙或端口问题测试机器间连通性开放所需端口训练loss不收敛学习率或batch size问题对比单卡与多卡配置按卡数缩放学习率通信超时某张卡异常或掉线检查所有节点状态替换故障节点5. 学习路径与阶段性目标建议5.1 第一阶段把单卡玩明白别急着上分布式。先在一张卡上把模型训练、推理、显存管理这些基本操作搞熟。具体目标包括能独立完成一个模型的训练脚本编写、能看懂nvidia-smi的各个指标含义、能通过调整batch size和精度设置来控制显存占用、能用profiler工具定位性能瓶颈。这个阶段大概需要一到两个月取决于你之前的工程基础。有后端经验的人会快一些因为很多工程思维是相通的。重点是把GPU编程的心智模型建立起来——理解host和device的区别、理解同步和异步操作、理解显存分配和释放的机制。5.2 第二阶段分布式入门与调优单卡玩明白之后开始接触分布式。先从单机多卡开始把DDP跑通理解数据并行的工作原理。然后学习模型并行和张量并行的基本概念知道什么情况下需要用这些策略。这个阶段的核心目标是建立“并行策略选择”的判断力。给你一个模型和一批硬件你能大致判断出该用什么并行方式、需要多少卡、预期效率如何。这个判断力来自大量的实操和踩坑没有捷径。调优方面重点掌握几个关键工具PyTorch Profiler用来定位计算瓶颈、NCCL日志用来排查通信问题、Nsight Systems用来做系统级性能分析。这些工具用熟了排查问题的效率会有质的提升。5.3 第三阶段平台化与自动化当你对单机和分布式的训练都熟悉之后就该考虑平台化的问题了。怎么让团队里的其他人也能高效地使用这些资源怎么管理几十上百张卡的集群怎么自动化训练流程这个阶段涉及的技术栈会更广容器化、Kubernetes、任务调度、监控告警、存储优化。你不再只是一个“会用GPU的人”而是一个“能建设GPU基础设施的人”。这也是AI-Infra工程师和普通算法工程师拉开差距的地方。我个人的体会是平台化阶段最难的不是技术而是平衡。平衡性能和易用性、平衡灵活性和稳定性、平衡成本和效率。这些权衡没有标准答案需要根据团队的具体情况来判断。6. 一些关于工具选型的个人看法工具选型这件事我的核心原则是优先选社区活跃、文档完善、出问题能搜到答案的。AI这个领域变化太快今天流行的框架明天可能就没人维护了。选一个有大厂背书、有活跃社区的工具能帮你省下大量排查问题的时间。训练框架方面PyTorch目前是事实标准这个没什么好纠结的。分布式加速库的选择要看场景DeepSpeed的ZeRO系列在显存优化上做得很好适合大模型训练Megatron-LM在张量并行和流水线并行上更成熟适合超大规模模型FSDP是PyTorch原生的方案集成度最好适合中小规模场景。推理引擎方面vLLM这两年在社区的热度很高它的PagedAttention机制确实解决了显存碎片化的问题吞吐表现很好。TensorRT在延迟敏感的场景下仍然是首选但灵活性差一些模型结构变了就要重新编译。TGI是HuggingFace生态的推理方案跟Transformers库集成得最好适合快速原型验证。监控方面Prometheus加Grafana是经典组合配合DCGM Exporter可以采集GPU指标。日志收集用Loki或者ELK都行看团队习惯。关键是监控指标要覆盖全面不能只看GPU利用率显存、温度、功耗、通信带宽这些都要监控。提示工具选型不要追求“最新最热”要追求“最稳最熟”。一个你完全掌握的旧工具比一个你半懂不懂的新工具效率高得多。7. 关于这个系列后续的规划这一篇主要聊了AI-Infra的整体图景和入门阶段的关键事项。后续我打算按这个顺序往下写第二篇聚焦分布式训练的具体实现把数据并行、模型并行、流水线并行的代码层面细节讲透第三篇聊推理优化包括量化、蒸馏、算子融合这些实用技术第四篇讲平台建设怎么从零搭一个支持多用户的训练平台。每篇都会保持这个风格——少讲概念多讲实操把踩过的坑和验证过的方案分享出来。AI-Infra这个方向内容很多一篇文章讲不完分几篇慢慢聊。如果你也在往这个方向走希望这些经验能帮你少走点弯路。有问题欢迎交流我知道的都会说不知道的咱们一起查。