iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Bugfix] Fix elastic EP scaling deadlock (#56610)

Browse Source 
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Matt
2026-09-12 23:09:19 +00:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 7ee8a6dd01
commit 72d4d83157
4 changed files with 106 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/distributed/device_communicators/pynccl.py
+31
View File
@@ -3,6 +3,7 @@
# ===================== import region =====================
import contextlib
import threading
import torch
@@ -26,6 +27,33 @@ from vllm.utils.torch_utils import current_stream
logger = init_logger(__name__)
_warmup_tls = threading.local()
def _warmup_deferred() -> bool:
return getattr(_warmup_tls, "deferred", False)
@contextlib.contextmanager
def defer_comm_warmup_on_rocm():
"""Build communicators on this thread without their warm-up on ROCm, where
the warm-up can deadlock if built next to a live engine whose existing
communicators still have collectives in flight; instead the warm-up is
deferred to when the engine is quiet (during commit).
"""
from vllm.platforms import current_platform
if not current_platform.is_rocm():
yield
return
prev = _warmup_deferred()
_warmup_tls.deferred = True
try:
yield
finally:
_warmup_tls.deferred = prev
_NCCL_SYMM_OPS_REGISTERED = False
_NCCL_SUSPEND_MEM = 0x01
@@ -153,6 +181,9 @@ class PyNcclCommunicator:
self.world_size, self.unique_id, self.rank
)
if _warmup_deferred():
return
stream = current_stream()
# A small all_reduce for warmup.
data = torch.zeros(1, device=device)
vllm/distributed/elastic_ep/elastic_execute.py
+13 -2
View File
@@ -50,6 +50,7 @@ from vllm.model_executor.warmup.flashinfer_autotune_cache import (
sync_flashinfer_autotune_cache,
)
from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
from vllm.platforms import current_platform
from vllm.utils import is_moe_layer
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
from vllm.v1.worker.dp_utils import skip_dp_coordination
@@ -283,7 +284,9 @@ class ElasticEPScalingExecutor:
self._prepare_eplb_communicator(get_standby_eplb_group())
if new_dp_size > old_dp_size:
self.transfer_weights(old_dp_size, new_dp_size)
self._warm_target_groups(get_standby_dp_group(), standby_ep_group)
# On ROCm this deadlocks while serving; warm_and_capture warms these at commit.
if not current_platform.is_rocm():
self._warm_target_groups(get_standby_dp_group(), standby_ep_group)
if new_dp_size > old_dp_size and self._can_reuse_fused_moe_kernel():
target_world_group = get_standby_world_group()
assert target_world_group is not None
@@ -577,7 +580,11 @@ class ElasticEPScalingExecutor:
mapping = self.receive_expert_mapping()
self.worker.model_runner.setup_eplb_from_mapping(mapping)
if not self._can_reuse_fused_moe_kernel():
# warm_and_capture's dummy run warms the deferred dp/ep groups.
self.warm_and_capture()
elif current_platform.is_rocm():
# Reuse path skips warm_and_capture; warm the deferred groups here.
self._warm_target_groups(get_dp_group(), get_ep_group())
self._perform_eplb_reshuffle(async_op=True)
if is_existing_worker:
self._start_group_cleanup(retired_groups)
@@ -590,6 +597,9 @@ class ElasticEPScalingExecutor:
retired_groups = self.switch_and_prepare()
if not self._can_reuse_fused_moe_kernel():
self.warm_and_capture()
elif current_platform.is_rocm():
# Reuse path skips warm_and_capture; warm the deferred groups here.
self._warm_target_groups(get_dp_group(), get_ep_group())
self._start_group_cleanup(retired_groups)
def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None:
@@ -646,7 +656,8 @@ class ElasticEPScalingExecutor:
expert_weights=expert_weights,
)
torch.accelerator.synchronize()
self._warm_target_groups(get_dp_group(), get_ep_group())
if not current_platform.is_rocm():
self._warm_target_groups(get_dp_group(), get_ep_group())
if self._can_reuse_fused_moe_kernel():
sync_flashinfer_autotune_cache(self.worker.model_runner, get_world_group())
vllm/distributed/elastic_ep/standby_state.py
+23 -14
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm
from vllm.distributed.parallel_state import (
_init_stateless_group,
_node_count,
@@ -78,27 +79,35 @@ def create_standby_groups(
)
standby_dp_ranks = all_ranks.transpose(1, 3).reshape(-1, new_dp_size).unbind(0)
standby_dp_ranks = [x.tolist() for x in standby_dp_ranks]
_STANDBY_DP = _init_stateless_group(
standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store
)
standby_ep_ranks = (
all_ranks.transpose(1, 2).reshape(-1, new_dp_size * tp_size).unbind(0)
)
standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
_STANDBY_EP = _init_stateless_group(
standby_ep_ranks, "ep", master_ip, backend, coord_store, use_all2all=use_all2all
)
# Deferred to commit so the warm-up runs while the engine is paused.
with defer_comm_warmup_on_rocm():
_STANDBY_DP = _init_stateless_group(
standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store
)
if enable_eplb:
_STANDBY_EPLB = _init_stateless_group(
standby_ep_ranks = (
all_ranks.transpose(1, 2).reshape(-1, new_dp_size * tp_size).unbind(0)
)
standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
_STANDBY_EP = _init_stateless_group(
standby_ep_ranks,
"eplb",
"ep",
master_ip,
backend,
coord_store=coord_store,
coord_store,
use_all2all=use_all2all,
)
if enable_eplb:
_STANDBY_EPLB = _init_stateless_group(
standby_ep_ranks,
"eplb",
master_ip,
backend,
coord_store=coord_store,
)
def pop_standby_groups() -> dict:
"""Return all standby groups and clear the standby state."""
vllm/distributed/parallel_state.py
+39 -22
View File
@@ -1919,6 +1919,20 @@ def init_distributed_environment(
_INNER_DP_WORLD = _WORLD
def _elastic_join_warmup_ctx() -> AbstractContextManager[Any]:
"""Defer communicator warm-up when joining a live elastic-EP cluster on ROCm.
The existing ranks defer their matching warm-up in create_standby_groups(),
so a joining worker must too or the warm-up collective has no peers. Both
sides warm up at commit.
"""
from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm
if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
return defer_comm_warmup_on_rocm()
return nullcontext()
def initialize_model_parallel(
tensor_model_parallel_size: int = 1,
pipeline_model_parallel_size: int = 1,
@@ -2101,13 +2115,14 @@ def initialize_model_parallel(
group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0)
group_ranks = [x.tolist() for x in group_ranks]
if enable_elastic_ep:
_DP = _init_stateless_group(
group_ranks,
"dp",
parallel_config.data_parallel_master_ip,
backend,
coord_store=coord_store,
)
with _elastic_join_warmup_ctx():
_DP = _init_stateless_group(
group_ranks,
"dp",
parallel_config.data_parallel_master_ip,
backend,
coord_store=coord_store,
)
else:
_DP = init_model_parallel_group(
group_ranks, get_world_group().local_rank, backend, group_name="dp"
@@ -2130,14 +2145,15 @@ def initialize_model_parallel(
group_ranks = [x.tolist() for x in group_ranks]
use_all2all = parallel_config.use_all2all
if enable_elastic_ep:
_EP = _init_stateless_group(
group_ranks,
"ep",
parallel_config.data_parallel_master_ip,
backend,
coord_store=coord_store,
use_all2all=use_all2all,
)
with _elastic_join_warmup_ctx():
_EP = _init_stateless_group(
group_ranks,
"ep",
parallel_config.data_parallel_master_ip,
backend,
coord_store=coord_store,
use_all2all=use_all2all,
)
else:
_EP = init_model_parallel_group(
group_ranks,
@@ -2155,13 +2171,14 @@ def initialize_model_parallel(
assert _EPLB is None, "EPLB group is already initialized"
if config.parallel_config.enable_eplb:
if enable_elastic_ep:
_EPLB = _init_stateless_group(
group_ranks,
"eplb",
parallel_config.data_parallel_master_ip,
backend,
coord_store=coord_store,
)
with _elastic_join_warmup_ctx():
_EPLB = _init_stateless_group(
group_ranks,
"eplb",
parallel_config.data_parallel_master_ip,
backend,
coord_store=coord_store,
)
else:
_EPLB = init_model_parallel_group(
group_ranks,