生成式人工智能/大语言模型

扩展 TensorRT-LLM 中的专家并行度:大规模 EP 的设计与实现

DeepSeek-V3 / R1 等模型采用大规模细粒度混合专家模型 (MoE) 架构,大幅提升了开源模型的质量。Llama 4 和 Qwen3 等新发布的开源模型的设计原则也采用了类似的大规模细粒度 MoE 架构。但大规模 MoE 模型为推理系统带来了新的挑战,如高显存需求和专家间负载失衡等。

之前,我们介绍过突破 DeepSeek R1 模型低延迟极限的 TensorRT-LLM 优化措施,多 Token 预测 (MTP) 的实现与优化以及提高 DeepSeek R1 吞吐量性能的优化措施

DeepSeek 团队还分享了优化此类大规模专家并行 (EP) 模型 (如 DeepEPEPLB) 的宝贵经验与实践。此外,DeepSeek 团队在这份[1]技术报告中详细阐述了具体设计考虑因素。除此之外,社区中也有在其他推理引擎中实现大规模 EP 的优秀实践,例如 SGLang 团队的这个项目[2]

这篇技术博客共分为上中下三篇,将介绍支持 TensorRT-LLM 中端到端大规模 EP 的详细设计与实现,主要包含以下内容:

  • 如何使用 NVIDIA 多节点 NVLink (MNNVL) 硬件特性来实现高性能通信内核。
  • 如何设计和实现在线专家负载均衡器,以动态平衡专家负载分布并适应在线流量模式的变化。我们将展示:
  • 证明此类优化措施必要性的经验数据分析。
  • 在线流量数据统计模块的实现。
  • 复制 / 放置策略的设计与实现。
  • 用于平衡多个 GPU 间在线工作负载的 MoE 权重负载 / 重新分配器。
  • 为适应专家负载均衡器需求而对 MoE 路由器和计算模块进行的必要修改。
  • 一些证明当前 TensorRT-LLM 中实现效果的初步数据。

未来的技术博客还将涵盖以下主题:

  • 对 TensorRT-LLM 大规模 EP 实现的性能调优和优化的介绍。
  • 如何在不使用 MNNVL 的情况下,为 Hopper 和其他 NVIDIA GPU 实现高效的大规模 EP 支持。
  • 使用大规模 EP 并获得性能提升的最佳实践。
  • 如何将大规模 EP 与其他系统优化技术相结合。

虽然本技术博客主要讨论 TensorRT-LLM,但我们相信其核心理念和实现方法也可用于其他推理引擎在 NVIDIA GPU 上的推理性能。此外,我们希望借助社区的力量,探索如何更好地将当前 TensorRT-LLM 大规模 EP 实现模块化,使其更容易被社区复用。

最后,本博客包含针对 Grace Blackwell 机架式系统的详细实现方式,例如使用 Grace Blackwell 机架式系统跨 GPU 连接的通信组件,以及使用 Grace CPU 与 Blackwell GPU 间高带宽 C2C 连接的 MoE 权重加载 / 重新分配模块等。但整体设计原则和软件架构仍适用于非此 NVIDIA GPU 系统。为了便于扩展到其他非此系统,我们有意识地关注了设计和实现的通用性。这些更改应能与现有其他组件轻松组合。

引入大规模 EP 的初衷

引入大规模 EP(本文中指 EP > 8)主要基于以下系统考量:

  • 我们希望通过提高聚合显存带宽来加载专家权重,从而降低执行延迟。
  • 我们希望通过增加有效批处理大小充分利用 GPU 算力。

需注意,当端到端 (E2E) 执行时间主要由 MoE GroupGEMM 计算主导时,引入大规模 EP 可显著提升性能。但若端到端执行时间未被 MoE GroupGEMM 计算主导,引入大规模 EP 提升的性能有限。

系统设计中不存在“免费的午餐”。当 EP 规模增大到超过 8(有时甚至不到 8)时,由于 MoE 模型的稀疏执行特性,会自动触发 EP 级别的负载失衡问题。

以下是一些基于特定数据集的经验观察(所有分析均使用DeepSeek R1 模型32 GPU 上进行):

对一个机器翻译数据集的观察结果

首先,我们将概述各层的整体失衡问题:

1. rank 0 发送到所有 rank(包括 rank 0)的 Token 数(对应解码迭代 1950)及所有 MoE

如图 1 所示,在第 36 层的 MoE 中,从 rank 0 发送到 rank 13 的 Token 数明显更多。

如果我们放大第 36 层的 MoE 并记录其激活专家 rank 的分布,可以清楚地看到有一个 rank 被更频繁地激活:

2. 36 层每个专家 rank 接收的 Token 数量

如果我们将数据展平以查看每个专家接收的 Token 数量,可以发现有一些专家比其他专家更活跃:

3. 36 层每个专家接收的 Token 数量

