Windows ROCm (RX 9070 XT) fix: two crashes on a clean venv_rocm install — vector_quantize_pytorch dist_nn import + torchao/quantization

[EmmanuelWorrell]

Title

README.md

Post

Posting this in case it saves someone else the troubleshooting time — hit two separate crashes getting ACE-Step 1.5 running on Windows + ROCm, both reproducible on a completely clean venv_rocm following the official requirements-rocm.txt install path. Neither is specific to my card; both trace back to gaps in AMD's Windows ROCm torch build itself.

Setup: RX 9070 XT, Windows 11, ROCm SDK 7.2.26024, torch 2.9.1+rocmsdk20260116, Python 3.12.

Crash 1: vector_quantize_pytorch — cannot import group from torch.distributed

ImportError: cannot import name 'group' from 'torch.distributed'
  ... vector_quantize_pytorch/lookup_free_quantization.py, line 15
      from torch.distributed import nn as dist_nn

vector_quantize_pytorch unconditionally imports torch.distributed.nn at module load, even though ACE-Step only needs it for single-GPU inference. The Windows ROCm torch build ships an incomplete torch.distributed (nn/group aren't there), so the import fails and takes model loading down with it. The package's own is_distributed() / maybe_distributed_mean() helpers already guard their usage of dist/dist_nn, so it's safe to just let the import fail quietly:

try:
    from torch.distributed import nn as dist_nn
except (ImportError, AttributeError):
    dist_nn = None

Saw the identical trace reported on a 9060 XT in #644, so this isn't card-specific.

Crash 2: torchao required for quantization — and can't actually be installed on Windows ROCm

ImportError: torchao is required for quantization but is not installed. Please install torchao to use quantization features.
  ... acestep/handler.py, line 656

ACE-Step auto-picks a quantization mode for some GPU memory tiers (mine, 16GB-class, was one), and handler.py treats missing torchao as fatal instead of falling back.

Important: don't just pip install torchao — it makes things worse. torchao's float8 path eagerly imports torch.distributed._functional_collectives, which needs torch._C._distributed_c10d — a native compiled component that's simply missing from the Windows ROCm torch build. This crashes even earlier than the original error (at diffusers import time):

ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package

There's no Windows ROCm build of torchao — it's built around CUDA kernels. The working fix is to catch the missing import in handler.py and just disable quantization:

if self.quantization is not None:
    try:
        import torchao
        assert compile_model, "Quantization requires compile_model to be True"
    except ImportError:
        print("[ACE-Step] torchao unavailable - disabling quantization and continuing without it.")
        self.quantization = None

Saw several people hit the missing-torchao error and try to force-install it in #55, #84, #124 — none mentioned the deeper native torch._C._distributed_c10d failure that causes on Windows ROCm specifically, so flagging that explicitly here.

Patch script

Wrote a small PowerShell script that applies both fixes to a venv_rocm + repo checkout. It's idempotent (detects and skips anything already patched) and backs up originals to <filename>.bak:

# from the ACE-Step-1.5 project root
Unblock-File -Path .\apply_rocm_patches.ps1   # first run only if downloaded via browser
.\apply_rocm_patches.ps1

Re-run it after recreating venv_rocm, after pip install -r requirements-rocm.txt, or after a git pull (fix #2 lives in handler.py, so an update can overwrite it).

Full script + writeup: [link to your gist/repo]

Happy to open this as a proper PR if a maintainer points me at the right approach for the quantization-tier auto-selection logic — the patch above is a workaround (quantization is disabled entirely on ROCm, not made to work), not a real fix for that path.
apply_rocm_patches.zip

[EmmanuelWorrell]

Title

Update: third Windows ROCm fix (LoRA/LoKr preprocessing) + correction on the quantization patch

Post

Follow-up to my earlier post on Windows ROCm crashes (vector_quantize_pytorch dist_nn + torchao/quantization). Two updates since then: a third fix needed specifically for LoRA/LoKr training, and a correction on the quantization patch that's worth flagging so nobody wastes time on it unnecessarily.

Setup, unchanged: RX 9070 XT, Windows 11, ROCm SDK 7.2.26024, torch 2.9.1+rocmsdk20260116, Python 3.12.

Correction: the quantization/torchao patch may no longer be needed

After pulling latest main, the fatal ImportError I originally hit in acestep/handler.py no longer reproduces the same way — that logic appears to have moved to acestep/core/generation/handler/init_service_setup.py (_validate_quantization_setup), and auto-tier quantization selection doesn't seem to trigger it under current main. I'm leaving the old patch in my combined script as a no-op fallback for anyone on an older checkout, but if you're on a recent main, you likely don't need it at all. Worth double-checking your own logs before assuming you need this one.

Still true and worth repeating: don't manually pip install torchao. It doesn't just fail to help — it actively breaks things worse. torchao's float8 code path eagerly imports torch.distributed._functional_collectives, which requires torch._C._distributed_c10d — a native compiled component that's simply absent from the Windows ROCm torch build. This crashes even earlier than the original error, at diffusers import time:

ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package

New: LoRA/LoKr dataset preprocessing crash

Hit this while trying to build a training dataset (Preprocess to Tensors step):

ImportError: TorchCodec is required for load_with_torchcodec. Please install torchcodec to use this function.
  ... torchaudio/_torchcodec.py, line 84, in load_with_torchcodec
  ... acestep/training/dataset_builder_modules/preprocess_audio.py, line 6, in load_audio_stereo
      audio, sr = torchaudio.load(audio_path)

Same underlying cause as the other torchcodec-related issues on ROCm: torchaudio.load() routes through torchcodec by default, and torchcodec isn't available on AMD ROCm at all (it's built around CUDA-specific dependencies). ACE-Step already has a soundfile fallback for audio export on ROCm elsewhere in the codebase (see the merged PR #1242 for XPU/ROCm export handling) — this same fallback pattern just wasn't applied to this particular load path in the training preprocessor.

Fix — acestep/training/dataset_builder_modules/preprocess_audio.py:

import torch
import torchaudio
import soundfile as sf


def load_audio_stereo(audio_path: str, target_sample_rate: int, max_duration: float):
    """Load audio, resample, convert to stereo, and truncate."""
    try:
        audio, sr = torchaudio.load(audio_path)
    except Exception:
        data, sr = sf.read(audio_path, always_2d=True, dtype="float32")
        audio = torch.from_numpy(data.T)
    if sr != target_sample_rate:
        resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
        audio = resampler(audio)
    if audio.shape[0] == 1:
        audio = audio.repeat(2, 1)
    elif audio.shape[0] > 2:
        audio = audio[:2, :]
    max_samples = int(max_duration * target_sample_rate)
    if audio.shape[1] > max_samples:
        audio = audio[:, :max_samples]
    return audio, sr

With this in place, LoRA and LoKr training both ran end-to-end on this card — training completed, export worked, and the adapter loaded correctly for generation (see note below on a confusing-but-harmless log warning you'll likely see).

Side note: confusing (but harmless) LoKr load warning

If you train and load a LoKr adapter, you'll likely see this in the log and think it failed:

WARNING - No adapter names discovered from decoder; LoRA registry will be empty.
INFO - LoRA adapter 'lokr_weights.safetensors' loaded from ... (adapters=[], targets=0)

This is a false alarm. That message comes from the PEFT-LoRA-specific adapter registry (registry_builder.py), which LoKr adapters never populate by design — LoKr scaling goes through a separate LyCORIS-specific path (_toggle_lokr → set_multiplier) that bypasses this registry entirely. Confirm via the LoKr multiplier set to <value> line elsewhere in the log instead — a real load failure would show LoKr adapter type set but no _lycoris_net found for scale, not the registry warning. Might be worth a small log-message tweak upstream so this doesn't read as an error to anyone training LoKr for the first time.

Updated patch script

All three fixes (dist_nn, quantization fallback as a no-op-safe legacy patch, preprocess_audio) are now in one idempotent script — detects what's already applied/not applicable and skips accordingly, backs up originals before touching anything: [link to your gist/repo]

Unblock-File -Path .\apply_rocm_patches.ps1   # first run only if downloaded via browser
.\apply_rocm_patches.ps1

Final Patches for Ace Step.zip
README.md