Gemma 4: sanitize() duplicates 'model.' prefix, all weights load as zero

[raullenchai]

Bug

Gemma 4 models (E2B, E4B, 31B) load with all-zero weights on mlx-vlm 0.4.3. The model produces only <pad> tokens.

Root Cause

In mlx_vlm/models/gemma4/gemma4.py, the sanitize() method has:

if new_key.startswith("language_model."):
    rest = new_key[len("language_model."):]
    new_key = "language_model.model." + rest

But the safetensors weights already have the full prefix language_model.model.:

language_model.model.embed_tokens.weight
language_model.model.layers.0.self_attn.q_proj.weight
...

So sanitize transforms language_model.model.embed_tokens.weight → language_model.model.model.embed_tokens.weight, which doesn't match any model parameter. The weights silently fail to load and everything is zero.

Fix

if new_key.startswith("language_model.model."):
    pass  # already correct
elif new_key.startswith("language_model."):
    rest = new_key[len("language_model."):]
    new_key = "language_model.model." + rest

Reproduction

from mlx_vlm import load, generate

model, processor = load("mlx-community/gemma-4-e4b-it-8bit")
prompt = processor.apply_chat_template(
    [{"role": "user", "content": "Hello"}],
    tokenize=False, add_generation_prompt=True,
)
output = generate(model, processor, prompt=prompt, max_tokens=10, verbose=False)
print(output.text)  # All <pad> tokens

Environment

• mlx-vlm 0.4.3
• mlx-lm 0.31.1
• macOS 15.5, M3 Ultra
• Models tested: gemma-4-e4b-it-4bit, gemma-4-e4b-it-8bit (both from mlx-community)

[raullenchai]

Additional finding: the issue is deeper than just the sanitize prefix.

Even after patching sanitize to fix the model.model. duplication, all weights are uint32 dtype with value 0. The 8-bit quantized weights (stored as packed integers + scales + biases in safetensors) are not being loaded into QuantizedLinear layers — they end up as raw uint32 tensors with no dequantization.

# After loading gemma-4-e4b-it-8bit:
lm.model.embed_tokens.weight.dtype  # mlx.core.uint32
lm.model.embed_tokens.weight[0,:5]  # [0, 0, 0, 0, 0]
lm.model.layers[0].self_attn.q_proj.weight.dtype  # mlx.core.uint32

This suggests the weight-to-module mapping in the model class doesn't match the safetensors key structure for quantized Gemma 4 models. The model creates standard nn.Linear layers but the weights have weight/scales/biases triplets that need nn.QuantizedLinear.

[raullenchai]

Update: Found a second, independent issue.

After patching the sanitize prefix bug and force-reloading weights, the linear layer scales load correctly (non-zero). However, embed_tokens.scales is all zeros in the safetensors file itself for mlx-community/gemma-4-e4b-it-8bit:

raw = mx.load("model-00001-of-00002.safetensors")
raw["language_model.model.embed_tokens.scales"]  # shape=(262144, 40), ALL ZEROS
raw["language_model.model.layers.0.mlp.down_proj.scales"]  # shape=(2560, 160), non-zero ✓

So there are two bugs:

1. sanitize() prefix duplication (code bug, fixable) — causes all weights to silently mismatch for non-MLX-format models
2. MLX-format models skip sanitize (code bug, fixable) — weight keys don't match model params even though keys look correct
3. Embedding quantization scales are zero in mlx-community quantized models (model/quantization bug) — the quantizer may not have handled the embedding layer correctly

Bug 1+2 are in mlx-vlm code. Bug 3 is in the mlx-community quantization pipeline.

[Blaizzy]

Hey

Thanks for reaching out!

Do you mind trying the main branch and seen if the issue persists