值得注意的是,这种失衡问题在多次迭代中非常稳定,如下图所示:

4. 36 层每个专家在 50 个解码步骤内接收的 Token 总数,本地 batch size=256

显然,图 4 中的热门专家与图 3 中仅包含单次解码迭代数据的专家相同。我们还对本地 batch size=1(对应单次请求)进行了基于持续时间的分析,观察到类似的模式:

5. 36 层每个专家在 400 次解码迭代内接收的 Token 总数,本地 batch size=1

综上所述,针对该机器翻译数据集的研究结果可总结为:

  • 某些层中存在一些热点,部分 EP 所在 GPU 的负载可能远高于其他 EP。
  • 其原因可能是最热门专家或多个热门专家位于同一 rank。
  • 路由 Token 的分布可能在数十至数百个迭代步骤甚至更多迭代步骤内保持一致。
  • 在单个请求的执行中,不同迭代步之间也存在相同的热门专家。

另一个实际问题是上述观察结果在其他数据集上是否会发生显著变化。因此,我们对 GSM8K 数据集进行了类似的分析。

对 GSM8K 数据集的观察结果

6. rank 0 发送到所有 rank Token 数(对应第 1950 个迭代步)及所有 MoE

如图 6 所示,与图 1 相比,GSM8K 数据集中的热门层变成了第 57 层而非第 36 层。那么 GSM8K 数据集中第 36 层的具体情况如何?

7. EP rank 0 发送到其他 EP rank Token 数(仍以迭代 1950MoE 36 层为例)

从图 7 可以清楚地看到,工作负载失衡于不同数据集(图 2 所示)中观察到的情况不同。在图 8 中可以观察到在 GSM8K 数据集上,工作负载的失衡在多次迭代中也相对稳定。这与之前的机器翻译数据集相同。

