ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: install rocm_sdk_core and rocm_sdk_libraries_custom alongside torch

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

24
install.ps1
View file

@ -1317,15 +1317,25 @@ shell.Run cmd, 0, False
$amdWheelBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.radeon.com/rocm/windows" }
if ($ROCmVersion -and $ROCmVersion -match '^7\.2') {
$amdRelBase = "$amdWheelBase/rocm-rel-7.2.1"
$ROCmTorchWheelUrl = "$amdRelBase/torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl"
$ROCmTorchVisionUrl = "$amdRelBase/torchvision-0.24.1+rocm7.2.1-cp312-cp312-win_amd64.whl"
$ROCmTorchAudioUrl = "$amdRelBase/torchaudio-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl"
$ROCmAllWheelUrls = @(
"$amdRelBase/rocm_sdk_core-7.2.1-py3-none-win_amd64.whl",
"$amdRelBase/rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl",
"$amdRelBase/torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl",
"$amdRelBase/torchvision-0.24.1+rocm7.2.1-cp312-cp312-win_amd64.whl",
"$amdRelBase/torchaudio-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl"
)
$ROCmTorchWheelUrl = $ROCmAllWheelUrls[2]
$TorchIndexUrl = $null
} elseif ($ROCmVersion -and $ROCmVersion -match '^7\.1') {
$amdRelBase = "$amdWheelBase/rocm-rel-7.1.1"
$ROCmTorchWheelUrl = "$amdRelBase/torch-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl"
$ROCmTorchVisionUrl = "$amdRelBase/torchvision-0.24.0+rocmsdk20251116-cp312-cp312-win_amd64.whl"
$ROCmTorchAudioUrl = "$amdRelBase/torchaudio-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl"
$ROCmAllWheelUrls = @(
"$amdRelBase/rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl",
"$amdRelBase/rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl",
"$amdRelBase/torch-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl",
"$amdRelBase/torchvision-0.24.0+rocmsdk20251116-cp312-cp312-win_amd64.whl",
"$amdRelBase/torchaudio-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl"
)
$ROCmTorchWheelUrl = $ROCmAllWheelUrls[2]
$TorchIndexUrl = $null
}
if ($ROCmTorchWheelUrl) {
@ -1430,7 +1440,7 @@ shell.Run cmd, 0, False
} elseif ($ROCmTorchWheelUrl) {
Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
substep "installing PyTorch (AMD ROCm $ROCmVersion)..."
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --no-cache-dir $ROCmTorchWheelUrl $ROCmTorchVisionUrl $ROCmTorchAudioUrl }
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --no-cache-dir @ROCmAllWheelUrls }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" $torchInstallExit)

36
studio/install_python_stack.py
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