AMD YES!

在Windows上跑深度学习是在还是太难受了，不过考虑到我们有世界上最好用的Linux发行版——WSL2，我们可以在WSL2上再搭一套环境。

而且ROCm再Linux上的支持也好一些，像triton等功能windows上根本装不上。

环境准备

驱动

正常安装 AMD Software: Adrenalin Edition

WSL2

wsl --install -d Ubuntu-24.04

随后在WSL内部安装librocdxg：

curl -sLO https://github.com/ROCm/librocdxg/releases/download/v1.2.2/rocdxg-roct_1.2.2_amd64.deb
sudo dpkg -i rocdxg-roct_1.2.2_amd64.deb

需要注意，官方版兼容矩阵不含RDNA2，因此我们使用社区fork来覆盖部分文件：

git clone https://github.com/the-zucc/librocdxg-rdna2.git
cd librocdxg-rdna2

export win_sdk='/mnt/c/Program Files (x86)/Windows Kits/10/Include/10.0.26100.0'
mkdir -p build && cd build
cmake .. -DWIN_SDK="${win_sdk}/shared"
make -j"$(nproc)"
sudo make install
sudo ldconfig

必要环境变量

可以写入~/.bashrc：

export HSA_ENABLE_DXG_DETECTION=1
export HSA_OVERRIDE_GFX_VERSION=10.3.0     # 让设备以 gfx1030 身份出现（gfx103X wheel 的内核都是按 gfx1030 编译的）
export HSA_ENABLE_SDMA=0
export HSA_ENABLE_PEER_SDMA=0
export LIBROCDXG_ALLOC_USER_QUEUE_FROM_UMD=1

uv

在终端中执行：

curl -LsSf https://astral.sh/uv/install.sh | sh

Python环境

首先，编写pyproject.toml

project]
name = "foo"
version = "0.1.0"
requires-python = ">=3.12,<3.13"

dependencies = [
    "torch==2.12.0+rocm7.14.0a20260612",
    "triton==3.7.0+gitb4e20bbe.rocm7.14.0a20260612",
    "transformers>=5.6.0",

    "rocm==7.14.0a20260612",
    "rocm-sdk-core==7.14.0a20260612",
    "rocm-sdk-libraries-gfx103X-all==7.14.0a20260612",
]

[[tool.uv.index]]
name = "amd"
url = "https://rocm.nightlies.amd.com/v2-staging/gfx103X-all/"
explicit = true

[[tool.uv.index]]
name = "tencent"
url = "http://mirrors.cloud.tencent.com/pypi/simple"
default = true

[tool.uv.sources]
torch = { index = "amd" }
triton = { index = "amd" }
rocm = { index = "amd" }
rocm-sdk-core = { index = "amd" }
rocm-sdk-libraries-gfx103X-all = { index = "amd" }

[tool.uv]
package = false

随后，把wheel包安装到本地并验证:

uv sync

验证

source .venv/bin/activate
rocminfo

输出里面可能会警告Warning: Windows driver is old, please update it.这个不用管他，升级到最新版也会报，不影响执行的。

如果你在这一步看不到你的显卡，检查一下前一步的环境变量设置是否对了。

随后，我们可以使用uv run python来检查torch是否工作正常：

import torch
>>> print(torch.__version__, torch.cuda.get_device_name(0))
2.12.0+rocm7.14.0a20260612 AMD Radeon RX 6800S
>>> a = torch.Tensor([1]).cuda()
>>> print(a+a)
tensor([2.], device='cuda:0')