hotstep-v1: soft-distilled + calibrated pooled-v4

README.md

@@ -0,0 +1,155 @@
+---
+license: other
+license_name: minimax-music3-terms
+license_link: https://huggingface.co/MiniMaxAI/MiniMax-Music3
+base_model: Mothersuperior/open-rvq-encoder-minimax-music3-169m-pooled-v4
+datasets:
+- Mothersuperior/minimax-music3-rvq-distill-corpus-8k
+tags:
+- minimax-music3
+- rvq
+- audio-encoder
+- reverse-distillation
+- knowledge-distillation
+---
+
+# open-rvq-encoder-minimax-music3-169m — soft-distilled + calibrated (hotstep-v1)
+
+[Mothersuperior/open-rvq-encoder-minimax-music3-169m-pooled-v4](https://huggingface.co/Mothersuperior/open-rvq-encoder-minimax-music3-169m-pooled-v4)
+further-trained with a **pure soft-distillation objective** (soft cross-entropy against the teacher
+top-50 distributions shipped in the 8k corpus) plus **learnable per-head logit temperatures**, then
+selected on a **real-audio gate** that no other release in this lineage measures. Architecture,
+config, and file format are identical to the base — this is a drop-in replacement.
+
+Three things in this release are useful even if you never load the weights:
+
+## 1. Corpus consumption fixes (apply these before training on the 8k corpus)
+
+**uint16 wrap in `probs.npz` semantic ids.** The semantic head's `idx` values were stored as raw
+LM token ids (`code + 151675`) in a uint16, which silently wraps them mod 65536. They arrive as
+`code + 20603`. Subtract 20603 and mask anything outside `[0, 16384)` (the EOS token wraps to −5).
+Verified: after unwrapping, the sampled ground-truth code lands inside the teacher top-50 for
+~100% of frames; consumed naively, the ids are garbage and soft-distillation losses are meaningless.
+
+**Stitched-timeline latent alignment.** The corpus audio is stitched from 200-frame DiT windows at
+a 100-frame / **345-latent** hop (documented in the corpus card), while the naive uniform mapping
+(`latent = frame * 441 // 128` ≈ 344.53/100 frames) drifts ~1 latent per 800 frames. Measured
+effect: semantic top-1 against ground-truth codes decays from 0.41 at track start to 0.06 by frame
+2400 under the uniform rule. The correct frame→latent mapping for this corpus is:
+
+```python
+def stitched_lat(i):                      # frame index -> DAV latent index
+    k = 0 if i < 125 else (i - 25) // 100 # chunk ownership (from local frame 25)
+    return 345 * k + ((i - 100 * k) * 441) // 128 + 1
+```
+
+The `+1` is a constant global offset, verified empirically as best at every probed position on
+every probed record (top-1 flat at 0.4–0.6 across track positions once applied). The corpus's own
+DAV latents are exact — recomputing them from the shipped FLAC through the official `dav.pth`
+reproduces them at cosine 1.0 — the drift lives in the audio-vs-codes timeline of chunked
+rendering. **Real audio has no such drift**: the uniform mapping remains correct at inference.
+
+## 2. Calibrated logits (folded)
+
+The released v4/pooled-v4 readout logits are muP-sharp and uncalibrated: harmless for
+argmax/top-K use (what the reference adapter and the ComfyUI/diffusers nodes do), but
+cross-entropy and softmax probabilities computed from them are meaningless (semantic soft-CE ~59
+nats at temperature 1; per-frame logit std ~45). This release trained **learnable per-head
+temperatures** jointly with the model (initialised from a line search: ~16 semantic, ~2 acoustic)
+and **folded them into the head weights** at export. Argmax and top-K ordering are unchanged;
+`softmax(logits)` and CE are now meaningful out of the box. The final folded values are in
+`calibration.json`.
+
+## 3. The weights
+
+Training: pure soft-CE vs teacher top-50 (all 8 heads; depth decoder teacher-forced), 20,000 steps
+at effective batch 64 (16 × 4 grad-accum), lr 1e-4 cosine with 200-step warmup, bf16, single
+RTX 5090, on the drift-corrected + unwrapped 8k corpus (records < 8 s filtered; 7,726 train / 394
+holdout by salted-hash split). Trainer included in [`training/`](training/).
+
+### In-domain holdout (8k corpus, drift-corrected loader, calibrated CE)
+
+| holdout metric | pooled-v4 (start) | this release (step 18500) | change |
+|---|---|---|---|
+| semantic soft-CE (vs teacher) | 3.299 | 3.039 | −7.9% |
+| acoustic soft-CE | 4.618 | 4.413 | −4.4% |
+| semantic hard-CE | 2.336 | 2.154 | −7.8% |
+| acoustic hard-CE | 3.925 | 3.763 | −4.1% |
+| semantic top-1 | 41.9% | 42.2% | +0.3 pt |
+| acoustic top-1 | 16.8% | 17.6% | +0.8 pt |
+
+In-domain argmax barely moves — the interesting result is below.
+
+### Real-audio gate (13 commercial studio recordings, frozen official LM/depth decoder base-CE)
+
+The metric nobody else in the lineage reports: export codes from real recordings, measure the
+frozen official model's cross-entropy on them. Lower = the encoder's codes look more like codes
+the official model would emit. All prior rows reproduced on the same harness:
+
+| encoder | params | sem CE | sem t1 | ac CE | ac t1 |
+|---|---|---|---|---|---|
+| Serveurperso community 41M | 41M | 8.447 | 5.8% | 6.763 | 2.3% |
+| Mothersuperior pooled-v3 | 41M | 7.261 | 8.5% | 6.347 | 4.5% |
+| SimpleTuner v4 | 169M | 7.158 | 8.9% | 5.638 | 5.8% |
+| Mothersuperior pooled-v4 | 169M | 6.838 | 10.9% | 5.373 | 7.5% |
+| **this release** | **169M** | **5.866** | **16.6%** | **4.725** | **10.9%** |
+
+The largest single-step improvement in the lineage on this metric — and it happens at the
+argmax level (temperature folding cannot cause it), despite near-flat in-domain argmax
+accuracy. Soft distillation appears to act as regularisation against synthetic-audio overfit,
+transferring to real recordings. Caveat: semantic unique-code usage on the real-audio set drops
+~14% vs the base (6,490 vs 7,515 over 64k frames; acoustic unchanged) — mild concentration,
+far from collapse, but stated for transparency.
+
+Ear-verified on frame-matched replay renders of five real commercial recordings across four
+genres (indie rock, pop, alt-metal, electronic): an audible improvement over pooled-v4 on all
+of them, though subtler to the ear than the pooled-v4 → v4 step was — consistent with
+diminishing returns from a synthetic-only corpus at this size. Best in the lineage by both
+metric and ear; the remaining real-audio fidelity gap likely needs real-audio objectives
+(consistency training, LM-prior regularisation, reconstruction through the frozen generator)
+rather than more synthetic epochs.
+
+## Usage
+
+Identical to the base model. Load with SimpleTuner's
+[`minimax_music3_reference_adapter.py`](https://huggingface.co/SimpleTuner/open-rvq-encoder-minimax-music3/blob/main/minimax_music3_reference_adapter.py):
+
+```python
+from minimax_music3_reference_adapter import MiniMaxMusic3ReferenceAdapter
+adapter = MiniMaxMusic3ReferenceAdapter.from_files(
+    "rvq_encoder.safetensors", "rvq_encoder_config.json", "dav.pth")
+codes = adapter.predict_codes(waveform, sample_rate)   # [frames, 8]
+```
+
+`dav.pth` is the official DAV encoder from
+[MiniMaxAI/MiniMax-Music3](https://huggingface.co/MiniMaxAI/MiniMax-Music3).
+Strict `load_state_dict` — the checkpoint keys match the base exactly.
+
+## Further training
+
+[`training/rvq_distill_train.py`](training/rvq_distill_train.py) is the self-contained trainer
+used for this release: 8k-corpus zip layout, both consumption fixes above built in, warm start
+from any encoder in this lineage, pure-soft or mixed soft/hard loss, learnable temperatures with
+fold-at-save, holdout eval + best-checkpoint selection. Depends only on torch/numpy/safetensors
+plus the reference adapter module for the model classes.
+
+## Credits
+
+This encoder exists because of a chain of community work, each stage building on the last:
+
+- **[MiniMax](https://huggingface.co/MiniMaxAI/MiniMax-Music3)** — MiniMax Music 3 and the DAV
+  encoder. All weights and training data derive from it; use is subject to its terms.
+- **[Serveurperso](https://github.com/ServeurpersoCom)** — the original community proof: a 41M
+  single-GPU encoder demonstrating that exact token agreement was not required, and the first
+  independent derivation of the stitched-timeline alignment contract.
+- **[bghira / SimpleTuner](https://huggingface.co/SimpleTuner/open-rvq-encoder-minimax-music3)** —
+  the encoder architecture (v1–v4, causal depth decoder), the reference adapter, the trainer, and
+  the [reverse-distillation trace corpus](https://huggingface.co/datasets/bghira/minimax-music3-rvq-reverse-distillation).
+- **[Mothersuperior](https://huggingface.co/Mothersuperior)** — the
+  [8k distillation corpus](https://huggingface.co/datasets/Mothersuperior/minimax-music3-rvq-distill-corpus-8k)
+  with teacher top-50 distributions (which make this release's objective possible), and the
+  pooled-v3/pooled-v4 fine-tunes this release warm-starts from.
+- **[HOT-Step CPP](https://github.com/scragnog/HOT-Step-CPP)** — this release: the soft-distillation
+  training run, the corpus consumption fixes, logit calibration, and the real-audio gate.
+
+Use is subject to the MiniMax Music 3 model terms and the source datasets' terms.

calibration.json

@@ -0,0 +1 @@
+{"folded_taus": [8.8834, 1.3853, 1.4031, 1.3873, 1.4341, 1.4484, 1.4719, 1.4799]}

mup_base_shapes.bsh

@@ -0,0 +1,517 @@
+# This is a base shape file encoded in yaml
+# - `null` indicates a dimension is "finite", i.e. a non-"width" dimension
+# - a number indicates the base dimension of an "infinite" dimension, i.e. some notion of "width"
+blocks.0.conv1.bias:
+- 128
+blocks.0.conv1.weight:
+- 128
+- 128
+- null
+blocks.0.conv2.bias:
+- 128
+blocks.0.conv2.weight:
+- 128
+- 128
+- null
+blocks.0.norm.bias:
+- 128
+blocks.0.norm.weight:
+- 128
+blocks.1.conv1.bias:
+- 128
+blocks.1.conv1.weight:
+- 128
+- 128
+- null
+blocks.1.conv2.bias:
+- 128
+blocks.1.conv2.weight:
+- 128
+- 128
+- null
+blocks.1.norm.bias:
+- 128
+blocks.1.norm.weight:
+- 128
+blocks.2.conv1.bias:
+- 128
+blocks.2.conv1.weight:
+- 128
+- 128
+- null
+blocks.2.conv2.bias:
+- 128
+blocks.2.conv2.weight:
+- 128
+- 128
+- null
+blocks.2.norm.bias:
+- 128
+blocks.2.norm.weight:
+- 128
+conv_in.bias:
+- 128
+conv_in.weight:
+- 128
+- null
+- null
+depth_decoder.context_projection.weight:
+- null
+- 128
+depth_decoder.heads.0.bias:
+- null
+depth_decoder.heads.0.weight:
+- null
+- null
+depth_decoder.heads.1.bias:
+- null
+depth_decoder.heads.1.weight:
+- null
+- null
+depth_decoder.heads.2.bias:
+- null
+depth_decoder.heads.2.weight:
+- null
+- null
+depth_decoder.heads.3.bias:
+- null
+depth_decoder.heads.3.weight:
+- null
+- null
+depth_decoder.heads.4.bias:
+- null
+depth_decoder.heads.4.weight:
+- null
+- null
+depth_decoder.heads.5.bias:
+- null
+depth_decoder.heads.5.weight:
+- null
+- null
+depth_decoder.heads.6.bias:
+- null
+depth_decoder.heads.6.weight:
+- null
+- null
+depth_decoder.layers.0.k_proj.bias:
+- null
+depth_decoder.layers.0.k_proj.weight:
+- null
+- null
+depth_decoder.layers.0.linear1.bias:
+- null
+depth_decoder.layers.0.linear1.weight:
+- null
+- null
+depth_decoder.layers.0.linear2.bias:
+- null
+depth_decoder.layers.0.linear2.weight:
+- null
+- null
+depth_decoder.layers.0.norm1.bias:
+- null
+depth_decoder.layers.0.norm1.weight:
+- null
+depth_decoder.layers.0.norm2.bias:
+- null
+depth_decoder.layers.0.norm2.weight:
+- null
+depth_decoder.layers.0.out_proj.bias:
+- null
+depth_decoder.layers.0.out_proj.weight:
+- null
+- null
+depth_decoder.layers.0.q_proj.bias:
+- null
+depth_decoder.layers.0.q_proj.weight:
+- null
+- null
+depth_decoder.layers.0.v_proj.bias:
+- null
+depth_decoder.layers.0.v_proj.weight:
+- null
+- null
+depth_decoder.layers.1.k_proj.bias:
+- null
+depth_decoder.layers.1.k_proj.weight:
+- null
+- null
+depth_decoder.layers.1.linear1.bias:
+- null
+depth_decoder.layers.1.linear1.weight:
+- null
+- null
+depth_decoder.layers.1.linear2.bias:
+- null
+depth_decoder.layers.1.linear2.weight:
+- null
+- null
+depth_decoder.layers.1.norm1.bias:
+- null
+depth_decoder.layers.1.norm1.weight:
+- null
+depth_decoder.layers.1.norm2.bias:
+- null
+depth_decoder.layers.1.norm2.weight:
+- null
+depth_decoder.layers.1.out_proj.bias:
+- null
+depth_decoder.layers.1.out_proj.weight:
+- null
+- null
+depth_decoder.layers.1.q_proj.bias:
+- null
+depth_decoder.layers.1.q_proj.weight:
+- null
+- null
+depth_decoder.layers.1.v_proj.bias:
+- null
+depth_decoder.layers.1.v_proj.weight:
+- null
+- null
+depth_decoder.norm.bias:
+- null
+depth_decoder.norm.weight:
+- null
+depth_decoder.position:
+- null
+- null
+- null
+depth_decoder.prior_embeddings.0.weight:
+- null
+- null
+depth_decoder.prior_embeddings.1.weight:
+- null
+- null
+depth_decoder.prior_embeddings.2.weight:
+- null
+- null
+depth_decoder.prior_embeddings.3.weight:
+- null
+- null
+depth_decoder.prior_embeddings.4.weight:
+- null
+- null
+depth_decoder.prior_embeddings.5.weight:
+- null
+- null
+depth_decoder.prior_embeddings.6.weight:
+- null
+- null
+heads.0.bias:
+- null
+heads.0.weight:
+- null
+- 128
+norm_out.bias:
+- 128
+norm_out.weight:
+- 128
+position:
+- null
+- null
+- 128
+transformer.0.k_proj.bias:
+- 128
+transformer.0.k_proj.weight:
+- 128
+- 128
+transformer.0.linear1.bias:
+- 512
+transformer.0.linear1.weight:
+- 512
+- 128
+transformer.0.linear2.bias:
+- 128
+transformer.0.linear2.weight:
+- 128
+- 512
+transformer.0.norm1.bias:
+- 128
+transformer.0.norm1.weight:
+- 128
+transformer.0.norm2.bias:
+- 128
+transformer.0.norm2.weight:
+- 128
+transformer.0.out_proj.bias:
+- 128
+transformer.0.out_proj.weight:
+- 128
+- 128
+transformer.0.q_proj.bias:
+- 128
+transformer.0.q_proj.weight:
+- 128
+- 128
+transformer.0.v_proj.bias:
+- 128
+transformer.0.v_proj.weight:
+- 128
+- 128
+transformer.1.k_proj.bias:
+- 128
+transformer.1.k_proj.weight:
+- 128
+- 128
+transformer.1.linear1.bias:
+- 512
+transformer.1.linear1.weight:
+- 512
+- 128
+transformer.1.linear2.bias:
+- 128
+transformer.1.linear2.weight:
+- 128
+- 512
+transformer.1.norm1.bias:
+- 128
+transformer.1.norm1.weight:
+- 128
+transformer.1.norm2.bias:
+- 128
+transformer.1.norm2.weight:
+- 128
+transformer.1.out_proj.bias:
+- 128
+transformer.1.out_proj.weight:
+- 128
+- 128
+transformer.1.q_proj.bias:
+- 128
+transformer.1.q_proj.weight:
+- 128
+- 128
+transformer.1.v_proj.bias:
+- 128
+transformer.1.v_proj.weight:
+- 128
+- 128
+transformer.2.k_proj.bias:
+- 128
+transformer.2.k_proj.weight:
+- 128
+- 128
+transformer.2.linear1.bias:
+- 512
+transformer.2.linear1.weight:
+- 512
+- 128
+transformer.2.linear2.bias:
+- 128
+transformer.2.linear2.weight:
+- 128
+- 512
+transformer.2.norm1.bias:
+- 128
+transformer.2.norm1.weight:
+- 128
+transformer.2.norm2.bias:
+- 128
+transformer.2.norm2.weight:
+- 128
+transformer.2.out_proj.bias:
+- 128
+transformer.2.out_proj.weight:
+- 128
+- 128
+transformer.2.q_proj.bias:
+- 128
+transformer.2.q_proj.weight:
+- 128
+- 128
+transformer.2.v_proj.bias:
+- 128
+transformer.2.v_proj.weight:
+- 128
+- 128
+transformer.3.k_proj.bias:
+- 128
+transformer.3.k_proj.weight:
+- 128
+- 128
+transformer.3.linear1.bias:
+- 512
+transformer.3.linear1.weight:
+- 512
+- 128
+transformer.3.linear2.bias:
+- 128
+transformer.3.linear2.weight:
+- 128
+- 512
+transformer.3.norm1.bias:
+- 128
+transformer.3.norm1.weight:
+- 128
+transformer.3.norm2.bias:
+- 128
+transformer.3.norm2.weight:
+- 128
+transformer.3.out_proj.bias:
+- 128
+transformer.3.out_proj.weight:
+- 128
+- 128
+transformer.3.q_proj.bias:
+- 128
+transformer.3.q_proj.weight:
+- 128
+- 128
+transformer.3.v_proj.bias:
+- 128
+transformer.3.v_proj.weight:
+- 128
+- 128
+transformer.4.k_proj.bias:
+- 128
+transformer.4.k_proj.weight:
+- 128
+- 128
+transformer.4.linear1.bias:
+- 512
+transformer.4.linear1.weight:
+- 512
+- 128
+transformer.4.linear2.bias:
+- 128
+transformer.4.linear2.weight:
+- 128
+- 512
+transformer.4.norm1.bias:
+- 128
+transformer.4.norm1.weight:
+- 128
+transformer.4.norm2.bias:
+- 128
+transformer.4.norm2.weight:
+- 128
+transformer.4.out_proj.bias:
+- 128
+transformer.4.out_proj.weight:
+- 128
+- 128
+transformer.4.q_proj.bias:
+- 128
+transformer.4.q_proj.weight:
+- 128
+- 128
+transformer.4.v_proj.bias:
+- 128
+transformer.4.v_proj.weight:
+- 128
+- 128
+transformer.5.k_proj.bias:
+- 128
+transformer.5.k_proj.weight:
+- 128
+- 128
+transformer.5.linear1.bias:
+- 512
+transformer.5.linear1.weight:
+- 512
+- 128
+transformer.5.linear2.bias:
+- 128
+transformer.5.linear2.weight:
+- 128
+- 512
+transformer.5.norm1.bias:
+- 128
+transformer.5.norm1.weight:
+- 128
+transformer.5.norm2.bias:
+- 128
+transformer.5.norm2.weight:
+- 128
+transformer.5.out_proj.bias:
+- 128
+transformer.5.out_proj.weight:
+- 128
+- 128
+transformer.5.q_proj.bias:
+- 128
+transformer.5.q_proj.weight:
+- 128
+- 128
+transformer.5.v_proj.bias:
+- 128
+transformer.5.v_proj.weight:
+- 128
+- 128
+transformer.6.k_proj.bias:
+- 128
+transformer.6.k_proj.weight:
+- 128
+- 128
+transformer.6.linear1.bias:
+- 512
+transformer.6.linear1.weight:
+- 512
+- 128
+transformer.6.linear2.bias:
+- 128
+transformer.6.linear2.weight:
+- 128
+- 512
+transformer.6.norm1.bias:
+- 128
+transformer.6.norm1.weight:
+- 128
+transformer.6.norm2.bias:
+- 128
+transformer.6.norm2.weight:
+- 128
+transformer.6.out_proj.bias:
+- 128
+transformer.6.out_proj.weight:
+- 128
+- 128
+transformer.6.q_proj.bias:
+- 128
+transformer.6.q_proj.weight:
+- 128
+- 128
+transformer.6.v_proj.bias:
+- 128
+transformer.6.v_proj.weight:
+- 128
+- 128
+transformer.7.k_proj.bias:
+- 128
+transformer.7.k_proj.weight:
+- 128
+- 128
+transformer.7.linear1.bias:
+- 512
+transformer.7.linear1.weight:
+- 512
+- 128
+transformer.7.linear2.bias:
+- 128
+transformer.7.linear2.weight:
+- 128
+- 512
+transformer.7.norm1.bias:
+- 128
+transformer.7.norm1.weight:
+- 128
+transformer.7.norm2.bias:
+- 128
+transformer.7.norm2.weight:
+- 128
+transformer.7.out_proj.bias:
+- 128
+transformer.7.out_proj.weight:
+- 128
+- 128
+transformer.7.q_proj.bias:
+- 128
+transformer.7.q_proj.weight:
+- 128
+- 128
+transformer.7.v_proj.bias:
+- 128
+transformer.7.v_proj.weight:
+- 128
+- 128

rvq_encoder.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8e36bd36d58f467823f3e1e8c1fc4f9d3311e5527dd49768e8e7cfed68c8e31
+size 676055160

rvq_encoder_config.json

@@ -0,0 +1,34 @@
+{
+  "codebook_vocab_sizes": [
+    16384,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024
+  ],
+  "conv_dilations": [
+    1,
+    3,
+    9
+  ],
+  "d_model": 1088,
+  "depth_decoder": true,
+  "depth_decoder_dim": 512,
+  "depth_decoder_dropout": 0.1,
+  "depth_decoder_ff_mult": 4,
+  "depth_decoder_heads": 8,
+  "depth_decoder_layers": 2,
+  "dropout": 0.1,
+  "ff_mult": 4,
+  "latent_channels": 128,
+  "max_position_embeddings": 128,
+  "mup": true,
+  "mup_attention_multiplier": 8.0,
+  "mup_output_mult": 1.0,
+  "mup_readout_zero_init": true,
+  "num_heads": 17,
+  "num_layers": 8
+}

training/rvq_distill_train.py

@@ -0,0 +1,471 @@
+r"""Further-train a SimpleTuner-class RVQ encoder on the published distillation corpora.
+
+Warm-starts from a released checkpoint (default: Mothersuperior pooled-v4, 169M)
+and trains with SOFT DISTILLATION against the teacher top-50 distributions the
+corpus ships — the axis the published runs (hard-label CE, 12 epochs) never used.
+Checkpoints regularly; final selection happens OFFLINE on our real-audio gate
+(export_codes_v4.py --encoder-file <ckpt> + 1-step lm_sft_train2 [Base] CE),
+which nobody else measures.
+
+Corpus support:
+  * Mothersuperior/minimax-music3-rvq-distill-corpus-8k (PRIMARY — ships
+    precomputed fp16 DAV latents frame-major [F,128] @ 86.1328 Hz, plus
+    top-50 PROBS [T,8,50]). Records: <id>.{flac,codes.npy,probs.npz,vae.npy,json}
+    inside data/shard-NNN.zip; manifest.jsonl has align_ok/probs_ok flags.
+    Warm-up convention: codes/probs row 0 is un-emitted; frame i <-> row i+1.
+  * bghira/minimax-music3-rvq-reverse-distillation: NOT yet wired — ships
+    top-50 LOGITS but NO DAV latents, so it needs a ChunkedDAV precache pass
+    first (see docs/plans/2026-08-18-encoder-training-plan.md). Respect its
+    deterministic dataset_split when added.
+
+Key geometry: every 128-frame window spans EXACTLY 441 DAV latents
+(bounds[i] = floor(i*441/128); the +128 difference is always 441), so batches
+need no padding — which matters because GroupNorm(1, C) in the conv stack
+normalises over length and padding would shift its statistics.
+
+The depth decoder is TEACHER-FORCED here (ground-truth c0..c6 as priors, causal
+mask, head i read from position i+1) — the published forward() feeds its own
+argmax chain and is inference-only.
+
+Usage (single GPU, 5090):
+  python rvq_distill_train.py ^
+    --corpus M:\HOT-Step-CPP\_corpora\mm3-rvq-distill-8k ^
+    --encoder-dir M:\HOT-Step-CPP\_experiments\open-rvq-pooled-v4 ^
+    --out M:\HOT-Step-CPP\_experiments\rvq-train\pv4-softdistill-r1 ^
+    --steps 20000 --batch 16 --grad-accum 4 --lr 1e-4 --vram-frac 0.85
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import io
+import json
+import math
+import os
+import shutil
+import sys
+import time
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+
+ADAPTER_DIR_DEFAULT = r"M:\HOT-Step-CPP\_experiments\open-rvq-v4"
+WINDOW = 128
+LATENTS_PER_WINDOW = 441  # exact for every start index (441/128 ratio)
+TOPK = 50
+SEM_TOKEN_WRAP = 151_675 % 65_536  # = 20603; see Corpus8k.__getitem__
+
+
+def load_adapter_module(adapter_dir: str):
+    sys.path.insert(0, adapter_dir)
+    import minimax_music3_reference_adapter as ref
+
+    return ref
+
+
+# ---------------------------------------------------------------- corpus ----
+
+def holdout_of(record_id: str, frac: float, salt: str = "hotstep-holdout-v1") -> bool:
+    digest = hashlib.sha1(f"{salt}:{record_id}".encode()).digest()
+    return (int.from_bytes(digest[:8], "big") / 2**64) < frac
+
+
+def stitched_lat(i: int) -> int:
+    """Frame index -> latent index in the corpus's RENDERED audio.
+
+    The 8k corpus audio is stitched from 200-frame DiT windows at a 100-frame /
+    345-LATENT hop (README: 'constant stitched timeline'; non-first-chunk
+    ownership from local frame 25). 100 frames nominally map to 344.53 latents,
+    so the uniform i*441//128 mapping drifts ~1 latent per 800 and top-1 decays
+    to ~0 by frame 2400 — measured, then eliminated by this rule. The +1 is a
+    constant global offset, verified best at every position on every record
+    probed (t1 0.4-0.6 flat vs position). Real audio has NO such drift — the
+    uniform rule stays correct at inference/export time.
+    """
+    k = 0 if i < 125 else (i - 25) // 100
+    return 345 * k + ((i - 100 * k) * 441) // 128 + 1
+
+
+def pool_matrix(bounds: list[int]) -> torch.Tensor:
+    """Local copy of the adapter's build_pool_matrix so DataLoader workers can
+    unpickle this Dataset without the runtime-sys.path adapter import."""
+    origin = bounds[0]
+    local = [b - origin for b in bounds]
+    pool = torch.zeros((len(local) - 1, local[-1]), dtype=torch.float32)
+    for i, (a, b) in enumerate(zip(local[:-1], local[1:])):
+        pool[i, a:b] = 1.0 / (b - a)
+    return pool
+
+
+class Corpus8k(Dataset):
+    """Random 128-frame windows from the Mothersuperior 8k distill corpus."""
+
+    def __init__(self, root: str, records: list[dict], seed: int, fixed_windows: bool = False):
+        self.root = Path(root)
+        self.records = records
+        self.seed = seed
+        self.fixed_windows = fixed_windows
+        self._zips: dict[str, zipfile.ZipFile] = {}
+
+    def __len__(self):
+        return len(self.records)
+
+    def _zip(self, shard: str) -> zipfile.ZipFile:
+        z = self._zips.get(shard)
+        if z is None:
+            z = zipfile.ZipFile(self.root / "data" / shard)
+            self._zips[shard] = z
+        return z
+
+    def __getitem__(self, index: int):
+        rec = self.records[index]
+        z = self._zip(rec["shard"])
+        rid = rec["id"]
+
+        def read(ext):
+            return z.read(f"{rid}/{rid}{ext}")
+
+        meta = json.loads(read(".json"))
+        warmup = int(meta.get("codes_warmup_frames", 1))
+        codes = np.load(io.BytesIO(read(".codes.npy")))          # [T_total, 8] int32
+        vae = np.load(io.BytesIO(read(".vae.npy")))              # [F_dav, 128] fp16 frame-major
+        probs_z = np.load(io.BytesIO(read(".probs.npz")))
+        p_idx, p_val = probs_z["idx"], probs_z["prob"]           # [T_total, 8, 50]
+
+        emitted = codes.shape[0] - warmup
+        # window start s needs stitched-timeline latents [stitched_lat(s),
+        # stitched_lat(s)+441) and frames [s, s+128) emitted. stitched_lat runs
+        # slightly FASTER than uniform (345 per 100 frames), so walk max_start
+        # down until its window fits (a handful of iterations at most).
+        max_start = emitted - WINDOW
+        while max_start >= 0 and stitched_lat(max_start) + LATENTS_PER_WINDOW > vae.shape[0]:
+            max_start -= 1
+        if max_start < 0:
+            # belt-and-braces: manifest filtering should prevent this; fall
+            # back to a neighbour rather than killing the DataLoader worker.
+            return self[(index + 1) % len(self.records)]
+        if self.fixed_windows:
+            start = (max_start // 2) if index % 2 else 0
+        else:
+            g = np.random.default_rng(
+                (self.seed * 1_000_003 + index) ^ int.from_bytes(os.urandom(4), "big")
+            )
+            start = int(g.integers(0, max_start + 1))
+
+        lat_start = stitched_lat(start)
+        latents = vae[lat_start : lat_start + LATENTS_PER_WINDOW].astype(np.float32)
+
+        bounds = [(start + i) * 441 // 128 for i in range(WINDOW + 1)]
+        pool = pool_matrix(bounds)                               # [128, 441] fp32
+
+        rows = slice(start + warmup, start + warmup + WINDOW)
+        target_codes = codes[rows].astype(np.int64)              # [128, 8]
+        t_idx = p_idx[rows].astype(np.int64)                     # [128, 8, 50]
+        t_val = p_val[rows].astype(np.float32)
+        # Head 0's ids were stored as raw LM token ids (code + 151675) and the
+        # uint16 dtype wrapped them mod 65536 -> code + 20603 (verified: own
+        # code lands in the unwrapped top-50 for ~100% of frames). Unwrap, and
+        # zero out anything outside the semantic vocab (e.g. EOS 151670 -> -5).
+        sem = t_idx[:, 0, :] - SEM_TOKEN_WRAP
+        bad = (sem < 0) | (sem >= 16384)
+        sem[bad] = 0
+        t_val[:, 0, :][bad] = 0.0
+        t_idx[:, 0, :] = sem
+        t_val = t_val / np.clip(t_val.sum(-1, keepdims=True), 1e-8, None)
+
+        return (
+            torch.from_numpy(latents),                           # [441, 128]
+            pool,                                                # [128, 441]
+            torch.from_numpy(target_codes),
+            torch.from_numpy(t_idx),
+            torch.from_numpy(t_val),
+        )
+
+
+def load_manifest(root: str, min_duration_s: float = 8.0) -> list[dict]:
+    # min_duration_s: a 128-frame window needs 5.12 s of emitted codes plus
+    # stitched-latent margin; 8 s is conservative and drops only 24/8120
+    # early-EOS outliers (min in corpus: 3.7 s).
+    records = []
+    with open(Path(root) / "manifest.jsonl", encoding="utf-8") as fh:
+        for line in fh:
+            rec = json.loads(line)
+            if rec.get("align_ok") and rec.get("probs_ok") and rec.get("duration_s", 0) >= min_duration_s:
+                records.append(rec)
+    return records
+
+
+# ----------------------------------------------------------------- model ----
+
+def encoder_trunk(model, latents: torch.Tensor, pool: torch.Tensor) -> torch.Tensor:
+    """conv stack + pooled transformer -> per-frame hidden [B, 128, d_model]."""
+    h = model.conv_in(latents.transpose(1, 2))
+    for block in model.blocks:
+        h = block(h)
+    h = torch.bmm(pool.to(h.dtype), h.transpose(1, 2))
+    h = h + model.position[:, : pool.shape[1]].to(h.dtype)
+    layers = model.transformer if isinstance(model.transformer, torch.nn.ModuleList) else model.transformer.layers
+    for layer in layers:
+        h = layer(h)
+    return model.norm_out(h)
+
+
+def depth_teacher_forced(dd, frame_context: torch.Tensor, codes: torch.Tensor) -> list[torch.Tensor]:
+    """Parallel teacher-forced depth pass. codes [B, F, 8] ground truth.
+    Sequence = [ctx, e0(c0), .., e6(c6)]; causal mask; head i reads position i+1."""
+    batch, frames, _ = frame_context.shape
+    parts = [dd.context_projection(frame_context).flatten(0, 1).unsqueeze(1)]
+    for i in range(7):
+        parts.append(dd.prior_embeddings[i](codes[..., i]).flatten(0, 1).unsqueeze(1))
+    hidden = dd._decode(torch.cat(parts, dim=1))                 # [B*F, 8, D]
+    return [dd.heads[i](hidden[:, i + 1]).view(batch, frames, -1) for i in range(7)]
+
+
+def soft_ce(logits: torch.Tensor, t_idx: torch.Tensor, t_val: torch.Tensor) -> torch.Tensor:
+    """-sum p_teacher * log q_student over the teacher's top-50 support."""
+    logq = F.log_softmax(logits.float(), dim=-1)
+    return -(t_val * logq.gather(-1, t_idx)).sum(-1).mean()
+
+
+def compute_losses(model, batch, device, acoustic_weight: float, hard_weight: float, log_taus=None):
+    latents, pool, codes, t_idx, t_val = (x.to(device, non_blocking=True) for x in batch)
+    hidden = encoder_trunk(model, latents, pool)
+    sem_logits = model.heads[0](hidden)                          # [B, F, 16384]
+    ac_logits = depth_teacher_forced(model.depth_decoder, hidden, codes)
+    if log_taus is not None:
+        # Released checkpoints carry uncalibrated (muP-sharp) readouts —
+        # harmless for argmax/top-K, fatal for CE. Learnable per-head
+        # temperature; folded into head weights at save time.
+        taus = log_taus.exp()
+        sem_logits = sem_logits / taus[0]
+        ac_logits = [ac_logits[i] / taus[i + 1] for i in range(7)]
+
+    losses = {"sem_soft": soft_ce(sem_logits, t_idx[..., 0, :], t_val[..., 0, :])}
+    ac = [soft_ce(ac_logits[i], t_idx[..., i + 1, :], t_val[..., i + 1, :]) for i in range(7)]
+    losses["ac_soft"] = torch.stack(ac).mean()
+
+    if hard_weight > 0:
+        losses["sem_hard"] = F.cross_entropy(sem_logits.flatten(0, 1).float(), codes[..., 0].flatten())
+        ach = [
+            F.cross_entropy(ac_logits[i].flatten(0, 1).float(), codes[..., i + 1].flatten())
+            for i in range(7)
+        ]
+        losses["ac_hard"] = torch.stack(ach).mean()
+
+    total = losses["sem_soft"] + acoustic_weight * losses["ac_soft"]
+    if hard_weight > 0:
+        total = total + hard_weight * (losses["sem_hard"] + acoustic_weight * losses["ac_hard"])
+    losses["total"] = total
+
+    with torch.no_grad():
+        losses["sem_t1"] = (sem_logits.argmax(-1) == codes[..., 0]).float().mean()
+        losses["ac_t1"] = torch.stack(
+            [(ac_logits[i].argmax(-1) == codes[..., i + 1]).float().mean() for i in range(7)]
+        ).mean()
+    return losses
+
+
+@torch.no_grad()
+def evaluate(model, loader, device, acoustic_weight, autocast_ctx, log_taus=None):
+    model.eval()
+    sums: dict[str, float] = {}
+    n = 0
+    for batch in loader:
+        with autocast_ctx():
+            losses = compute_losses(model, batch, device, acoustic_weight, hard_weight=1.0, log_taus=log_taus)
+        for k, v in losses.items():
+            sums[k] = sums.get(k, 0.0) + float(v)
+        n += 1
+    model.train()
+    return {k: v / max(n, 1) for k, v in sums.items()}
+
+
+# ------------------------------------------------------------------ main ----
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--corpus", default=r"M:\HOT-Step-CPP\_corpora\mm3-rvq-distill-8k")
+    ap.add_argument("--encoder-dir", default=r"M:\HOT-Step-CPP\_experiments\open-rvq-pooled-v4")
+    ap.add_argument("--encoder-file", default="rvq_encoder.safetensors")
+    ap.add_argument("--config-file", default="rvq_encoder_config.json")
+    ap.add_argument("--adapter-dir", default=ADAPTER_DIR_DEFAULT)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--steps", type=int, default=20000)
+    ap.add_argument("--batch", type=int, default=16)
+    ap.add_argument("--grad-accum", type=int, default=4)
+    ap.add_argument("--lr", type=float, default=1e-4)
+    ap.add_argument("--warmup", type=int, default=200)
+    ap.add_argument("--weight-decay", type=float, default=0.01)
+    ap.add_argument("--acoustic-weight", type=float, default=1.0)
+    ap.add_argument("--hard-weight", type=float, default=0.0,
+                    help="mix in hard-label CE alongside soft distillation")
+    ap.add_argument("--sem-tau", type=float, default=16.0,
+                    help="initial semantic logit temperature (line-searched on pooled-v4)")
+    ap.add_argument("--ac-tau", type=float, default=2.0,
+                    help="initial acoustic logit temperature")
+    ap.add_argument("--holdout-frac", type=float, default=0.05)
+    ap.add_argument("--eval-every", type=int, default=500)
+    ap.add_argument("--save-every", type=int, default=1000)
+    ap.add_argument("--workers", type=int, default=2)
+    ap.add_argument("--seed", type=int, default=1)
+    ap.add_argument("--vram-frac", type=float, default=0.85)
+    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    ap.add_argument("--smoke", action="store_true", help="2 optimizer steps on CPU-sized batch, then exit")
+    args = ap.parse_args()
+
+    torch.manual_seed(args.seed)
+    device = torch.device(args.device)
+    if device.type == "cuda":
+        torch.cuda.set_per_process_memory_fraction(args.vram_frac)
+
+    ref = load_adapter_module(args.adapter_dir)
+    from safetensors.torch import load_file, save_file
+
+    enc_dir = Path(args.encoder_dir)
+    config = ref.RVQEncoderConfig.from_dict(
+        json.loads((enc_dir / args.config_file).read_text(encoding="utf-8"))
+    )
+    model = ref.MiniMaxMusicRVQEncoder(config)
+    model.load_state_dict(load_file(str(enc_dir / args.encoder_file)), strict=True)
+    model.to(device).train()
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"[model] {n_params/1e6:.1f}M params, warm start from {enc_dir / args.encoder_file}")
+
+    records = load_manifest(args.corpus)
+    train_recs = [r for r in records if not holdout_of(r["id"], args.holdout_frac)]
+    hold_recs = [r for r in records if holdout_of(r["id"], args.holdout_frac)]
+    print(f"[data] {len(records)} usable records -> {len(train_recs)} train / {len(hold_recs)} holdout")
+
+    train_ds = Corpus8k(args.corpus, train_recs, args.seed)
+    hold_ds = Corpus8k(args.corpus, hold_recs, args.seed, fixed_windows=True)
+    loader_kw = dict(
+        batch_size=args.batch,
+        num_workers=args.workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=args.workers > 0,
+    )
+    train_loader = DataLoader(train_ds, shuffle=True, drop_last=True, **loader_kw)
+    hold_loader = DataLoader(hold_ds, shuffle=False, **loader_kw)
+
+    log_taus = torch.nn.Parameter(
+        torch.log(torch.tensor([args.sem_tau] + [args.ac_tau] * 7, dtype=torch.float32, device=device))
+    )
+    decay, no_decay = [], []
+    for name, p in model.named_parameters():
+        (no_decay if p.ndim <= 1 or "position" in name else decay).append(p)
+    opt = torch.optim.AdamW(
+        [
+            {"params": decay, "weight_decay": args.weight_decay},
+            {"params": no_decay, "weight_decay": 0.0},
+            {"params": [log_taus], "weight_decay": 0.0},
+        ],
+        lr=args.lr, betas=(0.9, 0.95),
+    )
+
+    def lr_at(step):
+        if step < args.warmup:
+            return args.lr * (step + 1) / args.warmup
+        t = (step - args.warmup) / max(args.steps - args.warmup, 1)
+        return args.lr * 0.5 * (1 + math.cos(math.pi * min(t, 1.0)))
+
+    use_bf16 = device.type == "cuda"
+
+    def autocast_ctx():
+        return torch.autocast("cuda", dtype=torch.bfloat16) if use_bf16 else torch.autocast("cpu", enabled=False)
+
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(enc_dir / args.config_file, out / "rvq_encoder_config.json")
+    (out / "run_args.json").write_text(json.dumps(vars(args), indent=1), encoding="utf-8")
+    log_path = out / "train_log.jsonl"
+
+    def save_ckpt(tag: str):
+        ck = out / tag
+        ck.mkdir(exist_ok=True)
+        state = {k: v.detach().to(torch.float32).cpu() for k, v in model.state_dict().items()}
+        # Fold the learned temperatures into the readout heads so the saved
+        # file is a plain calibrated encoder (argmax/top-K ordering unchanged,
+        # CE now meaningful). heads.0 = semantic; depth_decoder.heads.i = c(i+1).
+        taus = log_taus.detach().exp().cpu()
+        for key in ("heads.0.weight", "heads.0.bias"):
+            state[key] = state[key] / taus[0]
+        for i in range(7):
+            for suffix in ("weight", "bias"):
+                key = f"depth_decoder.heads.{i}.{suffix}"
+                state[key] = state[key] / taus[i + 1]
+        save_file(state, str(ck / "rvq_encoder.safetensors"))
+        shutil.copyfile(out / "rvq_encoder_config.json", ck / "rvq_encoder_config.json")
+        (ck / "calibration.json").write_text(
+            json.dumps({"folded_taus": [round(float(t), 4) for t in taus]}), encoding="utf-8"
+        )
+        print(f"[ckpt] {ck} (taus folded: {[round(float(t), 2) for t in taus]})")
+
+    best_hold = float("inf")
+    step = 0
+    t0 = time.time()
+    opt.zero_grad(set_to_none=True)
+    data_iter = iter(train_loader)
+    while step < args.steps:
+        for micro in range(args.grad_accum):
+            try:
+                batch = next(data_iter)
+            except StopIteration:
+                data_iter = iter(train_loader)
+                batch = next(data_iter)
+            with autocast_ctx():
+                losses = compute_losses(
+                    model, batch, device, args.acoustic_weight, args.hard_weight, log_taus=log_taus
+                )
+            (losses["total"] / args.grad_accum).backward()
+        if not torch.isfinite(losses["total"]):
+            raise FloatingPointError(f"non-finite loss at step {step}: {losses}")
+        torch.nn.utils.clip_grad_norm_(list(model.parameters()) + [log_taus], 1.0)
+        for group in opt.param_groups:
+            group["lr"] = lr_at(step)
+        opt.step()
+        opt.zero_grad(set_to_none=True)
+        step += 1
+
+        if step % 25 == 0 or args.smoke:
+            row = {k: round(float(v.detach()), 4) for k, v in losses.items()}
+            row.update(
+                step=step, lr=round(lr_at(step), 8), elapsed_s=round(time.time() - t0, 1),
+                taus=[round(float(t), 2) for t in log_taus.detach().exp()],
+            )
+            print(f"[{step}/{args.steps}] " + " ".join(f"{k}={v}" for k, v in row.items() if k != "step"))
+            with open(log_path, "a", encoding="utf-8") as fh:
+                fh.write(json.dumps(row) + "\n")
+
+        if args.smoke and step >= 2:
+            print("[smoke] OK — forward/backward/step ran clean")
+            return
+
+        if step % args.eval_every == 0:
+            ev = evaluate(model, hold_loader, device, args.acoustic_weight, autocast_ctx, log_taus=log_taus)
+            row = {("hold_" + k): round(float(v), 4) for k, v in ev.items()}
+            row["step"] = step
+            print(f"[eval @{step}] " + " ".join(f"{k}={v}" for k, v in row.items() if k != "step"))
+            with open(log_path, "a", encoding="utf-8") as fh:
+                fh.write(json.dumps(row) + "\n")
+            score = ev["sem_hard"] + args.acoustic_weight * ev["ac_hard"]
+            if score < best_hold:
+                best_hold = score
+                save_ckpt("best_holdout")
+
+        if step % args.save_every == 0:
+            save_ckpt(f"step{step:06d}")
+
+    save_ckpt("final")
+    print(f"[done] {args.steps} steps in {(time.time() - t0)/3600:.2f} h; best holdout hard-CE sum {best_hold:.4f}")
+    print("[next] run each kept checkpoint through the REAL-AUDIO gate: "
+          "export_codes_v4.py --encoder-file <ckpt>\\rvq_encoder.safetensors "
+          "--config-file <ckpt>\\rvq_encoder_config.json, then the 1-step lm_sft_train2 [Base] CE.")
+
+
+if __name__ == "__main__":
+    main()