diff --git a/install.ps1 b/install.ps1 index ce0211027e..cad8734bb8 100644 --- a/install.ps1 +++ b/install.ps1 @@ -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) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index a3e84d5885..8ffd9a1974 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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]==). +_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