[codex] Preserve real and padded THD sequence lengths

megatron/core/datasets/data_schedule.py

@@ -18,8 +18,10 @@
 )
 from megatron.core.packed_seq_params import (
     PackedSeqParams,
+    extend_thd_padding_before_cp_slice,
     get_thd_padding_kwargs,
     pad_sequence_for_thd,
+    resolve_thd_tail_padding_policy,
 )
 from megatron.core.pipeline_parallel.hybrid_cp_schedule import BalancedCPScheduler
 from megatron.core.process_groups_config import ProcessGroupCollection
@@ -29,7 +31,11 @@
 def _build_thd_padding_mask(
     cu_seqlens: torch.Tensor, cu_seqlens_padded: torch.Tensor
 ) -> torch.Tensor:
-    """Build a 1D THD padding mask from scheduler sequence metadata."""
+    """Build a 1D THD padding mask from packed-sequence metadata.
+
+    True marks physical slots that carry no valid tokens, covering both
+    inter-sequence gaps and the padding tail of the final sequence.
+    """
     assert cu_seqlens.dim() == 1
     assert cu_seqlens_padded.dim() == 1
     assert cu_seqlens.numel() == cu_seqlens_padded.numel()
@@ -467,8 +473,9 @@ def _get_scheduler_max_real_num_seqs(config) -> Optional[int]:
     if max_num_seqs < 1:
         raise ValueError(f"thd_max_packed_sequences must be >= 1, got {max_num_seqs}.")
 
