diff --git a/stable_audio_3/data/utils.py b/stable_audio_3/data/utils.py index 9b5f68a..80b5d63 100644 --- a/stable_audio_3/data/utils.py +++ b/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. diff --git a/stable_audio_3/interface/diffusion_cond.py b/stable_audio_3/interface/diffusion_cond.py index f365122..7c60260 100644 --- a/stable_audio_3/interface/diffusion_cond.py +++ b/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}") diff --git a/stable_audio_3/loading_utils.py b/stable_audio_3/loading_utils.py index 775ef0a..20dc8ff 100644 --- a/stable_audio_3/loading_utils.py +++ b/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) diff --git a/stable_audio_3/models/blocks.py b/stable_audio_3/models/blocks.py index 32b6441..2b88fe7 100644 --- a/stable_audio_3/models/blocks.py +++ b/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. diff --git a/stable_audio_3/models/transformer.py b/stable_audio_3/models/transformer.py index 72989d6..637a6a2 100644 --- a/stable_audio_3/models/transformer.py +++ b/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 diff --git a/stable_audio_3/utils/__init__.py b/stable_audio_3/utils/__init__.py new file mode 100644 index 0000000..8e05a40 --- /dev/null +++ b/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", +] diff --git a/stable_audio_3/utils/device.py b/stable_audio_3/utils/device.py new file mode 100644 index 0000000..d730365 --- /dev/null +++ b/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() diff --git a/tests/test_mps_training_smoke.py b/tests/test_mps_training_smoke.py new file mode 100644 index 0000000..7cabcd7 --- /dev/null +++ b/tests/test_mps_training_smoke.py @@ -0,0 +1,334 @@ +"""MPS (Apple Silicon) LoRA-training smoke test. + +Verifies the device-neutral AMP plumbing added for MPS training support: + + - stable_audio_3.utils.device: resolve_device / autocast_context / + make_grad_scaler on "mps" + - fp32 islands (RoPE, ExpoFourierFeatures) actually stay fp32 under an + active MPS autocast (the old @autocast("cuda", enabled=False) decorators + silently did nothing on MPS) + - a real (tiny) DiffusionTransformer with LoRA adapters runs 3 + forward/backward/AdamW steps on "mps" under fp16 autocast + GradScaler, + using underfit's rectified-flow signal-only masked-MSE loss shape; + LoRA params move, base params don't. + +Skipped entirely when MPS is unavailable. Runnable standalone: + python tests/test_mps_training_smoke.py +""" + +import os +import sys +from functools import partial + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +HAS_MPS = torch.backends.mps.is_available() +pytestmark = pytest.mark.skipif(not HAS_MPS, reason="MPS not available on this machine") + +DEVICE = "mps" + + +# --------------------------------------------------------------------------- +# underfit loss (real module when importable, faithful mirror otherwise) +# --------------------------------------------------------------------------- + + +def _underfit_loss_fns(): + try: + from underfit.training.loss import compute_masked_loss, compute_normalized_mse + + return compute_normalized_mse, compute_masked_loss + except ImportError: + pass + + # Mirror of underfit/training/loss.py (loss_normalization="none", + # mask_padding_attention=True → signal-only masked MSE). + def compute_normalized_mse( + pred, target, loss_mask, loss_normalization="none", loss_norm_eps=1e-6 + ): + return (pred - target) ** 2 + + def compute_masked_loss( + loss_full, loss_mask, mask_padding_attention, mask_loss_weight=0.0 + ): + signal = torch.where(loss_mask.unsqueeze(1), loss_full, 0.0) + signal_sum = signal.sum(dim=(1, 2)) + n_channels = loss_full.shape[1] + signal_count = loss_mask.sum(dim=1) * n_channels + per_sample_loss = signal_sum / (signal_count + 1e-8) + loss = per_sample_loss.mean() + signal_mean = signal.sum() / (signal_count.sum() + 1e-8) + return loss, signal_mean.detach(), torch.zeros(()) + + return compute_normalized_mse, compute_masked_loss + + +# --------------------------------------------------------------------------- +# Device helper probes +# --------------------------------------------------------------------------- + + +def test_resolve_device_prefers_mps_without_cuda(): + from stable_audio_3.utils.device import resolve_device + + expected = "cuda" if torch.cuda.is_available() else "mps" + assert resolve_device() == expected + assert resolve_device("cpu") == "cpu" + + +def test_autocast_context_fp16_on_mps(): + from stable_audio_3.utils.device import autocast_context + + a = torch.randn(8, 8, device=DEVICE) + with autocast_context(DEVICE, dtype=torch.float16): + assert torch.is_autocast_enabled(DEVICE) + out = a @ a + assert out.dtype == torch.float16 + + +def test_grad_scaler_on_mps_full_cycle(): + from stable_audio_3.utils.device import autocast_context, make_grad_scaler + + scaler = make_grad_scaler(DEVICE) + assert scaler.is_enabled(), "torch 2.13 GradScaler('mps') should be enabled" + + lin = torch.nn.Linear(4, 4).to(DEVICE) + opt = torch.optim.SGD(lin.parameters(), lr=0.1) + with autocast_context(DEVICE, dtype=torch.float16): + loss = lin(torch.randn(2, 4, device=DEVICE)).square().mean() + scaler.scale(loss).backward() + scaler.unscale_(opt) + scaler.step(opt) + scaler.update() + assert scaler.get_scale() > 0 + + +def test_fp32_islands_hold_under_mps_autocast(): + """RoPE + ExpoFourierFeatures must compute in fp32 inside MPS autocast.""" + from stable_audio_3.models.blocks import ExpoFourierFeatures + from stable_audio_3.models.transformer import RotaryEmbedding + + rope = RotaryEmbedding(16).to(DEVICE) + eff = ExpoFourierFeatures(32).to(DEVICE) + with torch.amp.autocast(DEVICE, dtype=torch.float16): + freqs, _ = rope.forward_from_seq_len(8) + feats = eff(torch.rand(4, device=DEVICE)) + # The einsum in RotaryEmbedding.forward is autocast-eligible; fp32 output + # proves the disable_autocast island took effect. + assert freqs.dtype == torch.float32 + assert feats.dtype == torch.float32 + + +# --------------------------------------------------------------------------- +# End-to-end: tiny DiT + LoRA, 3 training steps on MPS +# --------------------------------------------------------------------------- + + +def _build_tiny_dit_with_lora(): + from stable_audio_3.models.dit import DiffusionTransformer + from stable_audio_3.models.lora import LoRAParametrization, add_lora + + torch.manual_seed(0) + model = DiffusionTransformer( + io_channels=8, + embed_dim=64, + depth=2, + num_heads=2, + cond_token_dim=32, + global_cond_dim=16, + transformer_type="continuous_transformer", + diffusion_objective="rectified_flow", + ).to(DEVICE) + + lora_cfg = { + torch.nn.Linear: { + "weight": partial( + LoRAParametrization.from_linear, + rank=4, + lora_alpha=4.0, + adapter_type="lora", + ), + }, + } + add_lora(model, lora_cfg) + + lora_params, base_params = [], [] + for name, p in model.named_parameters(): + if "lora_" in name: + lora_params.append((name, p)) + else: + base_params.append((name, p)) + assert lora_params, "add_lora attached no LoRA parameters" + + # Mirror underfit/training/loop.py: only LoRA params train. + for _, p in base_params: + p.requires_grad_(False) + for _, p in lora_params: + p.data = p.data.float() + p.requires_grad_(True) + return model, lora_params, base_params + + +def test_lora_training_steps_on_mps(): + from stable_audio_3.utils.device import autocast_context, make_grad_scaler + + compute_normalized_mse, compute_masked_loss = _underfit_loss_fns() + model, lora_params, base_params = _build_tiny_dit_with_lora() + + lora_before = {n: p.detach().clone() for n, p in lora_params} + base_before = {n: p.detach().clone() for n, p in base_params} + + opt = torch.optim.AdamW([p for _, p in lora_params], lr=1e-2) + scaler = make_grad_scaler(DEVICE) + + # Fixed batch so loss trajectory is meaningful across steps. + torch.manual_seed(1) + B, C, T = 2, 8, 32 + x = torch.randn(B, C, T, device=DEVICE) + cross = torch.randn(B, 24, 32, device=DEVICE) + glob = torch.randn(B, 16, device=DEVICE) + t = torch.rand(B, device=DEVICE) * 0.8 + 0.1 + noise = torch.randn_like(x) + loss_mask = torch.ones(B, T, dtype=torch.bool, device=DEVICE) + + # rectified_flow noising (as in underfit/training/loop.py) + alphas, sigmas = (1 - t)[:, None, None], t[:, None, None] + noised = x * alphas + noise * sigmas + target = noise - x + + losses = [] + for _step in range(3): + with autocast_context(DEVICE, dtype=torch.float16): + out = model(noised, t, cross_attn_cond=cross, global_embed=glob) + mse_full = compute_normalized_mse(out, target, loss_mask) + loss, _sig, _pad = compute_masked_loss( + mse_full, loss_mask, mask_padding_attention=True + ) + opt.zero_grad() + scaler.scale(loss).backward() + scaler.step(opt) + scaler.update() + losses.append(loss.item()) + + assert all(x == x and abs(x) != float("inf") for x in losses), ( + f"non-finite loss: {losses}" + ) + assert losses[-1] != losses[0], f"loss did not change over 3 steps: {losses}" + # On a fixed batch with lr=1e-2 the loss should trend down. + assert losses[-1] < losses[0] * 1.05, f"loss did not decrease: {losses}" + + lora_changed = any( + not torch.equal(p.detach(), lora_before[n]) for n, p in lora_params + ) + assert lora_changed, "no LoRA parameter changed after 3 optimizer steps" + + for n, p in base_params: + assert torch.equal(p.detach(), base_before[n]), f"base param {n} changed" + + +def test_full_wrapper_training_step_on_mps(): + """The exact call shape underfit's loop uses, on the full SA3 wrapper: + conditioner(metadata, device) -> model(noised, t, cond=..., cfg_dropout_prob=...). + + Uses a NumberConditioner (no checkpoint downloads) and pretransform=None + (underfit's pre_encoded path never calls pretransform.encode). + """ + from stable_audio_3.models.conditioners import MultiConditioner, NumberConditioner + from stable_audio_3.models.diffusion import ( + ConditionedDiffusionModelWrapper, + DiTWrapper, + ) + from stable_audio_3.models.lora import LoRAParametrization, add_lora + from stable_audio_3.utils.device import autocast_context, make_grad_scaler + + compute_normalized_mse, compute_masked_loss = _underfit_loss_fns() + + torch.manual_seed(0) + dit = DiTWrapper( + diffusion_objective="rectified_flow", + io_channels=8, + embed_dim=64, + depth=2, + num_heads=2, + cond_token_dim=32, + global_cond_dim=32, + transformer_type="continuous_transformer", + ) + conditioner = MultiConditioner( + {"seconds_total": NumberConditioner(output_dim=32, min_val=0, max_val=512)} + ) + model = ConditionedDiffusionModelWrapper( + dit, + conditioner, + io_channels=8, + sample_rate=44100, + min_input_length=1, + diffusion_objective="rectified_flow", + pretransform=None, + cross_attn_cond_ids=["seconds_total"], + global_cond_ids=["seconds_total"], + ).to(DEVICE) + + add_lora( + model.model, + { + torch.nn.Linear: { + "weight": partial( + LoRAParametrization.from_linear, + rank=4, + lora_alpha=4.0, + adapter_type="lora", + ), + }, + }, + ) + lora_params = [p for n, p in model.named_parameters() if "lora_" in n] + assert lora_params + for p in model.parameters(): + p.requires_grad_(False) + for p in lora_params: + p.data = p.data.float() + p.requires_grad_(True) + + opt = torch.optim.AdamW(lora_params, lr=1e-2) + scaler = make_grad_scaler(DEVICE) + + torch.manual_seed(2) + B, C, T = 2, 8, 32 + metadata = [{"seconds_total": 30.0}, {"seconds_total": 47.5}] + diffusion_input = torch.randn(B, C, T, device=DEVICE) + loss_mask = torch.ones(B, T, dtype=torch.bool, device=DEVICE) + + losses = [] + for _step in range(3): + with autocast_context(DEVICE, dtype=torch.float16): + conditioning = model.conditioner(metadata, DEVICE) + t = torch.rand(B, device=DEVICE) * 0.8 + 0.1 + alphas, sigmas = (1 - t)[:, None, None], t[:, None, None] + noise = torch.randn_like(diffusion_input) + noised = diffusion_input * alphas + noise * sigmas + target = noise - diffusion_input + output = model(noised, t, cond=conditioning, cfg_dropout_prob=0.1) + mse_full = compute_normalized_mse(output, target, loss_mask) + loss, _sig, _pad = compute_masked_loss( + mse_full, loss_mask, mask_padding_attention=True + ) + opt.zero_grad() + scaler.scale(loss).backward() + scaler.step(opt) + scaler.update() + losses.append(loss.item()) + + assert all(x == x and abs(x) != float("inf") for x in losses), ( + f"non-finite loss: {losses}" + ) + assert len(set(losses)) > 1, f"loss frozen across steps: {losses}" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v", "-x"]))