Fix FP8 exclusions and MTP cache handling

config.json

@@ -130,8 +130,7 @@
     "modules_to_not_convert": [
       "lm_head",
       "model.word_embeddings",
-      "b_proj",
-      "mlp.gate",
+      "attention.b_proj",
       "model.layers.7.attention.g_proj",
       "model.layers.15.attention.g_proj",
       "model.layers.23.attention.g_proj",
@@ -251,4 +250,4 @@
   "vocab_size": 157184,
   "model_type": "bailing_hybrid",
   "torch_dtype": "bfloat16"
-}
+}

modeling_bailing_moe_v3.py

@@ -1272,6 +1272,7 @@ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         output_router_logits: Optional[bool] = None,
+        output_mtp_hidden_states: Optional[bool] = False,
         return_dict: Optional[bool] = None,
         **kwargs,
     ) -> Union[Tuple, MoeV3ModelOutputWithPast]:
@@ -1412,7 +1413,7 @@ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
 
         mtp_hidden_states = None
 
-        if mtp_layers:
+        if mtp_layers and output_mtp_hidden_states:
             for decoder_layer in mtp_layers:
                 input_ids, _ = roll_tensor(input_ids, shifts=-1, dims=-1)
                 inputs_embeds = self.word_embeddings(input_ids)
@@ -1424,10 +1425,10 @@ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
                         hidden_states,
                         attention_mask,
                         position_ids,
-                        past_key_values,
+                        None,
                         output_attentions,
                         output_router_logits,
-                        use_cache,
+                        False,
                         position_embeddings,
                     )
                 else:
@@ -1436,10 +1437,10 @@ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
                         hidden_states,
                         attention_mask=attention_mask,
                         position_ids=position_ids,
-                        past_key_value=past_key_values,
+                        past_key_value=None,
                         output_attentions=output_attentions,
                         output_router_logits=output_router_logits,
-                        use_cache=use_cache,
+                        use_cache=False,
                         position_embeddings=position_embeddings,
                     )
                 if mtp_hidden_states is None:
@@ -1450,9 +1451,6 @@ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
                 if output_hidden_states:
                     all_hidden_states += (hidden_states,)
 
-                if use_cache:
-                    next_decoder_cache = layer_outputs[2 if output_attentions else 1]
-
                 if output_attentions:
                     all_self_attns += (layer_outputs[1],)
 
@@ -1554,6 +1552,10 @@ 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
+        # 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
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         outputs = self.model(
             input_ids=input_ids,
@@ -1565,7 +1567,8 @@ class BailingMoeV3ForCausalLM(BailingMoeV3PreTrainedModel, GenerationMixin):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             output_router_logits=output_router_logits,
-            return_dict=return_dict,
+            output_mtp_hidden_states=compute_mtp_loss,
+            return_dict=model_return_dict,
             **kwargs,
         )
 
@@ -1580,7 +1583,7 @@ class BailingMoeV3ForCausalLM(BailingMoeV3PreTrainedModel, GenerationMixin):
             loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs)
 
         all_mtp_logits = None
-        if self.num_nextn_predict_layers > 0:
+        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):