fix(win32): scope ROCm workarounds to AMD hosts only

Three Codex-flagged issues where Windows ROCm workarounds incorrectly
applied to Windows CUDA (NVIDIA) machines:

main.py (P1): BNB_ROCM_VERSION was set unconditionally on all win32
hosts. On NVIDIA, bitsandbytes sees BNB_ROCM_VERSION and looks for a
ROCm DLL that doesn't exist, breaking bitsandbytes initialisation.
Fix: gate the block on HIP_PATH/ROCM_PATH being present (ROCm hosts only).

worker.py (P2): torchao stubs were seeded for all win32 runs, shadowing
real torchao on Windows CUDA and silently disabling torchao quantization
for NVIDIA users. Fix: gate on HIP_PATH/ROCM_PATH (win32 ROCm only).

install_python_stack.py (P1): _detect_windows_gfx_arch() only checked
shutil.which("hipinfo"), skipping the HIP_PATH/ROCM_PATH fallback that
the PowerShell installers use. On installs where the HIP SDK bin dir is
not on PATH, _ensure_rocm_torch() returned early without installing
ROCm wheels or bitsandbytes. Fix: mirror the env-var fallback.
This commit is contained in:
LeoBorcherding 2026-05-16 17:08:57 -05:00
commit 4e75d42e4a
3 changed files with 27 additions and 3 deletions

View file

@ -228,10 +228,24 @@ def _detect_rocm_version() -> tuple[int, int] | None:
def _detect_windows_gfx_arch() -> str | None:
"""Return the gcnArchName from hipinfo on Windows (e.g. 'gfx1200'), or None."""
"""Return the gcnArchName from hipinfo on Windows (e.g. 'gfx1200'), or None.
Resolves hipinfo via PATH first, then HIP_PATH\\bin and ROCM_PATH\\bin as
fallbacks -- the AMD HIP SDK installer sets these env vars but does not
always add the bin dir to the system PATH.
"""
import re
hipinfo = shutil.which("hipinfo")
if not hipinfo:
# Fallback: AMD HIP SDK sets HIP_PATH / ROCM_PATH even when bin isn't on PATH
for _env_var in ("HIP_PATH", "ROCM_PATH"):
_root = os.environ.get(_env_var)
if _root:
_candidate = os.path.join(_root, "bin", "hipinfo.exe")
if os.path.isfile(_candidate):
hipinfo = _candidate
break
if not hipinfo:
return None
try: