Support TE attention graphs for scheduler-packed THD

Adapt the fixed-capacity THD interface from dev #4359 and #5724 to main packing and graph scheduling. Preserve real and physical offsets, exclude padding from auxiliary losses, and keep replay bounded by the fixed source batch.

Signed-off-by: Deyu Fu <deyuf@nvidia.com>

Author: FDecaYed

docs/user-guide/features/cuda_graph.md

@@ -99,6 +99,18 @@ wired into custom training loops. The trade-off is that it requires more manual
 Megatron-LM's stock training loop already wires these calls in `megatron/training/training.py`,
 but custom training scripts must do the same work themselves.
 
+For scheduler-produced THD batches, use `--use-varlen-dataset`,
+`--sequence-packing-scheduler dp_balanced`, `--cuda-graph-modules attn`,
+`--pad-packed-seq-alignment max`, and an explicit `--thd-max-packed-sequences`. The metadata capacity includes one
+reserved dummy document for token padding. `--max-seqlen-per-dp-cp-rank` sets the
+fixed local token capacity; the source global batch size and micro-batch size 1
+must stay fixed. PP/VPP also requires `--cuda-graph-dynamic-microbatches`, which
+keeps graph inputs alive until their corresponding backward even when packing
+reduces the number of microbatches. Dynamic CP, EP communication overlap, delayed
+wgrad, and paged stash are outside this path. mHC wrappers and attention graph
+splitting require their separate graph integrations; these packing options do
+not enable mHC capture by themselves.
+
 ### Usage
 
 ```bash

megatron/core/datasets/data_schedule.py

@@ -16,7 +16,7 @@
     get_batch_and_global_seqlens,
     reroute_samples_to_dcp_ranks,
 )
-from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.packed_seq_params import PackedSeqParams, pad_packed_batch_to_capacity
 from megatron.core.pipeline_parallel.hybrid_cp_schedule import BalancedCPScheduler
 from megatron.core.process_groups_config import ProcessGroupCollection
 
@@ -361,6 +361,7 @@ def __init__(
         cp_size: int,
         dp_size: int,
         microbatch_group_size_per_vp_stage: Optional[int],
+        max_num_seqs: Optional[int] = None,
     ):
         """
         Args:
@@ -374,6 +375,7 @@ def __init__(
         self.cp_size = cp_size
         self.dp_size = dp_size
         self.microbatch_group_size_per_vp_stage = microbatch_group_size_per_vp_stage
+        self.max_num_seqs = max_num_seqs
 
     def get_required_sample_keys(self):
         """Return the required key of each batch."""
@@ -444,7 +446,9 @@ def get_groups_and_subsamples(self, sample_id_seqlens):
         single_microbatch = []
 
         for i in range(len(sample_id_seqlens)):
-            if sum_seqlen + sample_id_seqlens[i][1] <= self.max_seq_len_all_ranks:
+            if sum_seqlen + sample_id_seqlens[i][1] <= self.max_seq_len_all_ranks and (
+                self.max_num_seqs is None or len(single_microbatch) < self.max_num_seqs
+            ):
                 single_microbatch.append(i)
                 sum_seqlen += sample_id_seqlens[i][1]
             else:
@@ -720,6 +724,14 @@ def wrap_data_iterator(
             if config.virtual_pipeline_model_parallel_size is None
             else config.microbatch_group_size_per_vp_stage
         ),
+        # Reserve a document when token padding can append a dummy tail. Enforce
+        # the metadata bound before routing instead of falling back mid-pipeline.
+        max_num_seqs=(
+            config.thd_max_packed_sequences
+            - int(getattr(config, "pad_packed_seq_alignment", None) is not None)
+            if getattr(config, "thd_max_packed_sequences", None) is not None
+            else None
+        ),
     )
 
     (
@@ -745,6 +757,7 @@ def get_batch_on_this_rank_for_sequence_packing(
     mtp_on_this_rank: bool = False,
     vp_stage: Optional[int] = None,
     pg_collection: Optional[ProcessGroupCollection] = None,
+    config=None,
 ):
     """
     Get a batch of data for sequence packing.
@@ -921,5 +934,29 @@ def get_batch_on_this_rank_for_sequence_packing(
         max_seqlen_kv=max_seqlen,
     )
 
+    if config is not None and (
+        config.pad_packed_seq_alignment is not None
+        or getattr(config, "thd_max_packed_sequences", None) is not None
+    ):
+        alignment = config.pad_packed_seq_alignment
+        local_capacity = padding_mask.shape[-1]
+        if alignment == "max":
+            local_capacity = config.max_seqlen_per_dp_cp_rank
+        elif alignment is not None:
+            local_capacity = ((local_capacity + alignment - 1) // alignment) * alignment
+        tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = (
+            pad_packed_batch_to_capacity(
+                tokens,
+                labels,
+                loss_mask,
+                position_ids,
+                packed_seq_params,
+                padding_mask,
+                local_capacity=local_capacity,
+                cp_size=cp_group.size(),
+                max_num_seqs=getattr(config, "thd_max_packed_sequences", None),
+            )
+        )
+
     # "attention_mask" is not valid for sequence packing, so set it to None.
     return tokens, labels, loss_mask, None, position_ids, packed_seq_params, padding_mask

megatron/core/model_parallel_config.py

@@ -2,7 +2,7 @@
 
 import warnings
 from dataclasses import dataclass, field
-from typing import Callable, ContextManager, Literal, Optional
+from typing import Callable, ContextManager, Literal, Optional, Union
 
 import torch
 
@@ -11,6 +11,11 @@
 )
 
 
+def _parse_pad_packed_seq_alignment(value):
+    """Parse an integer packed-token alignment or the fixed-capacity value 'max'."""
+    return value if value == "max" else int(value)
+
+
 def resolve_tensor_parallel_weight_shards(
     tensor_model_parallel_size: int,
     tensor_parallel_num_weight_shards: Optional[int],
@@ -159,6 +164,16 @@ class ModelParallelConfig:
     dp_balanced: DP-balanced scheduler for sequence packing.
     """
 
+    pad_packed_seq_alignment: Optional[Union[int, Literal["max"]]] = field(
+        default=None, metadata={"argparse_meta": {"type": _parse_pad_packed_seq_alignment}}
+    )
+    """Pad scheduler THD tensors to a local token alignment, or to capacity with 'max'.
+
+    Padding occupies a separate physical document with no valid tokens. It is
+    excluded from loss, indexer training and MoE routing. Apply the same setting
+    in eager mode when comparing against CUDA graph execution.
+    """
+
     expert_model_parallel_size: int = 1
     """Distributes Moe Experts across sub data parallel dimension."""
 
@@ -545,6 +560,22 @@ def __post_init__(self):
         See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more
         details.
         """
+        if self.pad_packed_seq_alignment is not None:
+            self.pad_packed_seq_alignment = _parse_pad_packed_seq_alignment(
+                self.pad_packed_seq_alignment
+            )
+            assert self.sequence_packing_scheduler == "dp_balanced"
+            assert self.max_seqlen_per_dp_cp_rank is not None
+            if self.pad_packed_seq_alignment != "max":
+                assert 0 < self.pad_packed_seq_alignment <= self.max_seqlen_per_dp_cp_rank
+            if self.context_parallel_size > 1:
+                alignment = (
+                    self.max_seqlen_per_dp_cp_rank
+                    if self.pad_packed_seq_alignment == "max"
+                    else self.pad_packed_seq_alignment
+                )
+                assert alignment % 2 == 0, "Packed padding must preserve zigzag CP alignment."
+
         if self.sequence_parallel:
             if self.tensor_model_parallel_size <= 1:
                 raise ValueError("Cannot use sequence parallelism without tensor parallelism")

megatron/core/packed_seq_params.py

@@ -1,5 +1,5 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Optional
 
 import torch
@@ -73,6 +73,96 @@ def __post_init__(self):
             )
 
 
+def pad_packed_batch_to_capacity(
+    tokens,
+    labels,
+    loss_mask,
+    position_ids,
+    packed_seq_params,
+    padding_mask,
+    *,
+    local_capacity: int,
+    cp_size: int,
+    max_num_seqs: Optional[int] = None,
+):
+    """Pad already CP-sliced scheduler tensors without changing real document boundaries.
+
+    The existing mask supplies the local storage length on intermediate PP stages,
+    which do not own token or label tensors. Append a separate dummy document in
+    global THD coordinates; each CP rank appends its equal share of that document.
+    """
+    local_rows = padding_mask.shape[-1]
+    if local_rows > local_capacity:
+        raise ValueError(f"Packed local rows {local_rows} exceed capacity {local_capacity}.")
+    local_tail = local_capacity - local_rows
+    if cp_size > 1 and local_tail % 2:
+        raise ValueError("Packed padding tail must preserve the two-way zigzag CP partition.")
+
+    def pad_tensor(tensor, fill=0):
+        if tensor is None:
+            return None
+        if tensor.shape[-1] != local_rows:
+            raise ValueError("Packed token tensors and padding_mask must have the same length.")
+        return (
+            tensor
+            if local_tail == 0
+            else torch.nn.functional.pad(tensor, (0, local_tail), value=fill)
+        )
+
+    fields = {}
+    if local_tail:
+        for name in (
+            'cu_seqlens_q',
+            'cu_seqlens_kv',
+            'cu_seqlens_q_padded',
+            'cu_seqlens_kv_padded',
+        ):
+            offsets = getattr(packed_seq_params, name)
+            if offsets is None:
+                offsets = getattr(packed_seq_params, name.removesuffix('_padded'))
+            # The dummy document occupies physical storage but contains no real
+            # tokens. In particular, CSA's indexer auxiliary loss must ignore it.
+            tail = local_tail * cp_size if name.endswith('_padded') else 0
+            fields[name] = torch.cat((offsets, offsets[-1:] + tail))
+    if max_num_seqs is not None:
+        for name in (
+            'cu_seqlens_q',
+            'cu_seqlens_kv',
+            'cu_seqlens_q_padded',
+            'cu_seqlens_kv_padded',
+        ):
+            offsets = fields.get(name, getattr(packed_seq_params, name))
+            if offsets is None:
+                offsets = fields.get(
+                    name.removesuffix('_padded'),
+                    getattr(packed_seq_params, name.removesuffix('_padded')),
+                )
+            remaining = max_num_seqs + 1 - offsets.numel()
+            if remaining < 0:
+                raise ValueError("Packed documents exceed thd_max_packed_sequences.")
+            if remaining:
+                fields[name] = torch.cat((offsets, offsets[-1:].expand(remaining)))
+    static_capacity = local_capacity * cp_size if max_num_seqs is not None else 0
+    if fields or static_capacity:
+        packed_seq_params = replace(
+            packed_seq_params,
+            **fields,
+            max_seqlen_q=max(packed_seq_params.max_seqlen_q, local_tail * cp_size, static_capacity),
+            max_seqlen_kv=max(
+                packed_seq_params.max_seqlen_kv, local_tail * cp_size, static_capacity
+            ),
+            pad_between_seqs=True,
+        )
+    return (
+        pad_tensor(tokens),
+        pad_tensor(labels),
+        pad_tensor(loss_mask),
+        pad_tensor(position_ids),
+        packed_seq_params,
+        pad_tensor(padding_mask, True),
+    )
+
+
 def resolve_cp_group(
     static_cp_group: Optional[dist.ProcessGroup], packed_seq_params: PackedSeqParams = None
 ) -> Optional[dist.ProcessGroup]:

megatron/core/transformer/cuda_graphs.py

@@ -2249,9 +2249,16 @@ def __init__(
             "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using "
             "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True."
         )
+        if getattr(config, "sequence_packing_scheduler", None) is not None:
+            assert micro_batch_size == 1, "Scheduler THD graphs require micro_batch_size=1."
         self.model = model
         self.config = config
-        self.seq_length = seq_length
+        self.seq_length = (
+            config.max_seqlen_per_dp_cp_rank * config.context_parallel_size
+            if getattr(config, "sequence_packing_scheduler", None) is not None
+            and config.pad_packed_seq_alignment is not None
+            else seq_length
+        )
         self.micro_batch_size = micro_batch_size
         self.optimizers = optimizers
         self.pg_collection = pg_collection
@@ -2454,7 +2461,9 @@ def get_rotary_pos_emb(transformer_module, transformer_input):
                     )
                     if rotary_seq_len not in rotary_pos_emb_cache:
                         rotary_pos_emb_cache[rotary_seq_len] = transformer_module.rotary_pos_emb(
-                            rotary_seq_len
+                            rotary_seq_len,
+                            packed_seq=getattr(self.config, "sequence_packing_scheduler", None)
+                            is not None,
                         )
                     return rotary_pos_emb_cache[rotary_seq_len]
                 else:
@@ -2652,6 +2661,10 @@ def _get_cuda_graph_input_data(self):
             ), "If PP is not enabled, there should be only one model chunk."
             self.num_microbatches = 1
         else:
+            # Keep the source-batch bound when dp_balanced produces fewer packs.
+            # With a fixed source batch and MBS=1, each aligned packed schedule is
+            # a subsequence of this capture order, preserving the existing input
+            # buffer lifetimes without wrapping around the captured graph list.
             self.num_microbatches = get_num_microbatches()
 
         _, _, num_warmup_microbatches, _ = get_pp_rank_microbatches(

megatron/core/transformer/module.py

@@ -349,7 +349,17 @@ def _te_cuda_graph_replay(self, *args, **kwargs):
                 v, torch.Tensor
             ), "CUDA graph accepts only Tensor inputs."
 
-        cg_index = getattr(self, 'current_microbatch', 0) % len(self.cuda_graphs)
+        microbatch_idx = getattr(self, 'current_microbatch', 0)
+        if (
+            getattr(self.config, "cuda_graph_dynamic_microbatches", False)
+            and self.config.pipeline_model_parallel_size > 1
+        ):
+            if microbatch_idx >= len(self.cuda_graphs):
+                raise RuntimeError(
+                    "Packed microbatch count exceeds the captured source-batch bound; "
+                    "recapture CUDA graphs before increasing the source batch size."
+                )
+        cg_index = microbatch_idx % len(self.cuda_graphs)
         cudagraph_args, cudagraph_kwargs = self._get_te_cuda_graph_replay_args(*args, **kwargs)
 
         for hook, hook_args in self.cuda_graph_manual_hooks:

megatron/core/transformer/transformer_config.py

@@ -1252,6 +1252,20 @@ class TransformerConfig(ModelParallelConfig):
     cuda_graph_modules has no effect when cuda_graph_impl="none" and must be empty when
     cuda_graph_impl="full_iteration"."""
 
+    thd_max_packed_sequences: Optional[int] = None
+    """THD metadata capacity, including the dummy padding document.
+
+    Bounds real-document packing and fixes all four cumulative-offset tensor
+    shapes in eager and TE graph execution. Padding reserves one document slot.
+    """
+
+    cuda_graph_dynamic_microbatches: bool = False
+    """Allow TE graph replay when packing reduces the source microbatch count.
+
+    A fixed source global batch size bounds graph instances. Input slots stay
+    live until backward, independently of the number of packed microbatches.
+    """
+
     cuda_graph_modules: Union[str, CudaGraphModule, List[str], List[CudaGraphModule]] = "full"
     """Selects training capture coverage within per-layer CUDA graphs (local and
     transformer_engine implementations).
@@ -3562,6 +3576,27 @@ def _scope_to_str(s):
             "full_iteration",
         ], f"Invalid cuda graph implementation: {self.cuda_graph_impl}"
 
+        if self.thd_max_packed_sequences is not None:
+            assert self.thd_max_packed_sequences >= (
+                2 if self.pad_packed_seq_alignment is not None else 1
+            ), "THD metadata capacity must leave room for the dummy padding document."
+        if self.cuda_graph_dynamic_microbatches:
+            assert self.cuda_graph_impl == "transformer_engine"
+            assert self.sequence_packing_scheduler == "dp_balanced"
+            assert not self.overlap_moe_expert_parallel_comm and not self.delay_wgrad_compute
+            assert not self.moe_paged_stash
+        if (
+            self.sequence_packing_scheduler is not None
+            and self.cuda_graph_impl == "transformer_engine"
+        ):
+            assert self.pad_packed_seq_alignment in ("max", self.max_seqlen_per_dp_cp_rank)
+            assert self.thd_max_packed_sequences is not None
+            assert self.pipeline_model_parallel_size == 1 or self.cuda_graph_dynamic_microbatches
+            assert self.cuda_graph_modules == [
+                CudaGraphModule.attn
+            ], "Scheduler-produced THD CUDA graphs currently support attention capture only."
+            assert is_te_min_version("2.18.0"), "Packed TE graphs require TE >= 2.18.0."
+
         self.inference_cuda_graph_scope = normalize_inference_cuda_graph_scope(
             self.inference_cuda_graph_scope, self.cuda_graph_impl
         )