模型拆分

模型拆分是突破单卡显存限制的关键技术。当一个大型模型（如 Llama-3-70B 或 DeepSeek-V3）的显存占用超过单张显卡的容量（H100 为 80GB，而 70B 模型 FP16 权重就需 140GB+），必须进行模型切分。从分布式系统的视角来看，这是将一个巨大的计算图（DAG）进行跨节点/跨进程的解耦与流水线化。目前主流的拆分方式有两种：张量并行和流水线并行。

推理并行

推理场景的核心诉求是低延迟。vLLM 在分布式推理中主要采用张量并行，在单机多卡环境下通过 NVLink 的高带宽实现多卡协同计算。

张量并行原理

张量并行将模型每一层内部的巨大矩阵切开，分配到不同显卡上。以 Transformer 中的全连接层 $Y=XW$ 为例，有两种切分方式。

行并行将权重矩阵 $W$ 按行切分，每张卡计算一部分乘加，最后通过 All-Reduce 操作将各卡结果求和。列并行将权重矩阵 $W$ 按列切分，每张卡输出 $Y$ 的一部分通道，最后通过 All-Gather 拼接。

vLLM 使用列并行作为第一层，行并行作为第二层，这样中间结果无需通信，仅在输出层进行一次 All-Gather 或 All-Reduce。这种设计最小化了通信开销，充分利用了 NVLink 的 600GB/s 带宽。

底层实现

vLLM 的张量并行实现基于 SPMD 编程模型，每个显卡运行完全相同的代码，但处理不同的数据分片。在每一层计算结束时，通过 NCCL 进行跨卡的集合通信。

vLLM 通常为每张显卡 fork 一个独立进程，通过 Ray 或 PyTorch Distributed 管理。进程间通过共享内存交换控制信息，通过 NCCL 传输张量数据。为了抵消分布式环境下频繁 Launch Kernel 带来的 CPU 开销，vLLM 会使用 CUDA Graphs 将整条分布式路径录制下来。

分布式 PagedAttention

切分模型只是第一步，分布式场景下最难的是 KV Cache 的管理。vLLM 采用集中式调度，有一个 Master 节点负责计算每个请求的 Token 应该放在哪张卡的哪个物理块上。各卡在自己的显存空间里找到对应的局部 KV Cache 进行计算，然后通过 TP 进行结果同步。

流水线并行

当张量并行达到 8 卡（单机上限）依然放不下模型时，就需要跨机使用流水线并行。流水线并行将模型的层拆分到不同机器，例如 80 层 Transformer Block 可以拆分为：第 1-20 层放 0 号机，21-40 层放 1 号机，以此类推。数据像流水线一样流动，0 号机算完输出后通过网络传给 1 号机。

流水线并行的痛点是气泡问题，后端显卡在等待前端显卡计算时是空闲的。vLLM 较少在高性能场景下单独依赖流水线并行，通常采用张量并行加流水线并行的混合策略。

张量并行 vs 流水线并行

张量并行是算子内部矩阵的切分，通信频率极高（每一层都要同步），硬件要求极高带宽（如 NVLink），主要瓶颈是通信延迟。流水线并行是算子层与层之间的切分，通信频率较低（每组层算完才同步），可以使用 RoCE 或 InfiniBand，主要瓶颈是负载均衡与气泡。

单机多卡场景主要靠张量并行，利用 NVLink 把多张卡连成一张逻辑大显卡。多机多卡场景主要靠张量并行加流水线并行，利用网络让模型在不同机器间像流水线一样接力。

训练并行

训练场景的核心诉求是海量参数的显存吞吐平衡。DeepSpeed 通过 ZeRO 协议，将原本在每张卡上都冗余存储的参数、梯度、优化器状态打散，分布到整个集群中。

显存消耗的组成

在训练场景下，显存不只是存模型权重。一个 70B 模型训练时，显存占用远超其权重大小。模型权重（FP16）需要 $2\times P$ 字节，梯度同样需要 $2\times P$ 字节，优化器状态以 Adam 为例需要存储 FP32 的 Master Weights、Momentum 和 Variance，高达 $12\times P$ 字节。总计约 $16\times P$ 字节，远超模型权重大小。

ZeRO 优化的三阶段

ZeRO-1 分片优化器状态，每张卡只维护一部分优化器状态。当需要更新权重时，每张卡只更新自己负责的那部分参数。这几乎没有额外的通信开销，但能节省约 80% 的优化器相关显存。

ZeRO-2 进一步分片梯度。在反向传播计算出梯度后，每张卡只保留自己负责的那部分梯度，其余直接释放或发送给对应的节点。这类似于多线程编程中的分段锁或分区缓存。

ZeRO-3 是参数分片，每张卡在平时只存储 $1/N$ 的模型权重。前向计算时，当计算第 5 层时，所有卡通过 All-Gather 实时从其他卡拉取第 5 层的剩余权重，拼成完整层进行计算。计算结束后立即释放掉拉取过来的权重，显存只保留自己那 $1/N$。理论上你可以训练无限大的模型，只要集群的总显存够大。

ZeRO 与张量并行的对比

ZeRO 不像张量并行那样去切矩阵，而是像数据分片一样去切显存。张量并行产生的是 All-Reduce（高频、同步、数据量中等），ZeRO-3 产生的是 All-Gather（前向拉取）和 Reduce-Scatter（反向同步梯度）。

3D 并行

DeepSpeed 通常与 Megatron-LM 结合使用，形成 3D 并行架构。数据并行每张卡处理不同的数据 Batch，张量并行将单层内的矩阵横纵切分针对单机内 NVLink 优化，流水线并行将不同的层放在不同的机器上针对机间低带宽优化。这是将模型推向万卡规模的关键。

ZeRO-Offload

ZeRO-Offload 利用 CPU 内存和 NVMe 硬盘作为虚拟显存。将优化器状态的更新从 GPU 挪到 CPU 上执行，当 GPU 显存实在不够时，利用 cudaMemcpyAsync 将暂时不用的权重换出到系统内存。这实现了用时间换空间，训练速度下降 30-50%，但可将训练 1T 参数模型的 GPU 需求从 1024 张 A100 降至 128 张。