Add int8 dynamic quantization support (#147)

# Add int8 dynamic quantization
 
Closes #7
 
## Summary
 
This PR implements mixed-precision int8 dynamic quantization for
pocket-tts, as requested in issue #7. It adds a `quantize=True` flag to
`TTSModel.load_model()` that reduces runtime memory usage by ~48% and
improves inference speed by ~27% on x86 and up to 23% on ARM, with no
measurable impact on audio quality.
This commit is contained in:
nabil-tazi
2026-03-24 11:29:14 +01:00
committed by GitHub
parent 0dcf6cb442
commit 670e777260
8 changed files with 1506 additions and 512 deletions
+76
View File
@@ -0,0 +1,76 @@
# Quantization
Pocket TTS supports dynamic int8 quantization to reduce runtime memory usage and improve inference speed on x86 CPUs.
## Quick Start
### CLI
```bash
pocket-tts generate --quantize --text "Hello world"
pocket-tts serve --quantize
```
### Python API
```python
from pocket_tts import TTSModel
model = TTSModel.load_model(quantize=True)
voice_state = model.get_state_for_audio_prompt("alba")
audio = model.generate_audio(voice_state, "Hello world!")
```
## Installation
Quantization works out of the box on any supported PyTorch version (2.5+) using `torch.ao`.
For optimized performance, install `torchao` (requires torch 2.10+):
```bash
pip install pocket-tts[quantize]
```
The quantization module automatically selects the best available backend:
- **torchao** (torch 2.10+): optimized C++ kernels, faster on both ARM and x86
- **torch.ao** (torch 2.5-2.9): deprecated but functional fallback
## Performance
Benchmarks on the full eval paragraph across 8 voices, 5 isolated runs per config.
### x86 (FBGEMM, ubuntu-latest GitHub Actions runner)
| Config | Runtime Memory | RTS | Speedup vs Baseline |
| --------------------------- | -------------- | --------- | ------------------- |
| baseline | 450 MB | 3.17x | -- |
| **attention_ffn (default)** | **234 MB** | **4.04x** | **1.27x** |
| all | 206 MB | 4.01x | 1.26x |
### ARM (QNNPACK, Apple M4 MacBook Air, torchao backend)
| Config | Runtime Memory | RTS | Speedup vs Baseline |
| --------------------------- | -------------- | --------- | ------------------- |
| baseline | 450 MB | 6.33x | -- |
| **attention_ffn (default)** | **234 MB** | **7.76x** | **1.23x** |
With the `torch.ao` fallback (torch 2.5-2.9), ARM performance is ~16% slower than baseline rather than faster. Upgrading to torch 2.10+ with torchao is recommended for ARM users.
## Quality
Quantization has no measurable impact on speech quality:
- **WER (Word Error Rate)**: WER delta for the default attention_ffn config is −0.022 ±0.032 - the ± range crosses zero, meaning the delta is indistinguishable from measurement noise.
- **Subjective listening**: no audible difference across all 8 voices
## What gets quantized
When `quantize=True`, the following layer groups in the FlowLM transformer are quantized to int8:
| Group | Params | Description |
| ----------- | ------ | ------------------------------------------------ |
| `attention` | ~25M | Q/K/V/output projections in transformer layers |
| `ffn` | ~50M | Feed-forward linear layers in transformer layers |
The flow matching network (`flow_net`, ~7M params) and the Mimi VAE decoder (convolutional) remain in float32.