8. EP rank 0 发送到所有 rank Token 总数(MoE 57 层,50 个解码步骤内,本地 batch size=256

如果我们将每个 GPU EP 层面的数据展平为专家层面,可以得到下图。

9. 57 层的每个专家在 50 个解码步骤内接收的 Token 总数,本地 batch size=256

单个请求中也存在类似的失衡模式。

10. 单次请求下第 57 层的每个专家在 400 个解码步骤内接收的 Token 总数

如果使用另一个请求,我们仍然可以观察到专家失衡问题。虽然热门专家可能不同,但有一些是共同的(在此示例中是专家 10)。

11. 单次请求下第 57 层每个专家在 400 个解码步骤内接收的 Token 总数

通过对两个数据集的数据分析,我们得出以下结论:

  • EP 级别工作负载失衡问题在多个数据集的大规模 EP 推理中较为常见。且 EP 失衡的严重程度可能因层而异。此外,EP 失衡的问题具有数据集敏感性。
  • EP rank 级别失衡问题可能由某个最热门的专家或多个热门专家长期占据同一 EP rank 引起。
  • EP rank 失衡分布在数十到数百次迭代中相对稳定。
  • 尽管 EP rank 失衡分布在时间维度上具有稳定性,但不同请求的 EP 失衡分布可能有所不同。

这些发现可指导我们对 TensorRT-LLM 大规模 EP 实现的设计考量:

  • 设计时需考虑 EP 失衡问题以确保端到端的性能。
  • 基于实时在线请求流量的在线 EP 负载均衡器(而非仅实现离线 EP 负载均衡器)对确保 EP 均衡器的稳健性至关重要。
  • 可运用 EP rank 失衡分布的时间维度稳定性,以高效的方式将 MoE 权重重新分配至不同 EP rank。

下一部分将详细说明整体设计。

高层次设计介绍

根据引入大规模 EP 的初衷部分的详细分析与研究,可以明确观察到 EP 中的专家失衡是大规模 EP 的常见模式。这种 EP 失衡会通过以下方式显著降低整体系统性能:

  • 热门 EP rank 将消耗更多显存(用于激活值),这会限制推理过程中调度的有效最大批处理大小。
  • 更多数据将从热门 EP rank 被发送和接收。

这些问题将导致系统级拥塞效应,即热门 EP rank 将延迟整体端到端执行。

为确保大规模 EP 能稳定运行,需通过精心设计尽可能减少 EP 失衡问题。整体设计如下:

12. TensorRT-LLM 大规模 EP 的高层次设计

此设计同时包含 CPU 和 GPU 两侧逻辑:

  • CPU 侧
  • 使用复制与放置算法(复制与放置计算组件)实现更均衡的 EP 策略。这些算法是经典算法,更适合 CPU 计算。此外,将此计算卸载至 CPU 可减少对 GPU 的干扰。未来可探索基于机器学习的算法,并可能需要额外设计考量。复制与放置计算组件将生成放置信息,该信息将被 GPU 路由逻辑和 CPU 更新权重与放置组件共同使用。由 GPU 上运行的统计组件生成的统计数据将被用作复制与放置计算组件的输入。
  • 编排流程(更新权重与放置组件)将 MoE 权重从 CPU 内存更新并重新加载到 GPU 设备显存。该组件还将使用由复制与放置计算组件生成的放置信息。我们的可扩展设计允许通过 MNNVL 或 NIC 从远程 GPU 显存重新加载 MoE 权重。
  • GPU 侧
  • 这是推理的主要执行工作流。我们在设计中引入了以下新的 GPU 组件:
  • EP 通信内核,在图 11 中为分发合并组件。
  • 在线流量数据统计采集器(统计组件)。该组件采集统计数据复制与放置计算组件使用。
  • MoE 路由逻辑(路由组件)。该组件将 Token 发送至激活的专家,并且需要进行调整以支持 MoE 权重的动态放置。它使用复制与放置计算组件生成的放置信息
  • MoE 计算逻辑 (MoE 组件) 也需进行相应调整。
  • CPU 和 GPU 组件之间需要仔细同步,以确保整个执行过程的有效性,尤其是为了避免卡顿以及无效或次优执行。

我们为更新权重与放置组件提供了两种设计方案:

  • 批量方案
  • 在此方案中,当 MoE 权重重新分配逻辑启动时,当前服务实例上的推理过程将不得不暂停,直至 MoE 权重重新分配过程完成。我们估计这可能导致 0.5 1 的在线服务暂停,最坏情况下会引发请求超时。此类超时或暂停可通过系统级措施来缓解,例如将请求传送至其他服务实例或通过请求重试来应对。
  • 分层方案
  • 在此方案中,MoE 权重的重新分配是分层逐步进行,即在每次解码迭代中,只有某些层(可配置)会受到 MoE 权重重新分配的影响。这种设计需要多次迭代才能再平衡所有层的 MoE 权重。我们预计此方法对用户体验几乎没有影响。
13. 分层 MoE 权重重新分配示例

在当前系统中,我们选择采用分层方案以尽量减少对在线用户体验的影响。批量方案应更易于实现,但本文将不再讨论。为了正确实现分层方案,需仔细评估不同底层硬件的性能以确定具体实现方案。图 14 展示了系统节点中不同硬件组件的通信带宽。

14. 系统高层次拓扑结构

以 DeepSeek R1 模型为例,采用 FP4 精度时,每个 MoE 专家占用 24MiB 显存空间。每层包含 256 个专家,总共包含 58 个 MoE 层加 1 个 MTP 层。因此,为实现 EP 平衡所需重新分配的 MoE 权重最大总量为 348GiB。每个节点为每个 Grace CPU 提供 480GB LPDDR5X 显存。在 NUMA 域内,总计可提供 960GB Host 显存。一个节点可在其 CPU Host 显存中完整承载如 DeepSeek R1 LLM 等模型的全部 MoE 权重。基于此,MoE 权重重新分配可通过将对应的 MoE 权重从 CPU Host 显存移动至 GPU 设备显存来实现。

假设我们将 50ms 的跨 Token 延迟 (ITL) 作为主要延迟约束。通过粗略估算,可以计算出在每次解码迭代中,可从 MoE 权重池(可保存在 Grace CPU 显存或另一节点上的 GPU 显存中)移动到 Blackwell GPU(用于实际 MoE 推理)的专家权重数量为:

15. 在以下 50ms ITL 限制下,每次迭代理论上需要更新的专家数量(使用不同硬件作为存储完整 MoE 权重的池)

基于此分析,若依赖每个节点上的 Grace CPU 内存来存储 MoE 权重池,则每次解码迭代中,最多可将 300 个专家的权重重新分配至同一节点上的每个 GPU。假设目标是在 5 次解码迭代内完成整个模型 MoE 权重再平衡,以下为具体用例研究:

  • 用例 1(专家分配均衡,不进行专家复制)
  • 64 个 GPU,每个 GPU 分配 4 个专家
  • 58 层,每个 GPU 分配 232 个专家
  • 每次迭代需要 47 次专家更新,所有方法均可满足延迟目标。
  • 用例 2(专家分配均衡并进行复制)
  • 64 或 72 个 GPU,每个 GPU 分配 5 个专家
  • 58 层,每个 GPU 分配 290 个专家
  • 每次迭代需要 58 次专家更新,所有方法均可满足延迟目标。
  • 用例 3(专家分配均衡并进行复制)
    • 36 个 GPU,每个 GPU 分配 8 个专家
  • 58 层,每个 GPU 分配 464 个专家
  • 每次迭代需要 93 次专家更新,所有方法均可满足延迟目标。

综上所述,根据理论分析,采用 Grace CPU 内存作为存储完整大小 MoE 权重的池,应能使我们在 5 次解码迭代内实现 EP(专家并行)的再平衡。如果将要求放宽至 10 次或以上迭代,系统实现将变得更加灵活。

接下来我们将介绍大规模 EP 系统的详细实现方式。

EP 通信内核

我们评估了多种实现大规模 EP 所需 EP 通信内核的途径,包括 DeepEP、其他解决方案以及重新开发一种方法。

当前的技术决策是:

  • 对于其他系统(如 Hopper),我们选择直接集成 DeepEP 并进行一些可能的增强。

考虑因素:

  • DeepEP 是由 DeepSeek 团队完成的一项出色成果。我们在启动 TensorRT-LLM 大规模 EP 工作时,最初把重点放在 Grace Blackwell 机架式系统上。我们选择实现自己的定制 EP 通信内核,因为这更便于引入需要 Grace Blackwell 机架式系统功能的优化措施。
  • 当我们开始在 Hopper 上启用大规模 EP 工作时,我们得出的结论是 DeepEP 可以适应并满足我们在该平台上的需求。

我们也在积极评估将通信内核整合为单一解决方案以简化系统架构的可能性,并将持续向社区更新进展。接下来,我们将进一步探讨自定义 EP 通信内核实现中引入的优化措施。

在系统中引入 EP 通信内核的初衷

在解码阶段与预填充解码 (PD) 分离的场景中,我们观察到批处理大小可能不会很大,因此延迟成为一个重要考虑因素。在此背景下,我们非常需要实现与 CUDA graph 的兼容。NCCL 是一个优秀的 GPU 通信库,为我们提供了高效的通信内核和基本操作。目前,其 Send 和 Recv 操作在调用 ncclSend / ncclRecv 时,需要显式指定数据大小。但在大规模专家并行 (large-EP) 场景中,待传输的数据大小根据模型在每次迭代中的输出动态确定。当前 NCCL 通信接口需要同步将通信大小发回 CPU,并以对应数据大小从 CPU 发起 NCCL 调用。这将破坏 CUDA graph 兼容性。这一限制迫使我们开发与 CUDA graph 兼容,且能直接从 GPU 显存接受通信大小的高性能通信内核。我们还希望这些内核能够充分利用 MNNVL 的显存带宽。

EP 通信内核的实现

我们的内核采用与 NCCL 的 LL128 原语类似的通信方法。由于这种方法在延迟和带宽之间取得了良好的平衡,因此非常适合 LLM 推理。我们的自定义内核可直接从 GPU 显存读取通信大小并兼容 CUDA graph,即使数据大小在不同运行中变化也不例外。

我们的实现方式是使用 CUDA 的驱动程序 API 通过 MNNVL 建立点对点 (P2P) 缓冲区作为工作区。每个 GPU 都可以访问其他 GPU 的工作区。工作区被划分为多个通道,每个通道分配给远程 GPU 作为写入缓冲区。这些写入缓冲区以 FIFO 方式使用,通过标志同步 FIFO 状态以避免数据损坏。详细信息请参见 PR 3504

EP负载均衡器

TensorRT-LLM 采取了一系列实现 EP 负载均衡的功能,包括几个关键组件:

Python 接口

Python 接口层提供了一个对用户友好的 PyTorch / Python 原生接口,用于访问 MoE负载均衡的实现,例如用于 GPU / CPU 同步逻辑的 Python 包装器 (wrapper)、在线数据统计采集。

C++ 扩展

C++ 扩展作为 PyTorch / Python 接口与 C++ / CUDA 核心实现之间的桥梁。

Host 逻辑核心实现

Host 侧核心逻辑实现了以下关键部分:

  • 负载平衡算法
  • 复制算法
  • 放置算法
  • MoE 权重更新的编排逻辑
  • MoE 权重更新逻辑

GPU 逻辑核心实现

GPU 核心逻辑包含以下组件:

  • 在线流量统计采集
  • 为降低 CPU-GPU 往返同步开销,我们选择在 GPU 侧实现在线流量统计逻辑。
  • 专家路由逻辑
  • 需增强 MoE 路由逻辑以适应动态 EP 平衡影响。

已实现 GPU / CPU 同步组件。详情请参见 PR 4384PR 4495

基于这些核心实用程序,TensorRT-LLM 中提供了两种版本的 EP 负载均衡器:离线 EP 负载均衡器和在线 EP 负载均衡器。

在线 EP 负载均衡器

对于生产部署需求,推荐使用在线 EP 负载均衡器,因为它能够动态适应在线流量模式变化,从而更能保障性能。

但在线 EP 负载均衡器面临以下挑战。

首先,负载均衡引入了动态专家放置。单个专家的位置可能根据当前工作负载发生变化。例如,如果最初分配到 rank 0 的专家 0 和专家 1 均成为热门专家,负载均衡策略可能会将其重新分配到不同 rank,与冷门专家并列,这就需要及时更新权重数据。

我们旨在使在线负载均衡器能够快速响应请求模式的变化,并调整专家分配以避免负载失衡问题。重要的是,我们不希望平衡过程中干扰在线推理执行过程,也不希望采用“批量”策略更新权重。

大型 MoE 模型 (如 DeepSeek R1) 解码阶段的批处理大小通常较小,这使得 CUDA graph 成为一种有效的加速方法,尤其是在需要高单位用户 TPS 的情况下。这种优势在 Grace Blackwell 机架式系统等平台上表现得更为明显,因此我们希望整个负载均衡机制与 CUDA graph 兼容。

为避免无效化预捕获的 CUDA graph,我们通过将新专家权重写入相同显存位置的方式进行就地权重更新,而非替换张量指针。这确保了模型引擎中权重张量的地址保持不变。

在此设计中,每个专家槽位作为容器用于存储专家权重,与特定专家解耦。专家槽位的数量必须大于或等于专家总数,以确保每个专家始终至少有一个可用槽位。热门专家可能占用多个槽位。每个槽位通过 SlotId 进行标记。

由于 MoE 模型的路由逻辑输出 ExpertId(而非 SlotId),我们维护了一个从 ExpertId 到 SlotId 的路由表,该表由负载均衡策略定期更新。负载均衡路由模块使用当前路由表(专家复制信息和槽位)将每个 Token 映射到合适的专家槽位。

为实现权重更新的非阻塞处理并避免“停止世界 (Stop-The-World)”现象,我们采用分层更新的策略。在某一层的前向传播完成后,我们在其下一次前向传播开始前执行该层的权重平衡操作;若本次迭代中发生权重更新,则同一层的下一次前向传播需等待更新完成后再进行。

由于前向执行通常由单个 Python 线程调用一系列 PyTorch 操作驱动,我们将权重更新例程卸载到后台 C++ 线程。Python 侧仅初始化专家槽位,并在共享 Host 显存中注册专家权重。

在前向执行期间,我们在 MoE 计算的前后插入轻量化锁定 / 解锁内核,以及用于统计数据采集和将 SlotId 分配到 ExpertId 的内核。这些内核必须短且支持重叠,以尽量减少性能影响。只要 CPU 权重更新线程能及时完成工作,锁定 / 解锁操作将非常短暂。除路由内核外,所有内核均为轻量级且可轻松与不同 CUDA 流中的前向内核重叠,而路由内核是主要优化重点。

我们使用 MNNVL 进行专家调度和合并时的跨 GPU 通信。专家权重存储在 Host 显存中,并通过 C2C 传输至 GPU 显存以支持异步更新。多线程 Host 复制引擎管理此过程,自动检测 NUMA 拓扑并选择最优 CPU 核,实现与模型前向传递的完全异步。

对于不支持 C2C 但带有 PCIe 的服务器,若需跨节点通信,网络和权重更新可能会争夺 PCIe 带宽,需进行额外调优和设计考量。我们尚未为 PCIe 服务器实现复制引擎,该功能已列入未来任务清单。

离线 EP 负载均衡器

在线 EP 均衡器更适合部署在生产环境,可及时响应在线流量变化。但离线 EP 均衡器提供了一种适合性能研究 / 调试和验证的轻量级方式。您可以参考此 PR 了解离线 EP 负载均衡器的详细实现方式。此外,我们还提供了一个工具用于采集专家激活分布统计数据,可作为推导 EP 负载均衡放置策略的输入。您可参考此文档了解更多细节,以及如何通过端到端方式运行离线 EP 负载均衡器。

端到端评估

EP 负载均衡器的作用

  • 离线 EP 负载均衡器

如图 1 所示,在机器翻译数据集上,MoE 第 36 层存在严重的专家负载失衡问题,因此我们使用该层展示 EP 负载均衡器的效果。我们仍在 32 个 GPU 上运行 DeepSeek-R1 并采用 32 路专家并行。

16. 36 层各接收 rankx 轴)和迭代次数(y 轴)的传送 Token 数(未启用 EPLB
17. 36 层各专家(x 轴)和迭代次数(y 轴)的传送 Token 数(未启用 EPLB

图 16 显示了 50 次迭代中各接收 rank 收到的 Token 数,这可以代表每个 rank 的工作负载。rank 13 接收的 Token 数量明显多于其他所有 rank,且这种失衡的工作负载分布在整个迭代过程中几乎保持不变。图 17 将工作负载细分至专家。显然,rank 13 上的两个热门专家导致该 rank 上的压力过大。

基于上述统计数据,我们可以执行离线 EPLB。可行的策略之一是保持 32 路专家并行,同时将每个 rank 的专家槽位从 8 个增加到 9 个。这将产生 32 个冗余专家和总共 288 个专家槽位。图 18 和图 19 显示了 EPLB 后的传送 Token 数。显然,每个 rank 的 Token 分布更加平衡,且不再存在热门专家。

18. 36 层各接收 rankx 轴)和迭代次数(y 轴)的传送 Token 数(启用 EPLB,每个 rank 9 个专家槽位,EP 32 路)
19. 36 层各专家(x 轴)和迭代次数(y 轴)的传送 Token 数(启用 EPLB,每个 rank 9 个专家槽位,EP 32 路)

另一项 EPLB 策略是保持每 rank 8 个专家槽位,同时将专家并行度提升至 36 路。该策略同样产生 32 个冗余专家和总计 288 个专家槽位。如图 20 和 21 所示,工作负载在各 rank 或专家槽位之间也实现了平衡。

20. 36 层各接收 rankx 轴)和迭代次数(y 轴)的传送 Token 数(启用 EPLB,每个 rank 8 个专家槽位,EP 36 路)
21. 36 层各专家(x 轴)和迭代次数(y 轴)的传送 Token 数(启用 EPLB,每个 rank 8 个专家槽位,EP 36 路)

