fix: guard conv-weight PyTorch->MLX transpose with layout detection

[inureyes]

Summary

Several loaders transposed conv weights from PyTorch to MLX channel-last layout unconditionally while sanitizing a checkpoint. Pre-converted mlx-community checkpoints already store these weights channel-last, so the unconditional transpose double-converts and corrupts the shape. For mlx-community/gemma-4-e4b-it-qat-4bit this turned the audio subsample conv weight [128,3,3,1] into [128,3,1,3], which MLX conv2d rejected (input C_in=1 vs weight C_in=3), aborting the server. This adds a shared layout-detection guard 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.

What changed

• src/loading/mod.rs: new shared helpers conv2d_weight_is_channel_last(shape) and conv1d_weight_is_channel_last(shape) (both pub(crate), #[must_use]).
• src/loading/vlm_special.rs: should_transpose_phi3_patch_embedding now delegates to conv2d_weight_is_channel_last (behavior-preserving); flatten_phi4mm_patch_embedding gates only the transpose and keeps the reshape/flatten.
• src/loading/vlm_gemma.rs: sanitize_gemma4_audio_weights gates the subsample conv2d on !conv2d_weight_is_channel_last and the Conformer depthwise lconv1d on !conv1d_weight_is_channel_last (the confirmed crash site).
• src/vision/detection/rt_detr_v2/sanitize.rs: maybe_transpose_conv gated on !conv2d_weight_is_channel_last.
• src/loading/vlm_nemotron_h_nano_omni.rs: sanitize_audio_weights gates the 4D conv2d branch and the depthwise conv1d (keyed on the depthwise name); pointwise conv1d keeps the existing unconditional transpose (see note below).
• Tests added beside each touched sanitizer: helper unit tests, PyTorch-layout-is-transposed, MLX-layout-is-skipped, idempotency (conv2d 4D and depthwise conv1d 3D), and a Gemma 4 audio regression asserting [128,3,3,1], [32,3,3,128], [1024,5,1] are preserved.

Predicates and shape reasoning

conv2d_weight_is_channel_last(shape) returns true (already MLX [out, kH, kW, in], skip transpose) when shape.len() == 4 && shape[0] >= shape[1] && shape[0] >= shape[2] && shape[1] == shape[2]. This is the validated heuristic already used by should_transpose_phi3_patch_embedding: in channel-last layout the output-channel count leads and dominates both spatial dims, and the kernel dims are square. PyTorch [out, in, kH, kW] breaks dim1 == dim2 (e.g. [128,1,3,3] has dim1=1 != dim2=3) or out >= dim1 once in-channels exceed out-channels. Confirmed: Gemma [128,3,3,1] and [32,3,3,128] are detected as MLX (skip); PyTorch [128,1,3,3] and [32,128,3,3] are detected as PyTorch (transpose). It is idempotent: transposing [128,1,3,3] yields [128,3,3,1], which the predicate then recognizes as channel-last and leaves alone.

conv1d_weight_is_channel_last(shape) returns true when shape.len() == 3 && shape[2] == 1 && shape[1] > 1. This is a depthwise-only predicate. For a depthwise kernel the in-channel count per group is 1, so MLX channel-last [out, kW, 1] carries the 1 in the trailing dim while PyTorch [out, 1, kW] carries it in the middle dim. Confirmed: the Gemma lconv1d MLX form [1024,5,1] is detected as channel-last (skip) and its PyTorch form [1024,1,5] as PyTorch (transpose), idempotently. The predicate deliberately 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 are not separable from shape alone.

Sites gated vs left untouched

Gated (were unconditional): the Gemma 4 audio conv2d and depthwise conv1d, the phi4mm patch-embedding transpose, the rt_detr_v2 conv transpose, and the nemotron conv2d plus depthwise conv1d. Left untouched (already guarded or out of scope, per minimal blast radius): gemma3n_needs_conv_transpose (its own shape[1] > shape[2] guard), and the siglip/pixtral/internvl/minicpmo conv handling, which lives inside the vision::encoders::* model constructors rather than the loader sanitizers, so delegating it here would not be behavior-preserving and would widen scope beyond loaders/sanitizers.

Nemotron audio: maintainer verification note

Only the Gemma 4 audio checkpoint was verified locally. For the nemotron Parakeet sound encoder, the upstream mlx-vlm sanitize_audio_weights transposes 3D/4D *.weight unconditionally (no layout guard), so mlx-community nemotron checkpoints are expected to be PyTorch-layout and the transpose is currently always needed. The Parakeet convolution module mixes a depthwise conv1d ([out, 1, kW]) with two pointwise conv1d layers ([out, in, 1]), whose channel-last and PyTorch shape signatures collide. The depthwise weight is layout-detectable and is guarded; the pointwise weights keep the unconditional transpose to preserve current behavior and match upstream. A fully general channel-last nemotron audio checkpoint would need maintainer verification with a real checkpoint to settle the pointwise conv1d layout. The existing nemotron sanitize tests ([4,8,3] conv1d and [4,1,3,3] conv2d both transposed) remain green.

Build and test note

This dev host is Linux + CUDA (GB10) with no MLX build artifacts; a cold MLX C++ compile exceeds the watchdog and the quantized audio path may not be runnable here, so no compiling cargo command was run. Only cargo fmt --check (formatting, no compile) was run and passes. The fix is device-independent weight-sanitize shape logic verified by close reading of the diff and the shape arithmetic. Build, unit-test execution, and runtime validation of the audio generation path defer to CI (cargo-deny + cargo-fmt) and the maintainer's Apple Silicon make verify.

Closes #428

[inureyes]

Implementation Review Summary

Intent

Stop loaders from double-transposing conv weights for already-channel-last mlx-community checkpoints (issue #428, confirmed gemma-4-e4b-it-qat-4bit audio crash) by gating each unconditional PyTorch->MLX transpose on a shared shape-based layout guard.

Reviewed statically with extra rigor: cargo is unavailable on this host, so compile-level correctness (imports, pub(crate) paths, types, borrows, call-site type-checking) was verified by close reading, and rustfmt --edition 2024 --check was run on every touched file (passes). A second independent automated reviewer supplied supplementary input.

Findings Addressed

• RT-DETRv2 raw-HF stem conv corrupted by the layout guard (HIGH, regression introduced by this PR). maybe_transpose_conv in rt_detr_v2/sanitize.rs runs only inside sanitize(), which model.rs calls only when needs_sanitize() has already proven the checkpoint is raw HuggingFace (PyTorch) layout. The guard there adds no protection against fix: guard conv-weight PyTorch→MLX transpose with layout detection (Gemma4 audio crash + same-pattern loaders) #428 (the needs_sanitize marker check already prevents the double-transpose on mlx-community checkpoints) and actively breaks the documented raw-HF path: the ResNet-vd stem first conv is 3x3 over 3-channel RGB (num_channels=3, embedding_size=64), so its PyTorch shape [32,3,3,3] has in==kH==kW and conv2d_weight_is_channel_last returns true, wrongly skipping the transpose and leaving the stem in OIHW. Fixed by reverting maybe_transpose_conv to its original unconditional 4D transpose (with a doc comment explaining why no guard belongs there), removing the two new tests that asserted the wrong skip/idempotency behavior, keeping the correct PyTorch-transpose test, and adding a regression test that the ambiguous [32,3,3,3] stem is still transposed. rustfmt clean.

Remaining Items

• Nemotron pointwise conv1d layout (MEDIUM, informational, pre-existing). The pointwise conv1d ([out,in,1] PyTorch vs [out,1,in] MLX) is shape-indistinguishable from MLX depthwise [out,kW,1], so the PR leaves it on the unconditional transpose. This is byte-for-byte identical to main (pre-PR all 3D weights transposed unconditionally) and matches the upstream mlx-vlm sanitizer, so the PR introduces no regression. Whether mlx-community nemotron checkpoints could ever ship pointwise already-channel-last is a pre-existing question the author already flagged for maintainer verification with a real Parakeet checkpoint. No code change required here.

Verification

• All stated requirements implemented (all six unguarded sites from the issue table addressed; shared conv2d_weight_is_channel_last / conv1d_weight_is_channel_last helpers added in loading/mod.rs, pub(crate), #[must_use])
• No placeholder/mock code remaining
• Integrated into project code flow (helpers are called from every gated sanitizer; should_transpose_phi3_patch_embedding delegates to the helper, behavior-preserving against its existing test)
• Project conventions followed (rustfmt passes on all touched files; siglip/pixtral/internvl/minicpmo inline guards correctly left untouched as out-of-scope vision-encoder constructors with identical logic, no inconsistency)
• Existing modules reused where applicable (predicate consolidation; no reimplementation)
• No unintended structural changes (loaders/sanitizers only; no inference-core op changes)
• Tests pass (cannot execute cargo test on this Linux+CUDA host with no MLX build artifacts; defers to CI / maintainer Apple Silicon make verify. Tests read correctly: both layouts, idempotency, and the Gemma 4 regression shapes [128,3,3,1]/[32,3,3,128]/[1024,5,1] are asserted with correct expectations.)

Predicate-correctness verdict

PASS. conv2d guard: [128,3,3,1]->skip, [32,3,3,128]->skip, [128,1,3,3]->transpose, [32,128,3,3]->transpose, idempotent. conv1d depthwise guard: [1024,5,1]->skip, [1024,1,5]->transpose, [4,8,3]->transpose, idempotent. The depthwise key scoping in nemotron is correct (pointwise keys lack depthwise). The known heuristic blind spot (in==kH==kW conv2d misread as channel-last) materializes ONLY at the RT-DETR stem; phi3/phi4mm patch embeds (in=3, kernel=14) and the nemotron subsampling conv2ds (in=1 with kernel 3, or 1x1 pointwise with in>1) do not collide.

Regression verdict

After the fix: PASS. Every gated site still transposes genuine PyTorch-layout weights. The only real regression (RT-DETR stem) is fixed. phi3 delegation is behavior-preserving; nemotron pointwise and gemma3n are unchanged from main.

[inureyes]

PR Finalization

CHANGELOG

Added an [Unreleased] section above [v0.3.3] in CHANGELOG.md with a single Fixed bullet describing the double-transpose crash, the confirmed Gemma 4 shape corruption, the four affected sanitizers, the two predicates, and the #428 reference.

Docs

No user-facing API, CLI, or UI change. No docs page added or required beyond the changelog entry.

Coverage sanity (static read, no cargo run)

All gated sites have full test coverage:

• conv2d_weight_is_channel_last / conv1d_weight_is_channel_last helpers: unit tests in src/loading/tests.rs (MLX-detected, PyTorch-detected, non-4D/non-3D rejected).
• Gemma 4 audio sanitizer (sanitize_gemma4_audio_weights): three tests in src/loading/vlm_gemma_tests.rs (channel-last preserved, PyTorch transposed, idempotent).
• phi4mm patch-embed (flatten_phi4mm_patch_embedding): covers both layouts and already-flat in src/loading/vlm_special_tests.rs.
• Nemotron audio (sanitize_audio_weights): six new tests in src/loading/vlm_nemotron_h_nano_omni.rs (skip channel-last depthwise, transpose PyTorch depthwise, idempotent, pointwise always-transposes, skip channel-last conv2d, conv2d idempotent).
• RT-DETRv2 (maybe_transpose_conv): the per-weight guard was intentionally not applied here because needs_sanitize already gates the entire pipeline to raw-HF checkpoints. Two new tests document PyTorch-layout and ambiguous-RGB-stem behavior.

No gap found. No test was added in this finalization pass.

Formatting

All eight changed files pass rustfmt --edition 2024 --check via the pinned 1.93.1 toolchain. No reformatting needed.

No cargo commands run

Only rustfmt --check (format check, no compile) was invoked.

Not merged

Push only. No merge, no label changes.