diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 0038851ac5a..191786ddca3 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -61,7 +61,7 @@ def should_free_input(name, is_moe, config, num_local_experts): return False enable_deepep = ( config.moe_token_dispatcher_type == "flex" - and config.moe_flex_dispatcher_backend == "deepep" + and config.moe_flex_dispatcher_backend in ("deepep", "deepepv2") ) enable_hybridep = ( config.moe_token_dispatcher_type == "flex" @@ -494,7 +494,7 @@ def build_transformer_layer_callables(layer: TransformerLayer): is_moe = isinstance(layer.mlp, MoELayer) enable_deepep = ( layer.config.moe_token_dispatcher_type == "flex" - and layer.config.moe_flex_dispatcher_backend == "deepep" + and layer.config.moe_flex_dispatcher_backend in ("deepep", "deepepv2") ) enable_hybridep = ( layer.config.moe_token_dispatcher_type == "flex" diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index e0f7ab9030c..0ce27c7f21f 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -8,16 +8,32 @@ from megatron.core.utils import internal_api try: - from deep_ep import Buffer from deep_ep.utils import EventHandle, EventOverlap +except ImportError: + try: + from deep_ep import EventHandle, EventOverlap + except ImportError: + EventHandle = None + EventOverlap = None + +try: + from deep_ep import Buffer HAVE_DEEP_EP = True except ImportError: HAVE_DEEP_EP = False +try: + from deep_ep import ElasticBuffer + + HAVE_DEEP_EP_V2 = True +except ImportError: + HAVE_DEEP_EP_V2 = False + import torch _buffer = None +_elastic_buffer = None def get_hidden_bytes(x: torch.Tensor) -> int: @@ -71,6 +87,32 @@ def get_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int): return _buffer +def get_elastic_buffer( + group: torch.distributed.ProcessGroup, num_max_tokens_per_rank: int, hidden: int, num_topk: int +): + """Get or create a DeepEP v2 elastic buffer for all-to-all communication.""" + global _elastic_buffer + + num_bytes = ElasticBuffer.get_buffer_size_hint( + group, num_max_tokens_per_rank=num_max_tokens_per_rank, hidden=hidden, num_topk=num_topk + ) + + if ( + _elastic_buffer is None + or _elastic_buffer.group != group + or _elastic_buffer.num_bytes < num_bytes + or _elastic_buffer.num_max_tokens_per_rank < num_max_tokens_per_rank + ): + _elastic_buffer = ElasticBuffer( + group, + num_bytes=num_bytes, + num_max_tokens_per_rank=num_max_tokens_per_rank, + hidden=hidden, + num_topk=num_topk, + ) + return _elastic_buffer + + class FusedDispatch(torch.autograd.Function): """Fused dispatch operation for MoE routing combining computation and communication.""" @@ -270,6 +312,184 @@ def set_deepep_num_sms(num_sms): set_deepep_num_sms = None +class DeepepV2Dispatch(torch.autograd.Function): + """Dispatch operation using the DeepEP v2 ElasticBuffer backend.""" + + @staticmethod + def forward( + ctx, + buffer, + x, + token_indices, + token_probs, + num_experts, + num_max_tokens_per_rank, + expert_alignment, + num_sms, + async_finish=False, + allocate_on_comm_stream=False, + ): + """Forward pass of dispatch using the DeepEP v2 ElasticBuffer backend.""" + # Capture the current stream for the communication stream to wait on when + # DeepEP v2 allocates output tensors on the communication stream. + previous_event = buffer.capture() if async_finish and allocate_on_comm_stream else None + # Process the dispatch and keep the handle for the subsequent combine call. + recv_x, recv_token_indices, recv_token_probs, handle, event = buffer.dispatch( + x, + topk_idx=token_indices, + topk_weights=token_probs, + num_experts=num_experts, + num_max_tokens_per_rank=num_max_tokens_per_rank, + expert_alignment=expert_alignment, + num_sms=num_sms, + previous_event=previous_event, + async_with_compute_stream=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + + if async_finish: + event.current_stream_wait() + + ctx.buffer = buffer + ctx.handle = handle + ctx.num_sms = handle.num_sms + ctx.async_finish = async_finish + ctx.allocate_on_comm_stream = allocate_on_comm_stream + tokens_per_expert = torch.tensor(handle.num_recv_tokens_per_expert_list) + + return (recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, handle) + + @staticmethod + def backward( + ctx, grad_output, grad_token_indices, grad_token_probs, grad_tokens_per_expert, grad_handle + ): + """Backward pass of dispatch using the DeepEP v2 ElasticBuffer backend.""" + # The backward pass of dispatch is a combine over the dispatch handle. + previous_event = ( + ctx.buffer.capture() if ctx.async_finish and ctx.allocate_on_comm_stream else None + ) + grad_x, grad_token_probs, event = ctx.buffer.combine( + grad_output.contiguous(), + handle=ctx.handle, + topk_weights=grad_token_probs.float(), + num_sms=ctx.num_sms, + previous_event=previous_event, + async_with_compute_stream=ctx.async_finish, + allocate_on_comm_stream=ctx.allocate_on_comm_stream, + ) + if ctx.async_finish: + event.current_stream_wait() + return None, grad_x, None, grad_token_probs, None, None, None, None, None, None + + +class DeepepV2Combine(torch.autograd.Function): + """DeepEP v2 elastic combine with autograd support.""" + + @staticmethod + def forward(ctx, buffer, x, handle, num_sms, async_finish=False, allocate_on_comm_stream=False): + """Forward pass of DeepEP v2 elastic combine.""" + previous_event = buffer.capture() if async_finish and allocate_on_comm_stream else None + combined_x, combined_token_probs, event = buffer.combine( + x, + handle=handle, + num_sms=num_sms, + previous_event=previous_event, + async_with_compute_stream=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + if async_finish: + event.current_stream_wait() + + ctx.buffer = buffer + ctx.handle = handle + ctx.num_sms = handle.num_sms if num_sms == 0 else num_sms + ctx.async_finish = async_finish + ctx.allocate_on_comm_stream = allocate_on_comm_stream + return combined_x, combined_token_probs + + @staticmethod + def backward(ctx, grad_output, grad_combined_token_probs): + """Backward pass of DeepEP v2 elastic combine.""" + previous_event = ( + ctx.buffer.capture() if ctx.async_finish and ctx.allocate_on_comm_stream else None + ) + grad_x, _, _, _, event = ctx.buffer.dispatch( + grad_output.contiguous(), + handle=ctx.handle, + num_sms=ctx.num_sms, + previous_event=previous_event, + async_with_compute_stream=ctx.async_finish, + allocate_on_comm_stream=ctx.allocate_on_comm_stream, + ) + if ctx.async_finish: + event.current_stream_wait() + return None, grad_x, None, None, None, None + + +if HAVE_DEEP_EP_V2: + + def deepepv2_dispatch( + buffer, + x, + token_indices, + token_probs, + num_experts, + num_max_tokens_per_rank, + expert_alignment=1, + num_sms=0, + async_finish=False, + allocate_on_comm_stream=False, + ): + """Perform dispatch using the DeepEP v2 ElasticBuffer backend. + + Args: + buffer (ElasticBuffer): + DeepEP v2 buffer used for all-to-all communication. + x (torch.Tensor): + Input hidden states to dispatch. + token_indices (torch.Tensor): + Top-k expert indices for each token. + token_probs (torch.Tensor): + Top-k routing probabilities for each token. + num_experts (int): + Total number of experts across the communication group. + num_max_tokens_per_rank (int): + Maximum number of input tokens on each rank. + expert_alignment (int): + Alignment applied to per-expert token counts. + num_sms (int): + Number of SMs used by the dispatch API. + async_finish (bool): + Whether to use asynchronous communication completion. + allocate_on_comm_stream (bool): + Whether to allocate DeepEP output buffers on the communication stream. + """ + return DeepepV2Dispatch.apply( + buffer, + x.contiguous(), + token_indices, + token_probs, + num_experts, + num_max_tokens_per_rank, + expert_alignment, + num_sms, + async_finish, + allocate_on_comm_stream, + ) + + def deepepv2_combine( + buffer, x, handle, num_sms=0, async_finish=False, allocate_on_comm_stream=False + ): + """Perform DeepEP v2 elastic combine.""" + return DeepepV2Combine.apply( + buffer, x.contiguous(), handle, num_sms, async_finish, allocate_on_comm_stream + ) + +else: + deepepv2_dispatch = None + deepepv2_combine = None + + try: from deep_ep import HybridEPBuffer diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index e8f2044650e..ed18b6df511 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -19,8 +19,11 @@ from megatron.core.transformer.enums import CudaGraphModule from megatron.core.transformer.moe.fused_a2a import ( HYBRIDEP_TOKEN_ALIGNMENT, + deepepv2_combine, + deepepv2_dispatch, fused_combine, fused_dispatch, + get_elastic_buffer, hybrid_ep_combine, hybrid_ep_dispatch, set_deepep_num_sms, @@ -1282,7 +1285,10 @@ def __init__( "DeepEP is not installed. Please install DeepEP package from " "https://github.com/deepseek-ai/deepep." ) - set_deepep_num_sms(config.moe_deepep_num_sms) + if config.moe_deepep_num_sms is None: + set_deepep_num_sms(20) + else: + set_deepep_num_sms(config.moe_deepep_num_sms) def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): num_tokens = routing_map.shape[0] @@ -1461,6 +1467,112 @@ def get_restored_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> return hidden_states +class _DeepepV2Manager(_DeepepManager): + """ + A manager class for the DeepEP v2 ElasticBuffer backend. + + This keeps the original DeepEP backend isolated under "deepep", while "deepepv2" + uses the v2 dispatch/combine APIs. + """ + + def __init__( + self, + group: torch.distributed.ProcessGroup, + num_local_experts: int, + router_topk: int, + num_experts: int, + config: TransformerConfig, + ): + # Do not call _DeepepManager.__init__; v2-only images may not ship the v1 Buffer API. + self.group = group + self.num_local_experts = num_local_experts + self.config = config + + self.router_topk = router_topk + self.num_experts = num_experts + self.router_dtype = config.moe_router_dtype + self.capacity_factor = config.moe_expert_capacity_factor + self.permute_fusion = config.moe_permute_fusion + if config.moe_deepep_num_sms is None: + self.num_sms = 0 + else: + self.num_sms = config.moe_deepep_num_sms + + self.token_indices: Optional[torch.Tensor] = None + self.token_probs: Optional[torch.Tensor] = None + self.handle = None + self.buffer = None + + if deepepv2_dispatch is None: + raise ImportError( + "DeepEP v2 is not installed. Please install a DeepEP package that provides " + "ElasticBuffer." + ) + + def _get_buffer(self, hidden_states: torch.Tensor): + self.buffer = get_elastic_buffer( + self.group, + num_max_tokens_per_rank=hidden_states.shape[0], + hidden=hidden_states.shape[1], + num_topk=self.token_indices.shape[1], + ) + return self.buffer + + def dispatch( + self, + hidden_states: torch.Tensor, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + ) -> torch.Tensor: + # DeepEP v2 only supports float32 probs + if self.token_probs.dtype != torch.float32: + if self.token_probs.dtype in [torch.bfloat16, torch.float16]: + logger.warning( + "DeepEP v2 only supports float32 probs, please set --moe-router-dtype=fp32" + ) + self.token_probs = self.token_probs.float() + buffer = self._get_buffer(hidden_states) + hidden_states, dispatched_indices, dispatched_probs, num_tokens_per_expert, handle = ( + deepepv2_dispatch( + buffer, + hidden_states, + self.token_indices, + self.token_probs, + self.num_experts, + num_max_tokens_per_rank=hidden_states.shape[0], + expert_alignment=1, + num_sms=self.num_sms, + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + ) + self.handle = handle + self.tokens_per_expert = num_tokens_per_expert + self.dispatched_indices = dispatched_indices + self.dispatched_probs = dispatched_probs + + return hidden_states + + def combine( + self, + hidden_states: torch.Tensor, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + ) -> torch.Tensor: + hidden_states, _ = deepepv2_combine( + self.buffer, + hidden_states, + self.handle, + num_sms=self.num_sms, + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + self.handle = None + self.dispatched_indices = None + self.dispatched_probs = None + return hidden_states + + class MoEFlexTokenDispatcher(MoETokenDispatcher): """A flexible token dispatcher that abstracts the underlying tensor and expert parallelism. It uses a single communication group over all TP and EP ranks, @@ -1487,6 +1599,7 @@ def __init__( self.num_local_experts = num_local_experts self.local_expert_indices = local_expert_indices + self._comm_manager: _DispatchManager if self.config.moe_flex_dispatcher_backend == "deepep": assert self.tp_size * self.ep_size > 1, "DeepEP dispatcher requires TPxEP > 1" self._comm_manager = _DeepepManager( @@ -1497,6 +1610,16 @@ def __init__( config=self.config, ) self.cudagraph_attrs = ['_comm_manager.token_probs', '_comm_manager.token_indices'] + elif self.config.moe_flex_dispatcher_backend == "deepepv2": + assert self.tp_size * self.ep_size > 1, "DeepEP v2 dispatcher requires TPxEP > 1" + self._comm_manager = _DeepepV2Manager( + group=self.tp_ep_group, + num_local_experts=self.num_local_experts, + router_topk=self.tp_size * self.config.moe_router_topk, + num_experts=self.tp_size * self.config.num_moe_experts, + config=self.config, + ) + self.cudagraph_attrs = ['_comm_manager.token_probs', '_comm_manager.token_indices'] elif self.config.moe_flex_dispatcher_backend == "hybridep": self._comm_manager = _HybridEPManager( group=self.tp_ep_group, @@ -1508,7 +1631,8 @@ def __init__( else: raise ValueError( f"Invalid backend: {self.config.moe_flex_dispatcher_backend}" - "Please set --moe-flex-dispatcher-backend=deepep or " + "Please set --moe-flex-dispatcher-backend=deepep, " + "--moe-flex-dispatcher-backend=deepepv2 or " "--moe-flex-dispatcher-backend=hybridep" ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 2cd4490b0a6..10c8603dcac 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -874,10 +874,10 @@ class TransformerConfig(ModelParallelConfig): moe_enable_deepep: bool = False """[Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models.""" - moe_flex_dispatcher_backend: Literal['deepep', 'hybridep'] = "deepep" + moe_flex_dispatcher_backend: Literal['deepep', 'deepepv2', 'hybridep'] = "deepep" """[Experimental] The backend to use for flex token dispatcher. The default is "deepep". - Options are "deepep" and "hybridep". Currently only "hybridep" backend supports - the MNNVL case.""" + Options are "deepep", "deepepv2" and "hybridep". Currently only "hybridep" + backend supports the MNNVL case.""" moe_permute_fusion_into_hybridep: bool = False """Fuse token rearrangement ops during token dispatching for HybridEP.""" @@ -935,8 +935,8 @@ class TransformerConfig(ModelParallelConfig): moe_latent_size: Optional[int] = None """Latent projection dimension for MoE. If None, MoE latent projections are not used.""" - moe_deepep_num_sms: int = 20 - """Number of SMs to use for DeepEP.""" + moe_deepep_num_sms: Optional[int] = None + """Number of SMs to use for DeepEP. None uses v1's default or v2's theoretical default.""" moe_hybridep_num_sms: Optional[int] = None """Number of SMs to use for HybridEP. None uses the default from DeepEP. @@ -1705,7 +1705,7 @@ def __post_init__(self): if self.moe_enable_deepep: if self.moe_token_dispatcher_type != "flex": raise ValueError("DeepEP backend is only supported with flex token dispatcher.") - if self.moe_flex_dispatcher_backend == "hybridep": + if self.moe_flex_dispatcher_backend in ("deepepv2", "hybridep"): raise ValueError("Only one backend is supported for flex token dispatcher.") self.moe_flex_dispatcher_backend = "deepep" warnings.warn( @@ -1715,10 +1715,10 @@ def __post_init__(self): if self.moe_token_dispatcher_type == "flex": if self.moe_pad_expert_input_to_capacity and ( - self.moe_enable_deepep or self.moe_flex_dispatcher_backend == "deepep" + self.moe_enable_deepep or self.moe_flex_dispatcher_backend in ("deepep", "deepepv2") ): raise ValueError( - "Flex token dispatcher with deepep backend does not support " + "Flex token dispatcher with deepep/deepepv2 backend does not support " "moe_pad_expert_input_to_capacity" ) diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index f3d5e47a103..4c1837cd900 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -174,7 +174,7 @@ "mlp_chunks_for_prefill": 1, "moe_apply_probs_on_input": False, "moe_aux_loss_coeff": 0.0, - "moe_deepep_num_sms": 20, + "moe_deepep_num_sms": None, "moe_enable_deepep": False, "moe_expert_capacity_factor": None, "moe_expert_rank_capacity_factor": None, diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py index ddc8b313c19..e769dab664f 100644 --- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py +++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py @@ -465,6 +465,12 @@ def is_deep_ep_available(): return HAVE_DEEP_EP +def is_deep_ep_v2_available(): + from megatron.core.transformer.moe.fused_a2a import HAVE_DEEP_EP_V2 + + return HAVE_DEEP_EP_V2 + + def is_hybrid_ep_available(): from megatron.core.transformer.moe.fused_a2a import HAVE_HYBRIDEP @@ -636,9 +642,18 @@ def test_sequence_packing_thd_e2e_proxy_model(dispatcher): Utils.destroy_model_parallel() +def skip_if_flex_backend_unavailable(moe_flex_dispatcher_backend): + if moe_flex_dispatcher_backend == "deepep" and not is_deep_ep_available(): + pytest.skip("Deep EP is not available") + if moe_flex_dispatcher_backend == "deepepv2" and not is_deep_ep_v2_available(): + pytest.skip("Deep EP v2 is not available") + if moe_flex_dispatcher_backend == "hybridep" and not is_hybrid_ep_available(): + pytest.skip("Hybrid EP is not available") + + @pytest.mark.skipif( - not is_deep_ep_available() and not is_hybrid_ep_available(), - reason="Deep EP and Hybrid EP are not available", + not is_deep_ep_available() and not is_deep_ep_v2_available() and not is_hybrid_ep_available(), + reason="Deep EP, Deep EP v2 and Hybrid EP are not available", ) class TestFlexDispatcher: def setup_method(self, method): @@ -652,7 +667,7 @@ def teardown_method(self, method): @pytest.mark.internal @pytest.mark.parametrize("tp_size,ep_size", [(1, 8), (8, 1), (4, 2)]) @pytest.mark.parametrize("permute_fusion", permute_fusion_params) - @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "hybridep"]) + @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "deepepv2", "hybridep"]) @pytest.mark.parametrize("moe_permute_fusion_into_hybridep", [True, False]) def test_forward_backward( self, @@ -662,10 +677,7 @@ def test_forward_backward( moe_flex_dispatcher_backend, moe_permute_fusion_into_hybridep, ): - if moe_flex_dispatcher_backend == "deepep" and not is_deep_ep_available(): - pytest.skip("Deep EP is not available") - if moe_flex_dispatcher_backend == "hybridep" and not is_hybrid_ep_available(): - pytest.skip("Hybrid EP is not available") + skip_if_flex_backend_unavailable(moe_flex_dispatcher_backend) if moe_permute_fusion_into_hybridep: if permute_fusion or moe_flex_dispatcher_backend != "hybridep": pytest.skip( @@ -696,7 +708,7 @@ def test_forward_backward( @pytest.mark.timeout(120) @pytest.mark.parametrize("tp_size,ep_size", [(1, 8), (8, 1), (4, 2)]) @pytest.mark.parametrize("permute_fusion", permute_fusion_params) - @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "hybridep"]) + @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "deepepv2", "hybridep"]) @pytest.mark.parametrize("moe_permute_fusion_into_hybridep", [True, False]) def test_capacity_forward_backward( self, @@ -706,10 +718,7 @@ def test_capacity_forward_backward( moe_flex_dispatcher_backend, moe_permute_fusion_into_hybridep, ): - if moe_flex_dispatcher_backend == "deepep" and not is_deep_ep_available(): - pytest.skip("Deep EP is not available") - if moe_flex_dispatcher_backend == "hybridep" and not is_hybrid_ep_available(): - pytest.skip("Hybrid EP is not available") + skip_if_flex_backend_unavailable(moe_flex_dispatcher_backend) if moe_permute_fusion_into_hybridep: if permute_fusion or moe_flex_dispatcher_backend != "hybridep": pytest.skip( @@ -745,7 +754,7 @@ def test_capacity_forward_backward( @pytest.mark.timeout(120) @pytest.mark.parametrize("tp_size,ep_size", [(1, 8), (8, 1), (4, 2)]) @pytest.mark.parametrize("permute_fusion", [True]) - @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "hybridep"]) + @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "deepepv2", "hybridep"]) @pytest.mark.parametrize("moe_permute_fusion_into_hybridep", [True, False]) def test_router_padding_for_fp8_forward_backward( self, @@ -755,10 +764,7 @@ def test_router_padding_for_fp8_forward_backward( moe_flex_dispatcher_backend, moe_permute_fusion_into_hybridep, ): - if moe_flex_dispatcher_backend == "deepep" and not is_deep_ep_available(): - pytest.skip("Deep EP is not available") - if moe_flex_dispatcher_backend == "hybridep" and not is_hybrid_ep_available(): - pytest.skip("Hybrid EP is not available") + skip_if_flex_backend_unavailable(moe_flex_dispatcher_backend) if moe_permute_fusion_into_hybridep: if permute_fusion or moe_flex_dispatcher_backend != "hybridep": pytest.skip(