From 0681d42bbb5890e45f4bbe29a153c7401cadcb99 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 18 Oct 2026 20:34:10 +0900 Subject: [PATCH 1/3] fix(models): lift transformers 5 rope_parameters for every family transformers 5 saves rope_theta, rope_scaling and partial_rotary_factor only under a nested rope_parameters block. Only Qwen3 and GraniteMoeHybrid lifted it (#2439), so about 80 config structs silently ran at their serde default base and drifted on long prompts. There is no raw-Value choke point (families read config.json directly and VLMs nest the block under text_config), so the type is the choke point. The lifter moves to mlxcel_core::rope_parameters with an impl_rope_lifting_deserialize! macro, and every uncovered struct gets #[serde(remote = "Self")] plus one macro line. The lifter also reads the per-layer-type form (Gemma 3, Gemma 3n, OLMo 3), lifts partial_rotary_factor, keeps MRoPE mrope_section, and writes the scheme under type only so alias-both fields do not hit a duplicate-field error. DFlashConfig::from_json, the Idefics2 defaults and both vlm_qwen ad hoc lifts use it; JinaVlmTextConfig::from_json falls back to the nested base, and the Phi-4 text-field copy now carries rope_parameters. EXAONE-MoE and Mistral 4 now let a flat key win, and exaone4 accepts a missing rope_scaling. A table test covers every struct; all four new tests fail with the fix stashed. Transformers-5 copies of llama-3.2-1b-instruct-4bit and qwen3-30b-a3b-4bit decode a 20 KB greedy prompt byte-identically to the originals; main produces different output on the copies. Closes #2443 --- CHANGELOG.md | 1 + src/audio/qwen3_omni_moe/speech_config.rs | 9 + .../mlxcel-core/src/drafter/dflash/config.rs | 7 +- src/lib/mlxcel-core/src/lib.rs | 4 + src/lib/mlxcel-core/src/rope_parameters.rs | 166 ++++++++ .../mlxcel-core/src/rope_parameters_tests.rs | 159 ++++++++ src/loading/vlm_idefics2.rs | 3 + src/loading/vlm_qwen.rs | 42 +- src/loading/vlm_special.rs | 1 + src/models/afmoe.rs | 3 + src/models/apertus.rs | 3 + src/models/baichuan.rs | 3 + src/models/bailing_moe.rs | 3 + src/models/bailing_moe_linear.rs | 3 + src/models/bitnet.rs | 3 + src/models/cohere.rs | 3 + src/models/cohere2.rs | 3 + src/models/cohere2_moe.rs | 3 + src/models/dbrx.rs | 3 + src/models/deepseek.rs | 3 + src/models/deepseek_v2.rs | 3 + src/models/deepseek_v3.rs | 3 + src/models/deepseek_v32.rs | 3 + src/models/deepseek_v4.rs | 3 + src/models/dots1.rs | 3 + src/models/ernie4_5.rs | 3 + src/models/ernie4_5_moe.rs | 3 + src/models/ernie4_5_moe_vl.rs | 3 + src/models/exaone.rs | 3 + src/models/exaone4.rs | 7 + src/models/exaone_moe.rs | 13 +- src/models/falcon_h1.rs | 3 + src/models/falcon_ocr.rs | 3 + src/models/gemma.rs | 3 + src/models/gemma2.rs | 3 + src/models/gemma3.rs | 3 + src/models/gemma3n.rs | 3 + src/models/glm4.rs | 3 + src/models/glm4_moe.rs | 3 + src/models/glm4_moe_lite.rs | 3 + src/models/glm4v.rs | 3 + src/models/gpt_oss.rs | 3 + src/models/granite.rs | 3 + src/models/helium.rs | 3 + src/models/hunyuan_moe.rs | 3 + src/models/hunyuan_v1_dense.rs | 3 + src/models/hunyuan_vl.rs | 3 + src/models/internlm2.rs | 3 + src/models/internlm3.rs | 3 + src/models/iquestloopcoder.rs | 3 + src/models/jina_vlm.rs | 8 + src/models/kimi_linear.rs | 3 + src/models/klear.rs | 3 + src/models/llada2_moe/mod.rs | 3 + src/models/llama3.rs | 3 + src/models/llama4.rs | 3 + src/models/longcat_flash_ngram.rs | 3 + src/models/mimo.rs | 3 + src/models/minicpm.rs | 3 + src/models/minicpm3.rs | 3 + src/models/minimax.rs | 3 + src/models/minimax_m3_config.rs | 3 + src/models/mistral4.rs | 17 +- src/models/mixtral.rs | 3 + src/models/mllama/config.rs | 3 + src/models/molmo.rs | 3 + src/models/molmo2.rs | 3 + src/models/moondream2.rs | 3 + src/models/nanochat_config.rs | 3 + src/models/nemotron.rs | 3 + src/models/nemotron_nas.rs | 3 + src/models/olmo.rs | 3 + src/models/olmo2.rs | 3 + src/models/olmo3.rs | 3 + src/models/olmoe.rs | 3 + src/models/paddleocr_vl.rs | 3 + src/models/phi.rs | 3 + src/models/phi3.rs | 3 + src/models/phimoe.rs | 3 + src/models/phixtral.rs | 3 + src/models/qwen2_moe.rs | 3 + src/models/qwen2_vl.rs | 3 + src/models/qwen3_moe.rs | 3 + src/models/qwen3_next.rs | 3 + src/models/qwen3_vl.rs | 3 + src/models/qwen3_vl_moe.rs | 3 + src/models/recurrent_gemma.rs | 3 + src/models/rope_parameters_lift_tests.rs | 385 ++++++++++++++++++ src/models/rope_utils.rs | 83 +--- src/models/seed_oss.rs | 3 + src/models/smollm3.rs | 3 + src/models/stablelm.rs | 3 + src/models/starcoder2.rs | 3 + src/models/step3p5.rs | 3 + src/models/telechat3.rs | 3 + src/models/youtu_vl_lm_config.rs | 3 + 96 files changed, 1032 insertions(+), 116 deletions(-) create mode 100644 src/lib/mlxcel-core/src/rope_parameters.rs create mode 100644 src/lib/mlxcel-core/src/rope_parameters_tests.rs create mode 100644 src/models/rope_parameters_lift_tests.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 41e900ddf..01b363616 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Fixed +- Configs saved by transformers 5, which keep `rope_theta`, `rope_scaling` and `partial_rotary_factor` only under a nested `rope_parameters` block, now load with their real RoPE base and scaling in every family that reads the flat keys, not only Qwen3 and GraniteMoeHybrid (#2443). Before, about 80 config structs (Llama and everything sharing its args, Qwen2/Qwen3-MoE, Gemma, Phi, OLMo, DeepSeek, GLM, the VLM text configs, the Qwen3-Omni speech stack and the DFlash drafter) silently ran at the serde default base, decoding fluently on short prompts and drifting on long ones, and a Llama 3 config lost its `llama3` scaling. The per-layer-type form (Gemma 3, Gemma 3n, OLMo 3) and the MRoPE `mrope_section` are lifted too. A flat key still wins when both are present; EXAONE-MoE and Mistral 4 used to let the nested base win and now follow the same rule. - Checkpoints that load through the SentencePiece fallback (a `tokenizer.model` with no `tokenizer.json`) now encode their `"special": false` added tokens as the one id each was declared with, as `transformers` does, instead of splitting them into ordinary pieces (#1758). IQuest-Coder and IQuest-Coder Loop tool-calling prompts now carry ``, ``, `` and their closers as single ids, Baichuan-M1's `` turn marker and ERNIE-4.5-VL's `<|IMAGE_PLACEHOLDER|>` and `<|LOC_n|>` tokens do the same, and each added token's `lstrip`, `rstrip` and `single_word` flags are honored. `skip_special_tokens` still drops only the special ones. On this path `parse_special: false` (`/tokenize`, logit-bias text keys) now also recognizes non-special added tokens, as llama.cpp does, `added_tokens.json` is read when `added_tokens_decoder` is absent, and a `single_word` special token glued to a word (IQuest's `` in `de`) encodes as text, as in `transformers`. - On CUDA, a bf16 model on an fp16 TurboQuant KV cache (`--kv-cache-mode fp16+turbo4` and the other turbo modes) no longer aborts with `Mask type must promote to output type float32` when attention carries an array mask, which for Gemma 4 meant any prompt past the 1024-token sliding window (#2336). Additive masks now reach SDPA in the attention dtype, `result_type(q, k, v)`, instead of the query's dtype. - Mixed bf16/fp16 attention follow-ups to #2336 (#2422). On CUDA, the query-chunking gate for materializing SDPA now judges cuDNN eligibility and score size from the attention dtype, `result_type(q, k, v)`, instead of q's dtype, so a bf16 model on an fp16 TurboQuant KV cache, which attends in float32, chunks long prefills at 4 bytes per score instead of skipping the chunking at head_dim 128 or budgeting 2 bytes. The softcap attention composite (Gemma 2) casts its mask to and returns the attention dtype, as fast SDPA does, instead of q's and v's. Model-state snapshot donation checks the KV cache mode before the keys/values layout, so `--kv-cache-mode fp16+turbo4` logs `cache mode Turbo4Asym is not supported` instead of `has only one of keys/values`, and a populated symmetric `turbo4` layer is refused instead of being left out of the snapshot. Same-dtype attention is unchanged. diff --git a/src/audio/qwen3_omni_moe/speech_config.rs b/src/audio/qwen3_omni_moe/speech_config.rs index 975cb2cbf..64f158fb7 100644 --- a/src/audio/qwen3_omni_moe/speech_config.rs +++ b/src/audio/qwen3_omni_moe/speech_config.rs @@ -36,6 +36,7 @@ fn default_true() -> bool { /// Talker MoE decoder text config (`talker_config.text_config`). #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct TalkerTextConfig { pub hidden_size: usize, pub intermediate_size: usize, @@ -55,6 +56,8 @@ pub struct TalkerTextConfig { pub norm_topk_prob: bool, } +mlxcel_core::impl_rope_lifting_deserialize!(TalkerTextConfig); + fn default_shared_expert_intermediate_size() -> usize { 768 } @@ -63,6 +66,7 @@ fn default_shared_expert_intermediate_size() -> usize { /// (`talker_config.code_predictor_config`). All fields default per the /// reference dataclass. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct CodePredictorConfig { #[serde(default = "default_cp_layers")] pub num_hidden_layers: usize, @@ -86,6 +90,8 @@ pub struct CodePredictorConfig { pub num_code_groups: usize, } +mlxcel_core::impl_rope_lifting_deserialize!(CodePredictorConfig); + fn default_cp_layers() -> usize { 5 } @@ -179,6 +185,7 @@ fn default_codec_think_eos() -> i32 { /// code2wav codec vocoder config (`code2wav_config` at the config.json root). #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Code2WavConfig { #[serde(default = "default_c2w_hidden")] pub hidden_size: usize, @@ -206,6 +213,8 @@ pub struct Code2WavConfig { pub upsampling_ratios: Vec, } +mlxcel_core::impl_rope_lifting_deserialize!(Code2WavConfig); + fn default_c2w_hidden() -> usize { 1024 } diff --git a/src/lib/mlxcel-core/src/drafter/dflash/config.rs b/src/lib/mlxcel-core/src/drafter/dflash/config.rs index 10e54008b..199297db5 100644 --- a/src/lib/mlxcel-core/src/drafter/dflash/config.rs +++ b/src/lib/mlxcel-core/src/drafter/dflash/config.rs @@ -303,8 +303,11 @@ impl DFlashConfig { } } - serde_json::from_value(serde_json::Value::Object(flat)) - .map_err(|e| format!("Failed to parse DFlashConfig: {e}")) + // transformers 5 keeps `rope_theta` only under `rope_parameters` + // (#2443); without the lift the drafter would run at the default base. + let mut flat = serde_json::Value::Object(flat); + crate::rope_parameters::lift_rope_parameters(&mut flat); + serde_json::from_value(flat).map_err(|e| format!("Failed to parse DFlashConfig: {e}")) } /// Whether this checkpoint is an LFM2 / LFM2.5 DSpark drafter: a DFlash diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index 4250402b5..56caf6ea5 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -4046,6 +4046,10 @@ pub mod memory; // Currently: proportional RoPE used by Gemma 4 full-attention layers. pub mod rope_proportional; +// transformers 5 `rope_parameters` lift shared by every config struct that +// reads the flat `rope_theta` / `rope_scaling` keys (#2443). +pub mod rope_parameters; + // How many positions the sequence currently being prefilled will span, so a // model whose RoPE table is chosen from the whole prompt length can still make // that choice from inside a chunked prefill (Phi-3 / Phi-4 LongRoPE, #1358). diff --git a/src/lib/mlxcel-core/src/rope_parameters.rs b/src/lib/mlxcel-core/src/rope_parameters.rs new file mode 100644 index 000000000..161862149 --- /dev/null +++ b/src/lib/mlxcel-core/src/rope_parameters.rs @@ -0,0 +1,166 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Lift a transformers 5 `rope_parameters` block into the flat RoPE keys. +//! +//! transformers 5 saves `rope_theta`, `rope_scaling` and +//! `partial_rotary_factor` only under a nested `rope_parameters` block (for +//! example `{"rope_theta": 1000000, "rope_type": "default"}`) and drops the flat +//! keys. A config struct that reads the flat `rope_theta` with a serde default +//! would then silently run at the default base, decoding fluently at short +//! lengths and drifting at long ones (#2384, #2443). +//! +//! The lift lives here rather than in the model crate so that drafter configs +//! in this crate share it. The model crate re-exports it from +//! `models::rope_utils`. + +use serde::{Deserialize, Deserializer}; +use serde_json::{Map, Value}; + +/// Copy a nested `rope_parameters` block into the flat `rope_theta`, +/// `partial_rotary_factor` and `rope_scaling` keys, when the flat keys are +/// absent. +/// +/// A flat key that is present and non-null always wins over the nested one. +/// +/// `rope_scaling` is filled (unless the config already has a non-null one) +/// when the block names a non-default scheme or carries keys beyond the base +/// and the scheme name, such as the MRoPE `mrope_section` that transformers 4 +/// kept in `rope_scaling` and transformers 5 moved here. It is the block +/// without `rope_theta`, with the scheme under `type` only. The scheme is read +/// `type` first, then `rope_type`, the order `RopeScalingSpec` and upstream +/// use. `type` is the key every reader in the tree accepts, while carrying +/// both spellings would turn a `#[serde(alias = "type", alias = "rope_type")]` +/// field into a `duplicate field` error. +/// +/// The per-layer-type form transformers 5 writes for Gemma 3, Gemma 3n and +/// OLMo 3 (`{"full_attention": {...}, "sliding_attention": {...}}`) is read +/// too: the `full_attention` block supplies the keys above, and the +/// `sliding_attention` base becomes `rope_local_base_freq`. +/// +/// Upstream: ml-explore/mlx-lm#1898 (qwen3) and ml-explore/mlx-lm#1605 +/// (granitemoehybrid). +pub fn lift_rope_parameters(config: &mut Value) { + let Some(obj) = config.as_object_mut() else { + return; + }; + let Some(params) = obj.get("rope_parameters").and_then(Value::as_object) else { + return; + }; + let block_at = |key: &str| params.get(key).and_then(Value::as_object).cloned(); + let (main, sliding) = + if params.contains_key("full_attention") || params.contains_key("sliding_attention") { + ( + block_at("full_attention").unwrap_or_default(), + block_at("sliding_attention"), + ) + } else { + (params.clone(), None) + }; + + copy_number_if_absent(obj, "rope_theta", main.get("rope_theta")); + copy_number_if_absent( + obj, + "partial_rotary_factor", + main.get("partial_rotary_factor"), + ); + if let Some(sliding) = &sliding { + copy_number_if_absent(obj, "rope_local_base_freq", sliding.get("rope_theta")); + } + + let kind = main + .get("type") + .or_else(|| main.get("rope_type")) + .and_then(Value::as_str) + .map(str::to_string); + let has_extra_keys = main.keys().any(|key| { + !matches!( + key.as_str(), + "rope_theta" | "rope_type" | "type" | "partial_rotary_factor" + ) + }); + let has_scaling = obj.get("rope_scaling").is_some_and(|v| !v.is_null()); + if !has_scaling && (kind.as_deref().is_some_and(|k| k != "default") || has_extra_keys) { + let mut scaling = main; + scaling.remove("rope_theta"); + scaling.remove("rope_type"); + if let Some(kind) = kind { + scaling.insert("type".to_string(), Value::String(kind)); + } + obj.insert("rope_scaling".to_string(), Value::Object(scaling)); + } +} + +fn copy_number_if_absent(obj: &mut Map, key: &str, nested: Option<&Value>) { + if obj.get(key).is_none_or(Value::is_null) + && let Some(value) = nested.filter(|v| v.is_number()) + { + obj.insert(key.to_string(), value.clone()); + } +} + +/// Deserialize a config through [`lift_rope_parameters`]. +/// +/// A config struct whose `Deserialize` derive carries `#[serde(remote = +/// "Self")]` (which turns the derive into an inherent `Type::deserialize` +/// function instead of the trait impl) gets the trait impl from +/// [`impl_rope_lifting_deserialize!`](crate::impl_rope_lifting_deserialize), +/// which calls this. Every `serde_json::from_value` / `from_str` / +/// nested-field parse of that type then sees the lifted keys, so a new parse +/// site cannot miss the lift. +pub fn deserialize_lifting_rope_parameters<'de, D, T>( + deserializer: D, + parse: impl FnOnce(Value) -> Result, +) -> Result +where + D: Deserializer<'de>, +{ + let mut value = Value::deserialize(deserializer)?; + lift_rope_parameters(&mut value); + parse(value).map_err(serde::de::Error::custom) +} + +/// Implement `Deserialize` for config structs through +/// [`deserialize_lifting_rope_parameters`]. +/// +/// Each type must derive `Deserialize` with `#[serde(remote = "Self")]`: +/// +/// ```ignore +/// #[derive(Debug, Clone, Deserialize)] +/// #[serde(remote = "Self")] +/// pub struct ModelArgs { +/// #[serde(default = "default_rope_theta")] +/// pub rope_theta: f32, +/// } +/// mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); +/// ``` +#[macro_export] +macro_rules! impl_rope_lifting_deserialize { + ($($ty:ty),+ $(,)?) => {$( + impl<'de> ::serde::Deserialize<'de> for $ty { + fn deserialize>( + deserializer: D, + ) -> ::core::result::Result { + $crate::rope_parameters::deserialize_lifting_rope_parameters( + deserializer, + <$ty>::deserialize, + ) + } + } + )+}; +} + +#[cfg(test)] +#[path = "rope_parameters_tests.rs"] +mod tests; diff --git a/src/lib/mlxcel-core/src/rope_parameters_tests.rs b/src/lib/mlxcel-core/src/rope_parameters_tests.rs new file mode 100644 index 000000000..6058f619a --- /dev/null +++ b/src/lib/mlxcel-core/src/rope_parameters_tests.rs @@ -0,0 +1,159 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::lift_rope_parameters; +use serde::Deserialize; +use serde_json::json; + +#[test] +fn nested_theta_and_partial_factor_fill_absent_flat_keys() { + let mut config = json!({ + "rope_parameters": {"rope_theta": 10000.0, "partial_rotary_factor": 0.4, "rope_type": "default"} + }); + lift_rope_parameters(&mut config); + assert_eq!(config["rope_theta"], 10000.0); + assert_eq!(config["partial_rotary_factor"], 0.4); + assert!(config.get("rope_scaling").is_none()); +} + +#[test] +fn flat_keys_win_and_null_flat_keys_do_not() { + let mut config = json!({ + "rope_theta": 500000.0, + "partial_rotary_factor": null, + "rope_parameters": {"rope_theta": 123456.0, "partial_rotary_factor": 0.5} + }); + lift_rope_parameters(&mut config); + assert_eq!(config["rope_theta"], 500000.0); + assert_eq!(config["partial_rotary_factor"], 0.5); +} + +#[test] +fn llama3_block_becomes_rope_scaling_under_type_only() { + // The shape transformers 5.20 writes for a Llama 3.2 config. + let mut config = json!({ + "rope_parameters": { + "factor": 32.0, "high_freq_factor": 4.0, "low_freq_factor": 1.0, + "original_max_position_embeddings": 8192, "rope_type": "llama3", + "rope_theta": 500000.0 + } + }); + lift_rope_parameters(&mut config); + assert_eq!(config["rope_theta"], 500000.0); + let scaling = &config["rope_scaling"]; + assert_eq!(scaling["type"], "llama3"); + assert_eq!(scaling["factor"], 32.0); + assert_eq!(scaling["original_max_position_embeddings"], 8192); + assert!(scaling.get("rope_theta").is_none()); + // One spelling only, so an alias-both field does not see a duplicate. + assert!(scaling.get("rope_type").is_none()); + + #[derive(Deserialize)] + struct AliasBoth { + #[serde(alias = "type", alias = "rope_type")] + scaling_type: Option, + } + let parsed: AliasBoth = serde_json::from_value(scaling.clone()).unwrap(); + assert_eq!(parsed.scaling_type.as_deref(), Some("llama3")); +} + +#[test] +fn existing_rope_scaling_is_kept() { + let mut config = json!({ + "rope_scaling": {"type": "linear", "factor": 2.0}, + "rope_parameters": {"rope_theta": 1000000.0, "rope_type": "yarn", "factor": 4.0} + }); + lift_rope_parameters(&mut config); + assert_eq!( + config["rope_scaling"], + json!({"type": "linear", "factor": 2.0}) + ); + assert_eq!(config["rope_theta"], 1000000.0); +} + +#[test] +fn per_layer_type_block_lifts_full_and_sliding_bases() { + // The shape transformers 5.20 writes for a Gemma 3 text config. + let mut config = json!({ + "rope_parameters": { + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": {"rope_type": "linear", "factor": 8.0, "rope_theta": 1000000.0} + } + }); + lift_rope_parameters(&mut config); + assert_eq!(config["rope_theta"], 1000000.0); + assert_eq!(config["rope_local_base_freq"], 10000.0); + assert_eq!( + config["rope_scaling"], + json!({"type": "linear", "factor": 8.0}) + ); +} + +#[test] +fn default_mrope_block_still_becomes_rope_scaling() { + // transformers 5.20 Qwen2-VL / Qwen3-VL text configs: the scheme is + // "default" but `mrope_section` used to live in `rope_scaling`. + let mut qwen2_vl = json!({ + "rope_parameters": { + "type": "mrope", "mrope_section": [16, 24, 24], + "rope_theta": 1000000.0, "rope_type": "default" + } + }); + lift_rope_parameters(&mut qwen2_vl); + assert_eq!(qwen2_vl["rope_theta"], 1000000.0); + assert_eq!( + qwen2_vl["rope_scaling"], + json!({"type": "mrope", "mrope_section": [16, 24, 24]}) + ); + + let mut qwen3_vl = json!({ + "rope_parameters": { + "rope_type": "default", "mrope_section": [24, 20, 20], + "mrope_interleaved": true, "rope_theta": 5000000.0 + } + }); + lift_rope_parameters(&mut qwen3_vl); + assert_eq!( + qwen3_vl["rope_scaling"], + json!({"type": "default", "mrope_section": [24, 20, 20], "mrope_interleaved": true}) + ); +} + +#[test] +fn non_object_and_unusable_shapes_are_left_alone() { + let mut array = json!([1, 2]); + lift_rope_parameters(&mut array); + assert_eq!(array, json!([1, 2])); + + let mut no_block = json!({"rope_parameters": null, "rope_theta": 10.0}); + let before = no_block.clone(); + lift_rope_parameters(&mut no_block); + assert_eq!(no_block, before); + + // A non-numeric base is not copied, and a per-layer-type block with no + // `full_attention` entry leaves the flat keys alone. + let mut string_theta = + json!({"rope_parameters": {"rope_theta": "1e6", "rope_type": "default"}}); + lift_rope_parameters(&mut string_theta); + assert!(string_theta.get("rope_theta").is_none()); + assert!(string_theta.get("rope_scaling").is_none()); + + let mut sliding_only = json!({ + "rope_parameters": {"sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}} + }); + lift_rope_parameters(&mut sliding_only); + assert!(sliding_only.get("rope_theta").is_none()); + assert!(sliding_only.get("rope_scaling").is_none()); + assert_eq!(sliding_only["rope_local_base_freq"], 10000.0); +} diff --git a/src/loading/vlm_idefics2.rs b/src/loading/vlm_idefics2.rs index 701728619..60336624a 100644 --- a/src/loading/vlm_idefics2.rs +++ b/src/loading/vlm_idefics2.rs @@ -65,6 +65,9 @@ fn resolve_added_token_id(added_tokens: Option<&Value>, name: &str) -> i32 { /// deserializes into [`models::llama3::ModelArgs`]. Values already present in the /// checkpoint win over these defaults. fn normalize_text_config(text_config: &mut Value, full_config: &Value) -> Result<()> { + // Lift a transformers 5 `rope_parameters` block first, so the Mistral + // `rope_theta` default below does not shadow the checkpoint's base (#2443). + models::rope_utils::lift_rope_parameters(text_config); let obj = super::require_object_mut(text_config, "Idefics2 text_config")?; // HuggingFaceM4/idefics2-8b text_config == Mistral-7B-v0.1. Any field the diff --git a/src/loading/vlm_qwen.rs b/src/loading/vlm_qwen.rs index acbc6eb8c..554e1c22c 100644 --- a/src/loading/vlm_qwen.rs +++ b/src/loading/vlm_qwen.rs @@ -1075,12 +1075,13 @@ fn drop_extra_text_layers(weights: &mut mlxcel_core::weights::WeightMap, num_hid } } -/// Lift GLM-OCR's `rope_parameters` block into the fields `Glm4vTextConfig` -/// deserializes: `mrope_section` under `rope_scaling`, plus top-level -/// `partial_rotary_factor` and `rope_theta`. GLM-OCR nests all three under -/// `rope_parameters` (not `rope_scaling`) and omits the top-level scalars, so -/// feeding the config through unchanged silently yields half-width rotary with -/// the wrong sections. +/// Lift GLM-OCR's `rope_parameters.mrope_section` into `rope_scaling`, the +/// field `Glm4vTextConfig` deserializes it from, also when a `rope_scaling` +/// object already exists. GLM-OCR nests it under `rope_parameters` (not +/// `rope_scaling`), so feeding the config through unchanged silently yields the +/// wrong sections. The top-level `partial_rotary_factor` and `rope_theta` come +/// from the shared transformers 5 lift in `Glm4vTextConfig`'s `Deserialize` +/// (#2443). fn normalize_glm_ocr_rope(text_config: &mut serde_json::Value) -> Result<()> { let obj = super::require_object_mut(text_config, "GLM-OCR text_config")?; let Some(rp) = obj.get("rope_parameters").cloned() else { @@ -1094,21 +1095,6 @@ fn normalize_glm_ocr_rope(text_config: &mut serde_json::Value) -> Result<()> { rs_obj.insert("mrope_section".to_string(), section); } } - if obj - .get("partial_rotary_factor") - .and_then(|v| v.as_f64()) - .is_none() - && let Some(prf) = rp.get("partial_rotary_factor").cloned() - { - obj.insert("partial_rotary_factor".to_string(), prf); - } - if obj.get("rope_theta").and_then(|v| v.as_f64()).is_none() { - let theta = rp - .get("rope_theta") - .and_then(|v| v.as_f64()) - .unwrap_or(10000.0); - obj.insert("rope_theta".to_string(), serde_json::json!(theta)); - } Ok(()) } @@ -1233,21 +1219,19 @@ pub(crate) fn load_glm4v_moe(model_path: &Path) -> Result { let (_config_str, full_config) = read_sanitized_vlm_config(model_path)?; - // Text config: normalize `rope_theta` (it may live under `rope_parameters`) - // and inherit top-level quantization, then reuse the shared GLM-4 MoE args. + // Text config: lift `rope_parameters` through the shared transformers 5 + // lift (#2443), default a still-missing `rope_theta` (required by the + // GLM-4 MoE args) to 10000, and inherit top-level quantization, then reuse + // the shared GLM-4 MoE args. let mut text_config_val = full_config .get("text_config") .cloned() .ok_or_else(|| anyhow::anyhow!("Missing text_config in config.json"))?; + models::rope_utils::lift_rope_parameters(&mut text_config_val); { let obj = super::require_object_mut(&mut text_config_val, "GLM-4V MoE text_config")?; if obj.get("rope_theta").and_then(|v| v.as_f64()).is_none() { - let theta = obj - .get("rope_parameters") - .and_then(|rp| rp.get("rope_theta")) - .and_then(|v| v.as_f64()) - .unwrap_or(10000.0); - obj.insert("rope_theta".to_string(), serde_json::json!(theta)); + obj.insert("rope_theta".to_string(), serde_json::json!(10000.0)); } if obj.get("quantization_config").is_none() && obj.get("group_size").is_none() diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index b78cb61fb..3653f8f93 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -66,6 +66,7 @@ const PHI4_FUSED_TEXT_FIELDS: &[&str] = &[ "hidden_size", "num_key_value_heads", "rope_theta", + "rope_parameters", "rope_traditional", "partial_rotary_factor", "rope_scaling", diff --git a/src/models/afmoe.rs b/src/models/afmoe.rs index 4d9bfdb99..3c2cadeef 100644 --- a/src/models/afmoe.rs +++ b/src/models/afmoe.rs @@ -130,6 +130,7 @@ const AFMOE_KV_SNAPSHOT_NAMES: KvSnapshotNames = /// AFMoE `config.json`. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -215,6 +216,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + /// A `config.json` token-id field, which may be a single int or a list of ints. #[derive(Debug, Clone, Deserialize)] #[serde(untagged)] diff --git a/src/models/apertus.rs b/src/models/apertus.rs index 7c90ab439..f0f325336 100644 --- a/src/models/apertus.rs +++ b/src/models/apertus.rs @@ -46,6 +46,7 @@ use std::time::Instant; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -91,6 +92,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/baichuan.rs b/src/models/baichuan.rs index c53fdf497..fc6f65ae2 100644 --- a/src/models/baichuan.rs +++ b/src/models/baichuan.rs @@ -29,6 +29,7 @@ use std::path::Path; // Config. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct BaichuanConfig { pub model_type: String, pub vocab_size: usize, @@ -52,6 +53,8 @@ pub struct BaichuanConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(BaichuanConfig); + #[derive(Debug, Clone, Deserialize)] pub struct QuantizationConfig { pub group_size: i32, diff --git a/src/models/bailing_moe.rs b/src/models/bailing_moe.rs index a9f5f1958..2a1efda4a 100644 --- a/src/models/bailing_moe.rs +++ b/src/models/bailing_moe.rs @@ -120,6 +120,7 @@ use crate::models::switch_layers::{ /// `use_qk_norm`, `partial_rotary_factor` and `rotary_dim` are all absent from /// that file. Declaring any of them required would fail the whole parse. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -243,6 +244,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + /// A `config.json` token-id field, which may be a single int or a list of ints. /// /// Same shape as the per-family enums in [`crate::models::gpt_neox`] and diff --git a/src/models/bailing_moe_linear.rs b/src/models/bailing_moe_linear.rs index 3956cbd37..1468e1856 100644 --- a/src/models/bailing_moe_linear.rs +++ b/src/models/bailing_moe_linear.rs @@ -149,6 +149,7 @@ const GLA_CHUNK: i32 = 64; /// built. Widening the dense family's config to carry them would make a /// `bailing_moe` checkpoint parse fields no `bailing_moe` layer reads. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -281,6 +282,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + /// A `config.json` token-id field, which may be a single int or a list of ints. #[derive(Debug, Clone, Deserialize)] #[serde(untagged)] diff --git a/src/models/bitnet.rs b/src/models/bitnet.rs index 8b6362031..ec695c8d4 100644 --- a/src/models/bitnet.rs +++ b/src/models/bitnet.rs @@ -57,6 +57,7 @@ fn reject_without_bitlinear_kernel() -> Result<(), String> { // Config. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct BitNetConfig { pub model_type: String, pub hidden_size: usize, @@ -83,6 +84,8 @@ pub struct BitNetConfig { pub quantization_config: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(BitNetConfig); + #[derive(Debug, Clone, Deserialize)] pub struct BitNetQuantConfig { #[serde(default)] diff --git a/src/models/cohere.rs b/src/models/cohere.rs index 136d356dd..24ace571a 100644 --- a/src/models/cohere.rs +++ b/src/models/cohere.rs @@ -29,6 +29,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -54,6 +55,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/cohere2.rs b/src/models/cohere2.rs index 9d40f279c..0bd5c313c 100644 --- a/src/models/cohere2.rs +++ b/src/models/cohere2.rs @@ -33,6 +33,7 @@ use std::path::Path; // Configuration. /// Cohere2 model configuration #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Cohere2Config { pub model_type: String, pub hidden_size: usize, @@ -75,6 +76,8 @@ pub struct Cohere2Config { pub eos_token_id: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(Cohere2Config); + #[derive(Debug, Clone, Deserialize)] #[serde(untagged)] pub enum EosTokenId { diff --git a/src/models/cohere2_moe.rs b/src/models/cohere2_moe.rs index 0be036acf..5708e4df5 100644 --- a/src/models/cohere2_moe.rs +++ b/src/models/cohere2_moe.rs @@ -52,6 +52,7 @@ use std::path::Path; /// overrides `moe_gate_act`, `prefix_dense_intermediate_size` defaults to /// `intermediate_size`. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Cohere2MoeConfig { pub model_type: String, @@ -136,6 +137,8 @@ pub struct Cohere2MoeConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(Cohere2MoeConfig); + /// `eos_token_id` may be a single int or a list of ints (reused shape from /// `src/models/cohere2.rs`). #[derive(Debug, Clone, Deserialize)] diff --git a/src/models/dbrx.rs b/src/models/dbrx.rs index 70bd2e7b7..555ababd9 100644 --- a/src/models/dbrx.rs +++ b/src/models/dbrx.rs @@ -83,6 +83,7 @@ pub struct Quantization { /// have nothing to do with attention), so only the four fields that matter are /// declared here and serde drops the rest. #[derive(Debug, Clone, Deserialize, Default)] +#[serde(remote = "Self")] pub struct AttnConfig { #[serde(default)] pub kv_n_heads: Option, @@ -97,6 +98,8 @@ pub struct AttnConfig { pub rope_theta: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(AttnConfig); + /// Feed-forward / MoE geometry, nested under `ffn_config`. Carries the same /// `PretrainedConfig` noise as [`AttnConfig`]. #[derive(Debug, Clone, Deserialize, Default)] diff --git a/src/models/deepseek.rs b/src/models/deepseek.rs index 941b80cd4..6ef1e245b 100644 --- a/src/models/deepseek.rs +++ b/src/models/deepseek.rs @@ -35,6 +35,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -103,6 +104,8 @@ pub struct ModelArgs { pub bits: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_moe_layer_freq() -> usize { 1 } diff --git a/src/models/deepseek_v2.rs b/src/models/deepseek_v2.rs index 2991055d6..8df243df7 100644 --- a/src/models/deepseek_v2.rs +++ b/src/models/deepseek_v2.rs @@ -35,6 +35,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -115,6 +116,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/deepseek_v3.rs b/src/models/deepseek_v3.rs index 6fb6a311e..c36e34e81 100644 --- a/src/models/deepseek_v3.rs +++ b/src/models/deepseek_v3.rs @@ -33,6 +33,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct DeepSeekV3Config { #[serde(default = "default_model_type")] pub model_type: String, @@ -108,6 +109,8 @@ pub struct DeepSeekV3Config { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(DeepSeekV3Config); + #[derive(Debug, Clone, Deserialize)] pub struct QuantizationConfig { pub group_size: i32, diff --git a/src/models/deepseek_v32.rs b/src/models/deepseek_v32.rs index 704877c23..a55dc9173 100644 --- a/src/models/deepseek_v32.rs +++ b/src/models/deepseek_v32.rs @@ -61,6 +61,7 @@ pub struct RopeScaling { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -160,6 +161,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/deepseek_v4.rs b/src/models/deepseek_v4.rs index 08d5fcdf4..50d89c6a3 100644 --- a/src/models/deepseek_v4.rs +++ b/src/models/deepseek_v4.rs @@ -147,6 +147,7 @@ pub(crate) enum ScoringFunc { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -237,6 +238,8 @@ pub struct ModelArgs { pub quantization_config: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_model_type() -> String { "deepseek_v4".to_string() } diff --git a/src/models/dots1.rs b/src/models/dots1.rs index 451f8708c..9653cbb39 100644 --- a/src/models/dots1.rs +++ b/src/models/dots1.rs @@ -61,6 +61,7 @@ pub struct Quantization { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -113,6 +114,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_max_position_embeddings() -> usize { 32768 } diff --git a/src/models/ernie4_5.rs b/src/models/ernie4_5.rs index 10f1fb8f4..c2d62d269 100644 --- a/src/models/ernie4_5.rs +++ b/src/models/ernie4_5.rs @@ -28,6 +28,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -59,6 +60,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/ernie4_5_moe.rs b/src/models/ernie4_5_moe.rs index 648f1ee16..a4672e610 100644 --- a/src/models/ernie4_5_moe.rs +++ b/src/models/ernie4_5_moe.rs @@ -31,6 +31,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -85,6 +86,8 @@ pub struct ModelArgs { pub bits: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_rope_theta() -> f32 { 10000.0 } diff --git a/src/models/ernie4_5_moe_vl.rs b/src/models/ernie4_5_moe_vl.rs index 437cff35a..e14ab041a 100644 --- a/src/models/ernie4_5_moe_vl.rs +++ b/src/models/ernie4_5_moe_vl.rs @@ -93,6 +93,7 @@ impl IntOrPair { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Ernie45MoeVlTextConfig { pub hidden_size: usize, pub intermediate_size: usize, @@ -137,6 +138,8 @@ pub struct Ernie45MoeVlTextConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(Ernie45MoeVlTextConfig); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(default)] diff --git a/src/models/exaone.rs b/src/models/exaone.rs index e35f97e77..d64bd5653 100644 --- a/src/models/exaone.rs +++ b/src/models/exaone.rs @@ -29,6 +29,7 @@ use std::path::Path; // Config. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ExaOneConfig { pub model_type: String, pub hidden_size: usize, @@ -59,6 +60,8 @@ pub struct ExaOneConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ExaOneConfig); + #[derive(Debug, Clone, Deserialize)] pub struct QuantizationConfig { pub group_size: i32, diff --git a/src/models/exaone4.rs b/src/models/exaone4.rs index 544289e0a..42dc06937 100644 --- a/src/models/exaone4.rs +++ b/src/models/exaone4.rs @@ -30,6 +30,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -43,6 +44,9 @@ pub struct ModelArgs { pub rope_theta: f32, pub head_dim: usize, pub tie_word_embeddings: bool, + /// Absent from a transformers 5 save whose `rope_parameters` names the + /// default scheme (#2443); an empty map reads as no scaling. + #[serde(default)] pub rope_scaling: std::collections::HashMap, #[serde(default)] @@ -55,6 +59,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, @@ -569,6 +575,7 @@ fn compute_rope_freqs(args: &ModelArgs) -> Option> { let scaling = &args.rope_scaling; let rope_type = scaling .get("rope_type") + .or_else(|| scaling.get("type")) .and_then(|v| v.as_str()) .unwrap_or("default"); diff --git a/src/models/exaone_moe.rs b/src/models/exaone_moe.rs index 99f684af2..fa20dc10c 100644 --- a/src/models/exaone_moe.rs +++ b/src/models/exaone_moe.rs @@ -33,6 +33,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -89,6 +90,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, @@ -119,13 +122,9 @@ fn default_rope_theta() -> f32 { impl ModelArgs { pub fn get_rope_theta(&self) -> f32 { - // Check rope_parameters for rope_theta override - if let Some(ref params) = self.rope_parameters - && let Some(v) = params.get("rope_theta") - && let Some(f) = v.as_f64() - { - return f as f32; - } + // The `Deserialize` impl lifts `rope_parameters.rope_theta` into the + // flat field when the flat key is absent, so a flat key wins, matching + // every other family (#2443). self.rope_theta } diff --git a/src/models/falcon_h1.rs b/src/models/falcon_h1.rs index 851996fd2..2778f0a5f 100644 --- a/src/models/falcon_h1.rs +++ b/src/models/falcon_h1.rs @@ -77,6 +77,7 @@ pub struct Quantization { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -136,6 +137,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_head_dim() -> usize { 64 } diff --git a/src/models/falcon_ocr.rs b/src/models/falcon_ocr.rs index dcb55269a..fcf103448 100644 --- a/src/models/falcon_ocr.rs +++ b/src/models/falcon_ocr.rs @@ -69,6 +69,7 @@ use super::falcon_ocr_rope::{ /// scheme rather than the HF `hidden_size` / `num_hidden_layers` / /// `intermediate_size` names, so every field accepts both spellings. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct FalconOcrConfig { #[serde(alias = "hidden_size")] pub dim: usize, @@ -118,6 +119,8 @@ pub struct FalconOcrConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(FalconOcrConfig); + #[derive(Debug, Clone, Deserialize)] pub struct QuantConfig { #[serde(default)] diff --git a/src/models/gemma.rs b/src/models/gemma.rs index b5cf6e927..481ea70f8 100644 --- a/src/models/gemma.rs +++ b/src/models/gemma.rs @@ -30,6 +30,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -51,6 +52,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/gemma2.rs b/src/models/gemma2.rs index cabad13b3..27ec7c04f 100644 --- a/src/models/gemma2.rs +++ b/src/models/gemma2.rs @@ -33,6 +33,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -59,6 +60,8 @@ pub struct ModelArgs { pub tie_word_embeddings: bool, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/gemma3.rs b/src/models/gemma3.rs index 578029353..9a3188b43 100644 --- a/src/models/gemma3.rs +++ b/src/models/gemma3.rs @@ -46,6 +46,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] #[serde(default)] pub struct ModelArgs { pub model_type: String, @@ -71,6 +72,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + impl Default for ModelArgs { fn default() -> Self { // Defaults match Python mlx-vlm TextConfig for Gemma3 diff --git a/src/models/gemma3n.rs b/src/models/gemma3n.rs index 900b7a700..60717d116 100644 --- a/src/models/gemma3n.rs +++ b/src/models/gemma3n.rs @@ -38,6 +38,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct TextConfig { pub model_type: String, pub hidden_size: usize, @@ -68,6 +69,8 @@ pub struct TextConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(TextConfig); + impl TextConfig { pub fn get_intermediate_size(&self, layer_idx: usize) -> usize { match &self.intermediate_size { diff --git a/src/models/glm4.rs b/src/models/glm4.rs index 545c9ceb7..38fd2beca 100644 --- a/src/models/glm4.rs +++ b/src/models/glm4.rs @@ -29,6 +29,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -55,6 +56,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/glm4_moe.rs b/src/models/glm4_moe.rs index 14bbaf523..8034e1ece 100644 --- a/src/models/glm4_moe.rs +++ b/src/models/glm4_moe.rs @@ -42,6 +42,7 @@ fn default_one() -> usize { // Configuration. // Used by: GLM4 MoE, Solar Open #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -107,6 +108,8 @@ pub struct ModelArgs { pub quantization_config: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct QuantizationConfig { #[serde(default)] diff --git a/src/models/glm4_moe_lite.rs b/src/models/glm4_moe_lite.rs index e79e0ddf4..8a4acfde8 100644 --- a/src/models/glm4_moe_lite.rs +++ b/src/models/glm4_moe_lite.rs @@ -45,6 +45,7 @@ pub use glm4_moe_lite_sanitize::{kv_b_proj_geometry, sanitize_weights}; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -110,6 +111,8 @@ pub struct ModelArgs { pub num_nextn_predict_layers: usize, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_routed_scaling() -> f32 { 1.0 } diff --git a/src/models/glm4v.rs b/src/models/glm4v.rs index ecc572287..daa9a43ab 100644 --- a/src/models/glm4v.rs +++ b/src/models/glm4v.rs @@ -34,6 +34,7 @@ use serde::Deserialize; // Config. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Glm4vTextConfig { pub hidden_size: usize, pub num_hidden_layers: usize, @@ -62,6 +63,8 @@ pub struct Glm4vTextConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(Glm4vTextConfig); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(default)] diff --git a/src/models/gpt_oss.rs b/src/models/gpt_oss.rs index 673d07698..7e28f9c0c 100644 --- a/src/models/gpt_oss.rs +++ b/src/models/gpt_oss.rs @@ -41,6 +41,7 @@ use crate::distributed::pipeline::partial_loading::filter_weight_map; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -72,6 +73,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + /// Replace layer number with 0 for quantization config lookup fn regex_replace_layer(prefix: &str) -> String { // "model.layers.15.self_attn.q_proj" -> "model.layers.0.self_attn.q_proj" diff --git a/src/models/granite.rs b/src/models/granite.rs index e5aebb85b..b0f773f6c 100644 --- a/src/models/granite.rs +++ b/src/models/granite.rs @@ -34,6 +34,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -75,6 +76,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/helium.rs b/src/models/helium.rs index 3f66bde0b..4c352f545 100644 --- a/src/models/helium.rs +++ b/src/models/helium.rs @@ -68,6 +68,7 @@ use super::llama3::Llama3Model; /// source of the stop token) and the standard `quantization` block that an /// mlx-community conversion adds. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -115,6 +116,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + /// A `config.json` token-id field, which may be a single int or a list of ints. /// /// Same shape as the per-family enums in [`crate::models::gpt2`] and diff --git a/src/models/hunyuan_moe.rs b/src/models/hunyuan_moe.rs index 5f7a25b50..1f50b0785 100644 --- a/src/models/hunyuan_moe.rs +++ b/src/models/hunyuan_moe.rs @@ -31,6 +31,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -85,6 +86,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { pub alpha: f32, diff --git a/src/models/hunyuan_v1_dense.rs b/src/models/hunyuan_v1_dense.rs index 0751b8b79..7c419877e 100644 --- a/src/models/hunyuan_v1_dense.rs +++ b/src/models/hunyuan_v1_dense.rs @@ -28,6 +28,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -63,6 +64,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { pub alpha: f32, diff --git a/src/models/hunyuan_vl.rs b/src/models/hunyuan_vl.rs index f9abab601..3017524c4 100644 --- a/src/models/hunyuan_vl.rs +++ b/src/models/hunyuan_vl.rs @@ -40,6 +40,7 @@ use serde::Deserialize; // Config. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct HunyuanVlTextConfig { pub hidden_size: usize, pub num_hidden_layers: usize, @@ -66,6 +67,8 @@ pub struct HunyuanVlTextConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(HunyuanVlTextConfig); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(default)] diff --git a/src/models/internlm2.rs b/src/models/internlm2.rs index 73b391b70..52ca7f1cc 100644 --- a/src/models/internlm2.rs +++ b/src/models/internlm2.rs @@ -51,6 +51,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -89,6 +90,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/internlm3.rs b/src/models/internlm3.rs index e971c0de8..778d93c75 100644 --- a/src/models/internlm3.rs +++ b/src/models/internlm3.rs @@ -52,6 +52,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -93,6 +94,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/iquestloopcoder.rs b/src/models/iquestloopcoder.rs index 265d0fe81..e9d10d652 100644 --- a/src/models/iquestloopcoder.rs +++ b/src/models/iquestloopcoder.rs @@ -125,6 +125,7 @@ pub struct Quantization { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -170,6 +171,8 @@ pub struct ModelArgs { pub checkpoint_label: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + impl ModelArgs { pub fn head_dim(&self) -> usize { self.head_dim diff --git a/src/models/jina_vlm.rs b/src/models/jina_vlm.rs index 9ee96f780..49d82e853 100644 --- a/src/models/jina_vlm.rs +++ b/src/models/jina_vlm.rs @@ -193,9 +193,17 @@ impl JinaVlmTextConfig { .and_then(|l| l.get("eps")) .and_then(|v| v.as_f64()) .unwrap_or(d.rms_norm_eps as f64) as f32, + // A transformers 5 save keeps the base only under + // `rope_parameters`; a flat key still wins (#2443). rope_theta: text_config .get("rope_theta") .and_then(|v| v.as_f64()) + .or_else(|| { + text_config + .get("rope_parameters") + .and_then(|rp| rp.get("rope_theta")) + .and_then(|v| v.as_f64()) + }) .unwrap_or(d.rope_theta as f64) as f32, use_qk_norm: attn .and_then(|a| a.get("q_lnorm")) diff --git a/src/models/kimi_linear.rs b/src/models/kimi_linear.rs index 41446793f..ae62c56af 100644 --- a/src/models/kimi_linear.rs +++ b/src/models/kimi_linear.rs @@ -57,6 +57,7 @@ fn default_conv_kernel() -> usize { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct KimiLinearConfig { pub model_type: String, pub vocab_size: usize, @@ -105,6 +106,8 @@ pub struct KimiLinearConfig { pub eos_token_id: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(KimiLinearConfig); + fn default_one_usize() -> usize { 1 } diff --git a/src/models/klear.rs b/src/models/klear.rs index 1bd6db499..8a275869d 100644 --- a/src/models/klear.rs +++ b/src/models/klear.rs @@ -90,6 +90,7 @@ use crate::models::switch_layers::{SwitchGLU, fused_moe_enabled, moe_weighted_su /// Klear `config.json`. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -156,6 +157,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + /// A `config.json` token-id field, which may be a single int or a list of ints. /// /// `Klear-46B-A2.5B-Instruct` writes the list form, `[151645, 151643]`. diff --git a/src/models/llada2_moe/mod.rs b/src/models/llada2_moe/mod.rs index 46fe87f98..a45b80157 100644 --- a/src/models/llada2_moe/mod.rs +++ b/src/models/llada2_moe/mod.rs @@ -145,6 +145,7 @@ fn default_max_position_embeddings() -> usize { /// loads correctly. Derived values (`head_dim`, `rotary_dim`, `eos_token_id`) /// are resolved by the accessor methods below. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -211,6 +212,8 @@ pub struct ModelArgs { pub bits: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + impl ModelArgs { /// Per-head dimension, defaulting to `hidden_size / num_attention_heads`. pub fn head_dim(&self) -> usize { diff --git a/src/models/llama3.rs b/src/models/llama3.rs index c364ed277..6588f869e 100644 --- a/src/models/llama3.rs +++ b/src/models/llama3.rs @@ -30,6 +30,7 @@ use crate::models::rope_utils::RopeScalingKind; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -120,6 +121,8 @@ pub struct ModelArgs { pub checkpoint_label: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + /// Parse `rope_traditional`, mapping an explicit JSON `null` to `false`. /// /// See [`ModelArgs::rope_traditional`] for why this one field tolerates `null` diff --git a/src/models/llama4.rs b/src/models/llama4.rs index 47aded1a2..7370a7de1 100644 --- a/src/models/llama4.rs +++ b/src/models/llama4.rs @@ -216,6 +216,7 @@ impl Llama4Cache { // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct TextArgs { pub model_type: String, pub hidden_size: usize, @@ -258,6 +259,8 @@ pub struct TextArgs { pub bits: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(TextArgs); + fn default_rope_theta() -> f32 { 500000.0 } diff --git a/src/models/longcat_flash_ngram.rs b/src/models/longcat_flash_ngram.rs index 84ff3e3fd..67ca53f8b 100644 --- a/src/models/longcat_flash_ngram.rs +++ b/src/models/longcat_flash_ngram.rs @@ -44,6 +44,7 @@ pub struct QuantizationConfig { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct LongcatFlashNgramConfig { pub model_type: String, pub hidden_size: usize, @@ -88,6 +89,8 @@ pub struct LongcatFlashNgramConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(LongcatFlashNgramConfig); + fn default_zero_expert_type() -> String { "identity".to_string() } diff --git a/src/models/mimo.rs b/src/models/mimo.rs index 78e248ad2..57f80042a 100644 --- a/src/models/mimo.rs +++ b/src/models/mimo.rs @@ -28,6 +28,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -57,6 +58,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/minicpm.rs b/src/models/minicpm.rs index a023d282b..86cdc0908 100644 --- a/src/models/minicpm.rs +++ b/src/models/minicpm.rs @@ -29,6 +29,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -58,6 +59,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/minicpm3.rs b/src/models/minicpm3.rs index 58682a474..59e5101cf 100644 --- a/src/models/minicpm3.rs +++ b/src/models/minicpm3.rs @@ -31,6 +31,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -69,6 +70,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/minimax.rs b/src/models/minimax.rs index 37bcf45c1..eb33523b2 100644 --- a/src/models/minimax.rs +++ b/src/models/minimax.rs @@ -35,6 +35,7 @@ use std::path::Path; // ============================================================================ #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -60,6 +61,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_true() -> bool { true } diff --git a/src/models/minimax_m3_config.rs b/src/models/minimax_m3_config.rs index 159b064c8..78101475d 100644 --- a/src/models/minimax_m3_config.rs +++ b/src/models/minimax_m3_config.rs @@ -71,6 +71,7 @@ pub struct Quantization { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -149,6 +150,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_model_type() -> String { "minimax_m3".to_string() } diff --git a/src/models/mistral4.rs b/src/models/mistral4.rs index 1d5611f04..cf1715558 100644 --- a/src/models/mistral4.rs +++ b/src/models/mistral4.rs @@ -142,13 +142,16 @@ impl Mistral4Config { } pub fn rope_base(&self) -> f32 { - // First try rope_parameters.rope_theta, then fall back to rope_theta field - self.rope_parameters - .as_ref() - .and_then(|p| p.get("rope_theta")) - .and_then(|v| v.as_f64()) - .map(|v| v as f32) - .or(self.rope_theta) + // A flat `rope_theta` wins over `rope_parameters.rope_theta`, the same + // rule as the shared transformers 5 lift (#2443). + self.rope_theta + .or_else(|| { + self.rope_parameters + .as_ref() + .and_then(|p| p.get("rope_theta")) + .and_then(|v| v.as_f64()) + .map(|v| v as f32) + }) .unwrap_or(10000.0) } diff --git a/src/models/mixtral.rs b/src/models/mixtral.rs index abc30d5be..712c90dd2 100644 --- a/src/models/mixtral.rs +++ b/src/models/mixtral.rs @@ -31,6 +31,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -56,6 +57,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/mllama/config.rs b/src/models/mllama/config.rs index 735a997bb..7aa298172 100644 --- a/src/models/mllama/config.rs +++ b/src/models/mllama/config.rs @@ -27,6 +27,7 @@ use serde::Deserialize; /// replaced by gated cross-attention adapters that attend to the vision /// tower's features. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct MllamaTextConfig { #[serde(default = "default_text_model_type")] pub model_type: String, @@ -67,6 +68,8 @@ pub struct MllamaTextConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(MllamaTextConfig); + impl MllamaTextConfig { /// Head dimension (`hidden_size / num_attention_heads` unless overridden). pub fn head_dim(&self) -> usize { diff --git a/src/models/molmo.rs b/src/models/molmo.rs index 820231f60..6242242bf 100644 --- a/src/models/molmo.rs +++ b/src/models/molmo.rs @@ -52,6 +52,7 @@ pub use super::molmo2::{Molmo2Embedding, Quantization}; // from the reference `TextConfig` dataclass so the parser tolerates the flat // schema as well as a nested `text_config`. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct MolmoTextConfig { #[serde(default = "default_model_type")] pub model_type: String, @@ -88,6 +89,8 @@ pub struct MolmoTextConfig { pub tie_word_embeddings: bool, } +mlxcel_core::impl_rope_lifting_deserialize!(MolmoTextConfig); + fn default_model_type() -> String { "molmo".to_string() } diff --git a/src/models/molmo2.rs b/src/models/molmo2.rs index 0f45263e3..d0afc722d 100644 --- a/src/models/molmo2.rs +++ b/src/models/molmo2.rs @@ -32,6 +32,7 @@ use serde::Deserialize; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Molmo2TextConfig { #[serde(default = "default_model_type")] pub model_type: String, @@ -63,6 +64,8 @@ pub struct Molmo2TextConfig { pub tie_word_embeddings: bool, } +mlxcel_core::impl_rope_lifting_deserialize!(Molmo2TextConfig); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/moondream2.rs b/src/models/moondream2.rs index a048301af..00c825f8c 100644 --- a/src/models/moondream2.rs +++ b/src/models/moondream2.rs @@ -38,6 +38,7 @@ use mlxcel_core::{MlxArray, UniquePtr}; use serde::Deserialize; #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -71,6 +72,8 @@ pub struct ModelArgs { pub bos_token_id: i32, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_model_type() -> String { "moondream2".to_string() } diff --git a/src/models/nanochat_config.rs b/src/models/nanochat_config.rs index 2d42937c8..5fb88600c 100644 --- a/src/models/nanochat_config.rs +++ b/src/models/nanochat_config.rs @@ -18,6 +18,7 @@ use serde::Deserialize; /// nanochat `config.json`. Every optional key defaults to the d20 value. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -58,6 +59,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + /// `eos_token_id` may be a single int or a list of ints. #[derive(Debug, Clone, Deserialize)] #[serde(untagged)] diff --git a/src/models/nemotron.rs b/src/models/nemotron.rs index 38bd8fc25..f240cec6f 100644 --- a/src/models/nemotron.rs +++ b/src/models/nemotron.rs @@ -33,6 +33,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -70,6 +71,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(rename = "type", alias = "rope_type")] diff --git a/src/models/nemotron_nas.rs b/src/models/nemotron_nas.rs index 74ec7bf7d..789841988 100644 --- a/src/models/nemotron_nas.rs +++ b/src/models/nemotron_nas.rs @@ -106,6 +106,7 @@ pub struct BlockConfig { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct NemotronNASConfig { #[serde(default = "default_model_type")] pub model_type: String, @@ -158,6 +159,8 @@ pub struct NemotronNASConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(NemotronNASConfig); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/olmo.rs b/src/models/olmo.rs index 72d443149..7ab3adf11 100644 --- a/src/models/olmo.rs +++ b/src/models/olmo.rs @@ -29,6 +29,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub d_model: usize, @@ -54,6 +55,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_mlp_ratio() -> usize { 4 } diff --git a/src/models/olmo2.rs b/src/models/olmo2.rs index e5fce241b..6b4ca0092 100644 --- a/src/models/olmo2.rs +++ b/src/models/olmo2.rs @@ -29,6 +29,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -63,6 +64,8 @@ pub struct ModelArgs { pub tie_word_embeddings: bool, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(rename = "type", default)] diff --git a/src/models/olmo3.rs b/src/models/olmo3.rs index 96c5bfe55..c28d75dd1 100644 --- a/src/models/olmo3.rs +++ b/src/models/olmo3.rs @@ -32,6 +32,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct OLMo3Config { pub model_type: String, pub hidden_size: usize, @@ -70,6 +71,8 @@ pub struct OLMo3Config { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(OLMo3Config); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(rename = "type", default)] diff --git a/src/models/olmoe.rs b/src/models/olmoe.rs index a4499b69e..6fdc3934a 100644 --- a/src/models/olmoe.rs +++ b/src/models/olmoe.rs @@ -30,6 +30,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -67,6 +68,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/paddleocr_vl.rs b/src/models/paddleocr_vl.rs index 151c35643..892ca6289 100644 --- a/src/models/paddleocr_vl.rs +++ b/src/models/paddleocr_vl.rs @@ -36,6 +36,7 @@ use serde::Deserialize; // Config. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct PaddleOcrTextConfig { pub hidden_size: usize, pub num_hidden_layers: usize, @@ -60,6 +61,8 @@ pub struct PaddleOcrTextConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(PaddleOcrTextConfig); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(default)] diff --git a/src/models/phi.rs b/src/models/phi.rs index 1510c9a61..80cd974cd 100644 --- a/src/models/phi.rs +++ b/src/models/phi.rs @@ -31,6 +31,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -69,6 +70,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/phi3.rs b/src/models/phi3.rs index 685204108..c2f272b19 100644 --- a/src/models/phi3.rs +++ b/src/models/phi3.rs @@ -28,6 +28,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -72,6 +73,8 @@ pub struct ModelArgs { pub original_max_position_embeddings: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(rename = "type", default)] diff --git a/src/models/phimoe.rs b/src/models/phimoe.rs index c6337f327..e560f5e82 100644 --- a/src/models/phimoe.rs +++ b/src/models/phimoe.rs @@ -31,6 +31,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -75,6 +76,8 @@ pub struct ModelArgs { pub quantization_config: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/phixtral.rs b/src/models/phixtral.rs index 014a4df3b..f3bd2ef45 100644 --- a/src/models/phixtral.rs +++ b/src/models/phixtral.rs @@ -144,6 +144,7 @@ use crate::models::switch_layers::{SwitchLinear, gather_sort, moe_weighted_sum, /// reads only the alias and therefore falls back to a default for all four; see /// the module docs. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -201,6 +202,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/qwen2_moe.rs b/src/models/qwen2_moe.rs index e51ac40e4..0feabb6f6 100644 --- a/src/models/qwen2_moe.rs +++ b/src/models/qwen2_moe.rs @@ -30,6 +30,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -65,6 +66,8 @@ pub struct ModelArgs { pub head_dim: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + fn default_rope_theta() -> f32 { 1000000.0 } diff --git a/src/models/qwen2_vl.rs b/src/models/qwen2_vl.rs index 9b3a02537..e08c17c8b 100644 --- a/src/models/qwen2_vl.rs +++ b/src/models/qwen2_vl.rs @@ -29,6 +29,7 @@ use serde::Deserialize; // Config. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Qwen2VLConfig { pub hidden_size: usize, pub num_hidden_layers: usize, @@ -51,6 +52,8 @@ pub struct Qwen2VLConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(Qwen2VLConfig); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(default)] diff --git a/src/models/qwen3_moe.rs b/src/models/qwen3_moe.rs index beb28e5c6..453c7f48e 100644 --- a/src/models/qwen3_moe.rs +++ b/src/models/qwen3_moe.rs @@ -37,6 +37,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -78,6 +79,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/qwen3_next.rs b/src/models/qwen3_next.rs index cd93ff5f3..379f92a9c 100644 --- a/src/models/qwen3_next.rs +++ b/src/models/qwen3_next.rs @@ -62,6 +62,7 @@ pub struct Quantization { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Qwen3NextConfig { pub model_type: String, pub hidden_size: usize, @@ -118,6 +119,8 @@ pub struct Qwen3NextConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(Qwen3NextConfig); + fn default_rms_norm_eps() -> f32 { 1e-6 } diff --git a/src/models/qwen3_vl.rs b/src/models/qwen3_vl.rs index 453e24d88..2856085c3 100644 --- a/src/models/qwen3_vl.rs +++ b/src/models/qwen3_vl.rs @@ -33,6 +33,7 @@ use std::cell::RefCell; // Config. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Qwen3VLConfig { pub hidden_size: usize, pub num_hidden_layers: usize, @@ -57,6 +58,8 @@ pub struct Qwen3VLConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(Qwen3VLConfig); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(default)] diff --git a/src/models/qwen3_vl_moe.rs b/src/models/qwen3_vl_moe.rs index d3d272358..1f3e08569 100644 --- a/src/models/qwen3_vl_moe.rs +++ b/src/models/qwen3_vl_moe.rs @@ -38,6 +38,7 @@ use std::cell::RefCell; // Config. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Qwen3VLMoeConfig { pub hidden_size: usize, pub num_hidden_layers: usize, @@ -75,6 +76,8 @@ pub struct Qwen3VLMoeConfig { pub head_dim: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(Qwen3VLMoeConfig); + #[derive(Debug, Clone, Deserialize)] pub struct RopeScaling { #[serde(default)] diff --git a/src/models/recurrent_gemma.rs b/src/models/recurrent_gemma.rs index 678277c86..e5f1959db 100644 --- a/src/models/recurrent_gemma.rs +++ b/src/models/recurrent_gemma.rs @@ -45,6 +45,7 @@ pub struct Quantization { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct GriffinConfig { pub model_type: String, pub hidden_size: usize, @@ -84,6 +85,8 @@ pub struct GriffinConfig { pub eos_token_id: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(GriffinConfig); + fn default_rms_norm_eps() -> f32 { 1e-6 } diff --git a/src/models/rope_parameters_lift_tests.rs b/src/models/rope_parameters_lift_tests.rs new file mode 100644 index 000000000..cb2382549 --- /dev/null +++ b/src/models/rope_parameters_lift_tests.rs @@ -0,0 +1,385 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Every config struct that reads the flat `rope_theta` must also read a +//! transformers 5 `rope_parameters` block, and a flat key must still win +//! (#2443). +//! +//! The table holds one parse function per struct. Rather than a hand-written +//! minimal config per family, [`fill_required_fields`] grows a config from +//! the serde errors themselves: each `missing field` gets a placeholder, and a +//! placeholder of the wrong type is replaced by one matching the `expected` +//! clause of the next error. `rope_theta` and `rope_scaling` are never +//! fabricated, so a struct that requires either flat key and does not lift it +//! fails here instead of being papered over. + +use serde_json::{Value, json}; + +type ParseTheta = fn(Value) -> Result, String>; + +trait ThetaValue { + fn theta(&self) -> Option; +} + +impl ThetaValue for f32 { + fn theta(&self) -> Option { + Some(f64::from(*self)) + } +} + +impl ThetaValue for f64 { + fn theta(&self) -> Option { + Some(*self) + } +} + +impl ThetaValue for Option { + fn theta(&self) -> Option { + self.as_ref().and_then(ThetaValue::theta) + } +} + +impl ThetaValue for Value { + fn theta(&self) -> Option { + self.as_f64() + } +} + +/// A table row: the type name and a parse function returning its RoPE base. +/// The one-argument form reads the `rope_theta` field; the two-argument form +/// takes an accessor for families that resolve the base through a method. +macro_rules! case { + ($ty:ty) => { + case!($ty, |c| c.rope_theta.theta()) + }; + ($ty:ty, $get:expr) => { + ( + stringify!($ty), + (|value: Value| { + let get: fn(&$ty) -> Option = $get; + serde_json::from_value::<$ty>(value) + .map(|config| get(&config)) + .map_err(|e| e.to_string()) + }) as ParseTheta, + ) + }; +} + +fn cases() -> Vec<(&'static str, ParseTheta)> { + use crate::audio::qwen3_omni_moe::speech_config as omni; + use crate::models::*; + vec![ + // Shared Llama args: Llama, Qwen2, IQuestCoder, the Mistral3 fallback. + case!(llama3::ModelArgs), + // Qwen / Gemma. + case!(qwen2_moe::ModelArgs), + case!(qwen2_vl::Qwen2VLConfig), + case!(qwen3_moe::ModelArgs), + case!(qwen3_next::Qwen3NextConfig), + case!(qwen3_vl::Qwen3VLConfig), + case!(qwen3_vl_moe::Qwen3VLMoeConfig), + case!(gemma::ModelArgs), + case!(gemma2::ModelArgs), + case!(gemma3::ModelArgs), + case!(gemma3n::TextConfig), + case!(recurrent_gemma::GriffinConfig), + // Mistral / Phi / OLMo / Granite. + case!(mixtral::ModelArgs), + case!(phi::ModelArgs), + case!(phi3::ModelArgs), + case!(phimoe::ModelArgs), + case!(phixtral::ModelArgs), + case!(olmo::ModelArgs), + case!(olmo2::ModelArgs), + case!(olmo3::OLMo3Config), + case!(olmoe::ModelArgs), + case!(granite::ModelArgs), + // DeepSeek / GLM / Kimi. + case!(deepseek::ModelArgs), + case!(deepseek_v2::ModelArgs), + case!(deepseek_v3::DeepSeekV3Config), + case!(deepseek_v32::ModelArgs), + case!(deepseek_v4::ModelArgs), + case!(glm4::ModelArgs), + case!(glm4_moe::ModelArgs), + case!(glm4_moe_lite::ModelArgs), + case!(glm4v::Glm4vTextConfig), + case!(kimi_linear::KimiLinearConfig), + case!(dots1::ModelArgs), + // Other text families. + case!(afmoe::ModelArgs), + case!(apertus::ModelArgs), + case!(baichuan::BaichuanConfig), + case!(bailing_moe::ModelArgs), + case!(bailing_moe_linear::ModelArgs), + case!(bitnet::BitNetConfig), + case!(cohere::ModelArgs), + case!(cohere2::Cohere2Config), + case!(cohere2_moe::Cohere2MoeConfig), + case!(dbrx::AttnConfig), + case!(ernie4_5::ModelArgs), + case!(ernie4_5_moe::ModelArgs), + case!(exaone::ExaOneConfig), + case!(exaone4::ModelArgs), + case!(falcon_h1::ModelArgs), + case!(gpt_oss::ModelArgs), + case!(helium::ModelArgs), + case!(hunyuan_moe::ModelArgs), + case!(hunyuan_v1_dense::ModelArgs), + case!(internlm2::ModelArgs), + case!(internlm3::ModelArgs), + case!(iquestloopcoder::ModelArgs), + case!(klear::ModelArgs), + case!(llada2_moe::ModelArgs), + case!(llama4::TextArgs), + case!(longcat_flash_ngram::LongcatFlashNgramConfig), + case!(mimo::ModelArgs), + case!(minicpm::ModelArgs), + case!(minicpm3::ModelArgs), + case!(minimax::ModelArgs), + case!(minimax_m3::ModelArgs), + case!(nanochat::ModelArgs), + case!(nemotron::ModelArgs), + case!(nemotron_nas::NemotronNASConfig), + case!(seed_oss::ModelArgs), + case!(smollm3::ModelArgs), + case!(stablelm::ModelArgs), + case!(starcoder2::StarCoder2Config), + case!(step3p5::Step3p5Config), + case!(telechat3::ModelArgs), + // Families whose own resolver used to let the nested base win. + case!(exaone_moe::ModelArgs, |c| Some(f64::from( + c.get_rope_theta() + ))), + case!(mistral4::Mistral4Config, |c| Some(f64::from(c.rope_base()))), + // VLM text configs. + case!(ernie4_5_moe_vl::Ernie45MoeVlTextConfig), + case!(falcon_ocr::FalconOcrConfig), + case!(hunyuan_vl::HunyuanVlTextConfig), + case!(mllama::config::MllamaTextConfig), + case!(molmo::MolmoTextConfig), + case!(molmo2::Molmo2TextConfig), + case!(moondream2::ModelArgs), + case!(paddleocr_vl::PaddleOcrTextConfig), + case!(youtu_vl_lm::YoutuTextConfig), + ( + "jina_vlm::JinaVlmTextConfig", + (|value: Value| { + Ok(Some(f64::from( + jina_vlm::JinaVlmTextConfig::from_json(&value).rope_theta, + ))) + }) as ParseTheta, + ), + // Outside `src/models`. + case!(omni::TalkerTextConfig), + case!(omni::CodePredictorConfig), + case!(omni::Code2WavConfig), + ( + "mlxcel_core::drafter::dflash::DFlashConfig", + (|value: Value| { + mlxcel_core::drafter::dflash::DFlashConfig::from_json(&value) + .map(|c| Some(f64::from(c.rope_theta))) + }) as ParseTheta, + ), + ] +} + +/// The `` `name` `` right after `prefix` in a serde error message. +fn quoted_after<'a>(message: &'a str, prefix: &str) -> Option<&'a str> { + let rest = &message[message.find(prefix)? + prefix.len()..]; + rest.split('`').next() +} + +/// A placeholder matching the `expected ...` clause of a type error, or the +/// first listed variant of an `unknown variant` error. +fn placeholder_for(message: &str) -> Option { + if message.contains("unknown variant") { + let variants = &message[message.find("expected")?..]; + return quoted_after(variants, "`").map(|v| json!(v)); + } + let expected = &message[message.find("expected")?..]; + let value = if expected.contains("string") || expected.contains("variant identifier") { + json!("silu") + } else if expected.contains("bool") { + json!(false) + } else if expected.contains("sequence") || expected.contains("tuple") { + json!([]) + } else if expected.contains("map") || expected.contains("struct") { + json!({}) + } else if expected.contains("f32") || expected.contains("f64") || expected.contains("float") { + json!(1.0) + } else { + return None; + }; + Some(value) +} + +fn set_path(config: &mut Value, path: &[String], value: Value) { + let mut node = config; + for key in &path[..path.len() - 1] { + node = node.get_mut(key).expect("parent object exists"); + } + node[path.last().expect("non-empty path")] = value; +} + +fn get_path<'a>(config: &'a Value, path: &[String]) -> Option<&'a Value> { + path.iter().try_fold(config, |node, key| node.get(key)) +} + +/// Grow `config` until `parse` accepts it, filling required non-RoPE fields. +fn fill_required_fields(parse: ParseTheta, mut config: Value) -> Result { + let mut last: Option> = None; + let mut objects: Vec> = Vec::new(); + for _ in 0..400 { + let message = match parse(config.clone()) { + Ok(_) => return Ok(config), + Err(message) => message, + }; + if let Some(field) = quoted_after(&message, "missing field `") { + if matches!(field, "rope_theta" | "rope_scaling" | "rope_parameters") { + return Err(message); + } + // A field already present at the top level is missing from the + // object placeholder inserted most recently. + let path = if config.get(field).is_none() { + vec![field.to_string()] + } else { + let parent = objects + .iter() + .rev() + .find(|p| get_path(&config, p).is_some_and(|o| o.get(field).is_none())) + .ok_or_else(|| message.clone())?; + let mut path = parent.clone(); + path.push(field.to_string()); + path + }; + set_path(&mut config, &path, json!(1)); + last = Some(path); + continue; + } + let (Some(path), Some(value)) = (last.clone(), placeholder_for(&message)) else { + return Err(message); + }; + if get_path(&config, &path) == Some(&value) { + return Err(message); + } + if value.is_object() { + objects.push(path.clone()); + } + set_path(&mut config, &path, value); + } + Err("placeholder search did not converge".to_string()) +} + +#[test] +fn every_flat_rope_theta_config_lifts_rope_parameters_and_flat_wins() { + let mut failures = Vec::new(); + for (name, parse) in cases() { + let nested_only = json!({ + "rope_parameters": {"rope_theta": 123456.0, "rope_type": "default"} + }); + let mut config = match fill_required_fields(parse, nested_only) { + Ok(config) => config, + Err(e) => { + failures.push(format!("{name}: nested-only config does not parse: {e}")); + continue; + } + }; + let nested = parse(config.clone()); + if nested != Ok(Some(123456.0)) { + failures.push(format!("{name}: nested-only base {nested:?}, want 123456")); + } + config["rope_theta"] = json!(500000.0); + let both = parse(config); + if both != Ok(Some(500000.0)) { + failures.push(format!("{name}: flat + nested base {both:?}, want 500000")); + } + } + assert!(failures.is_empty(), "{}", failures.join("\n")); +} + +#[test] +fn llama3_nested_rope_type_llama3_resolves_llama3_scaling() { + // The `rope_parameters` transformers 5.20 writes for Llama 3.2 1B. + let args: crate::models::llama3::ModelArgs = serde_json::from_value(json!({ + "model_type": "llama", + "hidden_size": 2048, + "num_hidden_layers": 16, + "intermediate_size": 8192, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "head_dim": 64, + "rms_norm_eps": 1e-5, + "vocab_size": 128256, + "max_position_embeddings": 131072, + "rope_parameters": { + "factor": 32.0, "high_freq_factor": 4.0, "low_freq_factor": 1.0, + "original_max_position_embeddings": 8192, "rope_type": "llama3", + "rope_theta": 500000.0 + } + })) + .unwrap(); + assert_eq!(args.rope_theta, 500000.0); + assert!(matches!( + args.rope_scaling_kind(), + crate::models::rope_utils::RopeScalingKind::Llama3 { .. } + )); +} + +#[test] +fn per_layer_type_rope_parameters_fill_gemma3_bases_and_scaling() { + // transformers 5.20 writes Gemma 3 RoPE per layer type. + let parse = cases() + .into_iter() + .find(|(name, _)| *name == "gemma3::ModelArgs") + .unwrap() + .1; + let config = fill_required_fields( + parse, + json!({ + "rope_parameters": { + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": {"rope_type": "linear", "factor": 8.0, "rope_theta": 1000000.0} + } + }), + ) + .unwrap(); + let args: crate::models::gemma3::ModelArgs = serde_json::from_value(config).unwrap(); + assert_eq!(args.rope_theta, 1_000_000.0); + assert_eq!(args.rope_local_base_freq, 10_000.0); + let scaling = args.rope_scaling.expect("linear scaling lifted"); + assert_eq!(scaling["type"], "linear"); + assert_eq!(scaling["factor"], 8.0); +} + +#[test] +fn nested_partial_rotary_factor_reaches_phi() { + // transformers 5 also moves `partial_rotary_factor` under + // `rope_parameters`; Phi would otherwise rotate at its serde default + // (0.4), so the fixture uses a different value. + let parse = cases() + .into_iter() + .find(|(name, _)| *name == "phi::ModelArgs") + .unwrap() + .1; + let config = fill_required_fields( + parse, + json!({ + "rope_parameters": {"rope_theta": 10000.0, "partial_rotary_factor": 0.5, "rope_type": "default"} + }), + ) + .unwrap(); + let args: crate::models::phi::ModelArgs = serde_json::from_value(config).unwrap(); + assert_eq!(args.partial_rotary_factor, 0.5); +} diff --git a/src/models/rope_utils.rs b/src/models/rope_utils.rs index 3c8222694..0762401fa 100644 --- a/src/models/rope_utils.rs +++ b/src/models/rope_utils.rs @@ -747,79 +747,20 @@ pub(crate) fn printable_label(label: &str) -> String { } } -/// Copy a nested `rope_parameters` block into the flat `rope_theta` / -/// `rope_scaling` keys a family's `ModelArgs` reads, when the flat keys are -/// absent. +/// The transformers 5 `rope_parameters` lift, shared with `mlxcel-core` so the +/// drafter configs there use the same rule (#2443). See +/// [`mlxcel_core::rope_parameters`]. /// -/// transformers 5 saves `rope_theta` and `rope_scaling` only under -/// `rope_parameters` (for example `{"rope_theta": 1000000, "rope_type": -/// "default"}`) and drops the flat keys. A `ModelArgs` that reads the flat -/// `rope_theta` with a serde default would then silently run at the default -/// base (10000 for Qwen3), decoding fluently at short lengths and drifting at -/// long ones. A flat key that is present and non-null wins over the nested one. -/// A non-default `rope_type` is carried into `rope_scaling` (with the `type` key -/// [`RopeScalingSpec`] reads first) unless the config already has a non-null -/// `rope_scaling`. -/// -/// Upstream: ml-explore/mlx-lm#1898 (qwen3) and ml-explore/mlx-lm#1605 -/// (granitemoehybrid). Families that apply it at deserialization through -/// [`deserialize_lifting_rope_parameters`]: Qwen3 (and every text backbone or -/// pipeline stage that parses `qwen3::ModelArgs`) and GraniteMoeHybrid. -pub fn lift_rope_parameters(config: &mut serde_json::Value) { - use serde_json::Value; - - let Some(obj) = config.as_object_mut() else { - return; - }; - let Some(params) = obj - .get("rope_parameters") - .and_then(Value::as_object) - .cloned() - else { - return; - }; - if obj.get("rope_theta").is_none_or(Value::is_null) - && let Some(theta) = params.get("rope_theta").filter(|v| v.is_number()) - { - obj.insert("rope_theta".to_string(), theta.clone()); - } - let rope_type = params - .get("rope_type") - .or_else(|| params.get("type")) - .and_then(Value::as_str); - let has_scaling = obj.get("rope_scaling").is_some_and(|v| !v.is_null()); - if let Some(kind) = rope_type - && kind != "default" - && !has_scaling - { - let mut scaling = params.clone(); - scaling.remove("rope_theta"); - scaling.insert("type".to_string(), Value::String(kind.to_string())); - obj.insert("rope_scaling".to_string(), Value::Object(scaling)); - } -} - -/// Deserialize a config through [`lift_rope_parameters`]. -/// -/// A family whose `ModelArgs` derives `Deserialize` with -/// `#[serde(remote = "Self")]` (which turns the derive into an inherent -/// `ModelArgs::deserialize` function instead of the trait impl) implements the -/// trait by hand as -/// `deserialize_lifting_rope_parameters(deserializer, ModelArgs::deserialize)`. -/// Every `serde_json::from_value` / `from_str` / nested-field parse of that -/// type then sees the lifted keys, so a new parse site cannot miss the lift. -pub fn deserialize_lifting_rope_parameters<'de, D, T>( - deserializer: D, - parse: impl FnOnce(serde_json::Value) -> Result, -) -> Result -where - D: Deserializer<'de>, -{ - let mut value = serde_json::Value::deserialize(deserializer)?; - lift_rope_parameters(&mut value); - parse(value).map_err(serde::de::Error::custom) -} +/// Config structs apply it at deserialization with `#[serde(remote = "Self")]` +/// plus [`mlxcel_core::impl_rope_lifting_deserialize!`] (or a hand-written impl +/// calling [`deserialize_lifting_rope_parameters`], as Qwen3 and +/// GraniteMoeHybrid do), so every parse site of the type sees the lifted keys. +pub use mlxcel_core::rope_parameters::{deserialize_lifting_rope_parameters, lift_rope_parameters}; #[cfg(test)] #[path = "rope_utils_tests.rs"] mod tests; + +#[cfg(test)] +#[path = "rope_parameters_lift_tests.rs"] +mod rope_parameters_lift_tests; diff --git a/src/models/seed_oss.rs b/src/models/seed_oss.rs index 08f3638ba..e16a9f58e 100644 --- a/src/models/seed_oss.rs +++ b/src/models/seed_oss.rs @@ -48,6 +48,7 @@ use std::time::Instant; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -92,6 +93,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/smollm3.rs b/src/models/smollm3.rs index 7bfceb655..aacc0de9b 100644 --- a/src/models/smollm3.rs +++ b/src/models/smollm3.rs @@ -27,6 +27,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub hidden_size: usize, @@ -64,6 +65,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/stablelm.rs b/src/models/stablelm.rs index 38cfbf089..14770efd5 100644 --- a/src/models/stablelm.rs +++ b/src/models/stablelm.rs @@ -28,6 +28,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { pub model_type: String, pub vocab_size: usize, @@ -54,6 +55,8 @@ pub struct ModelArgs { pub eos_token_id: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + #[derive(Debug, Clone, Deserialize)] pub struct Quantization { pub group_size: i32, diff --git a/src/models/starcoder2.rs b/src/models/starcoder2.rs index 2f36d6b67..8b2b1184d 100644 --- a/src/models/starcoder2.rs +++ b/src/models/starcoder2.rs @@ -28,6 +28,7 @@ use std::path::Path; // Config. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct StarCoder2Config { pub model_type: String, pub hidden_size: usize, @@ -45,6 +46,8 @@ pub struct StarCoder2Config { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(StarCoder2Config); + #[derive(Debug, Clone, Deserialize)] pub struct QuantizationConfig { pub group_size: i32, diff --git a/src/models/step3p5.rs b/src/models/step3p5.rs index 594282b8b..ca93603fb 100644 --- a/src/models/step3p5.rs +++ b/src/models/step3p5.rs @@ -36,6 +36,7 @@ use std::path::Path; // Configuration. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct Step3p5Config { #[serde(default = "default_model_type")] pub model_type: String, @@ -126,6 +127,8 @@ pub struct Step3p5Config { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(Step3p5Config); + #[derive(Debug, Clone, Deserialize)] pub struct AttentionOtherSetting { pub num_attention_heads: usize, diff --git a/src/models/telechat3.rs b/src/models/telechat3.rs index b7a14d90d..9e28e0f49 100644 --- a/src/models/telechat3.rs +++ b/src/models/telechat3.rs @@ -87,6 +87,7 @@ fn default_max_position_embeddings() -> usize { } #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct ModelArgs { #[serde(default = "default_model_type")] pub model_type: String, @@ -134,6 +135,8 @@ pub struct ModelArgs { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(ModelArgs); + impl ModelArgs { pub fn num_kv_heads(&self) -> usize { self.num_key_value_heads.unwrap_or(self.num_attention_heads) diff --git a/src/models/youtu_vl_lm_config.rs b/src/models/youtu_vl_lm_config.rs index 8fdc6d215..97e53229e 100644 --- a/src/models/youtu_vl_lm_config.rs +++ b/src/models/youtu_vl_lm_config.rs @@ -37,6 +37,7 @@ use crate::models::deepseek_v3; /// MoE-related fields are kept for forward compatibility with possible larger /// variants but are not exercised by any published checkpoint. #[derive(Debug, Clone, Deserialize)] +#[serde(remote = "Self")] pub struct YoutuTextConfig { #[serde(default = "default_model_type")] pub model_type: String, @@ -124,6 +125,8 @@ pub struct YoutuTextConfig { pub quantization: Option, } +mlxcel_core::impl_rope_lifting_deserialize!(YoutuTextConfig); + #[derive(Debug, Clone, Deserialize)] pub struct QuantizationConfig { pub group_size: i32, From 1684f1515c477f3453f7c6c81ce5e7d4460ab492 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 18 Oct 2026 20:57:58 +0900 Subject: [PATCH 2/3] fix(models): lift DeepSeek-V4 main/compress rope_parameters Review follow-up for #2443. transformers 5.20 writes DeepSeek-V4 RoPE as {"main": {...}, "compress": {...}} with the compress-layer YaRN block under compress. The lifter only recognized full_attention / sliding_attention, so it treated this as a single block and synthesized an all-None rope_scaling out of the two sub-objects. A block whose values are all objects is now read as per-layer-type: full_attention or main supplies the base and scaling, sliding_attention gives rope_local_base_freq, compress gives compress_rope_theta and, when main names no scheme, rope_scaling. A "type": null key no longer hides a rope_type scheme. New lifter and DeepSeek-V4 ModelArgs tests pass, as do the affected family filters, clippy -D warnings for mlxcel and mlxcel-core, and the Llama 3.2 1B / Qwen3-30B-A3B transformers-5 copies, which still decode byte-identically to the originals. Refs #2443 --- CHANGELOG.md | 2 +- src/lib/mlxcel-core/src/rope_parameters.rs | 71 ++++++++++++------- .../mlxcel-core/src/rope_parameters_tests.rs | 49 +++++++++++++ src/models/rope_parameters_lift_tests.rs | 32 +++++++++ 4 files changed, 128 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 01b363616..0c38d3aa3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Fixed -- Configs saved by transformers 5, which keep `rope_theta`, `rope_scaling` and `partial_rotary_factor` only under a nested `rope_parameters` block, now load with their real RoPE base and scaling in every family that reads the flat keys, not only Qwen3 and GraniteMoeHybrid (#2443). Before, about 80 config structs (Llama and everything sharing its args, Qwen2/Qwen3-MoE, Gemma, Phi, OLMo, DeepSeek, GLM, the VLM text configs, the Qwen3-Omni speech stack and the DFlash drafter) silently ran at the serde default base, decoding fluently on short prompts and drifting on long ones, and a Llama 3 config lost its `llama3` scaling. The per-layer-type form (Gemma 3, Gemma 3n, OLMo 3) and the MRoPE `mrope_section` are lifted too. A flat key still wins when both are present; EXAONE-MoE and Mistral 4 used to let the nested base win and now follow the same rule. +- Configs saved by transformers 5, which keep `rope_theta`, `rope_scaling` and `partial_rotary_factor` only under a nested `rope_parameters` block, now load with their real RoPE base and scaling in every family that reads the flat keys, not only Qwen3 and GraniteMoeHybrid (#2443). Before, about 80 config structs (Llama and everything sharing its args, Qwen2/Qwen3-MoE, Gemma, Phi, OLMo, DeepSeek, GLM, the VLM text configs, the Qwen3-Omni speech stack and the DFlash drafter) silently ran at the serde default base, decoding fluently on short prompts and drifting on long ones, and a Llama 3 config lost its `llama3` scaling. The per-layer-type forms (Gemma 3, Gemma 3n and OLMo 3 `full_attention` / `sliding_attention`, DeepSeek-V4 `main` / `compress`) and the MRoPE `mrope_section` are lifted too. A flat key still wins when both are present; EXAONE-MoE and Mistral 4 used to let the nested base win and now follow the same rule. - Checkpoints that load through the SentencePiece fallback (a `tokenizer.model` with no `tokenizer.json`) now encode their `"special": false` added tokens as the one id each was declared with, as `transformers` does, instead of splitting them into ordinary pieces (#1758). IQuest-Coder and IQuest-Coder Loop tool-calling prompts now carry ``, ``, `` and their closers as single ids, Baichuan-M1's `` turn marker and ERNIE-4.5-VL's `<|IMAGE_PLACEHOLDER|>` and `<|LOC_n|>` tokens do the same, and each added token's `lstrip`, `rstrip` and `single_word` flags are honored. `skip_special_tokens` still drops only the special ones. On this path `parse_special: false` (`/tokenize`, logit-bias text keys) now also recognizes non-special added tokens, as llama.cpp does, `added_tokens.json` is read when `added_tokens_decoder` is absent, and a `single_word` special token glued to a word (IQuest's `` in `de`) encodes as text, as in `transformers`. - On CUDA, a bf16 model on an fp16 TurboQuant KV cache (`--kv-cache-mode fp16+turbo4` and the other turbo modes) no longer aborts with `Mask type must promote to output type float32` when attention carries an array mask, which for Gemma 4 meant any prompt past the 1024-token sliding window (#2336). Additive masks now reach SDPA in the attention dtype, `result_type(q, k, v)`, instead of the query's dtype. - Mixed bf16/fp16 attention follow-ups to #2336 (#2422). On CUDA, the query-chunking gate for materializing SDPA now judges cuDNN eligibility and score size from the attention dtype, `result_type(q, k, v)`, instead of q's dtype, so a bf16 model on an fp16 TurboQuant KV cache, which attends in float32, chunks long prefills at 4 bytes per score instead of skipping the chunking at head_dim 128 or budgeting 2 bytes. The softcap attention composite (Gemma 2) casts its mask to and returns the attention dtype, as fast SDPA does, instead of q's and v's. Model-state snapshot donation checks the KV cache mode before the keys/values layout, so `--kv-cache-mode fp16+turbo4` logs `cache mode Turbo4Asym is not supported` instead of `has only one of keys/values`, and a populated symmetric `turbo4` layer is refused instead of being left out of the snapshot. Same-dtype attention is unchanged. diff --git a/src/lib/mlxcel-core/src/rope_parameters.rs b/src/lib/mlxcel-core/src/rope_parameters.rs index 161862149..883e6b06c 100644 --- a/src/lib/mlxcel-core/src/rope_parameters.rs +++ b/src/lib/mlxcel-core/src/rope_parameters.rs @@ -44,10 +44,13 @@ use serde_json::{Map, Value}; /// both spellings would turn a `#[serde(alias = "type", alias = "rope_type")]` /// field into a `duplicate field` error. /// -/// The per-layer-type form transformers 5 writes for Gemma 3, Gemma 3n and -/// OLMo 3 (`{"full_attention": {...}, "sliding_attention": {...}}`) is read -/// too: the `full_attention` block supplies the keys above, and the -/// `sliding_attention` base becomes `rope_local_base_freq`. +/// The per-layer-type form (every value an object) is read too. For Gemma 3, +/// Gemma 3n and OLMo 3 (`{"full_attention": {...}, "sliding_attention": +/// {...}}`) the `full_attention` block supplies the keys above and the +/// `sliding_attention` base becomes `rope_local_base_freq`. For DeepSeek-V4 +/// (`{"main": {...}, "compress": {...}}`) `main` supplies the base, the +/// `compress` base becomes `compress_rope_theta`, and the `compress` YaRN block +/// becomes `rope_scaling`, which is where transformers 4 kept it. /// /// Upstream: ml-explore/mlx-lm#1898 (qwen3) and ml-explore/mlx-lm#1605 /// (granitemoehybrid). @@ -59,15 +62,19 @@ pub fn lift_rope_parameters(config: &mut Value) { return; }; let block_at = |key: &str| params.get(key).and_then(Value::as_object).cloned(); - let (main, sliding) = - if params.contains_key("full_attention") || params.contains_key("sliding_attention") { - ( - block_at("full_attention").unwrap_or_default(), - block_at("sliding_attention"), - ) - } else { - (params.clone(), None) - }; + let per_layer_type = + !params.is_empty() && params.values().all(|v| v.is_object() || v.is_null()); + let (main, sliding, compress) = if per_layer_type { + ( + block_at("full_attention") + .or_else(|| block_at("main")) + .unwrap_or_default(), + block_at("sliding_attention"), + block_at("compress"), + ) + } else { + (params.clone(), None, None) + }; copy_number_if_absent(obj, "rope_theta", main.get("rope_theta")); copy_number_if_absent( @@ -78,28 +85,42 @@ pub fn lift_rope_parameters(config: &mut Value) { if let Some(sliding) = &sliding { copy_number_if_absent(obj, "rope_local_base_freq", sliding.get("rope_theta")); } + if let Some(compress) = &compress { + copy_number_if_absent(obj, "compress_rope_theta", compress.get("rope_theta")); + } + + if obj.get("rope_scaling").is_some_and(|v| !v.is_null()) { + return; + } + if let Some(scaling) = scaling_block(main).or_else(|| compress.and_then(scaling_block)) { + obj.insert("rope_scaling".to_string(), Value::Object(scaling)); + } +} - let kind = main +/// The `rope_scaling` a nested block stands for, or `None` when it names the +/// default scheme and carries nothing beyond the base. +fn scaling_block(mut block: Map) -> Option> { + let kind = block .get("type") - .or_else(|| main.get("rope_type")) .and_then(Value::as_str) + .or_else(|| block.get("rope_type").and_then(Value::as_str)) .map(str::to_string); - let has_extra_keys = main.keys().any(|key| { + let has_extra_keys = block.keys().any(|key| { !matches!( key.as_str(), "rope_theta" | "rope_type" | "type" | "partial_rotary_factor" ) }); - let has_scaling = obj.get("rope_scaling").is_some_and(|v| !v.is_null()); - if !has_scaling && (kind.as_deref().is_some_and(|k| k != "default") || has_extra_keys) { - let mut scaling = main; - scaling.remove("rope_theta"); - scaling.remove("rope_type"); - if let Some(kind) = kind { - scaling.insert("type".to_string(), Value::String(kind)); - } - obj.insert("rope_scaling".to_string(), Value::Object(scaling)); + if !(kind.as_deref().is_some_and(|k| k != "default") || has_extra_keys) { + return None; } + block.remove("rope_theta"); + block.remove("rope_type"); + match kind { + Some(kind) => block.insert("type".to_string(), Value::String(kind)), + None => block.remove("type"), + }; + Some(block) } fn copy_number_if_absent(obj: &mut Map, key: &str, nested: Option<&Value>) { diff --git a/src/lib/mlxcel-core/src/rope_parameters_tests.rs b/src/lib/mlxcel-core/src/rope_parameters_tests.rs index 6058f619a..19f2f6977 100644 --- a/src/lib/mlxcel-core/src/rope_parameters_tests.rs +++ b/src/lib/mlxcel-core/src/rope_parameters_tests.rs @@ -157,3 +157,52 @@ fn non_object_and_unusable_shapes_are_left_alone() { assert!(sliding_only.get("rope_scaling").is_none()); assert_eq!(sliding_only["rope_local_base_freq"], 10000.0); } + +#[test] +fn null_type_key_does_not_hide_rope_type() { + let mut config = json!({ + "rope_parameters": {"type": null, "rope_type": "yarn", "factor": 4.0, "rope_theta": 1000000.0} + }); + lift_rope_parameters(&mut config); + assert_eq!( + config["rope_scaling"], + json!({"type": "yarn", "factor": 4.0}) + ); +} + +#[test] +fn deepseek_v4_main_and_compress_blocks_lift() { + // The shape transformers 5.20 writes for DeepSeek-V4: the flat bases stay, + // `rope_scaling` moves under `compress`. + let mut config = json!({ + "rope_theta": 10000.0, + "compress_rope_theta": 160000.0, + "rope_parameters": { + "main": {"rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 0.125}, + "compress": { + "type": "yarn", "factor": 16.0, "original_max_position_embeddings": 65536, + "beta_fast": 32, "beta_slow": 1, "rope_theta": 160000.0, "rope_type": "yarn", + "partial_rotary_factor": 0.125, "attention_factor": 1.0 + } + } + }); + lift_rope_parameters(&mut config); + assert_eq!(config["rope_theta"], 10000.0); + assert_eq!(config["compress_rope_theta"], 160000.0); + let scaling = &config["rope_scaling"]; + assert_eq!(scaling["type"], "yarn"); + assert_eq!(scaling["factor"], 16.0); + assert_eq!(scaling["original_max_position_embeddings"], 65536); + assert!(scaling.get("rope_type").is_none()); + + // With the flat bases dropped, both come from the nested blocks. + let mut nested_only = json!({ + "rope_parameters": { + "main": {"rope_type": "default", "rope_theta": 10000.0}, + "compress": {"rope_type": "yarn", "factor": 16.0, "rope_theta": 160000.0} + } + }); + lift_rope_parameters(&mut nested_only); + assert_eq!(nested_only["rope_theta"], 10000.0); + assert_eq!(nested_only["compress_rope_theta"], 160000.0); +} diff --git a/src/models/rope_parameters_lift_tests.rs b/src/models/rope_parameters_lift_tests.rs index cb2382549..02259723c 100644 --- a/src/models/rope_parameters_lift_tests.rs +++ b/src/models/rope_parameters_lift_tests.rs @@ -383,3 +383,35 @@ fn nested_partial_rotary_factor_reaches_phi() { let args: crate::models::phi::ModelArgs = serde_json::from_value(config).unwrap(); assert_eq!(args.partial_rotary_factor, 0.5); } + +#[test] +fn deepseek_v4_compress_yarn_reaches_rope_scaling() { + // transformers 5.20 keeps the flat DeepSeek-V4 bases but moves the + // compress-layer YaRN block under `rope_parameters.compress`. + let parse = cases() + .into_iter() + .find(|(name, _)| *name == "deepseek_v4::ModelArgs") + .unwrap() + .1; + let config = fill_required_fields( + parse, + json!({ + "rope_theta": 10000.0, + "compress_rope_theta": 160000.0, + "rope_parameters": { + "main": {"rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 0.125}, + "compress": { + "type": "yarn", "factor": 16.0, "original_max_position_embeddings": 65536, + "beta_fast": 32, "beta_slow": 1, "rope_theta": 160000.0, "rope_type": "yarn" + } + } + }), + ) + .unwrap(); + let args: crate::models::deepseek_v4::ModelArgs = serde_json::from_value(config).unwrap(); + assert_eq!(args.compress_rope_theta, 160000.0); + let scaling = args.rope_scaling.expect("compress YaRN lifted"); + assert_eq!(scaling.scaling_type.as_deref(), Some("yarn")); + assert_eq!(scaling.factor, Some(16.0)); + assert_eq!(scaling.original_max_position_embeddings, Some(65536)); +} From 51e261bb43a7faa53e11e6bd9c9943e46bbea812 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 19 Oct 2026 09:34:31 +0900 Subject: [PATCH 3/3] docs: add technical report for PR #2448 --- .../pr-2448-rope-parameters-lift.en.md | 72 +++++++++++++++++++ .../pr-2448-rope-parameters-lift.ko.md | 72 +++++++++++++++++++ 2 files changed, 144 insertions(+) create mode 100644 TECHNICAL_REPORTS/pr-2448-rope-parameters-lift.en.md create mode 100644 TECHNICAL_REPORTS/pr-2448-rope-parameters-lift.ko.md diff --git a/TECHNICAL_REPORTS/pr-2448-rope-parameters-lift.en.md b/TECHNICAL_REPORTS/pr-2448-rope-parameters-lift.en.md new file mode 100644 index 000000000..0868f5ca1 --- /dev/null +++ b/TECHNICAL_REPORTS/pr-2448-rope-parameters-lift.en.md @@ -0,0 +1,72 @@ +# PR #2448: Lift transformers 5 rope_parameters for every family + +**Date**: 2026-10-19 +**Status**: Implemented and verified on CUDA (GB10) with real checkpoints; PR open. The items under "Not addressed" remain open + +## Summary + +Closes #2443. transformers 5 saves `rope_theta`, `rope_scaling` and `partial_rotary_factor` only under a nested `rope_parameters` block. Only Qwen3 and GraniteMoeHybrid lifted it (#2439). Every other struct that reads the flat keys fell back to its serde default base without an error. This PR applies the lift at the config type, through one macro line per struct, for 86 structs plus the call sites that parse outside serde. + +## Problem + +- Config structs that read a flat `rope_theta` with a serde default (10000 for llama3, qwen3_moe, gemma2, phi3, olmo2) load a transformers 5 checkpoint without error, decode fluently on short prompts, and drift on long ones. A Llama 3 config also loses its `rope_type: "llama3"` scaling. +- On `main`, a transformers-5 copy of Llama 3.2 1B or Qwen3-30B-A3B produces garbage on a 20 KB prompt. The Llama copy degenerates into `and, and, ...` and the Qwen3 copy into `. The. The. ...`. +- mlx-lm 0.32.2 lifts only exaone_moe, gemma4, glm_moe_dsa, laguna, lfm2, mellum, ministral3, mistral4, muse_glimmer and qwen3_5, which mlxcel already handles, so no other family gets a fix from parity. +- EXAONE-MoE and Mistral 4 let the nested base win, which was inconsistent with the families that let a flat key win. +- `src/loading/vlm_qwen.rs` carried two ad hoc call-site lifts (GLM-OCR and GLM-4V MoE). + +## Technical decisions + +- **The type is the choke point.** The PR author looked for a shared parse helper that every site already calls and found none: about 125 family `load` functions read `config.json` directly, and VLMs nest the block under `text_config`. The issue rejects a top-level lift in `sanitize_config_json` for both reasons (it is a string replace that most family `load` functions skip, and it would miss `text_config`). +- **Shared lifter in `mlxcel-core`.** `lift_rope_parameters` / `deserialize_lifting_rope_parameters` move to `mlxcel_core::rope_parameters` (re-exported from `models::rope_utils`), so the DFlash drafter config can use them. An `impl_rope_lifting_deserialize!(Type)` macro gives each uncovered struct `#[serde(remote = "Self")]` plus one line. This covers 86 structs, including the VLM text configs and the three Qwen3-Omni speech configs. Pipeline and tensor-parallel sites parse these same types, so they are covered. No struct gains a field. +- **Lifter rules.** + - A flat non-null key wins. + - It lifts `partial_rotary_factor`. + - It reads the per-layer-type forms, meaning a block whose values are all objects. For Gemma 3 / Gemma 3n / OLMo 3, `full_attention` gives the base and scaling and `sliding_attention` gives `rope_local_base_freq`. For DeepSeek-V4, `main` gives the base and `compress` gives `compress_rope_theta` and the YaRN `rope_scaling`. + - It builds `rope_scaling` from a non-default scheme (Llama 3 `llama3` included) or from MRoPE extras like `mrope_section`. + - The scheme goes under `type` only (it reads `type` first, then `rope_type`), because the deepseek_v32 / deepseek_v4 / nemotron `alias = "type", alias = "rope_type"` fields reject a block that spells both. +- **Call sites outside serde use the shared lifter.** `DFlashConfig::from_json`, the Idefics2 Mistral defaults (which used to shadow the nested base) and both `vlm_qwen.rs` ad hoc lifts. `JinaVlmTextConfig::from_json` falls back to the nested base, and the Phi-4 text-field copy now carries `rope_parameters`. +- **Flat wins everywhere.** EXAONE-MoE and Mistral 4 now let a flat key win, as the issue specifies. transformers 5 itself gives the nested value priority, but the two only disagree when a config carries both with different values, and none of the reference checkpoints does. +- **`exaone4`.** It was the one reader that checked only `rope_type`, and it now accepts `type` too. Its `rope_scaling` is now `#[serde(default)]`, because a default-scheme transformers 5 save has none. + +## Changes + +- `src/lib/mlxcel-core/src/rope_parameters.rs` (new, 187 lines), `rope_parameters_tests.rs` (new, lifter unit tests), `lib.rs`: the shared lifter and macro. +- `src/models/rope_utils.rs`: lifter code removed and re-exported from `mlxcel-core`. +- About 85 config files under `src/models/` and `src/audio/qwen3_omni_moe/`: `#[serde(remote = "Self")]` and one `impl_rope_lifting_deserialize!` line each (typically 3 lines per file). +- `src/models/mistral4.rs`, `src/models/exaone_moe.rs`, `src/models/exaone4.rs`, `src/models/jina_vlm.rs`: flat-wins, `type` handling and fallback changes described above. +- `src/loading/vlm_qwen.rs`, `src/lib/mlxcel-core/src/drafter/dflash/config.rs`: call sites moved to the shared lifter. +- `src/models/rope_parameters_lift_tests.rs` (new, 417 lines): table-driven tests. + +## Verification + +- `rope_parameters_lift_tests::every_flat_rope_theta_config_lifts_rope_parameters_and_flat_wins` is table-driven over all 89 parse paths (the 86 lifted structs, Mistral 4, and the DFlash and Jina `from_json` paths). It checks that nested-only gives 123456 and flat 500000 plus nested gives 500000. Required non-RoPE fields are filled from serde's own errors, never `rope_theta` / `rope_scaling`. +- More tests: Llama 3 nested `llama3` resolves `RopeScalingKind::Llama3`, Gemma 3 per-layer-type, Phi `partial_rotary_factor`, and DeepSeek-V4 `compress` YaRN (added in review; not revert-run). The lifter unit tests are in `mlxcel-core`. +- With the fix stashed and only the tests kept, the first four tests all fail, and every one of the 89 rows reports a failure. +- Real checkpoints: transformers-5 copies of `~/models_archive/mlx/llama-3.2-1b-instruct-4bit` and `~/models_archive/mlx/qwen3-30b-a3b-4bit`, with only `config.json` rewritten (flat keys moved into a `rope_parameters` block identical to the one transformers 5.20.1 emits for that config) and the weights symlinked. Command: `MLXCEL_SDPA_DETERMINISTIC=1 gpu-lock run -- mlxcel generate -m -p "" -n 64 --temp 0 --show-reasoning`. On the branch, both copies match their originals byte for byte, and the originals match `main`. On `main` (self-parity checked), the Llama copy degenerates into `and, and, ...` and the Qwen3 copy into `. The. The. ...`. +- fmt, clippy (`--lib --tests -D warnings` for mlxcel and mlxcel-core), the affected-family test filters (647 tests) and the repo contract checks pass locally. + +## Changes during review + +The review found no CRITICAL or HIGH issues. Two follow-ups were applied, in commit `1684f151`: + +- The DeepSeek-V4 `main` / `compress` layout, which the first version turned into an all-`None` `rope_scaling`. +- A `"type": null` key that hid the `rope_type` scheme. + +## Not addressed + +- Vision-encoder configs (Pixtral, Llama 4 vision, etc.) and `mlxcel-xla`, which the issue does not list. +- DBRX is lifted at `attn_config`, where transformers 5 still keeps `rope_theta`, but its top-level `rope_parameters` is not copied there. + +## Statistics + +| Item | Value | +|------|-------| +| Files changed | 96 | +| Lines added | 1134 | +| Lines deleted | 116 | + +| Commit | Message | +|--------|---------| +| `0681d42b` | fix(models): lift transformers 5 rope_parameters for every family | +| `1684f151` | fix(models): lift DeepSeek-V4 main/compress rope_parameters | diff --git a/TECHNICAL_REPORTS/pr-2448-rope-parameters-lift.ko.md b/TECHNICAL_REPORTS/pr-2448-rope-parameters-lift.ko.md new file mode 100644 index 000000000..de8319da6 --- /dev/null +++ b/TECHNICAL_REPORTS/pr-2448-rope-parameters-lift.ko.md @@ -0,0 +1,72 @@ +# PR #2448: 모든 모델 계열에서 transformers 5 rope_parameters 끌어올리기 + +**작성일**: 2026-10-19 +**상태**: CUDA(GB10)에서 실제 체크포인트로 구현 및 검증 완료. PR은 열려 있음. "미해결 사항"의 항목은 계속 열려 있음 + +## 요약 + +#2443을 닫는다. transformers 5는 `rope_theta`, `rope_scaling`, `partial_rotary_factor`를 중첩된 `rope_parameters` 블록 안에만 저장한다. 지금까지 Qwen3와 GraniteMoeHybrid만 이를 끌어올렸고(#2439), 평평한(flat) 키를 읽는 나머지 struct는 오류 없이 serde 기본 base로 되돌아갔다. 이 PR은 설정 타입에서 끌어올리기를 적용한다. struct 86개에 매크로 한 줄씩 추가하고, serde 밖에서 파싱하는 호출 지점도 함께 처리한다. + +## 1. 문제 정의 + +- serde 기본값이 있는 flat `rope_theta`를 읽는 설정 struct(llama3, qwen3_moe, gemma2, phi3, olmo2는 10000)는 transformers 5 체크포인트를 오류 없이 로드하고, 짧은 프롬프트에서는 유창하게 디코딩하지만 긴 프롬프트에서는 어긋난다. Llama 3 설정은 `rope_type: "llama3"` 스케일링도 잃는다. +- `main`에서 Llama 3.2 1B 또는 Qwen3-30B-A3B의 transformers-5 사본은 20 KB 프롬프트에서 깨진 출력을 만든다. Llama 사본은 `and, and, ...`로, Qwen3 사본은 `. The. The. ...`로 퇴화한다. +- mlx-lm 0.32.2는 exaone_moe, gemma4, glm_moe_dsa, laguna, lfm2, mellum, ministral3, mistral4, muse_glimmer, qwen3_5만 끌어올리며 mlxcel은 이미 이들을 처리하므로, 다른 계열은 parity로는 수정을 얻지 못한다. +- EXAONE-MoE와 Mistral 4는 중첩된 base가 우선했는데, flat 키가 우선하는 다른 계열과 일관되지 않았다. +- `src/loading/vlm_qwen.rs`에 임시방편 호출 지점 끌어올리기가 두 곳 있었다(GLM-OCR, GLM-4V MoE). + +## 2. 기술적 선택과 그 이유 + +- **타입이 병목 지점이다.** PR 작성자는 모든 지점이 이미 호출하는 공유 파싱 헬퍼를 찾았으나 없었다. 약 125개의 계열별 `load` 함수가 `config.json`을 직접 읽고, VLM은 블록을 `text_config` 아래에 중첩한다. 이슈는 `sanitize_config_json`의 최상위 끌어올리기를 두 이유로 배제한다(대부분의 계열 `load` 함수가 건너뛰는 문자열 치환이며, `text_config`를 놓친다). +- **`mlxcel-core`의 공유 lifter.** `lift_rope_parameters` / `deserialize_lifting_rope_parameters`를 `mlxcel_core::rope_parameters`로 옮기고(`models::rope_utils`에서 re-export), DFlash drafter 설정도 사용할 수 있게 한다. `impl_rope_lifting_deserialize!(Type)` 매크로가 미처리 struct마다 `#[serde(remote = "Self")]`와 한 줄을 제공한다. VLM 텍스트 설정과 Qwen3-Omni speech 설정 3개를 포함해 86개 struct를 다룬다. 파이프라인과 tensor-parallel 지점도 같은 타입을 파싱하므로 포함된다. 어떤 struct에도 필드가 추가되지 않는다. +- **Lifter 규칙.** + - null이 아닌 flat 키가 우선한다. + - `partial_rotary_factor`를 끌어올린다. + - 레이어 종류별 형식, 즉 값이 모두 객체인 블록을 읽는다. Gemma 3 / Gemma 3n / OLMo 3는 `full_attention`이 base와 스케일링을, `sliding_attention`이 `rope_local_base_freq`를 준다. DeepSeek-V4는 `main`이 base를, `compress`가 `compress_rope_theta`와 YaRN `rope_scaling`을 준다. + - 기본이 아닌 스킴(Llama 3 `llama3` 포함) 또는 `mrope_section` 같은 MRoPE 추가 값으로 `rope_scaling`을 만든다. + - 스킴은 `type` 아래에만 넣는다(`type`을 먼저 읽고 `rope_type`을 읽는다). deepseek_v32 / deepseek_v4 / nemotron의 `alias = "type", alias = "rope_type"` 필드가 둘을 모두 적은 블록을 거부하기 때문이다. +- **serde 밖의 호출 지점도 공유 lifter를 쓴다.** `DFlashConfig::from_json`, Idefics2 Mistral 기본값(기존에는 중첩 base를 가렸다), 그리고 `vlm_qwen.rs`의 임시 끌어올리기 두 곳이다. `JinaVlmTextConfig::from_json`은 중첩 base로 폴백하고, Phi-4 텍스트 필드 복사는 이제 `rope_parameters`를 함께 옮긴다. +- **모든 곳에서 flat 우선.** 이슈의 지시대로 EXAONE-MoE와 Mistral 4도 flat 키가 우선한다. transformers 5 자체는 중첩 값에 우선순위를 주지만, 둘이 다른 값을 가진 설정에서만 차이가 나며 참조 체크포인트 중에는 그런 것이 없다. +- **`exaone4`.** `rope_type`만 확인하던 유일한 reader였으며 이제 `type`도 받아들인다. 기본 스킴으로 저장한 transformers 5 설정에는 `rope_scaling`이 없으므로 `#[serde(default)]`로 바꿨다. + +## 3. 변경 내용 + +- `src/lib/mlxcel-core/src/rope_parameters.rs`(신규, 187줄), `rope_parameters_tests.rs`(신규, lifter 단위 테스트), `lib.rs`: 공유 lifter와 매크로. +- `src/models/rope_utils.rs`: lifter 코드를 제거하고 `mlxcel-core`에서 re-export. +- `src/models/`와 `src/audio/qwen3_omni_moe/` 아래 약 85개 설정 파일: `#[serde(remote = "Self")]`와 `impl_rope_lifting_deserialize!` 한 줄(보통 파일당 3줄). +- `src/models/mistral4.rs`, `src/models/exaone_moe.rs`, `src/models/exaone4.rs`, `src/models/jina_vlm.rs`: 위에서 설명한 flat 우선, `type` 처리, 폴백 변경. +- `src/loading/vlm_qwen.rs`, `src/lib/mlxcel-core/src/drafter/dflash/config.rs`: 호출 지점을 공유 lifter로 이동. +- `src/models/rope_parameters_lift_tests.rs`(신규, 417줄): 테이블 기반 테스트. + +## 4. 검증 + +- `rope_parameters_lift_tests::every_flat_rope_theta_config_lifts_rope_parameters_and_flat_wins`는 89개 파싱 경로(끌어올린 struct 86개, Mistral 4, DFlash와 Jina의 `from_json` 경로)를 테이블로 순회한다. 중첩만 있으면 123456, flat 500000에 중첩이 있으면 500000인지 확인한다. RoPE가 아닌 필수 필드는 serde 자체 오류로 채우며 `rope_theta` / `rope_scaling`은 채우지 않는다. +- 추가 테스트: Llama 3 중첩 `llama3`가 `RopeScalingKind::Llama3`로 해석됨, Gemma 3 레이어 종류별, Phi `partial_rotary_factor`, DeepSeek-V4 `compress` YaRN(리뷰에서 추가, revert 실행은 하지 않음). lifter 단위 테스트는 `mlxcel-core`에 있다. +- 수정을 stash하고 테스트만 남기면 앞의 테스트 4개가 모두 실패하고, 89개 행 모두 실패를 보고한다. +- 실제 체크포인트: `~/models_archive/mlx/llama-3.2-1b-instruct-4bit`와 `~/models_archive/mlx/qwen3-30b-a3b-4bit`의 transformers-5 사본. `config.json`만 다시 쓰고(flat 키를 transformers 5.20.1이 해당 설정에 내보내는 것과 동일한 `rope_parameters` 블록으로 이동) 가중치는 심볼릭 링크로 연결했다. 명령: `MLXCEL_SDPA_DETERMINISTIC=1 gpu-lock run -- mlxcel generate -m -p "" -n 64 --temp 0 --show-reasoning`. 이 브랜치에서는 두 사본 모두 원본과 바이트 단위로 일치하고, 원본은 `main`과 일치한다. `main`(자기 일치 확인함)에서는 Llama 사본이 `and, and, ...`로, Qwen3 사본이 `. The. The. ...`로 퇴화한다. +- fmt, clippy(mlxcel과 mlxcel-core에 `--lib --tests -D warnings`), 영향받는 계열 테스트 필터(647개 테스트), 저장소 contract 검사가 로컬에서 통과한다. + +## 5. 리뷰 중 변경 + +리뷰에서 CRITICAL 또는 HIGH 이슈는 발견되지 않았다. 커밋 `1684f151`에서 후속 수정 두 건을 적용했다. + +- DeepSeek-V4 `main` / `compress` 레이아웃. 첫 버전은 이를 전부 `None`인 `rope_scaling`으로 만들었다. +- `rope_type` 스킴을 가리던 `"type": null` 키. + +## 6. 미해결 사항 + +- 비전 인코더 설정(Pixtral, Llama 4 vision 등)과 `mlxcel-xla`. 이슈에 나열되지 않았다. +- DBRX는 transformers 5가 여전히 `rope_theta`를 두는 `attn_config`에서 끌어올리지만, 최상위 `rope_parameters`는 그곳으로 복사되지 않는다. + +## 7. 변경 요약 + +| 항목 | 값 | +|------|----| +| 변경된 파일 | 96 | +| 추가된 줄 | 1134 | +| 삭제된 줄 | 116 | + +| 커밋 | 메시지 | +|------|--------| +| `0681d42b` | fix(models): lift transformers 5 rope_parameters for every family | +| `1684f151` | fix(models): lift DeepSeek-V4 main/compress rope_parameters |