vLLM 中的弹性专家并行(Elastic Expert Parallelism)

11 分钟阅读
Itay Alroy (NVIDIA), Yongji Wu (Sky Computing), Rui Qiao (Anyscale), Tyler Michael Smith (Red Hat), Moein Khazraee (NVIDIA), Omri Kahalon (NVIDIA), Tzu-Ling Kan (NVIDIA), Ron Tourgeman (NVIDIA)

专家并行(EP)是实现高吞吐量混合专家模型(MoE)服务的关键技术。广域专家并行(即 EP 跨越多个工作节点)可以最大化 KV 缓存容量,从而支持极高的并发量或极长的上下文。这对于强化学习工作负载(既需要长上下文又需要高吞吐量)以及代理(Agentic)工作负载(多轮对话会拉长上下文长度)尤为重要。

在 vLLM 中,与其他许多推理框架一样,EP 曾经是静态的:一旦部署开始,其服务容量就是固定的。如果请求量超过了该容量,vLLM 无法扩容以满足需求。如果需求下降,也无法缩容以减少 GPU 使用量和成本。唯一可行的方案是使用新配置进行全量重启,这不仅速度慢,还可能导致大量流量丢失。

弹性专家并行 (Elastic EP) 改变了这一现状。它允许 vLLM 在运行时重新配置工作节点数量,因此 MoE 部署可以随着需求变化进行扩容或缩容,且对服务的影响极小。

Elastic EP 通过增加或减少数据并行 (DP) 工作节点来进行扩展。在 vLLM 中,这将改变共享专家并行 (EP) 组的大小以及专家在工作节点间的分布方式,正如我们在背景部分所解释的那样。只需调用一个 API 即可实现:

curl -X POST https://:8000/scale_elastic_ep \
  -H "Content-Type: application/json" \
  -d '{"new_data_parallel_size": 8}'

此 API 调用将正在运行的部署从当前 DP 大小调整为 8 个工作节点。

