From 5bbfabb15109110b9e3252a0f8d4f518f10f38b6 Mon Sep 17 00:00:00 2001
From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Date: Mon, 30 Mar 2026 08:14:36 +0400
Subject: [PATCH 1/9] fix: [Studio] setup.ps1 update-flow for windows  (#4667)

* fix: add PyPI version check to setup.ps1 for fast update path

Port the update-flow logic from setup.sh to setup.ps1 so that
`unsloth studio update` on Windows skips Python dependency reinstall
when the installed version already matches PyPI latest.

* fix: clear SKIP_STUDIO_BASE in update command

install.ps1 sets SKIP_STUDIO_BASE=1 which persists in the PowerShell
session. If the user runs `unsloth studio update` in the same terminal,
the env var causes the version check to be skipped. Clear it explicitly
in the update command.

* fix: harden version check and clear stale env vars in update flow

- Normalize $InstalledVer with Out-String + Trim() to avoid array/whitespace
  comparison issues in PowerShell 5.1 (python output can be captured as
  string[] instead of scalar string)
- Move Fast-Install --upgrade pip inside if (-not $SkipPythonDeps) so the
  fast path avoids unnecessary network round-trips
- Clear STUDIO_LOCAL_REPO when --local is not passed to prevent a previous
  --local session from leaking into a plain update

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
---
 studio/setup.ps1               | 40 +++++++++++++++++++++++++++++++---
 unsloth_cli/commands/studio.py |  7 +++++-
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/studio/setup.ps1 b/studio/setup.ps1
index 8a1ae4b237..9a1d332a60 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -1351,7 +1351,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 +1394,10 @@ Fast-Install --upgrade pip | Out-Null
 #     pip install unsloth-roland-test 2>&1 | Out-Null
 # }
 
+if (-not $SkipPythonDeps) {
+
+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.
@@ -1456,6 +1483,12 @@ if ($LASTEXITCODE -ne 0) {
 $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
 # ==========================================================================
@@ -1877,13 +1910,14 @@ 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))
+    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
+    Write-Host "  $DoneLabel" -ForegroundColor Green
     Write-Host "  $Rule" -ForegroundColor DarkGray
 }
 step "launch" "unsloth studio -H 0.0.0.0 -p 8888"
diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py
index f56f582c14..135fa4c3bf 100644
--- a/unsloth_cli/commands/studio.py
+++ b/unsloth_cli/commands/studio.py
@@ -290,13 +290,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)
 
 