每个层和迭代可通过简单的指标(如标准差或失衡比率)测量负载失衡程度。基于所有 rank(或专家)传送 Token 数的情况下,失衡比率的计算方式为(最大值 − 均值) / 均值,表示最热门 rank(或专家)接收的过多工作负载。完美平衡负载的失衡比率为 0。

表 1 列出了上述情况下的标准差和失衡比率。每个数值均为各层各迭代指标的平均值。未启用 EPLB 时,负载失衡严重 —— 最热门 rank 接收的传送 Token 数平均比整体均值多 1.56 倍。EPLB 可有效降低负载失衡 —— 最热门 rank 接收的传送 Token 数平均仅比均值多约 0.11 倍。

1. 标准差与失衡比率(各层及各迭代指标的平均值)
  • 在线 EP 负载均衡器

在上一部分中,我们展示了离线 EP 负载均衡器的效果。接下来,我们将基于我们实现的在线 EP 负载均衡器,进一步分析在线环境下 EP 均衡的动态模式。我们仍以机器翻译数据集、DeepSeek R1 模型、第 36 层(如图 1 所示)为例来理解这一在线行为:

22. rank 0 发送到所有 rank Token 数(EP 32 路,本地 batch size=256256 个专家槽位(无复制),每个 rank 8 个专家)

