ARTICLE DETAIL

资讯详情

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

AI Infra架构实战:分层设计、组件选型与分布式训练推理优化指南

AI Infra架构实战:分层设计、组件选型与分布式训练推理优化指南 1. AI Infra架构到底在解决什么问题先把话说直白一点AI Infra人工智能基础设施架构本质上就是一套让AI模型能从实验室里跑通到在生产环境里稳定、高效、低成本地对外提供服务的工程体系。它跟传统后端架构最大的区别在于AI负载对计算资源的需求是极端不均匀的——训练阶段要吃满GPU、要处理海量数据吞吐、要保证多卡多机之间的通信不拖后腿推理阶段则对延迟极其敏感同时还要兼顾并发量和成本。这两类场景对架构的要求几乎是矛盾的所以AI Infra从来不是一套架构打天下而是分层、分场景、分阶段的一整套组合拳。我接触过不少团队一开始都觉得“不就是把模型部署上去吗”结果真到上线的时候发现单卡跑得好好的模型一上多卡就通信瓶颈离线测试延迟20ms线上并发一上来直接飙到2秒训练任务跑三天崩一次重启还得从头来。这些问题都不是模型本身的问题全是Infra架构没设计好。所以这篇文章我想从实际工程角度把AI Infra架构拆开讲清楚它分几层、每层核心组件是什么、选型时怎么权衡、实操中容易踩哪些坑。不管你是刚接触AI工程的后端开发还是正在搭建团队AI平台的负责人都能从中找到可以直接参考的东西。2. AI Infra架构的分层设计与核心思路2.1 为什么AI Infra必须分层传统Web服务的架构分层大家很熟悉接入层、应用层、数据层。AI Infra的分层逻辑不太一样它更像是一个“资源-调度-运行时-服务”的垂直栈。我习惯把它分成五层来看硬件资源层GPU/TPU/NPU、高速网络InfiniBand/RoCE、分布式存储资源调度层Kubernetes、Slurm、Ray Cluster负责把任务分配到合适的硬件上运行时框架层PyTorch、TensorFlow、JAX以及底层的CUDA、cuDNN、NCCL模型服务层Triton Inference Server、vLLM、TGI、TorchServe应用编排层Agent框架、RAG管线、业务API网关这么分的核心逻辑是每一层的变更频率和技术栈完全不同。硬件层可能两三年才换一代调度层跟着K8s生态走运行时框架跟着社区版本迭代服务层可能每个月都在换新方案。如果不分层任何一层的改动都会牵一发动全身维护成本会爆炸。2.2 训练架构与推理架构的分野这是AI Infra架构设计里最关键的一个决策点。训练和推理对架构的要求差异极大我列个表对比一下维度训练架构推理架构计算模式大批量、长周期小批量、低延迟通信需求极高AllReduce、参数同步较低主要是请求分发资源利用率目标吃满GPU算力降低单次请求成本容错要求支持断点续训支持弹性扩缩容典型调度器Slurm、RayK8s HPA网络要求InfiniBand/RoCE普通万兆即可很多团队初期想用一套K8s集群同时跑训练和推理结果发现训练任务把GPU占满之后推理服务直接饿死。我的建议是如果规模不大GPU少于32卡可以混布但要用MIG或时间片隔离如果规模上来了训练和推理集群物理隔离是更省心的选择。2.3 分布式架构的核心挑战AI Infra绕不开分布式。模型大到单卡放不下的时候就必须做分布式训练。目前主流的并行策略有三种数据并行DP每张卡持有完整模型副本处理不同数据批次梯度通过AllReduce同步。实现简单但模型必须能放进单卡显存。张量并行TP把单层内的矩阵运算切分到多卡适合Transformer这类结构规整的模型。通信量大通常限制在单机8卡内。流水线并行PP按层切分到不同卡上像流水线一样传递中间激活值。通信量小但有气泡bubble问题。实际生产中通常是3D并行组合使用TP放在机内NVLink带宽高PP跨机通信量小DP再叠加在外层。这个组合的选择直接决定了训练效率和硬件成本后面实操部分我会给出具体的参数计算。3. 核心组件选型与实操要点3.1 调度层K8s还是Slurm还是Ray这是被问得最多的问题。我的经验是看团队规模和任务类型SlurmHPC出身适合大规模训练任务排队调度对GPU拓扑感知好但生态偏封闭跟云原生结合麻烦。Kubernetes云原生标准生态丰富适合推理服务和中小规模训练。配合Volcano、Kueue等调度器可以支持gang scheduling。RayPython原生适合ML任务编排Ray Train/Ray Serve一套搞定训练和推理但大规模集群稳定性还需要打磨。我自己的选择逻辑是如果团队已经在用K8s推理服务直接上K8s没悬念训练任务如果超过64卡认真考虑Slurm或者K8sVolcano如果团队以算法工程师为主、运维力量薄弱Ray是上手最快的。注意不管选哪个调度器GPU拓扑感知一定要配好。同一台机器内的8卡通过NVLink互联跨机走网络如果调度器把需要频繁通信的TP组打散到不同机器上性能直接腰斩。3.2 推理服务引擎选型推理引擎这块这两年变化很快我按实际用过的顺序说Triton Inference ServerNVIDIA出品支持多框架多模型动态批处理做得好适合多模型混合部署的场景。缺点是配置偏复杂学习曲线陡。vLLM主打PagedAttention和连续批处理大语言模型推理吞吐量比朴素实现高好几倍。如果你的场景是LLM服务vLLM基本是首选。TGIText Generation InferenceHuggingFace出品跟Transformers生态结合紧密部署简单适合快速验证。TorchServePyTorch官方适合传统CV/NLP模型LLM支持一般。选型时重点看三个指标吞吐量tokens/s或requests/s、首token延迟TTFT、显存占用。这三个指标往往是互相矛盾的需要根据业务场景取舍。比如对话场景首token延迟最重要而离线批量生成场景吞吐量优先。3.3 存储与数据管线AI Infra里最容易被低估的就是存储。训练时数据加载跟不上GPU计算速度GPU利用率可能只有30%。核心要点训练数据放在本地NVMe或高性能分布式文件系统如Lustre、GPFS上不要用对象存储直接读。用WebDataset或LMDB等格式打包小文件避免海量小文件IO。数据预处理管线用DALI或torchdata做GPU加速把CPU解码的瓶颈消掉。我踩过的一个坑早期用NFS挂载训练数据100张卡同时读的时候NFS直接被打爆训练速度还不如单卡。后来换成每台机器本地缓存预取GPU利用率从40%拉到了85%以上。4. 实操过程与核心环节实现4.1 分布式训练环境搭建实录假设我们要在8台机器、每台8卡A100上训练一个70B参数的模型。先算显存账模型参数70B × 2字节FP16 140GB优化器状态Adam70B × 8字节 560GB梯度70B × 2字节 140GB合计约840GB单卡80GB放不下必须并行。采用TP8机内、PP4跨机、DP2的组合总共64卡。每卡分到的参数量140GB / 64 ≈ 2.2GB加上优化器状态和梯度约13GB再算上激活值单卡占用在40-60GB之间可行。具体配置步骤安装NCCL并验证多机通信nccl-tests跑all_reduce确认带宽达到预期A100InfiniBand应该在150GB/s以上。配置共享存储所有节点挂载同一份代码和数据。设置环境变量NCCL_IB_DISABLE0、NCCL_SOCKET_IFNAMEeth0、CUDA_DEVICE_MAX_CONNECTIONS1。用torchrun启动torchrun --nnodes8 --nproc_per_node8 --rdzv_backendc10d --rdzv_endpointnode0:29500 train.py。开启梯度检查点gradient checkpointing省显存代价是重计算增加约30%时间。4.2 推理服务部署与压测以vLLM部署一个13B模型为例python -m vllm.entrypoints.openai.api_server \ --model /models/llama-13b \ --tensor-parallel-size 2 \ --max-model-len 4096 \ --gpu-memory-utilization 0.9 \ --max-num-seqs 256关键参数说明tensor-parallel-size张量并行度13B模型2卡足够70B需要4-8卡。gpu-memory-utilizationKV Cache占用比例0.9意味着90%显存给模型和缓存留10%给其他。max-num-seqs最大并发序列数直接影响吞吐和延迟的平衡。压测用locust或wrk模拟并发请求重点观察P99延迟和吞吐量曲线。我实测下来13B模型在2卡A100上并发64时吞吐约1500 tokens/sP99首token延迟约800ms。并发再往上加延迟会快速上升这时候就需要加副本做水平扩展。4.3 监控与可观测性建设AI Infra的监控比普通服务复杂因为要同时看业务指标和硬件指标GPU指标利用率、显存占用、温度、功耗用DCGM Exporter采集到Prometheus。训练指标loss曲线、梯度范数、学习率、吞吐量samples/s。推理指标QPS、TTFT、TPOT每token输出时间、队列长度。通信指标NCCL带宽、AllReduce耗时。告警规则我建议设这几条GPU利用率持续低于50%说明有瓶颈、显存占用超过95%随时OOM、推理P99延迟超过SLA、训练loss出现NaN。这些规则帮我提前发现过好几次问题比如一次数据管线阻塞导致GPU空转监控告警之后十分钟就定位到了。5. 常见问题与排查技巧实录5.1 训练类问题速查现象可能原因排查方法解决手段loss不下降学习率过大/数据有问题检查数据标签、打印梯度范数调小学习率、清洗数据loss突然变NaN梯度爆炸/混合精度溢出开启梯度裁剪、检查loss scale加grad clip、调loss scaleGPU利用率低数据加载瓶颈看DataLoader耗时加worker、用DALI加速多机训练慢通信瓶颈跑nccl-tests看带宽检查网络配置、调整并行策略训练中断无法恢复checkpoint未保存检查保存逻辑定期保存断点续训5.2 推理类问题速查现象可能原因排查方法解决手段首token延迟高排队严重/模型太大看队列长度和batch组成加副本、用连续批处理吞吐上不去batch太小看GPU利用率开启动态批处理显存OOMKV Cache太大看max-model-len设置调小max-model-len或加卡输出乱码精度问题检查量化配置换回FP16或调量化参数5.3 几个血泪教训教训一不要忽视网络配置。有一次多机训练8台机器跑起来比单机还慢排查了两天才发现是其中一台机器的网卡协商速率掉到了1Gbps。后来我养成了习惯每次上集群先跑一遍ibstat和ethtool确认所有节点网络正常。教训二checkpoint策略要提前设计。训练任务跑了两天崩了结果发现checkpoint每24小时才存一次直接损失一天进度。现在我的标准是小模型每500步存一次大模型每1000步存一次同时保留最近3个checkpoint轮换。教训三推理服务的冷启动要预热。服务刚启动时前几个请求特别慢因为CUDA kernel要编译、显存要分配。解决办法是启动后先跑一批预热请求把常用shape的kernel都编译好再接入流量。教训四版本管理要严格。CUDA、cuDNN、NCCL、PyTorch之间的版本兼容性是个大坑。我现在的做法是用容器镜像把整个环境固化下来镜像tag里带上所有关键版本号比如pytorch2.1-cuda12.1-nccl2.19:v1避免“在我机器上能跑”的问题。6. 架构演进与规模化的几个关键决策6.1 从单机到多机的拐点什么时候该从单机扩展到多机我的经验判断标准是当单机8卡训练一个epoch的时间超过你能够接受的迭代周期时就该考虑多机了。但多机不是免费的通信开销会让线性加速比打折扣。实测下来8机64卡的训练效率通常是单机8卡的6-7倍不是8倍。所以如果单机能勉强跑优先优化单机效率数据管线、混合精度、梯度检查点实在不够再上多机。6.2 推理服务的弹性伸缩推理流量通常有明显波峰波谷固定副本数要么浪费资源要么扛不住峰值。基于K8s HPA的弹性伸缩是标配但AI推理的伸缩有个特殊问题模型加载慢。一个13B模型从磁盘加载到GPU可能要1-2分钟这期间新副本无法服务。解决办法有两个一是用模型缓存把模型预加载到节点本地SSD二是用KEDA基于队列长度做预测性扩容提前把副本拉起来。6.3 成本优化的几个实操手段AI Infra的成本大头是GPU。我总结的几个省钱手段混合精度训练FP16/BF16比FP32省一半显存速度还快精度损失可接受。量化推理INT8量化能把推理成本降低50-70%精度损失通常在1%以内。Spot实例训练任务用Spot实例能省60-70%成本配合checkpoint容错中断了也能恢复。模型蒸馏大模型蒸馏到小模型推理成本直接降一个数量级适合对精度要求不那么极致的场景。请求批处理把多个请求合并成一个batch推理GPU利用率能从30%拉到80%以上。这些手段不是孤立的实际项目中往往是组合使用。比如一个推荐系统训练用混合精度Spot实例推理用INT8量化动态批处理整体成本能压到最初的20%左右。6.4 未来值得关注的方向MoE混合专家架构正在改变AI Infra的很多假设。传统稠密模型推理时所有参数都参与计算MoE只激活部分专家同样参数量下计算量小很多。但MoE对Infra的挑战是专家分布和路由需要更复杂的调度策略。另外Agent架构的兴起让推理服务从“一问一答”变成“多轮工具调用”对状态管理和长上下文缓存提出了新要求。这些方向我在持续跟进有新实践再跟大家分享。我个人在实际操作中的体会是AI Infra没有银弹每个团队的业务场景、团队规模、成本预算都不一样照搬大厂的方案往往水土不服。最靠谱的做法是先把监控做扎实让数据告诉你瓶颈在哪然后针对性地优化。我见过太多团队一上来就追求“先进架构”结果基础的可观测性都没做好出了问题两眼一抹黑。先把单机跑稳再考虑分布式先把推理延迟压下来再考虑吞吐优化。这个顺序不能反。
返回列表