[dev] [follow-up] Qwen3.5 support: MoE aux loss padding_mask

examples/multimodal_dev/forward_step.py

@@ -227,9 +227,23 @@ def pack_or_pad_batch(
             cu_seqlens = list(accumulate(seqlens_list, initial=0))
             cu_seqlens_padded = list(accumulate(seqlens_padded_list, initial=0))
 
+            # padding_mask: True at collate-padded positions within each packed
+            # sample. Real tokens occupy [cu_seqlens_padded[i], +seqlens_list[i]);
+            # the tail up to cu_seqlens_padded[i+1] is padding. Consumed by MoE
+            # routing in megatron.core to exclude padded tokens from aux loss,
+            # z-loss, and expert-bias accumulation.
+            total_tokens_padded = cu_seqlens_padded[-1]
+            padding_mask_thd = torch.zeros(total_tokens_padded, dtype=torch.bool)
+            for i, real_seqlen in enumerate(seqlens_list):
+                pad_start = cu_seqlens_padded[i] + real_seqlen
+                pad_end = cu_seqlens_padded[i + 1]
+                if pad_end > pad_start:
+                    padding_mask_thd[pad_start:pad_end] = True
+
             packed_batch["input_ids"] = torch.concat(input_ids_list, dim=0).unsqueeze(0)
             packed_batch["labels"] = torch.concat(labels_list, dim=0).unsqueeze(0)
             packed_batch["loss_mask"] = torch.concat(loss_mask_list, dim=0).unsqueeze(0)
+            packed_batch["padding_mask"] = padding_mask_thd.unsqueeze(0)
             packed_batch["pixel_values"] = torch.concat(pixel_values_list)
             packed_batch["image_grid_thw"] = torch.concat(image_grid_thw_list)
             # cu_seqlens / cu_seqlens_padded need to reach non-source TP ranks
@@ -273,6 +287,10 @@ def pack_or_pad_batch(
         if divisible_by > 1:
             target_seqlens = math.ceil(target_seqlens / divisible_by) * divisible_by
 
+        # Capture real lengths before in-place padding so we can build a
+        # padding_mask for MoE routing (True at collate-padded positions).
+        real_seqlens = [s["input_ids"].shape[0] for s in batch]
+
         for sample in batch:
             sample["input_ids"] = F.pad(
                 sample["input_ids"], (0, target_seqlens - sample["input_ids"].shape[0]), value=0
@@ -291,6 +309,8 @@ def pack_or_pad_batch(
         padded_batch["loss_mask"] = torch.concat(
             [x["loss_mask"].unsqueeze(0) for x in batch], dim=0
         )
+        positions = torch.arange(target_seqlens).unsqueeze(0)
+        padded_batch["padding_mask"] = positions >= torch.tensor(real_seqlens).unsqueeze(1)
         padded_batch["pixel_values"] = torch.concat([x["pixel_values"] for x in batch])
         padded_batch["image_grid_thw"] = torch.concat([x["image_grid_thw"] for x in batch])
 
@@ -392,6 +412,7 @@ def forward_step(data_iterator, model):
         attention_mask=batch.get("attention_mask", None),
         labels=batch.get("labels", None),
         loss_mask=batch.get("loss_mask", None),
+        padding_mask=batch.get("padding_mask", None),
         pixel_values=pixel_values,
         image_grid_thw=batch.get("image_grid_thw", None),
         packed_seq_params=batch.get("packed_seq_params", None),

examples/multimodal_dev/models/base.py

@@ -202,6 +202,7 @@ def _cp_split_for_forward(
         attention_mask,
         position_ids,
         packed_seq_params,
+        padding_mask=None,
     ):
         """Apply CP split to model-forward inputs.
 
@@ -215,7 +216,10 @@ def _cp_split_for_forward(
         """
         cp_size = parallel_state.get_context_parallel_world_size()
         if cp_size <= 1:
-            return (decoder_input, input_ids, labels, loss_mask, attention_mask, position_ids)
+            return (
+                decoder_input, input_ids, labels, loss_mask,
+                attention_mask, position_ids, padding_mask,
+            )
         cp_rank = parallel_state.get_context_parallel_rank()
 
         if packed_seq_params is not None:
@@ -233,6 +237,8 @@ def _cp_split_for_forward(
                 labels = labels.index_select(1, idx)
             if loss_mask is not None:
                 loss_mask = loss_mask.index_select(1, idx)
+            if padding_mask is not None:
+                padding_mask = padding_mask.index_select(1, idx)
         else:
 
             def _split(t, seq_dim):
@@ -247,8 +253,12 @@ def _split(t, seq_dim):
             labels = _split(labels, 1)
             loss_mask = _split(loss_mask, 1)
             attention_mask = _split(attention_mask, 1)
+            padding_mask = _split(padding_mask, 1)
 
-        return (decoder_input, input_ids, labels, loss_mask, attention_mask, position_ids)
+        return (
+            decoder_input, input_ids, labels, loss_mask,
+            attention_mask, position_ids, padding_mask,
+        )
 
     @staticmethod
     def cp_split_loss_mask(loss_mask, packed_seq_params):
@@ -301,6 +311,7 @@ def forward(
         attention_mask: Tensor = None,
         labels: Tensor = None,
         loss_mask: Tensor = None,
+        padding_mask: Tensor = None,
         pixel_values: Tensor = None,
         image_grid_thw: Tensor = None,
         decoder_input: Tensor = None,
@@ -316,6 +327,11 @@ def forward(
             attention_mask: ``[B, S]`` attention mask (None in THD).
             labels: ``[B, S]`` target token IDs (``[1, T]`` in THD).
             loss_mask: ``[B, S]`` mask for loss (``[1, T]`` in THD).
+            padding_mask: ``[B, S]`` bool mask, True at collate-padded
+                positions (``[1, T]`` in THD). Forwarded to the language
+                decoder so MoE routing excludes padded tokens from aux
+                loss / z-loss / expert-bias accumulation. Distinct from
+                ``loss_mask``: only true padding, not SFT prompt tokens.
             pixel_values: Preprocessed image pixels.
             image_grid_thw: ``[num_images, 3]`` grid dimensions.
             decoder_input: Pre-computed decoder input (skip embed).
@@ -345,16 +361,18 @@ def forward(
             else:
                 decoder_input = text_embeddings
 
-        (decoder_input, input_ids, labels, loss_mask, attention_mask, position_ids) = (
-            self._cp_split_for_forward(
-                decoder_input=decoder_input,
-                input_ids=input_ids,
-                labels=labels,
-                loss_mask=loss_mask,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                packed_seq_params=packed_seq_params,
-            )
+        (
+            decoder_input, input_ids, labels, loss_mask,
+            attention_mask, position_ids, padding_mask,
+        ) = self._cp_split_for_forward(
+            decoder_input=decoder_input,
+            input_ids=input_ids,
+            labels=labels,
+            loss_mask=loss_mask,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            packed_seq_params=packed_seq_params,
+            padding_mask=padding_mask,
         )
 
         with self._thd_mrope_no_cp_override(packed_seq_params):
@@ -365,5 +383,6 @@ def forward(
                 decoder_input=decoder_input,
                 labels=labels,
                 loss_mask=loss_mask,
+                padding_mask=padding_mask,
                 packed_seq_params=packed_seq_params,
             )