Tighten and trim code comments
This commit is contained in:
Daniel Han
parent
d888184da4
commit
b50eb8bc71
10 changed files with 125 additions and 133 deletions
86
install.ps1
86
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 <str>`.
|
||||
# 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 <str>`.
|
||||
$_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"
|
||||
|
|
|
|||
16
install.sh
16
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 \
|
||||
|
|
|
|||
|
|
@ -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 - <distro>).lnk") left by pre-release dev builds.
|
||||
# Canonical name is "Unsloth Studio.lnk"; also sweep legacy distro-suffixed names
|
||||
# ("Unsloth Studio (WSL - <distro>).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 }
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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"(?<![A-Z0-9])" + re.escape(token) + r"(?![A-Z0-9])", name_upper):
|
||||
return token, True
|
||||
return "", False
|
||||
|
|
@ -2441,11 +2441,10 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
|
|||
# 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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
#!/usr/bin/env bash
|
||||
# Build CUDA llama.cpp for Studio GGUF *inference* into ~/.unsloth/llama.cpp
|
||||
# (resolver checks <dir>/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 <dir>/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; }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"(?<![A-Z0-9])" + re.escape(tok) + r"(?![A-Z0-9])", names_upper)
|
||||
|
|
@ -1010,8 +1009,8 @@ def _name_has_spark_token(names_upper):
|
|||
|
||||
@functools.lru_cache(maxsize = None)
|
||||
def is_dgx_spark():
|
||||
"""True only on DGX Spark / N1X Spark-class machines (gate: aarch64 + NVIDIA
|
||||
CUDA + known device-name token). UNSLOTH_FORCE_DGX_SPARK=1/0 forces on/off."""
|
||||
"""True only on DGX Spark / N1X (gate: aarch64 + CUDA + device-name token).
|
||||
UNSLOTH_FORCE_DGX_SPARK=1/0 forces on/off."""
|
||||
_force = os.environ.get("UNSLOTH_FORCE_DGX_SPARK")
|
||||
if _force == "1":
|
||||
return True
|
||||
|
|
@ -1034,8 +1033,8 @@ def is_dgx_spark():
|
|||
|
||||
@functools.lru_cache(maxsize = None)
|
||||
def _is_dgx_spark_no_cuda_init():
|
||||
"""Spark detection that never inits CUDA: reads device names via `nvidia-smi`,
|
||||
not torch, so PYTORCH_CUDA_ALLOC_CONF can still be set afterwards. Same
|
||||
"""Spark detection that never inits CUDA: reads names via `nvidia-smi`, not
|
||||
torch, so PYTORCH_CUDA_ALLOC_CONF can still be set afterwards. Same
|
||||
UNSLOTH_FORCE_DGX_SPARK override; False on any error."""
|
||||
_force = os.environ.get("UNSLOTH_FORCE_DGX_SPARK")
|
||||
if _force == "1":
|
||||
|
|
@ -1062,12 +1061,12 @@ def _is_dgx_spark_no_cuda_init():
|
|||
|
||||
|
||||
def patch_dgx_spark_caching_allocator_warmup():
|
||||
"""No-op `transformers.modeling_utils.caching_allocator_warmup` on Spark UMA.
|
||||
"""No-op `caching_allocator_warmup` on Spark UMA.
|
||||
|
||||
`cudaMemGetInfo()` undercounts free UMA memory, so HF's warmup
|
||||
`torch.empty(...)` raises `AcceleratorError: invalid argument` and aborts
|
||||
`torch.empty(...)` raises `AcceleratorError: invalid argument`, aborting
|
||||
bitsandbytes 4/8-bit loads. The warmup is only a speed hint, so skip it.
|
||||
Gated by `is_dgx_spark()`; idempotent (`_unsloth_spark_noop` marker).
|
||||
Gated by `is_dgx_spark()`; idempotent via `_unsloth_spark_noop` marker.
|
||||
"""
|
||||
if not is_dgx_spark():
|
||||
return
|
||||
|
|
@ -1089,7 +1088,7 @@ def patch_dgx_spark_caching_allocator_warmup():
|
|||
|
||||
def patch_dgx_spark_memory_config():
|
||||
"""Enable allocator `expandable_segments` on Spark UMA to cut fragmentation
|
||||
OOMs (accuracy-neutral; strict no-op off-Spark).
|
||||
OOMs (no-op off-Spark).
|
||||
|
||||
Appends to PYTORCH_CUDA_ALLOC_CONF only when absent; opt out with
|
||||
UNSLOTH_NO_EXPANDABLE_SEGMENTS=1. Must run before the first CUDA allocation,
|
||||
|
|
@ -1124,7 +1123,7 @@ def patch_dgx_spark_runtime_defaults():
|
|||
_frac = os.environ.get("UNSLOTH_SPARK_MEM_FRACTION")
|
||||
if _frac:
|
||||
try:
|
||||
# 0 would OOM every allocation; torch rejects > 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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue