Train LoRAs on AMD GPUs with Musubi Tuner

Musubi Tuner is a set of scripts for training LoRA (Low-Rank Adaptation) models with HunyuanVideo, Wan2.1/2.2, and many other image and video models.

The project documentation is excellent, but the default install commands assume NVIDIA + CUDA. This is an installation guide for AMD GPU / APU users.

I've trained character LoRAs for the Krea 2 image model on an AMD Ryzen Strix Halo (AI Max+ 395) running Fedora 44. I'm using the latest ROCm 7.14 here. In my tests, it was ~9% faster compared to ROCm 7.2 on identical tasks.

Other AMD chipsets, distros, WSL, etc. follow a similar setup.

This is only an installation/setup guide. For actual training, refer to Musubi's official docs. They're seriously great.

¶Prerequisites

• Python 3.10 - 3.12
• User should be in the video and/or render group. Check your distro's docs for how to do this.

¶Find your GPU architecture

Identify your GPU's gfx target. This determines which PyTorch device package to install. Use the table below, or refer to AMD's official PyTorch install guide.

¶Device tag reference

Run rocm-smi --showproductname, match the GFX Version to the table below, and note it down. You will need this while installing PyTorch.

Radeon (discrete GPU)

Ryzen AI (integrated / APU)

¶Set up Musubi Tuner

Clone the repo and create a venv:

git clone https://github.com/kohya-ss/musubi-tuner
cd musubi-tuner

python3 -m venv .venv
source .venv/bin/activate
pip install -U pip wheel

¶Install ROCm PyTorch

Replace device-gfx1151 with your GPU's device tag:

pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \
  "torch[device-gfx1151]==2.12.0+rocm7.14.0" \
  "torchvision[device-gfx1151]==0.27.0+rocm7.14.0"

¶Install the rest of the dependencies

We use --no-deps everywhere because PyPI's default torch is the CUDA build. Without it, pip might replace your ROCm torch with a CUDA version.

The versions below match Musubi's pyproject.toml at the time of writing. Check that file to confirm they're up to date before running.

pip install -e . --no-deps
pip install --no-deps accelerate==1.6.0 av==14.0.1 diffusers==0.32.1 \
  einops==0.7.0 huggingface-hub==0.34.3 opencv-python==4.10.0.84 \
  "pillow>=11.3.0" safetensors==0.4.5 toml==0.10.2 tqdm==4.67.1 \
  transformers==4.57.6 voluptuous==0.15.2 ftfy==6.3.1 easydict==1.13 \
  sentencepiece==0.2.1 psutil pyyaml numpy filelock fsspec \
  typing-extensions requests regex "tokenizers>=0.22.0,<=0.23.0" \
  importlib-metadata bitsandbytes urllib3 idna charset-normalizer certifi \
  wcwidth

Reinstall ROCm torch as a final safety measure (in case pip replaced it):

pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ --no-deps \
  "torch[device-gfx1151]==2.12.0+rocm7.14.0" \
  "torchvision[device-gfx1151]==0.27.0+rocm7.14.0"

Verify installation

python3 -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))"

You should see True and your AMD GPU name. The version string should end with +rocm7.14.0, not +cu.

¶Accelerate config

accelerate is the training launcher used by Musubi Tuner. It writes a config file that tells it your GPU setup and preferred precision.

accelerate config

Answer the questions as follows for a single-GPU AMD setup:

• Compute environment: This machine
• Machine type: No distributed training
• CPU only? NO
• Torch dynamo? NO
• DeepSpeed? NO
• GPU(s) by id: all (or 0 if you get fp16 mixed precision requires a GPU)
• NUMA efficiency? NO
• Mixed precision: bf16

¶Set environment variables

Export the following environment variables:

export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export TORCH_BLAS_PREFER_HIPBLASLT=1

• TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 enables AOTriton (ahead-of-time compiled kernels) for ROCm, needed for torch.compile and SDPA
• TORCH_BLAS_PREFER_HIPBLASLT=1 prefers hipBLASLt for matrix operations (faster on modern AMD GPUs)

If you plan to train LoRAs frequently, add the variables to your ~/.bashrc or ~/.zshrc because they're required for every training session:

echo 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1' >> ~/.bashrc
echo 'export TORCH_BLAS_PREFER_HIPBLASLT=1' >> ~/.bashrc
source ~/.bashrc

¶Training options

On AMD, use --sdpa for attention. FlashAttention, xformers, and SageAttention are (mostly) untested/unstable on ROCm.

You have two options depending on your VRAM:

¶Mode 1: Balanced (any AMD GPU)

Use this when you have limited VRAM (<48 GB) or want headroom for other applications.

--sdpa --compile --compile_cache_size_limit 32 --gradient_checkpointing

Other than the required --sdpa flag, we're using the following:

• --compile --compile_cache_size_limit 32 : JIT-compiles the model for faster training. Uses extra memory but significantly reduces iteration time
• --gradient_checkpointing : Slower, but uses less VRAM. Use on GPUs with less than 32 GB VRAM

