1. 为什么需要分布式推理
在大模型时代，单张GPU卡的内存容量和计算能力已经无法满足模型推理的需求。以Llama 2 70B为例，仅模型参数就需要140GB显存（按FP16计算），这远超目前任何单卡GPU的容量。分布式推理通过将计算任务拆分到多个设备上并行执行，主要解决以下问题：
- 显存墙限制：通过模型并行（TP/PP）将参数分散到多卡
- 计算效率瓶颈：利用数据并行（DP）提高吞吐量
- 延迟优化：专家并行（EP）等技术可降低特定场景的响应时间
实际案例：在A100 80G显卡上，单卡最多只能运行Llama 2 13B模型（26GB参数+KV缓存）。使用8卡Tensor Parallelism后，可稳定运行70B模型。
2. 分布式并行策略解析
2.1 张量并行（Tensor Parallelism）
2.1.1 工程实现机制
vLLM的TP实现核心位于vllm/distributed/parallel_state.py，关键数据结构如下：
python复制class GroupCoordinator:
def __init__(self, group_ranks, local_rank, backend="nccl"):
self.rank = dist.get_rank() # 全局rank
self.ranks = group_ranks # 组内所有rank列表
self.world_size = len(group_ranks)
self.local_rank = local_rank # 节点内局部rank
self.cpu_group = dist.new_group(backend='gloo')
self.device_group = dist.new_group(backend=backend)
通信流程示例：
- 前向传播时，每个GPU计算部分注意力头
- 通过
all-reduce操作聚合结果 - 输出张量在设备间同步
2.1.2 算法层实现
以Llama的注意力层为例（vllm/model_executor/models/llama.py）：
python复制class LlamaAttention(nn.Module):
def __init__(self, tp_size: int):
self.num_heads = config.num_attention_heads // tp_size
self.qkv_proj = ColumnParallelLinear(
hidden_size,
3 * hidden_size // tp_size,
gather_output=False)
def forward(self, x):
qkv = self.qkv_proj(x) # 各卡计算部分QKV
# ... 后续进行all-reduce通信
关键参数拆分原则：
- 注意力头数按TP维度均分
- 全连接层输出特征拆分为tp_size份
- 确保各卡计算量均衡
2.2 流水线并行（Pipeline Parallelism）
vLLM的PP实现特点：
- 按层划分模型阶段
- 使用微批次（micro-batch）提高设备利用率
- 通信模式为点对点（P2P）
典型配置示例：
python复制# 4卡PP配置
stage_mapping = {
0: range(0, 10), # 第1卡运行0-9层
1: range(10, 20), # 第2卡运行10-19层
...
}
避坑指南：PP需要仔细平衡各阶段计算时间，否则会产生气泡（bubble）开销。建议先用profiler分析各层耗时。
2.3 专家并行（Expert Parallelism）
适用于MoE架构的通信模式：
- 专家分布在不同设备
- 门控网络决定路由
- 使用
all-to-all通信交换token
实现示例（伪代码）：
python复制def forward(x):
gates = gating_network(x) # 计算路由权重
expert_inputs = all_to_all(gates, x) # 分发输入
expert_outputs = [expert(x) for x in expert_inputs]
return all_to_all(gates, expert_outputs) # 收集输出
2.4 数据并行（Data Parallelism）
vLLM的DP实现要点：
- 每卡持有完整模型副本
- 前向传播独立执行
- 反向传播时梯度all-reduce
性能优化技巧：
- 使用梯度累积（gradient accumulation）
- 采用ZeRO-3优化显存
- 重叠通信与计算
3. 混合并行实战配置
典型8卡混合并行配置示例：
| 并行策略 | 拆分维度 | 通信频率 | 适用场景 |
|---|---|---|---|
| TP=2 | 模型参数 | 高频 | 降低单卡显存压力 |
| PP=2 | 模型层 | 中频 | 超长序列处理 |
| DP=2 | 数据批次 | 低频 | 提高吞吐量 |
配置代码参考：
python复制parallel_config = {
"tensor_parallel_size": 2,
"pipeline_parallel_size": 2,
"data_parallel_size": 2
}
4. 性能调优经验
4.1 通信优化技巧
- 使用NCCL的
IB网络（InfiniBand） - 开启CUDA Graph减少启动开销
- 调整
all-reduce的group size
4.2 显存管理
- 激活检查点（activation checkpointing）
- 使用FP8量化
- 动态KV缓存压缩
4.3 典型问题排查
- 死锁问题：检查TP组内通信是否成环
- 显存溢出：减小micro-batch size
- 负载不均：用
torch.profiler分析各卡计算时间
5. 源码学习路线建议
-
初级路线：
vllm/distributed/device_communicators→ 通信原语vllm/model_executor/models/llama.py→ TP实现
-
进阶路线：
vllm/engine/llm_engine.py→ 调度逻辑vllm/worker/worker.py→ 执行单元
-
调试技巧：
- 设置
TORCH_DISTRIBUTED_DEBUG=DETAIL - 使用
nsight分析通信耗时
- 设置