From 17 s/step to 2.76 s/step — pre-computed encoder embeddings for ACE-Step 1.5 LoRA training

The starting point — and the honest research note

The model trained in this research is ACE-Step v1 (ACE Studio and StepFun, Apache 2.0), still available in the app as the legacy training stack. The models ExpSoft ACE-Step 1.5 generates with in version 1.2.0.0 — ACE-Step 1.5 and XL, MIT — use a different DiT (4B for XL) and a different text encoder (Qwen3-Embedding), so the figures below do not transfer to them. The v1 architecture is, in broad strokes, a flow-matching DiT conditioned by three frozen encoders:

• UMT5 for textual lyrics and prompt (a multilingual T5 variant).
• MERT for musical-content features extracted from reference audio.
• m-HuBERT for vocal/timbre features extracted from reference audio.

During training, each audio clip in your dataset gets passed through all three encoders at every step (or rather, would, if we did things the naive way). The encoders are frozen: their weights never change. The training signal only updates the DiT (or, for our purposes, the small LoRA matrices attached to the DiT’s attention/MLP projections).

Before going further, one honest research note belongs at the top: a LoRA on this architecture cannot clone the timbre of an individual singer’s voice. The reason is structural rather than a missing setting — the m-HuBERT and singer-conditioning encoders are frozen and were trained jointly with the base DiT on a vocabulary of timbres baked in at pre-training time. A LoRA touching only the DiT can re-weight the model toward styles, genres, instrumentations, languages, structures, even broad vocal qualities the base model already saw — but it can’t introduce a new singer embedding. That’s a known limitation of any LoRA fine-tune on a frozen-encoder architecture; it’s not a bug in the trainer. The product page now states plainly that LoRA training is not supported in version 1.2.0.0.

With that out of the way: the seven decisions that turn this from a 24-hour training session into a four-hour overnight session.

Decision 1 — pre-compute the encoder embeddings, once, at dataset build time

If the three encoders are frozen, their output for a given input never changes between training steps. Running them on every step is doing the same computation 8 000 times in a row. The cleaner approach is to run them once, at dataset preparation time, and cache the output tensors next to each audio clip as a sidecar file.

Concretely: a one-shot prepare step goes through each unique track once — 24 000 dataset rows built from 12 tracks need 12 passes — and saves the four encoder outputs (MERT and m-HuBERT hidden states, text hidden states and text mask) to one .pt sidecar per track under precomputed_embeddings/. The dataset is flagged as precomputed and the trainer picks the sidecars up automatically. Measured on an RTX 5080: 17 s/step with live encoders, 2.76 s/step with the sidecars.

Two implementation details that matter for repeatability:

• One prompt per item. Upstream picks a random recaption per item at every step; with precomputed embeddings the prompt is frozen to the one chosen at prepare time. Datasets without recaptions are not affected.
• One-pass dataset build with a progress bar. A novice user would otherwise see the first training step take an hour because of the latent encode-on-first-access pattern, with no clear feedback. Doing it as an explicit “Prepare dataset” phase with its own progress bar makes the latency budget honest.

This single change is the biggest lever. It accounts for nearly all of the ≈ 6× speed-up.

Decision 2 — bf16-true (not bf16-mixed)

The default precision recipe for training on Ampere/Ada/Blackwell GPUs is mixed precision: parameters and gradients live in fp32, the forward and backward passes happen in bf16, and a fp32 master copy of the optimiser state is kept on the side. It’s safe — it handles loss scaling, NaN avoidance, and tiny-gradient underflow gracefully — and it costs you about 2× the VRAM of pure bf16, because fp32 weights and bf16 weights both exist.

For the v1 DiT, the difference is enormous. The DiT weights themselves are:

• ~13.6 GB in bf16-mixed (fp32 master weights + bf16 active copy + bf16 grads, very roughly).
• ~6.8 GB in bf16-true (parameters, grads and optimiser state all in bf16, with the optimiser still doing its accumulation in bf16).

For a 16 GB GPU, that’s the difference between “fits with a tight batch size of 1” and “fits with comfortable head-room for activations and a small batch”. The cost of going bf16-true is reduced numerical safety: training can occasionally produce NaN gradients on outlier batches, and the optimiser is less forgiving of large learning rates. The mitigation is the watchdog of decision 3 and a learning rate of 1e-4 or lower.

