diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py
index 0f016473b6a..e3ae898b8c2 100644
--- a/megatron/core/datasets/data_schedule.py
+++ b/megatron/core/datasets/data_schedule.py
@@ -1,14 +1,68 @@
 # Copyright (c) 2025 NVIDIA CORPORATION.  All rights reserved.
 
-from typing import Any, List, Optional
+import enum
+from typing import Any, Dict, List, Optional, Type
 
 import torch
 
 from megatron.core import parallel_state
+from megatron.core.datasets.data_schedule_utils import (
+    broadcast_scalars,
+    broadcast_tensor,
+    broadcast_to_pp_group,
+    build_packed_microbatches,
+    create_data_iterator,
+    get_batch_and_global_seqlens,
+    reroute_samples_to_dcp_ranks,
+)
+from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices
+from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.pipeline_parallel.hybrid_cp_schedule import BalancedCPScheduler
 from megatron.core.process_groups_config import ProcessGroupCollection
 
 
+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."""
+    assert cu_seqlens.dim() == 1
+    assert cu_seqlens_padded.dim() == 1
+    assert cu_seqlens.numel() == cu_seqlens_padded.numel()
+
+    total_tokens = int(cu_seqlens_padded[-1].item())
+    if total_tokens == 0:
+        return torch.empty((0,), dtype=torch.bool, device=cu_seqlens.device)
+
+    num_sequences = cu_seqlens.numel() - 1
+    if num_sequences <= 0:
+        return torch.ones((total_tokens,), dtype=torch.bool, device=cu_seqlens.device)
+
+    positions = torch.arange(
+        total_tokens, dtype=cu_seqlens_padded.dtype, device=cu_seqlens_padded.device
+    )
+    seq_indices = torch.searchsorted(cu_seqlens_padded[1:].contiguous(), positions, right=True)
+
+    valid_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).clamp(min=0)
+    valid_ends = cu_seqlens_padded[:-1] + valid_lengths
+    return positions >= valid_ends[seq_indices]
+
+
+def _sanitize_thd_padding_values(batch: Dict[str, Any], padding_mask: torch.Tensor) -> None:
+    """Replace padded token-like slots with safe neutral values in-place."""
+    assert padding_mask.dim() == 1
+    pad_values = {'tokens': 0, 'labels': 0, 'loss_mask': 0.0, 'position_ids': 0}
+    for key, pad_value in pad_values.items():
+        tensor = batch.get(key)
+        if tensor is None:
+            continue
+        assert tensor.dim() == 1, f"{key} must be 1D before CP slicing, got {tensor.dim()}D"
+        assert tensor.numel() == padding_mask.numel(), (
+            f"{key} length ({tensor.numel()}) must match padding_mask length "
+            f"({padding_mask.numel()}) before CP slicing."
+        )
+        batch[key] = tensor.masked_fill(padding_mask, pad_value)
+
+
 class HybridCPDataLoaderWrapper:
     """
     A wrapper class that wraps around an existing data_iterator.
@@ -57,7 +111,6 @@ def get_global_seqlens(self, subsample_seqlens: torch.Tensor) -> List[int]:
         Gathers the sequence lengths of all subsamples from all DP ranks.
         Each DP rank loads the same number of microbatches but each microbatch
         may have a different number of subsamples.
-
         We find the number of subsamples each rank holds and then gather the
         sequence lengths of all subsamples from all ranks.
         """
@@ -299,3 +352,584 @@ def __next__(self) -> Any:
             batch, global_ids_this_rank, global_id_seqlens, sample_id_groups, offsets
         )
         return samples_this_rank_with_id, sample_id_groups
+
+
+class BasePackingScheduler:
+    """Base class for sequence packing schedulers."""
+
+    def __init__(
+        self,
+        max_seqlen_per_dp_cp_rank: int,
+        cp_size: int,
+        dp_size: int,
+        microbatch_group_size_per_vp_stage: Optional[int],
+    ):
+        """
+        Args:
+            max_seqlen_per_dp_cp_rank: The maximum sequence length per DPxCP rank.
+            cp_size: The context parallel size.
+            dp_size: The data parallel size.
+            microbatch_group_size_per_vp_stage: The microbatch group size per virtual
+            pipeline stage, only used when enabling VPP, otherwise None.
+        """
+        self.max_seqlen_per_dp_cp_rank = max_seqlen_per_dp_cp_rank
+        self.cp_size = cp_size
+        self.dp_size = dp_size
+        self.microbatch_group_size_per_vp_stage = microbatch_group_size_per_vp_stage
+
+    def get_required_sample_keys(self):
+        """Return the required key of each batch."""
+        raise NotImplementedError
+
+    def get_groups_and_subsamples(self, sample_id_seqlens):
+        """schedule the samples into groups"""
+        raise NotImplementedError
+
+    def run(
+        self,
+        data_iterator,
+        num_microbatches,
+        dp_group,
+        tp_group,
+        pp_group,
+        dp_cp_group,
+        dev,
+        config,
+    ):
+        """
+        Run the scheduler and return the new data_iterator.
+
+        Args:
+            data_iterator: The data iterator.
+            num_microbatches: The number of microbatches to fetch.
+            dp_group: Data parallel process group.
+            tp_group: Tensor parallel process group.
+            pp_group: Pipeline parallel process group.
+            dp_cp_group: Data parallel + context parallel process group.
+            dev: CUDA device.
+            config: Model parallel config.
+
+        Returns:
+            new_data_iterator: The new data iterator (or list for VPP).
+            num_micro_batches: Number of micro batches after scheduling.
+            seqlen_sum_this_global_batch: Total tokens for FLOPs calculation.
+            seqlen_squared_sum_this_global_batch: Sum of squared seqlens for FLOPs.
+        """
+        raise NotImplementedError
+
+
+class DpBalancedScheduler(BasePackingScheduler):
+    """Packs sequences in their original order until reaching the max limit of sequence length."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.max_seq_len_all_ranks = self.max_seqlen_per_dp_cp_rank * self.cp_size
+
+    def get_required_sample_keys(self):
+        """Return the required key of each batch."""
+        return [
+            "tokens",
+            "labels",
+            "loss_mask",
+            "position_ids",
+            "original_seq_len",  # Length of the original sequence length, should be a gpu tensor.
+            "padded_seq_len",  # Length of the padded sequence length, should be a gpu tensor.
+        ]
+
+    def get_groups_and_subsamples(self, sample_id_seqlens):
+        """
+        Packs sequences in their original order until reaching the max limit of sequence length.
+        """
+        sample_id_groups = []
+        packed_id_groups = []
+        sum_seqlen = 0
+        single_microbatch = []
+
+        for i in range(len(sample_id_seqlens)):
+            if sum_seqlen + sample_id_seqlens[i][1] <= self.max_seq_len_all_ranks:
+                single_microbatch.append(i)
+                sum_seqlen += sample_id_seqlens[i][1]
+            else:
+                packed_id_groups.append(single_microbatch)
+                single_microbatch = [i]
+                sum_seqlen = sample_id_seqlens[i][1]
+        if len(single_microbatch) > 0:
+            packed_id_groups.append(single_microbatch)
+
+        # we want the number of packed sequences to be multiple of dp_size
+        # so we move few samples from previous microbatch
+        # to the end of the microbatches if needed
+        num_packed_sequence = len(packed_id_groups)
+
+        # when enabling vpp, we want the number of packed sequences to be
+        # multiple of dp_size * microbatch_group_size_per_vp_stage
+        multiple = self.dp_size * (
+            self.microbatch_group_size_per_vp_stage
+            if self.microbatch_group_size_per_vp_stage is not None
+            else 1
+        )
+        if num_packed_sequence % multiple != 0:
+            remainder = num_packed_sequence % multiple
+            num_to_move = multiple - remainder
+            i = num_packed_sequence - 1
+            while num_to_move > 0:
+                assert i > 0, "Not enough samples to move"
+                if len(packed_id_groups[i]) > 1:
+                    seq_id = packed_id_groups[i].pop()
+                    packed_id_groups.append([seq_id])
+                    num_to_move -= 1
+                else:
+                    i -= 1
+
+        num_micro_batches = int(len(packed_id_groups) / self.dp_size)
+        for i in range(num_micro_batches):
+            sample_id_groups.append([])
+            for j in range(self.cp_size * self.dp_size):
+                seq_id = int(i * self.dp_size + j / self.cp_size)
+                sample_id_groups[i].append(packed_id_groups[seq_id])
+        return sample_id_groups
+
+    def run(
+        self,
+        data_iterator,
+        num_microbatches: int,
+        dp_group,
+        tp_group,
+        pp_group,
+        dp_cp_group,
+        dev: torch.device,
+        config,
+    ):
+        """
+        Run the complete scheduling pipeline.
+
+        Steps:
+            1. Fetch batches and gather global sequence lengths
+            2. Check required sample keys
+            3. Schedule samples into groups
+            4. Reroute samples to DCP ranks
+            5. Build packed microbatches
+            6. Calculate FLOPs info
+            7. Broadcast to PP group (for middle PP stages)
+            8. Broadcast to TP group (for non-TP-0 ranks)
+            9. Handle VPP if enabled
+
+        Args:
+            data_iterator: The data iterator.
+            num_microbatches: The number of microbatches to fetch.
+            dp_group: Data parallel process group.
+            tp_group: Tensor parallel process group.
+            pp_group: Pipeline parallel process group.
+            dp_cp_group: Data parallel + context parallel process group.
+            dev: CUDA device.
+            config: Model parallel config.
+
+        Returns:
+            new_data_iterator: The new data iterator (or list for VPP).
+            num_micro_batches: Number of micro batches after scheduling.
+            seqlen_sum_this_global_batch: Total tokens for FLOPs calculation.
+            seqlen_squared_sum_this_global_batch: Sum of squared seqlens for FLOPs.
+        """
+
+        total_dcp_gpus = dp_cp_group.size()
+
+        # Handle VPP: extract the correct data_iterator for this PP stage
+        if (
+            config.virtual_pipeline_model_parallel_size is not None
+            and config.virtual_pipeline_model_parallel_size > 1
+        ):
+            # if enable VPP, data_iterator is a list of data_iterators for each VPP stage,
+            # and only the first and last stage rank will have data_iterator,
+            # other stages will have None.
+            assert len(data_iterator) == config.virtual_pipeline_model_parallel_size
+            if pp_group.rank() == 0:
+                # the first stage
+                data_iterator = data_iterator[0]
+            elif pp_group.rank() == pp_group.size() - 1:
+                # the last stage
+                data_iterator = data_iterator[-1]
+            else:
+                data_iterator = None
+
+        # data_iterator is not None when TP rank 0, with PP stage 0 or -1.
+        if data_iterator is not None:
+            assert tp_group.rank() == 0 and (
+                pp_group.rank() == 0 or pp_group.rank() == pp_group.size() - 1
+            ), f"Only TP rank 0 and PP stage 0 or -1 should have data_iterator"
+
+            # Step 1: Fetch batches and gather global sequence lengths
+            batch, global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered = (
+                get_batch_and_global_seqlens(data_iterator, num_microbatches, dp_group)
+            )
+
+            # Step 2: Check required sample keys
+            for key in self.get_required_sample_keys():
+                assert (
+                    key in batch[0]
+                ), f"Batch missing required key {key}, provided keys: {batch[0].keys()}"
+
+            # Step 3: Schedule samples into groups
+            sample_id_groups = self.get_groups_and_subsamples(global_id_seqlens)
+
+            # Validate scheduling result
+            set_gbs = set()
+            for group in sample_id_groups:
+                for sub in group:
+                    set_gbs.update(sub)
+            assert len(set_gbs) == len(global_id_seqlens), (
+                f"set_gbs length: {len(set_gbs)} != "
+                f"global_id_seqlens length: {len(global_id_seqlens)}"
+            )
+
+            # Step 4: Reroute samples to DCP ranks
+            samples_this_rank_with_id = reroute_samples_to_dcp_ranks(
+                batch,
+                global_ids_this_rank,
+                global_id_seqlens,
+                sample_id_groups,
+                offsets,
+                dp_group,
+                tp_group,
+                dp_cp_group,
+                total_dcp_gpus,
+            )
+
+            dcp_rank = dp_cp_group.rank()
+            num_micro_batches = len(sample_id_groups)
+
+            grouped_samples = [
+                [
+                    samples_this_rank_with_id[sub_sample_id]
+                    for sub_sample_id in sample_id_groups[i][dcp_rank]
+                ]
+                for i in range(num_micro_batches)
+            ]
+
+            # Step 5: Build packed microbatches
+            new_samples = build_packed_microbatches(grouped_samples, dev)
+
+            # Step 6: Calculate FLOPs info
+            seqlen_sum_this_global_batch = float(sum(seqlens_gathered))
+            seqlen_squared_sum_this_global_batch = float(
+                sum(seqlen**2 for seqlen in seqlens_gathered)
+            )
+        else:
+            (
+                new_samples,
+                num_micro_batches,
+                seqlen_sum_this_global_batch,
+                seqlen_squared_sum_this_global_batch,
+            ) = (None, None, None, None)
+
+        # Step 7: Broadcast to PP group (for middle PP stages)
+        if tp_group.rank() == 0:
+            (
+                new_samples,
+                num_micro_batches,
+                seqlen_sum_this_global_batch,
+                seqlen_squared_sum_this_global_batch,
+            ) = broadcast_to_pp_group(
+                new_samples,
+                num_micro_batches,
+                seqlen_sum_this_global_batch,
+                seqlen_squared_sum_this_global_batch,
+                pp_group,
+                dev,
+            )
+
+        # Step 8: Broadcast to TP group (for non-TP-0 ranks)
+        (num_micro_batches, seqlen_sum_this_global_batch, seqlen_squared_sum_this_global_batch) = (
+            broadcast_scalars(
+                [
+                    num_micro_batches,
+                    seqlen_sum_this_global_batch,
+                    seqlen_squared_sum_this_global_batch,
+                ],
+                tp_group,
+                dev,
+            )
+        )
+        num_micro_batches = int(num_micro_batches)
+
+        # Step 9: create data_iterator and handle VPP if enabled
+        new_data_iterator = create_data_iterator(new_samples, pp_group, tp_group, config)
+
+        return (
+            new_data_iterator,
+            num_micro_batches,
+            seqlen_sum_this_global_batch,
+            seqlen_squared_sum_this_global_batch,
+        )
+
+
+class PackingSchedulerEnum(enum.Enum):
+    """Enum for supported sequence packing algorithms."""
+
+    DP_BALANCED = "dp_balanced"
+
+
+scheduler_map: Dict[PackingSchedulerEnum, Type[BasePackingScheduler]] = {
+    PackingSchedulerEnum.DP_BALANCED: DpBalancedScheduler
+}
+
+
+def wrap_data_iterator(
+    data_iterator, config, num_microbatches, pg_collection: Optional[ProcessGroupCollection] = None
+):
+    """
+    A wrapper function that wraps around an existing data_iterator
+    and return the num_micro_batches for sequence packing.
+
+    Args:
+        data_iterator: The original data_iterator to wrap around
+        config: The config object containing the max_seqlen_per_dp_cp_rank
+        dp_cp_group: Data parallel context parallel group.
+        pg_collection: The process group collection.
+    """
+
+    if pg_collection is None:
+        dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=True)
+        dp_group = parallel_state.get_data_parallel_group()
+        tp_group = parallel_state.get_tensor_model_parallel_group()
+        pp_group = parallel_state.get_pipeline_model_parallel_group()
+    else:
+        dp_cp_group = pg_collection.dp_cp
+        dp_group = pg_collection.dp
+        tp_group = pg_collection.tp
+        pp_group = pg_collection.pp
+    assert (
+        dp_cp_group is not None
+        and dp_group is not None
+        and tp_group is not None
+        and pp_group is not None
+    ), "dp_cp_group, dp_group, tp_group must not be None when using sequence packing"
+
+    dev = torch.cuda.current_device()
+    dp_size = dp_group.size()
+    cp_size = dp_cp_group.size() // dp_size
+
+    # Convert string to enum
+    scheduler_type = config.sequence_packing_scheduler
+    scheduler_type = PackingSchedulerEnum[scheduler_type.upper()]
+
+    scheduler = scheduler_map[scheduler_type](
+        config.max_seqlen_per_dp_cp_rank,
+        cp_size,
+        dp_size,
+        # When VPP is enabled, align num_micro_batches to this multiple.
+        (
+            None
+            if config.virtual_pipeline_model_parallel_size is None
+            else config.microbatch_group_size_per_vp_stage
+        ),
+    )
+
+    (
+        new_data_iterator,
+        num_micro_batches,
+        seqlen_sum_this_global_batch,
+        seqlen_squared_sum_this_global_batch,
+    ) = scheduler.run(
+        data_iterator, num_microbatches, dp_group, tp_group, pp_group, dp_cp_group, dev, config
+    )
+
+    return (
+        new_data_iterator,
+        num_micro_batches,
+        seqlen_sum_this_global_batch,
+        seqlen_squared_sum_this_global_batch,
+    )
+
+
+def get_batch_on_this_rank_for_sequence_packing(
+    data_iterator,
+    vpp_size: Optional[int] = None,
+    mtp_on_this_rank: bool = False,
+    vp_stage: Optional[int] = None,
+    pg_collection: Optional[ProcessGroupCollection] = None,
+):
+    """
+    Get a batch of data for sequence packing.
+    Args:
+        data_iterator (Iterator): The data iterator to get the batch from.
+        mtp_on_this_rank (bool): Whether to use multi-token prediction.
+        vp_stage (Optional[int]): The stage of the pipeline.
+    Returns:
+        tuple of (tokens, labels, loss_mask, attention_mask, position_ids,
+        packed_seq_params, padding_mask)
+    """
+
+    if pg_collection is None:
+        tp_group = parallel_state.get_tensor_model_parallel_group()
+        pp_group = parallel_state.get_pipeline_model_parallel_group()
+        cp_group = parallel_state.get_context_parallel_group()
+    else:
+        tp_group = pg_collection.tp
+        pp_group = pg_collection.pp
+        cp_group = pg_collection.cp
+
+    tp_src_rank = torch.distributed.get_process_group_ranks(tp_group)[0]
+
+    is_tp_rank_0 = tp_group.rank() == 0
+    is_first_stage = pp_group.rank() == 0 and (vp_stage is None or vp_stage == 0)
+    is_last_stage = pp_group.rank() == pp_group.size() - 1 and (
+        vp_stage is None or vp_stage == vpp_size - 1
+    )
+
+    is_first_or_last_stage = is_first_stage or is_last_stage
+    dev = torch.cuda.current_device()
+
+    # data_iterator should return a batch including the following keys.
+    batch_keys = ['cu_seqlens', 'cu_seqlens_padded', 'max_seqlen']
+    if is_first_stage:
+        batch_keys.append('tokens')
+        batch_keys.append('position_ids')
+    if is_last_stage:
+        batch_keys.append('labels')
+        batch_keys.append('loss_mask')
+
+    # Get a batch from data_iterator or create an emtpy batch.
+    if is_tp_rank_0:
+        assert data_iterator is not None
+        batch = next(data_iterator)
+        for key in batch_keys:
+            assert key in batch, f"{key} is missing in current batch."
+    else:
+        assert data_iterator is None, "Non TP 0 rank should not have data_iterator"
+        batch = {}
+
+    # Build padding_mask before CP slicing while tensors still have the full
+    # packed length represented by cu_seqlens_padded[-1].
+    if is_tp_rank_0:
+        batch['padding_mask'] = _build_thd_padding_mask(
+            batch['cu_seqlens'], batch['cu_seqlens_padded']
+        )
+        _sanitize_thd_padding_values(batch, batch['padding_mask'])
+
+    # Partition padding_mask for context parallel on every PP stage. Partition
+    # token-like tensors only on stages that own them.
+    if is_tp_rank_0:
+        cp_size = cp_group.size()
+        cp_rank = cp_group.rank()
+        # If cp_size == 1, no need to do further processing.
+        if cp_size > 1:
+            # Transformer Engine has a bug of cu_seqlens, we must treat cu_seqlens_padded as
+            # cu_seqlens to get the correct result.
+            # TODO: Revert this workaround once TE fixes the issue.
+            cu_seqlens = batch["cu_seqlens_padded"]
+            total_tokens = int(cu_seqlens[-1].item())
+            index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank)
+            cp_slice_keys = ['padding_mask']
+            if is_first_or_last_stage:
+                cp_slice_keys.extend(['tokens', 'position_ids', 'labels', 'loss_mask'])
+            for key in cp_slice_keys:
+                batch[key] = batch[key].index_select(0, index)
+
+    # Broadcast cu_seqlens_size because we need it to create placeholder for cu_seqlens and
+    # cu_seqlens_padded for non TP 0 ranks.
+    if is_tp_rank_0:
+        cu_seqlen_size = torch.tensor(batch['cu_seqlens'].size(0), dtype=torch.int32, device=dev)
+    else:
+        cu_seqlen_size = torch.empty(1, dtype=torch.int32, device=dev)
+    broadcast_tensor(cu_seqlen_size, tp_src_rank, tp_group)
+    cu_seqlen_size = cu_seqlen_size.item()
+
+    # Broadcast total_tokens because padding_mask is prepared on every PP stage.
+    # Tokens/labels/loss_mask/position_ids use the same length on stages that own them.
+    if is_tp_rank_0:
+        total_tokens = torch.tensor(batch['padding_mask'].size(0), dtype=torch.int32, device=dev)
+    else:
+        total_tokens = torch.empty(1, dtype=torch.int32, device=dev)
+    broadcast_tensor(total_tokens, tp_src_rank, tp_group)
+    total_tokens = total_tokens.item()
+
+    # Step1: Prepare "tokens", "position_ids" on all ranks.
+    if is_first_stage or mtp_on_this_rank:
+        if is_tp_rank_0:
+            assert batch['tokens'].dtype == torch.int64
+            assert batch['position_ids'].dtype == torch.int64
+            batch['tokens'] = batch['tokens'].view(1, total_tokens)
+            batch['position_ids'] = batch['position_ids'].view(1, total_tokens)
+        else:
+            batch['tokens'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev)
+            batch['position_ids'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev)
+    else:
+        # Non first stage rank doesn't need tokens and position_ids.
+        batch['tokens'] = None
+        batch['position_ids'] = None
+
+    # Step2: Prepare "labels", "loss_mask" on all ranks.
+    if is_last_stage:
+        if is_tp_rank_0:
+            assert batch['labels'].dtype == torch.int64
+            assert batch['loss_mask'].dtype == torch.float32
+            batch['labels'] = batch['labels'].view(1, total_tokens)
+            batch['loss_mask'] = batch['loss_mask'].view(1, total_tokens)
+        else:
+            batch['labels'] = torch.empty([1, total_tokens], dtype=torch.int64, device=dev)
+            batch['loss_mask'] = torch.empty([1, total_tokens], dtype=torch.float32, device=dev)
+    else:
+        # Non last stage rank doesn't need labels and loss_mask.
+        batch['labels'] = None
+        batch['loss_mask'] = None
+
+    # Step3: Prepare "padding_mask" on all TP ranks.
+    if is_tp_rank_0:
+        assert batch['padding_mask'].dtype == torch.bool
+        batch['padding_mask'] = batch['padding_mask'].view(1, total_tokens)
+    else:
+        batch['padding_mask'] = torch.empty([1, total_tokens], dtype=torch.bool, device=dev)
+
+    # Step4: Prepare "cu_seqlens", "cu_seqlens_padded", "max_seqlen" on all ranks.
+    if is_tp_rank_0:
+        assert batch['cu_seqlens'].dtype == torch.int32
+        assert batch['cu_seqlens_padded'].dtype == torch.int32
+        assert batch['cu_seqlens'].dim() == 1
+        assert batch['cu_seqlens_padded'].dim() == 1
+        if type(batch['max_seqlen']) == int:
+            batch['max_seqlen'] = torch.tensor(batch['max_seqlen'], dtype=torch.int32, device=dev)
+        else:
+            assert batch['max_seqlen'].dtype == torch.int32
+            assert batch['max_seqlen'].numel() == 1
+    else:
+        batch['cu_seqlens'] = torch.empty([cu_seqlen_size], dtype=torch.int32, device=dev)
+        batch['cu_seqlens_padded'] = torch.empty([cu_seqlen_size], dtype=torch.int32, device=dev)
+        batch['max_seqlen'] = torch.empty(1, dtype=torch.int32, device=dev)
+
+    # Broadcast batch inside TP group.
+    broadcast_tensor(batch['tokens'], tp_src_rank, tp_group)
+    broadcast_tensor(batch['position_ids'], tp_src_rank, tp_group)
+    broadcast_tensor(batch['labels'], tp_src_rank, tp_group)
+    broadcast_tensor(batch['loss_mask'], tp_src_rank, tp_group)
+    broadcast_tensor(batch['padding_mask'], tp_src_rank, tp_group)
+    broadcast_tensor(batch['cu_seqlens'], tp_src_rank, tp_group)
+    broadcast_tensor(batch['cu_seqlens_padded'], tp_src_rank, tp_group)
+    broadcast_tensor(batch['max_seqlen'], tp_src_rank, tp_group)
+
+    # Extract the data from batch after broadcasting.
+    tokens = batch['tokens']
+    position_ids = batch['position_ids']
+    labels = batch['labels']
+    loss_mask = batch['loss_mask']
+    padding_mask = batch['padding_mask']
+    cu_seqlens = batch['cu_seqlens']
+    cu_seqlens_padded = batch['cu_seqlens_padded']
+    max_seqlen = batch['max_seqlen'].item()
+
+    # Transformer Engine has a bug of cu_seqlens, we must treat cu_seqlens_padded as cu_seqlens to
+    # get the correct result.
+    # TODO: Revert this workaround once TE fixes the issue.
+    packed_seq_params = PackedSeqParams(
+        qkv_format="thd",
+        cu_seqlens_q=cu_seqlens_padded,
+        cu_seqlens_kv=cu_seqlens_padded,
+        cu_seqlens_q_padded=cu_seqlens_padded,
+        cu_seqlens_kv_padded=cu_seqlens_padded,
+        max_seqlen_q=max_seqlen,
+        max_seqlen_kv=max_seqlen,
+        local_cp_size=None,
+        cp_group=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
diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py
new file mode 100644
index 00000000000..86eb7d2e1fc
--- /dev/null
+++ b/megatron/core/datasets/data_schedule_utils.py
@@ -0,0 +1,540 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.  All rights reserved.
+
+from typing import Dict, List
+
+import torch
+
+from megatron.core.rerun_state_machine import RerunDataIterator
+
+
+def _unpack_batch(batch: List[Dict[str, torch.Tensor]]) -> List[Dict[str, torch.Tensor]]:
+    """
+    Unpacks the packed samples into a list of sub-samples.
+    Since each sub-sample may be routed to different DPxCP ranks,
+    we unpack the sample here to avoid unnecessarily transferring
+    the entire packed sample.
+
+    Two input shapes are accepted:
+
+      * **Pre-packed** (e.g. :class:`SFTDataset`): each sample carries a
+        ``cu_seqlens`` tensor and the tokens of multiple sub-samples
+        concatenated together. We slice them apart and synthesize
+        ``original_seq_len`` / ``padded_seq_len`` from the cu_seqlens deltas.
+
+      * **Already unpacked** (e.g. :class:`VarlenDataset`): each sample is a
+        single sub-sample that already carries ``padded_seq_len`` (and
+        usually ``original_seq_len``). We just normalize the leading batch
+        dimension introduced by the default collate_fn and return as-is.
+    """
+    # Short-circuit for datasets that already emit one sub-sample per index.
+    if batch and "padded_seq_len" in batch[0]:
+        for sample in batch:
+            for key in sample.keys():
+                if sample[key].ndim == 2 and sample[key].shape[0] == 1:
+                    # Drop the redundant batch dim added by collate_fn.
+                    sample[key] = sample[key].squeeze(0)
+            if "original_seq_len" not in sample:
+                sample["original_seq_len"] = sample["padded_seq_len"].clone()
+        return batch
+
+    batch_unpacked = []
+    dev = batch[0]["tokens"].device
+    original_seq_lens = []
+    padded_seq_lens = []
+    for sample in batch:
+        for key in sample.keys():
+            if len(sample[key].shape) == 2:
+                # squeeze the redundant batch dimension added by
+                # default collate_fn in pytorch dataloader
+                # we need a custom collate_fn for THD to avoid this
+                # current THD does not support micro_batch_size > 1 due to sft_dataset.py and
+                # data_loader in data_samples.py
+                sample[key] = sample[key].squeeze(0)
+        for sub_sample in range(sample["cu_seqlens"].shape[0] - 1):
+            sub_sample_dict = {}
+            start_idx = sample["cu_seqlens"][sub_sample]
+            end_idx = sample["cu_seqlens"][sub_sample + 1]
+            if end_idx - start_idx == 0:
+                continue
+            for key in ["tokens", "labels", "loss_mask", "position_ids"]:
+                sub_sample_dict[key] = sample[key][start_idx:end_idx]
+            # Since sft_dataset.py does not provide cu_seqlens_original,
+            # we assume original_seq_len equals padded_seq_len here.
+            # Ideally the dataset should define the pre-padding seq_len.
+            seq_len = (end_idx - start_idx).item()
+            original_seq_lens.append(seq_len)
+            padded_seq_lens.append(seq_len)
+            batch_unpacked.append(sub_sample_dict)
+
+    # Single H2D transfer for all seq lens
+    original_seq_lens_cuda = torch.tensor(original_seq_lens, device=dev)
+    padded_seq_lens_cuda = torch.tensor(padded_seq_lens, device=dev)
+    for i, sub_sample_dict in enumerate(batch_unpacked):
+        sub_sample_dict["original_seq_len"] = original_seq_lens_cuda[i : i + 1]
+        sub_sample_dict["padded_seq_len"] = padded_seq_lens_cuda[i : i + 1]
+
+    return batch_unpacked
+
+
+def _get_global_seqlens_and_ids(subsample_seqlens: torch.Tensor, dp_group):
+    """
+    Gathers the sequence lengths of all subsamples from all DP ranks and calculates global IDs.
+    """
+    # Collect the number of subsamples from all ranks
+    num_local_subsamples = subsample_seqlens.shape[0]
+    local_len = torch.tensor([num_local_subsamples], dtype=torch.int32).cuda()
+    dp_subsample_count = [torch.zeros_like(local_len) for _ in range(dp_group.size())]
+    torch.distributed.all_gather(dp_subsample_count, local_len, group=dp_group)
+
+    # Find the max number of subsamples across all ranks and pad subsample_seqlens to max length
+    dp_subsample_counts = torch.stack(dp_subsample_count, dim=0).cpu().view(-1)
+    max_sub_samples = int(dp_subsample_counts.max().item())
+
+    if num_local_subsamples < max_sub_samples:
+        subsample_seqlens_padded = torch.cat(
+            [
+                subsample_seqlens,
+                torch.zeros(max_sub_samples - num_local_subsamples, dtype=torch.int32).cuda(),
+            ],
+            dim=0,
+        )
+    else:
+        subsample_seqlens_padded = subsample_seqlens
+
+    # Gather the subsample_seqlens from all ranks
+    seqlens_gathered = [torch.empty_like(subsample_seqlens_padded) for _ in range(dp_group.size())]
+    torch.distributed.all_gather(seqlens_gathered, subsample_seqlens_padded, group=dp_group)
+
+    # Trim each seqlens_gathered to the length of the correct sample
+    for dp_rank, seqlen in enumerate(seqlens_gathered):
+        seqlens_gathered[dp_rank] = seqlen[: dp_subsample_counts[dp_rank]]
+
+    seqlens_gathered = torch.cat(seqlens_gathered, dim=0)
+    seqlens_gathered = seqlens_gathered.cpu().tolist()
+
+    # Calculate the offsets to assign unique global ID to each subsample.
+    csum = torch.cumsum(dp_subsample_counts, dim=0, dtype=torch.int32)
+    offsets = torch.cat([torch.zeros(1, dtype=torch.int32), csum], dim=0)
+
+    # Calculate global ID for each subsample
+    dp_rank = dp_group.rank()
+    global_ids = torch.arange(len(seqlens_gathered), dtype=torch.int32).cuda()
+
+    # Create a list of (global_id, seqlen) tuples for scheduling
+    global_id_seqlens = [(i, seqlens_gathered[i]) for i in range(len(global_ids))]
+
+    # Get the global IDs locally present on this rank
+    start_idx = offsets[dp_rank]
+    end_idx = offsets[dp_rank + 1]
+
+    global_ids_this_rank = global_ids[start_idx:end_idx]
+
+    return global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered
+
+
+def _pack_sequences(
+    samples: List, padded_lengths: torch.Tensor, original_lengths: torch.Tensor, dev: torch.device
+) -> Dict[str, torch.Tensor]:
+    """Pack multiple samples into a single packed sample."""
+
+    def _pack_tensors(tensors):
+        return torch.cat([t.reshape(-1) for t in tensors], dim=0)
+
+    tokens = _pack_tensors([sample["tokens"] for sample in samples])
+    labels = _pack_tensors([sample["labels"] for sample in samples])
+    loss_mask = _pack_tensors([sample["loss_mask"] for sample in samples])
+    position_ids = _pack_tensors([sample["position_ids"] for sample in samples])
+
+    new_sample = {}
+    new_sample["tokens"] = tokens
+    new_sample["labels"] = labels
+    new_sample["loss_mask"] = loss_mask
+    new_sample["position_ids"] = position_ids
+
+    padded_lengths = padded_lengths.to(device=dev, dtype=torch.int32, non_blocking=True).reshape(-1)
+    cu_seqlens_padded = torch.empty(padded_lengths.numel() + 1, device=dev, dtype=torch.int32)
+    cu_seqlens_padded[0] = 0
+    cu_seqlens_padded[1:] = torch.cumsum(padded_lengths, dim=0)
+    max_seqlen = torch.max(padded_lengths).to(dtype=torch.int32)
+
+    new_sample["cu_seqlens_padded"] = cu_seqlens_padded
+    new_sample["max_seqlen"] = max_seqlen
+
+    original_lengths = original_lengths.to(
+        device=dev, dtype=torch.int32, non_blocking=True
+    ).reshape(-1)
+    cu_seqlens = torch.empty(original_lengths.numel() + 1, device=dev, dtype=torch.int32)
+    cu_seqlens[0] = 0
+    cu_seqlens[1:] = torch.cumsum(original_lengths, dim=0).reshape(-1)
+    new_sample["cu_seqlens"] = cu_seqlens
+
+    return new_sample
+
+
+def broadcast_tensor(item, src_rank, group) -> None:
+    """Broadcast a tensor from src_rank to all ranks in the group."""
+    if item is not None:
+        torch.distributed.broadcast(item, src_rank, group=group)
+
+
+def broadcast_to_pp_group(
+    new_samples,
+    num_micro_batches,
+    seqlen_sum_this_global_batch,
+    seqlen_squared_sum_this_global_batch,
+    pp_group,
+    dev,
+):
+    """
+    Broadcast num_micro_batches, seqlen_sum_this_global_batch,
+    seqlen_squared_sum_this_global_batch and metadata to middle PP stages.
+    Before this broadcast, the new_samples on middle PP stages are None,
+    after this broadcast, the new_samples on middle PP stages contain the metadata but
+    without tokens, labels, loss_mask, position_ids.
+    """
+
+    pp_src_rank = torch.distributed.get_process_group_ranks(pp_group)[0]
+
+    if pp_group.size() > 2:
+        if pp_group.rank() == 0:
+            cu_seqlens_lengths = torch.tensor(
+                [sample["cu_seqlens"].numel() for sample in new_samples],
+                dtype=torch.float32,
+                device=dev,
+            )
+            cu_seqlens_padded_lengths = torch.tensor(
+                [sample["cu_seqlens_padded"].numel() for sample in new_samples],
+                dtype=torch.float32,
+                device=dev,
+            )
+            tensor_list = [
+                torch.tensor(
+                    [
+                        num_micro_batches,
+                        seqlen_sum_this_global_batch,
+                        seqlen_squared_sum_this_global_batch,
+                    ],
+                    dtype=torch.float32,
+                    device=dev,
+                )
+            ]
+            for sample in new_samples:
+                tensor_list.append(sample["max_seqlen"].reshape(1))
+            tensor_list.append(cu_seqlens_lengths)
+            tensor_list.append(cu_seqlens_padded_lengths)
+            for sample in new_samples:
+                tensor_list.append(sample["cu_seqlens"])
+                tensor_list.append(sample["cu_seqlens_padded"])
+            info_to_broadcast = torch.cat(tensor_list, dim=0).to(device=dev, dtype=torch.float32)
+            info_length_tensor = torch.tensor(
+                info_to_broadcast.shape[0], dtype=torch.int32, device=dev
+            )
+            broadcast_tensor(info_length_tensor, pp_src_rank, pp_group)
+            broadcast_tensor(info_to_broadcast, pp_src_rank, pp_group)
+        else:
+            info_length_tensor = torch.tensor(0, dtype=torch.int32, device=dev)
+            broadcast_tensor(info_length_tensor, pp_src_rank, pp_group)
+            info_to_broadcast = torch.empty(
+                info_length_tensor.item(), dtype=torch.float32, device=dev
+            )
+            broadcast_tensor(info_to_broadcast, pp_src_rank, pp_group)
+            if pp_group.rank() != pp_group.size() - 1:
+                # Middle PP stages receive the broadcasted info and unpack it.
+                # Cu-seqlens lengths are encoded explicitly so zero values inside
+                # the payload cannot be mistaken for tensor boundaries.
+                num_micro_batches = int(info_to_broadcast[0].item())
+                seqlen_sum_this_global_batch = info_to_broadcast[1].item()
+                seqlen_squared_sum_this_global_batch = info_to_broadcast[2].item()
+
+                cursor = 3
+                max_seqlens = info_to_broadcast[cursor : cursor + num_micro_batches]
+                cursor += num_micro_batches
+                cu_seqlens_lengths = info_to_broadcast[
+                    cursor : cursor + num_micro_batches
+                ].to(torch.int64)
+                cursor += num_micro_batches
+                cu_seqlens_padded_lengths = info_to_broadcast[
+                    cursor : cursor + num_micro_batches
+                ].to(torch.int64)
+                cursor += num_micro_batches
+
+                new_samples = []
+                for i in range(num_micro_batches):
+                    cu_seqlens_len = int(cu_seqlens_lengths[i].item())
+                    cu_seqlens_padded_len = int(cu_seqlens_padded_lengths[i].item())
+                    new_sample = {}
+                    new_sample["max_seqlen"] = max_seqlens[i].to(torch.int32)
+                    new_sample["cu_seqlens"] = info_to_broadcast[
+                        cursor : cursor + cu_seqlens_len
+                    ].to(torch.int32)
+                    cursor += cu_seqlens_len
+                    new_sample["cu_seqlens_padded"] = info_to_broadcast[
+                        cursor : cursor + cu_seqlens_padded_len
+                    ].to(torch.int32)
+                    cursor += cu_seqlens_padded_len
+                    new_samples.append(new_sample)
+
+    return (
+        new_samples,
+        num_micro_batches,
+        seqlen_sum_this_global_batch,
+        seqlen_squared_sum_this_global_batch,
+    )
+
+
+def broadcast_scalars(values: List, group, dev, dtype=torch.float32) -> List:
+    """
+    Broadcast scalar values from rank 0 to all ranks in the group.
+
+    Args:
+        values: List of scalar values to broadcast (only used on rank 0).
+        group: The process group to broadcast within.
+        dev: The device to use for the tensor.
+        dtype: The data type for the tensor.
+
+    Returns:
+        List of broadcasted values.
+    """
+    if group.size() <= 1:
+        return values
+
+    src_rank = torch.distributed.get_process_group_ranks(group)[0]
+    num_values = len(values)
+
+    if group.rank() == 0:
+        info_to_broadcast = torch.tensor(values, dtype=dtype, device=dev)
+    else:
+        info_to_broadcast = torch.zeros(num_values, dtype=dtype, device=dev)
+
+    broadcast_tensor(info_to_broadcast, src_rank, group)
+
+    if group.rank() != 0:
+        values = info_to_broadcast.cpu().tolist()
+
+    return values
+
+
+def create_data_iterator(new_samples, pp_group, tp_group, config):
+    """Handle virtual pipeline parallelism."""
+    if (
+        config.virtual_pipeline_model_parallel_size is not None
+        and config.virtual_pipeline_model_parallel_size > 1
+    ):
+        vpp_size = config.virtual_pipeline_model_parallel_size
+        if tp_group.rank() == 0:
+            if pp_group.rank() == 0 or pp_group.rank() == pp_group.size() - 1:
+                metadata = [
+                    {k: sample[k] for k in ["max_seqlen", "cu_seqlens", "cu_seqlens_padded"]}
+                    for sample in new_samples
+                ]
+                if pp_group.rank() == 0:
+                    new_data_iterator = [RerunDataIterator(iter(new_samples))] + [
+                        RerunDataIterator(iter(metadata)) for _ in range(vpp_size - 1)
+                    ]
+                else:
+                    new_data_iterator = [
+                        RerunDataIterator(iter(metadata)) for _ in range(vpp_size - 1)
+                    ] + [RerunDataIterator(iter(new_samples))]
+            else:
+                # on middle PP stages, the new_samples are the metadata
+                metadata = new_samples
+                new_data_iterator = [RerunDataIterator(iter(metadata)) for _ in range(vpp_size)]
+        else:
+            new_data_iterator = [None for _ in range(vpp_size)]
+    else:
+        new_data_iterator = RerunDataIterator(iter(new_samples)) if tp_group.rank() == 0 else None
+
+    return new_data_iterator
+
+
+def reroute_samples_to_dcp_ranks(
+    batch,
+    global_ids_this_rank,
+    global_id_seqlens,
+    sample_id_groups,
+    offsets,
+    dp_group,
+    tp_group,
+    dp_cp_group,
+    total_dcp_gpus,
+):
+    """
+    Reroutes the sub-samples to the correct rank after scheduling.
+
+    For each key in the batch dict, we perform an all-to-all communication
+    to transfer the data to the correct ranks.
+    """
+
+    def _gid_to_src_rank(gid: int) -> int:
+        dp_src_rank = torch.bucketize(gid, offsets[1:] - 1)
+        dcp_rank = (
+            torch.distributed.get_process_group_ranks(dp_group)[dp_src_rank] // tp_group.size()
+        ) % dp_cp_group.size()
+        return dcp_rank
+
+    gid2local_id = {int(gid): i for i, gid in enumerate(global_ids_this_rank)}
+    dcp_rank = dp_cp_group.rank()
+    dp_ranks = torch.distributed.get_process_group_ranks(dp_group)
+    dp_ranks = [(r // tp_group.size()) % dp_cp_group.size() for r in dp_ranks]
+
+    data_keys = batch[0].keys()
+
+    # Create the send plan
+    combined_sample_id_groups: List[List[int]] = [[] for _ in range(total_dcp_gpus)]
+    for d in range(total_dcp_gpus):
+        for sample_id_group in sample_id_groups:
+            combined_sample_id_groups[d].extend(sample_id_group[d])
+    for dest_rank in range(total_dcp_gpus):
+        combined_sample_id_groups[dest_rank].sort()
+
+    send_ids_sorted = [
+        gid for d in dp_ranks for gid in combined_sample_id_groups[d] if gid in global_ids_this_rank
+    ]
+
+    send_num_split = [0] * total_dcp_gpus
+    send_lens_split = [0] * total_dcp_gpus
+    for dest_rank in range(total_dcp_gpus):
+        if dest_rank in dp_ranks:
+            send_seq_lens = [
+                global_id_seqlens[gid][1]
+                for gid in combined_sample_id_groups[dest_rank]
+                if gid in global_ids_this_rank
+            ]
+            send_num_split[dest_rank] = len(send_seq_lens)
+            send_lens_split[dest_rank] = sum(send_seq_lens)
+        else:
+            send_lens_split[dest_rank] = 0
+
+    # Create the recv plan
+    recv_sample_id_groups = [[] for _ in range(total_dcp_gpus)]
+    for gid in combined_sample_id_groups[dcp_rank]:
+        src_rank = _gid_to_src_rank(gid)
+        recv_sample_id_groups[src_rank].append(gid)
+
+    recv_lens_split = [0] * total_dcp_gpus
+    for src_rank in range(total_dcp_gpus):
+        recv_lens_split[src_rank] = sum(
+            [global_id_seqlens[gid][1] for gid in recv_sample_id_groups[src_rank]]
+        )
+
+    recv_ids_sorted = [gid for d in range(total_dcp_gpus) for gid in recv_sample_id_groups[d]]
+    recv_counts = [len(recv_sample_id_groups[d]) for d in range(total_dcp_gpus)]
+
+    recv_samples = [{k: None for k in data_keys} for _ in range(sum(recv_counts))]
+
+    def _pack_sample_by_key(key: str) -> torch.Tensor:
+        flattened_tensors = []
+        for gid in send_ids_sorted:
+            t = batch[gid2local_id[gid]][key].to(torch.cuda.current_device(), non_blocking=True)
+            flattened_tensors.append(t.reshape(-1))
+        return (
+            torch.cat(flattened_tensors, dim=0)
+            if flattened_tensors
+            else torch.empty(1, device=torch.cuda.current_device(), dtype=batch[0][key].dtype)
+        )
+
+    def _unpack_sample_by_key(key: str, recv_tensor: torch.Tensor):
+        cursor = 0
+        for i, gid in enumerate(recv_ids_sorted):
+            sample_len = (
+                1 if key in ["original_seq_len", "padded_seq_len"] else global_id_seqlens[gid][1]
+            )
+            recv_samples[i][key] = recv_tensor[cursor : cursor + sample_len]
+            cursor += sample_len
+
+    for key in data_keys:
+        output_split_sizes, input_split_sizes = (
+            (recv_counts, send_num_split)
+            if key in ["original_seq_len", "padded_seq_len"]
+            else (recv_lens_split, send_lens_split)
+        )
+        send_tensor = _pack_sample_by_key(key)
+        recv_tensor_size = sum(output_split_sizes)
+        recv_tensor = torch.empty(
+            recv_tensor_size, device=torch.cuda.current_device(), dtype=send_tensor.dtype
+        )
+        torch.distributed.all_to_all_single(
+            output=recv_tensor,
+            input=send_tensor,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=dp_cp_group,
+        )
+        _unpack_sample_by_key(key, recv_tensor)
+
+    recv_sample_with_id = {recv_id: recv_samples[i] for i, recv_id in enumerate(recv_ids_sorted)}
+    return recv_sample_with_id
+
+
+def build_packed_microbatches(
+    grouped_samples: List[List[Dict[str, torch.Tensor]]], dev: torch.device
+) -> List[Dict[str, torch.Tensor]]:
+    """Build packed samples for each microbatch."""
+    num_micro_batches = len(grouped_samples)
+    seg_starts: List[int] = [0]
+    original_lens_tensors = []
+    padded_lens_tensors = []
+
+    for i in range(num_micro_batches):
+        samples = grouped_samples[i]
+        seg_starts.append(seg_starts[-1] + len(samples))
+        original_lens_tensors.extend([s["original_seq_len"].reshape(-1) for s in samples])
+        padded_lens_tensors.extend([s["padded_seq_len"].reshape(-1) for s in samples])
+
+    padded_lens_all_gpu = torch.cat(padded_lens_tensors, dim=0).to(dtype=torch.int32)
+    original_lens_all_gpu = torch.cat(original_lens_tensors, dim=0).to(dtype=torch.int32)
+
+    new_samples: List[Dict[str, torch.Tensor]] = []
+    for i in range(num_micro_batches):
+        samples = grouped_samples[i]
+        lens_padded = padded_lens_all_gpu[seg_starts[i] : seg_starts[i + 1]]
+        lens_original = original_lens_all_gpu[seg_starts[i] : seg_starts[i + 1]]
+        new_sample = _pack_sequences(samples, lens_padded, lens_original, dev)
+        new_samples.append(new_sample)
+
+    return new_samples
+
+
+def get_batch_and_global_seqlens(data_iterator, num_microbatches, dp_group):
+    """
+    Get the batch and global sequence lengths.
+    Each DP rank loads the same number of sequences, so we need to gather the sequence
+    lengths from all ranks then we can schedule the sequences into groups.
+    Args:
+        data_iterator: The data iterator.
+        num_microbatches: The number of microbatches.
+        dp_group: The data parallel group.
+
+    Returns:
+        batch: The batch.
+        global_id_seqlens: The global sequence lengths.
+        global_ids_this_rank: The global IDs locally present on this rank.
+    """
+
+    batch_list = [next(data_iterator) for _ in range(num_microbatches)]
+
+    batch = []
+    for item in batch_list:
+        if isinstance(item, dict):
+            batch.append(item)
+        elif isinstance(item, list):
+            batch.extend(item)
+        else:
+            raise ValueError(f"Invalid item type: {type(item)}")
+
+    # in sft_dataset.py, sequences are already packed before rescheduling,
+    # so we need to unpack them here and repack after rescheduling.
+    # This is only to adapt to the current megatron-lm sft_dataset.
+    # If you implement your own dataset, just have __getitem__ return List[Dict]
+    # and this step can be skipped.
+    batch = _unpack_batch(batch)
+
+    subsample_seqlens = torch.cat([sample["padded_seq_len"] for sample in batch]).to(
+        dtype=torch.int32, device=torch.cuda.current_device()
+    )
+
+    global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered = (
+        _get_global_seqlens_and_ids(subsample_seqlens, dp_group)
+    )
+
+    return batch, global_id_seqlens, global_ids_this_rank, offsets, seqlens_gathered
diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py
index 42146d1acd2..2fe168c6504 100644
--- a/megatron/core/datasets/gpt_dataset.py
+++ b/megatron/core/datasets/gpt_dataset.py
@@ -76,6 +76,26 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig):
     context_parallel_size: Optional[int] = None
     """The size of the context parallel group. Needed for padding in packed sequences."""
 
+    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__()
@@ -86,6 +106,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
diff --git a/megatron/core/datasets/readme.md b/megatron/core/datasets/readme.md
index 58721b7471b..40019642469 100644
--- a/megatron/core/datasets/readme.md
+++ b/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:
diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index b7de1013695..19764f15084 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -3391,3 +3391,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)
diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py
index 88bb070e105..1dca84dd393 100644
--- a/megatron/core/model_parallel_config.py
+++ b/megatron/core/model_parallel_config.py
@@ -59,7 +59,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 +69,12 @@ 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.
+    """
+
     expert_model_parallel_size: int = 1
     """Distributes Moe Experts across sub data parallel dimension."""
 
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index b2c23807bea..dc54748d07c 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -671,6 +671,9 @@ def forward_backward_no_pipelining(
         pg_collection.tp_dp_cp = parallel_state.get_tensor_and_data_parallel_group(
             with_context_parallel=True
         )
+        pg_collection.dp = parallel_state.get_data_parallel_group(
+            with_context_parallel=False, partial_data_parallel=False
+        )
 
     elif pg_collection is not None:
         assert hasattr(pg_collection, 'tp'), "pg_collection must have tp"
@@ -1001,6 +1004,9 @@ def forward_backward_pipelining_with_interleaving(
         pg_collection.tp_dp_cp = parallel_state.get_tensor_and_data_parallel_group(
             with_context_parallel=True
         )
+        pg_collection.dp = parallel_state.get_data_parallel_group(
+            with_context_parallel=False, partial_data_parallel=False
+        )
 
     elif p2p_communicator is not None and pg_collection is not None:
         model_type = get_model_type(model[0])
@@ -2160,6 +2166,9 @@ def forward_backward_pipelining_without_interleaving(
         pg_collection.tp_dp_cp = parallel_state.get_tensor_and_data_parallel_group(
             with_context_parallel=True
         )
+        pg_collection.dp = parallel_state.get_data_parallel_group(
+            with_context_parallel=False, partial_data_parallel=False
+        )
 
     elif p2p_communicator is not None and pg_collection is not None:
         assert hasattr(p2p_communicator, 'config'), "p2p_communicator must have a config"
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index ab84abd3a17..c273666ba10 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -2582,6 +2582,40 @@ def _scope_to_str(s):
                 self.attention_backend == AttnBackend.flash
             ), "Batch invariant mode only supports FlashAttention"
 
+        if self.sequence_packing_scheduler is not None:
+            # Check TE version.
+            if not HAVE_PACKAGING:
+                raise ImportError(
+                    "packaging is not installed. Please install it with `pip install packaging`."
+                )
+            # TODO: remove this after we fix the convergence issue with TE < 2.9.
+            if not (
+                is_te_min_version("2.9.0") or get_te_version() == PkgVersion("2.9.0.dev0+5b3092a")
+            ):
+                raise ValueError(
+                    "SFT sequence packing requires Transformer Engine >= 2.9.0 "
+                    f"but got {get_te_version()} (TE < 2.9.0 may have convergence issues)."
+                )
+
+            # Needed for passing variable sequences between pp stages.
+            self.variable_seq_lengths = True
+
+            # TODO(tailaim): add support for other dispatcher types
+            assert self.moe_token_dispatcher_type == "alltoall", (
+                f"sequence_packing only supports moe_token_dispatcher_type='alltoall', "
+                f"got '{self.moe_token_dispatcher_type}'"
+            )
+
+            supported_schedulers = ['dp_balanced']
+            if (
+                self.sequence_packing_scheduler is not None
+                and self.sequence_packing_scheduler not in supported_schedulers
+            ):
+                raise ValueError(
+                    f"Unsupported scheduler: {self.sequence_packing_scheduler}. "
+                    f"Available schedulers: {supported_schedulers}"
+                )
+
 
 @dataclass
 class MLATransformerConfig(TransformerConfig):
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 4be5d8b6776..f26b6082a8f 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -77,6 +77,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser):
     parser = _add_msc_args(parser)
     parser = _add_kitchen_quantization_arguments(parser)
     parser = _add_sft_args(parser)
+    parser = _add_varlen_dataset_args(parser)
 
     parser = _add_fault_injector_args(parser)
 
@@ -1202,13 +1203,6 @@ def validate_args(args, defaults={}):
     if args.rl_use_sequence_packing:
         args.consumed_train_bins = 0
 
-    # Support for variable sequence lengths across batches/microbatches.
-    # set it if the dataloader supports generation of variable sequence lengths
-    # across batches/microbatches. Due to additional communication overhead
-    # during pipeline parallelism, it should not be set if sequence length
-    # is constant during training.
-    args.variable_seq_lengths = False
-
     # Iteration-based training.
     # Skip these checks when skip_train is set: LR config is irrelevant.
     if args.train_iters and not args.skip_train:
@@ -1380,6 +1374,23 @@ def validate_args(args, defaults={}):
         assert args.dataloader_type == 'single', 'Hybrid context parallelism only supported with single dataloader type'
         assert args.calculate_per_token_loss, 'Hybrid context parallelism must be used with --calculate-per-token-loss'
 
+    # Support for variable sequence lengths across batches/microbatches.
+    # set it if the dataloader supports generation of variable sequence lengths
+    # across batches/microbatches. Due to additional communication overhead
+    # during pipeline parallelism, it should not be set if sequence length
+    # is constant during training.
+    args.variable_seq_lengths = False
+    if args.mock_data and args.sft and args.sft_mock_dataset_config_json is None:
+        args.sft_mock_dataset_config_json = json.dumps(
+            {
+                "mode": "distribution",
+                "type": "lognormal",
+                "min_seq_len": args.seq_length // 2,
+                "max_seq_len": args.seq_length,
+                "mean_seq_len": args.seq_length // 4 * 3,
+                "lognormal_sigma": 1.1,
+            }
+        )
     # disable async_tensor_model_parallel_allreduce when
     # model parallel memory optimization is enabled
     if (args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1) \
@@ -1493,6 +1504,47 @@ def validate_args(args, defaults={}):
     if args.ckpt_format == "fsdp_dtensor":
         assert args.use_megatron_fsdp, "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP."
 
+    # --use-varlen-dataset: independent of --sft. Cannot be combined with --sft
+    # because they are mutually-exclusive top-level dataset selectors that both
+    # drive the packed-sequence (THD) path.
+    if args.use_varlen_dataset:
+        assert not args.sft, (
+            "--use-varlen-dataset and --sft are mutually exclusive; both "
+            "select the packed-sequence dataset family. Pick one."
+        )
+        if args.varlen_sbhd_validation:
+            assert args.sequence_packing_scheduler is None, (
+                "--varlen-sbhd-validation does not use a sequence packing "
+                "scheduler; drop --sequence-packing-scheduler."
+            )
+            # SBHD validation is a real-data numerical-reference path only;
+            # MockVarlenDataset does not implement it.
+            assert not args.mock_data, (
+                "--varlen-sbhd-validation is not supported with --mock-data; "
+                "SBHD validation requires a real dataset."
+            )
+        else:
+            # VarlenDataset emits one unpacked sample per __getitem__; it
+            # relies on an upstream packing scheduler to group variable-length
+            # samples into THD batches. Auto-pick a default scheduler when
+            # the user did not request one explicitly:
+            # Otherwise fall back to ``dp_balanced`` (static packing).
+            if args.sequence_packing_scheduler is None:
+                args.sequence_packing_scheduler = 'dp_balanced'
+
+    # Packed-sequence buffer-size check. Placed after varlen scheduler
+    # auto-select so it validates the final resolved scheduler.
+    if args.sequence_packing_scheduler is not None:
+        args.variable_seq_lengths = True
+        assert args.max_seqlen_per_dp_cp_rank is not None, (
+            "--max-seqlen-per-dp-cp-rank must be set when using sequence packing"
+        )
+        total_cp_ranks = args.context_parallel_size
+        assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, (
+            f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) '
+            f'must be >= single sequence max length ({args.seq_length})'
+        )
+
     # Data blend checks
     assert args.mock_data + \
            bool(args.data_path) + \
@@ -2117,6 +2169,9 @@ def _add_network_size_args(parser):
         "persist_layer_norm",
         "bias_dropout_fusion",
         "apply_rope_fusion",
+        "max_seqlen_per_dp_cp_rank",
+        "hybrid_context_parallel",
+        "sequence_packing_scheduler",
     ]
     transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude)
     transformer_group = transformer_factory.build_group(parser, "transformer configuration")
@@ -2872,6 +2927,14 @@ def _add_distributed_args(parser):
                        'all layers will share the same communication type. Users can also '
                        'specify separated types for each layer like '
                        '--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p')
+    group.add_argument('--max-seqlen-per-dp-cp-rank', type=int, default=None,
+                       help='Maximum sequence length per CP rank. This is used to calculate the '
+                       'number of sub-samples assigned to each CP rank when using heterogeneous context parallel.')
+    group.add_argument('--hybrid-context-parallel', action='store_true', default=False,
+                       help='Enables hybrid context parallel. This is used to balance the workload '
+                       'of each CP rank when we use packed samples with variable sequence lengths. '
+                       'Requires --max-seqlen-per-dp-cp-rank to be set.')
+    group.add_argument('--sequence-packing-scheduler', type=str, default=None, choices=['dp_balanced'])
     group.add_argument('--fake-process-group', action='store_true', default=False,
                        help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \
                        This is quite useful for profiling memory usage of distributed training with just one GPU. \
@@ -3405,8 +3468,69 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser):
 def _add_sft_args(parser):
     group = parser.add_argument_group(title='sft')
     group.add_argument('--sft', action="store_true", help='Megatron SFT training')
-    group.add_argument('--sft-tokenizer-prompt-format', type=str, default="nemotron-h-aligned",
-                       help='SFT prompt format.')
+    group.add_argument(
+        '--sft-tokenizer-prompt-format',
+        type=str,
+        default="nemotron-h-aligned",
+        help='SFT prompt format.',
+    )
+    group.add_argument(
+        '--sft-mock-dataset-config-json',
+        type=str,
+        default=None,
+        help='This config provides the necessary information for the mock dataset. '
+        'Accepts either an inline JSON literal or a path to a JSON file containing '
+        'the same schema. You can either specify a CSV file that contains sequence lengths, '
+        'where each line stores the length of a sequence, for example: '
+        '{"mode":"file","path":"/path/to/file"}. Alternatively, you can specify a distribution '
+        '(currently only supporting lognormal distribution) along with the required parameters, '
+        'for example, {"mode":"distribution","type":"lognormal","min_seq_len":1024,'
+        '"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, where sigma controls '
+        'the variability of the lognormal distribution. '
+        'If not specified and --mock-data is set, defaults to a lognormal distribution with '
+        'min_seq_len=seq_length//2, max_seq_len=seq_length, mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.',
+    )
+    return parser
+
+
+def _add_varlen_dataset_args(parser):
+    group = parser.add_argument_group(title='varlen dataset')
+    group.add_argument(
+        '--use-varlen-dataset',
+        action="store_true",
+        help='Train with VarlenDataset, a variable-length packed (THD) dataset '
+        'that consumes instruction-tuning data from a HuggingFace Hub repo id, '
+        'a local parquet file, or a local jsonl file. Schema (alpaca / sharegpt '
+        '/ openai-messages) is auto-detected from the dataset columns. '
+        'Mutually exclusive with --sft. Auto-picks a sequence packing '
+        'scheduler when none is given: ``dp_balanced``. '
+        'Combine with --mock-data for a synthetic lognormal sequence-length '
+        'distribution; see --varlen-mock-dataset-config-json.',
+    )
+    group.add_argument(
+        '--varlen-sbhd-validation',
+        action="store_true",
+        help='Reference SBHD mode for THD numerical verification. When set, '
+        'VarlenDataset emits SBHD-style samples right-padded to '
+        '--seq-length (no cu_seqlens, no packing scheduler), so the run can '
+        'be compared against the THD path to validate correctness. '
+        'Incompatible with --sequence-packing-scheduler.',
+    )
+    group.add_argument(
+        '--varlen-mock-dataset-config-json',
+        type=str,
+        default=None,
+        help='Mock-dataset config for --use-varlen-dataset --mock-data. '
+        'Accepts either an inline JSON literal or a path to a JSON file containing '
+        'the same schema as --sft-mock-dataset-config-json: either '
+        '{"mode":"file","path":"/path/to/lengths.csv"}, '
+        '{"mode":"distribution","type":"lognormal","min_seq_len":1024,'
+        '"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, or '
+        '{"mode":"verification","data_path":"/prefix/of/IndexedDataset"}. '
+        'If not specified, defaults to a lognormal distribution with '
+        'min_seq_len=seq_length//2, max_seq_len=seq_length, '
+        'mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.',
+    )
     return parser
 
 def _add_logits_distillation_args(parser):
diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py
index 296acc97941..27c27393f82 100644
--- a/megatron/training/datasets/data_samplers.py
+++ b/megatron/training/datasets/data_samplers.py
@@ -39,14 +39,16 @@ def build_pretraining_data_loader(dataset, consumed_samples):
     is_eval = split in (Split.valid, Split.test)
     micro_batch_size = getattr(args, 'eval_micro_batch_size', args.micro_batch_size) if is_eval else args.micro_batch_size
     global_batch_size = getattr(args, 'eval_global_batch_size', args.global_batch_size) if is_eval else args.global_batch_size
-
     if split == Split.valid and args.full_validation:
         batch_sampler = MegatronFullValidationSampler(
             total_samples=len(dataset),
             data_parallel_rank=mpu.get_data_parallel_rank(),
             data_parallel_size=mpu.get_data_parallel_world_size())
     elif args.dataloader_type == 'single':
-        if args.hybrid_context_parallel:
+        if (
+            getattr(args, "hybrid_context_parallel", False)
+            and getattr(args, "sequence_packing_scheduler", None) is None
+        ):
             batch_sampler = HybridCPMegatronPretrainingSampler(
                 total_samples=len(dataset),
                 consumed_samples=consumed_samples,
@@ -55,7 +57,8 @@ def build_pretraining_data_loader(dataset, consumed_samples):
                 data_parallel_rank=mpu.get_data_parallel_rank(),
                 data_parallel_size=mpu.get_data_parallel_world_size())
         else:
-            # Megatron sampler
+            # Megatron sampler. Packing schedulers consume one microbatch at a
+            # time and form packed global batches themselves.
             batch_sampler = MegatronPretrainingSampler(
                 total_samples=len(dataset),
                 consumed_samples=consumed_samples,
@@ -97,9 +100,21 @@ def close_nvidia_fds():
     maybe_worker_init_fn = (
         worker_init_fn if args.num_workers > 0 else None
     )
-    # Torch dataloader.
-    if args.hybrid_context_parallel:
-        extra_kwargs = {"collate_fn": lambda x: x,}
+    # Identity collate for VarlenDataset and packing-scheduler paths;
+    # they emit one variable-length dict per sample, not stack-able by
+    # the default collate. --varlen-sbhd-validation is excluded: it bypasses
+    # packing and emits fixed-length [seq_length] samples that the default
+    # collate stacks normally.
+    if (
+        (
+            getattr(args, "use_varlen_dataset", False)
+            and not getattr(args, "varlen_sbhd_validation", False)
+        )
+        or getattr(args, "hybrid_context_parallel", False)
+        or getattr(args, "sequence_packing_scheduler", None) is not None
+        or getattr(args, "use_vanilla_collate_fn", False)
+    ):
+        extra_kwargs = {"collate_fn": lambda x: x}
     else:
         extra_kwargs = {}
     return torch.utils.data.DataLoader(
@@ -225,7 +240,6 @@ def __iter__(self):
                 global_batch_idx.extend(batch[start_idx[i]:end_idx[i]])
             yield global_batch_idx
 
-
 class MegatronFullValidationSampler:
     """Sampler for full validation that handles small datasets gracefully.
 
diff --git a/megatron/training/datasets/sft_dataset.py b/megatron/training/datasets/sft_dataset.py
index 9de5d2a52fe..0785cab5f80 100644
--- a/megatron/training/datasets/sft_dataset.py
+++ b/megatron/training/datasets/sft_dataset.py
@@ -1,15 +1,18 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 
-import atexit, json
+import atexit
 from collections import Counter
-from typing import Any, Dict, Optional
+import math
+from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
+import pandas as pd
 import torch
 
 from megatron.core.datasets.gpt_dataset import GPTDatasetConfig
 from megatron.core.datasets.megatron_dataset import LowLevelDataset, MegatronDataset
 from megatron.core.datasets.utils import Split
+from megatron.training.datasets.utils import load_json_arg
 
 IGNORE_INDEX = -100
 
@@ -61,6 +64,8 @@ def __init__(
         config: GPTDatasetConfig,
     ) -> None:
         super().__init__(dataset, dataset_path, indices, num_samples, index_split, config)
+        # Pre-calculate padding divisor to avoid redundant computation in get_item
+        self.padding_divisor = self._calculate_padding_divisor()
 
     @staticmethod
     def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int:
@@ -88,6 +93,26 @@ def _split_conversations(self, merged_conversations):
             split_conversations.append(current)
         return split_conversations
 
+    def _calculate_padding_divisor(self) -> int:
+        """
+            Calculate the divisor used for sequence padding.
+            tp_pad = tp_size * 2 if tp_size > 1 else 1
+            cp_pad = cp_size * 2 if cp_size > 1 else 1
+            cp_pad = cp_pad * dp_size if hybrid_cp else cp_pad
+            divisor = cp_pad * tp_pad
+        """
+        if self.config.hybrid_context_parallel:
+            # Hybrid CP: consider both CP and DP
+            cp_pad = self.config.data_parallel_size * self.config.context_parallel_size * 2
+        else:
+            # Standard CP: only consider CP
+            cp_pad = self.config.context_parallel_size * 2 if self.config.context_parallel_size > 1 else 1
+        tp_pad = self.config.sequence_parallel_size if self.config.sequence_parallel_size > 0 else 1
+        divisor = cp_pad * tp_pad
+        # TODO(tailaim): do we need to pad for FP8 execution?
+        # divisor = ((divisor + 15) // 16) * 16
+        return divisor
+
     def __getitem__(self, idx: int) -> Dict[str, Any]:
 
         tokenizer = self.config.tokenizer
@@ -124,12 +149,11 @@ def extend_with_padding(tokens, targets, positions, pad_len):
             assert not self.config.reset_position_ids
             pack_positions.extend(range(len(tokens_list)))
 
-            if self.config.context_parallel_size > 1:
-                pad_granularity = self.config.context_parallel_size * 2
-                mod_token_count = len(pack_tokens) % pad_granularity
-                if mod_token_count != 0:
-                    pad_len = pad_granularity - mod_token_count
-                    extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len)
+            pad_granularity = self.padding_divisor
+            mod_token_count = len(pack_tokens) % pad_granularity
+            if mod_token_count != 0:
+                pad_len = pad_granularity - mod_token_count
+                extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len)
 
             # TODO(duncan): Consider also padding to multiple of number of tokens here. This might
             # be needed for efficiency (and potentially set via command-line argument).