本文介绍了 vLLM 中的 Elastic EP(RFC #20323, PR #34861),包括扩容和缩容流程、vLLM 如何协调重配置与正在进行的请求执行、该功能如何与 EPLB 及 EP 通信后端交互,以及为什么这项工作对于 vLLM 正在发展的容错方向至关重要。文中还讨论了 NIXL EP(PR #35627)作为一种通信模型,其特别适合弹性重配置和容错。

运维人员须知 (TL;DR)

  • Elastic EP 允许 vLLM 在运行时通过更改 DP 大小来扩容或缩容 MoE 部署,而无需重启服务器。
  • 您可以使用 POST /scale_elastic_ep 触发调整大小;vLLM 会重新配置实时拓扑并根据需要重新分配专家。
  • 这种运行时重配置路径是 vLLM 实现容错服务的核心构建模块。
  • NIXL EP 可以显著减少缩放事件期间的重新初始化工作,并提供 EP 端的故障检测、上报和恢复能力。

背景:专家并行与数据并行注意力

在 MoE 模型中,注意力层保持密集,而大多数前馈层被替换为稀疏的专家层,将每个 Token 路由到选定的一组专家。在深入研究弹性扩展之前,了解 Elastic EP 所基于的两种并行策略非常有帮助。

数据并行 (DP) 注意力使用请求级并行:每个引擎核心处理不同的请求分片,并维护自己的 KV 缓存和调度器。这在 MLA 等架构中特别有用,因为张量并行 (TP) 会导致 KV 缓存跨 GPU 重复,从而浪费内存并限制批处理大小。

专家并行 (EP) 用于专家层。专家不再跨 GPU 进行分片,而是分布在不同的 GPU 上,Token 仅被分发到拥有所选专家的 GPU。

在 vLLM 中,注意力在每个 DP 工作节点上独立运行,而专家层在这些工作节点间共享一个 EP 组(EP 组大小为 DP x TP)。Elastic EP 在运行时更改 DP 工作节点的数量,从而相应地缩放 EP 组并在其中重新分配专家。

挑战:需要更改哪些状态?

在运行时缩放 DP 不仅仅是启动或终止进程的问题。EP 大小的变化会使多个运行时状态失效:

  • 分布式通信组。EP、DP 和全局组都嵌入了一组固定的 Rank。
  • 专家分配。当 EP 大小发生变化时,专家到 Rank 的映射关系会发生改变。
  • 模型权重。新 Rank 需要模型权重,现有 Rank 在重新分配后可能需要更新专家权重。
  • CUDA 图和编译状态。CUDA 图捕获和 torch.compile 都是基于拓扑发生变化时会失效的假设进行优化的。

因此,该实现将缩放视为一个协调的状态机。每个阶段都有明确的同步点,这些同步点必须与模型的前向执行安全共存。

扩容流程

DP=N 扩容到 DP=M(其中 M > N)比缩容更复杂,因为需要将新的 Rank 纳入正在运行的部署中。

1. 触发与请求处理

操作从 /scale_elastic_ep 开始。如果设置了 VLLM_ELASTIC_EP_DRAIN_REQUESTS=1,vLLM 首先会等待正在处理的工作排空,最多等待 drain_timeout 秒(默认 120 秒)。否则,将立即开始缩放。

2. 新引擎核心初始化

启动新的引擎核心工作节点依赖于 Ray DP 后端。在扩容期间,Ray DP 后端会在当前可用的 GPU 上启动目标 DP 大小所需的额外 DP 工作节点。新 Rank 接收当前的专家映射,并使用占位符权重初始化模型。然后,它们等待后续的传输和重配置阶段,将其纳入活跃拓扑中。

就绪状态通过两个阶段进行协调:一个信号允许现有 Rank 创建待机组,另一个信号则允许开始权重传输。

3. 待机通信组

一个关键的设计选择是,vLLM 不会立即拆除活跃的通信组。相反,现有 Rank 首先创建跨越目标 Rank 集合的待机组。这些组通过 StatelessGroupCoordinator 创建,它独立于 PyTorch 全局的 WORLD 状态。

这使得在切换前准备新配置成为可能,同时旧配置在此期间仍可执行前向计算。

使用 nixl_ep,此过渡可以是增量的:vLLM 不必拆除并重新创建所有 EP 端连接,而是可以通过 NIXL EP 的 connect_ranks() / disconnect_ranks() API 添加或移除 Rank,同时保持现有连接不受影响。

4. 专家映射与权重迁移

一旦待机组存在,我们就利用它们来广播当前的专家映射,并将非专家权重从现有 Rank 传输到新 Rank,传输工作尽可能均匀地分布在现有 Rank 上。Elastic EP 重用了 EPLB 用于专家权重移动的 GPU 到 GPU 发送/接收路径,但也将其扩展到了注意力层、归一化层、嵌入层和其他非专家权重,利用节点内的 NVLink 或节点间的 RDMA 等高速互联。

此阶段不移动专家权重。它们将在新拓扑激活后由 EPLB 传输。普通的 EPLB 活动在过渡期间会暂停,以免干扰重配置。

5. 切换

切换点是所有 Rank 停止使用旧拓扑并开始使用新拓扑的时刻。在此阶段,vLLM 会:

  1. 释放 CUDA 图并重置 torch.compile 状态。
  2. 将待机组提升为活跃的 EP、DP 和全局组。
  3. 销毁旧组。
  4. 为新的 EP 大小重新配置 MoE 模块。
  5. 预热模型,使 CUDA 图和编译路径与新设置匹配。

引擎协调状态(如运行标志、波次计数器和步骤计数器)在新的 DP 组中同步,以便每个 Rank 都能从一致的点恢复。

此时,新 Rank 已成为活跃 DP 组的一部分,可以参与前向计算并运行注意力,但它们尚不拥有专家。专家归属权将在随后的 EPLB 重洗牌中更新。

6. EPLB 重洗牌

随着新拓扑处于活跃状态,EPLB 在所有 M 个 Rank 上重新分配专家。这会更新专家映射并执行新布局所需的专家权重迁移。普通的 EPLB 操作在重洗牌完成后恢复。

缩容流程

DP=M 缩容到 DP=N 遵循与扩容相同的基本模式,但有一个重要区别:EPLB 重洗牌必须首先发生。即将被移除的 Rank 可能仍拥有专家权重,因此所有 M 个引擎核心首先参与重洗牌,将专家整合到 N 个幸存的 Rank 上,并将任何需要的专家权重从即将退出的 Rank 上迁移走。

跨 DP Rank 协调重配置步骤

一个细微的问题是 DP 引擎核心是异步运行的,因此它们接收重配置通知的时间可能略有不同。当某些 Rank 到达下一个 Elastic EP 阶段时,其他 Rank 可能已经开始了一个新的前向步骤。如果先行到达的 Rank 立即进行,组内就会出现一部分在重配置、一部分在执行前向计算的情况,从而导致部署死锁。

Elastic EP 使用两阶段屏障来处理这个问题。第一个屏障使用超时:如果未能及时完成,已到达的 Rank 会推断出某些对等节点已经进入了下一个引擎步骤,因此它们也会返回引擎循环再执行一次迭代,而不是独自推进。在下一次迭代中,一旦所有 Rank 到达相同的边界,第二个不带超时路径的屏障将让它们一起进入下一个阶段。

迈向容错性

Elastic EP 是容错的一个核心构建模块,因为它为 vLLM 提供了故障后所需的运行时重配置路径。如果一个 Rank 死亡,Elastic EP 提供了移除该 Rank、重新分配其专家、并在之后无需重启整个部署即可增加替换容量所需的缩容和扩容路径。这是 RFC #30112 中讨论的更广泛容错方向的一部分。

从高层次来看,恢复流程如下:

  1. 检测:通过健康检查或后端特定的故障信号检测故障。
  2. 缩容:移除故障 Rank 并重新分配其专家。
  3. 扩容:一旦有了替换容量,再次进行扩容。

NIXL EP 在这里也很有意义,因为它可以在 EP 端检测、报告和恢复故障,并在容量再次可用时重新连接替换 Rank。

后续步骤

Elastic EP 已经提供了核心的运行时重配置路径,但当前的实现范围仍然相当具体,并且有几个明显的后续方向:

  • 支持更丰富的并行配置。包括 tensor_parallel_size > 1 和其他并行配置。
  • 支持更多服务功能。当前实现将 api_server_count 限制为 1,且尚不支持 DBO 或 MoE 草稿/起草模型。
  • 减少重配置窗口。在重叠、预热成本、CUDA 图重新捕获和重用先前准备的状态方面仍有工作要做。
  • 将 Elastic EP 连接到自动缩放策略。运行时控制平面已经具备,但策略和编排是独立的工作(Dynamo, llm-d)。
  • 支持额外的 DP 后端。扩展操作目前依赖于 Ray DP 后端。

入门指南

启动并启用弹性专家并行 (Elastic EP)

下面的示例使用 DeepSeek-V2-Lite-Chat 作为 MoE 示例。当前实现针对的是 tensor_parallel_size=1、单 API 服务器且无 DBO 的 Ray DP 部署。

vllm serve deepseek-ai/DeepSeek-V2-Lite-Chat \
    --trust-remote-code \
    --tensor-parallel-size 1 \
    --data-parallel-size 2 \
    --data-parallel-backend ray \
    --api-server-count 1 \
    --enable-expert-parallel \
    --enable-elastic-ep \
    --enable-eplb \
    --eplb-config.num_redundant_experts 0 \
    --all2all-backend allgather_reducescatter \
    --gpu-memory-utilization 0.8

运行时扩容

使用 Ray DP 后端,增加容量可以简单到将另一个节点加入 Ray 集群;一旦 Ray 发现新 GPU,Elastic EP 就可以在运行时将部署扩展到这些 GPU 上。

例如,在新工作节点上:

ray start --address="${HEAD_NODE_IP}:6379"
curl -X POST https://:8000/scale_elastic_ep \
  -H "Content-Type: application/json" \
  -d '{"new_data_parallel_size": 16}'

缩容

curl -X POST https://:8000/scale_elastic_ep \
  -H "Content-Type: application/json" \
  -d '{"new_data_parallel_size": 8}'

使用 NIXL EP 作为通信后端

如果您想将 NIXL EP 与 Elastic EP 一起使用:

uv pip install nixl
 
vllm serve deepseek-ai/DeepSeek-V2-Lite-Chat \
    --trust-remote-code \
    --tensor-parallel-size 1 \
    --data-parallel-size 2 \
    --data-parallel-backend ray \
    --api-server-count 1 \
    --enable-expert-parallel \
    --enable-elastic-ep \
    --enable-eplb \
    --all2all-backend nixl_ep

请参阅 NIXL 仓库了解安装详情和传输配置。

参考文献

致谢

感谢所有为 vLLM 引入 Elastic EP 做出贡献的人。

  • Sky Computing: Yongji Wu
  • NVIDIA: Itay Alroy, Moein Khazraee, Omri Kahalon, Tzu-Ling Kan, Ron Tourgeman
  • Red Hat: Tyler Michael Smith
  • Anyscale: Rui Qiao
  • 更广泛的 vLLM 社区