From a2257f95b79287933defef5be40b74ce93053da3 Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Wed, 19 Aug 2026 05:13:06 +0300 Subject: [PATCH] [Elastic EP] Reduce eager-mode reconfiguration downtime (#51885) Signed-off-by: Itay Alroy --- tests/distributed/test_eplb_execute.py | 122 ------------- .../worker/test_gpu_model_runner_v2_eplb.py | 3 +- vllm/config/parallel.py | 2 +- .../distributed/elastic_ep/elastic_execute.py | 162 ++++++++++++------ vllm/distributed/elastic_ep/elastic_state.py | 49 ++---- vllm/distributed/eplb/eplb_communicator.py | 48 ++---- vllm/distributed/eplb/eplb_state.py | 71 +++++--- vllm/distributed/parallel_state.py | 14 +- vllm/v1/worker/gpu/eplb_utils.py | 2 - vllm/v1/worker/gpu/model_runner.py | 2 - vllm/v1/worker/gpu_model_runner.py | 20 +-- vllm/v1/worker/gpu_worker.py | 2 + 12 files changed, 194 insertions(+), 303 deletions(-) diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 4c9b98b62c..21fa057fd2 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -782,125 +782,3 @@ def test_rearrange_expert_weights_profile_mode(world_size): _test_rearrange_expert_weights_profile_mode, world_size, ) - - -def _test_nixl_deferred_init_worker( - env, - world_size: int, - num_layers: int, - num_local_experts: int, - num_logical_experts: int, -) -> None: - """Exercise NixlEplbCommunicator with defer_remote_setup=True (elastic EP path).""" - from vllm.distributed.eplb.eplb_communicator import NixlEplbCommunicator - - set_env_vars_and_device(env) - - vllm_config = VllmConfig() - vllm_config.parallel_config.tensor_parallel_size = world_size - - with set_current_vllm_config(vllm_config): - ensure_model_parallel_initialized( - tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1 - ) - - ep_group_coordinator = get_tp_group() - ep_group = ep_group_coordinator.cpu_group - ep_rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{ep_rank}") - - total_physical_experts = world_size * num_local_experts - hidden_sizes = [32, 64] - - redundancy_config = create_redundancy_config( - num_logical_experts, total_physical_experts - ) - old_indices = create_expert_indices_with_redundancy( - num_layers, - num_logical_experts, - total_physical_experts, - redundancy_config, - ) - - new_redundancy_config = create_redundancy_config( - num_logical_experts, total_physical_experts - ) - new_indices = create_expert_indices_with_redundancy( - num_layers, - num_logical_experts, - total_physical_experts, - new_redundancy_config, - ) - - expert_weights = create_expert_weights( - num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices - ) - - expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] - - communicator = NixlEplbCommunicator( - cpu_group=ep_group_coordinator.cpu_group, - all_expert_weights=expert_weights, - expert_buffer=expert_buffer, - defer_remote_setup=True, - ) - assert not communicator._remote_state_initialized - - rearrange_expert_weights_inplace( - old_indices, - new_indices, - expert_weights, - expert_buffer, - ep_group, - communicator, - ) - - assert communicator._remote_state_initialized - - local_ok = verify_expert_weights_after_shuffle( - expert_weights, - new_indices, - hidden_sizes, - ep_rank, - num_local_experts, - ) - - local_ok = ( - verify_redundant_experts_have_same_weights( - expert_weights, - new_indices, - hidden_sizes, - ep_rank, - world_size, - num_local_experts, - ) - and local_ok - ) - assert_verification_synced( - local_ok, - "Deferred NIXL init verification failed on at least one rank.", - ) - - -@pytest.mark.skipif(not has_nixl(), reason="NIXL is not available") -@pytest.mark.parametrize( - "world_size,num_layers,num_local_experts,num_logical_experts", - [(2, 2, 3, 4)], -) -def test_nixl_deferred_init( - world_size, - num_layers, - num_local_experts, - num_logical_experts, -): - """Test NixlEplbCommunicator with defer_remote_setup=True (elastic EP path).""" - - if torch.accelerator.device_count() < world_size: - pytest.skip(f"Need at least {world_size} GPUs to run the test") - distributed_run( - _test_nixl_deferred_init_worker, - world_size, - num_layers, - num_local_experts, - num_logical_experts, - ) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index ebb4beb2a5..637d078c73 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -159,13 +159,12 @@ def test_v2_setup_eplb_from_mapping_rebuilds_state(monkeypatch): runner = _make_runner(model=SimpleNamespace(is_moe=True)) mapping = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64) - mrv2.GPUModelRunner.setup_eplb_from_mapping(runner, mapping, 2) + mrv2.GPUModelRunner.setup_eplb_from_mapping(runner, mapping) assert runner.eplb_state is not None assert runner.eplb_state.built_from_mapping is True assert FakeEplbState.from_mapping_kwargs is not None assert FakeEplbState.from_mapping_kwargs["expanded_physical_to_logical"] is mapping - assert FakeEplbState.from_mapping_kwargs["num_valid_physical_experts"] == 2 def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index eb346ed4a3..d0118c6503 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -973,7 +973,7 @@ class ParallelConfig: if self.enable_eplb and self.eplb_config.communicator is None: # Prefer NIXL when available: zero-copy RDMA reads, compatible - # with both async EPLB and elastic EP (deferred remote setup). + # with both async EPLB and elastic EP. # Fallbacks: pynccl for elastic EP (stateless groups need it), # torch_gloo for static EP. torch_nccl is avoided because NCCL # is incompatible with async EPLB (multi-stream conflicts) and diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 99325a1706..6124d22b44 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -29,12 +29,15 @@ from vllm.distributed.elastic_ep.standby_state import ( create_standby_groups, get_standby_dp_group, get_standby_ep_group, + get_standby_eplb_group, pop_standby_groups, ) -from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator +from vllm.distributed.eplb.eplb_communicator import ( + EplbCommunicator, +) from vllm.distributed.parallel_state import ( + GroupCoordinator, _replace_active_groups, - get_eplb_group, ) from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator from vllm.logger import init_logger @@ -149,10 +152,12 @@ class ElasticEPScalingExecutor: self.worker_ref = weakref.ref(worker) self.reconfig_request = None self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {} + self._prepared_eplb_communicator: EplbCommunicator | None = None self._async_executor = ThreadPoolExecutor( max_workers=1, thread_name_prefix="ElasticEPAsync" ) self._async_future: Future[None] | None = None + self._group_cleanup_future: Future[None] | None = None @property def worker(self): @@ -207,12 +212,40 @@ class ElasticEPScalingExecutor: self._async_future = None future.result() + def _destroy_retired_groups( + self, groups: tuple[GroupCoordinator | None, ...] + ) -> None: + from vllm.platforms import current_platform + + current_platform.set_device(self.worker.device) + for group in groups: + if group is not None: + group.destroy() + + def _start_group_cleanup(self, groups: tuple[GroupCoordinator | None, ...]) -> None: + assert self._group_cleanup_future is None + self._group_cleanup_future = self._async_executor.submit( + self._destroy_retired_groups, groups + ) + + def _wait_for_group_cleanup(self) -> None: + if (future := self._group_cleanup_future) is not None: + self._group_cleanup_future = None + future.result() + + def shutdown(self) -> None: + try: + self._wait_for_group_cleanup() + finally: + self._async_executor.shutdown() + def load_model(self) -> None: self.worker.load_model(load_dummy_weights=True) - def create_standby_groups( + def prepare_reconfiguration( self, reconfig_request: ReconfigureDistributedRequest, use_all2all: bool ) -> None: + self._wait_for_group_cleanup() self.reconfig_request = reconfig_request new_dp_size = reconfig_request.new_data_parallel_size old_dp_size = get_dp_group().world_size @@ -227,8 +260,20 @@ class ElasticEPScalingExecutor: use_all2all=use_all2all, enable_eplb=parallel_config.enable_eplb, ) - if new_dp_size < old_dp_size: - self.stage_standby_moe_quant_methods() + self.stage_standby_moe_quant_methods() + 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(), get_standby_ep_group()) + + def _prepare_eplb_communicator(self, eplb_group) -> None: + assert eplb_group is not None + model_runner = self.worker.model_runner + eplb_state = model_runner.eplb_state + assert eplb_state is not None + self._prepared_eplb_communicator = eplb_state.create_communicator( + model_runner.model_config, eplb_group + ) def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None: standby_dp_group = get_standby_dp_group() @@ -276,6 +321,15 @@ class ElasticEPScalingExecutor: ) torch.accelerator.synchronize() + def _warm_target_groups(self, dp_group, ep_group) -> None: + assert dp_group is not None and ep_group is not None + stream = torch.Stream(device=dp_group.device) + with stream: + tensor = torch.zeros(1, dtype=torch.int32, device=dp_group.device) + for group in (dp_group, ep_group): + torch.distributed.all_reduce(tensor, group=group.device_group) + stream.synchronize() + def broadcast_expert_mapping(self) -> None: standby_dp_group = get_standby_dp_group() assert standby_dp_group is not None @@ -362,15 +416,22 @@ class ElasticEPScalingExecutor: torch.accelerator.empty_cache() def switch_and_remove(self) -> None: + # Removing ranks skipped preparation, so wait for prior cleanup here. + self._wait_for_group_cleanup() self._release_cuda_graphs() - _replace_active_groups(world=None, dp=None, ep=None, eplb=None, node_count=None) + retired_groups = _replace_active_groups( + world=None, dp=None, ep=None, eplb=None, node_count=None + ) + self._start_group_cleanup(retired_groups) + # Finish collective cleanup before this worker is shut down. + self._wait_for_group_cleanup() - def switch_and_prepare(self) -> None: + def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]: old_dp_size = get_dp_group().world_size old_ep_size = get_ep_group().world_size self._release_cuda_graphs() - _replace_active_groups(**pop_standby_groups()) + retired_groups = _replace_active_groups(**pop_standby_groups()) parallel_config = self.worker.vllm_config.parallel_config reconfig_request = self.reconfig_request @@ -453,7 +514,6 @@ class ElasticEPScalingExecutor: ) eplb_model_state.expert_load_pass = expanded_expert_load_pass eplb_model_state.expert_load_window = expanded_expert_load_window - eplb_state.num_valid_physical_experts = old_num_physical_experts else: assert pad_size < 0 eplb_model_state.expert_load_pass = eplb_model_state.expert_load_pass[ @@ -462,7 +522,6 @@ class ElasticEPScalingExecutor: eplb_model_state.expert_load_window = eplb_model_state.expert_load_window[ :, :, :num_physical_experts ] - eplb_state.num_valid_physical_experts = num_physical_experts model = self.worker.model_runner.get_model() model.expert_weights = [] @@ -485,18 +544,9 @@ class ElasticEPScalingExecutor: if getattr(module._quant_method, "wraps_legacy_quant_method", False): module._replace_quant_method(module._quant_method.old_quant_method) - eplb_model_state.expert_buffer = [ - torch.empty_like(w) for w in model.expert_weights[0] - ] - assert parallel_config.eplb_config.communicator is not None, ( - "EPLB communicator backend must be set by ParallelConfig" - ) - eplb_model_state.communicator = create_eplb_communicator( - group_coordinator=get_eplb_group(), - backend=parallel_config.eplb_config.communicator, - expert_weights=model.expert_weights, - expert_buffer=eplb_model_state.expert_buffer, - ) + assert self._prepared_eplb_communicator is not None + eplb_state.update_communicator(model_config, self._prepared_eplb_communicator) + self._prepared_eplb_communicator = None if ( self.worker.vllm_config.compilation_config.mode @@ -513,8 +563,12 @@ class ElasticEPScalingExecutor: compilation_counter.stock_torch_compile_count += 1 self.worker.model_runner.model.compile(fullgraph=True, backend=backend) + return retired_groups + def _perform_eplb_reshuffle( - self, rank_mapping: dict[int, int] | None = None + self, + rank_mapping: dict[int, int] | None = None, + async_op: bool = False, ) -> None: if get_ep_group().rank == 0: logger.info("[Elastic EP] Starting expert resharding...") @@ -522,21 +576,19 @@ class ElasticEPScalingExecutor: eplb_state = self.worker.model_runner.eplb_state assert eplb_state is not None - model_config = self.worker.model_runner.model_config - eplb_model_state = eplb_state.model_states[model_config.compute_hash()] is_async_enabled = eplb_state.is_async - eplb_state.is_async = False + run_async = async_op and is_async_enabled + eplb_state.is_async = run_async if rank_mapping is None: eplb_state.rearrange() else: eplb_state.rearrange(rank_mapping=rank_mapping) - # NOTE(yongji): check whether we need to synchronize here - torch.accelerator.synchronize() + if not run_async: + # Wait for non-blocking expert resharding copies before continuing + # the Elastic EP reconfiguration. + torch.accelerator.synchronize() # reset expert_rearrangement_step to ensure all ranks are synchronized eplb_state.expert_rearrangement_step = 0 - eplb_state.num_valid_physical_experts = ( - eplb_model_state.physical_to_logical_map.shape[1] - ) eplb_state.is_async = is_async_enabled # Start the async worker thread if it doesn't exist yet (idempotent). # This is needed for new workers after scale-up: they create EplbState @@ -544,25 +596,31 @@ class ElasticEPScalingExecutor: # groups aren't ready yet. eplb_state.start_async_loop() if get_ep_group().rank == 0: - logger.info("[Elastic EP] Expert resharding completed") + logger.info( + "[Elastic EP] Expert resharding %s", + "scheduled" if run_async else "completed", + ) def commit_scale_up(self, is_existing_worker: bool) -> None: if is_existing_worker: self.broadcast_expert_mapping() - self.switch_and_prepare() + retired_groups = self.switch_and_prepare() else: - mapping, _, num_valid_experts = self.receive_expert_mapping() - self.worker.model_runner.setup_eplb_from_mapping(mapping, num_valid_experts) - self._perform_eplb_reshuffle() + mapping = self.receive_expert_mapping() + self.worker.model_runner.setup_eplb_from_mapping(mapping) self.warm_and_capture() + self._perform_eplb_reshuffle(async_op=True) + if is_existing_worker: + self._start_group_cleanup(retired_groups) def commit_scale_down(self, new_dp_size: int, removing: bool) -> None: self.perform_scale_down_eplb_reshuffle(new_dp_size) if removing: self.switch_and_remove() else: - self.switch_and_prepare() + retired_groups = self.switch_and_prepare() self.warm_and_capture() + self._start_group_cleanup(retired_groups) def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: eplb_state = self.worker.model_runner.eplb_state @@ -578,7 +636,7 @@ class ElasticEPScalingExecutor: } self._perform_eplb_reshuffle(rank_mapping=rank_mapping) - def receive_weights(self) -> None: + def prepare_new_worker(self) -> None: dp_group = get_dp_group() assert isinstance(dp_group, StatelessGroupCoordinator) new_dp_size = dp_group.world_size @@ -618,19 +676,18 @@ class ElasticEPScalingExecutor: expert_weights=expert_weights, ) torch.accelerator.synchronize() + self._warm_target_groups(get_dp_group(), get_ep_group()) - def receive_expert_mapping(self) -> tuple[torch.Tensor, int, int]: + def receive_expert_mapping(self) -> torch.Tensor: dp_group = get_dp_group() assert isinstance(dp_group, StatelessGroupCoordinator) - physical_to_logical, num_local_physical_experts, num_logical_experts = ( - broadcast_expert_mapping( - physical_to_logical=None, - num_local_physical_experts=None, - num_logical_experts=None, - dp_group=dp_group, - src_rank=0, - device=self.worker.device, - ) + physical_to_logical, num_local_physical_experts, _ = broadcast_expert_mapping( + physical_to_logical=None, + num_local_physical_experts=None, + num_logical_experts=None, + dp_group=dp_group, + src_rank=0, + device=self.worker.device, ) num_moe_layers = physical_to_logical.shape[0] new_dp_size = get_dp_group().world_size @@ -642,13 +699,10 @@ class ElasticEPScalingExecutor: dtype=physical_to_logical.dtype, device=physical_to_logical.device, ) - old_num_physical_experts = physical_to_logical.shape[1] - expanded_physical_to_logical[:, :old_num_physical_experts] = physical_to_logical - return ( - expanded_physical_to_logical, - num_logical_experts, - old_num_physical_experts, + expanded_physical_to_logical[:, : physical_to_logical.shape[1]] = ( + physical_to_logical ) + return expanded_physical_to_logical def warmup_local_kernels(self) -> None: with set_current_vllm_config(self.worker.vllm_config): diff --git a/vllm/distributed/elastic_ep/elastic_state.py b/vllm/distributed/elastic_ep/elastic_state.py index 1b15ba8dc7..6718cd8ae8 100644 --- a/vllm/distributed/elastic_ep/elastic_state.py +++ b/vllm/distributed/elastic_ep/elastic_state.py @@ -30,12 +30,10 @@ WorkerType = Literal["existing", "new", "removing"] class ScaleUpExistingEngineState(enum.IntEnum): - CREATE_STANDBY_GROUPS = 0 - STAGE_QUANT_METHODS = 1 - TRANSFER_WEIGHTS = 2 - SYNC_KV_CACHE_MEMORY_SIZE = 3 - COMMIT_SCALE_UP = 4 # Blocks forward passes. - COMPLETE = 5 + PREPARE = 0 + SYNC_KV_CACHE_MEMORY_SIZE = 1 + COMMIT_SCALE_UP = 2 # Blocks forward passes. + COMPLETE = 3 class ScaleUpNewEngineState(enum.IntEnum): @@ -95,7 +93,7 @@ class ElasticEPScalingState: self.state = ( ScaleUpNewEngineState.PRE_KV_INIT if worker_type == "new" - else ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS + else ScaleUpExistingEngineState.PREPARE ) else: self.state = ( @@ -166,20 +164,8 @@ class ElasticEPScalingState: state = self.state assert self.old_dp_group is not None - if state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS: - if not self._create_standby_groups(): - return False - self.state = ScaleUpExistingEngineState.STAGE_QUANT_METHODS - return True - - elif state == ScaleUpExistingEngineState.STAGE_QUANT_METHODS: - if not self._execute_async("stage_standby_moe_quant_methods"): - return False - self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS - return True - - elif state == ScaleUpExistingEngineState.TRANSFER_WEIGHTS: - if not self._transfer_weights(): + if state == ScaleUpExistingEngineState.PREPARE: + if not self._prepare_workers(): return False self.state = ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE return True @@ -210,7 +196,7 @@ class ElasticEPScalingState: assert self.new_dp_group is not None and self.new_dp_store is not None if state == ScaleUpNewEngineState.PRE_KV_INIT: - self._collective_rpc("elastic_ep_execute", args=("receive_weights",)) + self._collective_rpc("elastic_ep_execute", args=("prepare_new_worker",)) self.engine_core.available_gpu_memory_for_kv_cache = ( ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1) ) @@ -243,7 +229,7 @@ class ElasticEPScalingState: assert self.old_dp_group is not None if state == ScaleDownRemainingEngineState.PREPARE: - if self._create_standby_groups(): + if self._prepare_workers(): self.state = ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN self._mark_ready_for_switch() return True @@ -328,29 +314,18 @@ class ElasticEPScalingState: self._prepare_future = None return True - def _create_standby_groups(self) -> bool: + def _prepare_workers(self) -> bool: assert self.old_dp_group is not None if not self._ensure_new_dp_group(): return False if not self._execute_async( - "create_standby_groups", + "prepare_reconfiguration", self.reconfig_request, self.new_parallel_config.use_all2all, ): return False if self.old_dp_group.rank() == 0: - logger.info("[Elastic EP] Created standby communication groups") - return True - - def _transfer_weights(self) -> bool: - assert self.reconfig_request is not None and self.old_dp_group is not None - old_dp_size = self.old_dp_group.size() - new_dp_size = self.reconfig_request.new_data_parallel_size - - if not self._execute_async("transfer_weights", old_dp_size, new_dp_size): - return False - if self.old_dp_group.rank() == 0: - logger.info("[Elastic EP] Transferred weights to new workers") + logger.info("[Elastic EP] Prepared reconfiguration") return True def _sync_kv_cache_memory_size(self) -> bool: diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 3f04b7bc91..14efd16177 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -248,7 +248,6 @@ class NixlEplbCommunicator(EplbCommunicator): cpu_group: ProcessGroup, all_expert_weights: Sequence[Sequence[torch.Tensor]], expert_buffer: Sequence[torch.Tensor], - defer_remote_setup: bool = False, ) -> None: """Create a NIXL-backed EPLB communicator. @@ -256,11 +255,6 @@ class NixlEplbCommunicator(EplbCommunicator): cpu_group: CPU process group for metadata exchange. all_expert_weights: Expert weight tensors for all MoE layers. expert_buffer: Pre-allocated receive buffer tensors. - defer_remote_setup: If True, postpone the collective - all-gather of NIXL agent metadata until the first - ``set_transfer_context`` call. Required for elastic EP - where ranks join asynchronously and cannot participate - in collectives at construction time. """ assert all_expert_weights, ( "NixlEplbCommunicator requires non-empty all_expert_weights." @@ -317,29 +311,17 @@ class NixlEplbCommunicator(EplbCommunicator): ] = {} self._cuda_device_id = int(self._device.index or 0) - self._remote_state_initialized = False self._init_step("buffers", self._init_registered_buffers) - if defer_remote_setup: - logger.info_once("NIXL EPLB: deferring remote agent setup (elastic EP).") - else: - self._init_remote_state() + self._init_remote_state() self._log_initialized() def _init_remote_state(self) -> None: """Exchange NIXL agent metadata and RDMA pointer info with all peers. This is a collective operation (uses ``all_gather_object`` twice). - Under elastic EP the call is deferred to the first - ``set_transfer_context`` invocation, where all ranks are - guaranteed to be synchronized. """ self._init_step("agents", self._init_remote_agents) self._init_step("send meta", self._exchange_remote_send_meta) - self._remote_state_initialized = True - - def _ensure_remote_state(self) -> None: - if not self._remote_state_initialized: - self._init_remote_state() @property def needs_profile_buffer_reservation(self) -> bool: @@ -373,7 +355,6 @@ class NixlEplbCommunicator(EplbCommunicator): pass def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None: - self._ensure_remote_state() assert not self._xfer_entries, ( f"set_transfer_context() called with {len(self._xfer_entries)} " f"pending transfers from layer {self._layer_idx}; " @@ -673,9 +654,8 @@ def create_eplb_communicator( Falls back to ``"torch_nccl"`` when *None*. Stateless (elastic EP) groups support ``"torch_nccl"``, ``"pynccl"``, and ``"nixl"``; ``"torch_nccl"`` is silently - promoted to ``"pynccl"``. ``"nixl"`` uses deferred remote - agent setup to avoid collective deadlocks during elastic - scaling. When tensors reside on CPU, ``"torch_gloo"`` or + promoted to ``"pynccl"``. When tensors reside on CPU, + ``"torch_gloo"`` or ``"torch_nccl"`` are used via the CPU process group. expert_weights: Expert weight tensors for *all* MoE layers. Shape ``(num_layers)(num_tensors_per_layer)``. @@ -729,22 +709,19 @@ def create_eplb_communicator( ) from exc is_stateless = isinstance(group_coordinator, StatelessGroupCoordinator) - if is_stateless: - if backend == "nixl": - pass # handled below with defer_remote_setup=True - elif backend not in ("torch_nccl", "pynccl"): + if is_stateless and backend != "nixl": + if backend not in ("torch_nccl", "pynccl"): raise ValueError( f"Elastic EP requires 'torch_nccl', 'pynccl', or 'nixl' " f"EPLB communicator (got '{backend}')." ) - else: - if backend == "torch_nccl": - logger.warning( - "Stateless elastic EP requires PyNCCL backend. " - "Forcing EPLB communicator to 'pynccl'." - ) - backend = "pynccl" - return _create_pynccl() + if backend == "torch_nccl": + logger.warning( + "Stateless elastic EP requires PyNCCL backend. " + "Forcing EPLB communicator to 'pynccl'." + ) + backend = "pynccl" + return _create_pynccl() if backend == "nixl": if not has_nixl(): @@ -761,7 +738,6 @@ def create_eplb_communicator( cpu_group=group_coordinator.cpu_group, all_expert_weights=expert_weights, expert_buffer=expert_buffer, - defer_remote_setup=is_stateless, ) except Exception as exc: raise RuntimeError( diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 5b2bb05ed0..0e3e38b253 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -37,6 +37,7 @@ from torch.distributed import ProcessGroup, all_reduce from vllm.config import ModelConfig, ParallelConfig from vllm.config.utils import compute_hash_cached from vllm.distributed.parallel_state import ( + GroupCoordinator, get_ep_group, get_eplb_group, get_node_count, @@ -289,14 +290,6 @@ class EplbState: """ CUDA device index for the async EPLB worker thread. """ - self.num_valid_physical_experts: int = 0 - """ - Number of valid physical experts. - This is the number of physical experts that are - actually mapped to logical experts. In elastic EP, - newly started EP ranks may not have physical experts - mapped yet. - """ if self.device.type == "cuda": self.cuda_device_index = self.device.index if self.cuda_device_index is None and torch.cuda.is_available(): @@ -499,7 +492,6 @@ class EplbState: num_unpadded_tokens_tensors=num_unpadded_tokens_tensors, ) self.model_states[model_config.compute_hash()] = model_state - self.num_valid_physical_experts = model.num_physical_experts def prepare_forward( self, @@ -766,28 +758,28 @@ class EplbState: # Map the physical expert load to global logical experts global_expert_load_windows = [] for eplb_model_state in self.model_states.values(): - expert_load_window = eplb_model_state.expert_load_window[ - :, :, : self.num_valid_physical_experts - ] + expert_load_window = eplb_model_state.expert_load_window + physical_to_logical = eplb_model_state.physical_to_logical_map + invalid_idx = eplb_model_state.model.num_logical_experts logical_expert_load_window = torch.zeros( self.expert_load_window_size, eplb_model_state.model.num_moe_layers, - eplb_model_state.model.num_logical_experts, + invalid_idx + 1, dtype=eplb_model_state.expert_load_window.dtype, device=eplb_model_state.expert_load_window.device, ) logical_expert_load_window.scatter_add_( dim=-1, - index=eplb_model_state.physical_to_logical_map[ - :, : self.num_valid_physical_experts - ] + index=physical_to_logical.masked_fill( + physical_to_logical < 0, invalid_idx + ) .unsqueeze(0) .expand_as(expert_load_window) .long(), src=expert_load_window, ) - global_expert_load_window = logical_expert_load_window.sum(dim=0) + global_expert_load_window = logical_expert_load_window[..., :-1].sum(dim=0) global_expert_load_windows.append(global_expert_load_window) # Perform all-reduce to get the expert load across all ranks for each model global_expert_load_windows = self._allreduce_list(global_expert_load_windows) @@ -960,7 +952,7 @@ class EplbState: Each pending result is acknowledged (consumed_event recorded) so the async worker can proceed, but the transferred weights are intentionally - NOT applied — a full synchronous rearrange is expected to follow. + NOT applied — a full rearrange is expected to follow. Ranks are kept in lockstep via _all_ranks_result_ready (all_reduce on the EP CPU group). The async worker's coordinated-stop collectives @@ -1060,7 +1052,6 @@ class EplbState: device: torch.device, parallel_config: ParallelConfig, expanded_physical_to_logical: torch.Tensor, - num_valid_physical_experts: int, ) -> "EplbState": eplb_state = cls( parallel_config=parallel_config, @@ -1070,12 +1061,24 @@ class EplbState: model=model, model_config=model_config, ) - eplb_state.num_valid_physical_experts = num_valid_physical_experts - eplb_model_state = eplb_state.model_states[model_config.compute_hash()] + eplb_state.update_mapping( + model_config, + expanded_physical_to_logical, + ) + + return eplb_state + + def update_mapping( + self, + model_config: ModelConfig, + expanded_physical_to_logical: torch.Tensor, + ) -> None: + eplb_model_state = self.model_states[model_config.compute_hash()] eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical) (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps( - expanded_physical_to_logical.cpu(), model.num_logical_experts + expanded_physical_to_logical.cpu(), + eplb_model_state.model.num_logical_experts, ) max_num_replicas = eplb_model_state.logical_to_physical_map.shape[-1] @@ -1087,13 +1090,31 @@ class EplbState: max_num_replicas - num_replicas, ), value=-1, - ).to(device) - logical_replica_count = logical_replica_count_cpu.to(device) + ).to(self.device) + logical_replica_count = logical_replica_count_cpu.to(self.device) eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map) eplb_model_state.logical_replica_count.copy_(logical_replica_count) - return eplb_state + def create_communicator( + self, model_config: ModelConfig, group_coordinator: GroupCoordinator + ) -> EplbCommunicator: + model_state = self.model_states[model_config.compute_hash()] + backend = self.parallel_config.eplb_config.communicator + assert backend is not None + return create_eplb_communicator( + group_coordinator, + backend, + model_state.model.expert_weights, + model_state.expert_buffer, + ) + + def update_communicator( + self, + model_config: ModelConfig, + communicator: EplbCommunicator, + ) -> None: + self.model_states[model_config.compute_hash()].communicator = communicator @dataclass diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 128db903d4..96a95a69fd 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1366,21 +1366,21 @@ def _replace_active_groups( ep: GroupCoordinator | None, eplb: GroupCoordinator | None, node_count: int | None, -) -> None: - """Destroy the current DP/EP/WORLD/EPLB groups and replace them. +) -> tuple[GroupCoordinator | None, ...]: + """Replace the active groups and return the groups they replaced. - Destruction is collective — all ranks in the old groups must call this - function together. Pass all-``None`` to tear down without replacement. + The caller must destroy the returned DP, EP, WORLD, and EPLB groups + collectively and in that order. Pass all-``None`` to remove the active + groups without replacement. """ global _WORLD, _DP, _EP, _EPLB, _NODE_COUNT - for group in (_DP, _EP, _WORLD, _EPLB): - if group is not None: - group.destroy() + old_groups = _DP, _EP, _WORLD, _EPLB _WORLD = world _DP = dp _EP = ep _EPLB = eplb _NODE_COUNT = node_count + return old_groups _TP: GroupCoordinator | None = None diff --git a/vllm/v1/worker/gpu/eplb_utils.py b/vllm/v1/worker/gpu/eplb_utils.py index 57a4e6db2c..24c0f0c4e1 100644 --- a/vllm/v1/worker/gpu/eplb_utils.py +++ b/vllm/v1/worker/gpu/eplb_utils.py @@ -144,7 +144,6 @@ class EPLBController: model: nn.Module, model_config: Any, expanded_physical_to_logical: torch.Tensor, - old_num_physical_experts: int, ) -> None: moe_model = get_mixture_of_experts_model(model) assert moe_model is not None @@ -155,6 +154,5 @@ class EPLBController: device=self.device, parallel_config=self.parallel_config, expanded_physical_to_logical=expanded_physical_to_logical, - num_valid_physical_experts=old_num_physical_experts, ) self._has_registered_models = True diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4f65769c95..3861bb5264 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1962,13 +1962,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): def setup_eplb_from_mapping( self, expanded_physical_to_logical: torch.Tensor, - old_num_physical_experts: int, ) -> None: self.eplb.setup_from_mapping( self.model, self.model_config, expanded_physical_to_logical, - old_num_physical_experts, ) ########### EPLB methods end ########### diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index a8d08a7bf2..ae11b2b602 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -3517,17 +3517,11 @@ class GPUModelRunner( def setup_eplb_from_mapping( self, expanded_physical_to_logical: torch.Tensor, - old_num_physical_experts: int, ) -> None: - assert self._moe_model is not None - - self.eplb_state = EplbState.from_mapping( - model=self._moe_model, - model_config=self.model_config, - device=self.device, - parallel_config=self.parallel_config, - expanded_physical_to_logical=expanded_physical_to_logical, - num_valid_physical_experts=old_num_physical_experts, + assert self.eplb_state is not None + self.eplb_state.update_mapping( + self.model_config, + expanded_physical_to_logical, ) def _pool( @@ -5488,11 +5482,7 @@ class GPUModelRunner( # MixtureOfExperts themselves. self._moe_model = get_mixture_of_experts_model(self.model) - if ( - self.parallel_config.enable_eplb - and not load_dummy_weights - and self._moe_model is not None - ): + if self._moe_model is not None and self.parallel_config.enable_eplb: logger.info_once( "EPLB is enabled for MoE part of model %s.", self.model_config.model, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index a3b00aaad2..958972f3c5 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1376,6 +1376,8 @@ class Worker(WorkerBase): if weight_transfer_engine := getattr(self, "weight_transfer_engine", None): weight_transfer_engine.shutdown() + self.elastic_ep_executor.shutdown() + # Release GPU resources held by the model runner so that memory # can be reclaimed when running in-process if model_runner := getattr(self, "model_runner", None):