minimal MPS-specific bf16 guard

[robbiemu]

Implemented a minimal MPS BF16 guard.

edit: I just wrapped up a long-format generation and listen test, and sadly this is not bit-perfect. This PR should be treated as a targeted MPS BF16 compatibility improvement, not a complete replacement for full FP32 MPS fallback. In my tests this is close enough to keep it from switching voices, but not identical renderings.

Reason:
MPS BF16 generation can diverge from CPU/CUDA BF16 at a narrow early-backbone decision boundary. In a focused diagnostic, CPU BF16 and CUDA BF16 produced an exact codebook1 tie, while MPS BF16 gave one candidate a one-step advantage. CPU FP32 preferred the same candidate as MPS BF16, but with a smaller gap, so this appears to be a BF16 backend boundary mismatch rather than a broad model failure. The divergence was traced to the first backbone block, and the smallest tested precision island that restored CPU/CUDA BF16 behavior was FP32 compute for only backbone.layers[0].mlp. summary

This PR keeps the bulk model in BF16 on MPS, but runs that one MLP path in FP32 and casts the output back to BF16 before returning to the normal residual path. This is intended to preserve MPS BF16 memory behavior while avoiding the observed first-layer parity issue.

Scope:
This does not change CPU or CUDA behavior, and the guard can be disabled with MISO_DISABLE_MPS_BF16_LAYER0_MLP_FP32=1 for A/B testing. Its a narrow fix:

MPS BF16 bulk path
+ scoped backbone.layer.00.mlp FP32 compute only
+ cast MLP output back to BF16 before residual

Changed:

• models.py: adds apply_mps_bf16_layer0_mlp_fp32_guard, patching only model.backbone.layers[0].mlp.forward.
• generator.py: auto-enables the guard only for device=mps + dtype=torch.bfloat16, with disable env var MISO_DISABLE_MPS_BF16_LAYER0_MLP_FP32=1, plus an explicit optional override.
• tests/test_mps_bf16_layer0_mlp_guard.py: focused unit coverage for FP32 compute, BF16 output, state_dict key preservation, layer-0-only patching, and auto-enable rules.

Local validation:

• rtk .venv/bin/python -m pytest tests/test_mps_bf16_layer0_mlp_guard.py -q passed: 3 passed
• rtk uv run --project . python -m pytest tests/test_mps_bf16_layer0_mlp_guard.py -q passed: 3 passed
• Tiny MPS BF16 runtime check passed: BF16 output on MPS, BF16 stored weights, finite output.
• Focused current-branch boundary validation passed with guard applied:
  • argmax_token = 316
  • logit[316] = 17.875
  • logit[1056] = 17.875
  • gap = 0.0
  • ties = 2

Results:
The overhead numbers are also important: the focused diagnostic I linked to above went from about 59.4s baseline to 60.5s with layer00_mlp_fp32_only, roughly +1.9%, and the extra storage cost for keeping FP32 copies of w1/w2/w3 is about 336 MiB.

The specific numerical boundary we isolated is resolved by the minimal MPS BF16 guard: it restores the observed tie behavior and keeps production outputs decode-valid. However, end-to-end production listening still exposes quality issues, including issues also present on CPU BF16, so this PR does not resolve the full generation/audio-quality problem. This PR should be treated as an experimental targeted MPS BF16 compatibility improvement, not a complete replacement for full FP32 MPS fallback.

Development note: ChatGPT was used to help plan diagnostics, interpret results, and draft this PR description.

probably: Fixes #1

[robbiemu]

There is also a usability branch I could add as a followup pr: https://github.com/robbiemu/MisoTTS/tree/mps-usability

here's what that accomplishes:

Changed:

• run_misotts.py: selects CUDA when available, then MPS, then CPU. This lets Apple Silicon users run the bundled example without editing the script.
• generator.py: adds _watermarker_device() and uses CPU for the watermarker only when the model device is MPS. SilentCipher’s MPS watermark encode path hit NotImplementedError: aten::unfold_backward inside torch.istft, so CPU avoids that unsupported MPS op. This keeps the decision at the generator integration point instead of changing load_watermarker() globally, which avoids changing behavior for the standalone watermark CLI or external users who import that helper directly.
• tests/test_mps_bf16_layer0_mlp_guard.py: adds focused coverage that the watermarker device maps MPS to CPU while leaving CPU and CUDA unchanged.
• README.md: updates the examples and deployment note so the docs no longer describe a CUDA/CPU-only path.