Update MOSS Audio Tokenizer v2 weight dtype loading methods

README.md

@@ -82,13 +82,30 @@ For production use with `trust_remote_code=True`, pin `revision` to a reviewed c
 
 `config.attention_implementation` controls whether transformer layers prefer `sdpa` or `flash_attention_2`.
 `config.compute_dtype` controls the non-quantizer autocast dtype and supports `fp32`, `bf16`.
+`config.codec_weight_dtype` controls encoder/decoder parameter dtype and defaults to `fp32`.
+The quantizer is always kept in fp32.
+
+GPU bf16 loading:
+
+```python
+import torch
+from transformers import AutoModel
+
+repo_id = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = AutoModel.from_pretrained(repo_id, trust_remote_code=True, low_cpu_mem_usage=True, codec_weight_dtype="bf16").eval().to(device)
+```
+Passing codec_weight_dtype="bf16" at load time avoids first materializing encoder/decoder weights as fp32 on GPU and then converting them to bf16, which would increase peak GPU memory.
+
+You can also switch an already loaded model:
 
 ```python
 model.set_attention_implementation("flash_attention_2")
 model.set_compute_dtype("bf16")
+model.set_codec_weight_dtype("bf16")  # encoder/decoder bf16, quantizer fp32
 ```
 
-The quantizer always runs in fp32.
+Avoid calling plain `model.to(torch.bfloat16)` on the whole codec; that also casts quantizer weights and can cause dtype mismatches or serious precision loss.
 
 ### Streaming
 

configuration_moss_audio_tokenizer.py

@@ -59,6 +59,9 @@ class MossAudioTokenizerConfig(PreTrainedConfig):
             `"flash_attention_2"`.
         compute_dtype (`str`, *optional*, defaults to `"fp32"`):
             Inference compute dtype for non-quantizer modules. Supported values are `"fp32"`, `"bf16"`.
+        codec_weight_dtype (`str`, *optional*, defaults to `"fp32"`):
+            Parameter dtype for encoder and decoder modules. The quantizer remains fp32 because it explicitly disables
+            autocast and performs numerically sensitive codebook operations in fp32.
         quantizer_type (`str`, *optional*, defaults to `"rlfq"`):
             Quantizer type. Options include `"rvq"`, `"spec_rvq"`, `"rlfq"`, `"random_prefix_rlfq"`.
         quantizer_kwargs (`dict`, *optional*):
@@ -95,6 +98,7 @@ class MossAudioTokenizerConfig(PreTrainedConfig):
     enable_channel_interleave: bool
     attention_implementation: str
     compute_dtype: str
+    codec_weight_dtype: str
     quantizer_type: str
     quantizer_kwargs: dict[str, Any]
 
@@ -110,6 +114,7 @@ class MossAudioTokenizerConfig(PreTrainedConfig):
         enable_channel_interleave: bool = True,
         attention_implementation: str = "sdpa",
         compute_dtype: str = "fp32",
+        codec_weight_dtype: str = "fp32",
         quantizer_type: str = "rlfq",
         quantizer_kwargs: dict[str, Any] | None = None,
         **kwargs,
@@ -125,6 +130,8 @@ class MossAudioTokenizerConfig(PreTrainedConfig):
             attention_implementation = kwargs.pop("attention_backend")
         if "codec_compute_dtype" in kwargs and compute_dtype == "fp32":
             compute_dtype = kwargs.pop("codec_compute_dtype")
+        if "codec_load_dtype" in kwargs and codec_weight_dtype == "fp32":
+            codec_weight_dtype = kwargs.pop("codec_load_dtype")
         reversed_decoder_kwargs = kwargs.pop("reversed_decoder_kwargs", None)
 
         # `version` is accepted for compatibility but not used in modeling.
@@ -136,6 +143,7 @@ class MossAudioTokenizerConfig(PreTrainedConfig):
         self.enable_channel_interleave = enable_channel_interleave
         self.attention_implementation = attention_implementation
         self.compute_dtype = compute_dtype
+        self.codec_weight_dtype = codec_weight_dtype
         # Default encoder configuration
         if encoder_kwargs is None:
             encoder_kwargs = [

modeling_moss_audio_tokenizer.py

@@ -62,6 +62,18 @@ except ImportError:
 
 SUPPORTED_ATTENTION_IMPLEMENTATIONS = {"sdpa", "flash_attention_2"}
 SUPPORTED_COMPUTE_DTYPES = {"fp32": None, "bf16": torch.bfloat16}
+SUPPORTED_CODEC_WEIGHT_DTYPES = {
+    "fp32": torch.float32,
+    "float32": torch.float32,
+    "bf16": torch.bfloat16,
+    "bfloat16": torch.bfloat16,
+}
+CANONICAL_CODEC_WEIGHT_DTYPES = {
+    "fp32": "fp32",
+    "float32": "fp32",
+    "bf16": "bf16",
+    "bfloat16": "bf16",
+}
 
 
 def resolve_compute_dtype(compute_dtype: str) -> torch.dtype | None:
@@ -72,6 +84,26 @@ def resolve_compute_dtype(compute_dtype: str) -> torch.dtype | None:
     return SUPPORTED_COMPUTE_DTYPES[compute_dtype]
 
 
+def canonicalize_codec_weight_dtype(codec_weight_dtype: str) -> str:
+    key = str(codec_weight_dtype).lower()
+    if key not in CANONICAL_CODEC_WEIGHT_DTYPES:
+        raise ValueError(
+            "Unsupported codec_weight_dtype="
+            f"{codec_weight_dtype!r}. Expected one of {sorted(CANONICAL_CODEC_WEIGHT_DTYPES)}."
+        )
+    return CANONICAL_CODEC_WEIGHT_DTYPES[key]
+
+
+def resolve_codec_weight_dtype(codec_weight_dtype: str) -> torch.dtype:
+    key = str(codec_weight_dtype).lower()
+    if key not in SUPPORTED_CODEC_WEIGHT_DTYPES:
+        raise ValueError(
+            "Unsupported codec_weight_dtype="
+            f"{codec_weight_dtype!r}. Expected one of {sorted(SUPPORTED_CODEC_WEIGHT_DTYPES)}."
+        )
+    return SUPPORTED_CODEC_WEIGHT_DTYPES[key]
+
+
 @contextmanager
 def disable_cuda_autocast():
     with torch.autocast(device_type="cuda", enabled=False):
@@ -715,7 +747,7 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
         )
 
     def _supports_flash_attention(self, device: torch.device, dtype: torch.dtype) -> bool:
