Add inter-document attention masking to GPTDataset

megatron/core/datasets/gpt_dataset.py

@@ -76,6 +76,10 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig):
     context_parallel_size: Optional[int] = None
     """The size of the context parallel group. Needed for padding in packed sequences."""
 
+    inter_document_masking: bool = False
+    """When True, return cu_seqlens marking document boundaries within each sample so
+    that attention is restricted to individual documents."""
+
     def __post_init__(self) -> None:
         """Do asserts and set fields post init"""
         super().__post_init__()
@@ -233,9 +237,9 @@ def __getitem__(self, idx: Optional[int]) -> Dict[str, torch.Tensor]:
         """
         if idx is None:
             # Batch padding sequence so the index does not matter
-            text, _ = self._query_document_sample_shuffle_indices(0)
+            text, _, document_lengths = self._query_document_sample_shuffle_indices(0)
         else:
-            text, _ = self._query_document_sample_shuffle_indices(idx)
+            text, _, document_lengths = self._query_document_sample_shuffle_indices(idx)
 
         text = torch.from_numpy(text).long()
         if self.config.add_extra_token_to_sequence:
@@ -279,32 +283,83 @@ def __getitem__(self, idx: Optional[int]) -> Dict[str, torch.Tensor]:
         if idx is None:
             loss_mask = torch.zeros_like(loss_mask)
 
-        if self.config.create_attention_mask:
-            return {
+        if self.config.inter_document_masking:
+            # document_lengths come from _query_document_sample_shuffle_indices
+            # which fetches sequence_length + add_extra_token_to_sequence tokens
+            # total. The extra token is appended to the last document part (used
+            # to produce the shifted labels), so subtract it before computing
+            # cu_seqlens which should index into the sequence_length-sized tokens
+            # tensor.
+            if self.config.add_extra_token_to_sequence:
+                document_lengths[-1] -= 1
+                if document_lengths[-1] == 0:
+                    document_lengths.pop()
+            # If the sample was padded (e.g., the last validation sample),
+            # fold the padding into the last document so cu_seqlens[-1]
+            # equals sequence_length.
+            shortfall = self.config.sequence_length - sum(document_lengths)
+            if shortfall > 0:
+                if document_lengths:
+                    document_lengths[-1] += shortfall
+                else:
+                    document_lengths.append(shortfall)
+            cu_seqlens = torch.tensor(numpy.cumsum([0] + document_lengths), dtype=torch.int32)
+
+            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+
+            # Reset position IDs per document.
+            position_ids = position_ids.clone()
+            for i in range(1, cu_seqlens.numel()):
+                start = cu_seqlens[i - 1].item()
+                end = cu_seqlens[i].item()
+                position_ids[start:end] = torch.arange(end - start, dtype=torch.long)
+
+            # Pad cu_seqlens to a fixed length so that default_collate can
+            # stack samples with different numbers of documents. Trailing
+            # entries are filled with sequence_length; the merge helper
+            # strips them later.
+            padded_cu_seqlens = torch.full(
+                (self.config.sequence_length + 1,), self.config.sequence_length, dtype=torch.int32
+            )
+            padded_cu_seqlens[: cu_seqlens.numel()] = cu_seqlens
+
+            result = {
+                "tokens": tokens,
+                "labels": labels,
+                "loss_mask": loss_mask,
+                "position_ids": position_ids,
+                "cu_seqlens": padded_cu_seqlens,
+                "max_seqlen": max_seqlen,
+            }
+        elif self.config.create_attention_mask:
+            result = {
                 "tokens": tokens,
                 "labels": labels,
                 "attention_mask": attention_mask,
                 "loss_mask": loss_mask,
                 "position_ids": position_ids,
             }
         else:
-            return {
+            result = {
                 "tokens": tokens,
                 "labels": labels,
                 "loss_mask": loss_mask,
                 "position_ids": position_ids,
             }
 
+        return result
+
     def _query_document_sample_shuffle_indices(
         self, idx: int
-    ) -> Tuple[numpy.ndarray, numpy.ndarray]:
+    ) -> Tuple[numpy.ndarray, numpy.ndarray, list]:
         """Get the text (token ids) and document ids for a given index
 
         Args:
             idx (int): The index into the dataset
 
         Returns:
-            Tuple[numpy.ndarray, numpy.ndarray]: The text ids and document ids
+            Tuple[numpy.ndarray, numpy.ndarray, list]: The text ids, document ids,
+                and per-document token counts (before any padding).
         """
         if self.shuffle_index is None:
             # NOTE(asolergi-nv): Lazy memmap the indexes
@@ -366,6 +421,8 @@ def _query_document_sample_shuffle_indices(
 
         length = sum(map(len, sample_parts))
 
+        document_lengths = [len(p) for p in sample_parts]
+
         # Pad the sample if necessary
         if length < (self.config.sequence_length + self.config.add_extra_token_to_sequence):
             sample_parts.append(
@@ -376,6 +433,7 @@ def _query_document_sample_shuffle_indices(
         return (
             numpy.concatenate(sample_parts, dtype=numpy.int64),
             numpy.array(document_ids, dtype=numpy.int64),
+            document_lengths,
         )
 
     def _build_document_sample_shuffle_indices(

megatron/core/utils.py

@@ -2038,7 +2038,7 @@ def is_submodule(module, parent_module, strict=True):
 
 def get_batch_on_this_tp_rank(
     batch: dict[str, torch.Tensor],
-    is_sft: bool,
+    has_cu_seqlens: bool,
     is_hybrid_cp: bool,
     create_attention_mask_in_dataloader: bool,
     broadcast_src_rank: int,
@@ -2073,8 +2073,8 @@ def get_batch_on_this_tp_rank(
         batch (dict[str, torch.Tensor]): The batch dict. On TP rank 0 this
             contains the actual data; on other ranks it is ignored (receive
             buffers are allocated internally).
-        is_sft (bool): Whether this is an SFT (supervised fine-tuning) run
-            using THD packed sequences.
+        has_cu_seqlens (bool): Whether the batch contains cu_seqlens and
+            max_seqlen metadata (e.g., SFT or --dataloader-inter-document-masking).
         is_hybrid_cp (bool): Whether hybrid context parallelism is enabled.
         create_attention_mask_in_dataloader (bool): Whether the dataloader
             creates an explicit attention mask tensor.
@@ -2131,7 +2131,7 @@ def _broadcast_cu_seqlens(cu_seqlens):
             _broadcast(batch['labels'])
             _broadcast(batch['loss_mask'])
             _broadcast(batch['position_ids'])
-            if is_sft or is_hybrid_cp:
+            if has_cu_seqlens or is_hybrid_cp:
                 _broadcast_cu_seqlens(batch['cu_seqlens'])
                 _broadcast(batch['max_seqlen'])
                 if cp_size > 1:
@@ -2147,7 +2147,7 @@ def _broadcast_cu_seqlens(cu_seqlens):
 
             _broadcast(batch['tokens'])
             _broadcast(batch['position_ids'])
-            if is_sft:
+            if has_cu_seqlens:
                 _broadcast_cu_seqlens(batch['cu_seqlens'])
                 _broadcast(batch['max_seqlen'])
                 if cp_size > 1:
@@ -2161,16 +2161,16 @@ def _broadcast_cu_seqlens(cu_seqlens):
 
             _broadcast(batch['labels'])
             _broadcast(batch['loss_mask'])
-            if is_sft:
+            if has_cu_seqlens:
                 _broadcast_cu_seqlens(batch['cu_seqlens'])
                 _broadcast(batch['max_seqlen'])
                 if cp_size > 1:
                     _broadcast_cu_seqlens(batch['cu_seqlens_padded'])
             if create_attention_mask_in_dataloader:
                 _broadcast(batch['attention_mask'])
 
-        elif is_sft:
-            # NOTE(asolergi-nv): Broadcast required THD metadata for SFT to intermediate stages
+        elif has_cu_seqlens:
+            # NOTE(asolergi-nv): Broadcast required THD metadata to intermediate stages.
             batch["tokens"] = None
             batch["labels"] = None
             batch["loss_mask"] = None
@@ -2202,7 +2202,7 @@ def _broadcast_cu_seqlens(cu_seqlens):
         attention_mask = None
         local_cp_size = None
 
-        if is_sft or is_hybrid_cp:
+        if has_cu_seqlens or is_hybrid_cp:
             max_seqlen = torch.empty(1, dtype=torch.int32, device=torch.cuda.current_device())
         if create_attention_mask_in_dataloader:
             attention_mask = torch.empty(
@@ -2242,7 +2242,7 @@ def _broadcast_cu_seqlens():
             _broadcast(labels)
             _broadcast(loss_mask)
             _broadcast(position_ids)
-            if is_sft or is_hybrid_cp:
+            if has_cu_seqlens or is_hybrid_cp:
                 cu_seqlens = _broadcast_cu_seqlens()
                 _broadcast(max_seqlen)
                 if cp_size > 1:
@@ -2258,7 +2258,7 @@ def _broadcast_cu_seqlens():
 
             _broadcast(tokens)
             _broadcast(position_ids)
-            if is_sft:
+            if has_cu_seqlens:
                 cu_seqlens = _broadcast_cu_seqlens()
                 _broadcast(max_seqlen)
                 if cp_size > 1:
@@ -2272,16 +2272,16 @@ def _broadcast_cu_seqlens():
 
             _broadcast(labels)
             _broadcast(loss_mask)
-            if is_sft:
+            if has_cu_seqlens:
                 cu_seqlens = _broadcast_cu_seqlens()
                 _broadcast(max_seqlen)
                 if cp_size > 1:
                     cu_seqlens_padded = _broadcast_cu_seqlens()
             if create_attention_mask_in_dataloader:
                 _broadcast(attention_mask)
 
-        elif is_sft:
-            # NOTE(asolergi-nv): Broadcast required THD metadata for SFT to intermediate stages
+        elif has_cu_seqlens:
+            # NOTE(asolergi-nv): Broadcast required THD metadata to intermediate stages.
             tokens = None
             labels = None
             loss_mask = None
@@ -2524,50 +2524,55 @@ def get_batch_on_this_cp_rank(
     is_hybrid_cp: bool,
     cp_group: Optional[torch.distributed.ProcessGroup] = None,
     hybrid_cp_group_func: Optional[Callable[[int], torch.distributed.ProcessGroup]] = None,
+    use_per_sequence_balancing: bool = False,
 ):
     """Dispatch batch partitioning across context-parallel ranks.
 
     Routes to the appropriate CP partitioning strategy based on the batch
     contents and parallelism mode:
+      - **Per-sequence zigzag**: When ``cu_seqlens`` is None, or when
+        ``use_per_sequence_balancing`` is True, delegates to
+        ``_get_batch_on_this_cp_rank_per_sequence_balancing``.
       - **Per-document zigzag**: When ``cu_seqlens`` is present and
         ``is_hybrid_cp`` is False, delegates to
         ``_get_batch_on_this_cp_rank_per_document_balancing``.
       - **Hybrid CP**: When ``cu_seqlens`` is present and ``is_hybrid_cp`` is
         True, creates a local hybrid CP group (via ``hybrid_cp_group_func``)
         and delegates to ``_get_batch_on_this_cp_rank_per_sequence_balancing``.
-      - **Per-sequence zigzag**: When ``cu_seqlens`` is None, delegates to
-        ``_get_batch_on_this_cp_rank_per_sequence_balancing``.
 
     Args:
         batch (Dict[str, Any]): Input batch tensors. Must contain a
             'cu_seqlens' key (may be None for pretraining).
         is_hybrid_cp (bool): Whether hybrid context parallelism is enabled.
         cp_group (Optional[torch.distributed.ProcessGroup]): Context-parallel
-            process group used for SFT and pretraining CP partitioning.
+            process group used for CP partitioning.
         hybrid_cp_group_func (Optional[Callable[[int], torch.distributed.ProcessGroup]]):
             Factory function that returns a hybrid CP process group for a given
             ``group_size``. Required when ``is_hybrid_cp`` is True.