Decision 3 — a NaN watchdog around the optimiser step

With bf16-true, divergence is a real risk. The training loop is watched by a small callback:

1. Every 50 steps it checks the LoRA parameters with torch.isnan().any() — cheap, about 98 M parameters next to a multi-billion-parameter base model.
2. At the first NaN or Inf it stops the run and prints a clear message naming the faulty parameter and the remedies, first of all a lower learning rate.
3. It is always on; there is no switch to disable it.

The failure it catches is real: with bf16-true, a learning rate of 5e-4 and the upstream Adam settings, the first tests diverged to NaN around step 700–1 000, and without the watchdog the run kept going for hours, rendering silent samples. With it, a diverged run ends within 50 steps.

Decision 4 — LoRA-only checkpoints (190 MB instead of 10 GB)

A natural inclination is to save the whole training state at each checkpoint: base DiT weights + LoRA adapters + optimiser state. For the v1 model, that’s about 10 GB per checkpoint. Saving every 500 steps over a 10 000-step run is 200 GB on disk for one run. That’s untenable on a personal machine.

The good news is that the base DiT weights never change during LoRA training. They’re frozen — the same as the encoders. The only state worth saving is the LoRA adapters (about 98 million parameters in the research runs, ~190 MB in bf16) and, optionally, the optimiser state restricted to those LoRA parameters.

ACE-Step’s checkpoint format makes this clean: we save only the LoRA state-dict. On disk this is about 190 MB per checkpoint. Twenty checkpoints over a run = 3.8 GB. That fits. The base DiT weights are loaded fresh from the cached HuggingFace download at load time, and the LoRA adapters are merged on top.

This decision is what makes the “keep many intermediate checkpoints” workflow practical. In the research runs that meant saving often, rendering one sample per checkpoint and discarding the bad ones. The LORA tab still offers Full / LoRA / Both checkpoints, but LoRA training is not supported in version 1.2.0.0.

Decision 5 — VRAM page-out for in-training sampling

One of the friction points of any LoRA training workflow is being able to see what your model is learning, mid-training. The natural test is to generate a few samples every N steps with the current LoRA state and listen to them. The problem is that this requires having two things in VRAM at the same time:

• The training graph: ~12 GB on a 16 GB GPU for the v1 model in bf16-true.
• The inference graph: ~8 GB for the same DiT in inference mode (no grads, no optimiser).

12 GB + 8 GB = 20 GB. Doesn’t fit. The naive workaround is to drop the training state to disk, swap the inference state in, run the sample, swap back. That’s correct but slow (tens of seconds per sample).

The approach used here is page-out to system RAM rather than disk: at a sample step, a callback saves a temporary LoRA folder, moves the transformer, the audio autoencoder and the Adam state to the CPU and empties the CUDA cache (about 7 GB freed), starts a separate render process that loads its own inference pipeline (about 8 GB), then moves everything back and resumes training. Measured cost: about 3 to 5 minutes per sample — 23 s of diffusion and about 2 s of decoding, the rest being page-out / page-in and the start of the render pipeline.

Samples can be rendered every N steps without ending the run. On a 5 000-step run, rendering every 500 steps (10 samples) adds about 50 minutes to a run of about 4 hours.

Decision 6 — saved training profiles instead of one-size defaults

Rather than a configuration copied from a tutorial written for another GPU, the LORA tab ships saved training profiles that set every knob at once: rank, target modules, optimiser, learning rate, precision, audio cap, batch size and accumulation. The two reference recipes use batch 1, accumulation 1, bf16-true, gradient checkpointing and an 80-second audio cap, and say what to lower first on an out-of-memory error.

Decision 7 — keep the frozen encoders off the GPU from the very first step

The upstream trainer still loads the three encoders, and on a 16 GB card that alone filled the GPU (15 of 16 GB) before step 1. Offloading them in Lightning’s on_fit_start hook was too late: Lightning moves the whole module to the GPU before that hook runs.

