diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py
index 8afdd6b0f39..b6a6a65dc3c 100644
--- a/megatron/core/datasets/data_schedule.py
+++ b/megatron/core/datasets/data_schedule.py
@@ -9,7 +9,6 @@
     align_sample_id_groups,
     broadcast_scalars,
     broadcast_tensor,
-    broadcast_to_pp_group,
     build_packed_microbatches,
     create_data_iterator,
     dcp_get_total_workload,
@@ -23,6 +22,7 @@
 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
+from megatron.core.transformer.multi_token_prediction import mtp_on_this_rank
 
 
 class BasePackingScheduler:
@@ -178,14 +178,17 @@ def run(
         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)
+            3. Strip data fields not needed by this PP stage
+            4. Schedule samples into groups
+            5. Reroute samples to DCP ranks
+            6. Build packed microbatches
+            7. Calculate FLOPs info
+            8. Broadcast scalars to TP group (for non-TP-0 ranks)
             9. Handle VPP if enabled
 
+        Note: There is no PP-group broadcast. In packed-sequence mode
+        is_dataset_built_on_rank returns True for every PP stage on TP rank 0
+
         Args:
             data_iterator: The data iterator.
             num_microbatches: The number of microbatches to fetch.
@@ -204,20 +207,20 @@ def run(
         """
 
         total_dcp_gpus = dp_cp_group.size()
+        is_first_pp = pp_group.rank() == 0
+        is_last_pp = pp_group.rank() == pp_group.size() - 1
+
+        mtp_on_this_pp = mtp_on_this_rank(config, ignore_virtual=True)
+        vpp_size = config.virtual_pipeline_model_parallel_size or 1
 
         # Handle VPP: extract the correct data_iterator for this PP stage.
         # When VPP is enabled, data_iterator is a list with one entry per VPP stage.
         # We only need one data_iterator to run the schedule (all VPP stages on the
         # same PP rank share the same underlying dataset), so pick the first non-None.
-        # Record which VPP stages had data so create_data_iterator knows which ones
-        # need full samples vs metadata only.
-        vpp_has_data = None
-        if (
-            config.virtual_pipeline_model_parallel_size is not None
-            and config.virtual_pipeline_model_parallel_size > 1
-        ):
-            assert len(data_iterator) == config.virtual_pipeline_model_parallel_size
-            vpp_has_data = [di is not None for di in data_iterator]
+        # Determine which VPP stages need full data based on pipeline position and MTP.
+        vpp_needs_data = None
+        if vpp_size > 1:
+            assert len(data_iterator) == vpp_size
             extracted = None
             for di in data_iterator:
                 if di is not None:
@@ -225,6 +228,19 @@ def run(
                     break
             data_iterator = extracted
 
+            # Only first VPP on first PP and last VPP on last PP need full data.
+            # MTP VPP stages also need full data (both tokens and labels).
+            # Middle VPP stages only need metadata (cu_seqlens, max_seqlen, etc.).
+            vpp_needs_data = [False] * vpp_size
+            if is_first_pp:
+                vpp_needs_data[0] = True
+            if is_last_pp:
+                vpp_needs_data[-1] = True
+            if mtp_on_this_pp:
+                for vp_i in range(vpp_size):
+                    if mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_i):
+                        vpp_needs_data[vp_i] = True
+
         # data_iterator is not None on TP rank 0 for PP stages that need data
         # (first stage, last stage, or any stage with MTP).
         if data_iterator is not None:
@@ -241,7 +257,24 @@ def run(
                     key in batch[0]
                 ), f"Batch missing required key {key}, provided keys: {batch[0].keys()}"
 
-            # Step 3: Schedule samples into groups
+            # Step 3: Strip data fields not needed by this PP stage to avoid
+            # unnecessary all-to-all communication. First PP needs tokens/position_ids,
+            # last PP needs labels/loss_mask. MTP stages need all four.
+            # NOTE: this assumes _unpack_batch produces only the six keys below
+            # (tokens, position_ids, labels, loss_mask, original_seq_len,
+            # padded_seq_len). Any custom dataset metadata key outside this set
+            # would be silently dropped here; extend keys_to_keep if needed.
+            keys_to_keep = {'original_seq_len', 'padded_seq_len'}
+            if is_first_pp or mtp_on_this_pp:
+                keys_to_keep.update(['tokens', 'position_ids'])
+            if is_last_pp or mtp_on_this_pp:
+                keys_to_keep.update(['labels', 'loss_mask'])
+            for sample in batch:
+                for key in list(sample.keys()):
+                    if key not in keys_to_keep:
+                        del sample[key]
+
+            # Step 4: Schedule samples into groups
             sample_id_groups = self.get_groups_and_subsamples(global_id_seqlens)
 
             # Validate scheduling result
@@ -254,7 +287,7 @@ def run(
                 f"global_id_seqlens length: {len(global_id_seqlens)}"
             )
 
-            # Step 4: Reroute samples to DCP ranks
+            # Step 5: Reroute samples to DCP ranks
             samples_this_rank_with_id = reroute_samples_to_dcp_ranks(
                 batch,
                 global_ids_this_rank,
@@ -270,12 +303,12 @@ def run(
             dcp_rank = dp_cp_group.rank()
             num_micro_batches = len(sample_id_groups)
 
-            # Step 5: Build packed microbatches
+            # Step 6: Build packed microbatches
             new_samples = build_packed_microbatches(
                 samples_this_rank_with_id, sample_id_groups, dcp_rank, dev, self.is_dynamic_cp
             )
 
-            # Step 6: Calculate FLOPs info
+            # Step 7: 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)
@@ -288,24 +321,7 @@ def run(
                 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,
-                is_dynamic_cp=self.is_dynamic_cp,
-            )
-
-        # Step 8: Broadcast to TP group (for non-TP-0 ranks)
+        # 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(
                 [
@@ -319,9 +335,9 @@ def run(
         )
         num_micro_batches = int(num_micro_batches)
 
-        # Step 9: create data_iterator and handle VPP if enabled
+        # Step 8: Broadcast to TP group and create data_iterator
         new_data_iterator = create_data_iterator(
-            new_samples, tp_group, config, vpp_has_data, self.is_dynamic_cp
+            new_samples, tp_group, config, vpp_needs_data, self.is_dynamic_cp
         )
 
         return (
@@ -551,7 +567,13 @@ def get_batch_on_this_rank_for_sequence_packing(
 
     if is_first_or_last_stage or mtp_on_this_rank:
         if is_tp_rank_0:
-            total_tokens = torch.tensor(batch['tokens'].size(0), dtype=torch.int32, device=dev)
+            # Use whichever data field is available (first stage has tokens, last has labels).
+            # Avoid `tokens or labels`: PyTorch tensors raise on truthiness when they have
+            # more than one element ("Boolean value of Tensor ... is ambiguous").
+            _data_field = batch.get('tokens')
+            if _data_field is None:
+                _data_field = batch.get('labels')
+            total_tokens = torch.tensor(_data_field.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)
diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py
index c59b1742c0a..0210a4bc0e3 100644
--- a/megatron/core/datasets/data_schedule_utils.py
+++ b/megatron/core/datasets/data_schedule_utils.py
@@ -5,7 +5,6 @@
 from math import ceil, log2
 from typing import Callable, Dict, List, Optional, Tuple
 
-import numpy as np
 import torch
 
 from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices
@@ -26,7 +25,13 @@ def get_cp_slice_for_thd(batch, cp_group):
     if cp_size <= 1:
         return
     cp_rank = cp_group.rank()
-    total_tokens = batch['tokens'].size(0)
+    # Use whichever data field is available to determine total_tokens
+    for _key in ['tokens', 'labels', 'loss_mask', 'position_ids']:
+        if _key in batch and batch[_key] is not None:
+            total_tokens = batch[_key].size(0)
+            break
+    else:
+        raise ValueError("Cannot determine total_tokens: no data field found in batch")
     # 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.
@@ -45,9 +50,11 @@ def _unpack_batch(batch: List[Dict[str, torch.Tensor]]) -> List[Dict[str, torch.
     the entire packed sample.
     """
     batch_unpacked = []
-    dev = batch[0]["tokens"].device
+    dev = batch[0]["cu_seqlens"].device
     original_seq_lens = []
     padded_seq_lens = []
+    # Determine which data fields exist in the batch
+    data_keys = [k for k in ["tokens", "labels", "loss_mask", "position_ids"] if k in batch[0]]
     for sample in batch:
         for key in sample.keys():
             if len(sample[key].shape) == 2:
@@ -63,7 +70,7 @@ def _unpack_batch(batch: List[Dict[str, torch.Tensor]]) -> List[Dict[str, torch.
             end_idx = sample["cu_seqlens"][sub_sample + 1]
             if end_idx - start_idx == 0:
                 continue
-            for key in ["tokens", "labels", "loss_mask", "position_ids"]:
+            for key in data_keys:
                 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.
@@ -151,16 +158,10 @@ def _pack_sequences(
     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
+    for key in ['tokens', 'labels', 'loss_mask', 'position_ids']:
+        if key in samples[0]:
+            new_sample[key] = _pack_tensors([sample[key] for sample in samples])
 
     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)
@@ -191,103 +192,6 @@ def broadcast_tensor(item, src_rank, group) -> 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,
-    is_dynamic_cp: bool = False,
-):
-    """
-    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:
-            tensor_list = [
-                torch.tensor(
-                    [
-                        num_micro_batches,
-                        seqlen_sum_this_global_batch,
-                        seqlen_squared_sum_this_global_batch,
-                    ],
-                    dtype=torch.float32,
-                ).cuda()
-            ]
-            for sample in new_samples:
-                tensor_list.append(sample["max_seqlen"].unsqueeze(0))
-
-            if is_dynamic_cp:
-                for sample in new_samples:
-                    tensor_list.append(sample["local_cp_size"].unsqueeze(0))
-
-            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).cuda()
-            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).cuda()
-            broadcast_tensor(info_length_tensor, pp_src_rank, pp_group)
-            info_to_broadcast = torch.empty(info_length_tensor.item(), dtype=torch.float32).cuda()
-            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
-                info_numpy = info_to_broadcast.cpu().numpy()
-                num_micro_batches = int(info_numpy[0])
-                seqlen_sum_this_global_batch = info_numpy[1]
-                seqlen_squared_sum_this_global_batch = info_numpy[2]
-                max_seqlens = info_to_broadcast[3 : 3 + num_micro_batches]
-                local_cp_sizes = (
-                    info_to_broadcast[3 + num_micro_batches : 3 + 2 * num_micro_batches]
-                    if is_dynamic_cp
-                    else None
-                )
-                cu_seqlens_list = []
-                cu_seqlens_padded_list = []
-                # cu_seqlens always starts with 0, and the other metadata values
-                # (num_micro_batches, seqlen_sum, seqlen_squared_sum, max_seqlens)
-                # are always positive, so we can use 0 as the delimiter to locate
-                # the start of each cu_seqlens / cu_seqlens_padded tensor.
-                # This avoids an extra broadcast for the lengths of cu_seqlens.
-                indices = np.where(info_numpy == 0)[0]
-                for i in range(num_micro_batches):
-                    cu_seqlens_list.append(info_to_broadcast[indices[i * 2] : indices[i * 2 + 1]])
-                    if i == num_micro_batches - 1:
-                        cu_seqlens_padded_list.append(info_to_broadcast[indices[i * 2 + 1] :])
-                    else:
-                        cu_seqlens_padded_list.append(
-                            info_to_broadcast[indices[i * 2 + 1] : indices[i * 2 + 2]]
-                        )
-
-                new_samples = []
-                for i in range(num_micro_batches):
-                    new_sample = {}
-                    new_sample["max_seqlen"] = max_seqlens[i].to(torch.int32)
-                    new_sample["cu_seqlens"] = cu_seqlens_list[i].to(torch.int32)
-                    new_sample["cu_seqlens_padded"] = cu_seqlens_padded_list[i].to(torch.int32)
-                    if is_dynamic_cp:
-                        new_sample["local_cp_size"] = local_cp_sizes[i].to(torch.int32)
-                    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.
@@ -321,21 +225,21 @@ def broadcast_scalars(values: List, group, dev, dtype=torch.float32) -> List:
 
 
 def create_data_iterator(
-    new_samples, tp_group, config, vpp_has_data=None, is_dynamic_cp: bool = False
+    new_samples, tp_group, config, vpp_needs_data=None, is_dynamic_cp: bool = False
 ):
     """Handle virtual pipeline parallelism.
 
     For VPP, each PP rank needs a list of data iterators (one per VPP stage).
-    VPP stages that originally had a data_iterator (indicated by vpp_has_data)
-    get full samples; others get metadata only (cu_seqlens, cu_seqlens_padded,
+    VPP stages that need full data (first/last pipeline stage, or MTP) get
+    full samples; others get metadata only (cu_seqlens, cu_seqlens_padded,
     max_seqlen).
 
     Args:
         new_samples: The packed samples after scheduling.
         tp_group: Tensor parallel process group.
         config: Model parallel config.
-        vpp_has_data: A list of booleans (one per VPP stage) indicating which
-            VPP stages originally had a data_iterator. None if VPP is disabled.
+        vpp_needs_data: A list of booleans (one per VPP stage) indicating which
+            VPP stages need full samples (data fields). None if VPP is disabled.
     """
     if (
         config.virtual_pipeline_model_parallel_size is not None
@@ -346,14 +250,19 @@ def create_data_iterator(
             metadata_keys = ["max_seqlen", "cu_seqlens", "cu_seqlens_padded"]
             if is_dynamic_cp:
                 metadata_keys.append("local_cp_size")
-            metadata = [
-                {k: sample[k] for k in metadata_keys if k in sample} for sample in new_samples
-            ]
             new_data_iterator = []
             for i in range(vpp_size):
-                if vpp_has_data is not None and vpp_has_data[i]:
-                    new_data_iterator.append(RerunDataIterator(iter(new_samples)))
+                if vpp_needs_data is not None and vpp_needs_data[i]:
+                    # Give each data-carrying VPP stage its own shallow-copied
+                    # sample dicts.
+                    samples_copy = [dict(sample) for sample in new_samples]
+                    new_data_iterator.append(RerunDataIterator(iter(samples_copy)))
                 else:
+                    # Create independent metadata dicts to avoid shared-reference mutation
+                    metadata = [
+                        {k: sample[k] for k in metadata_keys if k in sample}
+                        for sample in new_samples
+                    ]
                     new_data_iterator.append(RerunDataIterator(iter(metadata)))
         else:
             new_data_iterator = [None for _ in range(vpp_size)]
diff --git a/megatron/core/datasets/readme.md b/megatron/core/datasets/readme.md
index 64d0e3dd5bd..af3bccdebf4 100644
--- a/megatron/core/datasets/readme.md
+++ b/megatron/core/datasets/readme.md
@@ -219,7 +219,6 @@ Phase 1        │  (once per global      │   ────────►   
           ├─ get_groups_and_subsamples()              └─ PackedSeqParams(...)
           ├─ reroute_samples_to_dcp_ranks()  [utils]
           ├─ build_packed_microbatches()     [utils]
-          ├─ broadcast_to_pp_group()         [utils]
           ├─ broadcast_scalars()             [utils]
           └─ create_data_iterator()          [utils]
 ```
@@ -249,7 +248,6 @@ Utility functions consumed by the schedulers above:
 | `get_batch_and_global_seqlens()` | Fetch `num_microbatches` batches from the data iterator and all-gather sequence lengths across DP ranks. |
 | `reroute_samples_to_dcp_ranks()` | All-to-all communication to transfer sub-samples to their scheduled DP×CP rank. |
 | `build_packed_microbatches()` | Concatenate sub-samples within each microbatch group and produce `cu_seqlens`. |
-| `broadcast_to_pp_group()` | Broadcast packed samples and metadata from the first/last PP stage to middle stages. |
 | `broadcast_scalars()` | Broadcast scalar values (e.g. `num_microbatches`, FLOPs stats) across a process group. |
 | `broadcast_tensor()` | Broadcast a single tensor within a process group. |
 | `create_data_iterator()` | Wrap packed sample lists into a data iterator; handles VPP stage splitting. |
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 7bce9d96d2c..a7f7151316d 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -17,7 +17,7 @@
     RotaryEmbedding,
 )
 from megatron.core.models.common.language_module.language_module import LanguageModule
-from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group
 from megatron.core.pipeline_parallel.fine_grained_activation_offload import (
     FineGrainedActivationOffloadingInterface as off_interface,
 )
@@ -664,6 +664,7 @@ def _postprocess(
                 self._decoder_hidden_states_cache = hidden_states
             else:
                 # In training/eval, use the utility function for processing MTP loss/scaling.
+                mtp_cp_group = resolve_cp_group(self.pg_collection.cp, packed_seq_params)
                 hidden_states = process_mtp_loss(
                     hidden_states=hidden_states,
                     labels=labels,
@@ -674,7 +675,7 @@ def _postprocess(
                     is_training=self.training,
                     compute_language_model_loss=self.compute_language_model_loss,
                     config=self.config,
-                    cp_group=self.pg_collection.cp,
+                    cp_group=mtp_cp_group,
                     packed_seq_params=packed_seq_params,
                     scale_logits_fn=self._scale_logits if self.config.use_mup else None,
                 )
diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py
index 322f12a4122..b1b4275fee1 100644
--- a/megatron/core/packed_seq_params.py
+++ b/megatron/core/packed_seq_params.py
@@ -64,3 +64,17 @@ def __post_init__(self):
                 .to(torch.int32)
                 .unsqueeze(0)  # Add a batch dimension
             )
+
+
+def resolve_cp_group(
+    static_cp_group: dist.ProcessGroup, packed_seq_params: PackedSeqParams = None
+) -> dist.ProcessGroup:
+    """Return the dynamic CP group from packed_seq_params when available, else the static one.
+
+    Dynamic CP assigns a per-microbatch CP group that may differ from the
+    process-group stored at model construction time.  This helper centralises
+    the resolution logic used by GPTModel, GatedDeltaNet, and MTP layers.
+    """
+    if packed_seq_params is not None and packed_seq_params.cp_group is not None:
+        return packed_seq_params.cp_group
+    return static_cp_group
diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
index f9b923632f5..93ed577bc05 100644
--- a/megatron/core/ssm/gated_delta_net.py
+++ b/megatron/core/ssm/gated_delta_net.py
@@ -19,7 +19,7 @@
 from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.jit import jit_fuser
-from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.ssm.mamba_context_parallel import (
     _all_to_all_cp2hp,
@@ -56,6 +56,12 @@
 
 logger = logging.getLogger(__name__)
 
+# Triton's autotune key for causal_conv1d includes cdiv(total_tokens, 1024).
+# Dynamic CP causes total_tokens to vary per microbatch, triggering repeated
+# autotuning.  Aligning to this boundary collapses most variations into a
+# small number of buckets.
+_CONV_PAD_ALIGNMENT = 4096
+
 
 @dataclass
 class GatedDeltaNetSubmodules:
@@ -138,6 +144,22 @@ def __init__(
         self.qk_dim_local_tp = self.qk_dim // self.tp_size
         self.v_dim_local_tp = self.v_dim // self.tp_size
 
+        # GDN uses head-parallel CP: each CP rank handles a slice of heads.
+        # The static cp_size (== max dynamic cp_size) must evenly divide the
+        # per-TP head counts so that every possible runtime cp_size also divides.
+        num_key_heads_per_tp = self.num_key_heads // self.tp_size
+        num_value_heads_per_tp = self.num_value_heads // self.tp_size
+        assert num_key_heads_per_tp % self.cp_size == 0, (
+            f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) "
+            f"to evenly divide num_key_heads per TP rank ({num_key_heads_per_tp}); "
+            f"all runtime dynamic cp_size values divide the static one and so will also divide."
+        )
+        assert num_value_heads_per_tp % self.cp_size == 0, (
+            f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) "
+            f"to evenly divide num_value_heads per TP rank ({num_value_heads_per_tp}); "
+            f"all runtime dynamic cp_size values divide the static one and so will also divide."
+        )
+
         # Input projection (hidden_states -> q, k, v, gate, beta, alpha)
         # TODO: for now, output gate is forced for GDN.
         # We may remove this restriction in the future.
@@ -289,8 +311,11 @@ def forward(
 
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
+        cp_group = resolve_cp_group(self.pg_collection.cp, packed_seq_params)
+        cp_size = cp_group.size()
+
         seq_len, batch, _ = hidden_states.shape
-        seq_len = seq_len * self.sp_size * self.cp_size
+        seq_len = seq_len * self.sp_size * cp_size
 
         if inference_context is not None:
             assert (
@@ -339,14 +364,14 @@ def forward(
 
         # CP All to All: CP to HP
         if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
-            unpacked_qkvzba = _unpack_sequence(qkvzba, cu_seqlens_q // self.cp_size, dim=0)
+            unpacked_qkvzba = _unpack_sequence(qkvzba, cu_seqlens_q // cp_size, dim=0)
             outputs = []
             for qkvzba_i in unpacked_qkvzba:
                 qkvzba_i = tensor_a2a_cp2hp(
                     qkvzba_i,
                     seq_dim=0,
                     head_dim=-1,
-                    cp_group=self.pg_collection.cp,
+                    cp_group=cp_group,
                     split_sections=[
                         self.qk_dim_local_tp,
                         self.qk_dim_local_tp,
@@ -363,7 +388,7 @@ def forward(
                 qkvzba,
                 seq_dim=0,
                 head_dim=-1,
-                cp_group=self.pg_collection.cp,
+                cp_group=cp_group,
                 split_sections=[
                     self.qk_dim_local_tp,
                     self.qk_dim_local_tp,
@@ -382,10 +407,10 @@ def forward(
         qkv, gate, beta, alpha = torch.split(
             qkvzba,
             [
-                (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // self.cp_size,
-                self.v_dim_local_tp // self.cp_size,
-                self.num_value_heads // self.tp_size // self.cp_size,
-                self.num_value_heads // self.tp_size // self.cp_size,
+                (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size,
+                self.v_dim_local_tp // cp_size,
+                self.num_value_heads // self.tp_size // cp_size,
+                self.num_value_heads // self.tp_size // cp_size,
             ],
             dim=-1,
         )
@@ -402,16 +427,13 @@ def forward(
             self.v_dim_local_tp,
         ]
         conv1d_weight = get_parameter_local_cp(
-            self.conv1d.weight,
-            dim=0,
-            cp_group=self.pg_collection.cp,
-            split_sections=qkv_channels_split_sections,
+            self.conv1d.weight, dim=0, cp_group=cp_group, split_sections=qkv_channels_split_sections
         )
         conv1d_bias = (
             get_parameter_local_cp(
                 self.conv1d.bias,
                 dim=0,
-                cp_group=self.pg_collection.cp,
+                cp_group=cp_group,
                 split_sections=qkv_channels_split_sections,
             )
             if self.conv_bias
@@ -426,36 +448,47 @@ def forward(
                 stride=self.conv1d.stride,
                 padding=self.conv1d.padding,
                 dilation=self.conv1d.dilation,
-                groups=self.conv_dim_local_tp // self.cp_size,
+                groups=self.conv_dim_local_tp // cp_size,
             )
             qkv = self.act_fn(conv_out[..., :seq_len])
             qkv = qkv.transpose(1, 2)  # b, d, s -> b, s, d
         else:
             assert self.activation in ["silu", "swish"]
+            _orig_seq = qkv.shape[1]
+            _pad_n = -_orig_seq % _CONV_PAD_ALIGNMENT
+            _conv_input = qkv
+            _conv_cu_seqlens = cu_seqlens_q
+            if _pad_n > 0:
+                _conv_input = torch.nn.functional.pad(qkv, (0, 0, 0, _pad_n))
+                # cu_seqlens_q is None in non-packed-sequence mode; only the
+                # last-segment offset needs to grow to cover the padding tail.
+                if cu_seqlens_q is not None:
+                    _conv_cu_seqlens = cu_seqlens_q.clone()
+                    _conv_cu_seqlens[-1] += _pad_n
             qkv, _ = causal_conv1d(
-                x=qkv,  # FLA conv1d accepts [b, s, d] format input
+                x=_conv_input,  # FLA conv1d accepts [b, s, d] format input
                 weight=conv1d_weight.squeeze(1),  # d, 1, w -> d, w
                 bias=conv1d_bias,
                 activation=self.activation,
                 initial_state=None,
                 output_final_state=False,
-                cu_seqlens=cu_seqlens_q,
+                cu_seqlens=_conv_cu_seqlens,
             )
+            if _pad_n > 0:
+                qkv = qkv[:, :_orig_seq, :]
         nvtx_range_pop(suffix="conv1d")
 
         # Prepare QKV tensors (split, reshape, L2 norm, repeat_interleave, contiguous)
         nvtx_range_push(suffix="prepare_qkv_for_gated_delta_rule")
         query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule(
-            qkv, gate, beta, alpha, batch, seq_len
+            qkv, gate, beta, alpha, batch, seq_len, cp_size
         )
         nvtx_range_pop(suffix="prepare_qkv_for_gated_delta_rule")
 
         # Calculate g and beta
         nvtx_range_push(suffix="g_and_beta")
-        A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp)
-        dt_bias_local_cp = get_parameter_local_cp(
-            self.dt_bias, dim=0, cp_group=self.pg_collection.cp
-        )
+        A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group)
+        dt_bias_local_cp = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group)
         g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta)
         nvtx_range_pop(suffix="g_and_beta")
 
@@ -488,15 +521,11 @@ def forward(
             unpacked_norm_out = _unpack_sequence(norm_out, cu_seqlens_q, dim=0)
             outputs = []
             for norm_out_i in unpacked_norm_out:
-                norm_out_i = tensor_a2a_hp2cp(
-                    norm_out_i, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp
-                )
+                norm_out_i = tensor_a2a_hp2cp(norm_out_i, seq_dim=0, head_dim=-1, cp_group=cp_group)
                 outputs.append(norm_out_i)
             norm_out = torch.cat(outputs, dim=0)
         else:
-            norm_out = tensor_a2a_hp2cp(
-                norm_out, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp
-            )
+            norm_out = tensor_a2a_hp2cp(norm_out, seq_dim=0, head_dim=-1, cp_group=cp_group)
 
         # Output projection
         nvtx_range_push(suffix="out_proj")
@@ -518,16 +547,14 @@ def _apply_gated_norm(self, x, gate):
         return y
 
     @jit_fuser
-    def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len):
+    def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len, cp_size):
         """
         Prepare query, key, value, gate, beta, alpha tensors for gated delta rule.
         Fuses split, reshape, L2 norm, repeat_interleave, and contiguous operations.
         """
         # Split qkv into query_key and value
         query_key, value = torch.split(
-            qkv,
-            [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size],
-            dim=-1,
+            qkv, [2 * self.qk_dim_local_tp // cp_size, self.v_dim_local_tp // cp_size], dim=-1
         )
 
         # Reshape query_key and value
@@ -539,8 +566,7 @@ def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_l
             query_key = l2norm(query_key.contiguous())
 
         # Split query and key
-        split_size = self.qk_dim_local_tp // self.key_head_dim // self.cp_size
-        query, key = torch.split(query_key, [split_size, split_size], dim=2)
+        query, key = query_key.chunk(2, dim=2)
 
         # Expand query and key if needed (grouped query attention)
         if self.num_value_heads // self.num_key_heads > 1:
diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py
index 707c7d7690f..313cc8056b9 100644
--- a/megatron/core/ssm/mamba_mixer.py
+++ b/megatron/core/ssm/mamba_mixer.py
@@ -23,7 +23,7 @@
     tensor_masked_update,
     tensor_merge,
 )
-from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.ssm.ops.causal_conv1d_triton import causal_conv1d_update
 from megatron.core.ssm.ops.mamba_ssm import selective_state_update
@@ -452,10 +452,10 @@ def forward(
                     out, out_bias = self._decode(hidden_states, conv_state, ssm_state)
                     return out, out_bias
 
-        # Dynamic CP group support
         _orig_cp_group = self.cp.cp_group
-        if packed_seq_params is not None and packed_seq_params.cp_group is not None:
-            self.cp.set_context_parallel_group(packed_seq_params.cp_group)
+        _resolved_cp_group = resolve_cp_group(_orig_cp_group, packed_seq_params)
+        if _resolved_cp_group is not _orig_cp_group:
+            self.cp.set_context_parallel_group(_resolved_cp_group)
 
         zxBCdt, _ = self.in_proj(hidden_states)
 
@@ -474,7 +474,8 @@ def forward(
 
         out, out_bias = self.out_proj(y)
 
-        self.cp.set_context_parallel_group(_orig_cp_group)
+        if _resolved_cp_group is not _orig_cp_group:
+            self.cp.set_context_parallel_group(_orig_cp_group)
         return out, out_bias
 
     def _dynamic_inference(self, hidden_states: torch.Tensor, context: DynamicInferenceContext):
diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py
index 6c6d5b07a75..47f0d93ea4e 100644
--- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py
+++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py
@@ -356,11 +356,6 @@ def get_query_key_value_tensors(
         assert (
             hidden_states.ndim == 3
         ), f"hidden_states should be 3D, [s, b, h], got {hidden_states.ndim}D"
-        if packed_seq_params is not None:
-            assert (
-                packed_seq_params.local_cp_size is None
-            ), "dynamic context parallel is not supported with MLA yet and is planned for future. \
-            Please disable dynamic context parallel."
 
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
@@ -725,6 +720,14 @@ def forward(
             inference_context is None and inference_params is None
         ), "Inference is not supported for AbsorbedMLA"
 
+        # Set the right cp group for dynamic-cp. Mirrors Attention.forward:
+        # downstream RoPE uses self.pg_collection.cp, which must point at this
+        # microbatch's dynamic CP group. Restored before every return.
+        _orig_cp_group = self.pg_collection.cp
+        if packed_seq_params is not None and packed_seq_params.local_cp_size is not None:
+            assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode"
+            self.pg_collection.cp = packed_seq_params.cp_group
+
         # =====================
         # Query, Key, and Value
         # =====================
@@ -806,6 +809,7 @@ def forward(
         # =================
         output, bias = self.linear_proj(core_attn_out)
 
+        self.pg_collection.cp = _orig_cp_group
         return output, bias
 
     def backward_dw(self) -> NoReturn:
diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py
index 0c15ac5ac94..d4ee7783b73 100644
--- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py
+++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py
@@ -245,6 +245,15 @@ def forward(
             inference_context is None and inference_params is None
         ), "Inference is not supported for DSv4HybridAttention."
 
+        # Set the right cp group for dynamic-cp. Mirrors Attention.forward:
+        # both QKV RoPE and the post-attention inverse RoPE use
+        # self.pg_collection.cp, which must point at this microbatch's dynamic
+        # CP group. Restored before every return.
+        _orig_cp_group = self.pg_collection.cp
+        if packed_seq_params is not None and packed_seq_params.local_cp_size is not None:
+            assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode"
+            self.pg_collection.cp = packed_seq_params.cp_group
+
         # =====================
         # Query, Key, and Value
         # =====================
@@ -384,6 +393,7 @@ def forward(
             output, bias = self.linear_proj(core_attn_out)
         output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out])
 
+        self.pg_collection.cp = _orig_cp_group
         return output, bias
 
 
@@ -496,11 +506,6 @@ def get_query_key_value_tensors(
         assert (
             hidden_states.ndim == 3
         ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D"
-        if packed_seq_params is not None:
-            assert (
-                packed_seq_params.local_cp_size is None
-            ), "dynamic_context_parallel is not supported with MLA yet and is planned for future. \
-            Please disable dynamic_context_parallel."
 
         assert (
             inference_context is None and inference_params is None
diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py
index ab50f8a9067..c667ae960ea 100644
--- a/megatron/core/transformer/multi_latent_attention.py
+++ b/megatron/core/transformer/multi_latent_attention.py
@@ -334,6 +334,14 @@ def forward(
         if self.config.cache_mla_latents:
             self.prepare_for_absorption()
 
+        # Set the right cp group for dynamic-cp. Mirrors Attention.forward:
+        # downstream RoPE uses self.pg_collection.cp, which must point at this
+        # microbatch's dynamic CP group. Restored before every return.
+        _orig_cp_group = self.pg_collection.cp
+        if packed_seq_params is not None and packed_seq_params.local_cp_size is not None:
+            assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode"
+            self.pg_collection.cp = packed_seq_params.cp_group
+
         # =====================
         # Query, Key, and Value
         # =====================
@@ -461,6 +469,7 @@ def forward(
             output, bias = self.linear_proj(core_attn_out)
         output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out])
 
+        self.pg_collection.cp = _orig_cp_group
         return output, bias
 
 
@@ -662,11 +671,6 @@ def get_query_key_value_tensors(
         assert (
             hidden_states.ndim == 3
         ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D"
-        if packed_seq_params is not None:
-            assert (
-                packed_seq_params.local_cp_size is None
-            ), "dynamic_context_parallel is not supported with MLA yet and is planned for future. \
-            Please disable dynamic_context_parallel."
 
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 4ceb7050f1d..6a7f1d73374 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -17,7 +17,7 @@
 from megatron.core.extensions.transformer_engine import HAVE_TE
 from megatron.core.fp8_utils import get_fp8_context
 from megatron.core.models.backends import BackendSpecProvider, LocalSpecProvider
-from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group
 from megatron.core.pipeline_parallel.utils import is_vp_last_stage
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel import (
@@ -351,51 +351,66 @@ class MTPLossLoggingHelper:
 
     @staticmethod
     def save_loss_to_tracker(
-        loss: torch.Tensor,
+        loss_sum: torch.Tensor,
+        num_tokens: torch.Tensor,
         layer_number: int,
         num_layers: int,
         reduce_group: Optional[torch.distributed.ProcessGroup] = None,
         avg_group: Optional[torch.distributed.ProcessGroup] = None,
     ):
-        """Save the mtp loss for logging.
+        """Save the mtp loss sum and token count for logging.
+
+        Stores raw sums so that the global per-token loss can be computed
+        correctly after all-reduce, even when token counts differ across
+        ranks (e.g. Dynamic CP) or microbatches.
+
         Args:
-            loss (torch.Tensor): The loss tensor.
+            loss_sum (torch.Tensor): Sum of per-element losses on this rank.
+            num_tokens (torch.Tensor): Number of valid tokens on this rank.
             layer_number (int): Layer index of the loss.
             num_layers (int): The number of total layers.
-            reduce_group (torch.distributed.ProcessGroup): The group for reducing the loss.
-            mean_group (torch.distributed.ProcessGroup): The group for averaging the loss.
+            reduce_group (torch.distributed.ProcessGroup): The group for sum-reducing losses.
+            avg_group (torch.distributed.ProcessGroup): The group for sum-reducing before averaging.
         """
-        # Skip mtp loss logging if layer_number is None.
         if layer_number is None:
             return
 
         tracker = MTPLossLoggingHelper.tracker
-        if "values" not in tracker:
-            tracker["values"] = torch.zeros(num_layers, device=torch.cuda.current_device())
-        tracker["values"][layer_number] += loss.detach()
+        if "loss_sums" not in tracker:
+            tracker["loss_sums"] = torch.zeros(num_layers, device=torch.cuda.current_device())
+            tracker["num_tokens"] = torch.zeros(num_layers, device=torch.cuda.current_device())
+        tracker["loss_sums"][layer_number] += loss_sum.detach()
+        tracker["num_tokens"][layer_number] += num_tokens.detach()
         tracker["reduce_group"] = reduce_group
         tracker["avg_group"] = avg_group
 
     def clean_loss_in_tracker():
         """Clear the mtp losses."""
         tracker = MTPLossLoggingHelper.tracker
-        tracker["values"].zero_()
+        if "loss_sums" in tracker:
+            tracker["loss_sums"].zero_()
+            tracker["num_tokens"].zero_()
         tracker["reduce_group"] = None
         tracker["avg_group"] = None
 
     def reduce_loss_in_tracker():
-        """Collect and reduce the mtp losses across ranks."""
+        """Collect and reduce the mtp losses across ranks.
+
+        Packs loss sums and token counts into a single tensor for one
+        all-reduce, then computes per-token loss.  This produces correct
+        weighted-average results even when ranks hold different numbers
+        of tokens (e.g. Dynamic CP with variable CP sizes).
+        """
         tracker = MTPLossLoggingHelper.tracker
-        if "values" not in tracker:
+        if "loss_sums" not in tracker:
             return
-        values = tracker["values"]
-        # Reduce mtp losses across ranks.
-        if tracker.get('reduce_group') is not None:
-            torch.distributed.all_reduce(values, group=tracker.get('reduce_group'))
-        if tracker.get('avg_group') is not None:
-            torch.distributed.all_reduce(
-                values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG
-            )
+        packed = torch.cat([tracker["loss_sums"], tracker["num_tokens"]])
+        for group_key in ('reduce_group', 'avg_group'):
+            group = tracker.get(group_key)
+            if group is not None:
+                torch.distributed.all_reduce(packed, group=group)
+        loss_sums, num_tokens = packed.chunk(2)
+        tracker["values"] = loss_sums / num_tokens.clamp(min=1)
 
     def track_mtp_metrics(loss_scale, iteration, writer, wandb_writer=None, total_loss_dict=None):
         """Track the Multi-Token Prediction (MTP) metrics for logging."""
@@ -736,12 +751,9 @@ def output_layer_for_mtp(input_: Tensor, **kwargs):
             mtp_loss = compute_language_model_loss(mtp_labels, mtp_logits)
         mtp_loss = loss_mask * mtp_loss
         if is_training:
-            # Safe divide without sync: mask numerator when num_tokens==0, divide by clamp(min=1)
-            mtp_loss_for_log = (
-                torch.sum(mtp_loss) * (num_tokens > 0).to(mtp_loss.dtype)
-            ) / num_tokens.clamp(min=1)
             MTPLossLoggingHelper.save_loss_to_tracker(
-                mtp_loss_for_log,
+                torch.sum(mtp_loss),
+                num_tokens,
                 mtp_layer_number,
                 config.mtp_num_layers,
                 avg_group=parallel_state.get_data_parallel_group(with_context_parallel=True),
@@ -991,20 +1003,13 @@ def _get_embeddings(
                 sequence length, b is the batch size, and h is the hidden size.
             packed_seq_params (PackedSeqParams): Parameters for packed sequence processing.
         """
-        # Calc logits for the current Multi-Token Prediction (MTP) layers.
+        cp_group = resolve_cp_group(self.cp_group, packed_seq_params)
+
         input_ids, _ = roll_tensor(
-            input_ids,
-            shifts=-1,
-            dims=-1,
-            cp_group=self.cp_group,
-            packed_seq_params=packed_seq_params,
+            input_ids, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
         )
         position_ids, _ = roll_tensor(
-            position_ids,
-            shifts=-1,
-            dims=-1,
-            cp_group=self.cp_group,
-            packed_seq_params=packed_seq_params,
+            position_ids, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
         )
         # embedding
         decoder_input = embedding(input_ids=input_ids, position_ids=position_ids)
@@ -1417,8 +1422,7 @@ def forward(
         """
         assert context is None, "multi token prediction + cross attention is not yet supported."
         _orig_cp_group = self.cp_group
-        if packed_seq_params is not None and packed_seq_params.cp_group is not None:
-            self.cp_group = packed_seq_params.cp_group
+        self.cp_group = resolve_cp_group(self.cp_group, packed_seq_params)
         input_ids, position_ids, decoder_input, hidden_states = self._get_embeddings(
             input_ids=input_ids,
             position_ids=position_ids,
diff --git a/megatron/training/training.py b/megatron/training/training.py
index cf1cdda2685..3eebbc0360e 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -2546,7 +2546,10 @@ def training_log(
 
     # Log MTP metrics.
     if args.mtp_num_layers is not None:
-        mtp_loss_scale = 1 / get_num_microbatches()
+        # MTP tracker stores raw loss sums and token counts, so after reduction
+        # tracker["values"] already equals the per-token loss (loss_sum / num_tokens)
+        # aggregated across all ranks and microbatches. No further scaling needed.
+        mtp_loss_scale = 1.0
         MTPLossLoggingHelper.track_mtp_metrics(
             mtp_loss_scale, iteration, writer, wandb_writer, total_loss_dict
         )
diff --git a/megatron/training/utils.py b/megatron/training/utils.py
index 0e05c11e52a..e2417261e5f 100644
--- a/megatron/training/utils.py
+++ b/megatron/training/utils.py
@@ -561,6 +561,11 @@ def _broadcast(item):
         }
 
         def _broadcast_cu_seqlens(cu_seqlens):
+            if getattr(args, 'cuda_graph_impl', 'none') == 'full_iteration':
+                assert cu_seqlens is None, (
+                    "cu_seqlens is not supported with cuda_graph_impl=full_iteration"
+                )
+                return
             dev = torch.cuda.current_device()
             n = 0 if cu_seqlens is None else int(cu_seqlens.numel())
             n_tensor = torch.empty(1, dtype=torch.int64, device=dev).fill_(n)
@@ -642,6 +647,8 @@ def _broadcast_cu_seqlens(cu_seqlens):
         )
 
         def _broadcast_cu_seqlens():
+            if getattr(args, 'cuda_graph_impl', 'none') == 'full_iteration':
+                return None
             dev = torch.cuda.current_device()
 
             n = torch.empty((), dtype=torch.int64, device=dev)
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200.json
index ae215b3314a..8e0b6544b90 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200.json
@@ -433,7 +433,7 @@
         "step_interval": 1,
         "values": {
             "1": 10.91979,
-            "2": 10.92264,
+            "2": 10.92265,
             "3": 10.92219,
             "4": 10.91957,
             "5": 10.90846,
@@ -442,7 +442,7 @@
             "8": 10.92446,
             "9": 10.91998,
             "10": 10.92341,
-            "11": 10.91588,
+            "11": 10.91589,
             "12": 10.90792,
             "13": 10.92213,
             "14": 10.91569,
@@ -455,8 +455,8 @@
             "21": 10.90603,
             "22": 10.90322,
             "23": 10.90169,
-            "24": 10.8909,
-            "25": 10.8825,
+            "24": 10.89089,
+            "25": 10.88251,
             "26": 10.89359,
             "27": 10.8887,
             "28": 10.87619,
@@ -477,7 +477,7 @@
             "43": 10.84328,
             "44": 10.82898,
             "45": 10.84347,
-            "46": 10.83298,
+            "46": 10.83297,
             "47": 10.83911,
             "48": 10.82542,
             "49": 10.83132,
@@ -494,7 +494,7 @@
             "60": 10.77895,
             "61": 10.77556,
             "62": 10.76277,
-            "63": 10.77594,
+            "63": 10.77595,
             "64": 10.76136,
             "65": 10.7585,
             "66": 10.75798,
@@ -514,17 +514,17 @@
             "80": 10.67858,
             "81": 10.67147,
             "82": 10.65165,
-            "83": 10.63057,
+            "83": 10.63056,
             "84": 10.61714,
             "85": 10.60392,
             "86": 10.63183,
             "87": 10.62791,
-            "88": 10.62833,
+            "88": 10.62832,
             "89": 10.59789,
             "90": 10.59506,
             "91": 10.60606,
             "92": 10.58205,
-            "93": 10.55314,
+            "93": 10.55313,
             "94": 10.58516,
             "95": 10.57313,
             "96": 10.56963,
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200_2nd.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200_2nd.json
index 58d969d35d7..d280e191db7 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200_2nd.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_gb200_2nd.json
@@ -494,7 +494,7 @@
             "60": 10.77895,
             "61": 10.77556,
             "62": 10.76277,
-            "63": 10.77594,
+            "63": 10.77595,
             "64": 10.76136,
             "65": 10.7585,
             "66": 10.75798,
@@ -514,17 +514,17 @@
             "80": 10.67858,
             "81": 10.67147,
             "82": 10.65165,
-            "83": 10.63057,
+            "83": 10.63056,
             "84": 10.61714,
             "85": 10.60392,
             "86": 10.63183,
             "87": 10.62791,
-            "88": 10.62833,
+            "88": 10.62832,
             "89": 10.59789,
             "90": 10.59506,
             "91": 10.60606,
             "92": 10.58205,
-            "93": 10.55314,
+            "93": 10.55313,
             "94": 10.58516,
             "95": 10.57313,
             "96": 10.56963,
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json
index cc3963c29d9..48b68ba5823 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100.json
@@ -437,7 +437,7 @@
             "3": 10.93384,
             "4": 10.92739,
             "5": 10.90724,
-            "6": 10.91817,
+            "6": 10.91816,
             "7": 10.92486,
             "8": 10.92528,
             "9": 10.93457,
@@ -445,17 +445,17 @@
             "11": 10.91896,
             "12": 10.91863,
             "13": 10.92814,
-            "14": 10.91203,
+            "14": 10.91204,
             "15": 10.92041,
             "16": 10.92467,
             "17": 10.92235,
-            "18": 10.90719,
+            "18": 10.9072,
             "19": 10.91438,
             "20": 10.90506,
             "21": 10.91161,
             "22": 10.89778,
             "23": 10.90483,
-            "24": 10.88964,
+            "24": 10.88965,
             "25": 10.89765,
             "26": 10.88453,
             "27": 10.89849,
@@ -491,7 +491,7 @@
             "57": 10.78961,
             "58": 10.79824,
             "59": 10.78095,
-            "60": 10.77503,
+            "60": 10.77504,
             "61": 10.77627,
             "62": 10.7614,
             "63": 10.78392,
@@ -511,7 +511,7 @@
             "77": 10.69055,
             "78": 10.68188,
             "79": 10.66968,
-            "80": 10.67688,
+            "80": 10.67687,
             "81": 10.66904,
             "82": 10.65016,
             "83": 10.6267,
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100_2nd.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100_2nd.json
index 357c399d4b7..786d23d265f 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100_2nd.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph/golden_values_dev_dgx_h100_2nd.json
@@ -491,7 +491,7 @@
             "57": 10.78961,
             "58": 10.79824,
             "59": 10.78095,
-            "60": 10.77503,
+            "60": 10.77504,
             "61": 10.77627,
             "62": 10.7614,
             "63": 10.78392,
@@ -511,7 +511,7 @@
             "77": 10.69055,
             "78": 10.68188,
             "79": 10.66968,
-            "80": 10.67688,
+            "80": 10.67687,
             "81": 10.66904,
             "82": 10.65016,
             "83": 10.6267,
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph_1node/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph_1node/golden_values_dev_dgx_gb200.json
index d82c4eb4512..8dd231f1acc 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph_1node/golden_values_dev_dgx_gb200.json
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_scoped_cudagraph_1node/golden_values_dev_dgx_gb200.json
@@ -464,7 +464,7 @@
             "30": 10.93961,
             "31": 10.91806,
             "32": 10.92899,
-            "33": 10.92159,
+            "33": 10.92158,
             "34": 10.92688,
             "35": 10.91862,
             "36": 10.917,
@@ -475,13 +475,13 @@
             "41": 10.90546,
             "42": 10.88722,
             "43": 10.89763,
-            "44": 10.87484,
+            "44": 10.87485,
             "45": 10.88603,
             "46": 10.87926,
             "47": 10.87569,
             "48": 10.86052,
             "49": 10.86102,
-            "50": 10.84773,
+            "50": 10.84774,
             "51": 10.86037,
             "52": 10.84549,
             "53": 10.85137,
@@ -497,7 +497,7 @@
             "63": 10.80803,
             "64": 10.80094,
             "65": 10.78782,
-            "66": 10.78839,
+            "66": 10.78838,
             "67": 10.78222,
             "68": 10.76003,
             "69": 10.78043,
diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py
index 3eb02442fe9..651cb890f60 100644
--- a/tests/unit_tests/ssm/test_gated_delta_net.py
+++ b/tests/unit_tests/ssm/test_gated_delta_net.py
@@ -182,7 +182,9 @@ def test_jit_compiled_helpers(self):
         # which are normally wrapped by @jit_fuser (torch.compile).
         with torch._dynamo.config.patch(disable=True):
             query, key, value, gate_out, beta_out, alpha_out = (
-                gdn._prepare_qkv_for_gated_delta_rule(qkv, gate, beta, alpha, batch, seq_len)
+                gdn._prepare_qkv_for_gated_delta_rule(
+                    qkv, gate, beta, alpha, batch, seq_len, gdn.cp_size
+                )
             )
 
         assert query.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim)
diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py
index d594b46b373..bf929c374d4 100644
--- a/tests/unit_tests/test_sequence_packing.py
+++ b/tests/unit_tests/test_sequence_packing.py
@@ -396,6 +396,11 @@ def _create_single_sample(seq_len):
     config.microbatch_group_size_per_vp_stage = pp
     config.virtual_pipeline_model_parallel_size = vpp
     config.sequence_packing_scheduler = scheduler_type
+    # wrap_data_iterator -> mtp_on_this_rank reads these two config fields.
+    # A real TransformerConfig defaults both to None when MTP is unused, so
+    # mirror that here (this test does not exercise MTP).
+    config.pipeline_model_parallel_layout = None
+    config.mtp_num_layers = None
     if is_dynamic_cp:
         config.min_dynamic_context_parallel_size = 1
 
@@ -412,7 +417,11 @@ def _create_single_sample(seq_len):
 
     num_micro_batches_old = global_batch_size // micro_batch_size // dp_size
 
-    if is_tp_first and (is_pp_first or is_pp_last):
+    # In packed-sequence mode is_dataset_built_on_rank returns True for every
+    # PP stage on TP rank 0 (not just first/last), so all PP stages — including
+    # middle ones — own a data_iterator and run the scheduler. Middle stages
+    # later strip their samples down to metadata only.
+    if is_tp_first:
         # Seed torch RNG so CP siblings produce identical token values
         torch.manual_seed(42 + dp_rank)
         torch.cuda.manual_seed(42 + dp_rank)
@@ -431,7 +440,10 @@ def _create_single_sample(seq_len):
             elif is_pp_last:
                 data_iterator = [None for _ in range(vpp - 1)] + [data_iterator]
             else:
-                data_iterator = [None for _ in range(vpp)]
+                # Middle PP stage: no VPP sub-stage needs full data, but the
+                # scheduler still needs an iterator (slot 0) to derive the
+                # microbatch count and per-stage metadata.
+                data_iterator = [data_iterator] + [None for _ in range(vpp - 1)]
     try:
         # Call the function under test
         (
@@ -465,6 +477,16 @@ def _check_batch(batch_all, batch_keys):
             batch_keys = ["cu_seqlens", "max_seqlen", "cu_seqlens_padded"]
             if is_dynamic_cp:
                 batch_keys.append("local_cp_size")
+            # Per-stage data field stripping (see data_schedule.py): the first PP
+            # stage keeps tokens/position_ids, the last PP stage keeps labels/
+            # loss_mask. When pp==1 a single stage is both first and last, so it
+            # keeps all four. Middle stages carry metadata only.
+            stage_data_keys = []
+            if is_pp_first:
+                stage_data_keys += ["tokens", "position_ids"]
+            if is_pp_last:
+                stage_data_keys += ["labels", "loss_mask"]
+
             if vpp is not None and vpp > 1:
                 # check metadata for all stages (save batches to avoid re-consuming iterators)
                 all_stage_batches = []
@@ -476,21 +498,22 @@ def _check_batch(batch_all, 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)
+                    _check_batch(batch_all, batch_keys + stage_data_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_batch(batch_all, batch_keys + stage_data_keys)
 
-            # CHECK TOKEN SUM ON FIRST OR LAST PP RANK
+            # CHECK TOKEN SUM ON FIRST 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.
             # Skip for VPP: microbatch alignment may pad/duplicate samples,
             # changing the total token count.
-            if is_pp_first_or_last and (vpp is None or vpp <= 1):
+            # Only the first PP stage is checked: with per-stage data stripping the
+            # last PP stage (when pp>1) keeps labels/loss_mask and no longer carries
+            # 'tokens'. The dp/dp_cp all-reduce groups are disjoint per PP rank, so
+            # running this on the first PP stage alone is collective-safe.
+            if is_pp_first and (vpp is None or vpp <= 1):
                 dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=True)
                 cp_size = parallel_state.get_context_parallel_world_size()
                 cp_group = parallel_state.get_context_parallel_group()
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index a225cd376a6..d2b448e1618 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -22,7 +22,6 @@
 from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.parallel_state import get_context_parallel_group
-from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer.hyper_connection import learned_output_contract
 from megatron.core.transformer.multi_token_prediction import (
@@ -495,8 +494,10 @@ def test_forward_backward(self, tmp_path_dist_ckpt, tp, cp, full_recompute):
             labels=labels,
             loss_mask=loss_mask,
         )
+        # forward only fills raw loss_sums / num_tokens. Trigger the reduction
+        # so tracker["values"] (per-token loss across DP+CP) becomes available.
+        MTPLossLoggingHelper.reduce_loss_in_tracker()
         tracker = MTPLossLoggingHelper.tracker
-        mtp_loss_ref = None
         assert "values" in tracker
         mtp_loss_ref = tracker['values'].clone()
         MTPLossLoggingHelper.clean_loss_in_tracker()
@@ -547,14 +548,13 @@ def set_ckpt_path(ckpt_path):
                 labels=labels,
                 loss_mask=loss_mask,
             )
+            # reduce_loss_in_tracker performs sum-reduce of loss_sums and
+            # num_tokens across DP+CP, then computes sum/sum -- already the
+            # correct global per-token loss, no extra CP averaging needed.
+            MTPLossLoggingHelper.reduce_loss_in_tracker()
             tracker = MTPLossLoggingHelper.tracker
             assert "values" in tracker
             mtp_loss = tracker['values'].clone()
-            # Average MTP loss across CP ranks for comparison with reference
-            pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['cp'])
-            torch.distributed.all_reduce(
-                mtp_loss, group=pg_collection.cp, op=torch.distributed.ReduceOp.AVG
-            )
             MTPLossLoggingHelper.clean_loss_in_tracker()
             assert torch.allclose(output_ref, output, rtol=1e-03, atol=1e-03)
             assert torch.allclose(mtp_loss, mtp_loss_ref, rtol=1e-02, atol=1e-02)
@@ -651,7 +651,9 @@ def test_packed_sequences(self, tp, cp):
         assert output.shape[0] == 1  # batch size
         assert output.shape[1] == total_seq_length
 
-        # Verify MTP loss was computed
+        # Verify MTP loss was computed; reduce raw loss_sums/num_tokens into
+        # tracker["values"] (per-token loss) first.
+        MTPLossLoggingHelper.reduce_loss_in_tracker()
         tracker = MTPLossLoggingHelper.tracker
         assert "values" in tracker
         mtp_loss = tracker['values'].clone()
@@ -906,35 +908,40 @@ def teardown_method(self, method):
         MTPLossLoggingHelper.tracker = {}
 
     def test_save_loss_to_tracker(self):
-        """Test saving loss to tracker."""
-        # Create a dummy loss tensor
-        loss = torch.tensor(1.3)
+        """Test saving loss sum and token count to tracker."""
+        loss_sum = torch.tensor(1.3)
+        num_tokens = torch.tensor(5.0)
         layer_number = 2
         num_layers = self.num_layers
 
-        # Test saving loss
         MTPLossLoggingHelper.save_loss_to_tracker(
-            loss=loss, layer_number=layer_number, num_layers=num_layers
+            loss_sum=loss_sum,
+            num_tokens=num_tokens,
+            layer_number=layer_number,
+            num_layers=num_layers,
         )
 
-        # Verify tracker state
-        assert "values" in MTPLossLoggingHelper.tracker
-        assert MTPLossLoggingHelper.tracker["values"].shape == (num_layers,)
-        assert MTPLossLoggingHelper.tracker["values"][layer_number] == loss
+        # Tracker now stores raw loss sums and token counts; per-token loss
+        # is computed in reduce_loss_in_tracker.
+        assert "loss_sums" in MTPLossLoggingHelper.tracker
+        assert "num_tokens" in MTPLossLoggingHelper.tracker
+        assert MTPLossLoggingHelper.tracker["loss_sums"].shape == (num_layers,)
+        assert MTPLossLoggingHelper.tracker["num_tokens"].shape == (num_layers,)
+        assert MTPLossLoggingHelper.tracker["loss_sums"][layer_number] == loss_sum
+        assert MTPLossLoggingHelper.tracker["num_tokens"][layer_number] == num_tokens
         assert MTPLossLoggingHelper.tracker["reduce_group"] is None
         assert MTPLossLoggingHelper.tracker["avg_group"] is None
 
     def test_track_mtp_metrics(self):
         """Test tracking MTP metrics."""
-        # First save some losses
-        loss = torch.tensor(2.3)
+        loss_sum = torch.tensor(2.3)
+        num_tokens = torch.tensor(1.0)
         num_layers = self.num_layers
         for i in range(num_layers):
             MTPLossLoggingHelper.save_loss_to_tracker(
-                loss=loss, layer_number=i, num_layers=num_layers
+                loss_sum=loss_sum, num_tokens=num_tokens, layer_number=i, num_layers=num_layers
             )
 
-        # Create dummy writer and loss dict
         class DummyWriter:
             def add_scalar(self, name, value, iteration):
                 pass
@@ -949,7 +956,6 @@ def log(self, metrics, iteration):
         wandb_writer = DummyWandBWriter()
         total_loss_dict = {}
 
-        # Test tracking metrics
         MTPLossLoggingHelper.track_mtp_metrics(
             loss_scale=loss_scale,
             iteration=iteration,
@@ -958,13 +964,18 @@ def log(self, metrics, iteration):
             total_loss_dict=total_loss_dict,
         )
 
-        # Verify total_loss_dict is populated
+        # track_mtp_metrics reduces the tracker first, so per-layer log value
+        # equals (loss_sum / num_tokens) * loss_scale.
+        expected_per_token_loss = loss_sum / num_tokens
         for i in range(num_layers):
             assert f"mtp_{i + 1} loss" in total_loss_dict
-            assert total_loss_dict[f"mtp_{i + 1} loss"] == loss * loss_scale
+            assert torch.allclose(
+                total_loss_dict[f"mtp_{i + 1} loss"], expected_per_token_loss * loss_scale
+            )
 
-        # Verify tracker is cleaned
-        assert torch.all(MTPLossLoggingHelper.tracker["values"] == 0)
+        # Tracker raw sums are cleared by track_mtp_metrics.
+        assert torch.all(MTPLossLoggingHelper.tracker["loss_sums"] == 0)
+        assert torch.all(MTPLossLoggingHelper.tracker["num_tokens"] == 0)
         assert MTPLossLoggingHelper.tracker["reduce_group"] is None
         assert MTPLossLoggingHelper.tracker["avg_group"] is None
 
@@ -1125,8 +1136,10 @@ def test_forward_backward_mamba(self, tmp_path_dist_ckpt, tp, cp):
             labels=labels,
             loss_mask=loss_mask,
         )
+        # forward only fills raw loss_sums / num_tokens. Reduce them first so
+        # tracker["values"] (per-token loss across DP+CP) becomes available.
+        MTPLossLoggingHelper.reduce_loss_in_tracker()
         tracker = MTPLossLoggingHelper.tracker
-        mtp_loss_ref = None
         assert "values" in tracker
         mtp_loss_ref = tracker['values'].clone()
         MTPLossLoggingHelper.clean_loss_in_tracker()
@@ -1172,13 +1185,12 @@ def set_ckpt_path(ckpt_path):
                 labels=labels,
                 loss_mask=loss_mask,
             )
+            # reduce_loss_in_tracker already computes the cross-DP+CP per-token
+            # loss (sum/sum), no extra CP averaging needed.
+            MTPLossLoggingHelper.reduce_loss_in_tracker()
             tracker = MTPLossLoggingHelper.tracker
             assert "values" in tracker
             mtp_loss = tracker['values'].clone()
-            pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['cp'])
-            torch.distributed.all_reduce(
-                mtp_loss, group=pg_collection.cp, op=torch.distributed.ReduceOp.AVG
-            )
             MTPLossLoggingHelper.clean_loss_in_tracker()
             assert torch.allclose(output_ref, output, rtol=1e-03, atol=1e-03)
             assert torch.allclose(mtp_loss, mtp_loss_ref, rtol=1e-02, atol=1e-02)
@@ -1478,6 +1490,9 @@ def model_provider(
         )
         assert torch.isfinite(output).all(), f"Non-finite output (TP={tp})"
 
+        # Reduce raw loss_sums/num_tokens into tracker["values"] (per-token
+        # loss across DP+CP) before reading.
+        MTPLossLoggingHelper.reduce_loss_in_tracker()
         tracker = MTPLossLoggingHelper.tracker
         assert "values" in tracker, f"MTP loss not logged (TP={tp})"
         assert torch.isfinite(tracker['values']).all()