diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index e3d19e88d7f..2cb37e5d235 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -52,9 +52,12 @@ def get_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int): num_nvl_bytes = max( config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes ) - num_rdma_bytes = max( - config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes - ) + # Local-only EP groups do not need an RDMA buffer, and DeepEP builds + # without internode support may not expose RDMA size hints. + if group.size() > torch.cuda.device_count(): + num_rdma_bytes = max( + config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes + ) # Allocate buffer if not existed or not enough buffer # NOTES: the adaptive routing configuration of the network **must be off** @@ -276,6 +279,9 @@ def set_deepep_num_sms(num_sms): _hybrid_ep_buffer = None +# HybridEP dispatch/combine kernels use 64-token chunks for their public APIs. +HYBRIDEP_TOKEN_ALIGNMENT = 64 + def init_hybrid_ep_buffer( group: torch.distributed.ProcessGroup, diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 7f8b6a1c9a0..45aa652d453 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -18,6 +18,7 @@ ) from megatron.core.transformer.enums import CudaGraphModule from megatron.core.transformer.moe.fused_a2a import ( + HYBRIDEP_TOKEN_ALIGNMENT, fused_combine, fused_dispatch, hybrid_ep_combine, @@ -1049,16 +1050,46 @@ 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. + self._original_num_tokens: Optional[int] = None + self._padded_num_tokens: Optional[int] = None def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): num_tokens = routing_map.shape[0] - self.routing_map = routing_map.reshape(num_tokens, self.num_experts) - self.token_probs = probs.reshape(num_tokens, self.num_experts) + self._original_num_tokens = num_tokens + + padded_num_tokens = num_tokens + if self.config.sequence_packing_scheduler is not None: + # 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 self.config.sequence_packing_scheduler is not None and 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 + ) + probs = torch.cat([probs, probs.new_zeros((pad_rows, self.num_experts))], dim=0) + + self.routing_map = routing_map + self.token_probs = probs if self.moe_expert_rank_capacity_factor is not None: pad_multiple = get_align_size_for_quantization(self.config) budget = int( - routing_map.shape[0] + padded_num_tokens * self.config.moe_router_topk * self.moe_expert_rank_capacity_factor ) @@ -1066,7 +1097,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): self.num_permuted_tokens = budget # Compute the capacity for each expert at the drop_and_pad mode if self.drop_and_pad: - num_out_tokens = num_tokens * self.config.moe_router_topk + num_out_tokens = padded_num_tokens * self.config.moe_router_topk # Drop and pad the input to capacity. self.capacity = get_capacity( num_tokens=num_out_tokens, @@ -1095,6 +1126,11 @@ def dispatch( self.token_probs = self.token_probs.float() # downcast or upcast if self.config.fp8 or self.config.fp4: self.pad_multiple = get_align_size_for_quantization(self.config) + if self._padded_num_tokens is not None and hidden_states.shape[0] < self._padded_num_tokens: + pad_rows = self._padded_num_tokens - hidden_states.shape[0] + hidden_states = torch.cat( + [hidden_states, hidden_states.new_zeros((pad_rows, hidden_states.shape[-1]))], dim=0 + ) dispatched_hidden, self.dispatched_probs, _, tokens_per_expert, self.handle = ( hybrid_ep_dispatch( x=hidden_states, @@ -1137,12 +1173,20 @@ def combine( pad_multiple=self.pad_multiple, fused=self.config.moe_permute_fusion_into_hybridep, ) + if ( + self._padded_num_tokens is not None + and self._original_num_tokens is not None + and hidden_states.shape[0] > self._original_num_tokens + ): + hidden_states = hidden_states[: self._original_num_tokens] # Release the used handle/num_permuted_tokens which could change in each iteration. # For drop_and_pad mode, we don't need to reset the num_permuted_tokens and # num_dispatched_tokens, because their values never change. self.handle = None if not self.drop_and_pad: self.num_permuted_tokens = None + self._original_num_tokens = None + self._padded_num_tokens = None return hidden_states def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index a701d6bd268..d90ead62ea3 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -2851,10 +2851,9 @@ def _scope_to_str(s): # Needed for passing variable sequences between pp stages. self.variable_seq_lengths = True - # TODO(tailaim): add support for other dispatcher types - assert self.moe_token_dispatcher_type == "alltoall", ( - f"sequence_packing only supports moe_token_dispatcher_type='alltoall', " - f"got '{self.moe_token_dispatcher_type}'" + assert self.moe_token_dispatcher_type in ("alltoall", "flex"), ( + f"sequence_packing only supports moe_token_dispatcher_type in " + f"('alltoall', 'flex'), got '{self.moe_token_dispatcher_type}'" ) supported_schedulers = ['dp_balanced', 'default_dynamic_cp'] diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py index dead2c0c12d..af1c4f924f9 100644 --- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py +++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py @@ -2,16 +2,24 @@ import copy import dataclasses +import math from types import SimpleNamespace import pytest import torch from megatron.core import config, parallel_state -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules +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.transformer.moe.moe_layer import MoELayer from megatron.core.transformer.moe.moe_utils import get_capacity from megatron.core.transformer.moe.token_dispatcher import MoETokenDispatcher +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 @@ -461,6 +469,171 @@ 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() + + @pytest.mark.skipif( not is_deep_ep_available() and not is_hybrid_ep_available(), reason="Deep EP and Hybrid EP are not available", @@ -470,6 +643,7 @@ def setup_method(self, method): pass def teardown_method(self, method): + reset_hybrid_ep_buffer() Utils.destroy_model_parallel() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")