从图 22 可以看出,从迭代 1963 开始,由于 EPLB 已生效,原来的最热门 rank 13 不再是最热门 rank,原先发送到 rank 13 的工作负载已被重新分配到 rank 0 和 rank 1。

在图 22 中,仅通过在线 EPLB 进行了放置调整。如下图所示,若进一步引入专家复制,可进一步提升平衡效果:

23. rank 0 发送到所有 rank Token数(EP 32 路,本地 batch size=256288 个专家槽位(有复制),每个 rank 9 个专家)

显然,在使用 EPLB 时引入专家复制可进一步提升 EP 平衡效果。可设计更复杂的实验,观察在线 EPLB 在在线服务过程中的定期效果,以动态平衡 EP 工作负载。我们欢迎社区向我们反馈任何有趣的 EPLB 模式观察结果。

性能研究

请注意:本部分中展示的所有代表性工作负载均来自 DeepSeek R1 推理执行中提取的性能追踪。我们仍在对端到端的性能进行调优 / 优化,并将在未来的技术博客中进行讨论。

我们将通过一些具有代表性的工作负载说明大规模 EP 对性能的影响。

24. EP MoE GroupGEMM EP 通信的影响

24. EP MoE GroupGEMM EP 通信的影响

如图 24 所示,当 EP 大小从 4 增加到 72 时,MoE GroupGEMM 的计算时间有所减少,而 EP 通信时间 (在 EP4 / EP8 时使用 Reduce / Scatter,在 EP > 8 时使用 All2All) 几乎保持不变。当 EP 大小从 18 增加到 72 时,加速幅度有所下降。我们正在优化这个方面。

