diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py
index 42146d1acd2..92d6a00f371 100644
--- a/megatron/core/datasets/gpt_dataset.py
+++ b/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,8 +283,56 @@ 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,
@@ -288,23 +340,26 @@ def __getitem__(self, idx: Optional[int]) -> Dict[str, torch.Tensor]:
                 "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(
diff --git a/megatron/core/utils.py b/megatron/core/utils.py
index 326f95e5589..f99e0a4f409 100644
--- a/megatron/core/utils.py
+++ b/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,7 +2161,7 @@ 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:
@@ -2169,8 +2169,8 @@ def _broadcast_cu_seqlens(cu_seqlens):
             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,7 +2272,7 @@ 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:
@@ -2280,8 +2280,8 @@ def _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
 
 
diff --git a/megatron/elastification/pretrain_hybrid_flex.py b/megatron/elastification/pretrain_hybrid_flex.py
index c9f9a32d60a..967404c6298 100644
--- a/megatron/elastification/pretrain_hybrid_flex.py
+++ b/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,7 +222,7 @@ 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(
@@ -229,6 +230,9 @@ def get_batch(data_iterator, vp_stage=None):
         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),
     )
 
 
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 9764bb5f0b6..930168ef644 100644
--- a/megatron/training/arguments.py
+++ b/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',
diff --git a/megatron/training/datasets/fim_dataset.py b/megatron/training/datasets/fim_dataset.py
index 875f979c91b..4b5a32f16ba 100644
--- a/megatron/training/datasets/fim_dataset.py
+++ b/megatron/training/datasets/fim_dataset.py
@@ -101,14 +101,15 @@ def __init__(
             self.eod_tok_id,
         ) = fim_tokens_ids
 
-    def _query_document_sample_shuffle_indices(self, idx: int) -> Tuple[np.ndarray, np.ndarray]:
+    def _query_document_sample_shuffle_indices(self, idx: int) -> Tuple[np.ndarray, np.ndarray, list]:
         """Get the text (token ids) and document ids for a given index
 
         Args:
             idx (int): The index into the dataset
 
         Returns:
-            Tuple[np.ndarray, np.ndarray]: The text ids and document ids
+            Tuple[np.ndarray, np.ndarray, list]: The text ids, document ids,
+                and per-document token counts.
         """
         # Do the shuffle mapping
         idx = self.shuffle_index[idx]
@@ -179,7 +180,7 @@ def _query_document_sample_shuffle_indices(self, idx: int) -> Tuple[np.ndarray,
 
         assert sample.shape[0] == sample_len
 
-        return (np.array(sample, dtype=np.int64), np.array(document_ids, dtype=np.int64))
+        return (np.array(sample, dtype=np.int64), np.array(document_ids, dtype=np.int64), [sample_len])
 
     def _fim_permute_sequence(self, sequence, rate):
         return self._permute(
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 39ab4256ef0..08eea7786aa 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -2191,7 +2191,7 @@ def dummy_train_step(data_iterator):
     """Single dummy training step."""
     args = get_args()
     tp_rank = mpu.get_tensor_model_parallel_rank()
-    is_sft = getattr(args, 'sft', False)
+    has_cu_seqlens = getattr(args, 'sft', False) or getattr(args, 'dataloader_inter_document_masking', False)
     is_hybrid_cp = args.hybrid_context_parallel
 
     BATCH_KEYS = [
@@ -2214,7 +2214,7 @@ def dummy_train_step(data_iterator):
                 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=args.context_parallel_size,
diff --git a/pretrain_gpt.py b/pretrain_gpt.py
index bb9e06b71c9..11adbb773c2 100644
--- a/pretrain_gpt.py
+++ b/pretrain_gpt.py
@@ -98,6 +98,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = 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 args.dataloader_inter_document_masking
     create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader
     mtp_on_this_rank = mtp_on_this_rank_func(
         layout=config.pipeline_model_parallel_layout,
@@ -107,7 +108,11 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None):
     )
     is_hybrid_cp = args.hybrid_context_parallel
 
-    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 for _ in BATCH_KEYS]
 
     batch = {}
@@ -124,7 +129,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = 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=create_attention_mask_in_dataloader,
         cp_size=cp_size,
@@ -140,7 +145,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None):
     batch = flatten_batch_for_packed_sequences(batch)
 
     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,
             batch['cu_seqlens'],
@@ -159,6 +164,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None):
         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=args.dataloader_inter_document_masking and not is_sft,
     )
 
     # Return values in BATCH_KEYS order so callers can unpack into the fixed
