diff --git a/install.ps1 b/install.ps1 index 5228c7ba30..bf279b5850 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1311,23 +1311,31 @@ shell.Run cmd, 0, False # When the HIP SDK is present and Python 3.12 is in use, swap in the AMD wheel # URL and clear $TorchIndexUrl so the standard --index-url path is skipped. $ROCmTorchWheelUrl = $null + $ROCmTarballUrl = $null if ($HasROCm -and -not $SkipTorch) { $pyMajMin = if ($DetectedPython) { ($DetectedPython.Version -split '\.')[0..1] -join '.' } else { "" } 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') { $amdRelBase = "$amdWheelBase/rocm-rel-7.2.1" + # rocm tarball (14 KB) provides the 'rocm_sdk' Python namespace that + # torch/_rocm_init.py imports at startup. + $ROCmTarballUrl = "$amdRelBase/rocm-7.2.1.tar.gz" $ROCmAllWheelUrls = @( "$amdRelBase/rocm_sdk_core-7.2.1-py3-none-win_amd64.whl", + "$amdRelBase/rocm_sdk_devel-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] + $ROCmTorchWheelUrl = $ROCmAllWheelUrls[3] $TorchIndexUrl = $null } elseif ($ROCmVersion -and $ROCmVersion -match '^7\.1') { $amdRelBase = "$amdWheelBase/rocm-rel-7.1.1" + # rocm tarball (14 KB) provides the 'rocm_sdk' Python namespace that + # torch/_rocm_init.py imports at startup. + $ROCmTarballUrl = "$amdRelBase/rocm-0.1.dev0.tar.gz" $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", @@ -1440,12 +1448,17 @@ shell.Run cmd, 0, False } elseif ($ROCmTorchWheelUrl) { Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)" substep "installing PyTorch (AMD ROCm $ROCmVersion)..." - # Expand array to scalars — @array splatting requires & and doesn't - # work reliably inside scriptblocks passed to Invoke-InstallCommand. - $rw0 = $ROCmAllWheelUrls[0]; $rw1 = $ROCmAllWheelUrls[1] - $rw2 = $ROCmAllWheelUrls[2]; $rw3 = $ROCmAllWheelUrls[3] - $rw4 = $ROCmAllWheelUrls[4] - $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --no-deps $rw0 $rw1 $rw2 $rw3 $rw4 } + # Install the rocm namespace tarball first (provides the 'rocm_sdk' + # Python package that torch/_rocm_init.py imports at startup). + if ($ROCmTarballUrl) { + $tarballExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --no-deps $ROCmTarballUrl } + if ($tarballExit -ne 0) { + Write-Host "[WARN] ROCm namespace tarball install failed (exit $tarballExit) -- continuing" -ForegroundColor Yellow + } + } + # Install remaining SDK + torch wheels. @array splatting inside a + # scriptblock works in PS 5.1 because & $Command runs in-scope. + $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --no-deps @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 1a774e1f5c..39968ce484 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -70,7 +70,10 @@ _ROCM_WINDOWS_RELEASES: dict[tuple[int, int], tuple[str, list[str]]] = { (7, 2): ( "rocm-rel-7.2.1", [ + # rocm tarball provides the 'rocm_sdk' Python namespace package + "rocm-7.2.1.tar.gz", "rocm_sdk_core-7.2.1-py3-none-win_amd64.whl", + "rocm_sdk_devel-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", @@ -80,6 +83,8 @@ _ROCM_WINDOWS_RELEASES: dict[tuple[int, int], tuple[str, list[str]]] = { (7, 1): ( "rocm-rel-7.1.1", [ + # rocm tarball provides the 'rocm_sdk' Python namespace package + "rocm-0.1.dev0.tar.gz", "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", @@ -352,7 +357,6 @@ def _ensure_rocm_torch() -> None: pip_install( f"ROCm torch (Windows, {rel_tag})", "--force-reinstall", - "--no-cache-dir", "--no-deps", *wheel_urls, constrain = False, diff --git a/studio/setup.ps1 b/studio/setup.ps1 index bd6a6ac101..fdffaa30b8 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1855,6 +1855,7 @@ if ($HasNvidiaSmi) { # torch wheels instead of CPU-only PyTorch. $ROCmVersion = $script:ROCmVersion $ROCmTorchWheelUrls = $null +$ROCmTarballUrl = $null if ($HasROCm -and $CuTag -eq "cpu") { $pyVer = (& python --version 2>&1 | Out-String) -replace '[^0-9.]','' $pyMajMin = ($pyVer.Trim() -split '\.')[0..1] -join '.' @@ -1862,8 +1863,12 @@ if ($HasROCm -and $CuTag -eq "cpu") { if ($pyMajMin -eq "3.12" -and $ROCmVersion) { if ($ROCmVersion -match '^7\.2') { $rb = "$amdWheelBase/rocm-rel-7.2.1" + # rocm tarball (14 KB) provides the 'rocm_sdk' Python namespace that + # torch/_rocm_init.py imports at startup. + $ROCmTarballUrl = "$rb/rocm-7.2.1.tar.gz" $ROCmTorchWheelUrls = @( "$rb/rocm_sdk_core-7.2.1-py3-none-win_amd64.whl", + "$rb/rocm_sdk_devel-7.2.1-py3-none-win_amd64.whl", "$rb/rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl", "$rb/torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl", "$rb/torchvision-0.24.1+rocm7.2.1-cp312-cp312-win_amd64.whl", @@ -1871,6 +1876,9 @@ if ($HasROCm -and $CuTag -eq "cpu") { ) } elseif ($ROCmVersion -match '^7\.1') { $rb = "$amdWheelBase/rocm-rel-7.1.1" + # rocm tarball (14 KB) provides the 'rocm_sdk' Python namespace that + # torch/_rocm_init.py imports at startup. + $ROCmTarballUrl = "$rb/rocm-0.1.dev0.tar.gz" $ROCmTorchWheelUrls = @( "$rb/rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl", "$rb/rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl", @@ -1886,9 +1894,17 @@ $PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR if ($ROCmTorchWheelUrls) { substep "installing PyTorch (AMD ROCm $ROCmVersion)..." - $sw0 = $ROCmTorchWheelUrls[0]; $sw1 = $ROCmTorchWheelUrls[1] - $sw2 = $ROCmTorchWheelUrls[2]; $sw3 = $ROCmTorchWheelUrls[3]; $sw4 = $ROCmTorchWheelUrls[4] - $output = Fast-Install --force-reinstall --no-deps $sw0 $sw1 $sw2 $sw3 $sw4 | Out-String + # Install the rocm namespace tarball first (provides the 'rocm_sdk' Python + # package that torch/_rocm_init.py imports at startup). + if ($ROCmTarballUrl) { + $tarballOut = Fast-Install --force-reinstall --no-deps $ROCmTarballUrl | Out-String + if ($LASTEXITCODE -ne 0) { + Write-Host "[WARN] ROCm namespace tarball install failed -- continuing" -ForegroundColor Yellow + Write-Host $tarballOut -ForegroundColor Yellow + } + } + # Install remaining SDK + torch wheels using array splatting. + $output = Fast-Install --force-reinstall --no-deps @ROCmTorchWheelUrls | Out-String $torchInstallExit = $LASTEXITCODE if ($torchInstallExit -ne 0) { Write-Host "[WARN] AMD ROCm PyTorch install failed -- falling back to CPU" -ForegroundColor Yellow