metrale_config/parsers/
mistral.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Parser for a Mistral `params.json`, which the server uses for a checkpoint
4//! that has `params.json` and no `config.json`.
5//!
6//! Owner: config (model parsers).
7//! Invariants: none beyond the types.
8
9#![allow(unused_imports)]
10
11use anyhow::{Context, Result};
12use serde_json::Value;
13
14use super::super::{
15    LayerType, ModelConfig, QuantizationConfig, VisionConfig, default_conv_kernel, default_one,
16    default_one_f64, default_partial_rotary, default_rms_eps, default_rope_theta, finalize_config,
17    parse_quantization_config, parse_vision_config, validate_config,
18};
19
20pub fn parse_mistral_params(json: &str) -> Result<ModelConfig> {
21    let raw: serde_json::Value =
22        serde_json::from_str(json).context("Invalid JSON in params.json")?;
23
24    let dim = raw.get("dim").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
25    let n_heads = raw.get("n_heads").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
26    let n_kv_heads = raw
27        .get("n_kv_heads")
28        .and_then(|v| v.as_u64())
29        .unwrap_or(n_heads as u64) as usize;
30    let n_layers = raw.get("n_layers").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
31    let head_dim = raw.get("head_dim").and_then(|v| v.as_u64()).unwrap_or(128) as usize;
32    let vocab_size = raw.get("vocab_size").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
33    let hidden_dim = raw.get("hidden_dim").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
34    let rope_theta = raw
35        .get("rope_theta")
36        .and_then(|v| v.as_f64())
37        .unwrap_or(10000.0);
38    let norm_eps = raw.get("norm_eps").and_then(|v| v.as_f64()).unwrap_or(1e-6);
39
40    let kv_lora_rank = raw
41        .get("kv_lora_rank")
42        .and_then(|v| v.as_u64())
43        .unwrap_or(0) as usize;
44    let q_lora_rank = raw.get("q_lora_rank").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
45    let qk_nope_head_dim = raw
46        .get("qk_nope_head_dim")
47        .and_then(|v| v.as_u64())
48        .unwrap_or(0) as usize;
49    let qk_rope_head_dim = raw
50        .get("qk_rope_head_dim")
51        .and_then(|v| v.as_u64())
52        .unwrap_or(0) as usize;
53    let v_head_dim = raw
54        .get("v_head_dim")
55        .and_then(|v| v.as_u64())
56        .unwrap_or(head_dim as u64) as usize;
57
58    let moe = raw.get("moe");
59    let num_experts = moe
60        .and_then(|m| m.get("num_experts"))
61        .and_then(|v| v.as_u64())
62        .unwrap_or(0) as usize;
63    let num_experts_per_tok = moe
64        .and_then(|m| m.get("num_experts_per_tok"))
65        .and_then(|v| v.as_u64())
66        .unwrap_or(1) as usize;
67    let expert_hidden_dim = moe
68        .and_then(|m| m.get("expert_hidden_dim"))
69        .and_then(|v| v.as_u64())
70        .unwrap_or(0) as usize;
71    let num_shared_experts = moe
72        .and_then(|m| m.get("num_shared_experts"))
73        .and_then(|v| v.as_u64())
74        .unwrap_or(0) as usize;
75    let _shared_expert_size = if num_shared_experts > 0 {
76        expert_hidden_dim
77    } else {
78        0
79    };
80
81    let layer_types = vec![LayerType::FullAttention; n_layers];
82
83    // 2026-09-26: The template's fields are reset before `params.json` is applied, the
84    // linear-attention ones included: `validate_config` treats a non-zero
85    // `linear_num_key_heads` as a GDN model.
86    let mut config = ModelConfig::qwen3_next_80b_nvfp4();
87    config.num_hidden_layers = 0;
88    config.intermediate_size = 0;
89    config.vocab_size = 0;
90    config.num_attention_heads = 0;
91    config.num_key_value_heads = 0;
92    config.head_dim = 0;
93    config.num_experts = 0;
94    config.num_experts_per_tok = 1;
95    config.moe_intermediate_size = 0;
96    config.shared_expert_intermediate_size = 0;
97    config.mtp_num_hidden_layers = 0;
98    config.linear_num_key_heads = 0;
99    config.linear_key_head_dim = 0;
100    config.linear_num_value_heads = 0;
101    config.linear_value_head_dim = 0;
102    config.linear_conv_kernel_dim = 0;
103    config.partial_rotary_factor = if qk_rope_head_dim > 0 && head_dim > 0 {
104        qk_rope_head_dim as f64 / head_dim as f64
105    } else {
106        1.0
107    };
108    config.hidden_size = dim;
109    config.num_hidden_layers = n_layers;
110    config.intermediate_size = hidden_dim;
111    config.vocab_size = vocab_size;
112    config.num_attention_heads = n_heads;
113    config.num_key_value_heads = n_kv_heads;
114    config.head_dim = head_dim;
115    config.num_experts = num_experts;
116    config.num_experts_per_tok = num_experts_per_tok;
117    config.moe_intermediate_size = expert_hidden_dim;
118    config.shared_expert_intermediate_size = expert_hidden_dim;
119    config.layer_types = layer_types;
120    config.max_position_embeddings = raw
121        .get("max_position_embeddings")
122        .and_then(|v| v.as_u64())
123        .unwrap_or(8192) as usize;
124    config.rope_theta = rope_theta;
125    config.rms_norm_eps = norm_eps;
126    config.model_type = "mistral".to_string();
127    config.attn_gated = false;
128    config.kv_lora_rank = kv_lora_rank;
129    config.q_lora_rank = q_lora_rank;
130    config.qk_nope_head_dim = qk_nope_head_dim;
131    config.qk_rope_head_dim = qk_rope_head_dim;
132    config.v_head_dim = v_head_dim;
133
134    // 2026-09-26: `params.json` names the YaRN cutoffs `alpha` (read as `yarn_beta_slow`) and
135    // `beta` (read as `yarn_beta_fast`).
136    if let Some(yarn) = raw.get("yarn") {
137        config.yarn_factor = yarn.get("factor").and_then(|v| v.as_f64()).unwrap_or(0.0) as f32;
138        config.yarn_beta_slow = yarn.get("alpha").and_then(|v| v.as_f64()).unwrap_or(1.0) as f32;
139        config.yarn_beta_fast = yarn.get("beta").and_then(|v| v.as_f64()).unwrap_or(32.0) as f32;
140        config.yarn_original_max_position_embeddings = yarn
141            .get("original_max_position_embeddings")
142            .and_then(|v| v.as_u64())
143            .unwrap_or(8192) as usize;
144    }
145    if let Some(l4) = raw.get("llama_4_scaling") {
146        config.llama_4_scaling_beta = l4.get("beta").and_then(|v| v.as_f64()).unwrap_or(0.0) as f32;
147        config.llama_4_scaling_original_max_position_embeddings =
148            l4.get("original_max_position_embeddings")
149                .and_then(|v| v.as_u64())
150                .unwrap_or(8192) as usize;
151    }
152
153    config.tie_word_embeddings = raw
154        .get("tied_embeddings")
155        .and_then(|v| v.as_bool())
156        .unwrap_or(false);
157    config.eos_token_id = 2;
158    config.bos_token_id = 1;
159
160    finalize_config(&mut config, &raw)?;
161    Ok(config)
162}
metrale_config/parsers/mistral.rs

metrale_config/parsers/
mistral.rs
