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:
@@ -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.
|
||||
Reference in New Issue
Block a user