diff --git a/ci/conditional-test-policy.json b/ci/conditional-test-policy.json index f5e5b9c..0394c84 100644 --- a/ci/conditional-test-policy.json +++ b/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" ] }, diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 85c3fd5..289ef64 100644 --- a/ci/test-risk-map.json +++ b/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" ] diff --git a/obliteratus/architecture_profiles.py b/obliteratus/architecture_profiles.py index d835dba..1cb260c 100644 --- a/obliteratus/architecture_profiles.py +++ b/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 diff --git a/obliteratus/models/loader.py b/obliteratus/models/loader.py index b0eb208..c30c2c6 100644 --- a/obliteratus/models/loader.py +++ b/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 diff --git a/obliteratus/strategies/utils.py b/obliteratus/strategies/utils.py index 7e22e52..44acffa 100644 --- a/obliteratus/strategies/utils.py +++ b/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", diff --git a/tests/conditional/test_model_download_runtime.py b/tests/conditional/test_model_download_runtime.py index 223ad26..699de55 100644 --- a/tests/conditional/test_model_download_runtime.py +++ b/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) diff --git a/tests/test_loader_boundaries.py b/tests/test_loader_boundaries.py index c1f2871..0d60b43 100644 --- a/tests/test_loader_boundaries.py +++ b/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( diff --git a/tests/test_mistral_architecture_contracts.py b/tests/test_mistral_architecture_contracts.py new file mode 100644 index 0000000..264b3ab --- /dev/null +++ b/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="", eos_token=""), + 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="", eos_token="") + 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()))