Skip to content

Support for AMD ROCm compatible GPUs. #559

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cupertinomiranda wants to merge 1 commit into
base: main
Choose a base branch
from
+178 −8
Open

Support for AMD ROCm compatible GPUs. #559

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 24 additions & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,7 +216,7 @@ _Last updated: 2025-12-17 22:19 UTC_

Requirements:
- python >3.10
- Nvidia GPU with enough ram to do what you need
- Nvidia or AMD ROCm GPU with enough ram to do what you need
- python venv
- git

Expand Down Expand Up @@ -248,6 +248,29 @@ pip install --no-cache-dir torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --
pip install -r requirements.txt
```

Linux AMD:

Follow instructions in https://rocm.docs.amd.com to install amdgpu drivers and rocm support. Latest version 7.1.1 is confirmed to work. older versions might not work due to differences in amd-smi command.

```bash
git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
python3 -m venv venv
source venv/bin/activate
# install torch first - rocm7.1 compatible
pip install --pre torch torchvision torchaudio torchao --index-url https://download.pytorch.org/whl/nightly/rocm7.1
pip3 install -r requirements-amd.txt
```

AI Toolkit requires to have bitsandbytes installed. Perhaps this will change but for now the solution is to build it.
```bash
git clone https://github.com/bitsandbytes-foundation/bitsandbytes.git -b 0.48.2
cd bitsandbytes
# replace gfx1201 by the arch for your GPU. You can get it with `amd-smi static | grep gfx`
cmake -DCMAKE_HIP_COMPILER="/opt/rocm-7.1.1/lib/llvm/bin/clang++" -DCMAKE-DBNB_ROCM_ARCH="gfx1201" -DCOMPUTE_BACKEND=hip .
make -j32
pip install .
```

# AI Toolkit UI

Expand Down
39 changes: 39 additions & 0 deletions requirements-amd.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
torchao==0.10.0
safetensors
git+https://github.com/huggingface/diffusers@8600b4c10d67b0ce200f664204358747bd53c775
transformers==4.57.3
lycoris-lora==1.8.3
flatten_json
pyyaml
oyaml
tensorboard
kornia
invisible-watermark
einops
accelerate
toml
albumentations==1.4.15
albucore==0.0.16
pydantic
omegaconf
k-diffusion
open_clip_torch
timm
prodigyopt
controlnet_aux==0.0.10
python-dotenv
hf_transfer
lpips
pytorch_fid
optimum-quanto==0.2.4
sentencepiece
huggingface_hub
peft
gradio
python-slugify
opencv-python
pytorch-wavelets==1.3.0
matplotlib==3.10.1
setuptools==69.5.1
scipy==1.12.0
av==16.0.1
122 changes: 115 additions & 7 deletions ui/src/app/api/gpu/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,9 @@ export async function GET() {

// Check if nvidia-smi is available
const hasNvidiaSmi = await checkNvidiaSmi(isWindows);
const hasAmdSmi = await checkAMDSmi(isWindows);

if (!hasNvidiaSmi) {
if (!hasNvidiaSmi && !hasAmdSmi) {
return NextResponse.json({
hasNvidiaSmi: false,
gpus: [],
Expand All @@ -23,14 +24,22 @@ export async function GET() {
}

// Get GPU stats
const gpuStats = await getGpuStats(isWindows);
if (hasNvidiaSmi) {
const gpuStats = await getGpuStats(isWindows);
return NextResponse.json({
hasNvidiaSmi: true,
gpus: gpuStats,
});
} else {
const gpuStats = await getAMDGpuStats(isWindows);
return NextResponse.json({
hasNvidiaSmi: true,
gpus: gpuStats,
});
}

return NextResponse.json({
hasNvidiaSmi: true,
gpus: gpuStats,
});
} catch (error) {
console.error('Error fetching NVIDIA GPU stats:', error);
console.error('Error fetching GPU stats:', error);
return NextResponse.json(
{
hasNvidiaSmi: false,
Expand Down Expand Up @@ -58,6 +67,17 @@ async function checkNvidiaSmi(isWindows: boolean): Promise<boolean> {
return false;
}
}
async function checkAMDSmi(isWindows: boolean): Promise<boolean> {
try {
if (!isWindows) {
// Linux/macOS check
await execAsync('which amd-smi');
}
return true;
} catch (error) {
return false;
}
}

async function getGpuStats(isWindows: boolean) {
// Command is the same for both platforms, but the path might be different
Expand Down Expand Up @@ -121,3 +141,91 @@ async function getGpuStats(isWindows: boolean) {

return gpus;
}

// amdParseFloat and amdParseInt avoid errors when amd-smi entries
// contain the string "N/A".
function amdParseFloat(value) {
try {
const ret = parseFloat(value);
return ret;
} catch(error) {
return 0.0;
}
}

function amdParseInt(value) {
try {
const ret = parseInt(value);
return ret;
} catch(error) {
return 0;
}
}

async function getAMDGpuStats(isWindows: boolean) {
// Execute command
const command = 'amd-smi static --json && echo ";" && amd-smi metric --json';
Copy link

@ptelder ptelder Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This command will fail on systems with an inactive AMD iGPU in addition to an AMD discrete GPU.

The amd-smi metric JSON returns an empty usage attribute in gpu_data for an inactive iGPU looking like:

"usage": "N/A",

// Execute command
const { stdout } = await execAsync(command, {
env: { ...process.env, CUDA_DEVICE_ORDER: 'PCI_BUS_ID' },
});
var data = stdout.split(';');

var sdata = {};
var mdata = {};
try {
sdata = JSON.parse(data[0]);
mdata = JSON.parse(data[1]);
} catch (error) {
console.error('Failed to parse output of amd-smi returned json: ', error);
return [];
}

var gpus = sdata["gpu_data"].map(d => {
const i = amdParseInt(d["gpu"]);
const gpu_data = mdata["gpu_data"][i];
const mem_total = amdParseFloat(gpu_data["mem_usage"]["total_vram"]["value"]);
const mem_used = amdParseFloat(gpu_data["mem_usage"]["used_vram"]["value"]);
const mem_free = amdParseFloat(gpu_data["mem_usage"]["free_visible_vram"]["value"]);
const mem_utilization = ((1.0 - (mem_total - mem_free)) / mem_total) * 100;

return {
index: i,
name: d["asic"]["market_name"],
driverVersion: d["driver"]["version"],
temperature: amdParseInt(gpu_data["temperature"]["hotspot"]["value"]),
utilization: {
gpu: amdParseInt(gpu_data["usage"]["gfx_activity"]["value"]),
memory: mem_utilization,
},
memory: {
total: mem_total,
used: mem_used,
free: mem_free,
},
power: {
draw: amdParseFloat(gpu_data["power"]["socket_power"]["value"]),
limit: amdParseFloat(() => {
try {
if (d["limit"]["max_power"]) {
return d["limit"]["max_power"]["value"];
} else if (d["limit"]["ppt0"]["max_power_limit"]["value"]) {
return d["limit"]["ppt0"]["max_power_limit"]["value"];
}
} catch (error) {
return 0.0;
}
})
},
clocks: {
graphics: amdParseInt(gpu_data["clock"]["gfx_0"]["clk"]["value"]),
memory: amdParseInt(gpu_data["clock"]["mem_0"]["clk"]["value"]),
},
fan: {
speed: amdParseFloat(gpu_data["fan"]["usage"]["value"]),
}
};
});

return gpus;
}