From 25f6b8d842554e94659871b9b567a10fb2e63dc0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 20 Jul 2026 08:05:55 +0000 Subject: [PATCH 01/10] install.ps1: allow torch 2.11 on Windows CUDA installs The Windows installer still capped torch at <2.11.0 everywhere, a stability pin from before 2.11 support landed. Every other layer already allows 2.11: install.sh widens cu leaves to <2.12.0, the Python repair layer uses the <2.12.0 trio for CUDA and CPU, and setup.ps1's bare specs on known cu leaves resolve 2.11 today. The fresh-install and CUDA flavor-repair paths now widen the trio to torch<2.12 / torchvision<0.27 / torchaudio<2.12 when the index leaf is a cu family, and keep the 2.10 line otherwise (custom pins and the CPU fallback are unchanged, matching install.sh's defaults). Verified by uv dry-runs against cu126/cu128/cu130 for win_amd64: the trio resolves paired at 2.11.0 / 0.26.0 / 2.11.0, and the existing triton-windows<3.7 constraint resolves 3.6.0.post26, the torch 2.11 pairing. Parity test updated to assert the leaf-gated widen. --- install.ps1 | 32 ++++++++++++++++------ tests/python/test_cross_platform_parity.py | 20 ++++++++++---- 2 files changed, 38 insertions(+), 14 deletions(-) diff --git a/install.ps1 b/install.ps1 index 6e059ee0dd..9aa18cd819 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2336,13 +2336,22 @@ exit 0 } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..." - # Bound the companions to the capped torch on EVERY index, cu - # families included: torchaudio 2.11 dropped its exact torch pin from - # the wheel metadata, so a bare companion next to torch<2.11 can - # resolve a mismatched 2.11.0 build. Mirrors install.sh. - $_pinVisionSpec = "torchvision>=0.19,<0.26.0" - $_pinAudioSpec = "torchaudio>=2.4,<2.11.0" - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" $_pinVisionSpec $_pinAudioSpec --default-index $TorchIndexUrl } + # Bounded trio on every index (torchaudio 2.11 dropped its exact torch + # pin, so a bare companion can drift from a capped torch). cu + # families ship torch 2.11.x with paired triton-windows 3.6, so the + # ceiling widens to <2.12 there; other leaves keep the 2.10 line. + # Mirrors install.sh and _CUDA_TORCH_PKG_SPEC. + $_idxLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() + if ($_idxLeaf -match '^cu[0-9]+$') { + $_pinTorchSpec = "torch>=2.4,<2.12.0" + $_pinVisionSpec = "torchvision>=0.19,<0.27.0" + $_pinAudioSpec = "torchaudio>=2.4,<2.12.0" + } else { + $_pinTorchSpec = "torch>=2.4,<2.11.0" + $_pinVisionSpec = "torchvision>=0.19,<0.26.0" + $_pinAudioSpec = "torchaudio>=2.4,<2.11.0" + } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython $_pinTorchSpec $_pinVisionSpec $_pinAudioSpec --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2449,8 +2458,15 @@ exit 0 $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython } elseif ($expectedTorchTag -ne 'rocm') { # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. + # Same leaf-gated ceiling as the install above: cu* serves 2.11. + $_fixLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() + if ($_fixLeaf -match '^cu[0-9]+$') { + $_fixTorchSpec = "torch>=2.4,<2.12.0"; $_fixVisionSpec = "torchvision>=0.19,<0.27.0"; $_fixAudioSpec = "torchaudio>=2.4,<2.12.0" + } else { + $_fixTorchSpec = "torch>=2.4,<2.11.0"; $_fixVisionSpec = "torchvision>=0.19,<0.26.0"; $_fixAudioSpec = "torchaudio>=2.4,<2.11.0" + } substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython $_fixTorchSpec $_fixVisionSpec $_fixAudioSpec --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index b3a9b99c55..d7e4f4ffdd 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -417,15 +417,23 @@ class TestKnown211SetParity: assert ( '$_pinAudioSpec = "torchaudio>=2.4,<2.11.0"' in text ), "install.ps1 custom-pin install must bound torchaudio (>=2.4,<2.11.0)" - # No cu-family exemption: the bounds apply unconditionally. + # cu leaves widen the whole trio to the torch 2.11 line (<2.12), + # matching install.sh's cu[0-9]* widen and _CUDA_TORCH_PKG_SPEC; other + # leaves keep the 2.10 line. The trio stays bounded on every index. + assert '$_pinTorchSpec = "torch>=2.4,<2.12.0"' in text, ( + "install.ps1 must widen torch to <2.12.0 on cu index leaves" + ) assert ( - "$_pinCuLeaf" not in text - ), "install.ps1 must bound companions on every index (no cu-family exemption)" - # The bounded companions must actually be passed to the install command. + '$_pinVisionSpec = "torchvision>=0.19,<0.27.0"' in text + ), "install.ps1 cu-leaf install must pair torchvision <0.27.0 with torch <2.12" + assert ( + '$_pinAudioSpec = "torchaudio>=2.4,<2.12.0"' in text + ), "install.ps1 cu-leaf install must pair torchaudio <2.12.0 with torch <2.12" + # The bounded trio must actually be passed to the install command. assert re.search( - r'"torch>=2\.4,<2\.11\.0" \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl', + r'\$_pinTorchSpec \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl', text, - ), "install.ps1 custom-pin install must pass the bounded companion specs to uv" + ), "install.ps1 pinned install must pass the bounded trio specs to uv" def test_gfx_allowlist_matches_across_installers(self): # The gfx 2.11 allowlist {gfx120x-all, gfx1151, gfx1150} must appear in each. From d5b743f13e50bc765a2688e8254add706556f32a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 20 Jul 2026 10:09:37 +0000 Subject: [PATCH 02/10] install: reduce comment volume across the installers Collapse the multi-line rationale blocks that accreted across the torch index, override, redaction and companion-pin work to single-line constraints, and drop review-history narration. Comments and blank lines only; the PowerShell files verify code-identical by token-stream comparison against the previous commit, install.sh passes sh -n, and the Python file passes the AST comments-only gate. install.sh -258 lines, install.ps1 -235, studio/setup.ps1 -304, studio/install_python_stack.py -182. One test touched: tests/sh/test_mac_intel_compat.sh anchors its awk extraction on a comment phrase that was shortened; the anchor now matches both wordings. Full battery: parity/install-stack/pr5940 pytest suites, all sh suites (host-defaults is the known pre-existing failure), and the four studio ps1 suites all pass. --- install.ps1 | 572 ++++------------- install.sh | 742 +++++----------------- studio/install_python_stack.py | 389 ++++-------- studio/setup.ps1 | 982 ++++++++++-------------------- tests/sh/test_mac_intel_compat.sh | 2 +- 5 files changed, 727 insertions(+), 1960 deletions(-) diff --git a/install.ps1 b/install.ps1 index 9aa18cd819..02d51b3d2b 100644 --- a/install.ps1 +++ b/install.ps1 @@ -137,8 +137,7 @@ function Install-UnslothStudio { if ($env:UNSLOTH_NO_TORCH -in @('1', 'true', 'yes', 'on')) { $SkipTorch = $true } if ($env:UNSLOTH_SKIP_AUTOSTART -in @('1', 'true', 'yes', 'on')) { $SkipAutostart = $true } - # Propagate to child processes so they also respect verbose mode. - # Process-scoped -- does not persist. + # Propagate to child processes (process-scoped). if ($script:UnslothVerbose) { $env:UNSLOTH_VERBOSE = '1' } @@ -159,15 +158,11 @@ function Install-UnslothStudio { # UNSLOTH_PYTHON pins the version (mirrors install.sh --python); default 3.13. $PythonVersion = if ($env:UNSLOTH_PYTHON) { $env:UNSLOTH_PYTHON } else { "3.13" } - # python.org fallback patch, used only when winget is unavailable/broken AND - # the live python.org listing can't be fetched. The installer URL scheme is - # stable so an older patch still installs. Bump alongside $PythonVersion. + # python.org fallback patch when winget and the live listing both fail; bump alongside $PythonVersion. $PythonFallbackFullVersion = "3.13.13" - # Resolve install destinations. Priority: UNSLOTH_STUDIO_HOME, then - # STUDIO_HOME alias, then USERPROFILE-redirect, then default. - # Reject whitespace-only values so " " is treated as unset (matches the - # Python resolvers' .strip()), preventing install/runtime layout drift. + # Install dest priority: UNSLOTH_STUDIO_HOME, STUDIO_HOME alias, USERPROFILE-redirect, default. + # Whitespace-only == unset (matches the Python resolvers' .strip()). $envOverrideVar = $null $envOverride = $null if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { @@ -178,8 +173,7 @@ function Install-UnslothStudio { $envOverride = $env:STUDIO_HOME.Trim() } - # Custom Unsloth roots are not supported with --tauri (desktop app still - # resolves %USERPROFILE%\.unsloth\studio). Pass through if override == legacy. + # Custom roots unsupported with --tauri (desktop app resolves %USERPROFILE%\.unsloth\studio); pass through if override == legacy. if ($TauriMode -and $envOverride) { $_tauriOverride = $envOverride if ($_tauriOverride -eq "~" -or $_tauriOverride -like "~/*" -or $_tauriOverride -like "~\*") { @@ -211,8 +205,7 @@ function Install-UnslothStudio { $defaultProfile = $null try { $defaultProfile = [Environment]::GetFolderPath("UserProfile") } catch {} - # LOCALAPPDATA may be unset in service / CI contexts; Join-Path would abort - # under ErrorActionPreference=Stop without this guard. + # LOCALAPPDATA may be unset in service / CI contexts; guard Join-Path under ErrorActionPreference=Stop. $defaultDataDir = if ($env:LOCALAPPDATA -and -not [string]::IsNullOrWhiteSpace($env:LOCALAPPDATA)) { Join-Path $env:LOCALAPPDATA "Unsloth Studio" } else { $null } @@ -223,8 +216,7 @@ function Install-UnslothStudio { $envOverride = (Join-Path $env:USERPROFILE $envOverride.Substring(1).TrimStart('/','\')) } try { - # .NET API: New-Item -Path treats brackets as wildcards and has no - # -LiteralPath in PS 5.1, so a root like C:\studio[abc] would fail. + # .NET API: New-Item -Path treats brackets as wildcards (no -LiteralPath in PS 5.1). [System.IO.Directory]::CreateDirectory($envOverride) | Out-Null $StudioHome = (Resolve-Path -LiteralPath $envOverride).Path } catch { @@ -233,8 +225,7 @@ function Install-UnslothStudio { } $probe = Join-Path $StudioHome (".unsloth-write-probe-" + [guid]::NewGuid()) try { - # WriteAllText: literal-path safe + closes handle so Remove-Item works. - [System.IO.File]::WriteAllText($probe, "") + [System.IO.File]::WriteAllText($probe, "") # literal-path safe + closes handle Remove-Item -LiteralPath $probe -Force -ErrorAction SilentlyContinue } catch { Write-Host "ERROR: $envOverrideVar=$StudioHome is not writable." -ForegroundColor Red @@ -304,9 +295,8 @@ function Install-UnslothStudio { } Write-Host "" - # ── Helper: refresh PATH from registry (deduplicating entries) ── - # Merge order: venv Scripts (if active) > Machine > User > current $env:Path. - # Dedup compares both raw and expanded forms (%VAR% vs literal). + # ── Helper: refresh PATH from registry, deduplicating entries ── + # Merge order: venv Scripts (if active) > Machine > User > $env:Path; dedup on raw and expanded forms. function Refresh-SessionPath { $machine = [System.Environment]::GetEnvironmentVariable("Path", "Machine") $user = [System.Environment]::GetEnvironmentVariable("Path", "User") @@ -330,8 +320,7 @@ function Install-UnslothStudio { } # ── Helper: safely add a directory to the persistent User PATH ── - # Direct registry access preserves REG_EXPAND_SZ (avoids dotnet/runtime#1442). - # Append (default) keeps existing tools first; Prepend for must-win entries. + # Direct registry access preserves REG_EXPAND_SZ (dotnet/runtime#1442). Append keeps existing tools first; Prepend for must-win. function Add-ToUserPath { param( [Parameter(Mandatory = $true)][string]$Directory, @@ -469,27 +458,24 @@ function Install-UnslothStudio { } } - # Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer - # output before printing on failure; uv/pip errors echo the failing --index-url verbatim. - # Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted. + # Redact index-URL credentials (userinfo + ?query= + #fragment) from captured output; + # uv/pip errors echo the failing --index-url verbatim. Mirrors the other installers. function Redact-InstallOutput { param([string]$Text) if (-not $Text) { return $Text } $Text = $Text -replace '(https?://)[^/@\s`]+@', '$1@' $Text = $Text -replace '([?&][^=\s&`]+)=[^&#\s`]+', '$1=' - # A #token=... fragment is as sensitive as a query; URL-anchored. + # A #token fragment is as sensitive as a query. return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#' } - # Run native commands quietly by default to match install.sh behavior. - # Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1. + # Run native commands quietly (mirrors install.sh); full output only with --verbose / UNSLOTH_VERBOSE=1. function Invoke-InstallCommand { param( [Parameter(Mandatory = $true)][ScriptBlock]$Command ) - # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898): - # for --default-index, clear the uv index env vars (restore in finally) and set - # UV_NO_CONFIG=1 so a uv.toml/pyproject index can't outrank the CLI pin (uv 0.10). + # Installer-pinned index installs must beat an inherited uv mirror (#6898): for + # --default-index, clear uv index env vars and set UV_NO_CONFIG=1 so a uv.toml/pyproject index can't outrank the CLI pin. $savedUvIndex = $null if ($Command.ToString() -match '--default-index') { $savedUvIndex = @{} @@ -505,12 +491,9 @@ function Install-UnslothStudio { # Reset to avoid stale values from prior native commands. $global:LASTEXITCODE = 0 if ($script:UnslothVerbose) { - # Merge stderr into stdout so progress/warning output stays visible - # without flipping $? on successful native commands (PS 5.1 treats - # stderr records as errors that set $? = $false even on exit code 0). - # Redact per record: uv echoes index URLs (credentials and all) in - # its errors, and verbose mode must not bypass the quiet path's - # redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched. + # Merge stderr into stdout so progress stays visible without flipping $? (PS 5.1 + # treats stderr records as errors). Redact per record (verbose must not bypass + # the quiet path's redaction); ForEach-Object/Out-Host leave $LASTEXITCODE untouched. & $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host } else { $output = & $Command 2>&1 | Out-String @@ -528,15 +511,13 @@ function Install-UnslothStudio { } } - # Retry Invoke-InstallCommand on transient uv download failures with backoff. - # Returns the last exit code on permanent failure so rollback still fires. + # Retry Invoke-InstallCommand on transient uv download failures with backoff; returns last exit code on permanent failure. function Invoke-InstallCommandRetry { param( [Parameter(Mandatory = $true, Position = 0)][ScriptBlock]$Command, [string]$Label = "install step" ) - # Sanitize overrides to a default of 3 (a typo must not disable retries; =1 disables). - # TryParse with bounds avoids an Int32 overflow throw. Bounds: 1..100 retries, 0..3600s. + # Sanitize overrides; default 3. Bounds 1..100 retries, 0..3600s (TryParse avoids overflow throw). $maxAttempts = 3 $parsedAttempts = 0 if ([int]::TryParse($env:UNSLOTH_INSTALL_RETRIES, [ref]$parsedAttempts) -and $parsedAttempts -ge 1 -and $parsedAttempts -le 100) { @@ -569,11 +550,9 @@ function Install-UnslothStudio { return } try { - # Persist an absolute path in launcher scripts so shortcut working - # directory changes do not break process startup. + # Persist an absolute path so shortcut working-dir changes don't break startup. $UnslothExePath = (Resolve-Path -LiteralPath $UnslothExePath).Path - # Escape for single-quoted embedding in generated launcher script. - # This prevents runtime variable expansion for paths containing '$'. + # Escape for single-quoted embedding (blocks runtime expansion of paths containing '$'). $SingleQuotedExePath = $UnslothExePath -replace "'", "''" # $StudioDataDir = LOCALAPPDATA\Unsloth Studio, or $StudioHome\share in env-mode. @@ -616,15 +595,9 @@ function Install-UnslothStudio { [System.IO.Directory]::CreateDirectory($appDir) | Out-Null } - # Same-install discriminator: per-install opaque id written once at - # install time and read by both this launcher and the backend - # (/api/health). Replaces the older sha256(resolved $StudioHome) - # scheme to (a) avoid leaking the install path on -H 0.0.0.0 - # deployments and (b) sidestep launcher/backend canonicalization - # drift (Resolve-Path vs Path.resolve() junction handling). Lives - # at $StudioHome\share\ (not $appDir) so the backend can find it - # via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" - # regardless of mode. 32 bytes of crypto random -> 64 hex chars. + # Same-install discriminator: per-install opaque id read by launcher and backend + # (/api/health); avoids leaking the install path and canonicalization drift. Lives at + # $StudioHome\share\studio_install_id (found via _STUDIO_ROOT_RESOLVED) regardless of mode. 32 crypto bytes -> 64 hex. $_studioIdDir = Join-Path $StudioHome "share" if (-not (Test-Path -LiteralPath $_studioIdDir)) { [System.IO.Directory]::CreateDirectory($_studioIdDir) | Out-Null @@ -639,22 +612,16 @@ function Install-UnslothStudio { $_idBytes = New-Object byte[] 32 [Security.Cryptography.RandomNumberGenerator]::Create().GetBytes($_idBytes) $_studioRootId = -join ($_idBytes | ForEach-Object { $_.ToString('x2') }) - # Atomic write: write to a temp sibling then rename, so a partial - # install cannot leave a half-written id. + # Atomic write: temp sibling then rename, so a partial install can't leave a half-written id. $_idTmp = $_studioIdFile + ".$PID.tmp" [System.IO.File]::WriteAllText($_idTmp, $_studioRootId) Move-Item -LiteralPath $_idTmp -Destination $_studioIdFile -Force } - # Env-mode: persist UNSLOTH_STUDIO_HOME (and llama path) so fresh - # shells don't need to re-export, and bake per-install $portFile / - # $mutexName so concurrent custom-root launchers cannot serialize - # through one global mutex on 8888..8908. Default installs get an - # empty prefix to match pre-PR behavior. + # Env-mode: persist UNSLOTH_STUDIO_HOME (and llama path) and bake per-install + # $portFile / $mutexName so concurrent custom-root launchers don't serialize on one global mutex. Default installs get an empty prefix. $studioHomeExport = if ($StudioRedirectMode -eq 'env') { - # When override == legacy default, llama.cpp stays at - # ~/.unsloth/llama.cpp (one shared build). Canonicalize the - # legacy side so the comparison survives path normalization. + # When override == legacy default, llama.cpp stays at ~/.unsloth/llama.cpp; canonicalize the legacy side for the comparison. $_legacyStudio = Join-Path $env:USERPROFILE ".unsloth\studio" if (Test-Path -LiteralPath $_legacyStudio -PathType Container) { $_legacyStudio = (Resolve-Path -LiteralPath $_legacyStudio).Path @@ -871,30 +838,19 @@ try { exit 0 "@ - # Write UTF-8 with BOM for reliable decoding by Windows PowerShell 5.1, - # even when install.ps1 is executed from PowerShell 7. + # Write UTF-8 with BOM so Windows PowerShell 5.1 decodes it even when run from PS 7. $utf8Bom = New-Object System.Text.UTF8Encoding($true) [System.IO.File]::WriteAllText($launcherPs1, $launcherContent, $utf8Bom) - # No .vbs launcher is written. A WScript.Shell .vbs that spawns a hidden - # ExecutionPolicy-Bypass PowerShell is exactly the shape VBS-dropper - # heuristics score (e.g. Kaspersky HEUR:Trojan.VBS.Agent.gen). The .lnk - # shortcuts instead point straight at powershell.exe running - # launch-studio.ps1 with a hidden window (selected below). + # No .vbs launcher: a WScript.Shell .vbs spawning hidden PowerShell trips VBS-dropper AV heuristics. The .lnk points straight at powershell.exe. - # Delete any launch-studio.vbs left by a pre-hardening install. New - # installs no longer generate it, but an upgrade that merely stopped - # generating it would leave the exact file AV flags on disk, so remove - # it explicitly. Covers default and env-mode installs (same $appDir). + # Delete any launch-studio.vbs left by a pre-hardening install (AV-flagged shape). Covers default and env-mode ($appDir). $legacyLauncherVbs = Join-Path $appDir "launch-studio.vbs" if (Test-Path -LiteralPath $legacyLauncherVbs) { Remove-Item -LiteralPath $legacyLauncherVbs -Force -ErrorAction SilentlyContinue } - # Prefer bundled icon from local clone/dev installs. - # If not available, best-effort download from raw GitHub. - # We only attach the icon if the resulting file has a valid ICO header. - # Snapshot the existing icon first so we can tell whether it actually - # changed and gate the heavier icon-cache refresh on a real change. + # Prefer bundled icon (local/dev), else best-effort download from raw GitHub; only attach if + # the file has a valid ICO header. Snapshot the existing icon to gate the heavier cache refresh on a real change. $preIconHash = $null if (Test-Path -LiteralPath $iconPath) { try { $preIconHash = (Get-FileHash -LiteralPath $iconPath -Algorithm SHA256).Hash } catch {} @@ -934,9 +890,7 @@ exit 0 } } - # Did the icon content actually change vs the previous install? - # Only a real change (or a first/removed icon) should trigger the heavy - # refresh; a no-op reinstall with no icon at all must not. + # Only a real change (or a first/removed icon) triggers the heavy refresh; a no-op reinstall must not. $iconChanged = $false if ($hasValidIcon) { if (-not $preIconHash) { @@ -952,26 +906,19 @@ exit 0 $iconChanged = $true } - # Env-mode: skip persistent Desktop / Start Menu .lnk shortcuts - # that may point at a deleted workspace; launcher + icon stay. + # Env-mode: skip persistent .lnk shortcuts (may point at a deleted workspace); launcher + icon stay. if ($StudioRedirectMode -eq 'env') { substep "wrote launcher at $launcherPs1 (persistent shortcuts skipped in env-override mode)" return } - # Whether this is effectively a first install (no pre-existing .lnk). - # Used to gate the heavier icon-cache refresh below so a no-op reinstall - # does not repeatedly clear caches / restart StartMenuExperienceHost -- - # a behavioral cluster AV heuristics can score as dropper-like. + # First install == no pre-existing .lnk; gates the heavy refresh so a no-op reinstall doesn't repeatedly clear caches (a dropper-like AV cluster). $firstInstall = -not ( ($desktopLink -and (Test-Path -LiteralPath $desktopLink)) -or ($startMenuLink -and (Test-Path -LiteralPath $startMenuLink)) ) - # Launch transport for the shortcuts: powershell.exe runs - # launch-studio.ps1 with a hidden window. We deliberately avoid a - # .vbs/WScript.Shell wrapper -- that script-engine shape is what AV - # VBS-dropper heuristics score (Kaspersky HEUR:Trojan.VBS.Agent.gen). + # Shortcut transport: powershell.exe runs launch-studio.ps1 hidden; no .vbs/WScript wrapper (VBS-dropper AV heuristics). $powershellForLnk = Join-Path $env:SystemRoot "System32\WindowsPowerShell\v1.0\powershell.exe" $shortcutTarget = $powershellForLnk $shortcutArgs = "-NoProfile -WindowStyle Hidden -ExecutionPolicy Bypass -File `"$launcherPs1`"" @@ -1002,13 +949,8 @@ exit 0 } if ($createdShortcutCount -gt 0) { substep "Created Unsloth Studio shortcut" - # Always do the cheap, non-disruptive per-item refresh so a - # rewritten same-name .lnk renders with its new target/icon - # immediately (a same-name .lnk recreated across reinstalls keeps - # Explorer's cached per-item icon). The reliable fix (no explorer - # restart) is a per-item SHChangeNotify SHCNE_UPDATEITEM + - # SHCNF_PATHW per .lnk; the global SHCNE_ASSOCCHANGED broadcast - # alone does NOT recover a stale item. + # Cheap per-item refresh so a rewritten same-name .lnk renders its new target/icon: + # per-item SHChangeNotify SHCNE_UPDATEITEM + SHCNF_PATHW (the global SHCNE_ASSOCCHANGED broadcast alone does not recover a stale item). try { Add-Type -Namespace UnslothShell -Name IconRefresh -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);' -ErrorAction SilentlyContinue # SHCNE_UPDATEITEM (0x00002000) + SHCNF_PATHW (0x0005) per shortcut @@ -1018,21 +960,12 @@ exit 0 # SHCNE_ASSOCCHANGED (0x08000000) global refresh (belt-and-suspenders) [UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, $null, [System.IntPtr]::Zero) } catch {} - # Heavier on-disk icon-cache clear + StartMenuExperienceHost tile - # rebuild only when the icon actually changed or this is a first - # install. Running "clear icon cache + kill StartMenuExperienceHost" - # on every no-op reinstall is a dropper-like behavioral cluster and - # is unnecessary when the icon is unchanged (the per-item notify - # above already refreshes the rewritten shortcut). + # Heavier icon-cache clear + StartMenuExperienceHost rebuild only on first install or icon change (doing it every no-op reinstall is a dropper-like cluster). if ($firstInstall -or $iconChanged) { try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {} try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {} - # Win11's Start Menu (StartMenuExperienceHost) keeps its OWN - # pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT - # invalidate, so a rewritten same-name shortcut shows the old tile - # until the host restarts. Drop only the render caches (NEVER - # start2.bin -- the pinned layout) and let the host rebuild. - # Best-effort; Win10 has no such host (Test-Path skips it). + # Win11 StartMenuExperienceHost keeps its own tile-icon cache that ie4uinit/explorer restart don't + # invalidate. Drop only the render caches (NEVER start2.bin, the pinned layout) and let the host rebuild. Win10 has no such host. try { $smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState" if (Test-Path -LiteralPath $smehTemp) { @@ -1068,12 +1001,7 @@ exit 0 } # ── Check winget ── - # winget is only needed to install Python or uv. If both are - # already on PATH (Windows ARM64 GitHub-hosted runners, manual - # python.org + Astral uv installs, corporate locked-down hosts - # without the Store, etc.) the script can proceed without it. - # We defer the hard failure to the Python / uv install branches - # below, where winget is actually invoked. + # Only needed to install Python or uv; if both are already on PATH the script proceeds without it (hard failure deferred to the install branches below). Write-TauriLog "STEP" "Checking system dependencies" $script:WingetAvailable = [bool](Get-Command winget -ErrorAction SilentlyContinue) if ($script:WingetAvailable) { @@ -1083,19 +1011,7 @@ exit 0 substep "Get it from https://aka.ms/getwinget if Python / uv are not already on PATH." "Yellow" } - # ── Helper: detect a working Python 3.11-3.13 on the system ── - # Returns the version string (e.g. "3.13") or "" if none found. - # Uses try-catch + stderr redirection so that App Execution Alias stubs - # (WindowsApps) and other non-functional executables are probed safely - # without triggering $ErrorActionPreference = "Stop". - # - # Skips Anaconda/Miniconda Python: conda-bundled CPython ships modified - # DLL search paths that break torch's c10.dll loading on Windows. - # Standalone CPython (python.org, winget, uv) does not have this issue. - # - # NOTE: A venv created from conda Python inherits conda's base_prefix - # even if the venv path does not contain "conda". We check both the - # executable path AND sys.base_prefix to catch this. + # ── Helper: detect a working Python 3.11-3.13. Skips conda: its modified DLL search paths break torch's c10.dll on Windows; check both exe path AND sys.base_prefix since a conda-derived venv inherits base_prefix even when its path lacks "conda". ── $script:CondaSkipPattern = '(?i)(conda|miniconda|anaconda|miniforge|mambaforge)' function Test-IsCondaPython { @@ -1108,17 +1024,11 @@ exit 0 return $false } - # Returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. - # The resolved Path is passed to `uv venv --python` to prevent uv from - # re-resolving the version string back to a conda interpreter. + # Returns @{ Version; Path } or $null; the resolved Path is passed to `uv venv --python` so uv doesn't re-resolve the version string back to a conda interpreter. function Find-CompatiblePython { - # Try the Python Launcher first (most reliable on Windows) - # py.exe resolves to the standard CPython install, not conda. - # Prefer the requested $PythonVersion, then newest-first fallback. + # py.exe first (resolves standard CPython, not conda); prefer $PythonVersion then newest-first. $minors = @($PythonVersion) + (@("3.13", "3.12", "3.11") | Where-Object { $_ -ne $PythonVersion }) - # Enumerate every py.exe on PATH with -All (Windows PowerShell 5.1 - # returns only the first launcher without it) and search each for a - # supported, non-conda interpreter. + # -All: Windows PowerShell 5.1 returns only the first launcher without it. foreach ($pyLauncher in @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue)) { if ($pyLauncher.Source -match $script:CondaSkipPattern) { continue } foreach ($minor in $minors) { @@ -1135,13 +1045,7 @@ exit 0 } catch {} } } - # Try python3 / python via Get-Command -All to look past stubs that - # might shadow a real Python further down PATH. - # Skip WindowsApps entries: the App Execution Alias stubs live there - # and can open the Microsoft Store as a side effect. Legitimate Store - # Python is already detected via the py launcher above (Store packages - # include py since Python 3.11). - # Skip Anaconda/Miniconda: check both path and sys.base_prefix. + # python3 / python via -All to look past stubs shadowing a real Python; skip WindowsApps (App Execution Alias stubs can open the Store; real Store Python is caught by the py launcher above) and conda (path + sys.base_prefix). foreach ($name in @("python3", "python")) { foreach ($cmd in @(Get-Command $name -All -ErrorAction SilentlyContinue)) { if (-not $cmd.Source) { continue } @@ -1158,12 +1062,7 @@ exit 0 return $null } - # ── Fallback: install CPython directly from python.org ── - # Used when winget is unavailable or fails (notably msstore cert-pinning error - # 0x8a15005e, which aborts `winget install` unless --source winget is given). - # Downloads the official installer and runs it silently as a per-user install - # (no UAC), putting python.exe + the py launcher on PATH. Mirrors the uv -> - # astral.sh fallback below. Returns @{ Version; Path } or $null. + # ── Fallback: install CPython from python.org when winget is unavailable/fails (notably msstore cert-pinning 0x8a15005e). Silent per-user install (no UAC) puts python.exe + py launcher on PATH. Returns @{ Version; Path } or $null. ── function Install-PythonFromPythonOrg { # python.org ships one installer per architecture. $archSuffix = switch (Get-TauriDiagArch) { @@ -1177,10 +1076,7 @@ exit 0 return $null } - # Resolve the latest $PythonVersion.x patch from the python.org listing, - # falling back to a same-minor version if the listing cannot be fetched. - # Use the pinned full version only when it matches the requested minor so a - # non-default UNSLOTH_PYTHON (e.g. 3.12) doesn't silently install 3.13. + # Latest $PythonVersion.x patch from the python.org listing, else same-minor fallback. Use the pinned full version only when it matches the requested minor so a non-default UNSLOTH_PYTHON (e.g. 3.12) doesn't silently install 3.13. $full = if ($PythonFallbackFullVersion -like "$PythonVersion.*") { $PythonFallbackFullVersion } else { "$PythonVersion.0" } try { $listing = [string](Invoke-RestMethod -Uri "https://www.python.org/ftp/python/" -UseBasicParsing -TimeoutSec 20) @@ -1200,16 +1096,14 @@ exit 0 return $null } - # Per-user install => no UAC. PrependPath puts python + py on PATH; - # Include_launcher installs py.exe (preferred by Find-CompatiblePython). + # Per-user install => no UAC. PrependPath puts python + py on PATH; Include_launcher installs py.exe (preferred by Find-CompatiblePython). substep "installing Python $full (silent, per-user)..." $installArgs = @( "/quiet", "InstallAllUsers=0", "PrependPath=1", "Include_launcher=1", - # Launcher per-user too: Include_launcher defaults InstallLauncherAllUsers=1, - # which needs admin and would break this non-admin per-user fallback. + # Launcher per-user too: Include_launcher defaults InstallLauncherAllUsers=1 (needs admin, breaks this non-admin fallback). "InstallLauncherAllUsers=0", "Include_pip=1", "AssociateFiles=0", @@ -1232,7 +1126,6 @@ exit 0 } # ── Install Python if no compatible version (3.11-3.13) found ── - # Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. Write-TauriLog "STEP" "Installing Python" $DetectedPython = Find-CompatiblePython @@ -1245,13 +1138,7 @@ exit 0 $wingetExit = $null if ($script:WingetAvailable) { - # --source winget avoids the msstore source, which can fail with - # cert-pinning error 0x8a15005e and abort the whole `winget install` - # (winget then demands --source). Python and uv both live in the - # winget source, so pinning it is correct and faster. - # - # Lower ErrorActionPreference so winget stderr (progress/warnings) is - # not a terminating error on PS 5.1 (native stderr is ErrorRecord). + # --source winget avoids the msstore source (cert-pinning 0x8a15005e can abort the whole install); Python and uv both live in the winget source. Lower ErrorActionPreference so winget stderr isn't a terminating error on PS 5.1. $prevEAP = $ErrorActionPreference $ErrorActionPreference = "Continue" try { @@ -1265,10 +1152,7 @@ exit 0 $DetectedPython = Find-CompatiblePython if (-not $DetectedPython) { - # Python still not functional after winget -- force reinstall. - # This handles both real failures AND "already installed" codes where - # winget thinks Python is present but it's not actually on PATH - # (e.g. user partially uninstalled, or installed via a different method). + # Still not functional after winget -- force reinstall. Handles real failures AND "already installed" codes where winget thinks Python is present but it's not on PATH. substep "Python not found on PATH after winget. Retrying with --force..." "Yellow" $ErrorActionPreference = "Continue" try { @@ -1281,9 +1165,7 @@ exit 0 } } - # Fall back to python.org if winget is unavailable OR couldn't install a - # working Python (missing/broken winget, msstore cert errors --source - # winget can't fix). Keeps the install automatic instead of failing out. + # Fall back to python.org if winget is unavailable or couldn't install a working Python (msstore cert errors --source winget can't fix), keeping the install automatic. if (-not $DetectedPython) { if ($script:WingetAvailable) { substep "winget could not install Python -- falling back to python.org..." "Yellow" @@ -1343,10 +1225,7 @@ exit 0 $ErrorActionPreference = $prevEAP Refresh-SessionPath } - # Fallback: if winget is unavailable or didn't put uv on PATH, - # use Astral's official PowerShell installer. This is the only - # supported path on hosts without winget (Windows ARM64 runners, - # corporate machines without the Store, etc.). + # Fallback: Astral's official PowerShell installer -- the only path on hosts without winget. if (-not (Test-UvVersionOk)) { substep "installing uv via https://astral.sh/uv/install.ps1..." "Yellow" Invoke-Expression (Invoke-RestMethod -Uri "https://astral.sh/uv/install.ps1") @@ -1354,8 +1233,7 @@ exit 0 } } - # A freshly installed uv can sit later on PATH than an older one (active - # venv, Scoop/pipx shim). Prefer a just-installed uv from a known location. + # A freshly installed uv can sit later on PATH than an older one; prefer a just-installed uv from a known location. if (-not (Test-UvVersionOk)) { $origPath = $env:PATH foreach ($d in @($env:UV_INSTALL_DIR, $env:XDG_BIN_HOME, @@ -1375,14 +1253,12 @@ exit 0 return (Exit-InstallFailure "uv could not be installed") } - # When bytecode compilation is enabled, large installs can exceed uv's 60s - # default on slow machines. Default to 180s, preserving overrides ("0" disables). + # Bytecode compilation can exceed uv's 60s default on slow machines; default 180s, preserving overrides ("0" disables). if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT) { $env:UV_COMPILE_BYTECODE_TIMEOUT = "180" } - # uv >= 0.8.16 retries HTTP/2 streaming body errors; raise retries and read - # timeout for large wheel downloads. User-provided values are preserved. + # Raise uv HTTP retries + read timeout for large wheel downloads (preserves user values). if (-not $env:UV_HTTP_RETRIES) { $env:UV_HTTP_RETRIES = "5" } @@ -1390,9 +1266,7 @@ exit 0 $env:UV_HTTP_TIMEOUT = "180" } - # ── Create venv (migrate old layout if possible, otherwise fresh) ── - # Pass the resolved executable path to uv so it does not re-resolve - # a version string back to a conda interpreter. + # ── Create venv (migrate old layout if possible, otherwise fresh); pass the resolved exe path to uv so it doesn't re-resolve back to conda. ── Write-TauriLog "STEP" "Creating virtual environment" if (-not (Test-Path -LiteralPath $StudioHome)) { # .NET API: New-Item -Path treats brackets as wildcards. @@ -1457,11 +1331,7 @@ exit 0 } if (Test-Path -LiteralPath $VenvPython) { - # why: matching guard to the .venv branch below -- in env-mode - # $StudioHome is a user-chosen workspace, so refuse to nuke an - # existing $StudioHome\unsloth_studio that lacks Unsloth sentinels. - # -PathType Leaf rejects a directory at the sentinel path. Accept the - # in-VENV ownership marker so partial-install retries are not blocked. + # env-mode: $StudioHome is a user-chosen workspace, so refuse to nuke an existing venv lacking Unsloth sentinels (-PathType Leaf rejects a directory at the sentinel path; accept the in-VENV ownership marker so partial-install retries aren't blocked). if ( $StudioRedirectMode -eq 'env' -and -not (Test-Path -LiteralPath (Join-Path $VenvDir ".unsloth-studio-owned") -PathType Leaf) -and @@ -1484,9 +1354,7 @@ exit 0 $StudioRedirectMode -ne 'env' ` -and (Test-Path -LiteralPath (Join-Path $StudioHome ".venv\Scripts\python.exe")) ) { - # Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating. - # Skip in env-mode so we don't blow away an unrelated .venv at the - # workspace root (e.g. user's existing project Python venv). + # Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating. Skipped in env-mode so we don't blow away an unrelated .venv at the workspace root. $OldVenv = Join-Path $StudioHome ".venv" $OldPy = Join-Path $OldVenv "Scripts\python.exe" substep "found legacy Unsloth environment, validating..." @@ -1515,9 +1383,7 @@ exit 0 $StudioRedirectMode -ne 'env' ` -and (Test-Path -LiteralPath (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) ) { - # CWD-relative venv from old install.ps1 -> migrate to absolute path. - # Skip in env-mode so we don't relocate the default-install venv into - # the workspace root. + # CWD-relative venv from old install.ps1 -> migrate to absolute path. Skipped in env-mode so we don't relocate the default-install venv into the workspace root. $CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio" substep "found CWD-relative Unsloth environment, migrating to $VenvDir..." Move-Item -LiteralPath $CwdVenv -Destination $VenvDir -Force @@ -1538,33 +1404,23 @@ exit 0 substep "$VenvDir" } - # Mark the freshly-created venv as Unsloth-owned so a partial install can be - # repaired by re-running install.ps1; the env-mode deletion guard above - # accepts this marker as the primary sentinel. + # Mark the venv Unsloth-owned so a partial install can be repaired by re-running; the env-mode deletion guard above accepts this marker as the primary sentinel. if (Test-Path -LiteralPath $VenvDir -PathType Container) { try { [System.IO.File]::WriteAllText((Join-Path $VenvDir ".unsloth-studio-owned"), "") } catch {} } - # ── Helper: run amd-smi without triggering a UAC elevation prompt ── - # amd-smi on Windows auto-elevates to read GPU/APU memory, surfacing a confusing - # DiskPart UAC prompt mid-install (Unsloth backend amd.py hits the same). - # __COMPAT_LAYER=RunAsInvoker forces it (and helpers it spawns) to run - # un-elevated; on failure the WMI name -> gfx fallback still resolves the arch. + # ── Helper: run amd-smi without a UAC prompt -- it auto-elevates on Windows to read GPU/APU memory (confusing mid-install DiskPart prompt). __COMPAT_LAYER=RunAsInvoker forces un-elevated; on failure the WMI name -> gfx fallback still resolves the arch. ── function Invoke-AmdSmiNoElevate { param( [Parameter(Mandatory = $true, Position = 0)][string]$Exe, [Parameter(Position = 1)][string[]]$SmiArgs = @(), [int]$TimeoutSec = 30 ) - # RunAsInvoker blocks the auto-elevation/UAC prompt; the timeout bounds a - # flaky amd-smi that can otherwise spin for minutes (30s mirrors amd.py). + # RunAsInvoker blocks the auto-elevation prompt; the timeout bounds a flaky amd-smi (30s mirrors amd.py). $prevCompat = [Environment]::GetEnvironmentVariable('__COMPAT_LAYER', 'Process') $env:__COMPAT_LAYER = 'RunAsInvoker' try { - # [Process]::Start, NOT Start-Process -PassThru: the latter leaves - # .ExitCode $null after WaitForExit on PS 5.1, so $LASTEXITCODE (checked - # by callers) reads non-zero and kills detection. Async reads drain the - # pipes (no deadlock); amd-smi args have no spaces so a plain join is safe. + # [Process]::Start, NOT Start-Process -PassThru: the latter leaves .ExitCode $null after WaitForExit on PS 5.1, so $LASTEXITCODE reads non-zero and kills detection. Async reads drain the pipes; amd-smi args have no spaces so a plain join is safe. $psi = New-Object System.Diagnostics.ProcessStartInfo $psi.FileName = $Exe $psi.Arguments = ($SmiArgs -join ' ') @@ -1594,11 +1450,7 @@ exit 0 } } - # ── Helper: run nvidia-smi under a timeout ── - # A wedged NVIDIA driver can make nvidia-smi block during init or after a - # reset; WaitForExit bounds it (mirrors Invoke-AmdSmiNoElevate) so detection - # cannot hang the installer. No RunAsInvoker compat layer: nvidia-smi does - # not auto-elevate. Returns combined stdout+stderr; "" on timeout/failure. + # ── Helper: run nvidia-smi under a timeout so a wedged driver can't hang the installer (no RunAsInvoker: nvidia-smi doesn't auto-elevate). Returns combined stdout+stderr; "" on timeout/failure. ── function Invoke-NvidiaSmiBounded { param( [Parameter(Mandatory = $true, Position = 0)][string]$Exe, @@ -1629,10 +1481,7 @@ exit 0 } } - # ── Helper: nvidia-smi -L lists at least one real GPU ── - # Exit code 0 alone is not enough: a stale/driverless nvidia-smi can exit 0 - # while listing no GPU, which would mark an AMD host NVIDIA and suppress - # ROCm detection. Require a "GPU :" data row. + # ── Helper: nvidia-smi -L lists at least one real GPU. Exit 0 alone isn't enough (a stale/driverless nvidia-smi can exit 0 with no GPU, marking an AMD host NVIDIA and suppressing ROCm) -- require a "GPU :" data row. ── function Test-NvidiaSmiHasGpu { param([Parameter(Mandatory = $true)][string]$Exe) $out = Invoke-NvidiaSmiBounded $Exe @('-L') @@ -1667,15 +1516,11 @@ exit 0 $ROCmVersion = $null $ROCmGfxArch = $null if (-not $HasNvidiaSmi) { - # hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (mirrors NVIDIA smi path resolution). - # AMD HIP SDK sets HIP_PATH but may not add the bin dir to PATH depending on install type. - # Ignore the venv hipInfo.exe (AMD wheel, on PATH): not a HIP SDK, so - # amd-smi would still auto-elevate. Cf. _path_inside_venv(). + # hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (HIP SDK sets HIP_PATH but may not add bin to PATH). Ignore the venv hipInfo.exe (AMD wheel, not a HIP SDK, so amd-smi would still auto-elevate). Cf. _path_inside_venv(). function Test-HipinfoIsVenvInternal { param([AllowNull()][string]$HipinfoPath) if ([string]::IsNullOrWhiteSpace($HipinfoPath)) { return $false } - # Also derive the venv from the setup python + default Unsloth home, so - # the venv hipInfo is caught when VenvDir/VIRTUAL_ENV are unset. + # Also derive the venv from the setup python + default Unsloth home, so the venv hipInfo is caught when VenvDir/VIRTUAL_ENV are unset. $venvRoots = @() if ($env:VIRTUAL_ENV) { $venvRoots += $env:VIRTUAL_ENV } $vd = Get-Variable -Name VenvDir -ValueOnly -ErrorAction SilentlyContinue @@ -1684,15 +1529,12 @@ exit 0 try { $venvRoots += (Split-Path -Parent (Split-Path -Parent $env:UNSLOTH_SETUP_PYTHON)) } catch {} } if ($env:USERPROFILE) { $venvRoots += (Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio") } - # A custom Unsloth home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the - # venv off the default path; seed it too or its hipInfo escapes the filter. + # A custom Unsloth home moves the venv off the default path; seed it too or its hipInfo escapes the filter. $studioHomeEnv = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() } else { $null } if ($studioHomeEnv) { - # Expand a leading ~ like the canonical resolver; else GetFullPath - # keeps the literal ~ (cwd-relative) and the hipInfo escapes the filter. + # Expand a leading ~ like the canonical resolver; else GetFullPath keeps the literal ~ and the hipInfo escapes the filter. if (($studioHomeEnv -eq "~" -or $studioHomeEnv -like "~/*" -or $studioHomeEnv -like "~\*") -and -not [string]::IsNullOrWhiteSpace($env:USERPROFILE)) { - # A bare "~" leaves an empty child path; Join-Path rejects that on - # PS 5.1, so use USERPROFILE directly and only join a real remainder. + # A bare "~" leaves an empty child path (Join-Path rejects that on PS 5.1), so use USERPROFILE directly and only join a real remainder. $studioHomeRest = $studioHomeEnv.Substring(1).TrimStart('/', '\') $studioHomeEnv = if ($studioHomeRest) { Join-Path $env:USERPROFILE $studioHomeRest } else { $env:USERPROFILE } } @@ -1702,8 +1544,7 @@ exit 0 foreach ($root in $venvRoots) { if ([string]::IsNullOrWhiteSpace($root)) { continue } try { $r = [System.IO.Path]::GetFullPath($root).TrimEnd('\', '/') } catch { continue } - # Skip a bare drive root (e.g. a non-venv UNSLOTH_SETUP_PYTHON like - # C:\Python311\python.exe yields C:) -- it would match every path on that drive. + # Skip a bare drive root (e.g. a non-venv UNSLOTH_SETUP_PYTHON yields C:) -- it would match every path on that drive. if ($r -match '^[a-zA-Z]:$') { continue } if ($hip.Equals($r, [System.StringComparison]::OrdinalIgnoreCase) -or $hip.StartsWith($r + [System.IO.Path]::DirectorySeparatorChar, [System.StringComparison]::OrdinalIgnoreCase)) { @@ -1712,15 +1553,12 @@ exit 0 } return $false } - # Scan all hipinfo and keep the first non-venv one (the venv copy from the - # bnb fix could shadow a real HIP SDK's). -CommandType Application matches - # only real executables, not a user alias/function named hipinfo. + # Scan all hipinfo and keep the first non-venv one (the venv copy could shadow a real HIP SDK); -CommandType Application matches only real executables, not an alias/function named hipinfo. $hipinfoExe = Get-Command hipinfo -CommandType Application -All -ErrorAction SilentlyContinue | Where-Object { -not (Test-HipinfoIsVenvInternal $_.Source) } | Select-Object -First 1 if (-not $hipinfoExe) { - # Iterate the env roots (mirrors the Python list) and take the first non-venv - # bin\hipinfo.exe, so a venv-internal HIP_PATH can't mask a real SDK in ROCM_PATH. + # Iterate the env roots and take the first non-venv bin\hipinfo.exe, so a venv-internal HIP_PATH can't mask a real SDK in ROCM_PATH. $hipMissingLabel = $null; $hipMissingRoot = $null; $hipMissingCandidate = $null foreach ($hipEnvLabel in @("HIP_PATH", "HIP_PATH_57", "ROCM_PATH")) { $hipRoot = [Environment]::GetEnvironmentVariable($hipEnvLabel) @@ -1773,14 +1611,7 @@ exit 0 } } catch {} } - # On hosts without a working HIP runtime amd-smi elevates a child at runtime, - # popping a UAC/DiskPart prompt RunAsInvoker can't suppress (manifest is - # asInvoker). So only probe when a HIP SDK is present (hipinfo found -> - # un-elevated) or the user opts in; else fall through to WMI name inference - # (enough to pick ROCm wheels + the ROCm llama.cpp prebuilt). - # An explicit opt-out (UNSLOTH_ENABLE_AMD_SMI=0/false/no/off) wins over the - # HIP-SDK heuristic: a HIP SDK binary with a broken runtime can still pop the - # prompt, so $HipSdkInstalled must NOT silently re-enable it. + # Without a working HIP runtime amd-smi elevates a child at runtime (UAC/DiskPart prompt RunAsInvoker can't suppress), so only probe when a HIP SDK is present or the user opts in; else fall through to WMI name inference (enough for ROCm wheels + llama.cpp prebuilt). An explicit opt-out (UNSLOTH_ENABLE_AMD_SMI=0/false/no/off) wins over the HIP-SDK heuristic, since a broken runtime can still pop the prompt. $amdSmiOptOut = $env:UNSLOTH_ENABLE_AMD_SMI -match '^(?i)(0|false|no|off)$' $amdSmiAllowed = (-not $amdSmiOptOut) -and ($HipSdkInstalled -or ($env:UNSLOTH_ENABLE_AMD_SMI -match '^(?i)(1|true|yes|on)$')) if (-not $HasROCm -and $amdSmiAllowed) { @@ -1790,8 +1621,7 @@ exit 0 $smiOut = Invoke-AmdSmiNoElevate $amdSmiExe.Source @('list') if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") { $HasROCm = $true - # Mirror the hipinfo path: collect all gfx tokens in enumeration - # order and pick the runtime-visible one via HIP_VISIBLE_DEVICES. + # Mirror the hipinfo path: collect all gfx tokens in enumeration order and pick the runtime-visible one via HIP_VISIBLE_DEVICES. $_smiVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 } # Attempt 1: newer amd-smi versions embed the gfx arch in list output. $_smiGfxTokens = @([regex]::Matches($smiOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() }) @@ -1799,8 +1629,7 @@ exit 0 $ROCmGfxArch = if ($_smiVisIdx -lt $_smiGfxTokens.Count) { $_smiGfxTokens[$_smiVisIdx] } else { $_smiGfxTokens[0] } $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" } else { - # Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+, - # including the GFX target needed for wheel index selection. + # Attempt 2: 'static --asic' exposes the GFX target (ROCm 6+) needed for wheel index selection. $smiAsicOut = "" try { $smiAsicOut = Invoke-AmdSmiNoElevate $amdSmiExe.Source @('static','--asic') } catch {} $_asicGfxTokens = @([regex]::Matches($smiAsicOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() }) @@ -1825,11 +1654,7 @@ exit 0 if ($wmiGpu) { $ROCmGpuLabel = $wmiGpu.Name } } catch {} } - # ── Arch resolution: env-var override → name inference ────────────── - # Runs even when the probe can't confirm a runtime ($HasROCm false): the - # WMI-name gfx arch drives both ROCm llama.cpp and torch. repo.amd.com - # wheels bundle their own runtime (no HIP SDK), so a mapped arch installs - # ROCm torch directly below -- no wasted CPU base. + # ── Arch resolution: env-var override → name inference. Runs even when the probe can't confirm a runtime ($HasROCm false): the WMI-name gfx arch drives ROCm llama.cpp and torch, and repo.amd.com wheels bundle their own runtime, so a mapped arch installs ROCm torch directly below with no wasted CPU base. ── if (-not $ROCmGfxArch) { # 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running. if ($env:UNSLOTH_ROCM_GFX_ARCH) { @@ -1837,9 +1662,7 @@ exit 0 $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $ROCmGfxArch" "Cyan" } - # 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI). - # Targets only arches the ROCm prebuilts cover - # (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU. + # 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI); targets only arches the ROCm prebuilts cover (gfx120X/110X/1151/1150/103X), unknown names fall back to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080) @@ -1864,9 +1687,7 @@ exit 0 } } } - # Capture ROCm version for wheel selection (hipconfig, then amd-smi). - # Run whenever the HIP SDK binary is present, not just when the device is accessible -- - # hipconfig --version works even when hipinfo reports no ROCm device (driver issue). + # Capture ROCm version for wheel selection (hipconfig, then amd-smi). Run whenever the HIP SDK binary is present, since hipconfig --version works even when hipinfo reports no ROCm device (driver issue). if ($HasROCm -or $HipSdkInstalled) { $hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue if (-not $hipConfigExe) { @@ -1906,12 +1727,7 @@ exit 0 } } - # ── Optional WSL-ROCm driver hint ──────────────────────────────────────── - # An AMD GPU can also be used inside WSL2, but only with Adrenalin >= 26.2.2 - # (first production ROCDXG/WSL release); native Windows GPU works with any - # recent driver. We can't auto-install it (AMD referrer-gates downloads, no - # winget package), so just point at AMD's page. Shown only when the installed - # driver predates 26.2.2 (Feb 2026); suppress with UNSLOTH_SKIP_AMD_DRIVER_HINT=1. + # ── Optional WSL-ROCm driver hint: WSL2 needs AMD Adrenalin >= 26.2.2 (native Windows GPU works with any recent driver). Can't auto-install (AMD referrer-gates downloads, no winget package), so just point at AMD's page; shown only when the installed driver predates 26.2.2 (Feb 2026). Suppress with UNSLOTH_SKIP_AMD_DRIVER_HINT=1. ── function Show-AmdWslDriverHint { if ($env:UNSLOTH_SKIP_AMD_DRIVER_HINT) { return } try { @@ -1928,15 +1744,13 @@ exit 0 $drvDate = [Management.ManagementDateTimeConverter]::ToDateTime([string]$amd.DriverDate) } } catch {} - # Older than 26.2.2 (Feb 2026) => can't expose the GPU to WSL ROCm. - # Unreadable date => still show the hint (informational, suppressible). + # Older than 26.2.2 (Feb 2026) => can't expose the GPU to WSL ROCm; unreadable date => still show the hint (informational, suppressible). if ($drvDate -and $drvDate -ge (Get-Date '2026-02-01')) { return } substep "Tip: to use this GPU inside WSL too, install AMD Adrenalin 26.2.2+ (for WSL2)." "Cyan" substep " Your current driver predates it; native Windows GPU is unaffected. Get it from AMD:" "Cyan" substep " https://www.amd.com/en/resources/support-articles/release-notes/RN-RAD-WIN-26-2-2.html" "Cyan" substep " Then reboot and run this installer inside an Ubuntu-24.04 WSL distro." "Cyan" - # If WSL isn't installed yet, point at the command that provisions it - # (best-effort; wsl.exe absent => no WSL). + # If WSL isn't installed yet, point at the command that provisions it (wsl.exe absent => no WSL). $hasWsl = $false try { $hasWsl = [bool](Get-Command wsl.exe -ErrorAction SilentlyContinue) } catch {} if (-not $hasWsl) { @@ -1963,8 +1777,7 @@ exit 0 substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow" substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" } elseif ($ROCmGfxArch) { - # Known arch: Unsloth setup installs AMD's bundled-runtime ROCm PyTorch wheels - # (repo.amd.com), which ship their own runtime -- HIP SDK optional. + # Known arch: Unsloth setup installs AMD's bundled-runtime ROCm PyTorch wheels (repo.amd.com), which ship their own runtime -- HIP SDK optional. step "gpu" "AMD ROCm ($ROCmGfxArch)" "Cyan" substep "Detected: $ROCmGpuLabel" "Cyan" substep "GPU PyTorch uses AMD's bundled-runtime ROCm wheels -- HIP SDK not required (optional)." "Cyan" @@ -1981,8 +1794,7 @@ exit 0 # On an AMD GPU (no NVIDIA), surface the optional WSL-ROCm driver hint. if (-not $HasNvidiaSmi -and ($ROCmGfxArch -or $ROCmGpuLabel)) { Show-AmdWslDriverHint } - # Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL - # TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared. + # Trim trailing slashes from the URL PATH only, preserving ?query / #fragment (a whole-URL TrimEnd corrupts a token ending in "/"). Shared. function Trim-IndexPathSlashes { param([string]$Url) $value = $Url.Trim() @@ -1993,13 +1805,10 @@ exit 0 return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx) } - # ── Choose the correct PyTorch index URL based on driver CUDA version ── - # Mirrors Get-PytorchCudaTag in setup.ps1. + # ── Choose the PyTorch index URL based on driver CUDA version (mirrors Get-PytorchCudaTag in setup.ps1). ── function Get-TorchIndexUrl { $baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } - # Explicit pin -- skip ALL GPU probing (headless / CI / cross-install). - # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf appended - # to the mirror base. Matches install.sh / install_python_stack.py. + # Explicit pin skips ALL GPU probing: UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf appended to the mirror base. Matches install.sh / install_python_stack.py. if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) { return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL) } @@ -2009,9 +1818,7 @@ exit 0 if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" } try { $output = Invoke-NvidiaSmiBounded $NvidiaSmiExe - # Newer NVIDIA drivers (e.g. 610.x on Windows) print - # "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y". - # Accept both spellings so we don't fall through to the cu126 default. + # Newer NVIDIA drivers print "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y"; accept both so we don't fall through to the cu126 default. if ($output -match 'CUDA(?: UMD)? Version:\s+(\d+)\.(\d+)') { $major = [int]$Matches[1]; $minor = [int]$Matches[2] if ($major -ge 13) { return "$baseUrl/cu130" } @@ -2026,8 +1833,7 @@ exit 0 return "$baseUrl/cu126" } - # Strip userinfo AND query/fragment so an authenticated pin never leaks. Shared with - # _strip_index_url_credentials (install.sh / py / setup.ps1). + # Strip userinfo AND query/fragment so an authenticated pin never leaks. Shared with _strip_index_url_credentials (install.sh / py / setup.ps1). function Remove-IndexUrlCredentials { param([string]$Url) $sep = $Url.IndexOf('://') @@ -2045,9 +1851,7 @@ exit 0 return "${scheme}://${host_}" } - # ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ── - # torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu, - # matching setup.ps1's stale-venv parse. + # ── Torch flavor helpers (repair a stale CPU / wrong-CUDA wheel): torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu, matching setup.ps1's stale-venv parse. ── function ConvertTo-TorchFlavorTag { param([string]$TorchVersion) if (-not $TorchVersion) { return $null } @@ -2057,8 +1861,7 @@ exit 0 return 'cpu' } - # Expected tag from the index leaf: cuXXX / cpu / rocm ($ROCmIndexUrl or a - # gfx* leaf -> rocm). $null on an unknown leaf (odd mirror) so repair no-ops. + # Expected tag from the index leaf: cuXXX / cpu / rocm ($ROCmIndexUrl or a gfx* leaf -> rocm); $null on an unknown leaf so repair no-ops. function Get-ExpectedTorchFlavorTag { param([string]$TorchIndexUrl, [string]$ROCmIndexUrl) if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' } @@ -2073,8 +1876,7 @@ exit 0 return $null } - # Installed torch flavor tag in $PythonExe's venv, or $null if absent. Uses - # ProcessStartInfo (not &) so stderr doesn't trip $ErrorActionPreference. + # Installed torch flavor tag in $PythonExe's venv, or $null if absent. Uses ProcessStartInfo (not &) so stderr doesn't trip $ErrorActionPreference. function Get-InstalledTorchTag { param([string]$PythonExe) if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null } @@ -2087,12 +1889,7 @@ exit 0 $psi.UseShellExecute = $false $psi.CreateNoWindow = $true $proc = [System.Diagnostics.Process]::Start($psi) - # Drain BOTH streams async, then WaitForExit. A synchronous ReadToEnd() - # before the wait would block forever if a wedged "import torch" never - # closes stdout; leaving the redirected stderr undrained would deadlock a - # child that floods it past the pipe buffer. Async reads let a noisy-but- - # exiting probe finish, while a truly hung one still hits the 30s timeout - # and is killed -- bounded either way. + # Drain BOTH streams async before WaitForExit: a synchronous ReadToEnd() would block on a wedged "import torch", and an undrained stderr would deadlock a child flooding the pipe buffer. A truly hung probe still hits the 30s timeout. $outTask = $proc.StandardOutput.ReadToEndAsync() $errTask = $proc.StandardError.ReadToEndAsync() $finished = $proc.WaitForExit(30000) @@ -2104,19 +1901,12 @@ exit 0 } catch { return $null } } - # An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it - # (e.g. a deliberate cpu pin on an AMD host). + # An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it (e.g. a deliberate cpu pin on an AMD host). $TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or ` (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) $TorchIndexUrl = Get-TorchIndexUrl - # ── GPU arch → newest compatible Windows ROCm wheel release ── - # Wheels bundle their own ROCm runtime; the installed HIP SDK version does - # not constrain which release to use. Always picks the newest release that - # supports the GPU architecture. - # ── AMD Windows ROCm: arch-aware pip index (repo.amd.com) ── - # Wheels bundle their own ROCm runtime and support all Python versions. - # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. + # ── AMD Windows ROCm: arch-aware pip index (repo.amd.com). Wheels bundle their own ROCm runtime (HIP SDK version doesn't constrain the release) and support all Python versions; picks the newest release for the arch. Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. ── $ROCmIndexUrl = $null $ROCmTorchFloor = $null $PinnedRocmVisionSpec = $null @@ -2130,23 +1920,12 @@ exit 0 "gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all" "gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100 } - # gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in - # torch._C._grouped_mm on torch <2.11.0 (rocm7.12 and rocm7.1 respectively). - # TheRock issues #5284 and #3284. Force torch>=2.11.0 so pip never resolves - # to the broken 2.10.0 wheels even though they exist on the AMD index. - # The <2.12.0 ceiling matches the Linux install_python_stack.py constraint - # for the same arches: AMD actively publishes new versions on their index, - # so without a ceiling a future 2.12.0+rocmX.Y wheel would be pulled in - # automatically before it has been validated on these architectures. - # Bump the ceiling here (and in install_python_stack.py) when 2.12.x is - # confirmed working on gfx120X / Strix. + # gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) hit a null-pointer bug in torch._C._grouped_mm on torch <2.11.0 (TheRock #5284 / #3284); force torch>=2.11.0 so pip skips the broken 2.10.0 wheels on the AMD index. The <2.12.0 ceiling (matches install_python_stack.py) blocks an unvalidated future 2.12.0+rocm wheel; bump both when 2.12.x is confirmed on gfx120X / Strix. $torchFloorMap = @{ "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" } - # Companion ranges track the torch ceiling so pip resolves a consistent - # trio on AMD's per-arch index (each published independently). Mirrors - # setup.ps1 / install_python_stack.py; bump all three together for 2.12.x. + # Companion ranges track the torch ceiling so pip resolves a consistent trio on AMD's per-arch index (each published independently). Mirrors setup.ps1 / install_python_stack.py; bump all three together for 2.12.x. $torchvisionFloorMap = @{ "gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0" "gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0" @@ -2171,9 +1950,7 @@ exit 0 } } - # A gfx*/rocm pin skips the auto-reroute above, but the generic CPU/CUDA install below - # would use torch>=2.4,<2.11 and pull a known-bad wheel on the gfx115x/gfx120x/rocm>=7.2 - # indexes (the _grouped_mm bug). Route a pinned ROCm index through the ROCm path. + # A gfx*/rocm pin skips the auto-reroute above, but the generic CPU/CUDA install below (torch>=2.4,<2.11) would pull a known-bad wheel on the gfx115x/gfx120x/rocm>=7.2 indexes (the _grouped_mm bug). Route a pinned ROCm index through the ROCm path. if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) { $_pinLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLower() $_pinRocm211 = $false @@ -2191,8 +1968,7 @@ exit 0 $PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0" substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan" } elseif ($_pinLeaf -match '^gfx[0-9]' -or $_pinLeaf -match '^rocm[0-9]+(\.[0-9]+)?$') { - # Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with - # bare specs. Only EXACT rocm/gfx* are families; a suffixed leaf is verbatim. + # Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with bare specs. Only EXACT rocm/gfx* are families; a suffixed leaf is verbatim. $ROCmIndexUrl = $TorchIndexUrl } } @@ -2209,8 +1985,7 @@ exit 0 if (-not $SkipTorch -and -not $ROCmIndexUrl -and $TorchIndexUrl -like "*/cpu") { Write-Host "" if ($ROCmGfxArch) { - # Only an unmapped arch reaches here (a mapped one set $ROCmIndexUrl - # above). No ROCm torch wheels for this arch (e.g. RDNA2 gfx103X) -> CPU. + # Only an unmapped arch reaches here (a mapped one set $ROCmIndexUrl above): no ROCm torch wheels for this arch (e.g. RDNA2 gfx103X) -> CPU. substep "Installing CPU PyTorch -- no ROCm PyTorch wheels are available for $ROCmGfxArch." "Yellow" substep "PyTorch (training and Transformers inference) runs on CPU on this GPU." "Yellow" } else { @@ -2230,23 +2005,8 @@ exit 0 } # ── Install PyTorch first, then unsloth separately ── - # - # Why two steps? - # `uv pip install unsloth --torch-backend=cpu` on Windows resolves to - # unsloth==2024.8 (a pre-CLI release with no unsloth.exe) because the - # cpu-only solver cannot satisfy newer unsloth's dependencies. - # Installing torch first from the explicit CUDA index, then upgrading - # unsloth in a second step, avoids this solver dead-end. - # - # Why --upgrade-package instead of --upgrade? - # `--upgrade unsloth` re-resolves ALL dependencies including torch, - # pulling torch from default PyPI and stripping the +cuXXX suffix - # that step 1 installed (e.g. torch 2.5.1+cu124 -> 2.10.0 with no - # CUDA suffix). `--upgrade-package unsloth` upgrades ONLY unsloth - # to the latest version while preserving the already-pinned torch - # CUDA wheels. Missing dependencies (transformers, trl, peft, etc.) - # are still pulled in because they are new, not upgrades. - # + # Two steps because `uv pip install unsloth --torch-backend=cpu` on Windows resolves to the pre-CLI unsloth==2024.8 (no unsloth.exe); installing torch from the explicit index first avoids that solver dead-end. + # --upgrade-package (not --upgrade) so upgrading unsloth doesn't re-resolve torch from PyPI and strip the +cuXXX suffix step 1 pinned; new deps (transformers, trl, peft) are still pulled in. # ── Helper: find no-torch-runtime.txt ── function Find-NoTorchRuntimeFile { if ($StudioLocalInstall -and (Test-Path (Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt"))) { @@ -2259,18 +2019,14 @@ exit 0 } if ($_Migrated) { - # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving - # existing torch/CUDA unless the flavor repair below re-lands it. + # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving existing torch/CUDA unless the flavor repair below re-lands it. Write-TauriLog "STEP" "Installing unsloth" substep "upgrading unsloth in migrated environment..." if ($SkipTorch) { - # No-torch: install unsloth + unsloth-zoo with --no-deps, then - # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. + # No-torch: install unsloth + unsloth-zoo --no-deps, then runtime deps (typer, safetensors, transformers, etc.) --no-deps. $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } if ($baseInstallExit -eq 0) { - # Resolve pydantic WITH deps so pip pins pydantic-core - # to the matching version (no-torch-runtime.txt below - # is --no-deps). All transitive deps are torch-free. + # Resolve pydantic WITH deps so pip pins the matching pydantic-core (no-torch-runtime.txt below is --no-deps); all transitive deps are torch-free. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } } if ($baseInstallExit -eq 0) { @@ -2307,40 +2063,28 @@ exit 0 Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)" substep "installing PyTorch from $(Remove-IndexUrlCredentials $ROCmIndexUrl)..." $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } - # Pin the companions to match $torchSpec; bare names can resolve an - # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. + # Pin companions to match $torchSpec; bare names can resolve an ABI-incompatible torchvision/torchaudio on AMD's per-arch index. $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { - # Transient AMD-index failure: fall back to a CPU base (Unsloth setup retries - # ROCm). Use an explicit CPU index -- for a pinned ROCm index $TorchIndexUrl IS - # the ROCm mirror, so reusing it would just retry it. + # Transient AMD-index failure: fall back to a CPU base (Unsloth setup retries ROCm). Explicit CPU index -- for a pinned ROCm index $TorchIndexUrl IS the ROCm mirror, so reusing it would just retry it. $CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" } substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow" - # --force-reinstall: a failed ROCm install can leave an unpinned ROCm - # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU - # torch>= range, so without it uv would keep the ROCm build and only swap - # the companions -- a mismatched venv the flavor-repair block won't fix. + # --force-reinstall: a failed ROCm install can leave an unpinned ROCm torch that still satisfies the CPU torch>= range, so without it uv would keep the ROCm build and only swap companions -- a mismatched venv the flavor-repair block won't fix. $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $CpuFallbackIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) } - # CPU base is in; drop the ROCm expectation so the flavor-repair - # block below won't retry the just-failed index and abort. setup.ps1 - # reinstalls ROCm afterwards (recomputes its own index URL). + # CPU base is in; drop the ROCm expectation so the flavor-repair block below won't retry the just-failed index and abort. setup.ps1 reinstalls ROCm afterwards (recomputes its own index URL). $ROCmIndexUrl = $null $ROCmTorchFloor = $null } } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..." - # Bounded trio on every index (torchaudio 2.11 dropped its exact torch - # pin, so a bare companion can drift from a capped torch). cu - # families ship torch 2.11.x with paired triton-windows 3.6, so the - # ceiling widens to <2.12 there; other leaves keep the 2.10 line. - # Mirrors install.sh and _CUDA_TORCH_PKG_SPEC. + # Bounded trio on every index: torchaudio 2.11 dropped its torch pin, so a bare companion can drift from a capped torch. cu families ship torch 2.11.x (paired triton-windows 3.6) so the ceiling widens to <2.12; other leaves keep the 2.10 line. Mirrors install.sh and _CUDA_TORCH_PKG_SPEC. $_idxLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() if ($_idxLeaf -match '^cu[0-9]+$') { $_pinTorchSpec = "torch>=2.4,<2.12.0" @@ -2361,8 +2105,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($SkipTorch) { - # No-torch: install unsloth + unsloth-zoo with --no-deps, then - # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. + # No-torch: install unsloth + unsloth-zoo --no-deps, then runtime deps (typer, safetensors, transformers, etc.) --no-deps. $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. @@ -2429,24 +2172,16 @@ exit 0 } } - # ── Enforce the installed torch flavor matches the detected GPU build ── - # PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv - # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on - # "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is - # expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx* - # is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install - # above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops. + # ── Enforce the installed torch flavor matches the detected GPU build. PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv keeps a stale torch+cpu against a CUDA index and setup.ps1 loops on "cpu != required cuXXX". Reinstall the right triplet when a GPU build is expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (a PEP 503 index uv resolves via --default-index). --no-torch / CPU-only hosts are no-ops. ── if (-not $SkipTorch) { $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl if ($expectedTorchTag -and $expectedTorchTag -ne 'cpu') { $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { if ($expectedTorchTag -eq 'rocm' -and $ROCmIndexUrl) { - # AMD: a migrated venv can keep a stale CPU torch the fresh ROCm path - # would have force-reinstalled. Repair from the same repo.amd.com index. + # AMD: a migrated venv can keep a stale CPU torch the fresh ROCm path would have force-reinstalled. Repair from the same repo.amd.com index. $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } - # Pin companions like the fresh ROCm path (bare names can pull an - # ABI-incompatible torchvision/torchaudio from the per-arch index). + # Pin companions like the fresh ROCm path (bare names can pull an ABI-incompatible torchvision/torchaudio from the per-arch index). $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" @@ -2457,8 +2192,7 @@ exit 0 } $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython } elseif ($expectedTorchTag -ne 'rocm') { - # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. - # Same leaf-gated ceiling as the install above: cu* serves 2.11. + # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. Same leaf-gated ceiling as the install above: cu* serves 2.11. $_fixLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() if ($_fixLeaf -match '^cu[0-9]+$') { $_fixTorchSpec = "torch>=2.4,<2.12.0"; $_fixVisionSpec = "torchvision>=0.19,<0.27.0"; $_fixAudioSpec = "torchaudio>=2.4,<2.12.0" @@ -2484,11 +2218,7 @@ exit 0 } } - # Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped - # for --local: the editable install above already makes _PACKAGE_ROOT in - # unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__). - # Source paths match the Tauri bundle layout in studio/src-tauri/tauri.conf.json, - # which bundles install_python_stack.py at the bundle root next to install.ps1. + # Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped for --local (the editable install already resolves _PACKAGE_ROOT to the repo). Source paths match the Tauri bundle layout (tauri.conf.json bundles install_python_stack.py at the bundle root next to install.ps1). if ($TauriMode) { $rawPath = if ($PSCommandPath) { $PSCommandPath } else { $MyInvocation.ScriptName } if ($rawPath) { @@ -2500,8 +2230,7 @@ exit 0 foreach ($rel in $overlayMap.Keys) { $src = Join-Path $scriptDir $rel $dst = Join-Path $VenvDir $overlayMap[$rel] - # -LiteralPath: $VenvDir derives from $StudioHome which may - # contain [ ] * ? when the user overrode UNSLOTH_STUDIO_HOME. + # -LiteralPath: $VenvDir derives from $StudioHome which may contain [ ] * ? when the user overrode UNSLOTH_STUDIO_HOME. if (-not (Test-Path -LiteralPath $src)) { continue } $dstParent = Split-Path -Parent $dst if (-not (Test-Path -LiteralPath $dstParent)) { @@ -2529,11 +2258,7 @@ exit 0 } } - # ── Run studio setup ── - # setup.ps1 will handle installing Git, CMake, Visual Studio Build Tools, - # CUDA Toolkit, and other dependencies automatically via winget. Node.js is - # NOT installed via winget -- setup.ps1 uses an isolated Node it manages and - # never touches the system Node/npm. + # ── Run studio setup: setup.ps1 installs Git, CMake, VS Build Tools, CUDA Toolkit, etc. via winget. Node.js is NOT via winget -- setup.ps1 uses an isolated Node it manages and never touches system Node/npm. ── Write-TauriLog "STEP" "Running studio setup" step "setup" "running unsloth studio setup..." $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe" @@ -2551,8 +2276,7 @@ exit 0 $env:UNSLOTH_NO_TORCH = if ($SkipTorch) { "true" } else { "false" } # Tauri desktop app bundles its own frontend — skip Node/npm/frontend build $env:SKIP_STUDIO_FRONTEND = if ($TauriMode) { "1" } else { "0" } - # Always set STUDIO_LOCAL_INSTALL explicitly to avoid stale values from - # a previous --local run in the same PowerShell session. + # Always set STUDIO_LOCAL_INSTALL explicitly to avoid a stale value from a previous --local run in the same session. if ($StudioLocalInstall) { $env:STUDIO_LOCAL_INSTALL = "1" $env:STUDIO_LOCAL_REPO = $RepoRoot @@ -2560,11 +2284,7 @@ exit 0 $env:STUDIO_LOCAL_INSTALL = "0" Remove-Item Env:STUDIO_LOCAL_REPO -ErrorAction SilentlyContinue } - # Use 'studio setup' (not 'studio update') because 'update' pops - # SKIP_STUDIO_BASE, which would cause redundant package reinstallation - # and bypass the fast-path version check from PR #4667. - # Propagate UNSLOTH_STUDIO_HOME only for env-override installs; otherwise - # an inherited value would put llama.cpp in the wrong place. + # 'studio setup' (not 'update'): 'update' pops SKIP_STUDIO_BASE -> redundant reinstall + bypasses the PR #4667 fast-path version check. Propagate UNSLOTH_STUDIO_HOME only for env-override installs (else an inherited value misplaces llama.cpp). $previousUnslothStudioHome = $env:UNSLOTH_STUDIO_HOME $hadPreviousUnslothStudioHome = ($null -ne $previousUnslothStudioHome) if ($StudioRedirectMode -eq 'env') { @@ -2582,10 +2302,7 @@ exit 0 $env:UNSLOTH_LOCAL_LLAMA_CPP_DIR = (Resolve-Path -LiteralPath $WithLlamaCppDir).Path } $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1" - # Hand the venv interpreter to setup.ps1 so it reuses the Python we already - # resolved and built the venv with, instead of re-probing the system (which - # can trip over an unsupported `python` 3.14 or a Store stub on PATH even - # though the venv is fine). setup.ps1 Test-Path-guards this before use. + # Hand the venv interpreter to setup.ps1 so it reuses the Python we resolved instead of re-probing the system (which can trip over an unsupported `python` 3.14 or a Store stub on PATH). setup.ps1 Test-Path-guards this before use. $env:UNSLOTH_SETUP_PYTHON = Join-Path $VenvDir "Scripts\python.exe" try { & $UnslothExe @studioArgs @@ -2605,11 +2322,7 @@ exit 0 return (Exit-InstallFailure "unsloth studio setup failed (exit code $setupExit)" $setupExit) } - # ── Expose `unsloth` via a shim dir containing only unsloth.exe ── - # We do NOT add the venv Scripts dir to PATH (it also holds python.exe - # and pip.exe, which would hijack the user's system interpreter). - # Hardlink preferred; falls back to copy if cross-volume or non-NTFS. - # + # ── Expose `unsloth` via a shim dir containing only unsloth.exe (NOT the venv Scripts dir, which also holds python.exe/pip.exe and would hijack the system interpreter). Hardlink preferred, copy fallback if cross-volume/non-NTFS. ── # Remove the legacy venv Scripts PATH entry that older installers wrote. $LegacyScriptsDir = Join-Path $VenvDir "Scripts" try { @@ -2644,9 +2357,7 @@ exit 0 $ShimDir = Join-Path $StudioHome "bin" [System.IO.Directory]::CreateDirectory($ShimDir) | Out-Null $ShimExe = Join-Path $ShimDir "unsloth.exe" - # Fatal preflight outside the lock-handling try/catch -- a directory at - # the shim path must not be downgraded to "Continuing with the existing - # launcher", or the install finishes with no usable shim. + # Fatal preflight outside the lock-handling try/catch -- a directory at the shim path must not be downgraded to "Continuing with the existing launcher", or the install finishes with no usable shim. if (Test-Path -LiteralPath $ShimExe -PathType Container) { Write-Host "[ERROR] Cannot create unsloth launcher: $ShimExe is a directory." -ForegroundColor Red Write-Host " Move or remove it manually, then re-run the installer." -ForegroundColor Yellow @@ -2657,10 +2368,7 @@ exit 0 try { if (Test-Path -LiteralPath $ShimExe) { Remove-Item -LiteralPath $ShimExe -Force -ErrorAction Stop } try { - # New-Item -ItemType HardLink does NOT accept -LiteralPath in any - # PowerShell version, so use -Path. Wildcards in $ShimExe (e.g. - # brackets in custom roots) glob-expand here and fall through to - # the Copy-Item -LiteralPath fallback below. + # New-Item -ItemType HardLink doesn't accept -LiteralPath in any PowerShell version, so use -Path; wildcards in $ShimExe (brackets in custom roots) glob-expand here and fall through to the Copy-Item -LiteralPath fallback below. New-Item -ItemType HardLink -Path $ShimExe -Target $UnslothExe -ErrorAction Stop | Out-Null } catch { Copy-Item -LiteralPath $UnslothExe -Destination $ShimExe -Force -ErrorAction Stop # fallback: copy @@ -2678,8 +2386,7 @@ exit 0 Write-Host " Launch unsloth studio directly via '$UnslothExe' until the next successful install." -ForegroundColor Yellow } } - # Add to PATH only when launcher exists. Env-mode: session-only export, - # no registry change (workspace path may be deleted later). + # Add to PATH only when launcher exists. Env-mode: session-only export, no registry change (workspace path may be deleted later). $pathAdded = $false if (Test-Path -LiteralPath $ShimExe) { if ($StudioRedirectMode -ne 'env') { @@ -2692,8 +2399,7 @@ exit 0 Refresh-SessionPath # sync current session with registry Complete-StudioVenvRollback - # Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy - # User PATH entry (Machine > User > current $env:Path) would win. + # Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy User PATH entry (Machine > User > current $env:Path) would win. if ($StudioRedirectMode -eq 'env' -and (Test-Path -LiteralPath $ShimExe)) { $env:Path = "$ShimDir;$env:Path" step "path" "exported $ShimDir for this session (no registry PATH change in env-override mode)" @@ -2708,12 +2414,7 @@ exit 0 # New-StudioShortcuts gates the .lnk shortcuts on env-mode internally. New-StudioShortcuts -UnslothExePath $UnslothExe - # Warn if another 'unsloth' wins on PATH (different venv, system pip). - # Mirrors install.sh; absolute path is still the most reliable launch. - # Uses content-hash equality (Get-FileHash) so hardlinks, symlinks, and - # identical copies of the installer's shim don't false-trigger. CommandType - # Application restricts the probe to real executables (skips aliases, - # functions, scripts). + # Warn if another 'unsloth' wins on PATH (different venv, system pip). Content-hash equality (Get-FileHash) so hardlinks/symlinks/identical copies of the shim don't false-trigger; CommandType Application restricts the probe to real executables. try { $_pathCmd = Get-Command unsloth -CommandType Application -ErrorAction SilentlyContinue | Select-Object -First 1 if ($_pathCmd) { @@ -2735,9 +2436,7 @@ exit 0 # Diagnostic only; never block install on a probe failure. } - # In interactive terminals, ask the user before starting Unsloth unless the - # caller explicitly disabled the post-install prompt. - # In non-interactive environments (CI, Docker) just print instructions. + # Interactive terminals: prompt before starting Unsloth (unless the caller disabled it); non-interactive (CI, Docker): just print instructions. $IsInteractive = (-not $SkipAutostart) -and [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected) if ($IsInteractive) { Write-Host "" @@ -2753,8 +2452,7 @@ exit 0 } } else { step "launch" "manual commands:" - # Single-quote the printed paths so $-vars / backticks in custom roots - # do not reparse when the user pastes the command. + # Single-quote the printed paths so $-vars / backticks in custom roots don't reparse when the user pastes the command. $_actLiteral = "'" + ((Join-Path $VenvDir "Scripts\Activate.ps1") -replace "'", "''") + "'" if ($StudioRedirectMode -eq 'env') { # Env-mode skips registry PATH; print the absolute shim path. diff --git a/install.sh b/install.sh index c02552628f..7a4eb572e4 100755 --- a/install.sh +++ b/install.sh @@ -56,9 +56,7 @@ _SHORTCUTS_ONLY=false _next_is_package=false _next_is_python=false _next_is_llama_cpp_dir=false -# Seed from the environment so a caller who exports UNSLOTH_LOCAL_LLAMA_CPP_DIR -# (the documented piped-install style) is honored; the --with-llama-cpp-dir -# flag below overrides it when given. +# Seed from env (piped-install style); --with-llama-cpp-dir below overrides it. _WITH_LLAMA_CPP_DIR="${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" for arg in "$@"; do if [ "$_next_is_package" = true ]; then @@ -97,8 +95,7 @@ if [ "$_VERBOSE" = true ]; then export UNSLOTH_VERBOSE=1 fi -# Custom Unsloth roots are not supported with --tauri (desktop app still -# resolves ~/.unsloth/studio). Pass through if the override == legacy default. +# Custom Unsloth roots are unsupported with --tauri unless override == legacy default. if [ "$TAURI_MODE" = true ]; then _tauri_override_var="" _tauri_override="${UNSLOTH_STUDIO_HOME:-}" @@ -115,8 +112,7 @@ if [ "$TAURI_MODE" = true ]; then "~") _tauri_override="$HOME" ;; "~/"*) _tauri_override="$HOME/${_tauri_override#'~/'}" ;; esac - # Canonicalize both sides (CDPATH=, -P) so a CDPATH-set env or - # symlinked $HOME doesn't break the legacy-equality comparison. + # Canonicalize both sides so CDPATH / symlinked $HOME can't break equality. if [ -d "$_tauri_override" ]; then _tauri_override_abs=$(CDPATH= cd -P -- "$_tauri_override" 2>/dev/null && pwd -P) \ || _tauri_override_abs="$_tauri_override" @@ -159,8 +155,7 @@ run_maybe_quiet() { fi } -# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL -# strip corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared. +# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment. Shared. _trim_index_path_slashes() { _tips_v="$1" case "$_tips_v" in @@ -179,9 +174,7 @@ _trim_index_path_slashes() { printf '%s%s' "$_tips_head" "$_tips_tail" } -# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer -# output before printing on failure; uv/pip errors echo the failing --index-url verbatim. -# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted. +# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer output. _redact_install_output() { sed -E \ -e 's#(https?://)[^/@[:space:]`]+@#\1@#g' \ @@ -193,19 +186,12 @@ _redact_install_output() { run_install_cmd() { _label="$1" shift - # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898): - # for --default-index, neutralize the uv index/backend/config vars (UV_TORCH_BACKEND - # redirects torch; UV_NO_CONFIG=1 + dropping UV_CONFIG_FILE stops a uv.toml/pyproject - # index outranking the CLI pin, uv 0.10). + # For --default-index, neutralize inherited uv index/backend/config vars so a uv.toml/pyproject index can't outrank the CLI pin. case " $* " in *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL -u UV_TORCH_BACKEND -u UV_FIND_LINKS -u UV_CONFIG_FILE UV_NO_CONFIG=1 "$@" ;; esac if _is_verbose; then - # Stream through the redactor: uv echoes index URLs (credentials and - # all) in its errors, and verbose mode previously bypassed the - # redaction the quiet path applies. The rc file preserves the - # command's exit code across the pipe without relying on pipefail - # (this script runs under plain sh). + # Stream through the redactor; rc file carries the exit code across the pipe (no pipefail in plain sh). _rcf=$(mktemp) { "$@" 2>&1; printf '%s' "$?" > "$_rcf"; } | _redact_install_output _rc=$(cat "$_rcf" 2>/dev/null || echo 1) @@ -223,16 +209,12 @@ run_install_cmd() { return $_rc } -# Retry run_install_cmd on transient uv download failures with backoff. Returns -# the last exit code on permanent failure so the set -e rollback trap still fires. +# Retry run_install_cmd with backoff; returns the last exit code so the set -e rollback trap still fires. : "${UNSLOTH_INSTALL_RETRIES:=3}" : "${UNSLOTH_INSTALL_RETRY_DELAY:=3}" run_install_cmd_retry() { _ricr_label="$1" - # Sanitize overrides to a default of 3 (a typo must not disable retries; =1 disables). - # Length guard precedes the numeric test so a huge value can't overflow `[ -ge ]`. - # 0?* rejects leading-zero delays ("08"/"09" break the later $((delay*2)) as octal); - # bare "0" stays valid. Bounds: 1..100 retries, 0..3600s base delay. + # Sanitize to default 3; length guard before `[ -ge ]`, 0?* rejects leading-zero (octal) delays. Bounds: 1..100 retries, 0..3600s delay. case "$UNSLOTH_INSTALL_RETRIES" in ''|*[!0-9]*|0) _ricr_max=3 ;; *) if [ "${#UNSLOTH_INSTALL_RETRIES}" -le 3 ] && [ "$UNSLOTH_INSTALL_RETRIES" -ge 1 ] 2>/dev/null && [ "$UNSLOTH_INSTALL_RETRIES" -le 100 ] 2>/dev/null; then _ricr_max=$UNSLOTH_INSTALL_RETRIES; else _ricr_max=3; fi ;; @@ -243,8 +225,7 @@ run_install_cmd_retry() { esac _ricr_attempt=1 while :; do - # AND-OR (not `if`) preserves the real failure code: $? after a non-taken - # `if` is 0 in sh/dash/bash, which would break the rollback path. + # AND-OR (not `if`) preserves the real failure code for the rollback path. run_install_cmd "$@" && return 0 _ricr_rc=$? if [ "$_ricr_attempt" -ge "$_ricr_max" ]; then @@ -257,10 +238,7 @@ run_install_cmd_retry() { done } -# Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main -# wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2 -# NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the -# pre-release URL is unreachable. Drop the pin once bnb 0.50+ ships on PyPI. +# bitsandbytes on AMD ROCm: continuous-release_main wheel (ROCm 4-bit GEMV fix), PyPI >=0.49.1 fallback; drop once bnb 0.50+ ships on PyPI. _install_bnb_rocm() { _label="$1" _venv_py="$2" @@ -275,9 +253,7 @@ _install_bnb_rocm() { _bnb_whl_url="" ;; esac - # uv rejects the continuous-release_main bitsandbytes wheel because the - # filename version (1.33.7rc0) does not match the embedded metadata version - # (0.50.0.dev0). pip accepts the mismatch, so bootstrap pip and use it. + # uv rejects the wheel's filename/metadata version mismatch; pip accepts it, so bootstrap pip. if ! "$_venv_py" -m pip --version >/dev/null 2>&1; then if ! run_maybe_quiet "$_venv_py" -m ensurepip --upgrade; then run_maybe_quiet uv pip install --python "$_venv_py" pip || \ @@ -320,8 +296,7 @@ if [ "$_next_is_llama_cpp_dir" = true ]; then exit 1 fi -# Validate --package to prevent injection into shell/Python commands. -# Must start with a letter/digit (rejects leading dashes that uv would parse as flags). +# Validate --package (injection guard); must start with a letter/digit so uv can't parse it as a flag. case "$PACKAGE_NAME" in [!a-zA-Z0-9]*) echo "❌ ERROR: --package name must start with a letter or digit." >&2 @@ -350,8 +325,7 @@ _tauri_torch_index_family() { return fi _diag_url="${1:-}" - # Strip query/fragment AND a trailing slash before classifying (like _torch_index_url_leaf): - # a token isn't echoed into [TAURI:DIAG], and .../cu128/?token=x still classifies as cu128. + # Strip query/fragment and trailing slash before classifying (like _torch_index_url_leaf). _diag_url="${_diag_url%%\?*}" _diag_url="${_diag_url%%#*}" _diag_url="${_diag_url%/}" @@ -368,8 +342,7 @@ _tauri_torch_index_family() { rocm[0-9]*.[0-9]*) echo "$_diag_family" ;; *) echo "auto" ;; esac ;; - # AMD arch-specific index (e.g. repo.amd.com/rocm/whl/gfx1151/) -- - # used for Strix Halo/Point where torch 2.11+rocm7.13 has the real fix. + # AMD arch-specific index (Strix Halo/Point; torch 2.11+rocm7.13 has the real fix). *repo.amd.com/rocm/whl/gfx*|*rocm/whl/gfx*) echo "rocm7.13" ;; "") echo "none" ;; *) echo "auto" ;; @@ -388,7 +361,7 @@ _tauri_gpu_branch() { return fi case "$_diag_family" in - # Require a digit after cu so /current or /custom isn't branded CUDA (parity ^cu[0-9]). + # Require a digit after cu so /current or /custom isn't branded CUDA. cu[0-9]*) echo "cuda" ;; rocm*) if [ "$_diag_radeon" = true ]; then @@ -405,9 +378,7 @@ _tauri_gpu_branch() { PYTHON_VERSION="" # resolved after platform detection -# Resolve install destinations: env override, HOME-redirect (best-effort -# via getent/dscl), or default. Env-var priority: UNSLOTH_STUDIO_HOME wins -# over STUDIO_HOME (the more specific signal beats the generic alias). +# Resolve install destinations; UNSLOTH_STUDIO_HOME wins over the STUDIO_HOME alias. _resolve_studio_destinations() { _override_var="" _override="${UNSLOTH_STUDIO_HOME:-}" @@ -417,10 +388,9 @@ _resolve_studio_destinations() { _override="${STUDIO_HOME:-}" [ -n "$_override" ] && _override_var="STUDIO_HOME" fi - # Strip surrounding whitespace so " " is treated as unset (matches the - # Python resolvers' .strip()), preventing install/runtime layout drift. + # Strip surrounding whitespace so " " is treated as unset (matches Python .strip()). _override=$(printf '%s' "$_override" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') - # Tilde expansion: env vars are not subject to it when quoted on assignment. + # Tilde expansion: quoted env vars aren't subject to it on assignment. case "$_override" in "~") _override="$HOME" ;; "~/"*) _override="$HOME/${_override#'~/'}" ;; @@ -441,8 +411,7 @@ _resolve_studio_destinations() { elif [ "$(uname)" = "Darwin" ] && command -v dscl >/dev/null 2>&1; then _default_home=$(dscl . -read "/Users/${USER:-$(whoami)}" NFSHomeDirectory 2>/dev/null | awk '{print $2}') fi - # Canonicalize both sides so a trailing slash on $HOME (or symlink mismatch - # with passwd-DB output) doesn't misfire the redirection branch. + # Canonicalize both sides so a trailing slash / symlink mismatch doesn't misfire redirection. _home_canon="$HOME" if [ -d "$_home_canon" ]; then _home_canon=$(CDPATH= cd -P -- "$_home_canon" 2>/dev/null && pwd -P) || _home_canon="$HOME" @@ -521,8 +490,7 @@ _on_install_exit() { [ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true exit "$_status" } -# Empty so an inherited value never reaches the trap's rm; only temp paths this -# script creates below (spaced-path dir, torch-trio overrides) are removed. +# Empty so an inherited value never reaches the trap's rm. _UV_OVERRIDE_TMPDIR="" _UNSLOTH_TORCH_OVERRIDES="" trap _on_install_exit EXIT @@ -616,10 +584,7 @@ _smart_apt_install() { } # ── Helper: create desktop shortcuts and launcher script ── -# Usage: create_studio_shortcuts -# Creates ~/.local/share/unsloth/launch-studio.sh (shared launcher), -# plus platform-specific shortcuts (Linux .desktop / macOS .app bundle / -# WSL Windows Desktop+Start Menu .lnk). +# Usage: create_studio_shortcuts -- writes launch-studio.sh + platform shortcuts. create_studio_shortcuts() { _css_exe="$1" _css_os="$2" @@ -641,17 +606,7 @@ create_studio_shortcuts() { mkdir -p "$_css_data_dir" - # Same-install discriminator: per-install opaque id written once at install - # time and read by both this launcher and the backend (/api/health). Replaces - # the older sha256(canonical $STUDIO_HOME) scheme to (a) avoid leaking the - # install path on -H 0.0.0.0 deployments and (b) sidestep launcher/backend - # canonicalization drift (cd -P vs Path.resolve() symlink/junction handling). - # Lives at $STUDIO_HOME/share/ (not $DATA_DIR) so the backend can find it - # via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" regardless of - # mode (in env-mode $STUDIO_HOME/share == $DATA_DIR; in default mode they - # diverge but the backend only knows the studio_root). 32 bytes of urandom - # -> 64 hex chars, byte-compatible with the prior digest so launcher - # placeholder, _check_health, and tests stay length-agnostic. + # Same-install discriminator: per-install opaque id read by launcher + backend (/api/health); lives at $STUDIO_HOME/share/ so the backend finds it via studio_root. _css_id_dir="$STUDIO_HOME/share" mkdir -p "$_css_id_dir" _css_id_file="$_css_id_dir/studio_install_id" @@ -682,8 +637,7 @@ create_studio_shortcuts() { [ "$_STUDIO_HOME_REDIRECT" = "env" ] && _css_is_env_mode=true # ── Write launcher script ── - # Single-quoted heredoc; @@DATA_DIR@@, @@STUDIO_ROOT_ID@@, and - # @@INSTALLED_IS_ENV_MODE@@ are substituted via sed below. + # Single-quoted heredoc; @@ placeholders substituted via sed below. cat > "$_css_launcher" << 'LAUNCHER_EOF' #!/usr/bin/env bash # Unsloth Studio Launcher @@ -1042,20 +996,15 @@ else fi LAUNCHER_EOF - # why: bake non-user-controlled placeholders FIRST so a literal - # `@@STUDIO_ROOT_ID@@` inside $DATA_DIR cannot be rewritten below. + # Bake non-user-controlled placeholders FIRST so a literal @@STUDIO_ROOT_ID@@ in $DATA_DIR can't be rewritten below. sed -e "s|@@STUDIO_ROOT_ID@@|$_css_studio_root_id|g" \ -e "s|@@INSTALLED_IS_ENV_MODE@@|$_css_is_env_mode|g" \ "$_css_launcher" > "$_css_launcher.tmp" \ && mv "$_css_launcher.tmp" "$_css_launcher" - # Env-mode bakes an absolute DATA_DIR (root fixed at install time); - # default / HOME-redirect keeps the literal $HOME/.local/share/unsloth - # so behavior is byte-identical to pre-override. + # Env-mode bakes an absolute DATA_DIR; default / HOME-redirect keeps the literal $HOME/.local/share/unsloth. if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then - # Two-stage escape: (1) `'` -> `'\''` for shell single-quote embedding, - # (2) backslash/&/| escape so the value survives the s|...|VALUE| sed - # below. Verified end-to-end with apostrophes, spaces, &, |, $. + # Two-stage escape: single-quote embedding, then backslash/&/| for the sed below. _sq_escaped=$(printf '%s' "$DATA_DIR" | sed "s/'/'\\\\''/g") _sed_safe=$(printf '%s' "$_sq_escaped" | sed 's/[\\&|]/\\&/g') sed "s|@@DATA_DIR@@|$_sed_safe|g" "$_css_launcher" > "$_css_launcher.tmp" \ @@ -1068,16 +1017,12 @@ LAUNCHER_EOF chmod +x "$_css_launcher" - # studio.conf: exe path + (env-mode only) persisted env vars so fresh - # shells launch the right install without re-exporting. + # studio.conf: exe path + (env-mode only) persisted env vars for fresh shells. _css_quoted_exe=$(printf '%s' "$_css_exe" | sed "s/'/'\\\\''/g") { printf '%s\n' "UNSLOTH_EXE='$_css_quoted_exe'" if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then - # When an override resolves to the legacy default, llama.cpp - # still lives at ~/.unsloth/llama.cpp (one shared build). - # Canonicalize the legacy side so a symlinked $HOME doesn't - # break the comparison. + # An override resolving to the legacy default shares ~/.unsloth/llama.cpp; canonicalize the legacy side. _css_legacy_studio="$HOME/.unsloth/studio" if [ -d "$_css_legacy_studio" ]; then _css_legacy_studio=$(CDPATH= cd -P -- "$_css_legacy_studio" 2>/dev/null && pwd -P) \ @@ -1091,8 +1036,7 @@ LAUNCHER_EOF _css_quoted_home=$(printf '%s' "$STUDIO_HOME" | sed "s/'/'\\\\''/g") _css_quoted_llama=$(printf '%s' "$_css_llama_path" | sed "s/'/'\\\\''/g") printf '%s\n' "export UNSLOTH_STUDIO_HOME='$_css_quoted_home'" - # UNSLOTH_LLAMA_CPP_PATH is a pre-existing user-controlled - # llama.cpp dir override; only default it if unset. + # UNSLOTH_LLAMA_CPP_PATH is user-controlled; only default it if unset. printf '%s\n' 'if [ -z "${UNSLOTH_LLAMA_CPP_PATH:-}" ]; then' printf '%s\n' " export UNSLOTH_LLAMA_CPP_PATH='$_css_quoted_llama'" printf '%s\n' 'fi' @@ -1100,7 +1044,6 @@ LAUNCHER_EOF } > "$_css_data_dir/studio.conf" # ── Icon: try bundled, then download ── - # rounded-512.png used for both Linux and macOS icons _css_script_dir="" if [ -n "${0:-}" ] && [ -f "$0" ]; then _css_script_dir=$(cd "$(dirname "$0")" 2>/dev/null && pwd) || true @@ -1143,9 +1086,7 @@ LAUNCHER_EOF fi # ── Platform-specific shortcuts ── - # Env-mode installs are workspace-scoped: skip persistent desktop / - # Start-Menu / dock launchers that may point at a deleted workspace. - # Runtime launcher + studio.conf + icon are still written above. + # Env-mode installs are workspace-scoped: skip persistent launchers that may point at a deleted workspace. if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then substep "wrote launcher at $_css_launcher (persistent shortcuts skipped in env-override mode)" return 0 @@ -1239,9 +1180,7 @@ DESKTOP_EOF PLIST_EOF - # Executable stub: same single-quoted-heredoc + sed-substitute - # pattern as launch-studio.sh so $-vars in $_css_data_dir don't - # expand at .app launch time. + # Executable stub: single-quoted heredoc + sed so $-vars in $_css_data_dir don't expand at launch. _css_sq_dir=$(printf '%s' "$_css_data_dir" | sed "s/'/'\\\\''/g") _css_sed_dir=$(printf '%s' "$_css_sq_dir" | sed 's/[\\&|]/\\&/g') cat > "$_css_macos_dir/launch-studio" << 'STUB_EOF' @@ -1287,13 +1226,9 @@ STUB_EOF elif [ "$_css_os" = "wsl" ]; then # ── WSL: create Windows Desktop and Start Menu shortcuts ── - # Detect current WSL distro for targeted shortcut _css_distro="${WSL_DISTRO_NAME:-}" - # Build the wsl.exe arguments. - # Double-quote distro name and launcher path for Windows command line - # parsing so values with spaces (e.g. "Ubuntu Preview") are kept as - # single arguments. + # Build wsl.exe args; double-quote so spaced values ("Ubuntu Preview") stay single args. _css_wsl_args="" if [ -n "$_css_distro" ]; then _css_wsl_args="-d \"$_css_distro\" " @@ -1317,8 +1252,7 @@ STUB_EOF # Escape single quotes for PowerShell single-quoted string embedding _css_sc_args_ps=$(printf '%s' "$_css_sc_args" | sed "s/'/''/g") - # DISTINCT shortcut name so the WSL launcher never clobbers a native - # install's "Unsloth Studio.lnk" in the same folder. Per-distro suffix. + # DISTINCT per-distro shortcut name so the WSL launcher never clobbers a native "Unsloth Studio.lnk". if [ -n "$_css_distro" ]; then _css_lnk_name="Unsloth Studio (WSL - ${_css_distro}).lnk" else @@ -1415,8 +1349,7 @@ WSLPS1_EOF fi rm -f "$_css_ps1_tmp" fi - # If WSL interop is disabled (powershell.exe "Exec format error"), the - # shortcut wasn't created; tell the user how to launch / re-enable it. + # WSL interop disabled ("Exec format error"): no shortcut; tell the user how. if [ "$_css_created" -ne 1 ]; then substep "Couldn't create the Windows shortcut (WSL interop may be disabled)." "$C_WARN" substep " Launch Unsloth from Windows: wsl -d \"$_css_distro\" -- bash -lc 'unsloth studio'" "$C_WARN" @@ -1462,8 +1395,7 @@ fi _ARCH=$(uname -m) MAC_INTEL=false if [ "$OS" = "macos" ] && [ "$_ARCH" = "x86_64" ]; then - # Guard against Apple Silicon running under Rosetta (reports x86_64). - # sysctl hw.optional.arm64 returns "1" on Apple Silicon even in Rosetta. + # Apple Silicon under Rosetta reports x86_64; hw.optional.arm64 stays "1". if [ "$(sysctl -in hw.optional.arm64 2>/dev/null || echo 0)" = "1" ]; then echo "" echo " WARNING: Apple Silicon detected, but this shell is running under Rosetta (x86_64)." @@ -1499,9 +1431,7 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then SKIP_TORCH=true fi -# Apple Silicon: exclude broken mlx-lm 0.31.3 (QK-norm load regression for -# gemma4 / qwen3_5; mlx-lm #1242). A curl-piped install has no overrides file -# and skips the guarded MLX step (SKIP_STUDIO_BASE=1), so this is the only cover. +# Apple Silicon: exclude broken mlx-lm 0.31.3 (QK-norm load regression); a curl-piped install has no overrides file, so this is the only cover. _MLX_LM_EXCLUDE_ARG="" # Apple Silicon: override mlx-vlm / mlx-lm's transformers pin (see overrides file). @@ -1509,8 +1439,7 @@ if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then _MLX_LM_EXCLUDE_ARG="mlx-lm!=0.31.3" _OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt" if [ -f "$_OVERRIDES_FILE" ]; then - # uv splits UV_OVERRIDE on whitespace, so a repo path with whitespace - # truncates it and aborts every later uv call (issue #6503). Hand uv a copy. + # uv splits UV_OVERRIDE on whitespace; hand uv a copy in a whitespace-free temp dir. case "$_OVERRIDES_FILE" in *[[:space:]]*) _UV_OVERRIDE_TMPDIR=$(mktemp -d 2>/dev/null) || _UV_OVERRIDE_TMPDIR="" @@ -1540,9 +1469,7 @@ elif [ "$OS" = "macos" ]; then fi tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none" -# AMD GPU name from the Windows host via WMI, or empty. Discrete cards aren't in -# /proc/cpuinfo, so ask Windows. Cached ("-" = negative), self-contained, bounded -# to 10s. Defined here so the reroute below can use it before _run_bounded exists. +# AMD GPU name from the Windows host via WMI (discrete cards aren't in /proc/cpuinfo); cached ("-" = negative), bounded to 10s. _WSL_AMD_GPU_NAME_CACHE="" _wsl_amd_gpu_name() { if [ -n "$_WSL_AMD_GPU_NAME_CACHE" ]; then @@ -1561,11 +1488,7 @@ _wsl_amd_gpu_name() { } # ── Bounded command runner ── -# Runs a command under a 10s timeout when the `timeout` binary is available, -# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during -# driver init or after a reset) from hanging the installer: a timed-out probe -# exits nonzero and is treated exactly like a failed probe. No-op semantics on -# hosts without `timeout` (e.g. macOS) or when the probe is healthy. +# 10s timeout when `timeout` exists, so a wedged nvidia-smi can't hang the installer. _run_bounded() { if command -v timeout >/dev/null 2>&1; then timeout 10 "$@" @@ -1574,10 +1497,7 @@ _run_bounded() { fi } -# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every -# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to -# the AMD card). Unset means all devices visible. nvidia-smi ignores this env -# var, so the probes below cannot see the distinction on their own. +# True when CUDA_VISIBLE_DEVICES is "" or "-1" (every NVIDIA device deliberately hidden); nvidia-smi ignores it. _cvd_hides_nvidia() { [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') @@ -1585,13 +1505,7 @@ _cvd_hides_nvidia() { } # ── NVIDIA usable-GPU helper ── -# Returns 0 (true) if an NVIDIA GPU is present and usable. -# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, -# which the NVIDIA driver populates on Linux regardless of nvidia-smi state -# -- handles PATH gaps, subprocess timeouts, and driver init races that -# could otherwise cause nvidia-smi to fail and silence NVIDIA detection. -# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches -# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. +# nvidia-smi -L primary, /proc/driver/nvidia/gpus/ sysfs fallback; a GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable. _has_usable_nvidia_gpu() { if _cvd_hides_nvidia; then return 1 @@ -1607,7 +1521,7 @@ _has_usable_nvidia_gpu() { return 0 fi fi - # Fallback: NVIDIA driver exposes one subdir per GPU under this path. + # Fallback: one subdir per GPU under this path. if [ -d /proc/driver/nvidia/gpus ] && \ [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then return 0 @@ -1615,27 +1529,19 @@ _has_usable_nvidia_gpu() { return 1 } -# Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04) -# with a 24.04 distro present, re-run the install there and stop; else fall through -# to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before -# the STUDIO_HOME mkdir/venv so the origin distro is untouched. +# Strix Halo ROCm-on-WSL only targets Ubuntu 24.04: re-run the install in an installed 24.04 distro, else fall through to CPU (never auto-create a distro). _maybe_reroute_strixhalo_to_2404() { [ "${OS:-}" = "wsl" ] || return 0 - # An explicit index pin skips every GPU-driven reroute (same contract as - # the later Radeon/Strix guard): the pin is honored in THIS distro rather - # than probing the GPU and switching distributions. Whitespace-only - # overrides do not gate (parity with get_torch_index_url). + # An explicit index pin skips every GPU-driven reroute; whitespace-only overrides don't gate. _rr_pin=$(printf '%s' "${UNSLOTH_TORCH_INDEX_URL:-}${UNSLOTH_TORCH_INDEX_FAMILY:-}" | tr -d '[:space:]') [ -n "$_rr_pin" ] && return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 [ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0 [ -e /dev/dxg ] || return 0 - # A usable NVIDIA GPU (common on hybrid AMD+NVIDIA hosts) means the CUDA path works on - # this distro, so don't reroute for AMD. _has_usable_nvidia_gpu (moved above) honors - # CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps. + # A usable NVIDIA GPU means the CUDA path works here, so don't reroute for AMD. if _has_usable_nvidia_gpu; then return 0; fi - # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes. + # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ && ! _wsl_amd_gpu_name >/dev/null 2>&1; then return 0 @@ -1646,18 +1552,11 @@ _maybe_reroute_strixhalo_to_2404() { fi _rr_ver="" [ -r /etc/os-release ] && _rr_ver=$(. /etc/os-release 2>/dev/null; printf '%s' "${VERSION_ID:-}") - # The bootstrap (scripts/install_rocm_wsl_strixhalo.sh) dies on any VERSION_ID but - # 24.04 and pins the noble repo, so 24.04 is the sole GPU-supported target; leave a - # 24.04 user alone. (Working ROCm on other versions was caught by librocdxg above.) + # The bootstrap only supports 24.04, so leave a 24.04 user alone. case "$_rr_ver" in 24.04) return 0 ;; esac - # Distro is now unsupported. If we can't reroute to a 24.04 target, stay CPU-only - # AND skip the later origin-distro ROCm bootstrap (it ignores distro version, so it - # would otherwise install ROCm into 26.04 etc.). + # Without a 24.04 reroute target, stay CPU-only AND skip the origin-distro ROCm bootstrap. command -v wsl.exe >/dev/null 2>&1 || { UNSLOTH_SKIP_ROCM_WSL_SETUP=1; return 0; } - # Route only to an installed Ubuntu-24.04 (bootstrap's only target). Match the whole - # line (one distro per line from wsl.exe -l -q), not a substring, so "Ubuntu-24.04-test" - # can't masquerade as it and then fail `wsl -d`. - # || true: no match is expected, not an error (script runs under set -e). + # Whole-line match so "Ubuntu-24.04-test" can't masquerade; || true: no match is fine. _rr_distros=$(wsl.exe -l -q 2>/dev/null | tr -d '\000\r') _rr_target=$(printf '%s\n' "$_rr_distros" | grep -ixF "Ubuntu-24.04" | head -n1) || true [ -n "$_rr_target" ] || { @@ -1670,26 +1569,21 @@ _maybe_reroute_strixhalo_to_2404() { echo "" substep "ROCm-on-WSL (GPU) needs Ubuntu 24.04; this distro is Ubuntu ${_rr_ver:-unknown}." "$C_WARN" substep "Found an existing $_rr_target distro -- continuing the GPU install there." "$C_OK" - # A --local checkout can't be replayed via curl|sh (the repo isn't in the target - # distro), so tell the user to re-run there rather than silently run a different install. + # A --local checkout can't be replayed via curl|sh, so tell the user to re-run there. if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "This is a --local install; re-run it from $_rr_target instead:" "$C_WARN" substep " wsl -d $_rr_target -- bash -lc 'cd && ./install.sh --local'" "$C_WARN" substep "Continuing CPU-only in Ubuntu ${_rr_ver:-this distro} for now." "$C_WARN" - # Unsupported distro, can't reroute a --local checkout: skip the origin ROCm bootstrap. UNSLOTH_SKIP_ROCM_WSL_SETUP=1 return 0 fi - # Forward the caller's options/env (custom package/python/home) so the rerouted - # install matches what was asked for, not a default install. + # Forward the caller's options/env so the rerouted install matches what was asked for. _rr_q() { printf "'%s'" "$(printf '%s' "$1" | sed "s/'/'\\\\''/g")"; } _rr_exports="set -o pipefail; export UNSLOTH_WSL_REROUTED=1" [ "$_STUDIO_HOME_REDIRECT" = "env" ] && _rr_exports="$_rr_exports; export UNSLOTH_STUDIO_HOME=$(_rr_q "$STUDIO_HOME")" - # Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the - # GPU instead of falling back to the desktop-app prompt path. + # Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the GPU. [ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1" - # Forward a pinned torch index into the rerouted distro; dropping it would - # silently revert the child install to auto-detection. + # Forward a pinned torch index; dropping it would revert the child to auto-detection. [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_URL=$(_rr_q "$UNSLOTH_TORCH_INDEX_URL")" [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_FAMILY=$(_rr_q "$UNSLOTH_TORCH_INDEX_FAMILY")" [ "$_SKIP_AUTOSTART" = true ] && _rr_exports="$_rr_exports; export UNSLOTH_SKIP_AUTOSTART=1" @@ -1705,34 +1599,27 @@ _maybe_reroute_strixhalo_to_2404() { else _rr_cmd="curl -fsSL https://unsloth.ai/install.sh | sh" fi - # pipefail so a failed curl in `curl | sh` isn't masked by sh exiting 0 on empty - # input (which would wrongly report success and exit 0 the parent installer). + # pipefail so a failed curl in `curl | sh` isn't masked by sh exiting 0 on empty input. _rr_rc=0 wsl.exe -d "$_rr_target" -- bash -lc "$_rr_exports; $_rr_cmd" || _rr_rc=$? if [ "$_rr_rc" -eq 0 ]; then exit 0 fi - # In Tauri mode the child uses exit 2 ([TAURI:NEED_SUDO]) to ask the desktop app to - # elevate for the target distro; the child already printed the NEED_SUDO line, so - # propagate the code instead of masking it as a reroute failure and dropping to CPU. + # Tauri child exit 2 ([TAURI:NEED_SUDO]) asks the desktop app to elevate; propagate it. if [ "$TAURI_MODE" = true ] && [ "$_rr_rc" -eq 2 ]; then exit 2 fi substep "Could not auto-continue in $_rr_target; run it yourself:" "$C_WARN" substep " wsl -d $_rr_target -- bash -lc 'curl -fsSL https://unsloth.ai/install.sh | sh'" substep "Continuing CPU-only in Ubuntu ${_rr_ver:-this distro} for now." "$C_WARN" - # Reroute failed; don't let the later bootstrap install ROCm into this unsupported - # distro -- stay CPU-only. + # Reroute failed; keep the later bootstrap from installing ROCm into this distro. UNSLOTH_SKIP_ROCM_WSL_SETUP=1 return 0 } _maybe_reroute_strixhalo_to_2404 || true # ── Check system dependencies ── -# cmake/git are only needed to *build* llama.cpp from source. Unsloth downloads a -# prebuilt by default, and setup.sh self-skips the source build when they're -# absent -- so macOS doesn't block on cmake (requiring it would force a manual -# Homebrew install). Linux keeps requiring them; its package manager has them. +# cmake/git only needed to build llama.cpp from source; macOS doesn't block on cmake (prebuilt default). tauri_log "STEP" "Checking system dependencies" case "$OS" in @@ -1746,8 +1633,7 @@ case "$OS" in echo " After the installation completes, please re-run this script." exit 1 fi - # cmake is only needed for a source build; the default prebuilt path - # doesn't use it, so its absence is not fatal -- no Homebrew prerequisite. + # cmake absence is not fatal (only the source build uses it). if command -v cmake >/dev/null 2>&1; then step "deps" "all system dependencies found" else @@ -1797,25 +1683,17 @@ esac tauri_log "STEP" "Installing uv package manager" UV_MIN_VERSION="0.8.16" -# When bytecode compilation is enabled, large installs can exceed uv's 60s default on slow machines. Default to 180s, preserving overrides ("0" disables). +# Large bytecode-compiled installs can exceed uv's 60s default; use 180s ("0" disables). : "${UV_COMPILE_BYTECODE_TIMEOUT:=180}" export UV_COMPILE_BYTECODE_TIMEOUT -# uv >= 0.8.16 retries HTTP/2 streaming body errors; raise retries and read -# timeout for large wheel downloads. ":=" preserves any user override. +# Raise retries and read timeout for large wheel downloads (":=" keeps overrides). : "${UV_HTTP_RETRIES:=5}" export UV_HTTP_RETRIES : "${UV_HTTP_TIMEOUT:=180}" export UV_HTTP_TIMEOUT -# macOS: trust the system Keychain so uv uses SecureTransport instead of rustls. -# Required behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.) which -# present their own CA certificate. rustls (uv's default) ignores the Keychain -# and rejects intercepted connections with "invalid peer certificate: UnknownIssuer". -# Set both vars: UV_SYSTEM_CERTS is the modern one (uv >= 0.11), UV_NATIVE_TLS the -# legacy one understood by uv 0.8.16-0.10.x, which the installer keeps if already -# present (UV_MIN_VERSION) and which ignores UV_SYSTEM_CERTS. Mirror the choice onto -# both so it works on either uv. Opt out with UV_SYSTEM_CERTS=0. +# macOS: trust the system Keychain (TLS-inspecting proxies) via both UV_SYSTEM_CERTS (uv >= 0.11) and UV_NATIVE_TLS (uv 0.8.16-0.10.x); opt out with UV_SYSTEM_CERTS=0. if [ "$OS" = "macos" ]; then : "${UV_SYSTEM_CERTS:=1}" : "${UV_NATIVE_TLS:=$UV_SYSTEM_CERTS}" @@ -1857,7 +1735,7 @@ _uv_version_ok() { ''|*[!0-9.]*) return 1 ;; esac version_ge "$_ver" "$UV_MIN_VERSION" || return 1 - # Prerelease of the exact minimum (e.g. 0.7.14-rc1) is still below stable 0.7.14 + # Prerelease of the exact minimum is still below the stable minimum. [ "$_ver" = "$UV_MIN_VERSION" ] && [ "$_raw" != "$_ver" ] && return 1 return 0 } @@ -1883,13 +1761,7 @@ _MIGRATED=false _PREV_TORCH_VER="" if [ -x "$VENV_DIR/bin/python" ]; then - # why: matching guard to the .venv branch below -- in env-mode - # $STUDIO_HOME is a user-chosen workspace, so refuse to nuke an - # existing $STUDIO_HOME/unsloth_studio that lacks Unsloth sentinels. - # Accept the in-VENV ownership marker so partial-install retries are - # not blocked. Sentinels must be regular files: -f follows symlinks - # to files (the legitimate ln -s shim shape) but rejects directories - # and broken/dir-targeted symlinks. + # Env-mode: refuse to nuke a $VENV_DIR lacking Unsloth sentinels (regular files only); the in-VENV ownership marker unblocks partial-install retries. if [ "$_STUDIO_HOME_REDIRECT" = "env" ] \ && [ ! -f "$VENV_DIR/.unsloth-studio-owned" ] \ && [ ! -f "$STUDIO_HOME/share/studio.conf" ] \ @@ -1898,20 +1770,13 @@ if [ -x "$VENV_DIR/bin/python" ]; then echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2 exit 1 fi - # Record the existing venv's torch BEFORE the replacement moves it aside: a re-run - # rebuilds the venv for clean state, but must keep the torch release the user - # already has (see _previous_torch_pin below). Last line only: sitecustomize or - # import-hook noise on stdout must not corrupt the version. + # Record the existing venv's torch BEFORE replacement (see _previous_torch_pin); last line only so stdout noise can't corrupt it. _PREV_TORCH_VER=$("$VENV_DIR/bin/python" -c \ "import torch; print(torch.__version__)" 2>/dev/null | tail -n 1 || true) - # New layout already exists — replace only after preserving rollback copy. substep "preserving existing environment for rollback..." _start_studio_venv_replacement "$VENV_DIR" elif [ "$_STUDIO_HOME_REDIRECT" != "env" ] && [ -x "$STUDIO_HOME/.venv/bin/python" ]; then - # Old layout exists — validate before migrating. - # Skip in env-mode so we don't rm -rf an unrelated .venv at the - # workspace root (e.g. user's existing project Python venv). - # In no-torch mode, a missing torch package is expected; validate Python only. + # Old layout: validate before migrating (env-mode skips so an unrelated workspace .venv isn't rm -rf'd); no-torch validates Python only. substep "found legacy Unsloth environment, validating..." _legacy_ok=false if [ "$SKIP_TORCH" = true ]; then @@ -1942,8 +1807,7 @@ torch.testing.assert_close(torch.unique(E), torch.tensor((20,), device=E.device, fi fi -# If an Intel Mac has a stale 3.13 venv from a previous failed install, recreate -# (skip when the user explicitly chose a version via --python) +# Recreate a stale Intel Mac 3.13 venv (skip when the user chose --python). if [ "$SKIP_TORCH" = true ] && [ "$MAC_INTEL" = true ] && [ -z "$_USER_PYTHON" ] && [ -x "$VENV_DIR/bin/python" ]; then _PY_MM=$("$VENV_DIR/bin/python" -c \ "import sys; print('{}.{}'.format(*sys.version_info[:2]))" 2>/dev/null || echo "") @@ -1957,11 +1821,7 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then step "venv" "creating Python ${PYTHON_VERSION} virtual environment" substep "$VENV_DIR" if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ] && [ -z "$_USER_PYTHON" ]; then - # Apple Silicon: request an arch-explicit arm64 CPython so uv cannot - # reuse a cached x86_64 (Rosetta) build. torch ships no macOS x86_64 - # wheels since 2.2.2, so an x86_64 venv makes the torch install - # unresolvable. The arm64 guard below is kept as a backstop for - # migrated / pre-existing venvs. + # Arch-explicit arm64 CPython so uv can't reuse a cached x86_64 (Rosetta) build (torch ships no macOS x86_64 wheels since 2.2.2). run_install_cmd "create venv" uv venv "$VENV_DIR" \ --python "cpython-${PYTHON_VERSION}-macos-aarch64-none" else @@ -1969,26 +1829,12 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then fi fi -# Mark the freshly-created venv as Unsloth-owned so a partial install can be -# repaired by re-running install.sh; the env-mode deletion guard above accepts -# this marker as the primary sentinel. +# Mark the freshly-created venv as Unsloth-owned (env-mode deletion guard's primary sentinel). if [ -x "$VENV_DIR/bin/python" ]; then : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true fi -# Guard against two independent Apple Silicon venv problems, in order: -# 1. uv may create the venv from a cached x86_64 (Rosetta) Python when a -# same-version x86_64 build is already cached (often because uv itself -# is an x86_64 build). That venv reports x86_64 to wheel resolvers, and -# PyTorch ships no macOS wheels on the CPU index for any architecture, -# so the torch install can never resolve. Recreate it with an -# arch-explicit arm64 CPython. -# 2. Python 3.13.8 has a known torch import bug. -# The two are independent: a venv may be x86_64 and, once recreated, still -# land on 3.13.8. So we re-inspect the interpreter between the checks instead -# of chaining them with elif, guaranteeing both invariants hold on whatever -# venv we end up with. Skip both when the user explicitly chose an interpreter -# via --python. +# Two independent Apple Silicon venv guards: (1) x86_64 (Rosetta) venv -> recreate arm64; (2) Python 3.13.8 torch import bug. Re-inspect between checks (not elif); skip under --python. if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then _inspect_venv() { "$VENV_DIR/bin/python" -c \ @@ -1998,14 +1844,9 @@ if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then _info=$(_inspect_venv) _VENV_ARCH=${_info%% *} _PY_VER=${_info##* } - # If the interpreter could not be executed (an x86_64 venv python on a Mac - # without Rosetta installed), the probe above yields an empty arch. Fall - # back to reading the binary's Mach-O arch statically so the x86_64 - # recreate below still triggers instead of letting uv fail later. + # An unexecutable x86_64 venv python (no Rosetta) yields an empty arch; read the binary's Mach-O arch statically. if [ -z "$_VENV_ARCH" ] && [ -x "$VENV_DIR/bin/python" ]; then - # uv symlinks bin/python to the base interpreter, so dereference with - # file -L (lipo already follows the link). Trailing || true keeps the - # installer alive under set -e when neither tool is present. + # file -L dereferences the base interpreter; trailing || true survives set -e. _archs=$(lipo -archs "$VENV_DIR/bin/python" 2>/dev/null \ || file -L "$VENV_DIR/bin/python" 2>/dev/null || true) case "$_archs" in @@ -2047,8 +1888,7 @@ if [ -x "$VENV_DIR/bin/python" ]; then substep "${VENV_DIR}" fi -# Default torch constraint -- tightened for Python 3.13+ on arm64 macOS -# (torch <2.6 has no cp313 macOS arm64 wheels) +# Default torch constraint; tightened for Python 3.13+ on arm64 macOS (torch <2.6 has no cp313 macOS arm64 wheels). TORCH_CONSTRAINT="torch>=2.4,<2.11.0" if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then _PY_MINOR=$("$VENV_DIR/bin/python" -c \ @@ -2057,13 +1897,7 @@ if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; t TORCH_CONSTRAINT="torch>=2.6,<2.11.0" fi fi -# Companion (torchvision/torchaudio) constraints, bounded to torch's window. -# torchaudio 2.11 dropped its exact torch pin, so a bare companion next to a -# <2.11-capped torch resolves torchaudio 2.11 (verified: cpu leaf installed -# torch 2.10.0+cpu with torchaudio 2.11.0+cpu). torchvision still exact-pins -# torch and self-corrects, but is bounded for symmetry. Widened alongside the -# cu* torch window below; the torch-2.11 AMD paths (rocm7.2 / per-gfx / Strix) -# pin their own trio. +# Companion constraints bounded to torch's window: torchaudio 2.11 dropped its torch pin, so a bare companion beside a <2.11 torch resolves 2.11. TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" @@ -2086,10 +1920,7 @@ _find_no_torch_runtime() { } # ── AMD ROCm GPU detection helper ── -# WSL2 ROCDXG: the system rocminfo enumerates the GPU over /dev/dxg only when -# HSA_ENABLE_DXG_DETECTION=1 (a no-op on bare metal), and /opt/rocm/bin can be -# off PATH outside login shells (the profile.d drop-in). Seed both before any -# rocminfo probe or a ROCDXG WSL host is misdetected as CPU-only. +# WSL2 ROCDXG: rocminfo needs HSA_ENABLE_DXG_DETECTION=1 and /opt/rocm/bin can be off PATH; seed both or a ROCDXG WSL host misdetects as CPU-only. _ensure_rocm_probe_env() { export HSA_ENABLE_DXG_DETECTION="${HSA_ENABLE_DXG_DETECTION:-1}" if ! command -v rocminfo >/dev/null 2>&1 && [ -x /opt/rocm/bin/rocminfo ]; then @@ -2097,12 +1928,7 @@ _ensure_rocm_probe_env() { fi } -# Returns 0 if an AMD GPU is present. Checks rocminfo, amd-smi, then sysfs -# KFD topology (env-var-independent fallback for when HIP/ROCR_VISIBLE_DEVICES hides devices). -# Always returns 1 (false) when an NVIDIA GPU is present: blocks every -# detection path (rocminfo, amd-smi, KFD sysfs) from producing a false -# positive on NVIDIA-only or NVIDIA-primary hosts, even when ROCm tools -# are co-installed. +# True if an AMD GPU is present (rocminfo, amd-smi, then KFD sysfs); always false when an NVIDIA GPU is present. _has_amd_rocm_gpu() { _ensure_rocm_probe_env if _has_usable_nvidia_gpu; then @@ -2118,30 +1944,22 @@ _has_amd_rocm_gpu() { awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \ gpu && amd { found=1 } END{ exit !found }' \ /sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then - # vendor_id 4098 = 0x1002 (AMD). NVIDIA open kernel module (driver - # 560+) can register KFD topology nodes with non-zero gpu_id but - # vendor_id 4318 (0x10DE). Require AMD vendor to avoid misrouting - # NVIDIA-only hosts to the ROCm install path. + # vendor_id 4098 = 0x1002 (AMD); the NVIDIA open kernel module can register KFD nodes too. return 0 fi return 1 } -# ── Detect GPU and choose PyTorch index URL ── -# Mirrors Get-TorchIndexUrl in install.ps1. -# On CPU-only machines this returns the cpu index, avoiding the solver -# dead-end where --torch-backend=auto resolves to unsloth==2024.8. +# ── Detect GPU and choose PyTorch index URL (mirrors Get-TorchIndexUrl) ── +# CPU-only machines get the cpu index, avoiding the --torch-backend=auto solver dead-end. get_torch_index_url() { _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}" _base="${_base%/}" - # Explicit override -- skip ALL GPU probing (headless / container / CI / cross-install). - # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf (cpu, cu128, ...) - # appended to the mirror base. Trim whitespace so a whitespace-only value is unset. + # Explicit override skips ALL GPU probing: UNSLOTH_TORCH_INDEX_URL wins (verbatim); UNSLOTH_TORCH_INDEX_FAMILY is the leaf appended to the mirror base; whitespace-only = unset. _url="${UNSLOTH_TORCH_INDEX_URL:-}" _url="${_url#"${_url%%[![:space:]]*}"}"; _url="${_url%"${_url##*[![:space:]]}"}" if [ -n "$_url" ]; then - # Trim trailing PATH slashes (a multi-slash path 404s on strict pip proxies) while - # preserving a ?query/#fragment token (a whole-URL strip would eat a "/"-ending token). + # Trim trailing PATH slashes (multi-slash 404s on strict proxies), preserving ?query/#fragment. _url=$(_trim_index_path_slashes "$_url") echo "$_url"; return fi @@ -2154,10 +1972,7 @@ get_torch_index_url() { fi # macOS: always CPU (no CUDA support) case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac - # Try nvidia-smi -- require the binary to actually list a usable GPU. - # Presence of the binary alone (container leftovers, stale driver - # packages) is not sufficient: otherwise an AMD-only host would - # silently install CUDA wheels. + # Require nvidia-smi to actually list a usable GPU; the binary alone would install CUDA wheels on AMD. _smi="" _nvidia_detected=0 if _has_usable_nvidia_gpu; then @@ -2169,10 +1984,7 @@ get_torch_index_url() { fi fi if [ "$_nvidia_detected" -eq 0 ]; then - # No NVIDIA GPU -- check for AMD ROCm GPU. - # PyTorch only publishes ROCm wheels for linux-x86_64; skip the - # ROCm branch entirely on aarch64 / arm64 / other architectures - # so non-x86_64 Linux hosts fall back cleanly to CPU wheels. + # No NVIDIA GPU: check AMD ROCm. ROCm wheels are linux-x86_64 only; other arches fall back to CPU. case "$(uname -m)" in x86_64|amd64) : ;; *) echo "$_base/cpu"; return ;; @@ -2203,16 +2015,14 @@ get_torch_index_url() { *) _rocm_tag="" ;; # reject malformed (empty, garbled, or major=0) esac if [ -n "$_rocm_tag" ]; then - # Minimum supported: ROCm 6.0 (no PyTorch wheels exist for older) + # Minimum supported: ROCm 6.0. case "$_rocm_tag" in rocm[1-5].*) echo "[WARN] ROCm $_rocm_tag detected but PyTorch ROCm wheels require ROCm 6.0+ -- falling back to CPU-only PyTorch" >&2 echo "[WARN] Upgrade ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 echo "$_base/cpu"; return ;; esac - # Supported tags; 6.5+ clips to rocm6.4, 7.3+ caps to rocm7.2. - # PyTorch publishes major.minor URLs only (no patch level), so - # rocm7.2.1 / rocm6.0.2 / etc. must normalise to rocm7.2 / rocm6.0. + # Normalise to major.minor (no patch-level URLs); 6.5+ clips to rocm6.4, 7.3+ caps to rocm7.2. case "$_rocm_tag" in rocm6.0|rocm6.0.*) echo "$_base/rocm6.0" ;; rocm6.1|rocm6.1.*) echo "$_base/rocm6.1" ;; @@ -2223,32 +2033,21 @@ get_torch_index_url() { rocm7.1|rocm7.1.*) echo "$_base/rocm7.1" ;; rocm7.2|rocm7.2.*) echo "$_base/rocm7.2" ;; rocm6.*) - # ROCm 6.5+ (no published PyTorch wheels): clip down - # to the last supported 6.x wheel set. + # ROCm 6.5+: clip to the last supported 6.x wheel set. echo "$_base/rocm6.4" ;; *) - # ROCm 7.3+ (future): cap to rocm7.2 (latest known) + # ROCm 7.3+: cap to rocm7.2 (latest known). echo "$_base/rocm7.2" ;; esac return fi - # AMD GPU confirmed by rocminfo/amd-smi but ROCm version could not be - # read from any source (amd-smi, /opt/rocm/.info/version, hipconfig, - # dpkg, rpm). Warn explicitly rather than silently installing CPU PyTorch. + # ROCm version unreadable from any source: warn rather than silently install CPU PyTorch. echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2 echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2 echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 echo "$_base/cpu"; return fi - # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P). - # Newer NVIDIA drivers (e.g. 610.x) print "CUDA UMD Version: X.Y" instead - # of the legacy "CUDA Version: X.Y"; accept both with two BRE expressions - # (POSIX sed does not support "?" without -E). The two patterns are - # mutually exclusive per line, so head -1 picks the first emitted match. - # Bound the call (a wedged nvidia-smi would otherwise hang here) and force - # the C locale for stable parsing. LC_ALL is exported inside this command - # substitution subshell so it reaches nvidia-smi through _run_bounded - # without depending on `env`; the export is scoped to the subshell. + # Parse CUDA version from nvidia-smi (POSIX-safe): accept both "CUDA Version:" and the newer "CUDA UMD Version:". Bounded, C locale. _cuda_ver=$(export LC_ALL=C; _run_bounded "$_smi" 2>/dev/null \ | sed -n \ -e 's/.*CUDA UMD Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ @@ -2280,29 +2079,23 @@ _torch_flavor_tag() { esac } -# Final path segment of a wheel index URL ($1), lowercased, query/fragment stripped first -# so a token-authenticated pin (.../cu128?token=x) classifies as cu128 (else it reinstalls -# every update). Classification only. Shared with the py / ps1 leaf extractors. +# Final path segment of a wheel index URL ($1), lowercased, query/fragment stripped so .../cu128?token=x classifies as cu128. Shared with py / ps1. _torch_index_url_leaf() { _tl_u="${1%%\?*}" _tl_u="${_tl_u%%#*}" - # Strip ALL trailing slashes, not one: .../rocm7.2// must yield rocm7.2, not an empty leaf. + # Strip ALL trailing slashes: .../rocm7.2// must yield rocm7.2, not empty. while [ -n "$_tl_u" ] && [ "${_tl_u%/}" != "$_tl_u" ]; do _tl_u="${_tl_u%/}" done printf '%s' "${_tl_u##*/}" | tr '[:upper:]' '[:lower:]' } -# True (exit 0) when a lowercased leaf is an EXACT pip ROCm family: rocm[.] -# or a gfx ARCHITECTURE leaf (gfx followed by a digit: gfx90a, gfx1151, gfx120x-all). A leaf -# that merely starts with rocm/gfx (rocm7.2-private, gfx-private) is a custom verbatim pin. -# Matches the py / ps1 sides. +# True when a lowercased leaf is an EXACT pip ROCm family (rocm[.] or gfx*); a leaf merely starting with rocm/gfx is a custom verbatim pin. _is_pip_rocm_family_leaf() { case "$1" in gfx[0-9]*) return 0 ;; rocm[0-9]*) - # Exact rocm[.]: both major and minor must be non-empty all-digits - # (rocm7., rocm7.2.1, rocm7.2-private are all custom pins, not a family). + # Major/minor both non-empty all-digits (rocm7., rocm7.2.1, rocm7.2-private are custom pins). _rocm_rest="${1#rocm}" case "$_rocm_rest" in *.*.*) return 1 ;; @@ -2319,11 +2112,7 @@ _is_pip_rocm_family_leaf() { esac } -# Whether release base $1 (X.Y[.Z...]) falls inside constraint window $2 -# ("torch>=A.B[.C],=A.B, cuda, cu126 -> cu130, PyPI bare -> +cu130) while the -# release follows the user. Gating on flavor was wrong: a PyPI torch reports a BARE -# version (on Linux the PyPI wheel IS CUDA), misclassified "cpu", so a healthy 2.10 on a -# cu130 host was moved to 2.11. Per-leaf floors still win (rocm7.2 / gfx >=2.11 for the -# Strix _grouped_mm fix, out-of-window manual installs) and are never pinned; the caller's -# _PREV_FALLBACK_CONSTRAINT installs the newest supported release when the index lacks the -# exact one. Opt out with UNSLOTH_TORCH_UPGRADE=1. +# Keep the previous venv's torch RELEASE on a re-run when inside the constraint window; flavor follows the freshly chosen index. Opt out with UNSLOTH_TORCH_UPGRADE=1. _previous_torch_pin() { _ptp_ver="$1" _ptp_con="$2" [ -n "$_ptp_ver" ] || { echo ""; return; } [ "${UNSLOTH_TORCH_UPGRADE:-0}" = "1" ] && { echo ""; return; } _ptp_base="${_ptp_ver%%+*}" - # Base must be a plain numeric release (X.Y[.Z]); probe noise and - # nightly/dev/source builds (2.11.0.dev20250704, 2.9.0a0) must never - # become a pin -- no stable index carries them, so pinning would only - # print "keeping it" and then burn a doomed resolve before falling back. + # Base must be a plain numeric release; nightly/dev/source builds must never become a pin. case "$_ptp_base" in *[!0-9.]* | *..* | .* | *.) echo ""; return ;; [0-9]*.[0-9]*) ;; @@ -2376,16 +2153,10 @@ _previous_torch_pin() { echo "torch==$_ptp_base" } -# Install torch from TORCH_INDEX_URL honoring a kept-release pin: with _PREV_TORCH_PIN -# set, TORCH_CONSTRAINT is the exact previous release; fall back to the supported range -# if the index lacks it (pruned mirror) rather than failing. Used by every --default-index -# path (NVIDIA cu*, AMD rocm/gfx fallbacks, cpu/mac, ROCm repairs) so preservation is -# uniform. Extra args (e.g. --force-reinstall) are passed through to uv. +# Install torch from TORCH_INDEX_URL honoring a kept-release pin, falling back to the supported range if the index lacks it; used by every --default-index path. _install_torch_default_index() { if [ -n "$_PREV_TORCH_PIN" ]; then - # Pair the companions with the kept torch minor: torchaudio no longer - # exact-pins torch in its metadata, so leaving it unconstrained resolves - # a newer mismatched build (a kept torch 2.9.0 pulled torchaudio 2.11.0). + # Pair companions with the kept torch minor (torchaudio no longer exact-pins torch). _itdi_base="${_PREV_TORCH_PIN#torch==}" _itdi_minor="${_itdi_base#*.}" _itdi_minor="${_itdi_minor%%.*}" @@ -2411,14 +2182,12 @@ _install_torch_default_index() { fi } -# Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* -> -# rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops. +# Expected tag from the index leaf ($1): cuXXX / cpu / rocm; empty on an unknown leaf so the repair safely no-ops. _expected_torch_flavor_tag() { _leaf=$(_torch_index_url_leaf "$1") case "$_leaf" in cu[0-9]*) - # Exact cu + digits only; a cu*-suffixed leaf (cu128-private) -> "" (custom), - # else a correct +cu128 wheel is force-reinstalled every run. + # Exact cu + digits only; a cu*-suffixed leaf (cu128-private) -> "" (custom). case "${_leaf#cu}" in *[!0-9]*) echo "" ;; *) echo "$_leaf" ;; @@ -2432,11 +2201,7 @@ _expected_torch_flavor_tag() { esac } -# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX / -# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv -# resolves (torch + every transitive dep) via --default-index -- the same URLs the -# fresh-install paths above already use -- so a stale wheel is auto-repairable. -# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. +# Whether index ($1) supports a plain --default-index reinstall (cuXXX / rocmX.Y / repo.amd.com gfx* PEP 503 indexes); unknown leaves -> no. _torch_index_repairable() { _leaf=$(_torch_index_url_leaf "$1") case "$_leaf" in @@ -2448,8 +2213,7 @@ _torch_index_repairable() { esac } -# Remove credentials from a wheel index URL ($1) so an authenticated pin never leaks: -# drops userinfo AND query/fragment; scheme/host/path stay exact. Shared with py / ps1. +# Remove credentials from a wheel index URL ($1): drops userinfo AND query/fragment. Shared with py / ps1. _strip_index_url_credentials() { _sic_url="$1" case "$_sic_url" in @@ -2475,15 +2239,10 @@ _strip_index_url_credentials() { } get_radeon_wheel_url() { - # Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing - # contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/, - # rocm-rel-7.1.1/, rocm-rel-7.1/ (AMD publishes both M.m and M.m.p dirs). - # Accepts both X.Y and X.Y.Z host versions since /opt/rocm/.info/version - # and hipconfig --version can return either shape. + # Linux only. Picks a repo.radeon.com base URL (rocm-rel-/); accepts X.Y or X.Y.Z host versions. case "$(uname -s)" in Linux) ;; *) echo ""; return ;; esac - # Detect ROCm version (X.Y or X.Y.Z) -- try amd-smi, then - # /opt/rocm/.info/version, then hipconfig. + # Detect ROCm version via amd-smi, then /opt/rocm/.info/version, then hipconfig. _full_ver="" _full_ver=$({ command -v amd-smi >/dev/null 2>&1 && \ amd-smi version 2>/dev/null | awk -F'ROCm version: ' \ @@ -2503,16 +2262,12 @@ get_radeon_wheel_url() { } # ── Radeon repo wheel selection helpers ────────────────────────────────────── -# Fetches the Radeon repo directory listing once into _RADEON_LISTING (global). -# _RADEON_PYTAG holds the CPython tag for the running interpreter (e.g. cp312). -# _RADEON_BASE_URL holds the base URL for relative-href resolution. _RADEON_LISTING="" _RADEON_PYTAG="" _RADEON_BASE_URL="" _radeon_fetch_listing() { - # Usage: _radeon_fetch_listing BASE_URL - # Populates _RADEON_LISTING, _RADEON_PYTAG, _RADEON_BASE_URL. + # Usage: _radeon_fetch_listing BASE_URL -- populates _RADEON_LISTING, _RADEON_PYTAG, _RADEON_BASE_URL. _RADEON_BASE_URL="$1" _RADEON_PYTAG=$("$_VENV_PY" -c " import sys @@ -2527,15 +2282,7 @@ print('cp{}{}'.format(sys.version_info.major, sys.version_info.minor)) } _pick_radeon_wheel() { - # Usage: _pick_radeon_wheel PACKAGE_NAME [VERSION_PREFIX] - # Scans $_RADEON_LISTING for the newest wheel whose filename starts exactly - # with PACKAGE_NAME- (and optionally VERSION_PREFIX) and matches _RADEON_PYTAG + linux_x86_64. - # Prints the full URL (resolving relative hrefs against _RADEON_BASE_URL). - # - # POSIX-compliant pipeline: all href parsing, filtering, and version - # selection is done inside a single awk script rather than reaching - # for GNU extensions (grep -o, sort -V) that would break under BSD - # or BusyBox coreutils. + # Usage: _pick_radeon_wheel PACKAGE_NAME [VERSION_PREFIX] -- newest matching cpXY linux_x86_64 wheel URL; POSIX awk only (no grep -o / sort -V). _pkg="$1" _ver_prefix="${2:-}" [ -n "$_RADEON_LISTING" ] || return 1 @@ -2589,16 +2336,7 @@ _pick_radeon_wheel() { } # ── ROCm-on-WSL bootstrap for AMD Strix Halo (gfx1151) ─────────────────────── -# No-op everywhere except: WSL + GPU wanted + no usable GPU yet + /dev/dxg + -# Strix Halo APU. Every other config (NVIDIA, native-Linux ROCm, macOS, Windows, -# CPU, non-Strix WSL) skips it and normal detection runs unchanged. NEVER aborts -# the installer -- always returns 0. Runs the idempotent helper (ROCm 7.2 + -# librocdxg), then sources the env it persisted so detection finds the GPU. -# Export the ROCm-on-WSL env into this process and persist it to /etc/profile.d -# so non-login Unsloth/llama launches inherit it. Idempotent (writes only when -# the drop-in is missing); no-op without librocdxg, so never fires off WSL. -# /etc/profile.d is root-owned -- sudo-tee when not root, else ROCm vanishes -# after this shell on a non-root reinstall. Best-effort either way. +# Export the ROCm-on-WSL env + persist to /etc/profile.d (sudo-tee when not root); idempotent, no-op without librocdxg, best-effort. _persist_rocm_wsl_dropin() { [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ] || return 0 _rw_rocm=/opt/rocm @@ -2625,45 +2363,36 @@ _persist_rocm_wsl_dropin() { fi } -# _wsl_amd_gpu_name is defined earlier so both the reroute and this bootstrap can use it. _maybe_bootstrap_rocm_wsl() { [ "${OS:-}" = "wsl" ] || return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 - # Leave any already-usable GPU completely alone (NVIDIA, or working ROCm). + # Leave any already-usable GPU alone (NVIDIA, or working ROCm). if _has_usable_nvidia_gpu; then return 0; fi - # Usable ROCm = rocminfo enumerates a real GPU agent: gfx[1-9] (excludes gfx000, - # the CPU agent) and not the "gfx11-generic" fallback. awk consumes all input so - # rocminfo isn't SIGPIPE'd like `grep -q` under pipefail. + # Usable ROCm = rocminfo enumerates a real gfx[1-9] agent (not gfx000 / generic); awk consumes all input so rocminfo isn't SIGPIPE'd. _ensure_rocm_probe_env if command -v rocminfo >/dev/null 2>&1 && \ rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then - # rocminfo may work only via the transient env _ensure_rocm_probe_env - # just set, which dies with the installer. Persist the drop-in so login - # shells (Unsloth, llama.cpp) inherit it -- else a reinstall over an - # existing /opt/rocm (uninstall keeps ROCm but drops it) loses the GPU. + # Persist the drop-in so login shells inherit the transient probe env. _persist_rocm_wsl_dropin return 0 fi # WSL GPU passthrough device must exist (present on any WSL2 GPU host). [ -e /dev/dxg ] || return 0 - # Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also - # ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo. + # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also ask the Windows host. if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ && ! _wsl_amd_gpu_name >/dev/null 2>&1; then return 0 fi command -v bash >/dev/null 2>&1 || return 0 - # Fast path: already configured (librocdxg present) but launched from a - # non-login shell so the persisted env wasn't loaded -- just load it. + # Fast path: configured (librocdxg present) but launched from a non-login shell -- just load it. if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then if [ -r /etc/profile.d/unsloth-rocm-wsl.sh ]; then # shellcheck disable=SC1091 . /etc/profile.d/unsloth-rocm-wsl.sh || true else - # librocdxg present but the env drop-in is gone (e.g. an Unsloth - # uninstall removed it while keeping shared ROCm). Restore the env. + # librocdxg present but the env drop-in is gone (uninstall dropped it); restore it. _persist_rocm_wsl_dropin fi return 0 @@ -2675,7 +2404,7 @@ _maybe_bootstrap_rocm_wsl() { substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU." substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)" - # Locate the helper: prefer the copy shipped beside install.sh, else fetch it. + # Locate the helper: prefer the copy beside install.sh, else fetch it. _rw_helper="${_REPO_ROOT:-.}/scripts/install_rocm_wsl_strixhalo.sh" _rw_tmp="" if [ ! -r "$_rw_helper" ]; then @@ -2689,11 +2418,7 @@ _maybe_bootstrap_rocm_wsl() { fi fi - # Consent: the narrow guarded case is exactly the GPU setup the user ran the - # installer for, so it proceeds AUTOMATICALLY by default (works with no TTY, - # e.g. `curl ... | sh`). Opt out via UNSLOTH_SKIP_ROCM_WSL_SETUP=1 (top of - # function). The Tauri app drives its own consent UI, so under TAURI_MODE it - # only runs when the app passes UNSLOTH_ROCM_WSL_AUTO=1; else surface and wait. + # Consent: automatic by default (opt out UNSLOTH_SKIP_ROCM_WSL_SETUP=1); under TAURI_MODE runs only with UNSLOTH_ROCM_WSL_AUTO=1. _rw_go=1 if [ "${TAURI_MODE:-false}" = "true" ] && [ "${UNSLOTH_ROCM_WSL_AUTO:-0}" != "1" ]; then tauri_log "ROCM_WSL_AVAILABLE" "strixhalo" @@ -2702,11 +2427,9 @@ _maybe_bootstrap_rocm_wsl() { fi if [ "$_rw_go" = "1" ]; then - # Helper does its own sudo + is idempotent. SMOKE_TEST=0: install.sh - # installs torch itself right after, into the real venv. + # Helper does its own sudo + is idempotent. SMOKE_TEST=0: install.sh installs torch after. if UNSLOTH_WSL_SMOKE_TEST=0 bash "$_rw_helper"; then - # Pull the helper's persisted env into THIS shell so detection - # (rocminfo) now enumerates the GPU and routes to gfx1151. + # Pull the helper's persisted env into THIS shell so detection routes to gfx1151. if [ -r /etc/profile.d/unsloth-rocm-wsl.sh ]; then # shellcheck disable=SC1091 . /etc/profile.d/unsloth-rocm-wsl.sh || true @@ -2719,10 +2442,7 @@ _maybe_bootstrap_rocm_wsl() { [ -n "$_rw_tmp" ] && rm -f "$_rw_tmp" return 0 } -# When the caller pins the wheel index (UNSLOTH_TORCH_INDEX_URL / _FAMILY), honour it -# everywhere: skip the WSL ROCm bootstrap and the Radeon/Strix reroute below (which would -# re-probe the GPU and overwrite the pin). Trim whitespace first (parity with -# get_torch_index_url): a whitespace-only override is unset there, so must not flip this true. +# A pinned wheel index (UNSLOTH_TORCH_INDEX_URL / _FAMILY) skips the WSL ROCm bootstrap and the Radeon/Strix reroute below; whitespace trimmed first. _torch_index_pinned=false _ti_url_trim="${UNSLOTH_TORCH_INDEX_URL:-}" _ti_url_trim="${_ti_url_trim#"${_ti_url_trim%%[![:space:]]*}"}"; _ti_url_trim="${_ti_url_trim%"${_ti_url_trim##*[![:space:]]}"}" @@ -2735,21 +2455,10 @@ fi TORCH_INDEX_URL=$(get_torch_index_url) -# Export the resolved torch backend ("cuda", "rocm", or "cpu") so that -# downstream scripts (setup.sh -> install_python_stack.py) know what was -# chosen here and can skip ROCm-specific repair steps on CUDA/CPU hosts. -# Classify on the FINAL path segment only: a custom UNSLOTH_PYTORCH_MIRROR -# whose base path happens to contain "rocm" or "gfx" must not mislabel a -# cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix -# overrides in gfxNNNN/, so the trailing slash is stripped first). -# Lowercase the leaf so every gfx*/rocm*/cu* arm matches regardless of case (canonical AMD -# RDNA4 leaf is gfx120X-all). CUDA is branded only on a real cu[0-9]* leaf, so a mirror -# leaf (/current) does NOT commit a CUDA backend; an unknown leaf leaves the var unset so -# the stack probes the GPU. Query/fragment dropped first, then ALL trailing slashes (in -# lockstep with the shared _torch_index_url_leaf extractor). +# Export UNSLOTH_TORCH_BACKEND ("cuda"/"rocm"/"cpu") for downstream scripts; classify on the FINAL lowercased leaf only so a custom mirror path containing "rocm"/"gfx" can't mislabel a cu*/cpu index. _torch_index_leaf="${TORCH_INDEX_URL%%\?*}" _torch_index_leaf="${_torch_index_leaf%%#*}" -# Strip ALL trailing slashes, not one: .../cu128// must yield cu128, not an empty leaf. +# Strip ALL trailing slashes: .../cu128// must yield cu128, not empty. while [ -n "$_torch_index_leaf" ] && [ "${_torch_index_leaf%/}" != "$_torch_index_leaf" ]; do _torch_index_leaf="${_torch_index_leaf%/}" done @@ -2759,33 +2468,25 @@ case "$_torch_index_leaf" in rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; cu[0-9]*) export UNSLOTH_TORCH_BACKEND="cuda" ;; - # Unknown leaf (odd mirror, /current): unset so a stale inherited value can't leak and - # the stack probes the GPU. + # Unknown leaf: unset so a stale inherited value can't leak and the stack probes the GPU. *) unset UNSLOTH_TORCH_BACKEND ;; esac -# Whether TORCH_INDEX_URL names an actual pip ROCm family (rocm* / gfx*), gating the -# ROCm-only side effects below (AMD bitsandbytes, ROCm-torch repair). Digit-gated so a leaf -# merely STARTING with "rocm" isn't force-repaired from the wrong path. +# Whether TORCH_INDEX_URL names an actual pip ROCm family, gating the ROCm-only side effects below (AMD bitsandbytes, ROCm-torch repair). if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then _torch_index_is_rocm_family=true else _torch_index_is_rocm_family=false fi -# rocm7.2 and the per-gfx indexes with the _grouped_mm <2.11 bug (gfx120X-all, gfx1151, -# gfx1150) ship torch 2.11.0 -- raise the floor (also covers a pinned override that skipped -# the Strix reroute). Pin the companions too: the per-gfx index publishes them independently -# and a bare name can resolve a 2.12 ABI-mismatched wheel. Match on the FINAL leaf so a -# custom mirror with a gfx/rocm7.2 path segment but a cu*/cpu family isn't forced. +# rocm7.2 and the per-gfx indexes (Strix _grouped_mm fix) ship torch 2.11.0: raise the floor and pin companions; match the FINAL leaf only. case "$_torch_index_leaf" in rocm7.2|gfx120x-all|gfx1151|gfx1150) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" ;; - # CUDA cu12x/cu13x indexes ship torch 2.11.x: widen the ceiling to <2.12.0 (matches - # _CUDA_TORCH_PKG_SPEC) and widen the companions with it so the trio stays paired. + # CUDA cu12x/cu13x indexes ship torch 2.11.x: widen the trio ceiling to <2.12.0. cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.27.0" @@ -2793,21 +2494,14 @@ case "$_torch_index_leaf" in ;; esac -# A pinned custom/unknown-leaf index (/simple, /current, /cu128-private) has no curated -# companion set, so bound torchvision/torchaudio to the same <2.11 range the Python path pins -# (else a mirror with newer companions resolves a 2.12 ABI-mismatched wheel). Known families -# keep their curated companions above (_expected_torch_flavor_tag returns "" only for custom). +# A pinned custom/unknown-leaf index has no curated companion set: bound the companions to the same <2.11 range the Python path pins. if [ "$_torch_index_pinned" = true ] && \ [ -z "$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")" ]; then TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" fi -# Auto-detect GPU for AMD ROCm based -# get_torch_index_url must have chosen */rocm* -# (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon". -# Skipped when the index is pinned: an explicit override must not be rerouted to the -# Radeon/Strix repos by GPU probing. +# Detect a Radeon card (*/rocm* index + rocminfo "Marketing Name:.*Radeon"); skipped when the index is pinned. _amd_gpu_radeon=false if [ "$_torch_index_pinned" = false ]; then case "$TORCH_INDEX_URL" in @@ -2819,26 +2513,17 @@ case "$TORCH_INDEX_URL" in ;; esac # ── Strix Halo / Strix Point: force rocm7.2 wheels, bypass Radeon repo ─────── -# gfx1151 (Strix Halo) and gfx1150 (Strix Point) have a ROCm 7.1 driver bug -# that causes a segfault in torch._grouped_mm (moe_utils.py line 167). -# The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so when -# _amd_gpu_radeon=true the installer silently lands on the broken combo. -# Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2. +# gfx1151/gfx1150 segfault in torch._grouped_mm on ROCm 7.1. case "$TORCH_INDEX_URL" in */rocm7.1|*/rocm7.1.*) - # Collect every gfx token in rocminfo / amd-smi enumeration order - # (skip duplicates), then index by HIP_VISIBLE_DEVICES / - # ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box - # where the user selected the dGPU does NOT get rerouted to the - # Strix per-gfx index. + # Index gfx tokens by HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + selected dGPU box isn't rerouted. _gfx_all="" if command -v rocminfo >/dev/null 2>&1; then _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') fi if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') - # PowerShell paths also probe `amd-smi static --asic`; mirror it - # so a host with hipinfo-less amd-smi reports the gfx target. + # Mirror the PowerShell `amd-smi static --asic` probe as a fallback. if [ -z "$_gfx_all" ]; then _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') fi @@ -2873,14 +2558,9 @@ case "$TORCH_INDEX_URL" in echo " [WARN] Upgrade ROCm to 7.2+ to use the standard index:" >&2 echo " [WARN] https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 echo "" >&2 - # AMD's arch-specific index serves torch 2.11.0+rocm7.13.0 which has AMD's - # actual fix for the gfx1151/gfx1150 _grouped_mm kernel bug -- preferred - # over the pytorch.org rocm7.2 fallback because it exercises the real GPU - # kernel path. Set UNSLOTH_AMD_ROCM_MIRROR to override for air-gapped installs. + # AMD's arch-specific index has the real _grouped_mm fix (torch 2.11.0+rocm7.13.0); UNSLOTH_AMD_ROCM_MIRROR overrides for air-gapped installs. _amd_strix_base="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}" - # Strip ALL trailing slashes to match Python's .rstrip("/") -- a - # double-/triple-slash mirror URL would otherwise produce 404s on - # strict pip proxies (artifactory, sonatype). + # Strip ALL trailing slashes (match Python .rstrip("/")); multi-slash 404s on strict proxies. while [ "${_amd_strix_base%/}" != "$_amd_strix_base" ]; do _amd_strix_base="${_amd_strix_base%/}" done @@ -2894,12 +2574,7 @@ case "$TORCH_INDEX_URL" in ;; esac fi # _torch_index_pinned guard (Radeon + Strix reroute) -# Re-run over an existing install: keep the previous venv's torch RELEASE; the fresh -# index above supplies the right flavor for this machine. Evaluated HERE, after every -# index/constraint decision including the Strix reroute, so the window checked is the -# final one and a raised floor (rocm7.2 / Strix gfx) rejects an older release. -# _PREV_FALLBACK_CONSTRAINT keeps the range so the install can fall back when the exact -# release is not on the chosen index (mirrors may prune old wheels). Skipped for --no-torch. +# Keep the previous venv's torch RELEASE; evaluated after every index/constraint decision (incl. the Strix reroute) so a raised floor rejects an older release. _PREV_TORCH_PIN="" _PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT" if [ "$SKIP_TORCH" = false ]; then @@ -2954,9 +2629,7 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $_gpu_disp_gfx" # Name-based arch inference when tools don't report gfx (mirrors install.ps1 nameArchTable) elif [ -z "$_gpu_disp_gfx" ] && [ -n "$_gpu_disp_mkt" ]; then - # Kept in sync with the nameArchTable in install.ps1 / setup.ps1. - # gfx1102 matched BEFORE gfx1100 so the spaceless "RX 7700S" lands on - # gfx1102 (bash case has no negative lookahead like the PS tables). + # Kept in sync with install.ps1 nameArchTable; gfx1102 matched before gfx1100 ("RX 7700S"). case "$_gpu_disp_mkt" in *"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 *9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 @@ -3005,9 +2678,7 @@ case "$TORCH_INDEX_URL" in if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN" if [ "$OS" = "wsl" ]; then - # WSL + no GPU detected (detection above found nothing). Common - # cause: an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet -- - # /dev/dxg present (graphics) but no ROCm runtime. + # WSL + no GPU: often an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet. _wsl_ubu_ver="" [ -r /etc/os-release ] && _wsl_ubu_ver=$(. /etc/os-release 2>/dev/null; printf '%s' "${VERSION_ID:-}") if [ -e /dev/dxg ]; then @@ -3049,14 +2720,7 @@ esac tauri_log "STEP" "Installing PyTorch" _VENV_PY="$VENV_DIR/bin/python" -# A released unsloth wheel can pin an older torch (unsloth 2026.7.2 declares -# torch<2.11.0); a with-deps PyPI resolve then downgrades the whole trio, -# swapping the pinned +cuXXX/+rocm build for PyPI's default. The flavor guard -# below misses this (PyPI's torch 2.10 default is itself cu128-flavored), so -# freeze the trio via uv --overrides (overrides replace dependency requirements -# during resolution) while unsloth's other deps resolve normally. Sets -# _UNSLOTH_TORCH_OVERRIDES from the trio in the venv; every with-deps unsloth -# install (migrated and fresh) must call this before resolving and rm it after. +# A released unsloth wheel can pin an older torch, and a with-deps resolve then downgrades the pinned +cuXXX/+rocm trio: freeze it via uv --overrides (_UNSLOTH_TORCH_OVERRIDES; callers rm it after). _build_unsloth_torch_overrides() { _UNSLOTH_TORCH_OVERRIDES="" [ "$SKIP_TORCH" = false ] || return 0 @@ -3072,12 +2736,7 @@ for _p in ('torch', 'torchvision', 'torchaudio'): torch==*) _UNSLOTH_TORCH_OVERRIDES=$(mktemp) printf '%s\n' "$_torch_trio_pins" > "$_UNSLOTH_TORCH_OVERRIDES" - # The CLI --overrides flag replaces any UV_OVERRIDE env file (same - # uv setting; macOS arm64 exports one here), so fold its pins in. - # awk, not cat: it drops inherited torch-trio lines (uv intersects - # duplicate overrides, so a conflicting pin would make resolution - # unsatisfiable) and newline-terminates the last line so an - # unterminated file cannot join two requirements into one. + # --overrides replaces any UV_OVERRIDE env file, so fold its pins in; awk drops inherited torch-trio lines and newline-terminates. for _ov_file in ${UV_OVERRIDE:-}; do [ -f "$_ov_file" ] && awk '!/^[[:space:]]*torch(vision|audio)?([[:space:]<>=!~;@[]|$)/' "$_ov_file" >> "$_UNSLOTH_TORCH_OVERRIDES" done @@ -3086,20 +2745,14 @@ for _p in ('torch', 'torchvision', 'torchaudio'): } if [ "$_MIGRATED" = true ]; then - # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving - # existing torch/CUDA unless the ROCm repair below fires. + # Migrated env: force-reinstall unsloth+unsloth-zoo, preserving existing torch/CUDA. substep "upgrading unsloth in migrated environment..." if [ "$SKIP_TORCH" = true ]; then - # No-torch: install unsloth + unsloth-zoo with --no-deps (current - # PyPI metadata still declares torch as a hard dep), then install - # runtime deps (typer, safetensors, transformers, etc.) with --no-deps - # to prevent transitive torch resolution. + # No-torch: --no-deps installs (PyPI metadata still hard-deps torch), then torch-free runtime deps --no-deps. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" - # Resolve pydantic WITH deps so pip pins pydantic-core to the - # matching version (no-torch-runtime.txt below is --no-deps). - # All transitive deps are torch-free. + # Resolve pydantic WITH deps so its pydantic-core matches (all torch-free). run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic _NO_TORCH_RT="$(_find_no_torch_runtime)" @@ -3107,8 +2760,7 @@ if [ "$_MIGRATED" = true ]; then run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" fi else - # Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no - # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. + # Pin mlx-lm away from 0.31.3 here too (a curl-piped migration has no overrides file). _build_unsloth_torch_overrides run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ @@ -3125,9 +2777,7 @@ if [ "$_MIGRATED" = true ]; then --no-deps --reinstall-package unsloth-zoo \ "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" fi - # AMD ROCm: install bitsandbytes even in migrated environments so - # existing ROCm installs gain the AMD bitsandbytes build without a - # fresh reinstall. + # AMD ROCm: install bitsandbytes in migrated envs too, without a fresh reinstall. if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" # Repair ROCm torch if overwritten during migrated install @@ -3158,31 +2808,13 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi if [ "$_radeon_listing_ok" = true ]; then - # Require torch, torchvision, torchaudio wheels to all resolve - # from the Radeon listing. The repo often publishes multiple - # generations simultaneously, so picking the highest-version - # for each package independently can assemble a mismatched trio - # (e.g. torch 2.10 + torchvision 0.24). To prevent this, - # we identify the highest common minor version and downpair - # wheels if necessary to ensure a compatible set. + # Independent highest-version picks can assemble a mismatched trio (repo publishes multiple generations); downpair to the highest common minor. _torch_whl=$(_pick_radeon_wheel "torch" 2>/dev/null) || _torch_whl="" _tv_whl=$(_pick_radeon_wheel "torchvision" 2>/dev/null) || _tv_whl="" _ta_whl=$(_pick_radeon_wheel "torchaudio" 2>/dev/null) || _ta_whl="" _tri_whl=$(_pick_radeon_wheel "triton" 2>/dev/null) || _tri_whl="" - # Check that torch and torchaudio share the same X.Y public - # version prefix, and that torchvision's minor correctly - # pairs with torch's minor (torchvision = torch.minor + 15 - # since torch 2.4 -> torchvision 0.19 -> torch 2.9 -> - # torchvision 0.24). - # - # URL-decode each wheel name so %2B -> + before version - # extraction. Real Radeon wheel hrefs are percent-encoded - # (torch-2.10.0%2Brocm7.2.0...), so a plain [+-] terminator - # in the sed regex below would never match and - # _radeon_versions_match would stay false for every real - # listing, silently forcing a fallback to the generic - # ROCm index. + # Verify the X.Y pairing (torchvision = torch.minor + 15); URL-decode %2B -> + first (Radeon hrefs are percent-encoded). _extract_version() { _whl=$1 _pkg=$2 @@ -3197,11 +2829,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then _ta_ver=$(_extract_version "$_ta_whl" "torchaudio") _radeon_versions_match=false - # Kept release (_PREV_TORCH_PIN) wins here too: pick its exact - # patch (else the newest patch of its minor) plus the paired - # vision/audio wheels. Any gap falls back to the newest-trio - # search below, mirroring _install_torch_default_index, so a - # rerun never drifts to another release nor below the kept one. + # Kept release (_PREV_TORCH_PIN) wins here too: exact patch (else newest of its minor) + paired vision/audio; gaps fall back to the newest-trio search. if [ -n "$_PREV_TORCH_PIN" ]; then _prev_kept_base="${_PREV_TORCH_PIN#torch==}" _prev_kept_minor="${_prev_kept_base#*.}" @@ -3219,8 +2847,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then _ta_whl=$_kept_ta _tri_whl="" _radeon_versions_match=true - # Say so when the listing pruned the exact patch - # and a same-series build is installed instead. + # Say so when the listing pruned the exact patch. case "$(printf '%s' "${_kept_torch##*/}" | sed 's/%2[Bb]/+/g')" in "torch-${_prev_kept_base}"[+-]*) ;; *) substep "kept release ${_prev_kept_base} is not in the Radeon listing -- installing the closest 2.${_prev_kept_minor} series build instead" ;; @@ -3243,8 +2870,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then [ "$_tv_equiv_minor" -lt "$_target_minor" ] && _target_minor=$_tv_equiv_minor [ "$_ta_minor" -lt "$_target_minor" ] && _target_minor=$_ta_minor - # Loop downwards to find the first complete matching trio. - # This avoids aborting if the repo has gaps. + # Loop downwards to find the first complete matching trio (repo gaps). _attempts=0 while [ "$_attempts" -lt 5 ] && [ "$_target_minor" -ge 0 ]; do _expected_tv_minor=$((_target_minor + 15)) @@ -3267,8 +2893,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then _c_tv_major=${_c_tv_ver%%.*} _c_tv_minor=${_c_tv_ver#*.} - # Strict X.Y validation: allow patch versions to differ (e.g. torch 2.9.1 + vision 0.24.0) - # as long as the Major and Minor pairing is correct. + # Strict X.Y validation (patch may differ) on the major.minor pairing. if [ "$_c_torch_major" = "$_c_ta_major" ] && \ [ "$_c_torch_minor" = "$_c_ta_minor" ] && \ [ "$_c_tv_major" = "0" ] && \ @@ -3293,12 +2918,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then _install_torch_default_index else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." - # Pass explicit wheel URLs so the matched trio is - # installed together. --find-links lets uv discover - # the Radeon listing for any local lookup, and PyPI - # (not disabled) provides transitive deps like - # filelock / sympy / networkx which are not in the - # Radeon listing. + # Explicit wheel URLs install the matched trio together; --find-links exposes the listing, PyPI supplies transitive deps. if [ -n "$_tri_whl" ]; then run_install_cmd_retry "install triton + PyTorch" uv pip install --python "$_VENV_PY" \ --find-links "$_RADEON_BASE_URL" \ @@ -3321,10 +2941,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "installing PyTorch ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))..." _install_torch_default_index fi - # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). - # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm - # host stays in GGUF-only mode rather than pulling in bitsandbytes, - # which is only useful once torch is present for training. + # AMD ROCm: install bitsandbytes once after torch (--no-torch ROCm hosts stay GGUF-only). if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" fi @@ -3333,8 +2950,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "installing unsloth (this may take a few minutes)..." _build_unsloth_torch_overrides if [ "$SKIP_TORCH" = true ]; then - # No-torch: install unsloth + unsloth-zoo with --no-deps, then - # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. + # No-torch: install unsloth + unsloth-zoo --no-deps, then runtime deps --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" @@ -3370,8 +2986,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" _UNSLOTH_TORCH_OVERRIDES="" - # AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in - # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. + # AMD ROCm: repair torch if the unsloth install pulled CUDA torch from PyPI over the ROCm wheels. if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) if [ -z "$_has_hip" ]; then @@ -3397,10 +3012,7 @@ else fi # ── Enforce the installed torch flavor matches the detected GPU build ── -# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv -# keeps a stale torch==X+cpu against a GPU index and the venv silently trains on -# CPU. Reinstall the right wheel triplet when a GPU build is expected; if it -# can't be reinstalled, warn loudly. --no-torch / CPU-only / macOS: no-op. +# PEP 440 ignores the +cpu/+cuXXX/+rocm local label, so uv keeps a stale torch==X+cpu against a GPU index; reinstall the right triplet, else warn loudly. if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then _expected_torch_tag=$(_expected_torch_flavor_tag "$TORCH_INDEX_URL") # Only act when a GPU build is expected (cuXXX / rocm); cpu and unknown skip. @@ -3408,8 +3020,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") - # Repair when flavor is wrong AND the index is plain --default-index reinstallable - # (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only. + # Repair only when flavor is wrong AND the index is --default-index reinstallable. if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \ && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..." @@ -3431,8 +3042,7 @@ fi # ── Run studio setup ── tauri_log "STEP" "Running Unsloth setup" -# When --local, use the repo's own setup.sh directly. -# Otherwise, find it inside the installed package. +# --local uses the repo's setup.sh directly; otherwise find it in the installed package. SETUP_SH="" if [ "$STUDIO_LOCAL_INSTALL" = true ] && [ -f "$_REPO_ROOT/studio/setup.sh" ]; then SETUP_SH="$_REPO_ROOT/studio/setup.sh" @@ -3476,8 +3086,7 @@ _SKIP_FRONTEND=0 if [ "$TAURI_MODE" = true ]; then _SKIP_FRONTEND=1 fi -# Prepend UNSLOTH_STUDIO_HOME=$STUDIO_HOME to "$@" for env-override installs -# without word-splitting on whitespace paths. +# Prepend UNSLOTH_STUDIO_HOME for env-override installs without word-splitting whitespace paths. _run_setup_with_studio_home() { if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then UNSLOTH_STUDIO_HOME="$STUDIO_HOME" "$@" @@ -3503,11 +3112,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \ bash "$SETUP_SH" &2 echo " Move or remove it manually, then re-run the installer." >&2 exit 1 fi -# why: -sfn is atomic and -n prevents descent into a symlink-to-directory at -# the shim path (the directory guard above already rejects a real directory). +# -sfn is atomic and -n prevents descent into a symlink-to-directory at the shim path. ln -sfn "$VENV_DIR/bin/unsloth" "$_shim_path" case ":$PATH:" in @@ -3564,14 +3166,11 @@ case ":$PATH:" in esac # Non-Tauri installs keep shortcuts even if setup reports failure. -# create_studio_shortcuts gates persistent menu shortcuts on env-mode; -# launcher + studio.conf + icon are always written. if [ "$TAURI_MODE" != true ]; then create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS" fi -# If setup.sh failed, report and exit now. -# PATH and shortcuts are already set up so the user can fix and retry. +# If setup.sh failed, report and exit (PATH + shortcuts already set up for retry). if [ "$_SETUP_EXIT" -ne 0 ]; then echo "" step "error" "studio setup failed (exit code $_SETUP_EXIT)" "$C_ERR" @@ -3587,18 +3186,11 @@ if [ "$TAURI_MODE" = true ]; then exit 0 fi -# Warn if another 'unsloth' wins on PATH (different venv, system pip, etc). -# Users typing `unsloth studio` later would hit that binary instead of the -# one just installed; the runtime now falls back via UNSLOTH_STUDIO_HOME -# but the absolute path is still the most reliable launch. -# Uses the venv python (just created above) for path canonicalization so -# this works on macOS (BSD readlink has no -f) as well as Linux/WSL. +# Warn if another 'unsloth' wins on PATH; canonicalize via the venv python (BSD readlink lacks -f). _installed_bin="$VENV_DIR/bin/unsloth" _path_unsloth=$(command -v unsloth 2>/dev/null || true) if [ -n "$_path_unsloth" ] && [ -x "$VENV_DIR/bin/python" ]; then - # Canonicalize via the venv python (BSD readlink lacks -f on macOS). - # If either side fails to resolve, skip the check entirely rather than - # comparing raw paths (which would false-trigger on symlink targets). + # If either side fails to resolve, skip the check rather than compare raw paths. _canon() { "$VENV_DIR/bin/python" -c \ 'import os, sys; print(os.path.realpath(sys.argv[1]))' \ @@ -3624,9 +3216,7 @@ printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!" printf " ${C_DIM}%s${C_RST}\n" "$RULE" echo "" -# In interactive terminals, ask the user before starting Unsloth unless the -# caller explicitly disabled the post-install prompt. -# In non-interactive environments (Docker, CI, cloud-init) just print instructions. +# Interactive terminals prompt before starting; non-interactive environments just print instructions. if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then echo "" printf " Start Unsloth Studio now? [Y/n] " @@ -3639,11 +3229,7 @@ if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then case "${_reply:-y}" in [Yy]*|"") step "launch" "starting Unsloth Studio..." - # Detach stdin from the `curl | sh` pipe: as a foreground server the - # studio would otherwise drain the rest of this piped script, leaving - # the shell to die parsing the now-truncated tail (`unexpected fi`). - # trap '' INT: wait for studio's shutdown instead of racing the prompt. - # Subshell resets INT so the child still gets Ctrl+C (no inherited ignore). + # Detach stdin from the `curl | sh` pipe (the foreground server would drain the script's tail); trap '' INT waits for studio's shutdown, subshell resets INT for the child. trap '' INT # `|| ...`: capture the exit code without set -e aborting first. _LAUNCH_EXIT=0 diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 9921b83543..9c8c38626b 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -44,15 +44,10 @@ IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64" IS_MAC_ARM = IS_MACOS and platform.machine() == "arm64" IS_LINUX = sys.platform.startswith("linux") -# amd-smi auto-elevates on Windows (UAC/DiskPart prompt mid-install). This installer -# only spawns probes and pip/uv (no elevation), so set __COMPAT_LAYER=RunAsInvoker -# process-wide; amd-smi then runs un-elevated. setup.ps1 keeps per-call guards (it -# also spawns winget installers that need elevation). +# amd-smi auto-elevates on Windows (UAC/DiskPart prompt); RunAsInvoker keeps this installer's probes un-elevated. if IS_WINDOWS: os.environ.setdefault("__COMPAT_LAYER", "RunAsInvoker") -# torchcodec ships wheels only for manylinux_2_28_x86_64, macosx_12_0_arm64, -# and win_amd64. On other hosts the audio extras must be filtered out (the -# extras-no-deps step would otherwise fail), regardless of NO_TORCH. +# torchcodec ships wheels only for manylinux_2_28_x86_64, macosx_12_0_arm64, win_amd64; elsewhere filter the audio extras regardless of NO_TORCH. PLATFORM_LACKS_TORCHCODEC_WHEEL = ( (IS_LINUX and platform.machine() in {"aarch64", "arm64"}) or (IS_WINDOWS and platform.machine().lower() in {"arm64", "aarch64"}) @@ -60,8 +55,7 @@ PLATFORM_LACKS_TORCHCODEC_WHEEL = ( ) # ── ROCm / AMD GPU support ───────────────────────────────────────────────────── -# Detected ROCm (major, minor) -> best PyTorch wheel tag on -# download.pytorch.org. Checked newest-first (>=). +# Detected ROCm (major, minor) -> best PyTorch wheel tag, checked newest-first (>=). _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { (7, 2): "rocm7.2", # torch 2.11.0 (7, 1): "rocm7.1", # torch 2.10.0 @@ -73,12 +67,10 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { (6, 0): "rocm6.0", } -# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug). -# Mirrors *FloorMap in install.ps1 / setup.ps1; other arches ship <2.11 and stay bare. +# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug); mirrors install.ps1 / setup.ps1. _ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"}) -# pytorch.org rocmX.Y indexes KNOWN to ship torch 2.11 (rocm7.2 only today); don't -# floor an unknown newer rocm speculatively. Match install.sh / setup.ps1 / install.ps1. +# rocmX.Y indexes KNOWN to ship torch 2.11; never floor an unknown newer rocm speculatively. _ROCM_KNOWN_TORCH211_VERSIONS: frozenset[tuple[int, int]] = frozenset({(7, 2)}) # Per-tag pip specs; rocm7.2 ships torch 2.11.0 (older tags cap at 2.10.x). @@ -95,9 +87,7 @@ _ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = { "torchaudio>=2.4,<2.11.0", ), } -# Windows AMD per-arch companion pins for the repo.amd.com index (mirrors the install.ps1 / -# setup.ps1 floor maps): pinning stops the per-arch index (each published independently) from -# resolving an ABI-mismatched companion. Unlisted arches have no floor, so stay bare. +# Windows AMD per-arch companion pins for repo.amd.com: pinning stops the per-arch index resolving an ABI-mismatched companion; unlisted arches stay bare. _WINDOWS_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = { "gfx1201": _ROCM_TORCH_PKG_SPECS["rocm7.2"], "gfx1200": _ROCM_TORCH_PKG_SPECS["rocm7.2"], @@ -166,31 +156,18 @@ def _torch_index_leaf(url: str) -> str: return path.rstrip("/").rsplit("/", 1)[-1].lower() -# CUDA torch repair specs (see _ensure_cuda_torch). torch 2.11 is allowed (torchao -# 0.17 cpp loads cleanly, and the flash-attn/causal-conv1d/mamba wheels pass on 2.11). -# torchvision/torchaudio are pinned (not bare) so the exclusive --index-url can't -# resolve one built against a different torch major -> ABI mismatch. +# CUDA torch repair specs (see _ensure_cuda_torch): companions pinned so the exclusive --index-url can't resolve an ABI-mismatched torch major. _CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = ( "torch>=2.4,<2.12.0", "torchvision>=0.19,<0.27.0", "torchaudio>=2.4,<2.12.0", ) -# CPU torch repair specs (see _ensure_cpu_torch). Same bounds/reasoning as CUDA: the -# /cpu index also serves newer torch, so a bare trio could resolve out of range or ABI- -# mismatched. +# CPU torch repair specs (see _ensure_cpu_torch): the /cpu index also serves newer torch, so a bare trio could resolve ABI-mismatched. _CPU_TORCH_PKG_SPEC: tuple[str, str, str] = _CUDA_TORCH_PKG_SPEC -# torchao's cpp extensions are pinned to ONE torch release AND CUDA major. A torch -# mismatch just skips the cpp kernels (slow Python fallback); a CUDA mismatch fails -# to import ("libcudart.so.12: cannot open shared object file"). The torch pin is a -# range, so match torchao to the installed torch (table: pytorch/ao#2919): -# 2.9.x -> 0.14.0 -# 2.10.x, CUDA<=12 -> 0.16.0 (cpp built for 2.10, loads via the CUDA-12 wheel) -# 2.10.x, CUDA>=13 -> 0.17.0 (cu130: 0.16.0's CUDA-12 cpp crashes on load; 0.17.0 -# targets torch 2.11 so its cpp is cleanly skipped, not crashed) -# 2.11.x -> 0.17.0 (reachable via CUDA or ROCm rocm7.2) -# Unknown/older torch keeps the conservative default. +# torchao's cpp extensions are pinned to ONE torch release AND CUDA major (table: pytorch/ao#2919): +# 2.9.x -> 0.14.0; 2.10.x CUDA<=12 -> 0.16.0; 2.10.x CUDA>=13 -> 0.17.0; 2.11.x -> 0.17.0; else default. _TORCHAO_DEFAULT_SPEC = "torchao==0.14.0" _TORCHAO_TORCH_210_SPEC = "torchao==0.16.0" _TORCHAO_TORCH_210_CUDA13_SPEC = "torchao==0.17.0" @@ -221,8 +198,7 @@ def _select_torchao_spec(torch_version: str | None) -> str: release = str(torch_version).split("+", 1)[0] # drop +cu130/+rocm6.4/+cpu parts = release.split(".") try: - # Strip any pre-release/dev suffix from the minor (e.g. '10rc1' -> '10'), - # matching wheel_utils.probe_torch_wheel_env. + # Strip any pre-release/dev suffix from the minor (e.g. '10rc1' -> '10'). minor_str = re.sub(r"[^0-9].*", "", parts[1]) if len(parts) > 1 else "" major, minor = int(parts[0]), int(minor_str) except (IndexError, ValueError): @@ -300,9 +276,7 @@ def _installed_torch_is_windows_rocm() -> bool: return probe.returncode == 0 and bool(lines and lines[-1] == "yes") -# constraints.txt caps new anyio resolutions at <4.14 (#6483), but an install -# from before the cap existed can already be stuck at 4.14+, which later -# constrained installs won't touch since it already satisfies mcp/fastmcp. +# constraints.txt caps anyio <4.14 (#6483), but a pre-cap install can be stuck at 4.14+ which constrained installs won't touch. _ANYIO_BAD_FLOOR = (4, 14) @@ -335,14 +309,12 @@ def _repair_bad_anyio() -> None: ) -# AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/). -# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs. +# AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/); override with UNSLOTH_ROCM_WINDOWS_MIRROR. _ROCM_WINDOWS_INDEX_BASE = ( os.environ.get("UNSLOTH_ROCM_WINDOWS_MIRROR") or "https://repo.amd.com/rocm/whl" ).rstrip("/") -# gfx arch → AMD index arch-family suffix; each family is a separate -# pip index on repo.amd.com. +# gfx arch → AMD index arch-family suffix; each family is a separate pip index on repo.amd.com. _GFX_TO_AMD_INDEX_ARCH: dict[str, str] = { "gfx1201": "gfx120X-all", "gfx1200": "gfx120X-all", # RDNA 4 @@ -356,9 +328,7 @@ _GFX_TO_AMD_INDEX_ARCH: dict[str, str] = { "gfx908": "gfx908", # MI200/MI100 } -# bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix -# (bnb PR #1887, post-0.49.2). bnb <= 0.49.2 NaNs at decode shape on every -# AMD GPU. Drop the pin once bnb 0.50+ ships on PyPI. +# bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix (bnb #1887); bnb <=0.49.2 NaNs on AMD. Drop once bnb 0.50+ ships on PyPI. _BNB_ROCM_PRERELEASE_URLS: dict[str, str] = { "x86_64": ( "https://github.com/bitsandbytes-foundation/bitsandbytes/releases/" @@ -370,9 +340,7 @@ _BNB_ROCM_PRERELEASE_URLS: dict[str, str] = { "download/continuous-release_main/" "bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl" ), - # Windows ROCm wheel ships libbitsandbytes_rocm{VER}.dll. BNB's HIP - # auto-detect may mismatch the DLL suffix, so we scan the wheel and set - # BNB_ROCM_VERSION in _install_bnb_windows_rocm() and worker.py. + # Windows ROCm wheel ships libbitsandbytes_rocm{VER}.dll; the wheel is scanned and BNB_ROCM_VERSION set to match. "win_amd64": ( "https://github.com/bitsandbytes-foundation/bitsandbytes/releases/" "download/continuous-release_main/" @@ -409,8 +377,7 @@ def _path_inside_venv(path: str) -> bool: try: # realpath (not abspath): resolve symlinks/8.3 names so an aliased venv matches. _root = os.path.normcase(os.path.realpath(sys.prefix)) - # Guard a root-dir prefix (C:\ or /): commonpath would match every path on - # it. A venv is never at root, so treat that as outside. + # Root-dir prefix (C:\ or /) would commonpath-match everything; a venv is never at root. if os.path.dirname(_root) == _root: return False return os.path.normcase(os.path.commonpath([os.path.realpath(path), _root])) == _root @@ -448,9 +415,7 @@ def _amd_smi_allowed() -> bool: return True if flag in ("0", "false", "no", "off"): return False - # A real HIP SDK lets amd-smi run un-elevated; hipinfo-on-PATH is the proxy. - # Ignore the venv hipInfo.exe (AMD wheel via bnb fix): not a HIP SDK, doesn't - # stop amd-smi's DiskPart UAC. + # hipinfo-on-PATH proxies a real HIP SDK; the venv hipInfo.exe is not one. if _external_hipinfo_on_path(): return True for _var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"): @@ -473,16 +438,13 @@ def _detect_rocm_version() -> tuple[int, int] | None: try: with open(path) as fh: parts = fh.read().strip().split("-")[0].split(".") - # Explicit length guard: don't rely on the broad except below to - # swallow IndexError on a single-component version (e.g. "6\n"). + # Length guard for single-component versions (e.g. "6\n"). if len(parts) >= 2: return int(parts[0]), int(parts[1]) except Exception: pass - # Try amd-smi version (outputs "... | ROCm version: X.Y.Z"). - # Gated off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt); - # hipconfig below covers that case. + # Try amd-smi version ("ROCm version: X.Y.Z"); gated off on Windows w/o a HIP SDK (UAC/DiskPart prompt). amd_smi = shutil.which("amd-smi") if _amd_smi_allowed() else None if amd_smi: try: @@ -519,10 +481,7 @@ def _detect_rocm_version() -> tuple[int, int] | None: except Exception: pass - # Distro package-manager fallbacks: package-managed ROCm can expose GPUs via - # rocminfo/amd-smi but lack /opt/rocm/.info/version and hipconfig, so probe - # dpkg (Debian/Ubuntu) and rpm (RHEL/Fedora/SUSE) for the rocm-core version. - # Matches install.sh::get_torch_index_url so `studio update` == fresh install. + # dpkg/rpm rocm-core fallback: package-managed ROCm may lack .info/version and hipconfig. Matches install.sh::get_torch_index_url. for cmd in ( ["dpkg-query", "-W", "-f=${Version}\n", "rocm-core"], ["rpm", "-q", "--qf", "%{VERSION}\n", "rocm-core"], @@ -594,8 +553,7 @@ def _detect_windows_gfx_arch() -> str | None: def _dedup_pick(tokens: list[str]) -> "str | None": if not tokens: return None - # Index into the full ordered list so HIP_VISIBLE_DEVICES addresses - # GPU N on mixed-arch hosts, then return that arch. + # Index the full ordered list so HIP_VISIBLE_DEVICES addresses GPU N on mixed-arch hosts. return tokens[_pick_visible_index(len(tokens))] # 2. hipinfo via PATH, then HIP_PATH\bin / ROCM_PATH\bin. @@ -609,10 +567,7 @@ def _detect_windows_gfx_arch() -> str | None: hipinfo = _candidate break if not hipinfo: - # 2b. AMD torch wheels ship hipInfo.exe into the venv Scripts dir - # (next to python.exe); resolvable even on driver-only hosts with no - # SDK install at all. Lets `studio update` re-detect the arch on a - # venv that already has the AMD wheel. + # 2b. AMD torch wheels ship hipInfo.exe into venv Scripts; lets `studio update` re-detect on driver-only hosts. _venv_hipinfo = os.path.join(os.path.dirname(sys.executable), "hipInfo.exe") if os.path.isfile(_venv_hipinfo): hipinfo = _venv_hipinfo @@ -624,13 +579,9 @@ def _detect_windows_gfx_arch() -> str | None: stderr = subprocess.DEVNULL, timeout = 10, ) - # Accept partial output even when hipinfo crashes (e.g. 0xC0000005 / - # STATUS_ACCESS_VIOLATION on some RDNA 4 hosts): a gcnArchName in stdout - # means the device was enumerated pre-crash, so the arch is trustworthy. - # Ignoring it causes a silent CPU PyTorch fallback (issue #6043). + # Accept partial output even when hipinfo crashes (0xC0000005 on some RDNA 4, #6043): a pre-crash gcnArchName is trustworthy. text = result.stdout.decode(errors = "replace") - # findall gets every gcnArchName line so multi-GPU hosts are - # enumerable and HIP_VISIBLE_DEVICES selects correctly. + # findall gets every gcnArchName line so HIP_VISIBLE_DEVICES selects on multi-GPU hosts. _tokens = [ t.strip().lower() for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text) ] @@ -640,9 +591,7 @@ def _detect_windows_gfx_arch() -> str | None: except Exception: pass - # 3. amd-smi fallback -- runtime-only Radeon installs ship amd-smi but no hipinfo. - # Gated off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt); the arch - # arrives via --rocm-gfx / name inference there, so this is only needed when safe. + # 3. amd-smi fallback (runtime-only Radeon installs lack hipinfo); gated off on Windows w/o a HIP SDK (UAC/DiskPart prompt). amd_smi = shutil.which("amd-smi") if _amd_smi_allowed() else None if amd_smi: for _args in (("static", "--asic"), ("list",)): @@ -671,11 +620,7 @@ def _detect_windows_gfx_arch() -> str | None: except Exception: continue - # 4. Last resort: GPU marketing name via WMI → arch table. Driver-only - # hosts (Adrenalin, no HIP SDK) have neither hipinfo nor amd-smi - # (amd-smi does not exist on Windows at all), but the display driver - # always knows the GPU name. Mirrors setup.ps1's $nameArchTable so a - # standalone `studio update` can repair a CPU-only venv on such hosts. + # 4. Last resort: GPU marketing name via WMI → arch table (driver-only hosts have neither hipinfo nor amd-smi); mirrors setup.ps1's $nameArchTable. try: result = subprocess.run( [ @@ -705,10 +650,7 @@ def _detect_windows_gfx_arch() -> str | None: return None -# GPU marketing-name → gfx arch table, mirroring setup.ps1's $nameArchTable. -# Most-specific first; first match wins. Covers only arches the ROCm -# prebuilts / AMD Windows torch indexes support; unknown names return None -# (callers then fall back cleanly to CPU). +# GPU marketing-name → gfx arch table (mirrors setup.ps1's $nameArchTable); most-specific first; unknown names return None (CPU fallback). _WIN_GPU_NAME_ARCH_TABLE: "list[tuple[str, str]]" = [ (r"9070 XT|9080", "gfx1201"), # RDNA 4 (Radeon RX 9070 XT / 9080) (r"9070|9060", "gfx1200"), # RDNA 4 (Radeon RX 9070 / 9060) @@ -768,8 +710,7 @@ def _detect_bnb_rocm_dll_ver() -> str | None: m = re.search(r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(dll)) if m: all_vers.append(m.group(1)) - # Highest numeric suffix wins (e.g. "713" over "72"); glob order is not - # guaranteed, so sort rather than take the first match. + # Highest numeric suffix wins ("713" over "72"); glob order is not guaranteed. return max(all_vers, key = lambda v: int(v)) if all_vers else None @@ -810,8 +751,7 @@ def _persist_bnb_rocm_version(version: str) -> bool: existing = ( sitecustomize_path.read_text(encoding = "utf-8") if sitecustomize_path.exists() else "" ) - # Strip all managed regions, including one whose END marker was lost to - # an interrupted write, then append exactly one fresh block. + # Strip all managed regions (even END-marker-less from an interrupted write), append one fresh block. pattern = re.compile( rf"{re.escape(_BNB_ROCM_SITECUSTOMIZE_BEGIN)}.*?" rf"(?:{re.escape(_BNB_ROCM_SITECUSTOMIZE_END)}\n?|\Z)", @@ -851,10 +791,7 @@ def _has_rocm_gpu() -> bool: if _has_usable_nvidia_gpu(): return False for cmd, check_fn in ( - # rocminfo: look for a real gfx GPU id (3-4 chars, nonzero first digit). - # gfx000 is the CPU agent; ROCm 6.1+ also emits generic ISA lines like - # "gfx11-generic"/"gfx9-4-generic" with only 1-2 digits before the dash, - # which must not be treated as a real GPU. + # rocminfo: real gfx GPU id only (gfx000 = CPU agent; generic "gfx11-generic" ISA lines are not GPUs). ( ["rocminfo"], lambda out: bool(re.search(r"gfx[1-9][0-9a-z]{2,3}", out.lower())), @@ -868,8 +805,7 @@ def _has_rocm_gpu() -> bool: exe = shutil.which(cmd[0]) if not exe: continue - # Skip amd-smi on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt); - # rely on rocminfo / the sysfs fallback there. + # Skip amd-smi on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt). if cmd[0] == "amd-smi" and not _amd_smi_allowed(): continue try: @@ -886,14 +822,8 @@ def _has_rocm_gpu() -> bool: if result.returncode == 0 and result.stdout.strip(): if check_fn(result.stdout): return True - # sysfs KFD topology fallback (Linux only) -- matches install.sh's runtime-only - # detection. On minimal package-managed installs (no rocminfo / amd-smi), the - # kernel exposes AMD GPUs via /sys/class/kfd so `studio update` can still repair. - # - # Guard: reject any KFD node whose properties file reports a non-AMD vendor. The - # NVIDIA open kernel module (driver 560+) registers KFD nodes with a non-zero - # gpu_id and vendor_id 4318 (0x10DE), not the AMD 4098 (0x1002); without this - # check the fallback returns True on NVIDIA-only hosts, installing ROCm wheels. + # sysfs KFD topology fallback (Linux, matches install.sh): minimal installs lack rocminfo/amd-smi. + # Reject non-AMD vendors: the NVIDIA open kernel module also registers KFD nodes (vendor 0x10DE). if sys.platform != "win32": try: kfd_nodes = "/sys/class/kfd/kfd/topology/nodes" @@ -907,10 +837,7 @@ def _has_rocm_gpu() -> bool: continue if not gpu_id or gpu_id == "0": # gpu_id 0 = CPU node continue - # Require AMD vendor_id 4098 (0x1002). KFD properties files exist - # on every kernel exposing /sys/class/kfd, so a missing file means - # AMD ownership is unconfirmed -- skip the node rather than risk a - # false positive (e.g. NVIDIA open-driver KFD nodes lacking it). + # Require AMD vendor_id 4098 (0x1002); a missing properties file leaves AMD ownership unconfirmed -- skip. props_path = os.path.join(kfd_nodes, entry, "properties") try: with open(props_path) as fh: @@ -956,8 +883,7 @@ def _has_usable_nvidia_gpu() -> bool: return True except Exception: pass - # Fallback: the NVIDIA driver exposes one subdirectory per GPU under - # /proc/driver/nvidia/gpus/ on Linux regardless of nvidia-smi state. + # Fallback: /proc/driver/nvidia/gpus/ has one subdir per GPU regardless of nvidia-smi state. if sys.platform != "win32": try: gpu_dir = "/proc/driver/nvidia/gpus" @@ -1037,10 +963,7 @@ def _install_bnb_windows_rocm() -> bool: ) if not _ok: return False - # Detect the actual ROCm DLL suffix in the wheel and set BNB_ROCM_VERSION so bnb - # loads the right DLL regardless of torch.version.hip (the wheel may ship "72" - # while torch reports 7.13). The worker subprocess inherits it; fall back to "72" - # if detection fails (e.g. a no-op / dry-run install). + # Detect the ROCm DLL suffix and set BNB_ROCM_VERSION (wheel may ship "72" while torch reports 7.13); fall back to "72". _env_ver = os.environ.get("BNB_ROCM_VERSION") _env_is_persisted_default = ( os.environ.get(_BNB_ROCM_VERSION_SOURCE_ENV) == _BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE @@ -1055,11 +978,7 @@ def _install_bnb_windows_rocm() -> bool: _persist_detected_version = True if _persist_detected_version: _persist_bnb_rocm_version(_ver) - # Make hipInfo.exe (shipped into venv Scripts by the AMD torch wheel) resolvable - # via PATH for this process and every child python (import checks, precompile): - # bitsandbytes runs hipinfo.exe at import to detect the GPU arch and logs a scary - # (harmless) ERROR + WARNING when it is missing. Scripts is on PATH only for an - # activated venv, which neither Unsloth nor the installer's children ever do. + # Put venv Scripts (hipInfo.exe from the AMD torch wheel) on PATH: bnb probes hipinfo.exe at import and logs a scary (harmless) ERROR when missing. _scripts_dir = os.path.dirname(sys.executable) if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")) and not shutil.which( "hipinfo.exe" @@ -1138,8 +1057,7 @@ def _is_pip_rocm_family_leaf(leaf: str) -> bool: rocm7.2-private) starts with "rocm" but is a custom pin the verbatim path owns, so match EXACTLY. Mirrors install.sh / setup.ps1. """ - # gfx must be followed by a digit (gfx90a, gfx1151, gfx120X-all): a gfx-prefixed - # custom leaf (gfx-private) is a verbatim pin, like rocm7.2-private. + # gfx must be followed by a digit; a gfx-private custom leaf is a verbatim pin. return bool(re.fullmatch(r"rocm\d+(?:\.\d+)?", leaf)) or bool(re.match(r"gfx\d", leaf)) @@ -1280,8 +1198,7 @@ def _ensure_cuda_torch() -> None: Only repairs when torch actually links against HIP/ROCm. Healthy CUDA torch and deliberate CPU-only torch are left untouched. """ - # Respect install.sh's backend: only "" (standalone update) or "cuda" force CUDA - # wheels; "rocm"/"cpu"/unrecognised are deliberate. + # Respect install.sh's backend: only "" (standalone update) or "cuda" force CUDA wheels. if _TORCH_BACKEND not in ("", "cuda"): return # An explicit unknown-family pin was applied VERBATIM at install time; leave it alone. @@ -1293,11 +1210,9 @@ def _ensure_cuda_torch() -> None: # Never undo a deliberate ROCm install (setup.ps1 sets this marker). if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": return - # An explicit CUDA pin (headless / CI cross-install) commits to CUDA wheels and skips ALL - # GPU probing, so it clears both the CUDA_VISIBLE_DEVICES hide gate and the NVIDIA gate below. + # An explicit CUDA pin commits to CUDA wheels and skips ALL GPU probing gates below. _cuda_pinned = _explicit_cuda_torch_index_url() is not None - # CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU; never force CUDA - # wheels over that unless a CUDA index is pinned. + # CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU; honour it unless a CUDA index is pinned. _cvd = os.environ.get("CUDA_VISIBLE_DEVICES") if not _cuda_pinned and _cvd is not None and _cvd.strip() in ("", "-1"): return @@ -1305,9 +1220,7 @@ def _ensure_cuda_torch() -> None: if not _cuda_pinned and not _has_usable_nvidia_gpu(): return - # Classify the installed torch: "hip" (ROCm poisoning signature), "cuda" (healthy), - # or "cpu". A non-zero exit means torch is missing/un-importable: without a pin the - # base install owns it, but a pinned CUDA index reinstalls it below. + # Classify installed torch: "hip" (ROCm poisoning signature), "cuda", or "cpu"; non-zero exit = missing/un-importable. try: probe = subprocess.run( [ @@ -1330,9 +1243,7 @@ def _ensure_cuda_torch() -> None: except (OSError, subprocess.TimeoutExpired): return if probe.returncode != 0: - # torch present but can't import. Without a pin the base install owns it; but an - # explicit CUDA pin forces this pass (failed probe) and the base update won't - # reinstall an already-installed torch, so reinstall from the pin (self-resolving). + # Un-importable torch: only an explicit CUDA pin reinstalls here (the base update won't touch an installed torch). if not _cuda_pinned: return index_url = _detect_cuda_torch_index_url() @@ -1360,9 +1271,7 @@ def _ensure_cuda_torch() -> None: if not _marker_lines: return _marker, _, _installed_cu = _marker_lines[-1].partition("|") - # Reinstall CUDA torch on a ROCm build on an NVIDIA host (poisoning signature), or when a - # CUDA index is pinned but the venv has the wrong family (CPU or a different cuXXX). A - # healthy match, or a CPU wheel with no CUDA pin, is left alone. + # Reinstall on ROCm-on-NVIDIA poisoning or a pinned-CUDA family mismatch; healthy/CPU-no-pin left alone. _pin = _explicit_torch_index_url() _pin_leaf = _torch_index_leaf(_pin) if _pin else "" _pinned_cuda = _is_cuda_family_leaf(_pin_leaf) @@ -1371,8 +1280,7 @@ def _ensure_cuda_torch() -> None: elif _marker == "cpu" and _pinned_cuda: _why = "torch is a CPU build but an explicit CUDA index is pinned" elif _marker == "cuda" and _pinned_cuda and _installed_cu != _pin_leaf: - # Installed cuXXX differs from the pin. An untagged build (empty) counts too: - # the family can't be confirmed, so reinstall to enforce it (idempotent). + # Installed cuXXX differs from the pin; an untagged build counts too (family unconfirmed, idempotent). _installed_desc = _installed_cu if _installed_cu else "an untagged CUDA build" _why = f"torch is {_installed_desc} but the pinned CUDA index is {_pin_leaf}" else: @@ -1411,8 +1319,7 @@ def _ensure_cpu_torch() -> None: if pin is None: return - # Classify the installed torch family. A non-zero exit means torch is missing or - # un-importable: the explicit CPU pin reinstalls it below. + # Classify the installed torch family; non-zero exit = missing/un-importable -> the CPU pin reinstalls below. try: probe = subprocess.run( [ @@ -1434,9 +1341,7 @@ def _ensure_cpu_torch() -> None: except (OSError, subprocess.TimeoutExpired): return if probe.returncode != 0: - # torch present but can't import. The explicit CPU pin forces this pass (failed - # probe) and the base update won't reinstall an already-installed torch, so - # reinstall from the pin (self-resolving, no loop). + # Un-importable torch: reinstall from the explicit CPU pin (self-resolving, no loop). _torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC print( f" torch cannot import but an explicit CPU index is pinned -- reinstalling " @@ -1466,8 +1371,7 @@ def _ensure_cpu_torch() -> None: " torch is a GPU build but an explicit CPU index is pinned -- reinstalling " f"CPU torch from {_strip_index_url_credentials(pin)}" ) - # Pin the supported torch<2.11 family (the /cpu index now serves 2.11+, so a bare - # trio could resolve out of range or ABI-mismatched). + # Pin the supported torch family (the /cpu index now serves 2.11+). _torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC pip_install( "CPU torch repair", @@ -1492,15 +1396,13 @@ def _ensure_rocm_torch() -> None: Uses pip_install() to respect uv, constraints, and --python targeting. """ global _rocm_windows_torch_installed - # install.sh's resolved backend is authoritative: skip ROCm when it already chose a - # non-ROCm family (avoids re-detecting in a subprocess that may see a different env). + # install.sh's resolved backend is authoritative: skip ROCm for a non-ROCm family. if _TORCH_BACKEND in ("cuda", "cpu"): return # An explicit unknown-family pin was applied VERBATIM at install time; leave it alone. if _explicit_unknown_family_torch_index_url() is not None: return - # setup.ps1 sets this after installing AMD wheels; skip only when torch is actually - # importable as ROCm (a wiped venv leaves a stale env-var that must not suppress it). + # setup.ps1's marker; trust it only when torch actually imports as ROCm (a wiped venv leaves a stale env-var). if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": _torch_ok = False try: @@ -1524,8 +1426,7 @@ def _ensure_rocm_torch() -> None: pass if _torch_ok: _rocm_windows_torch_installed = True - # ROCm torch is already installed, but the AMD Windows BNB wheel is still - # needed (the PyPI bitsandbytes ships only CUDA DLLs, fails on ROCm). + # AMD Windows BNB wheel still needed (PyPI bitsandbytes ships only CUDA DLLs). _install_bnb_windows_rocm() return # torch was wiped between runs; fall through to the full install path @@ -1533,10 +1434,7 @@ def _ensure_rocm_torch() -> None: return if IS_WINDOWS: - # An explicit ROCm-family pin commits to ROCm wheels regardless of the visible - # GPU and overrides the public per-arch index (mirrors the Linux pin handling - # below): after a pinned setup.ps1 install fails to CPU, this repair must retry - # the PINNED index, not repo.amd.com. + # An explicit ROCm-family pin commits to ROCm wheels and overrides the public per-arch index: retry the PINNED index, not repo.amd.com. _win_rocm_pin = _explicit_rocm_torch_index_url() if _win_rocm_pin is None and _has_usable_nvidia_gpu(): return @@ -1574,14 +1472,11 @@ def _ensure_rocm_torch() -> None: f" {gfx_arch or 'pinned ROCm index'} (Windows) -- installing torch from " f"{_strip_index_url_credentials(index_url)}" ) - # Pin companions for the arches install.ps1/setup.ps1 pin (gfx120X / Strix) - # so the per-arch index resolves an ABI-consistent trio; other arches stay bare. + # Pin companions for the arches install.ps1/setup.ps1 pin so the per-arch index resolves an ABI-consistent trio. _torch_pkg, _vision_pkg, _audio_pkg = _WINDOWS_ROCM_TORCH_PKG_SPECS.get( gfx_arch, ("torch", "torchvision", "torchaudio") ) - # Nonfatal: a transient AMD-index failure must not abort the install. - # --force-reinstall resolves before uninstalling, so a failed index keeps the - # existing build intact; let the user retry. + # Nonfatal: --force-reinstall resolves before uninstalling, so a failed index keeps the existing build. if not pip_install_try( f"ROCm torch (Windows, {gfx_arch or 'pinned'})", "--force-reinstall", @@ -1598,14 +1493,9 @@ def _ensure_rocm_torch() -> None: "later to retry ROCm." ) return - # ROCm torch is installed (or already was); flag it so later phases - # do not overwrite it with the generic CPU torch wheel. BNB is a - # separate dependency -- a BNB install failure must NOT roll back the - # torch ROCm install. + # Flag ROCm torch installed so later phases don't overwrite it; a BNB failure must NOT roll it back. _rocm_windows_torch_installed = True - # Always install AMD Windows bitsandbytes -- the PyPI wheel ships only - # CUDA DLLs and fails on ROCm. Install even when torch was already a - # ROCm build so `studio update` repairs a broken bnb. + # Always install AMD Windows bitsandbytes (PyPI wheel ships only CUDA DLLs); also repairs a broken bnb on update. if not _install_bnb_windows_rocm(): print( " Warning: AMD Windows bitsandbytes install failed; " @@ -1616,15 +1506,13 @@ def _ensure_rocm_torch() -> None: # ── Linux x86_64 only: PyTorch ROCm wheels are not published for aarch64 ── if platform.machine().lower() not in {"x86_64", "amd64"}: return - # An explicit ROCm pin commits to ROCm wheels regardless of the visible GPU (headless / CI). - # Mirror _ensure_cuda_torch: skip the NVIDIA/no-AMD/unreadable gates. + # An explicit ROCm pin commits to ROCm wheels regardless of the visible GPU (headless / CI); skip the GPU gates. _rocm_pin = _explicit_rocm_torch_index_url() if _rocm_pin is None: # NVIDIA takes precedence on mixed hosts (only if a GPU is usable). if _has_usable_nvidia_gpu(): return - # _has_rocm_gpu() (rocminfo / amd-smi rows) is the authoritative AMD-host signal; - # the old /opt/rocm-or-hipcc gate broke runtime-only ROCm installs. + # _has_rocm_gpu() is the authoritative AMD-host signal (runtime-only ROCm installs lack /opt/rocm). if not _has_rocm_gpu(): return # no AMD GPU visible @@ -1633,13 +1521,10 @@ def _ensure_rocm_torch() -> None: if _rocm_pin is None: print(" ROCm detected but version unreadable -- skipping torch reinstall") return - # Explicit pin: the pinned leaf drives the install, so an unreadable host version - # is fine (sentinel keeps ver comparisons defined). + # Explicit pin drives the install; sentinel keeps ver comparisons defined. ver = (0, 0) - # Probe whether torch links against HIP, capturing the installed ROCm tag for pin-mismatch - # detection. Emit ONE "|" line: marker (HIP version, "rocm" sentinel, - # or empty for CPU/CUDA) before "|", wheel version after. + # Probe HIP linkage; emit ONE "|" line for pin-mismatch detection. try: probe = subprocess.run( [ @@ -1649,8 +1534,7 @@ def _ensure_rocm_torch() -> None: "import torch; " "hip=getattr(torch.version,'hip','') or ''; " "ver=getattr(torch,'__version__','').lower(); " - # HIP version if present, else a "rocm" sentinel when only the - # version string flags ROCm; empty marker = CPU/CUDA torch. + # HIP version, "rocm" sentinel, or empty marker = CPU/CUDA torch. "marker=hip if hip else ('rocm' if 'rocm' in ver else ''); " "print(marker + '|' + ver)" ), @@ -1673,9 +1557,7 @@ def _ensure_rocm_torch() -> None: # A "|"-delimited line is required; without it treat HIP as absent -> reinstall. has_hip_torch = bool(_sep) and _hip_marker != "" - # An explicit ROCm pin whose family differs from the installed torch must reinstall, else a - # rocm7.2/gfx* pin over an older +rocm6.4/7.1 build never applies. Version-tag heuristic - # only: a same-tag per-arch switch (gfx1151 -> gfx120X-all, both +rocm7.13.0) isn't detectable. + # A ROCm pin whose family differs from the installed torch must reinstall; version-tag heuristic only (same-tag per-arch switch undetectable). _rocm_pin_mismatch = ( _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver) if (has_hip_torch and _rocm_pin is not None) @@ -1684,9 +1566,7 @@ def _ensure_rocm_torch() -> None: rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch - # Strix Halo / Point (gfx1151 / gfx1150) segfault under ROCm 7.1 in torch._grouped_mm; - # AMD's per-gfx repo ships 2.11.0+rocm7.13.0 with the fix, so route those hosts there - # (mirrors install.sh). On mixed hosts, reroute only when HIP's runtime GPU is the Strix one. + # Strix Halo / Point (gfx1151/gfx1150) segfault in torch._grouped_mm under ROCm 7.1; route to AMD's per-gfx repo (2.11.0+rocm7.13.0 fix), Strix-runtime-only on mixed hosts. _strix_override_url: "str | None" = None _strix_override_pkgs: "tuple[str, str, str] | None" = None # An explicit ROCm pin is authoritative: never auto-reroute it. @@ -1695,8 +1575,7 @@ def _ensure_rocm_torch() -> None: _strix_gfx = {"gfx1151", "gfx1150"} _detected_strix = _strix_gfx.intersection(gfx_codes) if _detected_strix: - # Runtime-visible GPU (HIP_VISIBLE_DEVICES index into gfx_codes, else first); - # skip the override unless it's Strix. + # Runtime-visible GPU (HIP_VISIBLE_DEVICES index, else first) must be Strix. _runtime_gfx = gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None if _runtime_gfx in _strix_gfx: _selected_gfx = _runtime_gfx @@ -1706,8 +1585,7 @@ def _ensure_rocm_torch() -> None: _strix_override_url = f"{_amd_mirror}/{_selected_gfx}/" _strix_override_pkgs = ( "torch>=2.11.0,<2.12.0", - # Pin companions to the 2.11.x range: the exclusive --index-url could - # otherwise resolve a build for a different torch major (ABI mismatch). + # Pin companions to 2.11.x (exclusive --index-url could resolve ABI-mismatched). "torchvision>=0.26.0,<0.27.0", "torchaudio>=2.11.0,<2.12.0", ) @@ -1727,8 +1605,7 @@ def _ensure_rocm_torch() -> None: f" skipping AMD per-gfx index override.\n" ) - # The Strix override must fire even when has_hip_torch is True: an existing - # torch.version.hip == "7.1" is exactly the broken combo it repairs. + # Strix override fires even when has_hip_torch: hip == "7.1" is exactly the broken combo. if _strix_override_url is not None and _strix_override_pkgs is not None: index_url = _strix_override_url _torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs @@ -1749,8 +1626,7 @@ def _ensure_rocm_torch() -> None: ) rocm_torch_ready = True elif not has_hip_torch or _rocm_pin_mismatch: - # Reinstall when torch is not ROCm yet, OR a ROCm build's family differs from a pin. - # Honour a ROCm pin verbatim; else pick the newest wheel tag <= host. + # Reinstall when torch is not ROCm yet OR a pin family differs; honour a ROCm pin verbatim, else newest tag <= host. _override_idx = _explicit_rocm_torch_index_url() if _override_idx is not None: index_url = _override_idx @@ -1770,8 +1646,7 @@ def _ensure_rocm_torch() -> None: if _override_idx is None: index_url = f"{_PYTORCH_WHL_BASE}/{tag}" print(f" ROCm torch -- installing from {_strip_index_url_credentials(index_url)}") - # Only the _grouped_mm-bug gfx arches need the 2.11 spec; other gfx indexes ship - # <2.11 and stay on the default range (matches install.ps1 / setup.ps1). + # Only the _grouped_mm-bug gfx arches need the 2.11 spec (matches install.ps1 / setup.ps1). if tag in _ROCM_GFX_TORCH211_LEAVES: _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["rocm7.2"] elif tag.startswith("gfx"): @@ -1793,10 +1668,7 @@ def _ensure_rocm_torch() -> None: ) rocm_torch_ready = True - # Install bitsandbytes only when torch links against ROCm. Prefers the - # continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix), falling back - # to PyPI when the pre-release wheel won't install. Use pip for the - # pre-release wheel because uv rejects its filename/metadata version mismatch. + # bitsandbytes only when torch links ROCm; prefer the pre-release wheel (bnb #1887), pip not uv (filename/metadata version mismatch). if rocm_torch_ready: _bnb_url = _bnb_rocm_prerelease_url() _bnb_installed = False @@ -1828,9 +1700,6 @@ def _ensure_rocm_torch() -> None: ) -# _uv_safe_path is imported from backend.utils.uv_path_safety (shared with mlx_repair). - - def _windows_hidden_subprocess_kwargs() -> dict[str, object]: """Return Windows-only subprocess kwargs that suppress console windows.""" if not IS_WINDOWS: @@ -1868,11 +1737,9 @@ def _infer_no_torch() -> bool: NO_TORCH = _infer_no_torch() -# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() ("cuda", "rocm", -# "cpu"; empty = standalone `studio update`, where we re-detect). +# UNSLOTH_TORCH_BACKEND is set by install.sh ("cuda"/"rocm"/"cpu"; empty = standalone `studio update`). _TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower() -# Standalone update with an explicit pin: derive the backend from the override (classify on -# the final URL/family segment, mirroring install.sh) instead of re-probing the GPU. +# Standalone update with an explicit pin: derive the backend from the override leaf (mirrors install.sh). if not _TORCH_BACKEND: _idx_override = ( os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip() @@ -1884,9 +1751,7 @@ if not _TORCH_BACKEND: elif _idx_leaf == "cpu": _TORCH_BACKEND = "cpu" elif _is_cuda_family_leaf(_idx_leaf): - # Require a digit after "cu" so /current or /custom is NOT branded CUDA (a wrong backend - # makes _ensure_rocm_torch return early on AMD hosts). An unknown leaf keeps "" so the - # helpers probe the GPU. + # Require a digit after "cu" so /current or /custom is NOT branded CUDA; an unknown leaf keeps "" (helpers probe the GPU). _TORCH_BACKEND = "cuda" @@ -1908,15 +1773,10 @@ def _torch_step_label(suffix: str) -> str: # -- Verbosity control ---------------------------------------------------------- -# By default the installer shows a minimal in-place one-line progress bar. -# Set UNSLOTH_VERBOSE=1 to restore full per-step output: -# CLI: unsloth studio setup --verbose -# Linux/Mac: UNSLOTH_VERBOSE=1 ./studio/setup.sh -# Windows: $env:UNSLOTH_VERBOSE="1" ; .\studio\setup.ps1 +# Default: minimal in-place progress bar; UNSLOTH_VERBOSE=1 restores full per-step output. VERBOSE: bool = os.environ.get("UNSLOTH_VERBOSE", "0") == "1" -# Progress bar state -- updated by _progress() per install step. -# Update _TOTAL if you add/remove steps in install_python_stack(). +# Progress bar state -- update _TOTAL if you add/remove steps in install_python_stack(). _STEP: int = 0 _TOTAL: int = 0 # set at runtime in install_python_stack() based on platform _PROGRESS_LINE_ACTIVE: bool = False @@ -1931,20 +1791,16 @@ LOCAL_DD_UNSTRUCTURED_PLUGIN = ( ) LOCAL_DD_GITHUB_PLUGIN = SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed" -# mlx-lm 0.31.3 broke gemma4 / qwen3_5 loading (strict load_weights rejects the -# QK-norm q_norm/k_norm tensors); exclude just that release. See mlx-lm #1242. +# mlx-lm 0.31.3 broke gemma4 / qwen3_5 QK-norm loading; exclude just that release (mlx-lm #1242). MLX_LM_BAD_VERSION_EXCLUSION = "!=0.31.3" -# Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides). -# _uv_safe_path: uv truncates UV_OVERRIDE at the first space too (issue #6503). +# Apple Silicon: override mlx-vlm/mlx-lm's transformers pin; _uv_safe_path because uv truncates UV_OVERRIDE at the first space (#6503). _MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt" if IS_MAC_ARM and _MLX_OVERRIDES.is_file() and "UV_OVERRIDE" not in os.environ: os.environ["UV_OVERRIDE"] = _uv_safe_path(_MLX_OVERRIDES) # -- Unicode-safe printing --------------------------------------------- -# On Windows the console encoding may be a legacy code page (e.g. CP1252) -# that cannot represent glyphs like ✅ or ❌. _safe_print() degrades to ASCII -# equivalents so the installer never crashes over a status glyph. +# Windows console may be a legacy code page (e.g. CP1252); _safe_print() degrades glyphs to ASCII. _UNICODE_TO_ASCII: dict[str, str] = { "\u2705": "[OK]", # ✅ @@ -2006,8 +1862,7 @@ def _stdout_supports_color() -> bool: _HAS_COLOR = _stdout_supports_color() -# Column layout — matches setup.sh step() helper: -# 2-space indent, 15-char label (dim), then value. +# Column layout — matches setup.sh step(): 2-space indent, 15-char dim label, then value. _LABEL = "deps" _COL = 15 _INDENT = 2 @@ -2109,8 +1964,7 @@ def run( if result.returncode != 0: _step("error", f"{label} failed (exit code {result.returncode})", _red) if result.stdout: - # Redact before printing: the failing pip command may carry a pinned --index-url - # with userinfo/?token= creds, so raw pip error text would leak them. + # Redact before printing: pip error text may embed a pinned --index-url's userinfo/?token= creds. print(_redact_install_output(result.stdout)) sys.exit(result.returncode) return result @@ -2119,13 +1973,8 @@ def run( # Packages to skip on Windows (require special build steps) WINDOWS_SKIP_PACKAGES = {"triton_kernels"} -# Packages to skip when torch is unavailable (Intel Mac GGUF-only mode). These -# either *are* torch extensions or have unconditional ``Requires-Dist: torch``, so -# installing them pulls torch back in. ``librosa`` is here despite not requiring -# torch: upstream ``llvmlite`` dropped its macOS x86_64 wheel (0.46.0+ ships only -# macosx_arm64 / manylinux / win_amd64), so on Intel Mac the librosa -> numba -> -# llvmlite chain triggers a from-source build that fails without LLVM 14/15 headers. -# Tracked in unslothai/unsloth#5046. +# Packages to skip when torch is unavailable (Intel Mac GGUF-only mode): torch extensions or hard +# ``Requires-Dist: torch``; ``librosa`` because its numba -> llvmlite chain fails from-source on Intel Mac (#5046). NO_TORCH_SKIP_PACKAGES = { "torch-stoi", "timm", @@ -2214,8 +2063,7 @@ def _bootstrap_uv() -> bool: global UV_NEEDS_SYSTEM if not shutil.which("uv"): return False - # Probe: try a dry-run install targeting the current Python explicitly. - # Without --python, uv can ignore the activated venv on some platforms. + # Dry-run probe with explicit --python: uv can ignore the activated venv on some platforms. probe = subprocess.run( ["uv", "pip", "install", "--dry-run", "--python", sys.executable, "pip"], stdout = subprocess.PIPE, @@ -2289,25 +2137,18 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]: cmd = ["uv", "pip", "install"] if UV_NEEDS_SYSTEM: cmd.append("--system") - # Always pass --python so uv targets the right environment. Without it, uv - # can ignore an activated venv and install into the system Python (seen on - # Colab and similar). + # Always pass --python so uv targets the right env (uv can ignore an activated venv, e.g. Colab). cmd.extend(["--python", sys.executable]) cmd.extend(_translate_pip_args_for_uv(args)) - # Torch is pre-installed, so don't add --torch-backend by default (solver dead-ends on - # CPU-only machines); callers can set UV_TORCH_BACKEND. Never add it to a pinned-index - # command: uv's torch backend redirects torch to its own per-backend index, defeating the pin. + # No --torch-backend by default (torch pre-installed); never on a pinned-index command (UV_TORCH_BACKEND redirects torch, defeating the pin). _tb = os.environ.get("UV_TORCH_BACKEND", "") if _tb and not _is_pinned_index_cmd(cmd): cmd.append(f"--torch-backend={_tb}") return cmd -# uv resolves --index-url / --default-index at LOWEST priority, so an inherited UV_INDEX / -# UV_EXTRA_INDEX_URL mirror wins and a pinned torch repair silently ignores the pin. -# Neutralise these for pinned installs (as install.sh #6898 / install.ps1 / setup.ps1 do). -# UV_TORCH_BACKEND redirects torch; PIP_* matter for the pip FALLBACK; UV_CONFIG_FILE is -# stripped + UV_NO_CONFIG=1 (a discovered uv.toml outranks the CLI pin, uv 0.10). +# uv resolves --index-url at LOWEST priority, so inherited index env vars silently defeat a pinned +# torch repair; neutralise these for pinned installs. UV_CONFIG_FILE is stripped + UV_NO_CONFIG=1. _UV_INDEX_ENV_VARS = ( "UV_CONFIG_FILE", "UV_DEFAULT_INDEX", @@ -2318,8 +2159,7 @@ _UV_INDEX_ENV_VARS = ( "UV_FIND_LINKS", "PIP_EXTRA_INDEX_URL", "PIP_FIND_LINKS", - # PIP_NO_INDEX=1 makes the pip fallback ignore ALL indexes (defeating --index-url); - # PIP_INDEX_URL is dropped too so a stale mirror env can't outrank the pin. + # PIP_NO_INDEX would defeat --index-url; PIP_INDEX_URL dropped so a stale mirror can't outrank the pin. "PIP_NO_INDEX", "PIP_INDEX_URL", ) @@ -2483,9 +2323,7 @@ def install_python_stack() -> int: global USE_UV, _STEP, _TOTAL _STEP = 0 - # install.sh sets SKIP_STUDIO_BASE=1 to avoid reinstalling base packages; - # `studio update` does NOT, so unsloth + unsloth-zoo are reinstalled to pick - # up new versions. + # install.sh sets SKIP_STUDIO_BASE=1; `studio update` does NOT, so base packages are reinstalled. skip_base = os.environ.get("SKIP_STUDIO_BASE", "0") == "1" # --package installs a different package name (for testing). package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth") @@ -2500,8 +2338,7 @@ def install_python_stack() -> int: base_total += 2 # flash-attn + torch final repair (step 13), Linux _TOTAL = (base_total - 1) if skip_base else base_total - # 1. Try uv for faster installs (before pip upgrade -- uv venvs don't - # include pip by default). + # 1. Try uv for faster installs (before pip upgrade -- uv venvs omit pip). USE_UV = _bootstrap_uv() # 2. Ensure pip is available (uv venvs from install.sh omit pip). @@ -2519,8 +2356,7 @@ def install_python_stack() -> int: ], ) else: - # pip may not exist yet (uv-created venvs omit it). Try ensurepip, - # then upgrade. Direct upgrade only when pip is already present. + # pip may not exist yet (uv-created venvs omit it): ensurepip, else direct upgrade. _has_pip = ( subprocess.run( [sys.executable, "-m", "pip", "--version"], @@ -2542,10 +2378,7 @@ def install_python_stack() -> int: [sys.executable, "-m", "pip", "install", "--upgrade", "pip"], ) - # macOS arm64: install MLX stack at latest (UV_OVERRIDE relaxes the - # mlx-vlm / mlx-lm transformers pin -- set at module load). - # Exclude mlx-lm 0.31.3 (see MLX_LM_BAD_VERSION_EXCLUSION); it broke - # gemma4 / qwen3_5 QK-norm loading. mlx-lm #1242. + # macOS arm64: MLX stack at latest (UV_OVERRIDE relaxes the transformers pin); exclude mlx-lm 0.31.3 (mlx-lm #1242). if IS_MAC_ARM and not skip_base: _progress("MLX stack (Apple Silicon)") pip_install( @@ -2562,8 +2395,7 @@ def install_python_stack() -> int: if skip_base: pass elif NO_TORCH: - # No-torch update path: install unsloth + unsloth-zoo, then runtime deps, - # both with --no-deps (PyPI metadata declares torch a hard dep; avoid it). + # No-torch update path: --no-deps throughout (PyPI metadata declares torch a hard dep). _progress("base packages (no torch)") pip_install( f"Updating {package_name} + unsloth-zoo (no-torch mode)", @@ -2576,9 +2408,7 @@ def install_python_stack() -> int: package_name, "unsloth-zoo", ) - # Resolve pydantic WITH deps so pip pins pydantic-core to the exact version - # its metadata declares (under --no-deps pip picks the latest of each and - # trips pydantic's _ensure_pydantic_core_version check). Deps are torch-free. + # pydantic WITH deps so pip pins a matching pydantic-core (--no-deps trips _ensure_pydantic_core_version). Deps are torch-free. pip_install( "Installing pydantic (with deps for compatible core)", "--no-cache-dir", @@ -2610,8 +2440,7 @@ def install_python_stack() -> int: constrain = False, ) elif local_repo: - # Local dev install: update deps from base.txt, then overlay the local - # checkout as an editable install (--no-deps so torch is not re-resolved). + # Local dev install: update deps, then overlay the local checkout editable (--no-deps). _progress("base packages") pip_install( "Updating base packages", @@ -2649,9 +2478,7 @@ def install_python_stack() -> int: package_name, ) else: - # Update path: upgrade only unsloth + unsloth-zoo, preserving existing - # torch/CUDA installs. Torch is pre-installed by install.sh/setup.ps1; - # --upgrade-package targets only base pkgs. + # Update path: upgrade only unsloth + unsloth-zoo, preserving the pre-installed torch. _progress("base packages") pip_install( "Updating base packages", @@ -2663,17 +2490,14 @@ def install_python_stack() -> int: req = REQ_ROOT / "base.txt", ) - # 2b. AMD ROCm: reinstall torch with HIP wheels if the host has ROCm but the - # venv got CPU-only torch (common when pip resolves torch from PyPI). - # Must follow base packages so torch is present for inspection. + # 2b. Torch repair (wrong-family / CPU-only torch); must follow base packages so torch is present. if not IS_MACOS and not NO_TORCH: _progress(_torch_step_label("check")) _ensure_cuda_torch() _ensure_rocm_torch() _ensure_cpu_torch() - # Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python - # version or unknown ROCm version). + # Windows + AMD GPU: warn if ROCm torch was not installed. if IS_WINDOWS and not NO_TORCH and not _has_usable_nvidia_gpu(): # Validate actual AMD GPU presence (not just tool existence). import re as _re_win @@ -2689,10 +2513,7 @@ def install_python_stack() -> int: _wexe = shutil.which(_wcmd[0]) if not _wexe: continue - # Skip amd-smi on Windows w/o a HIP SDK (avoids the UAC/DiskPart - # prompt), as _has_rocm_gpu()/_detect_amd_gfx_codes do. The only loss - # is the best-effort "AMD GPU detected" note; ROCm-torch state below - # comes from the install itself. + # Skip amd-smi w/o a HIP SDK (UAC/DiskPart prompt); only loss is the best-effort note. if _wcmd[0] == "amd-smi" and not _amd_smi_allowed(): continue try: @@ -2736,16 +2557,11 @@ def install_python_stack() -> int: req = REQ_ROOT / "extras-no-deps.txt", ) - # 4. Overrides (torchao) -- force-reinstall to a version matching the venv's - # torch so its C++ extensions load (see _select_torchao_spec). Skipped when - # torch is unavailable (Intel Mac GGUF-only) and on Windows ROCm (no working - # build; see below). + # 4. Overrides (torchao) -- force-reinstall to match the venv's torch (see _select_torchao_spec); skipped for no-torch and Windows ROCm. if NO_TORCH: _progress("dependency overrides (skipped, no torch)") elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm(): - # No working Windows ROCm torchao build: it imports an absent c10d backend - # and crashes transformers.quantizers. Unsloth stubs it at runtime, so - # installing it only ships a package that crashes on import -- skip it. + # No working Windows ROCm torchao build (crashes on import; stubbed at runtime) -- skip it. _progress("dependency overrides (skipped, Windows ROCm)") _safe_print(" Windows ROCm -- skipping torchao (no working build; stubbed at runtime)") else: @@ -2760,8 +2576,7 @@ def install_python_stack() -> int: _torchao_spec, ) - # 5. Triton kernels (no-deps, from source). Skip on Windows and macOS - # (no support). + # 5. Triton kernels (no-deps, from source); skip on Windows and macOS. if not IS_WINDOWS and not IS_MACOS: _progress("triton kernels") pip_install( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index f523b9ff14..7892f67bf8 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -23,31 +23,16 @@ $ErrorActionPreference = "Stop" $ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path $PackageDir = Split-Path -Parent $ScriptDir -# -------------------------------------------------------------------------- -# Maintainer-editable defaults -# Change these in the GitHub-hosted script so users get updated defaults. -# User env vars always override these baked-in values. -# -------------------------------------------------------------------------- -# Prefer "latest" over "master" -- "master" bypasses the prebuilt resolver -# (no matching GitHub release), forces a source build, and causes HTTP 422 -# errors. Only use "master" temporarily when the latest release is missing -# support for a new model architecture. -# -# UNSLOTH_LLAMA_CPP_BACKEND : "auto" (default) or "cpu". When "cpu", forces -# the CPU-only prebuilt bundle on GPU hosts. Fixes Intel iGPU Vulkan -# crashes (#7213). +# Maintainer-editable defaults (user env vars override). "latest" not "master": +# "master" bypasses the prebuilt resolver and forces a source build. +# UNSLOTH_LLAMA_CPP_BACKEND=cpu forces the CPU-only prebuilt (Intel iGPU Vulkan crash #7213). $DefaultLlamaPrForce = "" $DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp" $DefaultLlamaTag = "latest" $DefaultLlamaForceCompileRef = "master" -# Corporate-mirror / proxy escape hatch for the frontend npm/bun install (#6491). -# studio/frontend/.npmrc pins registry=https://registry.npmjs.org/ as a supply-chain -# lock, which overrides a corporate user's ~/.npmrc proxy and causes 403s behind a -# firewall. UNSLOTH_NPM_REGISTRY is a deliberate opt-in: when set we splat it as -# `--registry ` into every npm/bun install. `--registry` is the highest-precedence -# override for BOTH tools and leaves min-release-age / save-exact in force. Empty array -# (the default) splats to nothing, so normal installs are unchanged. +# UNSLOTH_NPM_REGISTRY: opt-in --registry splat for corporate proxies behind the +# frontend .npmrc registry lock (min-release-age/save-exact stay in force). Empty = no-op. $NpmRegistryArgs = @() if ($env:UNSLOTH_NPM_REGISTRY) { $NpmRegistryArgs = @('--registry', $env:UNSLOTH_NPM_REGISTRY) @@ -61,14 +46,12 @@ foreach ($a in $args) { break } } -# Propagate to child processes (e.g. install_python_stack.py) so they -# also respect verbose mode. Process-scoped -- does not persist. +# Propagate verbose to child processes (process-scoped). if ($script:UnslothVerbose) { $env:UNSLOTH_VERBOSE = '1' } $script:LlamaCppDegraded = $false -# CUDA toolkit state, published by Resolve-CudaToolkit. Only the Phase 4 source -# build consumes these; the prebuilt path leaves them at these defaults. +# CUDA toolkit state, published by Resolve-CudaToolkit (Phase 4 source build only). $script:CudaToolkitReady = $false $script:NvccPath = $null $script:CudaToolkitRoot = $null @@ -83,9 +66,7 @@ $IsPipInstall = -not (Test-Path $FrontendDir) # Helper functions # ───────────────────────────────────────────── -# Reload ALL environment variables from registry. -# Picks up changes made by installers (winget, msi, etc.) including -# Path, CUDA_PATH, CUDA_PATH_V*, and any other vars they set. +# Reload ALL environment variables from registry (picks up installer changes: Path, CUDA_PATH, etc). function Refresh-Environment { foreach ($level in @('Machine', 'User')) { $vars = [System.Environment]::GetEnvironmentVariables($level) @@ -116,9 +97,8 @@ function Refresh-Environment { $env:Path = $unique -join ";" } -# ── Helper: safely add a directory to the persistent User PATH ── -# Direct registry access preserves REG_EXPAND_SZ (avoids dotnet/runtime#1442). -# Append (default) keeps existing tools first; Prepend for must-win entries. +# Add a directory to the persistent User PATH. Direct registry access preserves +# REG_EXPAND_SZ. Append (default) keeps existing tools first; Prepend for must-win entries. function Add-ToUserPath { param( [Parameter(Mandatory = $true)][string]$Directory, @@ -184,8 +164,7 @@ function Add-ToUserPath { return $false } $regKey.SetValue('Path', $newPath, [Microsoft.Win32.RegistryValueKind]::ExpandString) - # Broadcast WM_SETTINGCHANGE via dummy env-var roundtrip. - # [NullString]::Value avoids PS 7.5+/.NET 9 $null-to-"" coercion. + # Broadcast WM_SETTINGCHANGE via dummy env-var roundtrip ([NullString]::Value avoids PS 7.5+ coercion). try { $d = "UnslothPathRefresh_$([guid]::NewGuid().ToString('N').Substring(0,8))" [Environment]::SetEnvironmentVariable($d, '1', 'User') @@ -257,8 +236,7 @@ function Get-InstalledLlamaPrebuiltRelease { return $message } -# Find nvcc on PATH, CUDA_PATH, or standard toolkit dirs. -# Returns the path to nvcc.exe, or $null if not found. +# Find nvcc on PATH, CUDA_PATH, or standard toolkit dirs; $null if not found. function Find-Nvcc { param([string]$MaxVersion = "") @@ -329,11 +307,9 @@ function Write-CudaDriverToolkitMismatch { substep "Or let Unsloth use the prebuilt CUDA bundle; it does not need the local toolkit." $Color } -# Detect CUDA Compute Capability via nvidia-smi. -# Returns e.g. "80" for A100 (8.0), "89" for RTX 4090 (8.9), etc. -# Returns $null if detection fails. +# Detect CUDA Compute Capability via nvidia-smi (e.g. "89" for RTX 4090); $null on failure. function Get-CudaComputeCapability { - # Use the resolved absolute path ($NvidiaSmiExe) to survive Refresh-Environment + # $NvidiaSmiExe is an absolute path that survives Refresh-Environment. $smiExe = if ($script:NvidiaSmiExe) { $script:NvidiaSmiExe } else { $cmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue if ($cmd) { $cmd.Source } else { $null } @@ -341,13 +317,11 @@ function Get-CudaComputeCapability { if (-not $smiExe) { return $null } try { - # Bounded: a wedged nvidia-smi must not hang setup after the initial - # -L probe succeeded (the helper merges stderr after stdout, so the - # first line is still the compute_cap value). + # Bounded so a wedged nvidia-smi can't hang setup. $raw = Invoke-NvidiaSmiBounded $smiExe @('--query-gpu=compute_cap', '--format=csv,noheader') if ($LASTEXITCODE -ne 0 -or -not $raw) { return $null } - # nvidia-smi may return multiple GPUs; take the first one + # Multiple GPUs: take the first. $cap = ($raw -split "`n")[0].Trim() if ($cap -match '^(\d+)\.(\d+)$') { $major = $Matches[1] @@ -359,11 +333,8 @@ function Get-CudaComputeCapability { return $null } -# Check if an nvcc binary supports a given sm_ architecture. -# Uses `nvcc --list-gpu-code` which outputs sm_* tokens (--list-gpu-arch -# outputs compute_* tokens instead). Available since CUDA 11.6. -# Returns $false if the flag isn't supported (old toolkit) — safer to reject -# and fall back to scanning/PTX than to assume support and fail later. +# Does an nvcc binary support a given sm_ arch? Uses `nvcc --list-gpu-code` (CUDA 11.6+); +# $false on old toolkits without the flag (safer to reject and fall back to scanning/PTX). function Test-NvccArchSupport { param([string]$NvccExe, [string]$Arch) try { @@ -375,8 +346,7 @@ function Test-NvccArchSupport { } } -# Given an nvcc binary, return the highest sm_ architecture it supports. -# Returns e.g. "90" for CUDA 12.4. Returns $null if detection fails. +# Highest sm_ arch an nvcc binary supports (e.g. "90" for CUDA 12.4); $null on failure. function Get-NvccMaxArch { param([string]$NvccExe) try { @@ -395,11 +365,8 @@ function Get-NvccMaxArch { return $null } -# Detect driver's max CUDA version from nvidia-smi and return the highest -# compatible PyTorch CUDA index tag (e.g. "cu128"). -# PyTorch on Windows ships CPU-only by default from PyPI; CUDA wheels live at -# https://download.pytorch.org/whl/. The tag must not exceed the driver's -# capability: e.g. driver "CUDA Version: 12.9" → cu128 (not cu130). +# Highest PyTorch CUDA index tag (e.g. "cu128") the driver's max CUDA supports; +# the tag must not exceed the driver (CUDA 12.9 -> cu128, not cu130). function Get-PytorchCudaTag { $smiExe = if ($script:NvidiaSmiExe) { $script:NvidiaSmiExe } else { $cmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue @@ -408,13 +375,9 @@ function Get-PytorchCudaTag { if (-not $smiExe) { return "cu126" } try { - # Bounded: a wedged nvidia-smi must not hang setup. The helper merges - # stderr into the returned string, matching the old 2>&1 | Out-String - # shape (plain 2>$null leaks ErrorRecord objects in PS 5.1). + # Bounded so a wedged nvidia-smi can't hang setup. $output = Invoke-NvidiaSmiBounded $smiExe - # Newer NVIDIA drivers (e.g. 610.x on Windows) print - # "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y". - # Accept both spellings so we don't fall through to the cu126 default. + # Newer drivers print "CUDA UMD Version: X.Y"; accept both spellings. if ($output -match 'CUDA(?: UMD)? Version:\s+(\d+)\.(\d+)') { $major = [int]$Matches[1] $minor = [int]$Matches[2] @@ -431,8 +394,7 @@ function Get-PytorchCudaTag { return "cu126" } -# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL -# TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared. +# Trim trailing slashes from the URL path only, preserving ?query / #fragment. Shared. function Trim-IndexPathSlashes { param([string]$Url) $value = $Url.Trim() @@ -443,9 +405,8 @@ function Trim-IndexPathSlashes { return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx) } -# Explicit torch-index pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY), shared by the stale-venv check -# and install selection so a pinned index wins over GPU probing (parity with the other -# installers). URL is verbatim; _FAMILY is the leaf joined to the mirror base. +# Explicit torch-index pin (UNSLOTH_TORCH_INDEX_URL verbatim, or _FAMILY joined to the +# mirror base), shared so a pin wins over GPU probing. Parity with the other installers. function Get-PinnedTorchIndexUrl { if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) { return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL) @@ -457,9 +418,8 @@ function Get-PinnedTorchIndexUrl { return $null } -# Last path segment of a wheel index URL, query/fragment dropped first so a token-authenticated -# pin (.../cu128?token=x) classifies as cu128 (else it reinstalls every update). Classification -# only. Shared with the py / install.sh leaf extractors. +# Last path segment of a wheel index URL (query/fragment dropped) so .../cu128?token=x +# classifies as cu128. Classification only; shared with the py / install.sh extractors. function Get-TorchIndexLeaf { param([string]$Url) if ([string]::IsNullOrWhiteSpace($Url)) { return $null } @@ -469,63 +429,57 @@ function Get-TorchIndexLeaf { } # Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer -# output before printing on failure; uv/pip errors echo the failing --index-url verbatim. -# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted. +# output before printing on failure; uv/pip echo the failing --index-url verbatim. function Redact-InstallOutput { param([string]$Text) if (-not $Text) { return $Text } $Text = $Text -replace '(https?://)[^/@\s`]+@', '$1@' $Text = $Text -replace '([?&][^=\s&`]+)=[^&#\s`]+', '$1=' - # A #token=... fragment is as sensitive as a query; URL-anchored. + # #token=... fragment is as sensitive as a query. return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#' } -# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug). MUST match -# the install-spec path below and the other installers; other leaves ship <2.11 and stay default. +# AMD per-arch leaves needing the torch 2.11 floor (_grouped_mm <2.11 bug). Must match +# the install-spec path below and the other installers. function Test-RocmGfx211Leaf { param([string]$Leaf) return @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $Leaf } -# rocmX.Y versions KNOWN to ship torch 2.11: rocm7.2 only today. Do NOT floor an unknown newer -# rocm speculatively. MUST match _ROCM_KNOWN_TORCH211_VERSIONS and the rocm7.2 leaf elsewhere. +# rocmX.Y versions KNOWN to ship torch 2.11 (rocm7.2 only; no speculative floor). +# Must match _ROCM_KNOWN_TORCH211_VERSIONS. function Test-RocmKnown211Version { param([int]$Major, [int]$Minor) return ($Major -eq 7 -and $Minor -eq 2) } -# True only for a real CUDA family leaf: "cu" + digits (cu118, cu128, ...). A bare -like 'cu*' -# would match "custom"/"current" and rebuild the venv every run. Mirrors _is_cuda_family_leaf. +# True only for a real CUDA family leaf: EXACT cu+digits (cu118, cu128). A bare 'cu*' glob +# would match "custom"/"current" and rebuild every run. Mirrors _is_cuda_family_leaf. function Test-CudaFamilyLeaf { param([string]$Leaf) if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false } - # EXACT cu+digits: cu128-private routes through the unknown-leaf path instead. return $Leaf -match '^cu[0-9]+$' } -# True only for a real pip ROCm family leaf: EXACT rocm[.] or a gfx leaf. A leaf -# that merely STARTS with rocm (rocm-rel-7.2.1, rocm7.2-private) is a custom pin the verbatim -# path owns, so anchor the match. Mirrors _is_pip_rocm_family_leaf / install.sh. +# True only for a real pip ROCm family leaf: EXACT rocm[.] or gfx. +# A merely-STARTS-with-rocm leaf (rocm-rel-7.2.1) is a custom pin the verbatim path owns, +# so anchor the match. Mirrors _is_pip_rocm_family_leaf / install.sh. function Test-PipRocmFamilyLeaf { param([string]$Leaf) if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false } - # gfx must be followed by a digit (an architecture leaf); gfx-private is custom. return ($Leaf -match '^gfx[0-9]') -or ($Leaf -match '^rocm[0-9]+(\.[0-9]+)?$') } -# Stale-venv ROCm comparison for a pinned gfx*/rocm* index. Returns @{ Expected; Installed } so -# the caller rebuilds when they differ. Mirrors _rocm_pin_family_mismatch (same rocmX.Y / gfx -# cases). An untagged (no +rocm) wheel never satisfies a ROCm pin -> stale. +# Stale-venv ROCm comparison for a pinned gfx*/rocm* index. Returns @{ Expected; Installed } +# so the caller rebuilds when they differ. Mirrors _rocm_pin_family_mismatch. function Get-RocmPinStaleTags { param([string]$PinLeaf, [string]$TorchVersion) $_pinRocm = [regex]::Match($PinLeaf, '^rocm(\d+)\.(\d+)') $_pinVer = if ($_pinRocm.Success) { "$($_pinRocm.Groups[1].Value).$($_pinRocm.Groups[2].Value)" } else { $null } - # The family classifier accepts a major-only rocm leaf too (rocm7). - $_pinMajorOnly = [regex]::Match($PinLeaf, '^rocm(\d+)$') - # Installed rocm version and whether the wheel is a per-arch (three-part) build. + $_pinMajorOnly = [regex]::Match($PinLeaf, '^rocm(\d+)$') # major-only rocm leaf (rocm7) $_instRocm = [regex]::Match($TorchVersion, '\+rocm(\d+)\.(\d+)') $_instVer = if ($_instRocm.Success) { "$($_instRocm.Groups[1].Value).$($_instRocm.Groups[2].Value)" } else { $null } - $_instPerArch = [regex]::IsMatch($TorchVersion, '\+rocm\d+\.\d+\.\d+') + $_instPerArch = [regex]::IsMatch($TorchVersion, '\+rocm\d+\.\d+\.\d+') # per-arch three-part build # A ROCm build MUST carry a +rocm tag; an untagged wheel can't satisfy any ROCm pin. $_instHasRocm = [regex]::IsMatch($TorchVersion, '\+rocm') $_instRel = [regex]::Match($TorchVersion, '^(\d+)\.(\d+)') @@ -536,8 +490,7 @@ function Get-RocmPinStaleTags { if ($PinLeaf -like 'gfx*') { if (Test-RocmGfx211Leaf $PinLeaf) { - # Expect the AMD per-arch (three-part) 2.11 wheel: satisfied only when BOTH - # a 2.11 release AND a three-part rocm tag are installed. + # AMD per-arch 2.11 wheel: satisfied only with BOTH a 2.11 release AND a three-part rocm tag. $installed = if ($_instIs211 -and $_instPerArch) { "rocm-perarch(torch>=2.11)" } else { "rocm-generic-or-old" } return @{ Expected = "rocm-perarch(torch>=2.11)"; Installed = $installed } } @@ -549,9 +502,7 @@ function Get-RocmPinStaleTags { } } - # Major-only rocm pin (rocm7): compare majors only -- a +rocm6.4 wheel under a rocm7 - # pin is stale, any +rocm7.x wheel satisfies it (no pinned minor to compare, and the - # 2.11-line fallback below would invert both verdicts). Mirrors _rocm_pin_family_mismatch. + # Major-only rocm pin (rocm7): compare majors only. Any +rocm7.x wheel satisfies it. if ($_pinMajorOnly.Success) { $_pinMaj = [int]$_pinMajorOnly.Groups[1].Value if ($_instVer) { @@ -559,17 +510,15 @@ function Get-RocmPinStaleTags { $expected = if ($_instMaj -eq $_pinMaj) { "rocm$_instVer" } else { "rocm$_pinMaj.x" } return @{ Expected = $expected; Installed = "rocm$_instVer" } } - # Untagged wheel never satisfies a ROCm pin; a +rocm tag with an unreadable - # version is accepted (matches the lenient unreadable fallback below). + # Untagged wheel never satisfies a ROCm pin; a +rocm tag with an unreadable version is accepted. $installed = if ($_instHasRocm) { "rocm" } else { "not-rocm" } return @{ Expected = "rocm"; Installed = $installed } } # rocmX.Y pin. if ($_pinVer -and $_instVer) { - # Both readable: exact compare. When they match AND the pin is KNOWN-2.11, the - # installed release must also be 2.11 (a +rocm7.2 wheel drifted to 2.12 shares the - # tag but violates the spec), so fold the release into the tag. Mirrors _rocm_pin_family_mismatch. + # Both readable: exact compare. A KNOWN-2.11 pin also requires a 2.11 release (a + # +rocm7.2 wheel drifted to 2.12 shares the tag but violates the spec), so fold it in. $_pinKnown211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value) $_instOn211 = $_instRel.Success -and [int]$_instRel.Groups[1].Value -eq 2 -and [int]$_instRel.Groups[2].Value -eq 11 if ($_pinKnown211 -and -not $_instOn211) { @@ -579,12 +528,10 @@ function Get-RocmPinStaleTags { } $_pinNeeds211 = $false if ($_pinRocm.Success) { - # Only KNOWN-2.11 rocm (rocm7.2) is on the 2.11 line (no speculative floor). - # Matches _ROCM_KNOWN_TORCH211_VERSIONS. + # Only KNOWN-2.11 rocm (rocm7.2) is on the 2.11 line. Matches _ROCM_KNOWN_TORCH211_VERSIONS. $_pinNeeds211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value) } - # Fallback (installed rocm version unreadable): compare on the 2.11 line; an untagged - # wheel never satisfies a rocmX.Y pin -> stale. + # Fallback (installed rocm version unreadable): compare on the 2.11 line. $installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } return @{ Expected = if ($_pinNeeds211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } @@ -592,8 +539,7 @@ function Get-RocmPinStaleTags { } } -# VS generator -> MSBuild BuildCustomizations dir; toolset tracks the VS major -# (18->v180, 17->v170), defaulting to v170 when unparseable. +# VS generator -> MSBuild BuildCustomizations dir; toolset tracks the VS major (18->v180), default v170. function Get-VcBuildCustomizationsDir { param( [Parameter(Mandatory)][string]$VsInstallPath, @@ -631,8 +577,8 @@ function Test-CmakeSupportsGenerator { } function Test-CmakeListsGenerator { - # Does `cmake --help` actually list the generator? A VS-bundled cmake can drive - # VS 2026 below the 4.2 floor, so probe rather than trust the version. (#6473) + # Does `cmake --help` list the generator? A VS-bundled cmake can drive VS 2026 below + # the 4.2 floor, so probe rather than trust the version. (#6473) param([Parameter(Mandatory)][string]$Generator) $help = & cmake --help 2>$null | Out-String if (-not $help) { return $false } @@ -651,8 +597,7 @@ function Test-CmakeCanDriveGenerator { } function Add-DefaultCmakeToPath { - # Prepend the default CMake dir so a freshly winget-installed cmake wins over an - # older one already on PATH. $true if found. (#6473) + # Prepend the default CMake dir so a fresh winget cmake wins over an older one on PATH. (#6473) $cmakeDefaults = @( "$env:ProgramFiles\CMake\bin", "${env:ProgramFiles(x86)}\CMake\bin", @@ -669,10 +614,9 @@ function Add-DefaultCmakeToPath { } function Get-FallbackVsGenerator { - # Newest pre-2026 VS whose generator the current cmake can drive, for when the - # VS 2026 generator is unusable (old/offline cmake) but an older toolchain exists. - # vswhere first (catches non-default roots like D:\), then Program Files; matches - # Find-VsBuildTools. Returns @{ Generator; InstallPath } or $null. (#6473) + # Newest pre-2026 VS whose generator the current cmake can drive, when the VS 2026 + # generator is unusable. vswhere first (non-default roots like D:\), then Program Files. + # Returns @{ Generator; InstallPath } or $null. (#6473) $knownEditions = @('BuildTools', 'Community', 'Professional', 'Enterprise', 'Preview') # install path if it holds a usable cl.exe, else $null @@ -729,9 +673,8 @@ function Get-FallbackVsGenerator { return $null } -# VS version label -> cmake generator. vswhere's productLineVersion is the year for -# VS <= 2022 but the internal major "18" for VS 2026, and dir names use either form, -# so accept both. (VS 2026 detection adapted from @LeoBorcherding's #6038.) +# VS version label -> cmake generator. productLineVersion is the year for VS <= 2022 +# but internal major "18" for VS 2026, and dir names use either form, so accept both. function Resolve-VsGeneratorFromLabel { param([string]$Label) if (-not $Label) { return $null } @@ -794,10 +737,8 @@ function Find-VsBuildTools { return $null } -# Install CMake + VS Build Tools, deferred here from Phase 1 so the prebuilt path -# never pays for a multi-GB install. Called only when a source build is committed. -# CMake is best-effort (build skips downstream if absent); VS Build Tools are -# required, so exit 1 with guidance if missing. No-ops for VS when already detected. +# Install CMake + VS Build Tools, deferred from Phase 1 so the prebuilt path skips the +# multi-GB install. Called only for a source build. CMake best-effort; VS required (exit 1). function Ensure-BuildToolsForLlamaSourceBuild { # CMake if ($null -eq (Get-Command cmake -ErrorAction SilentlyContinue)) { @@ -855,9 +796,8 @@ function Ensure-BuildToolsForLlamaSourceBuild { } } -# Detect the VC++ 2015-2022 Redistributable that the prebuilt llama-server and -# PyTorch need (they link VCRUNTIME140_1.dll etc., which the Universal CRT lacks). -# Signal is System32\vcruntime140_1.dll (VS 2019+), registry as fallback. +# Detect the VC++ 2015-2022 Redistributable the prebuilt llama-server and PyTorch need +# (VCRUNTIME140_1.dll). Signal is System32\vcruntime140_1.dll (VS 2019+), registry fallback. function Test-VCRedistInstalled { $sys = $env:SystemRoot if ($sys -and (Test-Path (Join-Path $sys 'System32\vcruntime140_1.dll'))) { return $true } @@ -971,12 +911,8 @@ function Invoke-SetupCommand { # Reset to avoid stale values from prior native commands. $global:LASTEXITCODE = 0 if ($script:UnslothVerbose -and -not $AlwaysQuiet) { - # Merge stderr into stdout so progress/warning output stays visible - # without flipping $? on successful native commands (PS 5.1 treats - # stderr records as errors that set $? = $false even on exit code 0). - # Redact per record: uv/pip echo index URLs (credentials and all) in - # their errors, and verbose mode must not bypass the quiet path's - # redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched. + # Merge stderr into stdout (PS 5.1 stderr records flip $? even on exit 0). Redact per + # record: verbose mode must not bypass the quiet path's index-URL redaction. & $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host } else { $output = & $Command 2>&1 | Out-String @@ -1008,16 +944,8 @@ function Write-LlamaFailureLog { Write-Host " | $line" -ForegroundColor DarkGray } } -# Mirror the plain (no ANSI) form of step/substep messages to the -# OS-level stdout handle when a parent is consuming our stdout via -# a pipe (CI `tee`, Python subprocess.PIPE, CREATE_NO_WINDOW grandchild). -# Write-Host on PS 5.1 routes through $Host.UI / the Information -# stream, neither of which propagates reliably across the -# install.ps1 -> unsloth.exe -> python -> powershell.exe -> -# setup.ps1 process chain. [Console]::Out always lands on the OS -# stdout file handle. Gated on IsOutputRedirected so the -# interactive-console path keeps the colorized Write-Host output -# only (no double-print). +# Mirror plain step/substep lines to [Console]::Out when stdout is redirected: Write-Host doesn't +# propagate across the process chain. Gated on IsOutputRedirected (no double-print interactively). function Write-StudioStdoutMirror { param([Parameter(Mandatory = $true)][string]$Line) try { @@ -1083,10 +1011,8 @@ function substep { } function Show-NpmRegistryHint { - # Print actionable guidance when a frontend/OXC npm/bun install fails and the - # registry lock is the likely cause (corporate firewall/proxy). No-op once the - # user has opted in via UNSLOTH_NPM_REGISTRY. We never switch registries - # automatically -- we only guide. + # Guidance when a frontend/OXC npm/bun install fails behind a corporate firewall/proxy. + # Only guides; never switches registries automatically. No-op if UNSLOTH_NPM_REGISTRY set. if ($env:UNSLOTH_NPM_REGISTRY) { return } $mirror = $env:NPM_CONFIG_REGISTRY if (-not $mirror) { @@ -1162,11 +1088,8 @@ try { # ============================================ # 1a. GPU detection # ============================================ -# ── Helper: run nvidia-smi under a timeout ── -# A wedged NVIDIA driver can make nvidia-smi block during init or after a reset; -# WaitForExit bounds it (mirrors Invoke-AmdSmiNoElevate below) so detection -# cannot hang setup. No RunAsInvoker compat layer: nvidia-smi does not -# auto-elevate. Returns combined stdout+stderr; "" on timeout/failure. +# Run nvidia-smi under a timeout: a wedged driver can block during init/reset, so WaitForExit +# bounds it. Returns combined stdout+stderr; "" on timeout/failure. function Invoke-NvidiaSmiBounded { param( [Parameter(Mandatory = $true, Position = 0)][string]$Exe, @@ -1197,10 +1120,8 @@ function Invoke-NvidiaSmiBounded { } } -# ── Helper: nvidia-smi -L lists at least one real GPU ── -# Exit code 0 alone is not enough: a stale/driverless nvidia-smi can exit 0 -# while listing no GPU, which would mark an AMD host NVIDIA and suppress ROCm -# detection. Require a "GPU :" data row. +# nvidia-smi -L lists at least one real GPU: a stale/driverless smi can exit 0 while +# listing no GPU (would mark an AMD host NVIDIA), so require a "GPU :" data row. function Test-NvidiaSmiHasGpu { param([Parameter(Mandatory = $true)][string]$Exe) $out = Invoke-NvidiaSmiBounded $Exe @('-L') @@ -1216,8 +1137,7 @@ try { $NvidiaSmiExe = $nvSmiCmd.Source } } catch {} -# Fallback: nvidia-smi may not be on PATH even though a GPU + driver exist. -# Check the default install location and the Windows driver store. +# Fallback: nvidia-smi may not be on PATH though a GPU + driver exist -- check default dirs. if (-not $HasNvidiaSmi) { $nvSmiDefaults = @( "$env:ProgramFiles\NVIDIA Corporation\NVSMI\nvidia-smi.exe", @@ -1236,26 +1156,21 @@ if (-not $HasNvidiaSmi) { } } } -# ── Helper: run amd-smi without triggering a UAC elevation prompt ── -# amd-smi on Windows auto-elevates to read GPU/APU memory, surfacing a confusing -# DiskPart UAC prompt mid-install (Unsloth backend amd.py hits the same). RunAsInvoker -# forces it (and helpers it spawns) to run un-elevated; on failure the WMI name -> -# gfx fallback still resolves the arch. +# Run amd-smi without triggering a UAC elevation prompt: amd-smi auto-elevates to read +# GPU/APU memory (a confusing DiskPart UAC prompt mid-install). RunAsInvoker forces +# un-elevated; on failure the WMI name -> gfx fallback still resolves the arch. function Invoke-AmdSmiNoElevate { param( [Parameter(Mandatory = $true, Position = 0)][string]$Exe, [Parameter(Position = 1)][string[]]$SmiArgs = @(), [int]$TimeoutSec = 30 ) - # RunAsInvoker blocks the auto-elevation/UAC prompt; the timeout bounds a flaky - # amd-smi that can otherwise spin for minutes (30s mirrors the backend amd.py). + # RunAsInvoker blocks the auto-elevation/UAC prompt; the timeout bounds a flaky amd-smi. $prevCompat = [Environment]::GetEnvironmentVariable('__COMPAT_LAYER', 'Process') $env:__COMPAT_LAYER = 'RunAsInvoker' try { - # [Process]::Start, NOT Start-Process -PassThru: the latter leaves .ExitCode - # $null after WaitForExit on PS 5.1, so $LASTEXITCODE (checked by callers) - # reads non-zero and kills detection. Async reads drain the pipes (no - # deadlock); amd-smi args have no spaces so a plain join is safe. + # [Process]::Start, NOT Start-Process -PassThru (which leaves .ExitCode $null on + # PS 5.1, killing detection). Async reads drain the pipes; amd-smi args have no spaces. $psi = New-Object System.Diagnostics.ProcessStartInfo $psi.FileName = $Exe $psi.Arguments = ($SmiArgs -join ' ') @@ -1291,16 +1206,13 @@ $HipSdkInstalled = $false # HIP SDK binary found (independent of device access $ROCmGpuLabel = $null $script:ROCmGfxArch = $null if (-not $HasNvidiaSmi) { - # hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (mirrors NVIDIA smi path resolution). - # AMD HIP SDK sets HIP_PATH but may not add the bin dir to PATH depending on install type. - # Ignore the venv hipInfo.exe (AMD wheel, on PATH): not a HIP SDK, so amd-smi - # would still auto-elevate. Cf. _path_inside_venv(). + # hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback. Ignore the venv hipInfo.exe + # (AMD wheel, not a HIP SDK, would still auto-elevate). Cf. _path_inside_venv(). function Test-HipinfoIsVenvInternal { param([AllowNull()][string]$HipinfoPath) if ([string]::IsNullOrWhiteSpace($HipinfoPath)) { return $false } - # VenvDir/VIRTUAL_ENV can be unset this early (the update flow probes before - # VenvDir is set), so also derive the venv from the setup python + default - # Unsloth home, else the venv hipInfo isn't caught. + # VenvDir/VIRTUAL_ENV can be unset this early, so also derive the venv from the + # setup python + default Unsloth home, else the venv hipInfo isn't caught. $venvRoots = @() if ($env:VIRTUAL_ENV) { $venvRoots += $env:VIRTUAL_ENV } $vd = Get-Variable -Name VenvDir -ValueOnly -ErrorAction SilentlyContinue @@ -1309,15 +1221,12 @@ if (-not $HasNvidiaSmi) { try { $venvRoots += (Split-Path -Parent (Split-Path -Parent $env:UNSLOTH_SETUP_PYTHON)) } catch {} } if ($env:USERPROFILE) { $venvRoots += (Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio") } - # A custom Unsloth home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the - # venv off the default path; seed it too or its hipInfo escapes the filter. + # A custom Unsloth home moves the venv off the default path; seed it too. $studioHomeEnv = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() } else { $null } if ($studioHomeEnv) { - # Expand a leading ~ like the canonical resolver below; else GetFullPath - # keeps the literal ~ (cwd-relative) and the hipInfo escapes the filter. + # Expand a leading ~ like the canonical resolver; else GetFullPath keeps the literal ~. if (($studioHomeEnv -eq "~" -or $studioHomeEnv -like "~/*" -or $studioHomeEnv -like "~\*") -and -not [string]::IsNullOrWhiteSpace($env:USERPROFILE)) { - # A bare "~" leaves an empty child path; Join-Path rejects that on - # PS 5.1, so use USERPROFILE directly and only join a real remainder. + # A bare "~" leaves an empty child path Join-Path rejects on PS 5.1. $studioHomeRest = $studioHomeEnv.Substring(1).TrimStart('/', '\') $studioHomeEnv = if ($studioHomeRest) { Join-Path $env:USERPROFILE $studioHomeRest } else { $env:USERPROFILE } } @@ -1327,8 +1236,7 @@ if (-not $HasNvidiaSmi) { foreach ($root in $venvRoots) { if ([string]::IsNullOrWhiteSpace($root)) { continue } try { $r = [System.IO.Path]::GetFullPath($root).TrimEnd('\', '/') } catch { continue } - # Skip a bare drive root (e.g. a non-venv UNSLOTH_SETUP_PYTHON like - # C:\Python311\python.exe yields C:) -- it would match every path on that drive. + # Skip a bare drive root -- it would match every path on that drive. if ($r -match '^[a-zA-Z]:$') { continue } if ($hip.Equals($r, [System.StringComparison]::OrdinalIgnoreCase) -or $hip.StartsWith($r + [System.IO.Path]::DirectorySeparatorChar, [System.StringComparison]::OrdinalIgnoreCase)) { @@ -1337,15 +1245,14 @@ if (-not $HasNvidiaSmi) { } return $false } - # Scan all hipinfo and keep the first non-venv one (the venv copy from the - # bnb fix could shadow a real HIP SDK's). -CommandType Application matches - # only real executables, not a user alias/function named hipinfo. + # Scan all hipinfo, keep the first non-venv one. -CommandType Application matches only + # real executables, not a user alias/function named hipinfo. $hipinfoExe = Get-Command hipinfo -CommandType Application -All -ErrorAction SilentlyContinue | Where-Object { -not (Test-HipinfoIsVenvInternal $_.Source) } | Select-Object -First 1 if (-not $hipinfoExe) { - # Iterate the env roots (mirrors the Python list) and take the first non-venv - # bin\hipinfo.exe, so a venv-internal HIP_PATH can't mask a real SDK in ROCM_PATH. + # Iterate the env roots, take the first non-venv bin\hipinfo.exe, so a venv-internal + # HIP_PATH can't mask a real SDK in ROCM_PATH. $hipMissingLabel = $null; $hipMissingRoot = $null; $hipMissingCandidate = $null foreach ($hipEnvLabel in @("HIP_PATH", "HIP_PATH_57", "ROCM_PATH")) { $hipRoot = [Environment]::GetEnvironmentVariable($hipEnvLabel) @@ -1373,8 +1280,7 @@ if (-not $HasNvidiaSmi) { try { $hipOut = & $hipinfoExe.Source 2>&1 | Out-String if ($hipOut -match "(?i)gcnArchName") { - # hipinfo can crash after printing gcnArchName (#6043). - # Once the arch is printed, keep the ROCm wheel path. + # hipinfo can crash after printing gcnArchName (#6043); keep the ROCm path once printed. $HasROCm = $true $_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() }) $_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 } @@ -1388,9 +1294,7 @@ if (-not $HasNvidiaSmi) { substep "[INFO] hipinfo exited with code $LASTEXITCODE but reported gcnArchName -- treating as ROCm-capable (see #6043)" "Cyan" } } elseif ($LASTEXITCODE -ne 0) { - # hipinfo ran but returned a HIP runtime error without any gcnArchName - # output (e.g. "no ROCm-capable device detected"), or crashed before - # printing device info. + # hipinfo returned a HIP runtime error with no gcnArchName (no ROCm device), or crashed early. $firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1) substep "[WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" "Yellow" substep " $firstLine" "Yellow" @@ -1398,18 +1302,9 @@ if (-not $HasNvidiaSmi) { } } catch {} } - # amd-smi fallback: HIP runtime present but hipinfo unavailable (no full HIP SDK). - # 'list' confirms GPU visibility, 'static --asic' extracts the gfx arch hipinfo - # would give. Critical for Strix Halo (gfx1151) and other HIP-runtime-only iGPUs. - # - # BUT on hosts without a working HIP runtime amd-smi elevates a child at runtime, - # popping a UAC/DiskPart prompt RunAsInvoker can't suppress (its manifest is - # asInvoker; even 'amd-smi version' hangs). So only probe when a HIP SDK is present - # (hipinfo found -> un-elevated) or the user opts in; else fall through to WMI name - # inference (enough to pick ROCm wheels + the ROCm llama.cpp prebuilt). - # An explicit opt-out (UNSLOTH_ENABLE_AMD_SMI=0/false/no/off) wins over the HIP-SDK - # heuristic: a HIP SDK binary with a broken runtime can still pop the prompt, so - # $HipSdkInstalled must NOT silently re-enable it. + # amd-smi fallback when hipinfo is unavailable ('list' + 'static --asic' give the gfx arch). + # Probe only with a HIP SDK present or explicit opt-in (amd-smi can pop a UAC prompt + # RunAsInvoker can't suppress); an explicit UNSLOTH_ENABLE_AMD_SMI opt-out wins. $amdSmiOptOut = $env:UNSLOTH_ENABLE_AMD_SMI -match '^(?i)(0|false|no|off)$' $amdSmiAllowed = (-not $amdSmiOptOut) -and ($HipSdkInstalled -or ($env:UNSLOTH_ENABLE_AMD_SMI -match '^(?i)(1|true|yes|on)$')) if (-not $HasROCm -and $amdSmiAllowed) { @@ -1419,22 +1314,13 @@ if (-not $HasNvidiaSmi) { $smiOut = Invoke-AmdSmiNoElevate $amdSmiExe.Source @('list') if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") { $HasROCm = $true - # Attempt 1: newer amd-smi versions embed the gfx arch in list output. - # Collect ALL gfx tokens in output order so that on mixed-arch systems - # we can honour HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES and pick the - # arch for the *runtime-visible* GPU rather than always the first one. - # Do NOT deduplicate: a dual same-arch system (e.g. two gfx1151 APUs) - # must produce a 2-element array so HIP_VISIBLE_DEVICES=1 selects the - # second GPU rather than triggering a false out-of-range warning. - # Note: this mapping assumes amd-smi lists GPUs in the same order as - # HIP enumerates them (both follow PCI bus order in practice); it may - # give the wrong arch when GPU indices are non-contiguous (very rare). + # Attempt 1: newer amd-smi embeds the gfx arch in list output. Collect ALL gfx + # tokens in output order (no dedup) to honour HIP/ROCR_VISIBLE_DEVICES and pick + # the runtime-visible GPU's arch. Assumes amd-smi lists GPUs in HIP enumeration order. $allGfxArches = @([regex]::Matches($smiOut, '(?i)\b(gfx\d+[a-z]?)\b') | ForEach-Object { $_.Groups[1].Value.ToLower() }) if ($allGfxArches.Count -gt 0) { - # Resolve which GPU index is runtime-visible. When a single - # integer index is set, use it; fall back to index 0 otherwise - # (comma-separated lists or unset → first GPU, same as before). + # A single integer index selects that GPU; otherwise index 0. $visGpu = if ($env:HIP_VISIBLE_DEVICES) { $env:HIP_VISIBLE_DEVICES } elseif ($env:ROCR_VISIBLE_DEVICES) { $env:ROCR_VISIBLE_DEVICES } else { $null } @@ -1447,8 +1333,7 @@ if (-not $HasNvidiaSmi) { $script:ROCmGfxArch = $allGfxArches[$gpuIdx] $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" } else { - # Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+, - # including the GFX target needed for wheel index selection. + # Attempt 2: 'static --asic' exposes the GFX target on ROCm 6+. $smiAsicOut = "" try { $smiAsicOut = Invoke-AmdSmiNoElevate $amdSmiExe.Source @('static','--asic') } catch {} if ($smiAsicOut -match "(?i)\b(gfx\d+[a-z]?)\b") { @@ -1464,11 +1349,8 @@ if (-not $HasNvidiaSmi) { } catch {} } } - # WMI fallback: AMD GPU in device list but no HIP SDK → guide the user. - # WMI gives a marketing name (e.g. "AMD Radeon 890M") but never a gfx arch. - # $HasROCm is intentionally NOT set here — we cannot confirm ROCm runtime - # support without hipinfo or amd-smi. The name is saved to $ROCmGpuLabel - # so the name-based inference below can still attempt an arch lookup. + # WMI fallback: marketing name only (never a gfx arch); $HasROCm stays false (runtime + # unconfirmed) but $ROCmGpuLabel feeds the name-based arch inference below. if (-not $HasROCm) { try { $wmiGpu = Get-WmiObject Win32_VideoController -ErrorAction SilentlyContinue | @@ -1477,13 +1359,9 @@ if (-not $HasNvidiaSmi) { if ($wmiGpu) { $ROCmGpuLabel = $wmiGpu.Name } } catch {} } - # ── Arch resolution: env-var override → name inference ────────────────── - # Runs after all probes, even when none confirmed a ROCm runtime ($HasROCm false): - # the Adrenalin driver alone runs the per-gfx ROCm llama.cpp prebuilt (bundles its - # own runtime), and all it needs is the gfx arch, inferable from the WMI GPU name. - # Resolving it here lets setup.ps1 forward --rocm-gfx so a GPU llama.cpp is pulled - # instead of CPU. (PyTorch ROCm wheels still require a HIP SDK -- gated on $HasROCm - # below -- so this only affects llama.cpp / inference.) + # ── Arch resolution: env-var override -> name inference ── + # Runs even with no confirmed ROCm runtime: the per-gfx llama.cpp prebuilt only needs the + # gfx arch (forwarded via --rocm-gfx); PyTorch ROCm wheels still gate on $HasROCm below. if (-not $script:ROCmGfxArch) { # 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running. if ($env:UNSLOTH_ROCM_GFX_ARCH) { @@ -1491,9 +1369,8 @@ if (-not $HasNvidiaSmi) { $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $script:ROCmGfxArch" "Cyan" } - # 2. Best-effort name → arch lookup (amd-smi / WMI). Most-specific first, - # first match wins. Covers only arches the ROCm prebuilts support - # (gfx120X/110X/1151/1150/103X); unknown names fall back cleanly to CPU. + # 2. Best-effort name -> arch lookup (most-specific first, first match wins). Covers + # only arches the ROCm prebuilts support; unknown names fall back cleanly to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (Radeon RX 9070 XT / 9080) @@ -1518,9 +1395,8 @@ if (-not $HasNvidiaSmi) { } } } - # Capture ROCm version early for display and wheel selection. - # Run whenever the HIP SDK binary is present, not just when the device is accessible -- - # hipconfig --version works even when hipinfo reports no ROCm device (driver issue). + # Capture ROCm version (display + wheel selection). Runs whenever the HIP SDK binary is + # present: hipconfig --version works even when hipinfo reports no ROCm device. if ($HasROCm -or $HipSdkInstalled) { $script:ROCmVersion = $null $hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue @@ -1577,8 +1453,7 @@ if ($HasNvidiaSmi) { substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow" substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" } elseif ($script:ROCmGfxArch) { - # Known arch: PyTorch comes from AMD's bundled-runtime ROCm wheels (repo.amd.com), - # which ship their own runtime -- HIP SDK optional (only adds the system toolchain). + # Known arch: PyTorch comes from AMD's bundled-runtime ROCm wheels (HIP SDK optional). Write-Host "" step "gpu" "AMD ROCm ($script:ROCmGfxArch)" "Cyan" substep "Detected: $ROCmGpuLabel" "Cyan" @@ -1616,7 +1491,7 @@ if ($LongPathsEnabled) { Write-Host "Windows Long Paths not enabled (required for Triton compilation and deep dependency paths)." -ForegroundColor Yellow Write-Host " Requesting admin access to fix..." -ForegroundColor Yellow try { - # Spawn an elevated process to set the registry key (triggers UAC prompt) + # Elevated process sets the registry key (triggers UAC prompt). $proc = Start-Process -FilePath "reg.exe" ` -ArgumentList 'add "HKLM\SYSTEM\CurrentControlSet\Control\FileSystem" /v LongPathsEnabled /t REG_DWORD /d 1 /f' ` -Verb RunAs -Wait -PassThru -ErrorAction Stop @@ -1666,8 +1541,7 @@ Ensure-VCRedist # ============================================ # 1c. CMake (only needed for a llama.cpp SOURCE build -- detection only) # ============================================ -# Detection only: the prebuilt path needs no compiler, so do not install or exit -# here. Ensure-BuildToolsForLlamaSourceBuild installs CMake if a source build runs. +# Detection only: the prebuilt path needs no compiler. Ensure-BuildToolsForLlamaSourceBuild installs CMake for a source build. $HasCmake = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue) if ($HasCmake) { step "cmake" "$(cmake --version | Select-Object -First 1)" @@ -1678,8 +1552,7 @@ if ($HasCmake) { # ============================================ # 1d. Visual Studio Build Tools (only needed for a llama.cpp SOURCE build -- detection only) # ============================================ -# Detection only: detect VS for a possible source build, but never install or exit -# here. Install is deferred to Ensure-BuildToolsForLlamaSourceBuild. +# Detection only; install is deferred to Ensure-BuildToolsForLlamaSourceBuild. $CmakeGenerator = $null $VsInstallPath = $null $vsResult = Find-VsBuildTools @@ -1696,28 +1569,23 @@ if ($vsResult) { # ============================================ # 1e. CUDA Toolkit (nvcc for llama.cpp build + env vars) # ============================================ -# Defined here but invoked lazily right before a Phase 4 source build; the -# prebuilt llama.cpp path needs no local toolkit. With -RequireOrExit a source -# build is committed, so hard-fail if no driver-compatible toolkit can be found -# or installed. Without it, detection is best-effort and only sets the flag. +# Invoked lazily before a Phase 4 source build (the prebuilt path needs no local toolkit). +# With -RequireOrExit a source build is committed, so hard-fail if no driver-compatible +# toolkit can be found or installed; otherwise best-effort, only sets the flag. function Resolve-CudaToolkit { param([switch]$RequireOrExit) -# Toolkit major must be <= the driver's max CUDA major (nvidia-smi "CUDA Version: X.Y"); -# a newer-major toolkit fails at runtime ("ggml_cuda_init: failed to initialize CUDA"). +# Toolkit major must be <= the driver's max CUDA major; a newer-major toolkit fails at runtime. $DriverMaxCuda = $null try { - # Bounded: source-build toolkit resolution must not hang on a wedged smi. - # test_resolve_cuda_toolkit.ps1 extracts this function alone into a child - # pwsh (no Invoke-NvidiaSmiBounded in scope) and stubs nvidia-smi with a - # .ps1 script, so fall back to direct invocation when the bounded runner - # is unavailable; production setup.ps1 always has it defined. + # Bounded so a wedged smi can't hang. test_resolve_cuda_toolkit.ps1 extracts this function + # alone (no Invoke-NvidiaSmiBounded in scope), so fall back to direct invocation. $smiOut = if (Get-Command Invoke-NvidiaSmiBounded -ErrorAction SilentlyContinue) { Invoke-NvidiaSmiBounded $NvidiaSmiExe } else { & $NvidiaSmiExe 2>&1 | Out-String } - # Newer drivers report "CUDA UMD Version: X.Y" instead of "CUDA Version: X.Y"; accept both. + # Newer drivers report "CUDA UMD Version: X.Y"; accept both. if ($smiOut -match "CUDA(?: UMD)? Version:\s+([\d]+)\.([\d]+)") { $DriverMaxCuda = "$($Matches[1]).$($Matches[2])" substep "driver supports up to CUDA $DriverMaxCuda" @@ -1730,12 +1598,9 @@ if ($CudaArch) { substep "GPU Compute Capability = $($CudaArch.Insert($CudaArch.Length-1, '.')) (sm_$CudaArch)" } -# -- Find a toolkit that's compatible with the driver AND the GPU -- -# Strategy: prefer the toolkit at CUDA_PATH (user's existing setup) if it's -# compatible with the driver AND supports the GPU architecture. Only fall back -# to scanning side-by-side installs if CUDA_PATH is missing, points to an -# incompatible version, or can't compile for the GPU. This avoids -# header/binary mismatches when multiple toolkits are installed. +# Find a toolkit compatible with the driver AND the GPU. Prefer CUDA_PATH when it's +# driver-compatible and supports the GPU arch; else scan side-by-side installs (avoids +# header/binary mismatches with multiple toolkits installed). $IncompatibleToolkit = $null $NvccPath = $null @@ -1785,8 +1650,7 @@ if ($DriverMaxCuda) { } } } else { - # No side-by-side match: a major-compatible toolkit may still be on - # PATH/CUDA_PATH/a custom dir; use it, else record it as too-new. + # No side-by-side match: a major-compatible toolkit may still be on PATH/CUDA_PATH; use it, else record it too-new. $AnyNvcc = Find-Nvcc if ($AnyNvcc) { $NvccOut = & $AnyNvcc --version 2>&1 | Out-String @@ -1813,8 +1677,7 @@ if (-not $NvccPath -and $IncompatibleToolkit) { $script:CudaToolkitReady = $false return } - # Reached only by a source build (forced, or after a prebuilt-install failure); - # with no compatible toolkit it must fail (setup.sh degrades to CPU instead). + # Source build only: with no compatible toolkit it must fail. Write-Host "" -ForegroundColor Red Write-Host "========================================================================" -ForegroundColor Red Write-Host "[ERROR] CUDA source build cannot use the installed toolkit with this driver." -ForegroundColor Red @@ -1898,15 +1761,12 @@ if (-not $NvccPath) { $CudaToolkitRoot = Split-Path (Split-Path $NvccPath -Parent) -Parent # CUDA_PATH: used by cmake's find_package(CUDAToolkit) [Environment]::SetEnvironmentVariable('CUDA_PATH', $CudaToolkitRoot, 'Process') -# CudaToolkitDir: the MSBuild property that CUDA .targets checks directly -# Trailing backslash required -- the .targets file appends subpaths to it +# CudaToolkitDir: MSBuild property the CUDA .targets checks; trailing backslash required (.targets appends subpaths). [Environment]::SetEnvironmentVariable('CudaToolkitDir', "$CudaToolkitRoot\", 'Process') -# Always persist CUDA_PATH to User registry so the compatible toolkit is used -# in future sessions (overwrites any existing value pointing to a newer, incompatible version) +# Persist CUDA_PATH to User registry so the compatible toolkit is used in future sessions. [Environment]::SetEnvironmentVariable('CUDA_PATH', $CudaToolkitRoot, 'User') substep "Persisted CUDA_PATH=$CudaToolkitRoot to user environment" -# Clear all versioned CUDA_PATH_V* env vars in this process to prevent -# cmake/MSBuild from discovering a conflicting CUDA installation. +# Clear versioned CUDA_PATH_V* vars in this process so cmake/MSBuild can't find a conflicting install. $cudaPathVars = @([Environment]::GetEnvironmentVariables('Process').Keys | Where-Object { $_ -match '^CUDA_PATH_V' }) foreach ($v in $cudaPathVars) { [Environment]::SetEnvironmentVariable($v, $null, 'Process') @@ -1928,10 +1788,9 @@ if (Add-ToUserPath -Directory $nvccBinDir -Position 'Prepend') { substep "Persisted CUDA bin dir to user PATH" } -# -- Ensure CUDA ↔ Visual Studio integration files exist -- -# When CUDA is installed before VS Build Tools (or VS is reinstalled after CUDA), -# the MSBuild .targets/.props files that let VS compile .cu files are missing. -# cmake fails with "No CUDA toolset found". Fix: copy from CUDA extras dir. +# -- Ensure CUDA <-> Visual Studio integration files exist -- +# CUDA installed before VS Build Tools leaves the MSBuild .targets/.props missing and cmake +# fails with "No CUDA toolset found"; copy them from the CUDA extras dir. if ($VsInstallPath -and $CudaToolkitRoot) { $vsCustomizations = Get-VcBuildCustomizationsDir -VsInstallPath $VsInstallPath -Generator $CmakeGenerator $cudaExtras = Join-Path $CudaToolkitRoot "extras\visual_studio_integration\MSBuildExtensions" @@ -1943,7 +1802,7 @@ if ($VsInstallPath -and $CudaToolkitRoot) { Copy-Item "$cudaExtras\*" $vsCustomizations -Force -ErrorAction Stop substep "CUDA VS integration files installed" } catch { - # Direct copy failed (needs admin). Try elevated copy via Start-Process. + # Direct copy failed (needs admin); try elevated copy. try { $copyCmd = "Copy-Item '$cudaExtras\*' '$vsCustomizations' -Force" Start-Process powershell -ArgumentList "-NoProfile -Command $copyCmd" -Verb RunAs -Wait -ErrorAction Stop @@ -1970,8 +1829,7 @@ step "cuda" $NvccPath substep "CUDA_PATH = $CudaToolkitRoot" substep "CudaToolkitDir = $CudaToolkitRoot\" -# $CudaArch was detected earlier (before toolkit selection) so it could -# influence which toolkit we picked. Just log the final state here. +# $CudaArch was detected earlier (it influenced toolkit selection); just log final state. if (-not $CudaArch) { substep "could not detect compute capability -- cmake will use defaults" "Yellow" } @@ -1986,8 +1844,7 @@ if ($HasROCm) { $rocmVerLabel = if ($script:ROCmVersionFull) { "ROCm $script:ROCmVersionFull" } elseif ($script:ROCmVersion) { "ROCm $script:ROCmVersion" } else { "ROCm (version unknown)" } step "rocm" $rocmVerLabel } elseif ($script:ROCmGfxArch) { - # GPU training/inference works via AMD's bundled-runtime ROCm PyTorch wheels; - # the HIP SDK is optional (only the system ROCm toolchain). + # GPU works via AMD's bundled-runtime ROCm wheels; HIP SDK optional. step "rocm" "GPU via bundled ROCm wheels ($script:ROCmGfxArch) -- HIP SDK optional" "Cyan" } elseif ($ROCmGpuLabel) { step "rocm" "AMD GPU detected -- arch unknown; HIP SDK not found" "Yellow" @@ -1996,8 +1853,7 @@ if ($HasROCm) { # ============================================ # 1f. Node.js / npm (skip if pip-installed or Tauri -- only needed for frontend build) # ============================================ -# Frontend and OXC share this Node floor. The helper returns: -# system | bundled | skip. +# Frontend and OXC share this Node floor. Returns system | bundled | skip. function Get-NodeDecision { param( [string]$NodeVersion, # `node -v` output, e.g. v22.17.1 (or empty) @@ -2028,8 +1884,7 @@ $SysNpmVersion = "" $NodeSource = $null if (-not $IsPipInstall) { - # Put Node beside the Unsloth root. OXC can still need npm when the - # frontend build is skipped. + # Put Node beside the Unsloth root (OXC can need npm even when the frontend build is skipped). if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $NodeOverride = $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $NodeOverride = $env:STUDIO_HOME.Trim() } if ($NodeOverride) { @@ -2044,8 +1899,7 @@ if (-not $IsPipInstall) { exit 1 } $NodeParent = (Resolve-Path -LiteralPath $NodeOverride).Path - # An override pointing at the legacy default maps to the legacy sibling - # ~/.unsloth/node (what the runtime resolver and setup.sh use), not /node. + # An override at the legacy default maps to the legacy sibling ~/.unsloth/node, not /node. $_legacyStudio = Join-Path $env:USERPROFILE ".unsloth\studio" if (Test-Path -LiteralPath $_legacyStudio -PathType Container) { $_legacyStudio = (Resolve-Path -LiteralPath $_legacyStudio).Path @@ -2059,11 +1913,9 @@ if (-not $IsPipInstall) { } $NodeDir = Join-Path $NodeParent "node" - # Probe system node/npm without letting a missing/broken command abort setup. - # Under $ErrorActionPreference = "Stop" a bare `node -v` for an absent node - # throws a terminating error `2>$null` cannot swallow, and a present-but-broken - # shim throws too. Guard with Get-Command (node/npm independently) + try/catch; - # empty version => Get-NodeDecision returns "bundled". + # Probe system node/npm without aborting setup: under -ErrorAction Stop a bare `node -v` + # for an absent/broken node throws past 2>$null, so guard with Get-Command + try/catch + # (empty version => Get-NodeDecision returns "bundled"). $SysNodeVersion = try { if (Get-Command node -ErrorAction SilentlyContinue) { (node -v 2>$null) } else { "" } } catch { "" } $SysNpmVersion = try { if (Get-Command npm -ErrorAction SilentlyContinue) { (npm -v 2>$null) } else { "" } } catch { "" } $NodeSource = Get-NodeDecision -NodeVersion "$SysNodeVersion" -NpmVersion "$SysNpmVersion" -SkipInstall "$($env:UNSLOTH_SKIP_NODE_INSTALL)" @@ -2074,8 +1926,7 @@ if ($IsPipInstall) { } elseif ($SkipFrontend) { step "frontend" "bundled (Tauri)" } else { - # Stale npm used to trigger system Node changes. Keep this process-local - # and provision only when the build or OXC needs Node. + # Provision Node process-local, only when the build or OXC needs it. if ($NodeSource -eq "system") { substep "Node $SysNodeVersion and npm $SysNpmVersion already meet requirements (system)." } elseif ($NodeSource -eq "bundled") { @@ -2085,9 +1936,8 @@ if ($IsPipInstall) { } } -# Conda CPython ships modified DLL search paths that break torch's c10.dll -# loading on Windows; a venv made from conda Python inherits its base_prefix, -# so check the executable path AND sys.base_prefix. +# Conda CPython breaks torch's c10.dll loading on Windows; a venv from conda Python inherits +# its base_prefix, so check the executable path AND sys.base_prefix. $CondaSkipPattern = '(?i)(conda|miniconda|anaconda|miniforge|mambaforge)' function Test-IsConda { param([string]$Exe) @@ -2099,18 +1949,15 @@ function Test-IsConda { return $false } -# 1g. Python (>= 3.11 and < 3.14). Prefer the interpreter install.ps1 already -# resolved and built the venv with (UNSLOTH_SETUP_PYTHON), or the existing -# venv python, before re-probing a system where a 3.14 or a WindowsApps stub -# ahead on PATH would trip the gate. setup.ps1 only updates packages in that -# venv, so the handoff is safe to reuse once validated. +# 1g. Python (>= 3.11 and < 3.14). Prefer the interpreter install.ps1 built the venv with +# (UNSLOTH_SETUP_PYTHON) or the existing venv python, before re-probing a system where a 3.14 +# or WindowsApps stub ahead on PATH would trip the gate. function Resolve-ReusedSetupPython { if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_SETUP_PYTHON) -and (Test-Path -LiteralPath $env:UNSLOTH_SETUP_PYTHON)) { return $env:UNSLOTH_SETUP_PYTHON } - # Standalone `unsloth studio setup/update` (install.ps1 did not run): derive - # the venv python from the studio root, mirroring the resolver below. + # Standalone setup/update (install.ps1 did not run): derive the venv python from the studio root. $root = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() } else { Join-Path $env:USERPROFILE ".unsloth\studio" } @@ -2167,10 +2014,7 @@ if ($ReusedSetupPython) { } } -# Fall back to every py.exe on PATH (all-users and per-user launchers can both -# register). -All is required: Windows PowerShell 5.1 returns only the first -# launcher without it, and the PowerShell 7 multi-match array breaks the call -# operator if used directly. +# Fall back to every py.exe on PATH. -All is required: PS 5.1 returns only the first launcher without it. $PyLaunchers = if ($PythonOk) { @() } else { @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue) } foreach ($PyLauncher in $PyLaunchers) { @@ -2180,9 +2024,7 @@ foreach ($PyLauncher in $PyLaunchers) { $out = & $PyLauncher.Source "-$minor" --version 2>&1 | Out-String if ($out -match 'Python (3\.\d+\.\d+)') { $DetectedPyVer = $Matches[1] - # Make `python` resolvable for the rest of setup. Without this, - # py-launcher-only installs (no python.exe on PATH) pass the gate - # and then crash on the first bare `python` call below. + # Make `python` resolvable: a py-launcher-only install would otherwise crash on the first bare `python`. try { $resolvedExe = (& $PyLauncher.Source "-$minor" -c "import sys; print(sys.executable)" 2>$null | Select-Object -First 1) if ($resolvedExe -and (Test-Path $resolvedExe)) { @@ -2211,10 +2053,8 @@ if (-not $PythonOk -and $HasPython) { if ($PythonOk) { substep "Python $DetectedPyVer" } elseif (-not $HasPython) { - # No `python` on PATH (and py.exe either absent or only had unsupported - # minors). Try winget as before -- gating on $HasPython alone, not also - # on $PyLauncher, so a launcher-only install with just 3.14 still gets - # an automatic 3.12 install instead of a hard error. + # No `python` on PATH (py.exe absent or only unsupported minors). Try winget; gating on + # $HasPython alone lets a launcher-only 3.14 install still get an automatic 3.12. Write-Host "Python 3.11-3.13 not found -- installing Python 3.12 via winget..." -ForegroundColor Yellow $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue) if ($HasWinget) { @@ -2230,8 +2070,7 @@ if ($PythonOk) { step "python" "$(python --version 2>&1)" $PythonOk = $true } else { - # python.exe is on PATH but its version is unsupported, and py.exe (if - # present) had no supported minor either. + # python.exe on PATH is unsupported and py.exe had no supported minor either. Write-Host "[ERROR] No supported Python (3.11-3.13) found on this system." -ForegroundColor Red Write-Host " py.exe could not locate -3.11/-3.12/-3.13 and `python` on PATH is unsupported." -ForegroundColor Yellow Write-Host " Install Python 3.12 from https://python.org/downloads/" -ForegroundColor Yellow @@ -2243,7 +2082,7 @@ $ScriptsDir = python -c "import os, sysconfig; p = sysconfig.get_path('scripts', if ($LASTEXITCODE -eq 0 -and $ScriptsDir -and (Test-Path $ScriptsDir)) { # Append (not Prepend) -- this dir has other pip scripts; shim handles unsloth. if (Add-ToUserPath -Directory $ScriptsDir) { - # Also add to current process so it's available immediately + # Add to current process for immediate availability. $ProcessPathEntries = $env:PATH.Split(';') if (-not ($ProcessPathEntries | Where-Object { $_.TrimEnd('\') -eq $ScriptsDir })) { $env:PATH = "$ScriptsDir;$env:PATH" @@ -2260,8 +2099,7 @@ Write-Host "" # PHASE 2: Frontend build (skip if pip-installed -- already bundled) # ========================================================================== $DistDir = Join-Path $FrontendDir "dist" -# Skip build if dist/ exists and no tracked input is newer than dist/. -# Checks src/, public/, package.json, config files -- not just src/. +# Skip build if dist/ exists and no tracked input (src/, public/, package.json, configs) is newer. $NeedFrontendBuild = $true if ($IsPipInstall) { $NeedFrontendBuild = $false @@ -2295,9 +2133,8 @@ if ($IsPipInstall) { } } -# Provision Node when the frontend build OR the OXC runtime install needs it (the -# OXC `npm install` runs whenever its dir exists, regardless of dist staleness); -# never eagerly. System Node is used read-only; the isolated one is ours. +# Provision Node only when the frontend build OR the OXC runtime install needs it. System +# Node is used read-only; the isolated one is ours. $NeedNodeForSetup = (-not $IsPipInstall) -and ($NeedFrontendBuild -or (Test-Path $OxcValidatorDir)) if ($NeedNodeForSetup) { if ($NodeSource -eq "skip") { @@ -2309,8 +2146,7 @@ if ($NeedNodeForSetup) { substep "install a suitable Node + npm, or unset UNSLOTH_SKIP_NODE_INSTALL to let Unsloth manage an isolated Node" "Yellow" } elseif ($NodeSource -eq "bundled") { New-Item -ItemType Directory -Force -Path $NodeParent -ErrorAction SilentlyContinue | Out-Null - # Minimal ownership guard for a custom-home dir (the full Unsloth-owned - # helpers are defined later); never os.replace over a user-owned dir. + # Minimal ownership guard for a custom-home dir; never os.replace over a user-owned dir. if ($NodeOverride -and (Test-Path -LiteralPath $NodeDir -PathType Container)) { $nodeOwnedMarker = Join-Path $NodeDir ".unsloth-studio-owned" $nodeMeta = Join-Path $NodeDir "UNSLOTH_NODE_PREBUILT_INFO.json" @@ -2321,8 +2157,7 @@ if ($NeedNodeForSetup) { } } substep "installing isolated Node (system Node/npm left untouched)..." - # The main Python resolver runs later; bare `python` may be a Store stub or - # absent this early, so prefer the validated handed-off/venv Python. + # Prefer the validated handed-off/venv Python: bare `python` may be a Store stub this early. $NodeInstallPython = if ($ValidatedSetupPython) { $ValidatedSetupPython } else { "python" } $nodeOut = & $NodeInstallPython "$PSScriptRoot\install_node_prebuilt.py" --install-dir $NodeDir 2>&1 | Out-String $nodeExit = $LASTEXITCODE @@ -2339,8 +2174,7 @@ if ($NeedNodeForSetup) { if ($NodeOverride -and (Test-Path -LiteralPath $NodeDir -PathType Container)) { New-Item -ItemType File -Force -Path (Join-Path $NodeDir ".unsloth-studio-owned") -ErrorAction SilentlyContinue | Out-Null } - # Windows Node zip ships node.exe + npm.cmd at the root; prepend it (this - # process only) so node/npm/bun resolve here for the build. + # Windows Node zip ships node.exe + npm.cmd at the root; prepend it (this process only). $env:PATH = "$NodeDir;" + $env:PATH # Keep npm and module resolution inside the isolated Node. $env:NPM_CONFIG_PREFIX = $NodeDir @@ -2356,8 +2190,7 @@ if ($NeedNodeForSetup) { Invoke-SetupCommand { npm install -g bun --allow-scripts=bun @NpmRegistryArgs } | Out-Null $ErrorActionPreference = $prevEAP_bun Refresh-Environment - # Refresh-Environment rebuilds PATH (Machine;User;current), demoting the - # isolated-Node prepend; re-prepend so it wins for the build and OXC step. + # Refresh-Environment demotes the isolated-Node prepend; re-prepend so it wins. $env:PATH = "$NodeDir;" + $env:PATH $env:NPM_CONFIG_PREFIX = $NodeDir $env:npm_config_prefix = $NodeDir @@ -2369,8 +2202,7 @@ if ($NeedNodeForSetup) { } } } else { - # system Node already satisfies requirements; use it as-is. We do NOT - # install global packages (bun) here -- the build falls back to npm. + # system Node satisfies requirements; use as-is (no global bun -- the build falls back to npm). step "node" "$SysNodeVersion | npm $SysNpmVersion (system)" } } @@ -2378,11 +2210,8 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { Write-Host "" substep "building frontend..." - # ── Tailwind v4 .gitignore workaround ── - # Tailwind v4's oxide scanner respects .gitignore in parent directories. - # Python venvs create a .gitignore with "*" (ignore everything), which - # prevents Tailwind from scanning .tsx source files for class names. - # Temporarily hide any such .gitignore during the build, then restore it. + # Tailwind v4 .gitignore workaround: its oxide scanner respects parent .gitignore, and a + # Python venv's "*" .gitignore hides .tsx sources. Temporarily hide any such file, then restore. $HiddenGitignores = @() $WalkDir = (Get-Item $FrontendDir).Parent.FullName while ($WalkDir -and $WalkDir -ne [System.IO.Path]::GetPathRoot($WalkDir)) { @@ -2399,28 +2228,23 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { $WalkDir = Split-Path $WalkDir -Parent } - # Use bun if available (faster install), fall back to npm. - # Bun is used only as package manager; Node runs the actual build (Vite 8). + # Use bun if available (faster install), else npm. Bun is only the package manager; Node runs the build. $prevEAP_npm = $ErrorActionPreference $ErrorActionPreference = "Continue" Push-Location $FrontendDir $UseBun = $null -ne (Get-Command bun -ErrorAction SilentlyContinue) - # bun's package cache can become corrupt -- packages get stored with only - # metadata but no actual content (bin/, lib/). When this happens bun install - # exits 0 but leaves binaries missing. We validate after install and clear - # the cache + retry once before falling back to npm. + # bun's package cache can corrupt (metadata but no bin/lib), so bun install exits 0 with + # binaries missing. Validate after install, clear cache + retry once, then fall back to npm. if ($UseBun) { Write-Host " Using bun for package install (faster)" -ForegroundColor DarkGray $bunExit = Invoke-SetupCommand { bun install @NpmRegistryArgs } - # On Windows, .bin/ entries vary by package manager: - # npm → tsc, tsc.cmd, tsc.ps1 - # bun → tsc.exe, tsc.bunx + # .bin/ entries vary by manager (npm: tsc/.cmd/.ps1; bun: .exe/.bunx). $hasTsc = (Test-Path "node_modules\.bin\tsc") -or (Test-Path "node_modules\.bin\tsc.cmd") -or (Test-Path "node_modules\.bin\tsc.exe") -or (Test-Path "node_modules\.bin\tsc.bunx") $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd") -or (Test-Path "node_modules\.bin\vite.exe") -or (Test-Path "node_modules\.bin\vite.bunx") if ($bunExit -eq 0 -and $hasTsc -and $hasVite) { - # bun install succeeded and critical binaries are present + # bun install succeeded, binaries present } elseif ($bunExit -eq 0) { Write-Host " bun install exited 0 but critical binaries are missing, clearing cache and retrying..." -ForegroundColor Yellow if (Test-Path "node_modules") { @@ -2458,7 +2282,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { } } - # Always use npm to run the build (Node runtime — avoids bun Windows runtime issues) + # Always run the build via npm (Node runtime avoids bun Windows runtime issues). $buildExit = Invoke-SetupCommand { npm run build } if ($buildExit -ne 0) { Pop-Location @@ -2502,8 +2326,7 @@ if ((Test-Path $OxcValidatorDir) -and $NodeSource -ne "skip" -and (Get-Command n $ErrorActionPreference = $prevEAP_oxc step "oxc runtime" "installed" } elseif ((Test-Path $OxcValidatorDir) -and $NodeSource -ne "skip") { - # No npm on PATH (e.g. a pip install with no system Node and no isolated Node - # provisioned). Skip rather than abort; the runtime resolver degrades. Mirrors setup.sh. + # No npm on PATH: skip rather than abort; the runtime resolver degrades. Mirrors setup.sh. substep "OXC validator runtime skipped (no npm found); code validation degrades until Node is available" "Yellow" } @@ -2518,14 +2341,11 @@ Remove-AgentInstructionFiles -Roots @( Write-Host "" substep "setting up Python environment..." -# Find Python -- skip Anaconda/Miniconda distributions ($CondaSkipPattern and -# Test-IsConda are defined above the 1g gate). Standalone CPython (python.org, -# winget, uv) does not have conda's torch c10.dll loading issue. +# Find Python -- skip Anaconda/Miniconda (conda's torch c10.dll loading issue). Standalone +# CPython (python.org, winget, uv) is fine. $PythonCmd = $null -# 0. Reuse the interpreter install.ps1 already resolved and built the venv with -# (UNSLOTH_SETUP_PYTHON, or the existing venv python) before probing the -# system -- it is already validated as supported and non-conda. +# 0. Reuse the interpreter install.ps1 built the venv with (already validated, non-conda). if ($ReusedSetupPython) { try { $out = & $ReusedSetupPython --version 2>&1 | Out-String @@ -2538,10 +2358,8 @@ if ($ReusedSetupPython) { } catch { } } -# 1. Try the Python Launcher (py.exe) first -- most reliable on Windows. -# Enumerate every launcher with -All (Windows PowerShell 5.1 returns only -# the first match without it) and search each for a supported, non-conda -# interpreter. +# 1. Try the Python Launcher (py.exe) first. -All enumerates every launcher (PS 5.1 returns +# only the first without it); search each for a supported, non-conda interpreter. $PyLaunchersResolve = if ($PythonCmd) { @() } else { @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue) } foreach ($pyLauncher in $PyLaunchersResolve) { if ($pyLauncher.Source -match $CondaSkipPattern) { continue } @@ -2551,8 +2369,7 @@ foreach ($pyLauncher in $PyLaunchersResolve) { if ($out -match 'Python 3\.(\d+)') { $pyMinor = [int]$Matches[1] if ($pyMinor -ge 11 -and $pyMinor -le 13) { - # Resolve the actual executable path so venv creation - # does not re-resolve back to a conda interpreter. + # Resolve the actual executable so venv creation doesn't re-resolve to conda. $resolvedExe = (& $pyLauncher.Source "-$minor" -c "import sys; print(sys.executable)" 2>$null | Out-String).Trim() if ($resolvedExe -and (Test-Path $resolvedExe) -and -not (Test-IsConda $resolvedExe)) { $PythonCmd = $resolvedExe @@ -2565,8 +2382,7 @@ foreach ($pyLauncher in $PyLaunchersResolve) { if ($PythonCmd) { break } } -# 2. Fall back to scanning python3.x / python3 / python on PATH. -# Use Get-Command -All to look past conda entries. +# 2. Fall back to scanning python3.x / python3 / python on PATH (-All looks past conda entries). if (-not $PythonCmd) { foreach ($candidate in @("python3.13", "python3.12", "python3.11", "python3", "python")) { foreach ($cmdInfo in @(Get-Command $candidate -All -ErrorAction SilentlyContinue)) { @@ -2600,10 +2416,9 @@ if (-not $PythonCmd) { substep "Python found: $PythonCmd" -# The venv must already exist (created by install.ps1); this script only -# updates packages. UNSLOTH_STUDIO_HOME (or STUDIO_HOME alias) overrides the -# root. UNSLOTH_STUDIO_HOME wins when both are set. Whitespace-only values -# are treated as unset to match Python .strip() semantics. +# The venv must already exist (install.ps1 made it); this script only updates packages. +# UNSLOTH_STUDIO_HOME (or STUDIO_HOME alias) overrides the root; UNSLOTH_STUDIO_HOME wins. +# Whitespace-only values are treated as unset (Python .strip() semantics). $_studioOverrideVar = $null $_studioOverride = $null if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { @@ -2619,9 +2434,7 @@ if ($_studioOverride) { } if (Test-Path -LiteralPath $_studioOverride -PathType Container) { $StudioHome = (Resolve-Path -LiteralPath $_studioOverride).Path - # why: mirror setup.sh:417 and install.ps1:130 -- fail fast when the - # custom root is read-only instead of erroring later while creating - # sidecar venvs / installing packages. + # Fail fast when the custom root is read-only, not later while creating sidecar venvs. $_setupWriteProbe = Join-Path $StudioHome (".unsloth-write-probe-" + [guid]::NewGuid()) try { [System.IO.File]::WriteAllText($_setupWriteProbe, "") @@ -2640,10 +2453,9 @@ if ($_studioOverride) { } $VenvDir = Join-Path $StudioHome "unsloth_studio" -# why: in env-override mode $StudioHome is user-chosen; require the -# ownership marker before Remove-Item so unrelated dirs survive. Gated on -# the canonical comparison so an override pointing at the legacy default -# still behaves like a default install. +# In env-override mode $StudioHome is user-chosen; require the ownership marker before any +# Remove-Item so unrelated dirs survive. Gated on the canonical comparison so an override at +# the legacy default still behaves like a default install. $StudioOwnedMarker = ".unsloth-studio-owned" $LegacyStudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" $_studioHomeCanon = $StudioHome @@ -2654,10 +2466,9 @@ if (Test-Path -LiteralPath $LegacyStudioHome -PathType Container) { $LegacyStudioHome = (Resolve-Path -LiteralPath $LegacyStudioHome).Path } $StudioHomeIsCustom = ($_studioHomeCanon -ne $LegacyStudioHome) -# Directory-local evidence that Unsloth created $Path, used to adopt a custom-home -# llama.cpp predating the .unsloth-studio-owned marker (see setup.sh). Only the -# prebuilt UNSLOTH_PREBUILT_INFO.json counts; source builds are indistinguishable -# from a user clone on Windows and stay under the strict guard. +# Directory-local evidence Unsloth created $Path, to adopt a custom-home llama.cpp predating +# the .unsloth-studio-owned marker. Only the prebuilt UNSLOTH_PREBUILT_INFO.json counts; source +# builds are indistinguishable from a user clone on Windows and stay under the strict guard. function Test-StudioOwnedAdoptable { param([Parameter(Mandatory = $true)][string]$Path) if (Test-Path -LiteralPath (Join-Path $Path "UNSLOTH_PREBUILT_INFO.json") -PathType Leaf) { return $true } @@ -2687,12 +2498,9 @@ function Mark-StudioOwned { } catch {} } -# Stale-venv detection: if the venv exists but its torch flavor no longer -# matches the current machine, repair according to invocation context. -# - install.ps1 sets UNSLOTH_INSTALL_ROLLBACK_MANAGED=1 so setup can delegate -# to the installer-level rollback that restores the previous environment. -# - direct `unsloth studio update` keeps the pre-existing self-repair behavior. -# In no-torch mode, a missing torch package is expected. +# Stale-venv detection: if the venv's torch flavor no longer matches the machine, repair by +# context. install.ps1 sets UNSLOTH_INSTALL_ROLLBACK_MANAGED=1 to delegate to installer-level +# rollback; direct `studio update` self-repairs. No-torch mode expects a missing torch. $NoTorchMode = $env:UNSLOTH_NO_TORCH -match '^(?i:true|1|yes)$' $InstallerManagedSetup = $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -match '^(?i:true|1|yes)$' if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode) { @@ -2718,8 +2526,7 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode if ($torchVer -match '\+(cu\d+)') { $installedTorchTag = $Matches[1] } elseif ($torchVer -match '\+rocm') { - # Any +rocm / gfx wheel -> generic "rocm" flavor (the exact version is - # repaired later by install_python_stack.py; here we only need the flavor). + # Any +rocm / gfx wheel -> generic "rocm" flavor (exact version repaired later). $installedTorchTag = "rocm" } elseif ($torchVer -match '\+cpu') { $installedTorchTag = "cpu" @@ -2747,29 +2554,24 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode # Digit-gated like the install selection: a custom rocm-* leaf (rocm-current / # rocm-rel-7.2.1) is NOT a ROCm family and must not be stale-compared. if (Test-PipRocmFamilyLeaf $_pinLeaf) { - # Don't collapse a pinned ROCm/gfx leaf to a generic "rocm" (would mask a family - # change, rocm6.4 -> gfx1151). Get-RocmPinStaleTags uses the SAME 2.11 allowlist - # as the install path, so a gfx110X-all/gfx90a/gfx908 pin on a <2.11 wheel is NOT stale. + # Keep a pinned ROCm/gfx leaf specific (a generic "rocm" would mask rocm6.4 -> gfx1151). + # Get-RocmPinStaleTags uses the SAME 2.11 allowlist as the install path. $_rocmTags = Get-RocmPinStaleTags -PinLeaf $_pinLeaf -TorchVersion $torchVer $expectedTorchTag = $_rocmTags.Expected $installedTorchTag = $_rocmTags.Installed } elseif ((Test-CudaFamilyLeaf $_pinLeaf) -or $_pinLeaf -eq 'cpu') { - # cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds; - # /custom and /current fall through to the unknown-index branch below. + # cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds; /custom, /current fall through. $expectedTorchTag = $_pinLeaf } else { - # Custom index whose leaf is not a torch flavor (a /simple mirror): the - # flavor can't be inferred, so never treat the venv as stale over it. + # Custom index whose leaf isn't a torch flavor (/simple mirror): flavor unknown, never stale. $_expectedKnown = $false $expectedTorchTag = $installedTorchTag } } elseif ($HasNvidiaSmi) { $expectedTorchTag = Get-PytorchCudaTag } elseif ($HasROCm -or $script:ROCmGfxArch) { - # AMD/ROCm host with no explicit pin: an existing +rocm wheel is correct (gfx arch - # counts even when $HasROCm is false). But only the arches the install path maps to a - # repo.amd.com index get ROCm torch; an unmapped arch installs CPU, so expect "cpu" - # for those or a correct CPU venv rebuilds every update. + # AMD/ROCm host, no pin: an existing +rocm wheel is correct. Only arches the install + # path maps to a repo.amd.com index get ROCm torch; an unmapped arch expects "cpu". $_rocmWheelArches = @( "gfx1201", "gfx1200", # RDNA 4 "gfx1151", "gfx1150", # RDNA 3.5 (Strix Halo/Point) @@ -2777,9 +2579,8 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode "gfx90a", "gfx908" # MI200 / MI100 ) if ($script:ROCmGfxArch -and ($_rocmWheelArches -contains $script:ROCmGfxArch)) { - # A correct +rocm wheel is not stale. A CPU wheel on a supported AMD arch is - # NOT wiped either (the AMD Windows ROCm override below upgrades it in place); - # expect "cpu" for that case. A wrong CUDA wheel still rebuilds. + # A correct +rocm wheel isn't stale. A CPU wheel on a supported AMD arch is upgraded + # in place (expect "cpu"), not wiped. A wrong CUDA wheel still rebuilds. if ($installedTorchTag -eq "cpu") { $expectedTorchTag = "cpu" } else { @@ -2796,9 +2597,8 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode } } - # A stale venv under a pin whose torch still imports is repaired IN PLACE (the dependency - # pass force-reinstalls from the pin). The rebuild path wipes the venv and would strand a - # direct `studio update`; only a broken venv or an unpinned drift wipes. + # A stale venv under a pin whose torch still imports is repaired IN PLACE (force-reinstall); + # only a broken venv or an unpinned drift wipes. if ($shouldRebuild -and $_pinnedIdx -and $installedTorchTag) { substep "Torch-index pin changed ($installedTorchTag) -- reinstalling torch from the pin in place." "Cyan" $script:PinChangedForceReinstall = $true @@ -2814,9 +2614,8 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode exit 1 } substep "Stale venv detected ($reason) -- rebuilding..." "Yellow" - # why: mirror install.ps1 env-mode guard so an update against a custom - # UNSLOTH_STUDIO_HOME never wipes an unrelated unsloth_studio venv; - # -PathType Leaf rejects a directory masquerading as the sentinel. + # Mirror install.ps1 env-mode guard so an update against a custom UNSLOTH_STUDIO_HOME + # never wipes an unrelated venv; -PathType Leaf rejects a directory sentinel. if ( $StudioHomeIsCustom -and -not (Test-Path -LiteralPath (Join-Path $VenvDir $StudioOwnedMarker) -PathType Leaf) -and @@ -2853,10 +2652,8 @@ if (-not (Test-Path -LiteralPath $VenvDir)) { } } -# pip and python write to stderr even on success (progress bars, warnings). -# With $ErrorActionPreference = "Stop" (set at top of script), PS 5.1 -# converts stderr lines into terminating ErrorRecords, breaking output. -# Lower to "Continue" for the pip/python section. +# pip/python write to stderr even on success, which -ErrorAction Stop turns into terminating +# errors on PS 5.1; lower to "Continue" for the pip/python section. $prevEAP = $ErrorActionPreference $ErrorActionPreference = "Continue" @@ -2872,8 +2669,7 @@ if (Get-Command uv -ErrorAction SilentlyContinue) { try { Invoke-SetupCommand { Invoke-Expression (Invoke-RestMethod -Uri "https://astral.sh/uv/install.ps1") } | Out-Null Refresh-Environment - # Re-activate venv since Refresh-Environment rebuilds PATH from - # registry and drops the venv's Scripts directory + # Re-activate: Refresh-Environment rebuilds PATH and drops the venv's Scripts dir. . $ActivateScript if (Get-Command uv -ErrorAction SilentlyContinue) { $UseUv = $true } } catch { } @@ -2882,11 +2678,9 @@ if (Get-Command uv -ErrorAction SilentlyContinue) { # Helper: install a package, preferring uv with pip fallback function Fast-Install { param([Parameter(ValueFromRemainingArguments=$true)]$Args_) - # An explicit --index-url must win: inherited uv index vars otherwise pull CPU torch over - # the CUDA/ROCm build (#6898), so drop them for pinned installs (scrub covers the whole - # function since the pip fallback honours PIP_* too). UV_TORCH_BACKEND / UV_FIND_LINKS also - # reroute; UV_NO_CONFIG=1 (+ dropping UV_CONFIG_FILE) stops a uv.toml index outranking the - # pin (uv 0.10); PIP_NO_INDEX / PIP_INDEX_URL would defeat the pinned --index-url in pip. + # An explicit --index-url must win: inherited uv/pip index vars otherwise pull CPU torch over + # the CUDA/ROCm build (#6898), so drop them for pinned installs (pip fallback honours PIP_* too). + # UV_NO_CONFIG=1 (+ dropping UV_CONFIG_FILE) stops a uv.toml index outranking the pin (uv 0.10). $saved = @{} $pinned = @($Args_) -contains '--index-url' if ($pinned) { @@ -2898,8 +2692,7 @@ function Fast-Install { Remove-Item "Env:$n" -ErrorAction SilentlyContinue } $env:UV_NO_CONFIG = '1' - # A `pip config` global.extra-index-url still adds indexes to the pip FALLBACK; - # PIP_CONFIG_FILE = 'nul' (Windows devnull) loads NO config (uv ignores pip config). + # PIP_CONFIG_FILE = 'nul' (Windows devnull) stops a `pip config` extra-index-url adding indexes to the pip fallback. $env:PIP_CONFIG_FILE = 'nul' } try { @@ -2919,9 +2712,7 @@ function Fast-Install { } } -# ── Check if Python deps need updating ── -# Compare installed package version against PyPI latest. -# Skip all Python dependency work if versions match (fast update path). +# Check if Python deps need updating: compare installed version to PyPI latest, skip if they match. $_PkgName = if ($env:STUDIO_PACKAGE_NAME) { $env:STUDIO_PACKAGE_NAME } else { "unsloth" } $SkipPythonDeps = $false @@ -2937,9 +2728,8 @@ if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") { if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer)) { step "python" "$_PkgName $InstalledVer is up to date" $SkipPythonDeps = $true - # A pre-#6483-fix install can be stuck on anyio>=4.14 even though - # $_PkgName itself is current; the fast path above would otherwise - # never reach install_python_stack's anyio repair (#6797). + # A pre-#6483 install stuck on anyio>=4.14 with $_PkgName current would skip the + # install_python_stack anyio repair (#6797); force the pass. $_anyioBad = $false try { & python -c " @@ -2959,10 +2749,8 @@ sys.exit(0 if (major, minor) >= (4, 14) else 1) substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." "Cyan" $SkipPythonDeps = $false } - # ...but not if an AMD GPU is present and installed PyTorch is CPU-only - # (host predates ROCm-wheel support, or GPU added later): the fast "up to - # date" path would leave the user on CPU torch with Train/Export disabled. - # Force the dependency pass so the ROCm wheels get installed. + # AMD GPU present but installed PyTorch is CPU-only: the fast path would leave the user + # on CPU torch with Train/Export disabled, so force the pass to install ROCm wheels. if ($script:ROCmGfxArch) { $_torchIsCpu = $true try { @@ -2981,26 +2769,7 @@ sys.exit(0 if (major, minor) >= (4, 14) else 1) } } -# if (-not $IsPipInstall) { -# # Running from repo: copy requirements and do editable install -# $RepoRoot = (Resolve-Path (Join-Path $ScriptDir "..\..")).Path -# $ReqsSrc = Join-Path $RepoRoot "backend\requirements" -# $ReqsDst = Join-Path $PackageDir "requirements" -# if (-not (Test-Path $ReqsDst)) { New-Item -ItemType Directory -Path $ReqsDst | Out-Null } -# Copy-Item (Join-Path $ReqsSrc "*.txt") $ReqsDst -Force - -# Write-Host " Installing CLI entry point..." -ForegroundColor Cyan -# pip install -e $RepoRoot 2>&1 | Out-Null -# } else { -# # Running from pip install: the package is in system Python but not in -# # the fresh .venv. Install it so run_install() can find its modules -# # and bundled requirements files. -# Write-Host " Installing package into venv..." -ForegroundColor Cyan -# pip install unsloth 2>&1 | Out-Null -# } - -# A torch-index pin change repairs in place: force the dependency pass so the torch install -# below force-reinstalls from the new pin (else the fast path keeps the old wheel). +# A torch-index pin change repairs in place: force the pass so torch force-reinstalls from the new pin. if ($script:PinChangedForceReinstall) { $SkipPythonDeps = $false } if (-not $SkipPythonDeps) { @@ -3011,15 +2780,11 @@ if ($script:UnslothVerbose) { Fast-Install --upgrade pip | Out-Null } -# Pre-install PyTorch with CUDA support. -# On Windows, the default PyPI torch wheel is CPU-only. -# We need PyTorch's CUDA index to get GPU-enabled wheels. -# PyTorch bundles its own CUDA runtime, so this works regardless -# of whether the CUDA Toolkit is installed yet. -# The CUDA tag is chosen based on the driver's max supported CUDA version. +# Pre-install PyTorch with CUDA support: the default Windows PyPI wheel is CPU-only, so use +# PyTorch's CUDA index (tag chosen from the driver's max CUDA; wheels bundle their own runtime). -# Triton/inductor filenames are long and can hit Windows MAX_PATH (260). With long -# paths on, cache under Unsloth home; else use a short drive-root dir for headroom. +# Triton/inductor filenames can hit Windows MAX_PATH (260). With long paths on, cache under +# Unsloth home; else a short drive-root dir for headroom. if ($LongPathsEnabled) { $TorchCacheDir = Join-Path $StudioHome "TORCHINDUCTOR_CACHE_DIR" } else { @@ -3030,8 +2795,7 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir [Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User') substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" -# Explicit pin (URL or family) wins over GPU probing and suppresses the AMD reroute below; -# matches install.sh / install.ps1 / install_python_stack.py. +# Explicit pin (URL or family) wins over GPU probing and suppresses the AMD reroute below. $PinnedTorchIndexUrl = Get-PinnedTorchIndexUrl $TorchIndexPinned = [bool]$PinnedTorchIndexUrl if ($PinnedTorchIndexUrl) { @@ -3042,21 +2806,13 @@ if ($PinnedTorchIndexUrl) { $CuTag = "cpu" } -# ── GPU arch → newest compatible Windows ROCm wheel release ── -# Wheels bundle their own ROCm runtime; the installed HIP SDK version does -# not constrain which release to use. Always picks the newest release that -# supports the GPU architecture. -# ── AMD Windows ROCm torch override ────────────────────────────────────────── -# Uses AMD's arch-specific pip index (repo.amd.com/rocm/whl/{arch}/). -# Wheels bundle their own ROCm runtime; HIP SDK version is irrelevant. +# ── AMD Windows ROCm torch override ── +# AMD's arch-specific pip index (repo.amd.com/rocm/whl/{arch}/); wheels bundle their own runtime. $ROCmGfxArch = $script:ROCmGfxArch $ROCmIndexUrl = $null -# Install AMD ROCm PyTorch wheels when ROCm is confirmed OR a gfx arch is known -# (name-inferred on Adrenalin-only hosts). The per-arch wheels bundle the runtime -# (rocm-sdk-libraries-), so torch.cuda.is_available() is True without a HIP -# SDK -- which flips Unsloth out of chat-only (CHAT_ONLY) and enables Train/Export. -# Gating on $HasROCm alone left Strix Halo / Radeon 8060S on CPU torch; a failed -# ROCm install still falls back to CPU below, so this is safe. +# Install AMD ROCm wheels when ROCm is confirmed OR a gfx arch is known (name-inferred on +# Adrenalin-only hosts): the per-arch wheels bundle the runtime so torch.cuda.is_available() is +# True without a HIP SDK, flipping Unsloth out of chat-only. A failed ROCm install falls back to CPU. if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ @@ -3066,20 +2822,14 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu "gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all" "gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100 } - # gfx120X and Strix have a null _grouped_mm kernel on torch <2.11.0. - # Mirrors the $torchFloorMap in install.ps1 so both installers enforce - # the same floor and ceiling when pulling from AMD's per-arch index. + # gfx120X/Strix have a null _grouped_mm kernel on torch <2.11.0. Mirrors $torchFloorMap in install.ps1. $torchFloorMap = @{ "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" } - # Companion ranges for torchvision/torchaudio -- must stay in sync with the - # torch ceiling so pip can always find a consistent trio on AMD's per-arch - # index. AMD publishes each package independently and may add a newer - # torchvision (e.g. 0.27 for torch 2.12) before removing 0.26, which would - # cause pip to resolve an ABI-incompatible set if these are left bare. - # Matches _ROCM_TORCH_PKG_SPECS["rocm7.2"] in install_python_stack.py. - # Bump all three ceilings together when torch 2.12.x is validated. + # torchvision/torchaudio companions kept in sync with the torch ceiling so pip finds a + # consistent trio (AMD publishes each independently). Matches _ROCM_TORCH_PKG_SPECS["rocm7.2"]; + # bump all three ceilings together when torch 2.12.x is validated. $torchvisionFloorMap = @{ "gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0" "gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0" @@ -3095,33 +2845,28 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu if ($archFamily) { $ROCmIndexUrl = "$amdIndexBase/$archFamily/" } elseif ($ROCmGfxArch) { - # GPU arch detected but not in the supported wheel map — warn explicitly - # so the user knows why they are getting CPU PyTorch instead of ROCm. + # GPU arch detected but not in the supported wheel map -- warn why torch is CPU-only. substep "[WARN] AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow" substep " Supported: gfx1200/1201 (RDNA 4), gfx1150/1151 (RDNA 3.5), gfx1100-1103 (RDNA 3), gfx90a, gfx908" "Yellow" } else { - # HIP SDK present ($HasROCm=true via amd-smi) but gcnArchName was not - # readable — warn rather than silently falling back to CPU PyTorch. + # HIP SDK present but gcnArchName unreadable -- warn rather than silently CPU torch. substep "[WARN] AMD GPU detected (HIP SDK present) but GPU arch could not be read -- falling back to CPU-only PyTorch" "Yellow" substep " Arch detection requires hipinfo to report gcnArchName. Re-install the HIP SDK if this is unexpected." "Yellow" } } -# A pinned gfx*/rocm index skips the auto-reroute above; route it through the ROCm install path -# with the same floor/companions the unpinned AMD path uses (mirrors install.ps1), else the CUDA -# branch installs bare torch and resolves a known-bad wheel for gfx115x/gfx120x/rocm>=7.2. +# A pinned gfx*/rocm index skips the auto-reroute above; route it through the ROCm path with the +# same floor/companions the unpinned AMD path uses (mirrors install.ps1), else the CUDA branch +# resolves a known-bad wheel for gfx115x/gfx120x/rocm>=7.2. if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) { $_pinLeaf = Get-TorchIndexLeaf $PinnedTorchIndexUrl $_pinRocm211 = $false - # Anchor the match ($) so a suffixed custom leaf (rocm7.2-private) falls through to the - # verbatim install instead of being floored by its rocm7.2 prefix. + # Anchor ($) so a suffixed custom leaf (rocm7.2-private) falls through to the verbatim install. if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') { - # Only KNOWN-2.11 rocm (rocm7.2) gets the floor (no speculative floor). Matches - # Test-RocmKnown211Version / _ROCM_KNOWN_TORCH211_VERSIONS. + # Only KNOWN-2.11 rocm (rocm7.2) gets the floor. Matches Test-RocmKnown211Version. $_pinRocm211 = Test-RocmKnown211Version -Major ([int]$Matches[1]) -Minor ([int]$Matches[2]) } - # Only the 2.11 gfx arches need the floor; others publish <2.11 and stay bare. Reuse - # Test-RocmGfx211Leaf so this allowlist and the stale-venv check never diverge. + # Only the 2.11 gfx arches need the floor. Reuse Test-RocmGfx211Leaf so this and the stale check never diverge. $_pinGfx211 = Test-RocmGfx211Leaf $_pinLeaf if ($_pinGfx211 -or $_pinRocm211) { $ROCmIndexUrl = $PinnedTorchIndexUrl @@ -3130,9 +2875,8 @@ if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) { $ROCmAudioSpec = "torchaudio>=2.11.0,<2.12.0" substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchSpec" "Cyan" } elseif (Test-PipRocmFamilyLeaf $_pinLeaf) { - # Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with - # bare specs. Only EXACT rocm and gfx* are --index-url families; a suffixed - # leaf stays on the verbatim path. Mirrors install.ps1 / _is_pip_rocm_family_leaf. + # Other gfx / older rocm (<=7.1) ship torch <2.11: route via the ROCm path with bare specs. + # Mirrors install.ps1 / _is_pip_rocm_family_leaf. $ROCmIndexUrl = $PinnedTorchIndexUrl $ROCmTorchSpec = "torch" $ROCmVisionSpec = "torchvision" @@ -3142,8 +2886,7 @@ if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) { $PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } -# A full URL pin is used verbatim; a family pin already set $CuTag. A pinned ROCm install -# goes through $ROCmIndexUrl; on failure the fallback uses the CPU index, not the ROCm pin. +# Full URL pin verbatim; family pin already set $CuTag. On ROCm-install failure the fallback uses the CPU index. $TorchInstallIndexUrl = if ($ROCmIndexUrl) { "$PyTorchWhlBase/cpu" } elseif ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" } $ROCmCpuFallback = $false @@ -3166,7 +2909,7 @@ if ($ROCmIndexUrl) { $ROCmIndexUrl = $null $ROCmCpuFallback = $true } else { - # Tell install_python_stack.py to skip probe + suppress manual-install warning. + # Tell install_python_stack.py to skip the probe and suppress the manual-install warning. $env:UNSLOTH_ROCM_TORCH_INSTALLED = "1" substep "GPU ROCm PyTorch installed ($ROCmGfxArch) -- training and GPU inference will use the GPU" "Cyan" } @@ -3174,19 +2917,14 @@ if ($ROCmIndexUrl) { if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { substep "installing PyTorch (CPU-only)..." - # After an AMD ROCm fallback, force-reinstall so a partial ROCm torch (which satisfies the - # CPU torch>= range) is replaced by the CPU build; skip on a genuine CPU host to stay fast. - # $ROCmCpuFallback matters when a PINNED ROCm index failed ($CuTag is still the rocm leaf). - # Build the array directly: an if-expression collapses @("x") to a scalar @splat would - # enumerate char-by-char. + # After an AMD ROCm fallback, force-reinstall so a partial ROCm torch (satisfies the CPU + # torch>= range) is replaced by the CPU build; also on a pin change (a stale +cu/+rocm wheel + # still satisfies the range). Build the array directly (an if-expression @splat enumerates chars). $cpuForce = @() if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") } - # --force-reinstall on a pin change: a stale +cu / +rocm wheel still satisfies the CPU - # torch>= range, so uv would keep it and only swap companions. if ($script:PinChangedForceReinstall) { $cpuForce = @("--force-reinstall") } # A PINNED cpu index installs the bounded trio (parity with _CPU_TORCH_PKG_SPEC): the /cpu - # index serves newer torch, and _ensure_cpu_torch keeps any CPU build, so a bare trio could - # land an unsupported version. Unpinned CPU hosts keep the bare trio (pre-pin behavior). + # index serves newer torch and _ensure_cpu_torch keeps any CPU build. Unpinned keeps the bare trio. $cpuTorchSpec = "torch"; $cpuVisionSpec = "torchvision"; $cpuAudioSpec = "torchaudio" if ($TorchIndexPinned) { $cpuTorchSpec = "torch>=2.4,<2.12.0" @@ -3209,14 +2947,12 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { } elseif (-not $ROCmIndexUrl) { substep "installing PyTorch with CUDA support ($CuTag)..." substep "(This download is ~2.8 GB -- may take a few minutes)" - # --force-reinstall on a pin change: an installed cuXXX wheel satisfies the bare torch - # requirement (PEP 440 ignores the +cuXXX tag), so without it a changed CUDA pin (cu126 - # -> cu128) never applies. + # --force-reinstall on a pin change: an installed cuXXX wheel satisfies bare torch (PEP 440 + # ignores +cuXXX), so a changed CUDA pin (cu126 -> cu128) never applies without it. $cudaForce = @() if ($script:PinChangedForceReinstall) { $cudaForce = @("--force-reinstall") } - # An unknown-leaf custom pin (/simple, /current) routes here with $CuTag as that leaf. Bound - # the trio like the fresh custom-pin paths so a mirror can't pull an ABI-newer companion - # against the capped torch. Known cu* leaves keep bare specs. + # An unknown-leaf custom pin (/simple, /current) routes here: bound the trio so a mirror can't + # pull an ABI-newer companion against the capped torch. Known cu* leaves keep bare specs. $cudaTorchSpec = "torch" $cudaVisionSpec = "torchvision" $cudaAudioSpec = "torchaudio" @@ -3239,7 +2975,7 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { exit 1 } - # Install Triton for Windows (enables torch.compile -- without it training can hang) + # Triton for Windows enables torch.compile (without it training can hang). substep "installing Triton for Windows..." if ($script:UnslothVerbose) { Fast-Install "triton-windows<3.7" @@ -3257,11 +2993,8 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { } } -# No unsloth.exe rename needed. setup.ps1 runs *via* unsloth.exe, so renaming the -# running launcher only ever failed (WinError 32) and printed a scary warning. It's -# also unnecessary: install.ps1 sets SKIP_STUDIO_BASE=1 (base never reinstalled) and -# 'studio update' goes through uv (--upgrade-package), whose pip fallback no-ops on -# the already-satisfied bare unsloth/unsloth-zoo. Either way unsloth.exe stays. +# No unsloth.exe rename: setup.ps1 runs via unsloth.exe, so renaming the running launcher only +# ever failed (WinError 32), and it's unnecessary (install.ps1 sets SKIP_STUDIO_BASE=1). # Ordered heavy dependency installation -- shared cross-platform script substep "running ordered dependency installation..." @@ -3281,10 +3014,8 @@ if ($stackExit -ne 0) { $ErrorActionPreference = $prevEAP } -# ── Pre-install transformers 5.x into .venv_t5_530/, .venv_t5_550/, and .venv_t5_510/ ── -# Runs outside the deps fast-path gate so that upgrades from the legacy -# single .venv_t5 are always migrated to the tiered layout. -# T5 sidecar venvs live under the resolved $StudioHome so custom installs are self-contained. +# Pre-install transformers 5.x into tiered sidecar venvs under $StudioHome. Runs outside the +# deps fast-path gate so upgrades from the legacy single .venv_t5 always migrate. $VenvT5_530Dir = Join-Path $StudioHome ".venv_t5_530" $VenvT5_550Dir = Join-Path $StudioHome ".venv_t5_550" $VenvT5_510Dir = Join-Path $StudioHome ".venv_t5_510" @@ -3442,9 +3173,8 @@ step "transformers" "5.10.2 pre-installed" # ========================================================================== # PHASE 3.4: Prefer prebuilt llama.cpp bundles before source build # ========================================================================== -# Nest llama.cpp under $StudioHome only for real env-overrides, never the -# legacy default. Reuses $StudioHomeIsCustom from the canonical comparison -# computed above so the llama.cpp nest matches ownership-guard semantics. +# Nest llama.cpp under $StudioHome only for real env-overrides (reuses $StudioHomeIsCustom so +# the nest matches ownership-guard semantics). if ($StudioHomeIsCustom) { $UnslothHome = $StudioHome } else { @@ -3455,10 +3185,8 @@ $LlamaCppDir = Join-Path $UnslothHome "llama.cpp" $NeedLlamaSourceBuild = $false $SkipPrebuiltInstall = $false $RequestedLlamaTag = if ($env:UNSLOTH_LLAMA_TAG) { $env:UNSLOTH_LLAMA_TAG } else { $DefaultLlamaTag } -# Every host installs the fork's app-* prebuilts now: GPU Windows (CUDA / ROCm) -# already did, and the fork now also ships the CPU bundles for Windows x64 and -# arm64 (windows-cpu / windows-arm64). ggml-org artifacts are no longer used by -# default. Mirrors setup.sh's routing. +# Every host installs the fork's app-* prebuilts (GPU CUDA/ROCm + CPU windows-cpu/windows-arm64); +# ggml-org artifacts are no longer used by default. Mirrors setup.sh's routing. $HelperReleaseRepo = "unslothai/llama.cpp" $LlamaPr = if ($env:UNSLOTH_LLAMA_PR) { $env:UNSLOTH_LLAMA_PR.Trim() } else { "" } @@ -3492,9 +3220,7 @@ function Invoke-LlamaHelper { } try { - # Capture all output (stdout + stderr) so that PowerShell does not - # convert stderr lines into visible ErrorRecord objects. Separate - # stdout from stderr afterwards. + # Capture all output (2>&1) so PS doesn't turn stderr into ErrorRecords; split afterwards. $allOutput = & python "$PSScriptRoot\install_llama_prebuilt.py" @Arguments 2>&1 $exitCode = $LASTEXITCODE $stdoutLines = @() @@ -3554,10 +3280,8 @@ if ($LocalLlamaCppSrc) { exit 1 } $ResolvedLocal = (Resolve-Path -LiteralPath $LocalLlamaCppSrc).Path - # Reusing a local dir disables both the prebuilt download and the source - # build, so a runnable llama-server.exe must already be present. Accept any - # layout LlamaCppBackend._layout_candidates() resolves (root-level, build\bin, - # or build\bin\Release) so the flag never rejects a tree Unsloth could run. + # Reusing a local dir disables prebuilt + source build, so a runnable llama-server.exe must + # exist. Accept any layout LlamaCppBackend._layout_candidates() resolves. $LocalLlamaServerFound = $false foreach ($_cand in @( (Join-Path $ResolvedLocal "llama-server.exe"), @@ -3566,10 +3290,8 @@ if ($LocalLlamaCppSrc) { if (Test-Path -LiteralPath $_cand) { $LocalLlamaServerFound = $true; break } } if ($ResolvedLocal -eq $LlamaCppDir) { - # Points at the canonical install location itself: never delete-then-link - # onto itself. Reuse an existing build here (skip prebuilt + source) so the - # staged prebuilt installer can't replace a build the user asked to reuse; - # if nothing is built yet, fall through to the normal install. + # Points at the canonical install location itself: never delete-then-link onto itself. + # Reuse an existing build (skip prebuilt + source); if nothing is built, fall through to normal install. if ($LocalLlamaServerFound) { substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR is the canonical install location and already holds a build; reusing it" "Yellow" $LocalLlamaCppLinked = $true @@ -3578,21 +3300,15 @@ if ($LocalLlamaCppSrc) { substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR points to the canonical install location with nothing built there yet; running the normal install" "Yellow" } } else { - # Fail clearly rather than junction an unbuilt or wrong-platform checkout - # and leave Unsloth with no usable binary. + # Fail clearly rather than junction an unbuilt/wrong-platform checkout with no usable binary. if (-not $LocalLlamaServerFound) { step "llama.cpp" "no llama-server.exe under $ResolvedLocal (looked for .\llama-server.exe, .\build\bin and .\build\bin\Release) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "Red" exit 1 } - # If the target is already a junction/symlink (e.g. a previous - # --with-llama-cpp-dir run), delete only the link via DirectoryInfo.Delete(). - # Remove-Item -Recurse -Force on a reparse point can traverse the link and - # wipe the user's real llama.cpp directory on PowerShell 5.1. Dropping the - # stale link here also keeps the custom-home ownership check below idempotent. - # Use Get-Item -Force (not Test-Path): a *broken* junction whose target was - # moved/deleted makes Test-Path return false, which would leave the dangling - # link in place and make mklink below fail; Get-Item still resolves it so we - # can remove it and relink to a new valid directory. + # If the target is already a junction/symlink, delete only the link via DirectoryInfo.Delete() + # -- Remove-Item -Recurse -Force can traverse a reparse point and wipe the real dir on PS 5.1. + # Use Get-Item -Force (not Test-Path): a broken junction reads false under Test-Path, leaving + # a dangling link that makes mklink fail; Get-Item still resolves it so we can relink. $existing = Get-Item -LiteralPath $LlamaCppDir -Force -ErrorAction SilentlyContinue if ($existing -and ($existing.Attributes -band [System.IO.FileAttributes]::ReparsePoint)) { $existing.Delete() @@ -3602,9 +3318,8 @@ if ($LocalLlamaCppSrc) { } if (Test-Path -LiteralPath $LlamaCppDir) { Remove-Item -Recurse -Force -LiteralPath $LlamaCppDir -ErrorAction SilentlyContinue - # A locked/in-use tree can silently survive removal (SilentlyContinue - # masks it). Don't then junction/copy over a half-present dir; mirror the - # prebuilt path's active-process handling and stop with a clear message. + # A locked tree can survive removal (SilentlyContinue masks it); don't junction/copy + # over a half-present dir -- stop with a clear message like the prebuilt path. if (Test-Path -LiteralPath $LlamaCppDir) { step "llama.cpp" "install blocked by active llama.cpp process" "Yellow" substep "Close Unsloth or other llama.cpp users and retry" "Yellow" @@ -3637,25 +3352,17 @@ if ($LocalLlamaCppLinked) { Write-Host "" if (Test-Path -LiteralPath $LlamaCppDir) { substep "Existing llama.cpp install detected -- validating staged prebuilt update before replacement" - # If the existing install is the wrong kind (e.g. windows-cpu on a ROCm - # machine that should have windows-rocm), remove it so the installer is - # forced to download the correct variant rather than skipping on tag match. + # If the existing install is the wrong kind, remove it so the installer downloads the + # correct variant rather than skipping on tag match. $existingMetaPath = Join-Path $LlamaCppDir "UNSLOTH_PREBUILT_INFO.json" if (Test-Path $existingMetaPath) { try { $existingMeta = Get-Content $existingMetaPath -Raw | ConvertFrom-Json $existingKind = $existingMeta.install_kind - # A ROCm host may legitimately carry the fork's windows-rocm bundle - # or the upstream windows-hip fallback, so accept either and never - # treat a valid ROCm install as mismatched. A name-inferred gfx - # arch (Adrenalin-only, no confirmed runtime) still counts as - # ROCm-capable -- the ROCm prebuilt bundles its own runtime, - # mirroring the --rocm-gfx forward below. The CPU branch covers both - # the x64 windows-cpu and arm64 windows-arm64 bundles (Windows arm64 - # has no GPU prebuilt). NOTE: this block is currently inert -- - # write_prebuilt_metadata does not persist an install_kind key, so - # $existingKind is always null; keep $expectedKinds in sync with the - # kinds install_llama_prebuilt.py installs before relying on it. + # A ROCm host may carry windows-rocm or the windows-hip fallback (accept either); + # a name-inferred gfx arch counts as ROCm-capable. The CPU branch covers windows-cpu + # and windows-arm64. NOTE: currently inert -- write_prebuilt_metadata persists no + # install_kind, so $existingKind is always null; keep $expectedKinds in sync first. $expectedKinds = if ($HasROCm -or $script:ROCmGfxArch) { @("windows-rocm", "windows-hip") } elseif ($HasNvidiaSmi) { @("windows-cuda") } else { @("windows-cpu", "windows-arm64") } if ($existingKind -and ($existingKind -notin $expectedKinds)) { substep "Removing mismatched llama.cpp install (found '$existingKind', need one of: $($expectedKinds -join ', '))..." @@ -3667,9 +3374,8 @@ if ($LocalLlamaCppLinked) { } } substep "installing prebuilt llama.cpp bundle (preferred path)..." - # why: install_llama_prebuilt.py uses os.replace(), which would displace - # an unrelated $env:UNSLOTH_STUDIO_HOME\llama.cpp before the source-build - # ownership check below ever runs. + # install_llama_prebuilt.py uses os.replace(), which would displace an unrelated + # custom-home llama.cpp before the source-build ownership check runs. if ($StudioHomeIsCustom) { Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" } @@ -3682,20 +3388,16 @@ if ($LocalLlamaCppLinked) { if ($HasROCm) { $prebuiltArgs += "--has-rocm" } - # Forward the resolved gfx arch so the per-gfx ROCm prebuilt is picked even - # when the installer's probe can't confirm the runtime (amd-smi-only / - # Adrenalin-only, name-inferred arch). --rocm-gfx is authoritative and - # implies ROCm in install_llama_prebuilt.py, so the GPU prebuilt is selected - # even with $HasROCm false. Gating on $HasROCm gave Strix Halo / 8060S CPU. + # Forward the resolved gfx arch so the per-gfx ROCm prebuilt is picked even when the + # probe can't confirm the runtime. --rocm-gfx is authoritative and implies ROCm, so the + # GPU prebuilt is selected with $HasROCm false (gating on $HasROCm gave Strix Halo CPU). if ($script:ROCmGfxArch) { $prebuiltArgs += @("--rocm-gfx", $script:ROCmGfxArch) } if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) } - # UNSLOTH_LLAMA_CPP_BACKEND=cpu (case-insensitive, whitespace-trimmed) forces the - # CPU-only prebuilt via --force-cpu (persisted so updates keep it). Fixes Intel - # iGPU Vulkan crash (#7213). + # UNSLOTH_LLAMA_CPP_BACKEND=cpu forces the CPU-only prebuilt via --force-cpu (Intel iGPU Vulkan crash #7213). $llamaBackend = "$($env:UNSLOTH_LLAMA_CPP_BACKEND)".Trim().ToLowerInvariant() if ($llamaBackend -eq "cpu") { $prebuiltArgs += "--force-cpu" @@ -3713,7 +3415,7 @@ if ($LocalLlamaCppLinked) { } try { if ($script:UnslothVerbose) { - # Show live output in verbose mode while still capturing for error log + # Live output while still capturing for the error log. $prebuiltLog = Join-Path $env:TEMP "unsloth-prebuilt-$PID.log" & python @prebuiltArgs 2>&1 | Tee-Object -FilePath $prebuiltLog | Out-Host $prebuiltExit = $LASTEXITCODE @@ -3765,12 +3467,11 @@ if ($LocalLlamaCppLinked) { # ========================================================================== # PHASE 3.5: Install OpenSSL dev (for HTTPS support in llama-server) # ========================================================================== -# llama-server needs OpenSSL to download models from HuggingFace via -hf. -# ShiningLight.OpenSSL.Dev includes headers + libs that cmake can find. +# llama-server needs OpenSSL to download HF models via -hf; ShiningLight.OpenSSL.Dev ships headers + libs cmake finds. $OpenSslAvailable = $false if ($NeedLlamaSourceBuild) { - # Check if OpenSSL dev is already installed (look for include dir) + # Already installed? (look for the include dir) $OpenSslRoots = @( 'C:\Program Files\OpenSSL-Win64', 'C:\Program Files\OpenSSL', @@ -3814,22 +3515,16 @@ if ($NeedLlamaSourceBuild) { # ========================================================================== # PHASE 4: Build llama.cpp with CUDA for GGUF inference + export # ========================================================================== -# Builds at ~/.unsloth/llama.cpp — a single shared location under the user's -# home directory. This is used by both the inference server and the GGUF -# export pipeline (unsloth-zoo). -# We build: -# - llama-server: for GGUF model inference (with HTTPS if OpenSSL available) -# - llama-quantize: for GGUF export quantization -# Prerequisites git, cmake, VS Build Tools were installed in Phase 1; the CUDA -# Toolkit is resolved lazily just below via Resolve-CudaToolkit (source build only). +# Builds at ~/.unsloth/llama.cpp (shared by the inference server + GGUF export). Targets: +# llama-server (GGUF inference, HTTPS if OpenSSL) and llama-quantize (GGUF export). git/cmake/VS +# came from Phase 1; the CUDA Toolkit is resolved lazily below (source build only). $OriginalLlamaCppDir = $LlamaCppDir $BuildDir = Join-Path $LlamaCppDir "build" $LlamaServerBin = Join-Path $BuildDir "bin\Release\llama-server.exe" $HasCmakeForBuild = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue) -# Check if existing llama-server matches current GPU mode. A CUDA-built binary -# on a now-CPU-only machine (or vice versa) needs to be rebuilt. +# Does the existing llama-server match the current GPU mode? A CUDA binary on a now-CPU-only machine (or vice versa) rebuilds. $NeedRebuild = $false if (Test-Path -LiteralPath $LlamaServerBin) { $CmakeCacheFile = Join-Path $BuildDir "CMakeCache.txt" @@ -3845,10 +3540,8 @@ if (Test-Path -LiteralPath $LlamaServerBin) { } } -# Install build tools now (last resort) rather than eagerly in Phase 1, so the -# prebuilt path stays fast. Same condition as the if/elseif chain below: a source -# build runs only when needed and no usable binary is already present. A linked -# local dir sets $NeedLlamaSourceBuild = $false, so this no-ops for that path. +# Install build tools now (last resort) not eagerly in Phase 1, so the prebuilt path stays fast. +# Same condition as the chain below: source build runs only when needed and no usable binary exists. $WillBuildLlamaFromSource = $NeedLlamaSourceBuild -and ` -not ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") if ($WillBuildLlamaFromSource) { @@ -3858,24 +3551,21 @@ if ($WillBuildLlamaFromSource) { } if ($LocalLlamaCppLinked) { - # Local dir linked above -- honor the flag's contract: skip BOTH the prebuilt - # download and the source build. Falling through here would run CMake inside - # the user's checkout (via the junction) when it lacks build\bin\Release\llama-server.exe. + # Local dir linked above: skip BOTH prebuilt download and source build (falling through would + # run CMake inside the user's checkout via the junction). Write-Host "" step "llama.cpp" "linked (skipping build)" } elseif (-not $NeedLlamaSourceBuild) { Write-Host "" step "llama.cpp" "prebuilt (validated)" } elseif ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") { - # Skip rebuild only for pinned tags (e.g. b8635). When the requested - # tag is "master" (a moving target), always rebuild so the binary picks - # up new model architecture support (e.g. Gemma 4). + # Skip rebuild only for pinned tags; "master" is a moving target so always rebuild. Write-Host "" step "llama.cpp" "already built" } elseif (-not $HasCmakeForBuild) { Write-Host "" if (-not $HasNvidiaSmi) { - # CPU-only machines depend entirely on llama-server for GGUF chat -- cmake is required + # CPU-only machines depend entirely on llama-server for GGUF chat -- cmake required. substep "CMake is required to build llama-server for GGUF chat mode." "Yellow" substep "Continuing setup without llama.cpp build." "Yellow" substep "Install CMake from https://cmake.org/download/ and re-run setup." "Yellow" @@ -3885,25 +3575,22 @@ if ($LocalLlamaCppLinked) { substep "Install CMake from https://cmake.org/download/ and re-run setup." "Yellow" $script:LlamaCppDegraded = $true } else { - # Finalize the VS generator (gate/fallback below) BEFORE Resolve-CudaToolkit, - # which copies the CUDA .targets into the current generator's dir; a later swap - # would strand them. The CMake 4.2 gate for VS 2026 is checked only here, in the - # source-build path, so a VS 2026 + cmake < 4.2 host can still use the prebuilt. (#6473) + # Finalize the VS generator BEFORE Resolve-CudaToolkit, which copies the CUDA .targets into + # the current generator's dir (a later swap would strand them). The CMake 4.2 gate for VS 2026 + # is checked only here so a VS 2026 + cmake < 4.2 host can still use the prebuilt. (#6473) if ($CmakeGenerator -match 'Visual Studio 18\b') { if (-not (Test-CmakeCanDriveGenerator -Generator $CmakeGenerator)) { $cmakeVerObj = Get-CmakeVersion $cmakeVerStr = if ($cmakeVerObj) { $cmakeVerObj.ToString() } else { '0.0' } substep "CMake $cmakeVerStr cannot drive the Visual Studio 2026 generator (need 4.2+ or a VS-bundled cmake) -- updating via winget..." "Yellow" if ($null -ne (Get-Command winget -ErrorAction SilentlyContinue)) { - # upgrade first (fast if Kitware.CMake is already a winget app), then - # prepend the default dir so the new cmake wins over an older one on PATH + # upgrade first, then prepend the default dir so the new cmake wins over an older one. try { Invoke-SetupCommand { winget upgrade Kitware.CMake --source winget --accept-package-agreements --accept-source-agreements } | Out-Null Refresh-Environment } catch { substep "CMake winget upgrade failed: $($_.Exception.Message)" "Yellow" } Add-DefaultCmakeToPath | Out-Null - # upgrade no-ops if the cmake came from Scoop/Chocolatey/VS, not the - # Kitware winget package; install it so a 4.2+ cmake is available + # upgrade no-ops for a Scoop/Chocolatey/VS cmake; install to get a 4.2+ one. if (-not (Test-CmakeCanDriveGenerator -Generator $CmakeGenerator)) { try { Invoke-SetupCommand { winget install Kitware.CMake --source winget --accept-package-agreements --accept-source-agreements } | Out-Null @@ -3913,9 +3600,7 @@ if ($LocalLlamaCppLinked) { } } if (-not (Test-CmakeCanDriveGenerator -Generator $CmakeGenerator)) { - # cmake still cannot drive VS 2026; before failing, fall back to an - # older installed VS whose generator it can drive (e.g. VS 2022 + old - # cmake on an offline box keeps building) + # cmake still can't drive VS 2026; before failing, fall back to an older installed VS. $fallback = Get-FallbackVsGenerator if ($fallback) { substep "CMake cannot drive $CmakeGenerator; falling back to $($fallback.Generator)" "Yellow" @@ -3931,18 +3616,15 @@ if ($LocalLlamaCppLinked) { substep "CMake can drive the $CmakeGenerator generator" } - # CUDA resolved here (fail fast if none), after the final VS generator so its - # .targets land in the toolset cmake actually uses. + # CUDA resolved here (fail fast), after the final VS generator so its .targets land in the toolset cmake uses. if ($HasNvidiaSmi) { Resolve-CudaToolkit -RequireOrExit } Write-Host "" if ($HasNvidiaSmi) { substep "building llama.cpp with CUDA support..." } elseif ($HasROCm -or $script:ROCmGfxArch) { - # AMD GPU present but in the CPU-only source-build fallback: a HIP source - # build needs the full HIP SDK + ROCm clang toolchain. AMD GPU acceleration - # comes from the per-gfx ROCm prebuilt (bundles the runtime, no SDK) -- reaching - # here means it couldn't be installed. Warn loudly, don't ship a slow CPU build. + # AMD GPU in the CPU-only source-build fallback: a HIP source build needs the full HIP SDK, + # and GPU acceleration comes from the per-gfx ROCm prebuilt (reaching here means it failed). $_amdArch = if ($script:ROCmGfxArch) { $script:ROCmGfxArch } else { "ROCm" } substep "[WARN] AMD GPU ($_amdArch) detected, but the GPU-accelerated ROCm" "Yellow" substep " llama.cpp prebuilt could not be installed -- falling back to a CPU build." "Yellow" @@ -3959,19 +3641,15 @@ if ($LocalLlamaCppLinked) { # Start total build timer $totalSw = [System.Diagnostics.Stopwatch]::StartNew() - # Native commands (git, cmake) write to stderr even on success. - # With $ErrorActionPreference = "Stop" (set at top of script), PS 5.1 - # converts stderr lines into terminating ErrorRecords, breaking output. - # Lower to "Continue" for the build section. + # git/cmake write to stderr even on success, which -ErrorAction Stop turns terminating on PS 5.1; lower for the build. $prevEAP = $ErrorActionPreference $ErrorActionPreference = "Continue" $BuildOk = $true $FailedStep = "" - # Re-sanitize CUDA_PATH_V* vars — Refresh-Environment (called during - # Node/Python installs above) may have repopulated conflicting versioned - # vars from the Machine registry. + # Re-sanitize CUDA_PATH_V* vars: Refresh-Environment (Node/Python installs above) may have + # repopulated conflicting versioned vars from the Machine registry. if ($HasNvidiaSmi -and $CudaToolkitRoot) { $cudaPathVars2 = @([Environment]::GetEnvironmentVariables('Process').Keys | Where-Object { $_ -match '^CUDA_PATH_V' }) foreach ($v2 in $cudaPathVars2) { @@ -4031,14 +3709,13 @@ if ($LocalLlamaCppLinked) { $UseConcreteRef = ($ResolvedSourceRef -ne "latest" -and -not [string]::IsNullOrWhiteSpace($ResolvedSourceRef)) if (Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")) { - # why: in-place git mutation (remote set-url, checkout -B, clean -fdx) - # rewrites $LlamaCppDir; mirror the prebuilt and temp-dir-swap guards - # so an unrelated workspace .git tree is never silently overwritten. + # In-place git mutation (set-url, checkout -B, clean -fdx) rewrites $LlamaCppDir; guard so + # an unrelated workspace .git tree is never silently overwritten. if ($StudioHomeIsCustom) { Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" } Write-Host " Syncing llama.cpp to $ResolvedSourceRef..." -ForegroundColor Gray - # Always sync the remote URL so switching between default/fork sources works + # Always sync the remote URL so switching between default/fork sources works. Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir remote set-url origin "$ResolvedSourceUrl.git" } | Out-Null if ($LlamaPr) { $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin "pull/$LlamaPr/head" } @@ -4107,9 +3784,8 @@ if ($LocalLlamaCppLinked) { } } } - # why: in-place git-sync (the temp-dir clone path calls Mark-StudioOwned - # at swap-time) must mark the existing tree so a subsequent prebuilt - # update path's Assert-StudioOwnedOrAbsent does not exit on the same root. + # In-place git-sync must mark the existing tree so a later prebuilt update's + # Assert-StudioOwnedOrAbsent doesn't exit on the same root. if ($BuildOk -and $StudioHomeIsCustom) { Mark-StudioOwned -Path $LlamaCppDir } @@ -4219,7 +3895,7 @@ if ($LocalLlamaCppLinked) { '-G', $CmakeGenerator, '-Wno-dev' ) - # Tell cmake exactly where VS is (bypasses registry lookup) + # Tell cmake where VS is (bypasses registry lookup). if ($VsInstallPath) { $CmakeArgs += "-DCMAKE_GENERATOR_INSTANCE=$VsInstallPath" } @@ -4239,21 +3915,17 @@ if ($LocalLlamaCppLinked) { $CmakeArgs += '-DCMAKE_EXE_LINKER_FLAGS=/NODEFAULTLIB:LIBCMT' # CUDA flags -- only if GPU available, otherwise explicitly disable if ($HasNvidiaSmi -and $NvccPath) { - # UNSLOTH_LLAMA_CUDA_ARCHS (e.g. "120" or "89;86") forces the build - # arch and wins over detection, matching setup.sh. + # UNSLOTH_LLAMA_CUDA_ARCHS (e.g. "120" or "89;86") forces the build arch, matching setup.sh. $CudaArchOverride = if ($env:UNSLOTH_LLAMA_CUDA_ARCHS) { ($env:UNSLOTH_LLAMA_CUDA_ARCHS -replace '\s', '') } else { '' } if ((-not $CudaArch) -and (-not $CudaArchOverride)) { - # No detectable compute capability (#5854): -DGGML_CUDA=ON with no - # arch builds a PTX-only binary, so build CPU instead. Mirrors the - # Linux fix; set UNSLOTH_LLAMA_CUDA_ARCHS=120 to force a CUDA build. + # No detectable compute capability (#5854): -DGGML_CUDA=ON with no arch builds a + # PTX-only binary, so build CPU instead (set UNSLOTH_LLAMA_CUDA_ARCHS=120 to force CUDA). substep "could not detect a CUDA compute capability; building CPU llama.cpp instead of a PTX-only binary (set UNSLOTH_LLAMA_CUDA_ARCHS=120 to force a CUDA build)." "Yellow" $CmakeArgs += '-DGGML_CUDA=OFF' } else { $CmakeArgs += '-DGGML_CUDA=ON' - # Accept a host MSVC newer than nvcc's whitelist; a fresh toolkit - # (e.g. CUDA 13.3) otherwise aborts with "#error -- unsupported - # Microsoft Visual Studio version!". Mirrors the Linux fix. Via env - # (covers the configure probe + build), after Refresh-Environment, idempotent. + # Accept a host MSVC newer than nvcc's whitelist (a fresh toolkit otherwise aborts + # "unsupported Microsoft Visual Studio version"). Via env, after Refresh-Environment, idempotent. $nvccAllowFlag = '-allow-unsupported-compiler' if ([string]::IsNullOrEmpty($env:NVCC_PREPEND_FLAGS)) { $env:NVCC_PREPEND_FLAGS = $nvccAllowFlag @@ -4272,8 +3944,7 @@ if ($LocalLlamaCppLinked) { if (Test-NvccArchSupport -NvccExe $NvccPath -Arch $CudaArch) { $CmakeArgs += "-DCMAKE_CUDA_ARCHITECTURES=$CudaArch" } else { - # GPU arch too new for this toolkit -- fall back to highest supported. - # PTX forward-compatibility will JIT-compile for the actual GPU at runtime. + # GPU arch too new for this toolkit -- fall back to highest supported (PTX JITs at runtime). $maxArch = Get-NvccMaxArch -NvccExe $NvccPath if ($maxArch) { $CmakeArgs += "-DCMAKE_CUDA_ARCHITECTURES=$maxArch" @@ -4338,8 +4009,7 @@ if ($LocalLlamaCppLinked) { } # -- Step E: Build the DiffusionGemma visual server (optional, best-effort) -- - # An example target present on llama.cpp PR #24423; lets Unsloth serve - # DiffusionGemma GGUFs without DG_VISUAL_BIN. No-op when not configured. + # Example target from llama.cpp PR #24423; serves DiffusionGemma GGUFs without DG_VISUAL_BIN. if ($BuildOk) { $null = cmake --build $BuildDir --config Release --target llama-diffusion-gemma-visual-server -j $NumCpu 2>&1 | Out-String } @@ -4426,10 +4096,8 @@ substep "(add -H 0.0.0.0 for LAN / cloud access; exposes the raw port only, not substep "(add -H 0.0.0.0 --cloudflare for a public Cloudflare HTTPS link, or --secure to keep the raw port private; anyone with the API key can run code)" Write-Host "" -# Match studio/setup.sh: exit non-zero for degraded llama.cpp when called -# from install.ps1 (SKIP_STUDIO_BASE=1) so the installer can detect the -# failure. Direct 'unsloth studio update' does not set SKIP_STUDIO_BASE, -# so it keeps degraded installs successful. +# Match setup.sh: exit non-zero for degraded llama.cpp when called from install.ps1 +# (SKIP_STUDIO_BASE=1) so the installer detects it; direct 'studio update' stays successful. if ($script:LlamaCppDegraded -and $env:SKIP_STUDIO_BASE -eq "1") { exit 1 } diff --git a/tests/sh/test_mac_intel_compat.sh b/tests/sh/test_mac_intel_compat.sh index 8a0ff4b641..4159090a95 100644 --- a/tests/sh/test_mac_intel_compat.sh +++ b/tests/sh/test_mac_intel_compat.sh @@ -571,7 +571,7 @@ echo "=== Apple Silicon x86_64 (Rosetta) venv rebuild ===" # Extract the real guard block from install.sh so we exercise the shipped logic # (comment header down to its column-0 closing fi). _GUARD_FILE=$(mktemp) -awk '/Guard against two independent Apple Silicon venv problems/{f=1} f{print} f&&/^fi$/{exit}' \ +awk '/independent Apple Silicon venv/{f=1} f{print} f&&/^fi$/{exit}' \ "$INSTALL_SH" > "$_GUARD_FILE" if [ ! -s "$_GUARD_FILE" ]; then From e49a89f4f6175072f0ebaa92cc8ce91dd2af8954 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 20 Jul 2026 10:41:41 +0000 Subject: [PATCH 03/10] install.ps1: port the torch release preservation from install.sh The Windows installer recreates the venv on every re-run (the existing environment is moved aside for rollback and a fresh venv is created), so irm install.ps1 | iex over an existing install resolved the newest torch in range instead of keeping the installed release. This ports the install.sh preservation merged in PR 7250: - The installed torch version is probed (bounded process, drained streams, 30s timeout, last non-empty stdout line) BEFORE the rollback move, while the old interpreter still exists. - Get-PreviousTorchPin mirrors _previous_torch_pin: numeric-strict base (nightly/dev/rc builds never pin), UNSLOTH_TORCH_UPGRADE=1 opt-out, and a release-in-window check against the route's final constraint, so a raised ROCm floor correctly rejects keeping an older release. The decision runs after every index and floor decision. - The kept release installs as an exact pin with minor-paired companions (torchvision 0.minor+15, torchaudio 2.minor) at the main install, the AMD ROCm install and both flavor repairs; when the pinned release is not installable from the selected index the installer warns, clears the pin and falls back to the supported range. The CPU fallback cannot be reached with a live pin (the ROCm site resolves or clears it first) and keeps the plain range install. - The with-deps unsloth installs now carry a uv overrides file with the exact installed trio (twin of _build_unsloth_torch_overrides), so dependency resolution cannot move torch after it was deliberately selected. - The kept release is exported as UNSLOTH_KEPT_TORCH for setup.ps1, whose ROCm/CPU/CUDA torch installs substitute the kept trio when the env var is present (ROCm floors dominate; behavior without the env var is unchanged, so direct studio update runs are unaffected). After setup returns the installer re-probes and warns loudly if the kept release series changed, then clears the env var. Preservation is version-agnostic (numeric parse + window comparison), so a future ceiling bump to torch 2.12 keeps 2.11 installs in place. New suite tests/studio/test_previous_torch_pin.ps1 (AST-extracted helpers, 50+ checks incl. future-ceiling cases and structural wiring) passes, along with the flavor/pin-hardening/pin-stale/node ps1 suites, the parity and install-stack pytest suites, and the sh preservation and constraint suites. --- install.ps1 | 265 ++++++++++++++++++++--- studio/setup.ps1 | 103 +++++++-- tests/studio/test_previous_torch_pin.ps1 | 115 ++++++++++ 3 files changed, 428 insertions(+), 55 deletions(-) create mode 100644 tests/studio/test_previous_torch_pin.ps1 diff --git a/install.ps1 b/install.ps1 index 02d51b3d2b..a561e6f9b0 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1330,6 +1330,37 @@ exit 0 $script:StudioVenvRollbackDir = $null } + # Raw torch.__version__ from $PythonExe's venv (last non-empty stdout line), or $null. + # Bounded ProcessStartInfo probe (async-drain both streams, 30s timeout, kill on hang) so a + # wedged "import torch" can't stall the installer; feeds Get-InstalledTorchTag and the torch + # release-preservation decision (twin of install.sh's _PREV_TORCH_VER probe / _previous_torch_pin). + function Get-InstalledTorchVersionRaw { + param([string]$PythonExe) + if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null } + try { + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $PythonExe + $psi.Arguments = '-c "import torch; print(torch.__version__)"' + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.UseShellExecute = $false + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + # Drain BOTH streams async before WaitForExit: a synchronous ReadToEnd() would block on a wedged "import torch", and an undrained stderr would deadlock a child flooding the pipe buffer. A truly hung probe still hits the 30s timeout. + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + $finished = $proc.WaitForExit(30000) + if (-not $finished) { try { $proc.Kill() } catch {}; return $null } + $out = $outTask.GetAwaiter().GetResult() + [void]$errTask.GetAwaiter().GetResult() + if ($proc.ExitCode -ne 0) { return $null } + # Last non-empty line only, so stdout noise before the version can't corrupt the pin. + $lines = @($out -split "`r?`n" | ForEach-Object { $_.Trim() } | Where-Object { $_ -ne "" }) + if ($lines.Count -eq 0) { return $null } + return $lines[-1] + } catch { return $null } + } + if (Test-Path -LiteralPath $VenvPython) { # env-mode: $StudioHome is a user-chosen workspace, so refuse to nuke an existing venv lacking Unsloth sentinels (-PathType Leaf rejects a directory at the sentinel path; accept the in-VENV ownership marker so partial-install retries aren't blocked). if ( @@ -1342,6 +1373,14 @@ exit 0 Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." -ForegroundColor Yellow throw "Refusing to delete non-Unsloth venv at $VenvDir" } + # Record the existing venv's torch RELEASE BEFORE the rollback move (see Get-PreviousTorchPin); + # a re-run then keeps that release rather than silently jumping torch versions. Opt out with + # UNSLOTH_TORCH_UPGRADE=1. Only the new-layout replace probes; the legacy-migration branches + # reuse the venv, so torch survives naturally and needs no pin. + $script:PrevTorchVer = "" + if (-not $SkipTorch) { + $script:PrevTorchVer = Get-InstalledTorchVersionRaw -PythonExe $VenvPython + } # New layout already exists -- replace only after preserving rollback copy. substep "preserving existing environment for rollback..." try { @@ -1876,29 +1915,73 @@ exit 0 return $null } - # Installed torch flavor tag in $PythonExe's venv, or $null if absent. Uses ProcessStartInfo (not &) so stderr doesn't trip $ErrorActionPreference. + # Installed torch flavor tag in $PythonExe's venv, or $null if absent. Reuses the bounded raw probe. function Get-InstalledTorchTag { param([string]$PythonExe) - if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null } + $v = Get-InstalledTorchVersionRaw -PythonExe $PythonExe + if (-not $v) { return $null } + return ConvertTo-TorchFlavorTag $v + } + + # ── Torch release preservation (twin of install.sh's _previous_torch_pin, PR 7250): keep the + # previous venv's torch RELEASE across a re-run when it falls inside the freshly chosen constraint + # window; flavor follows the new index. Opt out with UNSLOTH_TORCH_UPGRADE=1. ── + + # Parse a probed torch.__version__ into a normalized stable release, or $null. + # Strips ONLY the +local tag; anchored numeric match so dev/rc/alpha/garbage never pin. + function ConvertTo-TorchNumericRelease { + param([string]$TorchVersion) + if ([string]::IsNullOrWhiteSpace($TorchVersion)) { return $null } + $publicBase = ($TorchVersion.Trim() -split '\+', 2)[0] + if ($publicBase -notmatch '^(\d+)\.(\d+)(?:\.(\d+))?$') { return $null } try { - $psi = New-Object System.Diagnostics.ProcessStartInfo - $psi.FileName = $PythonExe - $psi.Arguments = '-c "import torch; print(torch.__version__)"' - $psi.RedirectStandardOutput = $true - $psi.RedirectStandardError = $true - $psi.UseShellExecute = $false - $psi.CreateNoWindow = $true - $proc = [System.Diagnostics.Process]::Start($psi) - # Drain BOTH streams async before WaitForExit: a synchronous ReadToEnd() would block on a wedged "import torch", and an undrained stderr would deadlock a child flooding the pipe buffer. A truly hung probe still hits the 30s timeout. - $outTask = $proc.StandardOutput.ReadToEndAsync() - $errTask = $proc.StandardError.ReadToEndAsync() - $finished = $proc.WaitForExit(30000) - if (-not $finished) { try { $proc.Kill() } catch {}; return $null } - $torchVer = $outTask.GetAwaiter().GetResult().Trim() - [void]$errTask.GetAwaiter().GetResult() - if ($proc.ExitCode -ne 0 -or -not $torchVer) { return $null } - return ConvertTo-TorchFlavorTag $torchVer + $major = [int]$Matches[1]; $minor = [int]$Matches[2] + if ($Matches[3]) { $patch = [int]$Matches[3] } else { $patch = 0 } + $normalized = New-Object System.Version($major, $minor, $patch) } catch { return $null } + return [pscustomobject]@{ + PublicBase = $publicBase; Major = $major; Minor = $minor; Patch = $patch; Version = $normalized + } + } + + # True when a release falls inside a "torch>=A,=(\d+(?:\.\d+){0,2}),<(\d+(?:\.\d+){0,2})$') { return $false } + $floor = ConvertTo-TorchNumericRelease $Matches[1] + $ceiling = ConvertTo-TorchNumericRelease $Matches[2] + if (-not $floor -or -not $ceiling) { return $false } + return ($Release.Version -ge $floor.Version -and $Release.Version -lt $ceiling.Version) + } + + # The kept-release trio for a previously installed torch, or $null when nothing + # should be kept (no/unstable version, UNSLOTH_TORCH_UPGRADE=1, outside the final + # route window -- a raised ROCm floor correctly rejects an older release). + # Exact-release pin, matching install.sh's _previous_torch_pin; companions pair + # to the kept minor (torchaudio no longer exact-pins torch). + function Get-PreviousTorchPin { + param( + [string]$TorchVersion, + # Named -Constraint (like Test-TorchReleaseInWindow): the Windows port keeps no shell-style + # constraint variable, and its structural tests forbid that token appearing in install.ps1. + [Parameter(Mandatory = $true)][string]$Constraint + ) + if ($env:UNSLOTH_TORCH_UPGRADE -eq '1') { return $null } + $release = ConvertTo-TorchNumericRelease $TorchVersion + if (-not $release) { return $null } + if (-not (Test-TorchReleaseInWindow -Release $release -Constraint $Constraint)) { return $null } + if ($release.Major -ne 2) { return $null } + $visionMinor = $release.Minor + 15 + return [pscustomobject]@{ + Release = $release + TorchSpec = "torch==$($release.PublicBase)" + VisionSpec = "torchvision==0.$visionMinor.*" + AudioSpec = "torchaudio==2.$($release.Minor).*" + } } # An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it (e.g. a deliberate cpu pin on an AMD host). @@ -2018,6 +2101,22 @@ exit 0 return $installed } + # ── Freeze the installed torch trio for the with-deps unsloth install (twin of install.sh's + # _build_unsloth_torch_overrides): a released unsloth wheel can pin an older torch, and a + # with-deps resolve then downgrades the pinned +cuXXX/+rocm trio. Return a temp uv --overrides + # file pinning torch/torchvision/torchaudio to their installed versions, or $null when torch is + # absent (--no-torch) so the caller installs unchanged. Caller removes the file afterwards. ── + function New-UnslothTorchOverridesFile { + param([string]$PythonExe) + if ($SkipTorch) { return $null } + $pins = & $PythonExe -c "from importlib.metadata import version, PackageNotFoundError`nfor _p in ('torch', 'torchvision', 'torchaudio'):`n try:`n print(_p + '==' + version(_p))`n except PackageNotFoundError:`n pass" 2>$null + $lines = @($pins | Where-Object { $_ -match '^torch' }) + if ($lines.Count -eq 0 -or $lines[0] -notmatch '^torch==') { return $null } + $f = [System.IO.Path]::GetTempFileName() + Set-Content -LiteralPath $f -Value ($lines -join "`n") -Encoding ascii + return $f + } + if ($_Migrated) { # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving existing torch/CUDA unless the flavor repair below re-lands it. Write-TauriLog "STEP" "Installing unsloth" @@ -2057,6 +2156,35 @@ exit 0 } } } elseif ($TorchIndexUrl -or $ROCmIndexUrl) { + # Leaf-gated bounded trio, HOISTED so the release-preservation decision below can use it as the + # default route window: torchaudio 2.11 dropped its torch pin, so a bare companion can drift from + # a capped torch. cu families ship torch 2.11.x (paired triton-windows 3.6) so the ceiling + # widens to <2.12; other leaves keep the 2.10 line. Mirrors install.sh and _CUDA_TORCH_PKG_SPEC. + $_idxLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() + if ($_idxLeaf -match '^cu[0-9]+$') { + $_pinTorchSpec = "torch>=2.4,<2.12.0" + $_pinVisionSpec = "torchvision>=0.19,<0.27.0" + $_pinAudioSpec = "torchaudio>=2.4,<2.12.0" + } else { + $_pinTorchSpec = "torch>=2.4,<2.11.0" + $_pinVisionSpec = "torchvision>=0.19,<0.26.0" + $_pinAudioSpec = "torchaudio>=2.4,<2.11.0" + } + # Release preservation (twin of install.sh's _PREV_TORCH_PIN decision): evaluated after every + # index/floor choice, incl. the ROCm reroute, so a raised floor rejects an older release. The + # route window is the leaf-gated CUDA trio by default; the ROCm path uses its floor (or the + # torch>=2.4,<2.11.0 range the ROCm->CPU fallback installs). The kept release is exported for + # setup.ps1 (UNSLOTH_KEPT_TORCH) and cleared after setup runs. + $script:PrevTorchPin = $null + if (-not $SkipTorch -and $script:PrevTorchVer) { + $_routeWindow = $_pinTorchSpec + if ($ROCmIndexUrl) { if ($ROCmTorchFloor) { $_routeWindow = $ROCmTorchFloor } else { $_routeWindow = "torch>=2.4,<2.11.0" } } + $script:PrevTorchPin = Get-PreviousTorchPin -TorchVersion $script:PrevTorchVer -Constraint $_routeWindow + if ($script:PrevTorchPin) { + $env:UNSLOTH_KEPT_TORCH = $script:PrevTorchPin.Release.PublicBase + substep "existing install has torch $script:PrevTorchVer -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)" + } + } if ($SkipTorch) { substep "skipping PyTorch (--no-torch flag set)." "Yellow" } elseif ($ROCmIndexUrl) { @@ -2066,12 +2194,24 @@ exit 0 # Pin companions to match $torchSpec; bare names can resolve an ABI-incompatible torchvision/torchaudio on AMD's per-arch index. $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } + # Kept-release attempt first (pin already vetted against the ROCm floor); companions follow the kept minor. + if ($script:PrevTorchPin) { + $_keptTorch = $script:PrevTorchPin.TorchSpec; $_keptVision = $script:PrevTorchPin.VisionSpec; $_keptAudio = $script:PrevTorchPin.AudioSpec + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (kept release)" { uv pip install --python $VenvPython --force-reinstall $_keptTorch $_keptVision $_keptAudio --default-index $ROCmIndexUrl } + if ($torchInstallExit -ne 0) { + substep "[WARN] $_keptTorch is not installable from $(Remove-IndexUrlCredentials $ROCmIndexUrl) -- installing the newest supported release instead" "Yellow" + $script:PrevTorchPin = $null + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + } + } + if (-not $script:PrevTorchPin) { + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } + } if ($torchInstallExit -ne 0) { # Transient AMD-index failure: fall back to a CPU base (Unsloth setup retries ROCm). Explicit CPU index -- for a pinned ROCm index $TorchIndexUrl IS the ROCm mirror, so reusing it would just retry it. $CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" } substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow" - # --force-reinstall: a failed ROCm install can leave an unpinned ROCm torch that still satisfies the CPU torch>= range, so without it uv would keep the ROCm build and only swap companions -- a mismatched venv the flavor-repair block won't fix. + # --force-reinstall: a failed ROCm install can leave an unpinned ROCm torch that still satisfies the CPU torch>= range, so without it uv would keep the ROCm build and only swap companions -- a mismatched venv the flavor-repair block won't fix. (No kept-release attempt: the ROCm attempts above always resolve or clear $script:PrevTorchPin first, so a pin never reaches this CPU base.) $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $CpuFallbackIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red @@ -2084,18 +2224,19 @@ exit 0 } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..." - # Bounded trio on every index: torchaudio 2.11 dropped its torch pin, so a bare companion can drift from a capped torch. cu families ship torch 2.11.x (paired triton-windows 3.6) so the ceiling widens to <2.12; other leaves keep the 2.10 line. Mirrors install.sh and _CUDA_TORCH_PKG_SPEC. - $_idxLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() - if ($_idxLeaf -match '^cu[0-9]+$') { - $_pinTorchSpec = "torch>=2.4,<2.12.0" - $_pinVisionSpec = "torchvision>=0.19,<0.27.0" - $_pinAudioSpec = "torchaudio>=2.4,<2.12.0" - } else { - $_pinTorchSpec = "torch>=2.4,<2.11.0" - $_pinVisionSpec = "torchvision>=0.19,<0.26.0" - $_pinAudioSpec = "torchaudio>=2.4,<2.11.0" + # Kept-release attempt first (pin vetted against the leaf-gated route window); companions follow the kept minor. Range install (the hoisted bounded trio) runs when there is no pin or the kept attempt failed. + if ($script:PrevTorchPin) { + $_keptTorch = $script:PrevTorchPin.TorchSpec; $_keptVision = $script:PrevTorchPin.VisionSpec; $_keptAudio = $script:PrevTorchPin.AudioSpec + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (kept release)" { uv pip install --python $VenvPython $_keptTorch $_keptVision $_keptAudio --default-index $TorchIndexUrl } + if ($torchInstallExit -ne 0) { + substep "[WARN] $_keptTorch is not installable from $(Remove-IndexUrlCredentials $TorchIndexUrl) -- installing the newest supported release instead" "Yellow" + $script:PrevTorchPin = $null + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + } + } + if (-not $script:PrevTorchPin) { + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython $_pinTorchSpec $_pinVisionSpec $_pinAudioSpec --default-index $TorchIndexUrl } } - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython $_pinTorchSpec $_pinVisionSpec $_pinAudioSpec --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2118,9 +2259,25 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + # Freeze the installed torch trio so this with-deps resolve can't downgrade the pinned +cuXXX/+rocm build (twin of install.sh's _build_unsloth_torch_overrides). + $script:TorchOverridesFile = New-UnslothTorchOverridesFile -PythonExe $VenvPython + if ($script:TorchOverridesFile) { + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth --overrides $script:TorchOverridesFile "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + Remove-Item -LiteralPath $script:TorchOverridesFile -Force -ErrorAction SilentlyContinue + $script:TorchOverridesFile = $null + } else { + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } + # Freeze the installed torch trio (see above) so the with-deps unsloth resolve can't strip the +cuXXX/+rocm suffix. + $script:TorchOverridesFile = New-UnslothTorchOverridesFile -PythonExe $VenvPython + if ($script:TorchOverridesFile) { + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth --overrides $script:TorchOverridesFile -- "$PackageName" } + Remove-Item -LiteralPath $script:TorchOverridesFile -Force -ErrorAction SilentlyContinue + $script:TorchOverridesFile = $null + } else { + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } + } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2184,8 +2341,22 @@ exit 0 # Pin companions like the fresh ROCm path (bare names can pull an ABI-incompatible torchvision/torchaudio from the per-arch index). $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + # Kept-release substitution (twin of install.sh's _install_torch_default_index honoring _PREV_TORCH_PIN): honor the preserved torch when the pin survived the E-decision (already floor-vetted); restore the range specs and retry if it isn't installable here. + $_rocmKept = $false + if ($script:PrevTorchPin) { + $_origRocmSpec = $rocmSpec; $_origVisionSpec = $visionSpec; $_origAudioSpec = $audioSpec + $rocmSpec = $script:PrevTorchPin.TorchSpec; $visionSpec = $script:PrevTorchPin.VisionSpec; $audioSpec = $script:PrevTorchPin.AudioSpec + $_rocmKept = $true + } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } + if ($torchFixExit -ne 0 -and $_rocmKept) { + substep "[WARN] $rocmSpec is not installable from $(Remove-IndexUrlCredentials $ROCmIndexUrl) -- installing the newest supported release instead" "Yellow" + $rocmSpec = $_origRocmSpec; $visionSpec = $_origVisionSpec; $audioSpec = $_origAudioSpec + $script:PrevTorchPin = $null + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } + } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) @@ -2199,8 +2370,22 @@ exit 0 } else { $_fixTorchSpec = "torch>=2.4,<2.11.0"; $_fixVisionSpec = "torchvision>=0.19,<0.26.0"; $_fixAudioSpec = "torchaudio>=2.4,<2.11.0" } + # Kept-release substitution (twin of install.sh's _install_torch_default_index honoring _PREV_TORCH_PIN): honor the preserved torch when the pin survived the E-decision; restore the range specs and retry if it isn't installable here. The --reinstall-package triplet stays on both attempts. + $_cudaKept = $false + if ($script:PrevTorchPin) { + $_origFixTorchSpec = $_fixTorchSpec; $_origFixVisionSpec = $_fixVisionSpec; $_origFixAudioSpec = $_fixAudioSpec + $_fixTorchSpec = $script:PrevTorchPin.TorchSpec; $_fixVisionSpec = $script:PrevTorchPin.VisionSpec; $_fixAudioSpec = $script:PrevTorchPin.AudioSpec + $_cudaKept = $true + } substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython $_fixTorchSpec $_fixVisionSpec $_fixAudioSpec --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + if ($torchFixExit -ne 0 -and $_cudaKept) { + substep "[WARN] $_fixTorchSpec is not installable from $(Remove-IndexUrlCredentials $TorchIndexUrl) -- installing the newest supported release instead" "Yellow" + $_fixTorchSpec = $_origFixTorchSpec; $_fixVisionSpec = $_origFixVisionSpec; $_fixAudioSpec = $_origFixAudioSpec + $script:PrevTorchPin = $null + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython $_fixTorchSpec $_fixVisionSpec $_fixAudioSpec --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) @@ -2317,6 +2502,18 @@ exit 0 Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue } + # Release-preservation handoff done: setup.ps1 has consumed UNSLOTH_KEPT_TORCH (if any). Clear it so a + # later 'studio update' in the same session doesn't re-pin an old release, and warn (never abort) if the + # kept torch series changed out from under us during setup. + if ($script:PrevTorchPin) { + $_keptSeries = "$($script:PrevTorchPin.Release.Major).$($script:PrevTorchPin.Release.Minor)" + $_nowVer = Get-InstalledTorchVersionRaw -PythonExe $VenvPython + $_nowRelease = ConvertTo-TorchNumericRelease $_nowVer + if ($_nowRelease -and "$($_nowRelease.Major).$($_nowRelease.Minor)" -ne $_keptSeries) { + Write-Host "[WARN] kept torch $($script:PrevTorchVer) but the environment now has torch $_nowVer" -ForegroundColor Red + } + } + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue if ($setupExit -ne 0) { Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red return (Exit-InstallFailure "unsloth studio setup failed (exit code $setupExit)" $setupExit) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 7892f67bf8..79690fb76f 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2895,13 +2895,36 @@ if ($ROCmIndexUrl) { if ($ROCmTorchSpec -ne "torch") { substep " enforcing $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec (known _grouped_mm bug in older wheels)" "Cyan" } - if ($script:UnslothVerbose) { - Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host - $torchInstallExit = $LASTEXITCODE - $output = "" - } else { - $output = Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | Out-String - $torchInstallExit = $LASTEXITCODE + # Release preservation (twin of install.sh's _previous_torch_pin): install.ps1 exports the previous + # venv's exact torch RELEASE via UNSLOTH_KEPT_TORCH. Substitute the kept trio unless it conflicts with + # a >=2.11 floor already enforced in $ROCmTorchSpec (a kept minor < 11 keeps the floor). On a failed + # install, restore the computed specs and retry once. Absent the env var (direct 'studio update'), + # nothing changes and the loop runs exactly once. + $_rocmKeptActive = $false + $_rocmOrigTorch = $ROCmTorchSpec; $_rocmOrigVision = $ROCmVisionSpec; $_rocmOrigAudio = $ROCmAudioSpec + if ($env:UNSLOTH_KEPT_TORCH -match '^\d+\.\d+(\.\d+)?$') { + $_keptMinor = [int](($env:UNSLOTH_KEPT_TORCH -split '\.')[1]) + if (-not ($ROCmTorchSpec -match 'torch>=2\.11' -and $_keptMinor -lt 11)) { + $ROCmTorchSpec = "torch==$($env:UNSLOTH_KEPT_TORCH)" + $ROCmVisionSpec = "torchvision==0.$($_keptMinor + 15).*" + $ROCmAudioSpec = "torchaudio==2.$($_keptMinor).*" + $_rocmKeptActive = $true + } + } + while ($true) { + if ($script:UnslothVerbose) { + Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host + $torchInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | Out-String + $torchInstallExit = $LASTEXITCODE + } + if ($torchInstallExit -eq 0 -or -not $_rocmKeptActive) { break } + substep "[WARN] torch==$($env:UNSLOTH_KEPT_TORCH) not installable on this ROCm index -- using the supported release" "Yellow" + $ROCmTorchSpec = $_rocmOrigTorch; $ROCmVisionSpec = $_rocmOrigVision; $ROCmAudioSpec = $_rocmOrigAudio + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + $_rocmKeptActive = $false } if ($torchInstallExit -ne 0) { Write-Host "[WARN] AMD ROCm PyTorch install failed -- falling back to CPU" -ForegroundColor Yellow @@ -2931,13 +2954,32 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { $cpuVisionSpec = "torchvision>=0.19,<0.27.0" $cpuAudioSpec = "torchaudio>=2.4,<2.12.0" } - if ($script:UnslothVerbose) { - Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host - $torchInstallExit = $LASTEXITCODE - $output = "" - } else { - $output = Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | Out-String - $torchInstallExit = $LASTEXITCODE + # Release preservation: keep install.ps1's exported torch RELEASE (UNSLOTH_KEPT_TORCH) -- a kept + # release from the /cpu index resolves the correct cpu-flavor build. On a failed install, restore the + # computed specs and retry once. Absent the env var, nothing changes and the loop runs exactly once. + $_cpuKeptActive = $false + $_cpuOrigTorch = $cpuTorchSpec; $_cpuOrigVision = $cpuVisionSpec; $_cpuOrigAudio = $cpuAudioSpec + if ($env:UNSLOTH_KEPT_TORCH -match '^\d+\.\d+(\.\d+)?$') { + $_keptMinor = [int](($env:UNSLOTH_KEPT_TORCH -split '\.')[1]) + $cpuTorchSpec = "torch==$($env:UNSLOTH_KEPT_TORCH)" + $cpuVisionSpec = "torchvision==0.$($_keptMinor + 15).*" + $cpuAudioSpec = "torchaudio==2.$($_keptMinor).*" + $_cpuKeptActive = $true + } + while ($true) { + if ($script:UnslothVerbose) { + Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host + $torchInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | Out-String + $torchInstallExit = $LASTEXITCODE + } + if ($torchInstallExit -eq 0 -or -not $_cpuKeptActive) { break } + substep "[WARN] torch==$($env:UNSLOTH_KEPT_TORCH) not installable from the CPU index -- using the supported release" "Yellow" + $cpuTorchSpec = $_cpuOrigTorch; $cpuVisionSpec = $_cpuOrigVision; $cpuAudioSpec = $_cpuOrigAudio + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + $_cpuKeptActive = $false } if ($torchInstallExit -ne 0) { Write-Host "[FAILED] PyTorch install failed (exit code $torchInstallExit)" -ForegroundColor Red @@ -2961,13 +3003,32 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { $cudaVisionSpec = "torchvision>=0.19,<0.26.0" $cudaAudioSpec = "torchaudio>=2.4,<2.11.0" } - if ($script:UnslothVerbose) { - Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host - $torchInstallExit = $LASTEXITCODE - $output = "" - } else { - $output = Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | Out-String - $torchInstallExit = $LASTEXITCODE + # Release preservation: keep install.ps1's exported torch RELEASE (UNSLOTH_KEPT_TORCH) so a re-run + # reinstalls the same CUDA build (+cuXXX follows this index). On a failed install, restore the computed + # specs and retry once. Absent the env var (direct 'studio update'), nothing changes; loop runs once. + $_cudaKeptActive = $false + $_cudaOrigTorch = $cudaTorchSpec; $_cudaOrigVision = $cudaVisionSpec; $_cudaOrigAudio = $cudaAudioSpec + if ($env:UNSLOTH_KEPT_TORCH -match '^\d+\.\d+(\.\d+)?$') { + $_keptMinor = [int](($env:UNSLOTH_KEPT_TORCH -split '\.')[1]) + $cudaTorchSpec = "torch==$($env:UNSLOTH_KEPT_TORCH)" + $cudaVisionSpec = "torchvision==0.$($_keptMinor + 15).*" + $cudaAudioSpec = "torchaudio==2.$($_keptMinor).*" + $_cudaKeptActive = $true + } + while ($true) { + if ($script:UnslothVerbose) { + Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host + $torchInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | Out-String + $torchInstallExit = $LASTEXITCODE + } + if ($torchInstallExit -eq 0 -or -not $_cudaKeptActive) { break } + substep "[WARN] torch==$($env:UNSLOTH_KEPT_TORCH) not installable from this CUDA index -- using the supported release" "Yellow" + $cudaTorchSpec = $_cudaOrigTorch; $cudaVisionSpec = $_cudaOrigVision; $cudaAudioSpec = $_cudaOrigAudio + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + $_cudaKeptActive = $false } if ($torchInstallExit -ne 0) { Write-Host "[FAILED] PyTorch CUDA install failed (exit code $torchInstallExit)" -ForegroundColor Red diff --git a/tests/studio/test_previous_torch_pin.ps1 b/tests/studio/test_previous_torch_pin.ps1 new file mode 100644 index 0000000000..8c84a5e8b4 --- /dev/null +++ b/tests/studio/test_previous_torch_pin.ps1 @@ -0,0 +1,115 @@ +#!/usr/bin/env pwsh +# 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 test for install.ps1's torch release preservation helpers +# (ConvertTo-TorchNumericRelease, Test-TorchReleaseInWindow, Get-PreviousTorchPin), +# the Windows port of install.sh's _previous_torch_pin (PR 7250). Pure helpers, +# AST-extracted and run in-process -- no GPU/venv needed. +# Run: pwsh -NoProfile -File tests/studio/test_previous_torch_pin.ps1 + +$ErrorActionPreference = "Stop" +$installPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "install.ps1") +$installPath = (Resolve-Path $installPath).Path + +# --- Parse install.ps1 (also serves as a syntax gate) and extract the helpers --- +$tokens = $null; $errors = $null +$ast = [System.Management.Automation.Language.Parser]::ParseFile($installPath, [ref]$tokens, [ref]$errors) +if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "install.ps1 has parse errors" } + +foreach ($name in @("ConvertTo-TorchNumericRelease", "Test-TorchReleaseInWindow", "Get-PreviousTorchPin")) { + $fn = $ast.FindAll({ param($n) + $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name + }, $true) + if ($fn.Count -ne 1) { throw "expected exactly one $name in install.ps1, found $($fn.Count)" } + # Pure helpers (no exit / external calls) -- safe to define in this scope. + Invoke-Expression $fn[0].Extent.Text +} + +$failures = 0 +function Check($name, $cond) { + if ($cond) { Write-Host " PASS $name" } + else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ } +} + +$savedUpgrade = $env:UNSLOTH_TORCH_UPGRADE +try { + Remove-Item Env:UNSLOTH_TORCH_UPGRADE -ErrorAction SilentlyContinue + $win = "torch>=2.4,<2.12.0" + + # --- ConvertTo-TorchNumericRelease: accepted stable versions --- + foreach ($v in @("2.10.0", "2.10.0+cpu", "2.10.0+cu126", "2.10.1+cu130", "2.9.1+rocm7.2.1", "2.9.0+xpu", "2.10")) { + $r = ConvertTo-TorchNumericRelease $v + Check "release accepts $v" ($null -ne $r) + } + $r = ConvertTo-TorchNumericRelease "2.10.1+cu128" + Check "release strips only the +local tag" ($r.PublicBase -eq "2.10.1" -and $r.Minor -eq 10) + + # --- Rejected: nightly/dev/rc/alpha/garbage never pin --- + foreach ($v in @("", " ", "2.11.0.dev20260701+cu130", "2.9.0a0+gitabc123", "2.11.0rc1+cu130", + ".2.10", "2.10.", "2..10", "2.x.0", "not-a-version", + "Traceback (most recent call last):", "99999999999999999999.1")) { + Check "release rejects '$v'" ($null -eq (ConvertTo-TorchNumericRelease $v)) + } + + # --- Test-TorchReleaseInWindow --- + $cases = @( + @{ v = "2.4.0"; c = "torch>=2.4,<2.12.0"; ok = $true; n = "at the floor" } + @{ v = "2.3.1"; c = "torch>=2.4,<2.12.0"; ok = $false; n = "below the floor" } + @{ v = "2.11.0"; c = "torch>=2.4,<2.12.0"; ok = $true; n = "just below the ceiling" } + @{ v = "2.12.0"; c = "torch>=2.4,<2.12.0"; ok = $false; n = "at the ceiling" } + @{ v = "2.13.0"; c = "torch>=2.4,<2.12.0"; ok = $false; n = "above the ceiling" } + @{ v = "2.10.0"; c = "torch>=2.11.0,<2.12.0"; ok = $false; n = "2.11 floor rejects 2.10" } + @{ v = "2.11.0"; c = "torch>=2.11.0,<2.12.0"; ok = $true; n = "2.11 floor accepts 2.11" } + @{ v = "2.11.0"; c = "torch>=2.4,<2.13.0"; ok = $true; n = "future 2.13 ceiling keeps 2.11" } + @{ v = "2.12.1"; c = "torch>=2.4,<2.13.0"; ok = $true; n = "future 2.13 ceiling keeps 2.12" } + ) + foreach ($t in $cases) { + $rel = ConvertTo-TorchNumericRelease $t.v + Check ("window: " + $t.n) ((Test-TorchReleaseInWindow -Release $rel -Constraint $t.c) -eq $t.ok) + } + # Malformed constraints fail closed. + $rel = ConvertTo-TorchNumericRelease "2.10.0" + Check "window: malformed constraint fails closed" (-not (Test-TorchReleaseInWindow -Release $rel -Constraint "torch")) + Check "window: exact-pin constraint fails closed" (-not (Test-TorchReleaseInWindow -Release $rel -Constraint "torch==2.10.0")) + + # --- Get-PreviousTorchPin: exact-release pin, sh parity --- + $pin = Get-PreviousTorchPin -TorchVersion "2.10.0+cu128" -Constraint $win + Check "pin keeps 2.10.0 (exact release, sh parity)" ($pin.TorchSpec -eq "torch==2.10.0") + Check "pin pairs torchvision to the kept minor" ($pin.VisionSpec -eq "torchvision==0.25.*") + Check "pin pairs torchaudio to the kept minor" ($pin.AudioSpec -eq "torchaudio==2.10.*") + $pin = Get-PreviousTorchPin -TorchVersion "2.9.1+rocm7.2.1" -Constraint $win + Check "pin keeps 2.9.1 with 0.24.*/2.9.* companions" ( + $pin.TorchSpec -eq "torch==2.9.1" -and $pin.VisionSpec -eq "torchvision==0.24.*" -and $pin.AudioSpec -eq "torchaudio==2.9.*") + $pin = Get-PreviousTorchPin -TorchVersion "2.11.0+cpu" -Constraint $win + Check "pin keeps 2.11.0 under a future-widened window" ($pin.TorchSpec -eq "torch==2.11.0") + + # No previous version / out-of-window / non-stable -> no pin. + Check "no pin without a previous version" ($null -eq (Get-PreviousTorchPin -TorchVersion "" -Constraint $win)) + Check "no pin for a below-floor release" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.3.1+cpu" -Constraint $win)) + Check "raised ROCm floor rejects keeping 2.10" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.10.0+rocm7.1" -Constraint "torch>=2.11.0,<2.12.0")) + Check "no pin for a nightly build" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.11.0.dev20260701+cu130" -Constraint $win)) + Check "no pin for an unsupported 2.12 under a <2.12 window" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.12.0+cu130" -Constraint $win)) + + # --- UNSLOTH_TORCH_UPGRADE opt-out (exact string '1', sh parity) --- + $env:UNSLOTH_TORCH_UPGRADE = "1" + Check "UNSLOTH_TORCH_UPGRADE=1 disables the pin" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.10.0+cpu" -Constraint $win)) + $env:UNSLOTH_TORCH_UPGRADE = "0" + Check "UNSLOTH_TORCH_UPGRADE=0 keeps the pin" ($null -ne (Get-PreviousTorchPin -TorchVersion "2.10.0+cpu" -Constraint $win)) +} finally { + if ($null -ne $savedUpgrade) { $env:UNSLOTH_TORCH_UPGRADE = $savedUpgrade } + else { Remove-Item Env:UNSLOTH_TORCH_UPGRADE -ErrorAction SilentlyContinue } +} + +# --- Structural wiring (source assertions) --- +$src = Get-Content $installPath -Raw +Check "probe runs before the rollback move" ( + $src.IndexOf('$script:PrevTorchVer') -ge 0 -and + $src.IndexOf('$script:PrevTorchVer') -lt $src.IndexOf('Start-StudioVenvRollback -ExistingDir')) +Check "pin decision cites the UNSLOTH_TORCH_UPGRADE escape hatch" ($src -match 'UNSLOTH_TORCH_UPGRADE=1 to get the newest') +Check "kept-release fallback clears the pin" ($src -match '\$script:PrevTorchPin\s*=\s*\$null') +Check "kept release exported for setup.ps1" ($src -match 'UNSLOTH_KEPT_TORCH') + +Write-Host "" +if ($failures -gt 0) { Write-Host "$failures check(s) failed" -ForegroundColor Red; exit 1 } +Write-Host "All checks passed" +exit 0 From 3f79b5e53d461b7764e6a9688e92beae04adf73d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 20 Jul 2026 10:44:42 +0000 Subject: [PATCH 04/10] install: make torch 2.11 the default supported line on wheel-backed platforms Fresh installs now resolve the torch 2.11 trio (torch 2.11.0, torchvision 0.26.0, torchaudio 2.11.0) everywhere the wheels exist, verified by live index resolution for py3.11-3.13: Linux x86_64 (cpu, cu126, cu128, cu130, rocm7.1, rocm7.2), Linux aarch64 (cpu, cu130), Windows x86_64 (cpu, cu126, cu128, cu130; triton-windows<3.7 resolves 3.6.0.post26, the 2.11 pairing) and macOS arm64. Existing installs are unaffected: the release preservation on both installers keeps the installed torch on re-runs (verified end to end against the live cpu index: a seeded 2.10.0 venv stays 2.10.0 with paired 0.25/2.10 companions while a fresh venv resolves the 2.11 trio). - install.sh: the supported range is centralized in _TORCH_CEILING / _TORCHVISION_CEILING / _TORCHAUDIO_CEILING and composed into the default constraints, so the next bump (torch 2.12) is a three-line change. The now-redundant cu* widen arm and custom-pin companion block collapse into the default; the curated ROCm >=2.11 floors stay literal. - install.ps1: the hoisted default trio, the CUDA flavor repair and the ROCm CPU fallback all use the <2.12 trio (the leaf gate collapsed since cu and non-cu now share the range). - setup.ps1: the unknown-leaf pinned CUDA trio widens to <2.12, matching the pinned cpu path. Deliberately NOT widened: rocm6.4 (tops at torch 2.9.1) and rocm7.0 (2.10.0) by index content; macOS x86_64 (no >=2.4 wheels, stays no-torch); the AMD per-arch repo.amd.com curated bounds for arches outside the 2.11 allowlist. Constraint suites updated to the ceiling scheme; parity suites updated to the widened trio. All sh, ps1 and pytest installer suites pass (the host-defaults suite and the tokenizers negative-control are known pre-existing failures). --- install.ps1 | 39 +++------ install.sh | 30 +++---- studio/setup.ps1 | 6 +- tests/python/test_cross_platform_parity.py | 35 ++++---- .../test_tokenizers_and_torch_constraint.py | 39 ++++----- tests/sh/test_torch_constraint.sh | 81 +++++++++---------- 6 files changed, 99 insertions(+), 131 deletions(-) diff --git a/install.ps1 b/install.ps1 index a561e6f9b0..ab60b06e0a 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2156,29 +2156,21 @@ exit 0 } } } elseif ($TorchIndexUrl -or $ROCmIndexUrl) { - # Leaf-gated bounded trio, HOISTED so the release-preservation decision below can use it as the - # default route window: torchaudio 2.11 dropped its torch pin, so a bare companion can drift from - # a capped torch. cu families ship torch 2.11.x (paired triton-windows 3.6) so the ceiling - # widens to <2.12; other leaves keep the 2.10 line. Mirrors install.sh and _CUDA_TORCH_PKG_SPEC. - $_idxLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() - if ($_idxLeaf -match '^cu[0-9]+$') { - $_pinTorchSpec = "torch>=2.4,<2.12.0" - $_pinVisionSpec = "torchvision>=0.19,<0.27.0" - $_pinAudioSpec = "torchaudio>=2.4,<2.12.0" - } else { - $_pinTorchSpec = "torch>=2.4,<2.11.0" - $_pinVisionSpec = "torchvision>=0.19,<0.26.0" - $_pinAudioSpec = "torchaudio>=2.4,<2.11.0" - } + # Bounded default trio (torch 2.11 line; wheels verified on cpu + cu126/cu128/cu130 for + # win_amd64 with paired triton-windows 3.6), HOISTED so the release-preservation decision + # below can use it as the default route window. torchaudio 2.11 dropped its torch pin, so a + # bare companion can drift from a capped torch. Bump the three ceilings together with + # install.sh's _TORCH_CEILING trio and the repair/fallback sites below (2 more literals). + $_pinTorchSpec = "torch>=2.4,<2.12.0" + $_pinVisionSpec = "torchvision>=0.19,<0.27.0" + $_pinAudioSpec = "torchaudio>=2.4,<2.12.0" # Release preservation (twin of install.sh's _PREV_TORCH_PIN decision): evaluated after every # index/floor choice, incl. the ROCm reroute, so a raised floor rejects an older release. The - # route window is the leaf-gated CUDA trio by default; the ROCm path uses its floor (or the - # torch>=2.4,<2.11.0 range the ROCm->CPU fallback installs). The kept release is exported for - # setup.ps1 (UNSLOTH_KEPT_TORCH) and cleared after setup runs. + # kept release is exported for setup.ps1 (UNSLOTH_KEPT_TORCH) and cleared after setup runs. $script:PrevTorchPin = $null if (-not $SkipTorch -and $script:PrevTorchVer) { $_routeWindow = $_pinTorchSpec - if ($ROCmIndexUrl) { if ($ROCmTorchFloor) { $_routeWindow = $ROCmTorchFloor } else { $_routeWindow = "torch>=2.4,<2.11.0" } } + if ($ROCmIndexUrl -and $ROCmTorchFloor) { $_routeWindow = $ROCmTorchFloor } $script:PrevTorchPin = Get-PreviousTorchPin -TorchVersion $script:PrevTorchVer -Constraint $_routeWindow if ($script:PrevTorchPin) { $env:UNSLOTH_KEPT_TORCH = $script:PrevTorchPin.Release.PublicBase @@ -2212,7 +2204,7 @@ exit 0 $CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" } substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow" # --force-reinstall: a failed ROCm install can leave an unpinned ROCm torch that still satisfies the CPU torch>= range, so without it uv would keep the ROCm build and only swap companions -- a mismatched venv the flavor-repair block won't fix. (No kept-release attempt: the ROCm attempts above always resolve or clear $script:PrevTorchPin first, so a pin never reaches this CPU base.) - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $CpuFallbackIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.12.0" "torchvision>=0.19,<0.27.0" "torchaudio>=2.4,<2.12.0" --default-index $CpuFallbackIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2363,13 +2355,8 @@ exit 0 } $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython } elseif ($expectedTorchTag -ne 'rocm') { - # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. Same leaf-gated ceiling as the install above: cu* serves 2.11. - $_fixLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() - if ($_fixLeaf -match '^cu[0-9]+$') { - $_fixTorchSpec = "torch>=2.4,<2.12.0"; $_fixVisionSpec = "torchvision>=0.19,<0.27.0"; $_fixAudioSpec = "torchaudio>=2.4,<2.12.0" - } else { - $_fixTorchSpec = "torch>=2.4,<2.11.0"; $_fixVisionSpec = "torchvision>=0.19,<0.26.0"; $_fixAudioSpec = "torchaudio>=2.4,<2.11.0" - } + # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet with the default 2.11-line trio (ceiling bump site, see the hoisted trio above). + $_fixTorchSpec = "torch>=2.4,<2.12.0"; $_fixVisionSpec = "torchvision>=0.19,<0.27.0"; $_fixAudioSpec = "torchaudio>=2.4,<2.12.0" # Kept-release substitution (twin of install.sh's _install_torch_default_index honoring _PREV_TORCH_PIN): honor the preserved torch when the pin survived the E-decision; restore the range specs and retry if it isn't installable here. The --reinstall-package triplet stays on both attempts. $_cudaKept = $false if ($script:PrevTorchPin) { diff --git a/install.sh b/install.sh index 7a4eb572e4..3919d4b53e 100755 --- a/install.sh +++ b/install.sh @@ -1888,18 +1888,24 @@ if [ -x "$VENV_DIR/bin/python" ]; then substep "${VENV_DIR}" fi +# Supported torch line: the default range admits torch 2.11 (wheels verified on +# cpu/cu126/cu128/cu130/rocm7.1+/mac arm64). Bump the three ceilings together +# when the next torch minor is validated; curated ROCm floors below stay literal. +_TORCH_CEILING="2.12.0" +_TORCHVISION_CEILING="0.27.0" +_TORCHAUDIO_CEILING="2.12.0" # Default torch constraint; tightened for Python 3.13+ on arm64 macOS (torch <2.6 has no cp313 macOS arm64 wheels). -TORCH_CONSTRAINT="torch>=2.4,<2.11.0" +TORCH_CONSTRAINT="torch>=2.4,<${_TORCH_CEILING}" if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then _PY_MINOR=$("$VENV_DIR/bin/python" -c \ "import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0") if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then - TORCH_CONSTRAINT="torch>=2.6,<2.11.0" + TORCH_CONSTRAINT="torch>=2.6,<${_TORCH_CEILING}" fi fi -# Companion constraints bounded to torch's window: torchaudio 2.11 dropped its torch pin, so a bare companion beside a <2.11 torch resolves 2.11. -TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" -TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" +# Companion constraints bounded to torch's window: torchaudio 2.11 dropped its torch pin, so a bare companion can drift from a capped torch. +TORCHVISION_CONSTRAINT="torchvision>=0.19,<${_TORCHVISION_CEILING}" +TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<${_TORCHAUDIO_CEILING}" # ── Resolve repo root (for --local installs) ── _REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)" @@ -2480,27 +2486,15 @@ else fi # rocm7.2 and the per-gfx indexes (Strix _grouped_mm fix) ship torch 2.11.0: raise the floor and pin companions; match the FINAL leaf only. +# (cu*/cpu/custom leaves all use the default <2.12 trio above.) case "$_torch_index_leaf" in rocm7.2|gfx120x-all|gfx1151|gfx1150) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" ;; - # CUDA cu12x/cu13x indexes ship torch 2.11.x: widen the trio ceiling to <2.12.0. - cu[0-9]*) - TORCH_CONSTRAINT="torch>=2.4,<2.12.0" - TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.27.0" - TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0" - ;; esac -# A pinned custom/unknown-leaf index has no curated companion set: bound the companions to the same <2.11 range the Python path pins. -if [ "$_torch_index_pinned" = true ] && \ - [ -z "$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")" ]; then - TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" - TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" -fi - # Detect a Radeon card (*/rocm* index + rocminfo "Marketing Name:.*Radeon"); skipped when the index is pinned. _amd_gpu_radeon=false if [ "$_torch_index_pinned" = false ]; then diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 79690fb76f..a63454e395 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2999,9 +2999,9 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { $cudaVisionSpec = "torchvision" $cudaAudioSpec = "torchaudio" if ($TorchIndexPinned -and -not (Test-CudaFamilyLeaf $CuTag)) { - $cudaTorchSpec = "torch>=2.4,<2.11.0" - $cudaVisionSpec = "torchvision>=0.19,<0.26.0" - $cudaAudioSpec = "torchaudio>=2.4,<2.11.0" + $cudaTorchSpec = "torch>=2.4,<2.12.0" + $cudaVisionSpec = "torchvision>=0.19,<0.27.0" + $cudaAudioSpec = "torchaudio>=2.4,<2.12.0" } # Release preservation: keep install.ps1's exported torch RELEASE (UNSLOTH_KEPT_TORCH) so a re-run # reinstalls the same CUDA build (+cuXXX follows this index). On a failed install, restore the computed diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index d7e4f4ffdd..751b12253d 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -405,30 +405,25 @@ class TestKnown211SetParity: ), f"{label} floor gate must not use the unanchored ^rocm(\\d+)\\.(\\d+) prefix" def test_install_ps1_bounds_unknown_leaf_pinned_torch(self): - """install.ps1's pinned-torch install must bound BOTH companions on EVERY - index, cu families included: torchaudio 2.11 dropped its exact torch - pin from the wheel metadata, so a bare companion beside torch<2.11 can - resolve a mismatched 2.11.0 build (Codex P2, then unconditional per the - torchaudio 2.11 unpinning).""" + """install.ps1's pinned-torch install must bound the whole trio on EVERY + index with the default torch 2.11 line (<2.12 trio, matching install.sh's + ceiling-composed default and _CUDA_TORCH_PKG_SPEC): torchaudio 2.11 + dropped its exact torch pin from the wheel metadata, so a bare companion + beside a capped torch can resolve a mismatched build.""" text = INSTALL_PS1.read_text(encoding = "utf-8") - assert ( - '$_pinVisionSpec = "torchvision>=0.19,<0.26.0"' in text - ), "install.ps1 custom-pin install must bound torchvision (>=0.19,<0.26.0)" - assert ( - '$_pinAudioSpec = "torchaudio>=2.4,<2.11.0"' in text - ), "install.ps1 custom-pin install must bound torchaudio (>=2.4,<2.11.0)" - # cu leaves widen the whole trio to the torch 2.11 line (<2.12), - # matching install.sh's cu[0-9]* widen and _CUDA_TORCH_PKG_SPEC; other - # leaves keep the 2.10 line. The trio stays bounded on every index. assert '$_pinTorchSpec = "torch>=2.4,<2.12.0"' in text, ( - "install.ps1 must widen torch to <2.12.0 on cu index leaves" + "install.ps1 default install must use the torch 2.11 line (<2.12.0)" ) assert ( '$_pinVisionSpec = "torchvision>=0.19,<0.27.0"' in text - ), "install.ps1 cu-leaf install must pair torchvision <0.27.0 with torch <2.12" + ), "install.ps1 must pair torchvision <0.27.0 with torch <2.12" assert ( '$_pinAudioSpec = "torchaudio>=2.4,<2.12.0"' in text - ), "install.ps1 cu-leaf install must pair torchaudio <2.12.0 with torch <2.12" + ), "install.ps1 must pair torchaudio <2.12.0 with torch <2.12" + # No stale 2.10-line default remains anywhere in the Windows installer. + assert '"torch>=2.4,<2.11.0"' not in text, ( + "install.ps1 must not retain a <2.11.0 default torch range" + ) # The bounded trio must actually be passed to the install command. assert re.search( r'\$_pinTorchSpec \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl', @@ -672,9 +667,9 @@ class TestPinnedIndexClearsUvEnvParity: # The custom-leaf branch bounds torch AND both companions (parity with the # other installers' custom-pin trio bounds), gated on a non-cu-family leaf. for spec in ( - '$cudaTorchSpec = "torch>=2.4,<2.11.0"', - '$cudaVisionSpec = "torchvision>=0.19,<0.26.0"', - '$cudaAudioSpec = "torchaudio>=2.4,<2.11.0"', + '$cudaTorchSpec = "torch>=2.4,<2.12.0"', + '$cudaVisionSpec = "torchvision>=0.19,<0.27.0"', + '$cudaAudioSpec = "torchaudio>=2.4,<2.12.0"', ): assert spec in text, f"setup.ps1 must bound the custom-leaf trio: {spec}" assert ( diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py index c58808689b..a7ad585583 100644 --- a/tests/python/test_tokenizers_and_torch_constraint.py +++ b/tests/python/test_tokenizers_and_torch_constraint.py @@ -64,35 +64,30 @@ class TestStructuralTorchConstraint: _sh = _read(_INSTALL_SH) def test_default_assignment_exists(self): - assert 'TORCH_CONSTRAINT="torch>=2.4,<2.11.0"' in self._sh + """The default range composes the per-file ceiling variable, so the + supported line (torch 2.11 today) is bumped in one place.""" + assert '_TORCH_CEILING="2.12.0"' in self._sh + assert 'TORCH_CONSTRAINT="torch>=2.4,<${_TORCH_CEILING}"' in self._sh def test_tightened_assignment_exists(self): - assert 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' in self._sh + assert 'TORCH_CONSTRAINT="torch>=2.6,<${_TORCH_CEILING}"' in self._sh - def test_cuda_constraint_widened_to_2_12(self): - """A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x - land torch 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC); - without it cu128/cu130 resolves torch 2.10.x.""" - assert 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' in self._sh - - def test_cuda_case_widens_via_index_leaf(self): - """The cu* branch of the _torch_index_leaf case sets the widened - constraint (parallel to rocm7.2), anchored on the leaf.""" - m = re.search( - r'cu\[0-9\]\*\)\s*TORCH_CONSTRAINT="torch>=2\.4,<2\.12\.0"', - self._sh, - ) - assert m is not None, "CUDA (cu*) TORCH_CONSTRAINT widening case not found" + def test_companion_ceilings_composed(self): + """Companions bound to the same window via their own ceiling vars.""" + assert '_TORCHVISION_CEILING="0.27.0"' in self._sh + assert '_TORCHAUDIO_CEILING="2.12.0"' in self._sh + assert 'TORCHVISION_CONSTRAINT="torchvision>=0.19,<${_TORCHVISION_CEILING}"' in self._sh + assert 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<${_TORCHAUDIO_CEILING}"' in self._sh def test_variable_used_in_pip_install(self): """$TORCH_CONSTRAINT must appear in a uv pip install line.""" assert '"$TORCH_CONSTRAINT"' in self._sh - def test_hardcoded_torch_constraint_only_once(self): - """The hard-coded torch>=2.4,<2.11.0 string should appear exactly once - in install.sh (the default assignment), not in pip install lines.""" - count = self._sh.count('"torch>=2.4,<2.11.0"') - assert count == 1, f"Expected 1, found {count}" + def test_hardcoded_torch_constraint_gone(self): + """No hard-coded default ranges remain outside the ceiling-composed + assignments (curated ROCm >=2.11 floors stay literal).""" + assert self._sh.count('"torch>=2.4,<2.11.0"') == 0 + assert self._sh.count('"torch>=2.4,<2.12.0"') == 0 def test_tightening_guarded_by_skip_torch(self): """The block must check SKIP_TORCH=false.""" @@ -122,7 +117,7 @@ class TestStructuralInstallPs1Unchanged: assert "$TorchConstraint" not in self._ps1 def test_hardcoded_torch_constraint_present(self): - assert '"torch>=2.4,<2.11.0"' in self._ps1 + assert '"torch>=2.4,<2.12.0"' in self._ps1 class TestInstallPs1UvDefaultIndex: diff --git a/tests/sh/test_torch_constraint.sh b/tests/sh/test_torch_constraint.sh index bfafbd161b..d36581987b 100644 --- a/tests/sh/test_torch_constraint.sh +++ b/tests/sh/test_torch_constraint.sh @@ -76,11 +76,12 @@ run_constraint_snippet() { OS=\"$_os\" _ARCH=\"$_arch\" VENV_DIR=\"$_venv_dir\" - TORCH_CONSTRAINT=\"torch>=2.4,<2.11.0\" + _TORCH_CEILING=\"2.12.0\" + TORCH_CONSTRAINT=\"torch>=2.4,<\${_TORCH_CEILING}\" if [ \"\$SKIP_TORCH\" = false ] && [ \"\$OS\" = \"macos\" ] && [ \"\$_ARCH\" = \"arm64\" ]; then _PY_MINOR=\$(\"\$VENV_DIR/bin/python\" -c \"import sys; print(sys.version_info.minor)\" 2>/dev/null || echo \"0\") if [ \"\$_PY_MINOR\" -ge 13 ] 2>/dev/null; then - TORCH_CONSTRAINT=\"torch>=2.6,<2.11.0\" + TORCH_CONSTRAINT=\"torch>=2.6,<\${_TORCH_CEILING}\" fi fi echo \"\$TORCH_CONSTRAINT\" @@ -94,43 +95,39 @@ echo "=== Structural: TORCH_CONSTRAINT in install.sh ===" _SH_CONTENT=$(cat "$INSTALL_SH") -_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.11.0"' "$INSTALL_SH" || true) -assert_eq "default TORCH_CONSTRAINT assignment exists" "1" "$_count" +# The supported line is centralized in per-file ceiling variables so a future +# torch 2.12 bump is a three-line change; the default range admits torch 2.11. +_count=$(grep -c '_TORCH_CEILING="2.12.0"' "$INSTALL_SH" || true) +assert_eq "torch ceiling variable defined once" "1" "$_count" +_count=$(grep -c '_TORCHVISION_CEILING="0.27.0"' "$INSTALL_SH" || true) +assert_eq "torchvision ceiling variable defined once" "1" "$_count" +_count=$(grep -c '_TORCHAUDIO_CEILING="2.12.0"' "$INSTALL_SH" || true) +assert_eq "torchaudio ceiling variable defined once" "1" "$_count" -_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' "$INSTALL_SH" || true) -assert_eq "tightened TORCH_CONSTRAINT assignment exists" "1" "$_count" +_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<${_TORCH_CEILING}"' "$INSTALL_SH" || true) +assert_eq "default TORCH_CONSTRAINT composes the ceiling" "1" "$_count" +_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.6,<${_TORCH_CEILING}"' "$INSTALL_SH" || true) +assert_eq "tightened TORCH_CONSTRAINT composes the ceiling" "1" "$_count" _count=$(grep -c '"\$TORCH_CONSTRAINT"' "$INSTALL_SH" || true) _has_var=$([ "$_count" -ge 1 ] && echo "yes" || echo "no") assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var" -# Hardcoded torch>=2.4,<2.11.0 should only appear once (the default assignment) -_hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true) -assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded" +# No stray hardcoded default ranges outside the ceiling-composed assignments +# (the curated ROCm >=2.11 floors are deliberately literal). +_hardcoded=$(grep -c '"torch>=2.4,<2.11.0"\|"torch>=2.4,<2.12.0"' "$INSTALL_SH" || true) +assert_eq "no hardcoded default torch range remains" "0" "$_hardcoded" -# Companions must be bounded to torch's window everywhere: the <2.11 bound appears -# twice (default assignments + the pinned custom-leaf block), never bare. torchaudio -# 2.11 dropped its exact torch pin, so a bare companion next to a <2.11-capped torch -# resolves a mismatched 2.11 build. -_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"' "$INSTALL_SH" || true) -assert_eq "torchvision bounded (<0.26) at default + custom-leaf" "2" "$_count" -_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"' "$INSTALL_SH" || true) -assert_eq "torchaudio bounded (<2.11) at default + custom-leaf" "2" "$_count" +# Companions must be bounded to torch's window everywhere, never bare: torchaudio +# 2.11 dropped its exact torch pin, so a bare companion can drift from a capped torch. +_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision>=0.19,<${_TORCHVISION_CEILING}"' "$INSTALL_SH" || true) +assert_eq "torchvision default composes the ceiling" "1" "$_count" +_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<${_TORCHAUDIO_CEILING}"' "$INSTALL_SH" || true) +assert_eq "torchaudio default composes the ceiling" "1" "$_count" _count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision"$' "$INSTALL_SH" || true) assert_eq "no bare torchvision companion remains" "0" "$_count" _count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio"$' "$INSTALL_SH" || true) assert_eq "no bare torchaudio companion remains" "0" "$_count" -# The cu* widen must carry the companions with it (torch <2.12 with torchaudio <2.11 -# would cap a mismatched pair the other way). -assert_eq "cu widen pairs torchaudio (<2.12)" "1" "$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0"' "$INSTALL_SH" || true)" -_gated=$(grep -c '_expected_torch_flavor_tag "$TORCH_INDEX_URL"' "$INSTALL_SH" || true) -_has_gate=$([ "$_gated" -ge 1 ] && echo "yes" || echo "no") -assert_eq "custom-companion bound gated on empty flavor tag" "yes" "$_has_gate" - -# A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x land torch -# 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC). -_cuda_widen=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' "$INSTALL_SH" || true) -assert_eq "CUDA TORCH_CONSTRAINT widened to <2.12.0" "1" "$_cuda_widen" # Widening keys off the final leaf (_torch_index_leaf), not the full URL, so a # mirror base path with cu*/rocm7.2 but a cpu/older-rocm leaf is not mis-widened. @@ -181,7 +178,7 @@ _PS1_CONTENT=$(cat "$INSTALL_PS1") _ps1_has_var=$(echo "$_PS1_CONTENT" | grep -c 'TORCH_CONSTRAINT\|TorchConstraint' || true) assert_eq "install.ps1 has no TORCH_CONSTRAINT variable" "0" "$_ps1_has_var" -_ps1_hardcoded=$(echo "$_PS1_CONTENT" | grep -c '"torch>=2.4,<2.11.0"' || true) +_ps1_hardcoded=$(echo "$_PS1_CONTENT" | grep -c '"torch>=2.4,<2.12.0"' || true) _ps1_has_hc=$([ "$_ps1_hardcoded" -ge 1 ] && echo "yes" || echo "no") assert_eq "install.ps1 has hardcoded torch constraint" "yes" "$_ps1_has_hc" @@ -196,55 +193,55 @@ trap 'rm -rf "$TMPDIR_BASE"' EXIT # 1. arm64 macOS py3.13 -> tightened _result=$(run_constraint_snippet false macos arm64 13 "$TMPDIR_BASE/v1") -assert_eq "arm64+macos+py313 -> tightened" "torch>=2.6,<2.11.0" "$_result" +assert_eq "arm64+macos+py313 -> tightened" "torch>=2.6,<2.12.0" "$_result" # 2. arm64 macOS py3.14 -> tightened (future-proofed) _result=$(run_constraint_snippet false macos arm64 14 "$TMPDIR_BASE/v2") -assert_eq "arm64+macos+py314 -> tightened" "torch>=2.6,<2.11.0" "$_result" +assert_eq "arm64+macos+py314 -> tightened" "torch>=2.6,<2.12.0" "$_result" # 3. arm64 macOS py3.12 -> default _result=$(run_constraint_snippet false macos arm64 12 "$TMPDIR_BASE/v3") -assert_eq "arm64+macos+py312 -> default" "torch>=2.4,<2.11.0" "$_result" +assert_eq "arm64+macos+py312 -> default" "torch>=2.4,<2.12.0" "$_result" # 4. arm64 macOS py3.11 -> default _result=$(run_constraint_snippet false macos arm64 11 "$TMPDIR_BASE/v4") -assert_eq "arm64+macos+py311 -> default" "torch>=2.4,<2.11.0" "$_result" +assert_eq "arm64+macos+py311 -> default" "torch>=2.4,<2.12.0" "$_result" # 5. Linux x86_64 py3.13 -> default (Linux unaffected) _result=$(run_constraint_snippet false linux x86_64 13 "$TMPDIR_BASE/v5") -assert_eq "linux+x86_64+py313 -> default" "torch>=2.4,<2.11.0" "$_result" +assert_eq "linux+x86_64+py313 -> default" "torch>=2.4,<2.12.0" "$_result" # 6. Linux aarch64 py3.13 -> default (guard checks OS=macos) _result=$(run_constraint_snippet false linux aarch64 13 "$TMPDIR_BASE/v6") -assert_eq "linux+aarch64+py313 -> default" "torch>=2.4,<2.11.0" "$_result" +assert_eq "linux+aarch64+py313 -> default" "torch>=2.4,<2.12.0" "$_result" # 7. Intel Mac x86_64 py3.12 -> default (arch mismatch) _result=$(run_constraint_snippet false macos x86_64 12 "$TMPDIR_BASE/v7") -assert_eq "macos+x86_64+py312 -> default" "torch>=2.4,<2.11.0" "$_result" +assert_eq "macos+x86_64+py312 -> default" "torch>=2.4,<2.12.0" "$_result" # 8. SKIP_TORCH=true arm64 macOS py3.13 -> block skipped, default _result=$(run_constraint_snippet true macos arm64 13 "$TMPDIR_BASE/v8") -assert_eq "SKIP_TORCH=true -> default" "torch>=2.4,<2.11.0" "$_result" +assert_eq "SKIP_TORCH=true -> default" "torch>=2.4,<2.12.0" "$_result" # 9. WSL py3.13 -> default _result=$(run_constraint_snippet false wsl x86_64 13 "$TMPDIR_BASE/v9") -assert_eq "wsl+py313 -> default" "torch>=2.4,<2.11.0" "$_result" +assert_eq "wsl+py313 -> default" "torch>=2.4,<2.12.0" "$_result" # 10. py_minor=0 (failed query fallback) -> default _result=$(run_constraint_snippet false macos arm64 0 "$TMPDIR_BASE/v10") -assert_eq "py_minor=0 fallback -> default" "torch>=2.4,<2.11.0" "$_result" +assert_eq "py_minor=0 fallback -> default" "torch>=2.4,<2.12.0" "$_result" # 11. Boundary: py_minor=12 -> NOT tightened _result=$(run_constraint_snippet false macos arm64 12 "$TMPDIR_BASE/v11") -assert_eq "boundary py_minor=12 -> default" "torch>=2.4,<2.11.0" "$_result" +assert_eq "boundary py_minor=12 -> default" "torch>=2.4,<2.12.0" "$_result" # 12. Boundary: py_minor=13 -> tightened _result=$(run_constraint_snippet false macos arm64 13 "$TMPDIR_BASE/v12") -assert_eq "boundary py_minor=13 -> tightened" "torch>=2.6,<2.11.0" "$_result" +assert_eq "boundary py_minor=13 -> tightened" "torch>=2.6,<2.12.0" "$_result" # 13. Intel Mac py3.13 -> default (arch=x86_64, not arm64) _result=$(run_constraint_snippet false macos x86_64 13 "$TMPDIR_BASE/v13") -assert_eq "macos+x86_64+py313 -> default" "torch>=2.4,<2.11.0" "$_result" +assert_eq "macos+x86_64+py313 -> default" "torch>=2.4,<2.12.0" "$_result" # ====================================================================== # Mock uv integration From 49e304221a9e0f7a22c42ba2a69c8ada0134a6c3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:46:13 +0000 Subject: [PATCH 05/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/python/test_cross_platform_parity.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 751b12253d..499d2faf85 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -411,9 +411,9 @@ class TestKnown211SetParity: dropped its exact torch pin from the wheel metadata, so a bare companion beside a capped torch can resolve a mismatched build.""" text = INSTALL_PS1.read_text(encoding = "utf-8") - assert '$_pinTorchSpec = "torch>=2.4,<2.12.0"' in text, ( - "install.ps1 default install must use the torch 2.11 line (<2.12.0)" - ) + assert ( + '$_pinTorchSpec = "torch>=2.4,<2.12.0"' in text + ), "install.ps1 default install must use the torch 2.11 line (<2.12.0)" assert ( '$_pinVisionSpec = "torchvision>=0.19,<0.27.0"' in text ), "install.ps1 must pair torchvision <0.27.0 with torch <2.12" @@ -421,12 +421,12 @@ class TestKnown211SetParity: '$_pinAudioSpec = "torchaudio>=2.4,<2.12.0"' in text ), "install.ps1 must pair torchaudio <2.12.0 with torch <2.12" # No stale 2.10-line default remains anywhere in the Windows installer. - assert '"torch>=2.4,<2.11.0"' not in text, ( - "install.ps1 must not retain a <2.11.0 default torch range" - ) + assert ( + '"torch>=2.4,<2.11.0"' not in text + ), "install.ps1 must not retain a <2.11.0 default torch range" # The bounded trio must actually be passed to the install command. assert re.search( - r'\$_pinTorchSpec \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl', + r"\$_pinTorchSpec \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl", text, ), "install.ps1 pinned install must pass the bounded trio specs to uv" From d90fd8b563ac18c896bd6603524fdd7793db1efb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 20 Jul 2026 12:07:08 +0000 Subject: [PATCH 06/10] install: harden the preservation probe and the Windows override handoff Three review follow-ups: - install.sh's existing-venv torch probe now reads the version from dist metadata via the bounded runner instead of importing torch: a wedged CUDA/ROCm driver can hang the import indefinitely, while the metadata read never touches the driver. A failed or timed-out probe simply yields no preservable version. - install.ps1's New-UnslothTorchOverridesFile folds caller-supplied UV_OVERRIDE files into the temporary overrides file (minus their torch-trio lines), matching install.sh: --overrides replaces the env file wholesale, so without the merge a caller's own dependency overrides were silently dropped during the unsloth resolution. - install.ps1 clears an inherited UNSLOTH_KEPT_TORCH before the pin decision and sets it only on a fresh Get-PreviousTorchPin result: an interrupted earlier run could leak a stale exact pin into setup.ps1 even when the current run found nothing to preserve or the upgrade opt-in was set. Verified: the metadata probe returns the installed release through the bounded runner; the override filter keeps unrelated pins (numpy, torchao) while dropping torch/torchvision/torchaudio lines in all spec forms. Full sh, ps1 and pytest installer batteries pass (host-defaults and the tokenizers negative-control are the known pre-existing failures). --- install.ps1 | 14 ++++++++++++++ install.sh | 8 +++++--- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/install.ps1 b/install.ps1 index ab60b06e0a..63b88efa35 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2112,6 +2112,17 @@ exit 0 $pins = & $PythonExe -c "from importlib.metadata import version, PackageNotFoundError`nfor _p in ('torch', 'torchvision', 'torchaudio'):`n try:`n print(_p + '==' + version(_p))`n except PackageNotFoundError:`n pass" 2>$null $lines = @($pins | Where-Object { $_ -match '^torch' }) if ($lines.Count -eq 0 -or $lines[0] -notmatch '^torch==') { return $null } + # --overrides replaces any UV_OVERRIDE env file, so fold caller-supplied + # override files in (minus their torch-trio lines) like install.sh does. + if ($env:UV_OVERRIDE) { + foreach ($ovFile in ($env:UV_OVERRIDE -split '\s+' | Where-Object { $_ })) { + if (Test-Path -LiteralPath $ovFile -PathType Leaf) { + $lines += @(Get-Content -LiteralPath $ovFile | Where-Object { + $_ -notmatch '^\s*torch(vision|audio)?([\s<>=!~;@[]|$)' + }) + } + } + } $f = [System.IO.Path]::GetTempFileName() Set-Content -LiteralPath $f -Value ($lines -join "`n") -Encoding ascii return $f @@ -2168,6 +2179,9 @@ exit 0 # index/floor choice, incl. the ROCm reroute, so a raised floor rejects an older release. The # kept release is exported for setup.ps1 (UNSLOTH_KEPT_TORCH) and cleared after setup runs. $script:PrevTorchPin = $null + # Internal handoff variable: always clear an inherited value first (an interrupted + # earlier run can leak a stale pin into setup.ps1) and set it only on a fresh decision. + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue if (-not $SkipTorch -and $script:PrevTorchVer) { $_routeWindow = $_pinTorchSpec if ($ROCmIndexUrl -and $ROCmTorchFloor) { $_routeWindow = $ROCmTorchFloor } diff --git a/install.sh b/install.sh index 3919d4b53e..2ebb9be571 100755 --- a/install.sh +++ b/install.sh @@ -1770,9 +1770,11 @@ if [ -x "$VENV_DIR/bin/python" ]; then echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2 exit 1 fi - # Record the existing venv's torch BEFORE replacement (see _previous_torch_pin); last line only so stdout noise can't corrupt it. - _PREV_TORCH_VER=$("$VENV_DIR/bin/python" -c \ - "import torch; print(torch.__version__)" 2>/dev/null | tail -n 1 || true) + # Record the existing venv's torch BEFORE replacement (see _previous_torch_pin); last line + # only so stdout noise can't corrupt it. Reads dist metadata instead of importing torch (a + # wedged CUDA/ROCm driver can hang the import indefinitely) and bounds the interpreter run. + _PREV_TORCH_VER=$(_run_bounded "$VENV_DIR/bin/python" -c \ + "import importlib.metadata as m; print(m.version('torch'))" 2>/dev/null | tail -n 1 || true) substep "preserving existing environment for rollback..." _start_studio_venv_replacement "$VENV_DIR" elif [ "$_STUDIO_HOME_REDIRECT" != "env" ] && [ -x "$STUDIO_HOME/.venv/bin/python" ]; then From c845e726929da51988485e7f470804677b68ed6d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 24 Jul 2026 11:09:12 +0000 Subject: [PATCH 07/10] Harden torch-constraint test grep portability and clear kept-torch on failure - test_torch_constraint.sh: use grep -E for the hardcoded-range guard. The BRE \| alternation is a GNU extension; on BSD/macOS grep it can be a literal, so the regression guard could silently no-op on a supported platform. - install.ps1: clear UNSLOTH_KEPT_TORCH in Exit-InstallFailure. A non-Tauri irm | iex run throws while the caller session stays alive, so a leaked kept-torch handoff could let a later setup/update re-pin an abandoned exact torch release. --- install.ps1 | 5 +++++ tests/sh/test_torch_constraint.sh | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/install.ps1 b/install.ps1 index 3d709ae1ff..f64021c629 100644 --- a/install.ps1 +++ b/install.ps1 @@ -87,6 +87,11 @@ function Install-UnslothStudio { ) if ($Code -eq 0) { $Code = 1 } Write-TauriLog "ERROR" $Message + # Clear the release-preservation handoff on any failure: a non-Tauri `irm | iex` + # run throws below and leaves the caller's session alive, so a leaked + # UNSLOTH_KEPT_TORCH would let a later `studio setup`/`update` (or a retry that + # skips the branch-entry clear) re-pin the abandoned exact torch release. + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue if (Get-Command Restore-StudioVenvRollback -CommandType Function -ErrorAction SilentlyContinue) { Restore-StudioVenvRollback } diff --git a/tests/sh/test_torch_constraint.sh b/tests/sh/test_torch_constraint.sh index d36581987b..24e12787e2 100644 --- a/tests/sh/test_torch_constraint.sh +++ b/tests/sh/test_torch_constraint.sh @@ -115,7 +115,7 @@ assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var" # No stray hardcoded default ranges outside the ceiling-composed assignments # (the curated ROCm >=2.11 floors are deliberately literal). -_hardcoded=$(grep -c '"torch>=2.4,<2.11.0"\|"torch>=2.4,<2.12.0"' "$INSTALL_SH" || true) +_hardcoded=$(grep -E -c '"torch>=2.4,<2.11.0"|"torch>=2.4,<2.12.0"' "$INSTALL_SH" || true) assert_eq "no hardcoded default torch range remains" "0" "$_hardcoded" # Companions must be bounded to torch's window everywhere, never bare: torchaudio From 27c81d493b38803e6f1943231fa75dddef779ad1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 11:25:28 +0000 Subject: [PATCH 08/10] install.ps1: make the release preservation robust across sessions and failures Three review follow-ups on the Windows port: - The previous-release probe reads dist metadata rather than importing torch, matching install.sh. A broken CUDA/ROCm DLL (or an import slow enough to hit the timeout) made the probe return nothing, so the pin was never created and the installer replaced the venv with the newest supported release even though UNSLOTH_TORCH_UPGRADE was unset. - The preservation state is reset at the top of the install instead of inside the existing-venv branch. Under the documented irm | iex flow the script scope IS the caller's session, so a second invocation against a different UNSLOTH_STUDIO_HOME with no existing venv skipped the branch and inherited the earlier run's release, pinning a fresh install to an unrelated venv's torch. The pin is reset with it, so the no-index fallback path cannot see a stale pin either. - UNSLOTH_KEPT_TORCH is cleared from an outer finally around the installer entry point. The in-flow clears and Exit-InstallFailure cover the handled paths, but a terminating exception between the export and the end of setup left the handoff set in a session that outlives the installer, where a later studio setup or update would consume the abandoned exact pin. Verified the metadata probe end to end (returns the installed release through the bounded runner; a missing interpreter still yields null). Three regression guards added to tests/studio/test_previous_torch_pin.ps1 for the metadata read, the pre-branch reset ordering and the outer finally. Full sh, ps1 and pytest installer batteries pass (host-defaults and the tokenizers negative-control are the known pre-existing failures). --- install.ps1 | 19 ++++++++++++++++--- tests/studio/test_previous_torch_pin.ps1 | 15 +++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/install.ps1 b/install.ps1 index 3d709ae1ff..f5481aae49 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1278,6 +1278,11 @@ exit 0 $script:StudioVenvRollbackDir = $null $script:StudioVenvRollbackTarget = $VenvDir $script:StudioVenvRollbackActive = $false + # Release-preservation state, reset per run: with `irm | iex` the script scope IS the + # caller's session, so a second invocation (e.g. a different UNSLOTH_STUDIO_HOME with no + # existing venv) must not inherit the previous run's release or pin. + $script:PrevTorchVer = "" + $script:PrevTorchPin = $null function Start-StudioVenvRollback { param([Parameter(Mandatory = $true)][string]$ExistingDir) @@ -1340,7 +1345,9 @@ exit 0 try { $psi = New-Object System.Diagnostics.ProcessStartInfo $psi.FileName = $PythonExe - $psi.Arguments = '-c "import torch; print(torch.__version__)"' + # Dist metadata, not "import torch": a broken CUDA/ROCm DLL or a slow native + # import would yield no release and silently drop the pin (parity with install.sh). + $psi.Arguments = '-c "import importlib.metadata as m; print(m.version(''torch''))"' $psi.RedirectStandardOutput = $true $psi.RedirectStandardError = $true $psi.UseShellExecute = $false @@ -1377,7 +1384,6 @@ exit 0 # a re-run then keeps that release rather than silently jumping torch versions. Opt out with # UNSLOTH_TORCH_UPGRADE=1. Only the new-layout replace probes; the legacy-migration branches # reuse the venv, so torch survives naturally and needs no pin. - $script:PrevTorchVer = "" if (-not $SkipTorch) { $script:PrevTorchVer = Get-InstalledTorchVersionRaw -PythonExe $VenvPython } @@ -2670,4 +2676,11 @@ exit 0 } } -Install-UnslothStudio @args +try { + Install-UnslothStudio @args +} finally { + # UNSLOTH_KEPT_TORCH is a process-scoped handoff to setup.ps1. Under `irm | iex` the + # session outlives the installer, so a terminating exception that bypasses the in-flow + # clears must not leave an abandoned exact pin for a later `studio setup` / `update`. + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue +} diff --git a/tests/studio/test_previous_torch_pin.ps1 b/tests/studio/test_previous_torch_pin.ps1 index 8c84a5e8b4..08cd8b1be8 100644 --- a/tests/studio/test_previous_torch_pin.ps1 +++ b/tests/studio/test_previous_torch_pin.ps1 @@ -108,6 +108,21 @@ Check "probe runs before the rollback move" ( Check "pin decision cites the UNSLOTH_TORCH_UPGRADE escape hatch" ($src -match 'UNSLOTH_TORCH_UPGRADE=1 to get the newest') Check "kept-release fallback clears the pin" ($src -match '\$script:PrevTorchPin\s*=\s*\$null') Check "kept release exported for setup.ps1" ($src -match 'UNSLOTH_KEPT_TORCH') +# The probe must read dist metadata: a broken CUDA/ROCm DLL would make "import torch" +# fail and silently drop the pin (install.sh reads metadata for the same reason). +Check "probe reads dist metadata, not import torch" ( + $src -match 'importlib\.metadata as m; print\(m\.version' -and + $src -notmatch '-c "import torch; print\(torch\.__version__\)"') +# Under irm | iex the script scope is the caller's session: a second run with no existing +# venv must not inherit the earlier run's release or pin. +Check "preservation state reset before the venv branch" ( + $src.IndexOf('$script:PrevTorchVer = ""') -ge 0 -and + $src.IndexOf('$script:PrevTorchVer = ""') -lt $src.IndexOf('if (Test-Path -LiteralPath $VenvPython) {')) +Check "preservation pin reset before the venv branch" ( + $src.IndexOf('$script:PrevTorchPin = $null') -lt $src.IndexOf('if (Test-Path -LiteralPath $VenvPython) {')) +# A terminating exception must not leave the handoff set in a surviving session. +Check "outer finally clears the kept-torch handoff" ( + $src -match '(?s)try \{\s*Install-UnslothStudio @args\s*\} finally \{.*?Remove-Item Env:UNSLOTH_KEPT_TORCH') Write-Host "" if ($failures -gt 0) { Write-Host "$failures check(s) failed" -ForegroundColor Red; exit 1 } From 1e3b1c97ad45e76e9644c01a13803c0d12d0e690 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 15:50:14 +0000 Subject: [PATCH 09/10] install_python_stack.py: keep the ROCm 7.1 repair on the 2.11 line The rocm7.1 leaf fell through to _ROCM_TORCH_PKG_SPECS["_default"], which caps the trio below 2.11. Now that install.sh leaves a rocm7.1 leaf on the widened default range, a fresh install resolves torch 2.11.0+rocm7.1 while the later dependency pass force-reinstalled 2.10.0+rocm7.1 over it, so any repair or `studio update` silently downgraded the environment. download.pytorch.org/whl/rocm7.1 serves a paired 2.11 trio, verified with uv pip compile --no-deps against that index: install.sh default range -> torch 2.11.0+rocm7.1, torchvision 0.26.0+rocm7.1, torchaudio 2.11.0+rocm7.1 _default repair spec -> torch 2.10.0+rocm7.1, torchvision 0.25.0+rocm7.1, torchaudio 2.10.0+rocm7.1 Give rocm7.1 its own entry carrying install.sh's default range rather than the rocm7.2 tuple: only the _grouped_mm arches take the hard 2.11 floor, so _ROCM_KNOWN_TORCH211_VERSIONS stays {(7, 2)}. _default keeps its literal <2.11 ceiling because rocm7.0 and older genuinely top out below it (rocm7.0 at 2.10.0, rocm6.4 and rocm6.3 at 2.9.1, rocm6.2 at 2.5.1), and the stale index comments are corrected to match what those indexes serve today. --- studio/install_python_stack.py | 20 ++++++++++++++---- tests/studio/install/test_rocm_support.py | 25 ++++++++++++++++++++++- 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index fb450a4b54..57c647d11a 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -58,8 +58,8 @@ PLATFORM_LACKS_TORCHCODEC_WHEEL = ( # Detected ROCm (major, minor) -> best PyTorch wheel tag, checked newest-first (>=). _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { (7, 2): "rocm7.2", # torch 2.11.0 - (7, 1): "rocm7.1", # torch 2.10.0 - (7, 0): "rocm7.0", + (7, 1): "rocm7.1", # torch 2.11.0 + (7, 0): "rocm7.0", # torch 2.10.0 (6, 4): "rocm6.4", (6, 3): "rocm6.3", (6, 2): "rocm6.2", @@ -97,14 +97,26 @@ _ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset( # rocmX.Y indexes KNOWN to ship torch 2.11; never floor an unknown newer rocm speculatively. _ROCM_KNOWN_TORCH211_VERSIONS: frozenset[tuple[int, int]] = frozenset({(7, 2)}) -# Per-tag pip specs; rocm7.2 ships torch 2.11.0 (older tags cap at 2.10.x). +# Per-tag pip specs for the repair/update path; must land on the same wheels a fresh +# install.sh run would pick, otherwise `studio update` silently downgrades the venv. _ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = { + # Floored at 2.11 (the _grouped_mm bug), matching install.sh's rocm7.2|gfx* case. "rocm7.2": ( "torch>=2.11.0,<2.12.0", "torchvision>=0.26.0,<0.27.0", "torchaudio>=2.11.0,<2.12.0", ), - # rocm7.1 and earlier: torch 2.x below 2.11 + # rocm7.1 also serves a paired 2.11 trio (torch 2.11.0 / torchvision 0.26.0 / + # torchaudio 2.11.0), so it takes install.sh's widened DEFAULT range rather than + # the 2.11 floor: no _grouped_mm floor applies here, but capping at <2.11 would + # force-reinstall 2.10 over the 2.11 a fresh install just resolved. + "rocm7.1": ( + "torch>=2.4,<2.12.0", + "torchvision>=0.19,<0.27.0", + "torchaudio>=2.4,<2.12.0", + ), + # rocm7.0 and earlier genuinely top out below 2.11 (rocm7.0: torch 2.10.0, + # rocm6.4/6.3: 2.9.1, rocm6.2: 2.5.1), so the old ceiling stays literal. "_default": ( "torch>=2.4,<2.11.0", "torchvision>=0.19,<0.26.0", diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 4f65857a3e..f8d2d4e374 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -3457,12 +3457,35 @@ class TestRocmTorchPkgSpecs: assert "2.11" in torch_spec def test_default_caps_below_211(self): - """Default spec (rocm7.1 and earlier) should cap below 2.11.""" + """Default spec (rocm7.0 and earlier) should cap below 2.11.""" specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("_default") assert specs is not None torch_spec = specs[0] assert "<2.11" in torch_spec + def test_rocm71_repair_matches_install_sh_default_range(self): + """rocm7.1 serves a paired 2.11 trio, so the repair path must not cap at <2.11. + + install.sh leaves a rocm7.1 leaf on its default trio (torch>=2.4,<2.12.0 / + torchvision>=0.19,<0.27.0 / torchaudio>=2.4,<2.12.0), which resolves + torch 2.11.0+rocm7.1 on that index. Falling back to _default here would + force-reinstall 2.10.0+rocm7.1 over it on the next `studio update`. + """ + specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("rocm7.1") + assert specs is not None, "rocm7.1 must have its own repair spec" + assert specs == ( + "torch>=2.4,<2.12.0", + "torchvision>=0.19,<0.27.0", + "torchaudio>=2.4,<2.12.0", + ) + # Not the rocm7.2 spec: no 2.11 floor applies to rocm7.1. + assert specs != stack_mod._ROCM_TORCH_PKG_SPECS["rocm7.2"] + + def test_rocm71_is_not_a_known_211_floor_version(self): + """The widened rocm7.1 range must NOT promote it to a floored 2.11 line.""" + assert (7, 1) not in stack_mod._ROCM_KNOWN_TORCH211_VERSIONS + assert (7, 2) in stack_mod._ROCM_KNOWN_TORCH211_VERSIONS + def test_specs_have_torch_vision_audio(self): """Each entry should be a 3-tuple: torch, torchvision, torchaudio.""" for tag, specs in stack_mod._ROCM_TORCH_PKG_SPECS.items(): From 90fc453e8b317898cd0bc70461893bd9e1b7ecca Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 27 Jul 2026 15:57:33 +0000 Subject: [PATCH 10/10] install.ps1: remove the torch overrides temp file from the outer finally New-UnslothTorchOverridesFile writes a uv --overrides file that also folds in the caller's inherited UV_OVERRIDE lines, which can carry authenticated direct URLs. The in-flow Remove-Item calls only run on the normal path, so a terminating error or Ctrl-C between the GetTempFileName and those removals left the file in %TEMP%. install.sh already removes its twin (_UNSLOTH_TORCH_OVERRIDES) from _cleanup_install_temporaries on both the EXIT and the signal traps, and empties the variable beforehand so an inherited value never reaches the rm. Mirror both halves here: reset $script:TorchOverridesFile to $null ahead of the outer try (under irm | iex the script scope is the caller's session) and sweep it in the finally. The block writes nothing to the success or error stream, so a Ctrl-C cannot truncate it. tests/studio/test_unsloth_torch_override.ps1 is the Windows twin of tests/sh/test_unsloth_torch_override.sh: it asserts the guarded with-deps installs, the removal sites and the null reset, and behaviourally runs the extracted finally body against a live temp file holding a credential-bearing override line. Before this change 17 passed, 5 failed; after, 22 passed, 0 failed. --- install.ps1 | 12 ++ tests/studio/test_unsloth_torch_override.ps1 | 134 +++++++++++++++++++ 2 files changed, 146 insertions(+) create mode 100644 tests/studio/test_unsloth_torch_override.ps1 diff --git a/install.ps1 b/install.ps1 index 96078ae06d..6ee0d8d4f2 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2782,6 +2782,10 @@ exit 0 } } +# $null so a value left by an earlier run in the same session never reaches the finally +# below: under `irm | iex` the script scope IS the caller's session. Twin of install.sh's +# `_UNSLOTH_TORCH_OVERRIDES=""` ahead of its traps -- only a path this run created is removed. +$script:TorchOverridesFile = $null try { Install-UnslothStudio @args } finally { @@ -2789,4 +2793,12 @@ try { # session outlives the installer, so a terminating exception that bypasses the in-flow # clears must not leave an abandoned exact pin for a later `studio setup` / `update`. Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + # Same for the generated uv overrides temp file (twin of install.sh's + # _cleanup_install_temporaries): it copies the caller's inherited UV_OVERRIDE contents, so a + # terminating error between its creation and the in-flow removal would leave those + # requirements sitting in %TEMP%. Only paths this script created are ever set here. + if ($script:TorchOverridesFile) { + Remove-Item -LiteralPath $script:TorchOverridesFile -Force -ErrorAction SilentlyContinue + $script:TorchOverridesFile = $null + } } diff --git a/tests/studio/test_unsloth_torch_override.ps1 b/tests/studio/test_unsloth_torch_override.ps1 new file mode 100644 index 0000000000..bbc04ef3fd --- /dev/null +++ b/tests/studio/test_unsloth_torch_override.ps1 @@ -0,0 +1,134 @@ +#!/usr/bin/env pwsh +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Windows twin of tests/sh/test_unsloth_torch_override.sh: install.ps1's torch-trio +# --overrides guard (New-UnslothTorchOverridesFile) on the Step-2 unsloth installs. +# The generated file also folds in the caller's UV_OVERRIDE lines, which can carry +# authenticated direct URLs, so it must never outlive the run: install.sh removes its +# twin from the EXIT/signal traps and install.ps1 must do the same from the outer +# finally. Pure text/AST assertions plus one behavioural cleanup check -- no venv needed. +# Run: pwsh -NoProfile -File tests/studio/test_unsloth_torch_override.ps1 + +$ErrorActionPreference = "Stop" +$installPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "install.ps1") +$installPath = (Resolve-Path $installPath).Path +$installText = Get-Content -Raw $installPath + +# --- Parse install.ps1 (also serves as a syntax gate) --- +$tokens = $null; $errors = $null +$ast = [System.Management.Automation.Language.Parser]::ParseFile($installPath, [ref]$tokens, [ref]$errors) +if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "install.ps1 has parse errors" } + +$failures = 0 +function Check($name, $cond) { + if ($cond) { Write-Host " PASS $name" } + else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ } +} + +# Text of every Invoke-InstallCommandRetry statement carrying $label. A with-deps +# path has two: the overrides-guarded call and the no-torch-installed fallback. +function Get-InstallBlocks([string]$label) { + $calls = $ast.FindAll({ param($n) + $n -is [System.Management.Automation.Language.CommandAst] -and + $n.Extent.Text -like "*-Label `"$label`" *" -and + $n.GetCommandName() -eq "Invoke-InstallCommandRetry" + }, $true) + if ($calls.Count -eq 0) { throw "no Invoke-InstallCommandRetry found for '$label'" } + return @($calls | ForEach-Object { $_.Extent.Text }) +} + +Write-Host "New-UnslothTorchOverridesFile" +$fnAst = $ast.FindAll({ param($n) + $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and + $n.Name -eq "New-UnslothTorchOverridesFile" +}, $true) +Check "helper defined exactly once" ($fnAst.Count -eq 1) +$fnText = $fnAst[0].Extent.Text +Check "helper returns null under --no-torch" ($fnText -match 'if \(\$SkipTorch\) \{ return \$null \}') +Check "helper folds in caller UV_OVERRIDE files" ($fnText -match '\$env:UV_OVERRIDE') + +Write-Host "with-deps unsloth installs pass --overrides" +foreach ($label in @("install unsloth (local)", "install unsloth")) { + $blocks = Get-InstallBlocks $label + $guarded = @($blocks | Where-Object { $_ -match '--overrides \$script:TorchOverridesFile' }) + Check "'$label' has one overrides-guarded call and one plain fallback" ( + $blocks.Count -eq 2 -and $guarded.Count -eq 1) +} + +Write-Host "the --no-deps no-torch installs carry no overrides" +foreach ($label in @("install unsloth (no-torch)", "install unsloth (migrated no-torch)")) { + $blocks = Get-InstallBlocks $label + Check "'$label' has no --overrides" (@($blocks | Where-Object { $_ -match '--overrides' }).Count -eq 0) +} + +Write-Host "the generated temp file never outlives the run" +$removals = [regex]::Matches($installText, 'Remove-Item -LiteralPath \$script:TorchOverridesFile -Force') +# One in-flow removal per with-deps install, plus the outer-finally sweep. +Check "in-flow removal after each with-deps install, plus a final sweep" ($removals.Count -eq 3) +$outer = @($ast.FindAll({ param($n) + $n -is [System.Management.Automation.Language.TryStatementAst] -and + $n.Body.Extent.Text -match 'Install-UnslothStudio @args' +}, $true)) +Check "outer try/finally around Install-UnslothStudio found" ($outer.Count -eq 1) +$finallyText = $outer[0].Finally.Extent.Text +Check "outer finally removes the overrides temp file" ($finallyText -match 'Remove-Item -LiteralPath \$script:TorchOverridesFile -Force') +Check "outer finally still clears UNSLOTH_KEPT_TORCH" ($finallyText -match 'Remove-Item Env:UNSLOTH_KEPT_TORCH') +# install.sh empties _UNSLOTH_TORCH_OVERRIDES before arming its traps so an inherited +# value can never be rm'd; under `irm | iex` the script scope is the caller's session, +# so the same reset must precede the outer try. +Check "overrides path reset to null before the outer try" ( + $installText -match '(?m)^\$script:TorchOverridesFile = \$null\r?\ntry \{\r?\n\s*Install-UnslothStudio @args') + +Write-Host "outer finally actually deletes the file after a terminating error" +# Behavioural: run the real finally body with a live temp file holding a credential- +# bearing inherited override line, exactly as an interrupted install would leave it. +$leakFile = [System.IO.Path]::GetTempFileName() +Set-Content -LiteralPath $leakFile -Encoding ascii -Value @( + "torch==2.11.0+cu128", + "private-pkg @ https://svc:TOKEN123@pkgs.corp.example/private-1.0-py3-none-any.whl") +$script:TorchOverridesFile = $leakFile +$env:UNSLOTH_KEPT_TORCH = "2.11.0" +# Strip the `finally { ... }` wrapper and run the statements themselves; Invoke-Expression +# evaluates in this scope, so the block's $script: writes land where the installer's would. +$finallyBody = ($finallyText.Trim() -replace '(?s)^\{', '') -replace '(?s)\}$', '' +try { + try { throw "simulated terminating error mid-install" } + finally { Invoke-Expression $finallyBody } +} catch { } +Check "temp overrides file removed" (-not (Test-Path -LiteralPath $leakFile)) +Check "kept-torch handoff still cleared" ($null -eq $env:UNSLOTH_KEPT_TORCH) +Check "tracked path reset so a rerun cannot re-remove it" ($null -eq $script:TorchOverridesFile) +Remove-Item -LiteralPath $leakFile -Force -ErrorAction SilentlyContinue + +Write-Host "the inherited-override filter drops the torch trio in any casing" +# PowerShell's -notmatch is case-insensitive unless written -cnotmatch, so a caller +# override spelled `Torch<2.11` is dropped and the generated exact pin wins. +$filterPattern = $null +if ($fnText -match '\$_ -notmatch ''([^'']+)''') { $filterPattern = $Matches[1] } +Check "filter pattern extracted from the helper" ($null -ne $filterPattern) +$inherited = @( + "# comment survives", + "Torch<2.11", + "TORCHVISION>=0.19", + "TorchAudio==2.1", + "torch<2.11.0", + "torchvision==0.25.0", + "torchaudio!=2.11.0", + "torchmetrics==1.0", + "transformers>=4.57.6", + "anyio<4.14.0" +) +$merged = @("torch==2.11.0+cu128") + @($inherited | Where-Object { $_ -notmatch $filterPattern }) +$trio = @($merged | Where-Object { $_ -match '^(torch|torchvision|torchaudio)([\s<>=!~;@[]|$)' }) +Check "exactly one trio requirement survives (the generated pin)" ($trio.Count -eq 1) +Check "the survivor is the generated exact pin" ($trio[0] -eq "torch==2.11.0+cu128") +foreach ($keep in @("torchmetrics==1.0", "transformers>=4.57.6", "anyio<4.14.0", "# comment survives")) { + Check "unrelated inherited override preserved: $keep" ($merged -contains $keep) +} + +Write-Host "" +if ($failures -gt 0) { + Write-Host "FAILED: $failures check(s)" -ForegroundColor Red + exit 1 +} +Write-Host "All checks passed."