-        return HAS_FLASH_ATTN and device.type == "cuda" and dtype in {torch.float16, torch.bfloat16}
+        return HAS_FLASH_ATTN and device.type == "cuda" and dtype == torch.bfloat16
 
     def _get_backend_check_dtype(self, x: torch.Tensor) -> torch.dtype:
         if x.device.type != "cuda":
@@ -1778,6 +1810,9 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         self.attention_implementation = config.attention_implementation
         self.compute_dtype_name = config.compute_dtype
         self.compute_dtype = resolve_compute_dtype(config.compute_dtype)
+        self.codec_weight_dtype_name = canonicalize_codec_weight_dtype(
+            getattr(config, "codec_weight_dtype", "fp32")
+        )
 
         encoder_context_durations = [
             float(module_kwargs.get("context_duration", config.causal_transformer_context_duration))
@@ -1848,6 +1883,21 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
 
         self.post_init()
 
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        codec_weight_dtype = kwargs.pop("codec_weight_dtype", None)
+        explicit_torch_dtype = kwargs.get("torch_dtype", None)
+        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        if codec_weight_dtype is None:
+            codec_weight_dtype = getattr(model.config, "codec_weight_dtype", "fp32")
+        if explicit_torch_dtype is not None and canonicalize_codec_weight_dtype(codec_weight_dtype) != "fp32":
+            logger.warning(
+                "`torch_dtype` was passed while `codec_weight_dtype` is enabled. "
+                "Prefer leaving codec `torch_dtype` unset so the fp32 quantizer weights are loaded without precision loss."
+            )
+        model.set_codec_weight_dtype(codec_weight_dtype)
+        return model
+
     def _start_streaming(self, batch_size: int):
         """Start streaming mode for all modules."""
 
@@ -1936,6 +1986,9 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         if device.type == "cuda" and self.compute_dtype is not None:
             with torch.autocast(device_type="cuda", dtype=self.compute_dtype):
                 yield
+        elif device.type == "cpu" and self.compute_dtype is torch.bfloat16:
+            with torch.autocast(device_type="cpu", dtype=self.compute_dtype):
+                yield
         else:
             yield
 
@@ -1948,6 +2001,37 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
     def set_compute_dtype(self, compute_dtype: str) -> None:
         self.compute_dtype_name = compute_dtype
         self.compute_dtype = resolve_compute_dtype(compute_dtype)
+        self.config.compute_dtype = compute_dtype
+
+    def set_codec_weight_dtype(self, codec_weight_dtype: str) -> None:
+        codec_weight_dtype = canonicalize_codec_weight_dtype(codec_weight_dtype)
+        weight_dtype = resolve_codec_weight_dtype(codec_weight_dtype)
+
+        self.encoder.to(dtype=weight_dtype)
+        self.decoder.to(dtype=weight_dtype)
+
+        # Quantizer decode/encode intentionally disables autocast and builds fp32 intermediates.
+        # Keeping it fp32 avoids fp32-input/bf16-bias mismatches and preserves codebook numerics.
+        self.quantizer.to(dtype=torch.float32)
+
+        self.codec_weight_dtype_name = codec_weight_dtype
+        self.config.codec_weight_dtype = codec_weight_dtype
+        if codec_weight_dtype != "fp32" and self.compute_dtype is None:
+            self.set_compute_dtype(codec_weight_dtype)
+
+    def get_codec_dtype_summary(self) -> dict[str, str]:
+        def _first_param_dtype(module: nn.Module) -> str:
+            for param in module.parameters():
+                return str(param.dtype)
+            return "no_params"
+
+        return {
+            "encoder": _first_param_dtype(self.encoder),
+            "decoder": _first_param_dtype(self.decoder),
+            "quantizer": _first_param_dtype(self.quantizer),
+            "compute_dtype": self.compute_dtype_name,
+            "codec_weight_dtype": self.codec_weight_dtype_name,
+        }
 
     def _prepare_waveform_batch(
         self,
@@ -2117,6 +2201,9 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
 
         quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer)
         _, audio_codes, audio_codes_lengths = quantizer(encoder_hidden_states.float(), encoder_hidden_lengths, n_quantizers)
+        max_valid_length = int(audio_codes_lengths.max().item()) if audio_codes_lengths.numel() > 0 else 0
+        audio_codes = audio_codes[:, :, :max_valid_length]
+        encoder_hidden_states = encoder_hidden_states[:, :, :max_valid_length]
 
         return MossAudioTokenizerEncoderOutput(
             audio_codes=audio_codes,