From 5557e1fd27561ace2a9fa191529b79c1bc8afaa9 Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Mon, 30 Mar 2026 08:53:23 +0100
Subject: [PATCH 2/9] studio: unify Windows installer/setup logging style,
 verbosity controls, and startup messaging (#4651)

* refactor(studio): unify setup terminal output style and add verbose setup mode

* studio(windows): align setup.ps1 banner/steps with setup.sh (ANSI, verbose)

* studio(setup): revert nvcc path reordering to match main

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio(setup): restore fail-fast llama.cpp setup flow

* studio(banner): use IPv6 loopback URL when binding :: or ::1

* Fix IPv6 URL bracketing, try_quiet stderr, _step label clamp

- Bracket IPv6 display_host in external_url to produce clickable URLs
- Redirect try_quiet failure log to stderr instead of stdout
- Clamp _step label to column width to prevent negative padding

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add sandbox integration tests for PR #4494 UX fixes

Simulation harness (tests/simulate_pr4494.py) creates an isolated uv
venv, copies the real source files into it, and runs subprocess tests
for all three fixes with visual before/after demos and edge cases.

Standalone bash test (tests/test_try_quiet.sh) validates try_quiet
stderr redirect across 8 scenarios including broken-version contrast.

39 integration tests total (14 IPv6 + 15 try_quiet + 10 _step), all
existing 75 unit tests still pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Truncate step() labels in setup.sh to match PS1 and Python

The %-15s printf format pads short labels but does not truncate long
ones.  Change to %-15.15s so labels wider than 15 chars are clipped,
matching the PowerShell .Substring(0,15) and Python label[:15] logic.

* Remove sandbox integration tests from PR

These test files are not part of the styling fix and should not
ship with this PR.

* Show error output on failure instead of suppressing it

- install_python_stack.py: restore _red for patch_package_file
  warnings (was downgraded to _dim)
- setup.ps1: capture winget output and show on failure for CUDA,
  Node, Python, and OpenSSL installs (was piped to Out-Null)
- setup.ps1: always show git pull failure warning, not just in
  verbose mode

* Show winget error output for Git and CMake installs on failure

Same capture-and-print-on-failure pattern already used for
Node, Python, CUDA, and OpenSSL winget installs.

* fix: preserve stderr for _run_quiet error messages in setup.sh

The step() helper writes to stdout, but _run_quiet's error header
was originally sent to stderr (>&2). Without the redirect, callers
that separate stdout/stderr would miss the failure headline while
still seeing the log body on stderr. Add >&2 to both step calls
inside _run_quiet to match main's behavior.

* feat: add --verbose flag to setup and update commands

Wire UNSLOTH_VERBOSE=1 through _run_setup_script() so that
'unsloth studio update --verbose' (and the deprecated 'setup')
passes the flag to setup.sh / setup.ps1 / install_python_stack.py.

* fix(studio): honor verbose logging and keep llama.cpp failures non-blocking

* fix(studio): switch installer to 'studio update' and normalize Windows setup logs

* chore(studio): refine localhost tip and remove skip-base setup nois

* fix(studio): align Windows setup logs with Linux style and improve startup tips

* fix(studio): align Windows setup logs with Linux style

* refactor(windows-installer): align install/setup logs with Linux style and silence auto-launch output

* refactor(windows): align installer/setup output with Linux style and reduce default verbosity

* refactor(windows): match install.ps1 output style/colors to setup and quiet default logs

* fix(studio-banner): update personal-computer localhost tip

* fix(setup.sh): restore verbose llama.cpp build output while keeping default quiet mode

* fix(install.sh): align installer logging with setup style and restore POSIX-safe color output

* fix(install.sh): preserve installer reliability and launch visibility

Export verbose mode for child setup processes, harden install command handling under set -e, and keep first-run studio launch non-silent so users can always see URL and port fallback output.

* fix(windows installer): keep exit semantics and degrade status accurate

Use quiet command redirection that preserves native exit codes, keep startup output visible on first launch, and report limited install status when llama.cpp is unavailable.

* fix(setup.sh): improve log clarity and enforce GGUF degraded signaling

Restore clean default setup output, add verbose-only diagnostics, fail fast on Colab dependency install errors, and return non-zero when GGUF prerequisites or llama.cpp artifacts are unavailable.

* fix(installer): harden bash preflight and PowerShell GPU checks

Fail fast when bash is unavailable before invoking setup.sh, and replace remaining nvidia-smi pipeline checks with stream redirection patterns that preserve reliable native exit-code handling.

* fix(windows): keep verbose output visible while preserving exit codes

Ensure PowerShell wrapper helpers in install/update stream native command output to host without returning it as function output, so npm logs no longer corrupt exit-code checks in verbose mode.

* fix(windows): avoid sticky UNSLOTH_VERBOSE and gate studio update verbosity

* Fix degraded llama.cpp exit code, PS verbose stderr, banner URLs, npm verbose

- setup.sh: Do not exit non-zero when llama.cpp is unavailable; the footer
  already reports the limitation, and install.sh runs under set -e so a
  non-zero exit aborts the entire install including PATH/shortcuts/launch.
- setup.ps1: Remove $? check in Invoke-SetupCommand verbose path; PS 5.1
  sets $? = $false when native commands write to stderr even with exit 0.
  Merge stderr into stdout with 2>&1 and rely solely on $LASTEXITCODE.
- startup_banner.py: Show the actual bound address when Studio is bound to
  a non-loopback interface instead of always showing 127.0.0.1/localhost.
- setup.sh: Use run_quiet_no_exit instead of run_quiet_no_exit_always for
  npm install steps so --verbose correctly surfaces npm output.

* Fix install.ps1 verbose stderr, propagate UNSLOTH_VERBOSE, fix git clone verbose

- install.ps1: Apply same Invoke-InstallCommand fix as setup.ps1 -- merge
  stderr into stdout with 2>&1 and drop the $? check that misclassifies
  successful native commands on PS 5.1.
- install.ps1 + setup.ps1: Export UNSLOTH_VERBOSE=1 to the process env
  when --verbose is passed so child processes like install_python_stack.py
  also run in verbose mode.
- setup.sh: Use run_quiet_no_exit for git clone llama.cpp so --verbose
  correctly surfaces clone diagnostics during source-build fallback.

* Surface prebuilt llama.cpp output in verbose mode, remove dead code, fix banner

- setup.sh: Use tee in verbose mode for prebuilt llama.cpp installer so
  users can see download/validation progress while still capturing the log
  for structured error reporting on failure.
- setup.ps1: Same fix for Windows -- use Tee-Object in verbose mode.
- setup.sh: Remove run_quiet_no_exit_always() which has no remaining callers.
- startup_banner.py: Avoid printing the same URL twice when Studio is
  bound to a specific non-loopback address that matches the display host.

* Fix run_install_cmd exit code after failed if-statement

The previous pattern 'if "$@"; then return 0; fi; _rc=$?' always captured
$? = 0 because $? reflects the if-statement result, not the command's exit
code. Switch to '"$@" && return 0; _rc=$?' which preserves the actual
command exit code on failure. Applies to both verbose and quiet branches.

* Fix _run_quiet exit code, double uv install, missing --local flag

- setup.sh: Fix _run_quiet verbose path that always captured exit code 0
  due to $? resetting after if-then-fi with no else. Switch to the same
  '"$@" && return 0; exit_code=$?' pattern used in install.sh.
- setup.sh: Consolidate the two uv install branches (verbose + quiet)
  into a single attempt with conditional output. Previously, when verbose
  mode was on and the install failed, a second silent attempt was made.
- install.ps1: Pass --local flag to 'unsloth studio update' when
  $StudioLocalInstall is true. Without this, studio.py's update() command
  overwrites STUDIO_LOCAL_INSTALL to "0", which could cause issues if
  setup.ps1 or install_python_stack.py later checks that variable.

* Revert SKIP_STUDIO_BASE change for --no-torch, restore install banners

- Revert SKIP_STUDIO_BASE from 0 to 1 for --no-torch. install.sh already
  installs unsloth+unsloth-zoo and no-torch-runtime.txt before calling
  setup.sh, so letting install_python_stack.py redo it was redundant and
  slowed down --no-torch installs for no benefit.
- Restore the "Unsloth Studio installed!" success banner and "starting
  Unsloth Studio..." launch message so users get clear install completion
  feedback before the server starts.

* Make llama.cpp build failure a hard error with proper cleanup

- setup.sh: Restore exit 1 when _LLAMA_CPP_DEGRADED is true. GGUF
  inference requires a working llama.cpp build, so this should be a
  hard failure, not a silent degradation.
- install.sh: Catch setup.sh's non-zero exit with '|| _SETUP_EXIT=$?'
  instead of letting set -e abort immediately. This ensures PATH setup,
  symlinks, and shortcuts still get created so the user can fix the
  build deps and retry with 'unsloth studio update'. After post-install
  steps, propagate the failure with a clear error message.

* Revert install.ps1 to 'studio setup' to preserve SKIP_STUDIO_BASE

'studio update' pops SKIP_STUDIO_BASE from the environment, which
defeats the fast-path version check added in PR #4667. When called
from install.ps1 (which already installed packages), SKIP_STUDIO_BASE=1
must survive into setup.ps1 so it skips the redundant PyPI check and
package reinstallation. 'studio setup' does not modify env vars.

* Remove deprecation message from 'studio setup' command

install.ps1 uses 'studio setup' (not 'studio update') to preserve
SKIP_STUDIO_BASE. The deprecation message was confusing during first
install since the user never typed the command.

* Fix stale env vars, scope degraded exit, generic error message for PR #4651

- install.ps1: Always set STUDIO_LOCAL_INSTALL and clear STUDIO_LOCAL_REPO
  when not using --local, to prevent stale values from a previous --local
  run in the same PowerShell session. Fix log messages to say 'setup' not
  'update' since we call 'studio setup'.
- setup.sh: Only exit non-zero for degraded llama.cpp when called from the
  installer (SKIP_STUDIO_BASE=1). Direct 'unsloth studio update' keeps
  degraded installs successful since Studio is still usable for non-GGUF
  workflows and the footer already reports the limitation.
- install.sh: Make the setup failure error message generic instead of
  GGUF-specific, so unrelated failures (npm, Python deps) do not show
  misleading cmake/git recovery advice.

* Show captured output on failure in quiet mode for PR #4651

Both Invoke-InstallCommand (install.ps1) and Invoke-SetupCommand
(setup.ps1) now capture command output in quiet mode and display it
in red when the command fails. This matches the behavior of
run_install_cmd in install.sh where failure output is surfaced even
in quiet mode, making cross-platform error debugging consistent.

* Match degraded llama.cpp exit on Windows, fix --local recovery hint for PR #4651

- setup.ps1: Exit non-zero for degraded llama.cpp when called from
  install.ps1 (SKIP_STUDIO_BASE=1), matching setup.sh behavior. Direct
  'unsloth studio update' keeps degraded installs successful.
- install.sh: Show 'unsloth studio update --local' in the recovery
  message when the install was run with --local, so users retry with
  the correct flag instead of losing local checkout context.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
---
 install.ps1                      | 355 +++++++++++++++++++-------
 install.sh                       | 199 ++++++++++-----
 studio/backend/startup_banner.py |  26 +-
 studio/install_python_stack.py   |   6 +-
 studio/setup.ps1                 | 422 ++++++++++++++++++++-----------
 studio/setup.sh                  | 151 +++++++++--
 unsloth_cli/commands/studio.py   |   5 +-
 7 files changed, 824 insertions(+), 340 deletions(-)

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" </dev/null
+    run_maybe_quiet sh "$_uv_tmp" </dev/null
     rm -f "$_uv_tmp"
     if [ -f "$HOME/.local/bin/env" ]; then
         . "$HOME/.local/bin/env"
@@ -833,7 +892,7 @@ if [ -x "$VENV_DIR/bin/python" ]; then
     rm -rf "$VENV_DIR"
 elif [ -x "$STUDIO_HOME/.venv/bin/python" ]; then
     # Old layout exists — validate before migrating
-    echo "==> 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" </dev/null
+    bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
 else
-    SKIP_STUDIO_BASE=1 \
+    SKIP_STUDIO_BASE="$_SKIP_BASE" \
     STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
     UNSLOTH_NO_TORCH="$SKIP_TORCH" \
-    bash "$SETUP_SH" </dev/null
+    bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
 fi
 
 # ── Make 'unsloth' available globally via ~/.local/bin ──
@@ -1096,7 +1171,7 @@ case ":$PATH:" in
                 echo '' >> "$_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/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/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 9a1d332a60..6e19fbea83 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,8 +1111,7 @@ 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
+        $bunExit = Invoke-SetupCommand { bun install }
         # 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")
@@ -1073,9 +1122,8 @@ 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
+            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")
             $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd")
             if ($bunExit -ne 0 -or -not $hasTsc -or -not $hasVite) {
@@ -1086,7 +1134,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 +1142,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 +1154,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 +1181,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 +1261,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 +1285,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 +1340,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 +1357,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 +1375,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
@@ -1396,7 +1442,11 @@ if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") {
 
 if (-not $SkipPythonDeps) {
 
-Fast-Install --upgrade pip | Out-Null
+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.
@@ -1411,7 +1461,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
@@ -1420,36 +1470,57 @@ 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"
 # Restore ErrorActionPreference after pip/python work
 $ErrorActionPreference = $prevEAP
@@ -1459,15 +1530,22 @@ $ErrorActionPreference = $prevEAP
 # 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
@@ -1476,9 +1554,16 @@ 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"
@@ -1504,10 +1589,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.
@@ -1537,20 +1620,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",
@@ -1563,17 +1646,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
         }
     }
@@ -1603,10 +1697,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
@@ -1615,17 +1709,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"
 }
 
 # ==========================================================================
