diff --git a/install.ps1 b/install.ps1 index 345b8fa9a3..4dfd023af5 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1316,7 +1316,10 @@ shell.Run cmd, 0, False if ($pyMajMin -eq "3.12") { $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') { - $ROCmTorchWheelUrl = "$amdWheelBase/rocm-rel-7.2.1/torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl" + $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" $TorchIndexUrl = $null } if ($ROCmTorchWheelUrl) { @@ -1421,7 +1424,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 } + $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --no-cache-dir $ROCmTorchWheelUrl $ROCmTorchVisionUrl $ROCmTorchAudioUrl } 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 0be4607380..4b9159529d 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -63,9 +63,9 @@ _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_version_string) -_ROCM_WINDOWS_RELEASES: dict[tuple[int, int], tuple[str, str]] = { - (7, 2): ("rocm-rel-7.2.1", "2.9.1+rocm7.2.1"), +# Maps (major, minor) → (release_folder, torch_ver, torchvision_ver, torchaudio_ver) +_ROCM_WINDOWS_RELEASES: dict[tuple[int, int], tuple[str, str, str, 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"), } # bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix @@ -294,8 +294,8 @@ def _ensure_rocm_torch() -> None: return entry = next( ( - (rt, tv) - for (maj, mn), (rt, tv) in sorted( + v + for (maj, mn), v in sorted( _ROCM_WINDOWS_RELEASES.items(), reverse = True ) if ver >= (maj, mn) @@ -307,19 +307,19 @@ def _ensure_rocm_torch() -> None: f" No AMD Windows torch wheel for ROCm {ver[0]}.{ver[1]} -- skipping" ) return - rel_tag, torch_ver = entry - wheel_url = ( - f"{_ROCM_WINDOWS_WHEEL_BASE}/{rel_tag}/" - f"torch-{torch_ver}-cp312-cp312-win_amd64.whl" - ) - print( - f" ROCm {ver[0]}.{ver[1]} (Windows) -- installing torch from {wheel_url}" - ) + rel_tag, torch_ver, tv_ver, ta_ver = 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" + 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", - wheel_url, + torch_url, + tv_url, + ta_url, constrain = False, ) return