@@ -190,3 +214,171 @@ def extend_with_padding(tokens, targets, positions, pad_len):
             'cu_seqlens': cu_seqlens,
             'max_seqlen': max_seqlen,
         }
+
+
+class MockSFTLowLevelDataset:
+    """The low-level mock dataset for SFT
+
+    Args:
+        mode (str): Either 'file' or 'distribution'.
+        **kwargs: Additional arguments depending on mode.
+            For mode='file': path (str) - path to a CSV file with sequence lengths.
+            For mode='distribution': type (str), min_seq_len (int), max_seq_len (int),
+                mean_seq_len (int), and distribution-specific params (e.g. lognormal_sigma).
+    """
+
+    seed: int = 0
+    """The hard-coded random seed to use to set the NumPy RNG"""
+
+    size: int = 1000000
+    """The hard-coded number of sequence to generate"""
+
+    def __init__(self, mode: str, **kwargs) -> None:
+        np.random.seed(self.seed)
+
+        if mode == "file":
+            self.sequence_lengths = np.array(pd.read_csv(kwargs["path"])).flatten()
+            self.size = len(self.sequence_lengths)
+        elif mode == "distribution":
+            min_seq_len = kwargs["min_seq_len"]
+            max_seq_len = kwargs["max_seq_len"]
+            mean_seq_len = kwargs["mean_seq_len"]
+            if kwargs["type"] == "lognormal":
+                lognormal_sigma = kwargs["lognormal_sigma"]
+                self.sequence_lengths = self.generate_lognormal_samples(
+                    self.size, mean_seq_len, lognormal_sigma, min_seq_len, max_seq_len
+                )
+            else:
+                raise ValueError(f"Unsupported distribution type {kwargs['type']}")
+        else:
+            raise ValueError(f"Unsupported mode '{mode}', must be 'file' or 'distribution'")
+
+    def generate_lognormal_samples(self, size, mean, sigma, min_seq_len, max_seq_len):
+        mu = np.log(mean) - sigma**2 / 2
+        samples = np.random.lognormal(mu, sigma, size)
+        samples = np.clip(samples, min_seq_len, max_seq_len)
+        return samples.astype(int)
+
+    def __len__(self) -> int:
+        return self.size
+
+    def __getitem__(self, idx: int) -> List[np.ndarray]:
+        # the length of sample is 'length', but only length-1 elements are generated here,
+        # because an eod token will be appended at the end later in SFTDataset
+
+        length = self.sequence_lengths[idx % self.size]
+        sample = np.arange(1, length, dtype=np.int64)
+        return sample
+
+
+class MockSFTDataset(SFTDataset):
+    """The mock dataset used during SFT"""
+
+    def __init__(
+        self,
+        dataset: LowLevelDataset,
+        dataset_path: Optional[str],
+        indices: np.ndarray,
+        num_samples: Optional[int],
+        index_split: Split,
+        config: GPTDatasetConfig,
+    ) -> None:
+        super().__init__(dataset, dataset_path, indices, num_samples, index_split, config)
+
+    @staticmethod
+    def build_low_level_dataset(dataset_path: str, config: GPTDatasetConfig) -> LowLevelDataset:
+        if config.sft_mock_dataset_config_json is None:
+            mock_config = {
+                "mode": "distribution",
+                "type": "lognormal",
+                "min_seq_len": config.sequence_length // 2,
+                "max_seq_len": config.sequence_length,
+                "mean_seq_len": config.sequence_length // 4 * 3,
+                "lognormal_sigma": 1.1,
+            }
+        else:
+            mock_config = load_json_arg(config.sft_mock_dataset_config_json)
+        return MockSFTLowLevelDataset(**mock_config)
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+
+        tokenizer = self.config.tokenizer
+        pack_length = self.config.sequence_length
+        eod = tokenizer.eod
+        pad = tokenizer.pad
+
+        tokens = self.dataset[int(self.indices[idx % len(self.indices)])]
+
+        def extend_with_padding(tokens, targets, positions, pad_len):
+            tokens.extend([pad] * pad_len)
+            targets.extend([pad] * pad_len)
+            positions.extend(range(positions[-1] + 1, positions[-1] + 1 + pad_len))
+
+        # Convert tokens to list and add EOD
+        tokens_list = tokens.tolist()
+        if tokens_list[-1] != eod:
+            tokens_list.append(eod)
+        targets_list = list(tokens_list)
+
+        pack_tokens = list(tokens_list)
+        pack_targets = list(targets_list)
+        pack_positions = list(range(len(tokens_list)))
+        cu_seqlens = [0]
+
+        # Pad to padding_divisor alignment
+        if self.padding_divisor > 1:
+            mod_token_count = len(pack_tokens) % self.padding_divisor
+            if mod_token_count != 0:
+                pad_len = self.padding_divisor - mod_token_count
+                extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len)
+
+        # Record padded boundary after padding
+        cu_seqlens.append(len(pack_tokens))
+
+        # Handle any necessary truncation
+        if len(pack_tokens) >= pack_length + 1:  # +1 here to account for later alignment
+            max_body = pack_length - 1
+            pack_tokens = pack_tokens[:max_body]
+            pack_targets = pack_targets[:max_body]
+            pack_tokens.extend([eod, pad])
+            pack_targets.extend([eod, pad])
+            pack_positions = pack_positions[:pack_length + 1]
+            cu_seqlens[-1] = len(pack_tokens) - 1
+
+        # Handle any necessary padding
+        if len(pack_tokens) < pack_length + 1:  # +1 here to account for later alignment
+            pad_len = pack_length + 1 - len(pack_tokens)
+            extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len)
+            cu_seqlens[-1] = len(pack_tokens) - 1
+
+        assert len(pack_tokens) == pack_length + 1
+        assert len(pack_targets) == pack_length + 1
+        assert len(pack_positions) == pack_length + 1
+
+        # Align and convert to tensors
+        input_ids = torch.tensor(pack_tokens[:-1], dtype=torch.int64)
+        labels = torch.tensor(pack_targets[1:], dtype=torch.int64)
+        position_ids = torch.tensor(pack_positions[:-1], dtype=torch.int64)
+
+        # Loss mask
+        loss_mask = torch.ones(pack_length, dtype=torch.float32)
+        loss_mask[labels == pad] = 0.0
+        loss_mask[labels == IGNORE_INDEX] = 0.0
+
+        assert len(cu_seqlens) >= 2
+        cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32)
+        # Calculating max_seqlen here because of possible effects of truncation and padding
+        adjacent_diffs = cu_seqlens[1:] - cu_seqlens[:-1]
+        max_seqlen = adjacent_diffs.max()  # max_seqlen is a 0-D tensor
+
+        return {
+            'tokens': input_ids,
+            'labels': labels,
+            'loss_mask': loss_mask,
+            'position_ids': position_ids,
+            'cu_seqlens': cu_seqlens,
+            'max_seqlen': max_seqlen,
+        }
diff --git a/megatron/training/datasets/utils.py b/megatron/training/datasets/utils.py
new file mode 100644
index 00000000000..1fe6d7ef83e
--- /dev/null
+++ b/megatron/training/datasets/utils.py
@@ -0,0 +1,29 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+
+"""Shared utilities for training-side dataset helpers."""
+
+import json
+import os
+from typing import Any, Optional
+
+
+def load_json_arg(spec: Optional[str]) -> Optional[Any]:
+    """Parse a CLI JSON argument that may be either a JSON literal or a path
+    to a JSON file.
+
+    The argument is interpreted as a file path when ``spec`` points to an
+    existing regular file on the local filesystem; otherwise it is parsed as
+    a JSON literal string. Returns ``None`` when ``spec`` itself is ``None``,
+    so callers can use it transparently for optional CLI flags.
+
+    Used by the ``--sft-mock-dataset-config-json`` and
+    ``--varlen-mock-dataset-config-json`` flags, which both accept either an
+    inline JSON snippet or the path to a file containing the same JSON
+    document.
+    """
+    if spec is None:
+        return None
+    if os.path.isfile(spec):
+        with open(spec, "r") as f:
+            return json.load(f)
+    return json.loads(spec)
diff --git a/megatron/training/datasets/varlen_dataset.py b/megatron/training/datasets/varlen_dataset.py
new file mode 100644
index 00000000000..c2533f795bb
--- /dev/null
+++ b/megatron/training/datasets/varlen_dataset.py
@@ -0,0 +1,548 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+
+"""Variable-length packed (THD) dataset for SFT-style instruction data.
+
+This dataset is the entry point for the ``--use-varlen-dataset`` flag. It is
+independent of the ``--sft`` flag (no implicit coupling) but shares the same
+THD packing / cu_seqlens / dynamic-CP padding logic by extending the existing
+:class:`SFTDataset` family. The variable-length aspect is what matters here:
+samples have wildly different lengths and are packed into THD format for
+training throughput.
+
+Compared to :class:`SFTDataset`, this dataset adds:
+
+  * **Multi-source loading** — accepts HuggingFace Hub repo ids
+    (``owner/repo``), local ``.parquet`` files, and local ``.jsonl/.json``
+    files; the latter are read via pandas to sidestep pyarrow's per-chunk
+    JSON schema inference which fails when sample fields vary across rows.
+
+  * **Auto schema detection** — four input layouts are auto-detected by column
+    name. The three instruction-tuning layouts are normalized to the messages
+    list format expected by the parent ``SFTDataset.__getitem__``; the
+    ``pretrain-text`` fallback instead returns a raw string handled separately
+    in :meth:`VarlenDataset.__getitem__`:
+
+      * **openai-messages** — column ``messages`` (Llama post-training,
+        HuggingFaceH4/no_robots, ...)
+      * **sharegpt** — column ``conversations`` (OpenOrca, Vicuna, ...)
+      * **alpaca / dolly** — at least one of
+        ``instruction|prompt|query|question`` + one of
+        ``output|response|completion|answer``, plus optional context field
+        ``input|context``.
+      * **pretrain-text** — column ``text``; returns the raw string (no
+        messages list, no role masking), tokenized as plain pretraining text.
+
+  * **Mock variant** — :class:`MockVarlenDataset` mirrors
+    :class:`MockSFTDataset` end-to-end (synthetic lognormal sequence-length
+    distribution / fixed-length file / verification mode from an
+    ``IndexedDataset``), configured via
+    ``--varlen-mock-dataset-config-json``.
+
+Limitations (raise a clear ``ValueError`` instead of silently mishandling):
+
+  * Sample content/value must be a plain string — multi-modal content lists
+    (image+text parts) are not supported.
+  * Tree-structured (OpenAssistant oasst1) and preference (chosen/rejected)
+    datasets are out of scope.
+  * For HF Hub repos, only ``split="train"`` is loaded.
+"""
+
+import os
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
+
+import numpy as np
+import torch
+
+from megatron.core.datasets.gpt_dataset import GPTDatasetConfig
+from megatron.core.datasets.megatron_dataset import LowLevelDataset
+from megatron.core.datasets.utils import Split
+from megatron.training.datasets.sft_dataset import (
+    IGNORE_INDEX,
+    MockSFTDataset,
+    MockSFTLowLevelDataset,
+    SFTDataset,
+    SFTLowLevelDataset,
+)
+from megatron.training.datasets.utils import load_json_arg
+
+# Field-name synonyms (probed in order; first non-empty wins).
+_INSTRUCTION_FIELDS: Tuple[str, ...] = (
+    "instruction", "prompt", "query", "question",
+)
+_OUTPUT_FIELDS: Tuple[str, ...] = (
+    "output", "response", "completion", "answer",
+)
+# Supplementary user-turn context: Stanford Alpaca's "input", Dolly's "context".
+_EXTRA_INPUT_FIELDS: Tuple[str, ...] = ("input", "context")
+
+# ShareGPT "from" value -> chat-template "role". Unknown values fall back to
+# "user" so downstream tokenization does not crash on unfamiliar speakers.
+_SHAREGPT_ROLE_MAP: Dict[str, str] = {
+    "human": "user",
+    "user": "user",
+    "gpt": "assistant",
+    "assistant": "assistant",
+    "model": "assistant",
+    "chatgpt": "assistant",
+    "bing": "assistant",
+    "bard": "assistant",
+    "system": "system",
+    "tool": "tool",
+    "function": "tool",
+    "observation": "tool",
+}
+
+
+def _looks_like_hf_id(path: str) -> bool:
+    """Heuristic: does ``path`` look like an ``owner/repo`` HF dataset id?
+
+    True iff ``path`` contains ``/``, is not an absolute/relative file path,
+    and does not exist on the local filesystem.
+    """
+    if not path:
+        return False
+    if os.path.exists(path):
+        return False
+    if path.startswith(("/", "./", "../")):
+        return False
+    return "/" in path
+
+
+def _first_present(
+    sample: Dict[str, Any], fields: Iterable[str]
+) -> Optional[str]:
+    """Return the first non-empty string value among the given fields, or None."""
+    for f in fields:
+        v = sample.get(f)
+        if v in (None, ""):
+            continue
+        if not isinstance(v, str):
+            raise ValueError(
+                f"VarlenDataset: field '{f}' must be a string, "
+                f"got {type(v).__name__}."
+            )
+        return v
+    return None
+
+
+def _ensure_str_content(content: Any, where: str) -> str:
+    """Validate that a turn's content is a plain string (reject multi-modal lists)."""
+    if content is None:
+        return ""
+    if not isinstance(content, str):
+        raise ValueError(
+            f"VarlenDataset: {where} content must be a string, "
+            f"got {type(content).__name__}. Multi-modal datasets (e.g. "
+            "content as a list of image/text parts) are not supported."
+        )
+    return content
+
+
+def _alpaca_to_messages(sample: Dict[str, Any]) -> List[Dict[str, str]]:
+    """Convert an Alpaca/Dolly-style sample to a 3-turn messages list."""
+    instruction = _first_present(sample, _INSTRUCTION_FIELDS) or ""
+    extra_input = _first_present(sample, _EXTRA_INPUT_FIELDS) or ""
+    output = _first_present(sample, _OUTPUT_FIELDS) or ""
+    user_content = (
+        f"{instruction}\n\n{extra_input}" if extra_input else instruction
+    )
+    return [
+        {"role": "system", "content": ""},
+        {"role": "user", "content": user_content},
+        {"role": "assistant", "content": output},
+    ]
+
+
+def _sharegpt_to_messages(sample: Dict[str, Any]) -> List[Dict[str, str]]:
+    """Convert a ShareGPT ``conversations`` sample to a messages list.
+
+    Prepends an empty ``system`` turn unless the conversation already starts
+    with one, so ``SFTDataset._split_conversations`` treats the sample as a
+    single conversation.
+    """
+    conv = sample.get("conversations") or []
+    out: List[Dict[str, str]] = []
+    first_speaker = (conv[0].get("from") or "").lower() if conv else ""
+    if first_speaker != "system":
+        out.append({"role": "system", "content": ""})
+    for turn in conv:
+        speaker = (turn.get("from") or "").lower()
+        role = _SHAREGPT_ROLE_MAP.get(speaker, "user")
+        content = _ensure_str_content(turn.get("value"), f"sharegpt turn role={role}")
+        out.append({"role": role, "content": content})
+    return out
+
+
+def _messages_passthrough(sample: Dict[str, Any]) -> List[Dict[str, str]]:
+    """Pass through an OpenAI ``messages`` sample, ensuring a leading system turn.
+
+    Strips any keys other than ``role``/``content`` (e.g. ``name``,
+    ``tool_calls``) since they are not part of the chat-template input
+    expected by SFTTokenizer.
+    """
+    raw = list(sample.get("messages") or [])
+    if raw and raw[0].get("role") != "system":
+        raw = [{"role": "system", "content": ""}] + raw
+    out: List[Dict[str, str]] = []
+    for m in raw:
+        role = m.get("role") or "user"
+        content = _ensure_str_content(m.get("content"), f"messages turn role={role}")
+        out.append({"role": role, "content": content})
+    return out
+
+
+def _raw_text_loader(sample: Dict[str, Any]) -> str:
+    """Return the ``text`` column unchanged for pretrain-style packed runs.
+
+    Unlike the SFT schemas this returns a plain string (no messages list).
+    :class:`VarlenDataset.__getitem__` dispatches on the return type to pick
+    a tokenization path that skips chat templating and prompt masking.
+    """
+    text = sample.get("text") or ""
+    if not isinstance(text, str):
+        raise ValueError(
+            f"VarlenDataset (pretrain-text schema): 'text' must be a string, "
+            f"got {type(text).__name__}."
+        )
+    return text
+
+
+def _select_converter(
+    column_names: List[str],
+) -> Tuple[Callable[[Dict[str, Any]], Any], str]:
+    """Pick a sample converter based on dataset column names.
+
+    Priority (most explicit first): openai-messages > sharegpt > alpaca/dolly
+    > pretrain-text. ``pretrain-text`` is the fallback for datasets that
+    only carry a single ``text`` column (e.g. Dolma / OLMo midtraining
+    corpora) — long-context pretraining packed through the same THD path
+    as SFT.
+    """
+    cols = set(column_names)
+    if "messages" in cols:
+        return _messages_passthrough, "openai-messages"
+    if "conversations" in cols:
+        return _sharegpt_to_messages, "sharegpt"
+    has_instr = any(f in cols for f in _INSTRUCTION_FIELDS)
+    has_out = any(f in cols for f in _OUTPUT_FIELDS)
+    if has_instr and has_out:
+        return _alpaca_to_messages, "alpaca"
+    if "text" in cols:
+        return _raw_text_loader, "pretrain-text"
+    raise ValueError(
+        "VarlenDataset cannot infer schema from columns "
+        f"{sorted(cols)}. Supported schemas: "
+        f"alpaca/dolly ({'|'.join(_INSTRUCTION_FIELDS)} + "
+        f"{'|'.join(_OUTPUT_FIELDS)} [+ optional {'|'.join(_EXTRA_INPUT_FIELDS)}]), "
+        "sharegpt (conversations), openai-messages (messages), "
+        "pretrain-text (text)."
+    )
+
+
+class VarlenLowLevelDataset(SFTLowLevelDataset):
+    """Low-level loader: HF Hub repo / local parquet / local jsonl, normalized.
+
+    Dataset path interpretation:
+
+      * HF Hub repo id (e.g. ``Yukang/LongAlpaca-12k``) — contains ``/`` and
+        does not exist on the local filesystem; loaded via
+        ``datasets.load_dataset(path, split="train")``.
+      * Local ``.parquet`` — loaded via
+        ``datasets.load_dataset("parquet", data_files=path, split="all")``;
+        parquet's footer schema makes chunked loading safe.
+      * Otherwise local jsonl/json — loaded via pandas
+        ``read_json(lines=True)`` and wrapped in ``Dataset.from_pandas``.
+        We avoid ``datasets.load_dataset("json", ...)`` for local files
+        because its pyarrow-based JSON reader infers schema per parallel
+        chunk and fails with ``CastError`` when the union of fields varies
+        between rows (e.g. LongAlpaca-12k).
+
+    A per-sample converter is selected once at construction time based on
+    column names and applied at access time. The instruction-tuning schemas
+    convert to a messages list; the ``pretrain-text`` fallback returns the raw
+    string instead.
+    """
+
+    def __init__(self, dataset_path: str) -> None:
+        try:
+            from datasets import Dataset, load_dataset
+        except ImportError as exc:
+            raise ImportError(
+                "VarlenDataset requires the `datasets` library "
+                "(pip install datasets)."
+            ) from exc
+
+        if _looks_like_hf_id(dataset_path):
+            self.dataset = load_dataset(dataset_path, split="train")
+        elif dataset_path.endswith(".parquet"):
+            self.dataset = load_dataset(
+                "parquet", data_files=dataset_path, split="all"
+            )
+        else:
+            try:
+                import pandas as pd
+            except ImportError as exc:
+                raise ImportError(
+                    "VarlenDataset requires `pandas` to load local jsonl "
+                    "files (pip install pandas)."
+                ) from exc
+            df = pd.read_json(dataset_path, lines=True)
+            self.dataset = Dataset.from_pandas(df, preserve_index=False)
+
+        self._converter, self._schema_name = _select_converter(
+            list(self.dataset.column_names)
+        )
+
+    @property
+    def schema_name(self) -> str:
+        """Detected schema name: ``alpaca`` / ``sharegpt`` / ``openai-messages`` /
+        ``pretrain-text`` (the raw ``text``-column fallback)."""
+        return self._schema_name
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, idx: int) -> List[Dict[str, str]]:
+        return self._converter(self.dataset[idx])
+
+
+class VarlenDataset(SFTDataset):
+    """Variable-length single-sample SFT dataset for the packed-sequence path.
+
+    Each ``__getitem__`` returns **one tokenized conversation** in unpacked
+    form: ``tokens``/``labels``/``loss_mask``/``position_ids`` whose length
+    equals the sample's actual token count (padded to ``pad_granularity``,
+    NOT to ``sequence_length``), plus ``original_seq_len``/``padded_seq_len``
+    tensors that the upstream packing scheduler consumes directly via
+    :func:`get_batch_and_global_seqlens`.
+
+    This is the schema described in :class:`BasePackingScheduler.get_required_sample_keys`.
+    It deliberately skips the multi-conversation pre-packing that
+    :class:`SFTDataset.__getitem__` does, letting the upstream scheduler
+    pack variable-length samples across the DP×CP grid with no per-sample
+    padding waste.
+
+    Truncation: samples longer than ``config.sequence_length`` are truncated
+    on the right; an EOD token is appended if the truncation removed it.
+    """
+
+    def __init__(
+        self,
+        dataset: LowLevelDataset,
+        dataset_path: Optional[str],
+        indices: np.ndarray,
+        num_samples: Optional[int],
+        index_split: Split,
+        config: GPTDatasetConfig,
+    ) -> None:
+        super().__init__(dataset, dataset_path, indices, num_samples, index_split, config)
+
+    @staticmethod
+    def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int:
+        return len(low_level_dataset)
+
+    @staticmethod
+    def build_low_level_dataset(
+        dataset_path: str, config: GPTDatasetConfig
+    ) -> LowLevelDataset:
+        return VarlenLowLevelDataset(dataset_path)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        tokenizer = self.config.tokenizer
+        max_len = self.config.sequence_length
+        # HuggingFaceTokenizer returns None for ``pad`` when the underlying
+        # tokenizer has no explicit pad token (common for raw pretraining
+        # tokenizers like Qwen3). Fall back to eod for padding — irrelevant
+        # for loss because loss_mask zeros pad positions out.
+        eod = tokenizer.eod
+        pad = tokenizer.pad if tokenizer.pad is not None else eod
+        assert eod is not None, (
+            "VarlenDataset requires the tokenizer to expose an EOD/EOS token id."
+        )
+
+        # 1. Pull a single item from the low-level dataset. For SFT schemas
+        #    (alpaca / sharegpt / openai-messages) this is a messages list;
+        #    for the pretrain-text schema it is a raw string.
+        item = self.dataset[int(self.indices[idx % len(self.indices)])]
+
+        assert not self.config.reset_position_ids
+        assert not self.config.create_attention_mask and not self.config.reset_attention_mask
+
+        # 2. Tokenize. SFT schemas go through tokenize_conversation (chat
+        #    template + role-aware target masking); pretrain-text bypasses
+        #    chat templating and uses the plain ``tokenize`` interface,
+        #    treating every token as a target (no prompt masking).
+        if isinstance(item, str):
+            ids = list(tokenizer.tokenize(item))
+            tokens_list = ids
+            targets_list = list(ids)
+        else:
+            tokens, targets = tokenizer.tokenize_conversation(
+                item, return_target=True, add_generation_prompt=False
+            )
+            tokens_list = tokens.tolist()
+            targets_list = targets.tolist()
+
+        # 2b. Guard against an empty tokenization (e.g. a blank ``pretrain-text``
+        #     row where ``tokenizer.tokenize("")`` returns no ids). Represent it
+        #     as a single end-of-document token so the next-token shift still
+        #     yields a valid 1-token sample instead of raising on
+        #     ``tokens_list[-1]`` below or producing a zero-length sequence.
+        if len(tokens_list) == 0:
+            tokens_list = [eod, eod]
+            targets_list = [eod, eod]
+
+        # 3. Right-truncate to ``sequence_length + 1`` (we drop the last token
+        #    after the input/label shift below). Keep an EOD at the end so a
+        #    truncated assistant turn still has a valid stop token.
+        if len(tokens_list) > max_len + 1:
+            tokens_list = tokens_list[: max_len + 1]
+            targets_list = targets_list[: max_len + 1]
+            if tokens_list[-1] != eod:
+                tokens_list[-1] = eod
+                targets_list[-1] = eod
+
+        # 4. Ensure EOD is the last token (unconditional for short samples).
+        if tokens_list[-1] != eod:
+            tokens_list.append(eod)
+            targets_list.append(eod)
+
+        valid_len = len(tokens_list) - 1
+
+        # 5a. SBHD validation mode: right-pad to sequence_length + 1, drop
+        #     packing metadata, return shape [sequence_length]. Useful as a
+        #     numerical reference for THD path verification (no scheduler).
+        if self.config.varlen_sbhd_validation:
+            pad_len = max_len + 1 - len(tokens_list)
+            if pad_len > 0:
+                tokens_list.extend([pad] * pad_len)
+                targets_list.extend([pad] * pad_len)
+            assert len(tokens_list) == max_len + 1
+            input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64)
+            labels = torch.tensor(targets_list[1:], dtype=torch.int64)
+            loss_mask = torch.ones(max_len, dtype=torch.float32)
+            loss_mask[valid_len:] = 0.0  # mask the right-padded tail by position
+            loss_mask[labels == IGNORE_INDEX] = 0.0
+            return {
+                'tokens': input_ids,
+                'labels': labels,
+                'loss_mask': loss_mask,
+                'position_ids': torch.arange(max_len, dtype=torch.int64),
+            }
+
+        original_seq_len = len(tokens_list) - 1  # length after the shift below
+
+        # 5b. THD path: pad to pad_granularity (dp_size * cp_size * 2 * sp),
+        #     the minimum alignment required by CP slicing. We deliberately
+        #     do NOT pad to sequence_length — the upstream packing scheduler
+        #     will combine variable-length samples up to
+        #     max_seqlen_per_dp_cp_rank.
+        pad_granularity = self._calculate_padding_divisor()
+        mod = original_seq_len % pad_granularity
+        if mod != 0:
+            pad_len = pad_granularity - mod
+            tokens_list.extend([pad] * pad_len)
+            targets_list.extend([pad] * pad_len)
+        padded_seq_len = len(tokens_list) - 1
+
+        # 6. Apply the next-token shift.
+        input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64)
+        labels = torch.tensor(targets_list[1:], dtype=torch.int64)
+        position_ids = torch.arange(padded_seq_len, dtype=torch.int64)
+        loss_mask = torch.ones(padded_seq_len, dtype=torch.float32)
+        loss_mask[valid_len:] = 0.0  # mask the right-padded tail by position
+        loss_mask[labels == IGNORE_INDEX] = 0.0
+
+        return {
+            'tokens': input_ids,
+            'labels': labels,
+            'loss_mask': loss_mask,
+            'position_ids': position_ids,
+            # The packing scheduler consumes these directly; cu_seqlens /
+            # max_seqlen are produced downstream in _pack_sequences.
+            'original_seq_len': torch.tensor([original_seq_len], dtype=torch.int32),
+            'padded_seq_len': torch.tensor([padded_seq_len], dtype=torch.int32),
+        }
+
+
+class MockVarlenDataset(MockSFTDataset):
+    """Mock variable-length dataset for benchmarking the varlen path.
+
+    Uses :class:`MockSFTLowLevelDataset` for sequence-length sampling (lognormal
+    distribution / per-line CSV / IndexedDataset verification mode — same JSON
+    schema as ``--sft-mock-dataset-config-json``, just consumed via
+    ``--varlen-mock-dataset-config-json``).
+
+    Output shape mirrors :class:`VarlenDataset.__getitem__` (not the inherited
+    :meth:`MockSFTDataset.__getitem__`) so the mock and real-data paths
+    exercise exactly the same downstream pipeline:
+
+      * THD mode: emits **one unpacked sample** padded to ``pad_granularity``
+        with ``original_seq_len`` / ``padded_seq_len`` tensors. The upstream
+        scheduler packs across the DP×CP grid.
+
+    ``--varlen-sbhd-validation`` is intentionally not implemented for mock
+    data; it is guarded against in argument validation.
+    """
+
+    @staticmethod
+    def build_low_level_dataset(
+        dataset_path: str, config: GPTDatasetConfig
+    ) -> LowLevelDataset:
+        if config.varlen_mock_dataset_config_json is None:
+            mock_config = {
+                "mode": "distribution",
+                "type": "lognormal",
+                "min_seq_len": config.sequence_length // 2,
+                "max_seq_len": config.sequence_length,
+                "mean_seq_len": config.sequence_length // 4 * 3,
+                "lognormal_sigma": 1.1,
+            }
+        else:
+            mock_config = load_json_arg(config.varlen_mock_dataset_config_json)
+        return MockSFTLowLevelDataset(**mock_config)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        tokenizer = self.config.tokenizer
+        max_len = self.config.sequence_length
+        eod = tokenizer.eod
+        pad = tokenizer.pad if tokenizer.pad is not None else eod
+
+        # MockSFTLowLevelDataset returns ``length - 1`` token ids; append EOD
+        # to make the conversation end on a stop token, mirroring the real
+        # VarlenDataset path.
+        raw = self.dataset[int(self.indices[idx % len(self.indices)])]
+        tokens_list = raw.tolist()
+        tokens_list.append(eod)
+        # Mock data uses ``tokens == targets`` (no role masking).
+        targets_list = list(tokens_list)
+
+        # MockVarlenDataset only implements the THD (packed) path; SBHD
+        # validation is a real-data numerical-reference mode (guarded against
+        # --mock-data in validate_args).
+        # THD mode: unpacked single sample, pad to pad_granularity only.
+        if len(tokens_list) > max_len + 1:
+            tokens_list = tokens_list[: max_len - 1] + [eod]
+            targets_list = targets_list[: max_len - 1] + [eod]
+        original_seq_len = len(tokens_list) - 1
+
+        pad_granularity = self._calculate_padding_divisor()
+        mod = original_seq_len % pad_granularity
+        if mod != 0:
+            pad_len = pad_granularity - mod
+            tokens_list.extend([pad] * pad_len)
+            targets_list.extend([pad] * pad_len)
+        padded_seq_len = len(tokens_list) - 1
+
+        input_ids = torch.tensor(tokens_list[:-1], dtype=torch.int64)
+        labels = torch.tensor(targets_list[1:], dtype=torch.int64)
+        loss_mask = torch.ones(padded_seq_len, dtype=torch.float32)
+        loss_mask[original_seq_len:] = 0.0  # mask the right-padded tail by position
+        return {
+            'tokens': input_ids,
+            'labels': labels,
+            'loss_mask': loss_mask,
+            'position_ids': torch.arange(padded_seq_len, dtype=torch.int64),
+            'original_seq_len': torch.tensor([original_seq_len], dtype=torch.int32),
+            'padded_seq_len': torch.tensor([padded_seq_len], dtype=torch.int32),
+        }
diff --git a/megatron/training/training.py b/megatron/training/training.py
index d506086d07a..deacb0b6354 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -150,7 +150,7 @@ def set_startup_timestamps(program_start=None, main_entry=None):
 except ImportError:
     HAVE_FSDP2 = False
 
