From de9c6fa683fe972f1443fce6f9b3a96ab7240514 Mon Sep 17 00:00:00 2001 From: robbiemu Date: Thu, 4 Jun 2026 19:29:55 -0400 Subject: [PATCH] minimal MPS-specific guard --- generator.py | 27 +++++++- models.py | 24 +++++++ tests/test_mps_bf16_layer0_mlp_guard.py | 83 +++++++++++++++++++++++++ 3 files changed, 132 insertions(+), 2 deletions(-) create mode 100644 tests/test_mps_bf16_layer0_mlp_guard.py diff --git a/generator.py b/generator.py index 8441cee..64c29b6 100644 --- a/generator.py +++ b/generator.py @@ -8,7 +8,7 @@ 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 @@ -16,6 +16,7 @@ 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) diff --git a/models.py b/models.py index 914ceee..d279da7 100644 --- a/models.py +++ b/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 diff --git a/tests/test_mps_bf16_layer0_mlp_guard.py b/tests/test_mps_bf16_layer0_mlp_guard.py new file mode 100644 index 0000000..caefd5d --- /dev/null +++ b/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)