[model] fix: Dequantize FP8 weights during Ministral3 import and fix model init issues (#2895)

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

Author: cuichenx

src/megatron/bridge/models/ministral3/ministral3_bridge.py

@@ -32,6 +32,10 @@
 Reference: https://huggingface.co/mistralai/Ministral-3-3B-Base-2512
 """
 
+from typing import Mapping, Union
+
+import torch
+
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
 from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
 from megatron.bridge.models.conversion.param_mapping import (
@@ -170,6 +174,51 @@ def mapping_registry(self) -> MegatronMappingRegistry:
 
         return MegatronMappingRegistry(*mapping_list)
 
+    def maybe_modify_loaded_hf_weight(
+        self,
+        hf_param: Union[str, dict[str, str]],
+        hf_state_dict: Mapping[str, torch.Tensor],
+    ) -> Union[torch.Tensor, dict[str, torch.Tensor]]:
+        """Load HF weights and dequantize FP8 tensors on the fly.
+
+        Ministral-3-*-Instruct-2512 stores LM weights in FP8 (float8_e4m3fn) with
+        separate ``weight_scale_inv`` scalar tensors.  The true bfloat16 weight is::
+
+            w_bf16 = fp8_weight.to(bfloat16) * weight_scale_inv
+
+        This override applies dequantization transparently so that the bridge produces
+        correct Megatron checkpoints without a separate preprocessing step.
+        """
+        hf_weights = super().maybe_modify_loaded_hf_weight(hf_param, hf_state_dict)
+
+        if isinstance(hf_weights, dict):
+            # Compound params (QKV / GatedMLP): dequantize each component individually
+            return {
+                key: self._maybe_dequantize_fp8(tensor, hf_param[key], hf_state_dict)
+                for key, tensor in hf_weights.items()
+            }
+        return self._maybe_dequantize_fp8(hf_weights, hf_param, hf_state_dict)
+
+    @staticmethod
+    def _maybe_dequantize_fp8(
+        weight: torch.Tensor,
+        param_name: str,
+        hf_state_dict: Mapping[str, torch.Tensor],
+    ) -> torch.Tensor:
+        """Dequantize *weight* if it is stored as FP8.
+
+        Looks up ``param_name + "_scale_inv"`` in *hf_state_dict* and applies::
+
+            w_bf16 = weight.to(bfloat16) * scale_inv
+        """
+        if weight.dtype != torch.float8_e4m3fn:
+            return weight
+        scale_key = param_name + "_scale_inv"
+        if scale_key not in hf_state_dict:
+            return weight.to(torch.bfloat16)
+        scale_inv = hf_state_dict[scale_key].to(torch.bfloat16)
+        return weight.to(torch.bfloat16) * scale_inv
+
 
 # Register the bridge if Mistral3ForConditionalGeneration is available
 if HAS_MISTRAL3 and Mistral3ForConditionalGeneration is not None:

src/megatron/bridge/models/ministral3/modeling_ministral3.py

@@ -140,11 +140,10 @@ def _apply_preserve_inv_freq(fn):
                     # Apply the transformation (e.g., bfloat16 conversion)
                     result = original_apply(fn)
 
-                    # Restore inv_freq to FP32 (only move device, keep dtype)
+                    # Restore inv_freq to FP32 but on the correct device
                     if inv_freq_backup is not None:
-                        # Get the new device from the transformation
-                        new_device = fn(torch.tensor([1.0])).device
-                        pos_emb.inv_freq.data = inv_freq_backup.to(device=new_device)
+                        target_device = pos_emb.inv_freq.data.device
+                        pos_emb.inv_freq.data = inv_freq_backup.to(device=target_device)
 
                     return result
 
@@ -177,6 +176,9 @@ def _apply_preserve_inv_freq(fn):
         # Some config requires from HF vision tower
         self.config.spatial_merge_size = getattr(self.config.hf_config, "spatial_merge_size", 2)
         self.config.vision_feature_layer = getattr(self.config.hf_config, "vision_feature_layer", -1)
+        # HF's get_image_features accesses self.config.return_dict
+        if not hasattr(self.config, "return_dict"):
+            self.config.return_dict = True
 
     def set_input_tensor(self, input_tensor) -> None:
         """Set model chunk input tensor."""