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
+8
-2
@@ -188,11 +188,14 @@ def serve(
|
||||
help="Path to locally-saved model config .yaml file or model variant signature"
|
||||
),
|
||||
] = DEFAULT_VARIANT,
|
||||
quantize: Annotated[
|
||||
bool, typer.Option(help="Apply int8 quantization to reduce memory usage")
|
||||
] = False,
|
||||
):
|
||||
"""Start the FastAPI server."""
|
||||
|
||||
global tts_model, global_model_state
|
||||
tts_model = TTSModel.load_model(config)
|
||||
tts_model = TTSModel.load_model(config, quantize=quantize)
|
||||
|
||||
# Pre-load the voice prompt
|
||||
global_model_state = tts_model.get_state_for_audio_prompt(voice)
|
||||
@@ -236,6 +239,9 @@ def generate(
|
||||
max_tokens: Annotated[
|
||||
int, typer.Option(help="Maximum number of tokens per chunk.")
|
||||
] = MAX_TOKEN_PER_CHUNK,
|
||||
quantize: Annotated[
|
||||
bool, typer.Option(help="Apply int8 quantization to reduce memory usage")
|
||||
] = False,
|
||||
):
|
||||
"""Generate speech using Kyutai Pocket TTS."""
|
||||
log_level = logging.ERROR if quiet else logging.INFO
|
||||
@@ -248,7 +254,7 @@ def generate(
|
||||
logger.error("No input received from stdin.")
|
||||
raise typer.Exit(code=1)
|
||||
tts_model = TTSModel.load_model(
|
||||
config, temperature, lsd_decode_steps, noise_clamp, eos_threshold
|
||||
config, temperature, lsd_decode_steps, noise_clamp, eos_threshold, quantize=quantize
|
||||
)
|
||||
tts_model.to(device)
|
||||
|
||||
|
||||
@@ -33,6 +33,7 @@ from pocket_tts.modules import mimi_transformer
|
||||
from pocket_tts.modules.dummy_quantizer import DummyQuantizer
|
||||
from pocket_tts.modules.seanet import SEANetDecoder, SEANetEncoder
|
||||
from pocket_tts.modules.stateful_module import StatefulModule, increment_steps, init_states
|
||||
from pocket_tts.quantization import RECOMMENDED_CONFIG, apply_dynamic_int8
|
||||
from pocket_tts.utils.config import Config, load_config
|
||||
from pocket_tts.utils.utils import (
|
||||
PREDEFINED_VOICES,
|
||||
@@ -192,6 +193,7 @@ class TTSModel(nn.Module):
|
||||
lsd_decode_steps: int = DEFAULT_LSD_DECODE_STEPS,
|
||||
noise_clamp: float | int | None = DEFAULT_NOISE_CLAMP,
|
||||
eos_threshold: float = DEFAULT_EOS_THRESHOLD,
|
||||
quantize: bool = False,
|
||||
) -> Self:
|
||||
"""Load a pre-trained TTS model with specified configuration.
|
||||
|
||||
@@ -210,6 +212,11 @@ class TTSModel(nn.Module):
|
||||
is applied. Helps prevent extreme values in generation.
|
||||
eos_threshold: Threshold for end-of-sequence detection. Higher values
|
||||
make the model more likely to continue generating.
|
||||
quantize: If True, apply dynamic int8 quantization to the transformer's
|
||||
attention and FFN layers. Reduces runtime memory by ~48% and improves
|
||||
inference speed by ~27% on x86 (FBGEMM).
|
||||
No measurable impact on speech quality (WER unchanged).
|
||||
For optimized performance, install torchao: ``pip install pocket-tts[quantize]``
|
||||
|
||||
Returns:
|
||||
TTSModel: Fully initialized model with loaded weights on cpu, ready for
|
||||
@@ -227,8 +234,8 @@ class TTSModel(nn.Module):
|
||||
# Load with default settings
|
||||
model = TTSModel.load_model()
|
||||
|
||||
# Load with custom parameters
|
||||
model = TTSModel.load_model(variant="b6369a24", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0)
|
||||
# Load with int8 quantization
|
||||
model = TTSModel.load_model(quantize=True)
|
||||
```
|
||||
"""
|
||||
if str(config).endswith(".yaml"):
|
||||
@@ -241,6 +248,10 @@ class TTSModel(nn.Module):
|
||||
tts_model = TTSModel._from_pydantic_config_with_weights(
|
||||
config, temp, lsd_decode_steps, noise_clamp, eos_threshold
|
||||
)
|
||||
|
||||
if quantize:
|
||||
apply_dynamic_int8(tts_model.flow_lm, RECOMMENDED_CONFIG)
|
||||
|
||||
return tts_model
|
||||
|
||||
def _run_flow_lm_and_increment_step(
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
Dynamic int8 quantization for pocket-tts.
|
||||
|
||||
Uses torchao if available (torch 2.10+ with C++ extensions), otherwise
|
||||
falls back to torch.ao.quantization (deprecated but functional on torch 2.5-2.9).
|
||||
|
||||
Quantizes attention (Q/K/V/output projections) and FFN (linear1/linear2) layers
|
||||
in the FlowLM transformer. The flow matching network and Mimi VAE decoder
|
||||
remain in float32.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import platform
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Used by load_model(quantize=True) to specify which layer groups to quantize.
|
||||
RECOMMENDED_CONFIG = {"attention", "ffn"}
|
||||
|
||||
|
||||
def _get_backend():
|
||||
"""Detect the best available quantization backend.
|
||||
|
||||
Returns "torchao" if torchao is installed with working C++ extensions,
|
||||
otherwise returns "torch.ao".
|
||||
"""
|
||||
try:
|
||||
import importlib.util
|
||||
|
||||
if importlib.util.find_spec("torchao") is None:
|
||||
return "torch.ao"
|
||||
|
||||
import torchao
|
||||
|
||||
if hasattr(torchao, "_C") or not getattr(torchao, "_SKIPPED_CPP_EXTENSIONS", False):
|
||||
return "torchao"
|
||||
except ImportError:
|
||||
pass
|
||||
return "torch.ao"
|
||||
|
||||
|
||||
def _quantize_module_torchao(module: nn.Module):
|
||||
"""Apply int8 dynamic quantization using torchao."""
|
||||
from torchao.quantization import Int8DynamicActivationInt8WeightConfig, quantize_
|
||||
|
||||
quantize_(module, Int8DynamicActivationInt8WeightConfig())
|
||||
|
||||
|
||||
def _ensure_quantization_engine():
|
||||
"""Set the quantization engine for torch.ao (QNNPACK for ARM, FBGEMM for x86)."""
|
||||
if platform.machine() in ("arm64", "aarch64"):
|
||||
torch.backends.quantized.engine = "qnnpack"
|
||||
elif torch.backends.quantized.engine == "none":
|
||||
torch.backends.quantized.engine = "fbgemm"
|
||||
|
||||
|
||||
def apply_dynamic_int8(flow_lm: nn.Module, quantize_groups: set[str]) -> nn.Module:
|
||||
"""
|
||||
Apply dynamic int8 quantization to the specified layer groups of a FlowLM model.
|
||||
|
||||
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
|
||||
|
||||
Args:
|
||||
flow_lm: The FlowLM model (model.flow_lm)
|
||||
quantize_groups: Set of group keys to quantize.
|
||||
Valid keys: "attention", "ffn", "flow_net"
|
||||
|
||||
Returns:
|
||||
The quantized model (modified in-place).
|
||||
"""
|
||||
if not quantize_groups:
|
||||
logger.info("No quantization groups specified, returning model unchanged.")
|
||||
return flow_lm
|
||||
|
||||
backend = _get_backend()
|
||||
logger.info("Using quantization backend: %s", backend)
|
||||
|
||||
if backend == "torchao":
|
||||
_apply_torchao(flow_lm, quantize_groups)
|
||||
else:
|
||||
_apply_torch_ao(flow_lm, quantize_groups)
|
||||
|
||||
return flow_lm
|
||||
|
||||
|
||||
def _apply_torchao(flow_lm: nn.Module, quantize_groups: set[str]) -> None:
|
||||
"""Apply quantization using torchao backend."""
|
||||
if "flow_net" in quantize_groups:
|
||||
_quantize_module_torchao(flow_lm.flow_net)
|
||||
|
||||
for layer in flow_lm.transformer.layers:
|
||||
if "attention" in quantize_groups:
|
||||
_quantize_module_torchao(layer.self_attn)
|
||||
|
||||
if "ffn" in quantize_groups:
|
||||
wrapper1 = nn.Sequential(layer.linear1)
|
||||
wrapper2 = nn.Sequential(layer.linear2)
|
||||
_quantize_module_torchao(wrapper1)
|
||||
_quantize_module_torchao(wrapper2)
|
||||
layer.linear1 = wrapper1[0]
|
||||
layer.linear2 = wrapper2[0]
|
||||
|
||||
|
||||
def _apply_torch_ao(flow_lm: nn.Module, quantize_groups: set[str]) -> None:
|
||||
"""Apply quantization using deprecated torch.ao backend (fallback)."""
|
||||
from torch.ao.quantization import quantize_dynamic
|
||||
|
||||
_ensure_quantization_engine()
|
||||
|
||||
if "flow_net" in quantize_groups:
|
||||
quantize_dynamic(flow_lm.flow_net, {nn.Linear}, dtype=torch.qint8, inplace=True)
|
||||
|
||||
for layer in flow_lm.transformer.layers:
|
||||
if "attention" in quantize_groups:
|
||||
quantize_dynamic(layer.self_attn, {nn.Linear}, dtype=torch.qint8, inplace=True)
|
||||
|
||||
if "ffn" in quantize_groups:
|
||||
layer.linear1 = quantize_dynamic(
|
||||
nn.Sequential(layer.linear1), {nn.Linear}, dtype=torch.qint8
|
||||
)[0]
|
||||
layer.linear2 = quantize_dynamic(
|
||||
nn.Sequential(layer.linear2), {nn.Linear}, dtype=torch.qint8
|
||||
)[0]
|
||||
Reference in New Issue
Block a user