kaaaaan/t5gemma-sa3-int8-mlx

T5Gemma (Stable Audio 3) — int8 MLX quantization

An int8-quantized copy of the T5Gemma text encoder shipped in stabilityai/stable-audio-3-small-music, packaged for Apple MLX.

• 299 MB on disk — 1.9× smaller than the 563 MB fp16 encoder.
• Near-lossless: per-token cosine to the fp16 encoder = 0.99923 (mean), measured on real prompts.
• Audio-verified: feeding this encoder into the frozen Stable Audio 3 DiT gives prompt-adherence parity with fp16 by CLAP text-audio similarity (Δ = −0.0003 over 12 prompts × 2 seeds — see below).
• No retraining. Pure post-training quantization; identical architecture, identical tokenizer.

⚠️ You bring your own accepted licenses. These are quantized Stability/Google weights. See License below and the bundled NOTICE — You are responsible for your own compliance with the Stability AI Community License and the Gemma Terms of Use.

Also available: a framework-neutral (PyTorch/ONNX/CPU) build of the same encoder — kaaaaan/t5gemma-sa3-int8.

Why

70% of the encoder's 563 MB is the 256k-vocab embedding table; the transformer layers are only 30%. That table quantizes very cleanly (it's a lookup, not a compute path), so int8 recovers most of the size with negligible fidelity loss — far better than trying to distill a smaller foreign encoder (a different tokenizer breaks per-token alignment and caps quality well below usable).

Quantization scheme

mx.quantize, groupwise affine, int8, group size 64. All 2-D weight matrices (the embedding table + every linear projection) are quantized; 1-D params (RMSNorm weights, RoPE inverse frequencies) stay fp16.

scheme	size	mean cos vs fp16	min cos
fp16 (source)	563 MB	1.00000	1.00000
int8 g64 (this repo)	299 MB	0.99923	0.914

Validation

Two independent checks, both against the full-precision encoder.

1. Encoder fidelity (feature space). Per-token cosine of the int8 encoder's [256, 768] output vs the fp16 encoder, on 40 real prompts: mean 0.99923, min 0.914. (The T5Gemma forward used as reference was first validated to cosine 1.00000 against the encodings the original pipeline produces — so the reference is exact.)

2. Audio prompt-adherence (end-to-end). The feature test is necessary but not sufficient; what matters is whether generated audio still matches the prompt. So we generated audio with Stable Audio 3 small-music using the fp16 encoder vs this int8 encoder — same prompt, same seed, same DiT + decoder — and scored audio↔text similarity with CLAP (laion/clap-htsat-unfused), 12 prompts × 2 seeds:

encoder	mean CLAP adherence	retrieval@1
fp16 (563 MB)	0.4777	54%
int8 (299 MB)	0.4774	54%

Δ = −0.0003 (range ±0.06 — i.e. noise).

File format

t5gemma_int8_g64.safetensors. Quantized 2-D tensors are stored as three keys — {name}.q, {name}.scales, {name}.biases — ready for mx.quantized_matrix / mx.quantized_matmul / mx.dequantize. 1-D params are stored plain as fp16. Metadata: {"scheme": "mlx_int8_g64", "quant_bits": "8", "quant_group": "64"}.

Usage (MLX)

import mlx.core as mx
from safetensors import safe_open

BITS, GROUP = 8, 64
w = {}
with safe_open("t5gemma_int8_g64.safetensors", framework="numpy") as f:
    keys = set(f.keys())
    bases = {k[:-2] for k in keys if k.endswith(".q")}
    for base in bases:
        q      = mx.array(f.get_tensor(base + ".q"))
        scales = mx.array(f.get_tensor(base + ".scales"))
        biases = mx.array(f.get_tensor(base + ".biases"))
        # dequantize to fp16 (or keep packed and use mx.quantized_matmul in your layers)
        w[base] = mx.dequantize(q, scales, biases, group_size=GROUP, bits=BITS)
    for k in keys:  # 1-D params kept as fp16
        if not (k.endswith(".q") or k.endswith(".scales") or k.endswith(".biases")):
            w[k] = mx.array(f.get_tensor(k))

For best memory/speed, keep the big matrices packed and call mx.quantized_matmul(x, q, scales, biases, group_size=64, bits=8) inside your attention/MLP layers instead of dequantizing up front.

The encoder architecture is standard T5Gemma-b (hidden 768, 12 layers, 12 heads, head_dim 64, RoPE θ=10000, attn logit softcap 50, query_pre_attn_scalar 64, gated GELU, RMSNorm with (weight+1), embeddings scaled by √768). Use the tokenizer from stabilityai/stable-audio-3-small-music; pad to 256 tokens.

Drop-in replacement for the encoder Stable Audio 3 ships

The keys here are the flattened T5Gemma encoder names (embed_tokens.weight, layers.N.self_attn.q_proj.weight, norm.weight, …) — the same names a hand-written MLX T5Gemma forward uses. To swap it in:

1. Load the weights as above (dequantize, or keep packed for mx.quantized_matmul).
2. Run your T5Gemma-b forward on the tokenized prompt (tokenizer from the SA3 repo; pad to 256; embeddings scaled by √768) to get [1, 256, 768].
3. Feed the conditioner exactly as before. The result replaces the fp16 encoder's output 1:1 — no other pipeline change, since the architecture and tokenizer match.

⚠️ One gotcha if you assemble the cross-attention conditioning yourself: the SA3 conditioner zeros/learns the padded token positions. Substitute the encoder output only into the valid token rows; leave the padded rows as the conditioner produces them. Putting raw encoder output into all 256 positions fills ~95% of the sequence with noise and destroys prompt adherence. (If you instead swap the weights inside the conditioner's encoder module, padding is handled for you.)

Reproduce

quantize_export.py (included) regenerates this file from your own copy of the fp16 encoder:

python quantize_export.py --src /path/to/t5gemma_f16.safetensors        # int8, 299 MB
python quantize_export.py --src /path/to/t5gemma_f16.safetensors --bits 4  # int4, 158 MB

The int4 variant carries a fidelity risk — a few embedding rows overflow int4.

License

This is a derivative of Stability AI + Google Gemma weights and is governed by both upstream licenses. See the bundled NOTICE.

• Stability AI Community License — https://stability.ai/community-license-agreement (free for non-commercial use and for commercial use under US $1M annual revenue; above that, an enterprise license is required). Powered by Stability AI.
• Gemma Terms of Use — https://ai.google.dev/gemma/terms, subject to the Gemma Prohibited Use Policy — https://ai.google.dev/gemma/prohibited_use_policy

Modifications: int8 groupwise quantization only; no retraining or architectural change.