All three preserve full BF16 precision in attention, vision tower, embeddings, and norm layers. The 4-bit version matches the 8-bit on MBPP (86.5% both) and trails BF16 by only 1.2% on HumaneEval — well within benchmark variance.

Source model: https://huggingface.co/google/gemma-4-31B-it

(1) Introduction

The Objective: Get more throughput out of Gemma 4-31B on Apple Silicon without sacrificing the reasoning and multimodal capabilities that make it useful for agentic work and coding.

The Problem: BF16 gives you the full brain but demands ~58 GB of unified memory, and output is slow. Uniform quantizations are faster but degrade the hybrid attention mechanism and the deeply nested vision tower.

The Solution: Mixed Quantization. A surgical approach that compresses only the SwiGLU MLP layers — which make up 68% of the model's mass — to 8-bit, 6-bit, or 4-bit, while leaving every attention projection, the entire vision tower, embeddings, and norm layers at full BF16 precision. The resulting models load natively through mlx_vlm with zero modifications.

The result: 60% speedup at 4-bit with nearly identical benchmark scores, or 33–48% speedup at 8/6-bit if you want to be conservative.

Resulting sizes:

| Version | Size | Speed Gain | Reduction |
|---|---|---|---|
| BF16 baseline | ~58 GB | — | — |
| Mixed 8-bit | 40 GB | 33% | 31% |
| Mixed 6-bit | 35 GB | 48% | 40% |
| Mixed 4-bit | 30 GB | 60% | 48% |

(2) Architecture

Gemma 4-31B is a dense hybrid-attention model. It has 60 transformer layers following a repeating pattern:

[sliding_attention, sliding_attention, sliding_attention, sliding_attention, sliding_attention, full_attention] × 10 = 60 layers

The sliding attention layers use a 1024-token window with standard RoPE (θ=10,000) for local context. Every sixth layer uses full softmax attention with proportional RoPE (θ=1,000,000) for global reasoning across the 262K token window.

Additionally, the model includes:

A 27-layer vision tower for image/video comprehension
A multimodal embedding projector (embed_vision) that maps vision features into language model space
GeGLU MLPs in every layer for feature transformation

Critical Architectural Detail: Nested Vision Tower

Gemma 4's vision tower is interleaved inside every language model layer in the safetensors key structure. This means vision_tower.encoder.layers.N.* tensors appear alongside language_model.layers.N.* tensors in the same shard.

Identifying the Weight

Out of 31B parameters, the MLP projections (gate_proj, up_proj, down_proj) make up 68% of the model's mass — 39.5 GB of the 58.25 GB BF16 footprint. These GeGLU layers are the "knowledge lookup" / feature transformation matrices and are extremely tolerant of quantization.

The Glass Cannons: Layers Protected at BF16

Layer	Why Protected
Full attention (q_proj, k_proj, v_proj, o_proj, q_norm, k_norm)	Deep reasoning and cross-token attention. 16 layers of full softmax attention are the model's primary reasoning engine.
Sliding attention (q_proj, k_proj, v_proj, o_proj, q_norm, k_norm)	Local context tracking across 45 layers. Quantizing these degrades the model's ability to maintain coherent state.
Vision tower (all 27 encoder layers, nested in every LM layer)	Image and video comprehension. The nested structure makes this especially fragile — a single misclassified layer corrupts all vision understanding.
Embed vision projector	Bridges vision features into language model space. Quantizing this collapses the multimodal interface.
Embed tokens	Input embeddings. Quantizing these adds noise at the model's primary interface.
Norm layers (input_layernorm, post_attention_layernorm, pre_feedforward_layernorm, post_feedforward_layernorm)	Stability anchors. Gemma 4 uses four norm layers per layer — more than most architectures.
Layer scalars	Gemma's layer scaling mechanism. Critical for training stability.

The Crush Zone

Layer	Quantization	Rationale
MLP (gate_proj, up_proj, down_proj)	8-bit, 6-bit, or 4-bit	68% of model mass. GeGLU MLPs are the most quantization-tolerant component in a transformer. All the savings live here.

(3) The Quantization Script

The script uses MLX's native nn.quantize() with a class_predicate filter that targets only the three GeGLU projection matrices inside language_model.layers.N.mlp. Everything else passes through untouched.

python

import mlx.core as mx
import mlx.nn as nn
from mlx_vlm.utils import load
import json, os, gc, argparse

parser = argparse.ArgumentParser()
parser.add_argument("--bits", type=int, default=8, choices=[8, 6, 4])
parser.add_argument("--group-size", type=int, default=64)
args = parser.parse_args()

model_path = "/path/to/gemma-4-31B-it"
output_path = f"/path/to/gemma-4-31B-it-Mixed-{args.bits}bit"

# Lazy load — no eager evaluation
model, processor = load(model_path, lazy=True)

def gemma4_mlp_only_filter(path, module=None):
    """Target ONLY language model MLP projections. Protect everything else."""
    # Hard shield: vision tower (nested inside every layer)
    if "vision_tower" in path:
        return False
    # Hard shield: vision embedding projector
    if "embed_vision" in path:
        return False
    # Hard shield: embeddings
    if "embed_tokens" in path:
        return False
    # Hard shield: norms and scalars
    if "layernorm" in path or "layer_scalar" in path or path.endswith(".norm"):
        return False
    # Hard shield: attention
    if "self_attn" in path:
        return False
    # Target: MLP inside language_model only
    if "language_model" not in path or ".mlp." not in path:
        return False
    leaf = path.split('.')[-1]
    return leaf in ("gate_proj", "up_proj", "down_proj")

