Merge pull request #72 from Stability-AI/mlx-lora-training

MLX LoRA training for SA3 (Apple Silicon)

Author: Cortexelus

.github/workflows/mlx.yml

@@ -0,0 +1,32 @@
+name: MLX (Apple Silicon)
+
+# GitHub-hosted `macos-latest` runners are Apple Silicon (arm64), which the MLX
+# runtime requires. Scoped to optimized/mlx changes so it doesn't spend macOS
+# runner minutes on unrelated PRs.
+on:
+  pull_request:
+    paths:
+      - "optimized/mlx/**"
+      - ".github/workflows/mlx.yml"
+
+jobs:
+  test:
+    runs-on: macos-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install MLX deps
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r optimized/mlx/requirements.txt pytest
+      - name: Adapter math — all 9 LoRA/DoRA/BoRA types (weight-free)
+        run: pytest optimized/mlx/scripts/test_lora_merge.py -q
+      - name: Import + CLI smoke (no weights downloaded)
+        working-directory: optimized/mlx
+        run: |
+          python -c "import mlx.core; from models.defs import lora, training, demo_mlx, latent_dataset, lora_merge, sa3_pipeline; print('core imports OK')"
+          python scripts/lora_train_mlx.py --help > /dev/null
+          python scripts/pre_encode_mlx.py --help > /dev/null
+          echo "CLI smoke OK"

optimized/mlx/.gitignore

@@ -14,3 +14,6 @@ output/gradio/
 
 # Local LoRA library scanned by the gradio UI (user adapters, never committed)
 loras/
+
+# Trainer telemetry (written to cwd, underfit convention)
+loss_by_timestep.bin

optimized/mlx/README.md

