MLX engine for training + gradio (Apple Silicon)#10
Merged
Conversation
added 5 commits
July 15, 2026 12:07
Adds an alternate Apple-Silicon MLX engine alongside the default torch loop. engine=mlx shells out to the separate MLX trainer/gradio in the sibling stable-audio-3 checkout via a new underfit/backends/mlx_engine.py launcher. - defaults.ini: engine = torch (values torch|mlx) - lora_train.py / run_gradio.py: --engine flag (CLI > env UNDERFIT_ENGINE > defaults.ini > torch), dispatch to mlx_engine when engine=mlx - mlx_engine.py: path/name resolution (env-overridable), args->locked MLX trainer CLI mapping, stdout re-emit in dashboard-parseable format, gradio cmd - dashboard: Engine dropdown in New Finetune modal; server threads --engine onto both training and gradio launch commands, stores engine on the run Torch path is unchanged for engine=torch (default).
…-config list
The MLX trainer's --demo-config is a flat LIST of {prompt,cfg,seed,steps,...}
entries (its demo_cond equivalent), not underfit's whole training.demo dict.
Map demo_cond -> entries (cfg from demo_cfg_scales[0], steps from demo_steps),
skipping ARC entries (the MLX trainer finetunes the base model; ARC demos need
a weight-swap it doesn't do).
…efault
Map the SA3 template's training conventions onto the MLX trainer CLI: AdamW
betas/eps/weight_decay, InverseLR scheduler (+inv_gamma/power/warmup/final_lr),
model.diffusion.use_effective_length_for_schedule → --use-effective-length. Also
default adapter_type to 'lora' when the lora_config omits it (underfit's config-layer
default lora_config.get('adapter_type','lora'); the MLX trainer's own default is
dora-rows, which would otherwise mismatch a template without an explicit type).
…e npz) The MLX trainer now auto-downloads the base-model npz from HF, so underfit no longer needs a local base file. Pass --dit-weights only when UNDERFIT_MLX_BASE_WEIGHTS is set or a local base npz exists; otherwise omit it and let the trainer download. Removes the hard FileNotFoundError on a missing local base file.
…CUDA)
The dashboard modeled "devices" as an nvidia-smi CUDA-GPU list and hard-gated on
it — on a Mac it showed "No CUDA GPUs detected ... require an NVIDIA GPU" and left
Start disabled, even though the training loop already resolves cuda>mps>cpu and the
MLX engine exists.
server.py:
- _detect_platform() (torch-free: Darwin+arm64 -> apple, else nvidia count -> cuda, else cpu).
- _apple_gpu_info() — the Metal counterpart of the nvidia-smi query: chip/GPU name +
core count via "system_profiler SPDisplaysDataType", unified-memory total via
"sysctl hw.memsize"; _apple_gpu_mem_used_mb() via "vm_stat".
- /api/gpu returns platform + engines + (on apple) a single synthetic Metal device
{kind:apple, name, cores, used/total/free_mb}. CUDA path unchanged.
- Launch env: _cuda_env_prefix() omits CUDA_VISIBLE_DEVICES on apple/cpu;
_free_gpu_memory no-ops off-CUDA.
index.html:
- Store platform/engines; _noGpuHtml() keeps the exact CUDA message on cuda, softer
"CPU (slow)" otherwise; Apple renders a normal device tile (name + N-core / unified
memory) — no warning.
- Auto-select the single Apple device (Start/Launch enable without a click).
- Engine dropdown filtered/relabelled to the platform: torch -> "MPS (PyTorch)" + "MLX"
on apple (default MLX); mlx hidden on cuda/cpu (default torch).
All Apple/CPU branches guarded on platform -> the Linux/CUDA path is unchanged.
added 14 commits
July 15, 2026 16:19
The "base model not downloaded — run installer again" block in validateNewFt is the TORCH base checkpoint (modelMeta.registered). Skip it when engine=mlx: the MLX trainer brings its own base npz (auto-downloaded from HF / local in the MLX checkout), so the torch download isn't required. Added an onchange on #new-ft-engine so the error clears/appears when the engine is switched. Also shorten the Apple engine label "MPS (PyTorch)" -> "MPS" to save UI width. Dataset ENCODING still gates on the torch ckpt (validateNewDataset / launchEncoding) — wiring the MLX pre-encode into the dashboard encoding flow is a separate follow-up.
The Encode Dataset flow launched dataset_processing/pre_encode.py (torch, reads the torch base checkpoint) and gated on it being downloaded — so on Apple, where the torch backend isn't installed, dataset encoding was blocked with "base model not downloaded". On Apple, encode with the MLX SAME encoder instead: - mlx_engine.build_encode_cmd(input, output, model) -> pre_encode_mlx.py --audio-dir --output-dir --codec (small -> same-s, medium -> same-l). Encoder weights are the MLX SAME npz (auto-downloaded / local in the MLX checkout) — no torch checkpoint. - server.py /api/datasets/encode: dispatch to the MLX encoder when platform == apple (no CUDA_VISIBLE_DEVICES); torch path unchanged elsewhere. Output layout (npy + json + details.json) matches pre_encode.py, so the existing completion check works as-is. - index.html: skip the "not downloaded" gate (validateNewDataset + launchEncoding) when platform == apple. Known gap: the MLX encoder does not yet honor the per-file exclude list (encodes the full folder); logged when excludes are set.
pre_encode_mlx.py writes flat into --output-dir, but the dashboard expects latents in output_dir/latents/<model> (latent_dir) — where the torch pre_encode.py writes and where the encoding monitor looks for details.json + *.npy. Pointing the MLX encoder at output_dir left latent_dir empty, so the monitor saw a finished-but-empty encode and dropped the dataset. Point build_encode_cmd at latent_dir (+ mkdir it).
build_encode_cmd takes exclude_file and appends --exclude-file; the encode dispatch passes the run's exclude.txt. Removes the earlier 'mlx encoder does not honor the exclude list' note — per-file exclusions now work on the MLX path.
run_mlx_training now inherits the trainer's stdout/stderr instead of capturing and re-emitting a synthetic step line. The trainer emits the same tqdm output as underfit's torch loop, so writing it straight to the run log preserves tqdm's carriage-return progress updates (the dashboard collapses on them) and the grad_norm / lora_magnitude postfix (the metric charts parse it). Drops the translation loop and the now-unused _STEP_LINE_RE. _build_demo_entries passes ARC entries through (arc=true, cfg=1, steps=8 defaults) instead of skipping them - the MLX trainer renders them now. build_trainer_cmd forwards --gradient-clip-val (the dashboard passes 1.0) so the MLX trainer clips identically.
… model gate
(1) The /api/gradio gpu validation now only applies on CUDA. Apple/CPU have one
implicit accelerator (index 0) and nvidia-smi reports 0 GPUs, so the old
gpu >= _get_gpu_count() check rejected every Apple launch with "gpu must be
0--1". The MLX/MPS launch ignores the index anyway.
(2) The device-detail panel had no data on Apple: _query_gpu_mem returned {}
(no nvidia-smi) so the VRAM history never recorded ("Collecting VRAM data..."),
and _get_gpu_processes used nvidia-smi + /proc (empty on macOS → "No
processes"). Now _query_gpu_mem reports unified memory as GPU 0 (vm_stat), and
_apple_gpu_processes lists the managed training / gradio / encoding runs with
their whole process-tree RSS (ps/pgrep, no /proc).
(3) The "Base model not downloaded" popup (_warnUnregisteredModel) fired on
Apple for both MLX and MPS even though the MLX npz is present. Guarded with
_platform==='apple', matching the existing encode and validateNewFt guards.
(4) Resume now skips the torch-backend gate for engine=mlx, mirroring the
fresh-launch path — killing and reviving an MLX run no longer errors "requires
the sa3 backend".
When a base model's MLX weight pack (base + ARC + SAME codec + encoder + shared
t5gemma) isn't present, offer to download it instead of a dead-end "run
installer again" message.
- mlx_engine: shared pack helpers — mlx_pack_rel_paths / mlx_model_available /
mlx_missing_pack (files + approx GB) / build_mlx_download_cmd (runs the MLX
checkout's weights.ensure_local, the download source of truth, emitting
DL/OK/ALL_DONE lines for progress). mlx_root_or_none resolves without raising.
- server: /api/models annotates each model with mlx_available + mlx_download_gb
on Apple; POST /api/models/download starts a tracked background download and
GET /api/models/download?model= reports {state, done/total files, current
file, remaining GB}.
- index.html: New Finetune shows a "Download base + ARC (~X GB)" button when the
selected model's MLX pack is missing (engine=mlx on Apple). It downloads with
a live progress line, refreshes availability, and clears the gate — Next stays
disabled until the pack lands.
- setup.py: on Apple, the model phase pre-downloads the MLX packs into the
existing MLX checkout (reusing the same mlx_engine download) and skips the
torch packs an MLX-only Mac won't use. Degrades gracefully if the checkout /
venv isn't set up (weights then download lazily on first use).
_build_demo_entries dropped the per-demo `duration` field, so full-length demos collapsed to the crop length in the MLX trainer. Forward it (like seed / lora_strength / lora_interval_max) so each demo renders at its configured length.
… gate On Apple the engine dropdown always offered "MPS" (torch) even with no torch backend installed, so starting an MPS run greyed out Next with no warning — and only a name change surfaced "not downloaded — run installer again". - server: /api/gpu `engines` now reflects installed backends — mlx on Apple (its checkout brings its own stack) plus torch only where stable_audio_3 / stable_audio_tools is importable. An MLX-only Mac reports ["mlx"], so the UI omits the unusable MPS option entirely rather than showing a greyed dead-end. - index.html: _applyEngineOptions hides/disables the torch (MPS) option when torch isn't in _engines (it previously only handled mlx). - Fixed the greyed-silent gate itself: validateNewFt now flags a blocking error and updateCkptEstimate bails instead of overwriting it. The onchanges call both, so the size estimate was clobbering the "not downloaded" message, which then only reappeared on a name change. Any validation error shows immediately.
Launching an MLX gradio failed with "could not determine the --dit model": resolve_dit_model reads 'base_model' from the run's _model.json (which doesn't carry it) and the launch never passed --pretrained-name. The launch already knows base_model from the run record — forward it so the MLX gradio picks the --dit value (the torch path ignores --pretrained-name).
RF/Base demos are 50 steps (~8 min each at full length on MLX); ARC is 8 steps. For the MLX engine the suggested default now drops the RF demos and keeps only the ARC ones (all presets), falling back to RF only when the model has no ARC variant. _defaultNewFtDemos / _defaultDemoDurationLatents / the unconditional-prompt indices (_demoEmptyIdx) are engine-aware, and newFtGoStep4 rebuilds the default with the current engine on entry.
New Finetune → Demos shows a "Use SAME-S for demos (faster)" checkbox for SA3-medium only — default ON for Apple/MLX, off elsewhere. When on, demos decode with the small SAME-S autoencoder instead of the model's SAME-L (~1.6× faster, near-identical: the two share the latent space, corr 0.99). Wired payload.demo_decoder → training.demo.demo_decoder → the MLX trainer's --demo-decoder (mlx_engine); the torch loop keeps its own decoder.
MLX demos are ARC-only, so a preset yields half the demos of the RF+ARC layout (2→1, 4→2, 8→4). Relabel the preset buttons One/Two/Four on MLX to match the actual ARC-demo count; non-MLX keeps Two/Four/Eight (half RF + half ARC). The internal preset values (2/4/8) are unchanged — labels only.
_find_checkpoints included any *.safetensors, so the MLX ARC-demo scratch LoRA (.arc_demo_tmp_*.safetensors) showed up in the Checkpoints box and the Revive/launch checkpoint pickers. Skip dotfiles in the scan — hides that scratch (and any hidden temp) from every checkpoint list.
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.
Adds an engine selector (
torchdefault |mlx) so underfit can train + serve on Apple Silicon via MLX, alongside the existing torch (MPS/CUDA/CPU) path. The torch path is byte-for-byte unchanged whenengine=torch— this is purely additive.Pairs with the MLX trainer/runtime in Stability-AI/stable-audio-3#72; this PR is the underfit glue that drives it. Stacked on #9 (MPS training support) — base retargets to
mainwhen #9 merges.What it does (engine=mlx)
run_trainingshells out to the sibling stable-audio-3 MLX trainer (optimized/mlx/scripts/lora_train_mlx.py) in its own MLX venv, maps underfit's model/dataset/training config onto its CLI (adapter type/rank/alpha, timestep sampler, distribution shift +use_effective_length, AdamW betas/eps/weight-decay, InverseLR schedule, CFG-dropout, demos, resume offsets), and translates the trainer's stdout into the dashboard's progress + loss/lr-chart format. The base-model npz is auto-downloaded by the trainer, so no local base file is required.sa3_gradio.py) with the trained LoRA preloaded (--lora), instead of the torch gradio.Files
underfit/backends/mlx_engine.py(new) — torch-free launcher: MLX root/venv resolution, arg mapping, stdout→dashboard log translation, training + gradio runners.--engine {torch,mlx}inlora_train.py+run_gradio.py;engine = torchindefaults.ini(precedence: CLI >UNDERFIT_ENGINEenv > defaults.ini > torch).dashboard/server.py+dashboard/index.html— Engine dropdown + launch-command wiring.Config / paths
Overridable via env:
UNDERFIT_MLX_ROOT(default: the siblingstable-audio-3/optimized/mlxcheckout),UNDERFIT_MLX_PYTHON(default:<root>/.venv/bin/python),UNDERFIT_MLX_BASE_WEIGHTS(optional; otherwise the trainer downloads the base npz from HF).Verified (M4 Pro)
engine=mlxtraining end-to-end: 5 steps → checkpoint + demos + translated dashboard log lines, exit 0; losses match the direct MLX trainer.Δ=0, signal-only loss, effective-length) and the MLX-vs-torch forward+backward parity holds to the fp32 floor (see stable-audio-3#72).