# Underfit training conventions → MLX port reference

Working reference for the `mlx-lora-training` branch: every training convention and
default in [underfit](https://github.com/dada-bots/underfit) (the SA3 LoRA trainer),
inventoried from its source (2026-07-14), plus the gap analysis against the MLX
primitives this branch starts from (PR #51 by @betweentwomidnights). The goal:
the MLX runtime can run as an underfit backend — same behaviors, same defaults,
same checkpoint format. Trim/fold this file into README before the PR finalizes.

## 1. Adapter configuration

| Convention | Underfit value | Source |
|---|---|---|
| adapter_type (config-layer default) | `"lora"`; legacy `"dora"` → `"dora-rows"` | training/lora.py:34, resolve_adapter_type |
| adapter_type (product/dashboard default) | **`dora-rows`** | dashboard index.html:1336 |
| rank (config-layer default) | 8 | training/lora.py:32 |
| rank (product/dashboard default) | **16** | README.md:157 |
| alpha default | = rank | training/lora.py:33 |
| targets | every `nn.Linear` + `nn.Conv1d` under BOTH `model.model` (DiT) and `model.conditioner`, filtered by include/exclude | training/lora.py:53-64 |
| include/exclude | substring match + `[i-j]` numeric-range expansion; dashboard default EXCLUDES all conditioner + "other" modules (embedders/projections/convs) | dashboard index.html:5741-49 |
| adapter dropout | none | — |
| trainable dtype | LoRA params forced fp32; base stays fp16 (`base_precision`) | loop.py:344-347 |

## 2. Checkpoint format

- Keys: `{**get_lora_state_dict(model.model), **get_lora_state_dict(model.conditioner)}`
  → DiT keys come out `model.transformer.layers...` (inner `.model` on the DiT wrapper),
  conditioner keys `conditioners.seconds_total.embedder.embedding.1...`; per-key shape
  `<layer>.parametrizations.weight.<lora_index>.{lora_A,lora_B,magnitude,magnitude_r,magnitude_c,M_xs,U,V}`
  (index 0 for the finetune adapter).
- Tensors saved **fp16** (`save_lora_safetensors` casts `.half()`); DoRA magnitudes
  squeezed to 1-D on load (`prepare_dora_state_dict`).
- Metadata `lora_config` JSON: `rank, alpha, adapter_type, include, exclude` +
  injected `step` (int), `epoch` (int), `base_model` (str, from model_config).
- Filename: `{run_label}-step={step}-epoch={epoch}.safetensors` (run_label = run name
  with trailing 14-digit timestamp stripped); dir `<save_dir>/<name>/<uuid8>/checkpoints`.
- -xs: only `M_xs` (+ magnitudes) trained/saved; U/V bases precomputed OFFLINE
  (compute_svd.py) from fp32 base weights into a separate `.pt`
  (`{key: {"U": fp16, "V": fp16, "S": fp32, "shape"}}`, sign-canonicalized by max-abs
  U-column entry), loaded via `svd_bases_path` from the model registry; per-layer SVD
  at attach time is the fallback.
- Validation contract (lora_validate.py): `.safetensors` only, ≥1 parametrization key,
  `lora_config` present with `rank` + `adapter_type` (alpha:=rank, include/exclude:=[]
  back-filled), adapter type inferable from key fingerprints.

## 3. Loss (loop.py:561-613 + loss.py)

- Objective from the model: `rectified_flow` → `alphas, sigmas = 1−t, t`,
  `noised = x·α + noise·σ`, `target = noise − x` (`"v"` objective also exists:
  cos/sin schedule, `target = noise·α − x·σ`).
- MSE per-element; `loss_normalization` default `"none"` (timestep/sample/sample_channel
  variance-normalization modes exist, eps 1e-6, detached variance).
- Masking: `mask_loss_weight` default 0.0 and `mask_padding_attention` default False →
  **signal-only** loss; reduction is per-sample mean over valid positions×channels,
  then batch mean. Inpainting further restricts the mask to the regenerated region.
- No sigma/logSNR loss weighting. Runs under fp16 autocast; no explicit casts.

## 4. Timesteps

- Sampler default **`uniform`** (`torch.rand`); also logit_normal,
  trunc_logit_normal (left 0.075, flipped `1−t`), log_snr (mean −1.2, std 2.0),
  log_snr_uniform (min −6, max 5).
- Distribution shift comes from the MODEL config: SA3 models ship
  `{"type": "full", "min_length": 256, "max_length": 4096}` (base_shift 0.5,
  max_shift 1.15 defaults) — applied as `t = dist_shift.shift(t, seq_len)` with
  seq_len = the latent time dim (crop length) by default
  (`use_effective_length_for_schedule` False; when True: `ceil(seconds_total·sr/4096)`).

## 5. CFG dropout (dit.py:441-450)

- `cfg_dropout_prob` default **0.1**, applied INSIDE the model forward when
  `cfg_scale == 1.0`: per-sample bernoulli mask replaces the whole
  `cross_attn_cond` (prompt+seconds concat) with **zeros**; `global_cond` is NOT
  dropped. (Matches inference: the uncond branch zeroes cross_attn, keeps global.)

## 6. Optimizer / schedule / loop mechanics

| Convention | Underfit value |
|---|---|
| optimizer | AdamW, single param group (LoRA params only) |
| lr | **no default — must be supplied** (`--lr` or optimizer_configs) |
| betas/eps/weight_decay | torch defaults: (0.9, 0.999) / 1e-8 / 0.0 |
| LR schedule | none by default (InverseLR + any torch scheduler configurable) |
| grad clip | off by default (`gradient_clip_val = 0.0`); when set: clip_grad_norm_ |
| grad accumulation | none (defaults.ini key exists but inert) |
| EMA / validation / early stopping / save_top_k | none (inert keys) |
| precision | `"16-mixed"`: fp16 autocast + GradScaler; base fp16, LoRA fp32 |
| batch_size | 1 (defaults.ini) |
| seed | 42 |
| steps | step-driven; `max_steps` absolute global-step target (default ~1e10) |
| checkpoint cadence | every 1000 steps + final save; save BEFORE demos |
| resume offsets | explicit config > safetensors metadata `step/epoch` > filename `step=/epoch=` tokens > step//steps_per_epoch |
| logging | train/loss, train/lr, grad_norm + lora_magnitude (fp32, post-clip); `loss_by_timestep.bin` binary telemetry (step, t_mean, loss_mean) flushed every 10 steps |

## 7. Data pipeline

- **Pre-encode** (offline, once): whole files (no clip splitting), up to 600 s
  (aligned down to ds_ratio=4096 samples), 44.1 kHz stereo (mono→stereo by channel
  repeat, resample via torchaudio defaults, NO loudness/peak norm), encoder fp32,
  `chunked=True` for >30 s. Saved per file: `<stem>.npy` latents `[D,T]` fp32
  **unscaled** (training divides by pretransform scale) + `<stem>.json` with
  `seconds_total` (full duration, rounded 3dp), `seconds_start: 0`,
  `audio_samples`, `latent_shape`, latent-resolution 0/1 `padding_mask`, and all tags.
- **Training dataset** (PreEncodedDataset semantics): crop/pad every item to
  `latent_crop_length` — **1300** latents (sm-music/sfx, ≈120 s) / **4096** (medium,
  ≈380 s); `random_crop=True` (dashboard default) choosing a start within the valid
  region; short items padded with `silence.npy` latent if present else zeros, mask
  extended with zeros. Fixed-length batches + padding mask (no bucketing).
  **`seconds_total` stays the FULL song duration after cropping** (not recomputed).
- **Small datasets**: if `len(ds) < batch_size`, RandomSampler with replacement,
  `num_samples = batch_size*100` (~100 steps/epoch) — the core underfit workflow.
- **Prompts** (prompt_templates.py, per-sample at train time): tag-based by default
  (`use_tags`, tags→`Title:/Artist:/Album:/Genre:/Label:/Year:/Composer:/BPM:/Prompt:`
  joined ", "); augmentation: 50% shuffle-all vs 50% random-subset (1..n tags);
  optional trigger token prepended with prob **80%**; path/fixed prompt sources with
  weighted balance (tags 50/paths 50/fixed 0); `lyrics=""` always. Tag priority:
  JSON sidecar > `.txt` sidecar (whole content = `prompt`) > embedded ID3
  (`title,artist,album,genre,label,date,composer,bpm`).
- **Demos**: `demo_every` 1000 (template; loop default 0=off) + baseline at step 0;
  `demo_steps` 50 (ARC 8), cfg [7] (ARC 1), per-entry seed/duration/lora_strength/
  `lora_interval_max` (σ-interval demo path uses an explicit Euler loop);
  peak-normalized int16 → mp3 (`ffmpeg -q:a 0`), trimmed to seconds_total+5 s,
  `demo_<i>_<step:08d>.mp3` + json sidecar, idempotent per step.

## 8. Gap analysis: PR #51 primitives vs the above

Already aligned: adapter math for all 9 types (torch-parity tested), rf loss target
+ masked per-sample-then-mean reduction, the 5 timestep samplers + full/flux/logsnr
shifts with matching defaults, fp16 checkpoint tensors + `lora_config` metadata,
`.safetensors`-only trust boundary, fp32 adapter params over a fp16 base, SVD sign
canonicalization, `[i-j]` include/exclude expansion, whole-file chunked encoding
with validity masks.

To build (roughly in order):
1. **lora.py**: save with underfit key roots (`model.` + `conditioners.`) instead of
   bare DiT keys; adapter for the baked seconds-conditioner Linear
   (`cond.seconds_total_weight` ↔ `conditioners.seconds_total.embedder.embedding.1`);
   match include/exclude against checkpoint-convention names so dashboard filter
   strings work verbatim; config-layer defaults rank 8 / alpha=rank / lora;
   metadata `step/epoch/base_model` injection helpers; accept `svd_bases_path`
   `.pt`-equivalent (or keep recompute fallback, which underfit also has).
2. **conditioning + dropout**: training-time conditioning builder (T5Gemma frozen,
   prompt padding, seconds token; MLX runtime pieces exist) + per-sample
   cross_attn→zeros dropout at prob 0.1 (keep global_cond).
3. **trainer loop** (`lora_train`-equivalent CLI): AdamW (torch-default betas/eps/wd,
   lr required), no schedule, no clip, batch 1, seed 42, step-driven with
   checkpoint-every-1000 + final save, underfit filename convention, resume-offset
   ladder, loss/lr/grad_norm logging, `loss_by_timestep.bin`.
4. **dataset**: PreEncodedDataset-equivalent (npy+json pairs, latent_crop_length
   1300/4096 defaults, random_crop, silence/zero pad, mask handling, seconds_total
   NOT recomputed, oversample-with-replacement ×100 for tiny sets) +
   prompt_templates port (50/50 shuffle/subset, trigger 80%, display-name map).
5. **pre-encode CLI** on top of `audio_encoding.py` (600 s cap, fp32, unscaled fp32
   npy + json sidecar with the exact metadata fields).
6. **demos** (optional, last): pipeline-based demo step with underfit cadence/format.

## 9. Training-forward conventions discovered by parity testing (2026-07-14)

**Harness** (reusable, 2026-07-15): `scripts/parity_forward_torch.py` (sa3 venv —
builds the fp32 base model, injects shared latents/noise/timesteps, saves
prediction/loss/cross/global/scale) + `scripts/parity_forward_mlx.py` (MLX venv —
same forward, reports three bug-localizing PSNRs). Fresh run (sm-music base, B=3,
t={0.25,0.5,0.75}): **pretransform.scale=1.0 asserted against the real model**
(confirms the softnorm no-scale convention — a prime hidden-bug candidate),
DiT-only forward (torch cross/global injected → isolates the DiT) **79.0 dB**,
conditioning cross/global **88.5/77.6 dB**, end-to-end **79.0 dB**, loss torch
4.716201 vs MLX 4.716814 (**Δ 6.1e-4**). **Backward also checked**: ∂loss/∂noised
through the full DiT (torch cross/global injected) agrees at **77.5 dB** — the same
floor as the forward, so the backward sweep every adapter gradient is built from
matches too (the local adapter-param VJP was separately verified within-framework to
≤2.4e-7 MLX / ≤4.4e-7 torch during the DoRA reformulation). A convention/autograd bug
(RoPE/attn-mask/adaLN/local_add_cond layout/scale) reads ≪40 dB; ~78 dB is the
cross-framework fp32 reduction-order floor through 20 layers. Confirms fwd+bwd parity,
no bug.

Verified by a controlled forward (identical latents crop, numpy-seeded noise,
t=0.5, fixed prompt, fp32) through underfit's torch loop on MPS vs the MLX
trainer — final agreement **loss 3.9829 vs 3.9823, prediction PSNR 80.7 dB**
(bounded by the fp16 npz weights; torch MPS-vs-CPU agrees at 3e-6 rel):

- **Train on the BASE checkpoint** (`stabilityai/stable-audio-3-*-base`,
  rectified_flow), NOT the shipped ARC (rf_denoiser) weights that inference
  uses. `lora_train_mlx.py --dit-weights` + a 441-key npz conversion
  (strip `model.model.`, `.gamma`→`.weight`, conv (out,in,k)→(out,k,in),
  `to_local_embed.{0,2}`→`.seq.{0,2}`, conditioner keys → `cond.*`).
- **local_add_cond during training = all-ONES inpaint mask + zero context**
  (the models are the diffusion_cond_inpaint variant). The inference path's
  `local_add_cond=None` (≡ zeros) is an inference-only convention — using it
  in training changes the loss by ~4x (wrong conditioning regime).
- **Cross-attention runs over all 257 tokens at training time too**: the
  torch conditioner substitutes the learned padding_embedding into padded
  positions (exactly like the optimized runtimes) and the attention mask is
  effectively pass-through. Do NOT mask/slice to valid tokens.

## 10. MPS vs MLX benchmark (M4 Pro 48 GB, sm-music base, 30 steps,
## dora-rows r16 α16 = 141 layers / 9.2M params, batch 1, crop 256, fp16 base)

| | MPS (underfit loop, torch 2.13) | MLX (lora_train_mlx.py) |
|---|---|---|
| steady it/s | 0.668 (1498 ms/step) | **1.65 (2.47x faster)** |
| peak memory | 3.67 GiB driver-allocated | (not instrumented; comparable class) |
| loss regime | 1–43, spikes at low t | same (0.9–35, spikes at low t) |
| checkpoint | underfit-valid, 423 keys | underfit-valid, identical key set |

## 11. Dashboard integration (implements §6/§7 for the underfit dashboard)

- **Progress output is byte-identical to underfit's torch loop** (loop.py:516):
  one `tqdm` per epoch, `desc="Step N, Epoch E"`, `mininterval=0, miniters=1,
  file=sys.stdout`, postfix `{train/loss:.6f, train/lr:.3e, train/grad_norm:.6f,
  train/lora_magnitude:.6f}`. The dashboard collapses the progress bar on tqdm's
  `\r` and parses those four keys per step (server.py `_HISTORY_RE`/`_LR_RE`/
  `_GN_RE`/`_LM_RE`) — so `underfit.backends.mlx_engine.run_mlx_training` just
  inherits stdout (no line translation, which would break `\r` and drop the
  postfix). grad_norm/lora_magnitude are one global L2 over the adapter grads
  (post-clip) / params (pre-update), matching `_compute_grad_and_lora_norms`.
- **`--gradient-clip-val`** (dashboard passes 1.0) clips the global adapter-grad
  norm via `mlx.optimizers.clip_grad_norm` before the step; grad_norm is reported
  post-clip like torch. Default 0 = off.
- **ARC demos** (entries `arc=true`): the trained LoRA is merged into a fresh copy
  of the shipped rf_denoiser weights (`dit_<dit>_f16.npz`, auto-downloaded, or
  `--arc-weights`) and sampled with the pingpong integrator — underfit's
  train-on-base / demo-on-ARC story (its torch loop swaps the full ARC model in
  place; the MLX path loads it with the LoRA merged at load, once per demo round,
  freed after). The seconds conditioner is model-independent, so the trained
  `bundle.secs` is reused for both RF and ARC conditioning. RF/base entries keep
  the plain-Euler path. cfg/steps default to 1/8 for ARC, `demo_cfg_scales[0]`/
  `demo_steps` for RF.