ai-infra-introduction

AI Infra 是让大模型从实验室走向生产环境的技术底座。本文将从零开始帮你建立对这个领域的全景认知。 目录1. 什么是 AI Infra2. 为什么 AI Infra 越来越重要3. AI Infra 技术栈全景4. 硬件层算力基础5. 系统软件层让硬件跑起来6. 训练系统教会模型思考7. 推理系统让模型服务用户8. 性能工程榨干每一分算力9. AI Infra 工程师的能力模型总结1. 什么是 AI Infra打个比方如果大模型是一辆高性能赛车那么 AI Infra 就是赛道、加油站、维修团队和整套后勤保障系统——没有它们赛车哪儿也去不了。用更正式的语言来说AI InfraAI InfrastructureAI 基础设施是指支撑 AI 模型训练、推理和服务的全套技术体系涵盖从底层硬件到上层调度编排的完整技术栈。它要解决的核心问题只有一个如何高效、可靠、低成本地运行 AI 工作负载。1.1 AI Infra 的边界AI Infra 和 AI 算法是两个不同但紧密协作的领域 维度 AI 算法 AI Infra核心关注模型效果准确率、生成质量运行效率吞吐、延迟、成本典型工作设计网络结构、调参、训练策略GPU 编程、分布式训练、推理部署优化目标Loss 更低、效果更好速度更快、显存更省、扩展性更强代表技术Transformer、RLHF、MoECUDA、NCCL、vLLM、FlashAttention提示两者的交叉地带越来越多——比如 FlashAttention 既是算法创新IO 感知的精确注意力也是 Infra 优化CUDA Kernel 实现。理解 AI Infra 会让你对算法的工程落地有更深的认识。2. 为什么 AI Infra 越来越重要2.1 大模型时代的算力饥渴GPT-3 的训练需要约 3640 PF-days即 3640 petaflop/s 的算力运行一天而 GPT-4 的训练算力需求估计是 GPT-3 的数十倍以上。Llama 3 405B 的训练使用了 16384 张 H100 GPU持续训练数周。这意味着单张 GPU 的算力远远不够必须组建大规模集群多卡之间的通信开销可能比计算本身还耗时一次训练动辄花费数百万美元效率提升 10% 就是巨大的成本节省2.2 推理成本成为瓶颈训练是一次性投入但推理是持续性支出。当一个大模型被数百万用户同时使用时每个请求都需要 GPU 资源来做前向计算用户期望的响应延迟在毫秒到秒级推理成本可能占到 AI 服务总成本的 80% 以上关键点推理优化更高吞吐、更低延迟、更少显存直接决定了 AI 产品的商业可行性。2.3 人才缺口巨大目前行业对 AI Infra 工程师的需求远超供给。原因很简单这个领域需要同时理解硬件架构、系统编程、分布式系统和 AI 算法跨学科的知识壁垒很高但一旦掌握价值也非常大。3. AI Infra 技术栈全景AI Infra 可以按层次划分为四层从底向上依次是每一层都有独特的技术挑战和优化空间。接下来我们逐层展开。4. 硬件层算力基础4.1 GPUAI 计算的核心引擎为什么深度学习用 GPU 而不是 CPU因为深度学习的核心运算是矩阵乘法而矩阵乘法天然适合并行——一张 GPU 拥有数千个计算核心可以同时处理大量元素的运算。以 NVIDIA 的 GPU 为例核心架构演进如下架构年份代表产品关键特性Volta2017V100首次引入 Tensor CoreAmpere2020A100支持 TF32/BF16MIG 多实例Hopper2022H100FP8 支持Transformer EngineBlackwell2024B200第二代 Transformer EngineFP4核心概念Tensor Core是 NVIDIA GPU 中专门为矩阵运算设计的硬件单元。以 H100 为例其 FP16 Tensor Core 算力高达 989 TFLOPS远超普通 CUDA Core 的算力。4.2 互联技术多卡协同的关键单卡算力有上限多卡协同训练就需要高速互联NVLinkGPU 之间的直连通道。H100 的 NVLink 4.0 提供 900 GB/s 的双向带宽远超 PCIe 5.0 的 128 GB/sNVSwitch连接同一节点内多张 GPU 的交换芯片实现全对全互联InfiniBand跨节点的高速网络。NVIDIA NDR InfiniBand 提供 400 Gbps 带宽且延迟极低约 1 微秒级提示通信带宽往往是大规模训练的第一个瓶颈。当你发现训练扩展效率不高时首先应该排查的就是通信。4.3 显存HBM 的重要性GPU 的显存带宽直接影响模型的运行效率。现代 AI GPU 普遍采用 HBMHigh Bandwidth MemoryH100 SXM 配备 80GB HBM3带宽 3.35 TB/sA100 配备 80GB HBM2e带宽 2.0 TB/s很多 AI 工作负载尤其是推理是内存带宽瓶颈而非计算瓶颈这意味着显存带宽的提升有时比算力提升更重要。5. 系统软件层让硬件跑起来5.1 CUDAGPU 编程的基础语言CUDA 是 NVIDIA 推出的 GPU 并行计算平台和编程模型。可以把它理解为GPU 上的编程语言——你通过 CUDA 告诉 GPU 要做什么运算、怎么调度线程。一个简单的 CUDA 向量加法示例__global__voidvectorAdd(float*a,float*b,float*c,intn){intidxblockIdx.x*blockDim.xthreadIdx.x;if(idxn){c[idx]a[idx]b[idx];}}// 启动 Kernel256 个线程为一个 Blockintblocks(n255)/256;vectorAddblocks,256(d_a,d_b,d_c,n);CUDA 编程中最核心的概念是线程层级Thread最小的执行单元Block一组线程共享 Shared MemoryGrid所有 Block 的集合5.2 通信库NCCL当多张 GPU 需要协同计算时需要高效的通信原语。NCCLNVIDIA Collective Communications Library提供了 AllReduce、AllGather、ReduceScatter 等集合通信操作是分布式训练的通信基础。5.3 算子库高性能的基础运算不需要从头写 CUDANVIDIA 提供了丰富的算子库cuBLAS矩阵乘法GEMMcuDNN卷积、池化、归一化等深度学习算子cuFFT快速傅里叶变换6. 训练系统教会模型思考6.1 为什么需要分布式训练以 Llama 3 405B 为例模型参数量 4050 亿仅存储参数就需要约 810 GBFP16远超单张 GPU 的显存容量。更不用说训练时还需要保存梯度、优化器状态等中间数据总显存需求是参数量的 10-20 倍。所以大模型训练必须使用多卡甚至多节点。6.2 并行策略分布式训练有三种基本并行方式就像一个工厂里的三种分工方法数据并行Data Parallelism每张 GPU 持有完整的模型副本但各自处理不同的数据。训练完成后同步梯度。就像同一道菜谱多个厨师各做一份最后汇总味道反馈。代表技术PyTorch DDP、FSDP分片版数据并行适用场景模型能放进单卡显存模型并行Tensor Parallelism把模型的每一层切开分散到多张 GPU 上。就像一张大桌子太重了拆成几块分别搬。代表技术Megatron-LM适用场景单层参数量大如大模型的 Attention 和 FFN 层流水线并行Pipeline Parallelism把模型按层分段不同段放在不同 GPU 上数据像流水线一样依次通过各段。代表技术GPipe、PipeDream适用场景模型层数多适合按深度切分⚠️注意实际的大规模训练通常结合三种并行称为 3D 并行需要根据模型大小、集群规模和网络拓扑仔细调参。6.3 关键框架框架开发者核心特性PyTorch DDPMeta原生数据并行简单易用PyTorch FSDPMeta参数分片的数据并行省显存DeepSpeedMicrosoftZeRO 优化器灵活的并行组合Megatron-LMNVIDIA高性能的张量并行和流水线并行7. 推理系统让模型服务用户7.1 推理的核心挑战训练时追求的是吞吐量单位时间处理更多样本推理时既要高吞吐又要低延迟。大模型推理有一个独特的特点它是自回归的——每次只生成一个 token且每个 token 的生成都依赖之前的所有 token。这导致推理过程天然是串行的优化难度很大。7.2 关键优化技术KV Cache在自回归生成中每个 token 的 Attention 计算需要用到之前所有 token 的 Key 和 Value。把它们缓存起来避免重复计算。代价是显存占用会随序列长度线性增长。PagedAttention传统的 KV Cache 为每个请求预分配连续的显存块就像餐厅给每桌客人预留固定大小的区域不管实际用多少。PagedAttention 借鉴了操作系统虚拟内存分页的思想将 KV Cache 分成固定大小的 Page按需分配把显存利用率从约 60% 提升到接近 100%。量化用更低精度的数据类型FP8、INT8、INT4表示模型参数减少显存占用和计算量。例如 FP16 → INT8 量化可以将模型体积减半推理速度提升 1.5-2 倍精度损失通常在可接受范围内。7.3 主流推理引擎引擎核心技术特点vLLMPagedAttention开源标杆社区活跃SGLangRadixAttention前缀共享优化适合多轮对话TensorRT-LLMNVIDIA 自研 Kernel极致性能深度绑定 NVIDIA 硬件8. 性能工程榨干每一分算力8.1 为什么需要性能分析写出能跑的 CUDA Kernel 只是第一步写出跑得快的 Kernel 才是核心价值。性能分析Profiling帮你找到瓶颈在哪里——是计算不够快还是数据搬运太慢还是线程在等待8.2 Roofline 模型Roofline 模型是性能分析的核心思维框架。它把算子分为两类计算瓶颈Compute-bound算力不够用如大矩阵乘法访存瓶颈Memory-bound数据搬运速度跟不上计算速度如逐元素操作判断依据是算术强度Arithmetic Intensity每搬运 1 字节数据能做多少次浮点运算。Arithmetic Intensity FLOPs Bytes Accessed \text{Arithmetic Intensity} \frac{\text{FLOPs}}{\text{Bytes Accessed}}Arithmetic IntensityBytes AccessedFLOPs​当算术强度低于硬件的 “拐点” 时算子受限于内存带宽反之受限于计算能力。8.3 核心 Profiling 工具Nsight Systems系统级分析看 CPU/GPU 的时间线、Kernel 调度、通信开销——帮你找到哪里慢Nsight ComputeKernel 级分析看单个 Kernel 内部的计算效率、内存吞吐、Warp 占用率——帮你找到为什么慢提示先用 Nsight Systems 定位热点 Kernel再用 Nsight Compute 深入分析该 Kernel 的瓶颈是最高效的性能优化工作流。9. AI Infra 工程师的能力模型如果你想进入这个领域需要构建以下几层能力9.1 基础层C/C 编程CUDA 编程基于 C/C需要熟练掌握指针、内存管理计算机体系结构理解 CPU/GPU 的存储层级寄存器 → L1 → L2 → HBM → 主存Linux 系统大部分 AI Infra 工作在 Linux 环境下9.2 核心层CUDA 编程Thread/Block/Grid 模型、Shared Memory、Warp 调度分布式系统基础通信原语AllReduce、AllGather、一致性、容错PyTorch 内部机制Autograd、Module、分布式训练 API9.3 进阶层算子优化Tiling 策略、Memory Coalescing、Bank Conflict 消除并行策略设计根据模型和硬件特征选择最优的并行组合推理引擎原理调度器设计、KV Cache 管理、量化算法关键点AI Infra 是一个既需要广度又需要深度的领域。广度让你理解系统全貌深度让你解决实际的性能问题。建议先建立全景认知也就是这篇文章的目标再选择一个方向深入。 总结AI Infra是支撑大模型训练、推理和服务的完整技术体系从 GPU 硬件到调度编排大模型时代AI Infra 的重要性空前提升——训练需要大规模集群推理需要极致优化技术栈分为五层硬件层 → 系统软件层 → 训练系统层 → 推理系统层 → 应用与调度层核心技能包括 CUDA 编程、分布式训练、推理优化和性能分析这是一个人才缺口巨大、技术壁垒高但回报丰厚的领域 自我检验清单能用自己的话解释什么是 AI Infra 以及它和 AI 算法的区别能画出 AI Infra 技术栈的分层架构硬件层到应用层能说出 Tensor Core、NVLink、HBM 分别解决什么问题能区分数据并行、模型并行和流水线并行的适用场景能解释 KV Cache 和 PagedAttention 的基本原理能说出 Roofline 模型中计算瓶颈和访存瓶颈的判断依据能描述 Nsight Systems 和 Nsight Compute 的定位差异 参考资料NVIDIA H100 Tensor Core GPU ArchitectureNVIDIA CUDA C Programming GuidePyTorch Distributed OverviewvLLM: Easy, Fast, and Cheap LLM Serving with PagedAttentionEfficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMDeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion ParametersRoofline: An Insightful Visual Performance Model