Add optional int8 weight-only quantization

.gitignore

@@ -14,4 +14,7 @@ pretrained_models/
 
 # Eval
 data
-data/
+data/
+# Local venv / test output
+.venv/
+outputs/

README.md

@@ -79,6 +79,109 @@
 pip install -r requirements.txt
 ```
 
+### Installation on Apple Silicon (MPS)
+
+`flash_attn` has no macOS wheel, and `torchcodec` needs FFmpeg <= 7, so use the
+macOS requirement set instead (attention falls back to PyTorch SDPA, audio I/O
+to `soundfile`):
+
+```sh
+uv venv --python 3.12 .venv
+uv pip install --python .venv/bin/python -r requirements-mps.txt
+```
+
+The device and compute dtype are resolved once in `fireredtts3/utils/device.py`
+(CUDA -> MPS -> CPU) and can be overridden:
+
+| Variable | Values | Default |
+| --- | --- | --- |
+| `FIRERED_DEVICE` | `cuda` / `mps` / `cpu` | first available |
+| `FIRERED_DTYPE` | autocast dtype (activations) | `bfloat16` on CUDA, `float32` elsewhere |
+| `FIRERED_WEIGHT_DTYPE` | dtype the checkpoints are loaded in | as stored (fp32 officially) |
+| `FIRERED_QUANT` | `int8` weight-only quantization | `none` |
+
+`FIRERED_DTYPE` only casts activations, so it does **not** shrink resident
+memory; `FIRERED_WEIGHT_DTYPE` is the knob that does. Measured on an M4 Max
+(torch 2.8, macOS 26), Base zero-shot cloning of a 3-sentence paragraph over
+two reference voices, RTF as the range across 3 interleaved runs:
+
+| Weights | Resident weights | RTF | Speaker sim |
+| --- | --- | --- | --- |
+| `float32` (official checkpoints) | 11.42 GiB | 0.87-0.96 | 0.947 / 0.935 |
+| `bfloat16` | 5.71 GiB | **0.73-0.76** | 0.941 / 0.933 |
+
+**On MPS, bfloat16 weights are what you want:** half the memory *and* ~20%
+faster, at effectively unchanged speaker similarity. There are two ways to get
+them, and they are equivalent at inference time:
+
+**a) Cast on the fly** — keep the official fp32 checkpoints on disk and convert
+them while loading:
+
+```sh
+FIRERED_WEIGHT_DTYPE=bfloat16 python your_script.py
+```
+
+Nothing on disk changes (19 GB stays 19 GB), and you can go back to fp32 by
+dropping the variable. The cast is redone on every start.
+
+**b) Convert the checkpoints once** — halve them on disk too, then feed those
+weights with no variable set at all:
+
+```sh
+python scripts/convert_to_bf16.py          # in place: 19 GB -> 9.7 GB
+python your_script.py                      # loads bfloat16 by itself
+```
+
+Each tensor of the new file is verified against the fp32 original for exact
+bfloat16 rounding before the original is replaced, and the script skips
+components that are already converted. Loading gets faster too (3.3 s vs
+5.7-9.2 s here), since half as much is read from disk. Pass `--keep-fp32` to
+write `<dir>.bf16` alongside instead of replacing, and note the cast is one-way
+— recovering fp32 means re-downloading.
+
+Option (b) is also the way to move a converted model to another machine: copy
+`pretrained_models/` over, or run the script there after downloading. Either
+way, leaving `FIRERED_WEIGHT_DTYPE` unset means the loader keeps whatever the
+checkpoint stores, so nothing changes unless you ask for it.
+
+Autocast, by contrast, buys nothing on MPS (RTF 0.88-0.93 vs 0.85-0.91 over
+fp32 weights) — the DiT flow head, the real bottleneck, sits outside the
+autocast region — so autocast is off by default off CUDA.
+
+#### int8 quantization
+
+`FIRERED_QUANT=int8` applies weight-only int8 quantization to the transformer
+backbones. It needs `optimum-quanto`, imported lazily — the default
+`FIRERED_QUANT=none` requires nothing installed and does not touch the model:
+
+```sh
+pip install optimum-quanto
+FIRERED_QUANT=int8 FIRERED_WEIGHT_DTYPE=bfloat16 python your_script.py
+```
+
+int8 stacks on top of whatever the weights already are, so on a checkpoint
+already converted by `scripts/convert_to_bf16.py` the dtype variable is
+redundant — `FIRERED_QUANT=int8` alone gives the same result.
+
+Same benchmark as above:
+
+| Weights | Resident weights | RTF | Speaker sim |
+| --- | --- | --- | --- |
+| `float32`, no quant (default) | 11.42 GiB | 0.87-0.96 | 0.947 / 0.935 |
+| `bfloat16`, no quant | 5.71 GiB | 0.73-0.76 | 0.941 / 0.933 |
+| `bfloat16` + `int8` | **3.90 GiB** | 1.03-1.05 | 0.937 / 0.934 |
+
+Reach for int8 when memory is the binding constraint — 2.9x smaller than the
+default at ~40% higher latency than plain bf16. If you have the RAM, bf16
+without quantization is both faster and smaller than the fp32 default.
+
+The DiT flow head is deliberately left unquantized: it runs once per flow
+timestep with a CFG-doubled batch, so per-call dequantization dominates —
+quantizing every `nn.Linear` measured RTF 2.41 for only 8% more savings. `int4`
+is rejected with an explanatory error (measured 3x slower with a clear
+similarity drop, and quanto's int4 path conflicts with the
+`@torch.inference_mode()` decorators).
+
 ### Model Download
 
 Download the pretrained model from Hugging Face with the `hf` CLI:

fireredtts3/campp/campp.py

@@ -2,10 +2,13 @@
 import torchaudio
 import torchaudio.compliance.kaldi as kaldi
 from fireredtts3.campp.DTDNN import CAMPPlus
+from fireredtts3.utils.device import fft_device
 
 
 def extract_kaldi_mel(audio: torch.Tensor, audio_sr: int):
     audio = audio[:1]
+    # kaldi.fbank runs an FFT internally; fall back to CPU on devices without it
+    audio = fft_device(audio)
     if audio_sr != 16000:
         audio = torchaudio.functional.resample(audio, audio_sr, 16000)
         audio_sr = 16000

fireredtts3/core.py

@@ -381,7 +381,7 @@ def generate_tts(
         gen_audio_sr = None
         segments: List[torch.Tensor] = []
         for sent in sentences:
-            seg, seg_sr, _ = super().generate_tts(
+            seg, seg_sr = super().generate_tts(
                 prompt_text=prompt_text,
                 prompt_audio=prompt_audio,
                 prompt_audio_sr=prompt_audio_sr,

fireredtts3/llm/fireredtts3_base.py

@@ -11,6 +11,9 @@
 from fireredtts3.redae.redae import RedAE
 from fireredtts3.campp.campp import CamppEmbedding
 from fireredtts3.utils.utils import fix_seed
+from fireredtts3.utils.device import (
+    get_device, get_attn_implementation, get_weight_dtype, quantize_weights, autocast,
+)
 from fireredtts3.utils.text_tokenizer import (
     load_text_tokenizer, 
     MULTI_LANG_TAGS, MULTI_DIALECT_TAGS,
@@ -46,7 +49,7 @@
     "use_cache": True,
     "use_sliding_window": False,
     "vocab_size": 151936,
-    "attn_implementation": "flash_attention_2",
+    "attn_implementation": get_attn_implementation(),
 }
 
 
@@ -140,7 +143,7 @@ def _init_weights(self, module):
             module.rope_init()
 
     # Backbone Transformer AR wrapper
-    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
+    @autocast
     def _backbone_one_step(self, input_embeds: torch.Tensor, cache = None):
         outs = self.backbone_llm.forward(
             inputs_embeds=input_embeds,
@@ -162,7 +165,8 @@ def _flow_one_step(
         inference_cfg: float,
     ):
         # Compose input
-        x0 = torch.randn(1, self.patch_size, self.redae_dim, device=hist_latents.device)
+        x0 = torch.randn(1, self.patch_size, self.redae_dim,
+                         device=hist_latents.device, dtype=hist_latents.dtype)
 
         xt = torch.cat([hist_latents, x0], dim=1)   # History clean + current noise
         cond = torch.cat([
@@ -213,7 +217,7 @@ def generate(
         input_embeds = torch.cat([spk_embs_llm.unsqueeze(1), input_embeds, patch_prompt_latents], dim=1)
 
         # Prepare DiT decode
-        t_span = torch.linspace(0, 1, n_timesteps + 1).to(device)
+        t_span = torch.linspace(0, 1, n_timesteps + 1).to(device=device, dtype=input_embeds.dtype)
         t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) # (n_timesteps+1,)
         latents_gen = F.pad(prompt_latents, (0, 0, self.history_length, 0))
 
@@ -264,17 +268,19 @@ def generate(
 # RedAE + TextTokenizer + TTS3Core
 class FireRedTTS3Base(object):
     def __init__(self, pretrained_model_dir: str):
-        self.device = torch.device('cuda')
+        self.device = get_device()
         # RedAE
         redae_model_dir = os.path.join(pretrained_model_dir, 'redae')
         assert os.path.exists(redae_model_dir), f'{redae_model_dir} not found'
-        self.redae = RedAE.from_pretrained(redae_model_dir)
+        self.redae = RedAE.from_pretrained(redae_model_dir, dtype=get_weight_dtype())
         self.redae.to(self.device)
         # LLM-DiT 
         tts_model_dir = os.path.join(pretrained_model_dir, 'fireredtts3_base')
         assert os.path.exists(tts_model_dir), f'{tts_model_dir} not found'
-        self.tts_core = FireRedTTS3BaseCore.from_pretrained(tts_model_dir)
+        self.tts_core = FireRedTTS3BaseCore.from_pretrained(tts_model_dir, dtype=get_weight_dtype())
         self.tts_core.to(self.device)
+        # Optional int8 weight-only quantization (FIRERED_QUANT=int8)
+        quantize_weights(self.tts_core.backbone_llm, self.redae)
         # Text Tokenizer
         text_tok_dir = os.path.join(pretrained_model_dir, 'text_tokenizer')
         assert os.path.exists(text_tok_dir), f'{text_tok_dir} not found'
@@ -320,10 +326,10 @@ def generate(
         prompt_audio = self.redae.pad_to_multiple_of(prompt_audio, self.redae.downsample_rate*self.tts_core.patch_size)
         prompt_audio = prompt_audio.to(self.device)
         prompt_latents = self.redae.encode(prompt_audio, prompt_audio_sr)
-        prompt_latents = prompt_latents.to(torch.float32)
+        prompt_latents = prompt_latents.to(self.tts_core.dtype)
         # Spk emb
         spk_emb = self.spk_extractor.forward(prompt_audio, prompt_audio_sr)
-        spk_emb = spk_emb.to(self.device)
+        spk_emb = spk_emb.to(device=self.device, dtype=self.tts_core.dtype)
         # TTS
         if seed is not None:
             fix_seed(seed)

fireredtts3/llm/fireredtts3_instruct.py

@@ -17,6 +17,7 @@
 from fireredtts3.llm.dit import DiT
 from fireredtts3.redae.redae import RedAE
 from fireredtts3.utils.utils import fix_seed
+from fireredtts3.utils.device import get_device, get_weight_dtype, quantize_weights, autocast
 from fireredtts3.utils.text_tokenizer import load_text_tokenizer
 from fireredtts3.llm.fireredtts3_base import Qwen3_1_7B_ConfigDict
 from fireredtts3.utils.chatml import (
@@ -139,7 +140,7 @@ def _init_weights(self, module):
             module.rope_init()
 
     # Backbone Transformer AR wrapper
-    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
+    @autocast
     def _backbone_one_step(self, input_embeds: torch.Tensor, cache = None):
         outs = self.backbone_llm.model.forward(
             inputs_embeds=input_embeds,
@@ -160,7 +161,8 @@ def _flow_one_step(
         inference_cfg: float,
     ):
         # Compose input
-        x0 = torch.randn(1, self.patch_size, self.redae_dim, device=hist_latents.device)
+        x0 = torch.randn(1, self.patch_size, self.redae_dim,
+                         device=hist_latents.device, dtype=hist_latents.dtype)
 
         xt = torch.cat([hist_latents, x0], dim=1)   # History clean + current noise
         cond = backbone_cond.repeat_interleave(self.patch_size, dim=1)    # Correspond backbone cond
@@ -228,11 +230,11 @@ def generate(
                 latents_patch_out.reshape(-1).to(input_embeds),
             )
         else:
-            latents_out = torch.zeros(1, 0, self.config.redae_dim, device=device)
+            latents_out = torch.zeros(1, 0, self.config.redae_dim, device=device, dtype=input_embeds.dtype)
             latents_patch_out = None
 
         # Prepare DiT decode
-        t_span = torch.linspace(0, 1, n_timesteps + 1).to(device)
+        t_span = torch.linspace(0, 1, n_timesteps + 1).to(device=device, dtype=input_embeds.dtype)
         t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) # (n_timesteps+1,)
 
         # Init Backbone states
@@ -316,17 +318,19 @@ def generate(
 # RedAE + TextTokenizer + TTS3Core
 class FireRedTTS3Instruct(object):
     def __init__(self, pretrained_model_dir: str):
-        self.device = torch.device('cuda')
+        self.device = get_device()
         # RedAE
         redae_model_dir = os.path.join(pretrained_model_dir, 'redae')
         assert os.path.exists(redae_model_dir), f'{redae_model_dir} not found'
-        self.redae = RedAE.from_pretrained(redae_model_dir)
+        self.redae = RedAE.from_pretrained(redae_model_dir, dtype=get_weight_dtype())
         self.redae.to(self.device)
         # LLM-DiT 
         tts_model_dir = os.path.join(pretrained_model_dir, 'fireredtts3_instruct')
         assert os.path.exists(tts_model_dir), f'{tts_model_dir} not found'
-        self.tts_core = FireRedTTS3InstructCore.from_pretrained(tts_model_dir)
+        self.tts_core = FireRedTTS3InstructCore.from_pretrained(tts_model_dir, dtype=get_weight_dtype())
         self.tts_core.to(self.device)
+        # Optional int8 weight-only quantization (FIRERED_QUANT=int8)
+        quantize_weights(self.tts_core.backbone_llm, self.redae)
         # Text Tokenizer
         text_tok_dir = os.path.join(pretrained_model_dir, 'text_tokenizer')
         assert os.path.exists(text_tok_dir), f'{text_tok_dir} not found'
@@ -347,7 +351,7 @@ def _tokenize_audio(self, audio: torch.Tensor, audio_sr: int):
         audio = self.redae.pad_to_multiple_of(audio, self.redae.downsample_rate*self.tts_core.patch_size)
         audio = audio.to(self.device)
         latents = self.redae.encode(audio, audio_sr) * REDAE_SCALE
-        latents = latents.to(torch.float32)
+        latents = latents.to(self.tts_core.dtype)
         return latents
 
     # --- Inference Interface

fireredtts3/llm/rotary_embedding.py

@@ -1,8 +1,8 @@
 import torch
 from torch.nn import Module
-from torch.amp import autocast
 from torch import cat, stack, arange
 from einops import rearrange
+from fireredtts3.utils.device import disable_autocast
 
 
 class RotaryEmbedding(Module):
@@ -38,7 +38,7 @@ def forward_from_seq_len(self, seq_len):
         t = arange(seq_len, device = device)
         return self.forward(t)
 
-    @autocast('cuda', enabled = False)
+    @disable_autocast
     def forward(self, t, offset = 0):
         if t.ndim == 1:
             t = rearrange(t, 'n -> 1 n')
@@ -57,7 +57,7 @@ def rotate_half(x):
     return rearrange(x, '... d r -> ... (d r)')
 
 
-@autocast('cuda', enabled = False)
+@disable_autocast
 def apply_rotary_pos_emb(t, freqs, scale = 1):
     rot_dim, seq_len, orig_dtype = freqs.shape[-1], t.shape[-2], t.dtype