接下来,我们将通过一些具有代表性的工作负载了解 EPLB 对性能的影响。

25. EPLB 对性能的影响

如图 25 所示,当 EP 大小增加时,EPLB 能够在 MoE GroupGEMM 和 EP 通信上都实现大幅性能提升。

复现步骤

本部分中所述的复现步骤所需的代码和脚本已合并到主分支。

EP 负载均衡器的效果

请参见 EP 负载均衡器示例,了解如何复现离线 EP 负载均衡器的结果。

  • 第 1 步:运行推理和采集统计数据

如要生成负载再平衡所需的统计数据,请在目标数据集上运行模型,并在推理过程中统计传送的专家 ID。计数过程完成后,统计数据将被保存供进一步处理。

设置一些环境变量:

export MODEL_NAME=deepseek-ai/DeepSeek-R1

export MODEL_PATH=<YOUR_MODEL_PATH>

# Set the expert statistic data path

export EXPERT_STATISTIC_PATH=./expert_statistic

# Enable counting of routed expert IDs from iteration 100 to iteration 200

export EXPERT_STATISTIC_ITER_RANGE=100-200


按照基准测试文档准备数据集并将其保存为 ./dataset.json。

在准备好的数据集上运行 32 路专家并行推理。有关在 Slurm 上运行 trtllm-bench 的详细信息,请参见 LLM API MGMN 示例