In my benchmarks on Krea 2, measured on a Strix Halo with 31 images at 1024 px and 30 steps, I got:

• Speed: ~21 s/iteration
• VRAM usage: ~32 GB

¶Mode 2: Fast (large VRAM setups, recommended for Strix Halo)

On Strix Halo and other high-VRAM setups, we can trade high VRAM usage for better speed by skipping --gradient_checkpointing:

--sdpa --compile --compile_cache_size_limit 32

In my benchmarks on Krea 2, measured on a Strix Halo with 31 images at 1024 px and 30 steps, I got:

• Speed: ~15 s/iteration (~28% faster 🎉)
• VRAM usage: ~62 GB (~2x higher 🥹)

Optional but recommended: Warm up the Triton cache

On Strix Halo and other high-VRAM setups, you can run without a cache warmup. However, in my tests, the very first run in any mode without warming up lowers the speed and raises VRAM usage, causing an OOM even on 128 GB VRAM!

The warmup script below loads the Krea 2 DiT, compiles it with torch.compile, and runs one forward pass to warm the Triton cache. Adapt it to the models and folders you use:

#!/usr/bin/env python3
"""Warm up the Triton compile cache for torch.compile (bare metal, no container).

Run once per machine before training without --gradient_checkpointing.
~40-60 seconds (mostly model loading).
The Triton cache persists across reboots — you only need to run this once.
"""

import os
import sys
import torch
import argparse
import logging

logging.basicConfig(level=logging.INFO)

# Set environment variables
os.environ["TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL"] = "1"
os.environ["TORCH_BLAS_PREFER_HIPBLASLT"] = "1"

# Find repo root (script is in repo root)
REPO_ROOT = os.path.dirname(os.path.abspath(__file__))
os.chdir(REPO_ROOT)

from musubi_tuner.krea2.krea2_utils import load_krea2_dit, single_mmdit_large_wide
from musubi_tuner.utils.model_utils import compile_transformer

print("Loading DiT model...")
dit = load_krea2_dit(
    "models/dit/raw.safetensors",
    device="cuda",
    dtype=torch.bfloat16,
    config=single_mmdit_large_wide,
)

print("Compiling with torch.compile...")
args = argparse.Namespace(
    compile_backend="inductor",
    compile_mode="default",
    compile_dynamic=None,
    compile_fullgraph=False,
    compile_cache_size_limit=32,
)
compile_transformer(args, dit, [dit.blocks], disable_linear=False)

print("Running forward pass to warm Triton cache...")
patch, c = 2, 16
h, w = 64, 64  # 1024x1024 latent (1024/8/patch)
img = torch.randn(1, h * w, c * patch * patch, device="cuda", dtype=torch.bfloat16)
ctx = torch.randn(1, 77, 12, 2560, device="cuda", dtype=torch.bfloat16)
t = torch.rand(1, device="cuda", dtype=torch.bfloat16)
pos = torch.zeros(1, h * w + 77, 3, device="cuda")
mask = torch.ones(1, h * w + 77, device="cuda", dtype=torch.bool)
dit(img=img, context=ctx, t=t, pos=pos, mask=mask)

print("Triton cache warm. You can now train without --gradient_checkpointing.")

Save it as ./amd_compile_warmup.py and run:

python3 ./amd_compile_warmup.py

¶Start training

You're all set up now. I won't go into model downloads and training instructions. Refer to Musubi's official docs for that, they're great.

You might see a rocSHMEM Could not open libnuma warning when you start training. This is only relevant for multi-GPU setups. Single GPU / APU users can safely ignore it.

¶Training example

This is the command I ran to successfully train multiple character LoRAs for Krea 2. Adapt it to your needs:

accelerate launch --num_cpu_threads_per_process 1 \
  src/musubi_tuner/krea2_train_network.py \
  --dit models/dit/raw.safetensors \
  --vae models/vae/qwen_image_vae.safetensors \
  --dataset_config configs/your_dataset.toml \
  --sdpa --mixed_precision bf16 \
  --compile --compile_cache_size_limit 32 \
  --timestep_sampling krea2_shift --weighting_scheme none \
  --optimizer_type adamw8bit --learning_rate 1e-4 \
  --max_data_loader_n_workers 2 --persistent_data_loader_workers \
  --network_module networks.lora_krea2 --network_dim 32 --network_alpha 32 \
  --max_train_epochs 16 --save_every_n_epochs 1 --seed 42 \
  --output_dir output/your_run --output_name your_run

¶Troubleshooting

• torch.cuda.is_available() is False when verifying installation: Check that /dev/kfd and /dev/dri exist (ls /dev/kfd /dev/dri). Your user should be in video and/or render groups.
• Torch version shows +cu... (CUDA): Reinstall ROCm torch with --no-deps
• OOM during training: Add --gradient_checkpointing to your training run. See mode 1
• rocSHMEM Could not open libnuma warning: Harmless on single-GPU setups. Ignore it.