1. 从“调参侠”到“基建工”AI-Infra到底在修什么路干了三年算法模型结构改了几十版指标涨了不到两个点。后来换了个方向把训练数据管道的IO瓶颈解决掉同样的模型直接涨了四个点。这件事让我彻底想明白一个道理AI-Infra不是给算法打杂的它决定了算法能力的上限。很多人对AI-Infra的理解还停留在“帮算法同学装环境、配显卡”的层面。实际上这个领域覆盖的范围远比想象中宽。从最底层的GPU集群调度、分布式通信优化到中间层的训练框架适配、数据加载加速再到上层的实验管理、模型版本控制、推理服务部署每一层都有大量需要深耕的技术点。你可以把它理解成盖楼——算法工程师负责设计房间布局和装修风格AI-Infra工程师负责打地基、通水电、保证承重结构不出问题。地基没打好装修再漂亮也住不了人。这篇文章适合三类人看第一类是想从算法转Infra但不知道从哪下手的工程师第二类是刚入行做AI平台、训练框架相关工作的新人第三类是对大模型训练底层机制好奇、想了解“模型到底是怎么跑起来的”的开发者。我会从实际工作中遇到的问题出发把AI-Infra的核心版图拆开讲清楚包括每个方向到底在解决什么问题、需要掌握哪些硬技能、以及我在实践中踩过的那些坑。先给一个整体认知框架。AI-Infra大致可以分成四个层次硬件资源层GPU/TPU/NPU集群、高速网络、存储系统、调度编排层Kubernetes、Slurm、资源配额管理、训练加速层分布式并行策略、通信优化、混合精度、梯度检查点、工程效率层实验追踪、数据版本管理、CI/CD for ML、推理服务化。这四个层次不是孤立的很多时候一个问题会横跨好几层。比如训练吞吐上不去可能是数据加载的IO瓶颈工程效率层也可能是AllReduce通信没走RDMA训练加速层还可能是GPU被其他任务抢占导致算力碎片化调度编排层。排查这类问题需要你对整个栈都有基本认知。我见过太多团队在AI-Infra上走的弯路花大价钱买了A100集群结果数据管道还是用Python单线程读小文件GPU利用率常年不到30%或者盲目上马分布式训练通信开销比计算开销还大八卡并行还不如单卡跑得快。这些问题的根源都在于——没有从系统层面理解AI工作负载的特征。AI训练和传统Web服务有本质区别它是计算密集型通信密集型的混合负载对IO的吞吐和延迟极其敏感而且不同训练阶段的资源需求差异巨大。用管理微服务的那套思路来管AI训练任务大概率会翻车。2. 数据管道的隐形瓶颈为什么你的GPU在“饿肚子”2.1 从一次典型的训练卡顿说起去年帮一个团队排查训练慢的问题。他们的场景是图像分类ResNet50级别的模型八卡A100按理说应该跑得飞快。但实际训练速度只有理论峰值的四分之一nvidia-smi看GPU利用率像心电图一样上下跳——一会儿90%多一会儿掉到个位数。第一反应是通信问题但检查了NCCL配置和网络拓扑没发现异常。后来用PyTorch Profiler抓了一段trace真相大白GPU有超过40%的时间在等数据。具体原因很典型他们的数据管道用的是PyTorch原生DataLoadernum_workers设了8但每个worker都在从网络存储NFS上读JPEG小文件每个文件几KB到几十KB不等。NFS的元数据操作延迟本来就高加上小文件随机读IOPS直接打满。8个worker根本喂不饱8张A100。更麻烦的是他们的数据增强用的是PIL库做在线增强CPU单核性能成了瓶颈worker进程经常卡在图像解码和变换上。这个问题不是个例。我后来统计了一下在中小规模训练场景中数据管道瓶颈导致GPU利用率低于50%的情况占比超过六成。很多人把注意力全放在模型结构、学习率调优上却忽略了最基础的数据供给问题。2.2 数据管道的三层优化思路解决数据管道瓶颈我一般按三个层次来推进从易到难从低成本到高成本。第一层存储格式优化。把海量小文件合并成大文件这是性价比最高的做法。常用的方案有LMDB、RecordIO、WebDataset、以及HDF5。以WebDataset为例它把每个样本序列化成tar包每个tar包包含几千个样本读取时顺序IO吞吐量比随机读小文件高一个数量级。我实测过一个场景100万张ImageNet图片原始小文件读取吞吐约200MB/s转成WebDataset后顺序读达到2GB/s以上提升十倍。转换脚本也不复杂核心逻辑就是遍历文件列表按固定数量打包写入tar。import webdataset as wds import os source_dir /path/to/images output_pattern /path/to/shards/shard-%06d.tar files sorted(os.listdir(source_dir)) sink wds.ShardWriter(output_pattern, maxcount5000) for fname in files: with open(os.path.join(source_dir, fname), rb) as f: data f.read() sample { __key__: os.path.splitext(fname)[0], jpg: data, } sink.write(sample) sink.close()第二层并行加载与预取。PyTorch的DataLoader本身支持多进程加载但有几个参数需要仔细调。num_workers不是越大越好一般设为CPU核心数的50%到75%同时要配合pin_memoryTrue和prefetch_factor。另外如果用了IterableDataset要注意每个worker的数据分片逻辑避免重复读取。一个常见的坑是num_workers设得很大但每个worker都在做重CPU的数据增强结果CPU上下文切换开销反而拖慢了整体速度。这时候应该把数据增强卸载到GPU上做用NVIDIA DALI或者TorchVision的GPU transforms。第三层缓存与预计算。如果数据增强是确定性的比如固定的resize、归一化可以预先计算好增强后的结果存下来训练时直接读。代价是存储空间翻倍但换来的速度提升很可观。更激进的做法是把整个数据集缓存到本地NVMe SSD上训练时完全不碰网络存储。我们团队现在的标准做法是每个训练节点配2TB以上的本地NVMe训练开始前用rsync或Alluxio做数据预热训练过程中只读本地盘。2.3 一个容易被忽略的细节数据加载的随机性做数据管道优化时有一个问题经常被忽视shuffle的质量。为了追求读取速度很多人会用顺序读大buffer shuffle的方案。但如果shuffle buffer不够大模型看到的样本顺序仍然有很强的局部相关性影响收敛。我一般建议shuffle buffer至少设为数据集大小的10%如果内存允许越大越好。WebDataset的wds.shuffle就是干这个的配合wds.batched和wds.decode使用可以做到高效且随机性足够。还有一个坑是数据分片的均衡性。分布式训练时每个rank读到的数据量应该尽量相等。如果按文件分片而文件大小差异很大就会出现某些rank早早跑完、某些rank还在慢慢读的情况。解决方案是按样本数而非文件数分片或者用resampled dataset让每个epoch的样本分布更均匀。3. 分布式训练的通信账本AllReduce到底花了多少钱3.1 数据并行下的通信开销估算数据并行是最常用的分布式训练策略核心操作是梯度AllReduce。很多人知道AllReduce有开销但具体花多少、什么时候会成为瓶颈心里没数。这里给一个粗略的估算方法。假设模型参数量为P梯度用FP16存储2字节则每轮AllReduce需要传输的数据量约为2P字节发送接收。Ring AllReduce的通信量是2P(N-1)/NN为GPU数量当N较大时趋近于2P。通信时间 数据量 / 有效带宽。以A100为例NVLink带宽约600GB/sInfiniBand HDR约200Gb/s约25GB/s。如果模型是10B参数FP16梯度20GB在8卡NVLink下AllReduce时间约20GB/600GB/s≈33ms在跨节点IB下约20GB/25GB≈800ms。计算时间呢假设单卡算力300TFLOPSFP16模型计算量约6PD FLOPsP为参数量D为token数每轮迭代计算时间 6PD / (300e12 * N)。以10B模型、1M token的batch为例单卡计算时间约610e91e6/(300e12) 200ms。所以在NVLink下通信占比约14%可以接受但在跨节点IB下通信占比高达80%训练效率极低。这个估算说明一个关键结论小规模集群内用数据并行没问题但跨节点扩展时通信会成为主要瓶颈。这时候就需要引入更高级的并行策略。3.2 张量并行、流水并行与专家并行当模型大到单卡放不下或者数据并行的通信开销不可接受时就需要模型并行。主流方案有三种张量并行Tensor Parallelism把单个矩阵运算切分到多卡上每张卡计算一部分然后通过AllReduce汇总。Megatron-LM是典型实现。优点是计算和通信可以重叠缺点是通信频繁对带宽要求极高一般只在节点内NVLink上使用。流水并行Pipeline Parallelism把模型按层切分成多个stage每个stage放在不同卡上数据像流水线一样依次通过。GPipe和PipeDream是代表工作。优点是通信量小只在stage边界传activation缺点是存在流水线气泡bubble需要micro-batch来填充。1F1B调度策略可以把气泡率降到(N-1)/(MN-1)其中N是stage数M是micro-batch数。专家并行Expert Parallelism是MoEMixture of Experts架构的配套方案把不同的专家网络放在不同卡上每个token根据路由结果发送到对应的专家卡上计算。Switch Transformer和GShard是早期代表。核心挑战是负载均衡——如果路由不均某些专家卡会过载其他卡空闲。通常需要加auxiliary loss来鼓励均衡路由。实际训练大模型时这三种并行策略往往是组合使用的。比如GPT-3 175B的训练就用了8路张量并行16路流水并行数据并行的混合策略。组合的原则是张量并行放在节点内高带宽流水并行跨节点低带宽需求数据并行在最外层。3.3 通信优化的几个实操技巧除了选对并行策略还有一些工程技巧可以进一步压榨通信效率。梯度压缩用FP16或BF16代替FP32做AllReduce通信量直接减半。更激进的可以用1-bit Adam、PowerSGD等压缩算法但可能影响收敛需要谨慎评估。通信重叠把AllReduce和反向计算重叠起来。PyTorch的DDP已经内置了这个优化——它在反向传播过程中按bucket触发AllReduce而不是等所有梯度算完再通信。你可以通过调整bucket大小来平衡通信频率和单次通信量。bucket太小通信次数多启动开销大bucket太大重叠效果差。一般建议bucket大小在25MB到100MB之间。拓扑感知的通信组构建NCCL会自动检测网络拓扑但有时候需要手动指定。比如确保同一个NVLink域内的GPU分到同一个通信组跨节点通信走IB而不是绕道NVLink。可以通过NCCL_DEBUGINFO查看实际的通信路径确认没有走弯路。梯度累积如果显存不够大可以用梯度累积来模拟大batch。但要注意梯度累积会增加通信次数每个micro-batch都要AllReduce所以通常配合通信重叠使用。另一个技巧是只在最后一个micro-batch做AllReduce前面累积的梯度先不通信但这需要修改训练框架的底层逻辑。4. 集群调度与资源管理让每一块GPU都不闲着4.1 GPU共享与隔离的几种方案GPU是AI-Infra中最贵的资源如何提高利用率是永恒的话题。最简单的做法是一个任务独占一张卡但这样利用率往往很低——推理任务可能只用了几GB显存训练任务在数据加载阶段GPU空闲。于是就有了GPU共享的需求。目前主流的GPU共享方案有几种NVIDIA MPSMulti-Process Service允许多个进程共享同一张GPU的计算单元通过时间片轮转实现并发。优点是原生支持不需要修改代码缺点是隔离性差一个进程的显存溢出会影响其他进程而且MPS对CUDA Context的管理有额外开销。NVIDIA MIGMulti-Instance GPU是A100及以后架构支持的特性把一张物理GPU切分成多个独立的GPU实例每个实例有独立的显存、计算单元和缓存。隔离性最好但切分粒度固定比如A100可以切成7个1g.5gb实例不够灵活而且不是所有GPU型号都支持。时间片轮转调度Kubernetes的GPU插件默认就是这种模式多个Pod共享一张GPU通过CUDA的时间片调度实现并发。问题是显存不隔离一个Pod OOM会导致整张卡上的所有Pod崩溃。显存超卖软隔离一些公司自研的方案通过拦截CUDA API调用限制每个任务的显存使用量同时用cgroup限制CPU和内存。这种方案灵活但实现复杂需要维护CUDA hook库。我们团队的实际选择是训练任务用MIG做硬隔离保证关键任务的稳定性推理和开发任务用时间片共享提高利用率。MIG的配置通过NVIDIA的k8s device plugin管理在Pod的resource limit里指定nvidia.com/mig-1g.5gb这样的资源类型。4.2 调度策略从FIFO到抢占式调度Kubernetes默认的调度器对AI任务并不友好。它按Pod逐个调度不考虑任务之间的依赖关系和资源整体需求。一个常见的坑是一个分布式训练任务需要8张GPU但集群里只有7张空闲调度器会把Pod一个个挂起等有第8张卡释放时才全部启动。这期间那7张卡就被占着但没干活浪费严重。解决方案是使用Gang Scheduling成组调度。Volcano和Yunikorn都支持这个特性核心逻辑是一个任务的所有Pod要么全部调度成功要么全部不调度。Volcano的PodGroup资源就是干这个的你定义一个PodGroup指定minMember8调度器会等8个Pod都能满足时才一起分配资源。另一个重要特性是抢占式调度。高优先级的任务可以抢占低优先级任务的资源。比如一个紧急的线上模型训练任务可以抢占正在跑的离线实验任务。Volcano支持PriorityClass和PreemptAction配置好优先级后调度器会自动驱逐低优先级任务。但要注意抢占会导致低优先级任务被杀死所以需要配合checkpoint机制让被抢占的任务能从最近的checkpoint恢复。还有一个容易被忽略的点是拓扑感知调度。分布式训练对网络拓扑敏感同一个任务的Pod最好调度到同一台物理机或同一个交换机下。Volcano支持NUMA和网络拓扑感知可以在调度时考虑GPU之间的NVLink连接和节点间的网络带宽。配置方式是给节点打上拓扑标签然后在PodGroup里指定拓扑约束。4.3 资源配额与成本分摊在多团队共享集群的场景下资源配额管理是个绕不开的问题。没有配额先到先得容易导致资源被少数团队霸占配额太死又会导致资源闲置。我见过几种不同的做法静态配额每个团队分配固定的GPU数量不能超用。简单粗暴但利用率低。适合团队间信任度低、需要严格隔离的场景。弹性配额每个团队有保底配额空闲时可以借用其他团队的资源但对方需要时必须归还。这需要调度器支持资源回收和任务抢占。Volcano的capacity插件支持这种模式。按量计费配额结合成本核算每个团队有预算上限超过后需要审批。这需要一套完整的计量系统记录每个任务的GPU时长、存储用量、网络流量等。实现成本高但能有效抑制资源浪费。我们目前用的是弹性配额方案配合每周的资源使用报告。报告会列出每个团队的配额使用率、任务排队时间、GPU平均利用率等指标。数据透明之后大家自然会优化自己的任务减少不必要的资源占用。5. 实验管理与可复现性别让三个月后的自己骂现在的你5.1 实验追踪的最小可行方案AI研发和传统软件开发最大的区别在于实验结果的不确定性极高。同样的代码换个随机种子、换个数据版本、换个硬件环境结果可能完全不同。如果没有一套实验追踪机制三个月后你根本不知道哪个checkpoint对应哪组超参数更别提复现了。最小可行的实验追踪方案需要记录四类信息代码版本git commit hash、环境信息Python包版本、CUDA版本、驱动版本、超参数配置学习率、batch size、优化器参数等、训练指标loss曲线、评估指标、GPU利用率等。工具选择上WandB、MLflow、TensorBoard是主流。WandB的优点是开箱即用、可视化好缺点是数据存在云端有些团队对数据安全有顾虑。MLflow可以自托管但UI和协作功能弱一些。TensorBoard最轻量但只适合单机小规模实验。我们团队用的是MLflow自托管自研的看板核心指标推送到内部Prometheus用Grafana做可视化。关键是要把实验追踪做成自动化的不能依赖工程师手动记录。我的做法是在训练脚本的入口处加一个装饰器自动抓取git信息、环境变量、命令行参数生成一个唯一的experiment ID所有日志和checkpoint都带上这个ID。这样即使工程师忘了记录事后也能追溯。5.2 数据版本管理比代码版本更容易被忽视代码有git管理数据呢很多团队的数据集就是放在共享存储上的一个文件夹今天加点数据明天删点数据没有任何版本记录。结果就是三个月前的实验跑不出来了因为数据已经变了。数据版本管理的方案有几种DVCData Version Control用git来管理数据文件的元信息实际数据存在对象存储上通过hash来追踪版本。优点是和git工作流无缝集成缺点是对于超大文件比如几百GB的checkpoint性能一般。LakeFS提供了类似git的分支、提交、合并语义适合数据湖场景。Delta Lake和Iceberg是数据湖仓的方案支持ACID事务和时间旅行适合结构化数据。对于大多数AI团队我建议至少做到每次训练用的数据集打一个快照记录数据路径、文件列表和hash值。不需要复杂的版本管理工具一个manifest文件就够了。manifest里包含每个文件的路径、大小、md5训练时校验一遍确保数据没被意外修改。5.3 从实验到生产的最后一公里实验跑通了怎么上生产这是很多AI团队卡住的地方。实验环境是Jupyter Notebook生产环境是TF Serving或Triton中间隔着一堆手工操作导出模型、写推理脚本、配负载均衡、做A/B测试。每一步都可能出错。我的经验是从第一天就把训练和推理的接口对齐。训练框架导出模型时同时导出推理所需的预处理和后处理逻辑。比如用TorchScript或ONNX把整个pipeline包括tokenizer、模型、后处理打包成一个文件推理服务直接加载这个文件不需要额外写Python代码。Triton的ensemble功能支持这种模式把多个模型和预处理步骤串成一个pipeline。另一个关键是模型版本管理。每次上线的模型都要有唯一的版本号记录训练数据、代码版本、超参数、评估指标。推理服务支持多版本共存和灰度发布出问题时可以快速回滚。Kubernetes的Istio或Linkerd可以做流量切分把1%的流量导到新版本观察指标后再逐步扩大。6. 一线踩坑实录那些文档里不会写的教训6.1 显存泄漏的排查思路显存泄漏是AI-Infra中最难排查的问题之一。PyTorch的CUDA缓存机制让显存使用看起来只增不减nvidia-smi看到的显存占用往往比实际需要的高很多。真正的泄漏通常来自几个地方未释放的中间变量比如在循环里不断append tensor到list、CUDA graph未销毁、自定义CUDA kernel的显存未释放。排查方法用torch.cuda.memory_summary()看显存分配详情重点关注allocated和reserved的差值。如果allocated持续增长说明有变量没释放如果reserved远大于allocated说明缓存碎片化严重可以调torch.cuda.empty_cache()但注意这会影响性能。更精确的方法是torch.cuda.memory_snapshot()能看到每个显存块的分配栈直接定位到泄漏的代码行。一个常见的坑是DataLoader的worker进程持有CUDA tensor。如果Dataset的__getitem__里创建了CUDA tensor这个tensor会一直留在worker进程的显存里主进程看不到但显存确实被占了。解决方案是Dataset只返回CPU tensor在collate_fn或训练循环里再搬到GPU。6.2 分布式训练中的“幽灵”错误分布式训练最让人头疼的是那些只在多卡环境下出现、单卡跑不出来的错误。我遇到过几次典型情况NCCL超时训练跑着跑着卡住日志里出现NCCL timeout。原因可能是某张卡上的进程挂了其他卡在AllReduce时等不到数据。排查方法是设置NCCL_DEBUGINFO和NCCL_ASYNC_ERROR_HANDLING1让NCCL在检测到错误时主动abort而不是无限等待。另外TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC可以设置心跳超时默认是30分钟可以调短一些。随机种子不一致数据并行时每个rank的随机种子必须不同否则数据增强结果一样但模型初始化的种子必须相同否则梯度对不上。PyTorch的DistributedSampler会自动处理数据分片的种子但模型初始化的种子需要手动设置。我的做法是用torch.manual_seed(42 rank)初始化数据增强用torch.manual_seed(42)初始化模型。梯度为NaN混合精度训练时FP16的梯度容易溢出。PyTorch的GradScaler会自动缩放loss但如果缩放因子调得不好还是会出现NaN。排查方法是先关掉混合精度跑一遍确认模型本身没问题然后逐步调小初始缩放因子或者用BF16代替FP16BF16的动态范围更大不容易溢出。6.3 存储系统的选型与调优AI训练对存储的需求很特殊高吞吐、大容量、低延迟。训练时需要快速读取海量小文件checkpoint时需要高带宽写入大文件推理时需要低延迟加载模型。没有一种存储能同时满足所有需求通常需要组合使用。我们的方案是三层存储本地NVMe做训练数据的缓存和checkpoint的临时存储分布式文件系统如Lustre或GPFS做共享数据集和长期checkpoint存储对象存储如S3兼容存储做冷备和归档。数据预热流程是训练开始前从分布式文件系统把数据拉到本地NVMe训练过程中只读本地checkpoint先写本地再异步同步到分布式文件系统。调优方面有几个参数很关键文件系统的块大小对于大文件顺序读块大小设大一些比如1MB到4MB、客户端缓存Lustre的客户端缓存可以显著提升重复读的性能、网络协议RDMA over Converged Ethernet比TCP性能好很多但需要网卡和交换机支持。另外监控存储的IOPS和带宽使用率很重要瓶颈往往出现在你意想不到的地方。7. 给想入行AI-Infra的工程师几条实在建议如果你是从算法转Infra最大的优势是懂模型、懂训练流程知道算法同学真正需要什么。最大的劣势是系统功底可能不够对操作系统、网络、存储的理解停留在表面。我的建议是先补系统基础再深入某个垂直方向。系统基础包括Linux性能分析工具perf、ftrace、eBPF、网络协议栈TCP/IP、RDMA、存储原理文件系统、块存储、对象存储。这些不需要学到专家级别但至少要能看懂监控指标、定位常见问题。垂直方向的选择上训练框架和分布式并行是目前最缺人的方向门槛也最高需要懂CUDA、NCCL、集合通信算法。推理服务化门槛相对低一些但需求量大涉及模型编译优化TensorRT、TVM、服务编排Triton、KServe、性能调优。数据管道和存储是最容易被忽视但实际影响最大的方向适合细心、喜欢打磨细节的人。工具链方面Python是必须的但光会Python不够。C要能读懂和修改CUDA要能写简单的kernelBash要能写自动化脚本YAML要能配Kubernetes。不需要每样都精通但关键时候能顶上。最后说一个心态问题。AI-Infra的工作很多时候是“幕后英雄”——模型效果好大家夸算法训练挂了大家找Infra。这种落差感是真实存在的。但换个角度想Infra做得好算法才能跑得快、跑得稳、跑得省。每次看到GPU利用率从30%提到80%训练时间从一周缩到两天那种成就感不比调出一个新模型结构差。而且随着模型规模越来越大Infra的重要性只会越来越高。这个方向值得深耕。