several fixes for THD e2e

megatron/core/datasets/data_schedule.py

@@ -468,7 +468,7 @@ def _get_scheduler_max_real_num_seqs(config) -> Optional[int]:
         raise ValueError(f"thd_max_packed_sequences must be >= 1, got {max_num_seqs}.")
 
     if getattr(config, 'pad_packed_seq_alignment', None) is not None and getattr(
-        config, 'pad_packed_seq_by_appending_dummy_seq', True
+        config, 'pad_packed_seq_by_appending_dummy_seq', False
     ):
         if max_num_seqs < 2:
             raise ValueError(
@@ -778,7 +778,6 @@ def get_batch_on_this_rank_for_sequence_packing(
         max_seqlen_kv=max_seqlen,
         local_cp_size=local_cp_size,
         cp_group=cp_group,
-        pad_between_seqs=False,
     )
 
     # Pad the already-packed THD tensors at the end when requested. CUDA Graph
@@ -804,7 +803,7 @@ def get_batch_on_this_rank_for_sequence_packing(
                 target_len=target_len,
                 max_num_seqs=max_num_seqs,
                 pad_by_appending_dummy_seq=getattr(
-                    config, 'pad_packed_seq_by_appending_dummy_seq', True
+                    config, 'pad_packed_seq_by_appending_dummy_seq', False
                 ),
                 padding_mask=padding_mask,
                 cp_group=cp_group,

megatron/core/model_parallel_config.py

@@ -118,13 +118,13 @@ class ModelParallelConfig:
     tensors are padded to a multiple of N.
     """
 
-    pad_packed_seq_by_appending_dummy_seq: bool = True
+    pad_packed_seq_by_appending_dummy_seq: bool = False
     """Represent a THD packed-sequence padding tail by appending a dummy sequence.
 
-    When disabled, token-like tensors are still padded according to
-    pad_packed_seq_alignment, but cu_seqlens sequence boundaries are not extended
-    for the padding tail. CUDA Graph static-input padding may still pad the
-    cu_seqlens tensors to thd_max_packed_sequences + 1 entries.
+    By default, token-like tensors are padded according to
+    pad_packed_seq_alignment, cu_seqlens is unchanged, and the last
+    cu_seqlens_padded endpoint is extended over the padding tail. When enabled,
+    the tail is represented as a separate zero-valid-token dummy sequence.
     """
 
     expert_model_parallel_size: int = 1

megatron/core/packed_seq_params.py

@@ -159,6 +159,16 @@ def _append_dummy_seq(cu_seqlens: Optional[Tensor], dummy_end: int) -> Optional[
     return torch.cat((cu_seqlens, dummy), dim=0)
 
 
+def _replace_last_cu_seqlen(cu_seqlens: Optional[Tensor], padded_end: int) -> Optional[Tensor]:
+    """Return cu_seqlens with its final physical endpoint replaced."""
+    if cu_seqlens is None:
+        return None
+
+    result = cu_seqlens.clone()
+    result[-1] = int(padded_end)
+    return result
+
+
 def _round_up_to_alignment(value: int, alignment: int) -> int:
     assert alignment > 0, f"Packed sequence padding alignment must be > 0, got {alignment}."
     return ((value + alignment - 1) // alignment) * alignment
@@ -225,9 +235,14 @@ def _resolve_thd_padding_lengths(
             mask_device = candidate.device
             break
 
-    # Prefer THD metadata for the global packed length when it is available.
+    # The padded endpoint describes tensor storage. The unpadded endpoint is
+    # only the compact valid-token count when gaps exist between sequences.
     has_local_tensor = local_tensor_T is not None
-    if packed_seq_params.cu_seqlens_q is not None:
+    if packed_seq_params.cu_seqlens_q_padded is not None:
+        global_actual_T = int(packed_seq_params.cu_seqlens_q_padded[-1].item())
+        if mask_device is None:
+            mask_device = packed_seq_params.cu_seqlens_q_padded.device
+    elif packed_seq_params.cu_seqlens_q is not None:
         global_actual_T = int(packed_seq_params.cu_seqlens_q[-1].item())
         if mask_device is None:
             mask_device = packed_seq_params.cu_seqlens_q.device
@@ -335,7 +350,7 @@ def pad_sequence_for_thd(
     alignment: Optional[int] = None,
     target_len: Optional[int] = None,
     max_num_seqs: Optional[int] = None,
-    pad_by_appending_dummy_seq: bool = True,
+    pad_by_appending_dummy_seq: bool = False,
     padding_mask: Optional[Tensor] = None,
     cp_group: Optional[dist.ProcessGroup] = None,
     cp_size: Optional[int] = None,
@@ -371,7 +386,8 @@ def pad_sequence_for_thd(
         max_num_seqs: If set, pad cu_seqlens tensors to
             ``max_num_seqs + 1`` entries for static CUDA Graph inputs.
         pad_by_appending_dummy_seq: If true, represent the post-pack padding
-            tail as an extra dummy sequence in cu_seqlens metadata.
+            tail as an extra dummy sequence. Otherwise leave cu_seqlens unchanged
+            and extend the final cu_seqlens_padded endpoint over the tail.
         padding_mask: Existing bool padding mask for already-packed tokens,
             with True marking padding positions.
         cp_group: Context-parallel process group for resolving local/global
@@ -386,8 +402,10 @@ def pad_sequence_for_thd(
           stages, Megatron asks TE which packed rows this CP rank would receive
           and uses that row count as the local length instead of assuming equal
           division by CP size.
-        - When ``pad_by_appending_dummy_seq`` is true, the padding tail is also
-          represented as an ordinary dummy sequence in cu_seqlens metadata.
+        - By default, post-pack padding extends only the last padded sequence;
+          original cu_seqlens continue to describe valid-token counts.
+        - When ``pad_by_appending_dummy_seq`` is true, the padding tail is
+          represented as a separate zero-valid-token dummy sequence.
         - ``max_num_seqs`` pads all four cu_seqlens tensors; this is required
           by CUDA Graph replay because those tensors are graph inputs.
 
@@ -437,16 +455,40 @@ def pad_sequence_for_thd(
     cu_seqlens_q_padded = packed_seq_params.cu_seqlens_q_padded
     cu_seqlens_kv_padded = packed_seq_params.cu_seqlens_kv_padded
 
-    # Represent post-pack padding as a dummy sequence when requested.
+    # Cover post-pack padding by extending the last physical sequence, or by
+    # representing it as a separate zero-valid-token dummy sequence.
     target_cu_entries = None if max_num_seqs is None else max_num_seqs + 1
-    has_dummy_padding_seq = pad_by_appending_dummy_seq and global_target_len > global_actual_T
-    dummy_seq_len = global_target_len - global_actual_T if has_dummy_padding_seq else 0
+    has_padding_tail = global_target_len > global_actual_T
+    has_dummy_padding_seq = pad_by_appending_dummy_seq and has_padding_tail
 
     if has_dummy_padding_seq:
-        cu_seqlens_q = _append_dummy_seq(cu_seqlens_q, global_target_len)
-        cu_seqlens_kv = _append_dummy_seq(cu_seqlens_kv, global_target_len)
-        cu_seqlens_q_padded = _append_dummy_seq(cu_seqlens_q_padded, global_target_len)
-        cu_seqlens_kv_padded = _append_dummy_seq(cu_seqlens_kv_padded, global_target_len)
+        # None is TE's auto-detect mode. Keep real and physical boundaries
+        # separate unless the caller explicitly disables padding between sequences.
+        if (
+            packed_seq_params.pad_between_seqs is not False
+            and cu_seqlens_q is not None
+            and cu_seqlens_q_padded is not None
+        ):
+            cu_seqlens_q = _append_dummy_seq(cu_seqlens_q, int(cu_seqlens_q[-1].item()))
+            cu_seqlens_q_padded = _append_dummy_seq(cu_seqlens_q_padded, global_target_len)
+        else:
+            cu_seqlens_q = _append_dummy_seq(cu_seqlens_q, global_target_len)
+            cu_seqlens_q_padded = _append_dummy_seq(cu_seqlens_q_padded, global_target_len)
+        if (
+            packed_seq_params.pad_between_seqs is not False
+            and cu_seqlens_kv is not None
+            and cu_seqlens_kv_padded is not None
+        ):
+            cu_seqlens_kv = _append_dummy_seq(cu_seqlens_kv, int(cu_seqlens_kv[-1].item()))
+            cu_seqlens_kv_padded = _append_dummy_seq(cu_seqlens_kv_padded, global_target_len)
+        else:
+            cu_seqlens_kv = _append_dummy_seq(cu_seqlens_kv, global_target_len)
+            cu_seqlens_kv_padded = _append_dummy_seq(cu_seqlens_kv_padded, global_target_len)
+    elif has_padding_tail:
+        q_physical_cu = cu_seqlens_q_padded if cu_seqlens_q_padded is not None else cu_seqlens_q
+        kv_physical_cu = cu_seqlens_kv_padded if cu_seqlens_kv_padded is not None else cu_seqlens_kv
+        cu_seqlens_q_padded = _replace_last_cu_seqlen(q_physical_cu, global_target_len)
+        cu_seqlens_kv_padded = _replace_last_cu_seqlen(kv_physical_cu, global_target_len)
 
     # Pad cu_seqlens entry counts for static CUDA Graph inputs.
     if target_cu_entries is not None:
@@ -455,6 +497,12 @@ def pad_sequence_for_thd(
         cu_seqlens_q_padded = _pad_cu_seqlens(cu_seqlens_q_padded, target_cu_entries)
         cu_seqlens_kv_padded = _pad_cu_seqlens(cu_seqlens_kv_padded, target_cu_entries)
 
+    def _max_physical_seqlen(cu_seqlens: Optional[Tensor], fallback: int) -> int:
+        if cu_seqlens is None or cu_seqlens.numel() < 2:
+            return fallback
+        lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+        return max(fallback, int(lengths.max().item()))
+
     # Rebuild PackedSeqParams with the padded tensor and metadata shapes.
     padded_params = PackedSeqParams(
         qkv_format=packed_seq_params.qkv_format,
@@ -465,17 +513,17 @@ def pad_sequence_for_thd(
         max_seqlen_q=(
             global_target_len
             if target_cu_entries is not None
-            else max(packed_seq_params.max_seqlen_q, dummy_seq_len)
+            else _max_physical_seqlen(cu_seqlens_q_padded, packed_seq_params.max_seqlen_q)
         ),
         max_seqlen_kv=(
             global_target_len
             if target_cu_entries is not None
-            else max(packed_seq_params.max_seqlen_kv, dummy_seq_len)
+            else _max_physical_seqlen(cu_seqlens_kv_padded, packed_seq_params.max_seqlen_kv)
         ),
         local_cp_size=packed_seq_params.local_cp_size,
         cp_group=packed_seq_params.cp_group,
         total_tokens=local_target_len if target_cu_entries is None else None,
-        pad_between_seqs=False if has_dummy_padding_seq else packed_seq_params.pad_between_seqs,
+        pad_between_seqs=packed_seq_params.pad_between_seqs,
     )
 
     # True marks padded local token slots for routing/loss paths.

megatron/core/transformer/cuda_graphs.py

@@ -2344,7 +2344,7 @@ def _get_thd_varlen_max_num_microbatches(
         if max_num_seqs is not None:
             max_num_seqs = int(max_num_seqs)
             if getattr(self.config, 'pad_packed_seq_alignment', None) is not None and getattr(
-                self.config, 'pad_packed_seq_by_appending_dummy_seq', True
+                self.config, 'pad_packed_seq_by_appending_dummy_seq', False
             ):
                 max_num_seqs -= 1
 

megatron/core/transformer/moe/router.py

@@ -827,7 +827,15 @@ def routing(
                 router_replay=self.router_replay,
             )
 
-        # Apply token dropping to probs and routing_map.
+        # Padding tokens must not enter the dispatcher. Masking only auxiliary
+        # losses still lets padding alter expert capacity and grouped-GEMM shapes.
+        if padding_mask is not None:
+            valid_tokens = (~padding_mask).unsqueeze(-1)
+            probs = probs * valid_tokens
+            routing_map = routing_map & valid_tokens
+
+        # Apply token dropping after removing padding so padding tokens cannot
+        # consume expert capacity.
         if self.config.moe_expert_capacity_factor is not None:
             probs, routing_map = apply_router_token_dropping(
                 probs,

megatron/core/transformer/multi_token_prediction.py

@@ -371,8 +371,8 @@ def save_metrics_to_tracker(
         """Save normalized MTP loss and acceptance counts for logging.
 
         This compatibility path is used by tests and callers that already
-        computed a normalized per-layer loss. Dynamic-CP code should use
-        ``save_loss_to_tracker`` so loss is weighted by token counts.
+        computed a normalized per-layer loss. Dynamic-CP code uses
+        ``save_loss_to_tracker`` to normalize each local contribution safely.
         """
         if layer_number is None:
             return
@@ -402,11 +402,12 @@ def save_loss_to_tracker(
         reduce_group: Optional[torch.distributed.ProcessGroup] = None,
         avg_group: Optional[torch.distributed.ProcessGroup] = None,
     ):
-        """Save the mtp loss sum and token count for logging.
+        """Normalize and accumulate a local MTP loss 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.
+        MTP is normalized independently for each microbatch. The tracker
+        accumulates those normalized losses, then reduction combines the sums
+        across ranks. This intentionally preserves sequence-packing semantics
+        instead of changing the metric to global ``sum(loss) / sum(tokens)``.
 
         Args:
             loss_sum (torch.Tensor): Sum of per-element losses on this rank.
@@ -415,18 +416,17 @@ def save_loss_to_tracker(
             num_layers (int): The number of total layers.
             correct (Optional[torch.Tensor]): Number of correct MTP predictions.
             total (Optional[torch.Tensor]): Total number of MTP predictions.
-            reduce_group (torch.distributed.ProcessGroup): The group for sum-reducing losses.
-            avg_group (torch.distributed.ProcessGroup): The group for sum-reducing before averaging.
+            reduce_group (torch.distributed.ProcessGroup): Group for summing losses.
+            avg_group (torch.distributed.ProcessGroup): Group for averaging losses.
         """
         if layer_number is None:
             return
 
         tracker = MTPLossLoggingHelper.tracker
         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())
+        loss_sum = (loss_sum * (num_tokens > 0).to(loss_sum.dtype)) / num_tokens.clamp(min=1)
         tracker["loss_sums"][layer_number] += loss_sum.detach()
-        tracker["num_tokens"][layer_number] += num_tokens.detach()
         if correct is not None and total is not None:
             if "correct_values" not in tracker:
                 tracker["correct_values"] = torch.zeros(
@@ -485,7 +485,6 @@ def clean_loss_in_tracker():
         tracker = MTPLossLoggingHelper.tracker
         if "loss_sums" in tracker:
             tracker["loss_sums"].zero_()
-            tracker["num_tokens"].zero_()
         if "values" in tracker:
             tracker["values"].zero_()
         if "correct_values" in tracker:
@@ -499,21 +498,21 @@ def clean_loss_in_tracker():
     def reduce_loss_in_tracker():
         """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).
+        Each element is already a sum of normalized microbatch losses. Sum
+        reductions preserve additive groups, while the DP+CP average keeps the
+        legacy logging contract.
         """
         tracker = MTPLossLoggingHelper.tracker
         if "loss_sums" not in tracker:
             return
-        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)
+        values = tracker["loss_sums"]
+        if tracker.get('reduce_group') is not None:
+            torch.distributed.all_reduce(values, group=tracker['reduce_group'])
+        if tracker.get('avg_group') is not None:
+            torch.distributed.all_reduce(
+                values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG
+            )
+        tracker["values"] = values
 
     @staticmethod
     def track_mtp_metrics(loss_scale, iteration, writer, wandb_writer=None, total_loss_dict=None):

megatron/core/transformer/transformer_layer.py

@@ -1327,7 +1327,6 @@ def _reconstruct_packed_seq_params_from_kwargs(self, kwargs):
             cu_seqlens_kv_padded=kwargs.pop('cu_seqlens_kv_padded'),
             max_seqlen_q=max_seqlen,
             max_seqlen_kv=max_seqlen,
-            pad_between_seqs=False,
         )
         kwargs['packed_seq_params'] = packed_seq_params
 

megatron/training/datasets/varlen_dataset.py

@@ -406,10 +406,15 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
             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
+            # Keep physical padding separate from the LM loss mask: prompt
+            # tokens may be loss-masked but must still participate in MoE.
+            padding_mask = torch.zeros(max_len, dtype=torch.bool)
+            padding_mask[valid_len:] = True
             return {
                 'tokens': input_ids,
                 'labels': labels,
                 'loss_mask': loss_mask,
+                'padding_mask': padding_mask,
                 'position_ids': torch.arange(max_len, dtype=torch.int64),
             }
 

megatron/training/training.py

@@ -2954,10 +2954,10 @@ def training_log(
 
     # Log MTP metrics.
     if args.mtp_num_layers is not None:
-        # 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
+        # The tracker stores a sum of normalized microbatch losses.
+        # Sequence-packing schedulers may change the number of microbatches for
+        # this step, so use the scheduled count passed to training_log.
+        mtp_loss_scale = 1 / (num_microbatches or get_num_microbatches())
         MTPLossLoggingHelper.track_mtp_metrics(
             mtp_loss_scale, iteration, writer, wandb_writer, total_loss_dict
         )