From b50eb8bc7182e883ef1165848b4c8f9bda5c0888 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 08:25:40 +0000 Subject: [PATCH] Tighten and trim code comments --- install.ps1 | 86 ++++++++++---------- install.sh | 16 ++-- scripts/uninstall.ps1 | 22 ++--- scripts/uninstall.sh | 5 +- studio/backend/core/training/worker.py | 21 +++-- studio/backend/tests/test_spark_oom_guard.py | 11 ++- studio/scripts/provision_llama_cuda.sh | 30 +++---- studio/setup.sh | 22 ++--- unsloth/kernels/flex_attention.py | 10 +-- unsloth/models/_utils.py | 35 ++++---- 10 files changed, 125 insertions(+), 133 deletions(-) diff --git a/install.ps1 b/install.ps1 index b2486ad67d..eba41dfbe2 100644 --- a/install.ps1 +++ b/install.ps1 @@ -48,8 +48,7 @@ function Install-UnslothStudio { } } - # raw.githubusercontent.com ref for install assets (provision_llama_cuda.sh, .ico). - # UNSLOTH_INSTALL_REF overrides 'main' for pre-merge testing. + # raw.githubusercontent.com ref for install assets; UNSLOTH_INSTALL_REF overrides 'main' for pre-merge testing. function Get-UnslothInstallRef { if ($env:UNSLOTH_INSTALL_REF -and $env:UNSLOTH_INSTALL_REF.Trim()) { return $env:UNSLOTH_INSTALL_REF.Trim() } return 'main' @@ -97,8 +96,8 @@ function Install-UnslothStudio { if ($TauriMode) { exit $Code } - # -File ignores $LASTEXITCODE on plain return, so `exit` must carry the code; - # under `irm | iex` (no $PSCommandPath) `exit` would kill the user's shell. + # -File ignores $LASTEXITCODE on plain return, so `exit` carries the code; under + # `irm | iex` (no $PSCommandPath) `exit` would kill the user's shell, so set the var. if ($PSCommandPath) { exit $Code } @@ -1773,9 +1772,9 @@ shell.Run cmd, 0, False $TorchIndexUrl = Get-TorchIndexUrl # ===== Windows-on-ARM + NVIDIA GPU -> automatic WSL2 fallback (N1X "RTX Spark" / DGX Spark-class) ===== - # win_arm64 has no CUDA PyTorch/Triton wheel, so run the Linux installer inside WSL2 (full GPU) plus - # a Windows `unsloth` shim that forwards into it. x86_64 / ARM64-without-NVIDIA unaffected; the probe - # below keeps the native install if a win_arm64 CUDA torch wheel ever ships. + # win_arm64 has no CUDA PyTorch/Triton wheel, so run the Linux installer inside WSL2 (full GPU) plus a + # Windows `unsloth` shim forwarding into it; x86_64 / ARM64-without-NVIDIA unaffected, and the probe + # below keeps the native install if a win_arm64 CUDA wheel ever ships. # Opt out: UNSLOTH_NO_WSL_FALLBACK=1; pick distro with UNSLOTH_WSL_DISTRO. try { $_winArm64 = ([System.Runtime.InteropServices.RuntimeInformation]::OSArchitecture.ToString() -ieq 'Arm64') } catch { $_winArm64 = $false } # x64-emulated PS on ARM reports X64/AMD64; Win32_Processor.Architecture (12=ARM64) and machine-level @@ -1791,11 +1790,11 @@ shell.Run cmd, 0, False } $_nativeCudaTorchOk = $false if ($_winArm64 -and $HasNvidiaSmi -and (-not $SkipTorch)) { - # Probe with the SAME spec as the real install ("torch>=2.4,<2.11.0"): a bare `torch` probe - # could match an out-of-range wheel, skipping WSL only to fail the real pinned install. + # Probe the SAME spec as the real install ("torch>=2.4,<2.11.0"); a bare `torch` probe could + # match an out-of-range wheel, skipping WSL only to fail the real pinned install. $prevEapProbe = $ErrorActionPreference; $ErrorActionPreference = "Continue" - # --reinstall: an installed (e.g. CPU-only) torch mustn't satisfy the probe -- it must - # prove a native win_arm64 CUDA wheel exists on the index. + # --reinstall: an installed (e.g. CPU-only) torch mustn't satisfy the probe -- it must prove + # a native win_arm64 CUDA wheel exists on the index. $global:LASTEXITCODE = -1 try { & uv pip install --python $VenvPython --dry-run --reinstall "torch>=2.4,<2.11.0" --index-url $TorchIndexUrl *> $null @@ -1807,8 +1806,8 @@ 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 Tauri desktop app launches its backend from a Windows venv (resolve_backend_binary), - # not WSL, so a WSL-only install would start nothing -- send those users to the CLI installer. + # The Tauri desktop app launches its backend from a Windows venv (resolve_backend_binary), not + # WSL, so a WSL-only install would start nothing -- send those users to the CLI installer. 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) } @@ -1833,8 +1832,8 @@ shell.Run cmd, 0, False substep "reboot, then re-run: irm https://unsloth.ai/install.ps1 | iex" "Cyan" } # Deferred until reboot: restore any rolled-aside previous venv and signal not-complete. - # `exit 1` for -File (plain return exits 0); under `irm | iex` (no $PSCommandPath) return - # instead, since exit would kill the user's shell. + # `exit 1` for -File (plain return exits 0); under `irm | iex` (no $PSCommandPath) return, + # since exit would kill the user's shell. Restore-StudioVenvRollback $global:LASTEXITCODE = 1 if ($PSCommandPath) { exit 1 } @@ -1843,7 +1842,7 @@ shell.Run cmd, 0, False $distro = if ($env:UNSLOTH_WSL_DISTRO) { $env:UNSLOTH_WSL_DISTRO } else { "Ubuntu-24.04" } # For cmd-context uses (.cmd shim, copy-paste hints): wsl.exe rejects a QUOTED space-free name - # (WSL_E_DISTRO_NOT_FOUND on 2.x) yet splits a bare spaced one after -d, so quote ONLY when spaced. + # (WSL_E_DISTRO_NOT_FOUND on 2.x) but splits a bare spaced one after -d, so quote ONLY when spaced. $_distroArg = if ($distro -match '\s') { '"' + $distro + '"' } else { $distro } # Detect the distro by exit code (encoding-proof; wsl --list emits UTF-16 that PS mis-parses). $haveDistro = $false @@ -1851,15 +1850,15 @@ shell.Run cmd, 0, False 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" - # New distros install at the global default WSL version; force 2 so a host - # whose default is WSL1 doesn't get a GPU-less distro (fails only at torch.cuda). + # New distros install at the global default WSL version; force 2 so a WSL1-default + # host doesn't get a GPU-less distro (would fail only at torch.cuda). $global:LASTEXITCODE = -1 try { & wsl.exe --set-default-version 2 *> $null } catch {} try { & wsl.exe --install -d $distro --no-launch } catch {} } else { - # A PRE-EXISTING distro may be WSL1 (no GPU passthrough; would only fail at the final + # A PRE-EXISTING distro may be WSL1 (no GPU passthrough; would fail only at the final # torch.cuda check). Detect from inside (encoding-proof, unlike UTF-16 `wsl -l -v`) and - # convert in place -- `wsl --set-version` preserves files. Fresh installs default to WSL2. + # convert in place -- `wsl --set-version` preserves files. $_wsl2Probe = 'grep -qiE ''microsoft-standard|WSL2'' /proc/version 2>/dev/null || test -e /usr/lib/wsl/lib/libcuda.so' $_isWsl2 = $false $global:LASTEXITCODE = -1 @@ -1881,15 +1880,14 @@ shell.Run cmd, 0, False # Non-main ref: fetch + export THAT ref so the WSL venv gets the branch's setup.sh + patches # (else install.sh pulls PyPI unsloth). main == plain unsloth.ai/install.sh. $_instRef = Get-UnslothInstallRef - # UNSLOTH_WSL_LLAMA_DEFERRED=1: setup.sh skips its foreground CUDA llama.cpp build since we build - # it in the background (a DIRECT install.sh run in WSL doesn't set it). apt stderr stays visible - # (only stdout -> /dev/null) so network/repo failures are diagnosable. - # Forward UNSLOTH_NO_LLAMA_CUDA into WSL: it also skips the dispatch below, so unforwarded - # setup.sh would defer to a background builder that never starts (no llama-server). + # UNSLOTH_WSL_LLAMA_DEFERRED=1: setup.sh skips its foreground CUDA llama.cpp build since we build it + # in the background. apt stderr stays visible (only stdout -> /dev/null) so failures are diagnosable. + # Forward UNSLOTH_NO_LLAMA_CUDA into WSL: it also skips the dispatch below, so unforwarded setup.sh + # would defer to a background builder that never starts (no llama-server). $_fwdEnv = '' if ($env:UNSLOTH_NO_LLAMA_CUDA -eq '1') { $_fwdEnv = 'export UNSLOTH_NO_LLAMA_CUDA=1; ' } - # Forward a user Python pin: install.sh reads UNSLOTH_PYTHON, but a Windows env var - # isn't visible inside WSL unless bridged. Numeric-only guard (e.g. 3.12) = no injection. + # Forward a user Python pin (install.sh reads UNSLOTH_PYTHON, but Windows env vars don't cross + # into WSL unless bridged). Numeric-only guard (e.g. 3.12) prevents injection. if ($env:UNSLOTH_PYTHON -and ($env:UNSLOTH_PYTHON -match '^[0-9][0-9.]*$')) { $_fwdEnv += "export UNSLOTH_PYTHON=$($env:UNSLOTH_PYTHON); " } if ($_instRef -eq 'main') { $wslInstall = $_fwdEnv + '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 libcurl4-openssl-dev >/dev/null; curl -fsSL https://unsloth.ai/install.sh | sh' @@ -1897,7 +1895,7 @@ shell.Run cmd, 0, False $wslInstall = $_fwdEnv + '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 libcurl4-openssl-dev >/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; lower EAP so it doesn't abort under Stop. + # even though torch + unsloth + Studio install; lower EAP so it doesn't abort under Stop. $prevEapWsl = $ErrorActionPreference $ErrorActionPreference = "Continue" $global:LASTEXITCODE = -1 @@ -1918,9 +1916,9 @@ shell.Run cmd, 0, False & 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) } catch {} finally { $ErrorActionPreference = $prevEapChk } - # Self-heal web-server deps: a cut-short install.sh "studio deps" step leaves torch + unsloth - # but no fastapi/uvicorn/structlog/starlette (`unsloth studio` dies). Reinstall them unpinned - # (no huggingface-hub/transformers/datasets) so the verified GPU torch stack stays intact. + # Self-heal web-server deps: a cut-short install.sh "studio deps" step leaves torch + unsloth but + # no fastapi/uvicorn/structlog/starlette (`unsloth studio` dies). Reinstall them unpinned (no + # huggingface-hub/transformers/datasets) so the verified GPU torch stack stays intact. if ($torchOk) { $_studioPy = "/root/.unsloth/studio/unsloth_studio/bin/python" $_serverOk = $false @@ -1932,8 +1930,8 @@ shell.Run cmd, 0, False if (-not $_serverOk) { substep "Studio web-server deps incomplete (install.sh step cut short) -- installing them now..." "Cyan" # studio.txt minus the huggingface-hub pin; uv preferred, pip fallback. Bare names only: - # `>=` would become a redirection through PowerShell -> wsl.exe -> bash -lc, and - # latest-of-each satisfies the studio.txt minimums anyway. + # `>=` would become a redirection through PowerShell -> wsl.exe -> bash -lc, and latest-of-each + # satisfies the studio.txt minimums anyway. $_deps = 'typer fastapi uvicorn matplotlib pandas nest_asyncio pyjwt easydict addict structlog diceware ddgs cryptography httpx fastmcp sqlite-vec pymupdf python-docx' $_repair = 'PY=/root/.unsloth/studio/unsloth_studio/bin/python; UV="$(command -v uv 2>/dev/null || echo /root/.local/bin/uv)"; if [ -x "$UV" ] || command -v uv >/dev/null 2>&1; then "$UV" pip install --python "$PY" ' + $_deps + '; else "$PY" -m pip install ' + $_deps + '; fi' $prevEapR = $ErrorActionPreference; $ErrorActionPreference = "Continue" @@ -1946,8 +1944,8 @@ shell.Run cmd, 0, False if ($_serverOk) { substep "Studio web-server deps installed." "Green" } else { substep "(could not auto-install Studio server deps; 'unsloth studio' may fail to start)" "Yellow" } } - # The uv-managed venv ships no `pip`, but unsloth-zoo's exporter's check_pip() finds `uv pip` - # only when uv is on PATH. Seed pip so `save_pretrained_gguf` works regardless. + # The uv-managed venv ships no `pip`, but unsloth-zoo's check_pip() finds `uv pip` only + # when uv is on PATH. Seed pip so `save_pretrained_gguf` works regardless. $prevEapP = $ErrorActionPreference; $ErrorActionPreference = "Continue" try { & wsl.exe -d $distro --cd /root -u root -- $_studioPy -m pip --version *> $null @@ -1970,7 +1968,7 @@ shell.Run cmd, 0, False ) Set-Content -LiteralPath (Join-Path $shimDir "unsloth.cmd") -Value $shimLines -Encoding ASCII # Record the distro so the uninstaller can clean a custom UNSLOTH_WSL_DISTRO install - # without the env var set again. + # without the env var set. try { Set-Content -LiteralPath (Join-Path (Split-Path $shimDir -Parent) "wsl-distro.txt") -Value $distro -Encoding ASCII } catch {} # A fresh profile may have no HKCU 'Path'; null would make TrimEnd() throw. $userPath = [Environment]::GetEnvironmentVariable("Path", "User") @@ -2000,9 +1998,9 @@ shell.Run cmd, 0, False 'wsl.exe -d $distro --cd /root -u root -- bash -lic "unsloth studio -p 8888"' ) Set-Content -LiteralPath $launcher -Value $L -Encoding UTF8 - # Icon must live OUTSIDE %LOCALAPPDATA%: on WoA the sandboxed icon broker can't read a - # .ico under AppData\Local, so the shortcut renders BLANK; under the user profile it - # renders fine (verified on N1X). Only the icon moves. + # Icon must live OUTSIDE %LOCALAPPDATA%: on WoA the sandboxed icon broker can't read a .ico + # under AppData\Local, so the shortcut renders BLANK; under the user profile it renders + # fine (verified on N1X). Only the icon moves. $iconDir = Join-Path $env:USERPROFILE ".unsloth" New-Item -ItemType Directory -Force -Path $iconDir *> $null $icon = Join-Path $iconDir "unsloth.ico" @@ -2060,14 +2058,14 @@ shell.Run cmd, 0, False try { $_llamaUrl = "https://raw.githubusercontent.com/unslothai/unsloth/$(Get-UnslothInstallRef)/studio/scripts/provision_llama_cuda.sh" # Step 1: fetch the provision script + write a runner (base64 to dodge quoting layers). - # The runner restores PATH (non-login shells miss /usr/lib/wsl/lib nvidia-smi, so - # provision early-exits) and exports the env knobs below (Windows env vars don't cross - # into WSL). A runner FILE lets the detached launcher pass only space-free args, - # avoiding Start-Process mis-splitting `bash -lc `. + # The runner restores PATH (non-login shells miss /usr/lib/wsl/lib nvidia-smi, so provision + # early-exits) and exports the env knobs below (Windows env vars don't cross into WSL). A + # runner FILE lets the detached launcher pass only space-free args, avoiding Start-Process + # mis-splitting `bash -lc `. $_pathLine = 'export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/lib/wsl/lib:$PATH"' + "`n" $_jobsLine = if ($env:UNSLOTH_LLAMA_BUILD_JOBS) { "export UNSLOTH_LLAMA_BUILD_JOBS=$($env:UNSLOTH_LLAMA_BUILD_JOBS)`n" } else { "" } # Bridge UNSLOTH_LLAMA_TAG / UNSLOTH_LLAMA_PR pins into WSL, else the deferred build - # ignores them. sh-single-quoted (tags/PRs are simple tokens). + # ignores them. sh-single-quoted since tags/PRs are simple tokens. $_tagLine = if ($env:UNSLOTH_LLAMA_TAG) { "export UNSLOTH_LLAMA_TAG='$($env:UNSLOTH_LLAMA_TAG)'`n" } else { "" } $_prLine = if ($env:UNSLOTH_LLAMA_PR) { "export UNSLOTH_LLAMA_PR='$($env:UNSLOTH_LLAMA_PR)'`n" } else { "" } $_runner = "#!/usr/bin/env bash`n" + $_pathLine + $_jobsLine + $_tagLine + $_prLine + "exec bash /root/.unsloth/provision_llama_cuda.sh > /root/.unsloth/llama_cuda_build.log 2>&1`n" diff --git a/install.sh b/install.sh index 13d93531e7..2f8b3c0ff4 100755 --- a/install.sh +++ b/install.sh @@ -2648,10 +2648,9 @@ elif [ -n "$TORCH_INDEX_URL" ]; then --no-deps --reinstall-package unsloth-zoo \ "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" elif [ -n "${UNSLOTH_INSTALL_REF:-}" ] && [ "${UNSLOTH_INSTALL_REF}" != "main" ] && [ "$PACKAGE_NAME" = "unsloth" ]; then - # Pre-merge testing: install unsloth from a git ref (install.ps1 sets - # UNSLOTH_INSTALL_REF) so the branch's setup.sh + patches run. Name - # unsloth-zoo explicitly: it's not a base dep and SKIP_STUDIO_BASE skips - # base.txt, so otherwise it never installs. + # Pre-merge testing: install unsloth from a git ref (set by install.ps1) + # so the branch's setup.sh + patches run. Name unsloth-zoo explicitly -- + # not a base dep, and SKIP_STUDIO_BASE skips base.txt, so it never installs. substep "installing unsloth from git ref '$UNSLOTH_INSTALL_REF'..." run_install_cmd "install unsloth (@$UNSLOTH_INSTALL_REF)" uv pip install --python "$_VENV_PY" \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ @@ -2660,11 +2659,10 @@ elif [ -n "$TORCH_INDEX_URL" ]; then run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \ --upgrade-package unsloth -- "$PACKAGE_NAME" fi - # aarch64 + NVIDIA (DGX Spark / GB10 / N1X): unsloth's cuXXX extras are - # x86_64-oriented, so 4-bit QLoRA fails out of the box. aarch64 manylinux - # wheels work (verified on sm_121 via PTX JIT); best-effort, no wheel just - # keeps 16-bit LoRA / full finetuning. SKIP_TORCH gate: a --no-torch - # (GGUF-only) install must not let bitsandbytes drag torch back in. + # aarch64 + NVIDIA (DGX Spark / GB10 / N1X): unsloth's x86_64-oriented cuXXX + # extras break 4-bit QLoRA, but aarch64 manylinux wheels work (verified on + # sm_121 via PTX JIT). Best-effort: no wheel keeps 16-bit LoRA / full finetuning. + # SKIP_TORCH gate stops a --no-torch (GGUF-only) install dragging torch back in. if [ "$SKIP_TORCH" = false ] \ && { [ "$(uname -m)" = "aarch64" ] || [ "$(uname -m)" = "arm64" ]; } \ && command -v nvidia-smi >/dev/null 2>&1 \ diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 81f4dd58cb..7ff9ee3f66 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -341,8 +341,8 @@ function Uninstall-UnslothStudio { } # ── Remove desktop and Start Menu shortcuts ── - # Canonical name is "Unsloth Studio.lnk"; also sweep legacy distro-suffixed - # names ("Unsloth Studio (WSL - ).lnk") left by pre-release dev builds. + # Canonical name is "Unsloth Studio.lnk"; also sweep legacy distro-suffixed names + # ("Unsloth Studio (WSL - ).lnk") left by pre-release dev builds. _Step "Removing desktop and Start Menu shortcuts..." $shortcutDirs = @() try { $d = [Environment]::GetFolderPath("Desktop"); if ($d) { $shortcutDirs += $d } } catch { } @@ -423,11 +423,11 @@ function Uninstall-UnslothStudio { # ── Windows-on-Arm WSL-fallback artifacts ── # The ARM64+NVIDIA fallback puts Studio in WSL plus a native shim + launcher under - # %LOCALAPPDATA%\Unsloth (not "Unsloth Studio") with a PATH entry -- all missed above. + # %LOCALAPPDATA%\Unsloth (not "Unsloth Studio") with a PATH entry -- none caught above. _Step "Removing WSL-fallback artifacts (shim, launcher, PATH entry, WSL install)..." $unslothDir = if ($env:LOCALAPPDATA) { Join-Path $env:LOCALAPPDATA "Unsloth" } else { $null } - # wsl-distro.txt records a custom UNSLOTH_WSL_DISTRO install so it's cleanable without the - # env var set; read it BEFORE the directory is removed below. + # wsl-distro.txt records a custom UNSLOTH_WSL_DISTRO install so it's cleanable without the env + # var set; read it BEFORE the directory is removed below. $_recordedDistro = $null if ($unslothDir) { try { @@ -466,17 +466,17 @@ function Uninstall-UnslothStudio { try { # Probe candidates by exit code ('' = default distro) since `wsl --list` emits UTF-16 PS # mis-parses. rm runs FIRST (the kills could SIGKILL this shell) and drops the dangling - # /root/.local/bin/unsloth symlink. Scope STRICTLY to /root (where the fallback installs); + # /root/.local/bin/unsloth symlink. Scope STRICTLY to /root (the fallback's install dir); # /home/*/.unsloth may be another user's. The 8888 kill is gated on an Unsloth install - # existing (checked BEFORE rm deletes the marker) so an unrelated listener survives; pkill + # existing (checked BEFORE rm deletes the marker) so an unrelated listener survives. pkill # matches argv containing /root/.unsloth/ (not bare names that would hit a user's own - # llama-server), and the backslash + [h]-bracket in '/root/\.unslot[h]/' keep it from - # matching this command's own argv. + # llama-server); the backslash + [h]-bracket in '/root/\.unslot[h]/' keep it from matching + # this command's own argv. $_clean = '_had=0; if [ -d /root/.unsloth ] || [ -L /root/.local/bin/unsloth ]; then _had=1; fi; 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; if [ $_had -eq 1 ]; then fuser -k 8888/tcp 2>/dev/null; fi; pkill -9 -f ''/root/\.unslot[h]/'' 2>/dev/null; true' # Clean only distros with evidence of a fallback install: the wsl-distro.txt marker or an # explicit UNSLOTH_WSL_DISTRO. The broad candidate probe is only for legacy marker-less - # installs, which exist only on ARM64 -- on x86 it would delete distros this installer - # never touched (e.g. a ROCm-on-WSL Studio under /root). + # installs (ARM64 only); on x86 it would delete distros this installer never touched + # (e.g. a ROCm-on-WSL Studio under /root). $_cands = @() if ($env:UNSLOTH_WSL_DISTRO) { $_cands += $env:UNSLOTH_WSL_DISTRO } if ($_recordedDistro) { $_cands += $_recordedDistro } diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh index ab02c67974..e6ca561743 100755 --- a/scripts/uninstall.sh +++ b/scripts/uninstall.sh @@ -216,8 +216,7 @@ _remove_path "$HOME/.unsloth/studio" # by deleting it). No-op in env/custom mode (they nest under the custom root) and # when absent. A user-set UNSLOTH_LLAMA_CPP_PATH is intentionally kept. _remove_path "$HOME/.unsloth/llama.cpp" -# provision_llama_cuda.sh fetched by the WoA/Spark CUDA-build path (install.ps1 -# background build + direct-WSL setup.sh). No-op when absent. +# provision_llama_cuda.sh fetched by the WoA/Spark CUDA-build path. No-op when absent. _remove_path "$HOME/.unsloth/provision_llama_cuda.sh" _remove_path "$HOME/.unsloth/.cache" # llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging). @@ -278,7 +277,7 @@ case "$_os" in # receive trailing tokens as $args. WSL distro names are safe to # embed (no quotes/$/backtick). # shellcheck disable=SC2016 - # $env:APPDATA/$distro are PowerShell-side; $_wsl_distro is injected from shell. + # $env:APPDATA/$distro are PowerShell-side; $_wsl_distro is shell-injected. powershell.exe -NoProfile -Command '$distro = "'"$_wsl_distro"'"; $dirs = @( [Environment]::GetFolderPath("Desktop"), diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 77628e4a44..13d56e8c42 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -736,15 +736,15 @@ def _nvidia_classify_spark_unified_memory(props: Any) -> tuple[str, bool]: """Classify an NVIDIA device as Spark-class unified-memory or discrete. Returns ``(marker, is_unified)``; marker is ``"is_integrated"`` or the matched - device-name token, else ``""``. Spark-class parts (DGX Spark / GB10, N1X "RTX - Spark") share one memory pool with the OS, so like the ROCm APUs they need a + name token, else ``""``. Spark-class parts (DGX Spark / GB10, N1X "RTX Spark") + share one memory pool with the OS, so like the ROCm APUs they need a ``set_per_process_memory_fraction`` cap -- exhausting the pool can stall the box. ``is_integrated`` is authoritative on native Linux, but WSL2 paravirtualization masks it to 0 and renames the device (N1X reports ``JMJWOA-Generic-GPU``, - verified on hardware) -- hence the name-token fallback. Tokens mirror - ``_DGX_SPARK_DEVICE_TOKENS`` in ``unsloth/models/_utils.py`` (duplicated - because this guard runs before any ML import). + verified live) -- hence the name-token fallback. Tokens mirror + ``_DGX_SPARK_DEVICE_TOKENS`` in ``unsloth/models/_utils.py`` (duplicated since + this guard runs before any ML import). """ if getattr(props, "is_integrated", 0): return "is_integrated", True @@ -752,7 +752,7 @@ def _nvidia_classify_spark_unified_memory(props: Any) -> tuple[str, bool]: import re for token in ("GB10", "GB110", "JMJWOA", "N1X", "DGX SPARK"): - # Whole-token match so "GB10" does not match a discrete "GB100"/"GB10X". + # Whole-token match so "GB10" doesn't match discrete "GB100"/"GB10X". if re.search(r"(? # Discrete NVIDIA GPUs untouched. else: try: - # Set PYTORCH_CUDA_ALLOC_CONF before get_device_properties below inits the - # CUDA allocator -- the later `import unsloth` patch is too late for THIS - # worker process. CUDA-free nvidia-smi sniff (mirrors - # _is_dgx_spark_no_cuda_init), same append-don't-override and - # UNSLOTH_NO_EXPANDABLE_SEGMENTS opt-out as the library patch. + # Set PYTORCH_CUDA_ALLOC_CONF before get_device_properties below inits + # the allocator -- the later `import unsloth` patch is too late for THIS + # worker. CUDA-free nvidia-smi sniff (mirrors _is_dgx_spark_no_cuda_init), + # same append-don't-override and UNSLOTH_NO_EXPANDABLE_SEGMENTS opt-out. try: import platform as _plat diff --git a/studio/backend/tests/test_spark_oom_guard.py b/studio/backend/tests/test_spark_oom_guard.py index 4e61dfd3af..a0e13f0bd6 100644 --- a/studio/backend/tests/test_spark_oom_guard.py +++ b/studio/backend/tests/test_spark_oom_guard.py @@ -1,13 +1,12 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Unit tests for _nvidia_classify_spark_unified_memory (Spark OOM-guard classifier). +"""Tests for _nvidia_classify_spark_unified_memory (Spark OOM-guard classifier). -Two paths: (1) ``is_integrated`` device property (authoritative on native Linux), -(2) device-name token match -- needed because WSL2's GPU paravirtualization masks +Two paths: (1) ``is_integrated`` property (authoritative on native Linux), +(2) name-token match -- needed because WSL2 GPU paravirtualization masks ``is_integrated`` to 0 and renames the device (N1X reports ``JMJWOA-Generic-GPU``; -verified on hardware). Mirrors test_rocm_oom_guard.py for the ROCm/Strix-Halo -classifier the NVIDIA guard was modeled on. +verified live). Mirrors test_rocm_oom_guard.py, which the NVIDIA guard models. """ from __future__ import annotations @@ -20,7 +19,7 @@ from core.training.worker import _nvidia_classify_spark_unified_memory def _props(**kwargs) -> SimpleNamespace: - """Build a fake device-properties object with the given attributes.""" + """Fake device-properties object with the given attributes.""" return SimpleNamespace(**kwargs) diff --git a/studio/scripts/provision_llama_cuda.sh b/studio/scripts/provision_llama_cuda.sh index ef4c63bb50..1f36b8b90f 100644 --- a/studio/scripts/provision_llama_cuda.sh +++ b/studio/scripts/provision_llama_cuda.sh @@ -1,8 +1,8 @@ #!/usr/bin/env bash # Build CUDA llama.cpp for Studio GGUF *inference* into ~/.unsloth/llama.cpp -# (resolver checks /build/bin/llama-server). Idempotent, best-effort, always -# exits 0. Exists because no aarch64+CUDA prebuilt covers NVIDIA ARM hosts -# (DGX Spark / GB10, N1X "RTX" laptops). Platform gotchas handled: +# (resolver checks /build/bin/llama-server). Idempotent, best-effort, exits +# 0. Exists because no aarch64+CUDA prebuilt covers NVIDIA ARM hosts (DGX Spark / +# GB10, N1X "RTX" laptops). Platform gotchas handled: # * nvcc rejects gcc-15 -> force gcc-14 / g++-14 host compiler # * glibc >= 2.41 vs CUDA < 13.3 -> install CUDA 13.3 (rsqrt header clash) # * sm_121 (Blackwell) -> derive arch from the GPU's compute_cap @@ -13,7 +13,7 @@ LLAMA_DIR="${UNSLOTH_LLAMA_CPP_PATH:-$HOME/.unsloth/llama.cpp}" SERVER="$LLAMA_DIR/build/bin/llama-server" log() { printf ' - %s\n' "$*"; } -# CUDA shows up two ways: monolithic (libggml-cuda in ldd) or split (dlopen-ed +# Detect CUDA two ways: monolithic (libggml-cuda in ldd) or split (dlopen-ed # libggml-cuda.so* beside the binary, missed by ldd). CPU-only builds ship no # libggml-cuda.so, so its presence is the reliable signal. is_cuda_server() { @@ -39,12 +39,12 @@ SUDO=""; [ "$(id -u)" -ne 0 ] && SUDO="sudo" HAVE_APT=0; command -v apt-get >/dev/null 2>&1 && HAVE_APT=1 # 2. Base toolchain first, then gcc-14 (nvcc rejects gcc-15) in a SEPARATE apt -# transaction: gcc-14 is absent from default Ubuntu 22.04 / Debian 12 sources, -# which would abort a combined transaction and lose the base build tools too. +# transaction: gcc-14 is absent from default Ubuntu 22.04 / Debian 12 sources, so +# a combined transaction would abort and lose the base build tools too. if [ "$HAVE_APT" -eq 1 ]; then $SUDO apt-get update -y >/dev/null 2>&1 || true - # libcurl4-openssl-dev: -DLLAMA_CURL=ON needs it, and on the WSL deferred path - # setup.sh's GGUF dep install (which covers libcurl) was skipped. + # libcurl4-openssl-dev: -DLLAMA_CURL=ON needs it, and the WSL deferred path + # skips setup.sh's GGUF dep install that would otherwise provide libcurl. $SUDO apt-get install -y --no-install-recommends \ build-essential cmake git curl ca-certificates libcurl4-openssl-dev >/dev/null 2>&1 || true $SUDO apt-get install -y --no-install-recommends gcc-14 g++-14 >/dev/null 2>&1 || true @@ -89,7 +89,7 @@ fi CUDA_HOME="$(dirname "$(dirname "$NVCC")")" # CUDA + Linux dirs FIRST so the build uses Linux cmake/gcc/git, not Windows tools -# leaked in via WSL interop (/mnt/c); original PATH kept so nvidia-smi resolves. +# leaked in via WSL interop (/mnt/c). Keep original PATH so nvidia-smi resolves. export PATH="$CUDA_HOME/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:$PATH" export CUDAToolkit_ROOT="$CUDA_HOME" @@ -158,7 +158,7 @@ _cmake_configure() { -DLLAMA_CURL=ON >/dev/null 2>&1 } # A pre-existing build/ may carry a stale CMake cache (relocated dir: bad absolute -# paths + GGML_CUDA=OFF). Reuse it first (fast incremental); wipe only on failure. +# paths + GGML_CUDA=OFF). Reuse first (fast incremental); wipe only on failure. if ! _cmake_configure; then log "stale/incompatible CMake cache detected; wiping build dir for a clean CUDA configure" rm -rf build @@ -166,10 +166,10 @@ if ! _cmake_configure; then fi # Also builds the targets unsloth-zoo's GGUF exporter needs (llama-mtmd-cli, # llama-gguf-split). Jobs default to ~half the cores (full -j(nproc) CUDA builds -# trip thermal shutdowns on NVIDIA-ARM laptops like the N1X "RTX Spark") and are -# RAM-capped (~1.5 GB/nvcc job). Tune: UNSLOTH_LLAMA_BUILD_JOBS=N; re-runs resume. +# trip thermal shutdowns on NVIDIA-ARM laptops like the N1X "RTX Spark"), capped +# at ~1.5 GB/nvcc job. Tune: UNSLOTH_LLAMA_BUILD_JOBS=N; re-runs resume. _ncpu="$(nproc 2>/dev/null || echo 4)" -# Honor a valid positive-int override; ignore junk/0 (cmake reads -j0 as "all cores"). +# Honor a valid positive-int override; ignore junk/0 (cmake treats -j0 as all cores). if [ -n "${UNSLOTH_LLAMA_BUILD_JOBS:-}" ] && [ "${UNSLOTH_LLAMA_BUILD_JOBS}" -ge 1 ] 2>/dev/null; then JOBS="$UNSLOTH_LLAMA_BUILD_JOBS" else @@ -188,7 +188,7 @@ command -v nice >/dev/null 2>&1 && _NICE="nice -n 19" command -v ionice >/dev/null 2>&1 && _NICE="$_NICE ionice -c 3" _cmake_build() { # Only llama-server is REQUIRED: an old UNSLOTH_LLAMA_TAG pin may predate the - # helper targets, and those missing must not fail the whole provision. + # helper targets, whose absence must not fail the whole provision. $_NICE cmake --build build -j"$JOBS" --target llama-server >/dev/null 2>&1 } _cmake_build_extras() { @@ -200,7 +200,7 @@ _cmake_build_extras() { if ! _cmake_build; then # An interrupted build (thermal/power shutdown, common on this machine class) # can leave a half-linked libggml-cuda.so that breaks the resume link - # (undefined ggml_cuda_op_* refs); wipe and rebuild clean once. + # (undefined ggml_cuda_op_* refs); wipe and rebuild clean. log "build failed (likely interrupted/partial); wiping build dir and rebuilding clean" rm -rf build _cmake_configure || { log "cmake configure failed"; cd /; _restore_prev; exit 0; } diff --git a/studio/setup.sh b/studio/setup.sh index c8d9a630d8..44b790ffce 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -975,7 +975,7 @@ LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server" _NEED_LLAMA_SOURCE_BUILD=false _LLAMA_CPP_DEGRADED=false # Deferred != degraded: on WSL2 aarch64+NVIDIA install.ps1 builds the CUDA server -# in the background, so an absent server is success and must not trip the arm64 +# in the background, so an absent server is success -- must not trip the arm64 # CPU-prebuilt last-resort or exit 1. _LLAMA_CPP_DEFERRED=false _LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}" @@ -1164,7 +1164,7 @@ fi # install.ps1 builds the CUDA llama-server in the background; without nvcc, # section 9 could only make a slow CPU server that build discards. With nvcc we # fall through to section 9; opted out (UNSLOTH_NO_LLAMA_CUDA=1) the CPU build is -# kept as the only server. +# the only server. if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] \ && [ "$_LLAMA_FORCE_COMPILE" != "1" ] \ && [ -z "$_LLAMA_PR" ] \ @@ -1178,7 +1178,7 @@ if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] \ step "llama.cpp" "GGUF engine: CUDA build running in background (WSL aarch64 + NVIDIA)" "$C_WARN" substep "skipping slow CPU build; the background CUDA llama.cpp will provide the server" substep "(opt out / keep CPU build with UNSLOTH_NO_LLAMA_CUDA=1)" - # DEFERRED, not DEGRADED: DEGRADED would trigger the CPU-prebuilt last resort + exit 1. + # DEFERRED, not DEGRADED: DEGRADED triggers the CPU-prebuilt last resort + exit 1. _NEED_LLAMA_SOURCE_BUILD=false _LLAMA_CPP_DEFERRED=true fi @@ -1456,7 +1456,7 @@ else # glibc >= 2.41 + CUDA < 13.3: rsqrt/rsqrtf header clash fails # every .cu -> CPU fallback; only fix is CUDA >= 13.3. Diagnostic - # only (never changes flags or aborts). Checks the final _NVCC_VER + # only (never changes flags or aborts), against the final _NVCC_VER # (after the driver-compat swap above). _GLIBC_VER="$(getconf GNU_LIBC_VERSION 2>/dev/null | awk '{print $2}')" || _GLIBC_VER="" if [ -n "$_GLIBC_VER" ]; then @@ -1686,8 +1686,8 @@ fi # end _SKIP_GGUF_BUILD check # (provision_llama_cuda.sh) for native Linux. Best-effort: provision always exits # 0, and on failure the prior CPU/degraded state stands. # CUDA detection covers both layouts: monolithic (libggml-cuda in ldd) and split -# (dlopen-ed libggml-cuda.so* beside the binary, missed by ldd). CPU-only builds -# ship no libggml-cuda.so, so its presence is the signal. +# (dlopen-ed libggml-cuda.so* beside the binary, missed by ldd). Its presence is +# the signal -- CPU-only builds ship no libggml-cuda.so. _have_cuda_llama_server() { [ -x "$LLAMA_SERVER_BIN" ] || return 1 ldd "$LLAMA_SERVER_BIN" 2>/dev/null | grep -qi 'libggml-cuda' && return 0 @@ -1702,8 +1702,8 @@ if [ "$_HOST_SYSTEM" = "Linux" ] \ && nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}' \ && ! _have_cuda_llama_server; then # Under WSL this runs ONLY for a DIRECT `install.sh` run: install.ps1 sets - # UNSLOTH_WSL_LLAMA_DEFERRED=1 and builds in the background, but a direct run - # has no background builder, so provision here. + # UNSLOTH_WSL_LLAMA_DEFERRED=1 and builds in the background; a direct run has + # no background builder, so provision here. # Resolve provision_llama_cuda.sh: beside setup.sh, then local-dev repo, else # fetch from GitHub (so `curl | sh` works on an older wheel without it). _PROV_SH="" @@ -1721,13 +1721,13 @@ if [ "$_HOST_SYSTEM" = "Linux" ] \ if [ -n "$_PROV_SH" ]; then step "llama.cpp" "aarch64 + NVIDIA: provisioning CUDA toolkit + building CUDA llama.cpp for GGUF inference..." "$C_WARN" substep "(opt out with UNSLOTH_NO_LLAMA_CUDA=1; lower load with UNSLOTH_LLAMA_BUILD_JOBS=N)" - # UNSLOTH_LLAMA_CPP_PATH routes a custom STUDIO_HOME into $LLAMA_CPP_DIR; always exits 0. + # UNSLOTH_LLAMA_CPP_PATH routes a custom STUDIO_HOME into $LLAMA_CPP_DIR. UNSLOTH_LLAMA_CPP_PATH="$LLAMA_CPP_DIR" bash "$_PROV_SH" || true if _have_cuda_llama_server; then step "llama.cpp" "CUDA llama-server ready (aarch64 + NVIDIA)" _LLAMA_CPP_DEGRADED=false - # Claim ownership of the fresh $LLAMA_CPP_DIR or the next custom-STUDIO_HOME - # run's _assert_studio_owned_or_absent would abort on it. + # Claim ownership of the fresh $LLAMA_CPP_DIR, else the next custom-STUDIO_HOME + # run's _assert_studio_owned_or_absent aborts on it. if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then : > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true fi diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py index 71428b27b1..294a54e21d 100644 --- a/unsloth/kernels/flex_attention.py +++ b/unsloth/kernels/flex_attention.py @@ -27,9 +27,9 @@ torch_compile_options = { def _flex_is_dgx_spark(): - # CUDA-free copy of _utils._is_dgx_spark_no_cuda_init() (avoids a circular import). - # Runs at module import, before ._utils -- touching torch.cuda here would init the - # allocator before patch_dgx_spark_memory_config() can set PYTORCH_CUDA_ALLOC_CONF. + # CUDA-free copy of _utils._is_dgx_spark_no_cuda_init() (avoids circular import). + # Runs at module import, before ._utils -- touching torch.cuda here would init + # the allocator before patch_dgx_spark_memory_config() sets PYTORCH_CUDA_ALLOC_CONF. _force = os.environ.get("UNSLOTH_FORCE_DGX_SPARK") if _force == "1": return True @@ -49,7 +49,7 @@ def _flex_is_dgx_spark(): timeout = 5, ) names = (out.stdout or "").upper() - # Whole-token match so "GB10" does not match a discrete "GB100"/"GB10X". + # Whole-token match so "GB10" doesn't match discrete "GB100"/"GB10X". import re return any( @@ -60,7 +60,7 @@ def _flex_is_dgx_spark(): return False -# Spark's 48 SMs are under inductor's 68-SM is_big_gpu bar; max_autotune would only waste search time. +# Spark's 48 SMs are under inductor's 68-SM is_big_gpu bar; max_autotune just wastes search time. if _flex_is_dgx_spark(): torch_compile_options["max_autotune"] = False diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 8df76767a8..af96d1a85b 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -992,15 +992,14 @@ except: from transformers.modeling_utils import logger as transformers_logger -# ---- NVIDIA DGX Spark (GB10) / N1X "RTX Spark" unified-memory support ---- -# Device names vary ("NVIDIA GB10", "JMJWOA-Generic-GPU" on N1X); the -# aarch64 + CUDA gate keeps every Spark workaround a no-op elsewhere. +# NVIDIA DGX Spark (GB10) / N1X "RTX Spark" unified-memory support. +# Names vary ("NVIDIA GB10", "JMJWOA-Generic-GPU" on N1X); the aarch64 + CUDA +# gate keeps every Spark workaround a no-op elsewhere. _DGX_SPARK_DEVICE_TOKENS = ("GB10", "JMJWOA", "N1X", "DGX SPARK", "GB110") def _name_has_spark_token(names_upper): - # Whole-token match so "GB10" does NOT match "GB100"/"GB10X" -- a discrete - # Grace+Blackwell datacenter GPU must not be misread as a unified-memory Spark. + # Whole-token match so "GB10" doesn't match discrete "GB100"/"GB10X". import re return any( re.search(r"(? 1. Out-of-range = no cap. + # Out-of-range = no cap (0 OOMs everything; torch rejects > 1). _frac_val = float(_frac) if 0.0 < _frac_val <= 1.0: torch.cuda.set_per_process_memory_fraction(_frac_val) @@ -1133,10 +1132,10 @@ def patch_dgx_spark_runtime_defaults(): def patch_dgx_spark_dataloader_defaults(): - """Default `dataloader_pin_memory` to False on Spark UMA (accuracy-neutral). + """Default `dataloader_pin_memory` to False on Spark UMA. On one shared pool, pinning only reserves non-pageable RAM and adds a staging - copy (mirrors transformers' own use_cpu precedent). Wrapping the base + copy (mirrors transformers' use_cpu precedent). Wrapping the base `TrainingArguments.__post_init__` covers SFT + every TRL trainer in one idempotent patch. Opt out: UNSLOTH_SPARK_KEEP_PIN_MEMORY=1. No-op off-Spark. """ @@ -1153,7 +1152,7 @@ def patch_dgx_spark_dataloader_defaults(): return _orig_post_init = Base.__post_init__ - # *args/**kwargs: tolerate future InitVar parameters in __post_init__. + # *args/**kwargs: tolerate future InitVar params in __post_init__. def __post_init__(self, *args, **kwargs): _orig_post_init(self, *args, **kwargs) if getattr(self, "dataloader_pin_memory", None) is True: @@ -1731,7 +1730,7 @@ torch_compile_options = { "trace.enabled": UNSLOTH_COMPILE_DEBUG, "triton.cudagraphs": False, } -# Spark's 48 SMs are under inductor's 68-SM is_big_gpu bar; max_autotune would only waste search time. +# Spark's 48 SMs are under inductor's 68-SM is_big_gpu bar; max_autotune just wastes search time. if is_dgx_spark(): torch_compile_options["max_autotune"] = False