Feature request: support saving/loading quantized language model

[mscdex]

As far as I can tell there is no way to save/load a quantized TTSModel. Being able to do that would allow for a further reduction in file size and avoid redundant re-quantizing, resulting in faster loading times.

Is there a workaround for the existing behavior/logic that exists for 2.1.0 (that doesn't involve loading a model twice -- e.g. first loading the default model and then overwriting things with a quantized state_dict)?

[mscdex]

On second thought, it appears that the quantization can result in poor output for non-English languages, so I will have to stick to using the original models without quantization.

[gabrieldemarmiesse]

Indded we don't currently store pre-quantized weights, and it's true that it could be a nice addition. Let's keep this issue open and we'll let the community handle it :)

[Abhishek8108]

Hi @gabrieldemarmiesse — thanks for keeping this open. Happy to take this on. Wanted to align on a few design choices before writing code, as the area has some non-obvious tradeoffs.

What I tested first

Poked at both quantization backends in pocket_tts/quantization.py to understand what actually serialises:

Backend	State_dict structure	safetensors compat	torch.save + weights_only=True	Quantize time (151M-param test)
torchao	LinearActivationQuantizedTensor → AffineQuantizedTensor → PlainAQTTensorImpl	no — raw save_file fails ("invalid python storage" — subclasses need __tensor_flatten__ machinery)	yes, bit-identical roundtrip	~150ms
torch.ao (deprecated)	_packed_params (opaque tuple) + scale + zero_point per layer	no — rejected (tuple + torch.dtype non-tensor values)	yes, bit-identical roundtrip	~300ms

On a synthetic 151M-param transformer (similar shape to FlowLM): FP32 = 604 MB, quantized = 152 MB — 75% reduction for both backends. So the request's intuition is right: both file-size and load-time wins are real.

Where the field has converged

HuggingFace Transformers, vLLM (compressed-tensors), and MLX have all settled on the same pattern: a single loader that auto-detects quantization from a metadata file alongside the weights — no separate from_quantized() classmethod (AutoGPTQ shipped one historically, but no new library has adopted it). vLLM's docs are explicit: "pre-quantization of models is recommended as runtime quantization adds overhead to startup time" — exactly the request here.

Worth noting: torchao 0.15+ technically supports safetensors round-trip via __tensor_flatten__, but the machinery to walk subclass trees and serialise function references into the safetensors __metadata__ header lives in HF Transformers' integrations, not in safetensors itself. Vendoring it for pocket-tts is doable (~50 lines) but adds surface area to track against torchao's still-evolving subclass layout — PlainAQTTensorImpl, PlainLayout, and AffineQuantizedTensor are all marked deprecated as of 0.17 (pytorch/ao#2752).

Proposed design

Mirror the MLX / compressed-tensors directory-of-files convention:

my-model-q/
├── pocket_tts_quant.json    # manifest
└── weights.pt               # via torch.save (loaded with weights_only=True)

{
  "format_version": "1.0",
  "backend": "torchao",
  "torchao_version": "0.17.0",
  "torch_version": "2.10.0",
  "config_class": "Int8DynamicActivationInt8WeightConfig",
  "quantize_groups": ["attention", "ffn"],
  "source_config_path": "english.yaml",
  "source_config_sha256": "abc123…",
  "saved_at": "2026-05-10T11:23:00Z"
}

API surface (additive, no breakage):

# Save (new method on TTSModel)
model.save_quantized("./my-model-q")

# Load (one new optional kwarg on existing load_model — no new classmethod)
model = TTSModel.load_model(quantized_weights_path="./my-model-q")

Optional CLI:

pocket-tts export-quantized --language english --output ./my-model-q
pocket-tts generate --quantized-weights-path ./my-model-q --text "..."

Validation the manifest enables:

• Backend match — refuse to load a torchao-saved file under the deprecated torch.ao runtime (the tensor types differ; silent failure would be worse than an explicit error).
• Config SHA256 — catches "quantized with english.yaml, loaded with french_24l.yaml" — a real correctness pitfall otherwise.
• Format version — future format migrations stay safe; old loaders refuse newer formats with a clear message.
• Bit-identical roundtrip test in CI — torch.testing.assert_close(rtol=0, atol=0) on a fixed-seed forward pass guarantees quantize → save → load → infer matches quantize → infer exactly.

Three things I'd love your steer on before coding

1. Drop torch.ao backend for save/load entirely? It's deprecated and slated for removal after torch 2.10 per the upstream notice. Supporting two save formats is real maintenance. My instinct: torchao-only for save/load; the torch.ao runtime keeps its current re-quantize-on-load behaviour as the fallback. Acceptable, or do you want both paths supported?

2. Initial weights format: torch.save (.pt) or invest in safetensors-with-flatten machinery upfront? With weights_only=True the CVE-2025-32434 gap closes, and .pt works today across both backends. safetensors is more aligned with the wider ecosystem long-term but requires ~50 lines of subclass-walking logic that needs maintenance against torchao's evolving internals. My instinct: .pt for format_version: "1.0", document the safetensors migration as a future format bump. Reasonable, or strong preference for safetensors from day one?

3. CLI scope — add the export-quantized subcommand, or keep it Python-API only for now? export-quantized is the natural ergonomic for "quantize once locally, ship the artifact" (and would let you eventually publish a pre-quantized variant on Hub), but adds CLI surface to maintain.

Happy to do the work and write the tests — just want to be sure I'm building toward what you'd actually want to merge. Cheers!