minimal MPS-specific bf16 guard

generator.py

@@ -8,14 +8,15 @@
 import torch
 import torchaudio
 from huggingface_hub import hf_hub_download
-from models import MISO_TTS_8B_CONFIG, Model, ModelArgs
+from models import MISO_TTS_8B_CONFIG, Model, ModelArgs, apply_mps_bf16_layer0_mlp_fp32_guard
 from moshi_compat import patch_bitsandbytes_import_for_unquantized_layers
 from moshi.models import loaders
 from tokenizers.processors import TemplateProcessing
 from transformers import AutoTokenizer
 from watermarking import MISO_TTS_WATERMARK, load_watermarker, watermark
 
 DEFAULT_MISO_TTS_REPO_ID = "MisoLabs/MisoTTS"
+DISABLE_MPS_BF16_LAYER0_MLP_FP32_ENV = "MISO_DISABLE_MPS_BF16_LAYER0_MLP_FP32"
 patch_bitsandbytes_import_for_unquantized_layers()
 
 
@@ -44,6 +45,18 @@ def load_llama3_tokenizer():
     return tokenizer
 
 
+def _should_apply_mps_bf16_layer0_mlp_fp32_guard(
+    device: str,
+    dtype: torch.dtype,
+    enabled: Optional[bool],
+) -> bool:
+    if enabled is not None:
+        return enabled
+
+    disabled = os.environ.get(DISABLE_MPS_BF16_LAYER0_MLP_FP32_ENV, "").lower() in {"1", "true", "yes"}
+    return torch.device(device).type == "mps" and dtype == torch.bfloat16 and not disabled
+
+
 class Generator:
     def __init__(
         self,
@@ -199,6 +212,7 @@ def _load_model(
     config: ModelArgs,
     device: str,
     dtype: torch.dtype,
+    mps_bf16_layer0_mlp_fp32: Optional[bool] = None,
 ) -> Model:
     if os.path.isfile(model_path_or_repo_id):
         model_file = model_path_or_repo_id
@@ -224,6 +238,8 @@ def _load_model(
         raise FileNotFoundError(f"Could not resolve model checkpoint: {model_path_or_repo_id}")
 
     model.to(device=device, dtype=dtype)
+    if _should_apply_mps_bf16_layer0_mlp_fp32_guard(device, dtype, mps_bf16_layer0_mlp_fp32):
+        apply_mps_bf16_layer0_mlp_fp32_guard(model)
     model.eval()
     return model
 
@@ -232,7 +248,14 @@ def load_miso_8b(
     device: str = "cuda",
     model_path_or_repo_id: Optional[str] = None,
     dtype: torch.dtype = torch.bfloat16,
+    mps_bf16_layer0_mlp_fp32: Optional[bool] = None,
 ) -> Generator:
     source = model_path_or_repo_id or os.environ.get("MISO_TTS_8B_MODEL", DEFAULT_MISO_TTS_REPO_ID)
-    model = _load_model(source, MISO_TTS_8B_CONFIG, device=device, dtype=dtype)
+    model = _load_model(
+        source,
+        MISO_TTS_8B_CONFIG,
+        device=device,
+        dtype=dtype,
+        mps_bf16_layer0_mlp_fp32=mps_bf16_layer0_mlp_fp32,
+    )
     return Generator(model)

models.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 import contextlib
 import io
+import types
 from typing import Tuple
 
 import torch
@@ -103,6 +104,29 @@ def _masked_cross_entropy(logits, targets, mask, vocab_size):
     return (losses * weights).sum() / total, total
 
 
+def _linear_fp32(linear: nn.Linear, x: torch.Tensor) -> torch.Tensor:
+    bias = None if linear.bias is None else linear.bias.float()
+    return F.linear(x.float(), linear.weight.float(), bias)
+
+
+def apply_mps_bf16_layer0_mlp_fp32_guard(model: nn.Module) -> None:
+    """Run only the first backbone MLP matmuls in FP32 without changing state_dict keys."""
+    mlp = model.backbone.layers[0].mlp
+    if getattr(mlp, "_miso_mps_bf16_layer0_mlp_fp32_guard", False):
+        return
+
+    def forward_fp32_guard(self, x: torch.Tensor) -> torch.Tensor:
+        orig_dtype = x.dtype
+        h = self.activation(_linear_fp32(self.w1, x))
+        if self.w3 is not None:
+            h = h * _linear_fp32(self.w3, x)
+        h = _linear_fp32(self.w2, h)
+        return h.to(dtype=orig_dtype)
+
+    mlp.forward = types.MethodType(forward_fp32_guard, mlp)
+    mlp._miso_mps_bf16_layer0_mlp_fp32_guard = True
+
+
 @dataclass
 class ModelArgs:
     backbone_flavor: str

tests/test_mps_bf16_layer0_mlp_guard.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+
+from generator import _should_apply_mps_bf16_layer0_mlp_fp32_guard
+from models import apply_mps_bf16_layer0_mlp_fp32_guard
+
+
+class RecordingActivation(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.last_dtype = None
+
+    def forward(self, x):
+        self.last_dtype = x.dtype
+        return x
+
+
+class TinyMlp(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.w1 = nn.Linear(3, 4, bias=False)
+        self.w2 = nn.Linear(4, 3, bias=False)
+        self.w3 = nn.Linear(3, 4, bias=False)
+        self.activation = RecordingActivation()
+
+    def forward(self, x):
+        h = self.activation(self.w1(x))
+        h = h * self.w3(x)
+        return self.w2(h)
+
+
+class TinyLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mlp = TinyMlp()
+
+
+class TinyBackbone(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layers = nn.ModuleList([TinyLayer(), TinyLayer()])
+
+
+class TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.backbone = TinyBackbone()
+
+
+def test_layer0_mlp_guard_uses_fp32_compute_and_preserves_state_dict_keys():
+    model = TinyModel().to(dtype=torch.bfloat16)
+    before_keys = list(model.state_dict().keys())
+
+    apply_mps_bf16_layer0_mlp_fp32_guard(model)
+    out = model.backbone.layers[0].mlp(torch.ones(2, 3, dtype=torch.bfloat16))
+
+    assert out.dtype == torch.bfloat16
+    assert model.backbone.layers[0].mlp.activation.last_dtype == torch.float32
+    assert model.backbone.layers[0].mlp.w1.weight.dtype == torch.bfloat16
+    assert list(model.state_dict().keys()) == before_keys
+
+
+def test_layer0_mlp_guard_patches_only_first_backbone_layer():
+    model = TinyModel()
+
+    apply_mps_bf16_layer0_mlp_fp32_guard(model)
+
+    assert model.backbone.layers[0].mlp._miso_mps_bf16_layer0_mlp_fp32_guard
+    assert not hasattr(model.backbone.layers[1].mlp, "_miso_mps_bf16_layer0_mlp_fp32_guard")
+
+
+def test_mps_bf16_guard_auto_enable_rules(monkeypatch):
+    monkeypatch.delenv("MISO_DISABLE_MPS_BF16_LAYER0_MLP_FP32", raising=False)
+
+    assert _should_apply_mps_bf16_layer0_mlp_fp32_guard("mps", torch.bfloat16, None)
+    assert not _should_apply_mps_bf16_layer0_mlp_fp32_guard("cpu", torch.bfloat16, None)
+    assert not _should_apply_mps_bf16_layer0_mlp_fp32_guard("cuda", torch.bfloat16, None)
+    assert not _should_apply_mps_bf16_layer0_mlp_fp32_guard("mps", torch.float32, None)
+
+    monkeypatch.setenv("MISO_DISABLE_MPS_BF16_LAYER0_MLP_FP32", "1")
+    assert not _should_apply_mps_bf16_layer0_mlp_fp32_guard("mps", torch.bfloat16, None)
+    assert _should_apply_mps_bf16_layer0_mlp_fp32_guard("cpu", torch.bfloat16, True)
+    assert not _should_apply_mps_bf16_layer0_mlp_fp32_guard("mps", torch.bfloat16, False)