Repository navigation
fix(models): Lift transformers 5 rope_parameters for every family #2443
Description
Problem
transformers 5 saves rope_theta and rope_scaling only inside a nested rope_parameters block. #2439 (closes #2384) taught Qwen3 and GraniteMoeHybrid to lift it through deserialize_lifting_rope_parameters (src/models/rope_utils.rs:811). Every other config struct that reads a flat rope_theta with a serde default (10000 for llama3, qwen3_moe, gemma2, phi3, olmo2) loads a transformers-5 checkpoint without error, decodes fluently on short prompts, and drifts on long ones. A Llama 3 config also loses its rope_type: "llama3" scaling. 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 family below gets a fix from parity.
| Status | Config structs (file:line) |
|---|---|
| Lifted | qwen3.rs:57 (plus Mage-VL text config), granitemoehybrid.rs:119 |
| Own handling, flat wins | lfm2.rs:200; nested-native: ministral3.rs:244, gemma4.rs:186, glm_moe_dsa.rs:96, laguna.rs:94, mellum.rs:83, qwen3_5.rs:171, muse_glimmer_config.rs:85, cohere_compass_config.rs:88, embedding_gemma2_config.rs:96 |
| Own handling, nested wins (inconsistent) | exaone_moe.rs:121, mistral4.rs:144 |
| Ad hoc call-site lift | src/loading/vlm_qwen.rs:1078 (GLM-OCR), src/loading/vlm_qwen.rs:1236 (GLM-4V MoE) |
| Uncovered, shared args | llama3.rs:55 (Llama, Qwen2, IQuestCoder, Mistral3 llama fallback) |
| Uncovered, Qwen/Gemma | qwen2_moe.rs:50, qwen2_vl.rs:43, qwen3_moe.rs:56, qwen3_next.rs:100, qwen3_vl.rs:47, qwen3_vl_moe.rs:52, gemma.rs:45, gemma2.rs:50, gemma3.rs:62, gemma3.rs:87, gemma3n.rs:56, recurrent_gemma.rs:61 |
| Uncovered, Mistral/Phi/OLMo/Granite | mixtral.rs:44, phi.rs:66, phi3.rs:47, phimoe.rs:66, phixtral.rs:195, olmo.rs:45, olmo2.rs:54, olmo3.rs:46, olmoe.rs:43, granite.rs:60 |
| Uncovered, DeepSeek/GLM/Kimi | deepseek.rs:51, deepseek_v2.rs:103, deepseek_v3.rs:99, deepseek_v32.rs:148, deepseek_v4.rs:222, glm4.rs:42, glm4_moe.rs:56, glm4_moe_lite.rs:58, glm4v.rs:48, kimi_linear.rs:69, dots1.rs:90 |
| Uncovered, other text | afmoe.rs:172, apertus.rs:66, baichuan.rs:40, bailing_moe.rs:170, bailing_moe_linear.rs:210, bitnet.rs:71, cohere.rs:39, cohere2.rs:43, cohere2_moe.rs:71, dbrx.rs:97, ernie4_5.rs:50, ernie4_5_moe.rs:49, exaone.rs:40, exaone4.rs:43, falcon_h1.rs:93, gpt_oss.rs:56, helium.rs:103, hunyuan_moe.rs:55, hunyuan_v1_dense.rs:42, internlm2.rs:73, internlm3.rs:77, iquestloopcoder.rs:146, klear.rs:133, llada2_moe/mod.rs:170, llama4.rs:243, longcat_flash_ngram.rs:66, mimo.rs:45, minicpm.rs:49, minicpm3.rs:60, minimax.rs:49, minimax_m3_config.rs:95, nanochat_config.rs:40, nemotron.rs:61, nemotron_nas.rs:146, seed_oss.rs:69, smollm3.rs:46, stablelm.rs:39, starcoder2.rs:41, step3p5.rs:55, telechat3.rs:109 |
| Uncovered, VLM text configs | ernie4_5_moe_vl.rs:107, falcon_ocr.rs:91, hunyuan_vl.rs:56, jina_vlm.rs:86, mllama/config.rs:48, molmo.rs:77, molmo2.rs:57, moondream2.rs:59, paddleocr_vl.rs:50, youtu_vl_lm_config.rs:72 |
Uncovered, outside src/models |
src/audio/qwen3_omni_moe/speech_config.rs:47, src/audio/qwen3_omni_moe/talker.rs:82, src/lib/mlxcel-core/src/drafter/dflash/config.rs:98 |
Paths without a directory are under src/models/.
Proposed Solution
There is no raw-Value choke point. sanitize_config_json (src/models/sanitize.rs:2589) is a string replace that most family load functions skip: src/models has 125 direct config.json reads (for example llama3.rs:1237, qwen3_moe.rs:1197, gemma2.rs:517), and VLMs nest rope_parameters under text_config, which a top-level lift misses. A top-level lift in sanitize_config_json is rejected for both reasons.
The type is the choke point. Apply the #2439 pattern (#[serde(remote = "Self")] plus a hand Deserialize calling deserialize_lifting_rope_parameters) to every uncovered struct, through a one-line macro in rope_utils.rs such as impl_rope_lifting_deserialize!(ModelArgs). Move lift_rope_parameters into mlxcel-core and re-export it from rope_utils, so the drafter config can use it. Make exaone_moe and mistral4 let flat win, matching the lifter, and replace both vlm_qwen.rs call-site lifts with the shared lifter.
Implement everything in one PR, run the local gate once at the end, and keep pushes to a minimum.
Acceptance Criteria
- A table-driven test covers every struct in the uncovered rows plus
exaone_moeandmistral4: a nested-only config (rope_parameters: {"rope_theta": 123456.0, "rope_type": "default"}, no flat key) yields 123456, and flat 500000 with nested 123456 yields 500000. On current code the nested-only case fails. - A nested
rope_type: "llama3"block onllama3::ModelArgsyields llama3 scaling fromrope_scaling_kind(). - A transformers-5 copy (only
config.jsonrewritten) of~/models_archive/mlx/llama-3.2-1b-instruct-4bitand of~/models_archive/mlx/qwen3-30b-a3b-4bitdecodes the long prompt byte-identically to the original.
Verification
cargo fmt --all -- --check
cargo clippy --workspace --all-targets --features cuda -- -D warnings
cargo test --workspace --profile test-fast --features cuda -- --test-threads=1
# real checkpoint, prompt = CONTRIBUTING.md plus a question, as in #2439
MLXCEL_SDPA_DETERMINISTIC=1 gpu-lock run --tag rope-lift -- mlxcel generate -m <orig|t5-copy> -p "$PROMPT" -n 64 --temp 0
Activity
fix(models): lift DeepSeek-V4 main/compress rope_parameters
fix(models): lift transformers 5 rope_parameters for every family (#2448