kaaaaan/t5gemma-sa3-int8

T5Gemma (Stable Audio 3) — portable int8 encoder

An int8-quantized, encoder-only copy of the T5Gemma text encoder shipped in stabilityai/stable-audio-3-small-music, in a framework-neutral packed-safetensors format (no MLX/bitsandbytes/GPU required — load with plain safetensors + numpy or torch).

• 290 MB — vs the 1183 MB canonical bf16 encoder-decoder (~4× smaller), or ~2× vs the 563 MB encoder half.
• Encoder only. Stable Audio 3 uses only the encoder; the unused decoder half (620 MB) is dropped.
• Near-lossless: per-token cosine to the bf16 encoder = 0.9995 (mean), on real prompts.
• Audio-verified: through the frozen Stable Audio 3 DiT, this encoder gives CLAP prompt-adherence parity with the full-precision encoder (Δ ≈ 0 over 12 prompts × 2 seeds).
• No retraining — pure post-training quantization; identical architecture and 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.

There is also an MLX build for Apple Silicon — kaaaaan/t5gemma-sa3-int8-mlx.

Why int8 (and why not smaller)

70% of the encoder's size is the 256k-vocab embedding table. int8 quantizes it near-losslessly. int4 did not hold in my testing

scheme	size	mean cos vs bf16	min cos
bf16 encoder	563 MB	1.00000	1.00000
int8 (this repo)	290 MB	0.9995	0.933

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 bf16 encoder, on 40 real prompts: mean 0.9995, min 0.933. (The reproduction of the T5Gemma forward used for this 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 original encoder vs this int8 encoder — same prompt, same seed, same DiT + decoder (only the text encoder differs) — and scored audio↔text similarity with CLAP (laion/clap-htsat-unfused), over 12 prompts × 2 seeds:

encoder	mean CLAP adherence	retrieval@1
original	0.4777	54%
int8	0.4774	54%

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

Format

Symmetric int8, groupwise (group size 64) along the input dim. Each 2-D matrix W [out, in] is stored as two tensors:

• {name}.qweight — int8, shape [out, in]
• {name}.scales — fp16, shape [out, in/64] (per-(row, group) absmax / 127)

1-D params (RMSNorm weights, RoPE inverse frequencies) are stored plain as fp16 under their own name. Dequant: W ≈ qweight.float() * scales.repeat_interleave(64, -1). Metadata records format, group_size, bits, layout.

Usage (any framework — torch shown)

import torch
from safetensors import safe_open

GROUP = 64
weights = {}
with safe_open("t5gemma_encoder_int8.safetensors", framework="pt") as f:
    keys = set(f.keys())
    bases = {k[:-8] for k in keys if k.endswith(".qweight")}
    for base in bases:
        q = f.get_tensor(base + ".qweight").float()   # [out, in]
        s = f.get_tensor(base + ".scales").float()    # [out, in/GROUP]
        s_full = s.repeat_interleave(q.shape[1] // s.shape[1], dim=-1)
        weights[base] = (q * s_full).to(torch.float16)   # dequantized weight
    for k in keys:                                    # 1-D params: plain fp16
        if not (k.endswith(".qweight") or k.endswith(".scales")):
            weights[k] = f.get_tensor(k)

weights is now a standard {name: tensor} dict for the T5Gemma-b encoder (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). Key names match the upstream encoder with the model.encoder. prefix stripped. Use the tokenizer from stabilityai/stable-audio-3-small-music; pad to 256 tokens. (To keep memory low, dequant lazily per layer instead of all at once.)

Drop-in replacement for the encoder Stable Audio 3 ships

The cleanest swap is to load these weights into the text-conditioner's encoder and let the pipeline handle tokenization and padding as usual. This avoids the one real gotcha (see below).

import torch
from safetensors import safe_open

# `conditioner` is the SA3 MultiConditioner; its prompt encoder is a HF
# T5GemmaEncoderModel exposed at conditioner.conditioners["prompt"].model
prompt_enc = conditioner.conditioners["prompt"].model

GROUP = 64
state = {}
with safe_open("t5gemma_encoder_int8.safetensors", framework="pt") as f:
    keys = set(f.keys())
    bases = {k[:-8] for k in keys if k.endswith(".qweight")}
    for base in bases:
        q = f.get_tensor(base + ".qweight").float()
        s = f.get_tensor(base + ".scales").float()
        s_full = s.repeat_interleave(q.shape[1] // s.shape[1], dim=-1)
        state["encoder." + base] = q * s_full          # HF module uses 'encoder.' prefix
    for k in keys:
        if not (k.endswith(".qweight") or k.endswith(".scales")):
            state["encoder." + k] = f.get_tensor(k).float()

missing, unexpected = prompt_enc.load_state_dict(state, strict=False)
assert not missing and not unexpected           # keys map 1:1
prompt_enc.half()                               # optional

That's it — generate exactly as before. Verified: conditioning cosine 0.99972 vs the original encoder, padded positions identical. (Names strip model.encoder.; the conditioner's HF module wants an encoder. prefix, hence the rename above.)

If instead you precompute conditioning tensors yourself (the low-level conditioning_tensors= path), you MUST reproduce the conditioner's padding: it zeros/learns the padded positions, so feed the encoder output only into the valid token rows and leave the padded rows as the pipeline produces them. Substituting raw encoder output into all 256 positions puts noise in ~95% of the sequence and destroys prompt adherence. The load_state_dict route above sidesteps this entirely.

The architecture and tokenizer are unchanged, so nothing else in the pipeline needs to move.

Reproduce

python quantize_export.py --src /path/to/t5gemma-b-b-ul2/model.safetensors

Point --src at the upstream encoder-decoder safetensors (from stable-audio-3-small-music, which you have licensed). The script keeps the encoder half only and quantizes it.

License

A derivative of Stability AI + Google Gemma weights, governed by both upstream licenses. See the bundled NOTICE.

• Stability AI Community License — https://stability.ai/community-license-agreement (free for non-commercial and for commercial use under US $1M annual revenue). 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 + decoder dropped; no retraining or architectural change.