@@ -402,6 +408,7 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig:
         "data_parallel_size": args.data_parallel_size,
         "sequence_parallel_size": args.tensor_model_parallel_size * args.sequence_parallel,
         "hybrid_context_parallel": args.hybrid_context_parallel,
+        "inter_document_masking": args.dataloader_inter_document_masking,
     }
 
     # add FIM args to the config
diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py
index c2fe3bd510e..053040e656d 100644
--- a/pretrain_hybrid.py
+++ b/pretrain_hybrid.py
@@ -98,6 +98,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 args.dataloader_inter_document_masking
     create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader
     mtp_on_this_rank = mtp_on_this_rank_func(
         layout=config.pipeline_model_parallel_layout,
@@ -107,7 +108,11 @@ def get_batch(data_iterator, vp_stage=None):
     )
     is_hybrid_cp = args.hybrid_context_parallel
 
-    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 for _ in BATCH_KEYS]
 
     batch = {}
@@ -124,7 +129,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=create_attention_mask_in_dataloader,
         cp_size=cp_size,
@@ -140,7 +145,7 @@ def get_batch(data_iterator, vp_stage=None):
     batch = flatten_batch_for_packed_sequences(batch)
 
     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,
             batch['cu_seqlens'],
@@ -159,6 +164,7 @@ def get_batch(data_iterator, vp_stage=None):
         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=args.dataloader_inter_document_masking and not is_sft,
     )
 
     # Return values in BATCH_KEYS order so callers can unpack into the fixed
@@ -389,6 +395,7 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig:
         data_parallel_size=args.data_parallel_size,
         sequence_parallel_size=args.tensor_model_parallel_size * args.sequence_parallel,
         hybrid_context_parallel=args.hybrid_context_parallel,
+        inter_document_masking=args.dataloader_inter_document_masking,
     )
 
 
diff --git a/tests/unit_tests/data/test_get_batch.py b/tests/unit_tests/data/test_get_batch.py
index 27f8debe0a1..104acdd020a 100644
--- a/tests/unit_tests/data/test_get_batch.py
+++ b/tests/unit_tests/data/test_get_batch.py
@@ -2,13 +2,17 @@
 
 import os
 import sys
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
 
 from megatron.core import mpu
 from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator
-from megatron.core.utils import flatten_batch_for_packed_sequences
+from megatron.core.utils import (
+    _get_batch_on_this_cp_rank_per_sequence_balancing,
+    flatten_batch_for_packed_sequences,
+)
 from megatron.training.arguments import parse_args, validate_args
 from megatron.training.global_vars import destroy_global_vars, set_global_variables
 from pretrain_hybrid import get_batch
