fix: add rocm_sdk namespace tarball to Windows ROCm wheel installs
torch/_rocm_init.py calls `import rocm_sdk` at startup, which requires the rocm namespace tarball (rocm-*.tar.gz) in addition to the SDK wheel packages. This tarball was missing from both install.ps1 and setup.ps1, causing ModuleNotFoundError on first torch import. - Add rocm-0.1.dev0.tar.gz to ROCm 7.1.1 install (provides rocm_sdk namespace) - Add rocm-7.2.1.tar.gz + rocm_sdk_devel to ROCm 7.2.1 install - Install tarball in a dedicated step before main SDK/torch wheels - Switch to @array splatting in install.ps1 scriptblock for dynamic wheel count - Remove --no-cache-dir from Python-side ROCm wheel install (prevents ~2GB redownload)
This commit is contained in:
LeoBorcherding
parent
efcaccbbcf
commit
301d6c0aa7
3 changed files with 44 additions and 11 deletions
27
install.ps1
27
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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue