MPS training support (Apple Silicon)

stable_audio_3/data/utils.py

@@ -216,7 +216,7 @@ def compute_effective_seq_len_from_conditioning(
     conditioning: list,
     sample_rate: int,
     downsampling_ratio: int = 1,
-    device: str = "cuda"
+    device: str = None
 ) -> Optional[torch.Tensor]:
     """
     Compute effective sequence lengths from seconds_total in conditioning dicts.

stable_audio_3/interface/diffusion_cond.py

@@ -94,6 +94,8 @@ def generate_cond(
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+    elif torch.backends.mps.is_available():
+        torch.mps.empty_cache()
     gc.collect()
 
     print(f"Prompt: {prompt}")

stable_audio_3/loading_utils.py

@@ -7,6 +7,7 @@
     create_autoencoder_from_config,
     create_diffusion_cond_from_config,
 )
+from stable_audio_3.utils.device import resolve_device
 
 
 def copy_state_dict(model, state_dict):
@@ -61,9 +62,10 @@ def load_autoencoder(config_path: str, ckpt_path: str, device: str = "cpu"):
 def load_diffusion_cond(
     model_config,
     ckpt_path: str,
-    device: str = "cuda",
+    device: str = None,
     model_half: bool = False,
 ):
+    device = resolve_device(device)
     model = create_diffusion_cond_from_config(model_config)
     copy_state_dict(model, load_file(ckpt_path))
     model.to(device).eval().requires_grad_(False)

stable_audio_3/models/blocks.py

@@ -3,6 +3,8 @@
 from torch import nn
 from torch.nn.utils import weight_norm
 
+from ..utils.device import disable_autocast
+
 
 def get_activation(activation, channels=None) -> nn.Module:
     if activation == "elu":
@@ -54,7 +56,7 @@ def __init__(self, dim, min_freq=0.5, max_freq=10000.0):
         self.min_freq = min_freq
         self.max_freq = max_freq
 
-    @torch.amp.autocast("cuda",enabled=False)
+    @disable_autocast
     def forward(self, t):
         """
         t: [B] tensor.

stable_audio_3/models/transformer.py

@@ -72,6 +72,7 @@ def precompute_varlen_metadata(padding_mask: torch.Tensor):
     }
 
 from .utils import compile
+from ..utils.device import disable_autocast
 
 
 def _left_pad_to_match(emb, target_len):
@@ -161,19 +162,25 @@ def _sliding_window_chunked_halo_sdpa(q, k, v, w_left, w_right, chunk_size=_SLID
 
 def checkpoint(function, *args, **kwargs):
     kwargs.setdefault("use_reentrant", False)
-    # Preserve autocast context during recomputation to avoid dtype mismatches
+    # Preserve autocast context during recomputation to avoid dtype mismatches.
+    # Device-generic: torch.is_autocast_enabled() with no args only checks CUDA,
+    # which silently skips this on MPS/CPU autocast.
     if "context_fn" not in kwargs:
         from torch.amp import autocast
-        import functools
-        # Get current autocast state
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype('cuda')
-            def get_contexts():
-                return (
-                    autocast('cuda', dtype=dtype),
-                    autocast('cuda', dtype=dtype),
-                )
-            kwargs["context_fn"] = get_contexts
+        for _device_type in ("cuda", "mps", "cpu"):
+            try:
+                _ac_enabled = torch.is_autocast_enabled(_device_type)
+            except (RuntimeError, ValueError, TypeError):
+                _ac_enabled = False
+            if _ac_enabled:
+                dtype = torch.get_autocast_dtype(_device_type)
+                def get_contexts(_device_type=_device_type, dtype=dtype):
+                    return (
+                        autocast(_device_type, dtype=dtype),
+                        autocast(_device_type, dtype=dtype),
+                    )
+                kwargs["context_fn"] = get_contexts
+                break
     return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
 
 
@@ -273,7 +280,7 @@ def forward_from_seq_len(self, seq_len):
         t = torch.arange(seq_len, device = device)
         return self.forward(t)
 
-    @autocast("cuda", enabled = False)
+    @disable_autocast
     def forward(self, t):
         device = self.inv_freq.device
 
@@ -299,7 +306,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim = -1)
 
 
-@autocast("cuda", enabled = False)
+@disable_autocast
 def apply_rotary_pos_emb(t, freqs, scale = 1):
     out_dtype = t.dtype
 

stable_audio_3/utils/__init__.py

@@ -0,0 +1,15 @@
+from .device import (
+    autocast_context,
+    disable_autocast,
+    empty_device_cache,
+    make_grad_scaler,
+    resolve_device,
+)
+
+__all__ = [
+    "autocast_context",
+    "disable_autocast",
+    "empty_device_cache",
+    "make_grad_scaler",
+    "resolve_device",
+]

stable_audio_3/utils/device.py

@@ -0,0 +1,155 @@
+"""Device-neutral helpers for CUDA / MPS / CPU.
+
+Single-device training and inference should route device selection and AMP
+(autocast / GradScaler) construction through this module instead of hardcoding
+"cuda". Multi-GPU / DeepSpeed / Lightning-strategy code is out of scope and
+intentionally untouched.
+
+Empirically verified on torch 2.13.0 + macOS 15 (Apple Silicon):
+  - torch.amp.autocast("mps", dtype=torch.float16) works
+  - torch.amp.autocast("mps", dtype=torch.bfloat16) works (macOS 14+ required)
+  - torch.amp.GradScaler("mps") works (scale/unscale_/step/update)
+  - @autocast("cuda", enabled=False) does NOT disable an active MPS autocast,
+    so fp32 islands need the device-aware `disable_autocast` below.
+"""
+
+import functools
+
+import torch
+
+
+def resolve_device(preference=None) -> str:
+    """Best available device string: explicit preference > cuda > mps > cpu."""
+    if preference:
+        return preference
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+@functools.lru_cache(maxsize=None)
+def mps_bf16_supported() -> bool:
+    """bf16 on MPS needs macOS 14+; probe once with a tiny op."""
+    if not torch.backends.mps.is_available():
+        return False
+    try:
+        x = torch.ones(2, 2, device="mps", dtype=torch.bfloat16)
+        (x @ x).sum().item()
+        return True
+    except Exception:
+        return False
+
+
+def autocast_context(device_type=None, dtype=None, enabled=True):
+    """torch.amp.autocast with per-backend fixups.
+
+    - device_type may be a device string ("cuda:0", "mps") or None (auto-resolve).
+    - bf16 on MPS is downgraded to fp16 (with a printed note) when the OS
+      doesn't support it.
+    """
+    dt = torch.device(device_type or resolve_device()).type
+    if dt == "mps" and dtype is torch.bfloat16 and not mps_bf16_supported():
+        print(
+            "[device] bf16 autocast unsupported on this macOS/MPS build; "
+            "falling back to fp16 autocast",
+            flush=True,
+        )
+        dtype = torch.float16
+    if dt == "cpu":
+        # Training code historically used autocast("cuda"), a silent no-op on
+        # CPU-only hosts. Keep CPU runs in fp32 rather than newly enabling CPU
+        # fp16/bf16 autocast (slow and numerically different).
+        enabled = False
+    return torch.amp.autocast(dt, dtype=dtype, enabled=enabled)
+
+
+class NoOpGradScaler:
+    """Transparent stand-in when torch.amp.GradScaler doesn't support a backend."""
+
+    def scale(self, loss):
+        return loss
+
+    def unscale_(self, optimizer):
+        pass
+
+    def step(self, optimizer, *args, **kwargs):
+        return optimizer.step(*args, **kwargs)
+
+    def update(self, new_scale=None):
+        pass
+
+    def get_scale(self):
+        return 1.0
+
+    def is_enabled(self):
+        return False
+
+    def state_dict(self):
+        return {}
+
+    def load_state_dict(self, state_dict):
+        pass
+
+
+def make_grad_scaler(device_type=None, enabled=True):
+    """GradScaler for the given backend.
+
+    torch 2.13 supports GradScaler("mps") natively (verified). On CPU, or if
+    construction fails (older torch), returns a disabled/no-op scaler so
+    callers can use the scaler API unconditionally.
+    """
+    dt = torch.device(device_type or resolve_device()).type
+    if dt == "cpu":
+        # fp16 grad scaling is pointless on CPU; keep the API but disabled
+        # (matches the old GradScaler("cuda")-on-cpu auto-disable behavior).
+        enabled = False
+    try:
+        return torch.amp.GradScaler(dt, enabled=enabled)
+    except Exception as e:
+        if enabled:
+            print(
+                f"[device] torch.amp.GradScaler({dt!r}) unsupported "
+                f"({type(e).__name__}: {e}); using no-op scaler — fp16 loss "
+                "scaling disabled, watch for gradient underflow",
+                flush=True,
+            )
+        return NoOpGradScaler()
+
+
+def _first_tensor_device_type(args, kwargs):
+    for a in args:
+        if isinstance(a, torch.Tensor):
+            return a.device.type
+    for a in kwargs.values():
+        if isinstance(a, torch.Tensor):
+            return a.device.type
+    return None
+
+
+def disable_autocast(fn):
+    """Device-aware replacement for @autocast("cuda", enabled=False).
+
+    The cuda-pinned decorator does NOT disable an active MPS (or CPU) autocast
+    context, silently defeating fp32 islands (RoPE, Fourier timestep features)
+    on non-CUDA backends. This wrapper disables autocast for the device the
+    inputs actually live on.
+    """
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        dt = _first_tensor_device_type(args, kwargs) or "cuda"
+        with torch.amp.autocast(dt, enabled=False):
+            return fn(*args, **kwargs)
+
+    return wrapper
+
+
+def empty_device_cache(device_type=None):
+    """Release cached allocator memory on the active accelerator, if any."""
+    dt = torch.device(device_type or resolve_device()).type
+    if dt == "cuda" and torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    elif dt == "mps" and torch.backends.mps.is_available():
+        torch.mps.empty_cache()