fix(loading): route Mistral3-VLM mistral4 (MLA) text backbone to the Mistral4 loader (#423)#424
Conversation
…Mistral4 loader (#423) build_pixtral_family_context unconditionally parsed the VLM text backbone as models::llama3::ModelArgs and loaded it via Llama3Model::from_weights, which expects standard attention weights (q_proj). A Mistral3ForConditionalGeneration checkpoint whose text_config.model_type is mistral4 has an MLA + MoE backbone (q_a_proj / kv_a_proj_with_mqa / kv_b_proj, no q_proj), so the load failed with "Weight not found: model.layers.0.self_attn.q_proj.weight". This blocked loading mlx-community/Mistral-Small-4-119B-2603-4bit, the only public mistral4 checkpoint, which ships only as a VLM. Branch the text-backbone load on is_mistral4_config(&full_config), mirroring the text-only load_llama_family_from_weights path in src/loading/mod.rs. On the mistral4 branch, parse the raw text_config (with the top-level quantization block inherited when absent) as models::mistral4::Mistral4Config, run models::mistral4::sanitize_weights, and build models::Mistral4Model::from_weights wrapped as LoadedModel::Mistral4. The raw text_config is used rather than the Llama-massaged value because infer_llama_config_from_weights and apply_mistral_attention_head_override key off q_proj and can corrupt the MLA config. Weights are already stripped of the language_model. prefix by strip_language_model_prefix, so both loaders see the bare model.layers... namespace. The Llama path is unchanged for non-mistral4 Mistral3 VLMs (Pixtral and standard Mistral text backbones). text_config_value is still built via build_mistral_text_config for the scalar context fields (hidden_size, rms_norm_eps), which survive the Llama massaging on the mistral4 path. Real-model load validation of Mistral-Small-4-119B-2603-4bit is performed by the orchestrator.
inureyes
added
status:review
Real-model validation (Mistral-Small-4-119B-2603-4bit)Before this fix, loading the Mistral Small 4 VLM checkpoint aborted with With this fix, |
Implementation Review SummaryIntent
Findings AddressedNone. The implementation is correct and complete as submitted; no fixes were required. Verification
|
Mistral4Model overrode forward_with_embeddings but not embed_tokens, so it fell through to the LanguageModel trait default that returns None. PR #424 is the first to route a mistral4 (MLA + MoE) backbone into a Mistral3 VLM container, making it the first caller of embed_tokens on this model: VisionModule::get_input_embeddings (src/vision/mod.rs) calls it to build text-space embeddings before merging projected image features. With the default None, any image request to a Mistral4-backed Mistral3 VLM failed with "Text model must support embed_tokens for VLM". Text-only generation was unaffected because that path uses forward(input_ids) directly. Override embed_tokens to return the same self.embed_tokens.forward(input_ids) used by the no-input_embeddings branch of forward_with_embeddings and by the plain forward path, so the merge path and the forward path share one embedding table. Mirrors the llama3 backbone (src/models/llama3.rs).
Follow-up: fixed Mistral4 VLM image-path gap (HIGH)Deeper review found one issue that the text-only load validation could not surface, and it is now fixed on this branch (commit 29340f0). Finding (HIGH): Mistral4-backed Mistral3 VLM image requests fail at runtime. Fix: override Verification (narrow):
An image+text request against |
Correction on the image+text path (text-only is validated; image+text is NOT yet working)The
That size is Scope of this PR: it satisfies #423's required acceptance criteria (the checkpoint loads without the |
… supported-models The Mistral 3 VLM entry in supported-models.md now records that Mistral3 VLM routes Mistral4 (MLA+MoE) text backbones via the Mistral4 loader, that text generation is validated, and that image+text for the Mistral4 backbone is tracked in #425.
PR FinalizationTestsExisting coverage is complete for the checkpoint-free scope. The detection predicate and routing policy are covered by four tests in DocumentationUpdated the "Pixtral and Mistral 3 VLM wrappers" entry in Code quality
All checks passing. Not merged. |
…ng (#425) (#426) * fix(models): flatten mistral4 MoE tokens to 2D before SwitchGLU routing (#425) Mistral4MoE::forward passed a 3D `[batch, seq, hidden]` tensor and 3D `[batch, seq, top_k]` routing indices to the shared SwitchGLU::forward, which expects 2D `[n_tokens, hidden]` and `[n_tokens, top_k]`. SwitchGLU reads `indices_shape[0]` as `n_tokens` and `indices_shape[1]` as `top_k`, then decides `do_sort = n_tokens * top_k >= 64`. With 3D `[1, seq, top_k]` it misread `n_tokens=1` and `top_k=seq`. For `seq < 64` the `do_sort=false` branch limped along; for `seq >= 64` it entered gather_sort with the wrong dims and aborted with `[reshape] Cannot reshape array of size seq*hidden into (1,1,1)`. This is the 3D-vs-2D SwitchGLU contract violation behind #425: it fired on every real mistral4 prompt, on a 512-token chunked-prefill continuation chunk, and on any image-bearing multimodal prefill (e.g. the 3047-token merged text+image sequence, 3047 x 4096 = 12480512). It stayed latent because mistral4 was only ever run on a single-token warmup before #424 made the family loadable; the warmup's `[1, 1, hidden]` MoE forward succeeds. Fix mirrors qwen3_moe's proven flatten/reshape-back pattern: capture `orig_shape`, flatten `x` to `[n_tokens, hidden]` (prefill `[1, l, hidden]` -> `[l, hidden]`, decode `[1, 1, hidden]` -> `[1, hidden]`), run the existing routing math on the flat tensor so indices and scores come out correctly shaped `[n_tokens, top_k]`, dispatch experts and shared experts on the flat tensor, then reshape the result back to the original rank. Ops and eps are unchanged, so per-token routing is numerically identical, just correctly shaped. SwitchGLU, moe_weighted_sum, and all other models are untouched. This also corrects the decode path, which the old 3D shape was misrouting. Decode output may differ slightly from the pre-fix buggy path; that is the intended correction, not a regression. Validation: cargo fmt --check, cargo check --lib, and cargo clippy --lib --tests -- -D warnings all pass. Real-model 119B validation (chunked-prefill text plus image+text on mlx-community/Mistral-Small-4-119B-2603-4bit) is performed by the orchestrator. * docs(models): document Mistral 4 (Mistral Small 4) text and vision support (#425) Add a dedicated supported-models entry for the Mistral 4 (`mistral4`) MLA + MoE family now that it loads and runs end to end: the Mistral Small 4 119B checkpoint generates text and image-plus-text via the Mistral 3 VLM wrapper on the Mistral4 text backbone. Record the architecture values (q_lora_rank 1024, kv_lora_rank 256, split rope/nope head dims, 128 routed experts plus one shared with 4 active, Llama-4 attention scaling) and the validated checkpoint, and update the VLM wrapper note to reflect that image-plus-text on the Mistral4 backbone is now validated.
1 participant
Summary
Mistral3 VLM containers (
Mistral3ForConditionalGeneration) can wrap amistral4(MLA + MoE) text backbone, butbuild_pixtral_family_contextalways loaded the text backbone through the Llama3 loader, which fails because an MLA backbone has noq_proj. This routes themistral4case to the Mistral4 loader, unblockingmlx-community/Mistral-Small-4-119B-2603-4bit.Root cause
build_pixtral_family_context(src/loading/vlm_pixtral.rs) unconditionally parsedtext_configasmodels::llama3::ModelArgsand builtmodels::Llama3Model::from_weights, which expects standard attention weights (q_proj). For amistral4backbone the attention is MLA (q_a_proj/kv_a_proj_with_mqa/kv_b_proj, noq_proj), so loading failed withWeight not found: model.layers.0.self_attn.q_proj.weight. The text-only pathload_llama_family_from_weights(src/loading/mod.rs) already branches onis_mistral4_configand routes toMistral4Model::from_weights; the VLM context builder lacked the equivalent branch.What changed
src/loading/vlm_pixtral.rs: inbuild_pixtral_family_context, branch the text-backbone load onis_mistral4_config(&full_config). Themistral4branch parses the rawfull_config["text_config"](with the top-levelquantizationblock inherited viainherit_quantization_if_missingwhen absent) asmodels::mistral4::Mistral4Config, runsmodels::mistral4::sanitize_weights(&mut weights, &args, ""), and buildsmodels::Mistral4Model::from_weights(&weights, &args)wrapped asLoadedModel::Mistral4, mirroring the text-onlyload_llama_family_from_weightspath.text_configis used for model construction rather than the Llama-massaged value frombuild_mistral_text_config, becauseinfer_llama_config_from_weightsandapply_mistral_attention_head_overridekey offq_projand can corrupt the MLA config.language_model.prefix bystrip_language_model_prefixbefore this point, so both loaders see the baremodel.layers...namespace.mistral4Mistral3 VLMs (Pixtral and standard Mistral text backbones).text_config_valueis still built viabuild_mistral_text_configand used only for the scalar context fields (hidden_size,rms_norm_eps), which survive the Llama massaging on themistral4path.Test plan
cargo fmt --checkcargo check --lib --features metal,acceleratecargo clippy --lib --tests --features metal,accelerate -- -D warningscargo test --release --features metal,accelerate --lib loading::tests::is_mistral4(2 passed)A checkpoint-free unit test of the private routing branch is not practical (it needs a real model directory plus full model construction); the existing
is_mistral4_configdetection tests cover the predicate. Real-model load validation ofmlx-community/Mistral-Small-4-119B-2603-4bitis performed by the orchestrator.Closes #423