Gemma4 Pytorch Claude

gemma4-pt-claude

A clean, standalone PyTorch implementation of Google's **Gemma 4** model family — ported layer-by-layer from the official JAX reference with numerical verification at every step.

Text, vision, and audio. No mystery abstractions. Just PyTorch and matrix multiplies.

Supported Models

Variant	Layers	Embed	Heads	KV Heads	Attention	Notable Features
E2B	35	1536	8	1	4:1 local:global	PLI, KV sharing, vision, audio
E4B	42	2560	8	2	5:1 local:global	PLI, KV sharing, vision, audio
12B	48	3840	16	8 (1 global)	5:1 local:global	Encoder-free vision + audio, K=V global
31B	60	5376	32	16 (4 global)	5:1 local:global	K=V global, bidirectional vision
26B-A4B	30	2816	16	8 (2 global)	5:1 local:global	MoE (128 experts, top-8), K=V global, vision

All five use QK-norm and V-norm.

Encoder-free variants (12B)

The 12B (HF `gemma4_unified`) drops the multimodal towers entirely — there is no SigLIP-style vision encoder and no conformer:

• Vision: raw 48×48 pixel patches (3×3 groups of 16px patches merged) → LayerNorm → Linear → LayerNorm → factorised 2-D position embedding → LayerNorm → RMSNorm → Linear. The spatial pooling the tower models apply after their encoder happens here instead, before any projection.
• Audio: raw 16 kHz waveform chunked into 640-sample frames (40 ms), each frame projected by RMSNorm → Linear. No mel spectrogram.

The text decoder is the same `TransformerBlock` as every other variant. Config types are `EncoderFreeVisionConfig` / `EncoderFreeAudioConfig`; `Composer` and `preprocess_images` handle both families automatically, and the encoder-free audio path passes `audio_frames` where the tower path passes `audio_mel`.

Installation

pip install -e .

Core dependencies are `torch`, `sentencepiece`, `safetensor