atlas_core/config/parsers/
mistral.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Split out of `config.rs` for file-size budget. Parser for a model family.
4
5#![allow(unused_imports)]
6
7use anyhow::{Context, Result};
8use serde_json::Value;
9
10use super::super::{
11    LayerType, ModelConfig, QuantizationConfig, VisionConfig, default_conv_kernel, default_one,
12    default_one_f64, default_partial_rotary, default_rms_eps, default_rope_theta, finalize_config,
13    parse_quantization_config, parse_vision_config, validate_config,
14};
15
16pub fn parse_mistral_params(json: &str) -> Result<ModelConfig> {
17    let raw: serde_json::Value =
18        serde_json::from_str(json).context("Invalid JSON in params.json")?;
19
20    let dim = raw.get("dim").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
21    let n_heads = raw.get("n_heads").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
22    let n_kv_heads = raw
23        .get("n_kv_heads")
24        .and_then(|v| v.as_u64())
25        .unwrap_or(n_heads as u64) as usize;
26    let n_layers = raw.get("n_layers").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
27    let head_dim = raw.get("head_dim").and_then(|v| v.as_u64()).unwrap_or(128) as usize;
28    let vocab_size = raw.get("vocab_size").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
29    let hidden_dim = raw.get("hidden_dim").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
30    let rope_theta = raw
31        .get("rope_theta")
32        .and_then(|v| v.as_f64())
33        .unwrap_or(10000.0);
34    let norm_eps = raw.get("norm_eps").and_then(|v| v.as_f64()).unwrap_or(1e-6);
35
36    // MLA fields
37    let kv_lora_rank = raw
38        .get("kv_lora_rank")
39        .and_then(|v| v.as_u64())
40        .unwrap_or(0) as usize;
41    let q_lora_rank = raw.get("q_lora_rank").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
42    let qk_nope_head_dim = raw
43        .get("qk_nope_head_dim")
44        .and_then(|v| v.as_u64())
45        .unwrap_or(0) as usize;
46    let qk_rope_head_dim = raw
47        .get("qk_rope_head_dim")
48        .and_then(|v| v.as_u64())
49        .unwrap_or(0) as usize;
50    let v_head_dim = raw
51        .get("v_head_dim")
52        .and_then(|v| v.as_u64())
53        .unwrap_or(head_dim as u64) as usize;
54
55    // MoE config
56    let moe = raw.get("moe");
57    let num_experts = moe
58        .and_then(|m| m.get("num_experts"))
59        .and_then(|v| v.as_u64())
60        .unwrap_or(0) as usize;
61    let num_experts_per_tok = moe
62        .and_then(|m| m.get("num_experts_per_tok"))
63        .and_then(|v| v.as_u64())
64        .unwrap_or(1) as usize;
65    let expert_hidden_dim = moe
66        .and_then(|m| m.get("expert_hidden_dim"))
67        .and_then(|v| v.as_u64())
68        .unwrap_or(0) as usize;
69    let num_shared_experts = moe
70        .and_then(|m| m.get("num_shared_experts"))
71        .and_then(|v| v.as_u64())
72        .unwrap_or(0) as usize;
73    let _shared_expert_size = if num_shared_experts > 0 {
74        expert_hidden_dim
75    } else {
76        0
77    };
78
79    // All layers are attention+MoE for Mistral (no SSM)
80    let layer_types = vec![LayerType::FullAttention; n_layers];
81
82    // Start from a known-good template and override all fields.
83    let mut config = ModelConfig::qwen3_next_80b_nvfp4();
84    // Reset all fields to zero/empty before populating from params.json.
85    // CRITICAL: must reset ALL SSM fields inherited from the Qwen3 template,
86    // otherwise linear_num_key_heads > 0 misroutes this dense model onto the
87    // SSM/linear-attention dispatch path.
88    config.num_hidden_layers = 0;
89    config.intermediate_size = 0;
90    config.vocab_size = 0;
91    config.num_attention_heads = 0;
92    config.num_key_value_heads = 0;
93    config.head_dim = 0;
94    config.num_experts = 0;
95    config.num_experts_per_tok = 1;
96    config.moe_intermediate_size = 0;
97    config.shared_expert_intermediate_size = 0;
98    config.mtp_num_hidden_layers = 0;
99    config.linear_num_key_heads = 0;
100    config.linear_key_head_dim = 0;
101    config.linear_num_value_heads = 0;
102    config.linear_value_head_dim = 0;
103    config.linear_conv_kernel_dim = 0;
104    // MLA: RoPE is applied ONLY to the rope portion (qk_rope_head_dim dims per head).
105    // Our RoPE kernel rotates dims 0..rotary_dim-1, so we swap Q/K to [rope|nope]
106    // before RoPE, matching the kernel's expectation.
107    config.partial_rotary_factor = if qk_rope_head_dim > 0 && head_dim > 0 {
108        qk_rope_head_dim as f64 / head_dim as f64
109    } else {
110        1.0
111    };
112    config.hidden_size = dim;
113    config.num_hidden_layers = n_layers;
114    config.intermediate_size = hidden_dim;
115    config.vocab_size = vocab_size;
116    config.num_attention_heads = n_heads;
117    config.num_key_value_heads = n_kv_heads;
118    config.head_dim = head_dim;
119    config.num_experts = num_experts;
120    config.num_experts_per_tok = num_experts_per_tok;
121    config.moe_intermediate_size = expert_hidden_dim;
122    // Mistral has a shared expert with intermediate_size = expert_hidden_dim.
123    // The CUDA kernels have NULL guards for the shared expert slot, so it's safe
124    // to set this non-zero even during EP (NULL guard writes zeros for missing experts).
125    config.shared_expert_intermediate_size = expert_hidden_dim;
126    config.layer_types = layer_types;
127    config.max_position_embeddings = raw
128        .get("max_position_embeddings")
129        .and_then(|v| v.as_u64())
130        .unwrap_or(8192) as usize;
131    config.rope_theta = rope_theta;
132    config.rms_norm_eps = norm_eps;
133    config.model_type = "mistral".to_string();
134    config.attn_gated = false;
135    config.kv_lora_rank = kv_lora_rank;
136    config.q_lora_rank = q_lora_rank;
137    config.qk_nope_head_dim = qk_nope_head_dim;
138    config.qk_rope_head_dim = qk_rope_head_dim;
139    config.v_head_dim = v_head_dim;
140
141    // YaRN RoPE scaling — Mistral exposes this under `params.json::yarn` with
142    // its own naming: `alpha` is the low-rotation cutoff (HF `beta_slow`),
143    // `beta` is the high-rotation cutoff (HF `beta_fast`).
144    if let Some(yarn) = raw.get("yarn") {
145        config.yarn_factor = yarn.get("factor").and_then(|v| v.as_f64()).unwrap_or(0.0) as f32;
146        config.yarn_beta_slow = yarn.get("alpha").and_then(|v| v.as_f64()).unwrap_or(1.0) as f32;
147        config.yarn_beta_fast = yarn.get("beta").and_then(|v| v.as_f64()).unwrap_or(32.0) as f32;
148        config.yarn_original_max_position_embeddings = yarn
149            .get("original_max_position_embeddings")
150            .and_then(|v| v.as_u64())
151            .unwrap_or(8192) as usize;
152    }
153    // llama_4_scaling Q temperature multiplier (separate from YaRN).
154    if let Some(l4) = raw.get("llama_4_scaling") {
155        config.llama_4_scaling_beta = l4.get("beta").and_then(|v| v.as_f64()).unwrap_or(0.0) as f32;
156        config.llama_4_scaling_original_max_position_embeddings =
157            l4.get("original_max_position_embeddings")
158                .and_then(|v| v.as_u64())
159                .unwrap_or(8192) as usize;
160    }
161
162    // Detect tied embeddings
163    config.tie_word_embeddings = raw
164        .get("tied_embeddings")
165        .and_then(|v| v.as_bool())
166        .unwrap_or(false);
167    // Mistral uses standard BOS=1, EOS=2 (</s>)
168    config.eos_token_id = 2;
169    config.bos_token_id = 1;
170
171    finalize_config(&mut config, &raw)?;
172    Ok(config)
173}
atlas_core/config/parsers/mistral.rs

atlas_core/config/parsers/
mistral.rs
