fix(mistral3): preserve medium VLM checkpoint layout

nemo_automodel/components/models/mistral3_vlm/model.py

@@ -152,7 +152,7 @@ def __init__(self, config: PretrainedConfig):
                 except AttributeError:
                     pass
         super().__init__(config)
-        self.state_dict_adapter = Mistral3FP8StateDictAdapter.for_vlm_full()
+        self.state_dict_adapter = Mistral3FP8StateDictAdapter.for_vlm_full(config)
 
         # Lazy non-persistent buffer reinit. HF's Ministral3RotaryEmbedding /
         # PixtralRotaryEmbedding compute `inv_freq` in their __init__. Under

nemo_automodel/components/models/mistral3_vlm/state_dict_adapter.py

@@ -96,6 +96,28 @@ def _identity(k: str) -> str:
     return k
 
 
+_MISTRAL3P5_128B_NUM_HIDDEN_LAYERS = 88
+
+
+def _config_attr(config: Any | None, attr: str) -> Any:
+    if isinstance(config, dict):
+        return config.get(attr)
+    return getattr(config, attr, None)
+
+
+def _is_mistral3p5_128b_config(config: Any | None) -> bool:
+    text_config = _config_attr(config, "text_config")
+    return (
+        _config_attr(text_config, "model_type") == "ministral3"
+        and _config_attr(text_config, "num_hidden_layers") == _MISTRAL3P5_128B_NUM_HIDDEN_LAYERS
+    )
+
+
+def _uses_identity_vlm_layout(config: Any | None) -> bool:
+    """Return True for FP8 VLM checkpoints whose disk keys already match HF."""
+    return _is_mistral3p5_128b_config(config)
+
+
 # The runtime ``Mistral3ForConditionalGeneration`` puts body modules under
 # ``model.*`` while the on-disk checkpoint stores text weights under
 # ``language_model.model.*`` and non-text VLM components at top level.
