ARTICLE DETAIL

资讯详情

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

AI-Infra实战地图:从GPU集群到推理优化的工程师之路

AI-Infra实战地图:从GPU集群到推理优化的工程师之路 刚入行那会儿我总以为AI-Infra是属于大厂核心技术部门的高冷领域离自己很远。直到我负责的推荐模型上线前一周训练任务连续三个凌晨在GPU集群上崩溃我才真正意识到——AI项目的性能天花板从来不在算法而在基础设施。那段时间我几乎住在机房里跟NCCL通信超时、磁盘IO瓶颈、显存碎片化轮番较劲等真正把问题一个个摁下去才发现自己在AI基础设施上积累的经验比过去五年后端业务代码还要值钱。这篇文章是一线工程师的AI-Infra之路系列的第一章主要写给两类人一是像我一样从后端、运维、测试转过来刚接触大模型训练和推理基础设施的工程师二是算法团队里那个被逼着去排查环境的同学。我不打算堆概念而是用自己实操过的场景把AI-Infra的完整版图、每个环节的典型坑、以及亲测有效的排查思路串起来。你可以把它当成一张入门实战地图哪疼查哪。1. 算力的真相GPU集群不是买了卡就能跑1.1 GPU利用率到底骗了你多久先聊最刺激的GPU利用率。很多团队上线AI任务后第一件事就是看nvidia-smi发现利用率动不动就95%以上觉得算力用足了。但我要泼一盆凉水——这个数字是典型的报喜不报忧。我遇到过最离谱的例子某个8卡DDP训练任务nvidia-smi显示每张卡利用率都在90%以上但一个step要跑差不多1.8秒而算法同学在单卡小batch上估算的理论step时间应该只有0.6秒。差了三倍。问题出在哪用Nsight Systems拉一遍时间线就明白了真正的计算kernel执行只占了不到一半时间剩下全在等数据。nvidia-smi里的利用率统计的是GPU在某个时间片内是否有kernel在跑它不区分你是在做核心矩阵运算还是在做朴素的张量拷贝。数据加载慢、预处理在CPU端阻塞、或者H2D拷贝频繁都会把利用率垫高但实际训练吞吐低得可怜。所以从我这边给你三条实操建议训练阶段不要只看nvidia-smi务必用nsys或py-spy做一次profiling按时间占比排序找真正的热点。对PyTorch任务把DataLoader的num_workers从默认的0往上调通常吃到CPU核心数的一半以上才有明显改观。算算理论算力利用率 实际吞吐 / (GPU峰值算力 × 并行度)低于50%就要警觉别被面板上的数字自我感动。1.2 集群调度比你想象的更难伺候GPU集群调度是整个AI-Infra里最容易被低估的一块。做后端时我们用K8s调度CPU应用资源不够就等优先级高的插队一切都很标准。但GPU任务有两个特性让普通调度器直接抓瞎一是显存一旦分配出去即使任务闲置也无法共享给别的任务二是多卡训练要求多个GPU必须同时可用最好还在同一个节点或者同一个高带宽域里。我见到过最丑的排布任务A要4张卡任务B要2张卡任务C要2张卡调度器只按剩余显存足够来分配结果把任务A的四张卡拆到了两个节点上本来该走NVLink通信变成了走万兆网卡训练速度直接掉了六成。这种问题肉眼是看不出锅的得靠调度策略层面规避。在实际集群里我强烈建议把资源分配的最小单位从单卡提升到节点或GPU域。可以用K8s的nodeAffinity强制同一个job的Pod落在同一个节点或者直接引入支持gang scheduling成组调度的方案——比如Volcano、Kueue这类项目——保证一组Pod要么全起来要么一个都不起避免死等和碎片化。注意多卡训练里单节点内NVLink带宽通常在600GB/s以上而跨节点走IBInfiniBand一般也就200GB/s到400GB/s跨节点通信往往会成为瓶颈。调度策略的设计目标第一优先级永远是尽量把任务塞进一个节点。还有个小坑是显存管理。很多框架默认会做显存预分配例如PyTorch的缓存分配器一开始会占据几乎全部可用显存导致调度器以为没卡了。加一行export PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True或者用max_split_size_mb限制碎片能让显存利用率提升不少尤其适合多任务混布的场景。2. 数据流转的任督二脉存储与IO被忽略的隐形杀手2.1 你的数据集到底放在哪训练一个百亿参数模型单次epoch要读的数据量动辄几百GB甚至上TB。如果这些数据放在普通的NFS共享存储上刚开始小规模测试没问题一旦并行度上来NFS就成了全集群的公共瓶颈。我接手过一个训练任务数据放在基于机械盘的HDFS上每个epoch读数据要花掉将近45分钟而模型一个epoch的训练才20分钟。数据读取时间比训练时间还长这仗根本没法打。当时我们做了一次存储选型对比结论大概是这样存储方案吞吐特性适合场景我踩过的坑本地NVMe盘单机极高每节点独立数据准备、checkpoint临时存放跨节点数据拷贝麻烦宕机丢数据分布式文件系统JuiceFS、Alluxio等可水平扩展命中缓存后接近本地中大规模训练集共享访问元数据性能需要单独调优对象存储S3、OSS等带宽大但延迟偏高冷数据归档、预训练语料存放频繁小文件读取会慢到怀疑人生传统NFS简单易用小团队、小数据集吞吐上限明显并发一高就卡最终我们的方案是分层存储热数据当前epoch正在用的分片提前预热到每个节点的本地NVMe盘中间层用分布式文件系统做共享缓存冷数据全躺对象存储。效果立竿见影数据读取时间从45分钟压到了5分钟左右整体训练吞吐提升了差不多30%。2.2 DataLoader才是真正的隐藏瓶颈很多人以为存储换好了数据就快了其实另一个大头藏在你的代码里。PyTorch的DataLoader默认单进程加载数据如果你的预处理里有解码图片、做随机增强、tokenize文本这些操作CPU就变成了严重瓶颈GPU只能干等。我建议用这几板斧亲测有效调大num_workers一般设置为8到16具体看CPU核心数和数据复杂程度切不可贪多因为worker太多反而会增加进程切换和内存占用。开启persistent_workersTrue避免每个epoch都重新创建worker进程这个配置在PyTorch 1.9可用能省掉很大一部分启动开销。用prefetch_factor做流水线默认prefetch是2数据量大时可以加到4甚至8让数据加载和模型计算重叠起来。数据格式别再用一堆小文件尽量打成大文件格式。TensorFlow那边用TFRecordPyTorch生态可以用WebDataset或者直接做内存映射mmap的二进制格式小文件海量读会把IOPS打满而顺序读大文件几乎总能跑满带宽。我自己实践下来有个特别反直觉的经验与其频繁用shuffleTrue不如在生成数据集时离线做全局shuffle用确定性的方式打乱并切分好shard训练时再把shard顺序打乱。这样既保证了随机性又避免了在线shuffle对缓存命中率的破坏数据加载更稳定。3. 训练稳定性一场与随机性共舞的持久战3.1 崩溃的N种姿势与止损方案大模型训练和传统后端服务最大的区别在于训练是一个长时间、高成本、不可随意中断的过程。一个千卡规模的训练任务每小时的算力成本可能上万crash一次损失的不仅是时间还有无数电力。常见的崩溃姿势我总结成了一张表每一类我都踩过崩溃类型典型日志我的处理方式CUDA OOMCUDA out of memory降batch size、开梯度累积、检查是否有显存泄漏NCCL通信超时NCCL timeout先查网络IB链路、网卡驱动再查是否有慢节点拖后腿进程死锁日志停在某个barrier不动多半是数据分布不均或某个rank异常退出查NCCL初始化顺序NaN/Infloss突然变NaN查学习率、混合精度缩放、输入数据是否有异常值节点宕机某台机器心跳消失训练框架必须支持节点故障自动容错别指望人工盯着这里我特别想展开说下NCCL超时。它可能是AI-Infra里最折磨人的报错之一。一开始我以为是代码问题反复调torch.distributed的初始化折腾了两天没有任何进展。后来用nvidia-smi topo -m一看发现任务被调度到了跨节点的GPU组合PCIe拓扑还不对称NCCL默认选的通信路径绕了个大弯。解决方案也很粗暴但有效在启动命令里加NCCL_IB_DISABLE1先强制走TCP测试确认是网络拓扑问题后再调回IB模式并且用NCCL_P2P_LEVEL精确控制允许的通信层级。3.2 Checkpoint策略你的救命稻草训练到第80个epoch模型效果已经很好了结果第81个epoch开始loss飙升最后直接NaN——你想回滚到第80个epoch的权重发现没有存checkpoint。这种惨剧在团队里真实发生过。我的checkpoint策略是三级递进每N个step比如500步存一个latest.pt用来在异常退出后恢复到最近状态N的选择要平衡存储成本和恢复粒度。每个epoch结束存一个带epoch标记的checkpoint这是你的后悔药用于效果回退时选取最佳版本。定期清理策略只保留最近若干个epoch的checkpoint外加最优的top-K避免磁盘被写爆。还有两个容易忽略的点第一checkpoint不仅存模型权重还要存优化器状态Adam的一阶二阶矩、学习率调度器的进度、数据加载的offset和随机种子否则恢复训练后的loss曲线和之前对不上第二大模型单卡显存放不下整个模型时要保证能分片保存比如用torch.save配合state_dict的_metadata做分片或者直接用框架自带的分片checkpoint能力别指望单卡能完整加载。我用llama类模型做过对比完整checkpoint的保存耗时和存储占用大概是纯权重的3倍左右但这多出来的开销在故障恢复时都是真金白银。稳定的训练比省时省钱的checkpoint策略重要得多这个优先级一定要摆正。4. 推理上线的最后一公里服务化没那么简单4.1 离线训练到在线服务的思维切换训练跑通了只是万里长征的一半把模型部署成在线服务又是一个全新的战场。训练阶段你可以接受秒级的batch吞吐推理阶段用户可不答应几百毫秒的延迟。更关键的是训练和推理的优化目标是相反的训练追求吞吐最大化推理需要同时约束延迟、吞吐和成本。我上手推理服务的第一件事就是抛弃原来裸加载模型一个请求一个进程的天真想法。那时我们直接把PyTorch模型包了个FastAPI上线一看单卡QPS只有个位数延迟动不动就上秒级。后来才意识到两个致命问题一是动态shape导致每次推理都重新做kernel编译和显存分配二是没有做batchingGPU算力被小请求白白浪费。当时的解决方案是引入vLLM做推理引擎核心收益就是continuous batching连续批处理。传统的静态batching必须等一批请求都到齐才一起跑而continuous batching是边来边处理当一个sequence生成完了立刻插入新的请求GPU的空闲窗口被填得严严实实。同样一张A100从裸部署的个位数QPS提升到了接近40倍推理延迟还降了一个量级。4.2 量化与精度效果与速度的平衡木推理优化绕不开量化。大部分开源模型默认是FP16或BF16权重显存占用巨大推理速度也受限。把精度压到INT8或者INT4可以大幅降低显存占用和计算量但代价是精度损失。这里我分享一下我们的选型依据FP16/BF16精度无损显存占用大适合对效果极其敏感且GPU资源充足的场景。INT8W8A8精度损失通常小于1%推理速度提升约1.5到2倍是性价比最高的选择。INT4/FP8显存占用进一步下降可以单卡塞更大模型但量化敏感模型比如某些数学推理任务掉点明显需要做针对性评测。选型的核心方法是分任务评测别只看一两个测试case就拍板。我们用了一套业务侧真实的评测集包含几百条典型样本把FP16的结果作为baseline然后对比INT8和INT4的输出只有指标不掉才允许上线。上线之后还要做A/B观察线上真实流量下的效果波动。我还想提醒一件事量化之后一定要验证推理框架对量化算子的支持是否充分。有些框架的INT8算子只在特定GPU架构上有优化实现在别的架构上可能回退到未优化的实现速度反而更慢。这类问题就是典型的benchmark才能发现看文档永远猜不到。5. 可观测性把黑盒训练引擎变成透明系统5.1 指标四件套训练的可观测体系怎么搭没有观测就没有优化。AI-Infra到后期比拼的其实是谁能更快地定位问题。我每次接手一个新训练任务第一件事不是看代码而是搭三块板子资源监控、训练进度、效率分析。资源监控用DCGM导出GPU指标利用率、显存、温度、功耗、NVLink带宽配合Prometheus Grafana做可视化告警规则覆盖显存使用率突增GPU利用率长时间为0节点失联等核心场景。训练进度则靠框架自己暴露指标PyTorch这边可以挂TorchMetrics或者自定义callback把loss、学习率、吞吐samples/s实时推到监控系统。第三块是效率分析这个很多人忽略。把每个step的时间拆成数据加载时间、前向时间、反向时间、通信时间日常观测你能立刻判断瓶颈在哪。我搭过一个简单的profiling任务每500步用torch.profiler自动跑一次profile把结果落盘供事后分析配合告警基本能做到问题还没影响业务我就已经知道它要来了。5.2 线上推理服务的黄金指标推理服务的观测和训练又不一样。训练关注的是跑得快不快推理服务必须回答三个问题用户等不等得起、系统抗不抗压、模型效果崩没崩。我目前维护的一套推理服务指标包括延迟分位数重点看P50、P95、P99P99尤其重要因为长尾延迟才是用户体验差的根源。吞吐与排队长度QPS、并发数、排队请求数这几个数字配合起来能判断是否该扩容。GPU利用率与显存推理引擎经常出现GPU算力没吃满但延迟已经超标的情况说明batching策略或并发配置需要调。模型质量监控记录每次请求的输入摘要和输出质量分数用离线小模型做实时质量打分一旦质量分数均值下滑就触发告警——这个最容易被忽略但恰恰最重要毕竟线上效果loss才是真正的事故。我强烈建议推理服务从第一天上线就全量记录这些指标不要等出了问题再补。有一次模型悄悄退化用户反馈变笨了我们因为没有质量监控排查了整整两天才定位到是上游数据分布漂移所致。从那以后质量监控就成了推理服务上线的硬性门槛。6. 成本与效率AI-Infra工程师的终极KPI6.1 算力账单到底有多少水分聊完了技术必须聊聊钱。AI-Infra归根到底是一个成本中心你的价值体现在同样的训练效果用更少的钱跑出来。我自己做过一次全面的算力成本审计发现了很多看不见的水分空闲等待任务排队时GPU空转占了总账单的15%以上。低效训练GPU利用率低于50%的任务相当于一半的算力白买了。无节制重试同一个实验反复跑多次checkpoint没存好前面的进度全丢等于双倍成本。过度冗余推理服务为了扛峰值流量长期保持3倍冗余平峰时段大量算力闲置。解决成本水分我的核心手段是把成本可视化建立评估机制每个训练任务在启动时就绑定一个预估成本标签跑完自动计算实际成本每周拉一次成本报表按项目、按任务、按用户维度做对比设定GPU空闲超时回收策略超过15分钟没有实际计算的实例自动释放。这套机制跑了大半年我们整体的算力有效利用率提升了接近30%这在预算不增加的情况下等于白嫖了一批算力。6.2 规模化的几个判断题做到一定程度你会面临规模化的问题要不要上多机多卡要不要引入分布式训练框架我的建议是一切从你的模型多大、数据多少出发别为了用分布式而用分布式。如果单机8卡能跑就不要去碰多机模型超过单卡显存但单机内的多卡够用优先Tensor Parallelism或流水线并行只有模型大到单机完全装不下再上多机分布式。分布式训练框架的选型我踩过不少坑亲身对比下来DeepSpeedZeRO系列显存优化非常成熟适合中小团队用较少GPU训练大模型文档清晰上手快。Megatron-LMNVIDIA出品Tensor并行和序列并行支持完善适合超大模型但配置复杂学习曲线陡。PyTorch原生FSDP和PyTorch生态无缝衔接调试方便适合团队里对PyTorch更熟悉的同学。我给团队的通用建议是先跑通单机DDP作为baseline再按需引入高阶并行策略每一步都做性能复测。别一步到位上Megatron否则你会被一堆并行配置和通信原语搞得怀疑人生最后发现性能并没比简单方案快多少。写在第一章结尾的几点体会回头看完自己第一年的AI-Infra之路最大的感受是这个领域不像后端开发那样有明确的标准答案很多问题要靠现场感知和反复实验。没有人会在文档里告诉你这个任务应该把数据提前预热到本地盘那个模型量化后数学能力会崩这些几乎都是踩了坑、烧了钱才換来的肌肉记忆。如果你问我新人最该从哪开始我会说三件事先把profiling工具用熟你的定位能力会超过一半同行把观测体系搭起来让问题在发生之前就发出信号最后珍惜每一次线上事故那是成长最快的时候。这个系列的第一章就先写到这下一章我想专门聊聊大模型分布式训练里那些让人头秃的并行策略取舍以及我们在千卡集群上的调度实战。欢迎带着问题来交流毕竟AI-Infra这条路上没有一个人是孤军奋战。
返回列表