@@ -512,6 +516,173 @@ def test_flatten_batch_for_packed_sequences_padded_cu_seqlens(micro_batch_size,
     assert result['cu_seqlens'].shape[1] == expected_entries
 
 
+@pytest.mark.parametrize("tp_size", [1, 2, 4])
+@pytest.mark.parametrize("pp_size", [1, 2, 4])
+@pytest.mark.parametrize("cp_size", [1, 2, 4])
+@pytest.mark.parametrize("seq_length", [1024])
+def test_inter_document_masking_batch(tp_size, pp_size, cp_size, seq_length):
+    if tp_size * pp_size * cp_size > torch.cuda.device_count():
+        pytest.skip(
+            f"Skipping test because tp_size * pp_size * cp_size > torch.cuda.device_count() "
+            f"({tp_size * pp_size * cp_size} > {torch.cuda.device_count()})"
+        )
+
+    global_batch_size = int(os.environ.get("WORLD_SIZE", 1)) // (tp_size * pp_size * cp_size)
+    if global_batch_size < 1:
+        pytest.skip("Not enough ranks for the requested parallelism configuration")
+    args = initialize_test_environment(
+        tp_size,
+        pp_size,
+        cp_size,
+        seq_length,
+        micro_batch_size=1,
+        global_batch_size=global_batch_size,
+        sft=False,
+    )
+    args.dataloader_inter_document_masking = True
+
+    data_iterator = None
+    if mpu.get_tensor_model_parallel_rank() == 0:
+        data_iterator, _ = create_sft_data_iterator(seq_length)
+
+    (
+        attention_mask,
+        cu_seqlens,
+        cu_seqlens_padded,
+        hybrid_cp_group,
+        labels,
+        local_cp_size,
+        loss_mask,
+        max_seqlen,
+        position_ids,
+        tokens,
+    ) = get_batch(data_iterator)
+
+    is_first = mpu.is_pipeline_first_stage()
+    is_last = mpu.is_pipeline_last_stage()
+
+    # With CP > 1 and per-sequence balancing, sequence-dimension tensors
+    # are zigzag-partitioned to seq_length // cp_size while cu_seqlens
+    # and max_seqlen are left unchanged.
+    partitioned_seq_length = seq_length // cp_size
+
+    if pp_size == 1:
+        assert tokens is not None
+        assert labels is not None
+        assert loss_mask is not None
+        assert position_ids is not None
+        assert cu_seqlens is not None
+        assert max_seqlen is not None
+        assert attention_mask is None
+
+        assert tokens.shape[1] == partitioned_seq_length
+        assert labels.shape[1] == partitioned_seq_length
+        assert loss_mask.shape[1] == partitioned_seq_length
+        assert position_ids.shape[1] == partitioned_seq_length
+
+        assert cu_seqlens.dim() == 2
+        assert cu_seqlens.shape[0] == 1
+        assert cu_seqlens.dtype == torch.int32
+        assert cu_seqlens[0, 0].item() == 0
+        assert cu_seqlens[0, -1].item() == seq_length
+        assert cu_seqlens.shape[1] >= 2
+
+        assert max_seqlen.shape == (1,)
+        assert max_seqlen.dtype == torch.int32
+        assert 0 < max_seqlen.item() <= seq_length
+
+    elif is_first:
+        assert tokens is not None
+        assert position_ids is not None
+        assert labels is None
+        assert loss_mask is None
+        assert cu_seqlens is not None
+        assert max_seqlen is not None
+
+        assert tokens.shape[1] == partitioned_seq_length
+        assert position_ids.shape[1] == partitioned_seq_length
+
+        assert cu_seqlens.dim() == 2
+        assert cu_seqlens.dtype == torch.int32
+        assert cu_seqlens[0, 0].item() == 0
+        assert cu_seqlens[0, -1].item() == seq_length
+
+    elif is_last:
+        assert labels is not None
+        assert loss_mask is not None
+        assert tokens is None
+        assert position_ids is None
+        assert cu_seqlens is not None
+        assert max_seqlen is not None
+
+        assert labels.shape[1] == partitioned_seq_length
+        assert loss_mask.shape[1] == partitioned_seq_length
+
+        assert cu_seqlens.dim() == 2
+        assert cu_seqlens.dtype == torch.int32
+        assert cu_seqlens[0, 0].item() == 0
+        assert cu_seqlens[0, -1].item() == seq_length
+
+    else:
+        assert tokens is None
+        assert labels is None
+        assert loss_mask is None
+        assert position_ids is None
+        assert cu_seqlens is not None
+        assert max_seqlen is not None
+
+    Utils.destroy_model_parallel()
+
+
+@pytest.mark.parametrize("cp_size", [1, 2, 4])
+@pytest.mark.parametrize("seq_length", [16, 1024])
+def test_get_batch_on_this_cp_rank_per_sequence_balancing(cp_size, seq_length):
+    """Verify that per-sequence zigzag balancing selects the correct chunks.
+
+    Constructs a batch with tokens = range(seq_length) and checks that each
+    simulated CP rank receives the expected zigzag-interleaved chunks.
+    """
+    tokens = torch.arange(seq_length, dtype=torch.int64).unsqueeze(0)
+    cu_seqlens = torch.tensor([[0, seq_length // 2, seq_length]], dtype=torch.int32)
+    max_seqlen = torch.tensor([seq_length // 2], dtype=torch.int32)
+
+    for cp_rank in range(cp_size):
+        batch = {
+            'tokens': tokens.clone(),
+            'cu_seqlens': cu_seqlens.clone(),
+            'max_seqlen': max_seqlen.clone(),
+        }
+
+        mock_group = MagicMock()
+        with (
+            patch('torch.distributed.get_world_size', return_value=cp_size),
+            patch('torch.distributed.get_rank', return_value=cp_rank),
+        ):
+            result = _get_batch_on_this_cp_rank_per_sequence_balancing(batch, cp_group=mock_group)
+
+        if cp_size == 1:
+            assert torch.equal(result['tokens'], tokens)
+        else:
+            # The sequence is split into 2*cp_size equal chunks. This rank
+            # gets chunk cp_rank and chunk 2*cp_size - cp_rank - 1.
+            chunk_size = seq_length // (2 * cp_size)
+            chunk_0_start = cp_rank * chunk_size
+            chunk_1_start = (2 * cp_size - cp_rank - 1) * chunk_size
+            expected = torch.cat(
+                [
+                    tokens[0, chunk_0_start : chunk_0_start + chunk_size],
+                    tokens[0, chunk_1_start : chunk_1_start + chunk_size],
+                ]
+            ).unsqueeze(0)
+            assert torch.equal(
+                result['tokens'], expected
+            ), f"cp_rank={cp_rank}: expected {expected}, got {result['tokens']}"
+
+        # cu_seqlens and max_seqlen must be unchanged.
+        assert torch.equal(result['cu_seqlens'], cu_seqlens)
+        assert torch.equal(result['max_seqlen'], max_seqlen)
+
+
 def create_pretrain_data_iterator(
     seq_length: int = 1024, micro_batch_size: int = 1, create_attention_mask: bool = False
 ):
diff --git a/tests/unit_tests/data/test_gpt_dataset.py b/tests/unit_tests/data/test_gpt_dataset.py
index a2d25090fb8..26e773295ad 100644
--- a/tests/unit_tests/data/test_gpt_dataset.py
+++ b/tests/unit_tests/data/test_gpt_dataset.py
@@ -14,6 +14,7 @@
 from megatron.core.datasets.gpt_dataset import GPTDatasetConfig, MockGPTDataset
 from megatron.core.datasets.utils import compile_helpers
 from megatron.core.tokenizers import MegatronTokenizer
+from megatron.core.utils import _merge_cu_seqlens_across_micro_batch
 from tests.unit_tests.test_utilities import Utils
 
 _MOCK_VOCAB_SIZE = 8192
@@ -113,5 +114,81 @@ def test_mock_gpt_dataset():
     assert not torch.any(sample['loss_mask'])
 
 
+def test_inter_document_masking():
+    if torch.distributed.is_available():
+        Utils.initialize_distributed()
+        if torch.distributed.get_rank() == 0:
+            compile_helpers()
+        torch.distributed.barrier()
+    else:
+        compile_helpers()
+
+    tokenizer = MegatronTokenizer.from_pretrained(
+        metadata_path={"library": "null-text"}, vocab_size=_MOCK_VOCAB_SIZE
+    )
+
+    sequence_length = 1024
+
+    config = GPTDatasetConfig(
+        random_seed=1234,
+        sequence_length=sequence_length,
+        split="990,9,1",
+        reset_position_ids=False,
+        reset_attention_mask=False,
+        eod_mask_loss=False,
+        create_attention_mask=False,
+        tokenizer=tokenizer,
+        mid_level_dataset_surplus=0.005,
+        inter_document_masking=True,
+    )
+
+    datasets = BlendedMegatronDatasetBuilder(
+        MockGPTDataset, [100, 100, 100], lambda: True, config
+    ).build()
+
+    N = 20
+    for idx in range(N):
+        sample = datasets[0][idx]
+
+        assert "cu_seqlens" in sample
+        assert "max_seqlen" in sample
+        assert "attention_mask" not in sample
+
+        # Strip collation padding before validation.
+        cu_seqlens = _merge_cu_seqlens_across_micro_batch(
+            sample["cu_seqlens"].unsqueeze(0), sequence_length
+        )
+        max_seqlen = sample["max_seqlen"]
+        tokens = sample["tokens"]
+        position_ids = sample["position_ids"]
+
+        assert tokens.shape[0] == sequence_length
+        assert position_ids.shape[0] == sequence_length
+
+        assert cu_seqlens.dtype == torch.int32
+        assert cu_seqlens[0] == 0
+        assert cu_seqlens[-1] == sequence_length
+
+        # cu_seqlens must be strictly increasing.
+        diffs = cu_seqlens[1:] - cu_seqlens[:-1]
+        assert torch.all(diffs > 0), f"cu_seqlens not strictly increasing: {cu_seqlens}"
+
+        assert max_seqlen == diffs.max()
+
+        # Position IDs must reset to 0 at each document boundary.
+        for i in range(cu_seqlens.numel() - 1):
+            start = cu_seqlens[i].item()
+            end = cu_seqlens[i + 1].item()
+            expected = torch.arange(end - start, dtype=torch.long)
+            assert torch.equal(
+                position_ids[start:end], expected
+            ), f"position_ids mismatch in segment {i} [{start}:{end}]"
+
+    # Verify that None index zeros out loss_mask.
+    sample = datasets[0][None]
+    assert not torch.any(sample["loss_mask"])
+    assert "cu_seqlens" in sample
+
+
 if __name__ == "__main__":
     test_mock_gpt_dataset()