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

ACE-Step 1.5 is a 5-billion-parameter diffusion transformer (DiT) for music generation, published by ACE Studio and StepFun in 2024-2025 under an Apache 2.0 licence. The 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 repeats this in the legal/expectations section.

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: for each track_001.flac in the dataset, the prepare step writes a track_001.umt5.pt, a track_001.mert.pt and a track_001.hubert.pt. The shapes are small relative to raw audio: a typical 3-minute clip yields embeddings on the order of a few hundred KB per encoder. The disk overhead is roughly +10-20% on top of the audio dataset; the training-step cost goes from 17 s/step to about 5 s/step on the first pass, and the GPU memory previously held by the three encoders (about 4 GB combined in bf16) is now free.

Two implementation details that matter for repeatability:

• Cache invalidation by content hash. The prepare step hashes the audio bytes + the encoder version + the encoder config to derive a cache key. Change the encoder version or the audio file, the cache is recomputed; everything else, it’s reused. This is the difference between “fast on the second epoch” (every epoch re-reads the cached .pt files) and “fast forever” (the cache survives across training runs, dataset re-splits, model versions).
• 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. About 5× of the total speedup comes from here.

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 a 5B-parameter 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 comes from decisions 3 and 4 below.

Decision 3 — a NaN watchdog around the optimiser step

With bf16-true, occasional NaN gradients are real. The training loop wraps every optimiser step with a small watchdog:

1. After the backward pass, the watchdog checks torch.isnan(grad).any() across the LoRA parameter set (cheap — these are tiny tensors compared to the base model).
2. If any NaN is detected, the optimiser step is skipped: gradients are zeroed, no parameter update happens, and the offending batch is logged.
3. The training loop continues with the next batch. A counter tracks consecutive skipped steps; if it crosses a threshold (default 20), the run is paused and the user is notified — at that point something is structurally wrong with the data or the learning rate, not just one bad batch.

In practice, the NaN rate on a clean dataset with a sane learning rate is well under 1% of steps. The watchdog is cheap and the consequence of not having it is a single bad batch poisoning all subsequent steps via the optimiser’s momentum/variance state. With the watchdog, the cost of bf16-true is small and bounded.

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 a 5B-parameter 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 (a few million parameters, ~20 MB per layer set) 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. Users can save every 200 steps, browse them via the in-app sample generator (which loads the LoRA into the base model and produces a 30-second preview), and discard the bad ones.

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 ACE-Step 1.5 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 faster approach used here is page-out to system RAM rather than disk: the training optimiser state and the base DiT weights are .to("cpu")’d, the LoRA state is kept (it’s tiny), the inference state is .to("cuda")’d, the sample is generated (a 30-second clip takes ~30 seconds), then the swap is reversed. With pinned memory and a PCIe 4.0 ×16 link, each direction takes about 2-3 seconds. The whole sample cycle is ~45 seconds, compared to ~3 minutes for the disk-based version.

The user gets in-training samples every 500 steps without dropping out of the training loop. The continuous training run plus periodic sampling fits cleanly into the original 4-hour overnight envelope.

Decision 6 — VRAM-aware batch and gradient-accumulation auto-tuning

The setup phase detects the user’s GPU via nvidia-smi and the local CUDA build, reads the total VRAM, and proposes sane defaults for batch size and gradient accumulation steps. On 16 GB (RTX 4080 / 5080) the default is batch=1 with grad_accum=4, giving an effective batch of 4. On 24 GB (RTX 4090 / 3090) the default jumps to batch=2, grad_accum=2. On 12 GB (4070 Ti / 3080) it falls back to batch=1, grad_accum=8 and warns that training will be slower per wall-clock step but the same per effective-batch.

The point isn’t the specific numbers — they’re tuned per release and easy to override in the Settings tab. The point is that the user doesn’t see an OOM on first run because they tried the defaults from a tutorial written for a different GPU. A first-launch heuristic that’s right 90% of the time is much better than a one-size config and a half-page of post-mortem on why it crashed.

