From 02b86aee9e868e26756a71ac0dbd7c2cb194c1b4 Mon Sep 17 00:00:00 2001 From: 0xDELUXA Date: Fri, 30 Oct 2026 22:50:36 +0200 Subject: [PATCH] Support ROCm on Windows with ADLX VRAM readings --- diffsynth/core/device/__init__.py | 1 + diffsynth/core/device/rocm_windows.py | 88 +++++++++++++++++++++++++++ diffsynth/core/vram/layers.py | 4 +- docs/en/Pipeline_Usage/GPU_support.md | 2 + docs/en/Pipeline_Usage/Setup.md | 8 +++ docs/zh/Pipeline_Usage/GPU_support.md | 2 + docs/zh/Pipeline_Usage/Setup.md | 8 +++ pyproject.toml | 3 + 8 files changed, 115 insertions(+), 1 deletion(-) create mode 100644 diffsynth/core/device/rocm_windows.py diff --git a/diffsynth/core/device/__init__.py b/diffsynth/core/device/__init__.py index 889d6823a..a50e88944 100644 --- a/diffsynth/core/device/__init__.py +++ b/diffsynth/core/device/__init__.py @@ -1,2 +1,3 @@ from .npu_compatible_device import parse_device_type, parse_nccl_backend, get_available_device_type, get_device_name from .npu_compatible_device import IS_NPU_AVAILABLE, IS_CUDA_AVAILABLE +from .rocm_windows import IS_WINDOWS_ROCM, mem_get_info diff --git a/diffsynth/core/device/rocm_windows.py b/diffsynth/core/device/rocm_windows.py new file mode 100644 index 000000000..1c99deed0 --- /dev/null +++ b/diffsynth/core/device/rocm_windows.py @@ -0,0 +1,88 @@ +import atexit +import os +import sys +import threading +import time +import torch + + +IS_WINDOWS_ROCM = sys.platform == "win32" and torch.version.hip is not None + +# ADLX takes about half a millisecond per query and `check_free_vram` runs on every layer forward, +# so the VRAM used by other processes is refreshed at most this often (seconds). +OTHER_PROCESSES_REFRESH = 1.0 + +if torch.version.hip is not None: + # Without it, PyTorch SDPA has no flash or memory-efficient kernel on the GPUs AOTriton still + # marks experimental and falls back to the math path. Read at the first SDPA call. + os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1") + +_adlx = None +_adlx_lock = threading.Lock() +_other_processes = {} +_missing_adlx_tips_printed = False + + +def _adlx_session(): + global _adlx, _missing_adlx_tips_printed + with _adlx_lock: + if _adlx is None: + _adlx = {} + try: + from adlx import ADLX + except ImportError: + if not _missing_adlx_tips_printed: + print("On ROCm for Windows, torch only reports the VRAM used by this process. Install `amd-adlx` so that `vram_limit` also accounts for other programs and the desktop.") + _missing_adlx_tips_printed = True + return None + try: + helper = ADLX.ADLXHelper() + if helper.Initialize() == ADLX.ADLX_RESULT.ADLX_OK: + system = helper.GetSystemServices() + # UniqueId packs the PCI location as (bus << 8) | (device << 3) | function + gpus = {(gpu.UniqueId() >> 8) & 0xFF: gpu for gpu in system.GetGPUs()} + _adlx.update(helper=helper, perf=system.GetPerformanceMonitoringServices(), gpus=gpus) + atexit.register(_adlx_close) + except Exception as e: + print(f"ADLX is unavailable ({e}), falling back to the VRAM torch reports for this process.") + return _adlx or None + + +def _adlx_close(): + # ADLX crashes at interpreter exit unless every interface is released before Terminate. + with _adlx_lock: + helper = _adlx.pop("helper") + _adlx.clear() + helper.Terminate() + + +def _adlx_used_bytes(bus): + session = _adlx_session() + if session is None or bus not in session["gpus"]: + return None + with _adlx_lock: + metrics = session["perf"].GetCurrentGPUMetrics(session["gpus"][bus]) + used = metrics.GPUVRAM() + del metrics + return used * 1024 ** 2 + + +def mem_get_info(device=None): + """`torch.cuda.mem_get_info`, but for the whole GPU on ROCm for Windows. + + There torch only counts its own process, so the desktop and other programs look like free VRAM. + The VRAM they use comes from ADLX instead. + """ + free, total = torch.cuda.mem_get_info(device) + if not IS_WINDOWS_ROCM: + return free, total + used_by_this_process = total - free + bus = torch.cuda.get_device_properties(device).pci_bus_id + now = time.monotonic() + cached = _other_processes.get(bus) + if cached is None or now - cached[0] > OTHER_PROCESSES_REFRESH: + used = _adlx_used_bytes(bus) + if used is None: + return free, total + cached = _other_processes[bus] = (now, max(used - used_by_this_process, 0)) + return max(total - used_by_this_process - cached[1], 0), total diff --git a/diffsynth/core/vram/layers.py b/diffsynth/core/vram/layers.py index 80ae8589b..4c38e7acd 100644 --- a/diffsynth/core/vram/layers.py +++ b/diffsynth/core/vram/layers.py @@ -2,7 +2,7 @@ from typing import Union from .initialization import skip_model_initialization from .disk_map import DiskMap -from ..device import parse_device_type, get_device_name, IS_NPU_AVAILABLE +from ..device import parse_device_type, get_device_name, IS_NPU_AVAILABLE, mem_get_info class AutoTorchModule(torch.nn.Module): @@ -67,6 +67,8 @@ def check_free_vram(self): allocated = torch.mps.current_allocated_memory() recommended_max = torch.mps.recommended_max_memory() gpu_mem_state = (recommended_max - allocated, recommended_max) + elif self.computation_device_type == "cuda": + gpu_mem_state = mem_get_info(self.computation_device) else: device = self.computation_device if not IS_NPU_AVAILABLE else get_device_name() gpu_mem_state = getattr(torch, self.computation_device_type).mem_get_info(device) diff --git a/docs/en/Pipeline_Usage/GPU_support.md b/docs/en/Pipeline_Usage/GPU_support.md index d4a8f1493..e6012320e 100644 --- a/docs/en/Pipeline_Usage/GPU_support.md +++ b/docs/en/Pipeline_Usage/GPU_support.md @@ -12,6 +12,8 @@ All sample code provided by this project supports NVIDIA GPUs by default, requir AMD provides PyTorch packages based on ROCm, so most models can run without code changes. A small number of models may not be compatible due to their reliance on CUDA-specific instructions. +On Windows, `torch.cuda.mem_get_info` with ROCm only counts the VRAM used by the current process. When `amd-adlx` is installed, `vram_limit` uses the VRAM used by the whole GPU instead, as it does on NVIDIA GPUs. On ROCm, `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` is set by default so that PyTorch SDPA can use its flash and memory-efficient kernels on GPUs where AOTriton support is still experimental. + ### Apple Silicon On Apple Silicon devices, since VRAM and memory are unified, replace all `"cuda"` in the code with `"mps"` or `"cpu"`. diff --git a/docs/en/Pipeline_Usage/Setup.md b/docs/en/Pipeline_Usage/Setup.md index f6ade9f52..2d462624a 100644 --- a/docs/en/Pipeline_Usage/Setup.md +++ b/docs/en/Pipeline_Usage/Setup.md @@ -20,6 +20,7 @@ To keep the framework lightweight, the base installation only installs the neces * `[quant]`: For parameter quantization, enabling precisions such as NF4, INT8, NVFP4. * `[training]`: For distributed large-scale pretraining. * `[logger]`: To enable training loggers such as TensorBoard, SwanLab, etc. +* `[amd]`: For AMD GPUs on Windows, reading the VRAM used by the whole GPU through ADLX. * `[npu]`: For Ascend NPU devices with x86 architecture. * `[npu_aarch64]`: For Ascend NPU devices with aarch64/ARM architecture. * Dependencies of specific models @@ -45,6 +46,13 @@ You need to install the `torch` package with ROCm support. Taking ROCm 6.4 (as o pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm6.4 ``` +On Windows, install AMD's ROCm 10 `torch` packages instead (as of October 30, 2026), replacing `gfx1201` with your GPU's architecture, or with `all` if you are not sure. Then install the `[amd]` dependencies, which include `amd-adlx`, so that `vram_limit` accounts for the VRAM used by other programs and the desktop: + +```shell +pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ "torch[device-gfx1201]==2.14.0+rocm10.1.0" "torchvision[device-gfx1201]==0.29.0a0+rocm10.1.0" +pip install -e ".[amd]" +``` + ### Apple Silicon On Apple Silicon devices, no changes to the installation steps are needed. However, since VRAM and memory are unified, replace all `"cuda"` in the code with `"mps"` or `"cpu"`. diff --git a/docs/zh/Pipeline_Usage/GPU_support.md b/docs/zh/Pipeline_Usage/GPU_support.md index 4558fa0b2..aa44c8957 100644 --- a/docs/zh/Pipeline_Usage/GPU_support.md +++ b/docs/zh/Pipeline_Usage/GPU_support.md @@ -12,6 +12,8 @@ AMD 提供了基于 ROCm 的 torch 包，所以大多数模型无需修改代码即可运行，少数模型由于依赖特定的 cuda 指令无法运行。 +在 Windows 系统上，ROCm 版本的 `torch.cuda.mem_get_info` 只统计当前进程占用的显存。安装 `amd-adlx` 后，`vram_limit` 会改为使用整个 GPU 的显存占用，与 NVIDIA GPU 上的行为一致。在 ROCm 上，框架默认设置 `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1`，使 PyTorch SDPA 在 AOTriton 支持仍为实验性的 GPU 上也能使用 flash 和 memory-efficient 内核。 + ### Apple Silicon 在 Apple Silicon 设备上，由于显存和内存是统一的，请将代码中的 `"cuda"` 全部修改为 `"mps"` 或 `"cpu"`。 diff --git a/docs/zh/Pipeline_Usage/Setup.md b/docs/zh/Pipeline_Usage/Setup.md index 9ead662d5..467d9fdd4 100644 --- a/docs/zh/Pipeline_Usage/Setup.md +++ b/docs/zh/Pipeline_Usage/Setup.md @@ -20,6 +20,7 @@ pip install diffsynth * `[quant]`: 用于参数量化，启用 NF4、INT8、NVFP4 等精度 * `[training]`: 用于分布式大规模预训练 * `[logger]`: 用于启用 TensorBoard、SwanLab 等训练日志记录器 +* `[amd]`: 用于 Windows 系统上的 AMD GPU，通过 ADLX 读取整个 GPU 的显存占用 * `[npu]`: 用于 x86 架构的 Ascend NPU 设备 * `[npu_aarch64]`: 用于 aarch64/ARM 架构的 Ascend NPU 设备 * 特定模型的依赖 @@ -45,6 +46,13 @@ pip install diffsynth pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm6.4 ``` +在 Windows 系统上，请改为安装 AMD 提供的 ROCm 10 `torch` 包（本文更新于 2026 年 10 月 30 日），并将 `gfx1201` 替换为你的 GPU 架构，不确定时可使用 `all`。然后安装 `[amd]` 依赖（包含 `amd-adlx`），使 `vram_limit` 能够计入其他程序和桌面占用的显存： + +```shell +pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ "torch[device-gfx1201]==2.14.0+rocm10.1.0" "torchvision[device-gfx1201]==0.29.0a0+rocm10.1.0" +pip install -e ".[amd]" +``` + ### Apple Silicon 在 Apple Silicon 设备上，无需修改安装步骤。但由于显存和内存是统一的，请将代码中的 `"cuda"` 全部修改为 `"mps"` 或 `"cpu"`。 diff --git a/pyproject.toml b/pyproject.toml index 8986bbc7d..b881dae16 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,9 @@ logger = [ "swanlab", "wandb" ] +amd = [ + "amd-adlx; sys_platform == 'win32'" +] npu = [ "torch==2.7.1+cpu", "torch-npu==2.7.1",