Installer: GPU detection follow-ups after #6174 (poisoned venv repair, llama.cpp routing, probe bounds) (#6183)
* Installer: harden GPU detection follow-ups after #6174 Ports the NVIDIA-priority and /proc/driver/nvidia/gpus hardening from #6174 to the remaining pathways and adds recovery for already-poisoned venvs: - install_python_stack.py: add _ensure_cuda_torch so 'unsloth studio update' force-reinstalls CUDA torch when the venv carries a ROCm build on an NVIDIA Linux host (the pre-#6174 poisoning signature). Honors UNSLOTH_TORCH_BACKEND, UNSLOTH_ROCM_TORCH_INSTALLED, and CUDA_VISIBLE_DEVICES=-1/'' opt-outs; never touches healthy CUDA, deliberate CPU wheels, macOS, or Windows. - install_llama_prebuilt.py: detect_host gains the /proc NVIDIA fallback and skips ROCm probes when NVIDIA is usable; forwarded --rocm-gfx/--has-rocm overrides still win. - setup.sh: GPU summary classifies NVIDIA first through a timeout-bounded probe with the /proc fallback; AMD probes are bounded and gain a KFD vendor_id 4098 fallback; the llama.cpp source build only selects GGML_CUDA/GGML_HIP when the matching GPU is actually detected. - install.sh: bound both nvidia-smi calls with a 10s timeout (no behavior change when healthy or when the timeout binary is absent); classify the exported UNSLOTH_TORCH_BACKEND on the final index path segment so custom mirrors containing 'rocm'/'gfx' in their base path are not mislabeled. - install.ps1 + setup.ps1: NVIDIA probes now require a real 'GPU N:' row from nvidia-smi -L under a 10s bound instead of bare exit code 0; later CUDA version and compute_cap queries are bounded too. Tests: 3 new test files (50+ tests), suite at 788 passed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Resolve-CudaToolkit driver probe for extracted-function unit test tests/studio/test_resolve_cuda_toolkit.ps1 extracts Resolve-CudaToolkit alone into a child pwsh and stubs nvidia-smi with a .ps1 script. The bounded runner is not in scope there (and ProcessStartInfo cannot dispatch .ps1 stubs), so the DriverMaxCuda parse silently returned nothing and the major-mismatch scenarios failed. Fall back to direct invocation when Invoke-NvidiaSmiBounded is unavailable; production setup.ps1 always has it defined and keeps the 10s bound. * Treat CUDA_VISIBLE_DEVICES empty or -1 as hidden in NVIDIA-first guards The NVIDIA-first guards added in this branch only special-cased CUDA_VISIBLE_DEVICES=-1 at two setup.sh gates and ignored the empty-string form entirely, while the Python detector (install_llama_prebuilt.py) already treats both as hidden. On a mixed AMD+NVIDIA host steered to the AMD card via CUDA_VISIBLE_DEVICES, the guards suppressed the AMD probes, so setup.sh fell to a CPU llama.cpp build and install.sh picked CUDA wheels instead of ROCm. Move the policy into the helpers so every consumer agrees: - install.sh: new _cvd_hides_nvidia checked first in _has_usable_nvidia_gpu - studio/setup.sh: same via _setup_cvd_hides_nvidia; the two ad-hoc CUDA_VISIBLE_DEVICES=-1 gate conditions are now redundant and removed - studio/install_python_stack.py: _has_usable_nvidia_gpu returns False when CUDA_VISIBLE_DEVICES is set to or -1 (whitespace tolerated) Tests: 5 new sh scenarios (hidden via , -1, padded -1, visible device, and mixed host with hidden NVIDIA restoring the ROCm route) plus a pytest class covering all three implementations behaviourally. Addresses the review comment on the NVIDIA-first setup.sh block. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Retrigger CI after PyPI 503 outage during the previous run --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
28fe978549
commit
2db9fad4b5
10 changed files with 1439 additions and 74 deletions
|
|
@ -90,6 +90,18 @@ _PYTORCH_WHL_BASE = (
|
|||
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
|
||||
).rstrip("/")
|
||||
|
||||
# CUDA torch repair specs (see _ensure_cuda_torch). torchvision/torchaudio are
|
||||
# pinned to the torch<2.11 family rather than left bare: the install uses an
|
||||
# exclusive --index-url (no PyPI fallback), so a bare name could resolve a
|
||||
# torchvision built against a different torch major (e.g. 0.27 for torch 2.12)
|
||||
# and fail at runtime with an ABI mismatch. Same bounds as the _default ROCm
|
||||
# spec above, which targets the same torch family.
|
||||
_CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = (
|
||||
"torch>=2.4,<2.11.0",
|
||||
"torchvision>=0.19,<0.26.0",
|
||||
"torchaudio>=2.4,<2.11.0",
|
||||
)
|
||||
|
||||
# AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/).
|
||||
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs.
|
||||
_ROCM_WINDOWS_INDEX_BASE = (
|
||||
|
|
@ -654,7 +666,15 @@ def _has_usable_nvidia_gpu() -> bool:
|
|||
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.
|
||||
|
||||
CUDA_VISIBLE_DEVICES set to "" or "-1" hides every NVIDIA device (mixed
|
||||
AMD+NVIDIA hosts steering work to the AMD card); neither probe honours
|
||||
that env var, so check it first and report the GPU as not usable. Unset
|
||||
means all devices visible.
|
||||
"""
|
||||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||||
if cvd is not None and cvd.strip() in ("", "-1"):
|
||||
return False
|
||||
exe = shutil.which("nvidia-smi")
|
||||
if exe:
|
||||
try:
|
||||
|
|
@ -786,6 +806,139 @@ def _install_bnb_windows_rocm() -> bool:
|
|||
return True
|
||||
|
||||
|
||||
def _detect_cuda_torch_index_url() -> str:
|
||||
"""Return the pytorch.org CUDA wheel index URL for the host's NVIDIA driver.
|
||||
|
||||
Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update`
|
||||
repairs to the same wheel family a fresh `curl | sh` install would pick.
|
||||
Probes nvidia-smi (PATH, then /usr/bin/nvidia-smi) and parses both the
|
||||
legacy "CUDA Version:" and the newer "CUDA UMD Version:" spellings.
|
||||
Defaults to cu126 when nvidia-smi is missing or the version is unreadable
|
||||
(e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback).
|
||||
"""
|
||||
exe = shutil.which("nvidia-smi")
|
||||
if not exe and os.path.isfile("/usr/bin/nvidia-smi"):
|
||||
exe = "/usr/bin/nvidia-smi"
|
||||
tag = "cu126" # default when the driver CUDA version cannot be read
|
||||
if exe:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[exe],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.DEVNULL,
|
||||
text = True,
|
||||
timeout = 10,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
m = re.search(r"CUDA(?: UMD)? Version:\s*(\d+)\.(\d+)", result.stdout)
|
||||
if m:
|
||||
major, minor = int(m.group(1)), int(m.group(2))
|
||||
if major >= 13:
|
||||
tag = "cu130"
|
||||
elif major == 12 and minor >= 8:
|
||||
tag = "cu128"
|
||||
elif major == 12 and minor >= 6:
|
||||
tag = "cu126"
|
||||
elif major >= 12:
|
||||
tag = "cu124"
|
||||
elif major >= 11:
|
||||
tag = "cu118"
|
||||
else:
|
||||
tag = "cpu" # ancient driver: no usable CUDA wheels
|
||||
except Exception:
|
||||
pass
|
||||
return f"{_PYTORCH_WHL_BASE}/{tag}"
|
||||
|
||||
|
||||
def _ensure_cuda_torch() -> None:
|
||||
"""Repair a venv whose torch is a ROCm build on an NVIDIA host.
|
||||
|
||||
Counterpart to _ensure_rocm_torch. A venv poisoned by the pre-fix KFD
|
||||
gpu_id false positive (ROCm torch installed on an NVIDIA-only machine)
|
||||
keeps that broken torch on `studio update`, because a torch+rocm wheel
|
||||
satisfies the version constraint and nothing force-reinstalls it. This
|
||||
detects that exact case and reinstalls CUDA torch.
|
||||
|
||||
Only repairs when torch actually links against HIP/ROCm. Healthy CUDA
|
||||
torch and deliberate CPU-only torch are left untouched.
|
||||
"""
|
||||
# Respect an explicit backend choice from install.sh: only "" (standalone
|
||||
# `studio update`) or "cuda" should ever force CUDA wheels. "rocm"/"cpu"
|
||||
# (or any unrecognised value) are deliberate and must not be overridden.
|
||||
if _TORCH_BACKEND not in ("", "cuda"):
|
||||
return
|
||||
# No CUDA torch on macOS; Windows venv/torch lifecycle is owned by
|
||||
# install.ps1 (and the KFD poisoning bug is Linux-only), so skip both.
|
||||
if IS_MACOS or IS_WINDOWS or NO_TORCH:
|
||||
return
|
||||
# Never undo a deliberate ROCm install (setup.ps1 sets this marker).
|
||||
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
|
||||
return
|
||||
# CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU (for
|
||||
# example a mixed AMD+NVIDIA host that runs ROCm torch on the AMD card);
|
||||
# never force CUDA wheels over that choice.
|
||||
_cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||||
if _cvd is not None and _cvd.strip() in ("", "-1"):
|
||||
return
|
||||
# Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu()
|
||||
# covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent.
|
||||
if not _has_usable_nvidia_gpu():
|
||||
return
|
||||
|
||||
# Classify the installed torch: "hip" (ROCm build -- the poisoning
|
||||
# signature), "cuda" (healthy), or "cpu" (deliberate CPU wheel). A
|
||||
# non-zero exit means torch is missing or un-importable; the base install
|
||||
# step handles that, so leave it alone.
|
||||
try:
|
||||
probe = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
(
|
||||
"import torch; "
|
||||
"hip = getattr(torch.version, 'hip', '') or ''; "
|
||||
"cuda = getattr(torch.version, 'cuda', '') or ''; "
|
||||
"ver = getattr(torch, '__version__', '').lower(); "
|
||||
"print('hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'))"
|
||||
),
|
||||
],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.DEVNULL,
|
||||
timeout = 90,
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
return
|
||||
if probe.returncode != 0:
|
||||
return
|
||||
# Take the last non-empty stdout line: stray output from sitecustomize or
|
||||
# an import hook must not mask the marker (fail-closed either way).
|
||||
_marker_lines = [
|
||||
line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip()
|
||||
]
|
||||
if not _marker_lines or _marker_lines[-1] != "hip":
|
||||
return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is
|
||||
|
||||
index_url = _detect_cuda_torch_index_url()
|
||||
_torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC
|
||||
print(
|
||||
f" torch is a ROCm build on an NVIDIA host -- reinstalling "
|
||||
f"CUDA torch from {index_url}\n"
|
||||
f" (set UNSLOTH_TORCH_BACKEND=rocm to keep a deliberate ROCm torch "
|
||||
f"on a mixed AMD+NVIDIA host)"
|
||||
)
|
||||
pip_install(
|
||||
"CUDA torch repair",
|
||||
"--force-reinstall",
|
||||
"--no-cache-dir",
|
||||
_torch_pkg,
|
||||
_vision_pkg,
|
||||
_audio_pkg,
|
||||
"--index-url",
|
||||
index_url,
|
||||
constrain = False,
|
||||
)
|
||||
|
||||
|
||||
def _ensure_rocm_torch() -> None:
|
||||
"""Reinstall torch with ROCm wheels when the venv received CPU-only torch.
|
||||
|
||||
|
|
@ -1848,6 +2001,7 @@ def install_python_stack() -> int:
|
|||
# Must follow base packages so torch is present for inspection.
|
||||
if not IS_MACOS and not NO_TORCH:
|
||||
_progress(_torch_step_label("check"))
|
||||
_ensure_cuda_torch()
|
||||
_ensure_rocm_torch()
|
||||
|
||||
# Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python
|
||||
|
|
@ -2033,6 +2187,7 @@ def install_python_stack() -> int:
|
|||
# whichever intermediate step clobbered it.
|
||||
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
|
||||
_progress(_torch_step_label("final"))
|
||||
_ensure_cuda_torch()
|
||||
_ensure_rocm_torch()
|
||||
|
||||
# 14. Final check (silent; third-party conflicts are expected)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue