ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

AI-Infra实战指南:从GPU调度到分布式训练系统优化

AI-Infra实战指南:从GPU调度到分布式训练系统优化

我第一回正经接触AI基础设施,是在一个特别现实的问题面前——训练任务排队排到天亮,GPU一堆一堆空转,代码却老报OOM。那时候团队里都在卷模型结构、卷数据质量,没人正经想过算力、存储、网络调度这些东西,结果模型卷上去了,训练跑不动,跑得动也烧钱,最后全卡在“基础设施”这四个字上。

后来我慢慢摸到这个领域的门道,发现AI-Infra说白了就是一件事:让算力被高效、稳定、低成本地用起来。它横跨硬件、系统、网络、存储、调度、运维,是模型背后的地基。这篇文章是我的AI-Infra系列第一章,我打算从自己踩过的坑出发,聊聊这个领域到底在解决什么问题、一个一线工程师日常在做什么,以及如果你也想往这个方向走,可以从哪里入手。

这篇内容适合两类人:一是做算法、训练但被环境问题折磨得痛不欲生的同学,二是想从后端、运维、系统方向转型切入AI领域的工程师。我会尽量少讲空洞概念,多讲实际操作和选型背后的逻辑。

1. 先把“AI-Infra”这个词拆开看

1.1 AI-Infra不等于“装个GPU跑训练”

很多刚接触这个领域的朋友,以为AI-Infra就是给服务器插几块卡、装好驱动、配好CUDA,能跑通一个训练脚本就算完事。这个理解不能说错,但格局小了。你想想,单机单卡跑通一个demo,和在生产环境中让成百上千张GPU协同工作、同时支撑几十个团队的任务调度,是两个物种。

AI-Infra的完整拼图至少包括四块:算力资源层(GPU、CPU、内存、异构设备)、数据与存储层(数据集管理、缓存、高性能文件系统)、调度与编排层(任务排队、资源配额、弹性伸缩)、训练与推理运行时(分布式框架、容错、监控告警)。任何一个环节掉链子,模型训练就卡脖子。

我自己的体会是,AI-Infra更像是一个“翻译器”——把上层模型的需求翻译成下层硬件的指令,再把下层的故障和瓶颈翻译成上层能理解的指标。在传统后端,你关心的是QPS、RT、CPU利用率;在AI-Infra,你关心的是MFU(模型算力利用率)、显存占用、数据加载吞吐、NCCL通信效率。

1.2 为什么这两年它突然火了

如果你在两三年前问一个后端工程师什么是AI-Infra,他大概率会摇摇头。但今天这个岗位的招聘需求暴涨,核心原因是模型规模的增长速度远远超过了硬件和管理工具的进步速度。

举个例子,当你的模型参数从亿级涨到千亿级,单卡显存根本放不下,你需要做张量并行、流水线并行、数据并行各种组合,还要面对跨机通信延迟、负载不均、节点故障带来的训练中断。这时候如果基础设施层面没有一个好的设计,算法工程师再厉害也白搭——不是模型跑不起来,就是跑起来三天两头崩,每个epoch还没跑完,集群先炸了。

再加上GPU本身是重资产,一张卡的价格足够买一台不错的家用车,如果利用率上不去、任务排队排出一堆空档,烧的都是真金白银。所以“省钱”“提效”“稳定”这三个关键词,是AI-Infra存在的全部理由。

1.3 一线工程师与算法工程师的视角差异

我刚接触AI-Infra时有个错觉,觉得这活儿比算法简单,不就是搞搞运维吗。后来发现自己错得离谱。算法工程师关心的是“模型怎么收敛”,我关心的是“环境怎么不背锅”;算法工程师要“把GPU用满”,我要“让GPU别被某个任务占着茅坑不拉屎”;算法工程师说“我代码没问题”,我内心想说“你代码里那个DataLoader是不是有毛病”。

这其实是两种思维方式的碰撞。在传统软件工程里,你追求的是确定性和可预测性;在AI-Infra里,你面对的是一个充满随机性的系统——GPU之间通信有延迟抖动、数据读取有瞬时高峰、节点可能随时故障。这种随机性带来的不确定感,跟你在后端写CRUD接口是完全不同的体验。

