1. AI Infra 到底在解决什么问题先把话说直白一点AI Infra人工智能基础设施不是某一款软件也不是某一个框架而是一整套让 AI 模型能跑起来、跑得快、跑得稳、跑得省的工程体系。它横跨硬件、系统软件、调度平台、通信库、存储、推理引擎、可观测性等多个层面。很多人第一次听到这个词会下意识把它等同于“买几张显卡、装个驱动、跑个训练脚本”但真正在生产环境里摸爬滚打过的人都知道这只是冰山露出水面的那一角。我接触 AI Infra 是从一次很朴素的需求开始的团队要把一个视觉模型从单机搬到多机多卡上训练结果发现单卡能跑通的代码一上多机就各种卡死、掉卡、通信超时、显存溢出。那时候才意识到模型代码只是“应用层”真正决定你能不能把算力用起来的是底下那一整套基础设施。这也是为什么近两年“AI Infra”会成为一个独立且火热的方向——模型架构再漂亮没有配套的基础设施也落不了地。从职责边界上看AI Infra 大致要回答四个问题算力从哪来GPU、NPU、TPU、CPU 等异构硬件的选型、组网、驱动与运行时。算力怎么分集群调度、资源隔离、多租户、弹性伸缩、任务排队。算力怎么用得快分布式训练/推理的通信优化、并行策略、算子优化、显存管理。算力怎么用得稳故障自愈、监控告警、性能剖析、成本核算。这四个问题对应到具体技术上就是分布式架构、微服务架构、Agent 架构、指令集架构、Transformer 架构等一系列概念的组合落地。标题里给的关键词“AI infra、架构”看似宽泛但它恰恰点出了这个领域的本质AI Infra 的核心竞争力不在于你会不会调某个库而在于你能不能设计出一套合理的架构把异构资源、任务负载和业务需求三者对齐。这篇文章适合谁看如果你是刚入行的算法工程师想搞清楚自己训练脚本下面到底发生了什么如果你是后端或运维转过来的工程师想理解 AI 场景和传统 Web 服务的差异如果你是团队里负责选型和搭平台的人想少踩几个坑——那这篇内容应该能给你一些可以直接参考的东西。我会尽量用从业者之间聊天的口吻把架构拆开讲把参数算给你看把踩过的坑摊开说。2. AI Infra 的分层架构与设计取舍2.1 从下到上的五层结构聊架构最忌讳一上来就堆名词。我习惯把 AI Infra 从下到上分成五层这样每一层的职责和边界都清楚出了问题也容易定位到底是哪一层的事。层级名称核心职责典型组件/技术L1硬件与驱动层提供物理算力、互联带宽GPU/NPU、NVLink、RDMA、CUDA/ROCm 驱动L2资源抽象与运行时层把硬件抽象成可调度单元容器运行时、设备插件、CUDA 上下文管理L3集群调度层分配资源、编排任务Kubernetes、Slurm、自研调度器L4框架与通信层分布式训练/推理执行PyTorch、Megatron、DeepSpeed、NCCL、集合通信库L5平台与服务层面向用户的产品化能力训练平台、推理服务、Agent 平台、监控计费这个分层不是学术定义而是我在实际排查问题时脑子里的一张地图。比如“训练任务卡住不动”可能是 L4 的通信死锁也可能是 L3 调度把两个高带宽需求的任务塞到了同一台机器上抢网卡还可能是 L1 的某张卡降频了。有了分层排查就有了顺序而不是盲目重启。2.2 为什么是分层而不是一体化有人会问既然分层这么麻烦为什么不做一个大一统的系统答案在于变化频率不同。硬件每一两年换代调度策略可能每季度调整而业务侧的模型和接口可能每周都在变。如果全部耦合在一起任何一层的小改动都会牵一发动全身。这跟微服务架构的思路是一致的按变化频率和职责边界拆分让每一层可以独立演进。但 AI Infra 又和传统微服务有本质区别——传统微服务是无状态、可水平扩展的而 AI 任务往往是有状态的、强耦合硬件的、通信密集的。所以你不能照搬微服务那一套得做取舍。我个人的经验是调度层尽量通用化通信层尽量专用化。调度层用 Kubernetes 这类成熟方案能省很多事但通信层一定要针对你的网络拓扑比如是不是 RDMA、是不是多轨组网做定制否则带宽利用率可能连一半都跑不到。2.3 架构选型背后的三个核心权衡做 AI Infra 架构绕不开三个权衡我把它总结成一句话性能、成本、可维护性三者最多同时满足两个。追求极致性能用最贵的卡、最好的网络、最激进的并行策略代价是成本和维护复杂度飙升。追求低成本用混部、抢占式调度、异构卡池代价是性能抖动和故障率上升。追求可维护性用标准化组件、统一接口代价是可能牺牲一部分极限性能。举个具体例子。训练一个大模型时是否要用张量并行Tensor Parallelism张量并行能降低单卡显存压力但会引入大量通信。如果你的集群是 NVLink 全互联通信开销可以接受如果是普通以太网那通信就会成为瓶颈这时候可能更该用流水线并行Pipeline Parallelism配合梯度累积。这个决策不是拍脑袋而是要看你手里的网络带宽和模型规模算出来的。3. 核心组件拆解与实操要点3.1 分布式训练架构并行策略怎么选分布式训练是 AI Infra 里最硬核的部分。核心就三种并行方式我逐个说清楚它们适合什么场景。数据并行Data ParallelismDP是最常见的。每张卡持有一份完整模型副本喂不同的数据最后同步梯度。优点是实现简单PyTorch 的 DDP 开箱即用缺点是模型必须能塞进单卡显存。当模型参数量超过单卡容量时DP 就无能为力了。张量并行Tensor ParallelismTP把单个算子比如矩阵乘切到多张卡上。适合单层就很大的模型比如 Transformer 里的注意力层。缺点是通信极其频繁几乎每层都要 All-Reduce所以强烈建议只在 NVLink 这种高带宽互联的卡间使用。流水线并行Pipeline ParallelismPP把模型按层切成若干段每段放在不同卡上数据像流水线一样流过。优点是通信量小缺点是会有“气泡”bubble即某些卡在等待。为了减少气泡通常配合微批次micro-batch使用。实际生产中大模型训练基本是TP PP DP 的混合并行。这里给一个我常用的估算方法假设模型有 N 层用 P 路流水线并行微批次数为 M那么流水线气泡占比约为bubble_ratio ≈ (P - 1) / (M P - 1)举例P8M16则气泡占比 ≈ 7 / 23 ≈ 30%。这意味着有 30% 的算力在空转。要降低这个比例就得增大 M但 M 受限于显存。所以流水线并行度不是越大越好要结合显存和批次大小一起算。注意混合并行的并行度配置是个组合优化问题不要凭感觉设。我见过太多团队把 TP 设成 8 结果网络带宽不够训练速度还不如 TP2。3.2 推理服务架构从单机到集群训练是一次性的推理是长期在线的两者的架构思路完全不同。推理服务最关心的是延迟、吞吐和成本。单机推理很简单加载模型、起个 HTTP 服务就行。但生产环境要考虑批处理Batching把多个请求合并成一个批次送进 GPU能大幅提升吞吐。但批处理会增加延迟所以要设一个“等待窗口”比如最多等 10ms 凑批。模型并行大模型单卡放不下时推理也要做并行但推理的并行策略和训练不同更偏向张量并行。多副本负载均衡同一模型起多个副本前面挂负载均衡。这里要注意GPU 推理副本的负载均衡不能简单用轮询因为不同请求的计算量差异很大最好用基于队列长度的动态调度。我实测下来一个常见的坑是批处理窗口设得太大导致 P99 延迟爆炸。比如窗口设 50ms平均延迟可能只涨一点但尾部延迟会非常难看。建议从 5-10ms 起步根据业务 SLA 调整。3.3 调度层Kubernetes 在 AI 场景的适配Kubernetes 是调度层的事实标准但原生 K8s 对 AI 任务支持并不好主要问题是GPU 资源是整数分配的一张卡不能分给两个任务除非用 MPS 或 MIG。不支持 Gang Scheduling即一个分布式任务的所有 Pod 必须同时启动否则会死锁。拓扑感知弱不知道哪些卡在同一个 NVLink 域内。解决办法通常是引入扩展组件。比如用 Volcano 做 Gang Scheduling用设备插件暴露 GPU 拓扑用 MIG 做显存切分。这里给一个我常用的 GPU 资源声明示例resources: limits: nvidia.com/gpu: 2 nvidia.com/mig-1g.5gb: 1提示MIGMulti-Instance GPU能把一张 A100 切成 7 个实例适合推理小模型混部。但 MIG 实例之间不能做 NVLink 通信所以训练任务不要用 MIG。3.4 通信层集合通信是性能命门分布式训练的性能八成取决于通信。集合通信库如 NCCL提供了 All-Reduce、All-Gather、Reduce-Scatter 等原语。其中 All-Reduce 最常用它的实现有 Ring 和 Tree 两种算法。Ring All-Reduce所有卡围成一个环数据分块传递。带宽利用率高但延迟随卡数线性增长。Tree All-Reduce树形结构延迟低但带宽利用率不如 Ring。选择依据很简单卡数少、消息小用 Tree卡数多、消息大用 Ring。NCCL 会自动选择但你可以通过环境变量干预export NCCL_ALGORing export NCCL_PROTOSimple另外网络拓扑对通信影响巨大。如果跨机通信走的是普通以太网那跨机 All-Reduce 会成为瓶颈。这时候可以考虑梯度压缩、分层 All-Reduce机内先聚合再跨机等优化手段。4. 完整实操流程从零搭一个最小可用训练平台4.1 环境准备与依赖安装假设你手上有 4 台机器每台 8 张 GPU想搭一个能跑分布式训练的最小平台。我按实际顺序列一遍。第一步统一基础环境。所有机器的驱动版本、CUDA 版本、Python 版本必须一致否则会出现各种玄学问题。我习惯用 Docker 镜像锁定环境FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04 RUN apt-get update apt-get install -y python3.10 python3-pip RUN pip install torch2.1.0 --index-url https://download.pytorch.org/whl/cu121第二步配置免密登录和主机名解析。分布式训练需要节点间能互相 SSH且主机名要能解析。在/etc/hosts里把 4 台机器的 IP 和主机名写死比依赖 DNS 稳定得多。第三步验证 NCCL 通信。这一步千万别跳过我见过太多“训练慢”最后发现是 NCCL 走了错误网卡。用官方测试工具跑一遍./build/all_reduce_perf -b 8 -e 128M -f 2 -g 8看输出的 bus bandwidth如果远低于网卡理论带宽说明网络配置有问题。4.2 分布式训练任务启动用 PyTorch DDP 启动一个 4 机 32 卡的任务命令大致如下torchrun \ --nnodes4 \ --nproc_per_node8 \ --node_rank$RANK \ --master_addrnode0 \ --master_port29500 \ train.py这里$RANK是每台机器的编号从 0 到 3。master_addr是 rank 0 的主机名。代码里要初始化进程组import torch.distributed as dist dist.init_process_group(backendnccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) model torch.nn.parallel.DistributedDataParallel(model, device_ids[local_rank])注意backend一定要用nccl不要用gloo。gloo 是给 CPU 用的GPU 训练用 gloo 会慢到怀疑人生。4.3 监控与性能剖析平台搭起来只是开始能不能持续优化才是关键。我必装的监控有三样GPU 利用率用 DCGM 或 nvidia-smi 采集低于 70% 就说明有优化空间。通信耗时占比用 PyTorch Profiler 抓一段 trace看 All-Reduce 占了多少时间。显存占用曲线判断有没有内存泄漏或碎片化。PyTorch Profiler 的用法with torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CUDA], scheduletorch.profiler.schedule(wait1, warmup1, active3), on_trace_readytorch.profiler.tensorboard_trace_handler(./log) ) as prof: for step, data in enumerate(loader): train_step(data) prof.step()抓出来的 trace 丢进 TensorBoard 就能看到每个算子的耗时。我一般先看有没有“长尾算子”再看通信和计算有没有重叠。理想情况下通信应该被计算掩盖掉如果 trace 里通信和计算是串行的那说明没开overlap。5. 常见问题与排查技巧实录5.1 训练卡死与超时排查分布式训练最烦的就是卡死没有报错就是不走了。我整理了一个排查顺序表现象可能原因排查方法所有卡都停通信死锁检查是否有 rank 没进入集合通信部分卡停数据加载不均看各 rank 的 dataloader 耗时启动就挂端口冲突/主机名解析失败检查 master_port 和 /etc/hosts跑一会挂显存泄漏/OOM监控显存曲线开 memory snapshot一个经典死锁场景某个 rank 因为数据问题提前continue跳过了all_reduce其他 rank 就永远等在那里。解决办法是保证所有 rank 的集合通信调用次数完全一致哪怕数据为空也要参与通信。5.2 性能不达标的常见原因“为什么我的 32 卡还不如别人 16 卡快”这个问题我被问过无数次。常见原因按出现频率排序NCCL 走了错误网卡多网卡机器上NCCL 可能选了管理网而不是数据网。用NCCL_SOCKET_IFNAME指定。批处理太小GPU 没喂饱利用率上不去。适当增大 batch size。数据加载是瓶颈CPU 预处理跟不上 GPU。用num_workers多进程加载或提前把数据转成二进制格式。通信没重叠没开 gradient accumulation 的 overlap或没用bucket_cap_mb调优。提示bucket_cap_mb是 DDP 里一个很关键的参数控制梯度分桶大小。默认 25MB网络差的时候调小能提升重叠效果网络好的时候调大能减少通信次数。5.3 显存优化实战技巧显存不够是永恒的话题。除了加卡还有这些手段梯度检查点Gradient Checkpointing用计算换显存能省 50%-70% 显存代价是慢 20% 左右。混合精度AMPFP16/BF16 训练显存直接减半还能加速。ZeRO 优化DeepSpeed 的 ZeRO 把优化器状态、梯度、参数分片到多卡Stage 3 能省最多显存。及时释放中间变量用del加torch.cuda.empty_cache()但别频繁调用会拖慢速度。我个人的经验是先上 AMP再上梯度检查点最后才考虑 ZeRO。因为 ZeRO 会引入额外通信配置不当反而更慢。6. 架构演进方向与个人实践体会AI Infra 这个领域变化极快但有些趋势是明确的。一是异构化未来不会是 GPU 一家独大NPU、TPU 甚至专用芯片都会进来调度层要能屏蔽差异。二是云原生化训练和推理都在往容器化、Serverless 方向走弹性伸缩会成为标配。三是 Agent 化随着 Agent 架构兴起推理服务不再只是单次请求响应而是多轮、有状态、带工具调用的复杂流程这对基础设施提出了新要求。我自己在实际操作中的体会是不要追求一步到位搭一个“完美平台”而是从最小可用开始遇到问题再演进。我见过太多团队花半年搭平台结果业务需求早变了。先用最土的办法把任务跑起来用监控找到瓶颈再针对性优化这个节奏最稳。最后再分享一个小技巧把每次性能优化的前后数据记下来包括配置、参数、耗时。时间久了你会有自己的“调优直觉”这比任何文档都值钱。AI Infra 没有银弹只有一次次实测积累出来的经验。