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.

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:

Δ = −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.