fix(mtp): skip disabled loss and handle empty targets

modeling_bailing_moe_v3.py

@@ -1552,7 +1552,9 @@ class BailingMoeV3ForCausalLM(BailingMoeV3PreTrainedModel, GenerationMixin):
             output_router_logits if output_router_logits is not None else self.config.output_router_logits
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        compute_mtp_loss = labels is not None and self.num_nextn_predict_layers > 0
+        compute_mtp_loss = (
+            labels is not None and self.num_nextn_predict_layers > 0 and self.mtp_loss_scaling_factor > 0
+        )
         # MTP hidden states are kept outside the tuple output, so retain the structured output while
         # computing its loss.
         model_return_dict = return_dict or compute_mtp_loss
@@ -1584,10 +1586,8 @@ class BailingMoeV3ForCausalLM(BailingMoeV3PreTrainedModel, GenerationMixin):
 
         all_mtp_logits = None
         if compute_mtp_loss and outputs.mtp_hidden_states is not None:
-            mtp_hidden_states = outputs.mtp_hidden_states
             shift_labels_mtp = None
-            for i in range(self.num_nextn_predict_layers):
-                mtp_hidden_states = mtp_hidden_states[i]
+            for mtp_hidden_states in outputs.mtp_hidden_states:
                 mtp_logits = self.lm_head(mtp_hidden_states).float()
                 if all_mtp_logits is None:
                     all_mtp_logits = []
@@ -1597,8 +1597,20 @@ class BailingMoeV3ForCausalLM(BailingMoeV3PreTrainedModel, GenerationMixin):
                         shift_labels_mtp = labels.clone()
                     shift_labels_mtp, _ = roll_tensor(shift_labels_mtp, shifts=-1, dims=-1, fill_value=-100)
                     mtp_logits_ = mtp_logits.view(-1, self.config.vocab_size)
+                    mtp_loss_kwargs = dict(kwargs)
+                    if mtp_loss_kwargs.get("num_items_in_batch") is None:
+                        loss_labels = mtp_loss_kwargs.get("shift_labels")
+                        if loss_labels is None:
+                            loss_labels = shift_labels_mtp.view(-1)[1:]
+                        # Avoid a mean over zero targets after the causal shift.
+                        mtp_loss_kwargs["num_items_in_batch"] = (
+                            (loss_labels != mtp_loss_kwargs.get("ignore_index", -100)).sum().clamp_min(1)
+                        )
                     mtp_loss = self.loss_function(
-                        mtp_logits_, shift_labels_mtp.to(mtp_logits_.device).view(-1), self.config.vocab_size, **kwargs
+                        mtp_logits_,
+                        shift_labels_mtp.to(mtp_logits_.device).view(-1),
+                        self.config.vocab_size,
+                        **mtp_loss_kwargs,
                     )
                     if loss is not None:
                         loss += self.mtp_loss_scaling_factor * mtp_loss