metrale_config/
lib.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The typed model config (`ModelConfig`) parsed from a checkpoint's `config.json`
4//! or GGUF metadata, its capabilities, and the `METRALE_*` lever table.
5//!
6//! Owner: config.
7//! Invariants: none beyond the types.
8
9use anyhow::Result;
10use serde::Deserialize;
11
12/// 2026-09-26: Deserialize a u32 that may be JSON null; null gives 0.
13fn nullable_u32<'de, D: serde::Deserializer<'de>>(d: D) -> std::result::Result<u32, D::Error> {
14    Option::<u32>::deserialize(d).map(|v| v.unwrap_or(0))
15}
16
17/// 2026-09-26: `eos_token_id` as null (0), a scalar, or an array (its first element).
18///
19/// `parse_config` recovers the whole array into [`ModelConfig::eos_token_ids`]. A parser that
20/// wants another primary (`step3p7` takes the last element) rewrites the field before
21/// deserializing.
22fn eos_token_id_field<'de, D: serde::Deserializer<'de>>(
23    d: D,
24) -> std::result::Result<u32, D::Error> {
25    #[derive(Deserialize)]
26    #[serde(untagged)]
27    enum OneOrMany {
28        One(u32),
29        Many(Vec<u32>),
30    }
31    Ok(match Option::<OneOrMany>::deserialize(d)? {
32        None => 0,
33        Some(OneOrMany::One(v)) => v,
34        Some(OneOrMany::Many(v)) => v.first().copied().unwrap_or(0),
35    })
36}
37
38/// 2026-09-26: Which dtype ladder GLM-5.3's MoE router runs in. The glm5_next parser reads
39/// `moe_router_dtype` and uses `HfFp32` when it is absent.
40#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)]
41#[serde(rename_all = "snake_case")]
42pub enum Glm5NextRouterMode {
43    /// 2026-09-26: The router math in fp32 (`is_fp32`). The default.
44    #[default]
45    HfFp32,
46    /// 2026-09-26: The router's bf16 ladder (`router_bf16_ladder` in glm5next_mlp).
47    VllmBf16,
48}
49
50impl Glm5NextRouterMode {
51    /// 2026-09-26: Parse a `moe_router_dtype` value; an unrecognised one gives `None`, which
52    /// the glm5_next parser turns into an error.
53    pub fn from_config_str(s: &str) -> Option<Self> {
54        match s {
55            "float32" | "fp32" => Some(Self::HfFp32),
56            "bfloat16" | "bf16" => Some(Self::VllmBf16),
57            _ => None,
58        }
59    }
60
61    pub fn is_fp32(self) -> bool {
62        matches!(self, Self::HfFp32)
63    }
64}
65
66#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
67#[serde(rename_all = "snake_case")]
68pub enum LayerType {
69    FullAttention,
70    SlidingAttention,
71    LinearAttention,
72    /// 2026-09-26: A layer with no mixer, only the routed-expert FFN (Nemotron-H pattern `E`,
73    /// Puzzle `moe`).
74    Moe,
75    /// 2026-09-26: Sparse attention (`deepseek_sparse_attention`): a full-rank mixer whose
76    /// visible key set an indexer selects per query at run time.
77    SparseAttention,
78}
79
80impl LayerType {
81    pub fn is_attention(self) -> bool {
82        matches!(
83            self,
84            Self::FullAttention | Self::SlidingAttention | Self::SparseAttention
85        )
86    }
87
88    pub fn hf_name(self) -> &'static str {
89        match self {
90            Self::FullAttention => "full_attention",
91            Self::SlidingAttention => "sliding_attention",
92            Self::LinearAttention => "linear_attention",
93            Self::Moe => "moe",
94            Self::SparseAttention => "deepseek_sparse_attention",
95        }
96    }
97}
98
99/// 2026-09-26: The weight-quantization block of a checkpoint config (`quantization_config`,
100/// into which the server merges a sibling `hf_quant_config.json`), as read by
101/// `parse_quantization_config`. That parser returns `None` when the block declares no method,
102/// algorithm or ignore list.
103#[derive(Debug, Clone)]
104pub struct QuantizationConfig {
105    /// 2026-09-26: `quant_method` as written, `modelopt` for a ModelOpt sidecar that omits it,
106    /// or empty.
107    pub quant_method: String,
108    /// 2026-09-26: `quant_algo`, else `NVFP4`/`FP8` derived from `config_groups.group_0.weights`,
109    /// else empty.
110    pub quant_algo: String,
111    /// 2026-09-26: `format`, or empty.
112    pub format: String,
113    /// 2026-09-26: Module-path patterns from `ignore`, then those from `exclude_modules` not
114    /// already listed.
115    pub ignore_modules: Vec<String>,
116    /// 2026-09-28: What the block declares per linear layer, weights and activations: the
117    /// default precision every loader and kernel dispatch follows.
118    pub precision: precision_plan::DeclaredPrecisionPlan,
119}
120
121/// 2026-09-26: Whether GLM-5.3's vision tower is enabled for this process
122/// (`METRALE_GLM_VISION`, see [`glm_vision_enabled_from`]); off by default.
123///
124/// The glm5_next parser (which leaves `config.vision` `None` when off) and
125/// `Glm5NextWeightLoader::binds_vision_encoder` both call it, so they cannot disagree.
126pub fn glm_vision_enabled() -> bool {
127    glm_vision_enabled_from(std::env::var("METRALE_GLM_VISION").ok().as_deref())
128}
129
130/// 2026-09-26: The policy half of [`glm_vision_enabled`], taking the value instead of reading
131/// the environment, so tests need not set a variable. Only `1`, `true`, `yes` or `on`
132/// (trimmed, any case) enable; anything else, unset included, leaves the tower off.
133pub fn glm_vision_enabled_from(value: Option<&str>) -> bool {
134    matches!(
135        value.map(|v| v.trim().to_ascii_lowercase()).as_deref(),
136        Some("1" | "true" | "yes" | "on")
137    )
138}
139
140#[cfg(test)]
141mod glm_vision_gate_tests {
142    use super::glm_vision_enabled_from;
143
144    #[test]
145    fn only_an_explicit_affirmative_enables_the_tower() {
146        for off in [
147            None,
148            Some(""),
149            Some("0"),
150            Some("off"),
151            Some("false"),
152            Some("no"),
153            Some("yes please"),
154        ] {
155            assert!(!glm_vision_enabled_from(off), "{off:?} must not enable");
156        }
157        for on in ["1", "true", "TRUE", "Yes", " on "] {
158            assert!(glm_vision_enabled_from(Some(on)), "{on:?} must enable");
159        }
160    }
161}
162
163/// 2026-09-26: Vision tower configuration: `parse_vision_config` for the Qwen family, and the
164/// GLM-5.3 parser, which also fills the fields the Qwen parser leaves at their defaults.
165#[derive(Debug, Clone)]
166pub struct VisionConfig {
167    pub depth: usize,
168    pub hidden_size: usize,
169    pub num_heads: usize,
170    pub patch_size: usize,
171    pub temporal_patch_size: usize,
172    pub spatial_merge_size: usize,
173    pub intermediate_size: usize,
174    pub out_hidden_size: usize,
175    pub deepstack_visual_indexes: Vec<usize>,
176    /// 2026-09-26: The token that marks where vision embeddings are spliced into the text.
177    /// 0 means undeclared: the engine then uses `vision_encoder::IMAGE_PAD_TOKEN_ID`, and the
178    /// chat tokenizer the id of `<|image_pad|>`.
179    pub image_pad_token_id: u32,
180    /// 2026-09-26: The video-frame counterpart of [`Self::image_pad_token_id`], with the same
181    /// 0 rule (`VIDEO_PAD_TOKEN_ID`, `<|video_pad|>`).
182    pub video_pad_token_id: u32,
183    /// 2026-09-26: The vision area bound in pixels. The parsers leave it `None`; `serve_load`
184    /// sets it from `resolve_vision_max_pixels` before the model is built, and the vision
185    /// loaders pass it to the encoder.
186    pub max_pixels: Option<usize>,
187
188    pub hidden_act: String,
189    pub rms_norm_eps: f32,
190    pub attention_bias: bool,
191    pub projection_intermediate_size: usize,
192    pub swiglu_limit: f32,
193    pub rope_theta: f32,
194    pub image_mean: [f32; 3],
195    pub image_std: [f32; 3],
196    pub min_image_tokens: usize,
197    pub max_image_tokens: usize,
198    /// 2026-09-26: Patch tokens in 2x2-block-major order rather than raster; the preprocessor
199    /// reads it (`vision_preprocess`). The GLM-5.3 parser sets it.
200    pub block_major_patches: bool,
201}
202
203impl Default for VisionConfig {
204    fn default() -> Self {
205        Self {
206            depth: 0,
207            hidden_size: 0,
208            num_heads: 0,
209            patch_size: 0,
210            temporal_patch_size: 0,
211            spatial_merge_size: 0,
212            intermediate_size: 0,
213            out_hidden_size: 0,
214            deepstack_visual_indexes: Vec::new(),
215            image_pad_token_id: 0,
216            video_pad_token_id: 0,
217            max_pixels: None,
218            hidden_act: "gelu_pytorch_tanh".to_string(),
219            rms_norm_eps: 1e-6,
220            attention_bias: true,
221            projection_intermediate_size: 0,
222            swiglu_limit: 0.0,
223            rope_theta: 10_000.0,
224            image_mean: [0.5, 0.5, 0.5],
225            image_std: [0.5, 0.5, 0.5],
226            min_image_tokens: 0,
227            max_image_tokens: 0,
228            block_major_patches: false,
229        }
230    }
231}
232
233impl VisionConfig {
234    pub fn merger_input_size(&self) -> usize {
235        self.spatial_merge_size * self.spatial_merge_size * self.hidden_size
236    }
237}
238
239pub(crate) fn default_one() -> usize {
240    1
241}
242pub(crate) fn default_one_f64() -> f64 {
243    1.0
244}
245pub(crate) fn default_one_f32() -> f32 {
246    1.0
247}
248pub(crate) fn default_rope_theta() -> f64 {
249    10000.0
250}
251pub(crate) fn default_rms_eps() -> f64 {
252    1e-6
253}
254pub(crate) fn default_partial_rotary() -> f64 {
255    1.0
256}
257pub(crate) fn default_conv_kernel() -> usize {
258    4
259}
260
261pub mod activation_quantization;
262mod dispatch;
263mod factory;
264mod gguf;
265mod kv_completeness;
266#[cfg(test)]
267mod kv_completeness_tests;
268pub mod levers;
269mod methods;
270mod model_config;
271mod parsers;
272mod position_encoding;
273pub mod precision_plan;
274pub mod recipe_yaml;
275#[cfg(test)]
276mod tests;
277pub mod weight_quantization;
278
279pub use activation_quantization::{ActQuantFormat, ActivationQuantization, ProjFamily};
280pub use dispatch::parse_config;
281pub use gguf::{GgufConfigInputs, GgufMeta, config_from_gguf};
282pub use model_config::ModelConfig;
283pub use parsers::{
284    PEFT_SUPPORTED_TARGET_MODULES, PeftAdapterConfig, allow_partial_targets,
285    glm5_next_mtp_layer_index, parse_mistral_params, parse_peft_adapter_config,
286    parse_quantization_config,
287};
288pub(crate) use parsers::{
289    parse_deepseek_v4, parse_gemma4_params, parse_glm5_next, parse_kimi_k3, parse_laguna,
290    parse_longcat_ngram, parse_minimax_m2, parse_qwen4_exp, parse_step3p7, parse_vision_config,
291    sanitize_kimi_k3_eos,
292};
293pub use position_encoding::AttnPositionEncoding;
294pub(crate) use position_encoding::resolve_attn_position_encoding;
295pub use precision_plan::{DeclaredPrecisionPlan, LayerPrecision};
296pub use weight_quantization::{
297    Nvfp4Act, W4a4Downcast, WeightQuantPolicy, WeightQuantTier, WeightQuantization,
298};
299
300pub(crate) fn finalize_config(config: &mut ModelConfig, raw: &serde_json::Value) -> Result<()> {
301    if config.quantization_config.is_none() {
302        config.quantization_config = parse_quantization_config(raw)?;
303    }
304    validate_config(config)
305}
306
307/// 2026-09-26: Post-parse validation: `layer_types` and `num_attention_heads_per_layer`
308/// lengths against `num_hidden_layers`, and the SSM and Mamba-2 head fields.
309pub(crate) fn validate_config(config: &ModelConfig) -> Result<()> {
310    if !config.layer_types.is_empty() && config.layer_types.len() != config.num_hidden_layers {
311        anyhow::bail!(
312            "layer_types length ({}) doesn't match num_hidden_layers ({}) in config.json",
313            config.layer_types.len(),
314            config.num_hidden_layers,
315        );
316    }
317
318    if !config.num_attention_heads_per_layer.is_empty()
319        && config.num_attention_heads_per_layer.len() != config.num_hidden_layers
320    {
321        anyhow::bail!(
322            "num_attention_heads_per_layer length ({}) doesn't match num_hidden_layers ({}) in config.json",
323            config.num_attention_heads_per_layer.len(),
324            config.num_hidden_layers,
325        );
326    }
327
328    let has_ssm =
329        config.layer_types.contains(&LayerType::LinearAttention) || config.linear_num_key_heads > 0;
330    if has_ssm && config.linear_num_key_heads == 0 && config.mamba_num_heads == 0 {
331        anyhow::bail!(
332            "SSM model detected but linear_num_key_heads is 0 in config.json. \
333             This field is required for SSM/GDN layer initialization."
334        );
335    }
336
337    if config.mamba_num_heads > 0 {
338        if config.mamba_head_dim == 0 {
339            anyhow::bail!("mamba_head_dim must be greater than zero");
340        }
341        if config.ssm_state_size == 0 {
342            anyhow::bail!("ssm_state_size must be greater than zero");
343        }
344        if config.n_groups == 0 {
345            anyhow::bail!("n_groups must be greater than zero");
346        }
347        if !config.mamba2_d_inner().is_multiple_of(config.n_groups) {
348            anyhow::bail!("mamba_num_heads * mamba_head_dim must be divisible by n_groups");
349        }
350    }
351
352    Ok(())
353}
354
355pub mod capabilities;
metrale_config/lib.rs

metrale_config/
lib.rs
