From ad77ae6cae1a9e357b394437c3156f4ca884d641 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 8 Jun 2026 07:06:10 -0700 Subject: [PATCH] =?UTF-8?q?fix(install):=20address=20PR=20review=20(Codex?= =?UTF-8?q?=20+=20Gemini)=20=E2=80=94=20exit=20codes,=20over-broad=20unins?= =?UTF-8?q?tall,=20Spark=20allocator,=20provision=20robustness?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit install.ps1 (WoA WSL fallback): - report failure (non-zero) + restore the rolled-aside venv when the WSL GPU install fails (torch.cuda absent) or when WSL needs enabling+reboot, instead of returning success — so -File/Tauri callers don't see a broken install as complete - on WSL success, Complete-StudioVenvRollback so the previous-venv backup isn't orphaned - refuse under --tauri with a clear "use the CLI installer" message (the desktop launcher resolves a Windows-venv backend, which a WSL-only install can't provide) - reset $LASTEXITCODE before each wsl.exe / python probe (a stale 0 could mark WSL ready / torch OK if the native command fails to launch) - torch-availability probe: --reinstall so an already-installed CPU torch in a migrated venv isn't accepted as "satisfied" (would wrongly skip the WSL path) - treat a null HKCU PATH as empty (fresh profile) so shim PATH update can't throw - keep apt stderr visible inside WSL (only stdout -> /dev/null) for diagnosability scripts/uninstall.ps1: - scope WSL cleanup to /root (the fallback's install location); stop deleting /home/*/.unsloth, which could erase an unrelated WSL user's own Unsloth/cache studio/setup.sh: - direct (non-install.ps1) WSL installs now provision CUDA llama.cpp themselves instead of being left with no GGUF server: install.ps1 exports UNSLOTH_WSL_LLAMA_DEFERRED=1, and the aarch64+NVIDIA provision block runs under WSL only when that marker is absent - mark a provisioner-built llama.cpp as Studio-owned in custom-STUDIO_HOME mode so the next setup's _assert_studio_owned_or_absent doesn't abort - glibc>=2.41 check: also match a future major>2 (e.g. 3.0) studio/scripts/provision_llama_cuda.sh: - install base tools (cmake/git/curl) in their own apt transaction before the best-effort gcc-14/g++-14 (unavailable on Ubuntu 22.04 / Debian 12, where bundling them aborted the whole transaction and left no build tools) - back up an existing (e.g. CPU-only) llama.cpp before the destructive clone and restore it on clone failure, so a failed clone doesn't leave the user with no server - honor a pinned llama.cpp ref via UNSLOTH_LLAMA_TAG instead of always tracking main unsloth/models/_utils.py: - set PYTORCH_CUDA_ALLOC_CONF (expandable_segments) via a CUDA-free Spark detector (nvidia-smi, not torch.cuda.get_device_name) so it takes effect before CUDA/the caching allocator initialize — previously it was a silent no-op on auto-detected Spark Co-Authored-By: Claude Opus 4.8 --- install.ps1 | 57 +++++++++++++++++++++++--- scripts/uninstall.ps1 | 11 ++--- studio/scripts/provision_llama_cuda.sh | 39 +++++++++++++++--- studio/setup.sh | 20 ++++++--- unsloth/models/_utils.py | 46 ++++++++++++++++++--- 5 files changed, 145 insertions(+), 28 deletions(-) diff --git a/install.ps1 b/install.ps1 index 7768343239..f4d4599d30 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1504,8 +1504,12 @@ shell.Run cmd, 0, False # as the real install ("torch>=2.4,<2.11.0"): a bare `torch` probe can match an out-of-range # wheel on the index, a false positive that skips WSL then fails the real pinned install. $prevEapProbe = $ErrorActionPreference; $ErrorActionPreference = "Continue" + # --reinstall forces resolution from the index instead of accepting an already-installed + # (e.g. CPU-only) torch in a migrated venv as "satisfied" -- otherwise the probe could pass + # without proving a native win_arm64 CUDA wheel exists, wrongly skipping the WSL path. + $global:LASTEXITCODE = -1 try { - & uv pip install --python $VenvPython --dry-run "torch>=2.4,<2.11.0" --index-url $TorchIndexUrl *> $null + & uv pip install --python $VenvPython --dry-run --reinstall "torch>=2.4,<2.11.0" --index-url $TorchIndexUrl *> $null $_nativeCudaTorchOk = ($LASTEXITCODE -eq 0) } catch { $_nativeCudaTorchOk = $false } finally { $ErrorActionPreference = $prevEapProbe } if ($_nativeCudaTorchOk) { step "gpu" "native CUDA PyTorch now available for win_arm64 -- keeping native install" "Green" } @@ -1514,8 +1518,19 @@ shell.Run cmd, 0, False step "wsl" "Windows on ARM + NVIDIA, native CUDA unavailable -- routing GPU setup through WSL2" substep "no win_arm64 CUDA PyTorch/Triton yet; WSL2 delivers full GPU (DGX Spark / RTX Spark path)." "Yellow" + # The bundled desktop app passes --tauri and launches its backend from a Windows venv + # (resolve_backend_binary), not from WSL -- so a WSL-only install would report complete yet + # fail to start. Until the Tauri launcher can drive a WSL backend, send desktop-app users to + # the CLI installer rather than leaving them with a broken-looking app. + if ($TauriMode) { + return (Exit-InstallFailure "Windows-on-ARM + NVIDIA GPU needs the WSL2 GPU install, which the desktop app can't launch yet. Install from PowerShell instead: irm https://unsloth.ai/install.ps1 | iex" 1) + } + $wslReady = $false if (Get-Command wsl.exe -ErrorAction SilentlyContinue) { + # Reset first: if wsl.exe throws/fails to start, $LASTEXITCODE keeps its prior value + # (a stale 0 from an earlier command would wrongly mark WSL ready). + $global:LASTEXITCODE = -1 try { & wsl.exe --status *> $null; if ($LASTEXITCODE -eq 0) { $wslReady = $true } } catch {} } @@ -1532,12 +1547,18 @@ shell.Run cmd, 0, False substep "in an ADMINISTRATOR PowerShell run: wsl --install" "Cyan" substep "reboot, then re-run: irm https://unsloth.ai/install.ps1 | iex" "Cyan" } + # WSL2 must be enabled + the machine rebooted before anything can install. Restore any + # rolled-aside previous venv and signal not-complete so -File callers don't treat this + # deferred state as a successful install. + Restore-StudioVenvRollback + $global:LASTEXITCODE = 1 return } $distro = if ($env:UNSLOTH_WSL_DISTRO) { $env:UNSLOTH_WSL_DISTRO } else { "Ubuntu-24.04" } # Detect the distro by exit code (encoding-proof; wsl --list emits UTF-16 that PS mis-parses). $haveDistro = $false + $global:LASTEXITCODE = -1 try { & wsl.exe -d $distro -- true *> $null; if ($LASTEXITCODE -eq 0) { $haveDistro = $true } } catch {} if (-not $haveDistro) { substep "installing WSL distro '$distro' (first time only)..." "Cyan" @@ -1548,15 +1569,22 @@ shell.Run cmd, 0, False # setup.sh + unsloth patches (otherwise install.sh pulls released PyPI unsloth and the # branch never runs pre-merge). main is byte-identical to plain unsloth.ai/install.sh. $_instRef = Get-UnslothInstallRef + # UNSLOTH_WSL_LLAMA_DEFERRED=1 tells the inner setup.sh that install.ps1 will build the CUDA + # llama.cpp in the background after install -- so setup.sh skips its own foreground build. + # (A user who runs install.sh DIRECTLY inside WSL won't set it, so setup.sh provisions CUDA + # itself instead of leaving them with no GGUF server.) + # apt stderr is kept visible (only stdout -> /dev/null) so network/DNS/repo failures inside + # WSL are diagnosable rather than silently swallowed. if ($_instRef -eq 'main') { - $wslInstall = 'export DEBIAN_FRONTEND=noninteractive; apt-get update -y >/dev/null 2>&1; apt-get install -y build-essential cmake git curl pciutils >/dev/null 2>&1; curl -fsSL https://unsloth.ai/install.sh | sh' + $wslInstall = 'export DEBIAN_FRONTEND=noninteractive UNSLOTH_WSL_LLAMA_DEFERRED=1; apt-get update -y >/dev/null; apt-get install -y build-essential cmake git curl pciutils >/dev/null; curl -fsSL https://unsloth.ai/install.sh | sh' } else { - $wslInstall = 'export DEBIAN_FRONTEND=noninteractive; export UNSLOTH_INSTALL_REF=' + $_instRef + '; apt-get update -y >/dev/null 2>&1; apt-get install -y build-essential cmake git curl pciutils >/dev/null 2>&1; curl -fsSL https://raw.githubusercontent.com/unslothai/unsloth/' + $_instRef + '/install.sh | sh' + $wslInstall = 'export DEBIAN_FRONTEND=noninteractive UNSLOTH_WSL_LLAMA_DEFERRED=1; export UNSLOTH_INSTALL_REF=' + $_instRef + '; apt-get update -y >/dev/null; apt-get install -y build-essential cmake git curl pciutils >/dev/null; curl -fsSL https://raw.githubusercontent.com/unslothai/unsloth/' + $_instRef + '/install.sh | sh' } # install.sh may exit non-zero on the optional llama.cpp prebuilt step (no aarch64 prebuilt) # though torch + unsloth + Studio still install, so lower EAP so it doesn't abort under Stop. $prevEapWsl = $ErrorActionPreference $ErrorActionPreference = "Continue" + $global:LASTEXITCODE = -1 try { & wsl.exe -d $distro --cd /root -u root -- bash -lc $wslInstall $wslRc = $LASTEXITCODE @@ -1568,6 +1596,8 @@ shell.Run cmd, 0, False $torchOk = $false $prevEapChk = $ErrorActionPreference $ErrorActionPreference = "Continue" + # Reset first so a stale 0 from a prior command can't mark torch OK if this fails to launch. + $global:LASTEXITCODE = -1 try { & wsl.exe -d $distro --cd /root -u root -- /root/.unsloth/studio/unsloth_studio/bin/python -c "import torch,sys; sys.exit(0 if torch.cuda.is_available() else 3)" *> $null $torchOk = ($LASTEXITCODE -eq 0) @@ -1624,9 +1654,13 @@ shell.Run cmd, 0, False "wsl.exe -d $distro -u root -- /root/.unsloth/studio/unsloth_studio/bin/unsloth %*" ) Set-Content -LiteralPath (Join-Path $shimDir "unsloth.cmd") -Value $shimLines -Encoding ASCII + # A fresh Windows profile may have no HKCU 'Path' value at all -> $userPath is null + # and $userPath.TrimEnd() would throw, losing the shim. Treat null as empty. $userPath = [Environment]::GetEnvironmentVariable("Path", "User") + if (-not $userPath) { $userPath = "" } if (($userPath -split ';') -notcontains $shimDir) { - [Environment]::SetEnvironmentVariable("Path", ($userPath.TrimEnd(';') + ";" + $shimDir), "User") + $newUserPath = if ($userPath.Trim()) { $userPath.TrimEnd(';') + ";" + $shimDir } else { $shimDir } + [Environment]::SetEnvironmentVariable("Path", $newUserPath, "User") } $env:Path = $env:Path.TrimEnd(';') + ";" + $shimDir step "shim" "created native 'unsloth' command -> forwards to WSL '$distro'" "Green" @@ -1742,8 +1776,19 @@ shell.Run cmd, 0, False step "wsl" "WSL Studio install did not finish cleanly (torch.cuda not detected; inner exit $wslRc) -- see log above." "Yellow" substep "retry, or launch manually: wsl -d $distro -u root -- bash -lic 'unsloth studio -p 8888'" "Cyan" } - substep "GPU training + GGUF export run inside WSL. (GGUF *inference* additionally needs a CUDA llama.cpp build.)" "Yellow" - if ($torchOk) { $global:LASTEXITCODE = 0 } + if ($torchOk) { + # WSL GPU install succeeded. On this path the Windows venv is vestigial (everything + # runs in WSL), so drop the rolled-aside previous-venv backup instead of orphaning it. + Complete-StudioVenvRollback + substep "GPU training + GGUF export run inside WSL. (GGUF *inference* additionally needs a CUDA llama.cpp build.)" "Yellow" + $global:LASTEXITCODE = 0 + return + } + # WSL GPU install failed (torch.cuda unavailable). Restore any rolled-aside previous venv so + # a reinstall-over-existing isn't left worse off, and report non-zero so -File callers don't + # treat a broken install as success. + Restore-StudioVenvRollback + $global:LASTEXITCODE = 1 return } diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 5013e4a85b..46f9a87a77 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -374,11 +374,12 @@ function Uninstall-UnslothStudio { # `wsl --list` emits UTF-16 PowerShell mis-parses (empty list -> cleanup skipped), so probe a # candidate set by exit code instead ('' = default distro), which is encoding-proof. # rm runs FIRST (guaranteed) since the kills could SIGKILL this shell. Also rm the dangling - # ~/.local/bin/unsloth symlink (its target under ~/.unsloth is gone but the link still resolves - # on PATH). pkill patterns use the [x]-regex self-exclusion trick: '[u]nsloth_studio' keeps the - # shell's own argv from matching (no literal "unsloth_studio" substring) while real processes - # still match. Same for '[l]lama-server' (a dynamic port not covered by fuser -k 8888). - $_clean = 'rm -rf /root/.unsloth /home/*/.unsloth /root/llama-cuda /root/provision_llama_cuda.sh /root/llama_cuda_build.log 2>/dev/null; rm -f /root/.local/bin/unsloth /home/*/.local/bin/unsloth 2>/dev/null; fuser -k 8888/tcp 2>/dev/null; pkill -9 -f ''[u]nsloth_studio'' 2>/dev/null; pkill -9 -f ''[l]lama-server'' 2>/dev/null; true' + # /root/.local/bin/unsloth symlink (its target under /root/.unsloth is gone but the link still + # resolves on PATH). Scope STRICTLY to /root: the WoA fallback installs there (wsl -u root), so + # touching /home/*/.unsloth would erase an unrelated WSL user's own Unsloth/cache that this + # installer never created. pkill patterns use the [x]-regex self-exclusion trick: '[u]nsloth_studio' + # keeps the shell's own argv from matching while real processes still match. Same for '[l]lama-server'. + $_clean = 'rm -rf /root/.unsloth /root/llama-cuda /root/provision_llama_cuda.sh /root/llama_cuda_build.log 2>/dev/null; rm -f /root/.local/bin/unsloth 2>/dev/null; fuser -k 8888/tcp 2>/dev/null; pkill -9 -f ''[u]nsloth_studio'' 2>/dev/null; pkill -9 -f ''[l]lama-server'' 2>/dev/null; true' $_cands = @('', 'Ubuntu', 'Ubuntu-24.04', 'Ubuntu-22.04', 'Debian') if ($env:UNSLOTH_WSL_DISTRO) { $_cands = @($env:UNSLOTH_WSL_DISTRO) + $_cands } $_done = @{} diff --git a/studio/scripts/provision_llama_cuda.sh b/studio/scripts/provision_llama_cuda.sh index 5be96aa89b..a8c15fbbd0 100644 --- a/studio/scripts/provision_llama_cuda.sh +++ b/studio/scripts/provision_llama_cuda.sh @@ -42,11 +42,16 @@ fi SUDO=""; [ "$(id -u)" -ne 0 ] && SUDO="sudo" HAVE_APT=0; command -v apt-get >/dev/null 2>&1 && HAVE_APT=1 -# 2. Base toolchain. gcc-14 is required because nvcc rejects gcc-15. +# 2. Base toolchain (must succeed) THEN gcc-14 (best-effort, separate transaction). +# gcc-14 is preferred because nvcc rejects gcc-15, but it isn't in the default apt +# sources on Ubuntu 22.04 / Debian 12 -- installing it in the SAME transaction as +# cmake/git/curl would make apt abort the whole transaction there, leaving the box +# without the basic build tools needed to clone/configure llama.cpp. if [ "$HAVE_APT" -eq 1 ]; then $SUDO apt-get update -y >/dev/null 2>&1 || true $SUDO apt-get install -y --no-install-recommends \ - build-essential cmake git curl ca-certificates gcc-14 g++-14 >/dev/null 2>&1 || true + build-essential cmake git curl ca-certificates >/dev/null 2>&1 || true + $SUDO apt-get install -y --no-install-recommends gcc-14 g++-14 >/dev/null 2>&1 || true fi # 3. Locate nvcc; install the CUDA toolkit if missing. @@ -102,12 +107,34 @@ export CC="$HCC" CXX="$HCXX" CUDAHOSTCXX="$HCXX" CC_CAP="$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -1 | tr -d ' .')" if [ -n "$CC_CAP" ]; then CUDA_ARCH="$CC_CAP"; else CUDA_ARCH="native"; fi -# 6. Clone + build into ~/.unsloth/llama.cpp. +# 6. Clone + build into ~/.unsloth/llama.cpp. Honor a pinned llama.cpp ref +# (UNSLOTH_LLAMA_TAG, the same var setup.sh uses) so a provisioner-built tree matches +# the user's request instead of always tracking ggml-org main. mkdir -p "$(dirname "$LLAMA_DIR")" +_LLAMA_REF="${UNSLOTH_LLAMA_TAG:-}" if [ ! -d "$LLAMA_DIR/.git" ]; then - rm -rf "$LLAMA_DIR" - git clone --depth 1 https://github.com/ggml-org/llama.cpp "$LLAMA_DIR" >/dev/null 2>&1 \ - || { log "git clone failed"; exit 0; } + # Preserve any existing (e.g. CPU-only) llama.cpp so a FAILED clone doesn't leave the user + # with NO server -- restore it on clone failure. A successful clone makes it obsolete (the + # fresh CUDA build replaces it), so the backup is dropped then. + _LLAMA_BAK="" + if [ -e "$LLAMA_DIR" ]; then + _LLAMA_BAK="${LLAMA_DIR}.prev.$$" + rm -rf "$_LLAMA_BAK" 2>/dev/null + mv "$LLAMA_DIR" "$_LLAMA_BAK" 2>/dev/null || { rm -rf "$LLAMA_DIR" 2>/dev/null; _LLAMA_BAK=""; } + fi + _clone_ok=0 + if [ -n "$_LLAMA_REF" ]; then + git clone --depth 1 --branch "$_LLAMA_REF" https://github.com/ggml-org/llama.cpp "$LLAMA_DIR" >/dev/null 2>&1 && _clone_ok=1 + fi + if [ "$_clone_ok" -ne 1 ]; then + git clone --depth 1 https://github.com/ggml-org/llama.cpp "$LLAMA_DIR" >/dev/null 2>&1 && _clone_ok=1 + fi + if [ "$_clone_ok" -ne 1 ]; then + log "git clone failed" + [ -n "$_LLAMA_BAK" ] && mv "$_LLAMA_BAK" "$LLAMA_DIR" 2>/dev/null # restore previous server + exit 0 + fi + [ -n "$_LLAMA_BAK" ] && rm -rf "$_LLAMA_BAK" 2>/dev/null fi cd "$LLAMA_DIR" || exit 0 diff --git a/studio/setup.sh b/studio/setup.sh index 9d5053b6af..6ad42f0ff6 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1203,7 +1203,8 @@ else if [ -n "$_GLIBC_VER" ]; then _GLIBC_MAJ="${_GLIBC_VER%%.*}"; _GLIBC_MIN="${_GLIBC_VER#*.}"; _GLIBC_MIN="${_GLIBC_MIN%%.*}" _CU_MAJ="${_NVCC_VER%%.*}"; _CU_MIN="${_NVCC_VER#*.}"; _CU_MIN="${_CU_MIN%%.*}" - if [ "${_GLIBC_MAJ:-0}" -eq 2 ] 2>/dev/null && [ "${_GLIBC_MIN:-0}" -ge 41 ] 2>/dev/null \ + if { [ "${_GLIBC_MAJ:-0}" -gt 2 ] 2>/dev/null \ + || { [ "${_GLIBC_MAJ:-0}" -eq 2 ] 2>/dev/null && [ "${_GLIBC_MIN:-0}" -ge 41 ] 2>/dev/null; }; } \ && { [ "${_CU_MAJ:-0}" -lt 13 ] 2>/dev/null \ || { [ "${_CU_MAJ:-0}" -eq 13 ] 2>/dev/null && [ "${_CU_MIN:-0}" -lt 3 ] 2>/dev/null; }; }; then substep "CUDA toolkit ${_NVCC_VER} is incompatible with glibc ${_GLIBC_VER} (rsqrt/rsqrtf header clash)." "$C_ERR" @@ -1444,14 +1445,17 @@ _have_cuda_llama_server() { } if [ "$_HOST_SYSTEM" = "Linux" ] \ && { [ "$_HOST_MACHINE" = "aarch64" ] || [ "$_HOST_MACHINE" = "arm64" ]; } \ - && ! grep -qi microsoft /proc/version 2>/dev/null \ + && { ! grep -qi microsoft /proc/version 2>/dev/null || [ "${UNSLOTH_WSL_LLAMA_DEFERRED:-0}" != "1" ]; } \ && [ "${UNSLOTH_NO_LLAMA_CUDA:-0}" != "1" ] \ && command -v nvidia-smi >/dev/null 2>&1 \ && nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}' \ && ! _have_cuda_llama_server; then - # WSL2 is excluded above: there install.ps1 runs this in the background after - # setup, so doing it here would duplicate the work in the foreground. Native - # Linux (DGX Spark / GB10) only. + # Native Linux (DGX Spark / GB10) runs this. Under WSL it runs ONLY for a DIRECT + # `install.sh` invocation: when install.ps1 drives the WSL install it exports + # UNSLOTH_WSL_LLAMA_DEFERRED=1 and builds the CUDA llama.cpp in the background after + # setup, so this foreground build is skipped to avoid duplicating it. A user who runs + # install.sh themselves inside WSL has no background builder, so we provision here + # rather than leave them with no GGUF server. # Resolve provision_llama_cuda.sh: copy beside setup.sh, then local-dev repo, # else fetch from GitHub so `curl | sh` works on an older wheel without it. _PROV_SH="" @@ -1475,6 +1479,12 @@ if [ "$_HOST_SYSTEM" = "Linux" ] \ if _have_cuda_llama_server; then step "llama.cpp" "CUDA llama-server ready (aarch64 + NVIDIA)" _LLAMA_CPP_DEGRADED=false + # The provisioner just created $LLAMA_CPP_DIR. In custom-STUDIO_HOME mode the next + # setup/update runs _assert_studio_owned_or_absent on it, so claim ownership now or + # that assert would abort on a directory this installer made. + if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + : > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true + fi elif [ -f "$LLAMA_SERVER_BIN" ]; then substep "CUDA build unavailable; keeping existing (CPU) llama-server" "$C_WARN" else diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index f21863c6b2..25cc4d32a3 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -936,6 +936,38 @@ def is_dgx_spark(): return False +@functools.lru_cache(maxsize = None) +def _is_dgx_spark_no_cuda_init(): + """Spark detection that never initializes a CUDA context. + + `is_dgx_spark()` calls `torch.cuda.get_device_name()`, which lazily initializes CUDA + (and the caching allocator). Settings consumed at allocator-init time -- + `PYTORCH_CUDA_ALLOC_CONF` (expandable_segments) -- must be decided BEFORE that, so this + variant reads the GPU name from `nvidia-smi` (a separate process) instead of torch. + Honors the same UNSLOTH_FORCE_DGX_SPARK override. Falls back to False on any error. + """ + _force = os.environ.get("UNSLOTH_FORCE_DGX_SPARK") + if _force == "1": + return True + if _force == "0": + return False + try: + import platform + + if platform.machine().lower() not in ("aarch64", "arm64"): + return False + import subprocess + + out = subprocess.run( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output = True, text = True, timeout = 5, + ) + names = (out.stdout or "").upper() + return any(token in names for token in _DGX_SPARK_DEVICE_TOKENS) + except Exception: + return False + + def patch_dgx_spark_caching_allocator_warmup(): """No-op `transformers.modeling_utils.caching_allocator_warmup` on Spark UMA. @@ -975,13 +1007,15 @@ def patch_dgx_spark_memory_config(): fragmentation OOMs; headroom for larger models / longer sequences). Pure memory management: it never changes any computed value, so accuracy is unaffected. - Strictly no-op off-Spark (gated by `is_dgx_spark()`). Respects an existing - PYTORCH_CUDA_ALLOC_CONF (only appends `expandable_segments` when absent, never - overrides a user's setting) and an explicit opt-out - (UNSLOTH_NO_EXPANDABLE_SEGMENTS=1). Must run before the first CUDA allocation; - `import unsloth` precedes model load, so it is set in time for normal use. + Strictly no-op off-Spark. Respects an existing PYTORCH_CUDA_ALLOC_CONF (only appends + `expandable_segments` when absent, never overrides a user's setting) and an explicit + opt-out (UNSLOTH_NO_EXPANDABLE_SEGMENTS=1). Must run before the first CUDA allocation, + so it gates on the CUDA-free `_is_dgx_spark_no_cuda_init()` -- the regular + `is_dgx_spark()` calls `torch.cuda.get_device_name()`, which would initialize CUDA (and + the allocator) before this env var could take effect. `import unsloth` precedes model + load, so it is set in time for normal use. """ - if not is_dgx_spark(): + if not _is_dgx_spark_no_cuda_init(): return if os.environ.get("UNSLOTH_NO_EXPANDABLE_SEGMENTS") == "1": return