SoundGame Dev
f0f0a9b50e
- 共享引擎（app/src/main）：GestureEngine 6 手势识别、AudioEngine 三层音频 （ExoPlayer BGM 交叉淡入 / 对白 ducking / SoundPool 低延迟）、Game/GameEngine 状态机接口、MainActivity 全屏封面 + TalkBack - 《声命线》storyline：末世互动广播剧，3 幕 8 抉择点 4 结局 - 《声东击西》reflex：听声辨位反应，12 波战斗 + BGM 三档 BPM 变速 - 《暗室》escape：3x3 网格密室逃脱，钥匙→铁柜→保险丝→电箱→大门谜题链 - 文档：setup（从零搭建）/ architecture（架构）/ engine-api（引擎 API）/ asset-pipeline（资产流水线）/ games/*（三款游戏设计文档） - 资产管线：Qwen3-TTS 对白 + 程序合成音效/BGM + Mureka BGM 的 gen 脚本与台账， 全部 ogg 资产入库可复现 - env.d 环境脚本：android/tts 入库；sfx/mureka 含密钥不入库，仅提交 .example.sh 模板
228 lines
15 KiB
Markdown
228 lines
15 KiB
Markdown
# 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.
|