2. 算力池的搭建与调度:从裸机到集群

2.1 GPU选型背后的数学账

很多团队在GPU选型上特别随意,“哪个贵买哪个”“哪个新买哪个”,结果买回来一堆用不上的卡。其实选型必须回到自己的工作负载上来做匹配判断。

训练密集型任务(如大模型预训练)优先考虑H系列这类高算力卡,因为有大规模并行通信需求,卡间带宽和NVLink拓扑比单卡算力更关键;推理服务型任务则要权衡延迟与吞吐,A系列甚至消费级卡配合优秀的推理引擎也能撑起不错的业务;如果是CV、多模态这类任务,显存容量和带宽也要单独算账。

我在实际项目中总结过一个简单估算公式:整卡利用率 = 有效计算时间 / (有效计算时间 + 通信时间 + 等待时间)。如果你买卡只盯着算力峰值,忽略了集群网络带宽和存储性能,那大概率通信和等待时间会吃掉一半以上的性能空间。这也是为什么很多团队发现“换了更好的卡,训练反而没快多少”的原因。

2.2 从Kubernetes到GPU虚拟化:调度器的选择

当前AI-Infra领域最主流的算力调度方案,不用质疑,就是Kubernetes加上各种GPU扩展插件。Kubernetes本身不直接认识GPU,你得通过Device Plugin、调度器扩展和显存虚拟化工具这几层来补齐能力。

我在生产环境里用的是这么一套组合——Kubernetes负责资源编排,Generic Device Plugin暴露GPU资源,自定义调度器处理拓扑感知(比如一组训练任务尽量调度在同一个交换机下,减少跨机通信),并用显存虚拟化方案做显存超卖。为什么做超卖?因为不是每个任务都把显存卡得死死的,很多小任务实际用不到整卡,超卖就能把碎片化的显存利用起来,把集群整体利用率从40%拉到70%以上。

但也有个前提,超卖必须配合资源隔离和优先级抢占。我经历过一次P0事故,一个贪心任务把显存吃满,其他任务全部挤爆,整个集群的训练质量集体跳水。所以做超卖前,一定要先定好Quota和抢占策略,不能裸奔。

2.3 显存管理:OOM的终极解法

OOM(Out of Memory)是AI训练里最臭名昭著的错误之一。它背后往往不是“显存真的不够”,而是“显存管理太粗糙”。我见过的显存浪费有三个大坑:一是PyTorch的缓存分配器把显存留在手上不释放,二是DataLoader预取的Batch堆积在显存里,三是中间激活值(Activation)没有做内存换出。

解决第一个问题的常规手段是设置PYTORCH_CUDA_ALLOC_CONF环境变量,我建议直接用expandable_segments:True加max_split_size_mb组合参数。这个配置最直接的效果是减少显存碎片,实测下来能提升5%到10%的可用显存,而且不需要改代码。第二个问题要靠调整num_workers和prefetch_factor,别盲目拉大这两个参数,它们会直接影响显存水位;第三个问题在超大模型训练时几乎必须做——用Activation Checkpointing(重计算)或Offload方案,拿一部分计算换显存空间,这是业界标准操作。

3. 数据管线与存储加速:被忽视的性能刺客

3.1 数据加载为什么能把GPU饿死

很多团队在训练性能优化时,第一反应是调模型,却忽略了数据管线这个隐形瓶颈。GPU的计算速度快得惊人,但如果你喂数据的速度跟不上,GPU只能空转等待。我见过一个典型案例,某个团队买了8张高端卡,结果训练时GPU利用率长期在30%以下,排查半天发现是数据加载卡住了。

数据加载慢的原因通常有三层:底层存储的IOPS和带宽不够;中间缓存层没有做好热数据管理;上层数据预处理(解码、增强、打乱)太耗时。这三层如果都不处理,那GPU就是在做“吃饭等菜”的煎熬。

