diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 61bd7a6f94c..a5040df1734 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1057,9 +1057,8 @@ def __init__( self.moe_expert_rank_capacity_factor = self.config.moe_expert_rank_capacity_factor self.over_budget = torch.zeros(1, dtype=torch.bool, device='cuda') - # THD sequence packing can produce different token counts per rank. - # HybridEP dispatch expects equal per-rank input sizes, so metadata and - # hidden states are padded to the group-wide max and trimmed in combine. + # When runtime token equalization is enabled, HybridEP metadata and hidden + # states are padded to the group-wide max and trimmed again in combine. self._original_num_tokens: Optional[int] = None self._padded_num_tokens: Optional[int] = None @@ -1068,37 +1067,22 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): self._original_num_tokens = num_tokens padded_num_tokens = num_tokens - equalize_thd_token_counts = ( - self.config.sequence_packing_scheduler is not None - or self.config.moe_hybridep_pad_variable_tokens - ) - if equalize_thd_token_counts: - if self.config.sequence_packing_scheduler is not None and ( - torch.cuda.is_current_stream_capturing() or torch.compiler.is_compiling() - ): - # CUDA graph path: routing_map has already been padded to a static - # length upstream (CUDA graph + sequence packing implies - # cu_seqlens_q_padded -> max_seqlen_per_dp_cp_rank), so num_tokens - # is identical across the EP communication group. Skip the - # all_reduce + .item() during both dynamo tracing and stream - # capture, and use the local value directly. - padded_num_tokens = num_tokens - else: - # Use the actual tp_ep max so all ranks in the MoE communication - # group pass the same token count to HybridEP. - max_num_tokens_across_ep = torch.tensor( - [num_tokens], device=routing_map.device, dtype=torch.long - ) - torch.distributed.all_reduce( - max_num_tokens_across_ep, op=torch.distributed.ReduceOp.MAX, group=self.group - ) - padded_num_tokens = int(max_num_tokens_across_ep.item()) + if self.config.moe_hybridep_pad_variable_tokens: + # Use the actual tp_ep max so all ranks in the MoE communication + # group pass the same token count to HybridEP. + max_num_tokens_across_ep = torch.tensor( + [num_tokens], device=routing_map.device, dtype=torch.long + ) + torch.distributed.all_reduce( + max_num_tokens_across_ep, op=torch.distributed.ReduceOp.MAX, group=self.group + ) + padded_num_tokens = int(max_num_tokens_across_ep.item()) padded_num_tokens += -padded_num_tokens % HYBRIDEP_TOKEN_ALIGNMENT self._padded_num_tokens = padded_num_tokens routing_map = routing_map.reshape(num_tokens, self.num_experts) probs = probs.reshape(num_tokens, self.num_experts) - if equalize_thd_token_counts and padded_num_tokens > num_tokens: + if padded_num_tokens > num_tokens: pad_rows = padded_num_tokens - num_tokens routing_map = torch.cat( [routing_map, routing_map.new_zeros((pad_rows, self.num_experts))], dim=0 diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index bb01202b71c..96875ae45fc 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -899,10 +899,12 @@ class TransformerConfig(ModelParallelConfig): """Fuse token rearrangement ops during token dispatching for HybridEP.""" moe_hybridep_pad_variable_tokens: bool = False - """Pad uneven local token counts to the HybridEP group maximum before dispatch. + """Dynamically pad uneven local token counts to the HybridEP group maximum before dispatch. - This is needed when the frontend supplies locally packed THD inputs whose token counts - can differ across ranks, without using Megatron Core's sequence_packing_scheduler. + Enable this when local token counts can differ across ranks. When disabled, the caller must + guarantee equal token counts across the HybridEP communication group, for example by padding + THD inputs to a fixed maximum before dispatch. CUDA Graph inputs should be statically padded + upstream and leave this option disabled. """ moe_per_layer_logging: bool = False diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py index e769dab664f..95b8aa26b79 100644 --- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py +++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py @@ -1,25 +1,18 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import dataclasses -import math from types import SimpleNamespace import pytest import torch from megatron.core import config, parallel_state -from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices -from megatron.core.models.gpt.gpt_layer_specs import ( - get_gpt_layer_local_submodules, - get_gpt_layer_with_transformer_engine_spec, -) -from megatron.core.packed_seq_params import PackedSeqParams -from megatron.core.transformer.moe.fused_a2a import reset_hybrid_ep_buffer +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules +from megatron.core.transformer.moe.fused_a2a import HYBRIDEP_TOKEN_ALIGNMENT, reset_hybrid_ep_buffer from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules from megatron.core.transformer.moe.moe_utils import get_capacity -from megatron.core.transformer.moe.token_dispatcher import MoETokenDispatcher +from megatron.core.transformer.moe.token_dispatcher import MoETokenDispatcher, _HybridEPManager from megatron.core.transformer.spec_utils import get_submodules -from megatron.core.transformer.transformer_block import TransformerBlock from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.typed_torch import apply_module from megatron.core.utils import is_te_min_version @@ -85,6 +78,36 @@ def test_set_cudagraph_attr_supports_nested_paths(): assert dispatcher._comm_manager.routing_map is routing_map +def test_hybridep_variable_tokens_are_padded_to_group_max(monkeypatch): + manager = object.__new__(_HybridEPManager) + manager.config = SimpleNamespace(moe_hybridep_pad_variable_tokens=True) + manager.group = object() + manager.num_experts = 2 + manager.moe_expert_rank_capacity_factor = None + manager.drop_and_pad = False + + local_num_tokens = 65 + group_max_num_tokens = 129 + routing_map = torch.ones((local_num_tokens, manager.num_experts), dtype=torch.bool) + probs = torch.ones((local_num_tokens, manager.num_experts)) + + def fake_all_reduce(tensor, op, group): + assert op == torch.distributed.ReduceOp.MAX + assert group is manager.group + tensor.fill_(group_max_num_tokens) + + monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce) + manager.setup_metadata(routing_map, probs) + + expected_num_tokens = group_max_num_tokens + expected_num_tokens += -expected_num_tokens % HYBRIDEP_TOKEN_ALIGNMENT + assert manager._padded_num_tokens == expected_num_tokens + assert manager.routing_map.shape == (expected_num_tokens, manager.num_experts) + assert manager.token_probs.shape == (expected_num_tokens, manager.num_experts) + assert not manager.routing_map[local_num_tokens:].any() + assert not manager.token_probs[local_num_tokens:].any() + + class MoEModelTestContainer: def __init__( self, @@ -477,171 +500,6 @@ def is_hybrid_ep_available(): return HAVE_HYBRIDEP -def _round_up(value, divisor): - return value if divisor <= 1 else (value + divisor - 1) // divisor * divisor - - -def _get_thd_padded_seqlens(seqlens, cp_size, tp_size): - # This follows the runtime packed-sequence path used by the Moonlight script: - # per-sequence lengths must be CP partitionable, and the packed token count - # must be even for TP/SP slicing. - cp_divisor = 2 * cp_size if cp_size > 1 else 1 - padded_seqlens = [_round_up(seqlen, cp_divisor) for seqlen in seqlens] - total_seqlen = sum(padded_seqlens) - total_alignment = math.lcm(cp_divisor, tp_size) - padded_seqlens[-1] += _round_up(total_seqlen, total_alignment) - total_seqlen - return padded_seqlens - - -def _to_cu_seqlens(seqlens): - cu_seqlens = torch.empty(len(seqlens) + 1, dtype=torch.int32, device="cuda") - cu_seqlens[0] = 0 - cu_seqlens[1:] = torch.cumsum(torch.tensor(seqlens, dtype=torch.int32, device="cuda"), dim=0) - return cu_seqlens - - -def _make_thd_packed_seq_params(seqlens, cp_size, tp_size): - padded_seqlens = _get_thd_padded_seqlens(seqlens, cp_size, tp_size) - cu_seqlens_padded = _to_cu_seqlens(padded_seqlens) - max_seqlen = max(padded_seqlens) - # Match get_batch_on_this_rank_for_sequence_packing(): TE consumes padded - # cumulative lengths as both cu_seqlens and cu_seqlens_padded for THD. - return PackedSeqParams( - qkv_format="thd", - cu_seqlens_q=cu_seqlens_padded, - cu_seqlens_kv=cu_seqlens_padded, - cu_seqlens_q_padded=cu_seqlens_padded, - cu_seqlens_kv_padded=cu_seqlens_padded, - max_seqlen_q=max_seqlen, - max_seqlen_kv=max_seqlen, - ) - - -def _make_sharded_thd_hidden_states(seqlens, hidden_size, cp_size, tp_size, dtype): - padded_seqlens = _get_thd_padded_seqlens(seqlens, cp_size, tp_size) - padded_sequences = [] - for seqlen, padded_seqlen in zip(seqlens, padded_seqlens): - sequence = torch.randn(seqlen, hidden_size, device="cuda", dtype=dtype) - if padded_seqlen > seqlen: - sequence = torch.cat( - [ - sequence, - torch.zeros(padded_seqlen - seqlen, hidden_size, device="cuda", dtype=dtype), - ], - dim=0, - ) - padded_sequences.append(sequence) - - hidden_states = torch.cat(padded_sequences, dim=0) - if cp_size > 1: - cu_seqlens_padded = _to_cu_seqlens(padded_seqlens) - cp_rank = parallel_state.get_context_parallel_rank() - index = get_thd_partitioned_indices( - cu_seqlens_padded, hidden_states.shape[0], cp_size, cp_rank - ) - hidden_states = hidden_states.index_select(0, index) - - tp_rank = parallel_state.get_tensor_model_parallel_rank() - sequence_parallel_length = hidden_states.shape[0] // tp_size - hidden_states = hidden_states[ - tp_rank * sequence_parallel_length : (tp_rank + 1) * sequence_parallel_length - ] - return hidden_states.unsqueeze(1).contiguous().requires_grad_(True) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -@pytest.mark.skipif( - Utils.world_size % 8 != 0, reason="requires world size divisible by 8 for pp2/cp2/tp2/ep2/etp2" -) -@pytest.mark.internal -@pytest.mark.parametrize("dispatcher", ["alltoall", "deepep", "hybridep"]) -def test_sequence_packing_thd_e2e_proxy_model(dispatcher): - """Run packed THD attention + MoE forward/backward with major parallelisms enabled.""" - if not is_te_min_version("2.9.0"): - pytest.skip("SFT sequence packing requires Transformer Engine >= 2.9.0") - if dispatcher == "deepep" and not is_deep_ep_available(): - pytest.skip("Deep EP is not available") - if dispatcher == "hybridep" and not is_hybrid_ep_available(): - pytest.skip("Hybrid EP is not available") - - tp_size, pp_size, cp_size, ep_size, etp_size = 2, 2, 2, 2, 2 - Utils.initialize_model_parallel( - tensor_model_parallel_size=tp_size, - pipeline_model_parallel_size=pp_size, - context_parallel_size=cp_size, - expert_model_parallel_size=ep_size, - expert_tensor_parallel_size=etp_size, - ) - _set_random_seed(seed_=123, data_parallel_random_init=False) - - try: - spec = get_gpt_layer_with_transformer_engine_spec(num_experts=4, moe_grouped_gemm=False) - transformer_config = TransformerConfig( - num_layers=4, - hidden_size=1024, - ffn_hidden_size=2048, - moe_ffn_hidden_size=2048, - num_attention_heads=8, - tensor_model_parallel_size=tp_size, - pipeline_model_parallel_size=pp_size, - context_parallel_size=cp_size, - expert_model_parallel_size=ep_size, - expert_tensor_parallel_size=etp_size, - sequence_parallel=True, - sequence_packing_scheduler="dp_balanced", - max_seqlen_per_dp_cp_rank=1024, - cp_comm_type="p2p", - num_moe_experts=4, - moe_router_topk=2, - moe_router_load_balancing_type="aux_loss", - moe_token_dispatcher_type=( - "flex" if dispatcher in ("deepep", "hybridep") else dispatcher - ), - moe_flex_dispatcher_backend=( - dispatcher if dispatcher in ("deepep", "hybridep") else "deepep" - ), - moe_grouped_gemm=False, - moe_router_dtype="fp32", - params_dtype=torch.bfloat16, - pipeline_dtype=torch.bfloat16, - autocast_dtype=torch.bfloat16, - bf16=True, - add_bias_linear=False, - attention_dropout=0.0, - hidden_dropout=0.0, - use_cpu_initialization=True, - ) - transformer_block = TransformerBlock(transformer_config, spec).cuda().to(torch.bfloat16) - - torch.manual_seed(1000 + torch.distributed.get_rank()) - seqlens = [257, 509, 1021] - hidden_states = _make_sharded_thd_hidden_states( - seqlens, transformer_config.hidden_size, cp_size, tp_size, torch.bfloat16 - ) - packed_seq_params = _make_thd_packed_seq_params(seqlens, cp_size, tp_size) - - output = transformer_block( - hidden_states=hidden_states, attention_mask=None, packed_seq_params=packed_seq_params - ) - assert output.shape == hidden_states.shape - assert torch.isfinite(output).all() - - loss = output.float().square().mean() - loss.backward() - - assert hidden_states.grad is not None - assert hidden_states.grad.shape == hidden_states.shape - assert torch.isfinite(hidden_states.grad).all() - assert any( - param.grad is not None and torch.isfinite(param.grad).all() - for param in transformer_block.parameters() - if param.requires_grad - ) - finally: - reset_hybrid_ep_buffer() - 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")