# Single-pass quantization — only MLP gets touched
nn.quantize(model, group_size=args.group_size, bits=args.bits,
            class_predicate=gemma4_mlp_only_filter)

# Save config and processor
os.makedirs(output_path, exist_ok=True)
if hasattr(processor, "save_pretrained"):
    processor.save_pretrained(output_path)

with open(os.path.join(model_path, "config.json")) as f:
    config = json.load(f)
config["quantization"] = {
    "group_size": args.group_size,
    "bits": args.bits,
    "type": "mixed",
    "description": f"Gemma 4 MLP-only {args.bits}-bit. Attention, vision, embeddings, norms remain BF16."
}
with open(os.path.join(output_path, "config.json"), "w") as f:
    json.dump(config, f, indent=4)

# Flatten, destroy model tree, shard to SSD
flat_weights = dict(flatten_parameters(model.parameters()))
del model, processor
gc.collect()
mx.clear_cache()

# Sequential eval + shard (5 GB chunks)
current_shard, shard_size, shard_idx = {}, 0, 1
for name, tensor in flat_weights.items():
    if not isinstance(tensor, mx.array):
        continue
    mx.eval(tensor)
    current_shard[name] = tensor
    shard_size += tensor.nbytes
    if shard_size >= 5 * 1024**3:
        mx.save_safetensors(
            f"{output_path}/model-{shard_idx:05d}.safetensors", current_shard)
        current_shard, shard_size, shard_idx = {}, 0, shard_idx + 1
        gc.collect(); mx.clear_cache()
if current_shard:
    mx.save_safetensors(
        f"{output_path}/model-{shard_idx:05d}.safetensors", current_shard)

Full script: mixed_quantize_gemma4.py in the quantization-scripts repository.

(4) Deployment

Native Loading — No Hacks Required

These models load natively through mlx_vlm without any modifications to sanitize_weights or other framework internals. The surgical filter ensures that only language model MLP layers are quantized, and the output format follows MLX's native safetensors conventions exactly.

python

from mlx_vlm.utils import load

model, processor = load("your-username/gemma-4-31B-it-Mixed-8bit", lazy=True)

Note on config.json

The config.json declares a single "bits" value, but MLX reads per-layer quantization metadata dynamically from the safetensors tensor headers. The actual mixed-precision structure — which layers are quantized and which remain BF16 — lives in the safetensors files, not the config. Tooling that trusts only config.json will be confused; this is expected behavior for native MLX mixed quantization.

(5) Real-World Performance on Apple Silicon

Hardware: M5 Max, 128 GB Unified Memory (BF16, 8-bit, and 6-bit benchmarks). M1 Max, 64 GB Unified Memory (4-bit benchmark).

Baseline: gemma-4-31B-it BF16 (~58 GB)

Coding Benchmarks

Model: gemma-4-31B-it (BF16)
Benchmark         Accuracy   Correct   Total   Time(s)
------------------------------------------------------
HUMANEVAL            95.1%       156     164    1483.2
MBPP                 86.5%       173     200    2043.5
LIVECODEBENCH        74.0%        74     100    5799.9


Model: gemma-4-31B-it-Mixed-8bit
Benchmark         Accuracy   Correct   Total   Time(s)
------------------------------------------------------
HUMANEVAL            94.5%       155     164    1103.7
MBPP                 86.5%       173     200    1503.3
LIVECODEBENCH        73.0%        73     100    4076.3


Model: gemma-4-31B-it-Mixed-6bit
Benchmark         Accuracy   Correct   Total   Time(s)
------------------------------------------------------
HUMANEVAL            93.9%       154     164    1275.1
MBPP                 86.5%       173     200    1382.0
LIVECODEBENCH        72.0%        72     100    5259.5


Model: gemma-4-31B-it-Mixed-4bit  (M1 Max, 64 GB)
Benchmark         Accuracy   Correct   Total   Time(s)
------------------------------------------------------
HUMANEVAL            93.9%       154     164    3423.6
MBPP                 86.5%       173     200    3431.3
LIVECODEBENCH        70.0%        70     100    9461.1

(6) Conclusion

Mixed MLX quantization of Gemma 4-31B that compresses only the MLP layers — while preserving all attention, vision, and embedding layers at BF16 — delivers up to a 48% size reduction with no statistically significant intelligence loss.

The 8-bit version is the safe choice for production agentic work: at 40 GB, it matches BF16 on MBPP (86.5% both) and trails by only 0.6% on HumaneEval — well within benchmark variance. It completes LiveCodeBench 42% faster than BF16.

The 4-bit version is the sweet spot for memory-constrained machines: at 30 GB, it's deployable on an M1 Max with 64 GB of RAM and still matches the 8-bit on MBPP (86.5% both) and trails BF16 by only 1.2% on HumaneEval. The 4-bit version completes LiveCodeBench at 70.0% — a 4-point drop from BF16, but still strong absolute performance.

The hybrid attention architecture (sliding window + full softmax) is the key: by protecting both attention types in BF16, the model retains its full 262K context-window recall and reasoning capability. The MLP layers — which are just feature transformation matrices — absorb even 4-bit quantization with minimal penalty on coding tasks.

Notes

Fixed seed sampling was used on all benchmarks.
MBPP and LiveCodeBench were run on subsamples (200/500 and 100/1055 respectively).
The 4-bit benchmark was run on an M1 Max (64 GB), which explains the slower absolute times. Accuracy comparisons remain valid.
Test times include multitasking overhead and are not strictly comparable across hardware.
Vision capabilities were preserved but not formally benchmarked in this pass.

andrzejmontano/gemma-4-31B-it-Mixed-8bit

Gemma-4-31B-it-Mixed-Bit-MLX