README.md

---
license: gemma
base_model: google/gemma-4-12B
library_name: flax
pipeline_tag: image-text-to-text
tags:
  - jax
  - flax
  - nnx
  - gemma
  - gemma4
  - port
  - multimodal
  - encoder-free
---

# gemma4-jax — Flax NNX port of `google/gemma-4-12B` (full multimodal)

A from-scratch, faithful **Flax NNX** implementation of Google's Gemma 4 12B
"Unified" model — the **text decoder** *and* the encoder-free **vision + audio**
embedders — plus a HuggingFace `safetensors → NNX` weight converter.

- 💻 **Code / GitHub:** https://github.com/mlnomadpy/gemma4-jax
- 🧬 **Base model:** [`google/gemma-4-12B`](https://huggingface.co/google/gemma-4-12B)

This repository hosts the **port (code)**. It does **not** redistribute the
weights — at load time it reads the official `google/gemma-4-12B` safetensors
directly (the tensor names match; `nnx.Linear` kernels are just transposed). You
must accept the Gemma terms on the base-model page to download the weights.

## Why

Gemma 4 12B is not (yet) in the official `google-deepmind/gemma` JAX library.
This port makes the architecture an open, editable NNX module — so the global
attention layers, the encoder-free projectors, and the RoPE/norm details are a
clean swap point for research (e.g. linear-attention surgery on the 8 global
layers).

## What's implemented — verified exact to the parameter

| Component | State |
|---|---|
| Text decoder (48 layers, dual sliding/global attention) | ✅ |
| Dual attention: sliding GQA (hd 256, 8 KV) / global MQA (hd 512, 1 KV) | ✅ |
| `attention_k_eq_v` (V reuses pre-norm K on full layers, no RoPE on V) | ✅ |
| Per-head QK-norm, sandwich norm, `layer_scalar`, embed scaling, logit softcap | ✅ |
| Proportional / partial RoPE (zeroed-tail inv_freq) + default RoPE | ✅ |
| **Vision** (encoder-free: raw 48×48×3 patches → LN→Dense→LN→+2D-posemb→norm→proj) | ✅ |
| **Audio** (encoder-free: raw 640-sample frames → RMSNorm→proj) | ✅ |
| **Multimodal splice** (soft-token scatter + bidirectional-vision mask) | ✅ |
| `safetensors → NNX` converter (text + vision + audio) | ✅ |
| KV cache for fast decode | ❌ reference loop recomputes prefix |

**Param count matches the published checkpoint exactly:** text
`11,907,350,320` + multimodal `52,379,904` = **`11,959,730,224`** (0 diff).
The multimodal converter is verified against the real safetensors; text and
multimodal smoke tests pass (forward, causality drift 0, softcap bounds, splice).

## Non-obvious details baked in (vs. Gemma 2/3)

- RMSNorm is **plain `x·w`** (not `(1+w)`); eps *inside* the rsqrt; fp32 internals.
- Attention **`scaling = 1.0`** — magnitude set by per-head `q_norm`, not `1/√d`.
- **`k_eq_v`**: global layers have no `v_proj`; V reuses the K projection output
  (pre-norm, pre-RoPE) + a scale-free `v_norm`, and V is **not** rotated.
- **Proportional RoPE** on global layers: full-length inv_freq with only the
  first 64 frequencies nonzero (NoPE tail), base 1e6.
- Vision is **encoder-free**: no SigLIP; raw merged pixel patches project
  straight into the 3840-d decoder space. Audio likewise — no mel/conformer.

## Usage

```bash
pip install jax flax safetensors huggingface_hub tokenizers
git clone https://github.com/mlnomadpy/gemma4-jax && cd gemma4-jax
```

```python
import jax.numpy as jnp
from gemma4_jax.convert import unified_from_safetensors
from gemma4_jax.config import IMAGE_TOKEN_ID

# point at the official google/gemma-4-12B model.safetensors (accept terms first)
uni = unified_from_safetensors("path/to/model.safetensors")

# text
logits = uni.logits(input_ids)                       # [B, S, vocab]

# vision: pixel_values [B,P,6912], image_position_ids [B,P,2]
soft = uni.get_image_features(pixel_values, image_position_ids)   # [B,P,3840]
h = uni(input_ids, pixel_values=pixel_values, image_position_ids=image_position_ids)

# audio: input_features [B,T,640]
h = uni(input_ids, input_features=input_features)
```

The HF image processor (patchify + 3×3 pool + position ids) and audio feature
extractor are not ported — feed pre-patchified `pixel_values` / pre-framed
`input_features` exactly as the HF processors emit them.

## License

This **code** is a clean-room reimplementation. The **weights** it loads are
Google's Gemma 4, governed by the [Gemma Terms of Use](https://ai.google.dev/gemma/terms);
your use of the weights is subject to those terms and the Gemma Prohibited Use
Policy. See the [base model](https://huggingface.co/google/gemma-4-12B).