[codex] Propagate SBHD padding masks across pipeline stages

megatron/training/datasets/varlen_dataset.py

@@ -406,10 +406,15 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
             loss_mask = torch.ones(max_len, dtype=torch.float32)
             loss_mask[valid_len:] = 0.0  # mask the right-padded tail by position
             loss_mask[labels == IGNORE_INDEX] = 0.0
+            # Keep physical padding separate from the LM loss mask: prompt
+            # tokens may be loss-masked but must still participate in MoE.
+            padding_mask = torch.zeros(max_len, dtype=torch.bool)
+            padding_mask[valid_len:] = True
             return {
                 'tokens': input_ids,
                 'labels': labels,
                 'loss_mask': loss_mask,
+                'padding_mask': padding_mask,
                 'position_ids': torch.arange(max_len, dtype=torch.int64),
             }
 

megatron/training/utils/common_utils.py

@@ -530,9 +530,13 @@ def get_blend_and_blend_per_split(args):
     return blend, blend_per_split
 
 
-def get_batch_on_this_tp_rank(data_iterator, mtp_on_this_rank: bool = False):
+def get_batch_on_this_tp_rank(
+    data_iterator, mtp_on_this_rank: bool = False, needs_padding_mask: bool = False
+):
 
     args = get_args()
+    # Optional input structure must be identical across TP ranks and remain
+    # static during full-iteration CUDA Graph capture.
 
     def _broadcast(item):
         if item is not None:
@@ -555,6 +559,9 @@ def _broadcast(item):
                 if "attention_mask" not in data
                 else data["attention_mask"].cuda(non_blocking=True)
             ),
+            'padding_mask': (
+                data["padding_mask"].cuda(non_blocking=True) if needs_padding_mask else None
+            ),
             'position_ids': data["position_ids"].cuda(non_blocking=True),
             'cu_seqlens': (
                 None if "cu_seqlens" not in data else data["cu_seqlens"].cuda(non_blocking=True)
@@ -599,6 +606,7 @@ def _broadcast_cu_seqlens(cu_seqlens):
             _broadcast(batch['tokens'])
             _broadcast(batch['labels'])
             _broadcast(batch['loss_mask'])
+            _broadcast(batch['padding_mask'])
             _broadcast(batch['attention_mask'])
             _broadcast(batch['position_ids'])
             _broadcast_cu_seqlens(batch['cu_seqlens'])
@@ -607,6 +615,7 @@ def _broadcast_cu_seqlens(cu_seqlens):
 
         elif mpu.is_pipeline_first_stage():
             _broadcast(batch['tokens'])
+            _broadcast(batch['padding_mask'])
             _broadcast(batch['attention_mask'])
             _broadcast(batch['position_ids'])
             _broadcast_cu_seqlens(batch['cu_seqlens'])
@@ -618,8 +627,22 @@ def _broadcast_cu_seqlens(cu_seqlens):
             # to broadcast tokens and position_ids to all of the tensor parallel ranks on the last stage.
             _broadcast(batch['labels'])
             _broadcast(batch['loss_mask'])
+            _broadcast(batch['padding_mask'])
             _broadcast(batch['attention_mask'])
 
+        else:
+            # SBHD validation needs physical padding metadata on intermediate
+            # stages because those stages may also contain MoE layers.
+            _broadcast(batch['padding_mask'])
+            batch['tokens'] = None
+            batch['labels'] = None
+            batch['loss_mask'] = None
+            batch['attention_mask'] = None
+            batch['position_ids'] = None
+            batch['cu_seqlens'] = None
+            batch['max_seqlen'] = None
+            batch['local_cp_size'] = None
+
     else:
         if args.dynamic_context_parallel:
             seq_len = torch.tensor(0, dtype=torch.int32, device=torch.cuda.current_device())
@@ -631,6 +654,11 @@ def _broadcast_cu_seqlens(cu_seqlens):
         tokens = torch.empty(shape, dtype=torch.int64, device=torch.cuda.current_device())
         labels = torch.empty(shape, dtype=torch.int64, device=torch.cuda.current_device())
         loss_mask = torch.empty(shape, dtype=torch.float32, device=torch.cuda.current_device())
+        padding_mask = (
+            torch.empty(shape, dtype=torch.bool, device=torch.cuda.current_device())
+            if needs_padding_mask
+            else None
+        )
         if args.create_attention_mask_in_dataloader:
             shape_attention_mask = (
                 (args.micro_batch_size, 1, args.seq_length, args.seq_length)
@@ -676,6 +704,7 @@ def _broadcast_cu_seqlens():
             _broadcast(tokens)
             _broadcast(labels)
             _broadcast(loss_mask)
+            _broadcast(padding_mask)
             _broadcast(attention_mask)
             _broadcast(position_ids)
             cu_seqlens = _broadcast_cu_seqlens()
@@ -687,6 +716,7 @@ def _broadcast_cu_seqlens():
             loss_mask = None
 
             _broadcast(tokens)
+            _broadcast(padding_mask)
             _broadcast(attention_mask)
             _broadcast(position_ids)
             cu_seqlens = _broadcast_cu_seqlens()
@@ -703,12 +733,26 @@ def _broadcast_cu_seqlens():
 
             _broadcast(labels)
             _broadcast(loss_mask)
+            _broadcast(padding_mask)
             _broadcast(attention_mask)
 
+        else:
+            tokens = None
+            labels = None
+            loss_mask = None
+            attention_mask = None
+            position_ids = None
+            cu_seqlens = None
+            max_seqlen = None
+            local_cp_size = None
+
+            _broadcast(padding_mask)
+
         batch = {
             'tokens': tokens,
             'labels': labels,
             'loss_mask': loss_mask,
+            'padding_mask': padding_mask,
             'attention_mask': attention_mask,
             'position_ids': position_ids,
             'cu_seqlens': cu_seqlens,

pretrain_gpt.py

@@ -143,9 +143,11 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None):
 
     # TODO: this is pretty hacky, find a better way
     is_packed_sequence = args.sft or (args.use_varlen_dataset and not args.varlen_sbhd_validation)
+    needs_padding_mask = args.use_varlen_dataset and args.varlen_sbhd_validation
     if (
         not is_first_or_last_pipeline_stage(vp_stage)
         and not is_packed_sequence
+        and not needs_padding_mask
         and ((not mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage)))
     ):
         return None, None, None, None, None, None, None
@@ -154,6 +156,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None):
     batch = get_batch_on_this_tp_rank(
         data_iterator,
         mtp_on_this_rank=mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage),
