refactor

Signed-off-by: HaochenYuan <haocheny@nvidia.com>

Author: HaochenYuan

megatron/core/model_parallel_config.py

@@ -62,13 +62,6 @@ class ModelParallelConfig:
     each rank when using hybrid_context_parallel.
     """
 
-    sequence_packing_scheduler: Optional[Literal['dp_balanced', 'default_dynamic_cp']] = None
-    """
-    Scheduler for sequence packing.
-    dp_balanced: DP-balanced scheduler for sequence packing.
-    default_dynamic_cp: Dynamic-CP scheduler for packed sequence balancing.
-    """
-
     hybrid_context_parallel: bool = False
     """
     If true, enables hybrid context parallel. This is used to balance the workload of 

megatron/core/transformer/moe/token_dispatcher.py

@@ -1033,9 +1033,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.
+        # HybridEP dispatch expects equal per-rank input sizes. When requested,
+        # variable token counts 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
 
@@ -1044,11 +1043,7 @@ 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.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(
@@ -1063,7 +1058,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
 
         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

megatron/core/transformer/transformer_config.py

@@ -801,8 +801,8 @@ class TransformerConfig(ModelParallelConfig):
 
     moe_hybridep_pad_variable_tokens: bool = False
     """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 only when HybridEP inputs can have different token counts across ranks.
+    Leave disabled when inputs are already statically padded to equal token counts.
     """
 
     moe_per_layer_logging: bool = False

tests/unit_tests/transformer/moe/test_token_dispatcher.py

@@ -1,22 +1,16 @@
 # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
 
 import dataclasses
-import math
 
 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,
-    get_gpt_layer_with_transformer_engine_spec,
-)
-from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules
 from megatron.core.transformer.moe.fused_a2a import 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.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
@@ -425,173 +419,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:
-        import transformer_engine_torch as tex
-
-        cu_seqlens_padded = _to_cu_seqlens(padded_seqlens)
-        cp_rank = parallel_state.get_context_parallel_rank()
-        index = tex.thd_get_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",