fix thd cp convergence issue

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/models/gemma_vl/modeling_gemma3_vl.py

@@ -161,28 +161,54 @@ def forward(
 
         # CP slicing: slice embeddings, labels, loss_mask, position_ids, and attention_mask
         # This must happen AFTER vision-text merge so image token positions are correct
-        if self.config._pg_collection.cp.size() > 1:
-            # inputs_embeds is (T, B, D), need to transpose to (B, T, D) for get_batch_on_this_cp_rank
+        cp_size = self.config._pg_collection.cp.size()
+        if cp_size > 1:
+            cp_rank = self.config._pg_collection.cp.rank()
+
+            # (T, B, D) -> (B, T, D) for slicing
             if inputs_embeds is not None:
                 inputs_embeds = inputs_embeds.transpose(0, 1).contiguous()
 
-            cp_group = self.config._pg_collection.cp
-            cp_batch = get_batch_on_this_cp_rank(
-                {
-                    "decoder_input": inputs_embeds,
-                    "labels": labels,
-                    "loss_mask": loss_mask,
-                    "position_ids": position_ids,
-                    "attention_mask": attention_mask,
-                },
-                cp_group=cp_group,
-            )
-
-            inputs_embeds = cp_batch.get("decoder_input")
-            labels = cp_batch.get("labels")
-            loss_mask = cp_batch.get("loss_mask")
-            position_ids = cp_batch.get("position_ids")
-            attention_mask = cp_batch.get("attention_mask")
+            if packed_seq_params is not None and packed_seq_params.qkv_format == "thd":
+                import transformer_engine_torch as tex
+
+                cu_seqlens = packed_seq_params.cu_seqlens_q
+                cu_seqlens_padded = (
+                    packed_seq_params.cu_seqlens_q_padded
+                    if packed_seq_params.cu_seqlens_q_padded is not None
+                    else cu_seqlens
+                )
+                seq_len = inputs_embeds.size(1)
+
+                index = tex.thd_get_partitioned_indices(cu_seqlens_padded, seq_len, cp_size, cp_rank)
+
+                # Slice all tensors using THD indices
+                if inputs_embeds is not None:
+                    inputs_embeds = inputs_embeds.index_select(1, index)
+                if labels is not None:
+                    labels = labels.index_select(1, index)
+                if loss_mask is not None:
+                    loss_mask = loss_mask.index_select(1, index)
+                if position_ids is not None:
+                    position_ids = position_ids.index_select(1, index)
+            else:
+                cp_group = self.config._pg_collection.cp
+                cp_batch = get_batch_on_this_cp_rank(
+                    {
+                        "decoder_input": inputs_embeds,
+                        "labels": labels,
+                        "loss_mask": loss_mask,
+                        "position_ids": position_ids,
+                        "attention_mask": attention_mask,
+                    },
+                    cp_group=cp_group,
+                )
+
+                inputs_embeds = cp_batch.get("decoder_input")
+                labels = cp_batch.get("labels")
+                loss_mask = cp_batch.get("loss_mask")
+                position_ids = cp_batch.get("position_ids")
+                attention_mask = cp_batch.get("attention_mask")
 
             # Transpose back to (T, B, D)
             if inputs_embeds is not None:
@@ -198,7 +224,8 @@ def forward(
             runtime_gather_output=runtime_gather_output,
             packed_seq_params=packed_seq_params,
         )
-        return outputs
+        # Return both outputs and the CP-sliced loss_mask for consistent loss computation
+        return (outputs, loss_mask)
 
     def freeze(self, freeze_language_model: bool, freeze_vision_model: bool, freeze_vision_projection: bool):
         """Freeze model modules.

src/megatron/bridge/training/vlm_step.py

@@ -19,7 +19,7 @@
 import torch
 from megatron.core.models.gpt import GPTModel
 from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage
-from megatron.core.utils import get_batch_on_this_cp_rank, get_model_config
+from megatron.core.utils import get_model_config
 
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.losses import (
@@ -345,11 +345,10 @@ def _ceil_to_mult(n: int, mult: int) -> int:
         cu_seqlens = None
         max_seqlen = None
 
-    cp_batch = get_batch_on_this_cp_rank({"loss_mask": batch.get("loss_mask")}, cp_group=pg_collection.cp)
     return (
         (batch.get("tokens") if batch.get("tokens") is not None else batch.get("input_ids")),
         batch.get("labels"),
-        cp_batch.get("loss_mask"),
+        batch.get("loss_mask"),  # Full packed loss_mask, will be CP-sliced by model
         batch.get("attention_mask"),
         batch.get("position_ids"),
         cu_seqlens,
@@ -379,6 +378,7 @@ def forward_step(
     use_mtp = (getattr(config, "mtp_num_layers", None) or 0) > 0
 
     timers("batch-generator", log_level=2).start()
+    pg_collection = get_pg_collection(model)
     with straggler_timer(bdata=True):
         (
             tokens,
@@ -389,14 +389,15 @@ def forward_step(
             cu_seqlens,
             max_seqlen,
             visual_inputs,
-        ) = get_batch(data_iterator, state.cfg, use_mtp, pg_collection=get_pg_collection(model))
+        ) = get_batch(data_iterator, state.cfg, use_mtp, pg_collection=pg_collection)
     timers("batch-generator").stop()
 
     forward_args = {
         "input_ids": tokens,
         "position_ids": position_ids,
         "attention_mask": attention_mask,
         "labels": labels,
+        "loss_mask": loss_mask,  # Pass full loss_mask so model can slice it consistently with labels
     }
 
     if visual_inputs is not None:
@@ -423,7 +424,12 @@ def forward_step(
             loss_function = _create_loss_function(loss_mask, check_for_nan_in_loss, check_for_spiky_loss)
             return schedule_plan, loss_function
         else:
-            output_tensor = model(**forward_args)
+            model_output = model(**forward_args)
+            # Handle tuple return: (output_tensor, sliced_loss_mask) from VLM models with CPI'm
+            if isinstance(model_output, tuple):
+                output_tensor, loss_mask = model_output
+            else:
+                output_tensor = model_output
 
     loss_function = _create_loss_function(loss_mask, check_for_nan_in_loss, check_for_spiky_loss)