From 72d4d8315749c146cf19f047fecfec491911770d Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:09:19 -0700 Subject: [PATCH] [ROCm][Bugfix] Fix elastic EP scaling deadlock (#56610) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 4.8 (1M context) --- .../device_communicators/pynccl.py | 31 ++++++++++ .../distributed/elastic_ep/elastic_execute.py | 15 ++++- vllm/distributed/elastic_ep/standby_state.py | 37 ++++++----- vllm/distributed/parallel_state.py | 61 ++++++++++++------- 4 files changed, 106 insertions(+), 38 deletions(-) diff --git a/vllm/distributed/device_communicators/pynccl.py b/vllm/distributed/device_communicators/pynccl.py index 916d0308bf..ece4db82da 100644 --- a/vllm/distributed/device_communicators/pynccl.py +++ b/vllm/distributed/device_communicators/pynccl.py @@ -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) diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 343edd171f..19fc15b275 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -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()) diff --git a/vllm/distributed/elastic_ep/standby_state.py b/vllm/distributed/elastic_ep/standby_state.py index 1892f3e794..c9b12f447d 100644 --- a/vllm/distributed/elastic_ep/standby_state.py +++ b/vllm/distributed/elastic_ep/standby_state.py @@ -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.""" diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index b24ff64d4d..0466779c34 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -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,