+        needs_padding_mask=needs_padding_mask,
     )
 
     cu_seqlens = batch.pop('cu_seqlens', None)
@@ -199,6 +202,8 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None):
 
     # Pad the already-packed THD tensors at the end when requested. A configured
     # thd_max_packed_sequences also pads cu_seqlens to a fixed capacity in eager or graph mode.
+    # SBHD validation samples carry physical right-padding metadata. CP has
+    # already partitioned it with the other sequence-dimension tensors.
     padding_mask = batch.pop('padding_mask', None)
     if config.pad_packed_seq_alignment is not None and packed_seq_params is not None:
         tokens = batch.get('tokens', None)
@@ -400,7 +405,12 @@ def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False):
     config = core_transformer_config_from_args(args)
     if mpu.get_tensor_model_parallel_rank() != 0:
         return False
-    elif is_packed_sequence:
+    elif is_packed_sequence or (
+        getattr(args, 'use_varlen_dataset', False)
+        and getattr(args, 'varlen_sbhd_validation', False)
+    ):
+        # Packed THD and SBHD validation both need padding metadata on every
+        # pipeline stage so each MoE layer excludes physical padding.
         return True
     return is_first_or_last_pipeline_stage(vp_stage) or mtp_on_this_rank(
         config, ignore_virtual=False, vp_stage=vp_stage

tests/unit_tests/data/test_get_batch.py

@@ -2,6 +2,7 @@
 
 import os
 import sys
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -153,6 +154,63 @@ def create_sft_data_iterator(max_seq_length: int = 1024):
     return iter([batch]), num_real_tokens
 
 
+@pytest.mark.parametrize("tp_rank", [0, 1])
+@pytest.mark.parametrize("needs_padding_mask", [False, True])
+def test_get_batch_on_this_tp_rank_full_iteration_cudagraph_safe(
+    monkeypatch, tp_rank, needs_padding_mask
+):
+    """Optional batch inputs must not require host/device transfers during capture."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA Graph capture requires CUDA")
+
+    from megatron.training.utils import common_utils
+
+    torch.cuda.set_device(Utils.rank % torch.cuda.device_count())
+    device = torch.cuda.current_device()
+    shape = (1, 4)
+    data = {
+        "tokens": torch.ones(shape, dtype=torch.int64, device=device),
+        "labels": torch.ones(shape, dtype=torch.int64, device=device),
+        "loss_mask": torch.ones(shape, dtype=torch.float32, device=device),
+        "position_ids": torch.arange(shape[1], dtype=torch.int64, device=device).view(shape),
+    }
+    if needs_padding_mask:
+        data["padding_mask"] = torch.zeros(shape, dtype=torch.bool, device=device)
+
+    args = SimpleNamespace(
+        create_attention_mask_in_dataloader=False,
+        cuda_graph_impl="full_iteration",
+        dynamic_context_parallel=False,
+        micro_batch_size=shape[0],
+        pipeline_model_parallel_size=1,
+        seq_length=shape[1],
+        sft=False,
+    )
+    monkeypatch.setattr(common_utils, "get_args", lambda: args)
+    monkeypatch.setattr(common_utils.mpu, "get_tensor_model_parallel_rank", lambda: tp_rank)
+    monkeypatch.setattr(common_utils.mpu, "get_tensor_model_parallel_src_rank", lambda: 0)
+    monkeypatch.setattr(common_utils.mpu, "get_tensor_model_parallel_group", lambda: None)
+    monkeypatch.setattr(torch.distributed, "broadcast", lambda *args, **kwargs: None)
+
+    data_iterator = iter([data]) if tp_rank == 0 else None
+    marker = torch.zeros((), device=device)
+    marker.add_(1)
+    marker.zero_()
+    graph = torch.cuda.CUDAGraph()
+    torch.cuda.synchronize()
+    with torch.cuda.graph(graph):
+        marker.add_(1)
+        batch = common_utils.get_batch_on_this_tp_rank(
+            data_iterator, needs_padding_mask=needs_padding_mask
+        )
+    graph.replay()
+    torch.cuda.synchronize()
+
+    # Capture records CUDA work without executing it; one replay increments once.
+    assert marker.item() == 1
+    assert (batch["padding_mask"] is not None) == needs_padding_mask
+
+
 @pytest.mark.parametrize("tp_size", [1, 2, 4])
 @pytest.mark.parametrize("pp_size", [1, 2, 4])
 @pytest.mark.parametrize("cp_size", [1, 2, 4])

tests/unit_tests/data/test_varlen_dataset.py

@@ -569,12 +569,90 @@ def test_getitem_sbhd_pads_to_seq_length_and_masks_tail():
     ds = _make_varlen(["abc"], _make_config(tok, seq_length=8, sbhd=True))
     out = ds[0]
     # SBHD emits fixed [seq_length] samples with no packing metadata.
-    assert set(out) == {"tokens", "labels", "loss_mask", "position_ids"}
+    assert set(out) == {"tokens", "labels", "loss_mask", "padding_mask", "position_ids"}
     assert out["tokens"].numel() == 8
     loss_mask = out["loss_mask"].tolist()
+    padding_mask = out["padding_mask"].tolist()
     # tokens=[a,b,c,eod]: valid_len=3 -> first 3 kept (incl. real eod), rest masked.
     assert loss_mask[0:3] == [1.0, 1.0, 1.0]
     assert all(v == 0.0 for v in loss_mask[3:])
+    assert padding_mask[0:3] == [False, False, False]
+    assert all(padding_mask[3:])
+    assert out["padding_mask"].dtype == torch.bool
+
+
+def test_sbhd_padding_mask_is_partitioned_with_tokens(monkeypatch):
+    """CP zigzag slicing must select identical token and padding-mask positions."""
+    from megatron.core.utils import get_batch_on_this_cp_rank
+
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda group: 4)
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda group: 1)
+
+    tokens = torch.arange(16, dtype=torch.int64).view(1, 16)
+    padding_mask = tokens >= 10
+    batch = {"tokens": tokens.clone(), "padding_mask": padding_mask.clone()}
+    result = get_batch_on_this_cp_rank(batch, cp_group=object())
+
+    expected_indices = torch.tensor([2, 3, 12, 13])
+    assert torch.equal(result["tokens"], tokens.index_select(1, expected_indices))
+    assert torch.equal(result["padding_mask"], padding_mask.index_select(1, expected_indices))
+
+
+def test_sbhd_get_batch_returns_dataset_padding_mask(monkeypatch):
+    """The dataset padding mask must survive the pretrain_gpt batch handoff."""
+    import pretrain_gpt
+
+    padding_mask = torch.tensor([[False, False, True, True]], dtype=torch.bool)
+    source_batch = {
+        "tokens": torch.tensor([[1, 2, 0, 0]], dtype=torch.int64),
+        "labels": torch.tensor([[2, 0, 0, 0]], dtype=torch.int64),
+        "loss_mask": torch.tensor([[1.0, 1.0, 0.0, 0.0]]),
+        "padding_mask": padding_mask,
+        "attention_mask": None,
+        "position_ids": torch.arange(4, dtype=torch.int64).view(1, 4),
+    }
+    args = SimpleNamespace(
+        sequence_packing_scheduler=None,
+        sft=False,
+        use_varlen_dataset=True,
+        varlen_sbhd_validation=True,
+        dynamic_context_parallel=False,
+    )
+    config = SimpleNamespace(
+        virtual_pipeline_model_parallel_size=None, pad_packed_seq_alignment=None
+    )
+
+    monkeypatch.setattr(pretrain_gpt, "get_args", lambda: args)
+    monkeypatch.setattr(pretrain_gpt, "core_transformer_config_from_args", lambda _: config)
+    # Exercise an intermediate PP stage: it must not take the legacy early return.
+    monkeypatch.setattr(pretrain_gpt, "is_first_or_last_pipeline_stage", lambda _: False)
+    monkeypatch.setattr(pretrain_gpt, "mtp_on_this_rank", lambda *args, **kwargs: False)
+
+    def get_batch_on_this_tp_rank(*args, **kwargs):
+        assert kwargs["needs_padding_mask"] is True
+        return source_batch.copy()
+
+    monkeypatch.setattr(pretrain_gpt, "get_batch_on_this_tp_rank", get_batch_on_this_tp_rank)
+    monkeypatch.setattr(pretrain_gpt, "get_batch_on_this_cp_rank", lambda batch: batch)
+
+    *_, returned_padding_mask = pretrain_gpt.get_batch(iter(()))
+    assert torch.equal(returned_padding_mask, padding_mask)
+
+
+def test_sbhd_dataset_is_built_on_intermediate_pipeline_stage(monkeypatch):
+    """Every PP stage needs SBHD padding metadata for its local MoE layers."""
+    import pretrain_gpt
+
+    args = SimpleNamespace(use_varlen_dataset=True, varlen_sbhd_validation=True)
+    monkeypatch.setattr(pretrain_gpt, "get_args", lambda: args)
+    monkeypatch.setattr(
+        pretrain_gpt, "core_transformer_config_from_args", lambda _: SimpleNamespace()
+    )
+    monkeypatch.setattr(pretrain_gpt.mpu, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(pretrain_gpt, "is_first_or_last_pipeline_stage", lambda _: False)
+    monkeypatch.setattr(pretrain_gpt, "mtp_on_this_rank", lambda *args, **kwargs: False)
+
+    assert pretrain_gpt.is_dataset_built_on_rank() is True
 
 
 def test_mock_getitem_thd_keys_and_pad_fallback():
@@ -724,6 +802,8 @@ def test_sbhd_validation_dataloader_uses_default_collate():
         assert batch["tokens"].shape == (mbs, seq_len)
         assert batch["labels"].shape == (mbs, seq_len)
         assert batch["loss_mask"].shape == (mbs, seq_len)
+        assert batch["padding_mask"].shape == (mbs, seq_len)
+        assert batch["padding_mask"].dtype == torch.bool
     finally:
         destroy_global_vars()
         Utils.destroy_model_parallel()