cat > ./extra_llm_api_options.yaml <<EOF

enable_attention_dp: true

EOF

trtllm-llmapi-launch \

trtllm-bench --model ${MODEL_NAME} \

    --model_path ${MODEL_PATH} \

    throughput \

    --tp 32 \

    --ep 32 \

    --extra_llm_api_options ./extra_llm_api_options.yaml \

    --kv_cache_free_gpu_mem_fraction 0.75 \

    --backend pytorch \

    --dataset ./dataset.json \

    --warmup 0 \

    --eos_id -1

推理完成后,查看 $EXPERT_STATISTIC_PATH 中的导出统计文件。运行 examples/ep_load_balancer/report_load_statistics.py 脚本,以显示标准差和失衡比率指标:

python examples/ep_load_balancer/report_load_statistics.py --expert_statistic_path $EXPERT_STATISTIC_PATH

输出如下:

Load statistics:

           mean         std  imbalance-ratio

3        1024.0  187.955200         0.498043

4        1024.0  202.728516         0.537602

5        1024.0  209.339981         0.458676

...

58       1024.0  570.880676         2.461014

59       1024.0  341.339447         0.717498

60       1024.0  381.045471         1.119648

average  1024.0  491.651199         1.564272
  • 第 2 步:生成 EPLB 配置

使用提供的 examples/ep_load_balancer/generate_eplb_config.py 脚本,将采集到的统计数据转换为 EPLB 配置文件。指定目标专家并行度大小 (–ep_size) 和部署时使用的总槽位数 (–num_slots)。例如,如果我们选择保持每个 rank 8 个专家槽位,同时将专家并行度增加到 36 路,则应有 32 个冗余专家和总共 288 个专家槽位。

python examples/ep_load_balancer/generate_eplb_config.py \

    --ep_size 36 \

    --num_slots 288 \

    --expert_statistic_path $EXPERT_STATISTIC_PATH \

    --output_path ./moe_load_balancer.yaml

./moe_load_balancer.yaml 文件如下:

initial_global_assignments:

  3: [138, 81, 60, ..., 69, 250, 77]

  4: [24, 243, 72, ..., 90, 251, 52]

  5: [120, 162, 246, ..., 14, 192, 171]

  ...

  58: [67, 70, 160, ..., 212, 103, 125]

  59: [45, 142, 152, ..., 99, 205, 49]

  60: [34, 162, 119, ..., 234, 26, 129]

num_slots: 288

layer_updates_per_iter: 0
  • 第 3 步:使用 EPLB 配置运行推理

设置一些环境变量:

# Set a new expert statistic data path

export EXPERT_STATISTIC_PATH=./expert_statistic_eplb

# Enable counting of routed expert IDs from iteration 100 to iteration 200

export EXPERT_STATISTIC_ITER_RANGE=100-200

使用 EPLB 配置运行 36 路专家并行推理:

cat > ./extra_llm_api_options_eplb.yaml <<EOF

enable_attention_dp: true

moe_config:

  load_balancer: ./moe_load_balancer.yaml

EOF

trtllm-llmapi-launch \

trtllm-bench --model ${MODEL_NAME} \

    --model_path ${MODEL_PATH} \

    throughput \

    --tp 36 \

    --ep 36 \

    --extra_llm_api_options ./extra_llm_api_options_eplb.yaml \

    --kv_cache_free_gpu_mem_fraction 0.75 \

    --backend pytorch \

    --dataset ./dataset.json \

    --warmup 0 \

    --eos_id -1

再次运行 examples/ep_load_balancer/report_load_statistics.py 脚本:

python examples/ep_load_balancer/report_load_statistics.py --expert_statistic_path $EXPERT_STATISTIC_PATH

输出如下:

Load statistics:

           mean        std  imbalance-ratio

3        1024.0  37.612328         0.081947

4        1024.0  42.367714         0.093256

5        1024.0  42.623219         0.092623

...

58       1024.0  49.167507         0.113420

59       1024.0  44.529514         0.092314

60       1024.0  48.408348         0.101029

average  1024.0  53.976442         0.115378

注意:统计专家 ID 会大幅影响性能,因此在进行基准测试或生产环境运行时,请记得通过取消设置 EXPERT_STATISTIC_ITER_RANGE 来禁用该功能。

其他

  • 共享显存清理:为实现在线负载平衡,所有专家权重都存储在共享 Host 内存中。同一节点上的 4 个 rank 共享相同的专家权重以节省显存。通常,这些共享 Host 显存在进程退出时会被清理。但如果发生异常退出,它们可能无法被清理。在这种情况下,可能需要手动检查 /dev/shm 目录并删除  /dev/shm/moe_shared_* 中的任何文件(如存在)。

