MPS training support (Apple Silicon)#68
Open
Cortexelus wants to merge 2 commits into
Open
Conversation
added 2 commits
July 14, 2026 14:05
Makes the LoRA training path (the raw-PyTorch loop used by underfit) run
correctly on MPS. Two silent-failure classes fixed, not just device
plumbing:
- On CUDA-less builds, autocast('cuda') and GradScaler('cuda') are silent
no-ops — Mac 'training' ran with AMP and loss scaling secretly disabled.
- @autocast('cuda', enabled=False) does NOT disable an active MPS autocast,
so the RoPE and ExpoFourierFeatures fp32 islands were defeated under MPS
autocast (fp16 numerics where fp32 was intended).
New stable_audio_3/utils/device.py: resolve_device (cuda > mps > cpu),
autocast_context, make_grad_scaler (+ NoOpGradScaler fallback with a
printed note for older torch), device-aware disable_autocast, and
empty_device_cache. transformer.py/blocks.py fp32 islands now use
disable_autocast; gradient checkpointing re-enters autocast per device;
loading_utils/data device defaults resolve instead of assuming cuda;
interface empty_device_cache covers mps. CUDA behavior is bit-identical
(cuda checked first everywhere); the Lightning trainer stays CUDA-only by
design.
Verified on Apple Silicon (torch 2.13, macOS 15): fp16 and bf16 MPS
autocast and a full GradScaler('mps') cycle work; 6 smoke tests including
a real DiffusionTransformer + add_lora trained 3 steps on mps under
fp16 autocast + scaler with underfit's exact loss (loss decreasing, LoRA
params move, base bit-identical), and the full
ConditionedDiffusionModelWrapper path with underfit's call shape.
Cortexelus
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Apple Silicon (MPS) training support for the LoRA path. Companion PRs: dada-bots/underfit#9 (needs this merged first) and the stacked fast-forward PR based on this branch.
Two silent-failure classes fixed (not just device plumbing)
autocast("cuda")andGradScaler("cuda")are silent no-ops — Mac "training" ran with AMP and loss scaling secretly disabled.@autocast("cuda", enabled=False)does not disable an active MPS autocast, so the RoPE/ExpoFourierFeatures fp32 islands were silently defeated under MPS autocast.Changes
New
stable_audio_3/utils/device.py:resolve_device(cuda > mps > cpu),autocast_context,make_grad_scaler(+ no-op fallback with a printed note for older torch), device-awaredisable_autocast,empty_device_cache.transformer.py/blocks.pyfp32 islands usedisable_autocast; gradient checkpointing re-enters autocast per device; loading/data device defaults resolve instead of assuming cuda; interface cache clears cover mps. CUDA behavior is bit-identical (cuda checked first everywhere). The Lightning trainer remains CUDA-only by design; MPS training is the raw-PyTorch loop (underfit's).Verified on Apple Silicon (M4 Pro 48 GB, torch 2.13)
fp16 + bf16 MPS autocast and a full
GradScaler("mps")cycle; 6-part MPS smoke suite (incl. a realDiffusionTransformer+add_loratrained on mps: loss decreasing, LoRA params move, base bit-identical, and the full conditioned-wrapper call shape underfit uses); full 30-step underfit-loop runs on sm-music-base (0.668 it/s) and medium-base (0.254 it/s) at seq 256; demos (baseline + per-step, mp3 out); pre-encode--device mps; production crop lengths (sm@1300 16.3 s/step; medium@4096 ~20 s/step, 15.3 GiB peak → wants ≥32 GB, the raw loop has no grad checkpointing); config-path resume with metadata offsets. No MPS op fallbacks anywhere (fallback env not set — unsupported ops would have hard-errored).