Gemma-4-26B-A4B-NVFP4-mlx

Gemma-4-26B-A4B-NVFP4-mlx

An MLX transcode of NVIDIA's nvidia/gemma-4-26b-a4b-nvfp4 (a TensorRT Model Optimizer NVFP4 PTQ of Google's google/gemma-4-26b-a4b-it), repacked so it loads and runs on Apple Silicon via mlx-node (and other MLX runtimes that read the quantization block).

gemma-4-26b-a4b is a sparse Mixture-of-Experts model — 30 layers, 128 experts per layer, ~4B active parameters (A4B) out of ~26B total. NVIDIA's checkpoint ships the experts in a Blackwell FP4 layout (packed-U8 NVFP4 + weight_scale / weight_scale_2 / input_scale sidecars), which is not loadable by MLX runtimes as-is. This repo preserves NVIDIA's exact per-tensor bit allocation while re-expressing each quantized expert in an MLX-native NVFP4 the Metal kernels understand.

Quantization recipe

NVIDIA's NVFP4 recipe here is experts-only: only the MoE expert projections are 4-bit; everything else — attention, router, the per-layer dense MLP, all norms, the tied embedding/head, and the vision tower — stays bf16. This checkpoint reproduces that allocation exactly.

Notes:

• Experts-only NVFP4 mirrors NVIDIA's allocation. In an A4B MoE the experts hold the vast majority of the parameters, so quantizing just the experts to 4-bit captures almost all of the size win (~49 GB bf16 → ~19 GB) while leaving the small, precision-sensitive attention / router / dense-MLP paths in bf16 — exactly the modules NVIDIA lists under exclude_modules.
• The head is the bf16 tied embedding. tie_word_embeddings: true, and NVIDIA excludes lm_head, so logits are produced by the untied-nothing tied embedding matmul in bf16 — no separate (and no quantized) head is stored.
• NVFP4 experts are a faithful re-quantization of NVIDIA's. Each expert was dequantized from NVIDIA's NVFP4 (E2M1 × per-16 E4M3 block-scale × per-tensor F32) to bf16, then re-quantized to MLX NVFP4 (gs16). The dequantized experts match the bf16 -it reference at cosine 0.9955; loading a micro-scaled (biasless) NVFP4 expert needs an MLX runtime with a mode-aware batched-expert loader — on mlx-node that is the switch_glu NVFP4 path (gather_qmm mode nvfp4).
• Vision weights are carried through in bf16 but the multimodal path is not verified in this repo; this checkpoint is validated for text generation.

On-disk size is ~19 GB (bf16 reference is ~49 GB).

How it was made

1. Dequantize NVIDIA's NVFP4 experts → bf16, fuse gate‖up and stack the 128 experts into the reference layout experts.gate_up_proj [128, 1408, 2816] + down_proj [128, 2816, 704]; carry every non-expert tensor through unchanged (they are already bf16 in NVIDIA's checkpoint). Validated vs google/gemma-4-26b-a4b-it: cosine 0.9955.
2. mlx convert (bf16, no quantization) to sanitize into MLX layout (key remap model.language_model.* → language_model.model.*, split the fused experts into switch_glu.{gate,up,down}_proj, drop the tied lm_head, keep vision_tower.*).
3. Per-expert mx.quantize(..., mode="nvfp4", group_size=16, bits=4) on the switch_glu experts only; write quantization + quantization_config blocks (base {nvfp4,4,16} + per-expert overrides). All other tensors stay bf16 (no .scales → loaded dense).

Usage (mlx-node)

import { loadSession } from '@mlx-node/lm';

const session = await loadSession('./Gemma-4-26B-A4B-NVFP4-mlx');
for await (const ev of session.sendStream('Explain what a mixture-of-experts layer is.')) {
  if (!ev.done) process.stdout.write(ev.text);
}

Caveats

• Validated for text generation on mlx-node (loads + generates coherent output). Vision weights are carried in bf16 but the multimodal path is not verified here.
• Sparse MoE (128 experts, ~4B active): decode is expert-memory-bandwidth-bound. Performance is comparable to other ~26B-A4B MLX quants of this model; treat any numbers as indicative and thermally sensitive.

License & attribution

Governed by the Gemma Terms of Use, inherited from google/gemma-4-26b-a4b-it via nvidia/gemma-4-26b-a4b-nvfp4. This repository only re-expresses the weights in an MLX-loadable quantization; all model capability and training credit belongs to Google (and the NVFP4 PTQ to NVIDIA).