-from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper
+from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper, wrap_data_iterator
 from megatron.core.distributed import finalize_model_grads
 from megatron.core.enums import ModelType
 from megatron.core.inference.symmetric_memory import SymmetricMemoryManager
@@ -278,6 +278,7 @@ def print_datetime(string, override_timestamp=None):
 # never call ``update_*`` so the flag stays ``False`` and no collective fires.
 _seqlen_stats_in_iteration: Optional[torch.Tensor] = None
 _seqlen_stats_active: bool = False
+_seqlen_stats_are_global: bool = False
 
 
 def update_seqlen_stats_from_cu_seqlens(cu_seqlens):
@@ -297,7 +298,7 @@ def update_seqlen_stats_from_cu_seqlens(cu_seqlens):
     the all-reduce; BSHD callers that never invoke this function leave the
     flag at ``False`` and pay zero collective cost.
     """
-    global _seqlen_stats_in_iteration, _seqlen_stats_active
+    global _seqlen_stats_in_iteration, _seqlen_stats_active, _seqlen_stats_are_global
     if cu_seqlens is None or cu_seqlens.numel() < 2:
         return
     # Pin the accumulator to the current CUDA device when available so the
@@ -316,6 +317,21 @@ def update_seqlen_stats_from_cu_seqlens(cu_seqlens):
     _seqlen_stats_in_iteration[0] += seqlens.sum()
     _seqlen_stats_in_iteration[1] += (seqlens * seqlens).sum()
     _seqlen_stats_active = True
+    _seqlen_stats_are_global = False
+
+
+def set_seqlen_stats_in_iteration(total_real_tokens, seqlen_squared_sum):
+    """Seed per-iteration THD FLOPs stats that were already computed globally."""
+    global _seqlen_stats_in_iteration, _seqlen_stats_active, _seqlen_stats_are_global
+    if total_real_tokens is None or seqlen_squared_sum is None:
+        return
+    if _seqlen_stats_in_iteration is None:
+        device = torch.device(f'cuda:{torch.cuda.current_device()}') if torch.cuda.is_available() else 'cpu'
+        _seqlen_stats_in_iteration = torch.zeros(2, dtype=torch.float64, device=device)
+    _seqlen_stats_in_iteration[0] = float(total_real_tokens)
+    _seqlen_stats_in_iteration[1] = float(seqlen_squared_sum)
+    _seqlen_stats_active = True
+    _seqlen_stats_are_global = True
 
 
 def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float]]:
@@ -343,13 +359,15 @@ def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float
     replicated across TP/CP/PP); the world all-reduce therefore overcounts by a
     factor of ``TP * CP * PP``, which we divide out.
     """
-    global _seqlen_stats_in_iteration, _seqlen_stats_active
+    global _seqlen_stats_in_iteration, _seqlen_stats_active, _seqlen_stats_are_global
     if not _seqlen_stats_active:
         # BSHD path: never allocated the tensor; tell the caller to use the
         # closed-form defaults.
         return None, None
     t = _seqlen_stats_in_iteration
-    if torch.distributed.is_initialized() and mpu.model_parallel_is_initialized():
+    if _seqlen_stats_are_global:
+        dedup = 1
+    elif torch.distributed.is_initialized() and mpu.model_parallel_is_initialized():
         torch.distributed.all_reduce(t)
         tp_size = max(mpu.get_tensor_model_parallel_world_size(), 1)
         cp_size = max(mpu.get_context_parallel_world_size(), 1)
@@ -365,6 +383,7 @@ def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float
     # iterations reuse it without reallocating.
     t.zero_()
     _seqlen_stats_active = False
+    _seqlen_stats_are_global = False
     return total_real_tokens / dedup, seqlen_squared_sum / dedup
 
 
@@ -2248,6 +2267,20 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
                     if isinstance(optim_instance, DistributedOptimizer):
                         optim_instance._copy_main_params_to_param_buffer()
 
+        if getattr(config, "sequence_packing_scheduler", None) is not None:
+            (
+                data_iterator,
+                scheduled_num_microbatches,
+                total_real_tokens_in_batch,
+                seqlen_squared_sum_in_batch,
+            ) = wrap_data_iterator(data_iterator, config, get_num_microbatches())
+            set_seqlen_stats_in_iteration(
+                total_real_tokens_in_batch,
+                seqlen_squared_sum_in_batch,
+            )
+        else:
+            scheduled_num_microbatches = get_num_microbatches()
+
         # Forward pass.
         if save_activations_in_this_iteration:
             enable_activation_logging(model, args.save)
