Merge pull request #139 from elder-plinius/feat/117-mistral-architecture-contracts

feat: add tested Mistral 3 and 4 contracts

ci/conditional-test-policy.json

@@ -13,6 +13,14 @@
       "cache": "GitHub Actions cache keyed by repository, revision, runner OS, and Python version",
       "timeout_minutes": 20
     },
+    "mistral4_config": {
+      "repository": "mistralai/Mistral-Small-4-119B-2603",
+      "revision": "a11f36bebf709121056b1dbcc943d1c6afbe494d",
+      "trust_remote_code": false,
+      "scope": "configuration metadata only; model weights are not downloaded",
+      "cache": "GitHub Actions cache keyed with the model-runtime lane",
+      "timeout_minutes": 5
+    },
     "cuda_torch": {
       "version_source": "Exact base version from the locked CPU Torch package",
       "backend": "cu130",
@@ -72,6 +80,7 @@
       "prerequisites": "public HTTPS access to the pinned tiny model",
       "expected_cost": "under 20 runner-minutes and 100 MB download",
       "coverage_paths": [
+        "obliteratus/architecture_profiles.py",
         "obliteratus/abliterate.py",
         "obliteratus/auto_obliterate.py",
         "obliteratus/bayesian_optimizer.py",
@@ -79,6 +88,7 @@
         "obliteratus/evaluation/evaluator.py",
         "obliteratus/informed_pipeline.py",
         "obliteratus/lora_ablation.py",
+        "obliteratus/models/loader.py",
         "obliteratus/sweep.py"
       ]
     },

ci/test-risk-map.json

@@ -128,6 +128,7 @@
         "tests/test_config_properties.py",
         "tests/test_device_boundaries.py",
         "tests/test_loader_boundaries.py",
+        "tests/test_mistral_architecture_contracts.py",
         "tests/test_quant_dequant.py",
         "tests/test_mlx_backend_boundaries.py",
         "tests/test_model_profile.py",
@@ -330,6 +331,7 @@
         "tests/test_strategies.py",
         "tests/test_strategy_navigation_contracts.py",
         "tests/test_gemma4_support.py",
+        "tests/test_mistral_architecture_contracts.py",
         "tests/test_study_presets.py"
       ]
     }
@@ -361,7 +363,22 @@
       "risk": "architecture detection and projection-path contracts",
       "required_tests": [
         "tests/test_architecture_profiles.py",
-        "tests/test_gemma4_support.py"
+        "tests/test_gemma4_support.py",
+        "tests/test_mistral_architecture_contracts.py"
+      ],
+      "conditional_gates": [
+        "model-download-runtime"
+      ]
+    },
+    {
+      "path": "obliteratus/strategies/utils.py",
+      "risk_class": "mixed-runtime",
+      "risk": "architecture-specific layer, attention, FFN, and embedding navigation",
+      "required_tests": [
+        "tests/test_gemma4_support.py",
+        "tests/test_mistral_architecture_contracts.py",
+        "tests/test_strategies.py",
+        "tests/test_strategy_navigation_contracts.py"
       ],
       "conditional_gates": []
     },
@@ -443,9 +460,11 @@
       "risk": "model loading, cache, architecture, quantization, and device-map boundaries",
       "required_tests": [
         "tests/test_loader_boundaries.py",
+        "tests/test_mistral_architecture_contracts.py",
         "tests/test_offline_integration.py"
       ],
       "conditional_gates": [
+        "model-download-runtime",
         "cuda-runtime",
         "bitsandbytes-runtime"
       ]

obliteratus/architecture_profiles.py

@@ -81,7 +81,7 @@ class ArchitectureProfile:
 _MOE_MODEL_TYPES = {
     "mixtral", "qwen2_moe", "qwen3_moe", "deepseek_v2", "deepseek_v3",
     "dbrx", "grok", "jamba", "arctic", "olmoe", "switch_transformers",
-    "nllb_moe", "llama4",
+    "nllb_moe", "llama4", "mistral4",
 }
 
 # Patterns in model name that indicate MoE (fallback when model_type is ambiguous)
