fix: install rocm_sdk_core and rocm_sdk_libraries_custom alongside torch

The AMD Windows torch wheels declare rocm[libraries]==<ver> as a hard
dependency. Without installing rocm_sdk_core and rocm_sdk_libraries_custom
from the same AMD release folder, uv cannot resolve the dependency and
fails with 'No solution found'. Include all 5 wheels in one install call.
This commit is contained in:
LeoBorcherding 2026-05-06 18:08:01 -05:00
commit e53b543ef0
2 changed files with 38 additions and 22 deletions

View file

@ -63,19 +63,29 @@ _ROCM_WINDOWS_WHEEL_BASE = (
os.environ.get("UNSLOTH_ROCM_WINDOWS_MIRROR")
or "https://repo.radeon.com/rocm/windows"
).rstrip("/")
# Maps (major, minor) → (release_folder, torch_ver, torchvision_ver, torchaudio_ver)
_ROCM_WINDOWS_RELEASES: dict[tuple[int, int], tuple[str, str, str, str]] = {
# Maps (major, minor) → (release_folder, [wheel_filename, ...])
# Includes rocm_sdk_core and rocm_sdk_libraries_custom because the torch
# wheels declare them as hard dependencies (rocm[libraries]==<ver>).
_ROCM_WINDOWS_RELEASES: dict[tuple[int, int], tuple[str, list[str]]] = {
(7, 2): (
"rocm-rel-7.2.1",
"2.9.1+rocm7.2.1",
"0.24.1+rocm7.2.1",
"2.9.1+rocm7.2.1",
[
"rocm_sdk_core-7.2.1-py3-none-win_amd64.whl",
"rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl",
"torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl",
"torchvision-0.24.1+rocm7.2.1-cp312-cp312-win_amd64.whl",
"torchaudio-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl",
],
),
(7, 1): (
"rocm-rel-7.1.1",
"2.9.0+rocmsdk20251116",
"0.24.0+rocmsdk20251116",
"2.9.0+rocmsdk20251116",
[
"rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl",
"rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl",
"torch-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl",
"torchvision-0.24.0+rocmsdk20251116-cp312-cp312-win_amd64.whl",
"torchaudio-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl",
],
),
}
@ -316,19 +326,15 @@ def _ensure_rocm_torch() -> None:
f" No AMD Windows torch wheel for ROCm {ver[0]}.{ver[1]} -- skipping"
)
return
rel_tag, torch_ver, tv_ver, ta_ver = entry
rel_tag, wheel_files = entry
base = f"{_ROCM_WINDOWS_WHEEL_BASE}/{rel_tag}"
torch_url = f"{base}/torch-{torch_ver}-cp312-cp312-win_amd64.whl"
tv_url = f"{base}/torchvision-{tv_ver}-cp312-cp312-win_amd64.whl"
ta_url = f"{base}/torchaudio-{ta_ver}-cp312-cp312-win_amd64.whl"
wheel_urls = [f"{base}/{fn}" for fn in wheel_files]
print(f" ROCm {ver[0]}.{ver[1]} (Windows) -- installing torch from {base}/")
pip_install(
f"ROCm torch (Windows, {rel_tag})",
"--force-reinstall",
"--no-cache-dir",
torch_url,
tv_url,
ta_url,
*wheel_urls,
constrain = False,
)
return