Views
No views yet
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.| Layer group | NVIDIA (modelopt) | This checkpoint (MLX) |
|---|---|---|
MoE experts experts.{gate,up,down}_proj (128×/layer) | NVFP4 (E2M1, 2-level, gs16) | NVFP4 — 4-bit, group size 16 |
self_attn q/k/v/o + q_norm/k_norm | bf16 (excluded) | bf16 |
router.proj / router.scale / per_expert_scale | bf16 (excluded) | bf16 |
dense mlp.{gate,up,down}_proj (parallel to MoE) | bf16 (excluded) | bf16 |
embed_tokens (tied = lm_head) | bf16 (excluded) | bf16 |
vision_tower.* / embed_vision.* | bf16 (excluded) | bf16 |
all layernorms, final norm | bf16 | bf16 |
exclude_modules.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.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).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.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.*).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).1import { loadSession } from '@mlx-node/lm';
2
3const session = await loadSession('./Gemma-4-26B-A4B-NVFP4-mlx');
4for await (const ev of session.sendStream('Explain what a mixture-of-experts layer is.')) {
5 if (!ev.done) process.stdout.write(ev.text);
6}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).