The fix has three parts, each closing a distinct crash: offload the encoders in the earlier setup hook and detach them from the module tree so the move to the GPU skips them; once re-attached, turn their .to() / .cuda() calls into no-ops, because upstream moves the text encoder back to the GPU at every step; and cast their inputs to fp32, because the encoders stay in fp32 on the CPU while training runs in bf16-true. Without these three fixes, training this model on 16 GB was not possible.

What the research left out

• Multi-GPU training. Not attempted for ACE-Step LoRA training; every run in this note is single-GPU, on one RTX 5080.
• QLoRA / 8-bit base model. Would free memory at the risk of audible drift on the conditioning path; not tested. On the XL stack, upstream refuses to load a LoRA on a quantised model.
• Speaker-embedding LoRA (singer cloning). Already explained at the top — a structural limitation of the architecture, not a missing feature.

The twelve tabs of version 1.2.0.0

In version 1.2.0.0 the app is a song generator first; the training tabs stay for experimenters and are not supported.

1. WELCOME — what the app does, Quick Start, licences, requirements.
2. SETUP — one-click install of a portable Python 3.11, uv, Git, ACE-Step, PyTorch cu128 and the local lyrics engine, all inside the app folder.
3. GENERATE — a style description plus lyrics gives a song; XL (SFT), XL (Base), XL (Turbo), 2B (Turbo) and the older 2B.
4. REMIX — Cover, Inspire or Extend an existing track.
5. GALLERY — songs, LoRAs, datasets and lyrics, with right-click actions and one-click MP3.
6. LYRICS — a local model on the processor writes a first draft.
7. HELP — ready-to-paste lyrics templates.
8. DATASET — builds a training dataset (for experimenters, not supported).
9. LORA — the trainer this note is about (for experimenters, not supported).
10. DISK — what each folder takes and which ones are safe to clear.
11. SETTINGS — optional keys such as the Hugging Face token.
12. MORE TOOLS ? — version check and the ExpSoft products page.

Each tab has a single responsibility; there’s no global state machine and no “Run” button that does seven different things depending on which radio button was last selected.

Credits and license

Everything that makes this work belongs to the open-source communities behind it. ExpSoft’s contribution is the Windows-portable EasyInstaller around them and the engineering decisions described above. The pieces:

• ACE-Step (v1) — ACE Studio & StepFun, Apache 2.0. The v1 music DiT and the reference training scripts used in this research. ACE-Step 1.5 and XL, the models the app generates with in 1.2.0.0, are MIT. All credit for the architecture and weights belongs to the ACE Studio and StepFun teams.
• PyTorch — Meta + community, BSD-3. The training framework.
• Lightning — Lightning AI, Apache 2.0. The training loop scaffolding; the bf16-true precision plugin and the gradient-accumulation helpers in particular.
• uv + python-build-standalone — Astral; uv: MIT or Apache 2.0, python-build-standalone: MPL 2.0. Fast Python install on Windows without conda.
• Ollama — MIT. The local runner of the LYRICS tab’s lyrics model, on the processor.
• transformers + diffusers — HuggingFace, Apache 2.0. Model loading, tokenizers, scheduler implementations.

The C# host is .NET 8 WPF (MIT, Microsoft). The cross-app branding, the Welcome content and the engineering work described above are original ExpSoft contributions. None of this would exist without the ACE-Step team’s decision to release the model and the training scripts under an OSI-approved licence.

Take-aways

• For a frozen-encoder architecture, the biggest single training-speed lever is pre-computing the encoder outputs at dataset build time. Everything downstream gets faster by default, and VRAM previously held by encoders is freed.
• bf16-true with a NaN watchdog beats bf16-mixed for VRAM-constrained training. The watchdog cost is negligible; the VRAM saved unlocks a whole tier of consumer GPUs.
• LoRA-only checkpoint serialisation makes the “save often, browse often, prune later” workflow practical on a personal SSD.
• In-training sampling via VRAM page-out (not disk page-out) is the difference between “sample every 500 steps” being a real workflow vs a debugging-only luxury.
• Be honest about what a LoRA can and can’t do on a given architecture. Frozen-encoder architectures bound the space of fine-tunable behaviours; saying so up front saves users hours of frustration.
• A faster trainer is not yet a usable result: in version 1.2.0.0 LoRA training stays unsupported, because in the author’s tests it does not give usable results yet.