[main] support partial cuda graph under thd format

megatron/core/datasets/gpt_dataset.py

@@ -81,6 +81,26 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig):
     """When True, return cu_seqlens marking document boundaries within each sample so
     that attention is restricted to individual documents."""
 
+    sft_mock_dataset_config_json: Optional[str] = None
+    """This config provides the necessary information for the mock dataset."""
+
+    sequence_packing_scheduler: Optional[str] = None
+    """Scheduler for sequence packing and hybrid context parallel.
+    dp_balanced: DP-balanced scheduler for sequence packing.
+    """
+
+    varlen_mock_dataset_config_json: Optional[str] = None
+    """Mock-dataset config (same JSON schema as ``sft_mock_dataset_config_json``)
+    used by the ``--use-varlen-dataset`` path; kept separate so the varlen path
+    does not implicitly inherit SFT-specific knobs."""
+
+    varlen_sbhd_validation: bool = False
+    """When True, :class:`VarlenDataset.__getitem__` emits SBHD samples padded
+    to ``sequence_length`` (no ``cu_seqlens`` / ``original_seq_len`` /
+    ``padded_seq_len``), bypassing the packed-sequence path. Used to obtain a
+    SBHD reference run that mirrors the THD path's tokenization but skips all
+    packing — useful for THD numerical-correctness validation."""
+
     def __post_init__(self) -> None:
         """Do asserts and set fields post init"""
         super().__post_init__()
@@ -91,6 +111,12 @@ def __post_init__(self) -> None:
         assert self.reset_attention_mask is not None
         assert self.eod_mask_loss is not None
 
+        if self.varlen_sbhd_validation:
+            assert not self.hybrid_context_parallel, (
+                "--varlen-sbhd-validation is incompatible with "
+                "--hybrid-context-parallel (SBHD mode is not packed)."
+            )
+
         self.token_dtype_code = (
             None
             if self.tokenizer.vocab_size is None

megatron/core/datasets/readme.md

@@ -204,6 +204,28 @@ If the later training job does not specify `--global-batch-size` (which is neede
 
 `tools/prepare_cache.py` does not support `--mock-data`, `--sft`, `--fim-data`, or `--step-batch-size-schedule`.
 
+## Packing Scheduler
+
+The packing scheduler re-schedules variable-length sequences across DP×CP ranks to improve GPU utilization. It is built around the following modules:
+
+### `data_schedule`
+
+This module contains the high-level scheduling logic and entry points:
+
+- **`HybridCPDataLoaderWrapper`**: A wrapper class for hybrid context parallel (CP) scheduling. For every `__next__` call, it: (1) pulls a batch of packed samples from each DP rank, (2) gathers sequence lengths across the DP group, (3) schedules sub-samples using the `BalancedCPScheduler`, (4) reroutes sub-samples to the correct DPxCP ranks via all-to-all communication.
+
+- **`BasePackingScheduler`**: Abstract base class for packing schedulers. Defines the interface for `get_groups_and_subsamples()` (scheduling algorithm) and `run()` (full scheduling pipeline including fetch, schedule, reroute, pack, broadcast, and VPP handling).
+
+- **`DpBalancedScheduler`**: A concrete scheduler that packs sequences in their original order until reaching the max sequence length limit per DPxCP rank. Supports aligning the number of microbatches to DP size and VPP stage multiples.
+
+- **`wrap_data_iterator()`**: Top-level entry point that wraps an existing `data_iterator`. It creates the appropriate scheduler, runs the scheduling pipeline, broadcast metadata and new num_microbatches, returns a new data iterator along with the updated number of microbatches and FLOPs statistics.
+
+- **`get_batch_on_this_rank_for_sequence_packing()`**: Fetches and broadcasts a single packed microbatch for the current rank. Handles TP/PP broadcasting, constructs `PackedSeqParams` (with `cu_seqlens`, `max_seqlen`, `qkv_format=thd`), and optionally partitions sequences across CP ranks using Transformer Engine's `thd_get_partitioned_indices`.
+
+### `data_schedule_utils.py`
+
+This module contains the utility functions used by the schedulers.
+
 ## Fast DataLoader initialization
 
 Especially for large-scale runs, DataLoader initialization can take several minutes, since it involves opening and memory-mapping multiple files and can significantly stress the filesystem. To speed up this process, we have developed the following three optimizations, controlled by configuration flags:

megatron/core/extensions/transformer_engine.py

@@ -3393,3 +3393,24 @@ def set_save_original_input(module):
     from transformer_engine.pytorch.float8_tensor import Float8Tensor
 except ImportError:
     Float8Tensor = None
+
+
+def get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank):
+    """Get partitioned indices for THD format data in context parallel.
+
+    Args:
+        cu_seqlens: Cumulative sequence lengths tensor.
+        total_tokens: Total number of tokens.
+        cp_size: Context parallel world size.
+        cp_rank: Context parallel rank.
+
+    Returns:
+        Partitioned indices tensor.
+    """
+    assert is_te_min_version("1.10.0"), (
+        "Please update Transformer Engine to >= 1.10 to use "
+        "Context Parallel with THD format data"
+    )
+    import transformer_engine_torch as tex
+
+    return tex.thd_get_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank)

megatron/core/inference/contexts/dynamic_context.py

@@ -1771,6 +1771,7 @@ def apply_rotary_emb_query(
             cp_group=cp_group,
             mscale=mscale,
             mla_rotary_interleaved=config.multi_latent_attention,
+            max_seqlen=query_emb.size(0),
         )
         return query
 

megatron/core/model_parallel_config.py

@@ -2,11 +2,26 @@
 
 import warnings
 from dataclasses import dataclass, field
-from typing import Callable, ContextManager, Literal, Optional
+from typing import Callable, ContextManager, Literal, Optional, Union
 
 import torch
 
 
+def _parse_pad_packed_seq_alignment(value):
+    """Parse THD packed-sequence padding alignment.
+
+    Accepts ``"max"`` or a positive integer alignment.
+    """
+    if value == "max":
+        return value
+    try:
+        return int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(
+            "pad_packed_seq_alignment must be 'max' or a positive integer alignment."
+        ) from exc
+
+
 @dataclass
 class ModelParallelConfig:
     """Base configuration for Megatron Core
@@ -59,7 +74,7 @@ class ModelParallelConfig:
     can handle without overflowing the memory. Typically, a good starting point is to set this
     to maximum sequence length / context parallel size.
     This is used to calculate the number and length of sub-samples assigned to 
-    each rank when using hybrid_context_parallel.
+    each rank when sequence_packing_scheduler is not None.
     """
 
     hybrid_context_parallel: bool = False
@@ -69,6 +84,37 @@ class ModelParallelConfig:
     Please set max_seqlen_per_dp_cp_rank when using hybrid_context_parallel.
     """
 
+    sequence_packing_scheduler: Optional[Literal['dp_balanced']] = None
+    """
+    Scheduler for sequence packing and hybrid context parallel.
+    dp_balanced: DP-balanced scheduler for sequence packing.
+    """
+
+    pad_packed_seq_alignment: Optional[Union[int, Literal["max"]]] = field(
+        default=None,
+        metadata={
+            "argparse_meta": {
+                "arg_names": ["--pad-packed-seq-alignment"],
+                "type": _parse_pad_packed_seq_alignment,
+            }
+        },
+    )
+    """Pad THD packed sequence tensors after packing.
+
+    If set to ``max``, token-like tensors are padded to
+    max_seqlen_per_dp_cp_rank. If set to a positive integer N, token-like
+    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.
+
+    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. CUDA Graph static-input padding may still pad the
+    cu_seqlens tensors to thd_max_packed_sequences + 1 entries.
+    """
+
     expert_model_parallel_size: int = 1
     """Distributes Moe Experts across sub data parallel dimension."""
 
@@ -423,6 +469,28 @@ 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
+            )
+            if self.max_seqlen_per_dp_cp_rank is None:
+                raise ValueError(
+                    "max_seqlen_per_dp_cp_rank must be set when pad_packed_seq_alignment "
+                    "is enabled."
+                )
+            if self.pad_packed_seq_alignment != "max":
+                if self.pad_packed_seq_alignment <= 0:
+                    raise ValueError(
+                        "pad_packed_seq_alignment must be 'max' or a positive integer " "alignment."
+                    )
+                if self.pad_packed_seq_alignment > self.max_seqlen_per_dp_cp_rank:
+                    raise ValueError(
+                        "pad_packed_seq_alignment must not exceed "
+                        "max_seqlen_per_dp_cp_rank "
+                        f"({self.max_seqlen_per_dp_cp_rank}), got "
+                        f"{self.pad_packed_seq_alignment}."
+                    )
+
         if self.sequence_parallel:
             if self.tensor_model_parallel_size <= 1:
                 raise ValueError("Cannot use sequence parallelism without tensor parallelism")

megatron/core/models/common/embeddings/rope_utils.py

@@ -195,18 +195,22 @@ def _apply_rotary_pos_emb_thd(
     mscale: float = 1.0,
     cp_group: torch.distributed.ProcessGroup = None,
     multi_latent_attention: Optional[bool] = None,
+    max_seqlen: Optional[int] = None,
 ) -> Tensor:
-    """A baseline implementation of applying RoPE for `thd` format.
+    """Apply RoPE for `thd` format using pure CUDA ops (CUDA Graph compatible).
+
+    Replaces the original Python-loop + .tolist() implementation with vectorized
+    CUDA operations. No GPU->CPU syncs, compatible with CUDA Graph capture.
 
     Args:
-        t (Tensor): Input tensor T is of shape [t, h, d]
-        cu_seqlens(Tensor):  Cumulative sum of sequence lengths in a batch for `t`,
-        with shape [b + 1] and dtype torch.int32.
-        freqs (Tensor): Rotary Positional embedding tensor freq is of shape [max_s, 1, 1, d]
-        cp_group (torch.distributed.ProcessGroup): The context parallel group
+        t (Tensor): Input tensor of shape [total_tokens, h, d]
+        cu_seqlens (Tensor): Cumulative sequence lengths, shape [num_seqs + 1], int32.
+        freqs (Tensor): RoPE frequencies, shape [max_s, 1, 1, d] or [total_tokens, 1, 1, d]
+        cp_group: Context parallel group
+        max_seqlen: Global max sequence length for this packed batch when known.
 
     Returns:
-        Tensor: Shape [t, h, d]. The input tensor after applying RoPE.
+        Tensor: Shape [total_tokens, h, d]. Input with RoPE applied.
     """
     if multi_latent_attention is not None:
         warnings.warn(
@@ -219,53 +223,68 @@ def _apply_rotary_pos_emb_thd(
         raise ValueError("cp_group must be provided for THD format RoPE")
     cp_size = cp_group.size()
     cp_rank = cp_group.rank()
-    seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist()
-    sequence_splits = torch.split(t, seqlens)
-    total_seqlen = int(cu_seqlens[-1].item())
-    has_packed_freqs = freqs.dim() >= 1 and freqs.size(0) == total_seqlen
-
-    # Handle two different frequency tensor formats:
-    # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains positions for the whole packed
-    #    batch. Each sequence must therefore use its cu_seqlens offset when selecting the local CP
-    #    front/back slices. For example, with cu_seqlens=[0, 4, 8], cp_size=2, rank 0 should use
-    #    positions [0, 3, 4, 7], not [0, 3, 0, 3].
-    # 2. Otherwise: freqs contains only max sequence length positions. Each packed sequence should
-    #    reuse positions starting from 0, preserving the legacy THD behavior.
-    if has_packed_freqs:
-        # CASE 1: Exact mapping with offsets
-        local_freqs = []
-        for i, x in enumerate(sequence_splits):
-            # cu_seqlens[i] is the starting offset of this sequence in the original batch
-            seq_start_offset = cu_seqlens[i].item()
-            local_freqs.append(
-                _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset)
-            )
-        freqs = torch.cat(local_freqs, dim=0)
-        return _apply_rotary_pos_emb_bshd(
-            t.unsqueeze(1),
-            freqs,
-            rotary_interleaved=rotary_interleaved,
-            mla_rotary_interleaved=mla_rotary_interleaved,
-            mscale=mscale,
-        ).squeeze(1)
-
-    # CASE 2: Traditional mapping without offsets. Apply RoPE one sequence at a time so the second
-    # and later packed sequences do not look like continuations of the first sequence.
-    output = torch.empty_like(t)
-    output_offset = 0
-    for x in sequence_splits:
-        freq_slice = _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs)
-        output_slice = _apply_rotary_pos_emb_bshd(
-            x.unsqueeze(1),
-            freq_slice,
-            rotary_interleaved=rotary_interleaved,
-            mla_rotary_interleaved=mla_rotary_interleaved,
-            mscale=mscale,
-        ).squeeze(1)
-        output.narrow(0, output_offset, x.size(0)).copy_(output_slice)
-        output_offset += x.size(0)
 
-    return output
+    total_tokens = t.shape[0]
+    device = t.device
+
+    token_pos = torch.arange(total_tokens, device=device, dtype=torch.int64)
+
+    # `cu_seqlens` describes the global packed sequence. With CP, `t` is already
+    # CP-partitioned, so build a local cumulative-length view before assigning
+    # local tokens to packed sequences.
+    cu_seqlens_i64 = cu_seqlens.to(torch.int64)
+    global_seq_lens = cu_seqlens_i64[1:] - cu_seqlens_i64[:-1]
+    local_seq_lens = global_seq_lens // cp_size if cp_size > 1 else global_seq_lens
+    local_cu_seqlens = torch.zeros_like(cu_seqlens_i64)
+    local_cu_seqlens[1:] = torch.cumsum(local_seq_lens, dim=0)
+
+    # `searchsorted(..., right=True) - 1` returns the local sequence index. The
+    # clamp guards padded tokens that sit beyond the final real local token; they
+    # get a harmless frequency and are later masked out.
+    seq_idx = torch.searchsorted(local_cu_seqlens, token_pos, right=True) - 1
+    seq_idx = seq_idx.clamp(min=0, max=cu_seqlens.shape[0] - 2)
+
+    local_seq_start = local_cu_seqlens[seq_idx]
+    local_pos = token_pos - local_seq_start
+    local_seq_len = local_seq_lens[seq_idx]
+    global_seq_start = cu_seqlens_i64[seq_idx]
+
+    if cp_size > 1:
+        cp_seg = local_seq_len // 2
+        full_seqlen = local_seq_len * cp_size
+        is_first_half = local_pos < cp_seg
+        freq_pos = torch.where(
+            is_first_half,
+            cp_rank * cp_seg + local_pos,
+            full_seqlen - (cp_rank + 1) * cp_seg + (local_pos - cp_seg),
+        )
+    else:
+        freq_pos = local_pos.to(torch.int64)
+
+    assert max_seqlen is not None, (
+        "max_seqlen must be provided for THD RoPE so packed-frequency offset "
+        "detection does not silently depend on tensor shape heuristics."
+    )
+    exact_packed_freqs = freqs.dim() >= 1 and freqs.size(0) > max_seqlen
+    if exact_packed_freqs:
+        # `freqs` covers all positions across all sequences (used for non-1D
+        # RoPE / VLMs); shift by the per-sequence start offset so each token
+        # samples its absolute position. When `freqs` only spans one max-len
+        # sequence, no shift is needed.
+        freq_pos = freq_pos + global_seq_start
+
+    # Padded positions can sit outside the frequency table. Clamp them into
+    # range; downstream padding masks exclude those positions from the result.
+    freq_pos = freq_pos.clamp(min=0, max=freqs.shape[0] - 1)
+    freqs_packed = freqs[freq_pos]
+
+    return _apply_rotary_pos_emb_bshd(
+        t.unsqueeze(1),
+        freqs_packed,
+        rotary_interleaved=rotary_interleaved,
+        mla_rotary_interleaved=mla_rotary_interleaved,
+        mscale=mscale,
+    ).squeeze(1)
 
 
 def apply_rotary_pos_emb(
@@ -276,6 +295,7 @@ def apply_rotary_pos_emb(
     mscale: float = 1.0,
     cp_group: torch.distributed.ProcessGroup = None,
     mla_rotary_interleaved: bool = False,
+    max_seqlen: Optional[int] = None,
 ):
     """
     Reroute to the appropriate apply_rotary_pos_emb function depending on
@@ -343,6 +363,7 @@ def apply_rotary_pos_emb(
             mla_rotary_interleaved=mla_rotary_interleaved,
             mscale=mscale,
             cp_group=cp_group,
+            max_seqlen=max_seqlen,
         )