3.2 高性能存储方案怎么选

数据存储这块,我看到很多团队在用传统的NFS(网络文件系统)或者直接在本地磁盘上放数据。小规模训练问题不大,但一旦数据集到TB级别、训练节点数量增加,NFS很容易成为瓶颈。

业界目前比较成熟的方案大致分三类:一是并行文件系统,适合超大规模集群共享访问,性能和扩展性最好,但部署运维成本高,需要专门的团队维护;二是分布式对象存储,灵活且成本低,适合做数据湖和冷热分层,但训练直接读写对象存储延迟偏高,通常要和缓存配合;三是本地NVMe缓存加速,在每台训练节点上挂高速盘,把热数据预取到本地,这是见效最快、性价比最高的方案,也是我最推荐起步团队先做的。

我自己实践过的路径是:大数据集放在S3兼容的分布式存储里,训练前把常用数据集预热到各节点的本地NVMe盘上,训练过程中命中本地缓存,缓存不命中再回源分布式存储。这套组合能做到大部分数据读写在本地完成,训练速度提升立竿见影。

3.3 DataLoader调优的实战参数

有时候不需要上复杂的存储方案,光调DataLoader就能救回大量性能。我建议所有训练脚本都先检查这几个参数:num_workers是不是设成了0(这是默认值,很多人没改就直接上多机训练);pin_memory有没有打开(它能把数据从CPU固定内存直接传输到GPU,减少一次拷贝);persistent_workers这个参数是否开启(它能省掉每个epoch重复启动worker的开销)。

这三个参数组合起来效果惊人。我曾经帮一个团队排查训练速度慢的问题,发现他们num_workers=0跑了好几个星期,仅仅把这个参数调到CPU核心数的一半,训练吞吐直接翻倍。优化数据管线往往是AI-Infra性价比最高的操作——不用加一台机器,不用改一行模型代码,纯靠系统配置就能拿到几倍的速度提升。

4. 训练稳定性:分布式框架与容错机制

4.1 分布式训练框架选型:DP、PP、TP怎么组合

当你需要把模型做到单卡装不下,就要考虑分布式训练。主流框架比如PyTorch DDP(分布式数据并行)、DeepSpeed、Megatron-LM,它们各自的侧重点不一样。

DDP适合模型能放进单卡但想用多卡加速的场景,简单易用;DeepSpeed的ZeRO系列做了参数、梯度和优化器状态的分片,能大幅降低显存占用,是大模型训练的入门首选;Megatron-LM则更偏向超大模型的张量并行和流水线并行,适合千亿级甚至万亿级参数的训练场景。

我个人的建议是,不要一上来就追求最复杂的并行策略组合。先搞清楚你的瓶颈是显存、算力还是通信。显存不够,优先用ZeRO;通信成为瓶颈,再考虑张量并行、流水线并行这些切法;如果都不够,再上混合并行。从简到繁,每次只引入一个变量,出问题才好定位。

4.2 训练的“三高”问题:高波动、高故障、高成本

分布式训练最难搞的就是稳定性。我维护过的一个训练集群,高峰期几乎每天都有一个任务因为节点故障、网络抖动或者共享存储超时中断。训练中断的代价不仅仅是那段时间的算力白费,更重要的是大模型的训练有连续性,中断后即使从checkpoint恢复,也可能因为数据集顺序、随机数状态不一致而影响模型质量。

应对“三高”问题,我的核心思路是三点:一是快速失败,让问题尽早暴露,而不是等一个小时之后才报错;二是自动恢复,通过弹性调度把失败任务重新拉起,从最近的checkpoint续跑;三是checkpoint的高频化,现在训练框架普遍支持异步保存checkpoint,尽量打开,宁可多占点存储,不要拿训练进度冒险。

4.3 通信瓶颈如何定位:NCCL那张看不见的网

在多机训练中,GPU之间的通信是通过NCCL库调度的,通信效率往往决定了整个集群的训练效率。NCCL出问题的典型症状是训练吞吐忽高忽低,GPU利用率上不去,日志里出现大量通信超时或者hang住的迹象。

