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

[Elastic EP] Reduce eager-mode reconfiguration downtime (#51885)

Browse Source 
Signed-off-by: Itay Alroy <ialroy@nvidia.com>
This commit is contained in:
Itay Alroy
2026-08-19 02:13:06 +00:00
committed by GitHub
parent d36cc4254e
commit a2257f95b7
12 changed files with 194 additions and 303 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/distributed/test_eplb_execute.py
-122
View File
@@ -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,
)
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+1 -2
View File
@@ -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):
vllm/config/parallel.py
+1 -1
View File
@@ -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
vllm/distributed/elastic_ep/elastic_execute.py
+108 -54
View File
@@ -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):
vllm/distributed/elastic_ep/elastic_state.py
+12 -37
View File
@@ -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:
vllm/distributed/eplb/eplb_communicator.py
+12 -36
View File
@@ -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(
vllm/distributed/eplb/eplb_state.py
+46 -25
View File
@@ -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
vllm/distributed/parallel_state.py
+7 -7
View File
@@ -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
vllm/v1/worker/gpu/eplb_utils.py
-2
View File
@@ -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
vllm/v1/worker/gpu/model_runner.py
-2
View File
@@ -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 ###########
vllm/v1/worker/gpu_model_runner.py
+5 -15
View File
@@ -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,
vllm/v1/worker/gpu_worker.py
+2
View File
@@ -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):