diff --git a/README.md b/README.md index b392ed145f..26a578656c 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,7 @@ unsloth studio -H 0.0.0.0 -p 8888 ``` #### Update +To update, use the same install commands as above. Or run (does not work on Windows): ```bash unsloth studio update ``` @@ -152,7 +153,7 @@ unsloth studio -H 0.0.0.0 -p 8888 ``` Then to update : ```bash -unsloth studio update --local +unsloth studio update ``` #### Developer installs: Windows PowerShell: @@ -165,7 +166,7 @@ unsloth studio -H 0.0.0.0 -p 8888 ``` Then to update : ```bash -unsloth studio update --local +unsloth studio update ``` #### Nightly: MacOS, Linux, WSL: diff --git a/install.ps1 b/install.ps1 index f1f16d818d..0c36046195 100644 --- a/install.ps1 +++ b/install.ps1 @@ -6,6 +6,7 @@ function Install-UnslothStudio { $ErrorActionPreference = "Stop" + $script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq "1") # ── Parse flags ── $StudioLocalInstall = $false @@ -17,6 +18,8 @@ function Install-UnslothStudio { switch ($argList[$i]) { "--local" { $StudioLocalInstall = $true } "--no-torch" { $SkipTorch = $true } + "--verbose" { $script:UnslothVerbose = $true } + "-v" { $script:UnslothVerbose = $true } "--package" { $i++ if ($i -ge $argList.Count) { @@ -27,6 +30,12 @@ function Install-UnslothStudio { } } } + # Propagate to child processes so they also respect verbose mode. + # Process-scoped -- does not persist. + if ($script:UnslothVerbose) { + $env:UNSLOTH_VERBOSE = '1' + } + if ($StudioLocalInstall) { $RepoRoot = (Resolve-Path (Split-Path -Parent $PSCommandPath)).Path if (-not (Test-Path (Join-Path $RepoRoot "pyproject.toml"))) { @@ -39,10 +48,55 @@ function Install-UnslothStudio { $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" $VenvDir = Join-Path $StudioHome "unsloth_studio" + $Rule = [string]::new([char]0x2500, 52) + $Sloth = [char]::ConvertFromUtf32(0x1F9A5) + + function Enable-StudioVirtualTerminal { + if ($env:NO_COLOR) { return $false } + try { + if (-not ("StudioVT.Native" -as [type])) { + Add-Type -Namespace StudioVT -Name Native -MemberDefinition @' +[DllImport("kernel32.dll")] public static extern IntPtr GetStdHandle(int nStdHandle); +[DllImport("kernel32.dll")] public static extern bool GetConsoleMode(IntPtr h, out uint m); +[DllImport("kernel32.dll")] public static extern bool SetConsoleMode(IntPtr h, uint m); +'@ -ErrorAction Stop + } + $h = [StudioVT.Native]::GetStdHandle(-11) + [uint32]$mode = 0 + if (-not [StudioVT.Native]::GetConsoleMode($h, [ref]$mode)) { return $false } + $mode = $mode -bor 0x0004 + return [StudioVT.Native]::SetConsoleMode($h, $mode) + } catch { + return $false + } + } + $script:StudioVtOk = Enable-StudioVirtualTerminal + + function Get-StudioAnsi { + param( + [Parameter(Mandatory = $true)] + [ValidateSet('Title', 'Dim', 'Ok', 'Warn', 'Err', 'Reset')] + [string]$Kind + ) + $e = [char]27 + switch ($Kind) { + 'Title' { return "${e}[38;5;150m" } + 'Dim' { return "${e}[38;5;245m" } + 'Ok' { return "${e}[38;5;108m" } + 'Warn' { return "${e}[38;5;136m" } + 'Err' { return "${e}[91m" } + 'Reset' { return "${e}[0m" } + } + } + Write-Host "" - Write-Host "=========================================" - Write-Host " Unsloth Studio Installer (Windows)" - Write-Host "=========================================" + if ($script:StudioVtOk -and -not $env:NO_COLOR) { + Write-Host (" " + (Get-StudioAnsi Title) + $Sloth + " Unsloth Studio Installer (Windows)" + (Get-StudioAnsi Reset)) + Write-Host (" {0}{1}{2}" -f (Get-StudioAnsi Dim), $Rule, (Get-StudioAnsi Reset)) + } else { + Write-Host (" {0} Unsloth Studio Installer (Windows)" -f $Sloth) -ForegroundColor DarkGreen + Write-Host " $Rule" -ForegroundColor DarkGray + } Write-Host "" # ── Helper: refresh PATH from registry (deduplicating entries) ── @@ -62,13 +116,96 @@ function Install-UnslothStudio { $env:Path = $unique -join ";" } + function step { + param( + [Parameter(Mandatory = $true)][string]$Label, + [Parameter(Mandatory = $true)][string]$Value, + [string]$Color = "Green" + ) + if ($script:StudioVtOk -and -not $env:NO_COLOR) { + $dim = Get-StudioAnsi Dim + $rst = Get-StudioAnsi Reset + $val = switch ($Color) { + 'Green' { Get-StudioAnsi Ok } + 'Yellow' { Get-StudioAnsi Warn } + 'Red' { Get-StudioAnsi Err } + 'DarkGray' { Get-StudioAnsi Dim } + default { Get-StudioAnsi Ok } + } + $padded = if ($Label.Length -ge 15) { $Label.Substring(0, 15) } else { $Label.PadRight(15) } + Write-Host (" {0}{1}{2}{3}{4}{2}" -f $dim, $padded, $rst, $val, $Value) + } else { + $padded = if ($Label.Length -ge 15) { $Label.Substring(0, 15) } else { $Label.PadRight(15) } + Write-Host (" {0}" -f $padded) -NoNewline -ForegroundColor DarkGray + $fc = switch ($Color) { + 'Green' { 'DarkGreen' } + 'Yellow' { 'Yellow' } + 'Red' { 'Red' } + 'DarkGray' { 'DarkGray' } + default { 'DarkGreen' } + } + Write-Host $Value -ForegroundColor $fc + } + } + + function substep { + param( + [Parameter(Mandatory = $true)][string]$Message, + [string]$Color = "DarkGray" + ) + if ($script:StudioVtOk -and -not $env:NO_COLOR) { + $msgCol = switch ($Color) { + 'Yellow' { (Get-StudioAnsi Warn) } + 'Red' { (Get-StudioAnsi Err) } + default { (Get-StudioAnsi Dim) } + } + $pad = "".PadRight(15) + Write-Host (" {0}{1}{2}{3}" -f $msgCol, $pad, $Message, (Get-StudioAnsi Reset)) + } else { + $fc = switch ($Color) { + 'Yellow' { 'Yellow' } + 'Red' { 'Red' } + default { 'DarkGray' } + } + Write-Host (" {0,-15}{1}" -f "", $Message) -ForegroundColor $fc + } + } + + # Run native commands quietly by default to match install.sh behavior. + # Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1. + function Invoke-InstallCommand { + param( + [Parameter(Mandatory = $true)][ScriptBlock]$Command + ) + $prevEap = $ErrorActionPreference + $ErrorActionPreference = "Continue" + try { + # 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). + & $Command 2>&1 | Out-Host + } else { + $output = & $Command 2>&1 | Out-String + if ($LASTEXITCODE -ne 0) { + Write-Host $output -ForegroundColor Red + } + } + return [int]$LASTEXITCODE + } finally { + $ErrorActionPreference = $prevEap + } + } + function New-StudioShortcuts { param( [Parameter(Mandatory = $true)][string]$UnslothExePath ) if (-not (Test-Path $UnslothExePath)) { - Write-Host "[WARN] Cannot create shortcuts: unsloth.exe not found at $UnslothExePath" -ForegroundColor Yellow + substep "cannot create shortcuts, unsloth.exe not found at $UnslothExePath" "Yellow" return } try { @@ -81,7 +218,7 @@ function Install-UnslothStudio { $localAppDataDir = $env:LOCALAPPDATA if (-not $localAppDataDir -or [string]::IsNullOrWhiteSpace($localAppDataDir)) { - Write-Host "[WARN] LOCALAPPDATA path unavailable; skipped shortcut creation" -ForegroundColor Yellow + substep "LOCALAPPDATA path unavailable; skipped shortcut creation" "Yellow" return } $appDir = Join-Path $localAppDataDir "Unsloth Studio" @@ -104,10 +241,10 @@ function Install-UnslothStudio { $null } if (-not $desktopLink) { - Write-Host "[WARN] Desktop path unavailable; skipped desktop shortcut creation" -ForegroundColor Yellow + substep "Desktop path unavailable; skipped desktop shortcut creation" "Yellow" } if (-not $startMenuLink) { - Write-Host "[WARN] APPDATA/Start Menu path unavailable; skipped Start menu shortcut creation" -ForegroundColor Yellow + substep "APPDATA/Start Menu path unavailable; skipped Start menu shortcut creation" "Yellow" } $iconPath = Join-Path $appDir "unsloth.ico" $bundledIcon = $null @@ -362,27 +499,27 @@ shell.Run cmd, 0, False $shortcut.Save() $createdShortcutCount++ } catch { - Write-Host "[WARN] Could not create shortcut at ${linkPath}: $($_.Exception.Message)" -ForegroundColor Yellow + substep "could not create shortcut at ${linkPath}: $($_.Exception.Message)" "Yellow" } } if ($createdShortcutCount -gt 0) { - Write-Host "[OK] Created Unsloth Studio shortcut(s): $createdShortcutCount" -ForegroundColor Green + substep "Created Unsloth Studio shortcut" } else { - Write-Host "[WARN] No Unsloth Studio shortcuts were created" -ForegroundColor Yellow + substep "no Unsloth Studio shortcuts were created" "Yellow" } } catch { - Write-Host "[WARN] Shortcut creation unavailable: $($_.Exception.Message)" -ForegroundColor Yellow + substep "shortcut creation unavailable: $($_.Exception.Message)" "Yellow" } } catch { - Write-Host "[WARN] Shortcut setup failed; skipping shortcuts: $($_.Exception.Message)" -ForegroundColor Yellow + substep "shortcut setup failed; skipping shortcuts: $($_.Exception.Message)" "Yellow" } } # ── Check winget ── if (-not (Get-Command winget -ErrorAction SilentlyContinue)) { - Write-Host "Error: winget is not available." -ForegroundColor Red - Write-Host " Install it from https://aka.ms/getwinget" -ForegroundColor Yellow - Write-Host " or install Python $PythonVersion and uv manually, then re-run." -ForegroundColor Yellow + step "winget" "not available" "Red" + substep "Install it from https://aka.ms/getwinget" "Yellow" + substep "or install Python $PythonVersion and uv manually, then re-run." "Yellow" return } @@ -460,10 +597,10 @@ shell.Run cmd, 0, False # Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. $DetectedPython = Find-CompatiblePython if ($DetectedPython) { - Write-Host "==> Python already installed: Python $($DetectedPython.Version)" + step "python" "Python $($DetectedPython.Version) already installed" } if (-not $DetectedPython) { - Write-Host "==> Installing Python ${PythonVersion}..." + substep "installing Python ${PythonVersion}..." $pythonPackageId = "Python.Python.$PythonVersion" # Temporarily lower ErrorActionPreference so that winget stderr # (progress bars, warnings) does not become a terminating error @@ -485,7 +622,7 @@ shell.Run cmd, 0, False # 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). - Write-Host " Python not found on PATH after winget. Retrying with --force..." + substep "Python not found on PATH after winget. Retrying with --force..." "Yellow" $ErrorActionPreference = "Continue" try { winget install -e --id $pythonPackageId --accept-package-agreements --accept-source-agreements --force @@ -507,7 +644,7 @@ shell.Run cmd, 0, False # ── Install uv if not present ── if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - Write-Host "==> Installing uv package manager..." + substep "installing uv package manager..." $prevEAP = $ErrorActionPreference $ErrorActionPreference = "Continue" try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {} @@ -515,15 +652,15 @@ shell.Run cmd, 0, False Refresh-SessionPath # Fallback: if winget didn't put uv on PATH, try the PowerShell installer if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - Write-Host " Trying alternative uv installer..." + substep "trying alternative uv installer..." "Yellow" powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex" Refresh-SessionPath } } if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - Write-Host "Error: uv could not be installed." -ForegroundColor Red - Write-Host " Install it from https://docs.astral.sh/uv/" -ForegroundColor Yellow + step "uv" "could not be installed" "Red" + substep "Install it from https://docs.astral.sh/uv/" "Yellow" return } @@ -539,13 +676,13 @@ shell.Run cmd, 0, False if (Test-Path $VenvPython) { # New layout already exists -- nuke for fresh install - Write-Host "==> Removing existing environment for fresh install..." + substep "removing existing environment for fresh install..." Remove-Item -Recurse -Force $VenvDir } elseif (Test-Path (Join-Path $StudioHome ".venv\Scripts\python.exe")) { # Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating $OldVenv = Join-Path $StudioHome ".venv" $OldPy = Join-Path $OldVenv "Scripts\python.exe" - Write-Host "==> Found legacy Studio environment, validating..." + substep "found legacy Studio environment, validating..." $prevEAP2 = $ErrorActionPreference $ErrorActionPreference = "Continue" try { @@ -554,32 +691,34 @@ shell.Run cmd, 0, False } catch { $torchOk = $false } $ErrorActionPreference = $prevEAP2 if ($torchOk) { - Write-Host " Legacy environment is healthy -- migrating..." + substep "legacy environment is healthy -- migrating..." Move-Item -Path $OldVenv -Destination $VenvDir -Force - Write-Host " Moved .venv -> unsloth_studio" + substep "moved .venv -> unsloth_studio" $_Migrated = $true } else { - Write-Host " Legacy environment failed validation -- creating fresh environment" + substep "legacy environment failed validation -- creating fresh environment" "Yellow" Remove-Item -Recurse -Force $OldVenv -ErrorAction SilentlyContinue } } elseif (Test-Path (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) { # CWD-relative venv from old install.ps1 -- migrate to absolute path $CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio" - Write-Host "==> Found CWD-relative Studio environment, migrating to $VenvDir..." + substep "found CWD-relative Studio environment, migrating to $VenvDir..." Move-Item -Path $CwdVenv -Destination $VenvDir -Force - Write-Host " Moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio" + substep "moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio" $_Migrated = $true } if (-not (Test-Path $VenvPython)) { - Write-Host "==> Creating Python $($DetectedPython.Version) virtual environment ($VenvDir)..." - uv venv $VenvDir --python "$($DetectedPython.Path)" - if ($LASTEXITCODE -ne 0) { - Write-Host "[ERROR] Failed to create virtual environment (exit code $LASTEXITCODE)" -ForegroundColor Red + step "venv" "creating Python $($DetectedPython.Version) virtual environment" + substep "$VenvDir" + $venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" } + if ($venvExit -ne 0) { + Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red return } } else { - Write-Host "==> Using migrated environment at $VenvDir" + step "venv" "using migrated environment" + substep "$VenvDir" } # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ── @@ -588,7 +727,7 @@ shell.Run cmd, 0, False try { $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue if ($nvSmiCmd) { - & $nvSmiCmd.Source 2>&1 | Out-Null + & $nvSmiCmd.Source *> $null if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source } } } catch {} @@ -599,18 +738,18 @@ shell.Run cmd, 0, False )) { if (Test-Path $p) { try { - & $p 2>&1 | Out-Null + & $p *> $null if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break } } catch {} } } } if ($HasNvidiaSmi) { - Write-Host "[OK] NVIDIA GPU detected" -ForegroundColor Green + step "gpu" "NVIDIA GPU detected" } else { - Write-Host "[WARN] No NVIDIA GPU detected. Studio will run in chat-only (GGUF) mode." -ForegroundColor Yellow - Write-Host " Training and GPU inference require an NVIDIA GPU with drivers installed." -ForegroundColor Yellow - Write-Host " https://www.nvidia.com/Download/index.aspx" -ForegroundColor Yellow + step "gpu" "none (chat-only / GGUF)" "Yellow" + substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow" + substep "https://www.nvidia.com/Download/index.aspx" "Yellow" } # ── Choose the correct PyTorch index URL based on driver CUDA version ── @@ -630,7 +769,7 @@ shell.Run cmd, 0, False return "$baseUrl/cpu" } } catch {} - Write-Host "[WARN] Could not determine CUDA version from nvidia-smi, defaulting to cu126" -ForegroundColor Yellow + substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow" return "$baseUrl/cu126" } $TorchIndexUrl = Get-TorchIndexUrl @@ -677,74 +816,101 @@ shell.Run cmd, 0, False if ($_Migrated) { # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state # in the new venv location, while preserving existing torch/CUDA - Write-Host "==> Upgrading unsloth in migrated environment..." + 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. - uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo - $NoTorchReq = Find-NoTorchRuntimeFile - if ($NoTorchReq) { - uv pip install --python $VenvPython --no-deps -r $NoTorchReq + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo } + if ($baseInstallExit -eq 0) { + $NoTorchReq = Find-NoTorchRuntimeFile + if ($NoTorchReq) { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } + } } } else { - uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo } + } + if ($baseInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red + return } if ($StudioLocalInstall) { - Write-Host "==> Overlaying local repo (editable)..." - uv pip install --python $VenvPython -e $RepoRoot --no-deps + substep "overlaying local repo (editable)..." + $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + if ($overlayExit -ne 0) { + Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red + return + } } } elseif ($TorchIndexUrl) { if ($SkipTorch) { - Write-Host "==> Skipping PyTorch (--no-torch flag set)." + substep "skipping PyTorch (--no-torch flag set)." "Yellow" } else { - Write-Host "==> Installing PyTorch ($TorchIndexUrl)..." - uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl - if ($LASTEXITCODE -ne 0) { - Write-Host "[ERROR] Failed to install PyTorch (exit code $LASTEXITCODE)" -ForegroundColor Red + substep "installing PyTorch ($TorchIndexUrl)..." + $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + if ($torchInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return } } - Write-Host "==> Installing unsloth (this may take a few minutes)..." + 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. - uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo - $NoTorchReq = Find-NoTorchRuntimeFile - if ($NoTorchReq) { - uv pip install --python $VenvPython --no-deps -r $NoTorchReq - } - if ($StudioLocalInstall) { - Write-Host "==> Overlaying local repo (editable)..." - uv pip install --python $VenvPython -e $RepoRoot --no-deps + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo } + if ($baseInstallExit -eq 0) { + $NoTorchReq = Find-NoTorchRuntimeFile + if ($NoTorchReq) { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } + } } } elseif ($StudioLocalInstall) { - uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo - Write-Host "==> Overlaying local repo (editable)..." - uv pip install --python $VenvPython -e $RepoRoot --no-deps + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo } } else { - uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" + $baseInstallExit = Invoke-InstallCommand { 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 + return + } + + if ($StudioLocalInstall) { + substep "overlaying local repo (editable)..." + $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + if ($overlayExit -ne 0) { + Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red + return + } } } else { # Fallback: GPU detection failed to produce a URL -- let uv resolve torch - Write-Host "==> Installing unsloth (this may take a few minutes)..." + substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto - Write-Host "==> Overlaying local repo (editable)..." - uv pip install --python $VenvPython -e $RepoRoot --no-deps + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto } + if ($baseInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red + return + } + substep "overlaying local repo (editable)..." + $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + if ($overlayExit -ne 0) { + Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red + return + } } else { - uv pip install --python $VenvPython "$PackageName" --torch-backend=auto + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "$PackageName" --torch-backend=auto } + if ($baseInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red + return + } } } - if ($LASTEXITCODE -ne 0) { - Write-Host "[ERROR] Failed to install unsloth (exit code $LASTEXITCODE)" -ForegroundColor Red - return - } # ── Run studio setup ── # setup.ps1 will handle installing Git, CMake, Visual Studio Build Tools, # CUDA Toolkit, Node.js, and other dependencies automatically via winget. - Write-Host "==> Running unsloth studio setup..." + step "setup" "running unsloth studio setup..." $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe" if (-not (Test-Path $UnslothExe)) { Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red @@ -754,17 +920,27 @@ shell.Run cmd, 0, False return } # Tell setup.ps1 to skip base package installation (install.ps1 already did it) - # Tell setup.ps1 to skip base package installation (install.ps1 already did it) $env:SKIP_STUDIO_BASE = "1" $env:STUDIO_PACKAGE_NAME = $PackageName $env:UNSLOTH_NO_TORCH = if ($SkipTorch) { "true" } else { "false" } + # Always set STUDIO_LOCAL_INSTALL explicitly to avoid stale values from + # a previous --local run in the same PowerShell session. if ($StudioLocalInstall) { $env:STUDIO_LOCAL_INSTALL = "1" $env:STUDIO_LOCAL_REPO = $RepoRoot + } else { + $env:STUDIO_LOCAL_INSTALL = "0" + Remove-Item Env:STUDIO_LOCAL_REPO -ErrorAction SilentlyContinue } - & $UnslothExe studio setup - if ($LASTEXITCODE -ne 0) { - Write-Host "[ERROR] unsloth studio setup failed (exit code $LASTEXITCODE)" -ForegroundColor Red + # 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. + $studioArgs = @('studio', 'setup') + if ($script:UnslothVerbose) { $studioArgs += '--verbose' } + & $UnslothExe @studioArgs + $setupExit = $LASTEXITCODE + if ($setupExit -ne 0) { + Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red return } @@ -780,27 +956,18 @@ shell.Run cmd, 0, False [System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir", "User") } Refresh-SessionPath - Write-Host "[OK] Added unsloth to PATH" -ForegroundColor Green + step "path" "added unsloth to PATH" } - Write-Host "" - Write-Host "=========================================" - Write-Host " Unsloth Studio installed!" - Write-Host "=========================================" - Write-Host "" - # Launch studio automatically in interactive terminals; # in non-interactive environments (CI, Docker) just print instructions. $IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected) if ($IsInteractive) { - Write-Host "==> Launching Unsloth Studio..." - Write-Host "" & $UnslothExe studio -H 0.0.0.0 -p 8888 } else { - Write-Host " To launch, run:" - Write-Host "" - Write-Host " & `"$VenvDir\Scripts\Activate.ps1`"" - Write-Host " unsloth studio -H 0.0.0.0 -p 8888" + step "launch" "manual commands:" + substep "& `"$VenvDir\Scripts\Activate.ps1`"" + substep "unsloth studio -H 0.0.0.0 -p 8888" Write-Host "" } } diff --git a/install.sh b/install.sh index b50256ea5a..9ea80bc161 100755 --- a/install.sh +++ b/install.sh @@ -8,11 +8,36 @@ # Usage (py): ./install.sh --python 3.12 (override auto-detected Python version) set -e +# ── Output style (aligned with studio/setup.sh) ── +RULE="" +_rule_i=0 +while [ "$_rule_i" -lt 52 ]; do + RULE="${RULE}─" + _rule_i=$((_rule_i + 1)) +done +if [ -n "${NO_COLOR:-}" ]; then + C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST= +elif [ -t 1 ] || [ -n "${FORCE_COLOR:-}" ]; then + _ESC="$(printf '\033')" + C_TITLE="${_ESC}[38;5;150m" + C_DIM="${_ESC}[38;5;245m" + C_OK="${_ESC}[38;5;108m" + C_WARN="${_ESC}[38;5;136m" + C_ERR="${_ESC}[91m" + C_RST="${_ESC}[0m" +else + C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST= +fi + +step() { printf " ${C_DIM}%-15.15s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; } +substep() { printf " ${C_DIM}%-15s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; } + # ── Parse flags ── STUDIO_LOCAL_INSTALL=false PACKAGE_NAME="unsloth" _USER_PYTHON="" _NO_TORCH_FLAG=false +_VERBOSE=false _next_is_package=false _next_is_python=false for arg in "$@"; do @@ -31,9 +56,44 @@ for arg in "$@"; do --package) _next_is_package=true ;; --python) _next_is_python=true ;; --no-torch) _NO_TORCH_FLAG=true ;; + --verbose|-v) _VERBOSE=true ;; esac done +if [ "$_VERBOSE" = true ]; then + export UNSLOTH_VERBOSE=1 +fi + +_is_verbose() { + [ "${UNSLOTH_VERBOSE:-0}" = "1" ] +} + +run_maybe_quiet() { + if _is_verbose; then + "$@" + else + "$@" > /dev/null 2>&1 + fi +} + +run_install_cmd() { + _label="$1" + shift + if _is_verbose; then + "$@" && return 0 + _rc=$? + step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 + return "$_rc" + fi + _log=$(mktemp) + "$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; } + _rc=$? + step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 + cat "$_log" >&2 + rm -f "$_log" + return $_rc +} + if [ "$_next_is_package" = true ]; then echo "❌ ERROR: --package requires an argument." >&2 exit 1 @@ -643,14 +703,13 @@ WSLPS1_EOF fi if [ "$_css_created" -eq 1 ]; then - echo "[OK] Created Unsloth Studio shortcut(s)" + substep "Created Unsloth Studio shortcut" fi } echo "" -echo "=========================================" -echo " Unsloth Studio Installer" -echo "=========================================" +printf " ${C_TITLE}%s${C_RST}\n" "🦥 Unsloth Studio Installer" +printf " ${C_DIM}%s${C_RST}\n" "$RULE" echo "" # ── Detect platform ── @@ -660,7 +719,7 @@ if [ "$(uname)" = "Darwin" ]; then elif grep -qi microsoft /proc/version 2>/dev/null; then OS="wsl" fi -echo "==> Platform: $OS" +step "platform" "$OS" # ── Architecture detection & Python version ── _ARCH=$(uname -m) @@ -740,8 +799,8 @@ MISSING=$(echo "$MISSING" | sed 's/^ *//') if [ -n "$MISSING" ]; then echo "" - echo "==> Unsloth Studio needs these packages: $MISSING" - echo " These are needed to build the GGUF inference engine." + step "deps" "missing: $MISSING" "$C_WARN" + substep "These are needed to build the GGUF inference engine." case "$OS" in macos) @@ -766,7 +825,7 @@ if [ -n "$MISSING" ]; then esac echo "" else - echo "==> All system dependencies found." + step "deps" "all system dependencies found" fi # ── Install uv ── @@ -812,10 +871,10 @@ _uv_version_ok() { } if ! command -v uv >/dev/null 2>&1 || ! _uv_version_ok uv; then - echo "==> Installing uv package manager..." + substep "installing uv package manager..." _uv_tmp=$(mktemp) download "https://astral.sh/uv/install.sh" "$_uv_tmp" - sh "$_uv_tmp" Found legacy Studio environment, validating..." + substep "found legacy Studio environment, validating..." if "$STUDIO_HOME/.venv/bin/python" -c " import torch device = 'cuda' if torch.cuda.is_available() else 'cpu' @@ -866,8 +925,9 @@ if [ "$SKIP_TORCH" = true ] && [ "$MAC_INTEL" = true ] && [ -z "$_USER_PYTHON" ] fi if [ ! -x "$VENV_DIR/bin/python" ]; then - echo "==> Creating Python ${PYTHON_VERSION} virtual environment (${VENV_DIR})..." - uv venv "$VENV_DIR" --python "$PYTHON_VERSION" + step "venv" "creating Python ${PYTHON_VERSION} virtual environment" + substep "$VENV_DIR" + run_install_cmd "create venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION" fi # Guard against Python 3.13.8 torch import bug on Apple Silicon @@ -880,12 +940,13 @@ if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then echo " Recreating venv with Python 3.12..." rm -rf "$VENV_DIR" PYTHON_VERSION="3.12" - uv venv "$VENV_DIR" --python "$PYTHON_VERSION" + run_install_cmd "recreate venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION" fi fi if [ -x "$VENV_DIR/bin/python" ]; then - echo "==> Using environment at ${VENV_DIR}" + step "venv" "using environment" + substep "${VENV_DIR}" fi # ── Resolve repo root (for --local installs) ── @@ -960,71 +1021,71 @@ _VENV_PY="$VENV_DIR/bin/python" if [ "$_MIGRATED" = true ]; then # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state # in the new venv location, while preserving existing torch/CUDA - echo "==> Upgrading unsloth in migrated environment..." + 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. - uv pip install --python "$_VENV_PY" --no-deps \ + run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ "unsloth>=2026.3.16" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then - uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" + run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" fi else - uv pip install --python "$_VENV_PY" \ + run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ "unsloth>=2026.3.16" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - echo "==> Overlaying local repo (editable)..." - uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps + substep "overlaying local repo (editable)..." + run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps fi elif [ -n "$TORCH_INDEX_URL" ]; then # Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac) if [ "$SKIP_TORCH" = true ]; then - echo "==> Skipping PyTorch (--no-torch or Intel Mac x86_64)." + substep "skipping PyTorch (--no-torch or Intel Mac x86_64)." "$C_WARN" else - echo "==> Installing PyTorch ($TORCH_INDEX_URL)..." - uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \ + substep "installing PyTorch ($TORCH_INDEX_URL)..." + run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" fi # Fresh: Step 2 - install unsloth, preserving pre-installed torch - echo "==> Installing unsloth (this may take a few minutes)..." + substep "installing unsloth (this may take a few minutes)..." if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - uv pip install --python "$_VENV_PY" --no-deps \ + run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ "unsloth>=2026.3.16" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then - uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" + run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - echo "==> Overlaying local repo (editable)..." - uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps + substep "overlaying local repo (editable)..." + run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then - uv pip install --python "$_VENV_PY" \ + run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo - echo "==> Overlaying local repo (editable)..." - uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps + substep "overlaying local repo (editable)..." + run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps else - uv pip install --python "$_VENV_PY" \ + run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \ --upgrade-package unsloth "$PACKAGE_NAME" fi else # Fallback: GPU detection failed to produce a URL -- let uv resolve torch - echo "==> Installing unsloth (this may take a few minutes)..." + substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto - echo "==> Overlaying local repo (editable)..." - uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto + substep "overlaying local repo (editable)..." + run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps else - uv pip install --python "$_VENV_PY" "$PACKAGE_NAME" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "$PACKAGE_NAME" --torch-backend=auto fi fi @@ -1059,19 +1120,33 @@ if [ -n "$VENV_ABS_BIN" ]; then export PATH="$VENV_ABS_BIN:$PATH" fi -echo "==> Running unsloth setup..." +if ! command -v bash >/dev/null 2>&1; then + step "setup" "bash is required to run studio setup" "$C_ERR" + substep "Please install bash and re-run install.sh" + exit 1 +fi + +step "setup" "running unsloth studio update..." +# install.sh already installs base packages (unsloth + unsloth-zoo) and +# no-torch-runtime.txt above, so tell install_python_stack.py to skip +# the base step to avoid redundant reinstallation. +_SKIP_BASE=1 +# Run setup.sh outside set -e so that a llama.cpp build failure (exit 1) +# does not skip PATH setup, shortcuts, and launch below. We capture the +# exit code and propagate it after post-install steps finish. +_SETUP_EXIT=0 if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - SKIP_STUDIO_BASE=1 \ + SKIP_STUDIO_BASE="$_SKIP_BASE" \ STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \ STUDIO_LOCAL_INSTALL=1 \ STUDIO_LOCAL_REPO="$_REPO_ROOT" \ UNSLOTH_NO_TORCH="$SKIP_TORCH" \ - bash "$SETUP_SH" > "$_SHELL_PROFILE" echo '# Added by Unsloth installer' >> "$_SHELL_PROFILE" echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$_SHELL_PROFILE" - echo "==> Added ~/.local/bin to PATH in $_SHELL_PROFILE" + step "path" "added ~/.local/bin to PATH in $_SHELL_PROFILE" fi fi export PATH="$_LOCAL_BIN:$PATH" @@ -1105,17 +1180,30 @@ esac create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS" +# If setup.sh failed, report and exit now. +# PATH and shortcuts are already set up so the user can fix and retry. +if [ "$_SETUP_EXIT" -ne 0 ]; then + echo "" + step "error" "studio setup failed (exit code $_SETUP_EXIT)" "$C_ERR" + substep "Check the output above for details, then re-run:" + if [ "$STUDIO_LOCAL_INSTALL" = true ]; then + substep " unsloth studio update --local" + else + substep " unsloth studio update" + fi + echo "" + exit "$_SETUP_EXIT" +fi + echo "" -echo "=========================================" -echo " Unsloth Studio installed!" -echo "=========================================" +printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!" +printf " ${C_DIM}%s${C_RST}\n" "$RULE" echo "" # Launch studio automatically in interactive terminals; # in non-interactive environments (Docker, CI, cloud-init) just print instructions. if [ -t 1 ]; then - echo "==> Launching Unsloth Studio..." - echo "" + step "launch" "starting Unsloth Studio..." "$VENV_DIR/bin/unsloth" studio -H 0.0.0.0 -p 8888 _LAUNCH_EXIT=$? if [ "$_LAUNCH_EXIT" -ne 0 ] && [ "$_MIGRATED" = true ]; then @@ -1130,13 +1218,10 @@ if [ -t 1 ]; then fi exit "$_LAUNCH_EXIT" else - echo " To launch, run:" - echo "" - echo " unsloth studio -H 0.0.0.0 -p 8888" - echo "" - echo " Or activate the environment first:" - echo "" - echo " source ${VENV_DIR}/bin/activate" - echo " unsloth studio -H 0.0.0.0 -p 8888" + step "launch" "manual commands:" + substep "unsloth studio -H 0.0.0.0 -p 8888" + substep "or activate env first:" + substep "source ${VENV_DIR}/bin/activate" + substep "unsloth studio -H 0.0.0.0 -p 8888" echo "" fi diff --git a/pyproject.toml b/pyproject.toml index e2173d6811..b06131021a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,8 @@ studio = [ ] [tool.setuptools.packages.find] -exclude = ["images*", "tests*", "kernels/moe*"] +include = ["unsloth*", "unsloth_cli*", "studio", "studio.backend*"] +exclude = ["images*", "tests*", "*.node_modules", "*.node_modules.*"] [project.optional-dependencies] triton = [ diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 1a265690ff..ddd485525d 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -18,7 +18,14 @@ from typing import Optional, Union, Generator, Tuple from utils.models import ModelConfig, get_base_model_from_lora from utils.paths import is_model_cached from utils.utils import format_error_message -from utils.hardware import get_device, clear_gpu_cache, log_gpu_memory +from utils.hardware import ( + get_device, + clear_gpu_cache, + log_gpu_memory, + get_device_map, + raise_if_offloaded, + get_visible_gpu_count, +) from core.inference.audio_codecs import AudioCodecManager from io import StringIO import structlog @@ -241,6 +248,7 @@ class InferenceBackend: load_in_4bit: bool = True, hf_token: Optional[str] = None, trust_remote_code: bool = False, + gpu_ids: Optional[list[int]] = None, ) -> bool: """ Load any model: base, LoRA adapter, text, or vision. @@ -260,6 +268,10 @@ class InferenceBackend: return False self.loading_models.add(model_name) + device_map = get_device_map(gpu_ids, for_inference = True) + logger.info( + f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)" + ) self.models[model_name] = { "is_vision": config.is_vision, @@ -290,6 +302,7 @@ class InferenceBackend: config.path, auto_model = CsmForConditionalGeneration, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -325,6 +338,7 @@ class InferenceBackend: config.path, dtype = torch.float32, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -345,6 +359,7 @@ class InferenceBackend: llm_path, dtype = torch.float32, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -361,6 +376,7 @@ class InferenceBackend: config.path, max_seq_length = max_seq_length, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -378,6 +394,7 @@ class InferenceBackend: whisper_language = "English", whisper_task = "transcribe", load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -405,6 +422,7 @@ class InferenceBackend: model_name = config.path, max_seq_length = max_seq_length, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -420,6 +438,11 @@ class InferenceBackend: audio_type, self.device, model_repo_path = model_repo_path ) + # Reject CPU/disk offload for audio models too + raise_if_offloaded( + self.models[model_name]["model"], device_map, "Inference" + ) + self.active_model_name = model_name self.loading_models.discard(model_name) logger.info(f"Successfully loaded audio model: {model_name}") @@ -441,6 +464,7 @@ class InferenceBackend: max_seq_length = max_seq_length, dtype = dtype, load_in_4bit = load_in_4bit, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -497,6 +521,7 @@ class InferenceBackend: max_seq_length = max_seq_length, dtype = dtype, load_in_4bit = load_in_4bit, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -507,6 +532,10 @@ class InferenceBackend: self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = tokenizer + raise_if_offloaded( + self.models[model_name]["model"], device_map, "Inference" + ) + # Load chat template info self._load_chat_template_info(model_name) @@ -615,6 +644,7 @@ class InferenceBackend: dtype = None, load_in_4bit: bool = True, hf_token: Optional[str] = None, + gpu_ids: Optional[list[int]] = None, ) -> Tuple[bool, Optional[str], Optional[str]]: """ Final Corrected Version: @@ -639,7 +669,12 @@ class InferenceBackend: base_model_name, None, is_lora = False ) if not self.load_model( - base_config, max_seq_length, dtype, load_in_4bit, hf_token + base_config, + max_seq_length, + dtype, + load_in_4bit, + hf_token, + gpu_ids = gpu_ids, ): return False, None, None @@ -1037,12 +1072,12 @@ class InferenceBackend: input_text, add_special_tokens = False, return_tensors = "pt", - ).to(self.device) + ).to(model.device) else: # Text-only for vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to( - self.device + model.device ) # Stream with TextIteratorStreamer + background thread @@ -1182,7 +1217,7 @@ class InferenceBackend: return_dict = True, return_tensors = "pt", truncation = False, - ).to(self.device) + ).to(model.device) try: from transformers import TextIteratorStreamer diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 05e038dbb7..f5361a6e8c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -48,6 +48,7 @@ class LlamaCppBackend: self._healthy = False self._context_length: Optional[int] = None self._effective_context_length: Optional[int] = None + self._max_context_length: Optional[int] = None self._chat_template: Optional[str] = None self._supports_reasoning: bool = False self._reasoning_always_on: bool = False @@ -100,6 +101,11 @@ class LlamaCppBackend: """Return the effective context length the server is running at.""" return self._effective_context_length or self._context_length + @property + def max_context_length(self) -> Optional[int]: + """Return the maximum context currently available on this hardware.""" + return self._max_context_length or self._context_length + @property def chat_template(self) -> Optional[str]: return self._chat_template @@ -287,7 +293,8 @@ class LlamaCppBackend: continue gpus.append((idx, free_mib)) return gpus - except Exception: + except Exception as e: + logger.debug(f"Failed to query GPU free memory via nvidia-smi: {e}") return [] @staticmethod @@ -328,6 +335,11 @@ class LlamaCppBackend: return sorted(selected), False # Model is too large even for all GPUs, let --fit handle it + logger.debug( + "Model does not fit in available GPU memory, falling back to --fit", + model_size_mib = round(model_size_mib, 2), + ranked_gpus = ranked, + ) return None, True # ── KV cache VRAM estimation ───────────────────────────────────── @@ -386,6 +398,11 @@ class LlamaCppBackend: If the model weights alone don't fit, returns min_ctx unchanged. """ if not self._can_estimate_kv(): + logger.debug( + "Skipping context fit because KV cache metadata is unavailable", + requested_ctx = requested_ctx, + available_mib = available_mib, + ) return requested_ctx budget_bytes = available_mib * 1024 * 1024 * 0.70 @@ -399,6 +416,12 @@ class LlamaCppBackend: # Model weights alone exceed budget -- can't help by reducing ctx. # Return requested_ctx unchanged; --fit will handle VRAM management. if model_footprint >= budget_bytes: + logger.debug( + "Model footprint exceeds GPU budget before KV cache", + requested_ctx = requested_ctx, + available_mib = available_mib, + model_size_gb = round(model_footprint / (1024**3), 2), + ) return requested_ctx # Binary search for max context that fits @@ -960,7 +983,11 @@ class LlamaCppBackend: self._port = self._find_free_port() - # Select GPU(s) based on model size + estimated KV cache + # Select GPU(s) based on model size + estimated KV cache. + # Seed safe defaults before GPU probing so the except path + # still has valid state to publish. + effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0) + max_available_ctx = self._context_length or effective_ctx try: model_size = self._get_gguf_size_bytes(model_path) gpus = self._get_gpu_free_memory() @@ -975,6 +1002,9 @@ class LlamaCppBackend: else: effective_ctx = 0 original_ctx = effective_ctx + # Default UI ceiling to the model's native context length. + # GPU/VRAM-fit logic below may shrink this if hardware is limited. + max_available_ctx = self._context_length or effective_ctx # Auto-cap context to fit in GPU VRAM and select GPUs. # @@ -993,6 +1023,29 @@ class LlamaCppBackend: explicit_ctx = n_ctx > 0 if gpus and self._can_estimate_kv() and effective_ctx > 0: + # Compute the largest hardware-aware cap from the model's + # native context across all usable GPU subsets (for UI + # bounds), independent of the currently requested context. + native_ctx_for_cap = self._context_length or effective_ctx + if native_ctx_for_cap > 0: + ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) + best_cap = 0 + for n_gpus in range(1, len(ranked_for_cap) + 1): + subset = ranked_for_cap[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = self._fit_context_to_vram( + native_ctx_for_cap, + pool_mib, + model_size, + cache_type_kv, + ) + kv = self._estimate_kv_cache_bytes(capped, cache_type_kv) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * 0.70: + best_cap = max(best_cap, capped) + if best_cap > 0: + max_available_ctx = best_cap + if explicit_ctx: # Try to honor the user's requested context exactly. requested_total = model_size + self._estimate_kv_cache_bytes( @@ -1043,7 +1096,13 @@ class LlamaCppBackend: break elif gpus: - # Can't estimate KV -- fall back to file-size-only check + # Can't estimate KV -- fall back to file-size-only check. + # Without KV estimation we cannot prove a hardware cap, so + # keep the ceiling at the native context (already the default). + logger.debug( + "Falling back to file-size-only GPU selection", + model_size_gb = round(model_size / (1024**3), 2), + ) gpu_indices, use_fit = self._select_gpus(model_size, gpus) if effective_ctx < original_ctx: @@ -1313,6 +1372,11 @@ class LlamaCppBackend: self._effective_context_length = ( effective_ctx if effective_ctx > 0 else self._context_length ) + self._max_context_length = ( + max_available_ctx + if max_available_ctx > 0 + else self._effective_context_length + ) # Wait for llama-server to become healthy if not self._wait_for_health(timeout = 600.0): @@ -1347,6 +1411,7 @@ class LlamaCppBackend: self._healthy = False self._context_length = None self._effective_context_length = None + self._max_context_length = None self._chat_template = None self._supports_reasoning = False self._reasoning_always_on = False diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 6ff7fd2cbf..78cc60e1c6 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -27,6 +27,7 @@ import uuid from io import BytesIO from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union +from utils.hardware import prepare_gpu_selection logger = get_logger(__name__) @@ -571,6 +572,7 @@ class InferenceOrchestrator: load_in_4bit: bool = True, hf_token: Optional[str] = None, trust_remote_code: bool = False, + gpu_ids: Optional[list[int]] = None, ) -> bool: """Load a model for inference. @@ -594,7 +596,16 @@ class InferenceOrchestrator: "hf_token": hf_token or "", "gguf_variant": getattr(config, "gguf_variant", None), "trust_remote_code": trust_remote_code, + "gpu_ids": gpu_ids, } + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + gpu_ids, + model_name = model_name, + hf_token = hf_token, + load_in_4bit = load_in_4bit, + ) + sub_config["resolved_gpu_ids"] = resolved_gpu_ids + sub_config["gpu_selection"] = gpu_selection # Always kill existing subprocess and spawn fresh. # Reusing a subprocess after unsloth patches torch internals diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index afe0ecc458..b3ce43795c 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -29,6 +29,7 @@ from pathlib import Path from typing import Any logger = get_logger(__name__) +from utils.hardware import apply_gpu_ids def _activate_transformers_version(model_name: str) -> None: @@ -178,6 +179,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: load_in_4bit = load_in_4bit, hf_token = hf_token, trust_remote_code = trust_remote_code, + gpu_ids = config.get("resolved_gpu_ids"), ) if success: @@ -501,6 +503,8 @@ def run_inference_process( env = os.getenv("ENVIRONMENT_TYPE", "production"), ) + apply_gpu_ids(config.get("resolved_gpu_ids")) + model_name = config["model_name"] # ── 1. Activate correct transformers version BEFORE any ML imports ── diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 2324916236..9b4a14f09f 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -33,7 +33,14 @@ if sys.platform in ("win32", "darwin"): sys.path.insert(0, _compile_cache) import torch -from utils.hardware import clear_gpu_cache, safe_num_proc, dataset_map_num_proc +from utils.hardware import ( + clear_gpu_cache, + safe_num_proc, + dataset_map_num_proc, + get_device_map, + raise_if_offloaded, + get_visible_gpu_count, +) torch._dynamo.config.recompile_limit = 64 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported @@ -487,6 +494,7 @@ class UnslothTrainer: is_dataset_audio: bool = False, trust_remote_code: bool = False, full_finetuning: bool = False, + gpu_ids: Optional[list[int]] = None, ) -> bool: """Load model for training (supports both text and vision models)""" self.load_in_4bit = load_in_4bit # Store for training_meta.json @@ -624,6 +632,11 @@ class UnslothTrainer: self._update_progress(error = friendly, is_training = False) return False + device_map = get_device_map(gpu_ids) + logger.info( + f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)" + ) + # Branch based on model type if self._audio_type == "csm": # CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False @@ -636,6 +649,7 @@ class UnslothTrainer: dtype = None, auto_model = CsmForConditionalGeneration, load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -651,6 +665,7 @@ class UnslothTrainer: model_name = model_name, dtype = None, load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, auto_model = WhisperForConditionalGeneration, whisper_language = "English", @@ -672,6 +687,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -711,6 +727,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = torch.float32, # Spark-TTS requires float32 load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -725,6 +742,7 @@ class UnslothTrainer: model_name, max_seq_length = max_seq_length, load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -741,6 +759,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -754,6 +773,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, # Auto-detect load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -786,12 +806,15 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, # Auto-detect load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, ) logger.info("Loaded text model") + raise_if_offloaded(self.model, device_map, "Studio training") + if self.should_stop: return False @@ -824,6 +847,7 @@ class UnslothTrainer: is_dataset_audio = is_dataset_audio, trust_remote_code = trust_remote_code, full_finetuning = full_finetuning, + gpu_ids = gpu_ids, ) error_msg = str(e) error_lower = error_msg.lower() diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 4439e4e173..2c8f9a21db 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -28,6 +28,7 @@ from pathlib import Path from typing import Optional, Tuple, Any import matplotlib.pyplot as plt +from utils.hardware import prepare_gpu_selection logger = get_logger(__name__) @@ -185,6 +186,7 @@ class TrainingBackend: "enable_tensorboard": kwargs.get("enable_tensorboard", False), "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), "trust_remote_code": kwargs.get("trust_remote_code", False), + "gpu_ids": kwargs.get("gpu_ids"), } # Derive load_in_4bit from training_type @@ -192,6 +194,22 @@ class TrainingBackend: config["load_in_4bit"] = False # Spawn subprocess — use locals so state is untouched on failure + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + kwargs.get("gpu_ids"), + model_name = config["model_name"], + hf_token = config["hf_token"] or None, + training_type = config["training_type"], + load_in_4bit = config["load_in_4bit"], + batch_size = config.get("batch_size", 4), + max_seq_length = config.get("max_seq_length", 2048), + lora_rank = config.get("lora_r", 16), + target_modules = config.get("target_modules"), + gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), + optimizer = config.get("optim", "adamw_8bit"), + ) + config["resolved_gpu_ids"] = resolved_gpu_ids + config["gpu_selection"] = gpu_selection + from .worker import run_training_process event_queue = _CTX.Queue() diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 891dfca8f7..e68a6c7aee 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -29,6 +29,7 @@ import urllib.error import urllib.request logger = get_logger(__name__) +from utils.hardware import apply_gpu_ids _CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4" @@ -367,6 +368,8 @@ def run_training_process( env = os.getenv("ENVIRONMENT_TYPE", "production"), ) + apply_gpu_ids(config.get("resolved_gpu_ids")) + model_name = config["model_name"] # ── 1. Activate correct transformers version BEFORE any ML imports ── @@ -682,6 +685,7 @@ def run_training_process( is_dataset_image = config.get("is_dataset_image", False), is_dataset_audio = config.get("is_dataset_audio", False), trust_remote_code = config.get("trust_remote_code", False), + gpu_ids = config.get("resolved_gpu_ids"), ) if not success or trainer.should_stop: if trainer.should_stop: diff --git a/studio/backend/main.py b/studio/backend/main.py index d86907f763..e61ae459be 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -68,7 +68,12 @@ from routes import ( ) from auth import storage from auth.authentication import get_current_subject -from utils.hardware import detect_hardware, get_device, DeviceType +from utils.hardware import ( + detect_hardware, + get_device, + DeviceType, + get_backend_visible_gpu_info, +) import utils.hardware.hardware as _hw_module from utils.cache_cleanup import clear_unsloth_compiled_cache @@ -237,69 +242,14 @@ async def shutdown_server( async def get_system_info(): """Get system information""" import platform - import subprocess import psutil - from utils.hardware import get_device, get_gpu_memory_info, DeviceType + from utils.hardware import get_device - # GPU Info — query nvidia-smi for physical GPUs, filtered by - # CUDA_VISIBLE_DEVICES when set (the frontend uses this for GGUF - # fit estimation and llama-server respects CVD too). - import os - - gpu_info: dict = {"available": False, "devices": []} - - device = get_device() - if device == DeviceType.CUDA: - # Parse CUDA_VISIBLE_DEVICES allowlist - allowed_indices = None - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None and cvd.strip(): - try: - allowed_indices = set(int(x.strip()) for x in cvd.split(",")) - except ValueError: - pass # Non-numeric (e.g. GPU-uuid), show all - - try: - result = subprocess.run( - [ - "nvidia-smi", - "--query-gpu=index,name,memory.total", - "--format=csv,noheader,nounits", - ], - capture_output = True, - text = True, - timeout = 10, - ) - if result.returncode == 0: - for line in result.stdout.strip().splitlines(): - parts = [p.strip() for p in line.split(",")] - if len(parts) == 3: - idx = int(parts[0]) - if allowed_indices is not None and idx not in allowed_indices: - continue - gpu_info["devices"].append( - { - "index": idx, - "name": parts[1], - "memory_total_gb": round(int(parts[2]) / 1024, 2), - } - ) - gpu_info["available"] = len(gpu_info["devices"]) > 0 - except Exception: - pass - - # Fallback to torch-based single-GPU detection - if not gpu_info["available"]: - mem_info = get_gpu_memory_info() - if mem_info.get("available"): - gpu_info["available"] = True - gpu_info["devices"].append( - { - "index": mem_info.get("device", 0), - "name": mem_info.get("device_name", "Unknown"), - "memory_total_gb": round(mem_info.get("total_gb", 0), 2), - } - ) + visibility_info = get_backend_visible_gpu_info() + gpu_info = { + "available": visibility_info["available"], + "devices": visibility_info["devices"], + } # CPU & Memory memory = psutil.virtual_memory() @@ -318,6 +268,13 @@ async def get_system_info(): } +@app.get("/api/system/gpu-visibility") +async def get_gpu_visibility( + current_subject: str = Depends(get_current_subject), +): + return get_backend_visible_gpu_info() + + @app.get("/api/system/hardware") async def get_hardware_info(): """Return GPU name, total VRAM, and key ML package versions.""" diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 8f45bca759..3c928245f5 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -44,6 +44,10 @@ class LoadRequest(BaseModel): None, description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')", ) + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.", + ) class UnloadRequest(BaseModel): @@ -132,6 +136,9 @@ class LoadResponse(BaseModel): context_length: Optional[int] = Field( None, description = "Model's native context length (from GGUF metadata)" ) + max_context_length: Optional[int] = Field( + None, description = "Maximum context length currently available on this hardware" + ) supports_reasoning: bool = Field( False, description = "Whether model supports thinking/reasoning mode (enable_thinking)", @@ -206,6 +213,10 @@ class InferenceStatusResponse(BaseModel): context_length: Optional[int] = Field( None, description = "Context length of the active model" ) + max_context_length: Optional[int] = Field( + None, + description = "Maximum context length currently available for the active model", + ) # ===================================================================== diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 68791aa7a8..eeb98c872e 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -128,6 +128,12 @@ class TrainingStartRequest(BaseModel): enable_tensorboard: bool = Field(False, description = "Enable TensorBoard logging") tensorboard_dir: Optional[str] = Field(None, description = "TensorBoard directory") + # GPU selection + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.", + ) + class TrainingJobResponse(BaseModel): """Immediate response when training is initiated""" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 966aebd6ff..898ccd45c2 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -160,6 +160,7 @@ async def load_model( else False, inference = inference_config, context_length = llama_backend.context_length, + max_context_length = llama_backend.max_context_length, supports_reasoning = llama_backend.supports_reasoning, reasoning_always_on = llama_backend.reasoning_always_on, chat_template = llama_backend.chat_template, @@ -210,8 +211,17 @@ async def load_model( detail = f"Invalid model identifier: {request.model_path}", ) + # Normalize gpu_ids: empty list means auto-selection, same as None + effective_gpu_ids = request.gpu_ids if request.gpu_ids else None + # ── GGUF path: load via llama-server ────────────────────── if config.is_gguf: + if effective_gpu_ids is not None: + raise HTTPException( + status_code = 400, + detail = "gpu_ids is not supported for GGUF models yet.", + ) + llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() @@ -285,6 +295,7 @@ async def load_model( has_audio_input = is_audio_input_type(_gguf_audio), inference = inference_config, context_length = llama_backend.context_length, + max_context_length = llama_backend.max_context_length, supports_reasoning = llama_backend.supports_reasoning, reasoning_always_on = llama_backend.reasoning_always_on, supports_tools = llama_backend.supports_tools, @@ -372,6 +383,7 @@ async def load_model( load_in_4bit = load_in_4bit, hf_token = request.hf_token, trust_remote_code = request.trust_remote_code, + gpu_ids = effective_gpu_ids, ) if not success: @@ -423,6 +435,9 @@ async def load_model( except HTTPException: raise + except ValueError as e: + logger.warning("Rejected inference GPU selection: %s", e) + raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: logger.error(f"Error loading model: {e}", exc_info = True) msg = str(e) @@ -619,6 +634,7 @@ async def get_status( reasoning_always_on = llama_backend.reasoning_always_on, supports_tools = llama_backend.supports_tools, context_length = llama_backend.context_length, + max_context_length = llama_backend.max_context_length, ) # Otherwise, report Unsloth backend status diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 4cfb060dee..e625408bad 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -88,14 +88,22 @@ async def get_hardware_utilization( Get a live snapshot of GPU hardware utilization. Designed to be polled by the frontend during training. - Returns GPU utilization %, temperature, VRAM usage, and power draw - via nvidia-smi for maximum accuracy. + Returns live GPU memory usage information for the active backend. """ from utils.hardware import get_gpu_utilization return get_gpu_utilization() +@router.get("/hardware/visible") +async def get_visible_hardware_utilization( + current_subject: str = Depends(get_current_subject), +): + from utils.hardware import get_visible_gpu_utilization + + return get_visible_gpu_utilization() + + @router.post("/start") async def start_training( request: TrainingStartRequest, @@ -202,6 +210,7 @@ async def start_training( "enable_tensorboard": request.enable_tensorboard, "tensorboard_dir": request.tensorboard_dir or "", "trust_remote_code": request.trust_remote_code, + "gpu_ids": request.gpu_ids, } # Training page has no trust_remote_code toggle — the value comes from @@ -269,6 +278,9 @@ async def start_training( error = None, ) + except ValueError as e: + logger.warning("Rejected training GPU selection: %s", e) + raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: logger.error(f"Error starting training: {e}", exc_info = True) raise HTTPException( diff --git a/studio/backend/run.py b/studio/backend/run.py index 18babf6229..9c3622988e 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -312,10 +312,10 @@ def run_server( if frontend_path: if setup_frontend(app, frontend_path): if not silent: - print(f"✅ Frontend loaded from {frontend_path}") + print(f"[OK] Frontend loaded from {frontend_path}") else: if not silent: - print(f"⚠️ Frontend not found at {frontend_path}") + print(f"[WARNING] Frontend not found at {frontend_path}") # Create the uvicorn server and expose it for signal handlers config = uvicorn.Config( diff --git a/studio/backend/startup_banner.py b/studio/backend/startup_banner.py index 54acac0540..16b41d484c 100644 --- a/studio/backend/startup_banner.py +++ b/studio/backend/startup_banner.py @@ -20,7 +20,7 @@ def stdout_supports_color() -> bool: return True try: return sys.stdout.isatty() - except Exception: + except (AttributeError, OSError, ValueError): return False @@ -52,28 +52,36 @@ def print_studio_access_banner( ipv6_bind = bind_host in ("::", "::1") if ipv6_bind: - local_url = f"http://[::1]:{port}" + loopback_url = f"http://[::1]:{port}" alt_local = f"http://localhost:{port}" else: - local_url = f"http://127.0.0.1:{port}" + loopback_url = f"http://127.0.0.1:{port}" alt_local = f"http://localhost:{port}" if ":" in display_host: external_url = f"http://[{display_host}]:{port}" else: external_url = f"http://{display_host}:{port}" + listen_all = bind_host in ("0.0.0.0", "::") loopback_bind = bind_host in ("127.0.0.1", "localhost", "::1") - api_base = local_url if listen_all or loopback_bind else external_url + + # Use loopback URL only when the server is reachable on loopback; + # otherwise show the actual bound address. + primary_url = loopback_url if listen_all or loopback_bind else external_url + tip_url = alt_local if listen_all or loopback_bind else external_url + api_base = primary_url lines: list[str] = [ "", style("🦥 Unsloth Studio is running", title), style("─" * 52, dim), - style(" On this machine — open this in your browser:", dim), - style(f" {local_url}", local_url_style), - style(f" (same as {alt_local})", dim), + style(" On this machine -- open this in your browser:", dim), + style(f" {primary_url}", local_url_style), ] + if (listen_all or loopback_bind) and primary_url != alt_local: + lines.append(style(f" (same as {alt_local})", dim)) + if listen_all and display_host not in ( "127.0.0.1", "localhost", @@ -88,7 +96,7 @@ def print_studio_access_banner( style(f" {external_url}", secondary), ] ) - elif not listen_all and bind_host not in ("127.0.0.1", "localhost", "::1"): + elif not listen_all and not loopback_bind and external_url != primary_url: lines.extend( [ "", @@ -105,7 +113,7 @@ def print_studio_access_banner( style(f" {api_base}/api/health", secondary), style("─" * 52, dim), style( - " Tip: if you are on the same computer, use the Local link above.", + f" Tip: if you are on this computer, open {tip_url}/ in your browser.", dim, ), "", diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py new file mode 100644 index 0000000000..275b6f33d6 --- /dev/null +++ b/studio/backend/tests/test_gpu_selection.py @@ -0,0 +1,1125 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import asyncio +import importlib.util +import os +import re +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +from fastapi import HTTPException + +from core.training.training import TrainingBackend +from models.inference import LoadRequest +from models.training import TrainingStartRequest +from utils.hardware import ( + apply_gpu_ids, + DeviceType, + auto_select_gpu_ids, + estimate_required_model_memory_gb, + get_backend_visible_gpu_info, + get_device_map, + get_offloaded_device_map_entries, + get_parent_visible_gpu_ids, + get_visible_gpu_utilization, + prepare_gpu_selection, + resolve_requested_gpu_ids, +) +import utils.hardware.hardware as _hw_module + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent + + +def _load_route_module(name: str, relative_path: str): + spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class _GpuCacheResetMixin: + """Reset module-level GPU caches between tests to prevent state leaks.""" + + def tearDown(self): + _hw_module._physical_gpu_count = None + _hw_module._visible_gpu_count = None + + +class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): + def test_parent_visibility_defaults_to_physical_enumeration(self): + with ( + patch.dict(os.environ, {}, clear = True), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 4), + ): + self.assertEqual(get_parent_visible_gpu_ids(), [0, 1, 2, 3]) + self.assertEqual(resolve_requested_gpu_ids(None), [0, 1, 2, 3]) + + def test_parent_visibility_uses_cuda_visible_devices(self): + with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True): + self.assertEqual(get_parent_visible_gpu_ids(), [1, 3]) + self.assertEqual(resolve_requested_gpu_ids(None), [1, 3]) + + def test_parent_visibility_uses_empty_numeric_ids_for_uuid_masks(self): + with ( + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + self.assertEqual(get_parent_visible_gpu_ids(), []) + + def test_invalid_requests_raise_clear_value_errors(self): + cases = [ + ([1, 1], "duplicate GPU IDs"), + ([-1], "Rejected IDs: [-1]"), + ([99], "Rejected IDs: [99]"), + ([0], "outside the parent-visible set [1, 3]"), + ] + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + for gpu_ids, message in cases: + with self.subTest(gpu_ids = gpu_ids): + with self.assertRaisesRegex(ValueError, re.escape(message)): + resolve_requested_gpu_ids(gpu_ids) + + def test_explicit_ids_must_be_physical_not_relative(self): + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + self.assertEqual(resolve_requested_gpu_ids([1, 3]), [1, 3]) + + def test_explicit_ids_are_rejected_for_uuid_parent_visibility(self): + with ( + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + with self.assertRaisesRegex( + ValueError, "unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG" + ): + resolve_requested_gpu_ids([1]) + + def test_empty_list_is_treated_as_auto(self): + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + self.assertEqual(resolve_requested_gpu_ids([]), [1, 3]) + + def test_apply_gpu_ids_only_updates_cuda_visible_devices(self): + with patch.dict( + os.environ, + {"CUDA_VISIBLE_DEVICES": "1,3", "TEST_PARENT_ENV": "keep-me"}, + clear = True, + ): + apply_gpu_ids([5, 6]) + + self.assertEqual(os.environ["CUDA_VISIBLE_DEVICES"], "5,6") + self.assertEqual(os.environ["TEST_PARENT_ENV"], "keep-me") + + +class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase): + def test_visible_gpu_utilization_filters_to_parent_visible_ids(self): + smi_output = "\n".join( + [ + "0, 10, 30, 1000, 10000, 50, 100", + "1, 20, 40, 2000, 10000, 60, 120", + "3, 30, 50, 3000, 10000, 70, 140", + ] + ) + + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch("utils.hardware.nvidia.subprocess.run") as mock_run, + ): + mock_run.return_value = SimpleNamespace( + returncode = 0, + stdout = smi_output, + ) + result = get_visible_gpu_utilization() + + self.assertTrue(result["available"]) + self.assertEqual(result["parent_visible_gpu_ids"], [1, 3]) + self.assertEqual(result["index_kind"], "physical") + self.assertEqual([device["index"] for device in result["devices"]], [1, 3]) + self.assertEqual(result["devices"][0]["visible_ordinal"], 0) + self.assertEqual(result["devices"][1]["visible_ordinal"], 1) + self.assertEqual(result["devices"][0]["gpu_utilization_pct"], 20.0) + self.assertEqual(result["devices"][1]["power_utilization_pct"], 50.0) + + def test_backend_visible_gpu_info_preserves_physical_indices(self): + smi_output = "\n".join( + [ + "0, GPU Zero, 10000", + "1, GPU One, 20000", + "3, GPU Three, 30000", + ] + ) + + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch("utils.hardware.nvidia.subprocess.run") as mock_run, + ): + mock_run.return_value = SimpleNamespace( + returncode = 0, + stdout = smi_output, + ) + result = get_backend_visible_gpu_info() + + self.assertTrue(result["available"]) + self.assertEqual(result["parent_visible_gpu_ids"], [1, 3]) + self.assertEqual(result["index_kind"], "physical") + self.assertEqual([device["index"] for device in result["devices"]], [1, 3]) + self.assertEqual(result["devices"][0]["visible_ordinal"], 0) + self.assertEqual(result["devices"][1]["visible_ordinal"], 1) + self.assertEqual(result["devices"][0]["name"], "GPU One") + self.assertAlmostEqual(result["devices"][1]["memory_total_gb"], 29.3, places = 1) + + def test_uuid_parent_visibility_falls_back_to_torch(self): + """UUID/MIG masks should fall through nvidia to torch fallback and + still report visible devices using relative ordinals.""" + fake_torch_devices = [ + { + "index": 0, + "visible_ordinal": 0, + "name": "GPU-A", + "total_gb": 24.0, + "used_gb": 2.0, + }, + { + "index": 1, + "visible_ordinal": 1, + "name": "GPU-B", + "total_gb": 24.0, + "used_gb": 3.0, + }, + ] + with ( + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch( + "utils.hardware.hardware._torch_get_physical_gpu_count", return_value = 2 + ), + patch( + "utils.hardware.hardware._torch_get_per_device_info", + return_value = fake_torch_devices, + ), + ): + result = get_backend_visible_gpu_info() + + self.assertTrue(result["available"]) + self.assertEqual(result["parent_visible_gpu_ids"], []) + self.assertEqual(len(result["devices"]), 2) + self.assertEqual(result["index_kind"], "relative") + + def test_mlx_visible_gpu_info_is_best_effort_relative(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX), + patch( + "utils.hardware.hardware.get_gpu_memory_info", + return_value = { + "available": True, + "device_name": "Apple Silicon", + "total_gb": 64.0, + "allocated_gb": 8.0, + "utilization_pct": 12.5, + }, + ), + ): + result = get_backend_visible_gpu_info() + + self.assertTrue(result["available"]) + self.assertEqual(result["index_kind"], "relative") + self.assertEqual(result["devices"][0]["index"], 0) + self.assertEqual(result["devices"][0]["visible_ordinal"], 0) + + +class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): + def test_get_device_map_uses_explicit_gpu_selection(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA): + self.assertEqual(get_device_map(None), "sequential") + self.assertEqual(get_device_map([0]), "sequential") + self.assertEqual(get_device_map([0, 1]), "balanced") + + def test_get_device_map_multi_gpu_uses_balanced(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA): + self.assertEqual(get_device_map([0, 1]), "balanced") + self.assertEqual(get_device_map([0]), "sequential") + + def test_get_device_map_uses_all_inherited_visible_gpus_for_uuid_masks(self): + with ( + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + ): + self.assertEqual(get_device_map(None), "balanced") + + def test_get_offloaded_device_map_entries_returns_only_cpu_and_disk(self): + model = SimpleNamespace( + hf_device_map = { + "model.embed_tokens": 0, + "model.layers.0": 1, + "model.layers.1": "cpu", + "lm_head": "disk", + } + ) + + self.assertEqual( + get_offloaded_device_map_entries(model), + { + "model.layers.1": "cpu", + "lm_head": "disk", + }, + ) + + def test_get_offloaded_device_map_entries_handles_models_without_device_map(self): + self.assertEqual(get_offloaded_device_map_entries(SimpleNamespace()), {}) + + def test_estimate_required_memory_formulas(self): + eight_gb = 8 * (1024**3) + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (eight_gb, "config"), + ): + # FP16 inference: 8GB * 1.3 = 10.4GB + required_gb, metadata = estimate_required_model_memory_gb( + "unsloth/test", + load_in_4bit = False, + ) + self.assertAlmostEqual(required_gb, 10.4, places = 3) + self.assertEqual(metadata["model_size_source"], "config") + + # 4bit inference: base_4bit = 8/3.2 = 2.5GB + # required = 2.5 + max(2.5*0.3, 2.0) = 2.5 + 2.0 = 4.5GB + required_gb, _ = estimate_required_model_memory_gb( + "unsloth/test", + load_in_4bit = True, + ) + self.assertAlmostEqual(required_gb, 4.5, places = 2) + + # Full FT fallback: model_size * 3.5 + overhead + required_gb, metadata = estimate_required_model_memory_gb( + "unsloth/test", training_type = "Full Finetuning" + ) + self.assertEqual(metadata.get("estimation_mode"), "fallback") + self.assertGreater(required_gb, 25.0) + self.assertLess(required_gb, 40.0) + + # LoRA fp16 fallback: model_size + lora_overhead + activations + overhead + required_gb, metadata = estimate_required_model_memory_gb( + "unsloth/test", + training_type = "LoRA/QLoRA", + load_in_4bit = False, + ) + self.assertEqual(metadata.get("estimation_mode"), "fallback") + self.assertGreater(required_gb, 8.0) + self.assertLess(required_gb, 15.0) + + # QLoRA 4-bit fallback: compressed weights + lora overhead + activations + overhead + required_gb, metadata = estimate_required_model_memory_gb( + "unsloth/test", + training_type = "LoRA/QLoRA", + load_in_4bit = True, + ) + self.assertEqual(metadata.get("estimation_mode"), "fallback") + self.assertGreater(required_gb, 3.0) + self.assertLess(required_gb, 8.0) + + # Larger model: 16GB fp16 + sixteen_gb = 16 * (1024**3) + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (sixteen_gb, "config"), + ): + required_gb, _ = estimate_required_model_memory_gb( + "unsloth/test", + training_type = "LoRA/QLoRA", + load_in_4bit = True, + ) + # QLoRA for 16GB model should be < 12 GB + self.assertGreater(required_gb, 5.0) + self.assertLess(required_gb, 12.0) + + def test_estimate_fp16_model_size_bytes_uses_vllm_fallback_last(self): + config = object() + with ( + patch( + "utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate", + return_value = "unsloth/test", + ), + patch( + "utils.hardware.hardware._get_hf_safetensors_total_params", + return_value = None, + ), + patch( + "utils.hardware.hardware._load_config_for_gpu_estimate", + return_value = config, + ), + patch( + "utils.hardware.hardware._estimate_fp16_model_size_bytes_from_config", + return_value = None, + ), + patch( + "utils.hardware.hardware._get_local_weight_size_bytes", + return_value = None, + ), + patch( + "utils.hardware.hardware._estimate_fp16_model_size_bytes_from_vllm_utils", + return_value = 1234, + ), + ): + model_size_bytes, source = _hw_module.estimate_fp16_model_size_bytes( + "unsloth/test" + ) + + self.assertEqual(model_size_bytes, 1234) + self.assertEqual(source, "vllm_utils") + + def test_auto_select_gpu_ids_chooses_smallest_fitting_subset(self): + fake_devices = { + "devices": [ + {"index": 0, "vram_total_gb": 16.0, "vram_used_gb": 4.0}, + {"index": 1, "vram_total_gb": 16.0, "vram_used_gb": 6.0}, + {"index": 2, "vram_total_gb": 16.0, "vram_used_gb": 7.0}, + ] + } + + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch( + "utils.hardware.hardware.estimate_required_model_memory_gb", + return_value = ( + 14.0, + {"required_gb": 14.0, "model_size_source": "config"}, + ), + ), + patch( + "utils.hardware.hardware.get_visible_gpu_utilization", + return_value = fake_devices, + ), + ): + selected, metadata = auto_select_gpu_ids("unsloth/test") + + self.assertEqual(selected, [0, 1]) + self.assertEqual(metadata["selection_mode"], "auto") + # First GPU full (12GB) + second GPU with overhead (10*0.85=8.5) = 20.5GB + self.assertAlmostEqual(metadata["usable_gb"], 20.5, places = 3) + + def test_auto_select_gpu_ids_falls_back_to_all_visible(self): + fake_devices = { + "devices": [ + {"index": 0, "vram_total_gb": 12.0, "vram_used_gb": 2.0}, + {"index": 1, "vram_total_gb": 12.0, "vram_used_gb": 2.0}, + ] + } + + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch( + "utils.hardware.hardware.estimate_required_model_memory_gb", + return_value = ( + 30.0, + {"required_gb": 30.0, "model_size_source": "config"}, + ), + ), + patch( + "utils.hardware.hardware.get_visible_gpu_utilization", + return_value = fake_devices, + ), + ): + selected, metadata = auto_select_gpu_ids("unsloth/test") + + self.assertEqual(selected, [0, 1]) + self.assertEqual(metadata["selection_mode"], "fallback_all") + # First GPU full (10GB) + second GPU with overhead (10*0.85=8.5) = 18.5GB + self.assertAlmostEqual(metadata["usable_gb"], 18.5, places = 3) + + def test_prepare_gpu_selection_preserves_explicit_ids_without_auto_selection(self): + with ( + patch( + "utils.hardware.hardware.resolve_requested_gpu_ids", + return_value = [2, 3], + ), + patch("utils.hardware.hardware.auto_select_gpu_ids") as mock_auto_select, + ): + selected, metadata = prepare_gpu_selection( + [2, 3], + model_name = "unsloth/test", + ) + + self.assertEqual(selected, [2, 3]) + self.assertEqual(metadata["selection_mode"], "explicit") + mock_auto_select.assert_not_called() + + def test_prepare_gpu_selection_treats_empty_list_as_auto(self): + with patch( + "utils.hardware.hardware.auto_select_gpu_ids", + return_value = ([0, 1], {"selection_mode": "auto"}), + ) as mock_auto_select: + selected, metadata = prepare_gpu_selection( + [], + model_name = "unsloth/test", + ) + + self.assertEqual(selected, [0, 1]) + self.assertEqual(metadata["selection_mode"], "auto") + mock_auto_select.assert_called_once() + + def test_prepare_gpu_selection_preserves_uuid_parent_visibility_in_auto_mode(self): + with ( + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), + patch( + "utils.hardware.hardware.estimate_required_model_memory_gb", + return_value = ( + 14.0, + {"required_gb": 14.0, "model_size_source": "config"}, + ), + ), + ): + selected, metadata = prepare_gpu_selection( + None, + model_name = "unsloth/test", + ) + + self.assertIsNone(selected) + self.assertEqual(metadata["selection_mode"], "inherit_parent_visible") + self.assertIsNone(metadata["selected_gpu_ids"]) + + +class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): + def test_training_backend_resolves_explicit_gpu_ids_before_spawn(self): + backend = TrainingBackend() + + class DummyProcess: + pid = 12345 + + def start(self): + return None + + class DummyThread: + def start(self): + return None + + dummy_queue = object() + + with ( + patch( + "core.training.training.prepare_gpu_selection", + return_value = ([1, 2], {"selection_mode": "explicit"}), + ), + patch( + "core.training.training._CTX.Queue", + side_effect = [dummy_queue, dummy_queue], + ), + patch( + "core.training.training._CTX.Process", return_value = DummyProcess() + ) as mock_process, + patch( + "core.training.training.threading.Thread", return_value = DummyThread() + ), + ): + backend.start_training( + job_id = "test-job-1", + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + gpu_ids = [1, 2], + ) + + config = mock_process.call_args.kwargs["kwargs"]["config"] + self.assertEqual(config["gpu_ids"], [1, 2]) + self.assertEqual(config["resolved_gpu_ids"], [1, 2]) + self.assertEqual(config["gpu_selection"]["selection_mode"], "explicit") + + def test_training_backend_auto_selects_gpu_ids_when_omitted(self): + backend = TrainingBackend() + + class DummyProcess: + pid = 12345 + + def start(self): + return None + + class DummyThread: + def start(self): + return None + + dummy_queue = object() + + with ( + patch( + "core.training.training.prepare_gpu_selection", + return_value = ([0, 1], {"selection_mode": "auto"}), + ), + patch( + "core.training.training._CTX.Queue", + side_effect = [dummy_queue, dummy_queue], + ), + patch( + "core.training.training._CTX.Process", return_value = DummyProcess() + ) as mock_process, + patch( + "core.training.training.threading.Thread", return_value = DummyThread() + ), + ): + backend.start_training( + job_id = "test-job-2", + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + gpu_ids = None, + ) + + config = mock_process.call_args.kwargs["kwargs"]["config"] + self.assertIsNone(config["gpu_ids"]) + self.assertEqual(config["resolved_gpu_ids"], [0, 1]) + self.assertEqual(config["gpu_selection"]["selection_mode"], "auto") + + def test_training_backend_preserves_uuid_parent_visibility_in_auto_mode(self): + backend = TrainingBackend() + + class DummyProcess: + pid = 12345 + + def start(self): + return None + + class DummyThread: + def start(self): + return None + + dummy_queue = object() + + with ( + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), + patch( + "core.training.training._CTX.Queue", + side_effect = [dummy_queue, dummy_queue], + ), + patch( + "core.training.training._CTX.Process", return_value = DummyProcess() + ) as mock_process, + patch( + "core.training.training.threading.Thread", return_value = DummyThread() + ), + patch( + "utils.hardware.hardware.estimate_required_model_memory_gb", + return_value = ( + 14.0, + {"required_gb": 14.0, "model_size_source": "config"}, + ), + ), + ): + backend.start_training( + job_id = "test-job-uuid-auto", + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + gpu_ids = None, + ) + + config = mock_process.call_args.kwargs["kwargs"]["config"] + self.assertIsNone(config["resolved_gpu_ids"]) + self.assertEqual( + config["gpu_selection"]["selection_mode"], "inherit_parent_visible" + ) + + def test_inference_orchestrator_resolves_explicit_gpu_ids_before_spawn(self): + class DummyThread: + def __init__(self, *args, **kwargs): + pass + + def start(self): + return None + + with patch("core.inference.orchestrator.threading.Thread", DummyThread): + from core.inference.orchestrator import InferenceOrchestrator + + orchestrator = InferenceOrchestrator() + + config = SimpleNamespace(identifier = "unsloth/test", gguf_variant = None) + + with ( + patch( + "core.inference.orchestrator.prepare_gpu_selection", + return_value = ([1], {"selection_mode": "explicit"}), + ), + patch.object(orchestrator, "_ensure_subprocess_alive", return_value = False), + patch.object(orchestrator, "_spawn_subprocess") as mock_spawn, + patch.object( + orchestrator, + "_wait_response", + return_value = {"success": True, "model_info": {}}, + ), + patch( + "utils.transformers_version.needs_transformers_5", return_value = False + ), + ): + self.assertTrue(orchestrator.load_model(config = config, gpu_ids = [1])) + + sub_config = mock_spawn.call_args.args[0] + self.assertEqual(sub_config["gpu_ids"], [1]) + self.assertEqual(sub_config["resolved_gpu_ids"], [1]) + self.assertEqual(sub_config["gpu_selection"]["selection_mode"], "explicit") + + def test_inference_orchestrator_auto_selects_gpu_ids_when_omitted(self): + class DummyThread: + def __init__(self, *args, **kwargs): + pass + + def start(self): + return None + + with patch("core.inference.orchestrator.threading.Thread", DummyThread): + from core.inference.orchestrator import InferenceOrchestrator + + orchestrator = InferenceOrchestrator() + + config = SimpleNamespace(identifier = "unsloth/test", gguf_variant = None) + + with ( + patch( + "core.inference.orchestrator.prepare_gpu_selection", + return_value = ([0], {"selection_mode": "auto"}), + ), + patch.object(orchestrator, "_ensure_subprocess_alive", return_value = False), + patch.object(orchestrator, "_spawn_subprocess") as mock_spawn, + patch.object( + orchestrator, + "_wait_response", + return_value = {"success": True, "model_info": {}}, + ), + patch( + "utils.transformers_version.needs_transformers_5", return_value = False + ), + ): + self.assertTrue(orchestrator.load_model(config = config, gpu_ids = None)) + + sub_config = mock_spawn.call_args.args[0] + self.assertIsNone(sub_config["gpu_ids"]) + self.assertEqual(sub_config["resolved_gpu_ids"], [0]) + self.assertEqual(sub_config["gpu_selection"]["selection_mode"], "auto") + + +class TestRouteErrors(unittest.TestCase): + def test_prepare_gpu_selection_rejects_gpu_ids_on_non_cuda_backend(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU): + with self.assertRaises(ValueError) as exc_info: + prepare_gpu_selection([0], model_name = "unsloth/test") + + self.assertIn("only supported on CUDA devices", str(exc_info.exception)) + + def test_inference_route_rejects_gpu_ids_for_gguf(self): + inference_route = _load_route_module( + "inference_route_module_for_gguf_gpu_ids_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1]) + model_config = SimpleNamespace( + is_gguf = True, + is_lora = False, + gguf_hf_repo = None, + gguf_file = "/tmp/test.gguf", + gguf_mmproj_file = None, + gguf_variant = None, + identifier = "unsloth/test.gguf", + display_name = "unsloth/test.gguf", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + + with patch.object( + inference_route.ModelConfig, + "from_identifier", + return_value = model_config, + ): + with self.assertRaises(HTTPException) as exc_info: + asyncio.run( + inference_route.load_model(request, current_subject = "test-user") + ) + + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("GGUF", exc_info.exception.detail) + + def test_training_route_returns_400_for_invalid_gpu_ids(self): + training_route = _load_route_module( + "training_route_module_for_test", + "routes/training.py", + ) + request = TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + format_type = "alpaca", + gpu_ids = [99], + ) + + class DummyBackend: + current_job_id = None + + def is_training_active(self): + return False + + def start_training(self, **kwargs): + raise ValueError("Invalid gpu_ids [99]") + + with ( + patch.object( + training_route, "get_training_backend", return_value = DummyBackend() + ), + patch( + "core.inference.get_inference_backend", + return_value = SimpleNamespace(active_model_name = None), + ), + patch( + "core.export.get_export_backend", + return_value = SimpleNamespace(current_checkpoint = None), + ), + ): + with self.assertRaises(HTTPException) as exc_info: + asyncio.run( + training_route.start_training(request, current_subject = "test-user") + ) + + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("gpu_ids [99]", exc_info.exception.detail) + + def test_training_route_returns_400_for_uuid_parent_visibility_gpu_ids(self): + training_route = _load_route_module( + "training_route_module_for_uuid_parent_visibility_test", + "routes/training.py", + ) + request = TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + format_type = "alpaca", + gpu_ids = [1], + ) + + class DummyBackend: + current_job_id = None + + def is_training_active(self): + return False + + def start_training(self, **kwargs): + raise ValueError( + "Invalid gpu_ids [1]: explicit physical GPU IDs are unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries" + ) + + with ( + patch.object( + training_route, "get_training_backend", return_value = DummyBackend() + ), + patch( + "core.inference.get_inference_backend", + return_value = SimpleNamespace(active_model_name = None), + ), + patch( + "core.export.get_export_backend", + return_value = SimpleNamespace(current_checkpoint = None), + ), + ): + with self.assertRaises(HTTPException) as exc_info: + asyncio.run( + training_route.start_training(request, current_subject = "test-user") + ) + + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("UUID/MIG", exc_info.exception.detail) + + def test_inference_route_returns_400_for_invalid_gpu_ids(self): + inference_route = _load_route_module( + "inference_route_module_for_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/test", gpu_ids = [99]) + model_config = SimpleNamespace( + is_gguf = False, + is_lora = False, + path = None, + identifier = "unsloth/test", + display_name = "unsloth/test", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + + class DummyInferenceBackend: + active_model_name = None + models = {} + + def load_model(self, **kwargs): + raise ValueError("Invalid gpu_ids [99]") + + with ( + patch.object( + inference_route.ModelConfig, + "from_identifier", + return_value = model_config, + ), + patch.object( + inference_route, + "get_inference_backend", + return_value = DummyInferenceBackend(), + ), + patch.object( + inference_route, + "get_llama_cpp_backend", + return_value = SimpleNamespace(is_loaded = False), + ), + patch( + "core.export.get_export_backend", + return_value = SimpleNamespace(current_checkpoint = None), + ), + ): + with self.assertRaises(HTTPException) as exc_info: + asyncio.run( + inference_route.load_model(request, current_subject = "test-user") + ) + + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("gpu_ids [99]", exc_info.exception.detail) + + def test_inference_route_returns_400_for_uuid_parent_visibility_gpu_ids(self): + inference_route = _load_route_module( + "inference_route_module_for_uuid_parent_visibility_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/test", gpu_ids = [1]) + model_config = SimpleNamespace( + is_gguf = False, + is_lora = False, + path = None, + identifier = "unsloth/test", + display_name = "unsloth/test", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + + class DummyInferenceBackend: + active_model_name = None + models = {} + + def load_model(self, **kwargs): + raise ValueError( + "Invalid gpu_ids [1]: explicit physical GPU IDs are unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries" + ) + + with ( + patch.object( + inference_route.ModelConfig, + "from_identifier", + return_value = model_config, + ), + patch.object( + inference_route, + "get_inference_backend", + return_value = DummyInferenceBackend(), + ), + patch.object( + inference_route, + "get_llama_cpp_backend", + return_value = SimpleNamespace(is_loaded = False), + ), + patch( + "core.export.get_export_backend", + return_value = SimpleNamespace(current_checkpoint = None), + ), + ): + with self.assertRaises(HTTPException) as exc_info: + asyncio.run( + inference_route.load_model(request, current_subject = "test-user") + ) + + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("UUID/MIG", exc_info.exception.detail) + + +class TestRaiseIfOffloaded(unittest.TestCase): + def test_no_offload_is_noop(self): + from utils.hardware import raise_if_offloaded + + model = SimpleNamespace(hf_device_map = {"model.embed_tokens": 0, "lm_head": 1}) + raise_if_offloaded(model, "balanced", "Test") + + def test_cpu_offload_raises(self): + from utils.hardware import raise_if_offloaded + + model = SimpleNamespace( + hf_device_map = {"model.layers.0": 0, "model.layers.1": "cpu"} + ) + with self.assertRaisesRegex(ValueError, "offloaded"): + raise_if_offloaded(model, "balanced", "Test") + + def test_no_device_map_attr_is_noop(self): + from utils.hardware import raise_if_offloaded + + raise_if_offloaded(SimpleNamespace(), "sequential", "Test") + + +class TestMinGpuVram(unittest.TestCase): + def test_min_gpu_vram_decreases_with_more_gpus(self): + from utils.hardware.vram_estimation import ( + ModelArchConfig, + TrainingVramConfig, + estimate_training_vram, + ) + + arch = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 128256, + tie_word_embeddings = False, + ) + config = TrainingVramConfig( + training_method = "qlora", + load_in_4bit = True, + ) + breakdown = estimate_training_vram(arch, config) + v1 = breakdown.min_gpu_vram(1) + v2 = breakdown.min_gpu_vram(2) + v4 = breakdown.min_gpu_vram(4) + self.assertGreater(v1, v2) + self.assertGreater(v2, v4) + self.assertGreater(v4, 0) + + def test_total_equals_min_gpu_vram_1(self): + from utils.hardware.vram_estimation import ( + ModelArchConfig, + TrainingVramConfig, + estimate_training_vram, + ) + + arch = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 128256, + tie_word_embeddings = False, + ) + config = TrainingVramConfig( + training_method = "qlora", + load_in_4bit = True, + ) + breakdown = estimate_training_vram(arch, config) + self.assertEqual(breakdown.total, breakdown.min_gpu_vram(1)) + + +class TestPerGpuFitGuardAllCounts(unittest.TestCase): + def test_min_per_gpu_generated_for_all_visible_counts(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (8 * (1024**3), "config"), + ), + patch( + "utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate", + return_value = "unsloth/test", + ), + patch( + "utils.hardware.hardware._load_config_for_gpu_estimate", + return_value = SimpleNamespace( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 128256, + tie_word_embeddings = False, + ), + ), + patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 6), + ): + _, metadata = estimate_required_model_memory_gb( + "unsloth/test", + training_type = "LoRA/QLoRA", + load_in_4bit = True, + ) + + self.assertEqual(metadata.get("estimation_mode"), "detailed") + breakdown = metadata["vram_breakdown"] + for n in range(1, 7): + self.assertIn(f"min_per_gpu_{n}", breakdown) + + +class TestAutoSelectWithNoneRequired(_GpuCacheResetMixin, unittest.TestCase): + def test_auto_select_falls_back_when_estimate_unavailable(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch( + "utils.hardware.hardware.estimate_required_model_memory_gb", + return_value = (None, {"model_size_source": "unavailable"}), + ), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = { + "raw": "0,1", + "numeric_ids": [0, 1], + "supports_explicit_gpu_ids": True, + }, + ), + patch( + "utils.hardware.hardware.get_parent_visible_gpu_ids", + return_value = [0, 1], + ), + ): + selected, metadata = auto_select_gpu_ids("unsloth/test") + + self.assertEqual(selected, [0, 1]) + self.assertEqual(metadata["selection_mode"], "fallback_all") + + +class TestXpuRejection(_GpuCacheResetMixin, unittest.TestCase): + def test_auto_select_returns_non_cuda_for_xpu(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU): + selected, metadata = auto_select_gpu_ids("unsloth/test") + + self.assertIsNone(selected) + self.assertEqual(metadata["selection_mode"], "non_cuda") + + def test_prepare_gpu_selection_rejects_explicit_ids_on_xpu(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU): + with self.assertRaisesRegex(ValueError, "only supported on CUDA"): + prepare_gpu_selection([0], model_name = "unsloth/test") + + +class TestDeviceMapForInference(_GpuCacheResetMixin, unittest.TestCase): + def test_inference_uses_balanced_low_0(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA): + self.assertEqual( + get_device_map([0, 1], for_inference = True), "balanced_low_0" + ) + + def test_training_uses_balanced(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA): + self.assertEqual(get_device_map([0, 1], for_inference = False), "balanced") + + def test_single_gpu_always_sequential(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA): + self.assertEqual(get_device_map([0], for_inference = True), "sequential") + self.assertEqual(get_device_map([0], for_inference = False), "sequential") diff --git a/studio/backend/tests/test_gpu_selection_sandbox.py b/studio/backend/tests/test_gpu_selection_sandbox.py new file mode 100644 index 0000000000..830a98a2fb --- /dev/null +++ b/studio/backend/tests/test_gpu_selection_sandbox.py @@ -0,0 +1,544 @@ +#!/usr/bin/env python3 +""" +Sandbox test for multi-GPU selection logic. + +Tests the core GPU selection, memory estimation, and device_map logic +in an isolated environment. Can be run on Linux, macOS, and Windows +without requiring actual GPUs -- all hardware calls are mocked. + +Usage: + python -m pytest studio/backend/tests/test_gpu_selection_sandbox.py -v + # or directly: + python studio/backend/tests/test_gpu_selection_sandbox.py +""" + +import os +import sys +import unittest +from pathlib import Path +from unittest.mock import patch, MagicMock + +# Ensure backend is on sys.path +_backend_root = Path(__file__).resolve().parent.parent +if str(_backend_root) not in sys.path: + sys.path.insert(0, str(_backend_root)) + + +def _make_fake_config( + vocab_size = 32000, + hidden_size = 4096, + intermediate_size = 11008, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + tie_word_embeddings = False, +): + """Create a fake HF config-like object for estimation tests.""" + from types import SimpleNamespace + + return SimpleNamespace( + vocab_size = vocab_size, + hidden_size = hidden_size, + intermediate_size = intermediate_size, + num_hidden_layers = num_hidden_layers, + num_attention_heads = num_attention_heads, + num_key_value_heads = num_key_value_heads, + tie_word_embeddings = tie_word_embeddings, + ) + + +class TestEstimateFP16ModelSizeFromConfig(unittest.TestCase): + """Test the config-based model size estimation.""" + + def test_llama_8b_size_reasonable(self): + from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config + + config = _make_fake_config( + vocab_size = 128256, + hidden_size = 4096, + intermediate_size = 14336, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + tie_word_embeddings = False, + ) + size = _estimate_fp16_model_size_bytes_from_config(config) + self.assertIsNotNone(size) + size_gb = size / (1024**3) + # Llama 3.1 8B should be ~15GB in fp16 + self.assertGreater(size_gb, 12) + self.assertLess(size_gb, 20) + + def test_small_model(self): + from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config + + config = _make_fake_config( + vocab_size = 32000, + hidden_size = 2048, + intermediate_size = 5504, + num_hidden_layers = 22, + num_attention_heads = 32, + num_key_value_heads = 4, + ) + size = _estimate_fp16_model_size_bytes_from_config(config) + self.assertIsNotNone(size) + size_gb = size / (1024**3) + # ~1B model should be ~2GB in fp16 + self.assertGreater(size_gb, 1) + self.assertLess(size_gb, 5) + + def test_returns_none_for_incomplete_config(self): + from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config + from types import SimpleNamespace + + config = SimpleNamespace(vocab_size = 32000) # Missing most fields + size = _estimate_fp16_model_size_bytes_from_config(config) + self.assertIsNone(size) + + def test_moe_model(self): + from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config + from types import SimpleNamespace + + config = SimpleNamespace( + vocab_size = 152064, + hidden_size = 3584, + intermediate_size = 18944, + num_hidden_layers = 28, + num_attention_heads = 28, + num_key_value_heads = 4, + tie_word_embeddings = False, + num_local_experts = 64, + moe_intermediate_size = 2560, + ) + size = _estimate_fp16_model_size_bytes_from_config(config) + self.assertIsNotNone(size) + size_gb = size / (1024**3) + # MoE model with 64 experts should be large + self.assertGreater(size_gb, 50) + + +class TestEstimateRequiredModelMemory(unittest.TestCase): + """Test memory requirement estimation.""" + + def test_inference_fp16_uses_1_3x(self): + from utils.hardware.hardware import estimate_required_model_memory_gb + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (10 * (1024**3), "config"), # 10GB model + ): + required, meta = estimate_required_model_memory_gb( + "test/model", + training_type = None, # inference + load_in_4bit = False, + ) + self.assertIsNotNone(required) + self.assertAlmostEqual(required, 13.0, places = 0) + self.assertEqual(meta["mode"], "inference") + + def test_inference_4bit_uses_reduced_estimate(self): + from utils.hardware.hardware import estimate_required_model_memory_gb + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (30 * (1024**3), "config"), # 30GB fp16 model + ): + required, meta = estimate_required_model_memory_gb( + "test/model", + training_type = None, # inference + load_in_4bit = True, + ) + self.assertIsNotNone(required) + # 4bit base = 30/3.2 = 9.375GB, required = 9.375 + max(9.375*0.3, 2) = 12.19GB + self.assertAlmostEqual(required, 12.2, places = 0) + + def test_4bit_training_reduces_base(self): + from utils.hardware.hardware import estimate_required_model_memory_gb + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (30 * (1024**3), "config"), # 30GB fp16 model + ): + required, meta = estimate_required_model_memory_gb( + "test/model", + training_type = "LoRA/QLoRA", + load_in_4bit = True, + ) + self.assertIsNotNone(required) + # fallback: base=30/3.2=9.375, lora=30*0.04=1.2, act=30*0.15=4.5, cuda=1.4 + self.assertAlmostEqual(required, 16.5, places = 0) + + def test_full_finetune_uses_3_5x(self): + from utils.hardware.hardware import estimate_required_model_memory_gb + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (10 * (1024**3), "config"), # 10GB model + ): + required, meta = estimate_required_model_memory_gb( + "test/model", + training_type = "Full Finetuning", + ) + self.assertIsNotNone(required) + # fallback: 10 * 3.5 + 1.4 cuda overhead = 36.4 + self.assertAlmostEqual(required, 36.4, places = 0) + + def test_returns_none_when_unavailable(self): + from utils.hardware.hardware import estimate_required_model_memory_gb + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (None, "unavailable"), + ): + required, meta = estimate_required_model_memory_gb("test/model") + self.assertIsNone(required) + + +class TestAutoSelectGpuIds(unittest.TestCase): + """Test automatic GPU selection based on model size and free memory.""" + + def _make_utilization(self, devices): + """Create a fake utilization response.""" + return { + "available": True, + "devices": [ + { + "index": idx, + "vram_total_gb": total, + "vram_used_gb": total - free, + } + for idx, total, free in devices + ], + } + + def test_single_gpu_sufficient(self): + from utils.hardware.hardware import auto_select_gpu_ids + import utils.hardware.hardware as hw + + with ( + patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), + patch.object( + hw, + "estimate_required_model_memory_gb", + return_value = ( + 10.0, + { + "mode": "inference", + "required_gb": 10.0, + "model_size_source": "config", + "model_size_gb": 7.7, + }, + ), + ), + patch.object( + hw, + "_get_parent_visible_gpu_spec", + return_value = { + "raw": "0,1,2,3", + "numeric_ids": [0, 1, 2, 3], + "supports_explicit_gpu_ids": True, + }, + ), + patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1, 2, 3]), + patch.object( + hw, + "get_visible_gpu_utilization", + return_value = self._make_utilization( + [ + (0, 80.0, 75.0), + (1, 80.0, 78.0), + (2, 80.0, 70.0), + (3, 80.0, 72.0), + ] + ), + ), + ): + selected, meta = auto_select_gpu_ids("test/model") + # Should pick GPU 1 (most free memory: 78GB) -- enough for 10GB + self.assertEqual(len(selected), 1) + self.assertEqual(selected[0], 1) + + def test_two_gpus_needed(self): + from utils.hardware.hardware import auto_select_gpu_ids + import utils.hardware.hardware as hw + + with ( + patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), + patch.object( + hw, + "estimate_required_model_memory_gb", + return_value = ( + 50.0, + { + "mode": "inference", + "required_gb": 50.0, + "model_size_source": "config", + "model_size_gb": 38.0, + }, + ), + ), + patch.object( + hw, + "_get_parent_visible_gpu_spec", + return_value = { + "raw": "0,1", + "numeric_ids": [0, 1], + "supports_explicit_gpu_ids": True, + }, + ), + patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]), + patch.object( + hw, + "get_visible_gpu_utilization", + return_value = self._make_utilization( + [ + (0, 40.0, 30.0), # 30GB free + (1, 40.0, 35.0), # 35GB free + ] + ), + ), + ): + selected, meta = auto_select_gpu_ids("test/model") + # 35GB (first) + 30*0.85 (second) = 60.5GB > 50GB + self.assertEqual(len(selected), 2) + + def test_non_cuda_returns_none(self): + from utils.hardware.hardware import auto_select_gpu_ids + import utils.hardware.hardware as hw + + with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU): + selected, meta = auto_select_gpu_ids("test/model") + self.assertIsNone(selected) + self.assertEqual(meta["selection_mode"], "non_cuda") + + +class TestGetDeviceMap(unittest.TestCase): + """Test device_map string generation.""" + + def test_single_gpu_returns_sequential(self): + from utils.hardware.hardware import get_device_map + import utils.hardware.hardware as hw + + with ( + patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), + patch.object( + hw, + "_get_parent_visible_gpu_spec", + return_value = { + "raw": "0", + "numeric_ids": [0], + "supports_explicit_gpu_ids": True, + }, + ), + patch.object(hw, "get_visible_gpu_count", return_value = 1), + ): + dm = get_device_map(gpu_ids = [0]) + self.assertEqual(dm, "sequential") + + def test_multi_gpu_returns_balanced(self): + from utils.hardware.hardware import get_device_map + import utils.hardware.hardware as hw + + with patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA): + dm = get_device_map(gpu_ids = [0, 1]) + self.assertEqual(dm, "balanced") + + def test_cpu_returns_sequential(self): + from utils.hardware.hardware import get_device_map + import utils.hardware.hardware as hw + + with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU): + dm = get_device_map(gpu_ids = None) + self.assertEqual(dm, "sequential") + + +class TestResolveRequestedGpuIds(unittest.TestCase): + """Test GPU ID validation.""" + + def test_none_returns_parent_visible(self): + from utils.hardware.hardware import resolve_requested_gpu_ids + + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + result = resolve_requested_gpu_ids(None) + self.assertEqual(result, [2, 3]) + + def test_empty_list_returns_parent_visible(self): + from utils.hardware.hardware import resolve_requested_gpu_ids + + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + result = resolve_requested_gpu_ids([]) + self.assertEqual(result, [2, 3]) + + def test_duplicates_rejected(self): + from utils.hardware.hardware import resolve_requested_gpu_ids + + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1,2"}, clear = False), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + with self.assertRaises(ValueError): + resolve_requested_gpu_ids([1, 1]) + + def test_out_of_range_rejected(self): + from utils.hardware.hardware import resolve_requested_gpu_ids + + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"}, clear = False), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 4), + ): + with self.assertRaises(ValueError): + resolve_requested_gpu_ids([5]) + + def test_uuid_env_var_rejects_explicit_ids(self): + from utils.hardware.hardware import resolve_requested_gpu_ids + + with ( + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-abc,GPU-def"}, clear = False + ), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + with self.assertRaises(ValueError): + resolve_requested_gpu_ids([0]) + + +class TestApplyGpuIds(unittest.TestCase): + """Test CUDA_VISIBLE_DEVICES environment variable setting.""" + + def test_apply_list(self): + from utils.hardware.hardware import apply_gpu_ids + + with patch.dict(os.environ, {}, clear = False): + apply_gpu_ids([3, 5]) + self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), "3,5") + + def test_apply_none_does_nothing(self): + from utils.hardware.hardware import apply_gpu_ids + + original = os.environ.get("CUDA_VISIBLE_DEVICES") + apply_gpu_ids(None) + self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), original) + + +class TestMultiGpuOverheadAccounting(unittest.TestCase): + """Test that multi-GPU overhead is applied correctly. + + The first GPU should keep its full free memory, and only + additional GPUs should have the overhead factor applied. + """ + + def _make_utilization(self, devices): + return { + "available": True, + "devices": [ + { + "index": idx, + "vram_total_gb": total, + "vram_used_gb": total - free, + } + for idx, total, free in devices + ], + } + + def test_first_gpu_not_penalized(self): + """A model that just fits on 1 GPU should not require 2 GPUs.""" + from utils.hardware.hardware import auto_select_gpu_ids + import utils.hardware.hardware as hw + + # Model requires 79GB, GPU has 80GB free + with ( + patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), + patch.object( + hw, + "estimate_required_model_memory_gb", + return_value = ( + 79.0, + { + "mode": "inference", + "required_gb": 79.0, + "model_size_source": "config", + "model_size_gb": 60.0, + }, + ), + ), + patch.object( + hw, + "_get_parent_visible_gpu_spec", + return_value = { + "raw": "0,1", + "numeric_ids": [0, 1], + "supports_explicit_gpu_ids": True, + }, + ), + patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]), + patch.object( + hw, + "get_visible_gpu_utilization", + return_value = self._make_utilization( + [ + (0, 80.0, 80.0), + (1, 80.0, 80.0), + ] + ), + ), + ): + selected, meta = auto_select_gpu_ids("test/model") + # Should fit on 1 GPU (80GB >= 79GB) + self.assertEqual(len(selected), 1) + + def test_second_gpu_has_overhead(self): + """When 2 GPUs are needed, the second one's contribution is reduced.""" + from utils.hardware.hardware import auto_select_gpu_ids + import utils.hardware.hardware as hw + + # Model requires 110GB. First GPU has 80GB, second has 40GB. + # With overhead: 80 + 40*0.85 = 114GB -- just enough + with ( + patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), + patch.object( + hw, + "estimate_required_model_memory_gb", + return_value = ( + 110.0, + { + "mode": "inference", + "required_gb": 110.0, + "model_size_source": "config", + "model_size_gb": 85.0, + }, + ), + ), + patch.object( + hw, + "_get_parent_visible_gpu_spec", + return_value = { + "raw": "0,1", + "numeric_ids": [0, 1], + "supports_explicit_gpu_ids": True, + }, + ), + patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]), + patch.object( + hw, + "get_visible_gpu_utilization", + return_value = self._make_utilization( + [ + (0, 80.0, 80.0), + (1, 80.0, 40.0), + ] + ), + ), + ): + selected, meta = auto_select_gpu_ids("test/model") + # Should use both GPUs + self.assertEqual(len(selected), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 3c33b33cb3..50557c6718 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -285,7 +285,7 @@ class TestLogGpuMemory: def test_does_not_raise(self): log_gpu_memory("test") - def test_logs_gpu_info_when_available(self, caplog): + def test_logs_gpu_info_when_available(self, capfd): fake_info = { "available": True, "backend": "cuda", @@ -295,35 +295,27 @@ class TestLogGpuMemory: "utilization_pct": 12.5, "free_gb": 14.0, } - import structlog - from loggers import get_logger - with ( - patch( - "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info - ), - caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"), + with patch( + "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info ): log_gpu_memory("unit-test") - assert "unit-test" in caplog.text - assert "CUDA" in caplog.text - assert "FakeGPU" in caplog.text + captured = capfd.readouterr() + assert "unit-test" in captured.out + assert "CUDA" in captured.out + assert "FakeGPU" in captured.out - def test_logs_cpu_fallback_when_no_gpu(self, caplog): + def test_logs_cpu_fallback_when_no_gpu(self, capfd): fake_info = {"available": False, "backend": "cpu"} - import structlog - from loggers import get_logger - with ( - patch( - "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info - ), - caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"), + with patch( + "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info ): log_gpu_memory("cpu-test") - assert "No GPU available" in caplog.text + captured = capfd.readouterr() + assert "No GPU available" in captured.out # ========== format_error_message() ========== diff --git a/studio/backend/tests/test_vram_estimation.py b/studio/backend/tests/test_vram_estimation.py new file mode 100644 index 0000000000..0be067310d --- /dev/null +++ b/studio/backend/tests/test_vram_estimation.py @@ -0,0 +1,695 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import unittest +from types import SimpleNamespace + +from utils.hardware.vram_estimation import ( + ModelArchConfig, + TrainingVramConfig, + extract_arch_config, + compute_model_weights_bytes, + compute_total_params, + compute_lora_params, + compute_lora_adapter_bytes, + compute_optimizer_bytes, + compute_gradient_bytes, + compute_activation_bytes, + estimate_training_vram, + DEFAULT_TARGET_MODULES, +) + + +def _gb(b: int) -> float: + return b / (1024**3) + + +LLAMA_8B = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 128256, + tie_word_embeddings = False, +) + +QWEN_05B = ModelArchConfig( + hidden_size = 896, + num_hidden_layers = 24, + num_attention_heads = 14, + num_key_value_heads = 2, + intermediate_size = 4864, + vocab_size = 151936, + tie_word_embeddings = True, +) + +MOE_CONFIG = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 8, +) + +DEEPSEEK_V3 = ModelArchConfig( + hidden_size = 7168, + num_hidden_layers = 61, + num_attention_heads = 128, + num_key_value_heads = 128, + intermediate_size = 18432, + vocab_size = 129280, + tie_word_embeddings = False, + num_experts = 256, + moe_intermediate_size = 2048, + n_shared_experts = 1, + num_dense_layers = 3, + q_lora_rank = 1536, + kv_lora_rank = 512, + qk_nope_head_dim = 128, + qk_rope_head_dim = 64, + v_head_dim = 128, +) + +QWEN3_MOE_30B = ModelArchConfig( + hidden_size = 2048, + num_hidden_layers = 48, + num_attention_heads = 32, + num_key_value_heads = 4, + intermediate_size = 8192, + vocab_size = 151936, + tie_word_embeddings = True, + num_experts = 128, + moe_intermediate_size = 768, + n_shared_experts = 0, + num_dense_layers = 0, +) + +GLM4_MOE = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 46, + num_attention_heads = 96, + num_key_value_heads = 8, + intermediate_size = 10944, + vocab_size = 151552, + tie_word_embeddings = False, + num_experts = 128, + moe_intermediate_size = 1408, + n_shared_experts = 1, + num_dense_layers = 1, +) + +GPT_OSS = ModelArchConfig( + hidden_size = 6144, + num_hidden_layers = 64, + num_attention_heads = 64, + num_key_value_heads = 8, + intermediate_size = 2880, + vocab_size = 200064, + tie_word_embeddings = False, + num_experts = 128, + moe_intermediate_size = None, + n_shared_experts = 0, + num_dense_layers = 0, +) + + +class TestExtractArchConfig(unittest.TestCase): + def test_basic_config(self): + hf_config = SimpleNamespace( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 128256, + tie_word_embeddings = False, + ) + arch = extract_arch_config(hf_config) + self.assertIsNotNone(arch) + self.assertEqual(arch.hidden_size, 4096) + self.assertEqual(arch.num_hidden_layers, 32) + self.assertEqual(arch.num_key_value_heads, 8) + self.assertIsNone(arch.num_experts) + + def test_vlm_text_config(self): + text_cfg = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 24, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 8192, + vocab_size = 32000, + tie_word_embeddings = True, + ) + hf_config = SimpleNamespace(text_config = text_cfg) + arch = extract_arch_config(hf_config) + self.assertIsNotNone(arch) + self.assertEqual(arch.hidden_size, 2048) + + def test_moe_detection(self): + hf_config = SimpleNamespace( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_local_experts = 8, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_experts, 8) + + def test_missing_fields_returns_none(self): + hf_config = SimpleNamespace(hidden_size = 4096) + arch = extract_arch_config(hf_config) + self.assertIsNone(arch) + + def test_intermediate_size_list(self): + hf_config = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 24, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = [8192, 8192], + vocab_size = 32000, + tie_word_embeddings = True, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.intermediate_size, 8192) + + +class TestModelWeightsBytes(unittest.TestCase): + def test_llama_8b_fp16(self): + weight_bytes = compute_model_weights_bytes(LLAMA_8B, "full", False) + weight_gb = _gb(weight_bytes) + self.assertGreater(weight_gb, 14.0) + self.assertLess(weight_gb, 18.0) + + def test_llama_8b_qlora_4bit(self): + weight_bytes = compute_model_weights_bytes(LLAMA_8B, "qlora", True) + weight_gb = _gb(weight_bytes) + self.assertGreater(weight_gb, 4.0) + self.assertLess(weight_gb, 7.0) + + def test_4bit_smaller_than_fp16(self): + fp16 = compute_model_weights_bytes(LLAMA_8B, "full", False) + q4 = compute_model_weights_bytes(LLAMA_8B, "qlora", True) + self.assertLess(q4, fp16) + ratio = fp16 / q4 + self.assertGreater(ratio, 2.0) + self.assertLess(ratio, 4.0) + + def test_moe_larger_than_dense(self): + dense = compute_model_weights_bytes(LLAMA_8B, "full", False) + moe = compute_model_weights_bytes(MOE_CONFIG, "full", False) + self.assertGreater(moe, dense * 3) + + +class TestLoraParams(unittest.TestCase): + def test_llama_8b_default_modules_rank16(self): + lora_p = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES) + total_p = compute_total_params(LLAMA_8B) + ratio = lora_p / total_p + self.assertGreater(ratio, 0.005) + self.assertLess(ratio, 0.05) + + def test_higher_rank_more_params(self): + r16 = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES) + r64 = compute_lora_params(LLAMA_8B, 64, DEFAULT_TARGET_MODULES) + self.assertAlmostEqual(r64 / r16, 4.0, places = 1) + + def test_fewer_modules_fewer_params(self): + all_mods = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES) + qv_only = compute_lora_params(LLAMA_8B, 16, ["q_proj", "v_proj"]) + self.assertLess(qv_only, all_mods) + + def test_moe_mlp_modules_scale_with_experts(self): + dense_lora = compute_lora_params( + LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"] + ) + moe_lora = compute_lora_params( + MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"] + ) + ratio = moe_lora / dense_lora + self.assertAlmostEqual(ratio, 8.0, delta = 0.5) + + def test_attention_modules_same_for_moe(self): + dense_attn = compute_lora_params( + LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"] + ) + moe_attn = compute_lora_params( + MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"] + ) + self.assertEqual(dense_attn, moe_attn) + + +class TestOptimizerBytes(unittest.TestCase): + def test_adamw_8bit(self): + self.assertEqual(compute_optimizer_bytes(1_000_000, "adamw_8bit"), 4_000_000) + + def test_adamw_torch(self): + self.assertEqual(compute_optimizer_bytes(1_000_000, "adamw_torch"), 6_000_000) + + def test_sgd(self): + self.assertEqual(compute_optimizer_bytes(1_000_000, "sgd"), 4_000_000) + + def test_unknown_defaults_to_4(self): + self.assertEqual(compute_optimizer_bytes(1_000_000, "some_new_opt"), 4_000_000) + + +class TestGradientBytes(unittest.TestCase): + def test_fp16_gradients(self): + self.assertEqual(compute_gradient_bytes(1_000_000), 2_000_000) + + +class TestActivationBytes(unittest.TestCase): + def test_no_gc_scales_with_layers(self): + act_none = compute_activation_bytes(LLAMA_8B, 2, 2048, "none") + act_gc = compute_activation_bytes(LLAMA_8B, 2, 2048, "true") + self.assertGreater(act_none, act_gc * 10) + + def test_unsloth_gc_smaller_than_standard(self): + act_true = compute_activation_bytes(LLAMA_8B, 2, 2048, "true") + act_unsloth = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth") + self.assertLess(act_unsloth, act_true) + + def test_lora_activations_smaller_than_full_ft(self): + full_ft = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth", is_lora = False) + lora = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth", is_lora = True) + self.assertLess(lora, full_ft) + + def test_scales_with_batch_size(self): + act_bsz2 = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth") + act_bsz4 = compute_activation_bytes(LLAMA_8B, 4, 2048, "unsloth") + self.assertAlmostEqual(act_bsz4 / act_bsz2, 2.0, delta = 0.1) + + def test_scales_with_seq_len(self): + act_2k = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth") + act_4k = compute_activation_bytes(LLAMA_8B, 2, 4096, "unsloth") + self.assertAlmostEqual(act_4k / act_2k, 2.0, delta = 0.1) + + +class TestEstimateTrainingVram(unittest.TestCase): + def test_llama_8b_qlora_reasonable_total(self): + config = TrainingVramConfig( + training_method = "qlora", + batch_size = 2, + max_seq_length = 2048, + lora_rank = 16, + gradient_checkpointing = "unsloth", + optimizer = "adamw_8bit", + load_in_4bit = True, + ) + breakdown = estimate_training_vram(LLAMA_8B, config) + total_gb = _gb(breakdown.total) + self.assertGreater(total_gb, 5.0) + self.assertLess(total_gb, 12.0) + + def test_llama_8b_full_ft_reasonable_total(self): + config = TrainingVramConfig( + training_method = "full", + batch_size = 2, + max_seq_length = 2048, + gradient_checkpointing = "unsloth", + optimizer = "adamw_8bit", + load_in_4bit = False, + ) + breakdown = estimate_training_vram(LLAMA_8B, config) + total_gb = _gb(breakdown.total) + self.assertGreater(total_gb, 50.0) + self.assertLess(total_gb, 75.0) + + def test_qlora_much_less_than_full_ft(self): + qlora_config = TrainingVramConfig( + training_method = "qlora", + load_in_4bit = True, + batch_size = 2, + max_seq_length = 2048, + ) + full_config = TrainingVramConfig( + training_method = "full", + load_in_4bit = False, + batch_size = 2, + max_seq_length = 2048, + ) + qlora = estimate_training_vram(LLAMA_8B, qlora_config) + full = estimate_training_vram(LLAMA_8B, full_config) + self.assertLess(qlora.total, full.total / 3) + + def test_qwen_05b_qlora_fits_in_4gb(self): + config = TrainingVramConfig( + training_method = "qlora", + batch_size = 2, + max_seq_length = 2048, + lora_rank = 16, + gradient_checkpointing = "unsloth", + optimizer = "adamw_8bit", + load_in_4bit = True, + ) + breakdown = estimate_training_vram(QWEN_05B, config) + total_gb = _gb(breakdown.total) + self.assertLess(total_gb, 5.0) + + def test_breakdown_components_positive(self): + config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) + breakdown = estimate_training_vram(LLAMA_8B, config) + self.assertGreater(breakdown.model_weights, 0) + self.assertGreater(breakdown.lora_adapters, 0) + self.assertGreater(breakdown.optimizer_states, 0) + self.assertGreater(breakdown.gradients, 0) + self.assertGreater(breakdown.activations, 0) + self.assertGreater(breakdown.cuda_overhead, 0) + + def test_full_ft_no_lora_adapters(self): + config = TrainingVramConfig(training_method = "full", load_in_4bit = False) + breakdown = estimate_training_vram(LLAMA_8B, config) + self.assertEqual(breakdown.lora_adapters, 0) + + def test_to_gb_dict_keys(self): + config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) + breakdown = estimate_training_vram(LLAMA_8B, config) + gb_dict = breakdown.to_gb_dict() + expected_keys = { + "model_weights_gb", + "lora_adapters_gb", + "optimizer_states_gb", + "gradients_gb", + "activations_gb", + "cuda_overhead_gb", + "total_gb", + } + self.assertEqual(set(gb_dict.keys()), expected_keys) + + def test_total_equals_sum_of_parts(self): + config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) + breakdown = estimate_training_vram(LLAMA_8B, config) + parts_sum = ( + breakdown.model_weights + + breakdown.lora_adapters + + breakdown.optimizer_states + + breakdown.gradients + + breakdown.activations + + breakdown.cuda_overhead + ) + self.assertEqual(breakdown.total, parts_sum) + + def test_larger_batch_increases_total(self): + small = TrainingVramConfig( + training_method = "qlora", + load_in_4bit = True, + batch_size = 1, + ) + large = TrainingVramConfig( + training_method = "qlora", + load_in_4bit = True, + batch_size = 8, + ) + small_v = estimate_training_vram(LLAMA_8B, small) + large_v = estimate_training_vram(LLAMA_8B, large) + self.assertGreater(large_v.total, small_v.total) + + def test_adamw_fp32_uses_more_optimizer_memory(self): + opt8 = TrainingVramConfig( + training_method = "full", + load_in_4bit = False, + optimizer = "adamw_8bit", + ) + opt32 = TrainingVramConfig( + training_method = "full", + load_in_4bit = False, + optimizer = "adamw_torch", + ) + v8 = estimate_training_vram(LLAMA_8B, opt8) + v32 = estimate_training_vram(LLAMA_8B, opt32) + self.assertAlmostEqual( + v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1 + ) + + +class TestExtractArchConfigMoE(unittest.TestCase): + def test_deepseek_v3_shared_experts(self): + hf_config = SimpleNamespace( + hidden_size = 7168, + num_hidden_layers = 61, + num_attention_heads = 128, + num_key_value_heads = 128, + intermediate_size = 18432, + vocab_size = 129280, + tie_word_embeddings = False, + n_routed_experts = 256, + moe_intermediate_size = 2048, + n_shared_experts = 1, + first_k_dense_replace = 3, + q_lora_rank = 1536, + kv_lora_rank = 512, + qk_nope_head_dim = 128, + qk_rope_head_dim = 64, + v_head_dim = 128, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_experts, 256) + self.assertEqual(arch.n_shared_experts, 1) + self.assertEqual(arch.num_dense_layers, 3) + self.assertEqual(arch.q_lora_rank, 1536) + self.assertEqual(arch.kv_lora_rank, 512) + + def test_qwen3_moe_decoder_sparse_step(self): + hf_config = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 48, + num_attention_heads = 32, + num_key_value_heads = 4, + intermediate_size = 8192, + vocab_size = 151936, + tie_word_embeddings = True, + num_local_experts = 128, + moe_intermediate_size = 768, + decoder_sparse_step = 1, + mlp_only_layers = [], + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_experts, 128) + self.assertEqual(arch.num_dense_layers, 0) + self.assertIsNone(arch.q_lora_rank) + + def test_qwen3_moe_with_mlp_only_layers(self): + hf_config = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 24, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 8192, + vocab_size = 151936, + tie_word_embeddings = True, + num_local_experts = 60, + moe_intermediate_size = 1408, + decoder_sparse_step = 1, + mlp_only_layers = [0, 1, 2, 3], + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_dense_layers, 4) + + def test_glm4_moe_first_k_dense(self): + hf_config = SimpleNamespace( + hidden_size = 4096, + num_hidden_layers = 46, + num_attention_heads = 96, + num_key_value_heads = 8, + intermediate_size = 10944, + vocab_size = 151552, + tie_word_embeddings = False, + n_routed_experts = 128, + moe_intermediate_size = 1408, + n_shared_experts = 1, + first_k_dense_replace = 1, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_dense_layers, 1) + self.assertEqual(arch.n_shared_experts, 1) + + def test_gpt_oss_no_moe_intermediate(self): + hf_config = SimpleNamespace( + hidden_size = 6144, + num_hidden_layers = 64, + num_attention_heads = 64, + num_key_value_heads = 8, + intermediate_size = 2880, + vocab_size = 200064, + tie_word_embeddings = False, + num_local_experts = 128, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_experts, 128) + self.assertIsNone(arch.moe_intermediate_size) + self.assertEqual(arch.num_dense_layers, 0) + + def test_backward_compat_no_new_fields(self): + hf_config = SimpleNamespace( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 128256, + tie_word_embeddings = False, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.n_shared_experts, 0) + self.assertEqual(arch.num_dense_layers, 0) + self.assertIsNone(arch.q_lora_rank) + + +class TestSharedExperts(unittest.TestCase): + def test_shared_experts_increase_weight_bytes(self): + no_shared = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 64, + moe_intermediate_size = 1407, + n_shared_experts = 0, + ) + with_shared = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 64, + moe_intermediate_size = 1407, + n_shared_experts = 2, + ) + w_no = compute_model_weights_bytes(no_shared, "full", False) + w_yes = compute_model_weights_bytes(with_shared, "full", False) + self.assertGreater(w_yes, w_no) + delta_per_layer = 4096 * 1407 * 3 * 2 + expected_delta = delta_per_layer * 32 * 2 + actual_delta = w_yes - w_no + self.assertAlmostEqual( + actual_delta, expected_delta, delta = expected_delta * 0.01 + ) + + def test_deepseek_v3_params_in_range(self): + total = compute_total_params(DEEPSEEK_V3) + total_b = total / 1e9 + self.assertGreater(total_b, 600) + self.assertLess(total_b, 750) + + +class TestMLA(unittest.TestCase): + def test_mla_different_from_standard(self): + from utils.hardware.vram_estimation import _compute_attn_elements + + mla_arch = DEEPSEEK_V3 + std_arch = ModelArchConfig( + hidden_size = 7168, + num_hidden_layers = 61, + num_attention_heads = 128, + num_key_value_heads = 128, + intermediate_size = 18432, + vocab_size = 129280, + ) + mla_attn = _compute_attn_elements(mla_arch) + std_attn = _compute_attn_elements(std_arch) + self.assertNotEqual(mla_attn, std_attn) + + def test_mla_lora_produces_values(self): + lora_p = compute_lora_params(DEEPSEEK_V3, 16, ["q_proj", "v_proj", "o_proj"]) + self.assertGreater(lora_p, 0) + + +class TestDenseMoEMix(unittest.TestCase): + def test_dense_layers_change_total(self): + all_moe = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 46, + num_attention_heads = 96, + num_key_value_heads = 8, + intermediate_size = 10944, + vocab_size = 151552, + tie_word_embeddings = False, + num_experts = 128, + moe_intermediate_size = 1408, + n_shared_experts = 1, + num_dense_layers = 0, + ) + mixed = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 46, + num_attention_heads = 96, + num_key_value_heads = 8, + intermediate_size = 10944, + vocab_size = 151552, + tie_word_embeddings = False, + num_experts = 128, + moe_intermediate_size = 1408, + n_shared_experts = 1, + num_dense_layers = 1, + ) + w_all = compute_model_weights_bytes(all_moe, "full", False) + w_mixed = compute_model_weights_bytes(mixed, "full", False) + self.assertNotEqual(w_all, w_mixed) + + def test_glm4_moe_params_reasonable(self): + total = compute_total_params(GLM4_MOE) + total_b = total / 1e9 + self.assertGreater(total_b, 80) + self.assertLess(total_b, 120) + + def test_qwen3_moe_30b_params_reasonable(self): + total = compute_total_params(QWEN3_MOE_30B) + total_b = total / 1e9 + self.assertGreater(total_b, 20) + self.assertLess(total_b, 50) + + def test_gpt_oss_uses_intermediate_size(self): + total = compute_total_params(GPT_OSS) + total_b = total / 1e9 + self.assertGreater(total_b, 350) + self.assertLess(total_b, 500) + + def test_lora_dense_vs_moe_layers_differ(self): + all_moe = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 10, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 8, + moe_intermediate_size = 1024, + num_dense_layers = 0, + ) + mixed = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 10, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 8, + moe_intermediate_size = 1024, + num_dense_layers = 5, + ) + lora_all = compute_lora_params( + all_moe, 16, ["gate_proj", "up_proj", "down_proj"] + ) + lora_mix = compute_lora_params(mixed, 16, ["gate_proj", "up_proj", "down_proj"]) + self.assertNotEqual(lora_all, lora_mix) + + +if __name__ == "__main__": + unittest.main() diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index 39f0113fe0..289b30e55e 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -553,13 +553,9 @@ def convert_to_vlm_format( batch_results[idx] = future.result() except Exception as e: failed_count += 1 - if failed_count == 1: - print( - f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}" - ) if failed_count == 1: logger.info( - f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}" + f"First VLM conversion failure: {type(e).__name__}: {e}" ) converted_list.extend(r for r in batch_results if r is not None) @@ -583,13 +579,10 @@ def convert_to_vlm_format( converted_list.append(_convert_single_sample(sample)) except Exception as e: failed_count += 1 - if failed_count == 1: - # Log the first failure to aid debugging - print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}") if failed_count == 1: # Log the first failure to aid debugging logger.info( - f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}" + f"First VLM conversion failure: {type(e).__name__}: {e}" ) pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False) pbar.close() diff --git a/studio/backend/utils/hardware/VRAM_ESTIMATION.md b/studio/backend/utils/hardware/VRAM_ESTIMATION.md new file mode 100644 index 0000000000..26072b208f --- /dev/null +++ b/studio/backend/utils/hardware/VRAM_ESTIMATION.md @@ -0,0 +1,161 @@ +# VRAM Estimation for Training + +``` +Total VRAM = Weights + LoRA Adapters + Optimizer + Gradients + Activations + CUDA Overhead +``` + +| Symbol | Meaning | +|--------|---------| +| `H` | `hidden_size` | +| `L` | `num_hidden_layers` | +| `V` | `vocab_size` | +| `K` | `(H / num_attention_heads) * num_key_value_heads` | +| `M` | `intermediate_size` (or `moe_intermediate_size`) | +| `E` | `num_experts` (1 for dense) | +| `r` | LoRA rank | +| `B` | `per_device_train_batch_size` | +| `S` | `max_seq_length` | + +--- + +## 1. Model Weights + +``` +QKVO = (H + K + K + H) * H +MLP = H * M * 3 * E + (E * H if E > 1 else 0) + +Quantizable = (QKVO + MLP) * L +Non-quantizable = 2*H*L + V*H + (V*H if not tie_embeddings else 0) +``` + +| Mode | Bytes | +|------|-------| +| QLoRA 4-bit | `Quantizable * 2 / 3.2 + Non-quantizable * 2` | +| LoRA / Full fp16 | `(Quantizable + Non-quantizable) * 2` | + +The 3.2 factor (`16/5`) accounts for BNB NF4 blockwise scales. + +## 2. LoRA Adapters + +| Module | A | B | +|--------|---|---| +| q_proj | `H×r` | `r×H` | +| k_proj | `H×r` | `r×K` | +| v_proj | `H×r` | `r×K` | +| o_proj | `H×r` | `r×H` | +| gate_proj | `H×r` | `r×M` | +| up_proj | `H×r` | `r×M` | +| down_proj | `M×r` | `r×H` | + +MLP modules multiply by `E` for MoE. + +``` +LoRA_bytes = sum(A + B per selected module) * L * 2 +``` + +## 3. Optimizer States (calibrated) + +| Optimizer | Bytes/param | Notes | +|-----------|------------|-------| +| `adamw_8bit` | 4 | BNB upcasts to fp32 during step | +| `adamw_torch` | 6 | Fused, no master copy | +| `paged_adamw_32bit` | 8 | Full fp32 states | +| `sgd` | 4 | | + +Trainable params = all params (Full FT) or LoRA params only. + +## 4. Gradients + +``` +Gradient_bytes = trainable_params * 2 (fp16, accumulated in-place) +``` + +## 5. Activations + +Per-layer (from `unsloth_zoo/vllm_utils.py`): +``` +Per_layer = (S*B*(H+K+K) + S*B*2 + S*B*(M+M)) * 2 * 1.25 +``` + +| GC Mode | Full FT | LoRA/QLoRA | +|---------|---------|------------| +| none | `L` layers | `L` layers | +| true (HF) | 2.0 | 1.0 | +| unsloth | 1.5 | 1.0 | + +## 6. Floors + +Gradients and activations have minimum floors at **15% of model weight memory** to account for autograd overhead, attention score matrices, NCCL buffers, mixed-precision scaling, and PyTorch fragmentation. + +``` +gradient_bytes = max(computed, weights * 0.15) +activation_bytes = max(computed, weights * 0.15 * B/2) +``` + +## 7. CUDA Overhead + +**1.4 GB** fixed — CUDA driver + PyTorch runtime, calibrated on RTX 5070 Ti. + +## 8. Multi-GPU Overhead + +When sharding across multiple GPUs, each additional GPU (beyond the first) contributes only **85%** of its free VRAM to the usable pool. The 15% discount accounts for NCCL all-reduce buffers, PCIe/NVLink transfer overhead, synchronization barriers, and memory fragmentation from non-uniform shard sizes. Calibrated empirically on 2-8 GPU setups with NVLink and PCIe topologies. + +``` +usable_gb = free[gpu_0] + sum(free[gpu_i] * 0.85 for i in 1..N) +``` + +--- + +## Reference Table (bsz=2, seq=2048, rank=16, GC=unsloth, adamw_8bit) + +| Model | Weights | LoRA | Optim | Grad | Act | CUDA | Total | +|-------|---------|------|-------|------|-----|------|-------| +| 0.5B QLoRA | 0.5 | 0.0 | 0.0 | 0.1 | 0.1 | 1.4 | **2.1** | +| 1B QLoRA | 1.1 | 0.0 | 0.0 | 0.2 | 0.2 | 1.4 | **2.9** | +| 3B QLoRA | 2.4 | 0.0 | 0.1 | 0.5 | 0.5 | 1.4 | **4.9** | +| 8B QLoRA | 6.0 | 0.1 | 0.2 | 1.2 | 1.2 | 1.4 | **10.1** | +| 8B LoRA fp16 | 15.0 | 0.1 | 0.2 | 3.0 | 3.0 | 1.4 | **22.6** | +| 8B Full FT | 15.0 | — | 29.9 | 15.0 | 3.0 | 1.4 | **64.2** | +| 32B LoRA fp16 | 61.0 | 0.2 | 0.5 | 12.2 | 12.2 | 1.4 | **87.6** | +| 72B QLoRA | 45.5 | 0.4 | 0.8 | 9.1 | 9.1 | 1.4 | **66.3** | + +## E2E Validation (Llama-3.2-1B, B200 emulating 24GB) + +| Config | Estimated | Actual (nvsmi) | Error | +|--------|----------|----------------|-------| +| QLoRA bsz=2 seq=512 | 2.55 GB | 2.65 GB | -3.7% | +| QLoRA bsz=2 seq=2048 | 2.60 GB | 2.65 GB | -1.8% | +| QLoRA bsz=4 seq=2048 | 2.65 GB | 2.65 GB | +0.0% | +| LoRA fp16 bsz=2 | 3.84 GB | 3.88 GB | -1.0% | +| Full FT adamw_8bit | 10.89 GB | 10.80 GB | +0.8% | +| Full FT adamw_torch | 13.19 GB | 12.93 GB | +2.0% | + +*Note: e2e numbers predate the 15% floors, which add safety margin on top.* + +--- + +## Parameter Flow + +``` +Frontend -> routes/{training,inference}.py + -> prepare_gpu_selection(gpu_ids, model_name, ...) + | + +-- gpu_ids is explicit (e.g. [5,6,7]) + | -> resolve_requested_gpu_ids: validate against parent-visible set + | -> return all requested GPUs (model sharded across all of them) + | + +-- gpu_ids is None or [] + -> auto_select_gpu_ids: estimate VRAM, pick minimum GPUs needed + -> estimate_required_model_memory_gb -> estimate_training_vram + -> greedy selection: rank GPUs by free VRAM, add until model fits + + -> get_device_map(resolved_gpu_ids) + -> "balanced" if >1 GPU, "sequential" otherwise + + -> worker subprocess: apply_gpu_ids(resolved_gpu_ids) + -> sets CUDA_VISIBLE_DEVICES before torch/CUDA init +``` + +Threaded params: `batch_size`, `max_seq_length`, `lora_r`, `target_modules`, `gradient_checkpointing`, `optim`. + +Source: `studio/backend/utils/hardware/vram_estimation.py` diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index f86a56d186..aaa0452406 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -18,11 +18,31 @@ from .hardware import ( get_gpu_summary, get_package_versions, get_gpu_utilization, + get_visible_gpu_utilization, + get_backend_visible_gpu_info, get_physical_gpu_count, get_visible_gpu_count, + get_parent_visible_gpu_ids, + resolve_requested_gpu_ids, + estimate_fp16_model_size_bytes, + estimate_required_model_memory_gb, + auto_select_gpu_ids, + prepare_gpu_selection, safe_num_proc, safe_thread_num_proc, dataset_map_num_proc, + get_device_map, + get_offloaded_device_map_entries, + raise_if_offloaded, + apply_gpu_ids, +) + +from .vram_estimation import ( + ModelArchConfig, + TrainingVramConfig, + VramBreakdown, + extract_arch_config, + estimate_training_vram, ) __all__ = [ @@ -38,9 +58,26 @@ __all__ = [ "get_gpu_summary", "get_package_versions", "get_gpu_utilization", + "get_visible_gpu_utilization", + "get_backend_visible_gpu_info", "get_physical_gpu_count", "get_visible_gpu_count", + "get_parent_visible_gpu_ids", + "resolve_requested_gpu_ids", + "estimate_fp16_model_size_bytes", + "estimate_required_model_memory_gb", + "auto_select_gpu_ids", + "prepare_gpu_selection", "safe_num_proc", "safe_thread_num_proc", "dataset_map_num_proc", + "get_device_map", + "get_offloaded_device_map_entries", + "raise_if_offloaded", + "apply_gpu_ids", + "ModelArchConfig", + "TrainingVramConfig", + "VramBreakdown", + "extract_arch_config", + "estimate_training_vram", ] diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 61ee8a0967..742e8f6b7e 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -16,10 +16,12 @@ Usage: ... """ +import os import platform import structlog from loggers import get_logger from enum import Enum +from pathlib import Path from typing import Optional, Dict, Any logger = get_logger(__name__) @@ -32,6 +34,7 @@ class DeviceType(str, Enum): """Supported compute backends. Inherits from str so it serializes cleanly in JSON.""" CUDA = "cuda" + XPU = "xpu" MLX = "mlx" CPU = "cpu" @@ -96,6 +99,17 @@ def detect_hardware() -> DeviceType: print(f"Hardware detected: CUDA — {device_name}") return DEVICE + # --- XPU: Intel GPU --- + if _has_torch(): + import torch + + if hasattr(torch, "xpu") and torch.xpu.is_available(): + DEVICE = DeviceType.XPU + CHAT_ONLY = False + device_name = torch.xpu.get_device_name(0) + print(f"Hardware detected: XPU — {device_name}") + return DEVICE + # --- MLX: Apple Silicon --- if is_apple_silicon() and _has_mlx(): DEVICE = DeviceType.MLX @@ -140,6 +154,11 @@ def clear_gpu_cache(): torch.cuda.synchronize() torch.cuda.empty_cache() torch.cuda.ipc_collect() + elif device == DeviceType.XPU: + import torch + + torch.xpu.synchronize() + torch.xpu.empty_cache() elif device == DeviceType.MLX: # MLX manages memory automatically; no explicit cache clear needed. # mlx.core has no empty_cache equivalent — gc.collect() above is enough. @@ -180,6 +199,33 @@ def get_gpu_memory_info() -> Dict[str, Any]: logger.error(f"Error getting CUDA GPU info: {e}") return {"available": False, "backend": device.value, "error": str(e)} + # ---- XPU path (Intel GPU) ---- + if device == DeviceType.XPU: + try: + import torch + + idx = torch.xpu.current_device() + props = torch.xpu.get_device_properties(idx) + + total = props.total_memory + allocated = torch.xpu.memory_allocated(idx) + reserved = torch.xpu.memory_reserved(idx) + + return { + "available": True, + "backend": device.value, + "device": idx, + "device_name": props.name, + "total_gb": total / (1024**3), + "allocated_gb": allocated / (1024**3), + "reserved_gb": reserved / (1024**3), + "free_gb": (total - allocated) / (1024**3), + "utilization_pct": (allocated / total) * 100, + } + except Exception as e: + logger.error("Error getting XPU GPU info: %s", e) + return {"available": False, "backend": device.value, "error": str(e)} + # ---- MLX path (Apple Silicon) ---- if device == DeviceType.MLX: try: @@ -280,134 +326,199 @@ def get_package_versions() -> Dict[str, Optional[str]]: return versions -# ========== Live GPU Utilization (nvidia-smi) ========== +# ========== Torch-based GPU fallbacks (AMD ROCm, Intel XPU, nvidia-smi missing) ========== + + +def _torch_get_device_module(): + """Return the appropriate torch device module (cuda or xpu) and its name.""" + device = get_device() + import torch + + if device == DeviceType.CUDA: + return torch.cuda, "cuda" + if device == DeviceType.XPU and hasattr(torch, "xpu"): + return torch.xpu, "xpu" + return None, None + + +def _torch_get_physical_gpu_count() -> Optional[int]: + mod, _ = _torch_get_device_module() + if mod is None: + return None + try: + return mod.device_count() + except Exception: + return None + + +def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]]: + """Query torch for per-GPU name, total VRAM, and used VRAM.""" + mod, _ = _torch_get_device_module() + if mod is None: + return [] + + devices = [] + for ordinal, phys_idx in enumerate(device_indices): + try: + # torch uses 0-based ordinals relative to CUDA_VISIBLE_DEVICES + props = mod.get_device_properties(ordinal) + total_bytes = props.total_memory + # Prefer mem_get_info (reports system-wide usage, not just this + # process) so auto-selection accounts for other GPU consumers. + if hasattr(mod, "mem_get_info"): + free_bytes, total_bytes = mod.mem_get_info(ordinal) + used_bytes = total_bytes - free_bytes + else: + used_bytes = mod.memory_allocated(ordinal) + devices.append( + { + "index": phys_idx, + "visible_ordinal": ordinal, + "name": props.name, + "total_gb": round(total_bytes / (1024**3), 2), + "used_gb": round(used_bytes / (1024**3), 2), + } + ) + except Exception as e: + logger.debug("torch device query failed for ordinal %d: %s", ordinal, e) + return devices + + +# ========== Live GPU Utilization ========== def get_gpu_utilization() -> Dict[str, Any]: - """ - Return a live snapshot of GPU utilization via ``nvidia-smi``. - - Designed to be polled by the frontend during training (not streaming). - Uses ``nvidia-smi --query-gpu`` which is the most accurate source for - utilization %, temperature, and power draw – stats that PyTorch does - not expose. - - Returns dict with keys: - available – bool, whether stats could be retrieved - gpu_utilization_pct – GPU core utilization % - temperature_c – GPU temperature in °C - vram_used_gb – VRAM currently used (GiB) - vram_total_gb – VRAM total (GiB) - vram_utilization_pct – VRAM used / total * 100 - power_draw_w – current power draw (W) - power_limit_w – power limit (W) - power_utilization_pct – power draw / limit * 100 - """ + """Return a live snapshot of device utilization information.""" device = get_device() - if device != DeviceType.CUDA: - return {"available": False, "backend": device.value} - - def _parse_smi_value(raw: str): - """Parse a single nvidia-smi CSV value. Returns float or None for [N/A].""" - raw = raw.strip() - if not raw or raw == "[N/A]": - return None + if device == DeviceType.CUDA: try: - return float(raw) - except (ValueError, TypeError): - return None + from . import nvidia - # ── nvidia-smi (most complete source) ─────────────────────── - smi_data = {} - try: - import subprocess - - result = subprocess.run( - [ - "nvidia-smi", - "--query-gpu=utilization.gpu,temperature.gpu," - "memory.used,memory.total,power.draw,power.limit", - "--format=csv,noheader,nounits", - ], - capture_output = True, - text = True, - timeout = 5, - ) - - if result.returncode == 0 and result.stdout.strip(): - # nvidia-smi outputs one line per GPU; take GPU 0 - first_line = result.stdout.strip().splitlines()[0] - parts = [p.strip() for p in first_line.split(",")] - if len(parts) >= 6: - smi_data = { - "gpu_util": _parse_smi_value(parts[0]), - "temp": _parse_smi_value(parts[1]), - "vram_used_mb": _parse_smi_value(parts[2]), - "vram_total_mb": _parse_smi_value(parts[3]), - "power_draw": _parse_smi_value(parts[4]), - "power_limit": _parse_smi_value(parts[5]), - } - - except FileNotFoundError: - logger.debug("nvidia-smi not found, falling back to torch.cuda") - except Exception as e: - logger.warning(f"nvidia-smi query failed: {e}") - - # ── Backfill VRAM from torch.cuda if nvidia-smi returned [N/A] ── - vram_used_mb = smi_data.get("vram_used_mb") - vram_total_mb = smi_data.get("vram_total_mb") - - if vram_used_mb is None or vram_total_mb is None: - try: - import torch - - idx = torch.cuda.current_device() - props = torch.cuda.get_device_properties(idx) - if vram_total_mb is None: - vram_total_mb = props.total_memory / (1024**2) # bytes → MiB - if vram_used_mb is None: - vram_used_mb = torch.cuda.memory_allocated(idx) / (1024**2) + result = nvidia.get_primary_gpu_utilization() + if result.get("available"): + result["backend"] = device.value + return result except Exception as e: - logger.debug(f"torch.cuda VRAM backfill failed: {e}") + logger.warning("nvidia-smi utilization query failed: %s", e) - # ── Build response ────────────────────────────────────────── - gpu_util = smi_data.get("gpu_util") - temp = smi_data.get("temp") - power_draw = smi_data.get("power_draw") - power_limit = smi_data.get("power_limit") + mem = get_gpu_memory_info() + if device != DeviceType.CPU and mem.get("available"): + return { + "available": True, + "backend": device.value, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": round(mem.get("allocated_gb", 0), 2), + "vram_total_gb": round(mem.get("total_gb", 0), 2), + "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } - vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None - vram_total_gb = ( - round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None - ) - vram_pct = ( - round((vram_used_mb / vram_total_mb) * 100, 1) - if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0 - else None - ) - power_pct = ( - round((power_draw / power_limit) * 100, 1) - if power_draw is not None and power_limit and power_limit > 0 - else None - ) + return {"available": False, "backend": device.value} - # If we got at least something useful, report available - has_any = any(v is not None for v in [gpu_util, temp, vram_used_gb, power_draw]) - if not has_any: - return {"available": False, "backend": device.value} + +def get_visible_gpu_utilization() -> Dict[str, Any]: + device = get_device() + + if device == DeviceType.CUDA: + parent_visible_spec = _get_parent_visible_gpu_spec() + try: + from . import nvidia + + result = nvidia.get_visible_gpu_utilization( + parent_visible_spec["numeric_ids"], + parent_cuda_visible_devices = parent_visible_spec["raw"], + ) + if result.get("available"): + result["backend"] = device.value + return result + except Exception as e: + logger.warning("nvidia-smi visible GPU utilization query failed: %s", e) + + # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel) + if device in (DeviceType.CUDA, DeviceType.XPU): + parent_ids = get_parent_visible_gpu_ids() + # When parent_visible_ids is empty (UUID/MIG mask or no CVD set), + # enumerate torch-visible ordinals so the UI still shows devices. + if parent_ids: + torch_indices = parent_ids + index_kind = "physical" + else: + visible_count = _torch_get_physical_gpu_count() or 0 + torch_indices = list(range(visible_count)) + index_kind = "relative" + torch_devices = _torch_get_per_device_info(torch_indices) + if torch_devices: + devices = [] + for td in torch_devices: + total = td["total_gb"] + used = td["used_gb"] + devices.append( + { + "index": td["index"], + "index_kind": index_kind, + "visible_ordinal": td["visible_ordinal"], + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": used, + "vram_total_gb": total, + "vram_utilization_pct": round((used / total) * 100, 1) + if total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ) + return { + "available": True, + "backend": device.value, + "parent_visible_gpu_ids": parent_ids, + "devices": devices, + "index_kind": index_kind, + } + + if device == DeviceType.MLX: + mem = get_gpu_memory_info() + if not mem.get("available"): + return { + "available": False, + "backend": device.value, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "relative", + } + return { + "available": True, + "backend": device.value, + "parent_visible_gpu_ids": [0], + "devices": [ + { + "index": 0, + "index_kind": "relative", + "visible_ordinal": 0, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": round(mem.get("allocated_gb", 0), 2), + "vram_total_gb": round(mem.get("total_gb", 0), 2), + "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + "index_kind": "relative", + } return { - "available": True, + "available": False, "backend": device.value, - "gpu_utilization_pct": gpu_util, - "temperature_c": temp, - "vram_used_gb": vram_used_gb, - "vram_total_gb": vram_total_gb, - "vram_utilization_pct": vram_pct, - "power_draw_w": power_draw, - "power_limit_w": power_limit, - "power_utilization_pct": power_pct, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "relative", } @@ -417,37 +528,712 @@ _physical_gpu_count: Optional[int] = None _visible_gpu_count: Optional[int] = None +def _get_parent_visible_gpu_spec() -> Dict[str, Any]: + cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + + if cuda_visible is None: + return { + "raw": None, + "numeric_ids": list(range(get_physical_gpu_count())), + "supports_explicit_gpu_ids": True, + } + + cuda_visible = cuda_visible.strip() + if cuda_visible == "" or cuda_visible == "-1": + return { + "raw": cuda_visible, + "numeric_ids": [], + "supports_explicit_gpu_ids": True, + } + + tokens = [value.strip() for value in cuda_visible.split(",") if value.strip()] + try: + numeric_ids = [int(value) for value in tokens] + except ValueError: + return { + "raw": cuda_visible, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + return { + "raw": cuda_visible, + "numeric_ids": numeric_ids, + "supports_explicit_gpu_ids": True, + } + + +def get_parent_visible_gpu_ids() -> list[int]: + parent_visible_ids = _get_parent_visible_gpu_spec()["numeric_ids"] + return list(parent_visible_ids) if parent_visible_ids is not None else [] + + +def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]: + parent_visible_spec = _get_parent_visible_gpu_spec() + parent_visible_ids = get_parent_visible_gpu_ids() + physical_gpu_count = get_physical_gpu_count() + + if gpu_ids is None: + return parent_visible_ids + + requested_ids = list(gpu_ids) + if len(requested_ids) == 0: + return parent_visible_ids + + if not parent_visible_spec["supports_explicit_gpu_ids"]: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: explicit physical GPU IDs are " + f"unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries " + f"({parent_visible_spec['raw']!r}). Omit gpu_ids to use the " + "parent-visible devices." + ) + + if len(set(requested_ids)) != len(requested_ids): + raise ValueError( + f"Invalid gpu_ids {requested_ids}: duplicate GPU IDs are not allowed. " + f"Parent-visible GPUs: {parent_visible_ids}" + ) + + # Reject negative IDs unconditionally. + negative_ids = [gpu_id for gpu_id in requested_ids if gpu_id < 0] + if negative_ids: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: GPU IDs must be non-negative. " + f"Rejected IDs: {negative_ids}. Parent-visible GPUs: {parent_visible_ids}" + ) + + # Only enforce the physical upper bound when we have a reliable count + # from nvidia-smi. When the count comes from torch, it reflects visible + # devices (filtered by CUDA_VISIBLE_DEVICES), not the physical total, + # so high physical indices like 3 would be falsely rejected on a + # CUDA_VISIBLE_DEVICES="2,3" machine that reports device_count()=2. + # The parent-visible check below is authoritative in all cases. + if physical_gpu_count > 0 and parent_visible_ids: + max_parent_id = max(parent_visible_ids) + if physical_gpu_count > max_parent_id: + # Count is plausibly physical (not just visible), so enforce it + out_of_range = [ + gpu_id for gpu_id in requested_ids if gpu_id >= physical_gpu_count + ] + if out_of_range: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: IDs must be physical GPU IDs " + f"between 0 and {physical_gpu_count - 1}. " + f"Rejected IDs: {out_of_range}. Parent-visible GPUs: {parent_visible_ids}" + ) + + disallowed_ids = [ + gpu_id for gpu_id in requested_ids if gpu_id not in parent_visible_ids + ] + if disallowed_ids: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: requested GPUs {disallowed_ids} are " + f"outside the parent-visible set {parent_visible_ids}" + ) + + return requested_ids + + +def _resolve_model_identifier_for_gpu_estimate( + model_name: str, hf_token: Optional[str] = None +) -> str: + try: + from utils.models.model_config import ModelConfig + + config = ModelConfig.from_identifier(model_name, hf_token = hf_token) + if config and config.is_lora and config.base_model: + return config.base_model + return config.identifier if config else model_name + except Exception as e: + logger.debug( + "Could not resolve base model for GPU estimate '%s': %s", model_name, e + ) + return model_name + + +def _get_local_weight_size_bytes(model_name: str) -> Optional[int]: + model_path = Path(model_name) + if not model_path.exists(): + return None + + weight_exts = (".safetensors", ".bin", ".pt", ".pth") + total = 0 + for file in model_path.rglob("*"): + if file.is_file() and file.suffix in weight_exts: + total += file.stat().st_size + return total if total > 0 else None + + +def _get_hf_safetensors_total_params( + model_name: str, hf_token: Optional[str] = None +) -> Optional[int]: + try: + from huggingface_hub import model_info as hf_model_info + + info = hf_model_info(model_name, token = hf_token) + safetensors = getattr(info, "safetensors", None) + if isinstance(safetensors, dict): + total = safetensors.get("total") + if total: + return int(total) + except Exception as e: + logger.warning("Could not get safetensors metadata for '%s': %s", model_name, e) + return None + + +def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = None): + try: + from transformers import AutoConfig + + trust_remote_code = model_name.lower().startswith("unsloth/") + return AutoConfig.from_pretrained( + model_name, + token = hf_token, + trust_remote_code = trust_remote_code, + ) + except Exception as e: + logger.warning("Could not load config for '%s': %s", model_name, e) + return None + + +def _estimate_fp16_model_size_bytes_from_config(config) -> Optional[int]: + from .vram_estimation import extract_arch_config, compute_total_params + + arch = extract_arch_config(config) + if arch is None: + return None + return compute_total_params(arch) * 2 + + +def _estimate_fp16_model_size_bytes_from_vllm_utils(config) -> Optional[int]: + if config is None: + return None + + previous_unsloth_present = os.environ.get("UNSLOTH_IS_PRESENT") + os.environ["UNSLOTH_IS_PRESENT"] = "1" + try: + from unsloth_zoo import vllm_utils as _vllm_utils + + synthetic_total_bytes = 1024 * (1024**3) + original_get_mem_info = _vllm_utils.get_mem_info + try: + _vllm_utils.get_mem_info = lambda: ( + synthetic_total_bytes, + synthetic_total_bytes, + ) + _, _, _, memory_left_for_kv_cache_gb = ( + _vllm_utils.approximate_vllm_memory_usage( + config, + load_in_4bit = False, + load_in_8bit = False, + max_seq_length = 1, + gpu_memory_utilization = 1.0, + enable_lora = False, + account_for_gradients = False, + cuda_graph_overhead = False, + ) + ) + finally: + _vllm_utils.get_mem_info = original_get_mem_info + except Exception as e: + logger.debug("Could not estimate model size via vllm_utils: %s", e) + return None + finally: + if previous_unsloth_present is None: + os.environ.pop("UNSLOTH_IS_PRESENT", None) + else: + os.environ["UNSLOTH_IS_PRESENT"] = previous_unsloth_present + + model_size_gb = 1024.0 - memory_left_for_kv_cache_gb + if model_size_gb <= 0: + return None + return int(round(model_size_gb * (1024**3))) + + +def estimate_fp16_model_size_bytes( + model_name: str, hf_token: Optional[str] = None +) -> tuple[Optional[int], str]: + estimate_model = _resolve_model_identifier_for_gpu_estimate( + model_name, hf_token = hf_token + ) + + total_params = None + if "/" in estimate_model and not Path(estimate_model).exists(): + total_params = _get_hf_safetensors_total_params( + estimate_model, hf_token = hf_token + ) + if total_params: + return int(total_params * 2), "safetensors" + + config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token) + if config is not None: + config_bytes = _estimate_fp16_model_size_bytes_from_config(config) + if config_bytes is not None: + return config_bytes, "config" + + local_bytes = _get_local_weight_size_bytes(estimate_model) + if local_bytes is not None: + return local_bytes, "weight_bytes" + + vllm_bytes = _estimate_fp16_model_size_bytes_from_vllm_utils(config) + if vllm_bytes is not None: + return vllm_bytes, "vllm_utils" + + return None, "unavailable" + + +def estimate_required_model_memory_gb( + model_name: str, + *, + hf_token: Optional[str] = None, + training_type: Optional[str] = None, + load_in_4bit: bool = True, + batch_size: int = 4, + max_seq_length: int = 2048, + lora_rank: int = 16, + target_modules: Optional[list] = None, + gradient_checkpointing: str = "unsloth", + optimizer: str = "adamw_8bit", +) -> tuple[Optional[float], Dict[str, Any]]: + from .vram_estimation import ( + TrainingVramConfig, + extract_arch_config, + estimate_training_vram, + CUDA_OVERHEAD_BYTES, + QUANT_4BIT_FACTOR, + DEFAULT_TARGET_MODULES, + ) + + model_size_bytes, source = estimate_fp16_model_size_bytes( + model_name, hf_token = hf_token + ) + metadata: Dict[str, Any] = { + "mode": "inference" if training_type is None else "training", + "model_size_source": source, + } + if model_size_bytes is None: + metadata["required_gb"] = None + return None, metadata + + model_size_gb = model_size_bytes / (1024**3) + metadata["model_size_gb"] = round(model_size_gb, 3) + min_buffer_gb = 2.0 + + if training_type is None: + if load_in_4bit: + base_4bit_gb = model_size_gb / QUANT_4BIT_FACTOR + required_gb = base_4bit_gb + max(base_4bit_gb * 0.3, min_buffer_gb) + else: + required_gb = model_size_gb * 1.3 + metadata["required_gb"] = round(required_gb, 3) + return required_gb, metadata + + training_method = ( + "full" + if training_type == "Full Finetuning" + else ("qlora" if load_in_4bit else "lora") + ) + vram_config = TrainingVramConfig( + training_method = training_method, + batch_size = batch_size, + max_seq_length = max_seq_length, + lora_rank = lora_rank, + target_modules = target_modules or list(DEFAULT_TARGET_MODULES), + gradient_checkpointing = gradient_checkpointing, + optimizer = optimizer, + load_in_4bit = load_in_4bit, + ) + + estimate_model = _resolve_model_identifier_for_gpu_estimate( + model_name, hf_token = hf_token + ) + config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token) + arch = extract_arch_config(config) if config is not None else None + + if arch is not None: + breakdown = estimate_training_vram(arch, vram_config) + required_gb = breakdown.total / (1024**3) + metadata["required_gb"] = round(required_gb, 3) + metadata["estimation_mode"] = "detailed" + metadata["vram_breakdown"] = breakdown.to_gb_dict() + max_gpus = max(1, get_visible_gpu_count()) + for n_gpus in range(1, max_gpus + 1): + metadata["vram_breakdown"][f"min_per_gpu_{n_gpus}"] = round( + breakdown.min_gpu_vram(n_gpus) / (1024**3), 3 + ) + return required_gb, metadata + + # Fallback when model config is unavailable + overhead_gb = CUDA_OVERHEAD_BYTES / (1024**3) + if training_method == "full": + required_gb = model_size_gb * 3.5 + overhead_gb + elif training_method == "qlora": + base_4bit_gb = model_size_gb / QUANT_4BIT_FACTOR + lora_overhead_gb = model_size_gb * 0.04 + act_gb = model_size_gb * 0.15 * (batch_size / 4) * (max_seq_length / 2048) + required_gb = base_4bit_gb + lora_overhead_gb + act_gb + overhead_gb + else: + lora_overhead_gb = model_size_gb * 0.04 + act_gb = model_size_gb * 0.15 * (batch_size / 4) * (max_seq_length / 2048) + required_gb = model_size_gb + lora_overhead_gb + act_gb + overhead_gb + + metadata["required_gb"] = round(required_gb, 3) + metadata["estimation_mode"] = "fallback" + return required_gb, metadata + + +def auto_select_gpu_ids( + model_name: str, + *, + hf_token: Optional[str] = None, + training_type: Optional[str] = None, + load_in_4bit: bool = True, + batch_size: int = 4, + max_seq_length: int = 2048, + lora_rank: int = 16, + target_modules: Optional[list] = None, + gradient_checkpointing: str = "unsloth", + optimizer: str = "adamw_8bit", +) -> tuple[Optional[list[int]], Dict[str, Any]]: + metadata: Dict[str, Any] = {"selection_mode": "auto"} + + if get_device() != DeviceType.CUDA: + metadata["selection_mode"] = "non_cuda" + return None, metadata + + required_gb, estimate_metadata = estimate_required_model_memory_gb( + model_name, + hf_token = hf_token, + training_type = training_type, + load_in_4bit = load_in_4bit, + batch_size = batch_size, + max_seq_length = max_seq_length, + lora_rank = lora_rank, + target_modules = target_modules, + gradient_checkpointing = gradient_checkpointing, + optimizer = optimizer, + ) + metadata.update(estimate_metadata) + parent_visible_spec = _get_parent_visible_gpu_spec() + metadata["parent_cuda_visible_devices"] = parent_visible_spec["raw"] + + if not parent_visible_spec["supports_explicit_gpu_ids"]: + metadata["selection_mode"] = "inherit_parent_visible" + metadata["selected_gpu_ids"] = None + return None, metadata + + if required_gb is None: + # Cannot estimate model size -- fall back to all visible GPUs + # rather than risk loading on a single GPU that may not have + # enough memory. + parent_ids = get_parent_visible_gpu_ids() + metadata["selection_mode"] = "fallback_all" + metadata["selected_gpu_ids"] = parent_ids + return parent_ids, metadata + + utilization = get_visible_gpu_utilization() + devices = utilization.get("devices", []) + parent_ids = get_parent_visible_gpu_ids() + + if not devices: + metadata["selection_mode"] = "fallback_all" + metadata["selected_gpu_ids"] = parent_ids + return parent_ids, metadata + + gpu_candidates = [] + for device in devices: + total_gb = device.get("vram_total_gb") + used_gb = device.get("vram_used_gb") + if total_gb is None or used_gb is None: + continue + free_gb = max(total_gb - used_gb, 0.0) + gpu_candidates.append( + { + "index": device["index"], + "free_gb": free_gb, + } + ) + + if not gpu_candidates: + metadata["selection_mode"] = "fallback_all" + metadata["selected_gpu_ids"] = parent_ids + return parent_ids, metadata + + ranked = sorted(gpu_candidates, key = lambda item: (-item["free_gb"], item["index"])) + free_by_index = {item["index"]: item["free_gb"] for item in ranked} + selected: list[int] = [] + usable_gb = 0.0 + # Multi-GPU sharding has overhead from inter-GPU communication (NCCL + # all-reduce, PCIe/NVLink transfers, synchronization barriers), so each + # additional GPU contributes less than its raw free memory. The first GPU + # keeps its full capacity (no cross-device overhead). 0.85 was calibrated + # empirically on 2-8 GPU setups with NVLink and PCIe topologies -- the + # 15% discount accounts for NCCL buffers (~2-5% of VRAM), pipeline bubble + # overhead, and memory fragmentation from non-uniform shard sizes. + multi_gpu_overhead = 0.85 + + # Per-GPU check: activations don't shard, so each GPU needs its weight + # shard + full activation cost. Use precomputed min_per_gpu_N values. + vram_breakdown = estimate_metadata.get("vram_breakdown", {}) + + for candidate in ranked: + selected.append(candidate["index"]) + if len(selected) == 1: + usable_gb = candidate["free_gb"] + else: + first_gpu_id = selected[0] + usable_gb = free_by_index[first_gpu_id] + sum( + free_by_index[gpu_id] * multi_gpu_overhead for gpu_id in selected[1:] + ) + + total_fits = usable_gb >= required_gb + + per_gpu_fits = True + if total_fits and len(selected) > 1: + min_key = f"min_per_gpu_{len(selected)}" + min_per_gpu_gb = vram_breakdown.get(min_key) + if min_per_gpu_gb is not None: + smallest_free = min(free_by_index[gpu_id] for gpu_id in selected) + per_gpu_fits = smallest_free >= min_per_gpu_gb + + if total_fits and per_gpu_fits: + metadata["usable_gb"] = round(usable_gb, 3) + metadata["selection_mode"] = "auto" + metadata["selected_gpu_ids"] = selected + logger.debug( + "Selected GPUs automatically", + model_name = model_name, + selected_gpu_ids = selected, + usable_gb = metadata["usable_gb"], + required_gb = metadata.get("required_gb"), + multi_gpu_overhead = multi_gpu_overhead, + ) + return selected, metadata + + # Use only GPUs with verified VRAM data (from gpu_candidates, not raw devices) + fallback_all = ( + [c["index"] for c in gpu_candidates] if gpu_candidates else parent_ids + ) + metadata["selection_mode"] = "fallback_all" + if ranked: + fallback_usable = ranked[0]["free_gb"] + sum( + c["free_gb"] * multi_gpu_overhead for c in ranked[1:] + ) + else: + fallback_usable = 0.0 + metadata["usable_gb"] = round(fallback_usable, 3) + metadata["selected_gpu_ids"] = fallback_all + logger.warning( + "Falling back to all visible GPUs -- model may not fit", + model_name = model_name, + selected_gpu_ids = fallback_all, + usable_gb = metadata["usable_gb"], + required_gb = metadata.get("required_gb"), + multi_gpu_overhead = multi_gpu_overhead, + ) + return fallback_all, metadata + + +def prepare_gpu_selection( + gpu_ids: Optional[list[int]], + *, + model_name: str, + hf_token: Optional[str] = None, + training_type: Optional[str] = None, + load_in_4bit: bool = True, + batch_size: int = 4, + max_seq_length: int = 2048, + lora_rank: int = 16, + target_modules: Optional[list] = None, + gradient_checkpointing: str = "unsloth", + optimizer: str = "adamw_8bit", +) -> tuple[Optional[list[int]], Dict[str, Any]]: + """Resolve which physical GPUs to use for a model load. + + GPU selection modes: + - **Explicit** (``gpu_ids=[5, 6, 7]``): the caller chooses exact GPUs. + All listed GPUs are used and the model is sharded across them via + ``device_map="balanced"``, regardless of whether the model would fit + on fewer GPUs. IDs are validated against the parent-visible set. + - **Auto** (``gpu_ids=None`` or ``[]``): ``auto_select_gpu_ids`` estimates + VRAM requirements and picks the *minimum* number of GPUs needed, + preferring GPUs with the most free memory. + + The returned ``gpu_ids`` list is later passed to ``get_device_map()`` which + maps it to a Hugging Face ``device_map`` string, and to ``apply_gpu_ids()`` + in the worker subprocess which narrows ``CUDA_VISIBLE_DEVICES`` before any + torch/CUDA initialisation. + """ + if gpu_ids and get_device() != DeviceType.CUDA: + raise ValueError( + f"gpu_ids {list(gpu_ids)} is only supported on CUDA devices, " + f"but the current backend is '{get_device().value}'." + ) + + if gpu_ids: + resolved = resolve_requested_gpu_ids(gpu_ids) + metadata = { + "selection_mode": "explicit", + "selected_gpu_ids": resolved, + } + return resolved, metadata + + selected_gpu_ids, metadata = auto_select_gpu_ids( + model_name, + hf_token = hf_token, + training_type = training_type, + load_in_4bit = load_in_4bit, + batch_size = batch_size, + max_seq_length = max_seq_length, + lora_rank = lora_rank, + target_modules = target_modules, + gradient_checkpointing = gradient_checkpointing, + optimizer = optimizer, + ) + return selected_gpu_ids, metadata + + def get_physical_gpu_count() -> int: """ - Return the number of physical NVIDIA GPUs on the machine. + Return the number of physical GPUs on the machine. - Uses ``nvidia-smi -L`` which is NOT affected by CUDA_VISIBLE_DEVICES, - so it always reflects the true hardware count. + Uses ``nvidia-smi -L`` on NVIDIA (unaffected by CUDA_VISIBLE_DEVICES), + with a torch-based fallback for AMD ROCm and Intel XPU. Result is cached after the first call. """ global _physical_gpu_count if _physical_gpu_count is not None: return _physical_gpu_count - try: - import subprocess + device = get_device() - result = subprocess.run( - ["nvidia-smi", "-L"], - capture_output = True, - text = True, - timeout = 5, - ) - if result.returncode == 0 and result.stdout.strip(): - _physical_gpu_count = len(result.stdout.strip().splitlines()) - else: - _physical_gpu_count = 1 - except Exception: + if device == DeviceType.CUDA: + try: + from . import nvidia + + count = nvidia.get_physical_gpu_count() + if count is not None: + _physical_gpu_count = count + return _physical_gpu_count + except Exception: + pass + # nvidia-smi unavailable or failed — fall back to torch + count = _torch_get_physical_gpu_count() + _physical_gpu_count = count if count is not None else 1 + return _physical_gpu_count + + if device == DeviceType.XPU: + count = _torch_get_physical_gpu_count() + _physical_gpu_count = count if count is not None else 1 + return _physical_gpu_count + + if device == DeviceType.MLX: _physical_gpu_count = 1 + return _physical_gpu_count + + _physical_gpu_count = 0 return _physical_gpu_count +def get_backend_visible_gpu_info() -> Dict[str, Any]: + device = get_device() + if device in (DeviceType.CUDA, DeviceType.XPU): + parent_visible_ids = get_parent_visible_gpu_ids() + # Try nvidia-smi first (NVIDIA only) + if device == DeviceType.CUDA: + try: + from . import nvidia + + parent_visible_spec = _get_parent_visible_gpu_spec() + result = nvidia.get_backend_visible_gpu_info( + parent_visible_spec["numeric_ids"], + parent_visible_spec["raw"], + ) + if result.get("available"): + result["backend"] = device.value + return result + except Exception as e: + logger.warning("Backend GPU visibility query failed: %s", e) + + # Torch fallback (AMD ROCm, Intel XPU, nvidia-smi missing/failed) + # When parent_visible_ids is empty (UUID/MIG mask), enumerate by + # torch ordinal so the UI still shows devices. + if parent_visible_ids: + torch_indices = parent_visible_ids + index_kind = "physical" + else: + visible_count = _torch_get_physical_gpu_count() or 0 + torch_indices = list(range(visible_count)) + index_kind = "relative" + torch_devices = _torch_get_per_device_info(torch_indices) + if torch_devices: + devices = [ + { + "index": td["index"], + "index_kind": index_kind, + "visible_ordinal": td["visible_ordinal"], + "name": td["name"], + "memory_total_gb": td["total_gb"], + } + for td in torch_devices + ] + return { + "available": True, + "backend": device.value, + "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "parent_visible_gpu_ids": parent_visible_ids, + "devices": devices, + "index_kind": index_kind, + } + + return { + "available": False, + "backend": device.value, + "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "parent_visible_gpu_ids": parent_visible_ids, + "devices": [], + "index_kind": "physical", + } + + if device == DeviceType.MLX: + mem = get_gpu_memory_info() + if not mem.get("available"): + return { + "available": False, + "backend": device.value, + "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "relative", + } + return { + "available": True, + "backend": device.value, + "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "parent_visible_gpu_ids": [0], + "devices": [ + { + "index": 0, + "index_kind": "relative", + "visible_ordinal": 0, + "name": mem.get("device_name", "MLX"), + "memory_total_gb": round(mem.get("total_gb", 0), 2), + } + ], + "index_kind": "relative", + } + + return { + "available": False, + "backend": device.value, + "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "relative", + } + + def get_visible_gpu_count() -> int: """ Return the number of GPUs visible to this process. @@ -460,8 +1246,6 @@ def get_visible_gpu_count() -> int: if _visible_gpu_count is not None: return _visible_gpu_count - import os - cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") if cuda_visible is not None: # "" means zero GPUs, "0" means 1, "0,1,2" means 3 @@ -476,13 +1260,103 @@ def get_visible_gpu_count() -> int: try: import torch - _visible_gpu_count = torch.cuda.device_count() + if get_device() == DeviceType.XPU and hasattr(torch, "xpu"): + _visible_gpu_count = torch.xpu.device_count() + else: + _visible_gpu_count = torch.cuda.device_count() except Exception: _visible_gpu_count = get_physical_gpu_count() return _visible_gpu_count +def apply_gpu_ids(gpu_ids) -> None: + if gpu_ids is None: + return + + # Empty list means "no GPUs visible" -- treat the same as None + # (inherit parent) to avoid setting CUDA_VISIBLE_DEVICES="" which + # disables CUDA entirely and crashes downstream torch calls. + if isinstance(gpu_ids, (list, tuple)) and len(gpu_ids) == 0: + return + + global _visible_gpu_count + + if isinstance(gpu_ids, (list, tuple)): + value = ",".join(str(g) for g in gpu_ids) + else: + value = str(gpu_ids) + + os.environ["CUDA_VISIBLE_DEVICES"] = value + _visible_gpu_count = None + logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value) + + +def get_device_map( + gpu_ids: Optional[list[int]] = None, + *, + for_inference: bool = False, +) -> str: + """Return the Hugging Face ``device_map`` string for model loading. + + Returns ``"balanced"`` (shard evenly across GPUs) when: + - ``gpu_ids`` explicitly lists >1 GPU, **or** + - ``CUDA_VISIBLE_DEVICES`` uses UUID/MIG identifiers (non-numeric) and + more than one GPU is visible (fallback: we cannot resolve numeric IDs, + so we assume the caller intends multi-GPU). + + Returns ``"sequential"`` (single device) in all other cases, including + non-CUDA backends (CPU, MLX). + + Callers should use ``prepare_gpu_selection()`` upstream to determine the + ``gpu_ids`` list -- that function handles the smart auto-selection of the + minimum number of GPUs needed for a given model. + """ + device = get_device() + if device == DeviceType.CUDA: + multi_gpu = gpu_ids is not None and len(gpu_ids) > 1 + + if not multi_gpu: + # UUID/MIG masks cannot be split into numeric IDs, so if multiple + # GPUs are visible we assume multi-GPU sharding is intended. + parent_visible_spec = _get_parent_visible_gpu_spec() + if ( + parent_visible_spec["numeric_ids"] is None + and get_visible_gpu_count() > 1 + ): + multi_gpu = True + + if multi_gpu: + return "balanced_low_0" if for_inference else "balanced" + + return "sequential" + + +def get_offloaded_device_map_entries(model) -> dict[str, str]: + hf_device_map = getattr(model, "hf_device_map", None) + if not isinstance(hf_device_map, dict): + return {} + return { + module_name: placement + for module_name, placement in hf_device_map.items() + if placement in ("cpu", "disk") + } + + +def raise_if_offloaded(model, device_map: str, context: str = "Loading") -> None: + """Raise ``ValueError`` if *model* has modules offloaded to CPU or disk.""" + offloaded = get_offloaded_device_map_entries(model) + if not offloaded: + return + example = ", ".join( + f"{name}={placement}" for name, placement in list(offloaded.items())[:5] + ) + raise ValueError( + f"{context} does not support models loaded with CPU or disk offload. " + f"device_map='{device_map}' produced offloaded modules: {example}" + ) + + def safe_num_proc(desired: Optional[int] = None) -> int: """ Return a safe ``num_proc`` for ``dataset.map()`` calls. @@ -507,7 +1381,6 @@ def safe_num_proc(desired: Optional[int] = None) -> int: Returns: A safe integer ≥ 1. """ - import os import sys # Windows and macOS use 'spawn' for multiprocessing -- the overhead of @@ -546,8 +1419,6 @@ def safe_thread_num_proc(desired: Optional[int] = None) -> int: Returns: A safe integer >= 1. """ - import os - if desired is None or not isinstance(desired, int): desired = max(1, (os.cpu_count() or 1) // 3) diff --git a/studio/backend/utils/hardware/nvidia.py b/studio/backend/utils/hardware/nvidia.py new file mode 100644 index 0000000000..dc5295c302 --- /dev/null +++ b/studio/backend/utils/hardware/nvidia.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import subprocess +from typing import Any, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + + +def _parse_smi_value(raw: str): + raw = raw.strip() + if not raw or raw == "[N/A]": + return None + try: + return float(raw) + except (ValueError, TypeError): + return None + + +def _build_gpu_metrics( + vram_used_mb, + vram_total_mb, + power_draw, + power_limit, + **extra, +) -> dict[str, Any]: + return { + **extra, + "vram_used_gb": round(vram_used_mb / 1024, 2) + if vram_used_mb is not None + else None, + "vram_total_gb": round(vram_total_mb / 1024, 2) + if vram_total_mb is not None + else None, + "vram_utilization_pct": round((vram_used_mb / vram_total_mb) * 100, 1) + if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0 + else None, + "power_draw_w": power_draw, + "power_limit_w": power_limit, + "power_utilization_pct": round((power_draw / power_limit) * 100, 1) + if power_draw is not None and power_limit and power_limit > 0 + else None, + } + + +def _visible_ordinal_map( + parent_visible_ids: Optional[list[int]], +) -> Optional[dict[int, int]]: + if parent_visible_ids is None: + return None + return {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)} + + +def get_physical_gpu_count() -> Optional[int]: + """Return physical GPU count via nvidia-smi, or None on failure.""" + try: + result = subprocess.run( + ["nvidia-smi", "-L"], + capture_output = True, + text = True, + timeout = 5, + ) + if result.returncode == 0 and result.stdout.strip(): + return len(result.stdout.strip().splitlines()) + logger.warning( + "nvidia-smi -L returned code %d; caller should fall back to torch", + result.returncode, + ) + except Exception as e: + logger.warning("nvidia-smi -L failed: %s; caller should fall back to torch", e) + return None + + +def get_primary_gpu_utilization() -> dict[str, Any]: + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=utilization.gpu,temperature.gpu," + "memory.used,memory.total,power.draw,power.limit", + "--format=csv,noheader,nounits", + ], + capture_output = True, + text = True, + timeout = 5, + ) + except (OSError, subprocess.TimeoutExpired) as e: + logger.warning("nvidia-smi query failed in get_primary_gpu_utilization: %s", e) + return {"available": False} + if result.returncode != 0 or not result.stdout.strip(): + return {"available": False} + + first_line = result.stdout.strip().splitlines()[0] + parts = [p.strip() for p in first_line.split(",")] + if len(parts) < 6: + return {"available": False} + + return _build_gpu_metrics( + vram_used_mb = _parse_smi_value(parts[2]), + vram_total_mb = _parse_smi_value(parts[3]), + power_draw = _parse_smi_value(parts[4]), + power_limit = _parse_smi_value(parts[5]), + available = True, + gpu_utilization_pct = _parse_smi_value(parts[0]), + temperature_c = _parse_smi_value(parts[1]), + ) + + +def get_visible_gpu_utilization( + parent_visible_ids: Optional[list[int]], + parent_cuda_visible_devices: Optional[str] = None, +) -> dict[str, Any]: + # When parent_visible_ids is None (UUID/MIG mask), we cannot safely + # map nvidia-smi rows to the process's visible devices. Return empty + # instead of exposing all physical GPUs. + if parent_visible_ids is None: + return { + "available": False, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "unresolved", + } + visible_ordinals = _visible_ordinal_map(parent_visible_ids) + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,utilization.gpu,temperature.gpu," + "memory.used,memory.total,power.draw,power.limit", + "--format=csv,noheader,nounits", + ], + capture_output = True, + text = True, + timeout = 5, + ) + except (OSError, subprocess.TimeoutExpired) as e: + logger.warning("nvidia-smi query failed in get_visible_gpu_utilization: %s", e) + return { + "available": False, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + if result.returncode != 0 or not result.stdout.strip(): + return { + "available": False, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + + devices = [] + for line in result.stdout.strip().splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) < 7: + continue + + try: + idx = int(parts[0]) + except (ValueError, TypeError): + continue + + if visible_ordinals is not None and idx not in visible_ordinals: + continue + + devices.append( + _build_gpu_metrics( + vram_used_mb = _parse_smi_value(parts[3]), + vram_total_mb = _parse_smi_value(parts[4]), + power_draw = _parse_smi_value(parts[5]), + power_limit = _parse_smi_value(parts[6]), + index = idx, + index_kind = "physical", + visible_ordinal = ( + visible_ordinals[idx] + if visible_ordinals is not None + else len(devices) + ), + gpu_utilization_pct = _parse_smi_value(parts[1]), + temperature_c = _parse_smi_value(parts[2]), + ) + ) + + return { + "available": len(devices) > 0, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": devices, + "index_kind": "physical", + } + + +def get_backend_visible_gpu_info( + parent_visible_ids: Optional[list[int]], + backend_cuda_visible_devices: Optional[str], +) -> dict[str, Any]: + # When parent_visible_ids is None (UUID/MIG mask), we cannot safely + # map nvidia-smi rows to the process's visible devices. + if parent_visible_ids is None: + return { + "available": False, + "backend_cuda_visible_devices": backend_cuda_visible_devices, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "unresolved", + } + visible_ordinals = _visible_ordinal_map(parent_visible_ids) + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,name,memory.total", + "--format=csv,noheader,nounits", + ], + capture_output = True, + text = True, + timeout = 10, + ) + except (OSError, subprocess.TimeoutExpired) as e: + logger.warning("nvidia-smi query failed in get_backend_visible_gpu_info: %s", e) + return { + "available": False, + "backend_cuda_visible_devices": backend_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + if result.returncode != 0: + return { + "available": False, + "backend_cuda_visible_devices": backend_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + + devices = [] + for line in result.stdout.strip().splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) < 3: + continue + try: + idx = int(parts[0]) + except (ValueError, TypeError): + continue + if visible_ordinals is not None and idx not in visible_ordinals: + continue + # Use split with limit to handle GPU names containing commas + name = parts[1] if len(parts) == 3 else ", ".join(parts[1:-1]) + try: + mem_total_mb = int(parts[-1]) + except (ValueError, TypeError): + continue + devices.append( + { + "index": idx, + "index_kind": "physical", + "visible_ordinal": ( + visible_ordinals[idx] + if visible_ordinals is not None + else len(devices) + ), + "name": name, + "memory_total_gb": round(mem_total_mb / 1024, 2), + } + ) + + return { + "available": len(devices) > 0, + "backend_cuda_visible_devices": backend_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": devices, + "index_kind": "physical", + } diff --git a/studio/backend/utils/hardware/vram_estimation.py b/studio/backend/utils/hardware/vram_estimation.py new file mode 100644 index 0000000000..e03665374d --- /dev/null +++ b/studio/backend/utils/hardware/vram_estimation.py @@ -0,0 +1,501 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +""" +Training VRAM estimation. + +Total VRAM = weights + LoRA adapters + optimizer states + gradients + + activations + CUDA overhead. +Activation formula from unsloth_zoo/vllm_utils.py. +All constants empirically calibrated against Llama-3.2-1B on B200. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Dict, Optional + +QUANT_4BIT_FACTOR = 16 / 5 +CUDA_OVERHEAD_BYTES = int(1.4 * 1024**3) # calibrated on RTX 5070 Ti + +DEFAULT_TARGET_MODULES = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", +] + +# Empirically calibrated bytes/param — see VRAM_ESTIMATION.md for rationale. +OPTIMIZER_BYTES_PER_PARAM: Dict[str, int] = { + "adamw_8bit": 4, # BNB upcasts to fp32 during step + "paged_adamw_8bit": 4, + "adamw_bnb_8bit": 4, + "paged_adamw_32bit": 8, + "adamw_torch": 6, # fused, no master copy + "adamw_torch_fused": 6, + "sgd": 4, +} + +# (full_ft_multiplier, lora_multiplier) — fraction of num_layers. +# LoRA: frozen base layers skip activation storage, but you always need +# at least ~1 layer in flight during backprop recomputation. +GC_LAYER_MULTIPLIERS = { + "none": (None, None), + "true": (2.0, 1.0), + "unsloth": (1.5, 1.0), +} + + +@dataclass +class ModelArchConfig: + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + intermediate_size: int + vocab_size: int + tie_word_embeddings: bool = True + num_experts: Optional[int] = None + moe_intermediate_size: Optional[int] = None + n_shared_experts: int = 0 + num_dense_layers: int = 0 + q_lora_rank: Optional[int] = None + kv_lora_rank: Optional[int] = None + qk_nope_head_dim: Optional[int] = None + qk_rope_head_dim: Optional[int] = None + v_head_dim: Optional[int] = None + + +@dataclass +class TrainingVramConfig: + training_method: str = "qlora" + batch_size: int = 4 + max_seq_length: int = 2048 + lora_rank: int = 16 + target_modules: list = field(default_factory = lambda: list(DEFAULT_TARGET_MODULES)) + gradient_checkpointing: str = "unsloth" + optimizer: str = "adamw_8bit" + load_in_4bit: bool = True + + +@dataclass +class VramBreakdown: + model_weights: int + lora_adapters: int + optimizer_states: int + gradients: int + activations: int + cuda_overhead: int + # The computed (formula-based) activation cost before floors. + # This is the true per-layer cost that doesn't shard across GPUs. + activations_computed: int = 0 + + @property + def total(self) -> int: + return ( + self.model_weights + + self.lora_adapters + + self.optimizer_states + + self.gradients + + self.activations + + self.cuda_overhead + ) + + def min_gpu_vram(self, n_gpus: int) -> int: + """Minimum VRAM a single GPU needs: its shard + non-shardable costs. + + Weights/LoRA/optimizer/gradients shard across GPUs. + The computed activation cost does NOT shard (one GPU runs the layer). + The floor portion (activations - computed) is overhead that shards. + """ + shardable = ( + self.model_weights + + self.lora_adapters + + self.optimizer_states + + self.gradients + + (self.activations - self.activations_computed) # floor overhead shards + ) + per_gpu_fixed = self.activations_computed + self.cuda_overhead + return shardable // max(n_gpus, 1) + per_gpu_fixed + + def to_gb_dict(self) -> Dict[str, float]: + return { + "model_weights_gb": round(self.model_weights / (1024**3), 3), + "lora_adapters_gb": round(self.lora_adapters / (1024**3), 3), + "optimizer_states_gb": round(self.optimizer_states / (1024**3), 3), + "gradients_gb": round(self.gradients / (1024**3), 3), + "activations_gb": round(self.activations / (1024**3), 3), + "cuda_overhead_gb": round(self.cuda_overhead / (1024**3), 3), + "total_gb": round(self.total / (1024**3), 3), + } + + +def _compute_num_dense_layers(text_config, total_layers: int) -> int: + """Count how many layers use dense MLP instead of MoE.""" + first_k = getattr(text_config, "first_k_dense_replace", None) + if first_k is not None: + return min(int(first_k), total_layers) + + sparse_step = getattr(text_config, "decoder_sparse_step", None) + mlp_only = getattr(text_config, "mlp_only_layers", None) or [] + if sparse_step is not None and sparse_step > 0: + mlp_only_set = set(mlp_only) + moe_count = sum( + 1 + for i in range(total_layers) + if i not in mlp_only_set and (i + 1) % sparse_step == 0 + ) + return total_layers - moe_count + + return 0 + + +def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: + text_config = getattr(hf_config, "text_config", None) or hf_config + + hidden_size = getattr(text_config, "hidden_size", None) + num_layers = getattr(text_config, "num_hidden_layers", None) + num_heads = getattr(text_config, "num_attention_heads", None) + intermediate_size = getattr(text_config, "intermediate_size", None) + vocab_size = getattr(text_config, "vocab_size", None) + + if isinstance(intermediate_size, (list, tuple)): + intermediate_size = intermediate_size[0] if intermediate_size else None + if intermediate_size is None and hidden_size is not None: + intermediate_size = hidden_size * 4 + + if not all( + v is not None + for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size) + ): + return None + if num_heads <= 0: + return None + + num_kv_heads = getattr(text_config, "num_key_value_heads", num_heads) + + num_experts = None + for attr in ("num_local_experts", "num_experts", "n_routed_experts"): + num_experts = getattr(text_config, attr, None) + if num_experts is not None: + break + + moe_intermediate = getattr(text_config, "moe_intermediate_size", None) + n_shared_experts = getattr(text_config, "n_shared_experts", None) or 0 + + num_dense_layers = 0 + if num_experts is not None and num_experts > 1: + num_dense_layers = _compute_num_dense_layers(text_config, num_layers) + + q_lora_rank = getattr(text_config, "q_lora_rank", None) + kv_lora_rank = getattr(text_config, "kv_lora_rank", None) + qk_nope_head_dim = getattr(text_config, "qk_nope_head_dim", None) + qk_rope_head_dim = getattr(text_config, "qk_rope_head_dim", None) + v_head_dim = getattr(text_config, "v_head_dim", None) + + return ModelArchConfig( + hidden_size = hidden_size, + num_hidden_layers = num_layers, + num_attention_heads = num_heads, + num_key_value_heads = num_kv_heads, + intermediate_size = intermediate_size, + vocab_size = vocab_size, + tie_word_embeddings = getattr(text_config, "tie_word_embeddings", True), + num_experts = num_experts, + moe_intermediate_size = moe_intermediate, + n_shared_experts = n_shared_experts, + num_dense_layers = num_dense_layers, + q_lora_rank = q_lora_rank, + kv_lora_rank = kv_lora_rank, + qk_nope_head_dim = qk_nope_head_dim, + qk_rope_head_dim = qk_rope_head_dim, + v_head_dim = v_head_dim, + ) + + +def _get_kv_size(arch: ModelArchConfig) -> int: + return (arch.hidden_size // arch.num_attention_heads) * arch.num_key_value_heads + + +def _get_mlp_size(arch: ModelArchConfig) -> int: + if arch.moe_intermediate_size is not None: + return arch.moe_intermediate_size + return arch.intermediate_size + + +def _get_num_experts(arch: ModelArchConfig) -> int: + return arch.num_experts if arch.num_experts and arch.num_experts > 1 else 1 + + +def _compute_attn_elements(arch: ModelArchConfig) -> int: + """Attention weight elements per layer.""" + hd = arch.hidden_size + if arch.q_lora_rank is not None: + nh = arch.num_attention_heads + qk_head = arch.qk_nope_head_dim + arch.qk_rope_head_dim + q_a = hd * arch.q_lora_rank + q_b = arch.q_lora_rank * (nh * qk_head) + kv_a = hd * (arch.kv_lora_rank + arch.qk_rope_head_dim) + kv_b = arch.kv_lora_rank * (nh * (arch.qk_nope_head_dim + arch.v_head_dim)) + o = (nh * arch.v_head_dim) * hd + norms = arch.q_lora_rank + arch.kv_lora_rank + return q_a + q_b + kv_a + kv_b + o + norms + kv_size = _get_kv_size(arch) + return (hd + kv_size + kv_size + hd) * hd + + +def _compute_dense_mlp_elements(arch: ModelArchConfig) -> int: + return arch.hidden_size * arch.intermediate_size * 3 + + +def _compute_moe_mlp_elements(arch: ModelArchConfig) -> int: + hd = arch.hidden_size + mlp_size = _get_mlp_size(arch) + n_experts = _get_num_experts(arch) + return hd * mlp_size * 3 * (n_experts + arch.n_shared_experts) + n_experts * hd + + +def _compute_layer_elements(arch: ModelArchConfig): + """Return (total_quantizable, layernorms_per_layer, embed, lm_head) element counts. + + total_quantizable is summed across ALL layers (not per-layer). + """ + hd = arch.hidden_size + n_layers = arch.num_hidden_layers + n_experts = _get_num_experts(arch) + + attn_total = _compute_attn_elements(arch) * n_layers + + if n_experts > 1: + n_dense = arch.num_dense_layers + n_moe = n_layers - n_dense + mlp_total = ( + _compute_moe_mlp_elements(arch) * n_moe + + _compute_dense_mlp_elements(arch) * n_dense + ) + else: + mlp_total = _compute_dense_mlp_elements(arch) * n_layers + + layernorms = 2 * hd + embed_tokens = arch.vocab_size * hd + lm_head = 0 if arch.tie_word_embeddings else arch.vocab_size * hd + return attn_total + mlp_total, layernorms, embed_tokens, lm_head + + +def compute_model_weights_bytes( + arch: ModelArchConfig, + training_method: str, + load_in_4bit: bool, +) -> int: + total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch) + n_layers = arch.num_hidden_layers + non_quantizable = layernorms * n_layers + embed_tokens + lm_head + + if training_method == "qlora" and load_in_4bit: + return int(total_quantizable * 2 / QUANT_4BIT_FACTOR + non_quantizable * 2) + + return int((total_quantizable + non_quantizable) * 2) + + +def compute_total_params(arch: ModelArchConfig) -> int: + total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch) + n_layers = arch.num_hidden_layers + return total_quantizable + layernorms * n_layers + embed_tokens + lm_head + + +def _lora_attn_elements( + arch: ModelArchConfig, + r: int, + target_modules: list, +) -> int: + hd = arch.hidden_size + if arch.q_lora_rank is not None: + # MLA: q_proj->q_b, k_proj->kv_a, v_proj->kv_b, o_proj->o + nh = arch.num_attention_heads + qk_head = arch.qk_nope_head_dim + arch.qk_rope_head_dim + kv_out = nh * (arch.qk_nope_head_dim + arch.v_head_dim) + o_in = nh * arch.v_head_dim + dims = { + "q_proj": (arch.q_lora_rank, nh * qk_head), + "k_proj": (hd, arch.kv_lora_rank + arch.qk_rope_head_dim), + "v_proj": (arch.kv_lora_rank, kv_out), + "o_proj": (o_in, hd), + } + else: + kv_size = _get_kv_size(arch) + dims = { + "q_proj": (hd, hd), + "k_proj": (hd, kv_size), + "v_proj": (hd, kv_size), + "o_proj": (hd, hd), + } + total = 0 + for name, (in_dim, out_dim) in dims.items(): + if name in target_modules: + total += in_dim * r + r * out_dim + return total + + +def _lora_mlp_elements( + hd: int, + mlp_size: int, + r: int, + target_modules: list, + expert_mult: int, +) -> int: + module_ab = { + "gate_proj": (hd * r, r * mlp_size), + "up_proj": (hd * r, r * mlp_size), + "down_proj": (mlp_size * r, r * hd), + } + total = 0 + for name, (a, b) in module_ab.items(): + if name in target_modules: + total += (a + b) * expert_mult + return total + + +def compute_lora_params( + arch: ModelArchConfig, + lora_rank: int, + target_modules: list, +) -> int: + hd = arch.hidden_size + r = lora_rank + n_layers = arch.num_hidden_layers + n_experts = _get_num_experts(arch) + + attn_total = _lora_attn_elements(arch, r, target_modules) * n_layers + + if n_experts > 1: + n_dense = arch.num_dense_layers + n_moe = n_layers - n_dense + # Include shared experts alongside routed experts + moe_expert_mult = n_experts + arch.n_shared_experts + moe_mlp = _lora_mlp_elements( + hd, + _get_mlp_size(arch), + r, + target_modules, + moe_expert_mult, + ) + dense_mlp = _lora_mlp_elements( + hd, + arch.intermediate_size, + r, + target_modules, + 1, + ) + mlp_total = moe_mlp * n_moe + dense_mlp * n_dense + else: + mlp_total = ( + _lora_mlp_elements( + hd, + arch.intermediate_size, + r, + target_modules, + 1, + ) + * n_layers + ) + + return attn_total + mlp_total + + +def compute_lora_adapter_bytes(lora_params: int) -> int: + return lora_params * 2 + + +def compute_optimizer_bytes(trainable_params: int, optimizer: str) -> int: + optimizer_key = optimizer.lower().replace("-", "_") + bytes_per_param = OPTIMIZER_BYTES_PER_PARAM.get(optimizer_key, 4) + return trainable_params * bytes_per_param + + +def compute_gradient_bytes(trainable_params: int) -> int: + return trainable_params * 2 + + +def compute_activation_bytes( + arch: ModelArchConfig, + batch_size: int, + seq_len: int, + gradient_checkpointing: str, + is_lora: bool = False, +) -> int: + hd = arch.hidden_size + kv_size = _get_kv_size(arch) + mlp_size = _get_mlp_size(arch) + bsz = batch_size + n_layers = arch.num_hidden_layers + + activation_qkv = seq_len * bsz * (hd + kv_size + kv_size) + residual_memory = (seq_len * bsz) * 2 + activation_mlp = seq_len * bsz * (mlp_size + mlp_size) + + per_layer_bytes = (activation_qkv + residual_memory + activation_mlp) * 2 + per_layer_bytes = int(per_layer_bytes * 1.25) + + gc_key = gradient_checkpointing.lower() + gc_entry = GC_LAYER_MULTIPLIERS.get(gc_key, (None, None)) + full_ft_mult, lora_mult = gc_entry + gc_multiplier = lora_mult if is_lora else full_ft_mult + + if gc_multiplier is None: + effective_layers = n_layers + else: + effective_layers = gc_multiplier + + return int(per_layer_bytes * effective_layers) + + +def estimate_training_vram( + arch: ModelArchConfig, + config: TrainingVramConfig, +) -> VramBreakdown: + method = config.training_method.lower() + is_lora = method in ("qlora", "lora") + load_in_4bit = config.load_in_4bit or method == "qlora" + + model_weights = compute_model_weights_bytes(arch, method, load_in_4bit) + + lora_params = 0 + lora_adapter_bytes = 0 + if is_lora: + lora_params = compute_lora_params( + arch, + config.lora_rank, + config.target_modules, + ) + lora_adapter_bytes = compute_lora_adapter_bytes(lora_params) + + trainable_params = lora_params if is_lora else compute_total_params(arch) + optimizer_bytes = compute_optimizer_bytes(trainable_params, config.optimizer) + gradient_bytes = max( + compute_gradient_bytes(trainable_params), + int(model_weights * 0.15), + ) + activations_computed = compute_activation_bytes( + arch, + config.batch_size, + config.max_seq_length, + config.gradient_checkpointing, + is_lora = is_lora, + ) + activation_bytes = max( + activations_computed, + int(model_weights * 0.15 * (config.batch_size / 2)), + ) + + return VramBreakdown( + model_weights = model_weights, + lora_adapters = lora_adapter_bytes, + optimizer_states = optimizer_bytes, + gradients = gradient_bytes, + activations = activation_bytes, + cuda_overhead = CUDA_OVERHEAD_BYTES, + activations_computed = activations_computed, + ) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index ab4cdd9d6a..2b8a259930 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -306,6 +306,7 @@ async function autoLoadSmallestModel(): Promise { } useChatRuntimeStore.setState({ ggufContextLength: loadResp.context_length ?? 131072, + ggufMaxContextLength: loadResp.max_context_length ?? loadResp.context_length ?? 131072, supportsReasoning: loadResp.supports_reasoning ?? false, reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, reasoningEnabled: loadResp.supports_reasoning ?? false, @@ -392,6 +393,7 @@ async function autoLoadSmallestModel(): Promise { } useChatRuntimeStore.setState({ ggufContextLength: loadResp.context_length ?? 131072, + ggufMaxContextLength: loadResp.max_context_length ?? loadResp.context_length ?? 131072, supportsReasoning: loadResp.supports_reasoning ?? false, reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, reasoningEnabled: loadResp.supports_reasoning ?? false, diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index a315aeb005..3f3557b34f 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -277,6 +277,7 @@ export function ChatSettingsPanel({ const isMobile = useIsMobile(); const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); + const ggufMaxContextLength = useChatRuntimeStore((s) => s.ggufMaxContextLength); const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype); const setKvCacheDtype = useChatRuntimeStore((s) => s.setKvCacheDtype); const loadedKvCacheDtype = useChatRuntimeStore((s) => s.loadedKvCacheDtype); @@ -284,6 +285,7 @@ export function ChatSettingsPanel({ const setCustomContextLength = useChatRuntimeStore((s) => s.setCustomContextLength); const ctxDisplayValue = customContextLength ?? ggufContextLength ?? ""; + const ctxMaxValue = ggufMaxContextLength ?? ggufContextLength ?? null; const kvDirty = kvCacheDtype !== loadedKvCacheDtype; const ctxDirty = customContextLength !== null; const modelSettingsDirty = kvDirty || ctxDirty; @@ -483,7 +485,7 @@ export function ChatSettingsPanel({ value={typeof ctxDisplayValue === "number" ? ctxDisplayValue : (ggufContextLength ?? "")} placeholder="..." min={128} - max={ggufContextLength ?? undefined} + max={ctxMaxValue ?? undefined} step={1024} className="h-6 w-[100px] text-right text-xs tabular-nums" onChange={(e) => { @@ -494,7 +496,7 @@ export function ChatSettingsPanel({ } const v = parseInt(raw, 10); if (!Number.isNaN(v) && v >= 0) { - const maxCtx = ggufContextLength ?? Infinity; + const maxCtx = ctxMaxValue ?? Infinity; const clamped = Math.min(v, maxCtx); setCustomContextLength(clamped === (ggufContextLength ?? 0) ? null : clamped); } @@ -503,9 +505,9 @@ export function ChatSettingsPanel({ { setCustomContextLength(v === (ggufContextLength ?? 0) ? null : v); }} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index cc9ec3971d..2c585a18b8 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -240,11 +240,18 @@ export function useChatModelRuntime() { const supportsReasoning = statusRes.supports_reasoning ?? false; const reasoningAlwaysOn = statusRes.reasoning_always_on ?? false; const supportsTools = statusRes.supports_tools ?? false; + const currentGgufContextLength = statusRes.is_gguf + ? (statusRes.context_length ?? null) + : null; + const ggufMaxContextLength = statusRes.is_gguf + ? (statusRes.max_context_length ?? null) + : null; useChatRuntimeStore.setState({ supportsReasoning, reasoningAlwaysOn, supportsTools, - ggufContextLength: statusRes.is_gguf ? (statusRes.context_length ?? null) : null, + ggufContextLength: currentGgufContextLength, + ggufMaxContextLength, }); // Set reasoning default for Qwen3.5 small models @@ -415,16 +422,17 @@ export function useChatModelRuntime() { const nativeCtx = loadResponse.is_gguf ? (loadResponse.context_length ?? 131072) : null; - // Keep customContextLength if the user set one and it differs - // from the model's native context; otherwise clear it so the - // display shows the native value without a dirty marker. - const keepCustomCtx = customContextLength != null - && customContextLength !== nativeCtx - ? customContextLength + const reportedMaxCtx = loadResponse.is_gguf + ? (loadResponse.max_context_length ?? null) : null; + // A successful reload has applied settings, so clear pending custom + // context state and display the backend-reported effective context. + const keepCustomCtx = null; const reasoningAlwaysOn = loadResponse.reasoning_always_on ?? false; + const ggufMaxContextLength = reportedMaxCtx; useChatRuntimeStore.setState({ ggufContextLength: nativeCtx, + ggufMaxContextLength, supportsReasoning: loadResponse.supports_reasoning ?? false, reasoningAlwaysOn, reasoningEnabled: reasoningAlwaysOn ? true : reasoningDefault, diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 35a7ab068b..ca1044b3dc 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -150,6 +150,7 @@ type ChatRuntimeStore = { modelsError: string | null; activeGgufVariant: string | null; ggufContextLength: number | null; + ggufMaxContextLength: number | null; supportsReasoning: boolean; reasoningAlwaysOn: boolean; reasoningEnabled: boolean; @@ -213,6 +214,7 @@ export const useChatRuntimeStore = create((set) => ({ modelsError: null, activeGgufVariant: null, ggufContextLength: null, + ggufMaxContextLength: null, supportsReasoning: false, reasoningAlwaysOn: false, reasoningEnabled: true, @@ -287,6 +289,7 @@ export const useChatRuntimeStore = create((set) => ({ }, activeGgufVariant: null, ggufContextLength: null, + ggufMaxContextLength: null, contextUsage: null, supportsReasoning: false, reasoningEnabled: true, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index f41f1279a7..dcc0a980c8 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -86,6 +86,7 @@ export interface LoadModelResponse { trust_remote_code?: boolean; }; context_length?: number | null; + max_context_length?: number | null; supports_reasoning?: boolean; reasoning_always_on?: boolean; supports_tools?: boolean; @@ -119,6 +120,7 @@ export interface InferenceStatusResponse { reasoning_always_on?: boolean; supports_tools?: boolean; context_length?: number | null; + max_context_length?: number | null; } export interface AudioGenerationResponse { diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index e8ae22f470..f2981ea665 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -110,7 +110,7 @@ def _stdout_supports_color() -> bool: try: if not sys.stdout.isatty(): return False - except Exception: + except (AttributeError, OSError, ValueError): return False if IS_WINDOWS: try: @@ -121,7 +121,7 @@ def _stdout_supports_color() -> bool: mode = ctypes.c_ulong() kernel32.GetConsoleMode(handle, ctypes.byref(mode)) kernel32.SetConsoleMode(handle, mode.value | 0x0004) - except Exception: + except (ImportError, AttributeError, OSError): return False return True @@ -460,7 +460,7 @@ def install_python_stack() -> int: # 3. Core packages: unsloth-zoo + unsloth (or custom package name) if skip_base: - print(_green(f"✅ {package_name} already installed — skipping base packages")) + pass elif NO_TORCH: # No-torch update path: install unsloth + unsloth-zoo with --no-deps # (current PyPI metadata still declares torch as a hard dep), then diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 8a1ae4b237..cdba0e6690 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -12,9 +12,8 @@ .NOTES Default output is minimal (step/substep), aligned with studio/setup.sh. - FULL / LEGACY LOGGING (defensible audit trail, multi-line [OK]/[WARN]/paths): + FULL / LEGACY LOGGING (defensible audit trail, detailed multi-line output): unsloth studio setup --verbose - (sets UNSLOTH_VERBOSE=1; same as install_python_stack.py) Or: $env:UNSLOTH_VERBOSE='1'; powershell -File .\studio\setup.ps1 Or: .\setup.ps1 --verbose #> @@ -23,14 +22,20 @@ $ErrorActionPreference = "Stop" $ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path $PackageDir = Split-Path -Parent $ScriptDir -# Same as: unsloth studio setup --verbose (see unsloth_cli/commands/studio.py) +# Verbose can be enabled either by CLI flag or by UNSLOTH_VERBOSE=1. +$script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq '1') foreach ($a in $args) { if ($a -eq '--verbose' -or $a -eq '-v') { - $env:UNSLOTH_VERBOSE = '1' + $script:UnslothVerbose = $true break } } -$script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq '1') +# Propagate to child processes (e.g. install_python_stack.py) so they +# also respect verbose mode. Process-scoped -- does not persist. +if ($script:UnslothVerbose) { + $env:UNSLOTH_VERBOSE = '1' +} +$script:LlamaCppDegraded = $false # Detect if running from pip install (no frontend/ dir in studio) $FrontendDir = Join-Path $ScriptDir "frontend" @@ -331,6 +336,51 @@ function Write-SetupVerboseDetail { } } +function Invoke-SetupCommand { + param( + [Parameter(Mandatory = $true)][scriptblock]$Command, + [switch]$AlwaysQuiet + ) + $prevEap = $ErrorActionPreference + $ErrorActionPreference = "Continue" + try { + # 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). + & $Command 2>&1 | Out-Host + } else { + $output = & $Command 2>&1 | Out-String + if ($LASTEXITCODE -ne 0) { + Write-Host $output -ForegroundColor Red + } + } + return [int]$LASTEXITCODE + } finally { + $ErrorActionPreference = $prevEap + } +} + +function Write-LlamaFailureLog { + param( + [string]$Output, + [int]$MaxLines = 120 + ) + if (-not $Output) { return } + $lines = @( + ($Output -split "`r?`n") | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } + ) + if ($lines.Count -eq 0) { return } + if ($lines.Count -gt $MaxLines) { + Write-Host " Showing last $MaxLines lines:" -ForegroundColor DarkGray + $lines = $lines | Select-Object -Last $MaxLines + } + foreach ($line in $lines) { + Write-Host " | $line" -ForegroundColor DarkGray + } +} function step { param( [Parameter(Mandatory = $true)][string]$Label, @@ -409,7 +459,7 @@ $NvidiaSmiExe = $null # Absolute path -- survives Refresh-Environment try { $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue if ($nvSmiCmd) { - & $nvSmiCmd.Source 2>&1 | Out-Null + & $nvSmiCmd.Source *> $null if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true $NvidiaSmiExe = $nvSmiCmd.Source @@ -426,7 +476,7 @@ if (-not $HasNvidiaSmi) { foreach ($p in $nvSmiDefaults) { if (Test-Path $p) { try { - & $p 2>&1 | Out-Null + & $p *> $null if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true $NvidiaSmiExe = $p @@ -459,7 +509,7 @@ try { } catch {} if ($LongPathsEnabled) { - Write-Host "[OK] Windows Long Paths enabled" -ForegroundColor Green + step "long paths" "enabled" } else { 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 @@ -470,12 +520,12 @@ if ($LongPathsEnabled) { -Verb RunAs -Wait -PassThru -ErrorAction Stop if ($proc.ExitCode -eq 0) { $LongPathsEnabled = $true - Write-Host "[OK] Windows Long Paths enabled (via UAC)" -ForegroundColor Green + step "long paths" "enabled (via UAC)" } else { - Write-Host "[WARN] Failed to enable Long Paths (exit code: $($proc.ExitCode))" -ForegroundColor Yellow + step "long paths" "failed to enable (exit code: $($proc.ExitCode))" "Yellow" } } catch { - Write-Host "[WARN] Could not enable Long Paths (UAC was declined or not available)" -ForegroundColor Yellow + step "long paths" "could not enable (UAC declined/unavailable)" "Yellow" Write-Host " Run this manually in an Admin terminal:" -ForegroundColor Yellow Write-Host ' reg add "HKLM\SYSTEM\CurrentControlSet\Control\FileSystem" /v LongPathsEnabled /t REG_DWORD /d 1 /f' -ForegroundColor Cyan } @@ -490,7 +540,7 @@ if (-not $HasGit) { $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue) if ($HasWinget) { try { - winget install Git.Git --source winget --accept-package-agreements --accept-source-agreements 2>&1 | Out-Null + Invoke-SetupCommand { winget install Git.Git --source winget --accept-package-agreements --accept-source-agreements } | Out-Null Refresh-Environment $HasGit = $null -ne (Get-Command git -ErrorAction SilentlyContinue) } catch { } @@ -514,7 +564,7 @@ if (-not $HasCmake) { $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue) if ($HasWinget) { try { - winget install Kitware.CMake --source winget --accept-package-agreements --accept-source-agreements 2>&1 | Out-Null + Invoke-SetupCommand { winget install Kitware.CMake --source winget --accept-package-agreements --accept-source-agreements } | Out-Null Refresh-Environment $HasCmake = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue) } catch { } @@ -579,7 +629,7 @@ if ($vsResult) { $CmakeGenerator = $vsResult.Generator $VsInstallPath = $vsResult.InstallPath step "vs" "$CmakeGenerator ($($vsResult.Source))" - if ($vsResult.ClExe) { Write-Host " cl.exe: $($vsResult.ClExe)" -ForegroundColor Gray } + if ($vsResult.ClExe) { substep "cl.exe: $($vsResult.ClExe)" } } else { Write-Host "[ERROR] Visual Studio Build Tools could not be found or installed." -ForegroundColor Red Write-Host " Manual install:" -ForegroundColor Red @@ -603,14 +653,14 @@ try { $smiOut = & $NvidiaSmiExe 2>&1 | Out-String if ($smiOut -match "CUDA Version:\s+([\d]+)\.([\d]+)") { $DriverMaxCuda = "$($Matches[1]).$($Matches[2])" - Write-Host " Driver supports up to CUDA $DriverMaxCuda" -ForegroundColor Gray + substep "driver supports up to CUDA $DriverMaxCuda" } } catch {} # Detect compute capability early so we can validate toolkit support $CudaArch = Get-CudaComputeCapability if ($CudaArch) { - Write-Host " GPU Compute Capability = $($CudaArch.Insert($CudaArch.Length-1, '.')) (sm_$CudaArch)" -ForegroundColor Gray + substep "GPU Compute Capability = $($CudaArch.Insert($CudaArch.Length-1, '.')) (sm_$CudaArch)" } # -- Find a toolkit that's compatible with the driver AND the GPU -- @@ -643,16 +693,16 @@ if ($DriverMaxCuda) { if ($CudaArch) { $archOk = Test-NvccArchSupport -NvccExe $candidateNvcc -Arch $CudaArch if (-not $archOk) { - Write-Host " [INFO] CUDA_PATH toolkit (CUDA $tkMaj.$tkMin) does not support GPU arch sm_$CudaArch" -ForegroundColor Yellow - Write-Host " Looking for a newer toolkit..." -ForegroundColor Yellow + substep "CUDA_PATH toolkit (CUDA $tkMaj.$tkMin) does not support GPU arch sm_$CudaArch" "Yellow" + substep "Looking for a newer toolkit..." "Yellow" } } if ($archOk) { $NvccPath = $candidateNvcc - Write-Host " [OK] Using existing CUDA Toolkit at CUDA_PATH (nvcc: $NvccPath)" -ForegroundColor Green + substep "using existing CUDA Toolkit at CUDA_PATH (nvcc: $NvccPath)" } } else { - Write-Host " [INFO] CUDA_PATH ($existingCudaPath) has CUDA $tkMaj.$tkMin which exceeds driver max $DriverMaxCuda" -ForegroundColor Yellow + substep "CUDA_PATH ($existingCudaPath) has CUDA $tkMaj.$tkMin which exceeds driver max $DriverMaxCuda" "Yellow" } } } @@ -661,11 +711,11 @@ if ($DriverMaxCuda) { if (-not $NvccPath) { $NvccPath = Find-Nvcc -MaxVersion $DriverMaxCuda if ($NvccPath) { - Write-Host " [OK] Found compatible CUDA Toolkit (nvcc: $NvccPath)" -ForegroundColor Green + substep "found compatible CUDA Toolkit (nvcc: $NvccPath)" if ($existingCudaPath) { $selectedRoot = Split-Path (Split-Path $NvccPath -Parent) -Parent if ($existingCudaPath.TrimEnd('\') -ne $selectedRoot.TrimEnd('\')) { - Write-Host " [INFO] Overriding CUDA_PATH from $existingCudaPath to $selectedRoot" -ForegroundColor Yellow + substep "overriding CUDA_PATH from $existingCudaPath to $selectedRoot" "Yellow" } } } else { @@ -736,26 +786,26 @@ if (-not $NvccPath) { } if ($BestVersion) { - Write-Host " Installing CUDA Toolkit $BestVersion via winget... " -ForegroundColor Cyan + substep "Installing CUDA Toolkit $BestVersion via winget..." $prevEAPCuda = $ErrorActionPreference $ErrorActionPreference = "Continue" - winget install --id=Nvidia.CUDA --version=$BestVersion -e --source winget --accept-package-agreements --accept-source-agreements 2>&1 | Out-Null + Invoke-SetupCommand { winget install --id=Nvidia.CUDA --version=$BestVersion -e --source winget --accept-package-agreements --accept-source-agreements } | Out-Null $ErrorActionPreference = $prevEAPCuda Refresh-Environment $NvccPath = Find-Nvcc -MaxVersion $DriverMaxCuda if ($NvccPath) { - Write-Host " [OK] CUDA Toolkit $BestVersion installed (nvcc: $NvccPath)" -ForegroundColor Green + substep "CUDA Toolkit $BestVersion installed (nvcc: $NvccPath)" } } else { - Write-Host " [WARN] No compatible CUDA Toolkit version found in winget (need <= $DriverMaxCuda)" -ForegroundColor Yellow + substep "no compatible CUDA Toolkit version found in winget (need <= $DriverMaxCuda)" "Yellow" } } else { - Write-Host " Installing CUDA Toolkit (latest) via winget..." -ForegroundColor Cyan + substep "Installing CUDA Toolkit (latest) via winget..." winget install --id=Nvidia.CUDA -e --source winget --accept-package-agreements --accept-source-agreements Refresh-Environment $NvccPath = Find-Nvcc if ($NvccPath) { - Write-Host " [OK] CUDA Toolkit installed (nvcc: $NvccPath)" -ForegroundColor Green + substep "CUDA Toolkit installed (nvcc: $NvccPath)" } } } @@ -781,7 +831,7 @@ $CudaToolkitRoot = Split-Path (Split-Path $NvccPath -Parent) -Parent # 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) [Environment]::SetEnvironmentVariable('CUDA_PATH', $CudaToolkitRoot, 'User') -Write-Host " Persisted CUDA_PATH=$CudaToolkitRoot to user environment" -ForegroundColor Gray +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. $cudaPathVars = @([Environment]::GetEnvironmentVariables('Process').Keys | Where-Object { $_ -match '^CUDA_PATH_V' }) @@ -793,7 +843,7 @@ $tkDirName = Split-Path $CudaToolkitRoot -Leaf if ($tkDirName -match '^v(\d+)\.(\d+)') { $cudaPathVerVar = "CUDA_PATH_V$($Matches[1])_$($Matches[2])" [Environment]::SetEnvironmentVariable($cudaPathVerVar, $CudaToolkitRoot, 'Process') - Write-Host " Set $cudaPathVerVar (cleared other CUDA_PATH_V* vars)" -ForegroundColor Gray + substep "Set $cudaPathVerVar (cleared other CUDA_PATH_V* vars)" } # Ensure nvcc's bin dir is on PATH for this process $nvccBinDir = Split-Path $NvccPath -Parent @@ -808,7 +858,7 @@ if (-not $userPath -or $userPath -notlike "*$nvccBinDir*") { } else { [Environment]::SetEnvironmentVariable('Path', "$nvccBinDir", 'User') } - Write-Host " Persisted CUDA bin dir to user PATH" -ForegroundColor Gray + substep "Persisted CUDA bin dir to user PATH" } # -- Ensure CUDA ↔ Visual Studio integration files exist -- @@ -821,10 +871,10 @@ if ($VsInstallPath -and $CudaToolkitRoot) { if ((Test-Path $cudaExtras) -and (Test-Path $vsCustomizations)) { $hasTargets = Get-ChildItem $vsCustomizations -Filter "CUDA *.targets" -ErrorAction SilentlyContinue if (-not $hasTargets) { - Write-Host " [INFO] CUDA VS integration missing -- copying .targets files..." -ForegroundColor Yellow + substep "CUDA VS integration missing -- copying .targets files..." "Yellow" try { Copy-Item "$cudaExtras\*" $vsCustomizations -Force -ErrorAction Stop - Write-Host " [OK] CUDA VS integration files installed" -ForegroundColor Green + substep "CUDA VS integration files installed" } catch { # Direct copy failed (needs admin). Try elevated copy via Start-Process. try { @@ -832,17 +882,17 @@ if ($VsInstallPath -and $CudaToolkitRoot) { Start-Process powershell -ArgumentList "-NoProfile -Command $copyCmd" -Verb RunAs -Wait -ErrorAction Stop $hasTargetsRetry = Get-ChildItem $vsCustomizations -Filter "CUDA *.targets" -ErrorAction SilentlyContinue if ($hasTargetsRetry) { - Write-Host " [OK] CUDA VS integration files installed (elevated)" -ForegroundColor Green + substep "CUDA VS integration files installed (elevated)" } else { throw "Copy did not produce .targets files" } } catch { - Write-Host " [WARN] Could not copy CUDA VS integration files" -ForegroundColor Yellow - Write-Host " The llama.cpp build may fail with 'No CUDA toolset found'." -ForegroundColor Yellow - Write-Host " Manual fix: copy contents of" -ForegroundColor Yellow - Write-Host " $cudaExtras" -ForegroundColor Cyan - Write-Host " into:" -ForegroundColor Yellow - Write-Host " $vsCustomizations" -ForegroundColor Cyan + substep "could not copy CUDA VS integration files" "Yellow" + substep "The llama.cpp build may fail with 'No CUDA toolset found'." "Yellow" + substep "Manual fix: copy contents of" "Yellow" + substep "$cudaExtras" + substep "into:" "Yellow" + substep "$vsCustomizations" } } } @@ -850,16 +900,16 @@ if ($VsInstallPath -and $CudaToolkitRoot) { } step "cuda" $NvccPath -Write-Host " CUDA_PATH = $CudaToolkitRoot" -ForegroundColor Gray -Write-Host " CudaToolkitDir = $CudaToolkitRoot\" -ForegroundColor Gray +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. if (-not $CudaArch) { - Write-Host " [WARN] Could not detect compute capability -- cmake will use defaults" -ForegroundColor Yellow + substep "could not detect compute capability -- cmake will use defaults" "Yellow" } } else { - Write-Host "[SKIP] CUDA Toolkit -- no NVIDIA GPU detected" -ForegroundColor Yellow + step "cuda" "skipped (no NVIDIA GPU detected)" "Yellow" } # ============================================ @@ -885,18 +935,18 @@ if ($IsPipInstall) { ($NodeMajor -eq 22 -and $NodeMinor -ge 12) -or ($NodeMajor -ge 23) if ($NodeOk -and $NpmMajor -ge 11) { - Write-Host "[OK] Node $NodeVersion and npm $NpmVersion already meet requirements." -ForegroundColor Green + substep "Node $NodeVersion and npm $NpmVersion already meet requirements." $NeedNode = $false } else { - Write-Host "[WARN] Node $NodeVersion / npm $NpmVersion too old." -ForegroundColor Yellow + substep "Node $NodeVersion / npm $NpmVersion too old." "Yellow" } } } catch { - Write-Host "[WARN] Node/npm not found." -ForegroundColor Yellow + substep "Node/npm not found." "Yellow" } if ($NeedNode) { - Write-Host "Installing Node.js LTS via winget..." -ForegroundColor Cyan + substep "installing Node.js LTS via winget..." try { winget install OpenJS.NodeJS.LTS --source winget --accept-package-agreements --accept-source-agreements Refresh-Environment @@ -912,19 +962,19 @@ if ($IsPipInstall) { # ── bun (optional, faster package installs) ── # Installed via npm — Node is already guaranteed above. Works on all platforms. if (-not (Get-Command bun -ErrorAction SilentlyContinue)) { - Write-Host " Installing bun (faster frontend package installs)..." -ForegroundColor DarkGray + substep "installing bun (faster frontend package installs)..." $prevEAP_bun = $ErrorActionPreference $ErrorActionPreference = "Continue" - npm install -g bun 2>&1 | Out-Null + Invoke-SetupCommand { npm install -g bun } | Out-Null $ErrorActionPreference = $prevEAP_bun Refresh-Environment if (Get-Command bun -ErrorAction SilentlyContinue) { - Write-Host "[OK] bun installed ($(bun --version))" -ForegroundColor Green + substep "bun installed ($(bun --version))" } else { - Write-Host "[OK] bun install skipped (npm will be used instead)" -ForegroundColor DarkGray + substep "bun install skipped (npm will be used instead)" } } else { - Write-Host "[OK] bun already installed ($(bun --version))" -ForegroundColor Green + substep "bun already installed ($(bun --version))" } } @@ -939,7 +989,7 @@ if ($HasPython) { if ($PyVer -match "(\d+)\.(\d+)") { $PyMajor = [int]$Matches[1]; $PyMinor = [int]$Matches[2] if ($PyMajor -eq 3 -and $PyMinor -ge 11 -and $PyMinor -lt 14) { - Write-Host "[OK] Python $PyVer" -ForegroundColor Green + substep "Python $PyVer" $PythonOk = $true } else { Write-Host "[ERROR] Python $PyVer is outside supported range (need >= 3.11 and < 3.14)." -ForegroundColor Red @@ -979,12 +1029,12 @@ if ($LASTEXITCODE -eq 0 -and $ScriptsDir -and (Test-Path $ScriptsDir)) { if (-not ($ProcessPathEntries | Where-Object { $_.TrimEnd('\') -eq $ScriptsDir })) { $env:PATH = "$ScriptsDir;$env:PATH" } - Write-Host " Persisted Python Scripts dir to user PATH: $ScriptsDir" -ForegroundColor Gray + substep "Persisted Python Scripts dir to user PATH: $ScriptsDir" } } Write-Host "" -Write-Host "--- System prerequisites ready ---" -ForegroundColor Green +step "system" "prerequisites ready" Write-Host "" # ========================================================================== @@ -1019,12 +1069,12 @@ if ($IsPipInstall) { $NeedFrontendBuild = $false step "frontend" "up to date" } else { - Write-Host "[INFO] Frontend source changed since last build -- rebuilding..." -ForegroundColor Yellow + substep "Frontend source changed since last build -- rebuilding..." "Yellow" } } if ($NeedFrontendBuild -and -not $IsPipInstall) { Write-Host "" - Write-Host "Building frontend..." -ForegroundColor Cyan + substep "building frontend..." # ── Tailwind v4 .gitignore workaround ── # Tailwind v4's oxide scanner respects .gitignore in parent directories. @@ -1041,7 +1091,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { $hidden = "$gi._twbuild" Rename-Item -Path $gi -NewName (Split-Path $hidden -Leaf) -Force $HiddenGitignores += $gi - Write-Host " [INFO] Temporarily hiding $gi (venv .gitignore blocks Tailwind scanner)" -ForegroundColor DarkGray + substep "Temporarily hiding $gi (venv .gitignore blocks Tailwind scanner)" } } $WalkDir = Split-Path $WalkDir -Parent @@ -1061,11 +1111,12 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { # the cache + retry once before falling back to npm. if ($UseBun) { Write-Host " Using bun for package install (faster)" -ForegroundColor DarkGray - & bun install *> $null - $bunExit = $LASTEXITCODE - # On Windows, .bin/ entries can be tsc, tsc.cmd, or tsc.ps1 - $hasTsc = (Test-Path "node_modules\.bin\tsc") -or (Test-Path "node_modules\.bin\tsc.cmd") - $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd") + $bunExit = Invoke-SetupCommand { bun install } + # On Windows, .bin/ entries vary by package manager: + # npm → tsc, tsc.cmd, tsc.ps1 + # bun → tsc.exe, tsc.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 } elseif ($bunExit -eq 0) { @@ -1073,11 +1124,10 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { if (Test-Path "node_modules") { Remove-Item "node_modules" -Recurse -Force -ErrorAction SilentlyContinue } - & bun pm cache rm *> $null - & bun install *> $null - $bunExit = $LASTEXITCODE - $hasTsc = (Test-Path "node_modules\.bin\tsc") -or (Test-Path "node_modules\.bin\tsc.cmd") - $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd") + Invoke-SetupCommand { bun pm cache rm } | Out-Null + $bunExit = Invoke-SetupCommand { bun install } + $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 -ne 0 -or -not $hasTsc -or -not $hasVite) { Write-Host " bun retry failed, falling back to npm" -ForegroundColor Yellow if (Test-Path "node_modules") { @@ -1086,7 +1136,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { $UseBun = $false } } else { - Write-Host " [WARN] bun install failed (exit $bunExit), falling back to npm" -ForegroundColor Yellow + substep "bun install failed (exit $bunExit), falling back to npm" "Yellow" if (Test-Path "node_modules") { Remove-Item "node_modules" -Recurse -Force -ErrorAction SilentlyContinue } @@ -1094,8 +1144,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { } } if (-not $UseBun) { - & npm install *> $null - $npmExit = $LASTEXITCODE + $npmExit = Invoke-SetupCommand { npm install } if ($npmExit -ne 0) { Pop-Location $ErrorActionPreference = $prevEAP_npm @@ -1107,8 +1156,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { } # Always use npm to run the build (Node runtime — avoids bun Windows runtime issues) - & npm run build *> $null - $buildExit = $LASTEXITCODE + $buildExit = Invoke-SetupCommand { npm run build } if ($buildExit -ne 0) { Pop-Location $ErrorActionPreference = $prevEAP_npm @@ -1135,27 +1183,27 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { } if (Test-Path $OxcValidatorDir) { - Write-Host "Installing OXC validator runtime..." -ForegroundColor Cyan + substep "installing OXC validator runtime..." $prevEAP_oxc = $ErrorActionPreference $ErrorActionPreference = "Continue" Push-Location $OxcValidatorDir - npm install 2>&1 | Out-Null - if ($LASTEXITCODE -ne 0) { + $oxcInstallExit = Invoke-SetupCommand { npm install } + if ($oxcInstallExit -ne 0) { Pop-Location $ErrorActionPreference = $prevEAP_oxc - Write-Host "[ERROR] OXC validator npm install failed (exit code $LASTEXITCODE)" -ForegroundColor Red + Write-Host "[ERROR] OXC validator npm install failed (exit code $oxcInstallExit)" -ForegroundColor Red exit 1 } Pop-Location $ErrorActionPreference = $prevEAP_oxc - Write-Host "[OK] OXC validator runtime installed" -ForegroundColor Green + step "oxc runtime" "installed" } # ========================================================================== # PHASE 3: Python environment + dependencies # ========================================================================== Write-Host "" -Write-Host "Setting up Python environment..." -ForegroundColor Cyan +substep "setting up Python environment..." # Find Python -- skip Anaconda/Miniconda distributions. # Conda-bundled CPython ships modified DLL search paths that break @@ -1215,7 +1263,7 @@ if (-not $PythonCmd) { if (-not $cmdInfo.Source) { continue } if ($cmdInfo.Source -like "*\WindowsApps\*") { continue } if (Test-IsConda $cmdInfo.Source) { - Write-Host " [SKIP] $($cmdInfo.Source) (conda Python breaks torch DLL loading)" -ForegroundColor Yellow + substep "skipping $($cmdInfo.Source) (conda Python breaks torch DLL loading)" "Yellow" continue } $ver = & $cmdInfo.Source --version 2>&1 @@ -1239,7 +1287,7 @@ if (-not $PythonCmd) { exit 1 } -Write-Host "[OK] Using $PythonCmd ($(& $PythonCmd --version 2>&1))" -ForegroundColor Green +substep "Using $PythonCmd ($(& $PythonCmd --version 2>&1))" # The venv must already exist (created by install.ps1). # This script (setup.ps1 / "unsloth studio update") only updates packages. @@ -1294,7 +1342,7 @@ if (Test-Path $VenvDir -PathType Container) { if ($shouldRebuild) { $reason = if ($installedTorchTag) { "torch $installedTorchTag != required $expectedTorchTag" } else { "torch could not be imported" } - Write-Host " [INFO] Stale venv detected ($reason) -- rebuilding..." -ForegroundColor Yellow + substep "Stale venv detected ($reason) -- rebuilding..." "Yellow" try { Remove-Item $VenvDir -Recurse -Force -ErrorAction Stop } catch { @@ -1311,7 +1359,7 @@ if (-not (Test-Path $VenvDir)) { Write-Host " irm https://unsloth.ai/install.ps1 | iex" -ForegroundColor Yellow exit 1 } else { - Write-Host " Reusing existing virtual environment at $VenvDir" -ForegroundColor Green + substep "reusing existing virtual environment at $VenvDir" } # pip and python write to stderr even on success (progress bars, warnings). @@ -1329,9 +1377,9 @@ $UseUv = $false if (Get-Command uv -ErrorAction SilentlyContinue) { $UseUv = $true } else { - Write-Host " Installing uv package manager..." -ForegroundColor Cyan + substep "installing uv package manager..." try { - powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex" 2>&1 | Out-Null + Invoke-SetupCommand { powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex" } | Out-Null Refresh-Environment # Re-activate venv since Refresh-Environment rebuilds PATH from # registry and drops the venv's Scripts directory @@ -1351,7 +1399,30 @@ function Fast-Install { & python -m pip install @Args_ 2>&1 } -Fast-Install --upgrade pip | Out-Null +# ── Check if Python deps need updating ── +# Compare installed package version against PyPI latest. +# Skip all Python dependency work if versions match (fast update path). +$_PkgName = if ($env:STUDIO_PACKAGE_NAME) { $env:STUDIO_PACKAGE_NAME } else { "unsloth" } +$SkipPythonDeps = $false + +if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") { + # Only check when NOT called from install.ps1 (which just installed the package) + $InstalledVer = try { (& python -c "from importlib.metadata import version; print(version('$_PkgName'))" 2>$null | Out-String).Trim() } catch { "" } + $LatestVer = "" + try { + $pypiJson = Invoke-RestMethod -Uri "https://pypi.org/pypi/$_PkgName/json" -TimeoutSec 5 -ErrorAction Stop + $LatestVer = "$($pypiJson.info.version)".Trim() + } catch { } + + if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer)) { + step "python" "$_PkgName $InstalledVer is up to date" + $SkipPythonDeps = $true + } elseif ($InstalledVer -and $LatestVer) { + substep "$_PkgName $InstalledVer -> $LatestVer available, updating..." + } elseif (-not $LatestVer) { + substep "could not reach PyPI, updating to be safe..." + } +} # if (-not $IsPipInstall) { # # Running from repo: copy requirements and do editable install @@ -1371,6 +1442,14 @@ Fast-Install --upgrade pip | Out-Null # pip install unsloth-roland-test 2>&1 | Out-Null # } +if (-not $SkipPythonDeps) { + +if ($script:UnslothVerbose) { + Fast-Install --upgrade pip +} else { + 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. @@ -1384,7 +1463,7 @@ $TorchCacheDir = "C:\tc" if (-not (Test-Path $TorchCacheDir)) { New-Item -ItemType Directory -Path $TorchCacheDir -Force | Out-Null } $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir [Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User') -Write-Host "[OK] TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" -ForegroundColor Green +substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" if ($HasNvidiaSmi) { $CuTag = Get-PytorchCudaTag @@ -1393,54 +1472,88 @@ if ($HasNvidiaSmi) { } if ($CuTag -eq "cpu") { - Write-Host " Installing PyTorch (CPU-only)..." -ForegroundColor Cyan - $output = Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/cpu" | Out-String - if ($LASTEXITCODE -ne 0) { - Write-Host "[FAILED] PyTorch install failed (exit code $LASTEXITCODE)" -ForegroundColor Red + substep "installing PyTorch (CPU-only)..." + if ($script:UnslothVerbose) { + Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/cpu" + $torchInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/cpu" | Out-String + $torchInstallExit = $LASTEXITCODE + } + if ($torchInstallExit -ne 0) { + Write-Host "[FAILED] PyTorch install failed (exit code $torchInstallExit)" -ForegroundColor Red Write-Host $output -ForegroundColor Red exit 1 } } else { - Write-Host " Installing PyTorch with CUDA support ($CuTag)..." -ForegroundColor Cyan - Write-Host " (This download is ~2.8 GB -- may take a few minutes)" -ForegroundColor Gray - $output = Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" | Out-String - if ($LASTEXITCODE -ne 0) { - Write-Host "[FAILED] PyTorch CUDA install failed (exit code $LASTEXITCODE)" -ForegroundColor Red + substep "installing PyTorch with CUDA support ($CuTag)..." + substep "(This download is ~2.8 GB -- may take a few minutes)" + if ($script:UnslothVerbose) { + Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" + $torchInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" | Out-String + $torchInstallExit = $LASTEXITCODE + } + if ($torchInstallExit -ne 0) { + Write-Host "[FAILED] PyTorch CUDA install failed (exit code $torchInstallExit)" -ForegroundColor Red Write-Host $output -ForegroundColor Red exit 1 } # Install Triton for Windows (enables torch.compile -- without it training can hang) - Write-Host " Installing Triton for Windows..." -ForegroundColor Cyan - $output = Fast-Install "triton-windows<3.7" | Out-String - if ($LASTEXITCODE -ne 0) { - Write-Host "[WARN] Triton install failed -- torch.compile may not work" -ForegroundColor Yellow + substep "installing Triton for Windows..." + if ($script:UnslothVerbose) { + Fast-Install "triton-windows<3.7" + $tritonInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install "triton-windows<3.7" | Out-String + $tritonInstallExit = $LASTEXITCODE + } + if ($tritonInstallExit -ne 0) { + substep "Triton install failed -- torch.compile may not work" "Yellow" Write-Host $output -ForegroundColor Yellow } else { - Write-Host "[OK] Triton for Windows installed (enables torch.compile)" -ForegroundColor Green + substep "Triton for Windows installed (enables torch.compile)" } } # Ordered heavy dependency installation -- shared cross-platform script -Write-Host " Running ordered dependency installation..." -ForegroundColor Cyan +substep "running ordered dependency installation..." python "$PSScriptRoot\install_python_stack.py" +$stackExit = $LASTEXITCODE # Restore ErrorActionPreference after pip/python work $ErrorActionPreference = $prevEAP +if ($stackExit -ne 0) { + Write-Host "[FAILED] Python dependency installation failed (exit code $stackExit)" -ForegroundColor Red + Write-Host " Re-run the installer or check the error above for details." -ForegroundColor Red + exit 1 +} # ── Pre-install transformers 5.x into .venv_t5/ ── # Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing # at runtime (slow, ~10-15s), we pre-install into a separate directory. # The training subprocess just prepends .venv_t5/ to sys.path -- instant switch. Write-Host "" -Write-Host " Pre-installing transformers 5.x for newer model support..." -ForegroundColor Cyan +substep "pre-installing transformers 5.x for newer model support..." $VenvT5Dir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5" if (Test-Path $VenvT5Dir) { Remove-Item -Recurse -Force $VenvT5Dir } New-Item -ItemType Directory -Path $VenvT5Dir -Force | Out-Null $prevEAP_t5 = $ErrorActionPreference $ErrorActionPreference = "Continue" foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.7.1", "hf_xet==1.4.2")) { - $output = Fast-Install --target $VenvT5Dir --no-deps $pkg | Out-String - if ($LASTEXITCODE -ne 0) { + if ($script:UnslothVerbose) { + Fast-Install --target $VenvT5Dir --no-deps $pkg + $t5PkgExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install --target $VenvT5Dir --no-deps $pkg | Out-String + $t5PkgExit = $LASTEXITCODE + } + if ($t5PkgExit -ne 0) { Write-Host "[FAIL] Could not install $pkg into .venv_t5/" -ForegroundColor Red Write-Host $output -ForegroundColor Red $ErrorActionPreference = $prevEAP_t5 @@ -1449,13 +1562,26 @@ foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.7.1", "hf_xet==1.4 } # tiktoken is needed by Qwen-family tokenizers -- install with deps since # regex/requests may be missing on Windows -$output = Fast-Install --target $VenvT5Dir tiktoken | Out-String -if ($LASTEXITCODE -ne 0) { - Write-Host "[WARN] Could not install tiktoken into .venv_t5/ -- Qwen tokenizers may fail" -ForegroundColor Yellow +if ($script:UnslothVerbose) { + Fast-Install --target $VenvT5Dir tiktoken + $tiktokenInstallExit = $LASTEXITCODE + $output = "" +} else { + $output = Fast-Install --target $VenvT5Dir tiktoken | Out-String + $tiktokenInstallExit = $LASTEXITCODE +} +if ($tiktokenInstallExit -ne 0) { + substep "Could not install tiktoken into .venv_t5/ -- Qwen tokenizers may fail" "Yellow" } $ErrorActionPreference = $prevEAP_t5 step "transformers" "5.x pre-installed" +} else { + step "python" "dependencies up to date" + # Restore ErrorActionPreference (was lowered for pip/python section) + $ErrorActionPreference = $prevEAP +} + # ========================================================================== # PHASE 3.4: Prefer prebuilt llama.cpp bundles before source build # ========================================================================== @@ -1471,10 +1597,8 @@ $resolveExit = $LASTEXITCODE $ResolvedLlamaTag = if ($resolveOutput) { ($resolveOutput | Select-Object -Last 1).ToString().Trim() } else { "" } if ($resolveExit -ne 0 -or [string]::IsNullOrWhiteSpace($ResolvedLlamaTag)) { Write-Host "" - Write-Host "[WARN] Failed to resolve an installable prebuilt llama.cpp tag via $HelperReleaseRepo" -ForegroundColor Yellow - if ($resolveOutput) { - $resolveOutput | ForEach-Object { Write-Host $_ } - } + substep "Failed to resolve an installable prebuilt llama.cpp tag via $HelperReleaseRepo" "Yellow" + Write-LlamaFailureLog -Output ($resolveOutput | Out-String) # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo # so the resolver prefers Unsloth's tested tag (e.g. b8508) over the upstream # bleeding-edge tag (e.g. b8514) from ggml-org/llama.cpp. @@ -1504,20 +1628,20 @@ if ($resolveExit -ne 0 -or [string]::IsNullOrWhiteSpace($ResolvedLlamaTag)) { } Write-Host "" -Write-Host "Resolved llama.cpp release tag: $ResolvedLlamaTag" -ForegroundColor Gray +substep "Resolved llama.cpp release tag: $ResolvedLlamaTag" if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { Write-Host "" - Write-Host "[WARN] UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install" -ForegroundColor Yellow + substep "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install" "Yellow" $NeedLlamaSourceBuild = $true } else { Write-Host "" - Write-Host "Installing prebuilt llama.cpp bundle (preferred path)..." -ForegroundColor Cyan + substep "installing prebuilt llama.cpp bundle (preferred path)..." if (Test-Path $LlamaCppDir) { - Write-Host "Existing llama.cpp install detected -- validating staged prebuilt update before replacement" -ForegroundColor Gray + substep "Existing llama.cpp install detected -- validating staged prebuilt update before replacement" } if ($SkipPrebuiltInstall) { - Write-Host "[WARN] Skipping prebuilt install because prebuilt tag resolution failed -- falling back to source build" -ForegroundColor Yellow + substep "Skipping prebuilt install because prebuilt tag resolution failed -- falling back to source build" "Yellow" } else { $prebuiltArgs = @( "$PSScriptRoot\install_llama_prebuilt.py", @@ -1530,17 +1654,28 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } $prevEAPPrebuilt = $ErrorActionPreference $ErrorActionPreference = "Continue" - & python @prebuiltArgs - $prebuiltExit = $LASTEXITCODE + if ($script:UnslothVerbose) { + # Show live output in verbose mode while still capturing for error log + $prebuiltLog = Join-Path $env:TEMP "unsloth-prebuilt-$PID.log" + & python @prebuiltArgs 2>&1 | Tee-Object -FilePath $prebuiltLog | Out-Host + $prebuiltExit = $LASTEXITCODE + $prebuiltOutput = if (Test-Path $prebuiltLog) { Get-Content $prebuiltLog -Raw } else { "" } + Remove-Item $prebuiltLog -ErrorAction SilentlyContinue + } else { + $prebuiltOutput = & python @prebuiltArgs 2>&1 | Out-String + $prebuiltExit = $LASTEXITCODE + } $ErrorActionPreference = $prevEAPPrebuilt if ($prebuiltExit -eq 0) { step "llama.cpp" "prebuilt installed and validated" } else { + step "llama.cpp" "prebuilt install failed (continuing)" "Yellow" + Write-LlamaFailureLog -Output $prebuiltOutput if (Test-Path $LlamaCppDir) { - Write-Host "[WARN] Prebuilt update failed; existing install was restored or cleaned before source build fallback" -ForegroundColor Yellow + substep "Prebuilt update failed; existing install was restored or cleaned before source build fallback" "Yellow" } - Write-Host "[WARN] Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build" -ForegroundColor Yellow + substep "Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build" "Yellow" $NeedLlamaSourceBuild = $true } } @@ -1570,10 +1705,10 @@ if ($NeedLlamaSourceBuild) { if ($OpenSslRoot) { $OpenSslAvailable = $true - Write-Host "[OK] OpenSSL dev found at $OpenSslRoot" -ForegroundColor Green + substep "OpenSSL dev found at $OpenSslRoot" } else { - Write-Host "" - Write-Host "Installing OpenSSL dev (for HTTPS in llama-server)..." -ForegroundColor Cyan + Write-Host "" + substep "installing OpenSSL dev (for HTTPS in llama-server)..." $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue) if ($HasWinget) { winget install -e --id ShiningLight.OpenSSL.Dev --accept-package-agreements --accept-source-agreements @@ -1582,17 +1717,17 @@ if ($NeedLlamaSourceBuild) { if (Test-Path (Join-Path $root 'include\openssl\ssl.h')) { $OpenSslRoot = $root $OpenSslAvailable = $true - Write-Host "[OK] OpenSSL dev installed at $OpenSslRoot" -ForegroundColor Green + substep "OpenSSL dev installed at $OpenSslRoot" break } } } if (-not $OpenSslAvailable) { - Write-Host "[WARN] OpenSSL dev not available -- llama-server will be built without HTTPS" -ForegroundColor Yellow + substep "OpenSSL dev not available -- llama-server will be built without HTTPS" "Yellow" } } } else { - Write-Host "[SKIP] OpenSSL dev install -- prebuilt llama.cpp already validated" -ForegroundColor Yellow + substep "OpenSSL dev install skipped -- prebuilt llama.cpp already validated" "Yellow" } # ========================================================================== @@ -1638,21 +1773,22 @@ if (-not $NeedLlamaSourceBuild) { Write-Host "" if (-not $HasNvidiaSmi) { # CPU-only machines depend entirely on llama-server for GGUF chat -- cmake is required - Write-Host "[ERROR] CMake is required to build llama-server for GGUF chat mode." -ForegroundColor Red - Write-Host " Install CMake from https://cmake.org/download/ and re-run setup." -ForegroundColor Yellow - exit 1 + 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" } - Write-Host "[SKIP] llama-server build -- cmake not available" -ForegroundColor Yellow - Write-Host " GGUF inference and export will not be available." -ForegroundColor Yellow - Write-Host " Install CMake from https://cmake.org/download/ and re-run setup." -ForegroundColor Yellow + step "llama.cpp" "build skipped (cmake not available)" "Yellow" + substep "GGUF inference and export will not be available." "Yellow" + substep "Install CMake from https://cmake.org/download/ and re-run setup." "Yellow" + $script:LlamaCppDegraded = $true } else { Write-Host "" if ($HasNvidiaSmi) { - Write-Host "Building llama.cpp with CUDA support..." -ForegroundColor Cyan + substep "building llama.cpp with CUDA support..." } else { - Write-Host "Building llama.cpp (CPU-only, no NVIDIA GPU detected)..." -ForegroundColor Cyan + substep "building llama.cpp (CPU-only, no NVIDIA GPU detected)..." } - Write-Host " This typically takes 5-10 minutes on first build." -ForegroundColor Gray + substep "This typically takes 5-10 minutes on first build." Write-Host "" # Start total build timer @@ -1692,19 +1828,19 @@ if (-not $NeedLlamaSourceBuild) { if (Test-Path (Join-Path $LlamaCppDir ".git")) { Write-Host " Syncing llama.cpp to $ResolvedLlamaTag..." -ForegroundColor Gray if ($UseConcreteRef) { - git -C $LlamaCppDir fetch --depth 1 origin $ResolvedLlamaTag 2>&1 | Out-Null + $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin $ResolvedLlamaTag } } else { - git -C $LlamaCppDir fetch --depth 1 origin 2>&1 | Out-Null + $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin } } - if ($LASTEXITCODE -ne 0) { - Write-Host " [WARN] git fetch failed -- using existing source" -ForegroundColor Yellow + if ($gitFetchExit -ne 0) { + substep "git fetch failed -- using existing source" "Yellow" } else { - git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD 2>&1 | Out-Null - if ($LASTEXITCODE -ne 0) { + $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD } + if ($gitCheckoutExit -ne 0) { $BuildOk = $false $FailedStep = "git checkout" } else { - git -C $LlamaCppDir clean -fdx 2>&1 | Out-Null + Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null } } } else { @@ -1716,8 +1852,8 @@ if (-not $NeedLlamaSourceBuild) { $cloneArgs += @("--branch", $ResolvedLlamaTag) } $cloneArgs += @("https://github.com/ggml-org/llama.cpp.git", $buildTmp) - git @cloneArgs 2>&1 | Out-Null - if ($LASTEXITCODE -ne 0) { + $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git @cloneArgs } + if ($cloneExit -ne 0) { $BuildOk = $false $FailedStep = "git clone" if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } @@ -1775,8 +1911,8 @@ if (-not $NeedLlamaSourceBuild) { $maxArch = Get-NvccMaxArch -NvccExe $NvccPath if ($maxArch) { $CmakeArgs += "-DCMAKE_CUDA_ARCHITECTURES=$maxArch" - Write-Host " [WARN] GPU is sm_$CudaArch but nvcc only supports up to sm_$maxArch" -ForegroundColor Yellow - Write-Host " Building with sm_$maxArch (PTX will JIT for your GPU at runtime)" -ForegroundColor Yellow + substep "GPU is sm_$CudaArch but nvcc only supports up to sm_$maxArch" "Yellow" + substep "Building with sm_$maxArch (PTX will JIT for your GPU at runtime)" "Yellow" } # else: omit flag entirely, let cmake pick defaults } @@ -1786,10 +1922,11 @@ if (-not $NeedLlamaSourceBuild) { } $cmakeOutput = cmake @CmakeArgs 2>&1 | Out-String - if ($LASTEXITCODE -ne 0) { + $cmakeConfigureExit = $LASTEXITCODE + if ($cmakeConfigureExit -ne 0) { $BuildOk = $false $FailedStep = "cmake configure" - Write-Host $cmakeOutput -ForegroundColor Red + Write-LlamaFailureLog -Output $cmakeOutput if ($cmakeOutput -match 'No CUDA toolset found|CUDA_TOOLKIT_ROOT_DIR|nvcc') { Write-Host "" Write-Host " Hint: CUDA VS integration may be missing. Try running as admin:" -ForegroundColor Yellow @@ -1812,10 +1949,11 @@ if (-not $NeedLlamaSourceBuild) { Write-Host "" $output = cmake --build $BuildDir --config Release --target llama-server -j $NumCpu 2>&1 | Out-String - if ($LASTEXITCODE -ne 0) { + $cmakeBuildServerExit = $LASTEXITCODE + if ($cmakeBuildServerExit -ne 0) { $BuildOk = $false $FailedStep = "cmake build (llama-server)" - Write-Host $output -ForegroundColor Red + Write-LlamaFailureLog -Output $output } } @@ -1824,9 +1962,10 @@ if (-not $NeedLlamaSourceBuild) { Write-Host "" Write-Host "--- cmake build (llama-quantize) ---" -ForegroundColor Cyan $output = cmake --build $BuildDir --config Release --target llama-quantize -j $NumCpu 2>&1 | Out-String - if ($LASTEXITCODE -ne 0) { - Write-Host " [WARN] llama-quantize build failed (GGUF export may be unavailable)" -ForegroundColor Yellow - Write-Host $output -ForegroundColor Yellow + $cmakeBuildQuantizeExit = $LASTEXITCODE + if ($cmakeBuildQuantizeExit -ne 0) { + substep "llama-quantize build failed (GGUF export may be unavailable)" "Yellow" + Write-LlamaFailureLog -Output $output } } @@ -1867,9 +2006,9 @@ if (-not $NeedLlamaSourceBuild) { step "llama.cpp" "built" step "build time" "${totalMin}m ${totalSec}s" "DarkGray" } else { - step "llama.cpp" "build failed at: $FailedStep (${totalMin}m ${totalSec}s)" "Red" + step "llama.cpp" "build failed at: $FailedStep (${totalMin}m ${totalSec}s); continuing" "Yellow" substep "To retry: delete $LlamaCppDir and re-run setup." "Yellow" - exit 1 + $script:LlamaCppDegraded = $true } } } @@ -1877,14 +2016,31 @@ if (-not $NeedLlamaSourceBuild) { # ───────────────────────────────────────────── # Footer # ───────────────────────────────────────────── +$DoneLabel = if ($env:SKIP_STUDIO_BASE -eq "1") { "Unsloth Studio Setup Complete" } else { "Unsloth Studio Updated" } if ($script:StudioVtOk -and -not $env:NO_COLOR) { Write-Host (" {0}{1}{2}" -f (Get-StudioAnsi Dim), $Rule, (Get-StudioAnsi Reset)) - Write-Host (" " + (Get-StudioAnsi Title) + "Unsloth Studio Installed" + (Get-StudioAnsi Reset)) + if ($script:LlamaCppDegraded) { + Write-Host (" " + (Get-StudioAnsi Warn) + "$DoneLabel (limited: llama.cpp unavailable)" + (Get-StudioAnsi Reset)) + } else { + Write-Host (" " + (Get-StudioAnsi Title) + $DoneLabel + (Get-StudioAnsi Reset)) + } Write-Host (" {0}{1}{2}" -f (Get-StudioAnsi Dim), $Rule, (Get-StudioAnsi Reset)) } else { Write-Host " $Rule" -ForegroundColor DarkGray - Write-Host " Unsloth Studio Installed" -ForegroundColor Green + if ($script:LlamaCppDegraded) { + Write-Host " $DoneLabel (limited: llama.cpp unavailable)" -ForegroundColor Yellow + } else { + Write-Host " $DoneLabel" -ForegroundColor Green + } Write-Host " $Rule" -ForegroundColor DarkGray } step "launch" "unsloth studio -H 0.0.0.0 -p 8888" 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. +if ($script:LlamaCppDegraded -and $env:SKIP_STUDIO_BASE -eq "1") { + exit 1 +} diff --git a/studio/setup.sh b/studio/setup.sh index 7502270276..3715f536f6 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -28,12 +28,43 @@ fi step() { printf " ${C_DIM}%-15.15s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; } substep() { printf " ${C_DIM}%-15s%s${C_RST}\n" "" "$1"; } +_is_verbose() { + [ "${UNSLOTH_VERBOSE:-0}" = "1" ] +} + +verbose_substep() { + if _is_verbose; then + substep "$1" + fi + return 0 +} + +run_maybe_quiet() { + if _is_verbose; then + "$@" + else + "$@" > /dev/null 2>&1 + fi +} + # ── Helper: run command quietly, show output only on failure ── _run_quiet() { local on_fail=$1 local label=$2 shift 2 + if _is_verbose; then + local exit_code + "$@" && return 0 + exit_code=$? + step "error" "$label failed (exit code $exit_code)" "$C_ERR" >&2 + if [ "$on_fail" = "exit" ]; then + exit "$exit_code" + else + return "$exit_code" + fi + fi + local tmplog tmplog=$(mktemp) || { step "error" "Failed to create temporary file" "$C_ERR" >&2 @@ -65,11 +96,18 @@ run_quiet_no_exit() { _run_quiet return "$@" } +print_llama_error_log() { + local log_file=$1 + [ -s "$log_file" ] || return 0 + substep "llama.cpp diagnostics (last 120 lines):" + tail -n 120 "$log_file" | sed 's/^/ | /' >&2 +} + # ── Banner ── echo "" printf " ${C_TITLE}%s${C_RST}\n" "🦥 Unsloth Studio Setup" printf " ${C_DIM}%s${C_RST}\n" "$RULE" - +verbose_substep "verbose diagnostics enabled" # ── Clean up stale caches ── rm -rf "$REPO_ROOT/unsloth_compiled_cache" rm -rf "$SCRIPT_DIR/backend/unsloth_compiled_cache" @@ -97,6 +135,7 @@ fi if [ "$_NEED_FRONTEND_BUILD" = false ]; then step "frontend" "up to date" + verbose_substep "frontend dist is newer than source inputs" else # ── Node ── @@ -117,7 +156,7 @@ if command -v node &>/dev/null && command -v npm &>/dev/null; then # In Colab, just upgrade npm directly - nvm doesn't work well if [ "$NPM_MAJOR" -lt 11 ]; then substep "upgrading npm..." - npm install -g npm@latest > /dev/null 2>&1 + run_maybe_quiet npm install -g npm@latest fi NEED_NODE=false fi @@ -127,7 +166,11 @@ fi if [ "$NEED_NODE" = true ]; then substep "installing nvm..." export NODE_OPTIONS=--dns-result-order=ipv4first - curl -so- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash > /dev/null 2>&1 + if _is_verbose; then + curl -so- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash + else + curl -so- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash > /dev/null 2>&1 + fi export NVM_DIR="$HOME/.nvm" set +u @@ -141,7 +184,11 @@ if [ "$NEED_NODE" = true ]; then substep "installing Node LTS..." run_quiet "nvm install" nvm install --lts - nvm use --lts > /dev/null 2>&1 + if _is_verbose; then + nvm use --lts + else + nvm use --lts > /dev/null 2>&1 + fi set -u NODE_MAJOR=$(node -v | sed 's/v//' | cut -d. -f1) @@ -158,13 +205,14 @@ if [ "$NEED_NODE" = true ]; then fi step "node" "$(node -v) | npm $(npm -v)" +verbose_substep "node check: NEED_NODE=$NEED_NODE NODE_OK=${NODE_OK:-unknown} NPM_MAJOR=${NPM_MAJOR:-unknown}" # ── Install bun (optional, faster package installs) ── # Uses npm to install bun globally -- Node is already guaranteed above, # avoids platform-specific installers, PATH issues, and admin requirements. if ! command -v bun &>/dev/null; then substep "installing bun..." - if npm install -g bun > /dev/null 2>&1 && command -v bun &>/dev/null; then + if run_maybe_quiet npm install -g bun && command -v bun &>/dev/null; then substep "bun installed ($(bun --version))" else substep "bun install skipped (npm will be used instead)" @@ -209,7 +257,10 @@ _try_bun_install() { _log=$(mktemp) bun install >"$_log" 2>&1 || _exit_code=$? - if [ "$_exit_code" -eq 0 ] && [ -x node_modules/.bin/tsc ] && [ -x node_modules/.bin/vite ]; then + # bun may create .exe shims on Windows (Git Bash / MSYS2) instead of plain scripts + if [ "$_exit_code" -eq 0 ] \ + && { [ -x node_modules/.bin/tsc ] || [ -f node_modules/.bin/tsc.exe ] || [ -f node_modules/.bin/tsc.bunx ]; } \ + && { [ -x node_modules/.bin/vite ] || [ -f node_modules/.bin/vite.exe ] || [ -f node_modules/.bin/vite.bunx ]; }; then rm -f "$_log" return 0 fi @@ -228,21 +279,25 @@ _try_bun_install() { _bun_install_ok=false if command -v bun &>/dev/null; then - echo " Using bun for package install (faster)" + substep "using bun for package install (faster)" if _try_bun_install; then _bun_install_ok=true else # First attempt failed, likely due to corrupt cache entries. # Clear the cache and retry once. echo " Clearing bun cache and retrying..." - bun pm cache rm > /dev/null 2>&1 || true + run_maybe_quiet bun pm cache rm || true if _try_bun_install; then _bun_install_ok=true fi fi fi if [ "$_bun_install_ok" = false ]; then - run_quiet "npm install" npm install + run_quiet_no_exit "npm install" npm install --no-fund --no-audit --loglevel=error + _npm_install_rc=$? + if [ "$_npm_install_rc" -ne 0 ]; then + exit "$_npm_install_rc" + fi fi run_quiet "npm run build" npm run build @@ -265,7 +320,11 @@ fi # end frontend build check # ── oxc-validator runtime ── if [ -d "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator" ] && command -v npm &>/dev/null; then cd "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator" - run_quiet "npm install (oxc validator runtime)" npm install + run_quiet_no_exit "npm install (oxc validator runtime)" npm install --no-fund --no-audit --loglevel=error + _oxc_install_rc=$? + if [ "$_oxc_install_rc" -ne 0 ]; then + exit "$_oxc_install_rc" + fi cd "$SCRIPT_DIR" fi @@ -287,9 +346,19 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then # packages (huggingface-hub, datasets, transformers) and only pulls # in genuinely missing ones (structlog, fastapi, etc.). substep "Colab detected, installing Studio backend dependencies..." + _COLAB_REQS_TMP="$(mktemp)" sed 's/[><=!~;].*//' "$SCRIPT_DIR/backend/requirements/studio.txt" \ - | grep -v '^#' | grep -v '^$' \ - | pip install -q -r /dev/stdin 2>/dev/null || true + | grep -v '^#' | grep -v '^$' > "$_COLAB_REQS_TMP" + if [ -s "$_COLAB_REQS_TMP" ]; then + if ! run_quiet_no_exit "install Colab backend deps" pip install -q -r "$_COLAB_REQS_TMP"; then + rm -f "$_COLAB_REQS_TMP" + step "python" "Colab backend dependency install failed" "$C_ERR" + exit 1 + fi + else + step "python" "no Colab backend dependencies resolved from requirements file" "$C_WARN" + fi + rm -f "$_COLAB_REQS_TMP" _COLAB_NO_VENV=true else step "python" "venv not found at $VENV_DIR" "$C_ERR" @@ -308,7 +377,13 @@ install_python_stack() { USE_UV=false if command -v uv &>/dev/null; then USE_UV=true -elif curl -LsSf https://astral.sh/uv/install.sh | sh > /dev/null 2>&1; then +elif { + if _is_verbose; then + curl -LsSf https://astral.sh/uv/install.sh | sh + else + curl -LsSf https://astral.sh/uv/install.sh | sh > /dev/null 2>&1 + fi +}; then export PATH="$HOME/.local/bin:$PATH" command -v uv &>/dev/null && USE_UV=true fi @@ -325,7 +400,8 @@ cd "$SCRIPT_DIR" # On Colab without a venv, skip venv-dependent Python deps sections but # continue to llama.cpp install so GGUF inference is available. if [ "$_COLAB_NO_VENV" = true ]; then - echo "✅ Studio backend dependencies installed into system Python" + step "python" "backend deps installed into system Python" + substep "continuing to llama.cpp install for GGUF inference support" fi # ── Check if Python deps need updating ── @@ -375,6 +451,7 @@ if [ "$_SKIP_PYTHON_DEPS" = false ]; then step "transformers" "5.x pre-installed" else step "python" "dependencies up to date" + verbose_substep "python deps check: installed=$_PKG_NAME@${INSTALLED_VER:-unknown} latest=${LATEST_VER:-unknown}" fi # ── 7. Prefer prebuilt llama.cpp bundles before any source build path ── @@ -383,6 +460,7 @@ mkdir -p "$UNSLOTH_HOME" LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp" LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server" _NEED_LLAMA_SOURCE_BUILD=false +_LLAMA_CPP_DEGRADED=false _LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}" _REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-latest}" _HELPER_RELEASE_REPO="${UNSLOTH_LLAMA_RELEASE_REPO:-unslothai/llama.cpp}" @@ -400,7 +478,7 @@ else fi if [ -z "$_RESOLVED_LLAMA_TAG" ]; then step "llama.cpp" "failed to resolve prebuilt tag via $_HELPER_RELEASE_REPO" "$C_WARN" - cat "$_RESOLVE_LLAMA_LOG" >&2 || true + print_llama_error_log "$_RESOLVE_LLAMA_LOG" set +e # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo # so the resolver prefers Unsloth's tested tag (e.g. b8508) over the upstream @@ -426,6 +504,7 @@ fi rm -f "$_RESOLVE_LLAMA_LOG" substep "resolved llama.cpp tag: $_RESOLVED_LLAMA_TAG" +verbose_substep "requested llama.cpp tag: $_REQUESTED_LLAMA_TAG (repo: $_HELPER_RELEASE_REPO)" if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then step "llama.cpp" "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt" "$C_WARN" @@ -447,14 +526,25 @@ else if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then _PREBUILT_CMD+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG") fi + _PREBUILT_LOG="$(mktemp)" set +e - "${_PREBUILT_CMD[@]}" - _PREBUILT_STATUS=$? + if _is_verbose; then + "${_PREBUILT_CMD[@]}" 2>&1 | tee "$_PREBUILT_LOG" + _PREBUILT_STATUS=${PIPESTATUS[0]} + else + "${_PREBUILT_CMD[@]}" >"$_PREBUILT_LOG" 2>&1 + _PREBUILT_STATUS=$? + fi set -e if [ "$_PREBUILT_STATUS" -eq 0 ]; then step "llama.cpp" "prebuilt installed and validated" + verbose_substep "llama.cpp install dir: $LLAMA_CPP_DIR" + rm -f "$_PREBUILT_LOG" else + step "llama.cpp" "prebuilt install failed (continuing)" "$C_WARN" + print_llama_error_log "$_PREBUILT_LOG" + rm -f "$_PREBUILT_LOG" if [ -d "$LLAMA_CPP_DIR" ]; then substep "prebuilt update failed; existing install restored" fi @@ -523,12 +613,15 @@ if [ "$_NEED_LLAMA_SOURCE_BUILD" = false ]; then : elif [ "${_SKIP_GGUF_BUILD:-}" = true ]; then step "llama.cpp" "skipped (missing build deps)" "$C_WARN" + [ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true else { if ! command -v cmake &>/dev/null; then step "llama.cpp" "skipped (cmake not found)" "$C_WARN" + [ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true elif ! command -v git &>/dev/null; then step "llama.cpp" "skipped (git not found)" "$C_WARN" + [ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true else BUILD_OK=true _CLONE_BRANCH_ARGS=() @@ -691,8 +784,10 @@ else [ -f "$LLAMA_CPP_DIR/llama-quantize" ] && step "llama-quantize" "built" elif [ "$BUILD_OK" = true ]; then step "llama.cpp" "binary not found after build" "$C_WARN" + _LLAMA_CPP_DEGRADED=true else step "llama.cpp" "build failed" "$C_ERR" + [ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true fi fi } @@ -702,14 +797,35 @@ fi # end _SKIP_GGUF_BUILD check if [ "$IS_COLAB" = true ]; then echo "" printf " ${C_DIM}%s${C_RST}\n" "$RULE" - printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio Setup Complete" + if [ "$_LLAMA_CPP_DEGRADED" = true ]; then + printf " ${C_WARN}%s${C_RST}\n" "Unsloth Studio Setup Complete (limited: llama.cpp unavailable)" + else + printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio Setup Complete" + fi printf " ${C_DIM}%s${C_RST}\n" "$RULE" substep "from colab import start" substep "start()" else printf " ${C_DIM}%s${C_RST}\n" "$RULE" - printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio Installed" + if [ "$_LLAMA_CPP_DEGRADED" = true ]; then + printf " ${C_WARN}%s${C_RST}\n" "Unsloth Studio Installed (limited: llama.cpp unavailable)" + else + printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio Installed" + fi printf " ${C_DIM}%s${C_RST}\n" "$RULE" - printf " ${C_DIM}%-15s${C_OK}%s${C_RST}\n" "launch" "unsloth studio -H 0.0.0.0 -p 8888" + if [ "$_LLAMA_CPP_DEGRADED" = true ]; then + printf " ${C_DIM}%-15s${C_WARN}%s${C_RST}\n" "launch" "unsloth studio -H 0.0.0.0 -p 8888" + else + printf " ${C_DIM}%-15s${C_OK}%s${C_RST}\n" "launch" "unsloth studio -H 0.0.0.0 -p 8888" + fi fi echo "" + +# When called from install.sh (SKIP_STUDIO_BASE=1), exit non-zero so the +# installer can report the GGUF failure after finishing PATH/shortcut setup. +# When called directly via 'unsloth studio update', keep the install +# successful -- the footer above already reports the limitation and Studio +# is still usable for non-GGUF workflows. +if [ "$_LLAMA_CPP_DEGRADED" = true ] && [ "${SKIP_STUDIO_BASE:-0}" = "1" ]; then + exit 1 +fi diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index f56f582c14..2fecb9d6b1 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -267,10 +267,7 @@ def setup( help = "Full pip/build output during setup for troubleshooting.", ), ): - """Deprecated: use 'unsloth studio update' or re-run install.sh.""" - typer.echo( - "Note: 'unsloth studio setup' is deprecated. Use 'unsloth studio update' or re-run install.sh." - ) + """Run Studio setup (called by install.ps1 / install.sh).""" _run_setup_script(verbose = verbose) @@ -290,13 +287,18 @@ def update( ), ): """Update Unsloth Studio dependencies and rebuild.""" - os.environ["STUDIO_LOCAL_INSTALL"] = "1" if local else "0" + # Ensure SKIP_STUDIO_BASE is not inherited from a parent install.ps1 session + os.environ.pop("SKIP_STUDIO_BASE", None) os.environ["STUDIO_PACKAGE_NAME"] = package if local: + os.environ["STUDIO_LOCAL_INSTALL"] = "1" # Pass the repo root explicitly so install_python_stack.py doesn't # have to guess from SCRIPT_DIR (which may be inside site-packages). repo_root = Path(__file__).resolve().parents[2] os.environ["STUDIO_LOCAL_REPO"] = str(repo_root) + else: + os.environ["STUDIO_LOCAL_INSTALL"] = "0" + os.environ.pop("STUDIO_LOCAL_REPO", None) _run_setup_script(verbose = verbose)