@@ -91,6 +91,7 @@ _MOE_NAME_PATTERNS = [
     "gpt-oss", "kimi-k2", "glm-4.7",
     "step-3.5", "minimax-m2", "maverick", "scout",
     "mistral-large-3",
+    "mistral-small-4",
     "jamba", "olmoe", "arctic",
 ]
 
@@ -118,6 +119,7 @@ _LARGE_MOE_NAME_PATTERNS = [
     "kimi-k2",         # Kimi K2 (1T total)
     "-A22B",           # Qwen3-235B-A22B
     "mistral-large-3", # Mistral Large 3 (675B total)
+    "mistral-small-4", # Mistral Small 4 (119B total)
     "step-3.5",        # Step-3.5 Flash (large MoE)
     "minimax-m2",      # MiniMax-M2 (large MoE)
 ]
@@ -174,16 +176,22 @@ def detect_architecture(
 
     # ── Step 1: Extract info from config if available ────────────────
     if config is not None:
-        model_type = getattr(config, "model_type", "")
+        # Mistral 3 is a multimodal wrapper whose text_config may be either a
+        # dense Mistral backbone or a Mistral 4 MoE backbone.  Profile the text
+        # config when present without replacing the outer config used by the
+        # loader, tokenizer, and save path.
+        text_config = getattr(config, "text_config", None)
+        profile_config = text_config or config
+        model_type = getattr(profile_config, "model_type", "")
 
         # Check for MoE via config attributes
         for attr in _TOTAL_EXPERT_ATTRS:
-            val = getattr(config, attr, None)
+            val = getattr(profile_config, attr, None)
             if val is not None and val > 0:
                 is_moe = True
                 num_experts = max(num_experts, val)
         for attr in _ACTIVE_EXPERT_ATTRS:
-            val = getattr(config, attr, None)
+            val = getattr(profile_config, attr, None)
             if val is not None and val > 0:
                 is_moe = True
                 num_active_experts = max(num_active_experts, val)
@@ -194,20 +202,24 @@ def detect_architecture(
 
         # Extract layer/hidden info from config if not provided
         if num_layers == 0:
-            num_layers = getattr(config, "num_hidden_layers", 0)
+            num_layers = getattr(profile_config, "num_hidden_layers", 0)
         if hidden_size == 0:
-            hidden_size = getattr(config, "hidden_size", 0)
+            hidden_size = getattr(profile_config, "hidden_size", 0)
 
         # Rough param estimation
-        intermediate = getattr(config, "intermediate_size", hidden_size * 4)
-        vocab = getattr(config, "vocab_size", 32000)
+        intermediate = getattr(profile_config, "intermediate_size", hidden_size * 4)
+        moe_intermediate = getattr(profile_config, "moe_intermediate_size", intermediate)
+        vocab = getattr(profile_config, "vocab_size", 32000)
         if num_layers > 0 and hidden_size > 0:
             per_layer = 4 * hidden_size * hidden_size + 3 * hidden_size * intermediate
             if is_moe and num_experts > 0:
-                # MoE: multiply FFN part by num_experts
-                ffn_part = 3 * hidden_size * intermediate
+                # Mistral 4 and similar models expose a smaller routed-expert
+                # width separately from the shared/dense FFN width.
+                ffn_part = 3 * hidden_size * moe_intermediate
                 attn_part = 4 * hidden_size * hidden_size
                 per_layer = attn_part + ffn_part * num_experts
+                if getattr(profile_config, "n_shared_experts", 0) > 0:
+                    per_layer += 3 * hidden_size * intermediate
             embedding = 2 * vocab * hidden_size
             total_params_b = (per_layer * num_layers + embedding) / 1e9
 

obliteratus/models/loader.py

@@ -393,6 +393,15 @@ TASK_MODEL_MAP = {
 _IMAGE_TEXT_MODEL_TYPES = {
     "gemma4_unified",
     "gemma4",
+    "mistral3",
+    "mistral4",
+}
+
+_IMAGE_TEXT_ARCHITECTURES = {
+    "Gemma4ForConditionalGeneration",
+    "Gemma4UnifiedForConditionalGeneration",
+    "Mistral3ForConditionalGeneration",
+    "Mistral4ForCausalLM",
 }
 
 
@@ -405,14 +414,16 @@ def _select_model_class(task: str, config: AutoConfig):
     architectures = tuple(getattr(config, "architectures", None) or ())
     is_image_text = (
         model_type in _IMAGE_TEXT_MODEL_TYPES
-        or any("ForConditionalGeneration" in arch for arch in architectures)
-        and any("Gemma4" in arch for arch in architectures)
+        or any(arch in _IMAGE_TEXT_ARCHITECTURES for arch in architectures)
     )
     if task == "causal_lm" and is_image_text:
         if AutoModelForImageTextToText is None:
+            architecture_name = model_type or (architectures[0] if architectures else "unknown")
             raise RuntimeError(
-                "AutoModelForImageTextToText is required for Gemma 4 unified models. "
-                "Upgrade transformers to a version that provides it."
+                "AutoModelForImageTextToText is required for architecture "
+                f"{architecture_name!r}. "
+                "Upgrade transformers to a version that provides the matching "
+                "image-text model mapping."
             )
         return AutoModelForImageTextToText
 
@@ -545,11 +556,21 @@ def _estimate_model_memory_gb(config: AutoConfig, dtype: torch.dtype) -> float:
     if hidden == 0 or n_layers == 0:
         return 0.0
 
-    # For MoE models, the FFN is replicated per expert
-    num_experts = getattr(config, "num_local_experts", None) or getattr(config, "num_experts", 1)
+    profile_config = text_cfg or config
+    # For MoE models, the routed FFN is replicated per expert. Mistral 4 uses
+    # n_routed_experts plus a separate routed-expert width.
+    num_experts = (
+        getattr(profile_config, "num_local_experts", None)
+        or getattr(profile_config, "num_experts", None)
+        or getattr(profile_config, "n_routed_experts", 1)
+    )
+    moe_intermediate = getattr(profile_config, "moe_intermediate_size", intermediate)
 
-    # Per layer: attn (4 * hidden^2) + ffn (3 * hidden * intermediate * num_experts) + norms
-    per_layer = 4 * hidden * hidden + num_experts * 3 * hidden * intermediate
+    # Per layer: attention, routed experts, and any separately configured
+    # shared expert. Norms are negligible at this estimation scale.
+    per_layer = 4 * hidden * hidden + num_experts * 3 * hidden * moe_intermediate
+    if getattr(profile_config, "n_shared_experts", 0) > 0:
+        per_layer += 3 * hidden * intermediate
     # Embedding + LM head
     embedding = 2 * vocab * hidden
     total_params = per_layer * n_layers + embedding

obliteratus/strategies/utils.py

@@ -13,6 +13,8 @@ _LAYER_ATTR_PATHS: dict[str, list[str]] = {
     "gpt_neox": ["gpt_neox", "layers"],
     "llama": ["model", "layers"],
     "mistral": ["model", "layers"],
+    "mistral3": ["model", "language_model", "layers"],
+    "mistral4": ["model", "layers"],
     "gemma": ["model", "layers"],
     "gemma2": ["model", "layers"],
     "phi": ["model", "layers"],
@@ -50,12 +52,20 @@ _LAYER_ATTR_PATHS: dict[str, list[str]] = {
     "nemotron_h": ["backbone", "layers"],
 }
 
+_LAYER_ATTR_PATH_FALLBACKS: dict[str, list[list[str]]] = {
+    # Mistral3Model (without the conditional-generation head) exposes the
+    # same text backbone one level closer to the root.
+    "mistral3": [["language_model", "layers"]],
+}
+
 _ATTENTION_ATTR: dict[str, str] = {
     "gpt2": "attn",
     "gpt_neo": "attn.attention",
     "gpt_neox": "attention",
     "llama": "self_attn",
     "mistral": "self_attn",
+    "mistral3": "self_attn",
+    "mistral4": "self_attn",
     "gemma": "self_attn",
     "gemma2": "self_attn",
     "phi": "self_attn",
@@ -102,6 +112,8 @@ _FFN_ATTR: dict[str, str] = {
     "gpt_neox": "mlp",
     "llama": "mlp",
     "mistral": "mlp",
+    "mistral3": "mlp",
+    "mistral4": "mlp",
     "gemma": "mlp",
     "gemma2": "mlp",
     "phi": "mlp",
@@ -168,10 +180,21 @@ def get_layer_modules(handle: ModelHandle) -> nn.ModuleList:
     """Return the nn.ModuleList of transformer layers for this model."""
     arch = handle.architecture
     if arch in _LAYER_ATTR_PATHS:
-        obj = handle.model
-        for attr in _LAYER_ATTR_PATHS[arch]:
-            obj = getattr(obj, attr)
-        return obj
+        paths = [_LAYER_ATTR_PATHS[arch], *_LAYER_ATTR_PATH_FALLBACKS.get(arch, [])]
+        for path in paths:
+            obj = handle.model
+            try:
+                for attr in path:
+                    obj = getattr(obj, attr)
+            except AttributeError:
+                continue
+            if isinstance(obj, nn.ModuleList):
+                return obj
+        attempted = [".".join(path) for path in paths]
+        raise RuntimeError(
+            f"Cannot locate transformer layers for known architecture {arch!r}; "
+            f"expected one of {attempted}."
+        )
 
     # Fallback: walk the model looking for a ModuleList with the right length.
     # If num_layers is known, match exactly; otherwise find the largest ModuleList
@@ -227,6 +250,8 @@ def get_embedding_module(handle: ModelHandle) -> nn.Embedding:
     # Try common paths
     for path in [
         "transformer.wte",
+        "model.language_model.embed_tokens",
+        "language_model.embed_tokens",
         "model.embed_tokens",
         "gpt_neox.embed_in",
         "model.decoder.embed_tokens",

tests/conditional/test_model_download_runtime.py

@@ -7,12 +7,16 @@ import uuid
 import pytest
 import torch
 
-from obliteratus.models.loader import load_model
+from obliteratus.architecture_profiles import ArchitectureClass, detect_architecture
+from obliteratus.models.loader import _select_model_class, load_model
+from transformers import AutoConfig, AutoModelForImageTextToText
 
 
 pytestmark = [pytest.mark.network, pytest.mark.download]
 MODEL = "hf-internal-testing/tiny-random-gpt2"
 REVISION = "71034c5d8bde858ff824298bdedc65515b97d2b9"
+MISTRAL4_MODEL = "mistralai/Mistral-Small-4-119B-2603"
+MISTRAL4_REVISION = "a11f36bebf709121056b1dbcc943d1c6afbe494d"
 
 
 def test_pinned_tiny_model_download_inference_and_offline_cache(monkeypatch):
@@ -47,3 +51,23 @@ def test_pinned_tiny_model_download_inference_and_offline_cache(monkeypatch):
     missing = f"obliteratus/offline-missing-{uuid.uuid4().hex}"
     with pytest.raises(OSError):
         load_model(missing, revision=REVISION, device="cpu", local_files_only=True)
+
+
+def test_pinned_mistral4_config_resolves_composite_contract_without_remote_code():
+    config = AutoConfig.from_pretrained(
+        MISTRAL4_MODEL,
+        revision=MISTRAL4_REVISION,
+        trust_remote_code=False,
+    )
+
+    assert config.model_type == "mistral3"
+    assert config.architectures == ["Mistral3ForConditionalGeneration"]
+    assert config.text_config.model_type == "mistral4"
+    assert config.text_config.n_routed_experts == 128
+    assert config.text_config.num_experts_per_tok == 4
+    assert _select_model_class("causal_lm", config) is AutoModelForImageTextToText
+
+    profile = detect_architecture(MISTRAL4_MODEL, config=config)
+    assert profile.model_type == "mistral4"
+    assert profile.arch_class is ArchitectureClass.LARGE_MOE
+    assert (profile.num_experts, profile.num_active_experts) == (128, 4)

tests/test_loader_boundaries.py

@@ -492,6 +492,26 @@ def test_task_model_selection_and_gemma_contract(monkeypatch):
         loader._select_model_class("embedding", _config())
 
 
+def test_mistral3_load_uses_image_text_class_and_keeps_remote_code_disabled(
+    loader_boundary, monkeypatch,
+):
+    loader_boundary.config.model_type = "mistral3"
+    loader_boundary.config.architectures = ["Mistral3ForConditionalGeneration"]
+    image_text = SimpleNamespace(
+        from_pretrained=Mock(return_value=loader_boundary.model),
+    )
+    monkeypatch.setattr(loader, "AutoModelForImageTextToText", image_text)
+
+    handle = loader.load_model("local/mistral3", skip_snapshot=True)
+
+    assert handle.config is loader_boundary.config
+    assert handle.architecture == "mistral3"
+    assert loader.AutoConfig.from_pretrained.call_args.kwargs["trust_remote_code"] is False
+    assert image_text.from_pretrained.call_args.kwargs["trust_remote_code"] is False
+    assert loader.AutoTokenizer.from_pretrained.call_args.kwargs["trust_remote_code"] is False
+    loader_boundary.model_class.from_pretrained.assert_not_called()
+
+
 def test_model_handle_metadata_snapshot_restore_summary_and_cleanup(tmp_path):
     model = _model()
     nested = SimpleNamespace(

tests/test_mistral_architecture_contracts.py

@@ -0,0 +1,200 @@
+"""Offline Mistral 3/4 architecture and loader contracts."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from obliteratus.architecture_profiles import ArchitectureClass, detect_architecture
+from obliteratus.models import loader
+from obliteratus.models.loader import ModelHandle
+from obliteratus.strategies.utils import (
+    get_attention_module,
+    get_embedding_module,
+    get_ffn_module,
+    get_layer_modules,
+)
+
+
+class _MistralLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.self_attn = nn.Module()
+        self.mlp = nn.Module()
+
+
+class _Mistral3ConditionalModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Module()
+        self.model.language_model = nn.Module()
+        self.model.language_model.layers = nn.ModuleList(
+            [_MistralLayer(), _MistralLayer()]
+        )
+        self.model.language_model.embed_tokens = nn.Embedding(32, 16)
+
+
+class _BareMistral3Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.language_model = nn.Module()
+        self.language_model.layers = nn.ModuleList(
+            [_MistralLayer(), _MistralLayer()]
+        )
+        self.language_model.embed_tokens = nn.Embedding(32, 16)
+
+
+class _Mistral4CausalModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Module()
+        self.model.layers = nn.ModuleList([_MistralLayer(), _MistralLayer()])
+        self.model.embed_tokens = nn.Embedding(32, 16)
+
+
+def _mistral4_text_config():
+    return SimpleNamespace(
+        model_type="mistral4",
+        num_hidden_layers=36,
+        num_attention_heads=32,
+        hidden_size=4096,
+        intermediate_size=12288,
+        moe_intermediate_size=2048,
+        vocab_size=131072,
+        n_routed_experts=128,
+        n_shared_experts=1,
+        num_experts_per_tok=4,
+    )
+
+
+def _mistral3_config(text_config=None):
+    return SimpleNamespace(
+        model_type="mistral3",
+        architectures=["Mistral3ForConditionalGeneration"],
+        text_config=text_config or SimpleNamespace(
+            model_type="mistral",
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            hidden_size=16,
+            intermediate_size=64,
+            vocab_size=32,
+        ),
+    )
+
+
+def _handle(model: nn.Module, config, tokenizer=None) -> ModelHandle:
+    return ModelHandle(
+        model=model,
+        tokenizer=tokenizer or SimpleNamespace(pad_token="<pad>", eos_token="<eos>"),
+        config=config,
+        model_name="mistralai/synthetic-mistral",
+        task="causal_lm",
+    )
+
+
+def test_mistral_loader_uses_only_verified_image_text_mappings(monkeypatch):
+    causal = object()
+    classification = object()
+    image_text = object()
+    monkeypatch.setitem(loader.TASK_MODEL_MAP, "causal_lm", causal)
+    monkeypatch.setitem(loader.TASK_MODEL_MAP, "classification", classification)
+    monkeypatch.setattr(loader, "AutoModelForImageTextToText", image_text)
+
+    assert loader._select_model_class("causal_lm", _mistral3_config()) is image_text
+    assert loader._select_model_class(
+        "causal_lm",
+        SimpleNamespace(model_type="mistral4", architectures=["Mistral4ForCausalLM"]),
+    ) is image_text
+    assert loader._select_model_class("classification", _mistral3_config()) is classification
+    assert loader._select_model_class(
+        "causal_lm",
+        SimpleNamespace(model_type="unknown", architectures=["OtherForConditionalGeneration"]),
+    ) is causal
+
+
+def test_mistral_loader_fails_when_required_transformers_mapping_is_missing(monkeypatch):
+    monkeypatch.setattr(loader, "AutoModelForImageTextToText", None)
+
+    with pytest.raises(
+        RuntimeError,
+        match=r"AutoModelForImageTextToText.*mistral3.*Upgrade transformers",
+    ):
+        loader._select_model_class("causal_lm", _mistral3_config())
+
+
+def test_composite_profile_uses_text_backbone_without_misclassifying_mistral3():
+    dense = detect_architecture(
+        "mistralai/Mistral-Small-3.1-24B-Instruct-2503",
+        config=_mistral3_config(),
+    )
+    moe = detect_architecture(
+        "mistralai/Mistral-Small-4-119B-2603",
+        config=_mistral3_config(_mistral4_text_config()),
+    )
+
+    assert dense.model_type == "mistral"
+    assert dense.arch_class is ArchitectureClass.DENSE
+    assert not dense.is_moe
+    assert moe.model_type == "mistral4"
+    assert moe.arch_class is ArchitectureClass.LARGE_MOE
+    assert moe.is_moe
+    assert (moe.num_experts, moe.num_active_experts) == (128, 4)
+    assert moe.total_params_b >= 100
+
+
+def test_composite_memory_estimate_counts_routed_and_shared_experts():
+    config = _mistral3_config(_mistral4_text_config())
+
+    estimate_gb = loader._estimate_model_memory_gb(config, torch.bfloat16)
+
+    assert 200 < estimate_gb < 300
+
+
+def test_mistral_small_4_name_fallback_is_large_moe():
+    profile = detect_architecture("mistralai/Mistral-Small-4-119B-2603")
+
+    assert profile.arch_class is ArchitectureClass.LARGE_MOE
+    assert profile.is_moe
+
+
+@pytest.mark.parametrize(
+    "model",
+    [_Mistral3ConditionalModel(), _BareMistral3Model()],
+)
+def test_mistral3_navigation_preserves_outer_config_and_tokenizer(model):
+    config = _mistral3_config(_mistral4_text_config())
+    tokenizer = SimpleNamespace(pad_token="<pad>", eos_token="<eos>")
+    handle = _handle(model, config, tokenizer)
+    layers = get_layer_modules(handle)
+
+    assert handle.config is config
+    assert handle.tokenizer is tokenizer
+    assert handle.architecture == "mistral3"
+    assert (handle.num_layers, handle.num_heads, handle.hidden_size) == (36, 32, 4096)
+    assert len(layers) == 2
+    assert get_attention_module(layers[0], handle.architecture) is layers[0].self_attn
+    assert get_ffn_module(layers[0], handle.architecture) is layers[0].mlp
+    assert get_embedding_module(handle).embedding_dim == 16
+
+
+def test_direct_mistral4_navigation_uses_causal_wrapper_layout():
+    config = _mistral4_text_config()
+    handle = _handle(_Mistral4CausalModel(), config)
+    layers = get_layer_modules(handle)
+
+    assert handle.architecture == "mistral4"
+    assert len(layers) == 2
+    assert get_attention_module(layers[0], handle.architecture) is layers[0].self_attn
+    assert get_ffn_module(layers[0], handle.architecture) is layers[0].mlp
+    assert get_embedding_module(handle).num_embeddings == 32
+
+
+def test_known_mistral3_layout_mismatch_fails_with_attempted_paths():
+    with pytest.raises(
+        RuntimeError,
+        match=r"known architecture 'mistral3'.*model.language_model.layers.*language_model.layers",
+    ):
+        get_layer_modules(_handle(nn.Module(), _mistral3_config()))