AI System 体系
截止 2026 年,LLM、VLM 等模型的训练流程仍可以粗略划分为 Pretraining 和 Post-Training 两个阶段。后训练阶段已经衍生出较多技术路线,例如 SFT、DPO、PPO、GRPO、RLAIF 等。若从目标上理解:
Pretraining的目标是获得通用建模能力,依赖大规模语料与自监督学习。SFT的目标是让模型更贴近具体任务和指令格式。PPO、GRPO、RLAIF等强化学习或偏好优化方法,则更多是让模型输出更符合人类偏好、规则约束或奖励模型。
大型模型的训练与推理都高度受限于显存容量、显存带宽、计算吞吐和互连带宽。以一个 7B 模型为例,若仅考虑 FP16/BF16 权重,其静态参数大约需要 14GB 显存。推理时,除权重外,还要考虑 KV Cache。对 Transformer 来说,KV Cache 的大小近似与层数、KV 头数、head_dim、序列长度和精度成正比,因此在长上下文场景下,KV Cache 往往会成为主要瓶颈之一。
训练时的内存压力则更大。除了权重本身,还需要考虑梯度、优化器状态、可能存在的 FP32 master weight、激活值、临时 buffer 和通信 buffer。若以 AdamW + 混合精度训练 做工程估算,每个参数常可粗略按约 16 Byte/param 计算参数相关开销:
- 权重:
2B - 梯度:
2B FP32 master weight:4B- 一阶和二阶动量:
8B
这样一来,7B 模型仅参数相关状态就可能接近 112GB,还未包含 activation、临时 workspace 和通信 buffer。因此对 7B 及以上模型做全量训练时,多卡几乎是起步要求,而模型规模和序列长度继续上升后,互连带宽与通信效率会进一步成为瓶颈。
通信与互连
AI System 里需要区分几个不同层次的概念:
PCIe、NVLink、NVSwitch:主要描述单机内 GPU 之间的互连方式与拓扑。InfiniBand、RoCE:主要描述多机之间的高速网络互连。NCCL:是 NVIDIA 生态里常用的 GPU 集合通信库,用于实现AllReduce、AllGather、ReduceScatter等 collective。
这几个概念经常同时出现,但不属于同一层。NCCL 更像软件通信层,NVLink/NVSwitch/PCIe 更像机内硬件互连层,IB/RoCE 更像机间网络层。
在单机多卡场景下,消费级显卡通常缺少数据中心卡那样的高带宽互连能力,跨卡通信拓扑、驱动支持和稳定性也通常弱于 A100/H100 这类计算卡。对于通信密集型任务,例如 TP、部分 PP 切分或大规模推理中的 cache 同步,互连带宽会直接影响扩展效率。到了多机训练阶段,还常常需要依赖 IB/RoCE + GPUDirect RDMA 来降低跨机通信开销。
并行训练的基本动机
当单卡无法容纳模型或无法提供足够吞吐时,就需要在多卡甚至多机环境下进行切分。常见切分方式包括:
DP:数据并行。不同设备处理不同样本,每张卡通常保留完整模型。TP:张量并行。把一个算子的权重或计算切到多张卡上,通信压力较大,但可以突破单卡计算和显存上限。PP:流水线并行。按模型层划分 stage,不同 stage 放在不同设备上,核心问题是 bubble、调度与 stage 间通信。EP:专家并行,常用于MoE,把不同专家分散到不同设备上,再通过路由机制分发 token。
这些策略在真实系统里往往是混合使用的。训练系统常见的支持方式包括:
PyTorch DDPPyTorch FSDPDeepSpeed ZeROMegatron-LM
其中 ZeRO 更适合理解为一种分片和内存优化策略,而不是单独的推理框架;vLLM 则是推理框架,不应与 Megatron-LM、DeepSpeed 放在同一类中讨论。
训练与推理的系统差异
训练和推理虽然都运行在 GPU 集群上,但系统关注点并不一样:
- 训练更关注吞吐、扩展效率、参数同步、checkpoint、失败恢复、
MFU - 推理更关注首 token 时延、单 token 时延、吞吐、显存占用、请求调度和稳定性
训练时经常讨论 DDP/FSDP/ZeRO、TP/PP/EP、梯度同步、激活重计算等;推理时则更常讨论 Prefill/Decode、KV Cache、PagedAttention、Continuous Batching、量化和服务调度。
芯片架构
NVIDIA GPU
如果从 AI Infra 的角度理解 GPU,可以把关注点放在以下几个部分:
SM:执行计算的核心单元Tensor Core:矩阵乘加等张量计算的关键硬件HBM:高带宽显存L2 Cache / Shared Memory / Register:不同层级的片上或片外存储
模型训练与推理的很多优化,本质上都可以落回到两个问题:
- 计算是否足够密集,能否把算力吃满
- 数据搬运是否高效,是否被显存带宽或访存方式限制
因此很多算子优化,例如 FlashAttention,从根本上并不是“数学变了”,而是通过更好的分块和访存策略减少了中间结果读写,降低了 IO 开销。
DGX 系列
DGX 系列可以看作 NVIDIA 针对 AI 训练和推理场景提供的高密度整机方案。它的重要意义不只是“卡更多”,而在于:
- 提供了更适合多卡通信的拓扑
- 软硬件栈配套更完整
- 在规模化训练和推理部署中更容易获得稳定表现
从 AI System 角度看,整机设计会直接影响 TP/PP 的扩展效率、通信模式和容错方式。
通信网络
训练和推理规模上去之后,网络常常和显存一起成为瓶颈。一个更系统的拆法是:
- 机内通信:
PCIe、NVLink、NVSwitch - 机间通信:
InfiniBand、RoCE - 通信原语:
AllReduce、AllGather、ReduceScatter、Broadcast - 通信库:
NCCL
不同并行策略对通信模式的依赖不同:
DP更依赖梯度同步,常涉及AllReduceTP常涉及更频繁的张量切分和聚合,通信压力很大PP更依赖 stage 间激活和梯度传递EP则依赖 token 路由与专家间分发
深度学习框架与编译
PyTorch 生态
PyTorch 既是深度学习框架,也是很多 AI System 工作的入口。若从系统角度看,可以把它理解为:
- 前端:张量、算子、自动求导、模块抽象
- 中间层:图表示、调度、编译优化
- 后端:调用
CUDA、cuDNN、NCCL等底层库执行
当模型规模上升后,PyTorch 的工作就不再只是“写网络结构”,而是需要结合 DDP/FSDP、编译优化、内存优化和分布式运行时一起看。
CUDA 生态
CUDA 是 NVIDIA GPU 编程的基础。常见相关组件包括:
CUDA RuntimecuBLAScuDNNNCCLTensorRT
从学习路径上看,理解 CUDA 不一定一开始就要手写复杂 kernel,但至少需要建立这些概念:
- grid / block / thread 的执行模型
- global memory、shared memory、register 的层次
- 为什么访存模式会决定性能
- 为什么算子融合和 kernel fusion 会有效
CUDA 算子开发
CUDA 算子开发里,最常见的优化抓手通常集中在访存路径、线程组织、同步方式和 kernel 融合上;很多“性能差距巨大”的根源,往往并不是算得不够快,而是数据搬得不够对。
访存优化
在 GPU 的世界里,计算往往不是瓶颈,“把数据按时送到计算单元” 才是最难的。从冯诺伊曼机器的角度来看,GPU 的存储包含 Global Memory 和 Shared Memory,二者具有截然不同的访问效率,在算子层面的开发时需要注意一些铁律。
Global Memory(全局内存 / HBM)极大(几十 GB),但延迟极高(几百个时钟周期),对它的优化核心只有一个词:吞吐量(Throughput)。
合并访存(Coalesced Memory Access)是最核心的方案。GPU 是以 Warp(32 个线程)为单位执行的,当一个 Warp 去读写 Global Memory 时,硬件底层的内存控制器是以 32 字节、64 字节或 128 字节的块(Transaction)为单位拉取数据的。这就要求 32 个线程请求的内存地址必须是连续且对齐的,如果线程 0 读地址 0,线程 1 读地址 1,线程 31 读地址 31,硬件只需发一次车就能把数据全拉回来;如果地址是随机跳跃的(Stride 访问或间接索引),硬件可能需要发起 32 次独立的内存事务,带宽利用率直接暴跌 32 倍。
向量化访存(Vectorized Memory Access)的技巧在 vLLM、TensorRT-LLM 等推理框架的底层代码中极度常见。它的基本思想是不要让一个线程一次只读一个 float(4 字节),而是强转指针,一次性读取 float4(16 字节)。这能极大地减少内存请求指令的数量,更容易打满硬件的物理带宽上限。但是,数据结构的起始地址必须是其数据类型大小的整数倍;如果不对齐,强行做向量化读取会导致程序直接 Crash(段错误)或者退化为多次低效读取。
Shared Memory 极小(每个 SM 通常 100KB 左右),但就在计算单元旁边、延迟极低(几个时钟周期)。对它的优化核心也只有一个词:并发冲突(Conflicts)。使用它的核心是避免 Bank 冲突(Bank Conflicts)。Shared Memory 在物理上被划分成了 32 个独立的 Bank(存储体),每个 Bank 宽度通常是 4 字节,这 32 个 Bank 可以同时响应 32 个线程的读写请求。如果同一个 Warp 中的多个线程,试图访问同一个 Bank 中的不同物理地址,硬件就无法并发了,必须把这些请求串行化(Serialize),排队一个一个读,这就叫 Bank 冲突;然而,如果多个线程访问的是同一个 Bank 里的完全相同的地址,硬件会触发多播(Multicast)或广播(Broadcast),不仅不会冲突,反而一次就读完分发给所有人。
充分利用它的技巧之一是 Memory Padding(内存填充)。当我们在 Shared Memory 中存一个二维矩阵(比如 32 x 32)且需要按列读取时,极易发生严重的 Bank 冲突。经典的解法是在声明数组时,人为给每一行多加一个无用的元素,比如变成 32 x 33。这会导致原本落在同一个 Bank 上的列元素,在物理地址上发生错位,完美避开冲突,牺牲极其微小的内存空间,换来巨大的性能飞跃。
同步屏障 __syncthreads() 的细节在于,Shared Memory 是同一个 Block(线程块)内所有线程共享的,当一部分线程往里写数据,另一部分准备读时,必须在中间加上 __syncthreads()。否则,由于 Warp 执行进度不一致,读操作会拿到脏数据(Race Condition)。
NCCL
在 AI Infra 和分布式训练中,NCCL(NVIDIA Collective Communications Library)就是整个 GPU 集群的高速公路路网。
NCCL 提供了标准的集合通信(Collective Communication)和点对点通信能力。主要可以分为以下 6 大原语:
- Broadcast(广播)。一张卡(Root)把自己的数据原封不动地复制,发送给集群里的所有其他卡。在 DDP 训练初始化时,需要通过 Broadcast 将模型权重从主卡同步给所有卡。
- Reduce (规约)。所有卡把数据发给一张卡(Root),并在 Root 卡上进行数学运算(通常是求和 Sum、求最大值 Max 等)。较少单独使用,通常作为复杂操作的一环。
- All-Reduce (全规约)。所有卡的数据进行相加(或相乘),最后所有卡都得到了一份完整的最终计算结果。这是经典 DDP 和 TP(张量并行)中最核心的操作,每次反向传播结束算完梯度后,必须用它把所有卡的梯度加总求平均。
- All-Gather (全收集)。每张卡拿出一块切片数据,互相交换拼装,最后所有卡都得到了一份完整的全局数据。这是 ZeRO-3/FSDP 的核心,计算前,各卡通过 All-Gather 借来其他卡的参数切片,拼出完整模型。
- Reduce-Scatter (规约并打散)。所有卡的数据先进行相加(Reduce),但加完的结果不全给一个人,而是切成 N 块,每张卡只拿走属于自己的那一小块。这是 ZeRO-3 / FSDP 中同步梯度的核心,算完梯度后,不再需要每张卡都保留全量梯度(太占显存),而是只保留自己负责更新的那部分梯度。
- Send / Recv (点对点通信)。指定卡 A 发送数据,指定卡 B 接收数据。在 PP(流水线并行)中,跨机器、跨 GPU 传递中间层的 Hidden States 必须用到它。
All-Reduce 解析
在物理拓扑和 NCCL 底层,为了最大化榨干网络带宽,All-Reduce 通常被拆解为 2 个大阶段。
以业界最经典的 Ring All-Reduce(环形规约算法)为例。假设我们有 $N$ 张 GPU,要同步的数据总量为 $S$:
- 阶段一:Reduce-Scatter(规约并打散)
- 过程:NCCL 首先将庞大的梯度数据 $S$ 逻辑上切分为 $N$ 个数据块。所有的 GPU 连成一个环。在每个时钟周期,GPU 拿自己的某一个数据块发给右边的邻居,并接收左边邻居发来的数据块进行累加。
- 通信步数:需要绕环传递 $N-1$ 步。
- 阶段结果:此时,整个大集群的梯度确实被全部加总了,但没有任何一张卡拥有完整的全量结果。每张卡手里只有 $\frac{1}{N}$ 的、已经算完的最终梯度。
- 阶段二:All-Gather(全收集)
- 过程:既然每张卡现在手里都有 $\frac{1}{N}$ 的“真经”,大家只需要再绕着环互相抄一遍作业。GPU 再次把手里算好的最终结果传给右边邻居。
- 通信步数:同样需要绕环传递 $N-1$ 步。
- 阶段结果:所有卡都凑齐了 $N$ 个切片,每张卡都拥有了完整的全量梯度 $S$。
对于一次 All-Reduce 参数更新:
- 逻辑上:经历了 2 次复合通信操作(1 次
Reduce-Scatter+ 1 次All-Gather)。而ZeRO和FSDP之所以能够省显存,关键并不是简单“砍掉第二步”,而是把参数、梯度和优化器状态按shard分片存储,并尽量把All-Gather推迟到真正需要该层参数的时刻。 - 物理步骤上:经历了 $2(N-1)$ 步单步通信传递。
- 数据传输量:每张卡发送和接收的数据量都是 $2 \times \frac{N-1}{N} \times S$。当卡数 $N$ 足够大时,每张卡的通信量恒定在 $2S$ 左右。这意味着无论集群规模是一百张卡还是一万张卡,只要带宽足够,单张卡的通信负载不会随着节点数暴增而无限膨胀。
除了 Ring All-Reduce,在成千上万张卡的超大规模集群里(比如跨多台交换机时),NCCL 会抛弃环形算法,转而使用延迟更低的 Tree All-Reduce(双二叉树规约)。
All-Gather 解析
在 ZeRO-3 和 FSDP 中,没有任何一张 GPU 手里有完整的模型参数,它们只有“参数碎片”。为了进行正常的数学运算(前向和反向传播),GPU 必须在计算的那一瞬间,用 All-Gather 把碎片临时“拼凑”成完整的权重。
为了省显存,ZeRO-3/FSDP 把模型的 参数 (Parameters)、梯度 (Gradients) 和 优化器状态 (Optimizer States) 全部切成了 $N$ 份(假设有 $N$ 张 GPU)。这意味着,对于模型中的某一个 Linear 层(假设大小为 8GB),单张 GPU 显存里只存了 $\frac{1}{N}$ 的权重(比如 8 张卡,每张卡只存 1GB)。
现在,GPU 0 拿到了分配给自己的那一小批训练数据(Mini-batch),准备开始前向传播。
- 遇到的问题:当 GPU 0 准备计算第 1 层时,它发现自己手里只有第 1 层 12.5% 的权重矩阵,根本没法和输入数据做完整的矩阵乘法。
- All-Gather 登场:GPU 0 立刻向网卡发起一次 All-Gather 通信,大喊一声:“兄弟们,把你们手里第 1 层的碎片都发给我!” 同时,其他 7 张卡也在做同样的事情。
- 瞬间完整:通信完成后,GPU 0 的显存里短暂地拥有了完整的 8GB 第 1 层权重。它迅速完成矩阵乘法,算出激活用作下一层的输入。
- 阅后即焚 (Free):这是最关键的一步!算完第 1 层后,GPU 0 立刻把刚刚“借”来的那 87.5% 的权重从显存里删除(释放掉),只留下自己的那 12.5%。然后带着空出来的显存,继续去用 All-Gather 拼凑第 2 层的权重。
当损失函数算完,开始反向传播求梯度时,流程几乎一模一样:
- 为了算出第 1 层权重的梯度,GPU 0 再次需要第 1 层的完整权重矩阵。
- 它只能再做一次 All-Gather,重新把兄弟们的碎片借过来拼好。
- 算完梯度后,再次立刻释放(Free)借来的权重。
(注:算完梯度后,大家还会用我们上一轮聊过的 Reduce-Scatter 通信,把全局梯度规约求和,然后再次打散,每张卡只保留自己负责的那 $\frac{1}{N}$ 梯度准备更新优化器。)
这就是为什么 ZeRO-3/FSDP 对网络带宽(尤其是机内的 NVLink)要求极高,一层网络在前向和反向计算时,模型权重要在 GPU 之间反复横跳两次。由于这种频繁的 All-Gather 会导致 GPU 经常停下来等网卡传数据,现代的 Infra 框架(如 PyTorch FSDP)都引入了一种极其优雅的加速技术——预取(Prefetching)与计算重叠(Overlap)。
为了打破这个僵局,底层工程师引入了 CUDA 多流机制(Multi-Stream)和预取(Prefetching)。通过同时创建“计算流(Compute Stream)”和一条独立的“通信流(Communication Stream)”,GPU 可以在计算的同时预先拉取下一层的权重,并使用 CUDA Event 作为同步锁避免数据错误。
Triton 和 CUTLASS
Triton 更偏向于以张量为抽象来编写高性能 GPU kernel,适合快速做算子级原型和优化实验;CUTLASS 则更贴近 NVIDIA 提供的高性能矩阵运算模板库。二者都经常出现在 AI Infra 的算子优化语境里,但定位不完全相同:
Triton:更适合研究和快速实现自定义 kernelCUTLASS:更贴近高度优化的 GEMM/卷积模板化实现
分布式并行基础
并行计算
传统高性能计算里,OpenMPI 等工具更常用于 MPI 风格的多进程并行。它和大模型训练不是完全同一套生态,但很多并行与通信思想是一致的,例如:
- 进程间通信
- collective
- 同步与屏障
- 拓扑和数据分发
分布式计算
Ray 更适合 Python 生态下的分布式任务编排与服务组织。它的优势在于:
- 对任务和 Actor 抽象友好
- 易于快速构建实验系统和服务
- 在推理服务、离线处理、分布式实验中都很常见
需要注意的是,Ray worker 更准确地说是工作进程,而不是“等于一台机器”。
资源调度框架
在大模型场景下,常见的资源调度和编排方式包括:
Kubernetes:更常用于服务部署、资源编排、弹性扩缩容和平台化管理Slurm:在 HPC 和大规模训练场景里很常见,强调稳定、配额、队列和作业管理Ray:更像分布式运行时和任务编排层,不完全等同于集群调度器
从公开资料和工业实践来看,推理服务常与 K8s 结合,而训练任务则常见于 Slurm 或专门的训练平台上。
LLM 并行策略
TP、PP、DP 和 EP 是大模型分布式计算的四大基础维度(策略)。它们定义了“从哪个空间维度去切分计算任务”。
DP
DP 的核心思想:模型不动,切分数据。每张卡都有相同的模型,但吃进不同的数据批次(Batch)。
实现技术:
- PyTorch DDP:最基础的实现,每张卡完整复制模型、梯度和优化器。
- DeepSpeed ZeRO (1/2/3):由微软开发,是 DP 的显存优化版。它在逻辑上还是数据并行,但在物理显存上把模型状态“切碎”存放到不同卡上。
- PyTorch FSDP:Meta 贡献的 PyTorch 原生 ZeRO-3 实现。
TP
核心思想:切分矩阵,也就是“层内切分”。把模型中单层(比如一个巨大的 Linear 层或 Attention 矩阵)按行或按列切开,分给不同的 GPU 算,算完再拼起来。通信极其频繁,每次矩阵乘法都要通信,因此通常尽量限制在同一个节点(同一台物理机)内,严重依赖 NVLink 等高速互联带来的带宽。
实现技术:
- Megatron-LM:NVIDIA 开源的框架,是 TP 技术的绝对霸主和开创者。几乎所有大模型的 TP 底层逻辑都源自 Megatron。
- PyTorch DTensor:PyTorch 最近推出的原生 Device Mesh 分布式张量,试图降低 TP 的开发门槛。
- vLLM / TensorRT-LLM:推理框架中用于实现 TP 加速生成的底层支撑。
PP
核心思想:按层切分模型,也就是“层间切分”。比如一个 80 层的模型,GPU 0 负责 1-10 层,GPU 1 负责 11-20 层,数据像流水线一样在 GPU 之间接力传递。其通信量较小(只传递隐式状态),但容易产生“气泡”(Bubble,即某些 GPU 在等前面 GPU 算完时的闲置状态)。适合跨节点(跨机器)切分模型。
实现技术:
- GPipe / PipeDream:早期的学术界探索(引入了微批次 Micro-batch 概念)。
- Megatron-LM (1F1B 策略):NVIDIA 优化的交替前向-反向流水线,极大减少了显存占用和气泡。
- DeepSpeed Pipeline:DeepSpeed 框架内置的流水线引擎。
- PyTorch PiPPy:PyTorch 官方后来补充的流水线并行库。
EP
核心思想:专门针对 MoE(混合专家)架构模型(如 DeepSeek-V3/R1、Mixtral)。模型中的不同“专家”神经网络被放置在不同的 GPU 上。Token (数据) 会通过路由器(Router)被发送到对应的 GPU 去找对应的专家处理。它极度依赖网络通信(All-to-All 通信),因为 Token 需要在 GPU 之间飞来飞去寻找专家。
实现技术:
- DeepSpeed-MoE:目前应用最广泛的 MoE 并行训练框架。
- Megatron-Core (MoE):NVIDIA 针对 Hopper 架构优化的最新 EP 实现。
- Tutel:微软专门为 MoE 开发的高度优化的底层计算与通信库。
基础设施拓扑
显卡、机器之间的拓扑结构设计核心在于通信,在典型的现代 GPU 集群(如 NVIDIA DGX A100/H100 SuperPOD)中,通信带宽分为两个截然不同的等级:
- 机内通信:一台物理机内通常有 8 张 GPU,它们通过 NVLink / NVSwitch 连接。带宽极高(例如 H100 可达 900 GB/s),延迟极低。
- 机间通信:物理机与物理机之间,通过 InfiniBand (IB) 或 RoCE 交换机连接。带宽相对较低(通常是 400 Gbps,约合 50 GB/s),延迟较高。
对于训练场景,假设我们要训练一个千亿参数的大模型(如 Llama-3 70B),手头有 32 台 8 卡机器(共 256 张 GPU)。Infra 工程师会这样进行空间排布:
- 第一层切分,TP:一台机器内部的 8 张卡 (TP=8)。TP 在每一次矩阵乘法的前后都需要进行 All-Reduce 通信,频率极高,数据量大。如果跨机做 TP,网络带宽会被瞬间打爆。因此,TP 的大小通常不能超过单机卡数(最多 8)。
- 第二层切分,PP:机器与机器之间,比如前 4 台机器算前 20 层,后 4 台机器算后 20 层。PP 只需要在网络层的边界(Layer Boundaries)传递中间结果(Hidden States),不传递庞大的梯度和参数。它的通信量非常小,非常适合走相对较慢的 IB 网络。
- 第三层切分,DP / ZeRO:横跨所有不需要做 TP 和 PP 的机器。ZeRO 的 All-Gather 通信虽然也不小,但 Infra 框架(如 DeepSpeed/FSDP)可以做计算与通信的重叠(Overlap)。也就是说,GPU 一边在算当前层,网卡一边在后台拉取下一层的参数,从而掩盖了跨机通信的延迟。
推理(Inference)和训练的逻辑完全不同。训练追求的是高吞吐(Throughput),而推理往往更看重低延迟(Latency / TTFT)。为了让单次生成的延迟降到最低,推理时会极其激进地使用 TP。比如 70B 模型,通常直接用一台机器的 8 张卡做 TP=8。
Prefill-Decode 分离是目前大规模部署(如 Kimi、DeepSeek 线上服务)最火的架构:
- Prefill 节点(处理输入 Prompt):用户输入的长文本需要极大的并行计算力。我们会分配一批 GPU 专门负责“阅读理解”(Prefill),通常使用 TP=8,并配置极高的显存带宽。
- Decode 节点(生成逐个回答):逐字生成的阶段属于内存受限(Memory-bound)。我们会分配另一批不同的 GPU(甚至是低端一点的 GPU)专门做生成。
- 拓扑互联:Prefill 算完后,把庞大的 KV Cache 通过高速 RDMA 网络“飞传”(Transfer)给 Decode 节点。这要求集群的调度系统异常强大。
训练系统
正式训练
训练系统最核心的指标通常包括:
- 吞吐,例如
tokens/s MFU- 扩展效率
- 内存占用
- checkpoint 开销
理解训练系统时,不能只停留在“模型能不能跑起来”,还要关注:
- 多卡扩展后吞吐是否线性增长
- 通信是否成为瓶颈
- activation 是否占了主要内存
- 是否需要
checkpointing、recompute、ZeRO/FSDP
分布式训练中的数据并行
DDP、DeepSpeed ZeRO 和 FSDP 是分布式训练中的数据并行技术,其中,ZeRO 和 FSDP 可以视作 DDP 为了适应 LLM 训练而进化出的高级形态,与 DDP 相比提供了更灵活的数据切分方案。
PyTorch DDP
DDP 是最基础的并行训练技术,仅支持 DP(Data Parallel,数据并行),他在每张 GPU 上都完整复制一份一模一样的模型参数、梯度和优化器状态。训练时,把全量训练数据切分成多个小份(Mini-batches),分发给不同的 GPU。各算各的,然后通过网络通信(All-Reduce)把梯度加总平均,最后所有 GPU 同步更新自己的模型。
DeepSpeed ZeRO
ZeRO 是由微软开源的数据并行技术,与 DDP 相比,它引入了对优化器状态和模型参数的切分。微软发现,既然大家都在算同一份模型,为什么每张卡都要存一份完整的优化器状态和参数?ZeRO 的核心思想是“切片”(Sharding)。它把训练时的显存占用大头切成了 N 份,分散存在 N 张显卡上。需要用的时候,再通过网络把缺失的部分广播过来。ZeRO 提供了三个阶段:
- ZeRO-1:只切分优化器状态(Optimizer States)。
- ZeRO-2:切分优化器状态 + 梯度(Gradients)。
- ZeRO-3:切分优化器状态 + 梯度 + 模型参数(Parameters)。此时每张卡只存极小一部分模型,彻底解决大模型显存问题。
PyTorch FSDP
FSDP 是 PyTorch 原生的终极切片数据并行,相当于 PyTorch 亲自实现的 ZeRO-3。它同样将模型参数、梯度和优化器状态打散分布在各个 GPU 上。在 Forward(前向传播)和 Backward(反向传播)时,动态地进行 All-Gather 操作收集参数,算完立刻释放(Free)显存。
由于是 PyTorch 原生集成,它与 PyTorch 生态(如动态图、编译加速 torch.compile)的兼容性较好,代码侵入性也相对更低。在不少以 PyTorch 原生工作流为主的 Infra 框架中,FSDP 也常被拿来与外挂式的 DeepSpeed ZeRO-3 对比和取舍。
CPU/NVMe Offloading
当模型大到即使切成 N 份,把所有 GPU 的显存加起来都装不下,就需要使用 ZeRO-Offload(CPU 卸载)和 ZeRO-Infinity(NVMe 卸载)技术。在训练大模型(比如使用 Adam 优化器)时,真正吃显存的“大头”其实不是模型参数本身,而是优化器状态(Optimizer States)。对于一个参数,Adam 需要额外存两份 FP32 的历史状态(动量和方差),这通常占据了整个训练显存的 50% 以上。CPU Offload 基本的工作原理:
- 前向传播和反向传播依然在 GPU 上飞速进行,算出梯度。
- 卸载(Offload):算完梯度后,GPU 不在本地做更新,而是通过 PCIe 总线,把庞大的梯度数据“打包扔给” CPU 内存。
- CPU 兼职打工:CPU 内存里存着那极其庞大的优化器状态。此时,CPU 的计算核心(比如几十个核心的 EPYC 处理器)被唤醒,它在内存里根据传过来的梯度,慢条斯理地更新优化器状态和 FP32 的高精度权重。
- 回传(Reload):CPU 算完并更新好最新的 FP16/BF16 权重后,再通过 PCIe 传回给 GPU 显存,供下一轮前向传播使用。
当模型奔向超千亿、万亿规模时,几 TB 的 CPU 内存都不够用,DeepSpeed 又提供了 ZeRO-Infinity 将数据进一步卸载到 NVMe 固态硬盘。由于传统的硬盘读写要经过 Linux 操作系统的文件系统缓存,极其缓慢,于是 ZeRO-Infinity 底层引入了高度优化的异步 I/O 与 O_DIRECT 机制,尽量降低经过页缓存带来的额外开销。
天下没有免费的午餐。Offload 技术虽然极其省钱,但它触碰到了计算机硬件架构的逆鳞:PCIe 带宽太窄了。在物理层面上:
- GPU 内部的 HBM3 显存带宽高达 3000 GB/s。
- GPU 之间的 NVLink 带宽高达 900 GB/s。
- 而 CPU 和 GPU 之间那条 PCIe 4.0 x16 通道,理论极限只有区区 64 GB/s(实际可能只有 20-30 GB/s)。
这就导致了一个严重问题:如果 GPU 傻傻地等数据通过 PCIe 传到 CPU,再等 CPU 算完传回来,整个训练速度会被拖慢十几倍,GPU 的算力利用率(MFU)会跌到惨不忍睹的个位数。为了不让 PCIe 拖后腿,Infra 工程师再次掏出了我们之前聊过的多流机制(Multi-Stream),并在内存分配上做了手脚:
- 锁页内存 (Pinned / Page-locked Memory):普通的 CPU 内存数据,操作系统随时可能会把它们交换到虚拟内存里。GPU 如果要拉取,CPU 得先把它找出来放好,极其浪费时间。Infra 框架会在 CPU 内存中强行“圈出一块特区(Pinned Memory)”,规定操作系统绝对不准动它。这样 GPU 的 DMA 引擎就可以闭着眼睛以最高速跨越 PCIe 把数据抓过来。
- 极致的掩盖(Overlap):就像 FSDP 在显存之间玩戏法一样,Offload 也会在后台疯狂重叠。当 GPU 的 Tensor Core 正在计算第 $N$ 层的反向传播时,后台的 PCIe 通道正在偷偷把已经算完的第 $N+1$ 层梯度卸载给 CPU。只要 GPU 计算矩阵的时间,能够掩盖住 PCIe 传输数据的时间,这种“空间换时间”的魔法就不会严重拖慢整体速度。
RL Infra
RL 算法
LLM 的强化学习训练算法核心在于如何容纳更多的模型实例来提升数据的采样效率,需要理解不同算法的模型加载数量、通信模式是离线还是在线的。
PPO 是经典的在线 RL 算法,是大模型对齐的基石,但也是 Infra 压力最大的算法。PPO 通过一个奖励模型(Reward Model)给输出打分,再通过一个评论家模型(Critic)预测价值,最后更新策略模型(Actor)。PPO 在训练时通常需要同时加载 Actor、Ref、Reward、Critic 四个模型,对于 70B 级别的模型,这会导致严重的显存碎片和 OOM;Actor 负责生成(推理模式),Critic 负责打分(训练模式),Infra 需要在生成阶段和更新阶段之间频繁切换资源权重;由于模型分布在不同 GPU 上,参数同步和梯度累积的带宽需求极高。
DPO、IPO 和 KTO 统称为 Preference Optimization,它们将 RL 问题转化为了类似“分类”的监督学习问题。DPO 直接利用偏好数据对(Chosen/Rejected)更新模型,IPO 和 KTO 分别解决了 DPO 容易过拟合、需要成对数据的问题。这类算法只需要加载 Policy 和 Ref 两个模型,不需要在训练中动态生成文本(Rollout),直接读取离线数据集。
GRPO 由 DeepSeek 在相关工作中引入,并在其推理模型训练路线中得到使用。GRPO 取消了 Critic 模型,让 Actor 对同一个问题生成一组(Group)答案,通过这组答案的相对得分来计算 Advantage。相比 PPO,GRPO 更依赖高吞吐的推理,Infra 需要支持较大的 group_size(如 16 或 64),这对 vLLM 等推理后端的集成要求较高。由于去掉了 Critic,这类方法通常还能减少一部分训练时的显存与系统复杂度。
DPO 虽然简单,但目前学界发现 Online DPO(训练中动态生成新答案让模型对比)效果更好。这意味着 Infra 不能只做 SFT 架构,必须具备训练/推理混合调度的能力。
RL Infra 基础设施架构
一些 Tricks
在 RL 训练中,Rollout(采样生成答案)占据了 70%~80% 的时间,因此现在的最新研究有许多 Trick 用于加速采样效率。例如 DeepSeek 提出的 MTP 技术,能够增强模型的推理速度;投机解码通过自对齐,利用前一个 Epoch 的生成答案存入树形缓存来投机;通过 Ray 或 VLLM+DeepSpeed 来消解采样和训练环节之间存在的计算重叠等。
推理系统
指标视角
推理系统常见指标包括:
TTFT:time to first tokenTPOT:time per output token- 吞吐
- QPS
- 显存占用
- GPU 利用率
很多推理优化问题,本质上都需要用这些指标来判断,而不是只看“模型跑没跑通”。
PD 分离
Prefill-Decoding 分离是现在大型模型推理的重要优化手段。这是因为 LLM 的推理可以分为 Prefill 和 Decoding 两个极端。其中,Prefill 阶段会对用户输入完成嵌入、注意力计算并生成 KV Cache,这个阶段与序列长度密切相关,并且是 Compute-Bound;Decoding 阶段则利用已有的 KV Cache 逐个生成新 Token,整体更接近串行计算,因此往往是 Memory-Bound。Prefill 和 Decoding 呈现出明显不同的瓶颈特征,因此通过分离二者可以更高效地利用资源。
这两个阶段主要采用的优化技术包括:
- Prefill:长序列并行、FlashAttention、分块、算子融合、向量化
- Decode:KV cache、连续Batch、PagedAttention、Speculative Decoding
在将 Prefill 和 Decode 拆分到不同的物理机器后,Prefill 机器算完必须将庞大的 KVCache 通过 IB 网络传送到 Decoding 机器,对于超长的文本传输需要大量时间,这一开销被称为 KV Cache Transfer 开销。即使面对这种开销仍需要进行分离,核心的收益有:
- 提高利用率:拆分后,Prefill 机器可以专心致志地处理大规模的并发输入,把 Batch Size 开到最大,让矩阵乘法填满 GPU 的每一个计算周期;Decode 机器则专心做逐字生成,不用担心突然被巨大的 Prefill 任务打断,从而可以维持极其稳定的吐字速度。
- 流水线并行:在网络带宽充足的情况下,两类机器都可以实现充足的计算负载,使得整体吞吐量增长。
- 灵活的资源分配:在实际业务中,我们可能会发现 Decode 的压力更大,采用 PD 分离后可以部署 1 台 Prefill 机器 对应 4 台 Decode 机器。这种灵活的架构在传统的一体机模式下是根本做不到的。
跨 SM 的 PD 分离
传统的 PD 分离(跨机器)虽然解决了算力干扰,但付出了高昂的网络传输(KV Cache Transfer)代价。跨 SM 的 PD 分离,试图在单卡内部鱼与熊掌兼得。其核心逻辑是把单卡上的 SM 划分为两个“特区”。比如,左边 60 个 SM 专门负责跑用户的长文本 Prefill(狂吃 Tensor Core 算力);右边 72 个 SM 专门负责跑其他用户的 Decode 生成(狂吃 HBM 显存带宽)。这样,在同一张 GPU 上,所有的 SM 都共享一块巨大的 L2 Cache(H100 上有 50MB)和底层的全局显存(HBM)。
这意味着 Prefill SM 算出来的 KV Cache,完全不需要过 PCIe 或网卡,直接写在本地显存里;Decode SM 下一秒就能“零延迟”读取,从而完美避开了网络空泡。但是 NVIDIA 底层的 Thread Block Scheduler 默认是吃“大锅饭”的,它倾向于把一个大算子均匀打散到所有 132 个 SM 上以追求最高 Occupancy(占用率)。要实现跨 SM 分离,Infra 工程师必须手写极其底层的 Persistent Kernel(持久化内核) 或者利用 MPS(多进程服务),强行给算子“划地盘”。这不仅开发难度极高,而且一旦负载不均衡,部分 SM 就会处于极度浪费的空转状态。
跨 SM 的 AF 分离(Attention-FFN 微流水线)
如果说跨 SM 的 PD 分离是按“请求类型”切分,那 AF 分离就是按 “Transformer 内部的计算图” 切分,是对 GPU L1 Cache 极致压榨的神仙操作。
这一技术的痛点在于解决缓存颠簸(Cache Thrashing)。在一个标准的 Transformer 层里需要先做 Attention,再做 FFN(也就是 MLP 层)。其中 Attention 是重度访存操作,它会把海量的 KV Cache 从 HBM 搬进 SM 内部极速但极其有限的 L1 Cache/寄存器中;结果下一秒,FFN 开始了,它是密集的矩阵乘法,又会强行把巨大的 FFN 权重塞进同一个 L1 Cache,瞬间把刚才 Attention 辛苦搬进来的数据冲刷得一干二净;到了下一层,Attention 又得重新从慢速 HBM 里去捞数据。
跨 SM AF 分离的解法是让一批 SM 永远只做 Attention,另一批 SM 永远只做 FFN。这样,Attention SM 的 L1 Cache 里永远热乎乎地躺着 KV 数据;FFN SM 的 L1 Cache 里永远躺着模型权重。两者算完的中间结果(Hidden States),通过单卡内部极速的 L2 Cache 进行“微流水线接力(Micro-Pipeline)”。
算子层面
FlashAttention
FlashAttention 的关键意义在于优化 attention 的 IO 路径,通过分块与更高效的片上存储利用,减少中间矩阵写回显存的次数,因此在长序列场景和 attention 成本高的场景下效果尤其明显。
框架与调度层面
PagedAttention
PagedAttention 主要解决的是长上下文推理里 KV Cache 管理低效、显存碎片和大块连续分配困难的问题。
KV Cache
KV Cache 的核心价值在于避免 decode 阶段重复计算历史 token,但它也会带来显存占用、碎片管理和调度复杂度。理解它的首要问题,是先看清它的存在形式:
- 物理层面(单卡/分布式):在引入了 TP(张量并行)的情况下,KV Cache 是被切碎物理存储在多张 GPU 的显存(VRAM)里的。比如一个 8 卡 TP 组,每张卡只存 $1/8$ 的 Attention Head 的 KV Cache。
- 逻辑层面(全局化池化):在最前沿的 PD 分离架构中(如 Mooncake),Infra 工程师会构建一个全局分布式 KV Cache Pool(池)。除了 GPU 显存,还会利用 CPU 内存甚至极速 NVMe 固态硬盘,把全集群的 KV Cache 统管起来,实现跨机器的复用。
从生命周期的角度来看,一个完整多轮对话里 KV Cache 的流转:
- 第一轮对话(Turn 1):
- Prefill 节点接客:调度器把用户的系统提示词和 Prompt 1 发给 Prefill 机器(内部采用 TP 并行计算)。算完后,生成这部分文本的 KV_1。
- KV 转移与 Decode 生成:Prefill 节点通过高速网络(RDMA)把 KV_1 传给 Decode 机器。Decode 机器开始逐字生成回复(Response 1)。
- 自产自销:Decode 机器每生成一个字,都会在本地追加生成这个字的 KV Cache(记作 KV_R1)。
- 第一轮结束:此时,完整的历史 Cache (KV_1 + KV_R1) 存在于 Decode 机器上,同时可能被异步备份到了全局 CPU 内存池中。
- 第二轮对话(Turn 2):
- 当用户带着新的 Prompt 2 再次发问时,由于包含了上下文,整个输入实际上是:Prompt 1 + Response 1 + Prompt 2。
- 智能路由(Sticky Routing):调度器会检查哈希值(类似 SGLang 的 RadixAttention),发现前文的 Cache 已经存在了!
- 增量 Prefill:调度器将请求发给拥有(或能最快拉取)历史 KV 的 Prefill 机器。这台机器不需要重新计算前面的几十万字,它直接复用历史 Cache,只对最新的 Prompt 2 进行极其短暂的前向传播,计算出增量的 KV_2。
- 增量 Transfer:Prefill 机器把仅包含新问题的增量 KV_2 发给 Decode 机器。Decode 机器拼装好完整的 Cache,继续生成新的回答。
Continuous Batching
传统静态 batch 更适合离线、规则负载;在线推理里,请求长度和到达时间都不均匀,因此更需要 Continuous Batching 来动态拼批,提高 GPU 利用率。
Chunked Prefill
当输入特别长时,prefill 可能本身成为瓶颈。Chunked Prefill 可以把长输入切块处理,以平衡时延、显存和调度开销。
Megatron-LM 设计分析
通信优化
TP+PP+DP 三维并行
- 把通信极其频繁、数据量巨大的 TP(张量并行)锁死在单台机器内部(利用超高速 NVLink)。
- 把通信量极小(只传边界隐状态)的 PP(流水线并行)放在机器与机器之间(走相对较慢的 IB 网络)。
- 最后用 DP(数据并行) 铺满全局。通过这种三维组合,把网络带宽的压力降到了物理极限。
重叠通信计算
在时间维度上实现的流水线掩盖,让网卡和 GPU 尽量同时承担工作负载,避免任何一方的空转。
在 TP 中:当 GPU 正在计算某一块矩阵乘法时,后台的网络引擎同时在为下一块矩阵拉取数据(All-Gather)。只要计算的时间 $T_{compute}$ 大于通信的时间 $T_{comm}$,从外部看来,通信的时间就仿佛“消失”了,GPU 实现了近乎 100% 的运转。
分桶 AllReduce
这是针对 DP(数据并行)阶段梯度同步的碎片化打包技术。在反向传播(Backward)结束时,所有卡需要同步梯度(AllReduce)。如果等整个大模型(比如 70B)的所有梯度都算完再去同步,网络瞬间会拥堵瘫痪,且 GPU 在同步期间只能干等。
Infra 框架会在内存里建几个“桶”。算完最浅层的一小批梯度,就扔进桶里。一旦某个桶装满了(比如达到 25MB),就立刻触发 AllReduce 把这桶数据发走,而此时 GPU 还在继续算深层的梯度,使得梯度同步和反向传播计算完美重叠。
激活重计算
这其实是一个显存优化技术,但它对通信优化有极其重要的间接作用。在正常训练中,前向传播(Forward)的中间结果(激活值)需要保存在显存里,留给反向传播(Backward)算梯度用,极其吃显存。
Megatron 选择不保存这些庞大的激活值,而是在反向传播时重新算一遍(用算力换显存)。省下了海量显存后,Infra 工程师就可以大幅提高 Micro-batch Size。Batch Size 越大,GPU 每次计算的任务就越重(算力密度变高),就越容易在时间上掩盖住网络通信的延迟。
减少跨节点通信量
从算法底层动刀,直接砍掉需要传输的数据字节数。除了把 TP 放在机内,Megatron 后期还引入了一个神级操作——序列并行 (Sequence Parallelism, SP)。
在处理长文本时,原本 Dropout 和 LayerNorm 操作在 TP 中会产生冗余的 All-Reduce 通信。SP 技术把文本序列(Sequence Dimension)也沿着 TP 的维度切开,彻底消灭了这部分冗余通信,不仅省了带宽,还进一步省了显存。
vLLM 设计分析
理解 vLLM 时,比较值得关注的是:
- 它如何组织
KV Cache - 为什么引入
PagedAttention - 它如何支持高吞吐推理
- 它在长上下文和多请求场景下的优势是什么
如果后续继续深入,这一节可以进一步补:
Prefill与Decode的差异- scheduler 如何影响吞吐与时延
vLLM与SGLang、TensorRT-LLM的关注点差异
结语
AI System 的难点不只是模型本身,而是要同时理解模型、硬件、通信、编译、运行时和调度系统之间的耦合关系。很多表面上看像“模型问题”的现象,最终都可能落到:
- 显存不够
- 带宽不足
- 通信太慢
- 调度不合理
- 上下游链路没有把 GPU 喂满
如果把这些层次串起来,才算真正进入 AI Infra 的语境。