ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix installer selecting ROCm torch on NVIDIA Linux hosts (#6174)

Browse source 
* fix: prevent ROCm torch from installing on NVIDIA Linux hosts

NVIDIA's open kernel module (driver 560+) registers GPU topology nodes in
the KFD sysfs hierarchy with non-zero gpu_id values. The _has_amd_rocm_gpu
(install.sh) and _has_rocm_gpu (install_python_stack.py) sysfs fallbacks
previously treated any non-zero gpu_id as proof of an AMD GPU, so an
NVIDIA-only host with the open kernel driver was misrouted to the ROCm
install path, replacing the correctly-installed CUDA torch with ROCm wheels.

Fixes:

1. install.sh _has_amd_rocm_gpu sysfs fallback: require vendor_id 4098
   (AMD 0x1002) in the KFD node properties file before declaring an AMD
   GPU present. NVIDIA KFD nodes carry vendor_id 4318 (0x10DE) and are
   now skipped.

2. install_python_stack.py _has_rocm_gpu sysfs fallback: same vendor_id
   guard. Also preserves the existing fallback for older kernels that
   don't ship a properties file (trusts gpu_id alone there).

3. install.sh now exports UNSLOTH_TORCH_BACKEND ("cuda"/"rocm"/"cpu")
   immediately after get_torch_index_url() resolves the wheel family.
   install_python_stack.py reads this as _TORCH_BACKEND and short-circuits
   _ensure_rocm_torch() entirely on cuda/cpu hosts, providing a second
   layer of defense that is independent of subprocess GPU detection.

Tests: 9 new cases in TestHasRocmGpuKfdVendorGuard,
TestEnsureRocmTorch, and TestInstallShStructure cover all three changes.
Full test_rocm_support.py suite: 289 passed, 2 skipped, 0 failed.

Closes #6172

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: show actual torch backend in progress step labels

The 'ROCm torch check' and 'ROCm torch (final)' step labels were
hardcoded regardless of whether the installer was targeting CUDA, ROCm,
or CPU. On NVIDIA hosts they showed 'ROCm' even though no ROCm wheels
were being installed, which was misleading.

Add _torch_step_label(suffix) which reads UNSLOTH_TORCH_BACKEND (set by
install.sh) and formats the label as e.g. 'torch check (cuda)' or
'torch final (rocm)'. Falls back to live GPU detection for standalone
studio update runs that bypass install.sh.

* fix: make KFD sysfs vendor check conservative -- skip if no properties file

The previous implementation fell through to `return True` when the KFD
node's properties file was missing (OSError), intending to support older
kernels. But NVIDIA open driver KFD nodes can also lack a properties file
on some kernel versions, so the fallback still produced a false positive.

Change the `except OSError: pass` to `continue` so any node without a
readable properties file is skipped rather than trusted. KFD properties
files exist on every kernel version that actually exposes /sys/class/kfd,
so this does not regress real AMD GPU detection -- if the directory exists
at all, properties files will be present for genuine GPU nodes.

* fix: bulletproof NVIDIA vs AMD GPU detection

Four changes that together ensure ROCm torch can never be installed on an
NVIDIA host regardless of which detection path fires:

1. _has_rocm_gpu() (Python): NVIDIA guard at the top -- returns False
   immediately when _has_usable_nvidia_gpu() is True, blocking rocminfo,
   amd-smi, and KFD sysfs from producing a false positive even when ROCm
   tools are co-installed alongside the NVIDIA driver.

2. _has_amd_rocm_gpu() (install.sh): same NVIDIA guard -- calls
   _has_usable_nvidia_gpu first and returns 1 if it succeeds.

3. _has_usable_nvidia_gpu() (Python): adds /proc/driver/nvidia/gpus/
   sysfs fallback. The NVIDIA driver populates this directory on Linux
   regardless of nvidia-smi state, so a subprocess PATH gap, timeout, or
   driver initialisation race can no longer silence NVIDIA detection.

4. _has_usable_nvidia_gpu() (install.sh): same /proc/driver/nvidia/gpus
   fallback, tried after nvidia-smi -L rather than instead of it.

Together: NVIDIA wins at every decision point. If nvidia-smi works, it
confirms NVIDIA. If it fails, /proc/driver/nvidia confirms NVIDIA. If
somehow both fail, _has_rocm_gpu still checks NVIDIA first before any AMD
path runs.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: two KFD/proc-only corner cases from Codex review

1. KFD awk state not reset per node file (Ryzen+NVIDIA false positive):
   The awk glob processes all topology node properties files in one pass.
   Without FNR==1 reset, a Ryzen+NVIDIA host where an AMD CPU-agent node
   sets amd=1 (vendor_id 4098, gpu_id 0) can combine with a later NVIDIA
   node setting gpu=1 (gpu_id > 0), triggering found=1 before vendor_id
   4318 is seen. Added FNR==1{ gpu=0; amd=0 } to reset per file.

2. proc-only NVIDIA not reaching CUDA wheel selection:
   _has_usable_nvidia_gpu returning true via /proc/driver/nvidia fallback
   left _smi empty, so get_torch_index_url entered the AMD/CPU branch and
   selected CPU wheels despite NVIDIA being confirmed. Introduced
   _nvidia_detected flag (separate from _smi) so the AMD branch is skipped
   whenever NVIDIA is confirmed by any path, while _cuda_ver reads from
   _smi when available (with the existing cu126 fallback when _smi is absent).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Leo Borcherding 2026-06-10 23:00:11 -05:00 • committed by GitHub
commit bf2cd745b1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 366 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

115
studio/install_python_stack.py
View file

@ -557,7 +557,15 @@ def _persist_bnb_rocm_version(version: str) -> bool:
def _has_rocm_gpu() -> bool:
"""Return True only if an actual AMD GPU is visible (not just ROCm tools installed)."""
"""Return True only if an actual AMD GPU is visible (not just ROCm tools installed).
Always returns False when an NVIDIA GPU is present -- NVIDIA takes
priority on mixed hosts and prevents every detection path below
(rocminfo, amd-smi, KFD sysfs) from producing a false positive even
if ROCm tools are installed alongside the NVIDIA driver.
"""
if _has_usable_nvidia_gpu():
return False
for cmd, check_fn in (
# rocminfo: look for a real gfx GPU id (3-4 chars, nonzero first digit).
# gfx000 is the CPU agent; ROCm 6.1+ also emits generic ISA lines like
@ -598,6 +606,13 @@ def _has_rocm_gpu() -> bool:
# runtime-only detection. On minimal package-managed installs (no
# rocminfo / no amd-smi tools), the kernel exposes AMD GPUs via
# /sys/class/kfd so `studio update` can still detect and repair.
#
# Guard: reject any KFD node whose properties file reports a non-AMD
# vendor. With the NVIDIA open kernel module (driver 560+), NVIDIA GPUs
# can register KFD topology nodes with a non-zero gpu_id; those nodes
# have vendor_id 4318 (0x10DE) rather than the AMD value 4098 (0x1002).
# Without this check the fallback returns True on NVIDIA-only systems,
# causing _ensure_rocm_torch to install ROCm wheels on NVIDIA hardware.
if sys.platform != "win32":
try:
kfd_nodes = "/sys/class/kfd/kfd/topology/nodes"
@ -609,29 +624,61 @@ def _has_rocm_gpu() -> bool:
gpu_id = fh.read().strip()
except OSError:
continue
if gpu_id and gpu_id != "0": # gpu_id 0 = CPU node
return True
if not gpu_id or gpu_id == "0": # gpu_id 0 = CPU node
continue
# Require AMD vendor_id 4098 (0x1002) in the properties file.
# KFD properties files exist on every kernel that exposes
# /sys/class/kfd, so absence of the file means we cannot
# confirm AMD ownership -- skip the node rather than risk a
# false positive (e.g. NVIDIA open driver KFD nodes that
# lack a properties file on some kernel versions).
props_path = os.path.join(kfd_nodes, entry, "properties")
try:
with open(props_path) as fh:
props = fh.read()
except OSError:
continue # can't confirm vendor -- skip
if not re.search(r"\bvendor_id\s+4098\b", props):
continue
return True
except OSError:
pass
return False
def _has_usable_nvidia_gpu() -> bool:
"""Return True only when nvidia-smi exists AND reports at least one GPU."""
"""Return True when an NVIDIA GPU is present and usable.
Primary probe: nvidia-smi -L (subprocess).
Fallback: /proc/driver/nvidia/gpus/ sysfs (Linux only) -- handles the
case where nvidia-smi is present but the subprocess fails (PATH gap,
timeout, driver initialisation race). If either probe confirms an
NVIDIA GPU the function returns True so _has_rocm_gpu() is blocked.
"""
exe = shutil.which("nvidia-smi")
if not exe:
return False
try:
result = subprocess.run(
[exe, "-L"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
except Exception:
return False
return result.returncode == 0 and "GPU " in result.stdout
if exe:
try:
result = subprocess.run(
[exe, "-L"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
if result.returncode == 0 and "GPU " in result.stdout:
return True
except Exception:
pass
# Fallback: the NVIDIA driver exposes one subdirectory per GPU under
# /proc/driver/nvidia/gpus/ on Linux regardless of nvidia-smi state.
if sys.platform != "win32":
try:
gpu_dir = "/proc/driver/nvidia/gpus"
if os.path.isdir(gpu_dir) and os.listdir(gpu_dir):
return True
except OSError:
pass
return False
def _detect_amd_gfx_codes() -> list[str]:
@ -749,6 +796,13 @@ def _ensure_rocm_torch() -> None:
Uses pip_install() to respect uv, constraints, and --python targeting.
"""
global _rocm_windows_torch_installed
# install.sh sets UNSLOTH_TORCH_BACKEND to the resolved wheel family
# ("cuda", "rocm", "cpu"). Skip ROCm operations entirely when install.sh
# already selected a non-ROCm backend -- this is the authoritative signal
# and avoids re-running GPU detection in a subprocess that may see a
# different environment (different PATH, CUDA_VISIBLE_DEVICES, etc.).
if _TORCH_BACKEND in ("cuda", "cpu"):
return
# setup.ps1 sets this after installing AMD wheels; skip the probe only when
# torch is actually importable as ROCm. If the venv was wiped between runs,
# the stale env-var would suppress a needed reinstall.
@ -1088,6 +1142,29 @@ def _infer_no_torch() -> bool:
NO_TORCH = _infer_no_torch()
# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() so
# that this script knows which torch variant was selected without re-running
# GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown
# (standalone `unsloth studio update` runs, where we re-detect normally).
_TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower()
def _torch_step_label(suffix: str) -> str:
"""Return a progress label like 'torch check (cuda)' using the known backend.
Falls back to GPU detection when UNSLOTH_TORCH_BACKEND is not set (e.g.
standalone `unsloth studio update` runs that bypass install.sh).
"""
backend = _TORCH_BACKEND
if not backend:
if _has_usable_nvidia_gpu():
backend = "cuda"
elif _has_rocm_gpu():
backend = "rocm"
else:
backend = "cpu"
return f"torch {suffix} ({backend})"
# -- Verbosity control ----------------------------------------------------------
# By default the installer shows a minimal in-place one-line progress bar.
@ -1770,7 +1847,7 @@ def install_python_stack() -> int:
# venv got CPU-only torch (common when pip resolves torch from PyPI).
# Must follow base packages so torch is present for inspection.
if not IS_MACOS and not NO_TORCH:
_progress("ROCm torch check")
_progress(_torch_step_label("check"))
_ensure_rocm_torch()
# Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python
@ -1955,7 +2032,7 @@ def install_python_stack() -> int:
# Running the repair last ensures ROCm torch is in place at runtime,
# whichever intermediate step clobbered it.
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
_progress("ROCm torch (final)")
_progress(_torch_step_label("final"))
_ensure_rocm_torch()
# 14. Final check (silent; third-party conflicts are expected)