1. AI Infra 到底在解决什么问题聊 AI Infra 架构之前得先把一个误区掰正很多人以为 AI Infra 就是“买卡、装驱动、跑训练”这跟把云计算理解成“买服务器”一样离谱。AI Infra 的本质是一套围绕算力、数据、模型三要素构建的工程化底座它要解决的核心矛盾是——模型规模的增长速度远远快于单卡算力的增长速度。我拿实际数字说话。2018 年 BERT-Large 是 3.4 亿参数2020 年 GPT-3 是 1750 亿参数两年翻了 500 倍。而同期单张 GPU 的显存从 32GB 涨到 80GB只翻了 2.5 倍。这个剪刀差就是 AI Infra 存在的根本理由单卡装不下模型单机喂不饱训练单集群管不好任务。所以整个 AI Infra 架构的演进方向始终围绕三个问题展开——怎么把模型切开放进多卡、怎么让多卡之间高效通信、怎么让成百上千张卡稳定跑上几周不出事。从架构分层来看一套完整的 AI Infra 通常包含这么几层最底下是硬件资源层GPU/TPU/NPU、高速互联、存储往上是资源调度层K8s、Slurm、自研调度器再往上是训练/推理框架层PyTorch、Megatron、vLLM最上面是平台服务层实验管理、模型仓库、监控告警。每一层的架构选择都会影响最终的训练效率和成本。这篇文章适合三类人看一是刚转入 AI 工程方向的开发者想搞清楚整个技术栈的全貌二是已经在做模型训练但总被 OOM、通信瓶颈、任务失败折磨的算法工程师三是需要为团队做 AI 基础设施选型的技术负责人。我会尽量把每一层的架构决策逻辑讲透让你不仅知道“用什么”更知道“为什么这么选”。2. 算力层架构从单卡到千卡集群的设计逻辑2.1 单机多卡为什么先于分布式训练很多人一上来就想搞多机多卡结果连单机 8 卡的通信都没吃透。我的建议是先把单机 NVLink/NVSwitch 的拓扑搞清楚再往外扩。原因很简单单机内部的通信带宽比跨机网络高一个数量级。以一台典型的 8 卡 A100 服务器为例卡间通过 NVSwitch 互联双向带宽能到 600GB/s。而跨机的 InfiniBand HDR 网卡单口带宽是 200Gb/s也就是 25GB/s。差了 24 倍。这意味着什么如果你的模型并行策略把通信密集的部分放在了跨机链路上训练效率会断崖式下跌。所以架构设计的第一原则是通信最密集的操作尽量留在单机内部。具体到实践Tensor Parallel张量并行通常限制在单机 8 卡内因为它在每一层的前向和反向都要做 AllReduce通信频率极高。而 Pipeline Parallel流水线并行的通信只发生在层与层之间频率低得多可以跨机。2.2 分布式并行策略的组合逻辑现在训练一个大模型基本不会只用一种并行策略而是数据并行 张量并行 流水线并行的组合业界叫 3D 并行。我拿一个 1750 亿参数的模型举例算一下怎么切。假设你有 64 张 80GB 的 A100。模型参数用 FP16 存储需要 350GB 显存。优化器状态用 Adam需要存储一阶矩和二阶矩加上 FP32 的模型副本总共是参数的 12 倍也就是 4200GB。梯度又是 350GB。加起来远超单卡容量。切分方案是这样的先用 8 路张量并行把单层切开这样每张卡上的单层参数变成原来的 1/8。然后用 8 路流水线并行把层分组每组负责模型的一部分层。剩下的 64 / (8×8) 1也就是数据并行度为 1。如果卡更多比如 512 张那数据并行度就是 8。这里有个经验公式张量并行度不要超过单机卡数流水线并行度不要超过模型层数除以每个 stage 最少层数一般至少 2 层。数据并行度则是剩下的所有卡。这个组合不是拍脑袋定的是因为张量并行通信量最大所以要放在最快链路上流水线并行有气泡所以要控制 stage 数量数据并行通信最少所以可以放到最外层。2.3 显存优化的几个关键手段显存是训练大模型最紧的资源。除了并行切分还有几个手段必须掌握。激活重计算Activation Checkpointing前向传播时不保存中间激活值反向传播时重新算一遍。代价是多了 33% 的计算量收益是激活显存降低到原来的平方根级别。对于层数很深的模型这个收益非常可观。ZeROZero Redundancy Optimizer这是 DeepSpeed 提出的方案分三个阶段。ZeRO-1 只切优化器状态ZeRO-2 再切梯度ZeRO-3 连参数也切。切得越狠省得越多但通信量也越大。我的经验是如果跨机带宽有限用 ZeRO-1 或 ZeRO-2 就够了ZeRO-3 在带宽不足时反而会拖慢训练。混合精度训练用 FP16 或 BF16 做前向和反向用 FP32 做参数更新。BF16 比 FP16 的动态范围大不容易溢出现在基本是首选。但要注意某些算子对精度敏感比如 LayerNorm 和 Softmax需要保持 FP32。优化手段显存收益计算开销通信开销适用场景激活重计算高33%无层数深、激活占大头ZeRO-1中无低跨机带宽一般ZeRO-2中高无中跨机带宽较好ZeRO-3高无高带宽充足或单机混合精度中无无几乎所有场景3. 调度层架构让上千张卡不打架3.1 为什么 K8s 不是万能解很多团队第一反应是用 Kubernetes 管 GPU 集群。K8s 确实能管但原生 K8s 对 AI 训练任务的支持有几个硬伤。第一Gang Scheduling成组调度。一个分布式训练任务需要 64 个 Pod 同时起来才能开始但原生 K8s 是一个一个调度的。如果集群资源碎片化可能出现 63 个 Pod 起来了最后一个死活等不到资源前面 63 个就干耗着。这就需要 Volcano、Coscheduling 这类插件来补。第二拓扑感知调度。同样是 8 张卡分布在同一个 NVSwitch 域内和跨两台机器性能差好几倍。原生 K8s 不感知这个需要自己写调度器扩展。第三GPU 共享与隔离。一张 A100 80GB 跑一个小推理任务太浪费但原生 K8s 不支持显存切分。需要 MIGMulti-Instance GPU或者时间片共享方案。所以实际生产环境里常见的架构是K8s 自研调度器 设备插件的组合。调度器负责 Gang Scheduling 和拓扑感知设备插件负责 GPU 资源的上报和分配。3.2 任务队列与优先级设计集群资源永远不够用所以队列和优先级机制是必须的。我见过太多团队因为没做好这块导致重要任务被低优先级任务堵死。一个可落地的设计是多级反馈队列。把任务分成几个优先级在线推理最高交互式训练次之离线批量训练最低。高优先级任务可以抢占低优先级任务的资源被抢占的任务进入等待队列等资源释放后自动恢复。这里的关键是检查点Checkpoint机制。被抢占的任务必须能从中断点恢复否则抢占就没有意义。Checkpoint 的频率需要权衡太频繁影响训练速度太稀疏则恢复时浪费算力。我的经验是每 30 到 60 分钟存一次同时根据任务已运行时长动态调整——跑得越久的任务存得越勤。注意Checkpoint 的存储 IO 往往是瓶颈。建议用高吞吐的并行文件系统并且 Checkpoint 写入走独立网络通道不要和训练通信抢带宽。3.3 故障自愈与弹性训练千卡集群跑几周硬件故障是必然事件。GPU 掉卡、网卡抖动、内存 ECC 错误这些都会导致训练中断。架构上必须假设故障会发生而不是祈祷它不发生。弹性训练的思路是任务不绑定固定的卡数而是有一个最小卡数和最大卡数。当部分节点故障时任务缩容继续跑故障节点恢复后任务再扩容。这要求框架支持动态的并行度调整目前 PyTorch 的 Elastic 和 DeepSpeed 都有一定支持但实际用起来还是有坑尤其是并行度变化后 Checkpoint 的兼容性。故障检测需要多层次的探针GPU 层面看 ECC 错误计数和温度节点层面看心跳网络层面看 RDMA 重传率。任何一层异常都触发告警严重时自动隔离节点。我踩过的坑是只监控了 GPU 利用率结果网卡降速导致训练慢了一半却没告警白白浪费了两天。4. 框架层架构训练与推理的分野4.1 训练框架的核心抽象训练框架的架构演进本质是在易用性和性能之间找平衡。PyTorch 之所以赢了 TensorFlow 1.x就是因为它的动态图让调试变得直观。但动态图有性能损失所以后来又有了 TorchScript 和 torch.compile 来补性能。从架构上看一个训练框架要提供几个核心抽象张量Tensor、计算图Graph、自动微分Autograd、优化器Optimizer、分布式通信Collective。PyTorch 的架构是把这些做成模块化的组件用户可以自由组合。而 Megatron-LM 这类框架则是在 PyTorch 之上做了一层封装把 3D 并行的复杂性藏起来用户只需要改配置。我的建议是中小规模训练直接用 PyTorch FSDPFully Sharded Data Parallel就够了FSDP 本质上是 PyTorch 原生的 ZeRO-3集成度高调试方便。超大规模训练再上 Megatron 或 DeepSpeed因为它们的并行策略更灵活通信优化更极致。4.2 推理框架的性能博弈推理和训练的架构诉求完全不同。训练追求吞吐可以攒大批次慢慢算推理追求延迟用户等不了几秒钟。所以推理框架的架构核心是KV Cache 管理和批处理策略。KV Cache是 Transformer 推理的显存大头。每生成一个 token都要缓存之前所有 token 的 Key 和 Value。序列越长缓存越大。vLLM 提出的 PagedAttention 就是借鉴操作系统的虚拟内存分页思想把 KV Cache 切成固定大小的块按需分配大幅减少了显存碎片。连续批处理Continuous Batching是另一个关键。传统批处理要等一个批次里所有请求都完成才能处理下一批短请求被长请求拖累。连续批处理则是每生成一个 token 就检查有没有请求完成完成了就立刻腾出位置给新请求。这个策略能把 GPU 利用率从 30% 提到 70% 以上。推理框架核心特性适用场景显存效率vLLMPagedAttention、连续批处理高并发在线服务高TensorRT-LLM算子融合、量化极致延迟优化高TGI连续批处理、多模型通用服务中高llama.cppCPU/GPU 混合、量化边缘部署中4.3 量化与蒸馏的架构影响模型压缩技术会反过来影响架构设计。量化把 FP16 降到 INT8 或 INT4显存直接减半或减到四分之一但需要硬件支持对应的低精度算子。蒸馏用小模型学大模型的行为推理成本大幅降低但需要额外的训练流程。从架构角度看量化推理需要在框架层支持量化感知训练QAT或训练后量化PTQ。QAT 在训练时就模拟量化误差精度损失小但流程复杂PTQ 直接对训练好的模型做量化简单但精度可能掉点。我的经验是INT8 量化用 PTQ 基本够用INT4 最好上 QAT否则精度掉得厉害。5. 数据层架构被低估的瓶颈5.1 数据加载为什么总拖后腿训练时 GPU 利用率上不去十有八九是数据加载的问题。我见过太多团队花大价钱买 A100结果 GPU 利用率只有 40%因为数据管道喂不上。数据加载的架构要解决三个问题读取速度、解码速度、增强速度。读取速度靠存储解码速度靠 CPU增强速度靠 GPU 或专用硬件。任何一环慢了整个管道就堵了。一个优化过的数据管道架构是这样的数据预处理好之后存成WebDataset 格式把多个样本打包成 tar 文件放在高吞吐的并行文件系统上。训练时用多个 worker 并行读取每个 worker 负责解码和基础增强复杂的增强如 MixUp、CutMix放到 GPU 上做。这样 CPU 和 GPU 各司其职不会互相等。5.2 数据版本管理与血缘追踪模型训练的可复现性一半靠代码版本管理另一半靠数据版本管理。但数据版本比代码版本难管得多因为数据量大、变化频繁。实际可落地的方案是内容寻址存储。每个数据文件算一个哈希值数据集用文件哈希的列表来标识。这样即使文件路径变了只要内容没变哈希就不变。DVC、LakeFS 这类工具就是干这个的。血缘追踪则是记录“这个模型是用哪份数据、哪份代码、哪组参数训练出来的”。这在排查模型退化时特别有用。我遇到过模型效果突然下降查了半天代码没改最后发现是数据管道里某个预处理步骤的随机种子变了。提示数据版本管理不要追求一步到位。先从关键数据集开始用简单的哈希清单管理跑通了再上专业工具。6. 常见问题与排查技巧实录6.1 训练中断类问题速查现象可能原因排查方法解决手段NCCL 超时网络抖动或节点故障查 NCCL 日志、IB 计数器增大超时、隔离故障节点CUDA OOM批次太大或碎片打印显存快照减小批次、开重计算Loss 突增数据异常或学习率检查数据批次跳过异常批次、降学习率训练变慢降频或通信拥塞查 GPU 频率、网络带宽检查散热、调整拓扑Checkpoint 失败存储满或 IO 瓶颈查存储容量和 IO 延迟清理旧检查点、换存储6.2 几个反直觉的实操经验经验一不要迷信最大批次。很多人觉得批次越大训练越快其实批次大到一定程度收敛需要的步数会增加总时间反而变长。而且大批次对学习率敏感调不好容易发散。我的做法是从小批次开始逐步增大观察验证集损失找到拐点就停。经验二通信优化比计算优化更值钱。在分布式训练里计算优化能提升 10% 就不错了但通信优化搞好了能提升 50%。比如把 AllReduce 换成 Ring AllReduce把 FP32 通信换成 BF16 通信把跨机通信改成单机内通信这些收益都是立竿见影的。经验三监控要监控“慢”而不是“挂”。任务挂了会告警但任务变慢往往没人发现。我建议对每个训练任务记录 tokens/秒 或 samples/秒 的基线偏离基线 20% 就告警。这样能在问题恶化前就介入。经验四Checkpoint 要定期做恢复演练。存了 Checkpoint 不代表能恢复。我见过 Checkpoint 文件损坏、版本不兼容、恢复后精度对不上的各种情况。建议每周做一次恢复演练确保真出事时能顶上。6.3 成本优化的几个切入点AI Infra 的成本大头是 GPU。优化成本有几个方向提高利用率让 GPU 别闲着、用更便宜的卡推理用推理卡训练用训练卡、用竞价实例能容忍中断的任务用竞价成本能降 60% 到 70%。竞价实例的关键是快速检查点和快速恢复。如果恢复要 10 分钟那被抢占的损失就很大。我的做法是把 Checkpoint 存到内存文件系统或者高速 SSD恢复时直接加载把恢复时间压到 1 分钟以内。7. 架构选型的决策框架7.1 按规模选架构不同规模的团队架构选择完全不同。10 卡以下直接用单机多卡加 PyTorch DDP 就够了别搞复杂的调度系统。10 到 100 卡需要 K8s 加简单的队列管理训练用 FSDP 或 DeepSpeed。100 卡以上才需要自研调度器、拓扑感知、弹性训练这些高级特性。我见过 20 卡的小团队非要上全套 K8s Volcano 自研调度器结果运维成本比训练成本还高。架构要匹配规模过度设计是浪费。7.2 按任务类型选架构训练任务和推理任务的架构差异很大。训练任务重吞吐、容忍延迟、需要容错推理任务重延迟、要求稳定、需要弹性扩缩容。所以生产环境里训练集群和推理集群通常是分开的用不同的调度策略和硬件配置。如果资源有限必须混部那要做好资源隔离。训练任务会占满 GPU 和网络推理任务会被影响。可以用 MIG 把 GPU 切开或者用时间片轮转但都有性能损失。我的建议是能分开就分开混部的复杂度不值得。7.3 自研还是用开源这个问题我被问过无数次。我的判断标准是如果你的需求开源方案能满足 80%就用开源如果开源只能满足 50%才考虑自研。因为自研的维护成本极高而且容易陷入“造轮子”的陷阱。大部分团队的 AI Infra 需求用 PyTorch DeepSpeed K8s Prometheus 这套组合都能覆盖。真正需要自研的往往是超大规模团队有特殊的通信优化需求或者有特殊的硬件需要适配。中小团队把精力放在模型和业务上回报率更高。8. 我踩过的坑与最后几句实话做 AI Infra 这些年最大的体会是架构的复杂度应该由业务需求驱动而不是由技术炫技驱动。我见过太多团队为了“技术先进”上了各种复杂方案结果维护成本高企真正用于模型迭代的精力反而少了。另一个体会是监控和可观测性比什么都重要。训练任务跑在几百张卡上出问题是必然的。关键不是不出问题而是出问题能快速定位。我现在的习惯是任何新架构上线前先把监控埋点做好把关键指标GPU 利用率、通信带宽、显存占用、数据加载延迟都暴露出来再谈性能优化。最后一个建议保持对底层原理的理解。框架和工具会变但分布式训练的基本原理——通信、并行、容错——是不变的。理解了这些换什么框架都能快速上手。我见过太多人只会调库遇到框架不支持的场景就束手无策。底层原理才是真正的护城河。