diff --git a/src/models/qwen3_5.rs b/src/models/qwen3_5.rs index bd5a3efa..4508c6e2 100644 --- a/src/models/qwen3_5.rs +++ b/src/models/qwen3_5.rs @@ -2398,11 +2398,10 @@ pub fn sanitize_weights(mut weights: WeightMap, config: &Qwen35Config) -> Weight // Handle per-expert gate_proj/up_proj/down_proj naming variant. // Checkpoints that store experts under `experts.{e}.gate_proj.weight` - // instead of `experts.{e}.w1.weight` use this layout. The mapping - // mirrors the fused gate_up_proj path: gate->w1, up->w3, down->w2. - for (src_proj, dst_proj) in [("gate_proj", "w1"), ("up_proj", "w3"), ("down_proj", "w2")] { - // Skip if this target slot is already populated by the w1/w2/w3 pass above. - if weights.contains_key(format!("{}.{}.weight", base, dst_proj).as_str()) { + // instead of per-expert w1 weights use this layout. Keep the + // stacked names aligned with SwitchGLU::from_weights. + for proj in ["gate_proj", "up_proj", "down_proj"] { + if weights.contains_key(format!("{}.{}.weight", base, proj).as_str()) { continue; } @@ -2413,18 +2412,18 @@ pub fn sanitize_weights(mut weights: WeightMap, config: &Qwen35Config) -> Weight let mut e = 0; while let Some(w) = weights.remove(&format!( "model.layers.{}.mlp.experts.{}.{}.weight", - l, e, src_proj + l, e, proj )) { expert_weights.push(w); if let Some(s) = weights.remove(&format!( "model.layers.{}.mlp.experts.{}.{}.scales", - l, e, src_proj + l, e, proj )) { expert_scales.push(s); } if let Some(b) = weights.remove(&format!( "model.layers.{}.mlp.experts.{}.{}.biases", - l, e, src_proj + l, e, proj )) { expert_biases.push(b); } @@ -2433,16 +2432,16 @@ pub fn sanitize_weights(mut weights: WeightMap, config: &Qwen35Config) -> Weight if !expert_weights.is_empty() { let stacked = stack_arrays(&expert_weights, 0); - weights.insert(format!("{}.{}.weight", base, dst_proj), stacked); + weights.insert(format!("{}.{}.weight", base, proj), stacked); if !expert_scales.is_empty() { let stacked = stack_arrays(&expert_scales, 0); - weights.insert(format!("{}.{}.scales", base, dst_proj), stacked); + weights.insert(format!("{}.{}.scales", base, proj), stacked); } if !expert_biases.is_empty() { let stacked = stack_arrays(&expert_biases, 0); - weights.insert(format!("{}.{}.biases", base, dst_proj), stacked); + weights.insert(format!("{}.{}.biases", base, proj), stacked); } } } @@ -2481,31 +2480,6 @@ pub fn sanitize_weights(mut weights: WeightMap, config: &Qwen35Config) -> Weight } } - // 8. Rename switch_mlp.{gate_proj,up_proj,down_proj} → switch_mlp.{w1,w3,w2} - // Pre-quantized MoE models use gate_proj/up_proj/down_proj naming, - // but SparseMoeBlock expects w1/w2/w3 naming. - let rename_map = [ - ("switch_mlp.gate_proj.", "switch_mlp.w1."), - ("switch_mlp.up_proj.", "switch_mlp.w3."), - ("switch_mlp.down_proj.", "switch_mlp.w2."), - ]; - let keys_to_rename: Vec = weights - .keys() - .filter(|k| rename_map.iter().any(|(from, _)| k.contains(from))) - .cloned() - .collect(); - for key in keys_to_rename { - for (from, to) in &rename_map { - if key.contains(from) { - let new_key = key.replace(from, to); - if let Some(v) = weights.remove(&key) { - weights.insert(new_key, v); - } - break; - } - } - } - weights } diff --git a/src/models/qwen3_5_tests.rs b/src/models/qwen3_5_tests.rs index e61e9af4..f84a176b 100644 --- a/src/models/qwen3_5_tests.rs +++ b/src/models/qwen3_5_tests.rs @@ -21,7 +21,8 @@ //! that need a real Qwen 3.5 model and are gated behind hardware availability. use super::qwen3_5::{ - Qwen35Config, rebuild_with_zero_tail, sanitize_weights, zero_per_row_kv_tail, + Qwen35Config, rebuild_with_zero_tail, sanitize_weights, + zero_per_row_kv_tail, }; use mlxcel_core::dtype; use mlxcel_core::layers::KVCache; @@ -309,3 +310,34 @@ fn sanitize_weights_drops_lm_head_when_tied_embeddings() { ); assert!(sanitized.contains_key("model.embed_tokens.weight")); } + +#[test] +#[ignore = "requires serial MLX execution"] +fn sanitize_weights_stacks_per_expert_switch_proj_names_for_loader() { + let root = "model.layers.0.mlp.experts"; + let mut weights = WeightMap::new(); + for expert in 0..2 { + for proj in ["gate_proj", "up_proj", "down_proj"] { + weights.insert( + format!("{root}.{expert}.{proj}.weight"), + mlxcel_core::from_slice_f32(&[expert as f32; 8], &[2, 4]), + ); + } + } + + let mut config = make_tiny_config(); + config.num_experts = 2; + config.num_experts_per_tok = 1; + config.decoder_sparse_step = 1; + config.moe_intermediate_size = 2; + config.shared_expert_intermediate_size = 2; + + let sanitized = sanitize_weights(weights, &config); + + assert!(sanitized.contains_key("model.layers.0.mlp.switch_mlp.gate_proj.weight")); + assert!(sanitized.contains_key("model.layers.0.mlp.switch_mlp.up_proj.weight")); + assert!(sanitized.contains_key("model.layers.0.mlp.switch_mlp.down_proj.weight")); + assert!(!sanitized.contains_key("model.layers.0.mlp.switch_mlp.w1.weight")); + assert!(!sanitized.contains_key("model.layers.0.mlp.switch_mlp.w2.weight")); + assert!(!sanitized.contains_key("model.layers.0.mlp.switch_mlp.w3.weight")); +}