ARTICLE · INTELLIGENCE

战地情报 · 详情页

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

AI-Infra实践指南:容器化、GPU调度与数据IO优化全解析

AI-Infra实践指南:容器化、GPU调度与数据IO优化全解析 模型训练好的团队为什么会突然发现“GPU不够用”我接手某项目的第一周就碰到了这个状况——算法同事说训练跑不动资源同事说GPU经常白占着。两边对不上账问题不在这批卡而在于没有一套能支撑“拿到数据就能训练、训练完能顺利产出模型”的底层链路。这就是公司里开始反复提到AI-Infra时我作为一线工程师最直接的感受不是要造多么宏大的平台而是先让算力、数据、环境、监控这些最基础的东西不再互相拖后腿。这一篇是“一线工程师的AI-Infra之路”的第一章我把重点放在自己实际落地的那段经历上从画清AI-Infra的边界到用容器化解决环境冲突再到GPU调度、数据IO优化、可观测性搭建最后交付一条半自动化的训练流水线。章节里不会讲太多“架构师视角”的宏观规划每一步尽量给出当时的具体动作、踩过的坑和决策理由。如果你也是那个被迫从“写模型”转向“捯饬环境、盯GPU、查IO”的工程师这章应该能帮你少走一些弯路。1. AI-Infra到底管什么先画出边界再动手1.1 团队真实痛点而不是名词包装开始做AI-Infra之前我先把团队最近一个月的状态盘了一遍。模型训练总是要等人等环境装好等数据拷完等上一个任务跑完。更麻烦的是环境冲突有人升级了CUDA版本另一个人的旧训练脚本直接跑不了有人跑推理占满了显存训练任务只能卡在队列里。这些事表面看是“管理问题”实际全是基础设施能力不足。我当时就在白板上写了几类问题环境一致性问题、GPU资源分配问题、数据加载速度问题、任务状态不可见问题。每一类背后都有明确的技术选型和工程落地路径。AI-Infra对整个团队来说不是K8s集群拉起来就完事而是围绕“模型研发生命周期”把环境、算力、数据、流程串起来让算法工程师能专心抠模型结构而不是天天折腾运行环境。1.2 一线工程师的AI-Infra地图我之前也犯过一个错看到公司里有AI-Infra岗位就以为那是纯K8s运维。真正上手才发现AI-Infra覆盖的东西远不止容器和编排。它是立在“底层硬件”和“算法代码”之间的一整层工程能力至少要拆成四条线环境与镜像保证“昨天能跑的代码今天还能跑”算力调度GPU共享、排队、任务优先级、显存隔离数据管道存储选型、数据缓存、加载策略、IO吞吐可观测性GPU利用率、训练进度、任务状态、告警和日志这四个方向不需要全部一步到位但作为一线工程师第一件事是把地图画出来知道自己现在在哪个位置。我把第一章的核心目标定为训练任务从“手工裸奔”到“标准化容器半自动调度基础监控”这个范围足够聚焦也能最快产生价值。1.3 确定优先级的判断标准地图画完之后很容易贪心什么都想上。当时有其他工程师建议直接上K8s加全套调度平台被我按住了。判断优先级我用了两个简单标准最痛的点在哪里最小可用方案是什么。环境影响最大因为每个人都可能被坑GPU利用率只有30%左右因为大量时间浪费在等环境、等人、等数据上。最小可用方案就是先统一镜像再用简单的脚本做GPU排队和调度不必一上来就上大规模集群。先把“训不了、等半天”解决的问题搞定后面再加复杂度。2. 环境一致性是最早的痛从手工装环境到容器化基线2.1 为什么第一步选容器而不是虚拟化环境冲突这件事估计每个团队都遇到过。我加入某项目组时公共训练机上装了三套Python解释器、十几个版本的包每个人都在用自己的venv但GPU驱动和CUDA只有一个版本。一个同事需要A框架的新版本直接升级了公共环境的CUDA运行库另一个同事的代码瞬间跑崩。这种共享物理机模式天然就是不可维护的。容器化能解决的核心问题是环境隔离和可复现性。镜像里把CUDA、cuDNN、Python版本、依赖库全部固化跑训练只需要一条docker run命令谁都不会动到别人的环境。Kubernetes确实更完整但当时团队只有几台服务器K8s维护成本远大于收益。先用Docker加一台镜像仓库服务器基础设施复杂度最低也能达到目的。2.2 镜像构建时最容易踩的坑构建基础镜像我做了三次才稳定主要卡在几个细节上。第一个坑是CUDA基础镜像选择。一开始图省事直接用了带全部依赖的大镜像体积直奔十几个GB拉取一次能等十分钟。后来换成了runtime版基础镜像再手工装CUDA相关依赖镜像体积压到了5GB以内。注意训练阶段需要完整的CUDA工具链所以不能选runtime要用devel版本否则编译算子时会提示缺少nvcc。第二个坑是Python包版本的锁定习惯。早期镜像里都是requirements.txt只写包名结果一个月之后重新构建某依赖出了新版本行为变了原来能跑的训练开始报错。现在构建脚本里会把所有包用pip freeze锁死到精确版本一个版本都不能差。第三个坑是镜像tag管理。刚开始用latest当tag回滚的时候根本不知道当时跑的是哪个版本。后来强制要求每个镜像必须带唯一的构建ID和Git提交号比如train-base-v1.2.3-af4e2c1这样任何一次训练都能准确追溯到镜像内容。2.3 容器内GPU与驱动兼容细节容器里不是说挂上显卡就能用的NVIDIA的驱动是在宿主机上的容器内只有CUDA运行库二者必须通过特定工具对接。当时有同事直接docker run没加--gpus参数进去之后nvidia-smi都看不到卡还以为是镜像问题。正确做法是宿主机装好NVIDIA驱动和container toolkit容器启动时显式加上--gpus all或指定编号。这里还有一个隐藏坑宿主机驱动版本和容器内CUDA版本存在兼容关系。不是CUDA版本越高越好而是容器里的CUDA必须能匹配宿主机驱动支持的范围。我们统一记录了驱动的最大支持CUDA版本并在基础镜像的README里写明约束。2.4 数据卷挂载与权限管理的经验训练数据放在宿主机上容器里必须挂载才能访问这里权限容易出问题。容器默认以root运行出来的训练产物也是root权限宿主机上普通用户无法清理。我踩过几次之后统一在docker run里加--user参数让容器内进程用宿主机当前用户的UID和GID运行。数据目录的IO模式也要提前想清楚。训练大量小文件时直接挂载宿主目录通常比docker copy性能好因为省去了一层拷贝。我习惯把数据集放在宿主机SSD路径容器启动时用-v挂载同时只读挂载数据集避免训练脚本误写数据源。3. GPU调度与显存治理利用率从30%到70%的关键操作3.1 先摸清一张卡的显存占用习惯环境统一之后GPU利用率还是不高因为大家都挤在同一两台机器上谁先占谁用后到的人只能干等。要解决排队和共享先得搞清楚模型训练的显存占用模式。每个训练任务在启动时会申请显存运行过程中占用基本稳定。为了提升整卡利用率可以把显存较小的任务放入同一张卡。我们当时大多数模型跑在T4级别显卡上Batch Size调大后显存能用到25GB左右单任务无法占满一张卡的全部显存多任务共享一张卡是完全可行的。具体调度策略是每张卡最多同时跑两个训练任务并且互相之间通过环境变量限定可见GPU。只做显存隔离还不够因为同一张卡上并行跑多个任务会争抢算力。验证之后发现两个中等规模任务共享一张卡总过Token数比串行跑还高因为单任务很难把SM全部打满。3.2 用状态机脚本实现任务排队K8s上排队机制做起来很重但当时的服务器规模根本不需要。我写了一个轻量调度器每个训练任务启动前调用一个提交脚本脚本会查询当前哪些GPU卡空闲、哪些卡是半空闲状态再根据预设规则分配。核心逻辑像一个状态机。训练任务的执行脚本主要做三件事获取锁并分配GPU、执行训练、释放锁并更新状态。GPU状态存在一个共享目录下的JSON文件里提交脚本只做原子性的更新确保两个任务不会同时抢到同一张卡。排队等待这件事其实不需要自己重造轮子当时有人提出可以像银行叫号那样用一个FIFO队列文件。实际落地时用的是最简单的文件锁加队列列表比上消息队列轻量得多效果也足够了。3.3 显存泄漏的排查与预防显存泄漏是AI-Infra的经典场景。有个任务跑了一周显存占用随时间缓慢上涨最后直接OOM。排查方法很简单盯住nvidia-smi的输出对比内存占用随时间的变化曲线。我在调度脚本里加了一个自动巡检功能每隔一段时间记录一次每个进程的显存占用并把数据写入日志文件。如果发现某个任务显存增长趋势明显就先停止它再做检查避免波及同卡的其他任务。显存泄漏的根源通常是训练脚本中没有释放缓存或者验证集循环时缓存了上一轮的中间结果。PyTorch里torch.cuda.empty_cache能冻住释放但根本解法是检查代码里是否有循环引用。3.4 手动干预的灰度窗口期调度逻辑跑通之后我特意保留了一周的手动干预窗口期。因为自动调度再怎么设计也不如人在边上看着直观。灰度观察期间我会每天看调度日志发现某些任务总是容易被分配到同一张卡上原来是脚本判断显存空闲时没有考虑算力需求只是简单看显存剩余量。后来在状态机里加入预估时长、算力占用的组合评分分配结果才合理。4. 数据加载链路被低估的IO瓶颈与优化实录4.1 从“训练两个小时”到“前半小时在读数据”训练流程跑起来以后我习惯每天盯一次日志。某天发现一个视觉模型任务每次epoch耗时突然从25分钟涨到35分钟但GPU本身的算力没有下降排除了GPU问题之后把怀疑点转向了数据链路。一通排查下来问题出在数据读取上。数据集有几万个图片文件分散在几个目录里每次训练启动的时候要遍历并加载一遍。单机训练时这还好等到多卡训练时每张卡都要读同一份数据文件网络存储的吞吐就成了瓶颈。训练的前30分钟基本都在等数据到达显存GPU利用率自然会掉。4.2 存储选型本地SSD和网络存储的取舍刚开始做AI-Infra最直接的想法是找一台机器把数据全放上去然后网络挂载。真实体验是小团队小数据量时还好数据量超过几十GB之后网络传输带来的延迟和带宽竞争就很明显了。我的方案是分级存储。热数据放本地SSD也就是当前正在训练的数据集直接拷贝到各台机器上冷数据放集中存储。启动训练前有个同步脚本检查数据集版本和本地目录是否一致不一致就自动同步。同步采用增量方式只拉取变更部分。这样做好处明显每台机器的训练IO走本地盘不受网络影响GPU也能更稳定地吃满数据。4.3 DataLoader调优参数与缓存策略同一个模型训练DataLoader配置不合理也会造成严重的IO浪费。最经典的是num_workers设成0数据加载完全在主进程里串行做GPU只能空转等着。我根据CPU核数把num_workers调到了每个GPU对应8个进程加载速度肉眼可见地提升。prefetch_factor也要配合调整它是每个worker预先加载的batch数量。设置成2到4之间通常比较合适太大容易占内存太小仍然会等待。还有一个容易忽略的参数pin_memory开起来以后数据传输从分页内存搬到固定内存能显著减少Host到Device的拷贝开销。数据缓存策略上我把小而热的数据集直接做成LMDB或者WebDataset格式把几万个小文件合并成几个大文件。大文件顺序读取效率远高于随机小文件读取IO延迟从几十毫秒降到个位数毫秒。第一次做这个改造之后训练启动时间直接砍掉了将近一半。4.4 一条本地缓存与全局共享的混合路径用数据同步加本地缓存之后新的问题出现了多台机器上的数据集可能不一致某台机器同步了一半就断了。我设计了一个简单的一致性校验每台机器上的本地数据目录里放一个校验文件内容是数据集的版本Hash。训练启动时调度脚本会先检查Hash是否匹配不匹配就拒绝启动而不是稀里糊涂地拿旧数据开始训练。全局共享存储做兜底本地SSD做热路径中间加一层增量同步这套组合已经覆盖了大多数训练场景。后面如果数据量和训练规模再上一个量级再考虑并行文件系统也不迟。5. 可观测性把黑盒训练变成白盒监控5.1 可观测性三个维度日志、指标、追踪没有监控之前训练任务就像一个黑盒。某个任务卡了一中午只能登服务器看nvidia-smi和进程输出效率极低。可观测性建设不是简单地装个监控系统而是让每个训练任务的状态随时可见。我参照通用服务监控的思路拆成三层日志层负责记录训练脚本的标准输出、错误、异常和关键节点指标层定期收集GPU利用率、显存占用、吞吐量、loss曲线追踪层记录任务从提交到开始训练的时间线这样才能发现任务是在排队、同步数据、还是实际计算环节耗时最长。对于AI-Infra最重要的指标之一是GPU利用率的时间分布。注意不要只看纸面上的平均利用率要关注“时间平稳性”跑一个epoch时GPU利用率是否忽高忽低这能直接反映数据供给和计算是否匹配。5.2 极简方案脚本采集加轻量展示直接上完整Prometheus体系对于一个几台服务器的团队来说太重了我选择了一个极简但有效的方案每台机器上部署一个采集脚本定期把GPU状态、进程信息和训练日志尾部汇总写入时间序列文件。展示端用了轻量级看板工具按机器、按任务维度展示关键指标。这套东西没有用标准监控架构但胜在轻、快、零维护。每个训练脚本自己也会把loss和吞吐量写入结构化日志前端能直接看到实时曲线。如果你有资源直接上PrometheusGrafana也是完全可以的但初期我更建议先把手头数据可视化搞定再逐步完善告警。5.3 告警阈值设定的经验告警不在多而在准。常见的误区是GPU利用率低于某个值就报警但实际情况中模型刚启动时在同步数据、做预处理GPU利用率本来就会短暂下降这种报警只会制造噪音。我后来把告警分成三类第一类是任务异常卡死比如日志超过20分钟没有更新第二类是显存即将打满或超卖这种情况必须立即告警第三类是训练性能严重下滑通过对比过去一周同任务的平均吞吐量如果当前值低于平均值的70%就提示检查。这个阈值不是拍脑袋定的是跑了两个月以后根据历史数据分布设置的。5.4 从监控反推代码缺陷可观测性最大的价值在于反推代码问题。有一次新模型上线后监控面板显示GPU利用率周期性掉到零每半小时出现一次。我顺着时间轴查日志发现每半小时触发一次checkpoint保存保存操作在主进程里同步执行阻塞了训练流程。把保存操作挪到异步子进程之后周期性掉点就消失了。这种问题在纯黑盒状态下几乎不可能找到但有了按时间对齐的日志和指标定位起来就像看CT片子一样清楚。我建议每个训练脚本都定期输出进度同时把checkpoint时间、数据加载时间这些关键阶段记录到日志后续分析会非常省事。6. 团队协作的极简流水线第一章的交付成果6.1 算法工程师的工作流变化基础设施改造不是技术自嗨最终要看团队使用体验。改造前新同事入职第一周基本都在装环境、配依赖、问别人“这个包要怎么装”改造后拿到一份写好的部署文档照着docker run就能把训练跑起来。我开始给团队里每人整理了一份标准操作手册里面只写最常用的20条命令拉镜像、启动训练、查看日志、取消任务、查看GPU状态。算法工程师不再需要关心底层环境细节他们把模型代码交给统一入口的启动脚本即可。最明显的变化是“环境装了一天”这样的抱怨基本消失了。6.2 一个最简单的流水线训练到测试产出第一章的最终交付是一条极简流水线代码提交后自动触发镜像构建构建成功后打上新的版本号算法工程师可以一键提交训练任务。整个链路没有复杂的CI/CD平台只是用了Git的钩子加一个构建脚本再加前面说过的调度器。流水线的核心逻辑就两条版本可追溯操作可重复。每个训练任务都能查到对应镜像版本、代码提交号、数据集版本和运行参数这样出现问题时能够完全复现。如果以后规模变大再考虑上标准流水线平台现阶段这个方案足够稳定。6.3 我在这个阶段最想提醒的三件事第一不要一开始就追求大而全的架构。AI-Infra的复杂度是逐步累积的基础设施建设应该跟着实际痛点走而不是先把所有组件都搭起来再慢慢填内容。第二保持记录的习惯。每次踩坑、每次参数调整、每次优化前后对比都留在文档里。后面再排错时这些记录的价值比任何监控系统都大。第三关注人而不是只关注技术。基础设施要服务于使用它的人。如果团队里的算法同事用起来觉得繁琐那再先进的技术栈都等于零。我花了几天时间写使用文档和示例换来的是团队真正把这个流程用起来。第一章走到这里训练环境统一了GPU利用率有了统计任务调度有排队了数据加载瓶颈缓解了训练过程也不再是黑盒。后面章节里我打算接着写多机分布式训练、模型推理服务化、以及更细粒度的资源成本治理。如果你正在从算法转向基础设施这块或者团队刚开始重视AI-Infra希望这些一线实战经验能有点参考价值。也建议你从一个小痛点入手先跑通一条最小链路后面的事自然会慢慢清晰起来。
RELATED READING

延伸阅读

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