Skip to content

feat(deps): support ROCm on Windows through the rocm extra #9606

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

Closed
0xDELUXA wants to merge 1 commit into from
Closed

feat(deps): support ROCm on Windows through the rocm extra #9606

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
8 changes: 8 additions & 0 deletions docs/src/content/docs/start-here/manual.mdx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,14 @@ The following commands vary depending on the version of Invoke being installed a
--torch-backend=xpu
```
</TabItem>
<TabItem label="AMD GPU">
Do not use a torch backend. AMD publishes its Windows ROCm wheels on its own index, with the GPU kernels in a separate package per GPU target. In the next step, pin them for your GPU's target and point `uv` at that index instead:
```sh
uv pip install <PACKAGE_SPECIFIER>==<VERSION> "torch[device-<TARGET>]==2.13.0+rocm10.0.0" "torchvision[device-<TARGET>]==0.28.0+rocm10.0.0" --index https://stable.repo.amd.com/rocm/whl-next/ --index-strategy unsafe-best-match --python 3.12 --python-preference only-managed --force-reinstall
```

`<TARGET>` is `gfx1201` for the RX 9070 series and Radeon AI PRO R9700, `gfx1200` for the RX 9060 series, `gfx1100` for the RX 7900 series, `gfx1101` for the RX 7700 and 7800 series, `gfx1102` for the RX 7600 series, and `gfx1151` for Ryzen AI Max. For other GPUs, see AMD's [GPU target table](https://github.com/ROCm/TheRock/blob/main/RELEASES.md#gfx-target-lookup-table).
</TabItem>
<TabItem label="Other / no GPU">
Do not use a torch backend.
</TabItem>
Expand Down
14 changes: 9 additions & 5 deletions docs/src/content/docs/start-here/system-requirements.mdx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ Hardware requirements vary significantly depending on model and image output siz
The requirements below are rough guidelines for best performance. GPUs with less VRAM typically still work, if a bit slower. Follow the [Low VRAM Guide] to optimize performance.

- All Apple Silicon (M1, M2, etc) Macs work, but 16GB+ memory is recommended.
- AMD GPUs are supported on Linux only. The VRAM requirements are the same as Nvidia GPUs.
- AMD GPUs are supported on Linux, and on Windows 11 for RDNA3, RDNA3.5 and RDNA4 GPUs (Radeon RX 7000 and 9000 series, Radeon AI PRO R9700, Ryzen AI 300 and Ryzen AI Max). The VRAM requirements are the same as Nvidia GPUs.
- Intel Arc GPUs (Alchemist, Battlemage and newer) are supported on Windows and Linux `x86_64`. The VRAM requirements are the same as Nvidia GPUs.
- Linux ARM64 (`aarch64`) devices — e.g. Raspberry Pi 5, other SBCs, ARM servers — are supported in CPU-only mode. Local generation is slow without a GPU, but API-backed models (e.g. GPT Image, Gemini) work well.

Expand Down Expand Up @@ -140,14 +140,16 @@ If, after restarting the app, this doesn't improve your performance, either rest

### AMD

:::tip[Linux Only]{icon="linux"}
AMD GPUs are supported on Linux only, due to ROCm (the AMD equivalent of CUDA) support being Linux only.
:::

:::caution[Bumps Ahead]
While the application does run on AMD GPUs, there are occasional bumps related to spotty torch support.
:::

#### Windows

ROCm on Windows comes from AMD's own PyTorch wheels, which bundle the ROCm runtime, so there is no separate ROCm or HIP SDK install. You need Windows 11, an RDNA3, RDNA3.5 or RDNA4 GPU, and an up-to-date [AMD Software: Adrenalin Edition] driver. See AMD's [ROCm compatibility matrix] for details.

#### Linux

Run `rocm-smi` on your system's command line verify that drivers and ROCm are installed. If this command fails, or doesn't report versions, you will need to install them.

Go to the [ROCm Documentation] and carefully follow the instructions for your system to get everything installed.
Expand All @@ -170,4 +172,6 @@ An alternative to installing ROCm locally is to use a [ROCm docker container] to
[Intel Arc & Iris Xe Graphics driver]: https://www.intel.com/content/www/us/en/download/785597/intel-arc-iris-xe-graphics-windows.html
[Intel client GPU install guide]: https://dgpu-docs.intel.com/driver/client/overview.html
[ROCm Documentation]: https://rocmdocs.amd.com
[ROCm compatibility matrix]: https://rocm.docs.amd.com/en/docs-10.0.0/compatibility/compatibility-matrix.html
[AMD Software: Adrenalin Edition]: https://www.amd.com/en/support/download/drivers.html
[ROCm docker container]: https://rocmdocs.amd.com/en/latest/Deep_learning/Deep_learning.html#docker-containers
3 changes: 3 additions & 0 deletions invokeai/app/run_app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@
# keeps their value.
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

# Enables fused SDPA kernels on ROCm GPUs that AOTriton marks experimental. Ignored elsewhere.
os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")


def get_app():
"""Import the app and event loop. We wrap this in a function to more explicitly control when it happens, because
Expand Down
112 changes: 108 additions & 4 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,9 @@ dependencies = [
# reject torch>=2.12 on Windows/Linux CUDA, CPU and ARM64, where 2.12 has no known problem.
# The ROCm path is constrained where it can actually be targeted: the `rocm` extra below pins an
# exact version, and the manual docs carry the ROCm-specific caveat.
"torch>=2.7.0,<3.0; sys_platform != 'darwin'",
# Split so uv forks linux and win32, which take rocm torch from different indexes.
"torch>=2.7.0,<3.0; sys_platform == 'linux'",
"torch>=2.7.0,<3.0; sys_platform == 'win32'",
"torch>=2.7.0,<2.8.0; sys_platform == 'darwin'",
"torchsde", # diffusers needs this for SDE solvers, but it is not an explicit dep of diffusers
"torchvision",
Expand Down Expand Up @@ -138,10 +140,52 @@ dependencies = [
"torchvision==0.22.1; sys_platform == 'linux' and platform_machine == 'aarch64'",
]
"rocm" = [
# ROCm wheels are x86_64-linux-only; gate by platform so macOS/win/linux-aarch64 resolution is unaffected.
# PyTorch's ROCm wheels are x86_64-linux-only.
"torch==2.10.0+rocm7.1; sys_platform == 'linux' and platform_machine != 'aarch64'",
"torchvision==0.25.0+rocm7.1; sys_platform == 'linux' and platform_machine != 'aarch64'",
"triton-rocm==3.6.0; sys_platform == 'linux' and platform_machine != 'aarch64'",
# Windows wheels come from AMD's index. uv only applies index sources to direct requirements, so the
# packages that exist only there are listed explicitly.
"torch==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"torchvision==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx1100==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx1101==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx1102==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx1103==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx110x==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx1150==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx1151==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx1152==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx1153==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx115x==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx12-0==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx1200==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torch-device-gfx1201==2.13.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torchvision-device-gfx1100==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torchvision-device-gfx1101==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torchvision-device-gfx1102==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torchvision-device-gfx1103==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torchvision-device-gfx1150==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torchvision-device-gfx1151==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torchvision-device-gfx1152==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torchvision-device-gfx1153==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torchvision-device-gfx1200==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"amd-torchvision-device-gfx1201==0.28.0+rocm10.0.0; sys_platform == 'win32'",
"rocm==10.0.0; sys_platform == 'win32'",
"rocm-sdk-core==10.0.0; sys_platform == 'win32'",
"rocm-sdk-device-gfx1100==10.0.0; sys_platform == 'win32'",
"rocm-sdk-device-gfx1101==10.0.0; sys_platform == 'win32'",
"rocm-sdk-device-gfx1102==10.0.0; sys_platform == 'win32'",
"rocm-sdk-device-gfx1103==10.0.0; sys_platform == 'win32'",
"rocm-sdk-device-gfx1150==10.0.0; sys_platform == 'win32'",
"rocm-sdk-device-gfx1151==10.0.0; sys_platform == 'win32'",
"rocm-sdk-device-gfx1152==10.0.0; sys_platform == 'win32'",
"rocm-sdk-device-gfx1153==10.0.0; sys_platform == 'win32'",
"rocm-sdk-device-gfx1200==10.0.0; sys_platform == 'win32'",
"rocm-sdk-device-gfx1201==10.0.0; sys_platform == 'win32'",
"rocm-sdk-libraries==10.0.0; sys_platform == 'win32'",
# First release with a Windows ROCm binary.
"bitsandbytes>=0.50.2; sys_platform == 'win32'",
# linux_aarch64 fallback, see note under "cpu"
"torch==2.7.1; sys_platform == 'linux' and platform_machine == 'aarch64'",
"torchvision==0.22.1; sys_platform == 'linux' and platform_machine == 'aarch64'",
Expand Down Expand Up @@ -211,6 +255,22 @@ environments = ["sys_platform == 'win32' or sys_platform == 'darwin' or (sys_pla
# lockfile regenerated by a uv too old to honor this.
exclude-dependencies = ["notebook"]

# The Linux wheels of this version require a Linux-only triton; restate the Windows wheel's metadata.
[[tool.uv.dependency-metadata]]
name = "torch"
version = "2.13.0+rocm10.0.0"
requires-dist = [
"filelock",
"typing-extensions>=4.10.0",
"setuptools>=77.0.3",
"sympy>=1.13.3",
"networkx>=2.5.1",
"jinja2",
"fsspec>=0.8.5",
"rocm[libraries]==10.0.0",
"rocm-bootstrap",
]

[tool.uv.sources]
# Exclude linux_aarch64 from the PyTorch WHL indexes so that uv resolves torch /
# torchvision from PyPI on that platform. The indexes DO carry aarch64 torch wheels
Expand All @@ -219,20 +279,58 @@ exclude-dependencies = ["notebook"]
torch = [
{ index = "torch-cpu", extra = "cpu", marker = "sys_platform != 'linux' or platform_machine != 'aarch64'" },
{ index = "torch-cuda", extra = "cuda", marker = "sys_platform != 'linux' or platform_machine != 'aarch64'" },
{ index = "torch-rocm", extra = "rocm", marker = "sys_platform != 'linux' or platform_machine != 'aarch64'" },
{ index = "torch-rocm", extra = "rocm", marker = "sys_platform != 'win32' and (sys_platform != 'linux' or platform_machine != 'aarch64')" },
{ index = "torch-rocm-windows", extra = "rocm", marker = "sys_platform == 'win32'" },
{ index = "torch-xpu", extra = "xpu", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
{ index = "pypi", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
]
torchvision = [
{ index = "torch-cpu", extra = "cpu", marker = "sys_platform != 'linux' or platform_machine != 'aarch64'" },
{ index = "torch-cuda", extra = "cuda", marker = "sys_platform != 'linux' or platform_machine != 'aarch64'" },
{ index = "torch-rocm", extra = "rocm", marker = "sys_platform != 'linux' or platform_machine != 'aarch64'" },
{ index = "torch-rocm", extra = "rocm", marker = "sys_platform != 'win32' and (sys_platform != 'linux' or platform_machine != 'aarch64')" },
{ index = "torch-rocm-windows", extra = "rocm", marker = "sys_platform == 'win32'" },
{ index = "torch-xpu", extra = "xpu", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
{ index = "pypi", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
]
triton-rocm = [
{ index = "torch-rocm", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
]
amd-torch-device-gfx1100 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx1101 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx1102 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx1103 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx110x = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx1150 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx1151 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx1152 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx1153 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx115x = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx12-0 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx1200 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torch-device-gfx1201 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torchvision-device-gfx1100 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torchvision-device-gfx1101 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torchvision-device-gfx1102 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torchvision-device-gfx1103 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torchvision-device-gfx1150 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torchvision-device-gfx1151 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torchvision-device-gfx1152 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torchvision-device-gfx1153 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torchvision-device-gfx1200 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
amd-torchvision-device-gfx1201 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-core = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-device-gfx1100 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-device-gfx1101 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-device-gfx1102 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-device-gfx1103 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-device-gfx1150 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-device-gfx1151 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-device-gfx1152 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-device-gfx1153 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-device-gfx1200 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-device-gfx1201 = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
rocm-sdk-libraries = [{ index = "torch-rocm-windows", marker = "sys_platform == 'win32'" }]
triton-xpu = [
{ index = "torch-xpu", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
]
Expand All @@ -256,6 +354,12 @@ name = "torch-rocm"
url = "https://download.pytorch.org/whl/rocm7.1"
explicit = true

# Also mirrors PyPI packages, so it must stay explicit.
[[tool.uv.index]]
name = "torch-rocm-windows"
url = "https://stable.repo.amd.com/rocm/whl-next/"
explicit = true

[[tool.uv.index]]
name = "torch-xpu"
url = "https://download.pytorch.org/whl/xpu"
Expand Down
4 changes: 2 additions & 2 deletions scripts/check_pins.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,8 @@
#
# - darwin is deliberately empty: macOS uses MPS, for which PyTorch publishes no
# separate index, so the launcher installs the default PyPI wheels.
# - win32 has no rocm entry: PyTorch publishes no ROCm wheels for Windows, and
# the `rocm` extra in pyproject.toml is marked `sys_platform == 'linux'`.
# - win32 has no rocm entry: legacy (pre-6.14.0) installs do not support ROCm
# on Windows.
# - xpu is on both win32 and linux: PyTorch's XPU index publishes win_amd64 and
# linux-x86_64 wheels, which is exactly what the `xpu` extra's markers allow.
REQUIRED_BACKENDS: dict[str, set[str]] = {
Expand Down
Loading
Loading