diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 0066a60dd02..81006da401d 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -261,8 +261,13 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase): All2All communication based on DeepEP Low-Latency kernels. """ + _buffer: Any = None + _mask: torch.Tensor | None = None + _last_mask: torch.Tensor | None = None + def __init__(self, cpu_group, tcp_store_group=None): super().__init__(cpu_group, tcp_store_group) + self.support_fault_tolerance = False # TODO: set to True when FT is supported. def _make_all2all_kwargs( self, @@ -304,6 +309,7 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase): allow_nvlink_for_low_latency_mode=True, allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL, explicitly_destroy=True, + enable_shrink=self.support_fault_tolerance, ) return kwargs @@ -319,12 +325,30 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase): handle: deep_ep.Buffer = self.handle_cache.get_or_create( buffer_kwargs, deep_ep.Buffer ) + DeepEPLLAll2AllManager._buffer = handle return handle # DeepEP LL uses RDMA so no SMs are used for communication def max_sms_used(self) -> int | None: return 0 + def query_active_mask(self) -> torch.Tensor: + buf = DeepEPLLAll2AllManager._buffer + assert buf is not None + if DeepEPLLAll2AllManager._mask is None: + DeepEPLLAll2AllManager._mask = torch.zeros( + self.world_size, device="cuda", dtype=torch.int32 + ) + buf.low_latency_query_mask_buffer(DeepEPLLAll2AllManager._mask) + return DeepEPLLAll2AllManager._mask + + def query_fault(self) -> torch.Tensor: + current = self.query_active_mask() + if DeepEPLLAll2AllManager._last_mask is None: + DeepEPLLAll2AllManager._last_mask = torch.zeros_like(current) + has_fault = (current != DeepEPLLAll2AllManager._last_mask).any() + return has_fault + @dataclass class _NixlEPBufferState: @@ -341,6 +365,8 @@ class NixlEPAll2AllManager(All2AllManagerBase): _buffer: _NixlEPBufferState | None = None _lock = threading.RLock() + _mask: torch.Tensor | None = None + _last_mask: torch.Tensor | None = None def __init__(self, cpu_group, tcp_store_group=None): if tcp_store_group is None: @@ -350,6 +376,7 @@ class NixlEPAll2AllManager(All2AllManagerBase): store=dist.PrefixStore("nixl_ep", cpu_group.get_group_store()), ) super().__init__(cpu_group, tcp_store_group) + self.support_fault_tolerance = True self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS @@ -508,6 +535,25 @@ class NixlEPAll2AllManager(All2AllManagerBase): def max_sms_used(self) -> int | None: return 0 + def query_active_mask(self) -> torch.Tensor: + state = NixlEPAll2AllManager._buffer + assert state is not None + if NixlEPAll2AllManager._mask is None: + NixlEPAll2AllManager._mask = torch.zeros( + self.max_num_ep_ranks, device="cuda", dtype=torch.int32 + ) + state.buffer.query_mask_buffer(NixlEPAll2AllManager._mask) + return NixlEPAll2AllManager._mask[: state.active_ep_size] + + def query_fault(self) -> torch.Tensor: + current = self.query_active_mask() + last = NixlEPAll2AllManager._last_mask + if last is None or last.shape != current.shape: + NixlEPAll2AllManager._last_mask = torch.zeros_like(current) + last = NixlEPAll2AllManager._last_mask + has_fault = (current != last).any() + return has_fault + class FlashInferNVLinkTwoSidedManager(All2AllManagerBase): """ diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 0b4b81f93bb..6fd889daeb0 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -62,6 +62,8 @@ class All2AllManagerBase: in_the_same_node_as(tcp_store_group, source_rank=0) ) + self.support_fault_tolerance = False + def get_handle(self, kwargs): # get a handle for the all2all communication, # based on the kwargs. @@ -102,6 +104,13 @@ class All2AllManagerBase: # - raise a clear error if extra_tensors is not supported. raise NotImplementedError + def query_active_mask(self) -> torch.Tensor: + raise NotImplementedError + + def query_fault(self) -> torch.Tensor: + """Returns has_fault scalar.""" + raise NotImplementedError + def set_num_sms(self, num_sms: int): pass diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 1351e87b5b5..6af93bfde90 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -56,7 +56,7 @@ if current_platform.is_cuda_alike(): ) -def _get_ep_all2all_manager(eep_stage: bool = False) -> Any: +def get_ep_all2all_manager(eep_stage: bool = False) -> Any: if eep_stage: from vllm.distributed.elastic_ep.standby_state import get_standby_ep_group @@ -146,7 +146,7 @@ def maybe_make_prepare_finalize( "Detected DP deployment with no --enable-expert-parallel. " "Falling back to AllGather+ReduceScatter dispatch/combine." ) - all2all_manager = _get_ep_all2all_manager(eep_stage) + all2all_manager = get_ep_all2all_manager(eep_stage) return make_moe_prepare_and_finalize_naive_dp_ep( is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel, num_dispatchers=all2all_manager.world_size, @@ -155,7 +155,7 @@ def maybe_make_prepare_finalize( else: return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic) - all2all_manager = _get_ep_all2all_manager(eep_stage) + all2all_manager = get_ep_all2all_manager(eep_stage) prepare_finalize: FusedMoEPrepareAndFinalize | None = None diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 9b7581311e8..0937fad8f1f 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -939,7 +939,11 @@ class WorkerProc: converted to a FAILURE response. """ if isinstance(output, AsyncModelRunnerOutput): - output = output.get_output() + try: + output = output.get_output() + except Exception as e: + logger.exception("Error getting async model runner output") + output = e if isinstance(output, Exception): result = (WorkerProc.ResponseStatus.FAILURE, str(output)) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 022d598d6cb..84d989827c9 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -58,6 +58,7 @@ from vllm.logger import init_logger from vllm.lora.layers import LoRAMapping, LoRAMappingType from vllm.model_executor.layers.attention import Attention, MLAAttention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( RoutedExpertsCapturer, ) @@ -249,6 +250,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): async_output_copy_stream: torch.cuda.Stream, vocab_size: int, routed_experts: RoutedExpertsTensors | None = None, + check_ep_fault: bool = False, ): self._model_runner_output = model_runner_output self._invalid_req_indices = invalid_req_indices @@ -262,6 +264,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): self.vocab_size = vocab_size self._logprobs_tensors = logprobs_tensors self._routed_experts = routed_experts + self._has_fault: torch.Tensor | None = None # Initiate the copy on a separate stream, but do not synchronize it. default_stream = torch.cuda.current_stream() @@ -280,6 +283,9 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): if self._routed_experts is not None else None ) + if check_ep_fault: + has_fault = get_ep_all2all_manager().query_fault() + self._has_fault = has_fault.to("cpu", non_blocking=True) self.async_copy_ready_event.record() def get_output(self) -> ModelRunnerOutput: @@ -316,6 +322,14 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): output.routed_experts = self._routed_experts_cpu.tolists() del self._routed_experts + if self._has_fault is not None and self._has_fault.item(): + mask = get_ep_all2all_manager().query_active_mask() + raise RuntimeError( + "Fault detected in EP all2all communication: " + "one or more ranks timed out during dispatch/combine. " + f"Mask: {mask.cpu().tolist()}" + ) + return output @@ -445,6 +459,10 @@ class GPUModelRunner( self.device = device self.dtype = self.model_config.dtype + self.check_ep_fault = False + if parallel_config.data_parallel_size > 1 and self.model_config.is_moe: + self.check_ep_fault = get_ep_all2all_manager().support_fault_tolerance + self.kv_cache_dtype = kv_cache_dtype_str_to_dtype( cache_config.cache_dtype, self.model_config ) @@ -4695,6 +4713,7 @@ class GPUModelRunner( async_output_copy_stream=self._get_or_create_async_output_copy_stream(), vocab_size=self.input_batch.vocab_size, routed_experts=routed_experts_snapshot, + check_ep_fault=self.check_ep_fault, ) with record_function_or_nullcontext( "gpu_model_runner: set_async_sampled_token_ids"