Hugging Face's logo

google
/
gemma-4-12B-it

Any-to-Any
Transformers
Safetensors
gemma4_unified
image-text-to-text
Eval Results
Model card Files
xet
 Community
51

MLX 4-bit conversion of gemma-4-12B-it — with vision + audio embedder weights preserved

#27
by jokernifty - opened
Discussion
jokernifty
•
edited Jun 7

Hi all,

Shipped what I believe is the first MLX conversion of the 12B Unified model that keeps the encoder-free vision_embedder and audio_embedder weights intact (not just the text backbone):

https://huggingface.co/jokernifty/gemma-4-12B-it-mlx-4bit-multimodal

What works today

  • Text inference on Apple Silicon (via mlx_lm with a 4-line gemma4_unified.py alias — README has the snippet)
  • VisionEmbedder + AudioEmbedder ported to MLX, weights load cleanly, forward passes produce sensible outputs
  • Token-splicing forward pass (image_token_id / audio_token_id / video_token_id) implemented and tested with synthetic inputs

What's WIP

  • 🚧 Image preprocessor (resize → bucket to 70/140/280/560/1120 tokens → 48×48 patches → factorized 2D position IDs)
  • 🚧 Audio preprocessor (16 kHz waveform → 640-sample frames)
  • 🚧 Chat-template integration for <boi>…<eoi> / <boa>…<eoa> blocks

The companion package with the embedders + converter is at https://github.com/jokernifty/mlx-gemma4-unified (currently scaffolded — preprocessors next).

A few notes on what was tricky

  • mlx-lm doesn't yet register the gemma4_unified model_type — solved with a small alias that subclasses gemma4.Model and adds the encoder-free weight names to the sanitize skip list.
  • Standard mlx_lm.convert hits the macOS Metal 5-second command-buffer watchdog because the final eval+save of all quantized weights happens in one giant graph. Worked around with a layer-by-layer converter that calls mx.eval() after each transformer block (~1.2 s each on M-series).
  • The 262k × 3840 token embedding can't be quantized in one shot for the same watchdog reason; kept in bf16. Adds ~2 GB to the file size but standard practice for many mlx-community checkpoints.
  • MLX doesn't yet support boolean indexing, so the masked_scatter equivalent uses a numpy-host helper for the integer indices.

Would love a sanity check from anyone who's been working on gemma4_unified support in mlx-lm / mlx-vlm upstream — happy to coordinate or PR the alias module if it's useful. And of course feedback / issues / PRs on the embedder ports are very welcome.

Massive thanks to Google DeepMind for the open release

divinetribe

nice work on keeping the embedders, that is the harder road and the one worth taking.

one thing that might save you the alias entirely. upstream already registered gemma4_unified on main in #1349, merged 2026-06-05, the day before you posted. it remaps gemma4_unified straight to the gemma4 loader and drops the vision and audio names in sanitize(). so your alias and that remap are doing the same job.

the catch is it has not shipped. pypi is still on 0.31.3 from 2026-04-22 and i checked the released tarball, the remap is not in it, so pip install mlx-lm still stops with Model type gemma4_unified not supported. only main has it. that is probably why it looked unregistered from your side.

i came at it from the other end and just stripped the towers, relabeled model_type to gemma4 and left the 4bit weights untouched, which loads on the stock pypi release with nothing patched. i diffed the parameter names against the shard index to be sure, 1324 expected and 1324 shipped, nothing missing or extra. that only helps text only people though, yours is the one that actually keeps the modalities.

happy to sanity check the embedder ports if you want another pair of eyes. i have an mlx-lm port open right now for a different arch so i am in that code most days.

Sign up or log in to comment