-    if getattr(config, 'pad_packed_seq_alignment', None) is not None and getattr(
-        config, 'pad_packed_seq_by_appending_dummy_seq', True
+    if (
+        getattr(config, 'pad_packed_seq_alignment', None) is not None
+        and resolve_thd_tail_padding_policy(config) == 'append_dummy_seq'
     ):
         if max_num_seqs < 2:
             raise ValueError(
@@ -631,7 +638,9 @@ def get_batch_on_this_rank_for_sequence_packing(
         )
 
     cp_partition_mode = getattr(config, "cp_partition_mode", "zigzag")
+    tail_padding_policy = resolve_thd_tail_padding_policy(config)
     contiguous_cp_local_target_len = None
+    non_dummy_global_target_len = None
     pad_alignment = (
         getattr(config, 'pad_packed_seq_alignment', None) if config is not None else None
     )
@@ -666,17 +675,35 @@ def get_batch_on_this_rank_for_sequence_packing(
         )
         _sanitize_thd_padding_values(batch, batch['padding_mask'])
 
+        # In extend_last mode, the padding tail belongs to the final real
+        # sequence. Extend its global physical endpoint before CP slicing so
+        # both zigzag indices and contiguous rank origins see the padded layout.
+        if pad_alignment is not None and tail_padding_policy == 'extend_last':
+            batch['cu_seqlens_padded'], batch['max_seqlen'], non_dummy_global_target_len = (
+                extend_thd_padding_before_cp_slice(
+                    batch['cu_seqlens_padded'],
+                    batch['max_seqlen'],
+                    alignment=alignment,
+                    target_len=(
+                        target_len if target_len is not None else contiguous_cp_local_target_len
+                    ),
+                    cp_size=cp_group.size(),
+                    cp_partition_mode=cp_partition_mode,
+                )
+            )
+
     # Partition sequence tensors for context parallelism. Padding mask is needed
     # on every PP stage, while data tensors are only needed on first/last/MTP stages.
     if is_tp_rank_0:
         cp_slice_keys = ['padding_mask']
         if is_first_or_last_stage or mtp_on_this_rank:
             cp_slice_keys.extend(['tokens', 'position_ids', 'labels', 'loss_mask'])
-        partition_total_tokens = (
-            contiguous_cp_local_target_len * cp_group.size()
-            if contiguous_cp_local_target_len is not None
-            else None
-        )
+        if non_dummy_global_target_len is not None:
+            partition_total_tokens = non_dummy_global_target_len
+        elif contiguous_cp_local_target_len is not None:
+            partition_total_tokens = contiguous_cp_local_target_len * cp_group.size()
+        else:
+            partition_total_tokens = None
         get_cp_slice_for_thd(
             batch,
             cp_group,
@@ -823,12 +850,14 @@ def get_batch_on_this_rank_for_sequence_packing(
         local_cp_size=local_cp_size,
         cp_group=cp_group,
         cp_partition_mode=cp_partition_mode,
-        pad_between_seqs=False,
+        pad_between_seqs=True,
     )
 
-    # 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.
-    if pad_alignment is not None and packed_seq_params is not None:
+    # Dummy metadata is appended after CP slicing as an ordinary sequence.
+    # Non-dummy tensors and their final padded endpoint were extended before
+    # slicing; this call is then a no-op except for fixed-capacity cu_seqlens
+    # entries in eager or graph mode.
+    if pad_alignment is not None:
         tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = (
             pad_sequence_for_thd(
                 tokens,
@@ -839,9 +868,7 @@ def get_batch_on_this_rank_for_sequence_packing(
                 alignment=alignment,
                 target_len=target_len,
                 max_num_seqs=max_num_seqs,
-                pad_by_appending_dummy_seq=getattr(
-                    config, 'pad_packed_seq_by_appending_dummy_seq', True
-                ),
+                tail_padding_policy=tail_padding_policy,
                 padding_mask=padding_mask,
                 cp_group=cp_group,
             )

megatron/core/datasets/data_schedule_utils.py

@@ -34,7 +34,7 @@ def get_cp_slice_for_thd(
             before slicing. Existing cu_seqlens metadata is left unchanged.
     """
     cp_size = cp_group.size()
-    if cp_size <= 1:
+    if cp_size <= 1 and partition_total_tokens is None:
         return
     cp_rank = cp_group.rank()
     # Partition with padded cumulative lengths so CP slices match the THD
@@ -63,6 +63,8 @@ def get_cp_slice_for_thd(
             if pad_len > 0:
                 pad_value = True if key == 'padding_mask' else 0
                 batch[key] = torch.cat([batch[key], batch[key].new_full((pad_len,), pad_value)])
+    if cp_size <= 1:
+        return
     if cp_partition_mode == "contiguous":
         if total_tokens % cp_size != 0:
             raise RuntimeError(

megatron/core/model_parallel_config.py

@@ -118,13 +118,21 @@ class ModelParallelConfig:
     tensors are padded to a multiple of N.
     """
 
-    pad_packed_seq_by_appending_dummy_seq: bool = True
-    """Represent a THD packed-sequence padding tail by appending a dummy sequence.
+    thd_tail_padding_policy: Optional[Literal["append_dummy_seq", "extend_last"]] = None
+    """Policy for representing the THD packed-sequence padding tail.
 
-    When disabled, token-like tensors are still padded according to
-    pad_packed_seq_alignment, but cu_seqlens sequence boundaries are not extended
-    for the padding tail. When thd_max_packed_sequences is set, static-input
-    padding may still pad cu_seqlens tensors to that value + 1 entries.
+    - append_dummy_seq: cover the post-pack padding tail with an ordinary
+      dummy sequence appended to the cu_seqlens metadata. Existing
+      valid/physical gaps between real sequences are preserved. This is the
+      default behavior.
+    - extend_last: keep valid cu_seqlens boundaries unchanged and extend the
+      final padded boundary so the tail is physical padding of the last
+      sequence. With context parallelism the extension is applied to the
+      global metadata before CP slicing.
+    - None (default): treated as append_dummy_seq.
+
+    When thd_max_packed_sequences is set, cu_seqlens tensors are padded to
+    that value + 1 entries in both eager and CUDA Graph modes.
     """
 
     expert_model_parallel_size: int = 1
@@ -514,6 +522,12 @@ def __post_init__(self):
                     f"got {self.min_dynamic_context_parallel_size}"
                 )
 
+        if self.thd_tail_padding_policy not in (None, "append_dummy_seq", "extend_last"):
+            raise ValueError(
+                "thd_tail_padding_policy must be 'append_dummy_seq', 'extend_last', or None, "
+                f"got {self.thd_tail_padding_policy!r}."
+            )
+
         if self.pad_packed_seq_alignment is not None:
             self.pad_packed_seq_alignment = _parse_pad_packed_seq_alignment(
                 self.pad_packed_seq_alignment

megatron/core/models/hybrid/hybrid_block.py

@@ -169,7 +169,9 @@ def _reconstruct_packed_seq_params_from_kwargs(self, kwargs):
             cu_seqlens_kv_padded=kwargs.pop('cu_seqlens_kv_padded'),
             max_seqlen_q=max_seqlen,
             max_seqlen_kv=max_seqlen,
-            pad_between_seqs=False,
+            # This Python flag is baked into the captured graph and cannot vary
+            # between replay batches. Use the conservative THD-safe branch.
+            pad_between_seqs=True,
         )
         kwargs['packed_seq_params'] = packed_seq_params