Decision 7 — encoder offload during the first three training steps

This one is small but it solves a specific OOM that bites on tight VRAM budgets. The first three steps of any training run are particularly memory-heavy because:

• PyTorch is still allocating its caching allocator buckets and hasn’t yet learned the steady-state footprint.
• The autograd graph’s memory layout is being established and may briefly hold redundant tensors.
• CUDA kernels are being compiled JIT for the specific shapes seen for the first time.

For an otherwise-stable run, these three steps can be where you hit a 16.1 GB peak that the steady state never approaches again. The fix: during steps 0, 1 and 2, the three frozen encoders (if they’re in VRAM at all — which they shouldn’t be after decision 1, but the safety net exists) are explicitly .cpu()’d at the start of the step and only moved back if needed. From step 3 onward the steady state is well-behaved and the conservative offload is dropped.

This costs maybe 100 ms over the first three steps and saves a class of first-launch OOMs on machines right at the VRAM boundary.

What didn’t make the V1 cut

• Multi-GPU training. Pipeline parallelism for ACE-Step 1.5 LoRA isn’t shipped yet (the splitter is product-of-DiffusionForce’s engineering work and the conditioning graph for ACE-Step is meaningfully different from a pure image DiT). For now, single-GPU training is the supported path. Most users get their 4-hour overnight run on a single 4080/5080-class card and call it a day.
• QLoRA / 8-bit base model. Quantising the base DiT to 8-bit would free another ~3 GB. The trade-off is non-trivial numerical drift on the conditioning pathway, which on a music model can manifest as audible distortions in the sample previews. Left for a later release after more testing.
• Speaker-embedding LoRA (singer cloning). Already explained at the top — this is a structural limitation of the architecture, not a missing feature. The skill section about training expectations is upfront with users about it.

The 11-tab installer — why a single-purpose app needs 11 tabs

One thing first-time users sometimes ask is why the installer has 11 tabs. The short answer is that ACE-Step 1.5 is not a single tool — it’s an environment for an entire LoRA-training and inference workflow on Windows, and each tab is a workflow phase:

1. Welcome — licences, expectations, links.
2. Setup — Python via uv, venv, PyTorch CUDA build (matched to detected GPU series), ACE-Step repo clone, model weights from HuggingFace, FFmpeg essentials, Ollama for the local LLM tagger.
3. Inference — generate a clip from a prompt or a reference, with the trained LoRA optionally loaded.
4. Dataset — drag audio clips in, set per-track metadata, run the prepare phase (= run the encoders, write the sidecar .pt files).
5. Auto-tag — optional local LLM tagging via Ollama, for users who don’t want to write 200 prompts by hand.
6. Train — the actual training loop with live loss/grad/lr graphs and the in-training sample player.
7. Checkpoints — browse the saved LoRA checkpoints, generate a sample from any one of them.
8. Export — merge a LoRA into the base model for distribution, or export it as a standalone .safetensors.
9. Settings — per-run overrides for batch, grad-accum, lr, scheduler, sample frequency.
10. Logs — the full Python stdout/stderr for debugging.
11. More Tools — the cross-app discoverability tab embedded via WebView2.

The 11-tab layout is what lets the rest of the app stay simple. 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 — ACE Studio × StepFun, Apache 2.0. The 5B DiT music generation model and its reference training scripts. 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, MIT. Fast Python install on Windows without conda.
• Ollama — MIT. The local LLM runner used by the optional auto-tag tab.
• transformers + diffusers — HuggingFace, Apache 2.0. Model loading, tokenizers, scheduler implementations.
• onnxruntime — Microsoft, MIT. Used for some of the encoder inference paths on CPU page-out.
• FFmpeg — LGPL. Audio I/O for the dataset prepare phase.

The C# host is .NET 8 WPF (MIT, Microsoft). The cross-app branding, the Welcome content, the multilingual legal section, 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.