@@ -191,6 +191,70 @@ Sample run on **M4 Pro / 48 GB**:
 └───────────┴─────────┴─────────┴───────────┴───────────┴────────────┘
 ```
 
+## LoRA training
+
+Finetune SA3 on your own audio, entirely in MLX (no PyTorch) — a full training
+CLI that replicates [underfit](https://github.com/dada-bots/underfit)'s
+conventions, defaults, and checkpoint format (so it can run as underfit's
+Apple-Silicon backend). Three steps: **pre-encode → train → generate**. See
+`TRAINING_CONVENTIONS.md` for the complete convention inventory and the
+torch(MPS)-vs-MLX forward+backward parity results.
+
+**1. Pre-encode** your audio to SAME latents (once, offline — torch-free):
+
+```bash
+uv run python scripts/pre_encode_mlx.py \
+    --audio-dir ~/my-clips --output-dir ~/my-latents --codec same-s
+```
+
+Each file becomes `<stem>.npy` latents `[D, T]` + a `<stem>.json` sidecar
+(duration, padding mask, tags) — the exact format underfit's dataset uses.
+Use `same-s` for sm-music/sm-sfx, `same-l` for medium.
+
+**2. Train.** Training uses the **BASE** checkpoint (`stabilityai/stable-audio-3-*-base`,
+rectified_flow), *not* the shipped ARC weights inference uses. The base-model npz is
+**auto-downloaded** from HuggingFace on first run — no flag needed (override with
+`--dit-weights <path>`; regenerate one yourself from a `-base` repo with
+`scripts/export_base_npz.py`):
+
+```bash
+uv run python scripts/lora_train_mlx.py \
+    --dit sm-music --latents-dir ~/my-latents --lr 1e-4 --name my-lora \
+    --adapter-type dora-rows --rank 16 --max-steps 2000
+```
+
+Defaults mirror underfit (dora-rows / rank 16, AdamW, uniform sampler + the
+"full" distribution shift, CFG-dropout 0.1, signal-only masked rectified-flow
+loss, checkpoint every 1000 steps). To match the full dashboard template config,
+add `--timestep-sampler trunc_logit_normal --use-effective-length --beta2 0.95
+--weight-decay 0.01 --lr-scheduler inverse --lr-warmup 0.995`. Checkpoints land at
+`output/runs/<name>/<uuid>/checkpoints/<name>-step=S-epoch=E.safetensors` with
+`lora_config` metadata; resume with `--lora-ckpt-path <ckpt>`.
+
+Optional training-time demos (RF-Euler inference → mp3, underfit's on-disk
+format): `--demo-every 1000 --demo-config demos.json`, where `demos.json` is a
+list of `{"prompt", "cfg", "seed", "steps", "lora_strength"?, "lora_interval_max"?}`.
+
+**3. Generate** with your adapter — the checkpoint applies to the shipped ARC
+model at inference via `--lora` (see "Apply a LoRA finetune" above; also
+`scripts/sa3_gradio.py --lora <ckpt>` to open the web UI with it preloaded):
+
+```bash
+./sa3 --dit sm-music --decoder same-s --prompt "..." --out finetuned.wav \
+      --lora "output/runs/my-lora/"*"/checkpoints/my-lora-step=2000-epoch="*.safetensors
+```
+
+### Building your own loop
+
+The CLI is assembled from reusable pieces: `inject_trainable_lora` /
+`save_lora_checkpoint` / `apply_lora_checkpoint` (`models/defs/lora.py` — all 9
+LoRA/DoRA/BoRA/-XS types with the no-weight-materialization forward), the RF
+loss + samplers + distribution shift (`models/defs/training.py`), the
+pre-encoded dataset + prompt templates (`models/defs/latent_dataset.py`), and
+`encode_audio` (`models/defs/audio_encoding.py`). Checkpoints use the same
+safetensors keys + `lora_config` metadata as the PyTorch trainer, so adapters
+are interchangeable in either direction.
+
 ## Flag reference
 
 | Flag                  | Default  | Notes                                                                 |
@@ -221,6 +285,7 @@ sa3_mlx/
 ├── sa3                            ← shell wrapper (use this)
 ├── install.sh                     ← uv bootstrap (run once)
 ├── README.md
+├── TRAINING_CONVENTIONS.md         ← LoRA-training conventions + MPS-vs-MLX parity
 ├── requirements.txt
 ├── output/                        ← default landing zone for generated WAVs
 ├── scripts/
@@ -229,15 +294,26 @@ sa3_mlx/
 │   ├── examples.py                ← shared examples block (--help + post-install)
 │   ├── install.py                 ← install.sh's Python half (bundle picker)
 │   ├── test_all_configs.py        ← npz + CLI config sanity tests
-│   └── benchmark.py               ← wall-time + peak-RAM matrix across model × duration
+│   ├── benchmark.py               ← wall-time + peak-RAM matrix across model × duration
+│   ├── lora_train_mlx.py          ← LoRA training CLI (underfit conventions)
+│   ├── pre_encode_mlx.py          ← audio → SAME-latent pre-encode (torch-free)
+│   ├── sa3_gradio.py              ← web UI (invoked by ./sa3-gradio; --lora preload)
+│   ├── test_lora_merge.py         ← adapter-math tests (weight-free; run in CI)
+│   └── parity_forward_{torch,mlx}.py  ← torch(MPS)-vs-MLX fwd+bwd parity harness
 └── models/
     ├── defs/
     │   ├── sa3_pipeline.py        ← sampler + conditioner + unpatch
     │   ├── t5gemma_mlx.py         ← T5Gemma encoder + SentencePiece wrapper
     │   ├── dit_mlx.py             ← small DiT (sm-music + sm-sfx)
     │   ├── dit_mlx_medium.py      ← medium DiT (differential attention)
     │   ├── same_s_{encoder,decoder}.py    ← small codec
-    │   └── same_l_{encoder,decoder}.py    ← large codec
+    │   ├── same_l_{encoder,decoder}.py    ← large codec
+    │   ├── lora.py                ← trainable LoRA/DoRA/BoRA adapters (9 types)
+    │   ├── lora_merge.py          ← inference-time LoRA merge + per-step gating
+    │   ├── training.py            ← RF loss + timestep samplers + distribution shift
+    │   ├── latent_dataset.py      ← pre-encoded dataset + underfit prompt templates
+    │   ├── audio_encoding.py      ← waveform → SAME latents (pre-encode bridge)
+    │   └── demo_mlx.py            ← training-time RF-Euler demos → mp3
     └── mlx/                       ← .npz weights (auto-downloaded; ~8.4 GB total)
         ├── t5gemma_f16.npz                541 MB    text encoder + tokenizer
         ├── dit_sm-music_f16.npz           877 MB    DiT + conditioner baked in