@@ -2259,7 +2292,7 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
             forward_step_func=forward_step_func,
             data_iterator=data_iterator,
             model=model,
-            num_microbatches=get_num_microbatches(),
+            num_microbatches=scheduled_num_microbatches,
             seq_length=args.seq_length,
             micro_batch_size=args.micro_batch_size,
             decoder_seq_length=args.decoder_seq_length,
@@ -3528,6 +3561,9 @@ def trace_handler(p):
 
         # Completely skip iteration if needed.
         if (iteration + 1) in args.iterations_to_skip:
+            assert (
+                getattr(config, "sequence_packing_scheduler", None) is None
+            ), "Sequence packing scheduler is not supported in skip iteration mode"
             # Dummy train_step to fast forward train_data_iterator.
             dummy_train_step(train_data_iterator)
             if iteration == start_iteration:
@@ -3960,11 +3996,21 @@ def evaluate(
             # Don't care about timing during evaluation
             config.timers = None
             ft_integration.on_eval_step_start()
+            if getattr(config, "sequence_packing_scheduler", None) is not None:
+                try:
+                    (packed_data_iterator, scheduled_eval_num_microbatches, _, _) = (
+                        wrap_data_iterator(data_iterator, config, eval_num_microbatches)
+                    )
+                except StopIteration:
+                    break
+            else:
+                packed_data_iterator = data_iterator
+                scheduled_eval_num_microbatches = eval_num_microbatches
             loss_dicts = forward_backward_func(
                 forward_step_func=forward_step_func,
-                data_iterator=data_iterator,
+                data_iterator=packed_data_iterator,
                 model=model,
-                num_microbatches=eval_num_microbatches,
+                num_microbatches=scheduled_eval_num_microbatches,
                 seq_length=args.seq_length,
                 micro_batch_size=eval_micro_batch_size,
                 decoder_seq_length=args.decoder_seq_length,
diff --git a/pretrain_gpt.py b/pretrain_gpt.py
index 1884b728e05..07b0df163db 100644
--- a/pretrain_gpt.py
+++ b/pretrain_gpt.py
@@ -26,6 +26,7 @@
 from gpt_builders import gpt_builder
 from megatron.core import mpu
 from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
+from megatron.core.datasets.data_schedule import get_batch_on_this_rank_for_sequence_packing
 from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset
 from megatron.core.enums import ModelType
 from megatron.core.models.gpt import GPTModel
@@ -57,7 +58,8 @@
 from megatron.training.argument_utils import gpt_config_from_args, pretrain_cfg_container_from_args
 from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args
 from megatron.training.datasets.fim_dataset import GPTFIMDataset, GPTFIMDatasetConfig
-from megatron.training.datasets.sft_dataset import SFTDataset
+from megatron.training.datasets.sft_dataset import MockSFTDataset, SFTDataset
+from megatron.training.datasets.varlen_dataset import MockVarlenDataset, VarlenDataset
 from megatron.training.training import update_seqlen_stats_from_cu_seqlens
 from megatron.training.utils import get_blend_and_blend_per_split, is_first_or_last_pipeline_stage
 from model_provider import model_provider
@@ -94,6 +96,19 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None):
     args = get_args()
     config = core_transformer_config_from_args(args)
 
+    if args.sequence_packing_scheduler is not None:
+        return get_batch_on_this_rank_for_sequence_packing(
+            data_iterator,
+            vpp_size=config.virtual_pipeline_model_parallel_size,
+            mtp_on_this_rank=mtp_on_this_rank_func(
+                layout=config.pipeline_model_parallel_layout,
+                mtp_num_layers=config.mtp_num_layers,
+                ignore_virtual=False,
+                vp_stage=vp_stage,
+            ),
+            vp_stage=vp_stage,
+        )
+
     cp_size = args.context_parallel_size
     tp_rank = mpu.get_tensor_model_parallel_rank()
     is_sft = args.sft
@@ -280,43 +295,60 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa
     timers('batch-generator', log_level=2).start()
     with stimer(bdata=True):
         vp_stage = get_attr_wrapped_model(model, "vp_stage")
-        (
-            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, vp_stage)
-
-    packed_seq_params = None
-    if cu_seqlens is not None:
-        # cu_seqlens / cu_seqlens_padded carry the dataloader's batch dim (1, n).
-        # PackedSeqParams (and TE attention) expect 1-D, so squeeze before use.
-        cu_seqlens = cu_seqlens[0]
-        if cu_seqlens_padded is not None:
-            cu_seqlens_padded = cu_seqlens_padded[0]
-        # Use real (unpadded) cu_seqlens to feed the FLOPs accounting: varlen
-        # attention only computes work for real tokens within each chunk.
-        update_seqlen_stats_from_cu_seqlens(cu_seqlens)
-        cu_seqlens_for_params = (
-            cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens
-        )  # TODO(asolergi-nv): Currently there is a bug forcing cu_seqlens to be cu_seqlens_padded
-        packed_seq_params = PackedSeqParams(
-            qkv_format="thd",
-            cu_seqlens_q=cu_seqlens_for_params,
-            cu_seqlens_kv=cu_seqlens_for_params,
-            cu_seqlens_q_padded=cu_seqlens_padded,
-            cu_seqlens_kv_padded=cu_seqlens_padded,
-            max_seqlen_q=int(max_seqlen.item()),
-            max_seqlen_kv=int(max_seqlen.item()),
-            local_cp_size=int(local_cp_size.item()) if local_cp_size is not None else None,
-            cp_group=hybrid_cp_group,
-        )
+        batch = get_batch(data_iterator, vp_stage)
+
+        if len(batch) == 7:
+            (
+                tokens,
+                labels,
+                loss_mask,
+                attention_mask,
+                position_ids,
+                packed_seq_params,
+                padding_mask,
+            ) = batch
+        elif len(batch) == 6:
+            tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = batch
+            padding_mask = None
+        else:
+            (
+                attention_mask,
+                cu_seqlens,
+                cu_seqlens_padded,
+                hybrid_cp_group,
+                labels,
+                local_cp_size,
+                loss_mask,
+                max_seqlen,
+                position_ids,
+                tokens,
+            ) = batch
+
+            padding_mask = None
+            packed_seq_params = None
+            if cu_seqlens is not None:
+                # cu_seqlens / cu_seqlens_padded carry the dataloader's batch dim (1, n).
+                # PackedSeqParams (and TE attention) expect 1-D, so squeeze before use.
+                cu_seqlens = cu_seqlens[0]
+                if cu_seqlens_padded is not None:
+                    cu_seqlens_padded = cu_seqlens_padded[0]
+                # Use real (unpadded) cu_seqlens to feed the FLOPs accounting: varlen
+                # attention only computes work for real tokens within each chunk.
+                update_seqlen_stats_from_cu_seqlens(cu_seqlens)
+                cu_seqlens_for_params = (
+                    cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens
+                )  # TODO(asolergi-nv): Currently there is a bug forcing cu_seqlens to be cu_seqlens_padded
+                packed_seq_params = PackedSeqParams(
+                    qkv_format="thd",
+                    cu_seqlens_q=cu_seqlens_for_params,
+                    cu_seqlens_kv=cu_seqlens_for_params,
+                    cu_seqlens_q_padded=cu_seqlens_padded,
+                    cu_seqlens_kv_padded=cu_seqlens_padded,
+                    max_seqlen_q=int(max_seqlen.item()),
+                    max_seqlen_kv=int(max_seqlen.item()),
+                    local_cp_size=int(local_cp_size.item()) if local_cp_size is not None else None,
+                    cp_group=hybrid_cp_group,
+                )
 
     timers('batch-generator').stop()
 
@@ -326,7 +358,13 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa
                 args.overlap_moe_expert_parallel_comm
             ), "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan"
             schedule_plan = model.build_schedule_plan(
-                tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask
+                tokens,
+                position_ids,
+                attention_mask,
+                labels=labels,
+                loss_mask=loss_mask,
+                packed_seq_params=packed_seq_params,
+                padding_mask=padding_mask,
             )
             return schedule_plan, partial(loss_func, loss_mask, model=model)
         else:
@@ -337,6 +375,7 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa
                 labels=labels,
                 loss_mask=loss_mask,
                 packed_seq_params=packed_seq_params,
+                padding_mask=padding_mask,
             )
 
     # [ModelOpt]: model is needed to access ModelOpt distillation losses
@@ -399,6 +438,10 @@ 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,
+        "sft_mock_dataset_config_json": args.sft_mock_dataset_config_json,
+        "sequence_packing_scheduler": args.sequence_packing_scheduler,
+        "varlen_mock_dataset_config_json": args.varlen_mock_dataset_config_json,
+        "varlen_sbhd_validation": args.varlen_sbhd_validation,
     }
 
     # add FIM args to the config
@@ -437,8 +480,22 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None
 
     is_packed_sequence = False
     if args.sft:
-        dataset_type = SFTDataset
+        if args.mock_data:
+            dataset_type = MockSFTDataset
+        else:
+            dataset_type = SFTDataset
         is_packed_sequence = True  # SFT always uses packed sequence
+    elif args.use_varlen_dataset:
+        # Variable-length packed (THD) dataset, independent of --sft.
+        # Reuses SFTDataset's THD packing internally but is gated
+        # by its own top-level flag.
+        if args.mock_data:
+            dataset_type = MockVarlenDataset
+        else:
+            dataset_type = VarlenDataset
+        # SBHD validation mode runs the non-packed pipeline; THD mode
+        # is the packed-sequence path.
+        is_packed_sequence = not args.varlen_sbhd_validation
     else:
         if args.mock_data:
             dataset_type = MockGPTDataset
