From 7fbdce171c7dbb550047f2d1b0837ffdbc8cace1 Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Wed, 6 May 2026 20:37:44 -0500 Subject: [PATCH] fix: pass AMD torch install status via env var to suppress false warning setup.ps1 now sets UNSLOTH_ROCM_TORCH_INSTALLED=1 after a successful AMD wheel install. install_python_stack.py reads this at the top of _ensure_rocm_torch() to skip both the subprocess probe and the warning -- no re-import of torch needed, and the warning message now correctly says 'could not be auto-installed' rather than 'must be installed manually'. --- studio/install_python_stack.py | 10 ++++++++-- studio/setup.ps1 | 4 ++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 101978b362..cedbb8058d 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -286,6 +286,12 @@ def _ensure_rocm_torch() -> None: Uses pip_install() to respect uv, constraints, and --python targeting. """ global _rocm_windows_torch_installed + # setup.ps1 sets this env var when it successfully installs AMD wheels + # before calling install_python_stack.py, so we can skip the subprocess + # probe and avoid reinstalling what was just installed. + if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": + _rocm_windows_torch_installed = True + return if IS_MACOS: return @@ -1191,11 +1197,11 @@ def install_python_stack() -> int: if _win_amd_gpu and not _rocm_windows_torch_installed: _safe_print( _dim(" Note:"), - "AMD GPU detected on Windows. ROCm-enabled PyTorch must be", + "AMD GPU detected but ROCm PyTorch could not be auto-installed.", ) _safe_print( " " * 8, - "installed manually. See: https://docs.unsloth.ai/get-started/install-and-update/amd", + "Manual install may be required. See: https://docs.unsloth.ai/get-started/install-and-update/amd", ) # 3. Extra dependencies diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 21d99869e5..bd6a6ac101 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1894,6 +1894,10 @@ if ($ROCmTorchWheelUrls) { Write-Host "[WARN] AMD ROCm PyTorch install failed -- falling back to CPU" -ForegroundColor Yellow Write-Host $output -ForegroundColor Yellow $ROCmTorchWheelUrls = $null + } else { + # Signal to install_python_stack.py that AMD wheels are already installed + # so it skips the subprocess probe and suppresses the manual-install warning. + $env:UNSLOTH_ROCM_TORCH_INSTALLED = "1" } }