+        use_per_sequence_balancing (bool): When True, use per-sequence zigzag
+            even when ``cu_seqlens`` is present (e.g., for inter-document
+            masking where document lengths are not divisible by
+            ``2 * cp_size``).
 
     Returns:
         Dict[str, Any]: The batch with sequence-dimension tensors partitioned
         to this CP rank.
     """
 
-    if batch.get("cu_seqlens") is not None:  # NOTE(asolergi-nv): SFT & HybridCP case
-        if is_hybrid_cp:
-            assert (
-                batch['local_cp_size'] is not None
-            ), "local_cp_size is required for hybrid context parallel"
-            if batch['local_cp_size'].item() > 1:
-                hybrid_cp_group = hybrid_cp_group_func(group_size=batch['local_cp_size'].item())
-                batch = _get_batch_on_this_cp_rank_per_sequence_balancing(
-                    batch, cp_group=hybrid_cp_group
-                )
-                batch["hybrid_cp_group"] = hybrid_cp_group
-        else:
-            batch = _get_batch_on_this_cp_rank_per_document_balancing(batch, cp_group=cp_group)
-    else:  # NOTE(asolergi-nv): Pretrain case
+    if use_per_sequence_balancing or batch.get("cu_seqlens") is None:
         batch = _get_batch_on_this_cp_rank_per_sequence_balancing(batch, cp_group=cp_group)
+    elif is_hybrid_cp:
+        assert (
+            batch['local_cp_size'] is not None
+        ), "local_cp_size is required for hybrid context parallel"
+        if batch['local_cp_size'].item() > 1:
+            hybrid_cp_group = hybrid_cp_group_func(group_size=batch['local_cp_size'].item())
+            batch = _get_batch_on_this_cp_rank_per_sequence_balancing(
+                batch, cp_group=hybrid_cp_group
+            )
+            batch["hybrid_cp_group"] = hybrid_cp_group
+    else:
+        batch = _get_batch_on_this_cp_rank_per_document_balancing(batch, cp_group=cp_group)
     return batch
 
 

megatron/elastification/pretrain_hybrid_flex.py

@@ -178,6 +178,7 @@ def get_batch(data_iterator, vp_stage=None):
     cp_size = args.context_parallel_size
     tp_rank = mpu.get_tensor_model_parallel_rank()
     is_sft = args.sft
+    has_cu_seqlens = is_sft or getattr(args, 'dataloader_inter_document_masking', False)
     is_hybrid_cp = args.hybrid_context_parallel
     mtp_on_this_rank = mtp_on_this_rank_func(
         layout=config.pipeline_model_parallel_layout,
@@ -186,7 +187,7 @@ def get_batch(data_iterator, vp_stage=None):
         vp_stage=vp_stage,
     )
 
-    if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank and not is_sft:
+    if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank and not has_cu_seqlens:
         return None, None, None, None, None, None, None
 
     batch = {}
@@ -203,7 +204,7 @@ def get_batch(data_iterator, vp_stage=None):
         batch,
         broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(),
         broadcast_group=mpu.get_tensor_model_parallel_group(),
-        is_sft=is_sft,
+        has_cu_seqlens=has_cu_seqlens,
         is_hybrid_cp=is_hybrid_cp,
         create_attention_mask_in_dataloader=args.create_attention_mask_in_dataloader,
         cp_size=cp_size,
@@ -221,14 +222,17 @@ def get_batch(data_iterator, vp_stage=None):
     # Intermediate PP stage under SFT only needs THD metadata (matches the
     # pretrain_hybrid.py PP-SFT shortcut, collapsed to the flex 7-tuple shape).
     if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank:
-        assert is_sft
+        assert has_cu_seqlens
         return None, None, None, None, None, batch['cu_seqlens'], batch['max_seqlen']
 
     batch = get_batch_on_this_cp_rank(
         batch,
         is_hybrid_cp=is_hybrid_cp,
         cp_group=get_context_parallel_group(),
         hybrid_cp_group_func=get_hybrid_data_context_parallel_groups,
+        use_per_sequence_balancing=(
+            getattr(args, 'dataloader_inter_document_masking', False) and not is_sft
+        ),
     )
 
     cu_seqlens = batch.get('cu_seqlens')
@@ -474,6 +478,7 @@ def core_gpt_dataset_config_from_args(args):
         create_attention_mask=args.create_attention_mask_in_dataloader,
         object_storage_cache_path=args.object_storage_cache_path,
         mid_level_dataset_surplus=args.mid_level_dataset_surplus,
+        inter_document_masking=getattr(args, 'dataloader_inter_document_masking', False),
     )
 
 

megatron/training/arguments.py

@@ -1252,6 +1252,12 @@ def validate_args(args, defaults={}):
             'seq-length should be a multiple of 2 * context-parallel-size ' \
             'if context-parallel-size > 1.'
 
+    if getattr(args, 'dataloader_inter_document_masking', False):
+        # The dataset omits attention_mask when inter-document masking is
+        # enabled; disable the flag to avoid a TP broadcast mismatch.
+        if args.create_attention_mask_in_dataloader:
+            args.create_attention_mask_in_dataloader = False
+
     if args.seq_length is not None:
         assert args.encoder_seq_length is None
         args.encoder_seq_length = args.seq_length
@@ -2980,6 +2986,10 @@ def _add_data_args(parser):
     group.add_argument('--reset-attention-mask', action='store_true',
                        help='Reset self attention mask after '
                        'end-of-document token.')
+    group.add_argument('--dataloader-inter-document-masking', action='store_true',
+                       help='Return cu_seqlens marking document boundaries '
+                       'within each sample so that attention is restricted '
+                       'to individual documents.')
     group.add_argument('--eod-mask-loss', action='store_true',
                        help='Mask loss for the end of document tokens.')
     group.add_argument('--no-create-attention-mask-in-dataloader', action='store_false',