TensorRT LLM 中的并行策略

概述

当单个 GPU 无法满足大模型的容量或性能需求时，就需要通过并行化策略在多个 GPU 间分散计算。TensorRT LLM 提供了多种灵活的并行方案，开发者可根据模型特点和硬件配置选择最优策略组合。

何时采用并行化

满足以下任一条件时，需要采用多 GPU 并行化：

• 模型无法装入单个GPU的显存中，或
• 单个GPU无法达到所需的性能指标。

TensorRT LLM 支持多种并行策略，适用于单节点和多节点部署：

• 张量并行 (TP) - 在GPU间分片模型权重
• 流水线并行 (PP) - 在GPU间分布模型层
• 数据并行 (DP) - 为不同请求在GPU间复制模型
• 专家并行 (EP) - 在GPU间分布MoE模型的专家
• 上下文并行 (CP) - 在GPU间分布上下文处理
• 宽专家并行 (Wide-EP) - 面向大规模MoE模型的高级EP方案，具备负载均衡

并行策略概览

张量并行 (TP)

张量并行将模型权重分片到多个GPU上。每个GPU持有部分权重并处理相同的输入令牌，通过通信组合结果。

适用场景： 小批次大小、显存受限场景

流水线并行 (PP)

流水线并行将模型的不同层分布到多个GPU上。每个GPU处理层的子集，激活在GPU间传递。

适用场景： 无法装入单个GPU的大型模型

数据并行 (DP)

数据并行在多个GPU上复制完整模型。每个GPU独立处理不同请求。

适用场景： 大批次大小、高吞吐量场景

专家并行 (EP)

专家并行专为混合专家 (MoE) 模型设计，将不同专家分布在GPU间。

适用场景： 高专家数量的MoE模型

上下文并行 (CP)

上下文并行将长序列的处理分布到多个GPU上。

适用场景： 长上下文场景

宽专家并行 (Wide-EP)

宽专家并行是专家并行的高级形式，通过智能负载均衡和专家复制策略应对大规模MoE模型的工作负载不均衡问题。

适用场景： 大规模MoE模型（如DeepSeek-V3/R1、LLaMA4、Qwen3）

模块级并行指南

注意力模块

TensorRT LLM支持两种注意力模块策略：

• 张量并行 (TP) — 适用于小批次
• 数据并行 (DP) — 适用于大批次

张量并行 (TP)

• 注意力核前后的GEMM权重和注意力头数均匀分片到各GPU上。
• 例外情况：
  1. DeepSeek-R1：fused_A GEMM 不 分片。
  2. GQA / MQA / MLA：当 num_heads < tensor_parallel_size 时，KV缓存在所有GPU上复制。

数据并行 (DP)

• 所有GEMM权重在所有GPU上 复制。
• KV缓存被 分区，因为不同用户请求路由到不同的DP排名。

启用注意力并行

使用 trtllm-serve 部署或使用 trtllm-bench 进行基准测试时，创建名为 parallel_config.yaml 的YAML配置文件：

cat <<EOF > parallel_config.yaml
# TP-8
tensor_parallel_size: 8
enable_attention_dp: false    # 默认值
# DP-8
tensor_parallel_size: 8
enable_attention_dp: true
EOF

FFN模块

稠密模型

张量并行支持稠密模型的FFN层。

混合专家 (MoE)

MoE用多个专家替代单一的FFN。路由器为每个令牌选择top-k个专家并分派相应的隐藏状态。

TensorRT LLM支持三种MoE执行模式：

• TP - 每个专家的权重矩阵在所有GPU间分片。每个GPU看到所有令牌。
• EP - 每个专家的完整权重驻留在单个GPU上。每个GPU仅看到路由到其本地专家的令牌。
• 混合ETP - 每个GPU存储专家的子集 (EP) 并进一步分片这些权重 (TP)，平衡工作负载和核心效率。

启用MoE并行

使用 trtllm-serve 部署或使用 trtllm-bench 进行基准测试时，按如下方式创建 parallel_config.yaml 配置文件：

cat <<EOF > parallel_config.yaml
# 仅TP
tensor_parallel_size: 8
moe_tensor_parallel_size: 8

# 仅EP
tensor_parallel_size: 8
moe_expert_parallel_size: 8

# 混合 (TP-4 × EP-2)
tensor_parallel_size: 8      # 4 × 2
moe_tensor_parallel_size: 4
moe_expert_parallel_size: 2
EOF

`moe_tensor_parallel_size` 和 `moe_expert_parallel_size` 的乘积必须等于 `tensor_parallel_size`。

宽专家并行 (Wide-EP)

宽专家并行 (Wide-EP) 是TensorRT LLM针对大规模MoE模型推理的高级解决方案。它通过智能负载均衡和专家复制策略应对传统专家并行的挑战。

Wide-EP的动机

像DeepSeek-V3/R1、LLaMA4和Qwen3这样的大规模MoE模型采用细粒度专家设计，引入了新的挑战：

• 高显存需求 用于专家权重
• 固有的专家级工作负载不均衡 由于稀疏执行模式
• 通信开销 在分布式专家并行中
• 热专家问题 某些专家收到的令牌显著多于其他专家

Wide-EP的核心特性

1. 专家复制和负载均衡

Wide-EP引入了 专家槽 的概念，与特定专家解耦。这允许：

• 热专家的多个副本分布在不同GPU上
• 基于工作负载模式的动态专家放置
• 离线和在线负载均衡策略

2. 自定义EP通信核心

• 针对NVIDIA GB200多节点NVLink (MNNVL) 优化
• 用于专家分派和合并的高效全对全通信
• 相比传统EP的通信开销更低

3. 专家并行负载均衡器 (EPLB)

• 离线EPLB：基于历史工作负载统计的预计算专家放置
• 在线EPLB：适应实时流量模式的动态专家放置
• 层级权重重分布以最小化推理中断

架构概览

Wide-EP将 专家 和 槽 的概念分离：

• 专家：从模型角度的概念（例如，专家0、专家1等）
• 槽：从模型引擎角度的概念（例如，槽0、槽1等）

系统维护一个路由表，将专家ID映射到槽ID，可由负载均衡策略更新。

最佳实践

1. 从离线EPLB开始 用于已知工作负载模式的生产部署
2. 使用在线EPLB 用于动态工作负载或流量模式频繁变化的场景
3. 监控专家统计 以理解工作负载分布
4. 调整max_num_tokens 基于显存约束和EP大小
5. 使用代表性数据集测试 以验证负载均衡有效性

参考资源

• 技术博客：TensorRT LLM中的专家并行扩展
• DeepSeek-V3论文
• EPLB实现

详细实现示例和高级用法，参见：

• examples/wide_ep/：完整Wide-EP示例
• examples/wide_ep/ep_load_balancer/：负载均衡工具
• examples/wide_ep/slurm_scripts/：集群部署脚本