diff --git a/tests/unit_tests/data/test_varlen_dataset.py b/tests/unit_tests/data/test_varlen_dataset.py
new file mode 100644
index 00000000000..4dfc3db7061
--- /dev/null
+++ b/tests/unit_tests/data/test_varlen_dataset.py
@@ -0,0 +1,796 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Unit tests for :mod:`megatron.training.datasets.varlen_dataset`.
+
+These tests cover the schema-detection and message-normalization helpers and
+the :class:`VarlenLowLevelDataset` loader. The end-to-end SFTDataset packing
+behavior is exercised by the existing SFT test suite; here we focus on the
+varlen-specific contracts (auto-detect schema, normalize to messages,
+ValueError on unsupported shapes).
+"""
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+# Import via the public module path so this test gets discovered through the
+# regular pytest entry point. The functions under test are pure Python and do
+# not require torch.distributed.
+from megatron.training.datasets.sft_dataset import IGNORE_INDEX
+from megatron.training.datasets.varlen_dataset import (
+    MockVarlenDataset,
+    VarlenDataset,
+    VarlenLowLevelDataset,
+    _alpaca_to_messages,
+    _looks_like_hf_id,
+    _messages_passthrough,
+    _raw_text_loader,
+    _select_converter,
+    _sharegpt_to_messages,
+)
+
+# ----------------------------------------------------------------------------
+# _looks_like_hf_id heuristic
+# ----------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "path,expected",
+    [
+        ("Yukang/LongAlpaca-12k", True),
+        ("HuggingFaceH4/no_robots", True),
+        ("databricks/databricks-dolly-15k", True),
+        ("/tmp/foo.jsonl", False),
+        ("./local.jsonl", False),
+        ("../up.jsonl", False),
+        ("singlename", False),
+        ("", False),
+        (None, False),
+    ],
+)
+def test_looks_like_hf_id(path, expected):
+    assert _looks_like_hf_id(path) is expected
+
+
+# ----------------------------------------------------------------------------
+# Schema converters
+# ----------------------------------------------------------------------------
+
+
+def test_alpaca_canonical_with_input():
+    out = _alpaca_to_messages(
+        {"instruction": "Summarize.", "input": "Long passage", "output": "It says X."}
+    )
+    assert [m["role"] for m in out] == ["system", "user", "assistant"]
+    assert out[1]["content"] == "Summarize.\n\nLong passage"
+    assert out[2]["content"] == "It says X."
+
+
+def test_alpaca_without_input():
+    out = _alpaca_to_messages({"instruction": "Hi.", "output": "Hello."})
+    assert out[0] == {"role": "system", "content": ""}
+    assert out[1]["content"] == "Hi."
+    assert out[2]["content"] == "Hello."
+
+
+@pytest.mark.parametrize(
+    "instr_key,out_key",
+    [
+        ("prompt", "response"),
+        ("query", "answer"),
+        ("question", "completion"),
+        ("instruction", "answer"),
+    ],
+)
+def test_alpaca_field_synonyms(instr_key, out_key):
+    out = _alpaca_to_messages({instr_key: "Q?", out_key: "A."})
+    assert out[1]["content"] == "Q?"
+    assert out[2]["content"] == "A."
+
+
+def test_dolly_instruction_context_response():
+    """Dolly-15k: instruction + context + response, all via synonyms."""
+    out = _alpaca_to_messages(
+        {
+            "instruction": "Who wrote 1984?",
+            "context": "1984 was written in 1948.",
+            "response": "George Orwell.",
+        }
+    )
+    assert out[1]["content"] == "Who wrote 1984?\n\n1984 was written in 1948."
+    assert out[2]["content"] == "George Orwell."
+
+
+def test_sharegpt_human_gpt():
+    out = _sharegpt_to_messages(
+        {"conversations": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "hello"}]}
+    )
+    assert [m["role"] for m in out] == ["system", "user", "assistant"]
+    assert out[1]["content"] == "hi"
+
+
+def test_sharegpt_preserves_existing_system_turn():
+    out = _sharegpt_to_messages(
+        {
+            "conversations": [
+                {"from": "system", "value": "be terse"},
+                {"from": "human", "value": "hi"},
+                {"from": "gpt", "value": "hello"},
+            ]
+        }
+    )
+    assert [m["role"] for m in out] == ["system", "user", "assistant"]
+    assert out[0]["content"] == "be terse"
+
+
+@pytest.mark.parametrize(
+    "speaker,expected_role",
+    [
+        ("human", "user"),
+        ("user", "user"),
+        ("gpt", "assistant"),
+        ("assistant", "assistant"),
+        ("model", "assistant"),
+        ("chatgpt", "assistant"),
+        ("tool", "tool"),
+        ("function", "tool"),
+        ("alien", "user"),  # unknown speakers fall back to user
+    ],
+)
+def test_sharegpt_role_map(speaker, expected_role):
+    out = _sharegpt_to_messages({"conversations": [{"from": speaker, "value": "x"}]})
+    # First entry is the prepended system turn; second is the actual content.
+    assert out[1]["role"] == expected_role
+
+
+def test_messages_passthrough_prepends_system_when_missing():
+    out = _messages_passthrough(
+        {"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]}
+    )
+    assert [m["role"] for m in out] == ["system", "user", "assistant"]
+
+
+def test_messages_passthrough_keeps_existing_system():
+    out = _messages_passthrough(
+        {
+            "messages": [
+                {"role": "system", "content": "be terse"},
+                {"role": "user", "content": "hi"},
+                {"role": "assistant", "content": "ok"},
+            ]
+        }
+    )
+    assert [m["role"] for m in out] == ["system", "user", "assistant"]
+    assert out[0]["content"] == "be terse"
+
+
+def test_messages_passthrough_strips_extra_keys():
+    """OpenAI-style messages may carry ``name`` / ``tool_calls`` etc.;
+    chat-template input only wants ``role`` and ``content``."""
+    out = _messages_passthrough(
+        {
+            "messages": [
+                {"role": "user", "content": "hi", "name": "alice"},
+                {"role": "assistant", "content": "hi alice", "tool_calls": [{"function": "foo"}]},
+            ]
+        }
+    )
+    for m in out:
+        assert set(m.keys()) == {"role", "content"}
+
+
+# ----------------------------------------------------------------------------
+# Shape validation: reject multi-modal / non-string content
+# ----------------------------------------------------------------------------
+
+
+def test_messages_rejects_list_content():
+    with pytest.raises(ValueError, match="must be a string"):
+        _messages_passthrough(
+            {"messages": [{"role": "user", "content": [{"type": "image", "url": "x.png"}]}]}
+        )
+
+
+def test_alpaca_rejects_non_string_field():
+    with pytest.raises(ValueError, match="must be a string"):
+        _alpaca_to_messages({"instruction": ["a", "b"], "output": "x"})
+
+
+def test_sharegpt_rejects_list_value():
+    with pytest.raises(ValueError, match="must be a string"):
+        _sharegpt_to_messages({"conversations": [{"from": "human", "value": [1, 2, 3]}]})
+
+
+# ----------------------------------------------------------------------------
+# Pretrain-text schema
+# ----------------------------------------------------------------------------
+
+
+def test_raw_text_loader_returns_string():
+    """``text``-column samples are returned as plain strings (not messages)."""
+    out = _raw_text_loader({"text": "Once upon a time...", "id": "doc-1"})
+    assert isinstance(out, str)
+    assert out == "Once upon a time..."
+
+
+def test_raw_text_loader_handles_empty():
+    assert _raw_text_loader({"text": None}) == ""
+    assert _raw_text_loader({}) == ""
+
+
+def test_raw_text_rejects_non_string():
+    with pytest.raises(ValueError, match="must be a string"):
+        _raw_text_loader({"text": [1, 2, 3]})
+
+
+# ----------------------------------------------------------------------------
+# Schema selector priority
+# ----------------------------------------------------------------------------
+
+
+def test_select_converter_alpaca():
+    fn, name = _select_converter(["instruction", "output", "file"])
+    assert name == "alpaca"
+    assert fn is _alpaca_to_messages
+
+
+def test_select_converter_alpaca_via_synonyms():
+    fn, name = _select_converter(["prompt", "response"])
+    assert name == "alpaca"
+
+
+def test_select_converter_dolly_columns():
+    fn, name = _select_converter(["instruction", "context", "response", "category"])
+    assert name == "alpaca"
+
+
+def test_select_converter_sharegpt():
+    fn, name = _select_converter(["conversations", "id"])
+    assert name == "sharegpt"
+    assert fn is _sharegpt_to_messages
+
+
+def test_select_converter_messages():
+    fn, name = _select_converter(["messages"])
+    assert name == "openai-messages"
+    assert fn is _messages_passthrough
+
+
+def test_select_converter_priority_messages_over_alpaca():
+    # When both ``messages`` and alpaca-style columns are present, the more
+    # explicit ``messages`` schema wins.
+    fn, name = _select_converter(["messages", "instruction", "output"])
+    assert name == "openai-messages"
+
+
+def test_select_converter_unrecognized_columns():
+    with pytest.raises(ValueError, match="cannot infer schema"):
+        _select_converter(["foo", "bar"])
+
+
+def test_select_converter_alpaca_missing_output():
+    """Having an instruction column but no output column is not a match."""
+    with pytest.raises(ValueError, match="cannot infer schema"):
+        _select_converter(["instruction", "category"])
+
+
+def test_select_converter_pretrain_text():
+    fn, name = _select_converter(["text", "id"])
+    assert name == "pretrain-text"
+    assert fn is _raw_text_loader
+
+
+def test_select_converter_pretrain_text_with_metadata():
+    """Real corpora (e.g. Dolma) have ``text`` + ``url`` + ``metadata``."""
+    fn, name = _select_converter(["text", "url", "metadata", "id"])
+    assert name == "pretrain-text"
+
+
+def test_select_converter_alpaca_beats_pretrain_text():
+    """When both ``instruction``/``output`` and ``text`` are present (rare),
+    the alpaca schema is more specific and should win."""
+    fn, name = _select_converter(["text", "instruction", "output"])
+    assert name == "alpaca"
+
+
+def test_select_converter_messages_beats_pretrain_text():
+    fn, name = _select_converter(["text", "messages"])
+    assert name == "openai-messages"
+
+
+# ----------------------------------------------------------------------------
+# VarlenLowLevelDataset on local jsonl (no HF Hub network needed)
+# ----------------------------------------------------------------------------
+
+
+def _write_jsonl(tmp_path: Path, rows):
+    p = tmp_path / "data.jsonl"
+    with p.open("w") as f:
+        for row in rows:
+            f.write(json.dumps(row) + "\n")
+    return str(p)
+
+
+def test_low_level_loads_jsonl_alpaca(tmp_path):
+    pytest.importorskip("datasets")
+    pytest.importorskip("pandas")
+    path = _write_jsonl(
+        tmp_path,
+        [
+            {"instruction": "i1", "output": "o1"},
+            {"instruction": "i2", "output": "o2", "file": "extra"},
+        ],
+    )
+    ll = VarlenLowLevelDataset(path)
+    assert len(ll) == 2
+    assert ll.schema_name == "alpaca"
+    sample = ll[0]
+    assert [m["role"] for m in sample] == ["system", "user", "assistant"]
+    assert sample[1]["content"] == "i1"
+    assert sample[2]["content"] == "o1"
+
+
+def test_low_level_loads_jsonl_sharegpt(tmp_path):
+    pytest.importorskip("datasets")
+    pytest.importorskip("pandas")
+    path = _write_jsonl(
+        tmp_path,
+        [
+            {"conversations": [{"from": "human", "value": "q1"}, {"from": "gpt", "value": "a1"}]},
+            {"conversations": [{"from": "human", "value": "q2"}, {"from": "gpt", "value": "a2"}]},
+        ],
+    )
+    ll = VarlenLowLevelDataset(path)
+    assert len(ll) == 2
+    assert ll.schema_name == "sharegpt"
+    sample = ll[1]
+    # system prepended + 2 turns from the conversation
+    assert [m["role"] for m in sample] == ["system", "user", "assistant"]
+    assert sample[1]["content"] == "q2"
+
+
+def test_low_level_loads_jsonl_messages(tmp_path):
+    pytest.importorskip("datasets")
+    pytest.importorskip("pandas")
+    path = _write_jsonl(
+        tmp_path,
+        [
+            {
+                "messages": [
+                    {"role": "user", "content": "hi"},
+                    {"role": "assistant", "content": "hello"},
+                ]
+            }
+        ],
+    )
+    ll = VarlenLowLevelDataset(path)
+    assert ll.schema_name == "openai-messages"
+    sample = ll[0]
+    assert [m["role"] for m in sample] == ["system", "user", "assistant"]
+
+
+def test_low_level_jsonl_heterogeneous_columns(tmp_path):
+    """Real datasets often mix rows that have / lack an optional field. Our
+    pandas-based loader must accept the union schema without ``CastError``."""
+    pytest.importorskip("datasets")
+    pytest.importorskip("pandas")
+    rows = [{"instruction": "a", "output": "x"}] * 100 + [
+        {"instruction": "b", "output": "y", "file": "extra"}
+    ] * 100
+    path = _write_jsonl(tmp_path, rows)
+    ll = VarlenLowLevelDataset(path)
+    assert len(ll) == 200
+    # Both halves should normalize to the same messages structure.
+    assert [m["role"] for m in ll[0]] == ["system", "user", "assistant"]
+    assert [m["role"] for m in ll[150]] == ["system", "user", "assistant"]
+
+
+def test_low_level_rejects_unknown_schema(tmp_path):
+    pytest.importorskip("datasets")
+    pytest.importorskip("pandas")
+    path = _write_jsonl(tmp_path, [{"foo": "bar"}])
+    with pytest.raises(ValueError, match="cannot infer schema"):
+        VarlenLowLevelDataset(path)
+
+
+def test_low_level_loads_jsonl_pretrain_text(tmp_path):
+    """Pretrain-text corpora (Dolma / OLMo midtraining) typically have
+    ``text`` + extra fields like ``id`` / ``url`` / ``metadata``."""
+    pytest.importorskip("datasets")
+    pytest.importorskip("pandas")
+    path = _write_jsonl(
+        tmp_path,
+        [
+            {"text": "Doc one body...", "id": "1", "url": "https://x/1"},
+            {"text": "Doc two body...", "id": "2", "url": "https://x/2"},
+        ],
+    )
+    ll = VarlenLowLevelDataset(path)
+    assert ll.schema_name == "pretrain-text"
+    assert len(ll) == 2
+    # Each item is a raw string, NOT a messages list.
+    assert ll[0] == "Doc one body..."
+    assert ll[1] == "Doc two body..."
+
+
+# ----------------------------------------------------------------------------
+# VarlenDataset / MockVarlenDataset __getitem__ (fake tokenizer, no GPU)
+#
+# These bypass the heavy SFTDataset.__init__ and inject the minimal attributes
+# __getitem__ reads, so the EOD handling / position-based loss masking /
+# pad-to-divisor / packing-metadata contracts can be unit tested without a
+# real tokenizer or torch.distributed.
+# ----------------------------------------------------------------------------
+
+
+class _FakeTokenizer:
+    """Minimal tokenizer for exercising VarlenDataset.__getitem__.
+
+    ``tokenize`` maps each character to a non-zero id (so plain text never
+    collides with ``eod``/``pad``); ``tokenize("")`` returns ``[]`` to exercise
+    the empty-row guard. ``tokenize_conversation`` masks non-assistant turns
+    with ``IGNORE_INDEX`` in the targets.
+    """
+
+    def __init__(self, eod: int = 0, pad=None):
+        self._eod = eod
+        self._pad = pad
+
+    @property
+    def eod(self):
+        return self._eod
+
+    @property
+    def pad(self):
+        return self._pad
+
+    def tokenize(self, text):
+        return [ord(c) % 100 + 1 for c in text]  # always >= 1, never eod (0)
+
+    def tokenize_conversation(self, messages, return_target=True, add_generation_prompt=False):
+        tokens, targets = [], []
+        for m in messages:
+            ids = self.tokenize(m["content"])
+            tokens.extend(ids)
+            # Only assistant turns contribute to the loss; prompt is masked.
+            targets.extend(ids if m["role"] == "assistant" else [IGNORE_INDEX] * len(ids))
+        return (torch.tensor(tokens, dtype=torch.int64), torch.tensor(targets, dtype=torch.int64))
+
+
+def _make_config(tokenizer, seq_length=64, *, cp=1, dp=1, sp=1, sbhd=False):
+    return SimpleNamespace(
+        tokenizer=tokenizer,
+        sequence_length=seq_length,
+        reset_position_ids=False,
+        create_attention_mask=False,
+        reset_attention_mask=False,
+        varlen_sbhd_validation=sbhd,
+        data_parallel_size=dp,
+        context_parallel_size=cp,
+        sequence_parallel_size=sp,
+    )
+
+
+def _make_varlen(items, config):
+    ds = VarlenDataset.__new__(VarlenDataset)
+    ds.config = config
+    ds.dataset = items
+    ds.indices = np.arange(len(items))
+    return ds
+
+
+def _make_mock_varlen(token_arrays, config):
+    ds = MockVarlenDataset.__new__(MockVarlenDataset)
+    ds.config = config
+    ds.dataset = token_arrays  # each item exposes .tolist()
+    ds.indices = np.arange(len(token_arrays))
+    return ds
+
+
+def test_getitem_thd_pretrain_text_keys_and_shapes():
+    tok = _FakeTokenizer(eod=0, pad=7)
+    ds = _make_varlen(["hello world"], _make_config(tok, seq_length=64))
+    out = ds[0]
+    assert set(out) == {
+        "tokens",
+        "labels",
+        "loss_mask",
+        "position_ids",
+        "original_seq_len",
+        "padded_seq_len",
+    }
+    n = out["tokens"].numel()
+    assert out["labels"].numel() == n
+    assert out["loss_mask"].numel() == n
+    assert out["position_ids"].numel() == n
+    assert int(out["padded_seq_len"].item()) == n
+
+
+def test_getitem_thd_sft_prompt_is_masked():
+    tok = _FakeTokenizer(eod=0, pad=7)
+    messages = [
+        {"role": "system", "content": ""},
+        {"role": "user", "content": "question"},
+        {"role": "assistant", "content": "answer"},
+    ]
+    ds = _make_varlen([messages], _make_config(tok, seq_length=64))
+    out = ds[0]
+    # Prompt (user) tokens are IGNORE_INDEX in labels and must be masked out;
+    # assistant tokens must contribute to the loss.
+    labels = out["labels"]
+    loss_mask = out["loss_mask"]
+    assert torch.all(loss_mask[labels == IGNORE_INDEX] == 0.0)
+    assert loss_mask.sum() > 0  # assistant span still contributes
+
+
+def test_getitem_thd_pad_masked_by_position_keeps_real_eod():
+    """Regression: with pad falling back to eod, the real end-of-document EOD
+    target must stay in the loss (masked by position, not by value)."""
+    tok = _FakeTokenizer(eod=0, pad=None)  # pad falls back to eod
+    # cp=2 -> pad divisor = cp*2 = 4, so a 3-token doc gets a padding tail.
+    ds = _make_varlen(["abc"], _make_config(tok, seq_length=64, cp=2))
+    out = ds[0]
+    loss_mask = out["loss_mask"].tolist()
+    labels = out["labels"].tolist()
+    # tokens=[a,b,c,eod] padded to 4 -> labels=[b,c,eod,eod(pad)]
+    assert len(loss_mask) == 4
+    # index 2 is the real end-of-document EOD target -> kept (would be wrongly
+    # dropped by value-based ``labels == pad`` masking).
+    assert labels[2] == tok.eod and loss_mask[2] == 1.0
+    # index 3 is the appended pad -> masked.
+    assert loss_mask[3] == 0.0
+
+
+def test_getitem_thd_padded_to_divisor():
+    tok = _FakeTokenizer(eod=0, pad=7)
+    ds = _make_varlen(["abcde"], _make_config(tok, seq_length=64, cp=2))  # divisor 4
+    out = ds[0]
+    assert int(out["padded_seq_len"].item()) % 4 == 0
+
+
+def test_getitem_thd_empty_text_does_not_crash():
+    """A blank pretrain-text row tokenizes to [] -> must not crash and must
+    yield a valid (non-zero-length) sample."""
+    tok = _FakeTokenizer(eod=0, pad=7)
+    ds = _make_varlen([""], _make_config(tok, seq_length=64))
+    out = ds[0]
+    assert out["tokens"].numel() >= 1
+    assert out["labels"].numel() == out["tokens"].numel()
+    assert out["loss_mask"].numel() == out["tokens"].numel()
+
+
+def test_getitem_sbhd_pads_to_seq_length_and_masks_tail():
+    tok = _FakeTokenizer(eod=0, pad=None)
+    ds = _make_varlen(["abc"], _make_config(tok, seq_length=8, sbhd=True))
+    out = ds[0]
+    # SBHD emits fixed [seq_length] samples with no packing metadata.
+    assert set(out) == {"tokens", "labels", "loss_mask", "position_ids"}
+    assert out["tokens"].numel() == 8
+    loss_mask = out["loss_mask"].tolist()
+    # tokens=[a,b,c,eod]: valid_len=3 -> first 3 kept (incl. real eod), rest masked.
+    assert loss_mask[0:3] == [1.0, 1.0, 1.0]
+    assert all(v == 0.0 for v in loss_mask[3:])
+
+
+def test_mock_getitem_thd_keys_and_pad_fallback():
+    tok = _FakeTokenizer(eod=0, pad=None)  # exercise the eod fallback (no crash)
+    ds = _make_mock_varlen([np.array([1, 2, 3, 4], dtype=np.int64)], _make_config(tok, cp=2))
+    out = ds[0]
+    assert set(out) == {
+        "tokens",
+        "labels",
+        "loss_mask",
+        "position_ids",
+        "original_seq_len",
+        "padded_seq_len",
+    }
+    n = out["tokens"].numel()
+    assert out["labels"].numel() == n and out["loss_mask"].numel() == n
+    assert int(out["padded_seq_len"].item()) % 4 == 0
+
+
+# ----------------------------------------------------------------------------
+# THD handoff: _unpack_batch contract for VarlenDataset-style samples
+#
+# VarlenDataset already emits one unpacked sub-sample carrying ``padded_seq_len``,
+# so _unpack_batch must short-circuit (no cu_seqlens slicing) and only normalize
+# the collate batch dim. SFTDataset-style pre-packed samples (cu_seqlens, no
+# padded_seq_len) still take the slicing path.
+# ----------------------------------------------------------------------------
+
+
+def test_unpack_batch_short_circuits_for_varlen_samples():
+    from megatron.core.datasets.data_schedule_utils import _unpack_batch
+
+    # Two VarlenDataset-style samples, each already a single sub-sample with a
+    # leading batch dim (as added by the default collate_fn) and padded_seq_len.
+    batch = [
+        {
+            "tokens": torch.arange(4, dtype=torch.int64).view(1, 4),
+            "labels": torch.arange(4, dtype=torch.int64).view(1, 4),
+            "loss_mask": torch.ones(1, 4),
+            "position_ids": torch.arange(4, dtype=torch.int64).view(1, 4),
+            "padded_seq_len": torch.tensor([4], dtype=torch.int32),
+        },
+        {
+            "tokens": torch.arange(8, dtype=torch.int64).view(1, 8),
+            "labels": torch.arange(8, dtype=torch.int64).view(1, 8),
+            "loss_mask": torch.ones(1, 8),
+            "position_ids": torch.arange(8, dtype=torch.int64).view(1, 8),
+            "padded_seq_len": torch.tensor([8], dtype=torch.int32),
+            "original_seq_len": torch.tensor([8], dtype=torch.int32),
+        },
+    ]
+    out = _unpack_batch(batch)
+    # Short-circuit: same number of samples (no slicing into sub-samples).
+    assert len(out) == 2
+    # Leading collate batch dim dropped.
+    assert out[0]["tokens"].shape == (4,)
+    assert out[1]["tokens"].shape == (8,)
+    # Missing original_seq_len synthesized from padded_seq_len.
+    assert "original_seq_len" in out[0]
+    assert int(out[0]["original_seq_len"].item()) == 4
+    # Existing original_seq_len preserved.
+    assert int(out[1]["original_seq_len"].item()) == 8
+
+
+def test_unpack_batch_slices_prepacked_cu_seqlens_samples():
+    from megatron.core.datasets.data_schedule_utils import _unpack_batch
+
+    # SFTDataset-style pre-packed sample: two sub-sequences [0:3) and [3:5),
+    # described by cu_seqlens, NO padded_seq_len -> takes the slicing path.
+    batch = [
+        {
+            "tokens": torch.arange(5, dtype=torch.int64),
+            "labels": torch.arange(5, dtype=torch.int64),
+            "loss_mask": torch.ones(5),
+            "position_ids": torch.arange(5, dtype=torch.int64),
+            "cu_seqlens": torch.tensor([0, 3, 5], dtype=torch.int32),
+        }
+    ]
+    out = _unpack_batch(batch)
+    # One packed sample with two sub-sequences -> two unpacked samples.
+    assert len(out) == 2
+    assert out[0]["tokens"].numel() == 3
+    assert out[1]["tokens"].numel() == 2
+    assert int(out[0]["padded_seq_len"].item()) == 3
+    assert int(out[1]["padded_seq_len"].item()) == 2
+
+
+# ----------------------------------------------------------------------------
+# DataLoader collate selection (distributed; run under torch.distributed.run).
+#
+# Validates the build_pretraining_data_loader contract for the varlen paths:
+#   * --varlen-sbhd-validation emits fixed-length [seq_length] samples that the
+#     DEFAULT collate stacks into a [mbs, seq_length] batch.
+#   * The THD path (--use-varlen-dataset without SBHD) uses the identity collate
+#     (variable-length dicts are returned as a list, not stacked).
+# ----------------------------------------------------------------------------
+
+
+def _build_varlen_for_loader(items, config, num_samples):
+    from megatron.core.datasets.utils import Split
+
+    ds = VarlenDataset.__new__(VarlenDataset)
+    ds.config = config
+    ds.dataset = items
+    ds.indices = np.arange(len(items))
+    ds.num_samples = num_samples
+    ds.index_split = Split.train
+    return ds
+
+
+def _loader_args(*, use_varlen, sbhd, scheduler, mbs, gbs=None):
+    return SimpleNamespace(
+        dataloader_type='single',
+        micro_batch_size=mbs,
+        global_batch_size=mbs if gbs is None else gbs,
+        full_validation=False,
+        num_workers=0,
+        use_varlen_dataset=use_varlen,
+        varlen_sbhd_validation=sbhd,
+        sequence_packing_scheduler=scheduler,
+    )
+
+
+def test_sbhd_validation_dataloader_uses_default_collate():
+    from megatron.core import parallel_state
+    from megatron.training.datasets.data_samplers import build_pretraining_data_loader
+    from megatron.training.global_vars import destroy_global_vars, set_args
+    from tests.unit_tests.test_utilities import Utils
+
+    Utils.initialize_model_parallel(1, 1)
+    try:
+        tok = _FakeTokenizer(eod=0, pad=7)
+        seq_len, mbs = 16, 2
+        # One global batch needs micro_batch_size * data_parallel_size samples;
+        # size the dataset off the runtime DP world size so this passes under
+        # any --nproc-per-node (the CI default is 8 ranks -> dp=8).
+        dp = parallel_state.get_data_parallel_world_size()
+        n = mbs * dp * 4
+        cfg = _make_config(tok, seq_length=seq_len, sbhd=True)
+        ds = _build_varlen_for_loader(["hello world"] * n, cfg, num_samples=n)
+        set_args(_loader_args(use_varlen=True, sbhd=True, scheduler=None, mbs=mbs))
+        loader = build_pretraining_data_loader(ds, consumed_samples=0)
+        batch = next(iter(loader))
+        # Default collate stacks fixed-length SBHD samples into a tensor batch.
+        assert isinstance(batch, dict)
+        assert batch["tokens"].shape == (mbs, seq_len)
+        assert batch["labels"].shape == (mbs, seq_len)
+        assert batch["loss_mask"].shape == (mbs, seq_len)
+    finally:
+        destroy_global_vars()
+        Utils.destroy_model_parallel()
+
+
+def test_thd_dataloader_uses_identity_collate():
+    from megatron.core import parallel_state
+    from megatron.training.datasets.data_samplers import build_pretraining_data_loader
+    from megatron.training.global_vars import destroy_global_vars, set_args
+    from tests.unit_tests.test_utilities import Utils
+
+    Utils.initialize_model_parallel(1, 1)
+    try:
+        tok = _FakeTokenizer(eod=0, pad=7)
+        mbs = 2
+        dp = parallel_state.get_data_parallel_world_size()
+        n = mbs * dp * 4
+        cfg = _make_config(tok, seq_length=64, sbhd=False)
+        # Variable-length samples so identity collate is required.
+        variable = ["a", "abcdef", "xy", "qwerty"]
+        items = [variable[i % len(variable)] for i in range(n)]
+        ds = _build_varlen_for_loader(items, cfg, num_samples=n)
+        set_args(_loader_args(use_varlen=True, sbhd=False, scheduler="dp_balanced", mbs=mbs))
+        loader = build_pretraining_data_loader(ds, consumed_samples=0)
+        batch = next(iter(loader))
+        # Identity collate returns the raw list of per-sample dicts (unstacked).
+        assert isinstance(batch, list)
+        assert len(batch) == mbs
+        assert "padded_seq_len" in batch[0]
+    finally:
+        destroy_global_vars()
+        Utils.destroy_model_parallel()
+
+
+def test_packing_scheduler_dataloader_yields_microbatches():
+    from megatron.core import parallel_state
+    from megatron.training.datasets.data_samplers import build_pretraining_data_loader
+    from megatron.training.global_vars import destroy_global_vars, set_args
+    from tests.unit_tests.test_utilities import Utils
+
+    Utils.initialize_model_parallel(1, 1)
+    try:
+        tok = _FakeTokenizer(eod=0, pad=7)
+        mbs = 2
+        num_microbatches = 3
+        dp = parallel_state.get_data_parallel_world_size()
+        gbs = mbs * dp * num_microbatches
+        n = gbs * 2
+        cfg = _make_config(tok, seq_length=64, dp=dp, cp=1)
+        variable = ["a", "abcdef", "xy", "qwerty"]
+        items = [variable[i % len(variable)] for i in range(n)]
+        ds = _build_varlen_for_loader(items, cfg, num_samples=n)
+        set_args(
+            _loader_args(
+                use_varlen=True,
+                sbhd=False,
+                scheduler="dp_balanced",
+                mbs=mbs,
+                gbs=gbs,
+            )
+        )
+        loader = build_pretraining_data_loader(ds, consumed_samples=0)
+        batch = next(iter(loader))
+        # The packing scheduler calls next(data_iterator) num_microbatches times;
+        # each loader step must therefore be one local microbatch, not all
+        # local samples from the global batch.
+        assert isinstance(batch, list)
+        assert len(batch) == mbs
+        assert "padded_seq_len" in batch[0]
+    finally:
+        destroy_global_vars()
+        Utils.destroy_model_parallel()
diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py
index 69e2f95071d..f0e918b542b 100644
--- a/tests/unit_tests/models/test_hybrid_moe_model.py
+++ b/tests/unit_tests/models/test_hybrid_moe_model.py
@@ -318,6 +318,7 @@
     "use_transformer_engine_op_fuser": False,
     "moe_single_grouped_weight": False,
     "moe_single_grouped_bias": False,
+    "sequence_packing_scheduler": None,
 }
 # Fields to ignore entirely (ephemeral, environment-specific, very large).
 SKIP_FIELDS = set()
diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py
new file mode 100644
index 00000000000..b2fa3d69db9
--- /dev/null
+++ b/tests/unit_tests/test_sequence_packing.py
@@ -0,0 +1,517 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import random
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from megatron.core import parallel_state
+from megatron.core.datasets.data_schedule import (
+    _build_thd_padding_mask,
+    _sanitize_thd_padding_values,
+    get_batch_on_this_rank_for_sequence_packing,
+    wrap_data_iterator,
+)
+from megatron.core.rerun_state_machine import RerunDataIterator
+from megatron.training.global_vars import unset_global_variables
+from tests.unit_tests.test_utilities import Utils
+
+
+def test_scheduler_thd_padding_mask_from_cu_seqlens():
+    cu_seqlens = torch.tensor([0, 3, 5], dtype=torch.int32)
+    cu_seqlens_padded = torch.tensor([0, 4, 8], dtype=torch.int32)
+
+    padding_mask = _build_thd_padding_mask(cu_seqlens, cu_seqlens_padded)
+
+    assert torch.equal(
+        padding_mask, torch.tensor([False, False, False, True, False, False, True, True])
+    )
+
+
+def test_scheduler_sanitizes_thd_padding_values():
+    padding_mask = torch.tensor([False, False, True, False, True])
+    batch = {
+        'tokens': torch.tensor([11, 12, -1, 21, -1], dtype=torch.int64),
+        'labels': torch.tensor([12, 13, -1, 22, -1], dtype=torch.int64),
+        'loss_mask': torch.ones(5, dtype=torch.float32),
+        'position_ids': torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64),
+    }
+
+    _sanitize_thd_padding_values(batch, padding_mask)
+
+    assert torch.equal(batch['tokens'], torch.tensor([11, 12, 0, 21, 0]))
+    assert torch.equal(batch['labels'], torch.tensor([12, 13, 0, 22, 0]))
+    assert torch.equal(batch['loss_mask'], torch.tensor([1.0, 1.0, 0.0, 1.0, 0.0]))
+    assert torch.equal(batch['position_ids'], torch.tensor([0, 1, 0, 0, 0]))
+
+
+class MockVariableLengthSequencePackingDataIterator:
+    """
+    Mock data iterator for testing get_batch_on_this_rank_for_sequence_packing.
+
+    Generates variable-length (THD format) packed sequences with deterministic
+    data for verification across parallel ranks.
+    """
+
+    def __init__(
+        self,
+        total_seq_length: int,
+        sequence_lengths: list,
+        local_cp_size: int = None,
+        device: str = "cuda",
+        seed: int = 42,
+    ):
+        """
+        Args:
+            total_seq_length: Total length of packed sequences
+            sequence_lengths: List of individual sequence lengths (variable-length).
+                              If None, generates random variable lengths.
+            device: Device to create tensors on
+            seed: Random seed for reproducibility
+        """
+        self.total_seq_length = total_seq_length
+        self.sequence_lengths = sequence_lengths
+        self.local_cp_size = local_cp_size
+        self.device = device
+        self.seed = seed
+        assert (
+            sum(self.sequence_lengths) == total_seq_length
+        ), f"Sequence lengths sum {sum(self.sequence_lengths)} != total {total_seq_length}"
+
+    def __iter__(self):
+        """Interface for the data iterator."""
+        return self
+
+    def __next__(self):
+        """Generate a mock batch with variable-length THD format."""
+        dev = self.device
+        torch.manual_seed(self.seed)
+        torch.cuda.manual_seed(self.seed)
+
+        tokens = torch.randint(0, 16384, (self.total_seq_length,), dtype=torch.int64, device=dev)
+
+        # Create position_ids that reset for each sequence (THD format)
+        position_ids = []
+        for seq_len in self.sequence_lengths:
+            position_ids.extend(range(seq_len))
+        position_ids = torch.tensor(position_ids, dtype=torch.int64, device=dev)
+
+        # Labels are tokens shifted by 1 for easy verification
+        labels = tokens + 1
+
+        # Loss mask: 1.0 for all positions except padding (none here)
+        loss_mask = torch.ones(self.total_seq_length, dtype=torch.float32, device=dev)
+
+        # Create cu_seqlens for variable-length packed sequences
+        cu_seqlens = [0]
+        for seq_len in self.sequence_lengths:
+            cu_seqlens.append(cu_seqlens[-1] + seq_len)
+        cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=dev)
+        cu_seqlens_padded = cu_seqlens.clone()
+
+        max_seqlen = torch.tensor([max(self.sequence_lengths)], dtype=torch.int32, device=dev)
+
+        batch = {
+            "tokens": tokens,
+            "position_ids": position_ids,
+            "labels": labels,
+            "loss_mask": loss_mask,
+            "cu_seqlens": cu_seqlens,
+            "cu_seqlens_padded": cu_seqlens_padded,
+            "max_seqlen": max_seqlen,
+        }
+
+        if not (
+            parallel_state.is_pipeline_first_stage(ignore_virtual=True)
+            or parallel_state.is_pipeline_last_stage(ignore_virtual=True)
+        ):
+            batch["tokens"] = None
+            batch["position_ids"] = None
+            batch["labels"] = None
+            batch["loss_mask"] = None
+
+        if self.local_cp_size is not None:
+            batch["local_cp_size"] = torch.tensor(
+                [self.local_cp_size], dtype=torch.int32, device=dev
+            )
+
+        return batch
+
+
+def _gather_tensor_from_tp_group(tensor):
+    """Gather tensors from all TP ranks for comparison."""
+    assert tensor is not None, "Tensor should not be None"
+    tp_size = parallel_state.get_tensor_model_parallel_world_size()
+    gathered = [torch.zeros_like(tensor) for _ in range(tp_size)]
+    torch.distributed.all_gather(
+        gathered, tensor, group=parallel_state.get_tensor_model_parallel_group()
+    )
+    return gathered
+
+
+def _gather_tensor_from_all_ranks(tensor):
+    """Gather tensors from all PP ranks for comparison."""
+    assert tensor is not None, "Tensor should not be None"
+    if type(tensor) is int:
+        tensor = torch.tensor(tensor, dtype=torch.int32, device=torch.cuda.current_device())
+    gathered = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())]
+    torch.distributed.all_gather(gathered, tensor)
+    return gathered
+
+
+@pytest.mark.parametrize(
+    ("tp", "pp", "cp"),
+    [
+        (1, 1, 1),  # Basic case: no parallelism
+        (2, 1, 1),  # Tensor parallel only
+        (1, 2, 1),  # Pipeline parallel only
+        (2, 2, 1),  # TP + PP
+        (1, 1, 2),  # CP only
+        (2, 1, 2),  # TP + CP
+        (1, 2, 2),  # PP + CP
+        (1, 4, 1),  # Has middle pp stage
+    ],
+)
+def test_get_batch_on_this_rank_for_sequence_packing(tp, pp, cp):
+    """
+    Test get_batch_on_this_rank_for_sequence_packing function with variable-length THD format.
+
+    This test verifies:
+    1. TP ranks: All ranks within a TP group receive identical data after broadcast
+    2. PP ranks: Middle PP ranks have the same packed_seq_params as first/last stages
+    3. CP ranks: Data is correctly partitioned with proper shape and values
+    4. Variable-length (THD) format: Different sequence lengths are handled correctly
+    """
+    args = SimpleNamespace()
+    args.tensor_model_parallel_size = tp
+    args.pipeline_model_parallel_size = pp
+    args.context_parallel_size = cp
+    args.virtual_pipeline_model_parallel_size = None
+    args.data_parallel_size = 8 // (tp * pp * cp)
+    args.seq_length = 8192
+
+    # Skip invalid configurations
+    if args.data_parallel_size < 1:
+        raise ValueError(f"Invalid config: tp={tp}, pp={pp}, cp={cp} exceeds world size 8")
+
+    # Initialize model parallel
+    Utils.initialize_model_parallel(tp, pp, None, context_parallel_size=cp)
+
+    try:
+        # Create mock data iterator with variable-length sequences
+        # Only TP rank 0 needs the iterator; other TP ranks pass None
+        tp_rank = parallel_state.get_tensor_model_parallel_rank()
+        if tp_rank == 0:
+            # Use deterministic seed based on DP rank so same data within TP/PP/CP group
+            dp_rank = parallel_state.get_data_parallel_rank()
+            sequence_lengths = [1024, 2048, 512, 1536, 3072]
+            assert (
+                sum(sequence_lengths) == args.seq_length
+            ), f"Sequence lengths sum {sum(sequence_lengths)} != total {args.seq_length}"
+            data_iterator = iter(
+                MockVariableLengthSequencePackingDataIterator(
+                    total_seq_length=args.seq_length,
+                    sequence_lengths=sequence_lengths,  # Variable lengths, sum=8192
+                    seed=42 + dp_rank,  # Same seed within PP/CP group
+                )
+            )
+        else:
+            # Non-TP-rank-0 ranks don't need the iterator
+            data_iterator = None
+
+        # Call the function under test
+        result = get_batch_on_this_rank_for_sequence_packing(
+            data_iterator=data_iterator, mtp_on_this_rank=False, vp_stage=None
+        )
+
+        # Unpack the result. Scheduler THD always returns padding_mask.
+        tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask = (
+            result
+        )
+
+        # Get parallel state info
+        tp_rank = parallel_state.get_tensor_model_parallel_rank()
+        pp_rank = parallel_state.get_pipeline_model_parallel_rank()
+        cp_rank = parallel_state.get_context_parallel_rank()
+        is_first_stage = parallel_state.is_pipeline_first_stage(ignore_virtual=True)
+        is_last_stage = parallel_state.is_pipeline_last_stage(ignore_virtual=True)
+        is_first_or_last = is_first_stage or is_last_stage
+
+        assert padding_mask is not None
+        assert padding_mask.dtype == torch.bool
+        assert padding_mask.dim() == 2
+        assert padding_mask.size(0) == 1
+        assert not padding_mask.any(), "Mock data has no per-sequence padding."
+
+        # =====================================================================
+        # TEST 1: Verify data based on pipeline stage
+        # =====================================================================
+        if is_first_stage:
+            assert tokens is not None, "First stage should have tokens"
+            assert position_ids is not None, "First stage should have position_ids"
+            assert tokens.dim() == 2, "Tokens should be 2D (batch, seq)"
+            assert position_ids.dim() == 2, "Position IDs should be 2D (batch, seq)"
+            assert tokens.size(0) == 1, "batch should be 1 in THD format"
+            assert position_ids.size(0) == 1, "batch should be 1 in THD format"
+        else:
+            assert tokens is None, "Non-first stage should not have tokens"
+            assert position_ids is None, "Non-first stage should not have position_ids"
+
+        if is_last_stage:
+            assert labels is not None, "Last stage should have labels"
+            assert loss_mask is not None, "Last stage should have loss_mask"
+            assert labels.dim() == 2, "Labels should be 2D (batch, seq)"
+            assert loss_mask.dim() == 2, "Loss mask should be 2D (batch, seq)"
+            assert labels.size(0) == 1, "batch should be 1 in THD format"
+            assert loss_mask.size(0) == 1, "batch should be 1 in THD format"
+        else:
+            assert labels is None, "Non-last stage should not have labels"
+            assert loss_mask is None, "Non-last stage should not have loss_mask"
+
+        # =====================================================================
+        # TEST 2: Verify all ranks have consistent packed_seq_params
+        # =====================================================================
+        assert packed_seq_params is not None
+        assert packed_seq_params.qkv_format == "thd"
+
+        test_keys = [
+            "cu_seqlens_q",
+            "cu_seqlens_q_padded",
+            "max_seqlen_q",
+            "cu_seqlens_kv",
+            "cu_seqlens_kv_padded",
+            "max_seqlen_kv",
+        ]
+        for key in test_keys:
+            tensor = getattr(packed_seq_params, key)
+            assert tensor is not None
+            gathered_tensor = _gather_tensor_from_all_ranks(tensor)
+            for i in range(1, len(gathered_tensor)):
+                assert torch.equal(
+                    gathered_tensor[0], gathered_tensor[i]
+                ), f"Rank 0 and rank {i} have different {key}"
+
+        # =====================================================================
+        # TEST 3: Verify TP ranks receive identical data after broadcast
+        # =====================================================================
+        if tp > 1:
+            test_tensors = []
+            if is_first_stage:
+                test_tensors.extend([tokens, position_ids])
+            if is_last_stage:
+                test_tensors.extend([labels, loss_mask])
+
+            for tensor in test_tensors:
+                gathered_tensors = _gather_tensor_from_tp_group(tensor)
+                for i in range(1, tp):
+                    assert torch.equal(
+                        gathered_tensors[0], gathered_tensors[i]
+                    ), f"TP rank 0 and rank {i} have different data"
+
+        # =====================================================================
+        # TEST 4: Verify CP partitioning
+        # =====================================================================
+        if cp > 1:
+            # With CP, the sequence should be partitioned
+            expected_seq_len = args.seq_length // cp
+
+            if is_first_stage:
+                actual_seq_len = tokens.shape[1]
+                assert (
+                    actual_seq_len == expected_seq_len
+                ), f"CP partitioned tokens have wrong shape: {actual_seq_len} != {expected_seq_len}"
+
+            # Verify labels only if all CP ranks are at last stage
+            if is_last_stage:
+                actual_seq_len = labels.shape[1]
+                assert (
+                    actual_seq_len == expected_seq_len
+                ), f"CP partitioned labels have wrong shape: {actual_seq_len} != {expected_seq_len}"
+
+    finally:
+        Utils.destroy_model_parallel()
+        unset_global_variables()
+
+
+@pytest.mark.parametrize(
+    ("tp", "pp", "cp", "vpp", "scheduler_type"),
+    [
+        (1, 1, 8, None, "dp_balanced"),
+        (2, 1, 4, None, "dp_balanced"),
+        (2, 4, 1, None, "dp_balanced"),
+        (2, 2, 1, None, "dp_balanced"),
+        (1, 4, 1, 4, "dp_balanced"),
+    ],
+)
+def test_wrap_dataloader(tp, pp, cp, vpp, scheduler_type):
+    '''
+    Test wrap_dataloader function with different scheduler types.
+    '''
+    args = SimpleNamespace()
+    args.tensor_model_parallel_size = tp
+    args.pipeline_model_parallel_size = pp
+    args.context_parallel_size = cp
+    args.virtual_pipeline_model_parallel_size = None
+    args.data_parallel_size = 8 // (tp * pp * cp)
+    args.seq_length = 8192
+    args.max_seqlen_per_dp_cp_rank = 8192
+
+    # Skip invalid configurations
+    if args.data_parallel_size < 1:
+        raise ValueError(f"Invalid config: tp={tp}, pp={pp}, cp={cp} exceeds world size 8")
+
+    def _create_single_sample(seq_len):
+        # hard code the padding size to 16
+        pad_size = 16
+        seq_len_padded = ((seq_len + pad_size - 1) // pad_size) * pad_size
+        device = torch.device("cuda", torch.cuda.current_device())
+        tokens = torch.randint(0, 128, (seq_len_padded,), dtype=torch.int64, device=device)
+        labels = tokens + 1
+        position_ids = torch.arange(seq_len_padded, dtype=torch.int64, device=device)
+        loss_mask = torch.ones(seq_len_padded, dtype=torch.float32, device=device)
+        loss_mask[0:seq_len] = 1
+        loss_mask[seq_len:] = 0
+        cu_seqlens = torch.tensor([0, seq_len_padded], dtype=torch.int32, device=device)
+
+        return {
+            'tokens': tokens,
+            'labels': labels,
+            'loss_mask': loss_mask,
+            'position_ids': position_ids,
+            'cu_seqlens': cu_seqlens,
+        }
+
+    # Initialize model parallel
+    Utils.initialize_model_parallel(tp, pp, vpp, context_parallel_size=cp)
+
+    global_batch_size = 64
+    micro_batch_size = 1
+    nums = [random.randint(2048, args.seq_length) for _ in range(global_batch_size)]  # 64 sequences
+
+    config = SimpleNamespace()
+    config.max_seqlen_per_dp_cp_rank = args.max_seqlen_per_dp_cp_rank
+    config.microbatch_group_size_per_vp_stage = pp
+    config.virtual_pipeline_model_parallel_size = vpp
+    config.sequence_packing_scheduler = scheduler_type
+
+    dp_rank = parallel_state.get_data_parallel_rank()
+    dp_size = parallel_state.get_data_parallel_world_size()
+
+    pp_rank = parallel_state.get_pipeline_model_parallel_rank()
+    tp_rank = parallel_state.get_tensor_model_parallel_rank()
+
+    is_pp_first = pp_rank == 0
+    is_pp_last = pp_rank == pp - 1
+    is_pp_first_or_last = is_pp_first or is_pp_last
+    is_tp_first = tp_rank == 0
+
+    num_micro_batches_old = global_batch_size // micro_batch_size // dp_size
+
+    if is_tp_first and (is_pp_first or is_pp_last):
+        samples = [
+            _create_single_sample(num)
+            for num in nums[dp_rank * num_micro_batches_old : (dp_rank + 1) * num_micro_batches_old]
+        ]
+        data_iterator = RerunDataIterator(iter(samples))
+    else:
+        data_iterator = None
+
+    if is_tp_first:
+        if vpp is not None and vpp > 1:
+            if is_pp_first:
+                data_iterator = [data_iterator] + [None for _ in range(vpp - 1)]
+            elif is_pp_last:
+                data_iterator = [None for _ in range(vpp - 1)] + [data_iterator]
+            else:
+                data_iterator = [None for _ in range(vpp)]
+    try:
+        # Call the function under test
+        (
+            new_data_iterator,
+            num_micro_batches,
+            num_total_tokens_this_global_batch,
+            sequence_square_sum_this_global_batch,
+        ) = wrap_data_iterator(data_iterator, config, num_micro_batches_old)
+
+        # check the result
+        assert type(num_micro_batches) is int
+        assert (
+            type(num_total_tokens_this_global_batch) is float
+            or type(num_total_tokens_this_global_batch) is np.float32
+        )
+        assert (
+            type(sequence_square_sum_this_global_batch) is float
+            or type(sequence_square_sum_this_global_batch) is np.float32
+        )
+
+        def _check_batch(batch_all, batch_keys):
+            for batch in batch_all:
+                assert set(batch_keys) <= set(
+                    batch.keys()
+                ), f"batch keys: {set(batch.keys())} missing {set(batch_keys) - set(batch.keys())}"
+                for key in batch_keys:
+                    assert batch[key] is not None
+
+        if is_tp_first:
+            # CHECK KEYS
+            batch_keys = ["cu_seqlens", "max_seqlen", "cu_seqlens_padded"]
+            if vpp is not None and vpp > 1:
+                # check metadata for all stages (save batches to avoid re-consuming iterators)
+                all_stage_batches = []
+                for temp_data_iterator in new_data_iterator:
+                    stage_batch = [next(temp_data_iterator) for _ in range(num_micro_batches)]
+                    all_stage_batches.append(stage_batch)
+                    _check_batch(stage_batch, batch_keys)
+
+                # check for first or last stage on first or last pp rank
+                if is_pp_first_or_last:
+                    batch_all = all_stage_batches[0] if is_pp_first else all_stage_batches[-1]
+                    batch_keys += ["tokens", "position_ids", "labels", "loss_mask"]
+                    _check_batch(batch_all, batch_keys)
+            else:
+                # non-VPP: single iterator
+                batch_all = [next(new_data_iterator) for _ in range(num_micro_batches)]
+                if is_pp_first_or_last:
+                    batch_keys += ["tokens", "position_ids", "labels", "loss_mask"]
+                _check_batch(batch_all, batch_keys)
+
+            # CHECK TOKEN SUM ON FIRST OR LAST PP RANK
+            # Note: data_iterator is consumed by wrap_data_iterator, new_data_iterator is consumed above.
+            # Use `samples` for before-wrap, reuse `batch_all` from the check above for after-wrap.
+            if is_pp_first_or_last:
+                # Compute token sum before wrap
+                token_sum_before = torch.tensor(0, dtype=torch.int64, device='cuda')
+                for sample in samples:
+                    token_sum_before += sample['tokens'].long().sum()
+
+                # Compute token sum after wrap (batch_all already collected above with tokens)
+                token_sum_after = torch.tensor(0, dtype=torch.int64, device='cuda')
+                for batch in batch_all:
+                    token_sum_after += batch['tokens'].long().sum()
+
+                # Reduce sum across dp_cp group and verify equality
+                dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=False)
+                torch.distributed.all_reduce(
+                    token_sum_before, op=torch.distributed.ReduceOp.SUM, group=dp_cp_group
+                )
+                torch.distributed.all_reduce(
+                    token_sum_after, op=torch.distributed.ReduceOp.SUM, group=dp_cp_group
+                )
+
+                assert (
+                    token_sum_before == token_sum_after
+                ), f"Token sum mismatch: before={token_sum_before.item()}, after={token_sum_after.item()}"
+
+        else:
+            if vpp is not None and vpp > 1:
+                assert type(new_data_iterator) is list and len(new_data_iterator) == vpp
+                for data_iterator in new_data_iterator:
+                    assert data_iterator is None
+            else:
+                assert new_data_iterator is None
+
+    finally:
+        Utils.destroy_model_parallel()
+        unset_global_variables()