748 lines
30 KiB
Python
748 lines
30 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import gc
|
|
import weakref
|
|
from collections.abc import Iterable, Sequence
|
|
from concurrent.futures import Future, ThreadPoolExecutor
|
|
from dataclasses import replace
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.distributed import P2POp
|
|
|
|
from vllm.compilation.counter import compilation_counter
|
|
from vllm.compilation.cuda_graph import CUDAGraphWrapper
|
|
from vllm.compilation.wrapper import reset_compile_wrapper
|
|
from vllm.config import (
|
|
CompilationMode,
|
|
set_current_vllm_config,
|
|
)
|
|
from vllm.distributed import (
|
|
get_dp_group,
|
|
get_ep_group,
|
|
get_pcp_group,
|
|
get_tp_group,
|
|
)
|
|
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 (
|
|
EplbCommunicator,
|
|
)
|
|
from vllm.distributed.parallel_state import (
|
|
GroupCoordinator,
|
|
_replace_active_groups,
|
|
)
|
|
from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
|
|
from vllm.model_executor.layers.fused_moe.eep_reconfigure import (
|
|
make_eep_staged_quant_method,
|
|
)
|
|
from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
|
|
from vllm.utils import is_moe_layer
|
|
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
|
|
from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
|
|
from vllm.v1.worker.workspace import lock_workspace, unlock_workspace
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
|
|
FusedMoEMethodBase,
|
|
)
|
|
|
|
|
|
def batch_transfer_weights(
|
|
model: nn.Module,
|
|
is_sender: bool,
|
|
peer_rank: int,
|
|
dp_group: StatelessGroupCoordinator,
|
|
expert_weights: Sequence[Iterable[torch.Tensor]],
|
|
) -> None:
|
|
device_comm = dp_group.device_communicator
|
|
if device_comm is None:
|
|
raise ValueError("No device communicator found")
|
|
|
|
expert_weights_set = set()
|
|
for weight_group in expert_weights:
|
|
for weight in weight_group:
|
|
expert_weights_set.add(weight.data_ptr())
|
|
|
|
state_dict = model.state_dict()
|
|
all_params = []
|
|
|
|
for name, param in state_dict.items():
|
|
if name.endswith("expert_map") or name.find("._shared_experts") != -1:
|
|
continue
|
|
if param.data_ptr() not in expert_weights_set:
|
|
all_params.append(param.data)
|
|
|
|
assert len(all_params) > 0
|
|
p2p_ops = []
|
|
for param in all_params:
|
|
# PyNccl transfers flat memory and does not honor tensor strides.
|
|
transfer_param = param.contiguous()
|
|
op = object.__new__(P2POp)
|
|
op.op = torch.distributed.isend if is_sender else torch.distributed.irecv
|
|
op.tensor = transfer_param
|
|
op.group_peer = peer_rank
|
|
p2p_ops.append(op)
|
|
if transfer_param is not param:
|
|
device_comm.batch_isend_irecv(p2p_ops)
|
|
p2p_ops.clear()
|
|
if not is_sender:
|
|
param.copy_(transfer_param)
|
|
if p2p_ops:
|
|
device_comm.batch_isend_irecv(p2p_ops)
|
|
|
|
|
|
def broadcast_expert_mapping(
|
|
physical_to_logical: torch.Tensor | None,
|
|
num_local_physical_experts: int | None,
|
|
num_logical_experts: int | None,
|
|
dp_group: StatelessGroupCoordinator,
|
|
device: torch.device,
|
|
src_rank: int = 0,
|
|
) -> tuple[torch.Tensor, int, int]:
|
|
if dp_group.rank_in_group == src_rank:
|
|
assert physical_to_logical is not None
|
|
assert num_local_physical_experts is not None
|
|
assert num_logical_experts is not None
|
|
assert physical_to_logical.dtype == torch.int64
|
|
shape_tensor = torch.tensor(
|
|
list(physical_to_logical.shape), dtype=torch.int64, device="cpu"
|
|
)
|
|
metadata_tensor = torch.tensor(
|
|
[num_local_physical_experts, num_logical_experts],
|
|
dtype=torch.int64,
|
|
device="cpu",
|
|
)
|
|
else:
|
|
shape_tensor = torch.empty(2, dtype=torch.int64, device="cpu")
|
|
metadata_tensor = torch.empty(2, dtype=torch.int64, device="cpu")
|
|
|
|
shape_tensor = dp_group.tcp_store_group.broadcast(shape_tensor, src_rank)
|
|
metadata_tensor = dp_group.tcp_store_group.broadcast(metadata_tensor, src_rank)
|
|
|
|
if dp_group.rank_in_group != src_rank:
|
|
assert device is not None
|
|
physical_to_logical = torch.empty(
|
|
tuple(shape_tensor.tolist()),
|
|
dtype=torch.int64,
|
|
device=device,
|
|
)
|
|
|
|
assert physical_to_logical is not None
|
|
physical_to_logical = dp_group.broadcast(physical_to_logical, src_rank)
|
|
num_local_physical_experts = int(metadata_tensor[0].item())
|
|
num_logical_experts = int(metadata_tensor[1].item())
|
|
|
|
return physical_to_logical, num_local_physical_experts, num_logical_experts
|
|
|
|
|
|
class ElasticEPScalingExecutor:
|
|
def __init__(self, worker):
|
|
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):
|
|
worker = self.worker_ref()
|
|
if worker is None:
|
|
raise RuntimeError("Worker has been garbage collected")
|
|
return worker
|
|
|
|
def execute(self, execute_method: str, *args, **kwargs):
|
|
method = getattr(self, execute_method, None)
|
|
if method is None:
|
|
raise ValueError(f"Unknown execute method: {execute_method}")
|
|
return method(*args, **kwargs)
|
|
|
|
def start_async(self, execute_method: str, *args, **kwargs) -> str:
|
|
if self._async_future is not None:
|
|
raise RuntimeError("Another Elastic EP async method is active")
|
|
if args and isinstance(args[0], ReconfigureDistributedRequest):
|
|
self.reconfig_request = args[0]
|
|
dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank
|
|
done_key = f"eep_async/{execute_method}/{dp_rank}/{self.worker.rank}"
|
|
self._async_future = self._async_executor.submit(
|
|
self._run_async, execute_method, *args, **kwargs
|
|
)
|
|
self._async_future.add_done_callback(lambda _: self._mark_async_done(done_key))
|
|
return done_key
|
|
|
|
def _run_async(self, execute_method: str, *args, **kwargs) -> None:
|
|
from vllm.platforms import current_platform
|
|
|
|
self.worker.vllm_config.enable_trace_function_call_for_thread()
|
|
assert hasattr(self.worker, "device")
|
|
current_platform.set_device(self.worker.device)
|
|
with set_current_vllm_config(self.worker.vllm_config):
|
|
self.execute(execute_method, *args, **kwargs)
|
|
|
|
def _mark_async_done(self, done_key: str) -> None:
|
|
from vllm.distributed.utils import get_cached_tcp_store_client
|
|
|
|
assert self.reconfig_request is not None
|
|
get_cached_tcp_store_client(
|
|
self.reconfig_request.new_data_parallel_master_ip,
|
|
self.reconfig_request.coord_store_port,
|
|
).set(done_key, b"1")
|
|
|
|
def clear_async(self) -> None:
|
|
future = self._async_future
|
|
if future is None:
|
|
raise RuntimeError("No Elastic EP async method is active")
|
|
if not future.done():
|
|
raise RuntimeError("Elastic EP async method is not done")
|
|
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 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
|
|
parallel_config = self.worker.vllm_config.parallel_config
|
|
world_size = parallel_config.world_size
|
|
new_world_size_across_dp = world_size * new_dp_size
|
|
create_standby_groups(
|
|
new_dp_size=new_dp_size,
|
|
new_world_size_across_dp=new_world_size_across_dp,
|
|
master_ip=reconfig_request.new_data_parallel_master_ip,
|
|
coord_store_port=reconfig_request.coord_store_port,
|
|
use_all2all=use_all2all,
|
|
enable_eplb=parallel_config.enable_eplb,
|
|
)
|
|
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()
|
|
assert standby_dp_group is not None
|
|
# Broadcast old_dp_size to all workers in standby group
|
|
if standby_dp_group.rank_in_group < old_dp_size:
|
|
old_dp_size_tensor = torch.tensor(
|
|
[old_dp_size], dtype=torch.int64, device="cpu"
|
|
)
|
|
else:
|
|
old_dp_size_tensor = torch.empty(1, dtype=torch.int64, device="cpu")
|
|
old_dp_size_tensor = standby_dp_group.tcp_store_group.broadcast(
|
|
old_dp_size_tensor, 0
|
|
)
|
|
|
|
num_new_workers = new_dp_size - old_dp_size
|
|
dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank
|
|
|
|
# Sender-receiver pairing: the first new_workers % old_dp_size
|
|
# senders get (k+1) contiguous receivers, the rest get k
|
|
# receivers.
|
|
num_dst_per_sender = num_new_workers // old_dp_size
|
|
remainder = num_new_workers % old_dp_size
|
|
|
|
if dp_rank < remainder:
|
|
recv_begin = dp_rank * (num_dst_per_sender + 1)
|
|
recv_end = recv_begin + num_dst_per_sender + 1
|
|
else:
|
|
recv_begin = (
|
|
remainder * (num_dst_per_sender + 1)
|
|
+ (dp_rank - remainder) * num_dst_per_sender
|
|
)
|
|
recv_end = recv_begin + num_dst_per_sender
|
|
|
|
ranks_to_send = list(range(old_dp_size + recv_begin, old_dp_size + recv_end))
|
|
|
|
model = self.worker.model_runner.get_model()
|
|
for new_worker_rank in sorted(ranks_to_send):
|
|
batch_transfer_weights(
|
|
model=model,
|
|
is_sender=True,
|
|
peer_rank=new_worker_rank,
|
|
dp_group=standby_dp_group,
|
|
expert_weights=model.expert_weights,
|
|
)
|
|
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
|
|
model_config = self.worker.model_runner.model_config
|
|
eplb_state = self.worker.model_runner.eplb_state
|
|
assert eplb_state is not None
|
|
eplb_state.drain_async()
|
|
eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
|
|
physical_to_logical = eplb_model_state.physical_to_logical_map
|
|
num_physical_experts = physical_to_logical.shape[1]
|
|
num_local_physical_experts = num_physical_experts // get_ep_group().world_size
|
|
num_logical_experts = eplb_model_state.logical_replica_count.shape[1]
|
|
broadcast_expert_mapping(
|
|
physical_to_logical=physical_to_logical,
|
|
num_local_physical_experts=num_local_physical_experts,
|
|
num_logical_experts=num_logical_experts,
|
|
dp_group=standby_dp_group,
|
|
src_rank=0,
|
|
device=self.worker.device,
|
|
)
|
|
|
|
def _make_eep_moe_config(self, module, dp_group, ep_group):
|
|
parallel_config = self.worker.vllm_config.parallel_config
|
|
tp_size = get_tp_group().world_size
|
|
sp_size = tp_size if parallel_config.use_sequence_parallel_moe else 1
|
|
moe_parallel_config = FusedMoEParallelConfig.make(
|
|
tp_size_=tp_size,
|
|
pcp_size_=get_pcp_group().world_size,
|
|
dp_size_=dp_group.world_size,
|
|
sp_size_=sp_size,
|
|
vllm_parallel_config=parallel_config,
|
|
)
|
|
return replace(
|
|
module.moe_config,
|
|
num_experts=module.moe_config.num_local_experts * ep_group.world_size,
|
|
moe_parallel_config=moe_parallel_config,
|
|
)
|
|
|
|
def stage_standby_moe_quant_methods(self) -> None:
|
|
standby_dp_group = get_standby_dp_group()
|
|
standby_ep_group = get_standby_ep_group()
|
|
model = self.worker.model_runner.get_model()
|
|
moe_modules = [module for module in model.modules() if is_moe_layer(module)]
|
|
self._staged_moe_quant_methods.clear()
|
|
with set_current_vllm_config(self.worker.vllm_config):
|
|
for module in moe_modules:
|
|
staged_quant_method = make_eep_staged_quant_method(
|
|
module,
|
|
self._make_eep_moe_config(
|
|
module,
|
|
standby_dp_group,
|
|
standby_ep_group,
|
|
),
|
|
)
|
|
if staged_quant_method is not None:
|
|
self._staged_moe_quant_methods[module] = staged_quant_method
|
|
|
|
def _commit_staged_moe_quant_methods(self) -> None:
|
|
model = self.worker.model_runner.get_model()
|
|
moe_modules = [module for module in model.modules() if is_moe_layer(module)]
|
|
for module in moe_modules:
|
|
staged_quant_method = self._staged_moe_quant_methods.pop(module, None)
|
|
if staged_quant_method is None:
|
|
continue
|
|
assert staged_quant_method.moe_kernel is not None
|
|
module._replace_quant_method(staged_quant_method)
|
|
staged_quant_method.moe_kernel.prepare_finalize.on_commit()
|
|
self._staged_moe_quant_methods.clear()
|
|
|
|
def _release_cuda_graphs(self) -> None:
|
|
if isinstance(self.worker.model_runner.model, CUDAGraphWrapper):
|
|
wrapper = self.worker.model_runner.model
|
|
wrapper.concrete_cudagraph_entries = {}
|
|
|
|
elif isinstance(self.worker.model_runner.model, UBatchWrapper):
|
|
raise RuntimeError("DBO is not yet supported in elastic EP")
|
|
|
|
torch.compiler.reset()
|
|
with set_current_vllm_config(self.worker.vllm_config):
|
|
reset_compile_wrapper(self.worker.model_runner.get_model())
|
|
|
|
gc.collect()
|
|
torch.accelerator.synchronize()
|
|
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()
|
|
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) -> tuple[GroupCoordinator | None, ...]:
|
|
old_dp_size = get_dp_group().world_size
|
|
old_ep_size = get_ep_group().world_size
|
|
|
|
self._release_cuda_graphs()
|
|
retired_groups = _replace_active_groups(**pop_standby_groups())
|
|
|
|
parallel_config = self.worker.vllm_config.parallel_config
|
|
reconfig_request = self.reconfig_request
|
|
assert reconfig_request is not None
|
|
new_dp_size = reconfig_request.new_data_parallel_size
|
|
new_ep_size = get_ep_group().world_size
|
|
|
|
parallel_config.data_parallel_size = new_dp_size
|
|
if (
|
|
reconfig_request.new_data_parallel_rank
|
|
!= ReconfigureRankType.KEEP_CURRENT_RANK
|
|
):
|
|
parallel_config.data_parallel_rank = reconfig_request.new_data_parallel_rank
|
|
if (
|
|
reconfig_request.new_data_parallel_rank_local
|
|
!= ReconfigureRankType.KEEP_CURRENT_RANK
|
|
):
|
|
parallel_config.data_parallel_rank_local = (
|
|
reconfig_request.new_data_parallel_rank_local
|
|
)
|
|
parallel_config.data_parallel_master_ip = (
|
|
reconfig_request.new_data_parallel_master_ip
|
|
)
|
|
parallel_config.data_parallel_master_port = (
|
|
reconfig_request.new_data_parallel_master_port
|
|
)
|
|
|
|
# Reconfigure MoE modules with new EP size
|
|
moe_modules = [
|
|
module
|
|
for module in self.worker.model_runner.model.modules()
|
|
if is_moe_layer(module)
|
|
]
|
|
num_local_experts = moe_modules[0].moe_config.num_local_experts
|
|
assert all(
|
|
module.moe_config.num_local_experts == num_local_experts
|
|
for module in moe_modules
|
|
), "All MoE modules must have the same number of experts"
|
|
dp_group = get_dp_group()
|
|
ep_group = get_ep_group()
|
|
for module in moe_modules:
|
|
new_moe_config = self._make_eep_moe_config(module, dp_group, ep_group)
|
|
module._set_moe_config(new_moe_config)
|
|
|
|
# Update EPLB state
|
|
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()]
|
|
|
|
num_physical_experts = num_local_experts * new_ep_size
|
|
num_logical_experts = eplb_model_state.logical_replica_count.shape[1]
|
|
parallel_config.eplb_config.num_redundant_experts = (
|
|
num_physical_experts - num_logical_experts
|
|
)
|
|
old_physical_to_logical = eplb_model_state.physical_to_logical_map
|
|
num_moe_layers = old_physical_to_logical.shape[0]
|
|
num_local_experts = eplb_model_state.expert_load_pass.shape[1] // old_ep_size
|
|
if new_dp_size > old_dp_size:
|
|
expanded_physical_to_logical = torch.full(
|
|
(num_moe_layers, num_local_experts * new_ep_size),
|
|
-1,
|
|
dtype=old_physical_to_logical.dtype,
|
|
device=old_physical_to_logical.device,
|
|
)
|
|
expanded_physical_to_logical[:, : num_local_experts * old_ep_size] = (
|
|
old_physical_to_logical
|
|
)
|
|
eplb_model_state.physical_to_logical_map = expanded_physical_to_logical
|
|
|
|
old_num_physical_experts = eplb_model_state.expert_load_pass.shape[1]
|
|
pad_size = num_physical_experts - old_num_physical_experts
|
|
if new_dp_size > old_dp_size:
|
|
assert pad_size > 0
|
|
expanded_expert_load_pass = F.pad(
|
|
eplb_model_state.expert_load_pass, (0, pad_size), value=0
|
|
)
|
|
expanded_expert_load_window = F.pad(
|
|
eplb_model_state.expert_load_window, (0, pad_size), value=0
|
|
)
|
|
eplb_model_state.expert_load_pass = expanded_expert_load_pass
|
|
eplb_model_state.expert_load_window = expanded_expert_load_window
|
|
else:
|
|
assert pad_size < 0
|
|
eplb_model_state.expert_load_pass = eplb_model_state.expert_load_pass[
|
|
:, :num_physical_experts
|
|
]
|
|
eplb_model_state.expert_load_window = eplb_model_state.expert_load_window[
|
|
:, :, :num_physical_experts
|
|
]
|
|
|
|
model = self.worker.model_runner.get_model()
|
|
model.expert_weights = []
|
|
with set_current_vllm_config(self.worker.vllm_config):
|
|
model.set_eplb_state(
|
|
eplb_model_state.expert_load_pass,
|
|
eplb_model_state.logical_to_physical_map,
|
|
eplb_model_state.logical_replica_count,
|
|
)
|
|
eplb_state._propagate_shared_tensors(
|
|
model, eplb_model_state.num_unpadded_tokens_tensors
|
|
)
|
|
model.update_physical_experts_metadata(
|
|
num_physical_experts=num_physical_experts,
|
|
num_local_physical_experts=num_local_experts,
|
|
)
|
|
self._commit_staged_moe_quant_methods()
|
|
# Legacy modular methods need to be recreated for the new EP size.
|
|
for module in moe_modules:
|
|
if getattr(module._quant_method, "wraps_legacy_quant_method", False):
|
|
module._replace_quant_method(module._quant_method.old_quant_method)
|
|
|
|
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
|
|
== CompilationMode.STOCK_TORCH_COMPILE
|
|
):
|
|
# NOTE(yongji): when using stock torch.compile,
|
|
# torch.compile is triggered during GPUModelRunner's load_model()
|
|
# TODO(yongji):check do we need to re-trigger torch.compile here?
|
|
# any changes to the tensor shapes in execution should already
|
|
# be handled internally by torch.compile.
|
|
backend = self.worker.vllm_config.compilation_config.init_backend(
|
|
self.worker.vllm_config
|
|
)
|
|
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,
|
|
async_op: bool = False,
|
|
) -> None:
|
|
if get_ep_group().rank == 0:
|
|
logger.info("[Elastic EP] Starting expert resharding...")
|
|
|
|
eplb_state = self.worker.model_runner.eplb_state
|
|
assert eplb_state is not None
|
|
|
|
is_async_enabled = eplb_state.is_async
|
|
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)
|
|
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.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
|
|
# in setup_eplb_from_mapping() but don't start the thread there because
|
|
# groups aren't ready yet.
|
|
eplb_state.start_async_loop()
|
|
if get_ep_group().rank == 0:
|
|
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()
|
|
retired_groups = self.switch_and_prepare()
|
|
else:
|
|
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:
|
|
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
|
|
if eplb_state is not None:
|
|
eplb_state.drain_async()
|
|
parallel_config = self.worker.vllm_config.parallel_config
|
|
tp_size = parallel_config.tensor_parallel_size
|
|
old_ep_size = parallel_config.data_parallel_size * tp_size
|
|
new_ep_size = new_dp_size * tp_size
|
|
rank_mapping = {
|
|
old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1
|
|
for old_ep_rank in range(old_ep_size)
|
|
}
|
|
self._perform_eplb_reshuffle(rank_mapping=rank_mapping)
|
|
|
|
def prepare_new_worker(self) -> None:
|
|
dp_group = get_dp_group()
|
|
assert isinstance(dp_group, StatelessGroupCoordinator)
|
|
new_dp_size = dp_group.world_size
|
|
dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank
|
|
|
|
# Receive old_dp_size broadcasted during transfer_weights
|
|
old_dp_size_tensor = torch.empty(1, dtype=torch.int64, device="cpu")
|
|
old_dp_size_tensor = dp_group.tcp_store_group.broadcast(old_dp_size_tensor, 0)
|
|
old_dp_size = int(old_dp_size_tensor[0].item())
|
|
|
|
# Calculate which existing worker will send to this new worker
|
|
num_new_workers = new_dp_size - old_dp_size
|
|
new_worker_idx = dp_rank - old_dp_size
|
|
num_dst_per_sender = num_new_workers // old_dp_size
|
|
remainder = num_new_workers % old_dp_size
|
|
|
|
if new_worker_idx < remainder * (num_dst_per_sender + 1):
|
|
sender_rank = new_worker_idx // (num_dst_per_sender + 1)
|
|
else:
|
|
sender_rank = (
|
|
remainder
|
|
+ (new_worker_idx - remainder * (num_dst_per_sender + 1))
|
|
// num_dst_per_sender
|
|
)
|
|
|
|
model = self.worker.model_runner.get_model()
|
|
expert_weights = [
|
|
module.get_expert_weights()
|
|
for module in model.modules()
|
|
if is_moe_layer(module)
|
|
]
|
|
batch_transfer_weights(
|
|
model=model,
|
|
is_sender=False,
|
|
peer_rank=sender_rank,
|
|
dp_group=dp_group,
|
|
expert_weights=expert_weights,
|
|
)
|
|
torch.accelerator.synchronize()
|
|
self._warm_target_groups(get_dp_group(), get_ep_group())
|
|
|
|
def receive_expert_mapping(self) -> torch.Tensor:
|
|
dp_group = get_dp_group()
|
|
assert isinstance(dp_group, StatelessGroupCoordinator)
|
|
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
|
|
tp_size = self.worker.vllm_config.parallel_config.tensor_parallel_size
|
|
new_ep_size = new_dp_size * tp_size
|
|
expanded_physical_to_logical = torch.full(
|
|
(num_moe_layers, num_local_physical_experts * new_ep_size),
|
|
-1,
|
|
dtype=physical_to_logical.dtype,
|
|
device=physical_to_logical.device,
|
|
)
|
|
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):
|
|
kernel_warmup(self.worker, process_local_only=True)
|
|
|
|
def warm_and_capture(self) -> None:
|
|
# Must run on every DP sibling in lockstep: _dummy_run calls
|
|
# coordinate_batch_across_dp whenever data_parallel_size > 1
|
|
# (gpu_model_runner.py:3663), which deadlocks if any rank skips it.
|
|
|
|
# Save and clear block tables so the dummy MoE forward doesn't
|
|
# write dummy slot mappings into real KV-cache blocks.
|
|
multi_block_table = self.worker.model_runner.input_batch.block_table
|
|
saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
|
|
for bt in multi_block_table.block_tables:
|
|
saved_block_tables.append(
|
|
(bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
|
|
)
|
|
multi_block_table.clear()
|
|
|
|
# _ensure_workspace_size allocates a fresh tensor on grow, leaving
|
|
# any captured CUDA graph with a stale data pointer; drop graphs
|
|
# before re-warm so captures realign with the resized buffer.
|
|
self._release_cuda_graphs()
|
|
unlock_workspace()
|
|
|
|
# Grow the MoE workspace at max_num_tokens. compile_or_warm_up_model
|
|
# alone only exercises cudagraph-capture sizes and can leave the
|
|
# workspace too small for post-reshuffle routing. Use _dummy_run
|
|
# directly with skip_eplb=True so dummy routing doesn't pollute the
|
|
# just-rebalanced EPLB stats.
|
|
runner = self.worker.model_runner
|
|
runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)
|
|
self.worker.compile_or_warm_up_model()
|
|
|
|
lock_workspace()
|
|
|
|
for bt, (saved_gpu, saved_cpu) in zip(
|
|
multi_block_table.block_tables, saved_block_tables
|
|
):
|
|
bt.block_table.gpu.copy_(saved_gpu)
|
|
bt.block_table.cpu.copy_(saved_cpu)
|