MLX LoRA training for SA3 (Apple Silicon)#72
Merged
Conversation
Add pure-MLX LoRA, DoRA, BoRA, and XS adapter injection, checkpoint interoperability, fixed-strength inference support, waveform-to-SAME-latent encoding, SA3 timestep sampling, distribution shifting, rectified-flow loss, and focused parity tests.
Complete inventory of every training convention and default in underfit (adapter config, checkpoint key roots and metadata, loss/timesteps/CFG dropout, optimizer and loop mechanics, pre-encode and dataset semantics, prompt augmentation, demos), cross-referenced to source, plus the gap analysis against the MLX primitives this branch starts from and the build order. Working reference for the branch — fold into README before the PR finalizes.
…apter, defaults Checkpoint keys now use underfit's roots (model. prefix for DiT layers via a checkpoint_prefix on inject; conditioners.seconds_total.embedder. embedding.1 for the seconds conditioner) — verified key-for-key and shape-for-shape against a real underfit checkpoint. include/exclude filters match both bare runtime names and checkpoint names so dashboard filter strings work verbatim. New: underfit_lora_config() (dora-rows, rank 16, alpha=rank, the dashboard exclude list), inject_from_lora_config (config-layer fallbacks rank 8/alpha=rank/lora), TrainableSecondsEmbedder (bit-exact vs the inference embedder at init), load_trainable_lora_state for resume (strict=False, DoRA magnitude squeeze). 24 tests.
…ntics pre_encode_mlx.py: offline audio→latent CLI (whole files ≤600 s aligned to 4096, fp32 SAME encoders, chunked >30 s, no normalization; npy + json with seconds_total/seconds_start/audio_samples/latent_shape/padding_mask and sidecar tag extraction — underfit field parity). latent_dataset.py: PreEncodedDataset port (latent_crop_length crop with random_crop over the mask-valid region, silence/zero padding, seconds_total kept at FULL duration after cropping, min/max-length rejection redraw) + prompt_templates port (tag prompts with the 50/50 shuffle-vs-subset augmentation, trigger token at 80%, path/fixed sources, legacy fallback) + iterate_batches with the oversample-with-replacement x100 tiny-dataset behavior. 31 tests, all synthetic fixtures.
Step-driven LoRA training on Apple Silicon mirroring underfit's raw loop:
rectified-flow velocity target with signal-only masked MSE (per-sample
then mean), uniform timestep sampler + the SA3 'full' distribution shift
(256/4096), CFG dropout 0.1 (whole cross_attn to zeros per sample inside
the grad scope, global_cond kept), AdamW at torch defaults with required
--lr (weight decay explicitly 0.0 — MLX's default differs), fp32 adapters
over a frozen fp16 base, checkpoints every 1000 steps + final
({run}-step={s}-epoch={e}.safetensors with step/epoch/base_model
metadata), resume ladder (flags > metadata > filename), per-step
loss_by_timestep.bin telemetry, <save>/<name>/<uuid8>/checkpoints layout.
E2E verified on sm-music: pre-encode → 20 steps dora-rows r16 (140 DiT
layers = the underfit per-block set + conditioner, 9.2M trainable,
1.6 it/s at T=256) → checkpoint passes underfit's lora_validate contract →
resume restores 141 layers with metadata offsets → checkpoint loads
directly via sa3_mlx.py --lora (141 layers merged).
…ventions --dit-weights flag: train on the BASE (rectified-flow) checkpoint like underfit does — inference uses the ARC weights, training must not (loud warning otherwise). Conversion recipe for the HF *-base safetensors into the 441-key MLX npz documented in TRAINING_CONVENTIONS.md. Training forward now feeds the diffusion_cond_inpaint pure-generation conditioning (all-ones inpaint mask + zero context) instead of the inference path's zeros — verified against underfit's torch loop with a controlled forward: loss 3.9829 vs 3.9823, prediction PSNR 80.7 dB (fp16-weight bound; torch MPS-vs-CPU 3e-6 rel). Using inference-style zeros trains against the wrong conditioning regime (4x loss difference). Cross-attention correctly runs over all 257 padding-embedded tokens at training time (parity-tested; masking/slicing to valid tokens is wrong). Benchmark (M4 Pro, 30 steps, identical setup, full underfit loop on MPS via the mps-training/mps-support branches): MLX 1.65 it/s vs MPS 0.668 it/s — MLX 2.47x faster. Loss regimes match; both checkpoints pass underfit validation with identical key sets.
dora-rows/cols and the -xs variants (and lora-xs) no longer build the full [out,in] adapted weight per call. dora-rows is exactly a row-scale of the (base + low-rank) output with the norm computed in closed form: y = (x@W0.T + s·(x@A.T)@B.T) ⊙ m/rownorm + bias rownorm² = Σ_row W0² (cached const) + 2s·rowsum((W0@A.T)⊙B) + s²·rowsum((B@AAᵀ)⊙B) dora-cols is the input-feature-scale dual. One fp16 read of W0 through a rank-r matmul replaces ~5 fp32 full-matrix passes + an fp32 GEMM; the base matmul runs in native dtype. bora keeps the full-weight path (its nested col-norm has no rank-r expansion without storing W0²); SA3_LORA_NAIVE_DORA=1 restores the old path for ablation. Equivalence proven: outputs ≤9e-6 rel fp32, gradients ≤2.4e-7 wrt A/B/magnitude/M_xs vs the naive path; 55 tests. Layer micro-bench (12288x1536 dora-rows fwd+bwd): 7.2x faster. Full training step: sm-music 1.60→3.15 it/s, medium 0.52→~1.2 it/s, and peak training memory 9.8→5.5 GB (small) / 26.8→12.9 GB (medium) since the fp32 weight copies are gone.
…d-checkpoint Train step (loss+grad+optimizer update) is mx.compile'd with [bundle.state, optimizer.state, mx.random.state] capture (--no-compile for eager; optimizer.init up front so captured state is stable; the frozen-seconds branch hoisted out of the step for purity). T5Gemma prompt conditioning cached per exact prompt string (--no-t5-cache; bit-exact, ~100% hit rate in the tiny-dataset workflow). Wired memory limit raised to the device recommendation. --grad-checkpoint ports mlx-lm's grad_checkpoint onto the DiT blocks: bit-exact losses, −19%+ peak memory, ~1.4x step cost — for big crops on small Macs. Compiled-vs-eager losses agree to 8e-7 rel; save/resume verified under compile. Combined ablation (30 same-seed steps, crop 256): sm-music 1.60→3.78 it/s and medium 0.52→1.29 it/s at ≤3e-4 max loss deviation; vs the MPS/underfit loop that is 5.7x (small) and 5.1x (medium).
BoRA/BoRA-XS training forwards join the no-materialization path: y = (((x⊙β) @ W0.T + s·(((x⊙β)@A.T)@B.T)) ⊙ α) + bias α reuses the dora-rows closed-form rownorm; β's colnorm of the row-rescaled intermediate expands into three cheap terms, the W0²-weighted one via a cached native-dtype W0⊙W0 — the speed-mode memory cost (+1 copy of adapted weights, e.g. ~2.6 GiB for a full medium bora injection; the underfit default config adapts no convs and bora is rare, so usually 0). --bora-mode memory (trainer flag, threaded through inject_*) keeps the exact old full-weight path with no cache; SA3_LORA_NAIVE_DORA=1 still forces naive everywhere for ablation. Equivalence: fp32 forward ≤1.4e-6 abs, grads ≤8.5e-5 rel (incl. magnitude_r/magnitude_c/M_xs); fp16 within the shared tolerances with no loosening. Layer micro-bench (12288x1536 bora, fwd+bwd): 5.8x vs the full-weight path. 71 tests.
Add underfit-style demos to lora_train_mlx.py: baseline at step 0 then every --demo-every steps + a final render. models/defs/demo_mlx.py runs MLX inference on the base rectified_flow model + the trained LoRA — plain Euler velocity integration (x += dt*v, distinct from the distilled rf_denoiser pingpong), the same conditioning/CFG as sa3_mlx (uncond = cross->zeros, denoised-space guidance, local_add_cond=None), decode via SAME-S/SAME-L, peak-normalized int16 -> mp3 with a json sidecar, files demo_<i>_<step:08d>.mp3, idempotent per step. Per-entry cfg/seed/steps/duration; lora_strength and lora_interval_max (sigma gating) are applied by scaling lora_B (M_xs for -xs) pre-norm, which is exactly underfit's lora_strength semantics for lora/DoRA and touches no forward code. sa3_gradio.py: --lora PATH (repeatable, preloads a slot active-on-launch) and --port, so underfit can open the MLX gradio with a trained checkpoint loaded.
… betas/wd, InverseLR) The underfit SA3 training templates set conventions the trainer didn't honor: - use_effective_length_for_schedule=True: shift timesteps by the effective latent length ceil(int(seconds_total*44100)/4096) per sample, not the crop length (--use-effective-length). - AdamW betas [0.9,0.95] + weight_decay 0.01 (--beta1/--beta2/--eps/--weight-decay; MLX AdamW's decoupled wd matches torch AdamW). - InverseLR scheduler with warmup (--lr-scheduler inverse + --inv-gamma/--lr-power/ --lr-warmup/--lr-final): lr(step)=(1-warmup^(step+1))*max(final,lr*(1+step/inv_gamma)^-power), stepped per optimizer step (last_epoch=raw_step, fresh on resume like torch). At the template's warmup=0.995 the step-0 LR is ~5e-7, not the nominal 1e-4 — omitting it diverges immediately. Set via optimizer.learning_rate each step; the compiled step reads it through inputs=state (verified: exact InverseLR values, no recompile).
Reusable controlled-forward comparison: parity_forward_torch.py (sa3 venv) builds the fp32 base model, injects shared latents/noise/timesteps, replicates underfit's training-step forward, saves prediction/loss/cross/global + pretransform.scale; parity_forward_mlx.py (MLX venv) runs the same forward and reports three PSNRs that localize any mismatch — DiT-only (inject torch's cross/global → isolates the forward), conditioning (T5+seconds), end-to-end — plus asserts pretransform.scale==1.0. Fresh sm-music result: scale=1.0 (verified against the real model, not just assumed), DiT-only 79.0 dB, conditioning 88.5/77.6 dB, end-to-end 79.0 dB, loss Δ 6.1e-4. A convention bug would read ≪40 dB; ~79 dB is the cross-framework fp32 floor. No bug — the MLX and torch(MPS) training forwards agree to the floating-point limit.
Extend the parity harness to compare the full DiT backward sweep cross-framework: torch saves d loss/d noised (autograd w.r.t. the injected noised input, the same backward every interior/adapter gradient is built from); the MLX side computes mx.grad of the same loss (torch cross/global injected to isolate the DiT backward) and reports its PSNR. Result: 77.5 dB — the same fp32 floor as the forward (79 dB), so forward AND backward agree to the floating-point limit. The local adapter-param VJP was separately verified within-framework (DoRA reformulation, <=2.4e-7).
GitHub-hosted macos-latest runners are Apple Silicon, which MLX requires. Runs the weight-free adapter-math suite (all 9 LoRA/DoRA/BoRA types, scripts/ test_lora_merge.py) + a core-module import smoke + lora_train_mlx.py/ pre_encode_mlx.py --help. Scoped via paths filter to optimized/mlx changes so it doesn't spend macOS runner minutes on unrelated PRs. (test_all_configs.py needs the shipped npz weights and isn't CI-runnable — left out.)
The old 'LoRA training primitives' section predated the trainer and said there was 'no dataset or training CLI' — now stale. Replace it with the real end-to-end workflow: pre_encode_mlx.py → lora_train_mlx.py (train on the BASE ckpt via --dit-weights, underfit-default flags + the full-template flag set, demos) → generate with --lora. Keep a short 'build your own loop' note for the primitives, and list the training files + TRAINING_CONVENTIONS.md in the Files tree.
…onverter
Training uses the BASE checkpoint (not the shipped ARC weights). Ship the base
weights via HF so training works out of the box:
- weights.py: TRAINING_BASE manifest entries for dit_{sm-music,sm-sfx,medium}-base_f16.npz
(MLX/… in stabilityai/stable-audio-3-optimized), added to FLAT_MANIFEST for lazy
ensure_local download.
- lora_train_mlx.py: default to the base npz (ensure_local auto-download) when
--dit-weights is omitted — drop the old 'warn + train on ARC' fallback. --dit-weights
still overrides with a custom path.
- export_base_npz.py: torch-free converter (safetensors.numpy) from a *-base checkpoint
→ dit_<model>-base_f16.npz (DiT + baked conditioner). Validated to reproduce the
known-good sm-music/medium npz within fp16 (419/441 bit-identical, rest ≤1 fp16 ULP);
used to produce the sm-sfx base npz.
- README: training uses the auto-downloaded base npz (no flag needed).
added 9 commits
July 15, 2026 16:30
Print "Engine: MLX · device: Device(gpu, 0)" at the top of the encode run so the dashboard log (and CLI) shows unambiguously that the MLX SAME encoder on the Metal GPU is running — not the torch (MPS) pre_encode path.
…de parity) Add --exclude-file (text file of relpaths relative to --audio-dir, one per line) — same format/semantics as dataset_processing/pre_encode.py. Filters the discovered files in both main() (banner shows N excluded) and run(); the dashboard passes its exclude.txt through so per-file exclusions work on the MLX encode path.
Emit the same per-step tqdm line as underfit's torch loop (desc "Step N, Epoch E"; postfix train/loss, train/lr, train/grad_norm, train/lora_magnitude) instead of a plain print, so the dashboard collapses the progress bar and charts grad-norm / lora-magnitude with no translation layer. Both norms are one global L2 over the adapter grads (post-clip) and params (pre-update), matching loop.py _compute_grad_and_lora_norms. Add optional --gradient-clip-val (the dashboard passes 1.0) via mlx clip_grad_norm; grad_norm is post-clip like torch. Demos: entries tagged arc=true now render on the shipped rf_denoiser weights (train-on-base / demo-on-ARC). The trained LoRA is merged into a fresh copy of dit_<dit>_f16.npz (auto-downloaded, or --arc-weights) and sampled with the pingpong integrator; the seconds conditioner is model-independent so the trained embedder is reused. The ARC model is loaded once per demo round and freed after. RF/base entries keep the Euler path. Documents both in TRAINING_CONVENTIONS section 11.
Demo entries carry a `duration` (seconds) when they aren't at the crop length (the dashboard's full-length demos set it, crop-length ones omit it). The trainer was ignoring it and rendering every demo at the crop length. Now each demo computes its own T_lat from `duration` (falling back to the crop) for both the RF/base and ARC paths. For that to work on a model loaded at the training crop length, the DiT is now length-agnostic: the local-add-cond zeros in the text-to-audio path (local_add_cond=None) are built at the INPUT length instead of the baked self.T_lat. Identical for inference (which always loads T_lat == the gen length) and unused by training (which passes the all-ones inpaint cond), so this only unlocks reusing one loaded model across demo lengths. Both DiT sizes (dit_mlx_medium, dit_mlx) updated; attention is full (no mask) and RoPE is dynamic, so length is otherwise free. Verified: crop 128 + a 256-latent demo (via duration) render at 11.9 s / 23.8 s.
A full-length (4096-latent) demo is ~10 s/forward on MLX medium → ~8 min for a 50-step RF demo, and nothing was printed until it finished — so training looked stuck at the step-0 baseline. Now each demo prints its length + step count before running and shows a tqdm step bar (rf_euler_sample + pingpong_sample gain a `desc`). Makes clear it's progressing, not hung.
Each demo's completion line now shows how long generation took, e.g. "♪ demo 0 @ step 0: '...' → demo_0_00000000.mp3 (482s)" — useful given full-length demos run minutes each on MLX.
When --demo-decoder differs from the model's default (SA3-medium's SAME-L), print "demos: decoding with SAME-S instead of the model's SAME-L (faster demos)" so the console shows the swap. Both RF and ARC demos use the chosen decoder.
The temp LoRA merged into the ARC DiT was written to the run's checkpoints/ dir (as a dotfile), so it could surface in the dashboard's checkpoint list during the demo (or leak on a hard kill). Write it to the session dir instead (ckpt_dir.parent, not the scanned checkpoints/ subdir); it's still a dotfile and still removed after the demo.
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
2 participants
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.
Builds on #51 (@betweentwomidnights — the MLX LoRA training + audio-encoding primitives) and carries commits by @brxs. This completes them into a full MLX-native LoRA training runtime for SA3 that replicates underfit's conventions and can run as its Apple-Silicon backend — same behaviors, same defaults, same checkpoint format.
What's here (
optimized/mlx)scripts/lora_train_mlx.py— step-driven loop; AdamW (betas/eps/weight-decay + InverseLR warmup schedule); uniform / logit_normal / trunc_logit_normal / log_snr(_uniform) samplers; the SA3 "full" distribution shift with per-sample effective-length; CFG dropout 0.1; signal-only masked rectified-flow loss (per-sample-then-mean); checkpoint-every + underfit{run}-step=S-epoch=E.safetensorsfilename/metadata; resume-offset ladder;loss_by_timestep.bintelemetry;mx.compile'd step; per-prompt T5 conditioning cache;--grad-checkpoint.models/defs/lora.py— all 9 LoRA-family types with the no-weight-materialization DoRA/BoRA reformulation (--bora-mode speed|memory).models/defs/latent_dataset.py— PreEncodedLatentDataset (crop/pad, random-crop, oversample-with-replacement for tiny sets, underfit prompt-template augmentation;seconds_totalkept full after crop).scripts/pre_encode_mlx.py— torch-free audio→SAME-latent CLI, underfitnpy+jsonfield parity.models/defs/demo_mlx.py— training-time RF Euler+CFG inference on the base model + trained LoRA → peak-normalized mp3 + json (underfit on-disk format), baseline@0 + cadence + final, idempotent.scripts/sa3_gradio.py—--lorapreload (×3 slots, active on launch) +--port.scripts/parity_forward_{torch,mlx}.py— controlled fp32 forward+backward comparison vs the torch model.TRAINING_CONVENTIONS.md— the full underfit-convention inventory + parity results.Parity vs the torch (MPS) trainer — controlled fp32, sm-music base
pretransform.scale == 1.0verified against the real model (the softnorm no-scale convention was previously only assumed)Δ=0, signal-only loss identical, effective-length identicalFollow-ups (not blocking)
TRAINING_CONVENTIONS.mdinto the README; add a macOS-latest CI job for the MLX test suitespre_encode_mlx.pyinto underfit's dashboard pre-encode step forengine=mlx(today a dataset pre-encoded by either encoder feeds both engines, since thenpy+jsonformat is identical)Please preserve @betweentwomidnights and @brxs authorship on merge (a merge/rebase merge, or co-author trailers if squashed).