[training] fix: Route batches to standalone MTP stages (#4208)

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/training/gpt_step.py

@@ -21,9 +21,13 @@
 from megatron.core import parallel_state
 from megatron.core.models.gpt import GPTModel
 from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage
+from megatron.core.transformer.enums import LayerType
+from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
 from megatron.core.utils import (
     get_batch_on_this_cp_rank,
     get_model_config,
+    get_pg_rank,
+    get_pg_size,
     is_te_min_version,
     unwrap_model,
 )
@@ -57,6 +61,50 @@ def _middle_pp_stage_needs_batch(cfg: ConfigContainer) -> bool:
     return uses_custom_attention_mask or _uses_packed_sequence_metadata(cfg)
 
 
+def _layout_stage_has_mtp(layout, *, pp_rank: int, pp_size: int, vp_stage: int) -> bool:
+    """Return whether a parsed or raw pipeline layout stage owns MTP layers."""
+    if isinstance(layout, str):
+        layout = PipelineParallelLayerLayout.from_str(layout, pp_size)
+
+    if isinstance(layout, PipelineParallelLayerLayout):
+        stage_layout = layout.layout[pp_rank][vp_stage]
+    elif isinstance(layout, list):
+        stage_layout = layout[vp_stage * pp_size + pp_rank]
+    else:
+        return False
+
+    return any(
+        layer == "mtp" or layer == LayerType.mtp or getattr(layer, "name", None) == "mtp" for layer in stage_layout
+    )
+
+
+def _current_pp_stage_has_mtp(cfg: ConfigContainer, *, pg_collection) -> bool:
+    """Return whether the current PP/VPP stage owns the configured MTP block."""
+    model_cfg = getattr(cfg, "model", None)
+    layout = getattr(model_cfg, "pipeline_model_parallel_layout", None)
+    if layout is None:
+        return False
+
+    pp_group = getattr(pg_collection, "pp", None)
+    pp_rank = get_pg_rank(pp_group)
+    pp_size = get_pg_size(pp_group)
+    vp_stage = parallel_state.get_virtual_pipeline_model_parallel_rank()
+    if vp_stage is None:
+        vp_stage = 0
+
+    return _layout_stage_has_mtp(layout, pp_rank=pp_rank, pp_size=pp_size, vp_stage=vp_stage)
+
+
+def _current_pp_stage_needs_mtp_inputs(cfg: ConfigContainer, *, pg_collection, is_last: bool) -> bool:
+    """Return whether this stage needs token ids for MTP embedding lookup."""
+    model_cfg = getattr(cfg, "model", None)
+    layout = getattr(model_cfg, "pipeline_model_parallel_layout", None)
+    if layout is None:
+        return is_last
+
+    return _current_pp_stage_has_mtp(cfg, pg_collection=pg_collection)
+
+
 def _partition_packed_batch_for_cp(batch: dict[str, torch.Tensor], cp_size: int) -> dict[str, torch.Tensor]:
     """Partition THD/packed batches across context-parallel ranks.
 
@@ -105,7 +153,7 @@ def _partition_packed_batch_for_cp(batch: dict[str, torch.Tensor], cp_size: int)
 
 def get_batch_from_iterator(
     data_iterator: Iterable,
-    use_mtp: bool = False,
+    include_mtp_inputs: bool = False,
     skip_getting_attention_mask_from_dataset: bool = True,
     *,
     is_first_pp_stage: bool,
@@ -116,7 +164,7 @@ def get_batch_from_iterator(
 
     Args:
         data_iterator: The data iterator to get the batch from.
-        use_mtp: Whether Multi-Token Prediction layers are enabled.
+        include_mtp_inputs: Whether this PP stage needs Multi-Token Prediction input tensors.
         skip_getting_attention_mask_from_dataset: If set, the dataset will pass a None attention mask.
         include_full_batch_fields: Whether to include all standard training tensors regardless of PP stage.
 
@@ -143,7 +191,7 @@ def get_batch_from_iterator(
             required_host_keys.add("cu_seqlens_unpadded_argmin")
 
     if not include_full_batch_fields:
-        if is_first_pp_stage or use_mtp:
+        if is_first_pp_stage or include_mtp_inputs:
             required_device_keys.update(("tokens", "position_ids"))
         if is_last_pp_stage:
             required_device_keys.update(("labels", "loss_mask"))
@@ -191,13 +239,18 @@ def get_batch(
     is_last = is_pp_last_stage(pg_collection.pp)
     is_middle = (not is_first) and (not is_last)
     include_full_batch_fields = is_middle and _middle_pp_stage_needs_batch(cfg)
-    if is_middle and not include_full_batch_fields:
+    include_mtp_inputs = use_mtp and _current_pp_stage_needs_mtp_inputs(
+        cfg, pg_collection=pg_collection, is_last=is_last
+    )
+    if is_middle and not include_full_batch_fields and not include_mtp_inputs:
         return None, None, None, None, None, None, None, None, None, None
 
     batch = get_batch_from_iterator(
         data_iterator,
-        use_mtp,
-        getattr(cfg.dataset, "skip_getting_attention_mask_from_dataset", True),
+        include_mtp_inputs=include_mtp_inputs,
+        skip_getting_attention_mask_from_dataset=getattr(
+            cfg.dataset, "skip_getting_attention_mask_from_dataset", True
+        ),
         is_first_pp_stage=is_first,
         is_last_pp_stage=is_last,
         include_full_batch_fields=include_full_batch_fields,

tests/unit_tests/training/test_gpt_step.py

@@ -46,23 +46,25 @@ def __next__(self):
 
 
 class _MockProcessGroup:
+    def __init__(self, rank=0, size=1):
+        self._rank = rank
+        self._size = size
+
     def rank(self):
-        return 0
+        return self._rank
 
     def size(self):
-        return 1
+        return self._size
 
 
 class _MockPGCollection:
-    def __init__(self, cp_size=1):
-        self.pp = _MockProcessGroup()
+    def __init__(self, cp_size=1, pp_rank=0, pp_size=1):
+        self.pp = _MockProcessGroup(rank=pp_rank, size=pp_size)
         self._cp_size = cp_size
 
     @property
     def cp(self):
-        pg = _MockProcessGroup()
-        pg.size = lambda: self._cp_size
-        return pg
+        return _MockProcessGroup(size=self._cp_size)
 
 
 class _NoCudaTensor(torch.Tensor):
@@ -74,7 +76,14 @@ def _as_nocuda(tensor):
     return tensor.as_subclass(_NoCudaTensor)
 
 
-def _make_cfg(*, packed_sequence_specs=None, skip_getting_attention_mask_from_dataset=True):
+def _make_cfg(
+    *,
+    packed_sequence_specs=None,
+    skip_getting_attention_mask_from_dataset=True,
+    pipeline_model_parallel_layout=None,
+    pipeline_model_parallel_size=1,
+    mtp_num_layers=0,
+):
     cfg = type("Cfg", (), {})()
     cfg.dataset = type(
         "D",
@@ -84,6 +93,15 @@ def _make_cfg(*, packed_sequence_specs=None, skip_getting_attention_mask_from_da
             "skip_getting_attention_mask_from_dataset": skip_getting_attention_mask_from_dataset,
         },
     )()
+    cfg.model = type(
+        "M",
+        (),
+        {
+            "pipeline_model_parallel_layout": pipeline_model_parallel_layout,
+            "pipeline_model_parallel_size": pipeline_model_parallel_size,
+            "mtp_num_layers": mtp_num_layers,
+        },
+    )()
     return cfg
 
 
@@ -92,6 +110,15 @@ def _set_middle_pp_stage(monkeypatch):
     monkeypatch.setattr("megatron.bridge.training.gpt_step.is_pp_last_stage", lambda pg: False)
 
 
+def _set_last_pp_stage(monkeypatch):
+    monkeypatch.setattr("megatron.bridge.training.gpt_step.is_pp_first_stage", lambda pg: False)
+    monkeypatch.setattr("megatron.bridge.training.gpt_step.is_pp_last_stage", lambda pg: True)
+
+
+def _set_distributed_initialized(monkeypatch):
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+
+
 class _NoopTimer:
     def __call__(self, *args, **kwargs):
         return self
@@ -187,6 +214,133 @@ def test_middle_pp_stage_keeps_non_packed_fast_path(self, monkeypatch):
         assert result == (None, None, None, None, None, None, None, None, None, None)
         data_iterator.__next__.assert_not_called()
 
+    def test_middle_pp_stage_without_mtp_keeps_fast_path_when_mtp_enabled(self, monkeypatch):
+        """Global MTP does not force ordinary middle PP stages to load a batch."""
+        _set_middle_pp_stage(monkeypatch)
+        _set_distributed_initialized(monkeypatch)
+        data_iterator = MagicMock()
+
+        result = get_batch(
+            data_iterator,
+            _make_cfg(
+                pipeline_model_parallel_layout=[["embedding", "decoder"], ["decoder"], ["mtp"], ["loss"]],
+                pipeline_model_parallel_size=4,
+                mtp_num_layers=1,
+            ),
+            use_mtp=True,
+            pg_collection=_MockPGCollection(pp_rank=1, pp_size=4),
+        )
+
+        assert result == (None, None, None, None, None, None, None, None, None, None)
+        data_iterator.__next__.assert_not_called()
+
+    def test_standalone_mtp_middle_pp_stage_loads_tokens_and_position_ids(self, monkeypatch):
+        """A middle PP stage that owns MTP receives input ids for MCore MTP."""
+        _set_middle_pp_stage(monkeypatch)
+        _set_distributed_initialized(monkeypatch)
+        monkeypatch.setattr(
+            "megatron.bridge.training.gpt_step.get_batch_on_this_cp_rank",
+            lambda batch, is_hybrid_cp=False, cp_group=None, hybrid_cp_group_func=None: batch,
+        )
+        monkeypatch.setattr(
+            "megatron.bridge.training.gpt_step.parallel_state.get_virtual_pipeline_model_parallel_rank",
+            lambda: None,
+        )
+
+        tokens = _as_nocuda(torch.tensor([[1, 2, 3, 4]]))
+        labels = _as_nocuda(torch.tensor([[2, 3, 4, 5]]))
+        loss_mask = _as_nocuda(torch.ones(1, 4))
+        position_ids = _as_nocuda(torch.arange(4).unsqueeze(0))
+        batch = {
+            "tokens": tokens,
+            "labels": labels,
+            "loss_mask": loss_mask,
+            "attention_mask": None,
+            "position_ids": position_ids,
+        }
+
+        (
+            out_tokens,
+            out_labels,
+            out_loss_mask,
+            out_attention_mask,
+            out_position_ids,
+            out_cu_seqlens,
+            out_cu_seqlens_argmin,
+            out_max_seqlen,
+            out_cu_seqlens_unpadded,
+            out_cu_seqlens_unpadded_argmin,
+        ) = get_batch(
+            _Iterator(batch),
+            _make_cfg(
+                pipeline_model_parallel_layout=[["embedding", "decoder"], ["decoder"], ["mtp"], ["loss"]],
+                pipeline_model_parallel_size=4,
+                mtp_num_layers=1,
+            ),
+            use_mtp=True,
+            pg_collection=_MockPGCollection(pp_rank=2, pp_size=4),
+        )
+
+        assert torch.equal(out_tokens, tokens)
+        assert out_labels is None
+        assert out_loss_mask is None
+        assert out_attention_mask is None
+        assert torch.equal(out_position_ids, position_ids)
+        assert out_cu_seqlens is None
+        assert out_cu_seqlens_argmin is None
+        assert out_max_seqlen is None
+        assert out_cu_seqlens_unpadded is None
+        assert out_cu_seqlens_unpadded_argmin is None
+
+    def test_standalone_mtp_loss_stage_skips_mtp_inputs(self, monkeypatch):
+        """The loss-only final PP stage does not load token ids for standalone MTP."""
+        _set_last_pp_stage(monkeypatch)
+        _set_distributed_initialized(monkeypatch)
+        monkeypatch.setattr(
+            "megatron.bridge.training.gpt_step.get_batch_on_this_cp_rank",
+            lambda batch, is_hybrid_cp=False, cp_group=None, hybrid_cp_group_func=None: batch,
+        )
+        monkeypatch.setattr(
+            "megatron.bridge.training.gpt_step.parallel_state.get_virtual_pipeline_model_parallel_rank",
+            lambda: None,
+        )
+
+        tokens = _as_nocuda(torch.tensor([[1, 2, 3, 4]]))
+        labels = _as_nocuda(torch.tensor([[2, 3, 4, 5]]))
+        loss_mask = _as_nocuda(torch.ones(1, 4))
+        position_ids = _as_nocuda(torch.arange(4).unsqueeze(0))
+        batch = {
+            "tokens": tokens,
+            "labels": labels,
+            "loss_mask": loss_mask,
+            "attention_mask": None,
+            "position_ids": position_ids,
+        }
+
+        (
+            out_tokens,
+            out_labels,
+            out_loss_mask,
+            out_attention_mask,
+            out_position_ids,
+            *_,
+        ) = get_batch(
+            _Iterator(batch),
+            _make_cfg(
+                pipeline_model_parallel_layout=[["embedding", "decoder"], ["decoder"], ["mtp"], ["loss"]],
+                pipeline_model_parallel_size=4,
+                mtp_num_layers=1,
+            ),
+            use_mtp=True,
+            pg_collection=_MockPGCollection(pp_rank=3, pp_size=4),
+        )
+
+        assert out_tokens is None
+        assert torch.equal(out_labels, labels)
+        assert torch.equal(out_loss_mask, loss_mask)
+        assert out_attention_mask is None
+        assert out_position_ids is None
+
     def test_forward_common_passes_packed_seq_params_on_middle_pp_stage(self, monkeypatch):
         """Forward path must pass packed metadata on middle PP stages."""
         sentinel_packed_seq_params = object()