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

Compress PR comments to essentials (comment-only; AST/token-verified)

Browse source 
Comment-compression sweep over comments this PR added, mirroring the
sweep already done on main. No non-comment token changed: .py verified
by AST equality (docstrings normalized), .sh by non-comment-line
equality + bash -n, .ps1 by token-stream equality minus comments.
test_spark_oom_guard.py: 13 passed before and after.

Files touched:
- install.ps1
- install.sh
- scripts/uninstall.ps1
- studio/backend/core/training/worker.py
- studio/scripts/provision_llama_cuda.sh
- studio/setup.sh
- unsloth/kernels/flex_attention.py
- unsloth/models/_utils.py

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Daniel Han 2026-06-10 00:32:41 -07:00
commit 42e69031b9
8 changed files with 222 additions and 366 deletions
Download patch file Download diff file Expand all files Collapse all files

180
install.ps1
View file

@ -48,9 +48,8 @@ function Install-UnslothStudio {
}
}
# Git ref for fetching repo-versioned install assets (provision_llama_cuda.sh,
# the .ico) from raw.githubusercontent.com. Defaults to 'main' (unchanged for
# existing users); set UNSLOTH_INSTALL_REF to a branch to test pre-merge.
# Ref for fetching install assets (provision_llama_cuda.sh, the .ico) from
# raw.githubusercontent.com; 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'
@ -98,11 +97,9 @@ function Install-UnslothStudio {
if ($TauriMode) {
exit $Code
}
# File-based runs (powershell -File / .\install.ps1) exit 0 on a plain return no
# matter what $LASTEXITCODE says, so automation would treat a fatal failure as a
# completed install -- `exit` carries the code there. Under `irm | iex` there is
# no $PSCommandPath and `exit` would kill the user's shell, so fall through and
# let the caller return (the message is the signal).
# -File runs exit 0 on a plain return regardless of $LASTEXITCODE, so `exit`
# must carry the code there; under `irm | iex` (no $PSCommandPath) `exit`
# would kill the user's shell, so fall through.
if ($PSCommandPath) {
exit $Code
}
@ -1488,16 +1485,13 @@ 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 the GPU stack can't run natively. On ARM64 with an
# NVIDIA GPU and no installable native CUDA torch, route GPU setup through WSL2: enable/install WSL2,
# run the Linux installer there (full GPU), and add a Windows `unsloth` shim that forwards into WSL.
# Strictly gated: x86_64 and ARM64-without-NVIDIA are unaffected. Future-proof: if a win_arm64 CUDA
# torch wheel ships, the probe below passes and native install is kept automatically.
# Opt out with UNSLOTH_NO_WSL_FALLBACK=1; choose the distro with UNSLOTH_WSL_DISTRO.
# win_arm64 has no CUDA PyTorch/Triton wheel, so run the Linux installer inside WSL2 (full GPU) and
# add a Windows `unsloth` shim that forwards into it. x86_64 / ARM64-without-NVIDIA unaffected; if a
# win_arm64 CUDA torch wheel ever ships, the probe below keeps the native install automatically.
# Opt out: UNSLOTH_NO_WSL_FALLBACK=1; choose the distro with UNSLOTH_WSL_DISTRO.
try { $_winArm64 = ([System.Runtime.InteropServices.RuntimeInformation]::OSArchitecture.ToString() -ieq 'Arm64') } catch { $_winArm64 = $false }
# Under x64-emulated PowerShell on ARM, .NET OSArchitecture and $env:PROCESSOR_ARCHITECTURE report
# X64/AMD64; Win32_Processor.Architecture (12=ARM64) and machine-level PROCESSOR_ARCHITECTURE read
# the true arch. Additive: only turns $_winArm64 ON for genuine ARM64 hosts.
# x64-emulated PS on ARM reports X64/AMD64 via .NET and $env:; Win32_Processor.Architecture
# (12=ARM64) and machine-level PROCESSOR_ARCHITECTURE read the true arch. Only ever turns $_winArm64 ON.
if (-not $_winArm64) {
try { if ((@(Get-CimInstance Win32_Processor -ErrorAction Stop))[0].Architecture -eq 12) { $_winArm64 = $true } } catch {}
}
@ -1509,13 +1503,11 @@ shell.Run cmd, 0, False
}
$_nativeCudaTorchOk = $false
if ($_winArm64 -and $HasNvidiaSmi -and (-not $SkipTorch)) {
# Can a native CUDA torch wheel be resolved for this platform/index? Must use the SAME spec
# 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.
# Probe with the SAME spec as the real install ("torch>=2.4,<2.11.0"): a bare `torch` probe
# can match an out-of-range wheel, skipping WSL only to fail 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.
# --reinstall: an already-installed (e.g. CPU-only) torch must not satisfy the probe --
# it has to 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
@ -1527,24 +1519,21 @@ 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.
# The Tauri desktop app launches its backend from a Windows venv (resolve_backend_binary),
# not WSL, so a WSL-only install would report complete yet fail to start -- 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)
}
$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).
# Reset first: a stale 0 in $LASTEXITCODE would wrongly mark WSL ready if wsl.exe fails to start.
$global:LASTEXITCODE = -1
try { & wsl.exe --status *> $null; if ($LASTEXITCODE -eq 0) { $wslReady = $true } } catch {}
}
if (-not $wslReady) {
# Enabling WSL2 is a one-time operation that requires admin + a reboot.
$isAdmin = $false
try { $isAdmin = ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltinRole]::Administrator) } catch {}
step "wsl" "WSL2 isn't enabled yet -- one-time setup (needs admin + reboot)" "Yellow"
@ -1556,11 +1545,9 @@ 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. A plain return exits 0 for `-File` runs no
# matter what $LASTEXITCODE says, so exit explicitly there; under `irm | iex`
# ($PSCommandPath empty) exit would kill the user's shell, so return instead.
# Deferred until reboot: restore any rolled-aside previous venv and signal not-complete.
# A plain return exits 0 for -File runs regardless of $LASTEXITCODE, so `exit 1` there;
# under `irm | iex` ($PSCommandPath empty) exit would kill the user's shell, so return.
Restore-StudioVenvRollback
$global:LASTEXITCODE = 1
if ($PSCommandPath) { exit 1 }
@ -1568,11 +1555,9 @@ shell.Run cmd, 0, False
}
$distro = if ($env:UNSLOTH_WSL_DISTRO) { $env:UNSLOTH_WSL_DISTRO } else { "Ubuntu-24.04" }
# For cmd-context uses of the name (the generated .cmd shim, copy-paste hints):
# wsl.exe parses its raw command line itself, and a QUOTED space-free name
# ('wsl -d "Ubuntu-24.04"') fails with WSL_E_DISTRO_NOT_FOUND (verified live on
# 2.x) -- while a bare spaced name would split after -d. So quote ONLY when the
# name contains whitespace.
# For cmd-context uses (.cmd shim, copy-paste hints): wsl.exe rejects a QUOTED space-free
# name (WSL_E_DISTRO_NOT_FOUND, verified on 2.x) yet splits a bare spaced one after -d --
# so quote ONLY when the name contains whitespace.
$_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
@ -1582,12 +1567,10 @@ shell.Run cmd, 0, False
substep "installing WSL distro '$distro' (first time only)..." "Cyan"
try { & wsl.exe --install -d $distro --no-launch } catch {}
} else {
# A PRE-EXISTING distro may be WSL1, which has no GPU passthrough: the existence
# probe passes but the full install would only fail at the final torch.cuda
# check. Detect WSL1 up-front from inside the distro (kernel string + libcuda --
# encoding-proof, unlike parsing UTF-16 `wsl -l -v` output) and convert in place;
# `wsl --set-version` preserves the distro's files. Freshly installed distros
# are WSL2 (default version 2), so only the pre-existing case needs this.
# A PRE-EXISTING distro may be WSL1 (no GPU passthrough; would only fail at the final
# torch.cuda check). Detect from inside the distro (encoding-proof, unlike UTF-16
# `wsl -l -v`) and convert in place -- `wsl --set-version` preserves the files.
# Fresh installs default to WSL2, so only the pre-existing case needs this.
$_wsl2Probe = 'grep -qiE ''microsoft-standard|WSL2'' /proc/version 2>/dev/null || test -e /usr/lib/wsl/lib/libcuda.so'
$_isWsl2 = $false
$global:LASTEXITCODE = -1
@ -1606,20 +1589,14 @@ shell.Run cmd, 0, False
}
}
substep "installing Unsloth Studio inside WSL '$distro' with full GPU (this downloads PyTorch)..." "Cyan"
# For a non-main ref, fetch + export THAT ref so the WSL venv gets the branch's
# 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.
# 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 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.
# Forward the CUDA llama.cpp opt-out into WSL: without it the inner setup.sh would
# defer its build to a background builder this script then never starts (the same
# opt-out skips the dispatch below), leaving no llama-server and a misleading
# "building in background" footer. Forwarded, setup.sh keeps its own build instead.
# UNSLOTH_WSL_LLAMA_DEFERRED=1: setup.sh skips its foreground CUDA llama.cpp build because
# install.ps1 builds 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: the same opt-out 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; ' }
if ($_instRef -eq 'main') {
@ -1649,10 +1626,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 Studio's web-server deps: if install.sh's late "studio deps" step was cut short,
# torch + unsloth land but fastapi/uvicorn/structlog/starlette are missing and `unsloth studio`
# dies with ModuleNotFoundError. Reinstall those without pinning 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 without
# pinning huggingface-hub/transformers/datasets so the verified GPU torch stack stays intact.
if ($torchOk) {
$_studioPy = "/root/.unsloth/studio/unsloth_studio/bin/python"
$_serverOk = $false
@ -1663,10 +1639,9 @@ shell.Run cmd, 0, False
} catch {} finally { $ErrorActionPreference = $prevEapS }
if (-not $_serverOk) {
substep "Studio web-server deps incomplete (install.sh step cut short) -- installing them now..." "Cyan"
# Mirrors studio.txt minus the huggingface-hub pin (protected above); uv preferred,
# pip fallback. Bare names only -- a version spec's quotes get mangled through
# PowerShell -> wsl.exe -> bash -lc and `>=` becomes a redirection. uv resolves the
# latest of each, which satisfies the studio.txt minimums anyway.
# studio.txt minus the huggingface-hub pin; uv preferred, pip fallback. Bare names
# only: `>=` becomes 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'
$_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"
@ -1698,16 +1673,14 @@ shell.Run cmd, 0, False
New-Item -ItemType Directory -Force -Path $shimDir *> $null
$shimLines = @(
'@echo off',
# $_distroArg: quoted only if the name has spaces -- wsl.exe rejects a
# quoted space-free name (WSL_E_DISTRO_NOT_FOUND) but splits a bare spaced one.
# $_distroArg: pre-quoted only when spaced (wsl.exe quoting rule above).
"wsl.exe -d $_distroArg -u root -- /root/.unsloth/studio/unsloth_studio/bin/unsloth %*"
)
Set-Content -LiteralPath (Join-Path $shimDir "unsloth.cmd") -Value $shimLines -Encoding ASCII
# Record the distro for the uninstaller: a custom UNSLOTH_WSL_DISTRO install
# must be cleanable without the env var being set again at uninstall time.
# Record the distro so the uninstaller can clean a custom UNSLOTH_WSL_DISTRO
# install without the env var being set again.
try { Set-Content -LiteralPath (Join-Path (Split-Path $shimDir -Parent) "wsl-distro.txt") -Value $distro -Encoding ASCII } catch {}
# 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.
# A fresh profile may have no HKCU 'Path' at all; null would make TrimEnd() throw.
$userPath = [Environment]::GetEnvironmentVariable("Path", "User")
if (-not $userPath) { $userPath = "" }
if (($userPath -split ';') -notcontains $shimDir) {
@ -1735,11 +1708,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 Windows-on-ARM the shell's sandboxed
# icon-extraction broker can't read a standalone .ico under AppData\Local (it gets a
# redirected/virtualized view), so the shortcut renders BLANK -- while the identical
# file under the user profile renders fine (verified on N1X). Keep the shim/launcher
# in $appDir; only the icon needs the profile location.
# Icon must live OUTSIDE %LOCALAPPDATA%: on WoA the shell's sandboxed icon broker
# can't read a .ico under AppData\Local, so the shortcut renders BLANK -- the same
# file under the user profile 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"
@ -1774,44 +1745,37 @@ shell.Run cmd, 0, False
$sc.Save()
}
step "shortcuts" "created Desktop + Start Menu shortcuts (launch WSL Studio + open browser)" "Green"
# Nudge Explorer to pick up the new/changed shortcuts now: clear+rebuild the icon
# cache, then per-.lnk SHCNE_UPDATEITEM + a global SHCNE_ASSOCCHANGED. (The real
# blank-icon cause on WoA was the AppData\Local icon path, fixed above.)
# Nudge Explorer: clear+rebuild icon cache, per-.lnk SHCNE_UPDATEITEM, global
# SHCNE_ASSOCCHANGED. (The real WoA blank-icon cause was the icon path, fixed above.)
try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {}
try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {}
try {
if (-not ("UnslothShell.Notify" -as [type])) {
Add-Type -Namespace UnslothShell -Name Notify -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int eventId, uint flags, string item1, System.IntPtr item2);'
}
# Per-.lnk SHCNE_UPDATEITEM (0x00002000), SHCNF_PATHW (0x0005): force Explorer to
# re-read each shortcut's icon now (the global notify alone often misses existing .lnks).
# SHCNE_UPDATEITEM (0x00002000), SHCNF_PATHW (0x0005): the global notify alone often misses existing .lnks.
foreach ($lnk in $lnks) { try { [UnslothShell.Notify]::SHChangeNotify(0x00002000, 0x0005, $lnk, [System.IntPtr]::Zero) } catch {} }
# SHCNE_ASSOCCHANGED (0x08000000), SHCNF_IDLIST (0): flush global icon associations
# (item args unused for this event).
# SHCNE_ASSOCCHANGED (0x08000000), SHCNF_IDLIST (0): flush global icon associations.
[UnslothShell.Notify]::SHChangeNotify(0x08000000, 0, $null, [System.IntPtr]::Zero)
} catch {}
} catch {
substep "(could not create shortcuts: $($_.Exception.Message))" "Yellow"
}
# GGUF *inference* needs a CUDA-linked llama-server and no aarch64+CUDA prebuilt exists, so
# build one into ~/.unsloth/llama.cpp in the BACKGROUND: Studio + training are usable now and
# GGUF inference lights up minutes later. Best-effort; opt out with UNSLOTH_NO_LLAMA_CUDA=1.
# GGUF *inference* needs a CUDA llama-server (no aarch64+CUDA prebuilt exists), so build one
# into ~/.unsloth/llama.cpp in the BACKGROUND. Best-effort; opt out: UNSLOTH_NO_LLAMA_CUDA=1.
if ($env:UNSLOTH_NO_LLAMA_CUDA -ne '1') {
$prevEapL = $ErrorActionPreference; $ErrorActionPreference = "Continue"
try {
$_llamaUrl = "https://raw.githubusercontent.com/unslothai/unsloth/$(Get-UnslothInstallRef)/studio/scripts/provision_llama_cuda.sh"
# Step 1: fetch the provision script + write a small runner, shipped as base64 to
# dodge quoting layers. The runner (a) restores PATH so a non-login shell finds
# nvidia-smi (/usr/lib/wsl/lib) and apt -- else provision early-exits "no nvidia-smi";
# (b) caps build jobs from UNSLOTH_LLAMA_BUILD_JOBS (Windows env vars don't cross into
# WSL); (c) runs provision with logging. A runner FILE lets the detached launcher below
# pass only space-free args, avoiding Start-Process mis-splitting `bash -lc <str>`.
# 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 ->
# 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 llama.cpp pins into WSL: the provisioner honors UNSLOTH_LLAMA_TAG /
# UNSLOTH_LLAMA_PR, but Windows env vars don't cross into WSL on their own --
# without these exports a user's pin would be silently ignored by the
# deferred background build. sh-single-quoted (tags/PRs are simple tokens).
# Bridge UNSLOTH_LLAMA_TAG / UNSLOTH_LLAMA_PR pins into WSL too, else the deferred
# build silently ignores them. sh-single-quoted (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"
@ -1819,13 +1783,10 @@ shell.Run cmd, 0, False
$_fetchCmd = 'mkdir -p /root/.unsloth; if curl -fsSL "' + $_llamaUrl + '" -o /root/.unsloth/provision_llama_cuda.sh && [ -s /root/.unsloth/provision_llama_cuda.sh ]; then chmod +x /root/.unsloth/provision_llama_cuda.sh; echo ' + $_runnerB64 + ' | base64 -d > /root/.unsloth/run_llama_build.sh; chmod +x /root/.unsloth/run_llama_build.sh; echo PROV_FETCHED; else echo PROV_NOSCRIPT; fi'
$_fetchOut = & wsl.exe -d $distro --cd /root -u root -- bash -lc $_fetchCmd 2>$null
if ("$_fetchOut" -match 'PROV_FETCHED') {
# Step 2: anchor the build to a detached Windows process. A WSL-side `nohup &`
# doesn't survive -- WSL stops the VM when the launching session exits, killing
# the build. A persistent Windows-side wsl.exe (Start-Process, no -Wait) keeps the
# VM up for the whole build while install.ps1 returns. PS 5.1's Start-Process
# joins -ArgumentList with spaces WITHOUT quoting, so a spaced distro name would
# split after -d -- pass $_distroArg (pre-quoted only when spaced; wsl.exe
# rejects a quoted space-free name). All other tokens are space-free.
# Step 2: a detached Windows-side wsl.exe keeps the WSL VM up for the whole build
# (a WSL-side `nohup &` dies: WSL stops the VM when the launching session exits).
# PS 5.1 Start-Process joins -ArgumentList WITHOUT quoting, so pass $_distroArg
# (pre-quoted only when spaced); all other tokens are space-free.
Start-Process -WindowStyle Hidden -FilePath 'wsl.exe' -ArgumentList @('-d', $_distroArg, '--cd', '/root', '-u', 'root', '--', 'bash', '/root/.unsloth/run_llama_build.sh') | Out-Null
step "llama.cpp" "building CUDA llama.cpp for GGUF inference in the background (a few min); log: ~/.unsloth/llama_cuda_build.log" "Green"
} else {
@ -1838,16 +1799,15 @@ shell.Run cmd, 0, False
substep "retry, or launch manually: wsl -d $_distroArg -u root -- bash -lic 'unsloth studio -p 8888'" "Cyan"
}
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.
# Success: the Windows venv is vestigial here (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 (plain return exits 0 for -File; iex must not exit).
# Failed (torch.cuda unavailable): restore any rolled-aside previous venv and report
# non-zero (plain return exits 0 for -File; under iex `exit` would kill the caller's shell).
Restore-StudioVenvRollback
$global:LASTEXITCODE = 1
if ($PSCommandPath) { exit 1 }

22
install.sh
View file

@ -2360,11 +2360,10 @@ 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 branch testing: install unsloth from a git ref so its bundled
# setup.sh + Python patches are exercised (not yet on PyPI). install.ps1 sets
# UNSLOTH_INSTALL_REF; gated to the "unsloth" package and a non-"main" ref.
# unsloth-zoo is an optional extra (not a base dep) and SKIP_STUDIO_BASE skips
# the studio base.txt step, so name it explicitly or it never gets installed.
# Pre-merge testing: install unsloth from a git ref (install.ps1 sets
# UNSLOTH_INSTALL_REF) so the branch's setup.sh + patches run. unsloth-zoo
# is not a base dep and SKIP_STUDIO_BASE skips studio base.txt, so name it
# explicitly or it never gets installed.
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 \
@ -2373,14 +2372,11 @@ 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, native or WSL): the base unsloth
# package does not depend on bitsandbytes and the cuXXX extras that normally
# add it are x86_64-oriented, so 4-bit QLoRA fails with ModuleNotFoundError
# out of the box. bitsandbytes ships working aarch64 manylinux wheels
# (verified on sm_121 Blackwell via PTX JIT), so add it best-effort -- a
# platform without a wheel just keeps 16-bit LoRA / full finetuning.
# Gated on SKIP_TORCH: a --no-torch/UNSLOTH_NO_TORCH (GGUF-only) install must
# not have bitsandbytes drag torch back into the venv via its dependencies.
# aarch64 + NVIDIA (DGX Spark / GB10 / N1X): base unsloth lacks bitsandbytes
# (the 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.
if [ "$SKIP_TORCH" = false ] \
&& { [ "$(uname -m)" = "aarch64" ] || [ "$(uname -m)" = "arm64" ]; } \
&& command -v nvidia-smi >/dev/null 2>&1 \

48
scripts/uninstall.ps1
View file

@ -344,9 +344,8 @@ function Uninstall-UnslothStudio {
# %LOCALAPPDATA%\Unsloth (not "Unsloth Studio") with a PATH entry -- all missed by the cleanup above.
_Step "Removing WSL-fallback artifacts (shim, launcher, PATH entry, WSL install)..."
$unslothDir = if ($env:LOCALAPPDATA) { Join-Path $env:LOCALAPPDATA "Unsloth" } else { $null }
# The installer records its WSL distro in wsl-distro.txt so a custom
# UNSLOTH_WSL_DISTRO install is cleanable without the env var being set again
# at uninstall time. Read it BEFORE the directory is removed below.
# wsl-distro.txt records a custom UNSLOTH_WSL_DISTRO install so it is cleanable
# without the env var set; read it BEFORE the directory is removed below.
$_recordedDistro = $null
if ($unslothDir) {
try {
@ -378,36 +377,24 @@ function Uninstall-UnslothStudio {
} catch { }
_RemovePath $unslothDir
}
# The WoA shortcut icon lives under the user profile (the shell icon broker can't read a .ico
# under AppData\Local), so remove it here too.
# The WoA shortcut icon lives under the user profile (icon broker can't read AppData\Local).
if ($env:USERPROFILE) { _RemovePath (Join-Path $env:USERPROFILE ".unsloth\unsloth.ico") }
# Remove the Studio install inside each WSL distro (the real GPU install + any CUDA llama.cpp build).
if (Get-Command wsl.exe -ErrorAction SilentlyContinue) {
try {
# `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
# /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.
# The port-8888 kill is gated on an Unsloth install actually existing in the
# distro (checked BEFORE the rm deletes the marker): a probed distro with an
# unrelated listener on 8888 (Jupyter etc.) must not lose it. The process kill is
# scoped to argv referencing /root/.unsloth/ -- the fallback's install dir, which
# its Studio server, llama-server, and build runner all reference -- instead of
# bare name patterns that would also kill a user's own unrelated llama-server or
# a /home Studio in that distro. The backslash in '/root/\.unslot[h]/' keeps the
# pattern from matching this command's own argv (whose literal text contains the
# escaped form, not the resolved path) -- same idea as the [x]-bracket trick.
# `wsl --list` emits UTF-16 PS mis-parses, so probe candidates by exit code instead
# ('' = default distro). rm runs FIRST (the kills could SIGKILL this shell) and also
# drops the dangling /root/.local/bin/unsloth symlink. Scope STRICTLY to /root (where
# the fallback installs): /home/*/.unsloth may be an unrelated user's. The port-8888
# kill is gated on an Unsloth install existing (checked BEFORE rm deletes the marker)
# so an unrelated 8888 listener survives; pkill matches argv containing /root/.unsloth/
# rather than bare names that would kill a user's own llama-server, and 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'
# Scope the in-distro cleanup to evidence the WoA fallback actually
# installed there: the recorded wsl-distro.txt marker (written by
# install.ps1) or an explicit UNSLOTH_WSL_DISTRO. Only legacy
# marker-less fallback installs need the broad candidate probe, and
# those can only exist on ARM64 hosts -- on x86 machines the probe
# would reach into distros this installer never touched (e.g. an
# AMD ROCm-on-WSL Studio under /root) and delete them.
# 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 hosts -- 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 }
@ -440,9 +427,8 @@ function Uninstall-UnslothStudio {
Write-Host " `$env:UNSLOTH_STUDIO_HOME = 'C:\your\path'; irm https://raw.githubusercontent.com/unslothai/unsloth/main/scripts/uninstall.ps1 | iex"
}
# The distro-probe loop leaves $LASTEXITCODE from its last probe, which fails by design for
# absent distros -- reset it so a successful uninstall exits 0. Set the var rather than `exit 0`
# so `irm ... | iex` doesn't terminate the caller's shell.
# The distro probes leave a failing $LASTEXITCODE; reset it so success exits 0. Set the
# var rather than `exit 0` so `irm ... | iex` doesn't terminate the caller's shell.
$global:LASTEXITCODE = 0
}

53
studio/backend/core/training/worker.py
View file

@ -706,23 +706,17 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
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``: the signal that matched (``"is_integrated"`` or the matching
device-name token), else ``""``.
- ``is_unified``: ``True`` for Spark-class parts that share one memory pool
with the OS (DGX Spark / GB10, N1X "RTX Spark", Grace-Blackwell desksides)
— these need the same lower ``set_per_process_memory_fraction`` cap as the
ROCm APUs: exhausting the shared pool can stall the whole box instead of
raising a catchable OutOfMemoryError.
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
``set_per_process_memory_fraction`` cap -- exhausting the pool can stall the
box instead of raising a catchable OutOfMemoryError.
Classification priority:
1. ``is_integrated`` device property (authoritative on native Linux).
2. Device-name token match — WSL2's GPU paravirtualization masks
``is_integrated`` to 0 and renames the device (the N1X reports
``JMJWOA-Generic-GPU`` with ``is_integrated == 0``, verified on
hardware), so the property alone misses Spark-under-WSL. Tokens mirror
``_DGX_SPARK_DEVICE_TOKENS`` in ``unsloth/models/_utils.py`` (duplicated
because this guard runs before any ML import).
``is_integrated`` is authoritative on native Linux, but WSL2 paravirtualization
masks it to 0 and renames the device (the 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).
"""
if getattr(props, "is_integrated", 0):
return "is_integrated", True
@ -2225,25 +2219,18 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err)
# ── 1h. NVIDIA Spark-class unified-memory OOM guard ──
# Same failure mode as the ROCm APU guard above, NVIDIA flavor: Spark-class
# parts (DGX Spark / GB10, N1X "RTX Spark") share one memory pool with the
# OS, so over-allocation can stall the whole box instead of raising a
# catchable OutOfMemoryError. Cap the allocator at 0.80 like Strix Halo —
# the pool is shared with the host OS and page cache, so 20% headroom stays
# with the system. UNSLOTH_SPARK_MEM_FRACTION overrides the cap; any value
# outside (0, 1] disables the guard. Discrete NVIDIA GPUs are untouched
# (they already raise a graceful OOM). The generic OOM handler in the
# training loop surfaces the resulting OutOfMemoryError with remediation.
# NVIDIA flavor of the ROCm APU guard above: Spark-class parts share one
# memory pool with the OS, so over-allocation can stall the box instead of
# raising a catchable OutOfMemoryError. Cap at 0.80 like Strix Halo (20%
# headroom stays with the OS/page cache). UNSLOTH_SPARK_MEM_FRACTION
# overrides; outside (0, 1] disables. Discrete NVIDIA GPUs untouched.
else:
try:
# The Spark allocator config must be decided BEFORE this guard's first
# CUDA touch: get_device_properties below initializes the CUDA allocator,
# after which PYTORCH_CUDA_ALLOC_CONF changes are ignored -- and the later
# `import unsloth` (patch_dgx_spark_memory_config) would be too late for
# THIS worker process even though it is in time for a plain
# `import unsloth`. CUDA-free sniff via nvidia-smi device names (mirrors
# _is_dgx_spark_no_cuda_init), with the same append-don't-override and
# UNSLOTH_NO_EXPANDABLE_SEGMENTS opt-out semantics as the library patch.
# 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.
try:
import platform as _plat

89
studio/scripts/provision_llama_cuda.sh
View file

@ -1,14 +1,11 @@
#!/usr/bin/env bash
# Build a CUDA llama.cpp for Unsloth Studio GGUF *inference* into
# ~/.unsloth/llama.cpp (resolver checks <dir>/build/bin/llama-server).
# Idempotent, best-effort: safe to re-run, always exits 0.
#
# Needed because no aarch64+CUDA llama.cpp prebuilt exists for NVIDIA ARM hosts
# (DGX Spark / GB10, N1X "RTX" laptops). Handles the platform gotchas:
# * nvcc rejects gcc-15 -> force gcc-14 / g++-14 as the host compiler
# * glibc >= 2.41 vs CUDA < 13.3 -> install CUDA 13.3 (rsqrt header clash)
# * sm_121 (Blackwell) GPUs -> derive arch from the GPU's compute_cap
#
# 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:
# * 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
# Opt out with UNSLOTH_NO_LLAMA_CUDA=1 (handled by the caller).
set -uo pipefail
@ -16,10 +13,9 @@ LLAMA_DIR="${UNSLOTH_LLAMA_CPP_PATH:-$HOME/.unsloth/llama.cpp}"
SERVER="$LLAMA_DIR/build/bin/llama-server"
log() { printf ' - %s\n' "$*"; }
# CUDA-capable in two layouts: old monolithic (libggml-cuda is a direct ldd dep)
# or current split build (CUDA is a dlopen-ed backend libggml-cuda.so* beside the
# binary, not shown by ldd). ldd alone false-negatives on current llama.cpp; a
# CPU-only build has no libggml-cuda.so, so its presence is the reliable signal.
# CUDA shows up two ways: old monolithic (libggml-cuda in ldd) or current split
# build (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() {
[ -x "$1" ] || return 1
ldd "$1" 2>/dev/null | grep -qi 'libggml-cuda' && return 0
@ -42,16 +38,13 @@ 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 (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.
# 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.
if [ "$HAVE_APT" -eq 1 ]; then
$SUDO apt-get update -y >/dev/null 2>&1 || true
# libcurl4-openssl-dev: _cmake_configure forces -DLLAMA_CURL=ON, and on the WSL
# deferred path this script is the only build path -- setup.sh's GGUF dep install
# (which covers libcurl) was skipped, so configure would fail without the headers.
# 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.
$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
@ -95,9 +88,8 @@ if [ -z "$NVCC" ]; then
fi
CUDA_HOME="$(dirname "$(dirname "$NVCC")")"
# CUDA toolkit + Linux dirs FIRST so the build uses Linux cmake/gcc/git, not a
# Windows tool leaked into PATH via WSL interop (/mnt/c, also has spaces). Keep
# the original PATH after so nvidia-smi etc. still resolve.
# 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.
export PATH="$CUDA_HOME/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:$PATH"
export CUDAToolkit_ROOT="$CUDA_HOME"
@ -110,14 +102,12 @@ 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. 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.
# 6. Clone + build into ~/.unsloth/llama.cpp, honoring a UNSLOTH_LLAMA_TAG pin
# (same var setup.sh uses) instead of always tracking ggml-org main.
mkdir -p "$(dirname "$LLAMA_DIR")"
_LLAMA_REF="${UNSLOTH_LLAMA_TAG:-}"
# Preserve any existing (e.g. CPU-only) llama.cpp so a failed clone OR a failed CUDA
# build doesn't leave the user with NO server: the backup is restored on any failure
# exit and only dropped once a server from the fresh build is confirmed.
# Back up any existing (e.g. CPU-only) llama.cpp: restored on any failure exit,
# dropped only once the fresh build yields a server -- never leave NO server.
_LLAMA_BAK=""
_restore_prev() {
if [ -n "$_LLAMA_BAK" ] && [ -e "$_LLAMA_BAK" ]; then
@ -143,9 +133,8 @@ if [ ! -d "$LLAMA_DIR/.git" ]; then
_restore_prev
exit 0
fi
# Honor a llama.cpp PR pin (UNSLOTH_LLAMA_PR, the same var setup.sh supports)
# so a provisioned tree matches the user's request instead of silently building
# the default branch. Best-effort: a failed fetch keeps the default branch.
# Honor a UNSLOTH_LLAMA_PR pin (same var setup.sh supports); best-effort --
# a failed fetch keeps the default branch.
case "${UNSLOTH_LLAMA_PR:-}" in
''|*[!0-9]*) ;;
*)
@ -168,22 +157,17 @@ _cmake_configure() {
-DCMAKE_CUDA_HOST_COMPILER="$HCXX" \
-DLLAMA_CURL=ON >/dev/null 2>&1
}
# A pre-existing build/ may carry an incompatible CMake cache (e.g. the installer
# relocates a versioned build dir here, leaving stale absolute paths + GGML_CUDA=OFF),
# making CUDA configure fail. Try to reuse build/ first (fast incremental resume);
# only wipe and configure clean if that fails.
# 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.
if ! _cmake_configure; then
log "stale/incompatible CMake cache detected; wiping build dir for a clean CUDA configure"
rm -rf build
_cmake_configure || { log "cmake configure failed"; cd /; _restore_prev; exit 0; }
fi
# Build the full target set unsloth-zoo's GGUF exporter also needs (llama-mtmd-cli,
# llama-gguf-split) so one build serves both Studio inference and save_pretrained_gguf.
# Parallelism default = ~half the cores: much faster than a tiny -j4, but leaves
# thermal/power headroom -- a full -j(nproc) CUDA build trips shutdowns on
# thermally constrained NVIDIA-ARM laptops (e.g. N1X "RTX Spark"). Also cap by RAM
# (~1.5 GB per nvcc job) to avoid OOM. Tune with UNSLOTH_LLAMA_BUILD_JOBS=N (raise
# on a well-cooled box, lower if it still trips). Incremental: a re-run resumes.
# 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 (N1X "RTX Spark") -- and are
# RAM-capped (~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").
if [ -n "${UNSLOTH_LLAMA_BUILD_JOBS:-}" ] && [ "${UNSLOTH_LLAMA_BUILD_JOBS}" -ge 1 ] 2>/dev/null; then
@ -198,15 +182,13 @@ else
if [ "$_memjobs" -lt "$JOBS" ]; then JOBS="$_memjobs"; fi
fi
log "building with -j${JOBS} (cores=${_ncpu})"
# Lowest CPU + idle I/O priority so this background build keeps full speed when the
# box is idle but instantly yields to a foreground `unsloth studio` / training run.
# nice/ionice: full speed when idle, yields to foreground Studio/training runs.
_NICE=""
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 (mirrors setup.sh's source path): an older
# UNSLOTH_LLAMA_TAG pin may predate newer helper targets (llama-mtmd-cli,
# llama-gguf-split), and those missing must not fail the whole provision.
# 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.
$_NICE cmake --build build -j"$JOBS" --target llama-server >/dev/null 2>&1
}
_cmake_build_extras() {
@ -216,10 +198,9 @@ _cmake_build_extras() {
done
}
if ! _cmake_build; then
# An interrupted build (e.g. a thermal/power shutdown mid-compile, which this
# machine class is prone to) can leave a partially-linked libggml-cuda.so that
# then fails to link llama-server on resume (undefined ggml_cuda_op_* refs).
# Wipe build/ and rebuild clean once before giving up.
# An interrupted build (thermal/power shutdown -- this machine class is prone)
# can leave a half-linked libggml-cuda.so that breaks the resume link
# (undefined ggml_cuda_op_* refs); wipe and rebuild clean once.
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; }

72
studio/setup.sh
View file

@ -756,10 +756,9 @@ LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp"
LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server"
_NEED_LLAMA_SOURCE_BUILD=false
_LLAMA_CPP_DEGRADED=false
# Distinct from _LLAMA_CPP_DEGRADED: on WSL2 aarch64+NVIDIA with no nvcc, the CPU
# build is skipped because install.ps1 builds the real CUDA server in the background.
# A temporarily-absent server here is success, not failure, so it must not trip the
# arm64 CPU-prebuilt last-resort or the exit 1.
# Deferred != degraded: on WSL2 aarch64+NVIDIA install.ps1 builds the real CUDA
# server in the background, so a temporarily-absent server is success and must not
# trip the arm64 CPU-prebuilt last-resort or the exit 1.
_LLAMA_CPP_DEFERRED=false
_LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}"
_REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}"
@ -933,12 +932,10 @@ if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] && \
fi
# ── WSL2 aarch64 + NVIDIA, no nvcc yet: defer to the background CUDA build ──
# On Windows-on-ARM + NVIDIA, install.ps1 builds the real CUDA llama-server in the
# background after this install. Without nvcc yet the section-9 build can only make a
# slow CPU server that the background build throws away, so skip it on this exact path.
# Gated: WSL + aarch64/arm64 + NVIDIA GPU + nvcc missing + CUDA not opted out
# (UNSLOTH_NO_LLAMA_CUDA!=1) + no forced compile / PR pin. If nvcc is present we fall
# through to section 9; if opted out we keep the CPU build as the only server.
# install.ps1 builds the real CUDA llama-server in the background after install;
# 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.
if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] \
&& [ "$_LLAMA_FORCE_COMPILE" != "1" ] \
&& [ -z "$_LLAMA_PR" ] \
@ -952,8 +949,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)"
# Use DEFERRED, not DEGRADED: DEGRADED would trigger the CPU-prebuilt last
# resort + exit 1, but install.ps1's background build is the intended builder.
# DEFERRED, not DEGRADED: DEGRADED would trigger the CPU-prebuilt last resort + exit 1.
_NEED_LLAMA_SOURCE_BUILD=false
_LLAMA_CPP_DEFERRED=true
fi
@ -1195,10 +1191,9 @@ else
else
CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON"
# glibc >= 2.41 vs CUDA < 13.3: rsqrt/rsqrtf header clash makes every .cu
# fail "exception specification is incompatible" and the GPU build drops to
# CPU. No workaround but CUDA >= 13.3. Diagnostic only: never changes flags
# or aborts, so it cannot regress any platform.
# glibc >= 2.41 + CUDA < 13.3: rsqrt/rsqrtf header clash fails every .cu
# ("exception specification is incompatible") -> CPU fallback; only fix is
# CUDA >= 13.3. Diagnostic only -- never changes flags or aborts.
_GLIBC_VER="$(getconf GNU_LIBC_VERSION 2>/dev/null | awk '{print $2}')" || _GLIBC_VER=""
if [ -n "$_GLIBC_VER" ]; then
_GLIBC_MAJ="${_GLIBC_VER%%.*}"; _GLIBC_MIN="${_GLIBC_VER#*.}"; _GLIBC_MIN="${_GLIBC_MIN%%.*}"
@ -1424,19 +1419,14 @@ fi # end _SKIP_GGUF_BUILD check
# ── aarch64 + NVIDIA (DGX Spark / GB10 / N1X "RTX Spark"): provision a CUDA
# llama.cpp when the source build above could not (no CUDA toolkit found) ──
# No aarch64+CUDA prebuilt exists, and the source build above only emits a CUDA
# server when nvcc is already present (a fresh Spark ships only driver + nvidia-smi,
# so it falls back to CPU). The Windows/WSL path closes this gap via
# provision_llama_cuda.sh; mirror it here so native-Linux Spark gets the same.
#
# Gated + additive: only on Linux aarch64/arm64 + NVIDIA GPU with no CUDA server
# yet (opt out via UNSLOTH_NO_LLAMA_CUDA=1). x86_64, ROCm, Metal, CPU-only ARM, and
# ARM hosts that already built CUDA are unaffected. Best-effort: provision always
# exits 0; on failure the prior CPU/degraded state stands for the fallback below.
# CUDA-capable in two layouts: old monolithic (libggml-cuda is a direct ldd dep) or
# split build (dlopen-ed backend libggml-cuda.so* beside the binary, not in ldd). ldd
# alone false-negatives; a CPU-only build has no libggml-cuda.so, so its presence is
# the reliable signal.
# No aarch64+CUDA prebuilt exists and a fresh Spark ships only driver + nvidia-smi,
# so the build above fell back to CPU; mirror the Windows/WSL fix
# (provision_llama_cuda.sh) for native Linux. Gated to Linux aarch64 + NVIDIA with
# no CUDA server yet (opt out: UNSLOTH_NO_LLAMA_CUDA=1); best-effort -- provision
# always exits 0 and on failure the prior CPU/degraded state stands.
# CUDA detection covers both layouts: old monolithic (libggml-cuda in ldd) and
# split build (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).
_have_cuda_llama_server() {
[ -x "$LLAMA_SERVER_BIN" ] || return 1
ldd "$LLAMA_SERVER_BIN" 2>/dev/null | grep -qi 'libggml-cuda' && return 0
@ -1450,14 +1440,11 @@ if [ "$_HOST_SYSTEM" = "Linux" ] \
&& 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
# 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.
# Under WSL this runs ONLY for a DIRECT `install.sh` run: install.ps1 exports
# UNSLOTH_WSL_LLAMA_DEFERRED=1 and builds CUDA llama.cpp in the background, but
# a direct run has no background builder, so provision here.
# 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.
# else fetch from GitHub (so `curl | sh` works on an older wheel without it).
_PROV_SH=""
if [ -f "$SCRIPT_DIR/scripts/provision_llama_cuda.sh" ]; then
_PROV_SH="$SCRIPT_DIR/scripts/provision_llama_cuda.sh"
@ -1473,15 +1460,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)"
# Builds into $LLAMA_CPP_DIR (via UNSLOTH_LLAMA_CPP_PATH so a custom
# STUDIO_HOME lands where setup.sh validates); always exits 0.
# UNSLOTH_LLAMA_CPP_PATH routes a custom STUDIO_HOME into $LLAMA_CPP_DIR; always exits 0.
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
# 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.
# 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.
if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then
: > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true
fi
@ -1489,9 +1474,8 @@ if [ "$_HOST_SYSTEM" = "Linux" ] \
substep "CUDA build unavailable; keeping existing (CPU) llama-server" "$C_WARN"
else
substep "CUDA build unavailable and no llama-server present; see $LLAMA_CPP_DIR build output" "$C_WARN"
# No server at all (e.g. the provisioner replaced a previous build and then
# failed): mark degraded so the arm64 CPU-prebuilt last resort below and the
# installer failure exit fire instead of reporting a working install.
# No server at all: mark degraded so the arm64 CPU-prebuilt last resort
# and the failure exit fire instead of reporting a working install.
_LLAMA_CPP_DEGRADED=true
fi
fi

13
unsloth/kernels/flex_attention.py
View file

@ -27,12 +27,10 @@ torch_compile_options = {
def _flex_is_dgx_spark():
# Inlined CUDA-free copy of _utils._is_dgx_spark_no_cuda_init() (kept local to
# avoid a circular import). Spark = aarch64 + a Spark device name via nvidia-smi.
# Must NOT touch torch.cuda: this runs at module import, and vision.py imports
# ..kernels before ._utils -- a device-name query here would initialize the CUDA
# allocator before patch_dgx_spark_memory_config() can set PYTORCH_CUDA_ALLOC_CONF
# on the very Spark hosts this check targets.
# Local 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 on the very Spark hosts this targets.
_force = os.environ.get("UNSLOTH_FORCE_DGX_SPARK")
if _force == "1":
return True
@ -57,8 +55,7 @@ def _flex_is_dgx_spark():
return False
# Spark's 48 SMs are below inductor's 68-SM is_big_gpu threshold, so max_autotune
# is already skipped; disabling it just avoids a wasted compile-time search.
# Spark's 48 SMs are under inductor's 68-SM is_big_gpu bar; max_autotune would only waste search time.
if _flex_is_dgx_spark():
torch_compile_options["max_autotune"] = False

103
unsloth/models/_utils.py
View file

@ -992,20 +992,16 @@ except:
from transformers.modeling_utils import logger as transformers_logger
# ---- NVIDIA DGX Spark (GB10) / N1X "RTX Spark" (Blackwell unified-memory) support ----
# Shared detector for Spark-class UMA machines, which report varying device names
# ("NVIDIA GB10" on DGX Spark, "JMJWOA-Generic-GPU" on the N1X laptop). The
# aarch64 + CUDA gate keeps every Spark workaround a strict no-op elsewhere.
# ---- NVIDIA DGX Spark (GB10) / N1X "RTX Spark" unified-memory support ----
# Device names vary ("NVIDIA GB10" on DGX Spark, "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")
@functools.lru_cache(maxsize = None)
def is_dgx_spark():
"""True only on a DGX Spark / N1X Spark-class machine.
Gate: aarch64 + NVIDIA CUDA + a known Spark device-name token. Overridable for
testing via UNSLOTH_FORCE_DGX_SPARK=1 (force on) / =0 (force off).
"""
"""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."""
_force = os.environ.get("UNSLOTH_FORCE_DGX_SPARK")
if _force == "1":
return True
@ -1028,14 +1024,10 @@ def is_dgx_spark():
@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.
"""
"""Spark detection that never initializes CUDA: reads device names via
`nvidia-smi` instead of torch, so allocator-init-time settings
(PYTORCH_CUDA_ALLOC_CONF) can still be set after calling it. Same
UNSLOTH_FORCE_DGX_SPARK override; False on any error."""
_force = os.environ.get("UNSLOTH_FORCE_DGX_SPARK")
if _force == "1":
return True
@ -1063,14 +1055,10 @@ 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.
HF sizes a GPU pre-allocation from `cudaMemGetInfo()` to warm the caching
allocator. On Spark unified memory `cudaMemGetInfo` undercounts free memory
(reclaimable buffer cache is reported unavailable), so the warmup
`torch.empty(...)` raises `AcceleratorError: invalid argument` and aborts any
runtime-quantized (bitsandbytes 4/8-bit) load. The warmup is only a speed hint,
so skipping it on Spark merely forgoes a minor warmup while letting loads
succeed. No-op on every non-Spark platform (gated by `is_dgx_spark()`).
Idempotent: re-applying is a no-op (marked via `_unsloth_spark_noop`).
`cudaMemGetInfo()` undercounts free memory on Spark unified memory, so HF's
warmup `torch.empty(...)` raises `AcceleratorError: invalid argument` and
aborts 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).
"""
if not is_dgx_spark():
return
@ -1091,21 +1079,13 @@ def patch_dgx_spark_caching_allocator_warmup():
def patch_dgx_spark_memory_config():
"""Memory-efficiency default for Spark UMA (accuracy-neutral, gated).
"""Enable allocator `expandable_segments` on Spark UMA to cut fragmentation
OOMs (accuracy-neutral; strict no-op off-Spark).
Enables the CUDA caching allocator's `expandable_segments` mode so segments can
grow in virtual address space instead of fragmenting the shared unified-memory
pool -- more of the pool stays usable for weights/activations (fewer
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. 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.
Appends to PYTORCH_CUDA_ALLOC_CONF only when absent; opt out with
UNSLOTH_NO_EXPANDABLE_SEGMENTS=1. Must run before the first CUDA allocation,
hence the CUDA-free `_is_dgx_spark_no_cuda_init()` gate -- `is_dgx_spark()`
would initialize the allocator before the env var could take effect.
"""
if not _is_dgx_spark_no_cuda_init():
return
@ -1120,21 +1100,14 @@ def patch_dgx_spark_memory_config():
def patch_dgx_spark_runtime_defaults():
"""Spark UMA runtime defaults (accuracy-neutral, gated, env-overridable).
"""Spark UMA runtime defaults (no-op off-Spark; env-overridable).
- `UNSLOTH_DISABLE_DOUBLE_BUFFER=1`: unsloth-zoo's gradient-checkpointing
double-buffer is enabled via a `torch.cuda.mem_get_info` free-memory check
that UNDERCOUNTS on UMA, and it stages an extra GPU buffer to overlap a
host<->device copy that is physically free on a shared pool. Default it off
on Spark (`setdefault`, so a user can still force it back on). Must be set
before unsloth-zoo initializes gradient checkpointing -- `import unsloth`
precedes that, so this is in time.
- `set_per_process_memory_fraction`: OPT-IN safety valve. On Spark UMA an
over-allocation can wedge the box (untracked UMA allocations may never trip
a catchable OOM). If the user sets `UNSLOTH_SPARK_MEM_FRACTION=<0..1>`, cap
the caching allocator so it raises OutOfMemoryError early. Default unset ->
NO cap (no capacity loss); purely opt-in.
Strict no-op off-Spark.
- UNSLOTH_DISABLE_DOUBLE_BUFFER=1 (setdefault): zoo's grad-checkpointing
double-buffer gates on a mem_get_info check that UNDERCOUNTS on UMA, and
its extra staging buffer is pure waste on a shared pool.
- UNSLOTH_SPARK_MEM_FRACTION=<0..1> (opt-in, default NO cap): caps the
allocator so over-allocation raises OutOfMemoryError early instead of
wedging the box (untracked UMA allocations may never trip a catchable OOM).
"""
if not is_dgx_spark():
return
@ -1142,8 +1115,7 @@ def patch_dgx_spark_runtime_defaults():
_frac = os.environ.get("UNSLOTH_SPARK_MEM_FRACTION")
if _frac:
try:
# Only (0, 1] is a usable cap: 0 would make EVERY allocation OOM
# and values > 1 are rejected by torch. Out-of-range = no cap.
# 0 would OOM every allocation; torch rejects > 1. Out-of-range = no cap.
_frac_val = float(_frac)
if 0.0 < _frac_val <= 1.0:
torch.cuda.set_per_process_memory_fraction(_frac_val)
@ -1152,17 +1124,12 @@ def patch_dgx_spark_runtime_defaults():
def patch_dgx_spark_dataloader_defaults():
"""On Spark UMA, default `dataloader_pin_memory` to False (accuracy-neutral).
"""Default `dataloader_pin_memory` to False on Spark UMA (accuracy-neutral).
Page-locked host memory exists to speed host->device DMA; on unified memory
there is no separate device memory, so pinning only reserves non-pageable RAM
from the shared pool and adds a staging copy -- pure waste. Mirrors
transformers' own `if self.use_cpu: self.dataloader_pin_memory = False`
precedent. Wraps the base `TrainingArguments.__post_init__`, so SFT + every
TRL trainer (whose configs call `super().__post_init__()`) are covered with
one idempotent patch. Only flips the library default `True`; opt out with
`UNSLOTH_SPARK_KEEP_PIN_MEMORY=1`. Strict no-op off-Spark; never changes any
computed value, so accuracy is unaffected.
With one shared memory pool, pinning only reserves non-pageable RAM and adds
a staging copy (mirrors transformers' own 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.
"""
if not is_dgx_spark():
return
@ -1177,8 +1144,7 @@ def patch_dgx_spark_dataloader_defaults():
return
_orig_post_init = Base.__post_init__
# Forward *args/**kwargs so a future TrainingArguments (or a subclass) that
# adds InitVar parameters to __post_init__ keeps working through the wrapper.
# *args/**kwargs: tolerate future InitVar parameters in __post_init__.
def __post_init__(self, *args, **kwargs):
_orig_post_init(self, *args, **kwargs)
if getattr(self, "dataloader_pin_memory", None) is True:
@ -1744,8 +1710,7 @@ torch_compile_options = {
"trace.enabled": UNSLOTH_COMPILE_DEBUG,
"triton.cudagraphs": False,
}
# Spark's 48 SMs are below inductor's 68-SM is_big_gpu threshold, so max_autotune
# is already skipped; disabling it just avoids a wasted compile-time search.
# Spark's 48 SMs are under inductor's 68-SM is_big_gpu bar; max_autotune would only waste search time.
if is_dgx_spark():
torch_compile_options["max_autotune"] = False