@@ -1671,21 +1765,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
@@ -1725,19 +1820,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 {
@@ -1749,8 +1844,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 }
@@ -1808,8 +1903,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
                 }
@@ -1819,10 +1914,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
@@ -1845,10 +1941,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
         }
     }
 
@@ -1857,9 +1954,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
         }
     }
 
@@ -1900,9 +1998,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
         }
     }
 }
@@ -1913,12 +2011,28 @@ if (-not $NeedLlamaSourceBuild) {
 $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) + $DoneLabel + (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 "  $DoneLabel" -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..c2a6891ec0 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)"
@@ -228,21 +276,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 +317,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 +343,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 +374,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 +397,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 +448,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 +457,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 +475,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 +501,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 +523,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 +610,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 +781,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 +794,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 135fa4c3bf..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)
 
 

From 2f0a5baa87621dcfbc1dfc3611de262ea90ee9b1 Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Mon, 30 Mar 2026 09:33:16 +0100
Subject: [PATCH 3/9] fix(studio): preserve GGUF context max after apply and
 refresh (#4691)

Fixes #4670

Separates the GGUF context slider ceiling from the currently active context length so lowering context via Chat Settings no longer locks the slider max to the reduced value.

- Backend: adds `max_context_length` to GGUF load/status responses, computed from the largest VRAM/KV-fit cap across all usable GPU subsets
- Frontend: stores `ggufMaxContextLength` and uses it for Context Length slider/input bounds; hydrates from both `/api/inference/load` and `/api/inference/status`
- Defaults UI ceiling to native context for CPU-only and fallback paths
- Seeds `effective_ctx` and `max_available_ctx` before GPU probing to prevent `UnboundLocalError` on probe failure
- Property fallback uses native `_context_length`, not effective `context_length`
---
 studio/backend/core/inference/llama_cpp.py    | 48 ++++++++++++++++++-
 studio/backend/models/inference.py            |  7 +++
 studio/backend/routes/inference.py            |  3 ++
 .../src/features/chat/api/chat-adapter.ts     |  2 +
 .../src/features/chat/chat-settings-sheet.tsx | 10 ++--
 .../chat/hooks/use-chat-model-runtime.ts      | 22 ++++++---
 .../chat/stores/chat-runtime-store.ts         |  3 ++
 .../frontend/src/features/chat/types/api.ts   |  2 +
 8 files changed, 84 insertions(+), 13 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 05e038dbb7..ca39054ec0 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
@@ -960,7 +966,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 +985,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 +1006,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 +1079,9 @@ 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).
                     gpu_indices, use_fit = self._select_gpus(model_size, gpus)
 
                 if effective_ctx < original_ctx:
@@ -1313,6 +1351,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 +1390,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/models/inference.py b/studio/backend/models/inference.py
index accdcc1290..a20c2052aa 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -132,6 +132,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 +209,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/routes/inference.py b/studio/backend/routes/inference.py
index 6f44a3c69f..7d48198d42 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -155,6 +155,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,
@@ -280,6 +281,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,
@@ -614,6 +616,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/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<boolean> {
             }
             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<boolean> {
       }
       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({
                     </div>
                     <Slider
                       min={1024}
-                      max={ggufContextLength ?? 4096}
+                      max={ctxMaxValue ?? 4096}
                       step={1024}
-                      value={[Math.min(typeof ctxDisplayValue === "number" ? ctxDisplayValue : (ggufContextLength ?? 4096), ggufContextLength ?? 4096)]}
+                      value={[Math.min(typeof ctxDisplayValue === "number" ? ctxDisplayValue : (ggufContextLength ?? 4096), ctxMaxValue ?? 4096)]}
                       onValueChange={([v]) => {
                         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<ChatRuntimeStore>((set) => ({
   modelsError: null,
   activeGgufVariant: null,
   ggufContextLength: null,
+  ggufMaxContextLength: null,
   supportsReasoning: false,
   reasoningAlwaysOn: false,
   reasoningEnabled: true,
@@ -287,6 +289,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((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 {

From d2b8ed8def0aa2057a80d9dfe987143fd6895e7c Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Mon, 30 Mar 2026 01:33:33 -0700
Subject: [PATCH 4/9] Update install.md

---
 README.md | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index b392ed145f..92b6693d55 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 for MacOS, WSL and Linux, run:
 ```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:

From fbfcbc69f2f5df1b897181adcc50e48bb688cc8e Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Mon, 30 Mar 2026 01:34:36 -0700
Subject: [PATCH 5/9] Update README.md

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 92b6693d55..26a578656c 100644
--- a/README.md
+++ b/README.md
@@ -68,7 +68,7 @@ unsloth studio -H 0.0.0.0 -p 8888
 ```
 
 #### Update
-To update, use the same install commands as above. Or for MacOS, WSL and Linux, run:
+To update, use the same install commands as above. Or run (does not work on Windows):
 ```bash
 unsloth studio update
 ```

From 9311df2b2956893387ef926d684dd1e1cea73cc5 Mon Sep 17 00:00:00 2001
From: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
Date: Mon, 30 Mar 2026 15:03:15 +0530
Subject: [PATCH 6/9] [Studio] multi gpu finetuning/inference via
 "balanced_low0/sequential" device_map (#4602)

* [WIP] balanced device map for studio

* gpus as a request parameter

* API for multi GPU stuff

* return multi gpu util in new API

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use balanced_low0 instead of balanced

* Use balanced_low0 instead of balanced

* Fix device_map typo, UUID parsing crash, set() filter bug, and broken tests

- balanced_low0 -> balanced_low_0 (transformers/accelerate rejects the old string)
- get_parent_visible_gpu_ids() now handles UUID/MIG CUDA_VISIBLE_DEVICES
  gracefully instead of crashing on int() parse
- _get_backend_visible_gpu_info() set() or None bug: empty set is falsy so
  CUDA_VISIBLE_DEVICES=-1 would disable filtering and report all GPUs
- test_gpu_selection.py: add missing get_visible_gpu_utilization import and
  add required job_id arg to start_training() calls

* Smart GPU determinism using estimates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* disallow gpu selection for gguf for now

* cleanup

* Slightly larger baseline

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Treat empty list as auto

* Verbose logging/debug

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Cleanup and revert unnecessary deletions

* Cleanup excessive logs and guard against disk/cpu offload

* auth for visibility API. cleanup redundant imports. Adjust QLoRA estimate

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* support for non cuda gpus

* Fix multi-GPU auto-selection memory accounting

The multi_gpu_factor was applied uniformly to all GPUs including the
first one, which unfairly penalizes single-GPU capacity when
transitioning to multi-GPU. This created a discontinuity where a model
that barely fits 1 GPU would suddenly require 2 GPUs because the first
GPU's free memory was discounted by 20%.

Now the first GPU keeps its full free memory, and only additional GPUs
have an overhead factor (0.85) applied to account for inter-GPU
communication and sharding overhead. This gives more accurate
auto-selection and avoids unnecessary multi-GPU for models that
comfortably fit on one device.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add sandbox tests for multi-GPU selection logic

24 tests covering model size estimation, memory requirements, automatic
GPU selection, device map generation, GPU ID validation, and multi-GPU
overhead accounting. All tests use mocks so they run without GPUs on
Linux, macOS, and Windows.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix reviewer findings: 4bit inference estimate, fallback, GGUF gpu_ids, retry

1. 4-bit inference now uses reduced memory estimate (model_size/3 + buffer)
   instead of the FP16 1.3x multiplier. This prevents over-sharding
   quantized models across unnecessary GPUs.

2. When model size estimation fails, auto_select_gpu_ids now falls back to
   all visible GPUs instead of returning None (which could default to
   single-GPU loading for an unknown-size model).

3. GGUF inference route now treats gpu_ids=[] as auto-selection (same as
   None) instead of rejecting it as an unsupported explicit request.

4. Training retry path for "could not get source code" now preserves the
   gpu_ids parameter so the retry lands on the same GPUs.

5. Updated sandbox tests to cover the new 4-bit inference estimate branch.

* Remove accidentally added unsloth-zoo submodule

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix UUID/MIG visibility and update test expectations

1. nvidia.py: When CUDA_VISIBLE_DEVICES uses UUID/MIG tokens, the
   visibility APIs now return "unresolved" with empty device lists instead
   of exposing all physical GPUs. This prevents the UI from showing GPUs
   that the backend process cannot actually use.

2. test_gpu_selection.py: Updated test expectations to match the new
   multi-GPU overhead accounting (first GPU at full capacity, 0.85x for
   additional GPUs) and 4-bit inference memory estimation formula.
   All 60 tests now pass.

* Add CPU/disk offload guard to audio inference path

The audio model loading branch returned before the common
get_offloaded_device_map_entries() check, so audio models loaded with a
multi-GPU device_map that spilled layers to CPU/disk would be accepted
instead of rejected. Now audio loads also verify no modules are offloaded.

* Improve VRAM requirement estimates

* Replace balanced_low_0 with balanced

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refine calculations for slightly easier nums

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* adjust estimates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use nums instead of obj to avoid seralisation error

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden nvidia-smi parsing and fix fallback GPU list

1. nvidia.py: Wrap int() casts for GPU index and memory in try/except
   so MIG slices, N/A values, or unexpected nvidia-smi output skip the
   unparseable row instead of aborting the entire GPU list.

2. nvidia.py: Handle GPU names containing commas by using the last
   field as memory instead of a fixed positional index.

3. hardware.py: fallback_all now uses gpu_candidates (GPUs with verified
   VRAM data) instead of raw devices list, which could include GPUs
   with null VRAM that were excluded from the ranking.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* consolidate raise_if_offload

* Improve MoE support. Guard against nvidia-smi failures

* Improve MoE support. Guard against nvidia-smi failures

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix shared-expert LoRA undercount, torch VRAM fallback, and apply_gpu_ids edge case

1. vram_estimation.py: compute_lora_params now includes shared experts
   (n_shared_experts) alongside routed experts when computing MoE LoRA
   adapter parameters. Previously only n_experts were counted, causing
   the estimator to undercount adapter, optimizer, and gradient memory
   for DeepSeek/GLM-style models with shared experts.

2. hardware.py: _torch_get_per_device_info now uses mem_get_info (which
   reports system-wide VRAM usage) instead of memory_allocated (which
   only reports this process's PyTorch allocations). This prevents
   auto-selection from treating a GPU as mostly free when another
   process is consuming VRAM. Falls back to memory_allocated when
   mem_get_info is unavailable.

3. hardware.py: apply_gpu_ids([]) now returns early instead of setting
   CUDA_VISIBLE_DEVICES="" which would disable CUDA entirely. Empty
   list inherits the parent visibility, same as None.

4. hardware.py: Upgraded fallback_all GPU selection log from debug to
   warning so operators are notified when the model likely will not fit
   in available VRAM.

* Guard nvidia-smi subprocess calls against OSError and TimeoutExpired

get_visible_gpu_utilization and get_backend_visible_gpu_info now catch
OSError (nvidia-smi not found) and TimeoutExpired internally instead
of relying on callers to wrap every invocation. Returns the standard
available=False sentinel on failure so the torch-based fallback in
hardware.py can take over.

* Guard get_primary_gpu_utilization and reset GPU caches between tests

1. nvidia.py: get_primary_gpu_utilization now catches OSError and
   TimeoutExpired internally, matching the pattern already used in
   get_visible_gpu_utilization and get_backend_visible_gpu_info. All
   three nvidia-smi callers are now self-contained.

2. test_gpu_selection.py: Added _GpuCacheResetMixin that resets the
   module-level _physical_gpu_count and _visible_gpu_count caches in
   tearDown. Applied to all test classes that exercise GPU selection,
   device map, or visibility functions. This prevents stale cache
   values from leaking between tests and causing flaky results on
   machines with real GPUs.

* Fix nvidia-smi fallback regression and physical GPU count validation

1. hardware.py: get_gpu_utilization, get_visible_gpu_utilization, and
   get_backend_visible_gpu_info now check result.get("available") before
   returning the nvidia-smi result. When nvidia-smi is unavailable or
   returns no data (e.g., containers without nvidia-smi, UUID/MIG masks),
   the functions fall through to the torch-based fallback instead of
   returning an empty result. This fixes a regression where the internal
   exception handling in nvidia.py prevented the caller's except block
   from triggering the fallback.

2. hardware.py: resolve_requested_gpu_ids now separates negative-ID
   validation from physical upper-bound validation. The physical count
   check is only enforced when it is plausibly a true physical count
   (i.e., higher than the largest parent-visible ID), since
   torch.cuda.device_count() under CUDA_VISIBLE_DEVICES returns the
   visible count, not the physical total. The parent-visible-set check
   remains authoritative in all cases. This prevents valid physical IDs
   like [2, 3] from being rejected as "out of range" when nvidia-smi is
   unavailable and CUDA_VISIBLE_DEVICES="2,3" makes torch report only
   2 devices.

* Fix UUID/MIG torch fallback to enumerate devices by ordinal

When CUDA_VISIBLE_DEVICES uses UUID or MIG identifiers,
get_parent_visible_gpu_ids() returns [] because the tokens are
non-numeric. The torch fallback in get_visible_gpu_utilization() and
get_backend_visible_gpu_info() previously passed that empty list to
_torch_get_per_device_info(), getting nothing back.

Now both functions detect the empty-list case and fall back to
enumerating torch-visible ordinals (0..device_count-1) with
index_kind="relative". This means the UI and auto-selection still
see real device data in Kubernetes, MIG, and Slurm-style UUID
environments where nvidia-smi output cannot be mapped to physical
indices.

Updated test_uuid_parent_visibility to verify the new torch fallback
path returns available=True with relative ordinals.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add type hint for gpu_ids parameter in InferenceOrchestrator.load_model

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
---
 studio/backend/core/inference/inference.py    |   45 +-
 studio/backend/core/inference/llama_cpp.py    |   23 +-
 studio/backend/core/inference/orchestrator.py |   11 +
 studio/backend/core/inference/worker.py       |    4 +
 studio/backend/core/training/trainer.py       |   26 +-
 studio/backend/core/training/training.py      |   18 +
 studio/backend/core/training/worker.py        |    4 +
 studio/backend/main.py                        |   81 +-
 studio/backend/models/inference.py            |    4 +
 studio/backend/models/training.py             |    6 +
 studio/backend/routes/inference.py            |   13 +
 studio/backend/routes/training.py             |   16 +-
 studio/backend/tests/test_gpu_selection.py    | 1125 ++++++++++++++++
 .../tests/test_gpu_selection_sandbox.py       |  544 ++++++++
 studio/backend/tests/test_utils.py            |   32 +-
 studio/backend/tests/test_vram_estimation.py  |  695 ++++++++++
 .../backend/utils/hardware/VRAM_ESTIMATION.md |  161 +++
 studio/backend/utils/hardware/__init__.py     |   37 +
 studio/backend/utils/hardware/hardware.py     | 1143 +++++++++++++++--
 studio/backend/utils/hardware/nvidia.py       |  279 ++++
 .../backend/utils/hardware/vram_estimation.py |  501 ++++++++
 21 files changed, 4541 insertions(+), 227 deletions(-)
 create mode 100644 studio/backend/tests/test_gpu_selection.py
 create mode 100644 studio/backend/tests/test_gpu_selection_sandbox.py
 create mode 100644 studio/backend/tests/test_vram_estimation.py
 create mode 100644 studio/backend/utils/hardware/VRAM_ESTIMATION.md
 create mode 100644 studio/backend/utils/hardware/nvidia.py
 create mode 100644 studio/backend/utils/hardware/vram_estimation.py

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 ca39054ec0..f5361a6e8c 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -293,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
@@ -334,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 ─────────────────────────────────────
@@ -392,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
@@ -405,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
@@ -1082,6 +1099,10 @@ class LlamaCppBackend:
                     # 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:
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 67908d8617..c18f18a743 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -67,7 +67,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
@@ -230,69 +235,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()
@@ -311,6 +261,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 a20c2052aa..aabfba9b3a 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):
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 7d48198d42..1a94256059 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -206,8 +206,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()
 
@@ -369,6 +378,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:
@@ -420,6 +430,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)
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/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/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,
+    )

From a0bca759f362fb1294d41792f77b6c7c3c85749e Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 30 Mar 2026 02:40:29 -0700
Subject: [PATCH 7/9] Fix editable install scanning 6,500+ node_modules dirs
 (#4697)

* fix: scope packages.find to prevent node_modules namespace scanning

The packages.find section had no include filter, so setuptools'
find_namespace_packages discovered all directories as potential Python
packages -- including the 6,557 directories inside
studio/frontend/node_modules/ after the frontend build step.

This caused the editable install overlay step to run 20,000+ glob
operations across 6,619 "packages", which on fast NVMe takes ~5s but
on slower disks can take 7+ minutes.

Adding an explicit include filter scopes discovery to only the packages
we actually ship (unsloth, unsloth_cli, studio, studio.backend), dropping
from 6,619 to 58 discovered packages and the editable build time from
5.4s to 1.2s.

Also removes the broken kernels/moe exclude (used "/" instead of "."
notation so it never matched) and adds a node_modules exclude as a
safety net.

* fix: use precise node_modules exclude patterns

Use "*.node_modules" and "*.node_modules.*" instead of "*.node_modules*"
to avoid accidentally excluding valid packages that might contain
"node_modules" as a substring in their name.
---
 pyproject.toml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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 = [

From 6d83ad9a2834fc84cb39bd2a2fe2ba0ceb8d8262 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 30 Mar 2026 06:40:47 -0700
Subject: [PATCH 8/9] fix(studio): avoid UnicodeEncodeError on Windows cp1252
 consoles (#4699)

* fix(studio): replace unicode emoji in print() to avoid cp1252 crash on Windows

On Windows the default console encoding is cp1252 which cannot encode
unicode emoji like U+2705 or U+26A0. bare print() calls with these
characters cause a UnicodeEncodeError at runtime.

- run.py: replace emoji with ASCII status prefixes [OK] and [WARNING]
- format_conversion.py: remove duplicate print() that mirrors the
  logger.info() call on the next line, and drop the emoji from the
  log message since loggers handle encoding separately

* fix(studio): apply same emoji/print cleanup to parallel VLM conversion path

The parallel URL-based conversion logic has the same duplicate print()
with emoji that was fixed in the sequential path. Remove the bare
print() and drop the emoji from the logger.info() call.

* Treat install_python_stack.py failure as fatal in setup.ps1

On Linux/Mac, setup.sh runs under set -euo pipefail so a non-zero
exit from install_python_stack.py aborts the installer. On Windows,
setup.ps1 had no exit code check -- if the Python script crashed
(eg from the cp1252 UnicodeEncodeError), the installer silently
continued past the dependency loop and reported success. Studio
would then fail at launch with ModuleNotFoundError for structlog,
fastapi, and other deps that were never installed.

Capture $LASTEXITCODE and exit 1 if the dependency installer fails,
matching the error handling pattern already used for PyTorch install.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
---
 studio/backend/run.py                              |  4 ++--
 studio/backend/utils/datasets/format_conversion.py | 11 ++---------
 studio/setup.ps1                                   |  6 ++++++
 3 files changed, 10 insertions(+), 11 deletions(-)

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/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/setup.ps1 b/studio/setup.ps1
index 6e19fbea83..7cea60371a 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -1522,8 +1522,14 @@ if ($CuTag -eq "cpu") {
 # Ordered heavy dependency installation -- shared cross-platform script
 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

From 34272a796f3a33dc98a75c36e1370956280f6273 Mon Sep 17 00:00:00 2001
From: Etherll <61019402+Etherll@users.noreply.github.com>
Date: Mon, 30 Mar 2026 19:58:33 +0200
Subject: [PATCH 9/9] Fix/bun windows bin detection (#4703)

* fix(studio): detect bun .exe shims in Windows binary check

* Update setup.sh

* add .bunx checking
---
 studio/setup.ps1 | 12 +++++++-----
 studio/setup.sh  |  5 ++++-
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/studio/setup.ps1 b/studio/setup.ps1
index 7cea60371a..cdba0e6690 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -1112,9 +1112,11 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) {
     if ($UseBun) {
         Write-Host "   Using bun for package install (faster)" -ForegroundColor DarkGray
         $bunExit = Invoke-SetupCommand { bun install }
-        # 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")
+        # 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) {
@@ -1124,8 +1126,8 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) {
             }
             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")
-            $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd")
+            $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") {
diff --git a/studio/setup.sh b/studio/setup.sh
index c2a6891ec0..3715f536f6 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -257,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