排查通信瓶颈,我看三个地方:第一是物理拓扑,GPU之间是否通过NVLink直连、跨机走的是InfiniBand还是RoCE,拓扑决定了理论带宽上限;第二是环与树的通信模式,NCCL会自己探测拓扑并选择最优算法,但有时候需要手动设置NCCL_P2P_LEVEL或NCCL_SOCKET_IFNAME来纠正它选错网络;第三是竞争问题,多个任务同时跑在同一批节点上,网络流量互相干扰,这时需要通过拓扑感知调度把任务错开。

我印象最深的一次事故,是在一个混合网络环境里(同时有InfiniBand和以太网),NCCL默认走了慢速网络导致训练速度骤降。后来通过检查IB端口状态和ibstatus命令确认了物理链路没有问题,再挨个排查环境变量,最后锁定是NCCL_SOCKET_IFNAME指向了错误的网卡,改成正确的接口后训练速度立刻恢复正常。那个排查过程相当痛苦,但确实让我把NCCL的底层逻辑吃透了。

5. 实操复盘:我在生产中踩过的那些坑

5.1 案例一:GPU利用率暴跌,罪魁祸首是磁盘IO

有个业务团队找我说,训练任务突然变得特别慢,GPU利用率只有20%不到。我先看了监控面板,发现CPU使用率也不高,但iowait特别高,磁盘读延迟飙到几百毫秒。于是判断大概率是存储层出了问题。

进一步排查发现,他们在训练中途加入了一个新的数据集预处理逻辑,会在每个epoch开始时重新扫描几千个小文件,而对象存储在大量小文件场景下的读取性能极差,每次扫描都耗费大量时间在元数据操作上,导致GPU长时间空闲等数据。最终的处理方案是:把预处理后的数据打包成几个大文件,放到本地NVMe缓存里,同时把num_workers调高来掩盖剩余的低频IO,GPU利用率直接从20%拉到了85%以上。

这个案例其实特别典型,很多人遇到训练慢第一反应就是调模型、调超参,很少有人想到去看看存储和IO。我的经验是,遇到训练性能问题,先看数据管线,再看网络,最后才怀疑模型代码——这个排查顺序能在90%的场景下迅速命中问题。

5.2 案例二:一次OOM引发的调度策略重构

有次集群里一个训练任务频繁OOM,我上机器一看,发现这哥们儿的代码里有个显存泄漏的Bug,跑的时间越长占用越大,到了某个临界点就爆掉。当时集群还没有做显存限额,一个任务就能吃光整张卡,导致同卡其他任务全部被挤下线。

这个事故让我下决心重做调度策略。首先是给每个任务强制设置显存上限,超过配额自动触发回收,防止一个任务拖垮全局;其次是引入优先级队列,重要任务可以插队,但不允许抢占正在运行的高优任务;最后是给每个节点配置了显存使用的水位告警,超过80%就提前预警。这套策略上线之后,虽然偶尔还是有OOM,但影响范围被控制住了,不再是一崩全崩的局面。

我觉得做AI-Infra这事儿,最重要的不是做出多炫酷的功能,而是先把边界划清楚、把影响面控制住。基础设施的职责是“不让单点故障变成全局灾难”,这一点在AI训练场景里尤为突出。

5.3 常见问题速查表:一线排查手册

结合我自己的实操经验,给读者整理一份排查手册,这张表可以说是我自己的“救火指南”:

症状可能原因排查命令/工具处理建议
GPU利用率低数据加载慢、CPU瓶颈nvidia-smi dmon、htop、iostat优先调DataLoader参数,检查存储IO
训练速度忽高忽低NCCL通信不稳定、网络拥塞ibstat、perf query检查NCCL环境变量和网络链路质量
莫名其妙OOM显存碎片、内存泄漏、激活值堆积nvidia-smi、torch.cuda内存快照配置显存分配器、限制任务显存限额、检查代码
Checkpoint恢复后效果变差数据集顺序变化、随机种子不一致对比日志中的配置和数据集路径固定全局采样顺序和随机种子,校验配置
多任务互相干扰资源配额缺失、优先级未设置查看调度器事件和资源请求建立资源池和抢占策略,设置Quota