@@ -161,14 +183,19 @@ def __init__(
         self._not_fp8_prefixes = tuple(not_fp8_prefixes)
 
     @classmethod
-    def for_vlm_full(cls) -> "Mistral3FP8StateDictAdapter":
+    def for_vlm_full(cls, config: Any | None = None) -> "Mistral3FP8StateDictAdapter":
         """Full-VLM path for Mistral3ForConditionalGeneration checkpoints.
 
-        The runtime module keeps VLM body modules under ``model.*`` but the
-        checkpoint stores text weights under ``language_model.model.*`` and
-        non-text component names at top level. The **LM head** has one extra
-        quirk: the model exposes it at the top level (``lm_head.weight``) while
-        the checkpoint nests it (``language_model.lm_head.weight``).
+        Mistral3 FP8 VLM checkpoints have two observed body-key layouts. The
+        Mistral-Medium-3.5 128B checkpoint already stores keys in the same
+        layout as HF's VLM ``state_dict()`` (``model.language_model.*`` /
+        ``model.vision_tower.*`` / ``model.multi_modal_projector.*``). Newer
+        Ministral/Devstral-style checkpoints store text weights under
+        ``language_model.model.*`` and non-text component names at top level.
+
+        The **LM head** has one extra quirk in the nested layout: the model
+        exposes it at the top level (``lm_head.weight``) while the checkpoint
+        nests it (``language_model.lm_head.weight``).
         Tied checkpoints (Ministral-3) never serialize the head, so the head
         translation is a harmless no-op there; untied checkpoints (Devstral-24B)
         rely on it to find the head during the DCP load.
@@ -183,6 +210,8 @@ def for_vlm_full(cls) -> "Mistral3FP8StateDictAdapter":
             "model.multi_modal_projector",
             # "lm_head" already in _NON_QUANTIZED_SUFFIXES via suffix match.
         )
+        if _uses_identity_vlm_layout(config):
+            return cls(layout_name="vlm_full_identity", not_fp8_prefixes=not_fp8)
         return cls(
             native_to_hf=_vlm_full_native_to_hf,
             hf_to_native=_vlm_full_hf_to_native,

tests/unit_tests/models/mistral3_vlm/test_model.py

@@ -153,6 +153,16 @@ def test_state_dict_adapter_for_vlm_full(self, patched_super_init):
         assert isinstance(m.state_dict_adapter, Mistral3FP8StateDictAdapter)
         assert m.state_dict_adapter._layout_name == "vlm_full"
 
+    def test_state_dict_adapter_uses_identity_layout_for_mistral_medium_35(self, patched_super_init):
+        qc = {"quant_method": "fp8"}
+        cfg = SimpleNamespace(
+            text_config=SimpleNamespace(model_type="ministral3", num_hidden_layers=88),
+            quantization_config=qc,
+        )
+        m = Mistral3FP8VLMForConditionalGeneration(cfg)
+        assert isinstance(m.state_dict_adapter, Mistral3FP8StateDictAdapter)
+        assert m.state_dict_adapter._layout_name == "vlm_full_identity"
+
 
 class TestInitRegistersRotaryHooks:
     def test_hook_registered_on_modules_with_inv_freq(self, patched_super_init):

tests/unit_tests/models/mistral3_vlm/test_state_dict_adapter.py

@@ -14,6 +14,8 @@
 
 """Unit tests for the Mistral3 FP8 VLM state-dict adapter."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -164,6 +166,29 @@ def test_lm_head_key_is_remapped_both_ways(self):
         # Round-trips back to the model name.
         assert a._hf_to_native(a._native_to_hf("lm_head.weight")) == "lm_head.weight"
 
+    def test_mistral_medium_35_uses_identity_body_layout(self):
+        # Mistral-Medium-3.5 128B stores full-VLM keys in the same layout as
+        # HF's runtime state_dict. Remapping the body to language_model.model.*
+        # makes DCP request keys that are absent from that checkpoint.
+        cfg = SimpleNamespace(text_config=SimpleNamespace(model_type="ministral3", num_hidden_layers=88))
+        a = Mistral3FP8StateDictAdapter.for_vlm_full(cfg)
+        assert a._layout_name == "vlm_full_identity"
+        for key in (
+            "model.language_model.embed_tokens.weight",
+            "model.language_model.layers.0.self_attn.q_proj.weight",
+            "model.vision_tower.patch_conv.weight",
+            "model.multi_modal_projector.linear_1.weight",
+            "lm_head.weight",
+        ):
+            assert a._native_to_hf(key) == key
+            assert a._hf_to_native(key) == key
+
+    def test_smaller_mistral3_configs_keep_nested_body_layout(self):
+        cfg = SimpleNamespace(text_config=SimpleNamespace(model_type="ministral3", num_hidden_layers=36))
+        a = Mistral3FP8StateDictAdapter.for_vlm_full(cfg)
+        assert a._layout_name == "vlm_full"
+        assert a._native_to_hf("model.language_model.embed_tokens.weight") == "language_model.model.embed_tokens.weight"
+
 
 # --------------------------------------------------------------------------- #
 # from_hf                                                                     #
@@ -311,6 +336,16 @@ def test_lm_head_remapped_to_nested_name_no_scale(self):
         # Original dtype preserved (no FP8 cast for the non-quantized head).
         assert out["language_model.lm_head.weight"].dtype == torch.bfloat16
 
+    def test_mistral_medium_35_quantization_keeps_identity_body_keys(self):
+        a = Mistral3FP8StateDictAdapter.for_vlm_full(
+            {"text_config": {"model_type": "ministral3", "num_hidden_layers": 88}}
+        )
+        w_key = "model.language_model.layers.0.self_attn.q_proj.weight"
+        out = a.to_hf({w_key: torch.zeros(2, 2, dtype=torch.bfloat16)}, quantization=True)
+        assert w_key in out
+        assert "language_model.model.layers.0.self_attn.q_proj.weight" not in out
+        assert w_key + "_scale_inv" in out
+
     def test_exclude_key_regex(self):
         a = self._adapter()
         out = a.to_hf(