Gemma 4 is not loading

[luiz290788]

I'm trying to run gemma 4 with the latest release (b153) and I'm getting the following error:

[rednaxela5950]

I'm getting the same error.

[abbasmousavi]

Same error here when trying to run gemma-4-26b-a4b-it-8bit

Error: Mismatched parameter model.embed_tokens.weight in Gemma4Model.Gemma4ModelInternal.QuantizedEmbedding shape. Actual [262144, 704], expected [262144, 352]

[rednaxela5950]

I tried FakeRockert543/gemma-4-26b-a4b-it-MLX-8bit instead and got it to start.

[solderzzc]

Should be working now.
I just merged the code to main.

[michaelstingl]

Edit (2026-04-06): cleaner fix — router stays at 8-bit, fix now fully contained in Load.swift. Root cause below unchanged.

Reproduced on b153, b160, b188 with mlx-community/gemma-4-26b-a4b-it-4bit @ 8bcfa0d (2026-04-03).

Note: the HF repo was re-uploaded on 2026-04-03 (b86b3e2 → 8bcfa0d): weights replaced in all 3 shards, config.json expanded from 157 to ~972 lines adding the per-layer quantization section. Code tested against the earlier SHA would not have hit this path.

The -it-4bit variant uses mixed precision (MLP + router at 8-bit, experts at 4-bit) to fit Gemma 4 26B MoE in ~15 GB — the only variant small enough for 32 GB Macs. @abbasmousavi's comment above reports a similar shape mismatch on gemma-4-26b-a4b-it-8bit (embed_tokens [262144, 704] vs [262144, 352]), which looks like the same prefix-mismatch class of bug.

Root cause

Two issues stacked:

1. config.json has per-layer overrides for MLP (gate_proj/down_proj/up_proj) and router.proj at 8-bit under keys prefixed with language_model..
2. Gemma4.sanitize() strips the language_model. prefix from weights, so leafModules() paths (model.layers.0.mlp.gate_proj) never match perLayerQuantization keys (language_model.model.layers.0.mlp.gate_proj). PerLayerQuantization.quantization(layer:) falls through to the global 4-bit default → QuantizedLinear created with wrong bits → shape mismatch on load.

The existing moduleUpdates workaround in sanitize() doesn't take effect because self.update(modules: ModuleChildren.unflattened(...)) throws UpdateError.unexpectedStructure(key: "layers", ...) on Gemma 4's nested module tree, so the router mismatch surfaces first.

Fix

Libraries/MLXLMCommon/Load.swift — extend the per-layer lookup with a fallback chain (prefix + .experts. stripping for Gemma 4's router remap):

if let perLayerQuantization {
    let dict = perLayerQuantization.perLayerQuantization
    var candidates = [path, "language_model.\(path)"]
    if path.contains(".experts.") {
        let flat = path.replacingOccurrences(of: ".experts.", with: ".")
        candidates.append(contentsOf: [flat, "language_model.\(flat)"])
    }
    for key in candidates {
        if let opt = dict[key] {
            switch opt {
            case .skip: return nil
            case .quantize(let q): return q.asTuple
            }
        }
    }
    return perLayerQuantization.quantization?.asTuple
}

Libraries/MLXLLM/Models/Gemma4.swift — delete the moduleUpdates block in sanitize(). With the fallback above, the framework's standard quantize() resolves the correct per-layer bits directly from config.json, so the manual replacement is no longer needed (and never worked anyway due to the update(modules:) crash).

Result

Gemma 4 loads and generates. ~16 tok/s, 20.9 GB on M5 32GB. Qwen2.5-3B still loads (no regression).

---

Analysis and patch drafted with Claude Code assistance.

[luiz290788]

Yeah, I confirmed that the issue is still happening in b188

[solderzzc]

I'm uploading code for a fix, running HomeSec Benchmark for accuracy evaluation. Currently it's testing gemma-4-26b-a4b-it-8bit.

It's triggered with ./run_benchmark.sh select 3 for real workload tests.

[solderzzc]

Thanks for the detailed analysis and for sharing this patch, @michaelstingl! Your write-up was spot on in identifying the root cause of the shape mismatch.

We actually ran into the exact same issue trying to get the Gemma 4 MoE models running over the last couple of days! We landed on a very similar fix which was reviewed and merged into our main branch today (April 6th). Here is how we tackled it on our end:

Like you noted, the HF -4bit and -8bit variants use language_model.-prefixed keys in config.json, which breaks Load.swift's per-layer quantization.

Our Fix:

Load.swift: We added the language_model. fallback chain to properly capture the MLP 8-bit overrides (gate_proj, down_proj, up_proj).
Gemma4.swift (sanitize): Instead of expanding the Load.swift fallback to strip .experts. for the router remap, we implemented an alternative fix inside sanitize(). We manually dequantize the router weights back to dense float32 natively. The router is tiny (~0.1% of params), so the memory impact is negligible. Because the .scales and .biases are stripped during this dequantization step, Load.swift's quantize() pass naturally skips the router entirely, side-stepping the .experts. path mismatch completely.

Results: Ran both the 4-bit and 8-bit variants through our 96-test HomeSec benchmark suite to ensure stability across edge contexts (M5 32GB). https://github.com/SharpAI/DeepCamera/tree/master/skills/analysis/home-security-benchmark , this is being used with ./run_benchmark.sh Option 3.
Both variants load and generate flawlessly with no shape mismatch crashes, and importantly, no regressions in existing 4-bit models:

4-bit (~15.5 GB Memory Demand): 87/96 passed, ~7.0 tok/s decode
8-bit (~20.9 GB Memory Demand): 90/96 passed, ~6.4 tok/s decode

Thanks again for digging so deeply into this and sharing your findings. We really appreciate the community collaboration and your detailed documentation of the config.json mismatch!

[solderzzc]

Let me close this issue since the fix was committed. Reopen otherwise.

[luiz290788]

It loads the model and it works but it struggles when I try to use it with opencode.