6. AI-Infra的进阶思考:从“能用”到“好用”

6.1 可观测性建设:眼睛先亮了,手才能快

做AI-Infra,如果没有一套完整的可观测体系,本质上就是盲人摸象。我见过很多团队,连GPU平均利用率、任务排队时长这种最基础的数据都是手算的,出了问题只能靠经验和运气去猜。这绝对不行。

可观测性要从三个视角同时建设:资源视角(GPU/CPU/内存/网络的实时使用率)、任务视角(每个训练任务的进度、吞吐、失败率)、框架视角(模型训练过程的各种指标,如Loss曲线、梯度范数)。三个视角的数据要串起来看,才能快速定位“是任务太贪婪吃了资源,还是资源不够拖累了任务”。

工具链上,我建议不要上来就搭自研平台——先直接用Prometheus采集指标,配合Grafana做可视化,用Alertmanager做告警,基本够用。等到团队规模大了,再考虑引入更专业的一体化平台。

6.2 成本治理:单位算力产出的极限压缩

由于GPU资源稀缺且昂贵,AI-Infra工程师有一个算法工程师不懂的特殊使命:让每一分钱花得值。成本治理有两个层面:宏观上,做集群的水位管理,提高整体利用率,减少空闲浪费;微观上,做任务的资源画像,搞清楚每个任务到底需要多少算力、多少显存,避免盲目申请资源。

我见过一个团队,所有训练任务都按最大规格申请资源,但实际消耗只有申请的一半甚至更少。这等于一半的GPU在空转。做了任务资源画像之后,把申请量压缩到合理范围,集群整体能承载的任务数直接翻倍,相当于白赚了一批算力。这件事不需要任何高级技术,只需要一个愿意和算法团队沟通的基建工程师。

6.3 如何从零开始学AI-Infra:我的路线建议

看完了前面的内容,如果你对这个方向感兴趣,想转岗或者入门,我根据自己的经历给一条学习路线。

第一步,把单机训练玩透——在一台带GPU的机器上装好驱动、跑通一个PyTorch训练脚本,理解GPU显存、CUDA、DataLoader这些基础概念。第二步,上手多机多卡训练——用DDP跑一个ResNet级别的模型,感受一下分布式训练带来的速度和通信开销。第三步,了解调度系统——学会用Kubernetes部署和调度GPU任务,理解Device Plugin、调度器扩展的原理。第四步,深入分布式训练框架——用DeepSpeed或Megatron训练一个模型,熟悉ZeRO、张量并行和流水线并行的实践场景。

这几步走完,你基本就具备了一线AI-Infra工程师的实战能力。然后再去折腾存储选型、网络调优、可观测性等进阶话题。这个过程不需要你数学多好,但需要你有折腾的耐心和定位问题的嗅觉——这些恰恰是工程师最值钱的能力。

写在最后的一个真实感悟

我做AI-Infra这段时间,最大的心得是:这个岗位从来不拼谁的title光鲜,而是拼在故障发生时,谁能保持冷静、按照正确顺序把问题一层层剥开。GPU会烧、网络会抖、存储会慢、代码会有Bug,这些都不恐怖,恐怖的是你连问题在哪一层都说不清楚。

如果这篇文章能帮你建立起对AI-Infra的整体认知,哪怕只在遇到训练故障时多了一个排查方向,我觉得就值了。后续章节我打算接着聊更多细节——比如Kubernetes调度器的二次开发、NCCL通信调优的实战记录、大模型训练checkpoint策略的设计……

关于第一章节,我个人最后非常想说的是:不要被“AI-Infra”这个看起来很宏大的词唬住。剥开外壳,它就是一群工程师在做一件朴实的事——让模型训练的每一秒钟、每一张卡、每一字节带宽都花在刀刃上。这个领域的大门对踏实的工程师永远敞开,你随时可以进来。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进