diff --git a/CHANGELOG.md b/CHANGELOG.md index c707cc53..c726b6d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). +## [Unreleased] + +### Fixed +- **Double-transpose crash on mlx-community conv checkpoints (Gemma 4 audio, phi4mm patch-embed, nemotron audio, RT-DETRv2).** Several weight-sanitizer functions transposed conv weights from PyTorch `[out, in, kH, kW]` to MLX channel-last `[out, kH, kW, in]` unconditionally. Pre-converted mlx-community checkpoints already store these weights in channel-last order, so the unconditional transpose double-converted them and produced a corrupted shape. The confirmed crash: loading `mlx-community/gemma-4-e4b-it-qat-4bit` turned the audio subsample conv weight `[128, 3, 3, 1]` into `[128, 3, 1, 3]`, which MLX conv2d rejected because the input C_in=1 did not match the weight C_in=3. All four affected sanitizers now check the tensor shape before transposing: `conv2d_weight_is_channel_last` (already-MLX `[out, kH, kW, in]` skips; PyTorch `[out, in, kH, kW]` transposes) and `conv1d_weight_is_channel_last` (depthwise-only; MLX `[out, kW, 1]` skips; PyTorch `[out, 1, kW]` transposes). Both predicates are idempotent. Resolves #428. + ## [v0.3.3] - 2026-06-23 ### Added diff --git a/src/loading/mod.rs b/src/loading/mod.rs index 6ee8778a..5af4cc08 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -88,6 +88,53 @@ pub(super) fn parse_model_config(config_str: &str) -> Resul serde_json::from_str(config_str).map_err(|e| anyhow::anyhow!("Failed to parse config: {}", e)) } +/// Layout-detection guard for conv weights loaded from a checkpoint. +/// +/// Several loaders transpose conv weights from PyTorch to MLX channel-last +/// layout while sanitizing a checkpoint. Pre-converted `mlx-community` +/// checkpoints already store these weights channel-last, so an unconditional +/// transpose double-converts and corrupts the shape (see issue #428). These +/// helpers detect the layout from the shape so the transpose fires only for +/// genuine PyTorch-layout weights, which is safe in both directions: a true +/// PyTorch weight is still transposed and an already-MLX weight is left as-is. +/// +/// Returns `true` when a 4D conv2d weight is already in MLX channel-last +/// `[out, kH, kW, in]` layout (skip the transpose). PyTorch layout is +/// `[out, in, kH, kW]`. +/// +/// Heuristic: in channel-last layout the output-channel count is the leading +/// dim and dominates both spatial dims, and the two kernel dims are equal +/// (square kernels are the norm for these patch/subsample/backbone convs). +/// PyTorch layout puts `in` in dim 1, which generally breaks `dim1 == dim2` +/// (e.g. `[128, 1, 3, 3]` has `dim1=1 != dim2=3`) or `out >= dim1` once the +/// in-channel count exceeds the output count. This mirrors the validated +/// `should_transpose_phi3_patch_embedding` predicate. +#[must_use] +pub(crate) fn conv2d_weight_is_channel_last(shape: &[i32]) -> bool { + shape.len() == 4 && shape[0] >= shape[1] && shape[0] >= shape[2] && shape[1] == shape[2] +} + +/// Returns `true` when a 3D depthwise conv1d weight is already in MLX +/// channel-last `[out, kW, in]` layout (skip the transpose). PyTorch layout is +/// `[out, in, kW]`. +/// +/// This predicate is valid only for depthwise kernels, where the in-channel +/// count per group is `1`. In that case channel-last `[out, kW, 1]` carries the +/// `1` in the trailing dim, while PyTorch `[out, 1, kW]` carries it in the +/// middle dim. Detecting `shape[2] == 1 && shape[1] > 1` therefore distinguishes +/// the two depthwise layouts unambiguously: the confirmed Gemma 4 audio kernel +/// `[1024, 5, 1]` is recognized as already-MLX (skip), and its PyTorch form +/// `[1024, 1, 5]` as needing the transpose. +/// +/// It does NOT generalize to pointwise conv1d (`kernel == 1`): PyTorch pointwise +/// `[out, in, 1]` and MLX depthwise `[out, kW, 1]` share the same shape +/// signature and cannot be told apart from shape alone. Callers that may receive +/// pointwise weights must scope this guard to depthwise keys. +#[must_use] +pub(crate) fn conv1d_weight_is_channel_last(shape: &[i32]) -> bool { + shape.len() == 3 && shape[2] == 1 && shape[1] > 1 +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum Qwen35VlmKind { Dense, diff --git a/src/loading/tests.rs b/src/loading/tests.rs index d0a51b2a..4283ed96 100644 --- a/src/loading/tests.rs +++ b/src/loading/tests.rs @@ -13,8 +13,9 @@ // limitations under the License. use super::{ - Qwen35VlmKind, model_path_str, parse_eos_token_ids, qwen35_vlm_kind, read_eos_token_ids, - require_qwen35_vlm_kind, resolve_model_dir, + Qwen35VlmKind, conv1d_weight_is_channel_last, conv2d_weight_is_channel_last, model_path_str, + parse_eos_token_ids, qwen35_vlm_kind, read_eos_token_ids, require_qwen35_vlm_kind, + resolve_model_dir, }; use crate::model_metadata::{ DirectoryLoadRoute, ModelCapabilities, ModelKind, ModelLoadPolicy, WeightLoadRoute, @@ -347,3 +348,29 @@ fn model_load_policy_handles_mistral3_mistral4_wrapper_config() { ); assert_eq!(policy.weight_route, Some(WeightLoadRoute::LlamaFamily)); } + +#[test] +fn conv2d_weight_is_channel_last_detects_mlx_layout() { + // MLX channel-last [out, kH, kW, in]: out dominates and kernels are square. + assert!(conv2d_weight_is_channel_last(&[128, 3, 3, 1])); + assert!(conv2d_weight_is_channel_last(&[32, 3, 3, 128])); + // PyTorch [out, in, kH, kW]: in breaks dim1 == dim2 (or out >= dim1). + assert!(!conv2d_weight_is_channel_last(&[128, 1, 3, 3])); + assert!(!conv2d_weight_is_channel_last(&[32, 128, 3, 3])); + // Non-4D shapes are never treated as conv2d channel-last. + assert!(!conv2d_weight_is_channel_last(&[128, 3, 3])); + assert!(!conv2d_weight_is_channel_last(&[4, 4])); +} + +#[test] +fn conv1d_weight_is_channel_last_detects_depthwise_mlx_layout() { + // MLX depthwise [out, kW, in=1]: trailing in dim is 1. + assert!(conv1d_weight_is_channel_last(&[1024, 5, 1])); + // PyTorch depthwise [out, in=1, kW]: middle in dim is 1, kW > 1. + assert!(!conv1d_weight_is_channel_last(&[1024, 1, 5])); + // A regular conv1d [out, in, kW] with in, kW > 1 is treated as PyTorch. + assert!(!conv1d_weight_is_channel_last(&[4, 8, 3])); + // Non-3D shapes are never treated as depthwise conv1d channel-last. + assert!(!conv1d_weight_is_channel_last(&[1024, 5, 1, 1])); + assert!(!conv1d_weight_is_channel_last(&[4, 4])); +} diff --git a/src/loading/vlm_gemma.rs b/src/loading/vlm_gemma.rs index 47bc632d..d14a3737 100644 --- a/src/loading/vlm_gemma.rs +++ b/src/loading/vlm_gemma.rs @@ -280,17 +280,25 @@ fn sanitize_gemma4_audio_weights(weights: &mut mlxcel_core::weights::WeightMap) mlxcel_core::array_shape(w) }; - // Conv2d: PyTorch [out, in, kH, kW] -> MLX [out, kH, kW, in] + // Conv2d: PyTorch [out, in, kH, kW] -> MLX [out, kH, kW, in]. + // Skip already-channel-last checkpoints to avoid double-converting the + // shape (issue #428): the mlx-community gemma4 weights ship channel-last. if key.contains("subsample_conv_projection") && key.contains("conv.weight") && shape.len() == 4 + && !crate::loading::conv2d_weight_is_channel_last(&shape) { let w = weights.remove(&key).unwrap(); let transposed = mlxcel_core::transpose_axes(&w, &[0, 2, 3, 1]); weights.insert(key, transposed); } - // Conv1d depthwise: PyTorch [out, in, kW] -> MLX [out, kW, in] - else if key.contains("depthwise_conv1d.weight") && shape.len() == 3 { + // Conv1d depthwise: PyTorch [out, in, kW] -> MLX [out, kW, in]. + // The Conformer lconv1d is depthwise (in == 1), so the channel-last + // form is [out, kW, 1]; skip it when already converted (issue #428). + else if key.contains("depthwise_conv1d.weight") + && shape.len() == 3 + && !crate::loading::conv1d_weight_is_channel_last(&shape) + { let w = weights.remove(&key).unwrap(); let transposed = mlxcel_core::transpose_axes(&w, &[0, 2, 1]); weights.insert(key, transposed); diff --git a/src/loading/vlm_gemma_tests.rs b/src/loading/vlm_gemma_tests.rs index 2dda9c02..8479e250 100644 --- a/src/loading/vlm_gemma_tests.rs +++ b/src/loading/vlm_gemma_tests.rs @@ -14,7 +14,7 @@ use super::{ gemma3n_language_model_prefix, gemma3n_metadata, gemma3n_needs_conv_transpose, - sanitize_gemma3n_weights, + sanitize_gemma3n_weights, sanitize_gemma4_audio_weights, }; use mlxcel_core::dtype; use mlxcel_core::weights::WeightMap; @@ -109,3 +109,101 @@ fn sanitize_gemma3n_weights_strips_model_prefix_and_transposes_conv_weights() { vec![8, 3, 3, 16] ); } + +fn audio_conv_shape(weights: &WeightMap, key: &str) -> Vec { + mlxcel_core::array_shape(weights.get(key).unwrap()) +} + +const LAYER0_CONV: &str = "audio_tower.subsample_conv_projection.layer0.conv.weight"; +const LAYER1_CONV: &str = "audio_tower.subsample_conv_projection.layer1.conv.weight"; +const LCONV1D: &str = "audio_tower.layers.0.lconv1d.depthwise_conv1d.weight"; + +#[test] +fn sanitize_gemma4_audio_weights_preserves_channel_last_checkpoint() { + // The mlx-community/gemma-4-e4b-it-qat-4bit checkpoint stores these conv + // weights already in MLX channel-last layout. The sanitizer must leave them + // untouched so the audio subsample conv receives a C_in=1 weight matching + // the C_in=1 input (regression for issue #428). + let mut weights = WeightMap::new(); + weights.insert( + LAYER0_CONV.to_string(), + mlxcel_core::ones(&[128, 3, 3, 1], dtype::FLOAT32), + ); + weights.insert( + LAYER1_CONV.to_string(), + mlxcel_core::ones(&[32, 3, 3, 128], dtype::FLOAT32), + ); + weights.insert( + LCONV1D.to_string(), + mlxcel_core::ones(&[1024, 5, 1], dtype::FLOAT32), + ); + + sanitize_gemma4_audio_weights(&mut weights); + + assert_eq!(audio_conv_shape(&weights, LAYER0_CONV), vec![128, 3, 3, 1]); + assert_eq!(audio_conv_shape(&weights, LAYER1_CONV), vec![32, 3, 3, 128]); + assert_eq!(audio_conv_shape(&weights, LCONV1D), vec![1024, 5, 1]); +} + +#[test] +fn sanitize_gemma4_audio_weights_transposes_pytorch_layout() { + // Synthetic PyTorch-layout weights: conv2d [out, in, kH, kW] and depthwise + // conv1d [out, in=1, kW]. Both must be transposed to channel-last. + let mut weights = WeightMap::new(); + weights.insert( + LAYER0_CONV.to_string(), + mlxcel_core::ones(&[128, 1, 3, 3], dtype::FLOAT32), + ); + weights.insert( + LCONV1D.to_string(), + mlxcel_core::ones(&[1024, 1, 5], dtype::FLOAT32), + ); + + sanitize_gemma4_audio_weights(&mut weights); + + // conv2d [128, 1, 3, 3] --transpose[0,2,3,1]--> [128, 3, 3, 1]. + assert_eq!(audio_conv_shape(&weights, LAYER0_CONV), vec![128, 3, 3, 1]); + // depthwise conv1d [1024, 1, 5] --transpose[0,2,1]--> [1024, 5, 1]. + assert_eq!(audio_conv_shape(&weights, LCONV1D), vec![1024, 5, 1]); +} + +#[test] +fn sanitize_gemma4_audio_weights_is_idempotent() { + // Running the sanitizer twice must equal running it once: the first pass + // converts the PyTorch-layout weights to channel-last, and the second pass + // must detect that layout and leave them unchanged. + let mut once = WeightMap::new(); + once.insert( + LAYER0_CONV.to_string(), + mlxcel_core::ones(&[128, 1, 3, 3], dtype::FLOAT32), + ); + once.insert( + LCONV1D.to_string(), + mlxcel_core::ones(&[1024, 1, 5], dtype::FLOAT32), + ); + sanitize_gemma4_audio_weights(&mut once); + + let mut twice = WeightMap::new(); + twice.insert( + LAYER0_CONV.to_string(), + mlxcel_core::ones(&[128, 1, 3, 3], dtype::FLOAT32), + ); + twice.insert( + LCONV1D.to_string(), + mlxcel_core::ones(&[1024, 1, 5], dtype::FLOAT32), + ); + sanitize_gemma4_audio_weights(&mut twice); + sanitize_gemma4_audio_weights(&mut twice); + + assert_eq!( + audio_conv_shape(&once, LAYER0_CONV), + audio_conv_shape(&twice, LAYER0_CONV) + ); + assert_eq!( + audio_conv_shape(&once, LCONV1D), + audio_conv_shape(&twice, LCONV1D) + ); + // And the idempotent result is the channel-last layout. + assert_eq!(audio_conv_shape(&twice, LAYER0_CONV), vec![128, 3, 3, 1]); + assert_eq!(audio_conv_shape(&twice, LCONV1D), vec![1024, 5, 1]); +} diff --git a/src/loading/vlm_nemotron_h_nano_omni.rs b/src/loading/vlm_nemotron_h_nano_omni.rs index 589d76c9..74f556dd 100644 --- a/src/loading/vlm_nemotron_h_nano_omni.rs +++ b/src/loading/vlm_nemotron_h_nano_omni.rs @@ -88,15 +88,48 @@ fn filter_out_audio_weights(weights: WeightMap, drop_audio: bool) -> WeightMap { .collect() } +/// Decide whether the `sound_encoder` conv weights are in PyTorch `[O, I, K]` +/// layout (need transposing to MLX channel-last) or are already channel-last +/// (must be left untouched). +/// +/// The Parakeet conv module mixes depthwise and pointwise Conv1d whose +/// channel-last vs PyTorch shape signatures are mirror images (depthwise +/// channel-last `[O, kW>1, 1]` vs pointwise channel-last `[O, 1, in>1]`), so no +/// single per-weight shape test can classify both: the earlier per-weight guard +/// left pointwise on an unconditional transpose, which corrupted a channel-last +/// checkpoint's pointwise weights and crashed the audio conv (#428). Decide the +/// whole tower's layout ONCE from an unambiguous depthwise weight: PyTorch +/// depthwise is `[O, 1, kW]` (axis 1 == 1) while channel-last depthwise is +/// `[O, kW, 1]` (axis 2 == 1). mlx-community checkpoints (e.g. +/// Nemotron-3-Nano-Omni) ship the encoder already channel-last, so this returns +/// false and every `sound_encoder` conv transpose is skipped. +fn nemotron_audio_needs_conv_transpose(weights: &WeightMap) -> bool { + for (key, value) in weights.iter() { + if key.starts_with("sound_encoder.encoder.") + && key.contains("depthwise") + && key.ends_with(".weight") + { + let shape = mlxcel_core::array_shape(value); + if shape.len() == 3 { + return shape[1] == 1; + } + } + } + // No depthwise weight found: preserve the historical unconditional transpose. + true +} + /// Mirror upstream `sanitize_audio_weights` from /// https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/nemotron_h_nano_omni/audio.py: /// - drop `sound_encoder.encoder.feature_extractor.*` (training-only), /// - drop `*.num_batches_tracked` BatchNorm scratch state, -/// - transpose Conv1d weights from PyTorch `[O, I, K]` to MLX -/// channel-last `[O, K, I]`, -/// - transpose Conv2d weights from PyTorch `[O, I, kh, kw]` to MLX -/// channel-last `[O, kh, kw, I]`. +/// - transpose Conv1d `[O, I, K]` -> MLX `[O, K, I]` and Conv2d +/// `[O, I, kH, kW]` -> MLX `[O, kH, kW, I]`, but ONLY when the checkpoint is +/// in PyTorch layout. A checkpoint already stored channel-last (detected once +/// via [`nemotron_audio_needs_conv_transpose`]) is left untouched so the +/// pointwise conv1d weights are not corrupted by a double transpose (#428). fn sanitize_audio_weights(weights: WeightMap) -> WeightMap { + let needs_transpose = nemotron_audio_needs_conv_transpose(&weights); let mut out = WeightMap::new(); for (key, value) in weights { if key.starts_with("sound_encoder.encoder.feature_extractor.") { @@ -105,17 +138,16 @@ fn sanitize_audio_weights(weights: WeightMap) -> WeightMap { if key.ends_with(".num_batches_tracked") { continue; } - if key.starts_with("sound_encoder.encoder.") && key.ends_with(".weight") { + if needs_transpose && key.starts_with("sound_encoder.encoder.") && key.ends_with(".weight") + { let shape = mlxcel_core::array_shape(&value); match shape.len() { 3 => { - let transposed = mlxcel_core::transpose_axes(&value, &[0, 2, 1]); - out.insert(key, transposed); + out.insert(key, mlxcel_core::transpose_axes(&value, &[0, 2, 1])); continue; } 4 => { - let transposed = mlxcel_core::transpose_axes(&value, &[0, 2, 3, 1]); - out.insert(key, transposed); + out.insert(key, mlxcel_core::transpose_axes(&value, &[0, 2, 3, 1])); continue; } _ => {} @@ -553,10 +585,10 @@ mod tests { #[test] fn sanitize_audio_weights_transposes_conv_weights() { let mut weights = WeightMap::new(); - // Conv1d weight in PyTorch layout: [out=4, in=8, kernel=3]. + // Conv1d depthwise in PyTorch layout: [out=4, in=1, kernel=3]. weights.insert( "sound_encoder.encoder.layers.0.conv.depthwise_conv.weight".into(), - ones(&[4, 8, 3]), + ones(&[4, 1, 3]), ); // Conv2d weight in PyTorch layout: [out=4, in=1, kh=3, kw=3]. weights.insert( @@ -570,8 +602,8 @@ mod tests { .get("sound_encoder.encoder.layers.0.conv.depthwise_conv.weight") .unwrap(), ); - // After transpose(0,2,1): [4, 3, 8]. - assert_eq!(conv1d_shape, vec![4, 3, 8]); + // After transpose(0,2,1): [4, 3, 1]. + assert_eq!(conv1d_shape, vec![4, 3, 1]); let conv2d_shape = mlxcel_core::array_shape( sanitized @@ -581,4 +613,105 @@ mod tests { // After transpose(0,2,3,1): [4, 3, 3, 1]. assert_eq!(conv2d_shape, vec![4, 3, 3, 1]); } + + fn audio_shape(weights: &WeightMap, key: &str) -> Vec { + mlxcel_core::array_shape(weights.get(key).unwrap()) + } + + const DEPTHWISE: &str = "sound_encoder.encoder.layers.0.conv.depthwise_conv.weight"; + const POINTWISE: &str = "sound_encoder.encoder.layers.0.conv.pointwise_conv1.weight"; + const SUBSAMPLE_CONV2D: &str = "sound_encoder.encoder.subsampling.layers.0.weight"; + + #[test] + fn sanitize_audio_weights_skips_channel_last_depthwise_conv1d() { + // A depthwise conv1d already in MLX channel-last [out, kW, in=1] layout + // must be left untouched (issue #428). + let mut weights = WeightMap::new(); + weights.insert(DEPTHWISE.into(), ones(&[1024, 5, 1])); + let sanitized = sanitize_audio_weights(weights); + assert_eq!(audio_shape(&sanitized, DEPTHWISE), vec![1024, 5, 1]); + } + + #[test] + fn sanitize_audio_weights_transposes_pytorch_depthwise_conv1d() { + // PyTorch depthwise [out, in=1, kW] -> MLX [out, kW, 1]. + let mut weights = WeightMap::new(); + weights.insert(DEPTHWISE.into(), ones(&[1024, 1, 5])); + let sanitized = sanitize_audio_weights(weights); + assert_eq!(audio_shape(&sanitized, DEPTHWISE), vec![1024, 5, 1]); + } + + #[test] + fn sanitize_audio_weights_depthwise_conv1d_is_idempotent() { + let mut once = WeightMap::new(); + once.insert(DEPTHWISE.into(), ones(&[1024, 1, 5])); + let once = sanitize_audio_weights(once); + + let mut twice = WeightMap::new(); + twice.insert(DEPTHWISE.into(), ones(&[1024, 1, 5])); + let twice = sanitize_audio_weights(sanitize_audio_weights(twice)); + + assert_eq!( + audio_shape(&once, DEPTHWISE), + audio_shape(&twice, DEPTHWISE) + ); + assert_eq!(audio_shape(&twice, DEPTHWISE), vec![1024, 5, 1]); + } + + #[test] + fn sanitize_audio_weights_channel_last_pointwise_left_untouched() { + // Regression (#428): in a channel-last checkpoint (depthwise [O, kW, 1]) + // the pointwise conv1d [O, 1, in] must NOT be transposed. The earlier + // per-weight guard transposed the pointwise and crashed the audio conv. + // The depthwise anchors the checkpoint-level layout detection. + let mut weights = WeightMap::new(); + weights.insert(DEPTHWISE.into(), ones(&[1024, 9, 1])); + weights.insert(POINTWISE.into(), ones(&[2048, 1, 1024])); + let sanitized = sanitize_audio_weights(weights); + assert_eq!(audio_shape(&sanitized, DEPTHWISE), vec![1024, 9, 1]); + assert_eq!(audio_shape(&sanitized, POINTWISE), vec![2048, 1, 1024]); + } + + #[test] + fn sanitize_audio_weights_transposes_pytorch_pointwise_conv1d() { + // PyTorch checkpoint (depthwise [O, 1, kW]): pointwise [O, in, 1] is + // transposed to channel-last [O, 1, in]. + let mut weights = WeightMap::new(); + weights.insert(DEPTHWISE.into(), ones(&[1024, 1, 9])); + weights.insert(POINTWISE.into(), ones(&[2048, 1024, 1])); + let sanitized = sanitize_audio_weights(weights); + assert_eq!(audio_shape(&sanitized, POINTWISE), vec![2048, 1, 1024]); + } + + #[test] + fn sanitize_audio_weights_skips_channel_last_conv2d() { + // Conv2d already channel-last [out, kH, kW, in] must be left untouched; + // a channel-last depthwise anchors the checkpoint-level layout detection. + let mut weights = WeightMap::new(); + weights.insert(DEPTHWISE.into(), ones(&[1024, 3, 1])); + weights.insert(SUBSAMPLE_CONV2D.into(), ones(&[32, 3, 3, 1])); + let sanitized = sanitize_audio_weights(weights); + assert_eq!(audio_shape(&sanitized, SUBSAMPLE_CONV2D), vec![32, 3, 3, 1]); + } + + #[test] + fn sanitize_audio_weights_conv2d_is_idempotent() { + // A PyTorch depthwise anchors detection; after the first pass it is + // channel-last, so the second pass skips and the result is stable. + let mut once = WeightMap::new(); + once.insert(DEPTHWISE.into(), ones(&[1024, 1, 3])); + once.insert(SUBSAMPLE_CONV2D.into(), ones(&[4, 1, 3, 3])); + let once = sanitize_audio_weights(once); + + let mut twice = WeightMap::new(); + twice.insert(DEPTHWISE.into(), ones(&[1024, 1, 3])); + twice.insert(SUBSAMPLE_CONV2D.into(), ones(&[4, 1, 3, 3])); + let twice = sanitize_audio_weights(sanitize_audio_weights(twice)); + + assert_eq!( + audio_shape(&once, SUBSAMPLE_CONV2D), + audio_shape(&twice, SUBSAMPLE_CONV2D) + ); + assert_eq!(audio_shape(&twice, SUBSAMPLE_CONV2D), vec![4, 3, 3, 1]); + } } diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index e8c75406..12829086 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -733,7 +733,14 @@ fn flatten_phi4mm_patch_embedding( return mlxcel_core::copy(weight); } - let transposed = mlxcel_core::transpose_axes(weight, &[0, 2, 3, 1]); + // Transpose only genuine PyTorch-layout weights; channel-last checkpoints + // are already `[out, kH, kW, in]` and must skip it (issue #428). The + // reshape/flatten below runs unconditionally either way. + let transposed = if crate::loading::conv2d_weight_is_channel_last(&shape) { + mlxcel_core::copy(weight) + } else { + mlxcel_core::transpose_axes(weight, &[0, 2, 3, 1]) + }; mlxcel_core::reshape(&transposed, &[shape[0], shape[1] * shape[2] * shape[3]]) } @@ -1177,13 +1184,11 @@ pub(super) fn rewrite_phi3_weight_key(key: &str) -> Option { } pub(super) fn should_transpose_phi3_patch_embedding(shape: &[i32]) -> bool { - if shape.len() != 4 { - return false; - } - let out_ch = shape[0]; - let dim1 = shape[1]; - let dim2 = shape[2]; - !(out_ch >= dim1 && out_ch >= dim2 && dim1 == dim2) + // A 4D weight that is not already MLX channel-last is PyTorch layout and + // must be transposed. Delegates to the shared layout guard (issue #428); + // for a non-4D shape `conv2d_weight_is_channel_last` is `false`, so this + // returns `false` (no transpose), preserving the original behavior. + shape.len() == 4 && !crate::loading::conv2d_weight_is_channel_last(shape) } fn remap_phi3_weights(raw_weights: WeightMap) -> WeightMap { diff --git a/src/loading/vlm_special_tests.rs b/src/loading/vlm_special_tests.rs index 960c1db4..3e7dc330 100644 --- a/src/loading/vlm_special_tests.rs +++ b/src/loading/vlm_special_tests.rs @@ -13,14 +13,15 @@ // limitations under the License. use super::{ - cap_molmo2_vit_num_layers, dequantize_moondream3_weight, inherit_quantization_if_missing, - llama4_mm_tokens_per_image, llama4_quantization_params, llama4_token_ids, llama4_vision_prefix, - minicpmv4_6_text_weights, molmo2_max_crops, moondream3_text_config_value, - moondream3_vision_config_value, parse_molmo2_vit_layers, phi3_num_crops, - phi4_siglip_text_config_value, phi4mm_text_config_value, phi4mm_vision_config_value, - remap_minicpmo_text_weights, remap_minicpmv4_6_weights, rewrite_molmo2_weight_key, - rewrite_moondream3_weight_key, rewrite_phi3_weight_key, rewrite_phi4_siglip_weight_key, - rewrite_phi4mm_vision_key, should_transpose_phi3_patch_embedding, + cap_molmo2_vit_num_layers, dequantize_moondream3_weight, flatten_phi4mm_patch_embedding, + inherit_quantization_if_missing, llama4_mm_tokens_per_image, llama4_quantization_params, + llama4_token_ids, llama4_vision_prefix, minicpmv4_6_text_weights, molmo2_max_crops, + moondream3_text_config_value, moondream3_vision_config_value, parse_molmo2_vit_layers, + phi3_num_crops, phi4_siglip_text_config_value, phi4mm_text_config_value, + phi4mm_vision_config_value, remap_minicpmo_text_weights, remap_minicpmv4_6_weights, + rewrite_molmo2_weight_key, rewrite_moondream3_weight_key, rewrite_phi3_weight_key, + rewrite_phi4_siglip_weight_key, rewrite_phi4mm_vision_key, + should_transpose_phi3_patch_embedding, }; use mlxcel_core::dtype; use mlxcel_core::weights::WeightMap; @@ -204,6 +205,25 @@ fn phi3_patch_embedding_transpose_detection_matches_layout_expectations() { assert!(!should_transpose_phi3_patch_embedding(&[1024, 196])); } +#[test] +fn flatten_phi4mm_patch_embedding_flattens_both_layouts_to_same_shape() { + // PyTorch layout [out, in, kH, kW]: transpose to channel-last, then flatten. + let pytorch = mlxcel_core::ones(&[1024, 3, 14, 14], dtype::FLOAT32); + let flat = flatten_phi4mm_patch_embedding(&pytorch); + assert_eq!(mlxcel_core::array_shape(&flat), vec![1024, 3 * 14 * 14]); + + // Already channel-last [out, kH, kW, in]: skip the transpose (issue #428), + // flatten to the same [out, in*kH*kW] shape without double-converting. + let channel_last = mlxcel_core::ones(&[1024, 14, 14, 3], dtype::FLOAT32); + let flat = flatten_phi4mm_patch_embedding(&channel_last); + assert_eq!(mlxcel_core::array_shape(&flat), vec![1024, 14 * 14 * 3]); + + // Non-4D weights are copied through unchanged. + let already_flat = mlxcel_core::ones(&[1024, 196], dtype::FLOAT32); + let out = flatten_phi4mm_patch_embedding(&already_flat); + assert_eq!(mlxcel_core::array_shape(&out), vec![1024, 196]); +} + #[test] fn phi3_num_crops_prefers_preprocessor_then_config_then_default() { assert_eq!( diff --git a/src/vision/detection/rt_detr_v2/sanitize.rs b/src/vision/detection/rt_detr_v2/sanitize.rs index 9a9c3ec9..31a0cd12 100644 --- a/src/vision/detection/rt_detr_v2/sanitize.rs +++ b/src/vision/detection/rt_detr_v2/sanitize.rs @@ -183,6 +183,12 @@ pub fn sanitize(weights: WeightMap) -> WeightMap { /// PyTorch conv weights are `(out, in, kH, kW)`; MLX wants NHWC /// `(out, kH, kW, in)`. Transpose any 4D `*.conv.weight`. +/// +/// This runs only inside [`sanitize`], which the caller invokes only when +/// [`needs_sanitize`] has already determined the checkpoint is raw +/// HuggingFace (PyTorch) layout, so an unconditional transpose is correct and +/// a layout guard here would wrongly skip stem convs whose shape is ambiguous +/// (e.g. a `[out, 3, 3, 3]` RGB 3x3 stem conv). fn maybe_transpose_conv(key: &str, value: UniquePtr) -> UniquePtr { if key.ends_with(".conv.weight") { let shape = mlxcel_core::array_shape(&value); @@ -313,4 +319,30 @@ mod tests { ); assert!(needs_sanitize(&hf_map)); } + + #[test] + fn maybe_transpose_conv_transposes_pytorch_layout() { + // PyTorch stem conv [out=64, in=3, kH=7, kW=7]; in != kernel makes the + // layout unambiguous. transpose[0,2,3,1]: [64, 7, 7, 3]. + let value = mlxcel_core::ones(&[64, 3, 7, 7], mlxcel_core::dtype::FLOAT32); + let out = maybe_transpose_conv("vision.backbone.embedder.0.conv.weight", value); + assert_eq!(mlxcel_core::array_shape(&out), vec![64, 7, 7, 3]); + + // PyTorch conv [out=8, in=16, kH=3, kW=3] where in > out also resolves + // to PyTorch layout. transpose[0,2,3,1]: [8, 3, 3, 16]. + let value = mlxcel_core::ones(&[8, 16, 3, 3], mlxcel_core::dtype::FLOAT32); + let out = maybe_transpose_conv("vision.backbone.embedder.0.conv.weight", value); + assert_eq!(mlxcel_core::array_shape(&out), vec![8, 3, 3, 16]); + } + + #[test] + fn maybe_transpose_conv_transposes_ambiguous_rgb_stem() { + // The ResNet-vd stem first conv is a 3x3 over 3-channel RGB input, so + // its PyTorch shape [out, in=3, kH=3, kW=3] has in == kH == kW. This + // runs only after needs_sanitize confirmed raw-HF layout, so it must + // still be transposed [0,2,3,1] even though the shape looks square. + let value = mlxcel_core::ones(&[32, 3, 3, 3], mlxcel_core::dtype::FLOAT32); + let out = maybe_transpose_conv("vision.backbone.embedder.0.conv.weight", value); + assert_eq!(mlxcel_core::array_shape(&out), vec![32, 3, 3, 3]); + } }