Fix Request: Tensor shape mismatch in audio_tower causing complete fallback from VLM to LLM on gemma-4-e4b-it-qat-OptiQ-4bit
Hi mlx-community team,
Thank you for providing the high-quality mlx-optiq mixed-precision quants! While benchmarking the gemma-4-e4b-it-qat-OptiQ-4bit model locally using the MLX framework, I noticed that the engine consistently fails to load it as a Vision-Language Model (VLM), forcing a complete fallback to a text-only Language Model (LLM).
This regression disables the model's native multimodal (vision/audio) capabilities and prevents downstream serving architectures (such as oMLX) from utilizing speculative multi-token prediction (VLM MTP) via the companion 74M assistant/drafter model.
Error Logs / Traces:
During the model initialization phase, the runtime catches a shape mismatch on the convolutional projection layer within the audio tower:
Code snippet
WARNING - EnginePool - VLM loading failed for gemma-4-e4b-it-qat-OptiQ-4bit, falling back to LLM:
Expected shape (128, 3, 3, 1) but received shape (128, 1, 3, 3) for parameter audio_tower.subsample_conv_projection.layer0.conv.weight
INFO - EnginePool - Successfully loaded gemma-4-e4b-it-qat-OptiQ-4bit as LLM (fallback from VLM)
Root Cause Analysis:
The native Google Gemma-4 QAT base incorporates an integrated audio and vision processing architecture. During the mlx-optiq quantization process or the subsequent safetensors export pass, the weights for audio_tower.subsample_conv_projection.layer0.conv.weight appear to have been inadvertently transposed/flattened into a channel-first layout (128, 1, 3, 3) instead of preserving the layout expected by native MLX VLM parsing implementations (128, 3, 3, 1).
Because the shapes do not match on this specific parameter tensor, inference engines discard the entire multimodal computational graph (VLMBatchedEngine) entirely.
Request:
Could you please re-run the quantization export pipeline or provide an updated set of weights that corrects this dimension layout translation for the audio_tower weights? Having the VLM identity fully intact will allow us to restore both its native multimodal understanding and unlock full VLM Speculative MTP acceleration.
Thanks again for your awesome work on local Apple Silicon inference tooling!
Thanks for the precise report. Your analysis was correct.
Confirmed and fixed. The sidecar was storing audio_tower conv weights in PyTorch channel-first layout instead of MLX channel-last. OptiQ's builder copied the towers raw from the bf16 base without mlx-vlm's conv sanitize(), so the audio convs shipped in the wrong layout. It went unnoticed because OptiQ's own path is vision and text only and never loads the audio tower.
This model's optiq_vision.safetensors has been re-uploaded with the corrected layout, so just re-pull and the audio tower will load. The sibling gemma-4-e2b-it-qat-OptiQ-4bit had the same issue and is fixed too. The build-time fix is in mlx-optiq 0.2.9.
Thanks again. Let us know if VLM plus speculative MTP loads cleanly on your side now.
Thanks for the fast turnaround and the precise root-cause breakdown!
The PyTorch NCHW vs MLX NHWC channel mismatch perfectly explains the layout rejection we saw during the audio_tower initialization. Skipping conv_sanitize() during the raw tensor migration is a classic slip that's easy to miss on vision-only runs.
I can confirm that after pulling the re-uploaded optiq_vision.safetensors and updating to mlx-optiq 0.2.9, the complete VLM pipeline plus speculative MTP loads flawlessly on oMLX now. No more layout errors, and no engine fallback.
Telemetry on a 16GB UMA Mac mini (M4) looks incredible over the vlm_mtp path:
100% Token Acceptance: The 0.16B assistant hit a 100% acceptance rate (tokens_per_round=4.00), pushing token generation speeds up to 48.8 - 60.4+ tok/s.
Tight Memory Footprint: Absolute peak memory ceiling capped at 8.19 GB, leaving plenty of headroom on a 16GB machine.
This build-time fix completely unblocks the Gemma-4 architecture for ultra-fast, local multimodal pipelines on everyday hardware.
Outstanding work on this fix!