拓展思考

我们高度认可 DeepSeek 团队的系统创新。他们将大规模 EP 支持引入内部推理系统,并以开放态度与社区分享工程洞察,这种做法和精神极其宝贵并且大大提升了推理系统设计性能。同时需要指出的是,系统设计与优化(如大规模 EP)并无万能解决方案。根据当前性能分析,当计划应用大规模 EP 时,需考虑以下因素:

  • MoE GroupGEMM 计算时间是否为端到端的性能瓶颈?
    • 大规模 EP 主要通过减轻专家权重负载压力缩短 MoE GroupGEMM 的执行时间,从而提升 MoE GroupGEMM 层的计算强度。如果在工作负载设置中 MoE GroupGEMM 计算不是瓶颈,那么大规模 EP 可能帮助不大。
  • 延迟限制
    • 大规模 EP 主要在存在严格延迟限制时发挥作用,尤其是在具有更大显存容量的 GPU 上。对于显存容量较小的 GPU,在延迟约束较弱的场景下,大规模 EP 仍可通过提升并发度和实现更高的 Token / 秒 / GPU 数发挥作用。
    • 可用的硬件规格
      • 大规模 EP 的优化配置取决于 GPU 规格,包括显存带宽、容量、GPU 间带宽及算力等。这些因素决定了是否采用大规模 EP 以及理想的并行度。
      • 系统复杂性和生产部署限制
        • 在没有故障容错保证的情况下,大规模 EP 可能会增加在线系统的故障率。当某些大规模 EP 服务实例故障时,即使可以通过集群级协调将流量传送到其他运行中的服务实例,大规模 EP 单实例部署所需的大量 GPU 也会增加系统级部署难度。

        未来,我们计划总结并分享更多大规模 EP 技术部署的最佳实践。

        请根据自身判断决定是否将大规模 EP 集成到系统中,以及在使用时确定何种 EP 规模和具体部署设置最适合您的需求。

        当前 TensorRT-LLM 大规模 EP 的实现方式尚不完善,仍存在已知的限制(欢迎至 TensorRT-LLM 社区帮助进行改进)。例如,我们需要:

        • 更广泛的平台支持
          • 扩展对其他 GPU 硬件的支持。
          • 目前大型 EP 支持仅覆盖 NVFP4 数据精度,需要开展进一步的工作将支持范围扩展到 FP8 和 INT8 / INT4 数据精度。
          • 性能
            • 进一步的性能调优和优化。
            • 进行更多接近生产流量的工作负载验证。我们期望TensorRT-LLM社区在这方面的反馈,以帮助我们基于更具体的工作负载校准 TensorRT-LLM 大规模 EP 实现。
            • 对与其他推理核心功能的组合进行全面验证,例如分离服务、推测性解码、验证更多 MoE 模型系列等。
                • 易用性
                  • 易于定制
                    • 我们认为大规模 EP 可至少分为两层:
                      • 由推理引擎开发者开发的内核层。该层包含自定义 EP 通信内核、CPU 与 GPU 之间的同步逻辑以及 MoE 权重重新分配逻辑。
                      • 策略层,可由推理引擎开发者与机器学习研究人员共同开发。该层包含用于以不同方法收集在线流量统计数据的工具,以及专家最佳复制与放置算法。
                    • 基于此理解,我们计划使接近策略层的组件更易于社区用户扩展和自定义,以鼓励涌现出更好的想法。
                  • 我们希望能够基于用户输入的部署要求 (ISL / OSL、延迟约束、硬件规格等),自动推荐最佳 EP 设置。
                • 容错性
                  • 由于大规模 EP 部署方案可能导致在线部署系统故障率增加,这可能会增加与 NVIDIA GPU 上端到端大语言模型 (LLM) 推理系统的多个组件进行跨层交互的需求,包括低级通信内核、集群级编排器和调度器等。我们正在与 NVIDIA 的多个工程团队积极合作,推动相关进展。

                  我们认为当前实现方式可作为一种合理的端到端大规模 EP 实现方案,并鼓励社区尝试新思路和性能验证。我们期待着社区能够提供反馈,帮助我们在该领域快速进步。我们正在积极追踪 GitHub 问题中的 TensorRT-LLM 大规模 EP 执行情况,以保证对社区的公开透明。

                  致谢

                  上述大规模 EP 工作是团队又一项巨大成果,涉及内核级优化、运行时增强以及系统性能分析和调优。我们无法逐一感谢每位贡献者,但我们为这支具有奉献精神的工程师团队感到骄傲,正是他们的集体专业知识推动了 TensorRT-LLM 性能的大幅提升。

                  此次团队合作为我们提供了提高大语言模型推理中 GPU 利用率的宝贵洞察。希望本文介绍的技术和经验能够帮助开发者社区在关键 LLM 推理应用中更加充分地利用 NVIDIA GPU 的性能。

                  标签