Merge branch 'main' into pip

.github/dependabot.yml

@@ -0,0 +1,40 @@
+---
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    groups:
+      actions:
+        patterns: ["*"]
+
+  - package-ecosystem: "pip"
+    directories:
+      - "/"
+      - "/studio/backend/plugins/data-designer-unstructured-seed"
+      - "/studio/backend/requirements"
+      - "/unsloth/kernels/moe"
+    schedule:
+      interval: "weekly"
+    open-pull-requests-limit: 10
+    groups:
+      pip:
+        patterns: ["*"]
+
+  - package-ecosystem: "bun"
+    directory: "/studio/frontend"
+    schedule:
+      interval: "weekly"
+    groups:
+      bun-frontend:
+        patterns: ["*"]
+
+  - package-ecosystem: "npm"
+    directory: "/studio/backend/core/data_recipe/oxc-validator"
+    schedule:
+      interval: "weekly"
+    groups:
+      npm-oxc-validator:
+        patterns: ["*"]
+...

build.sh

@@ -29,7 +29,22 @@ _restore_gitignores() {
 }
 trap _restore_gitignores EXIT
 
-npm install
+# Use bun for install if available (faster), fall back to npm.
+_install_ok=false
+if command -v bun &>/dev/null; then
+    if bun install; then
+        _install_ok=true
+    else
+        echo "⚠ bun install failed, falling back to npm"
+        rm -rf node_modules
+    fi
+fi
+if [ "$_install_ok" != "true" ]; then
+    if ! npm install; then
+        echo "❌ ERROR: package install failed" >&2
+        exit 1
+    fi
+fi
 npm run build       # outputs to studio/frontend/dist/
 
 _restore_gitignores

install.ps1

@@ -5,8 +5,9 @@
 function Install-UnslothStudio {
     $ErrorActionPreference = "Stop"
 
-    $VenvName = "unsloth_studio"
     $PythonVersion = "3.13"
+    $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
+    $VenvDir = Join-Path $StudioHome "unsloth_studio"
 
     Write-Host ""
     Write-Host "========================================="
@@ -31,6 +32,275 @@ function Install-UnslothStudio {
         $env:Path = $unique -join ";"
     }
 
+    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
+            return
+        }
+        try {
+            # Persist an absolute path in launcher scripts so shortcut working
+            # directory changes do not break process startup.
+            $UnslothExePath = (Resolve-Path $UnslothExePath).Path
+            # Escape for single-quoted embedding in generated launcher script.
+            # This prevents runtime variable expansion for paths containing '$'.
+            $SingleQuotedExePath = $UnslothExePath -replace "'", "''"
+
+            $localAppDataDir = $env:LOCALAPPDATA
+            if (-not $localAppDataDir -or [string]::IsNullOrWhiteSpace($localAppDataDir)) {
+                Write-Host "[WARN] LOCALAPPDATA path unavailable; skipped shortcut creation" -ForegroundColor Yellow
+                return
+            }
+            $appDir = Join-Path $localAppDataDir "Unsloth Studio"
+            $launcherPs1 = Join-Path $appDir "launch-studio.ps1"
+            $launcherVbs = Join-Path $appDir "launch-studio.vbs"
+            $desktopDir = [Environment]::GetFolderPath("Desktop")
+            $desktopLink = if ($desktopDir -and $desktopDir.Trim()) {
+                Join-Path $desktopDir "Unsloth Studio.lnk"
+            } else {
+                $null
+            }
+            $startMenuDir = if ($env:APPDATA -and $env:APPDATA.Trim()) {
+                Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs"
+            } else {
+                $null
+            }
+            $startMenuLink = if ($startMenuDir -and $startMenuDir.Trim()) {
+                Join-Path $startMenuDir "Unsloth Studio.lnk"
+            } else {
+                $null
+            }
+            if (-not $desktopLink) {
+                Write-Host "[WARN] Desktop path unavailable; skipped desktop shortcut creation" -ForegroundColor Yellow
+            }
+            if (-not $startMenuLink) {
+                Write-Host "[WARN] APPDATA/Start Menu path unavailable; skipped Start menu shortcut creation" -ForegroundColor Yellow
+            }
+            $iconPath = Join-Path $appDir "unsloth.ico"
+            $bundledIcon = $null
+            if ($PSScriptRoot -and $PSScriptRoot.Trim()) {
+                $bundledIcon = Join-Path $PSScriptRoot "studio\frontend\public\unsloth.ico"
+            }
+            $iconUrl = "https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/unsloth.ico"
+
+            if (-not (Test-Path $appDir)) {
+                New-Item -ItemType Directory -Path $appDir -Force | Out-Null
+            }
+
+            $launcherContent = @"
+`$ErrorActionPreference = 'Stop'
+`$basePort = 8888
+`$maxPortOffset = 20
+`$timeoutSec = 60
+`$pollIntervalMs = 1000
+
+function Test-StudioHealth {
+    param([Parameter(Mandatory = `$true)][int]`$Port)
+    try {
+        `$url = "http://127.0.0.1:`$Port/api/health"
+        `$resp = Invoke-RestMethod -Uri `$url -TimeoutSec 1 -Method Get
+        return (`$resp -and `$resp.status -eq 'healthy' -and `$resp.service -eq 'Unsloth UI Backend')
+    } catch {
+        return `$false
+    }
+}
+
+function Get-CandidatePorts {
+    # Fast path: only probe base port + currently listening ports in range.
+    `$ports = @(`$basePort)
+    try {
+        `$maxPort = `$basePort + `$maxPortOffset
+        `$listening = Get-NetTCPConnection -State Listen -ErrorAction Stop |
+            Where-Object { `$_.LocalPort -ge `$basePort -and `$_.LocalPort -le `$maxPort } |
+            Select-Object -ExpandProperty LocalPort
+        `$ports = (@(`$basePort) + `$listening) | Sort-Object -Unique
+    } catch {
+        Write-Host "[DEBUG] Get-NetTCPConnection failed: `$(`$_.Exception.Message). Falling back to full port scan." -ForegroundColor DarkGray
+        # Fallback when Get-NetTCPConnection is unavailable/restricted.
+        for (`$offset = 1; `$offset -le `$maxPortOffset; `$offset++) {
+            `$ports += (`$basePort + `$offset)
+        }
+    }
+    return `$ports
+}
+
+function Find-HealthyStudioPort {
+    foreach (`$candidate in (Get-CandidatePorts)) {
+        if (Test-StudioHealth -Port `$candidate) {
+            return `$candidate
+        }
+    }
+    return `$null
+}
+
+# If Studio is already healthy on any expected port, just open it and exit.
+`$existingPort = Find-HealthyStudioPort
+if (`$existingPort) {
+    Start-Process "http://localhost:`$existingPort"
+    exit 0
+}
+
+`$launchMutex = [System.Threading.Mutex]::new(`$false, 'Local\UnslothStudioLauncher')
+`$haveMutex = `$false
+try {
+    try {
+        `$haveMutex = `$launchMutex.WaitOne(0)
+    } catch [System.Threading.AbandonedMutexException] {
+        `$haveMutex = `$true
+    }
+    if (-not `$haveMutex) {
+        # Another launcher is already running; wait for it to bring Studio up
+        `$deadline = (Get-Date).AddSeconds(`$timeoutSec)
+        while ((Get-Date) -lt `$deadline) {
+            `$port = Find-HealthyStudioPort
+            if (`$port) { Start-Process "http://localhost:`$port"; exit 0 }
+            Start-Sleep -Milliseconds `$pollIntervalMs
+        }
+        exit 0
+    }
+
+    `$powershellExe = Join-Path `$env:SystemRoot 'System32\WindowsPowerShell\v1.0\powershell.exe'
+    `$studioExe = '$SingleQuotedExePath'
+    `$studioCommand = '& "' + `$studioExe + '" studio -H 0.0.0.0 -p ' + `$basePort
+    `$launchArgs = @(
+        '-NoExit',
+        '-NoProfile',
+        '-ExecutionPolicy',
+        'Bypass',
+        '-Command',
+        `$studioCommand
+    )
+
+    try {
+        `$proc = Start-Process -FilePath `$powershellExe -ArgumentList `$launchArgs -WorkingDirectory `$env:USERPROFILE -PassThru
+    } catch {
+        `$msg = "Could not launch Unsloth Studio terminal.`n`nError: `$(`$_.Exception.Message)"
+        try {
+            Add-Type -AssemblyName System.Windows.Forms -ErrorAction Stop
+            [System.Windows.Forms.MessageBox]::Show(`$msg, 'Unsloth Studio') | Out-Null
+        } catch {}
+        exit 1
+    }
+
+    `$browserOpened = `$false
+    `$deadline = (Get-Date).AddSeconds(`$timeoutSec)
+    while ((Get-Date) -lt `$deadline) {
+        `$healthyPort = Find-HealthyStudioPort
+        if (`$healthyPort) {
+            Start-Process "http://localhost:`$healthyPort"
+            `$browserOpened = `$true
+            break
+        }
+        if (`$proc.HasExited) { break }
+        Start-Sleep -Milliseconds `$pollIntervalMs
+    }
+    if (-not `$browserOpened) {
+        if (`$proc.HasExited) {
+            `$msg = "Unsloth Studio exited before becoming healthy. Check terminal output for errors."
+        } else {
+            `$msg = "Unsloth Studio is still starting but did not become healthy within `$timeoutSec seconds. Check the terminal window for the selected port and open it manually."
+        }
+        try {
+            Add-Type -AssemblyName System.Windows.Forms -ErrorAction Stop
+            [System.Windows.Forms.MessageBox]::Show(`$msg, 'Unsloth Studio') | Out-Null
+        } catch {}
+    }
+} finally {
+    if (`$haveMutex) { `$launchMutex.ReleaseMutex() | Out-Null }
+    `$launchMutex.Dispose()
+}
+exit 0
+"@
+
+            # Write UTF-8 with BOM for reliable decoding by Windows PowerShell 5.1,
+            # even when install.ps1 is executed from PowerShell 7.
+            $utf8Bom = New-Object System.Text.UTF8Encoding($true)
+            [System.IO.File]::WriteAllText($launcherPs1, $launcherContent, $utf8Bom)
+            $vbsContent = @"
+Set shell = CreateObject("WScript.Shell")
+cmd = "powershell -NoProfile -ExecutionPolicy Bypass -WindowStyle Hidden -File ""$launcherPs1"""
+shell.Run cmd, 0, False
+"@
+            # WSH handles UTF-16LE reliably for .vbs files with non-ASCII paths.
+            Set-Content -Path $launcherVbs -Value $vbsContent -Encoding Unicode -Force
+
+            # Prefer bundled icon from local clone/dev installs.
+            # If not available, best-effort download from raw GitHub.
+            # We only attach the icon if the resulting file has a valid ICO header.
+            $hasValidIcon = $false
+            if ($bundledIcon -and (Test-Path $bundledIcon)) {
+                try {
+                    Copy-Item -Path $bundledIcon -Destination $iconPath -Force
+                } catch {
+                    Write-Host "[DEBUG] Error copying bundled icon: $($_.Exception.Message)" -ForegroundColor DarkGray
+                }
+            } elseif (-not (Test-Path $iconPath)) {
+                try {
+                    Invoke-WebRequest -Uri $iconUrl -OutFile $iconPath -UseBasicParsing
+                } catch {
+                    Write-Host "[DEBUG] Error downloading icon: $($_.Exception.Message)" -ForegroundColor DarkGray
+                }
+            }
+
+            if (Test-Path $iconPath) {
+                try {
+                    $bytes = [System.IO.File]::ReadAllBytes($iconPath)
+                    if (
+                        $bytes.Length -ge 4 -and
+                        $bytes[0] -eq 0 -and
+                        $bytes[1] -eq 0 -and
+                        $bytes[2] -eq 1 -and
+                        $bytes[3] -eq 0
+                    ) {
+                        $hasValidIcon = $true
+                    } else {
+                        Remove-Item $iconPath -Force -ErrorAction SilentlyContinue
+                    }
+                } catch {
+                    Write-Host "[DEBUG] Error validating or removing icon: $($_.Exception.Message)" -ForegroundColor DarkGray
+                    Remove-Item $iconPath -Force -ErrorAction SilentlyContinue
+                }
+            }
+
+            $wscriptExe = Join-Path $env:SystemRoot "System32\wscript.exe"
+            $shortcutArgs = "//B //Nologo `"$launcherVbs`""
+
+            try {
+                $wshell = New-Object -ComObject WScript.Shell
+                $createdShortcutCount = 0
+                foreach ($linkPath in @($desktopLink, $startMenuLink)) {
+                    if (-not $linkPath -or [string]::IsNullOrWhiteSpace($linkPath)) { continue }
+                    try {
+                        $shortcut = $wshell.CreateShortcut($linkPath)
+                        $shortcut.TargetPath = $wscriptExe
+                        $shortcut.Arguments = $shortcutArgs
+                        $shortcut.WorkingDirectory = $appDir
+                        $shortcut.Description = "Launch Unsloth Studio"
+                        if ($hasValidIcon) {
+                            $shortcut.IconLocation = "$iconPath,0"
+                        }
+                        $shortcut.Save()
+                        $createdShortcutCount++
+                    } catch {
+                        Write-Host "[WARN] Could not create shortcut at ${linkPath}: $($_.Exception.Message)" -ForegroundColor Yellow
+                    }
+                }
+                if ($createdShortcutCount -gt 0) {
+                    Write-Host "[OK] Created Unsloth Studio shortcut(s): $createdShortcutCount" -ForegroundColor Green
+                } else {
+                    Write-Host "[WARN] No Unsloth Studio shortcuts were created" -ForegroundColor Yellow
+                }
+            } catch {
+                Write-Host "[WARN] Shortcut creation unavailable: $($_.Exception.Message)" -ForegroundColor Yellow
+            }
+        } catch {
+            Write-Host "[WARN] Shortcut setup failed; skipping shortcuts: $($_.Exception.Message)" -ForegroundColor Yellow
+        }
+    }
+
     # ── Check winget ──
     if (-not (Get-Command winget -ErrorAction SilentlyContinue)) {
         Write-Host "Error: winget is not available." -ForegroundColor Red
@@ -180,20 +450,59 @@ function Install-UnslothStudio {
         return
     }
 
-    # ── Create venv (skip if it already exists and has a valid interpreter) ──
+    # ── Create venv (migrate old layout if possible, otherwise fresh) ──
     # Pass the resolved executable path to uv so it does not re-resolve
     # a version string back to a conda interpreter.
-    $VenvPython = Join-Path $VenvName "Scripts\python.exe"
+    if (-not (Test-Path $StudioHome)) {
+        New-Item -ItemType Directory -Path $StudioHome -Force | Out-Null
+    }
+
+    $VenvPython = Join-Path $VenvDir "Scripts\python.exe"
+    $_Migrated = $false
+
+    if (Test-Path $VenvPython) {
+        # New layout already exists -- nuke for fresh install
+        Write-Host "==> 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..."
+        $prevEAP2 = $ErrorActionPreference
+        $ErrorActionPreference = "Continue"
+        try {
+            & $OldPy -c "import torch; A = torch.ones((2,2)); B = A + A" 2>$null | Out-Null
+            $torchOk = ($LASTEXITCODE -eq 0)
+        } catch { $torchOk = $false }
+        $ErrorActionPreference = $prevEAP2
+        if ($torchOk) {
+            Write-Host "   Legacy environment is healthy -- migrating..."
+            Move-Item -Path $OldVenv -Destination $VenvDir -Force
+            Write-Host "   Moved .venv -> unsloth_studio"
+            $_Migrated = $true
+        } else {
+            Write-Host "   Legacy environment failed validation -- creating fresh environment"
+            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..."
+        Move-Item -Path $CwdVenv -Destination $VenvDir -Force
+        Write-Host "   Moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio"
+        $_Migrated = $true
+    }
+
     if (-not (Test-Path $VenvPython)) {
-        if (Test-Path $VenvName) { Remove-Item -Recurse -Force $VenvName }
-        Write-Host "==> Creating Python $($DetectedPython.Version) virtual environment (${VenvName})..."
-        uv venv $VenvName --python "$($DetectedPython.Path)"
+        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
             return
         }
     } else {
-        Write-Host "==> Virtual environment ${VenvName} already exists, skipping creation."
+        Write-Host "==> Using migrated environment at $VenvDir"
     }
 
     # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ──
@@ -267,15 +576,26 @@ function Install-UnslothStudio {
     #   CUDA wheels.  Missing dependencies (transformers, trl, peft, etc.)
     #   are still pulled in because they are new, not upgrades.
     #
-    Write-Host "==> Installing PyTorch ($TorchIndexUrl)..."
-    uv pip install --python $VenvPython torch torchvision torchaudio --index-url $TorchIndexUrl
-    if ($LASTEXITCODE -ne 0) {
-        Write-Host "[ERROR] Failed to install PyTorch (exit code $LASTEXITCODE)" -ForegroundColor Red
-        return
-    }
+    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..."
+        uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.11" unsloth-zoo
+    } elseif ($TorchIndexUrl) {
+        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
+            return
+        }
 
-    Write-Host "==> Installing unsloth (this may take a few minutes)..."
-    uv pip install --python $VenvPython --upgrade-package unsloth unsloth
+        Write-Host "==> Installing unsloth (this may take a few minutes)..."
+        uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11"
+    } else {
+        # Fallback: GPU detection failed to produce a URL -- let uv resolve torch
+        Write-Host "==> Installing unsloth (this may take a few minutes)..."
+        uv pip install --python $VenvPython "unsloth>=2026.3.11" --torch-backend=auto
+    }
     if ($LASTEXITCODE -ne 0) {
         Write-Host "[ERROR] Failed to install unsloth (exit code $LASTEXITCODE)" -ForegroundColor Red
         return
@@ -285,7 +605,7 @@ function Install-UnslothStudio {
     # 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..."
-    $UnslothExe = Join-Path $VenvName "Scripts\unsloth.exe"
+    $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe"
     if (-not (Test-Path $UnslothExe)) {
         Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red
         Write-Host "        Expected: $UnslothExe" -ForegroundColor Yellow
@@ -293,12 +613,16 @@ function Install-UnslothStudio {
         Write-Host "        Try re-running the installer or see: https://github.com/unslothai/unsloth?tab=readme-ov-file#-quickstart" -ForegroundColor Yellow
         return
     }
+    # Tell setup.ps1 to skip base package installation (install.ps1 already did it)
+    $env:SKIP_STUDIO_BASE = "1"
     & $UnslothExe studio setup
     if ($LASTEXITCODE -ne 0) {
         Write-Host "[ERROR] unsloth studio setup failed (exit code $LASTEXITCODE)" -ForegroundColor Red
         return
     }
 
+    New-StudioShortcuts -UnslothExePath $UnslothExe
+
     Write-Host ""
     Write-Host "========================================="
     Write-Host "   Unsloth Studio installed!"
@@ -311,12 +635,11 @@ function Install-UnslothStudio {
     if ($IsInteractive) {
         Write-Host "==> Launching Unsloth Studio..."
         Write-Host ""
-        $UnslothExe = Join-Path $VenvName "Scripts\unsloth.exe"
         & $UnslothExe studio -H 0.0.0.0 -p 8888
     } else {
         Write-Host "  To launch, run:"
         Write-Host ""
-        Write-Host "    .\${VenvName}\Scripts\activate"
+        Write-Host "    & `"$VenvDir\Scripts\Activate.ps1`""
         Write-Host "    unsloth studio -H 0.0.0.0 -p 8888"
         Write-Host ""
     }

install.sh

@@ -1,11 +1,35 @@
 #!/bin/sh
 # Unsloth Studio Installer
-# Usage (curl): curl -fsSL https://raw.githubusercontent.com/unslothai/unsloth/main/install.sh | sh
-# Usage (wget): wget -qO- https://raw.githubusercontent.com/unslothai/unsloth/main/install.sh | sh
+# Usage (curl):  curl -fsSL https://unsloth.ai/install.sh | sh
+# Usage (wget):  wget -qO- https://unsloth.ai/install.sh | sh
+# Usage (local): ./install.sh --local   (install from local repo instead of PyPI)
+# Usage (test):  ./install.sh --package roland-sloth  (install a different package name)
 set -e
 
-VENV_NAME="unsloth_studio"
+# ── Parse flags ──
+STUDIO_LOCAL_INSTALL=false
+PACKAGE_NAME="unsloth"
+_next_is_package=false
+for arg in "$@"; do
+    if [ "$_next_is_package" = true ]; then
+        PACKAGE_NAME="$arg"
+        _next_is_package=false
+        continue
+    fi
+    case "$arg" in
+        --local) STUDIO_LOCAL_INSTALL=true ;;
+        --package) _next_is_package=true ;;
+    esac
+done
+
+if [ "$_next_is_package" = true ]; then
+    echo "❌ ERROR: --package requires an argument." >&2
+    exit 1
+fi
+
 PYTHON_VERSION="3.13"
+STUDIO_HOME="$HOME/.unsloth/studio"
+VENV_DIR="$STUDIO_HOME/unsloth_studio"
 
 # ── Helper: download a URL to a file (supports curl and wget) ──
 download() {
@@ -89,6 +113,441 @@ _smart_apt_install() {
     fi
 }
 
+# ── Helper: create desktop shortcuts and launcher script ──
+# Usage: create_studio_shortcuts <unsloth_exe> <os>
+# Creates ~/.local/share/unsloth/launch-studio.sh (shared launcher),
+# plus platform-specific shortcuts (Linux .desktop / macOS .app bundle).
+# Skipped on WSL (no native desktop).
+create_studio_shortcuts() {
+    _css_exe="$1"
+    _css_os="$2"
+
+    # Skip on WSL -- no native desktop environment
+    if [ "$_css_os" = "wsl" ]; then
+        return 0
+    fi
+
+    # Validate exe
+    if [ ! -x "$_css_exe" ]; then
+        echo "[WARN] Cannot create shortcuts: unsloth not found at $_css_exe"
+        return 0
+    fi
+
+    # Resolve absolute path
+    _css_exe_dir=$(cd "$(dirname "$_css_exe")" && pwd)
+    _css_exe="$_css_exe_dir/$(basename "$_css_exe")"
+
+    _css_data_dir="$HOME/.local/share/unsloth"
+    _css_launcher="$_css_data_dir/launch-studio.sh"
+    _css_icon_png="$_css_data_dir/unsloth-studio.png"
+    _css_gem_png="$_css_data_dir/unsloth-gem.png"
+
+    mkdir -p "$_css_data_dir"
+
+    # ── Write launcher script ──
+    # The launcher is Bash (not POSIX sh).
+    # We write it with a placeholder and substitute the exe path via sed.
+    cat > "$_css_launcher" << 'LAUNCHER_EOF'
+#!/usr/bin/env bash
+# Unsloth Studio Launcher
+# Auto-generated by install.sh -- do not edit manually.
+set -euo pipefail
+
+DATA_DIR="$HOME/.local/share/unsloth"
+
+# Read exe path from config written at install time.
+# Sourcing is safe: the config file is written by install.sh, not user input.
+if [ -f "$DATA_DIR/studio.conf" ]; then
+    . "$DATA_DIR/studio.conf"
+fi
+if [ -z "${UNSLOTH_EXE:-}" ] || [ ! -x "${UNSLOTH_EXE:-}" ]; then
+    echo "Error: UNSLOTH_EXE not set or not executable. Re-run the installer." >&2
+    exit 1
+fi
+
+BASE_PORT=8888
+MAX_PORT_OFFSET=20
+TIMEOUT_SEC=60
+POLL_INTERVAL_SEC=1
+LOG_FILE="$DATA_DIR/studio.log"
+LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u).lock"
+
+# ── HTTP GET helper (supports curl and wget) ──
+_http_get() {
+    _url="$1"
+    if command -v curl >/dev/null 2>&1; then
+        curl -fsS --max-time 1 "$_url" 2>/dev/null
+    elif command -v wget >/dev/null 2>&1; then
+        wget -qO- --timeout=1 "$_url" 2>/dev/null
+    else
+        return 1
+    fi
+}
+
+# ── Health check ──
+_check_health() {
+    _port=$1
+    _resp=$(_http_get "http://127.0.0.1:$_port/api/health") || return 1
+    case "$_resp" in
+        *'"status"'*'"healthy"'*'"service"'*'"Unsloth UI Backend"'*) return 0 ;;
+        *'"service"'*'"Unsloth UI Backend"'*'"status"'*'"healthy"'*) return 0 ;;
+    esac
+    return 1
+}
+
+# ── Port scanning ──
+_candidate_ports() {
+    echo "$BASE_PORT"
+    _max_port=$((BASE_PORT + MAX_PORT_OFFSET))
+    if command -v ss >/dev/null 2>&1; then
+        ss -tlnH 2>/dev/null | awk '{print $4}' | grep -oE '[0-9]+$' | \
+            awk -v lo="$BASE_PORT" -v hi="$_max_port" '$1 >= lo && $1 <= hi && $1 != lo {print}' || true
+    elif command -v lsof >/dev/null 2>&1; then
+        lsof -iTCP -sTCP:LISTEN -nP 2>/dev/null | awk '{print $9}' | grep -oE '[0-9]+$' | \
+            awk -v lo="$BASE_PORT" -v hi="$_max_port" '$1 >= lo && $1 <= hi && $1 != lo {print}' || true
+    else
+        _offset=1
+        while [ "$_offset" -le "$MAX_PORT_OFFSET" ]; do
+            echo $((BASE_PORT + _offset))
+            _offset=$((_offset + 1))
+        done
+    fi
+}
+
+_find_healthy_port() {
+    for _p in $(_candidate_ports | sort -un); do
+        if _check_health "$_p"; then
+            echo "$_p"
+            return 0
+        fi
+    done
+    return 1
+}
+
+# ── Check if a port is busy ──
+_is_port_busy() {
+    _port=$1
+    if command -v ss >/dev/null 2>&1; then
+        ss -tlnH 2>/dev/null | awk '{print $4}' | grep -qE "[.:]$_port$"
+    elif command -v lsof >/dev/null 2>&1; then
+        lsof -iTCP:"$_port" -sTCP:LISTEN -nP >/dev/null 2>&1
+    else
+        return 1
+    fi
+}
+
+# ── Find a free port in range ──
+_find_launch_port() {
+    _offset=0
+    while [ "$_offset" -le "$MAX_PORT_OFFSET" ]; do
+        _candidate=$((BASE_PORT + _offset))
+        if ! _is_port_busy "$_candidate"; then
+            echo "$_candidate"
+            return 0
+        fi
+        _offset=$((_offset + 1))
+    done
+    return 1
+}
+
+# ── Open browser ──
+_open_browser() {
+    _url="$1"
+    if [ "$(uname)" = "Darwin" ] && command -v open >/dev/null 2>&1; then
+        open "$_url"
+    elif command -v xdg-open >/dev/null 2>&1; then
+        xdg-open "$_url" >/dev/null 2>&1 &
+    else
+        echo "Open in your browser: $_url" >&2
+    fi
+}
+
+# ── Spawn terminal with studio command ──
+_spawn_terminal() {
+    _cmd="$1"
+    _os=$(uname)
+    if [ "$_os" = "Darwin" ]; then
+        # Escape backslashes and double-quotes for AppleScript string
+        _cmd_escaped=$(printf '%s' "$_cmd" | sed 's/\\/\\\\/g; s/"/\\"/g')
+        osascript -e "tell application \"Terminal\" to do script \"$_cmd_escaped\"" >/dev/null 2>&1 && return 0
+    else
+        for _term in gnome-terminal konsole xfce4-terminal mate-terminal lxterminal xterm; do
+            if command -v "$_term" >/dev/null 2>&1; then
+                case "$_term" in
+                    gnome-terminal) "$_term" -- sh -c "$_cmd" & return 0 ;;
+                    konsole)        "$_term" -e sh -c "$_cmd" & return 0 ;;
+                    xterm)          "$_term" -e sh -c "$_cmd" & return 0 ;;
+                    *)              "$_term" -e sh -c "$_cmd" & return 0 ;;
+                esac
+            fi
+        done
+    fi
+    # Fallback: background with log
+    echo "No terminal emulator found; running in background. Logs: $LOG_FILE" >&2
+    nohup sh -c "$_cmd" >> "$LOG_FILE" 2>&1 &
+    return 0
+}
+
+# ── Atomic directory-based single-instance guard ──
+_acquire_lock() {
+    if mkdir "$LOCK_DIR" 2>/dev/null; then
+        echo "$$" > "$LOCK_DIR/pid"
+        return 0
+    fi
+
+    # Lock dir exists -- check if owner is still alive
+    _old_pid=$(cat "$LOCK_DIR/pid" 2>/dev/null || true)
+    if [ -n "$_old_pid" ] && kill -0 "$_old_pid" 2>/dev/null; then
+        # Another launcher is running; wait for it to bring Studio up
+        _deadline=$(($(date +%s) + TIMEOUT_SEC))
+        while [ "$(date +%s)" -lt "$_deadline" ]; do
+            _port=$(_find_healthy_port) && {
+                _open_browser "http://localhost:$_port"
+                exit 0
+            }
+            sleep "$POLL_INTERVAL_SEC"
+        done
+        echo "Timed out waiting for other launcher (PID $_old_pid)" >&2
+        exit 0
+    fi
+
+    # Stale lock -- reclaim
+    rm -rf "$LOCK_DIR"
+    mkdir "$LOCK_DIR" 2>/dev/null || return 1
+    echo "$$" > "$LOCK_DIR/pid"
+}
+
+_release_lock() {
+    rm -rf "$LOCK_DIR"
+}
+
+# ── Main ──
+# Fast path: already healthy
+_port=$(_find_healthy_port) && {
+    _open_browser "http://localhost:$_port"
+    exit 0
+}
+
+_acquire_lock
+trap '_release_lock' EXIT INT TERM
+
+# Post-lock re-check (handles race with another launcher)
+_port=$(_find_healthy_port) && {
+    _open_browser "http://localhost:$_port"
+    exit 0
+}
+
+# Find a free port in range
+_launch_port=$(_find_launch_port) || {
+    echo "No free port found in range ${BASE_PORT}-$((BASE_PORT + MAX_PORT_OFFSET))" >&2
+    exit 1
+}
+
+# Launch studio in a terminal
+_launch_cmd=$(printf '%q ' "$UNSLOTH_EXE" studio -H 0.0.0.0 -p "$_launch_port")
+_launch_cmd=${_launch_cmd% }
+_spawn_terminal "$_launch_cmd"
+
+# Poll for health
+_deadline=$(($(date +%s) + TIMEOUT_SEC))
+while [ "$(date +%s)" -lt "$_deadline" ]; do
+    _port=$(_find_healthy_port) && {
+        _open_browser "http://localhost:$_port"
+        exit 0
+    }
+    sleep "$POLL_INTERVAL_SEC"
+done
+
+echo "Unsloth Studio did not become healthy within ${TIMEOUT_SEC}s." >&2
+echo "Check logs at: $LOG_FILE" >&2
+exit 1
+LAUNCHER_EOF
+
+    chmod +x "$_css_launcher"
+
+    # Write the exe path to a separate conf file sourced by the launcher.
+    # Using single-quote wrapping with the standard '\'' escape for any
+    # embedded apostrophes. This avoids all sed metacharacter issues.
+    _css_quoted_exe=$(printf '%s' "$_css_exe" | sed "s/'/'\\\\''/g")
+    printf '%s\n' "UNSLOTH_EXE='$_css_quoted_exe'" > "$_css_data_dir/studio.conf"
+
+    # ── Icon: try bundled, then download ──
+    # favicon.png (small, for Linux) and unsloth-gem.png (large, for macOS icns)
+    _css_script_dir=""
+    if [ -n "${0:-}" ] && [ -f "$0" ]; then
+        _css_script_dir=$(cd "$(dirname "$0")" 2>/dev/null && pwd) || true
+    fi
+
+    # Try to find favicon.png from installed package (site-packages) or local repo
+    _css_found_favicon=""
+    _css_found_gem=""
+    _css_venv_dir=$(dirname "$(dirname "$_css_exe")")
+    # Check site-packages
+    for _sp in "$_css_venv_dir"/lib/python*/site-packages/unsloth/studio/frontend/public; do
+        if [ -f "$_sp/favicon.png" ]; then
+            _css_found_favicon="$_sp/favicon.png"
+        fi
+        if [ -f "$_sp/unsloth-gem.png" ]; then
+            _css_found_gem="$_sp/unsloth-gem.png"
+        fi
+    done
+    # Check local repo (when running from clone)
+    if [ -z "$_css_found_favicon" ] && [ -n "$_css_script_dir" ] && [ -f "$_css_script_dir/studio/frontend/public/favicon.png" ]; then
+        _css_found_favicon="$_css_script_dir/studio/frontend/public/favicon.png"
+    fi
+    if [ -z "$_css_found_gem" ] && [ -n "$_css_script_dir" ] && [ -f "$_css_script_dir/studio/frontend/public/unsloth-gem.png" ]; then
+        _css_found_gem="$_css_script_dir/studio/frontend/public/unsloth-gem.png"
+    fi
+
+    # Copy or download favicon.png
+    if [ -n "$_css_found_favicon" ]; then
+        cp "$_css_found_favicon" "$_css_icon_png" 2>/dev/null || true
+    elif [ ! -f "$_css_icon_png" ]; then
+        download "https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/favicon.png" "$_css_icon_png" 2>/dev/null || true
+    fi
+    # Copy or download unsloth-gem.png (for macOS icns)
+    if [ -n "$_css_found_gem" ]; then
+        cp "$_css_found_gem" "$_css_gem_png" 2>/dev/null || true
+    elif [ ! -f "$_css_gem_png" ]; then
+        download "https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/unsloth-gem.png" "$_css_gem_png" 2>/dev/null || true
+    fi
+
+    # Validate PNG header (first 4 bytes: \x89PNG)
+    _css_validate_png() {
+        [ -f "$1" ] || return 1
+        _hdr=$(od -An -tx1 -N4 "$1" 2>/dev/null | tr -d ' ')
+        [ "$_hdr" = "89504e47" ]
+    }
+    if [ -f "$_css_icon_png" ] && ! _css_validate_png "$_css_icon_png"; then
+        rm -f "$_css_icon_png"
+    fi
+    if [ -f "$_css_gem_png" ] && ! _css_validate_png "$_css_gem_png"; then
+        rm -f "$_css_gem_png"
+    fi
+
+    # ── Platform-specific shortcuts ──
+    _css_created=0
+
+    if [ "$_css_os" = "linux" ]; then
+        # ── Linux: .desktop file ──
+        _css_app_dir="$HOME/.local/share/applications"
+        mkdir -p "$_css_app_dir"
+
+        _css_desktop="$_css_app_dir/unsloth-studio.desktop"
+        # Escape backslashes and double-quotes for .desktop Exec= field
+        _css_exec_escaped=$(printf '%s' "$_css_launcher" | sed 's/\\/\\\\/g; s/"/\\"/g')
+        _css_icon_escaped=$(printf '%s' "$_css_icon_png" | sed 's/\\/\\\\/g; s/"/\\"/g')
+        cat > "$_css_desktop" << DESKTOP_EOF
+[Desktop Entry]
+Version=1.0
+Type=Application
+Name=Unsloth Studio
+Comment=Launch Unsloth Studio
+Exec="$_css_exec_escaped"
+Icon=$_css_icon_escaped
+Terminal=false
+StartupNotify=true
+Categories=Development;Science;
+DESKTOP_EOF
+        chmod +x "$_css_desktop"
+
+        # Copy to ~/Desktop if it exists
+        if [ -d "$HOME/Desktop" ]; then
+            cp "$_css_desktop" "$HOME/Desktop/unsloth-studio.desktop" 2>/dev/null || true
+            chmod +x "$HOME/Desktop/unsloth-studio.desktop" 2>/dev/null || true
+            # Mark as trusted so GNOME/Nautilus allows launching via double-click
+            if command -v gio >/dev/null 2>&1; then
+                gio set "$HOME/Desktop/unsloth-studio.desktop" metadata::trusted true 2>/dev/null || true
+            fi
+        fi
+
+        # Best-effort update database
+        update-desktop-database "$_css_app_dir" 2>/dev/null || true
+        _css_created=1
+
+    elif [ "$_css_os" = "macos" ]; then
+        # ── macOS: .app bundle ──
+        _css_app="$HOME/Applications/Unsloth Studio.app"
+        _css_contents="$_css_app/Contents"
+        _css_macos_dir="$_css_contents/MacOS"
+        _css_res_dir="$_css_contents/Resources"
+        mkdir -p "$_css_macos_dir" "$_css_res_dir"
+
+        # Info.plist
+        cat > "$_css_contents/Info.plist" << 'PLIST_EOF'
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+    <key>CFBundleIdentifier</key>
+    <string>ai.unsloth.studio</string>
+    <key>CFBundleName</key>
+    <string>Unsloth Studio</string>
+    <key>CFBundleDisplayName</key>
+    <string>Unsloth Studio</string>
+    <key>CFBundleExecutable</key>
+    <string>launch-studio</string>
+    <key>CFBundleIconFile</key>
+    <string>AppIcon</string>
+    <key>CFBundlePackageType</key>
+    <string>APPL</string>
+    <key>CFBundleVersion</key>
+    <string>1.0</string>
+    <key>CFBundleShortVersionString</key>
+    <string>1.0</string>
+    <key>LSMinimumSystemVersion</key>
+    <string>10.15</string>
+    <key>NSHighResolutionCapable</key>
+    <true/>
+</dict>
+</plist>
+PLIST_EOF
+
+        # Executable stub
+        cat > "$_css_macos_dir/launch-studio" << STUB_EOF
+#!/bin/sh
+exec "$HOME/.local/share/unsloth/launch-studio.sh" "\$@"
+STUB_EOF
+        chmod +x "$_css_macos_dir/launch-studio"
+
+        # Build AppIcon.icns from unsloth-gem.png (2240x2240)
+        if [ -f "$_css_gem_png" ] && command -v sips >/dev/null 2>&1 && command -v iconutil >/dev/null 2>&1; then
+            _css_tmpdir=$(mktemp -d 2>/dev/null)
+            if [ -d "$_css_tmpdir" ]; then
+                _css_iconset="$_css_tmpdir/AppIcon.iconset"
+                mkdir -p "$_css_iconset"
+                _css_icon_ok=true
+                for _sz in 16 32 128 256 512; do
+                    _sz2=$((_sz * 2))
+                    sips -z "$_sz" "$_sz" "$_css_gem_png" --out "$_css_iconset/icon_${_sz}x${_sz}.png" >/dev/null 2>&1 || _css_icon_ok=false
+                    sips -z "$_sz2" "$_sz2" "$_css_gem_png" --out "$_css_iconset/icon_${_sz}x${_sz}@2x.png" >/dev/null 2>&1 || _css_icon_ok=false
+                done
+                if [ "$_css_icon_ok" = "true" ]; then
+                    iconutil -c icns "$_css_iconset" -o "$_css_res_dir/AppIcon.icns" 2>/dev/null || true
+                fi
+                rm -rf "$_css_tmpdir"
+            fi
+        fi
+        # Fallback: copy PNG as icon
+        if [ ! -f "$_css_res_dir/AppIcon.icns" ] && [ -f "$_css_icon_png" ]; then
+            cp "$_css_icon_png" "$_css_res_dir/AppIcon.icns" 2>/dev/null || true
+        fi
+
+        # Touch so Finder indexes it
+        touch "$_css_app"
+
+        # Symlink on Desktop
+        if [ -d "$HOME/Desktop" ]; then
+            ln -sf "$_css_app" "$HOME/Desktop/Unsloth Studio" 2>/dev/null || true
+        fi
+        _css_created=1
+    fi
+
+    if [ "$_css_created" -eq 1 ]; then
+        echo "[OK] Created Unsloth Studio shortcut(s)"
+    fi
+}
+
 echo ""
 echo "========================================="
 echo "   Unsloth Studio Installer"
@@ -224,32 +683,197 @@ if ! command -v uv >/dev/null 2>&1 || ! _uv_version_ok uv; then
     export PATH="$HOME/.local/bin:$PATH"
 fi
 
-# ── Create venv (skip if it already exists and has a valid interpreter) ──
-if [ ! -x "$VENV_NAME/bin/python" ]; then
-    [ -e "$VENV_NAME" ] && rm -rf "$VENV_NAME"
-    echo "==> Creating Python ${PYTHON_VERSION} virtual environment (${VENV_NAME})..."
-    uv venv "$VENV_NAME" --python "$PYTHON_VERSION"
-else
-    echo "==> Virtual environment ${VENV_NAME} already exists, skipping creation."
+# ── Create venv (migrate old layout if possible, otherwise fresh) ──
+mkdir -p "$STUDIO_HOME"
+
+_MIGRATED=false
+
+if [ -x "$VENV_DIR/bin/python" ]; then
+    # New layout already exists — nuke for fresh install
+    rm -rf "$VENV_DIR"
+elif [ -x "$STUDIO_HOME/.venv/bin/python" ]; then
+    # Old layout exists — validate before migrating
+    echo "==> Found legacy Studio environment, validating..."
+    if "$STUDIO_HOME/.venv/bin/python" -c "
+import torch
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+A = torch.ones((10, 10), device=device)
+B = torch.ones((10, 10), device=device)
+C = torch.ones((10, 10), device=device)
+D = A + B
+E = D @ C
+torch.testing.assert_close(torch.unique(E), torch.tensor((20,), device=E.device, dtype=E.dtype))
+" >/dev/null 2>&1; then
+        echo "✅ Legacy environment is healthy — migrating..."
+        mv "$STUDIO_HOME/.venv" "$VENV_DIR"
+        echo "   Moved ~/.unsloth/studio/.venv → $VENV_DIR"
+        _MIGRATED=true
+    else
+        echo "⚠️  Legacy environment failed validation — creating fresh environment"
+        rm -rf "$STUDIO_HOME/.venv"
+    fi
 fi
 
+if [ ! -x "$VENV_DIR/bin/python" ]; then
+    echo "==> Creating Python ${PYTHON_VERSION} virtual environment (${VENV_DIR})..."
+    uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
+else
+    echo "==> Using migrated environment at ${VENV_DIR}"
+fi
+
+# ── Resolve repo root (for --local installs) ──
+_REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)"
+
+# ── Detect GPU and choose PyTorch index URL ──
+# Mirrors Get-TorchIndexUrl in install.ps1.
+# On CPU-only machines this returns the cpu index, avoiding the solver
+# dead-end where --torch-backend=auto resolves to unsloth==2024.8.
+get_torch_index_url() {
+    _base="https://download.pytorch.org/whl"
+    # macOS: always CPU (no CUDA support)
+    case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac
+    # Try nvidia-smi
+    _smi=""
+    if command -v nvidia-smi >/dev/null 2>&1; then
+        _smi="nvidia-smi"
+    elif [ -x "/usr/bin/nvidia-smi" ]; then
+        _smi="/usr/bin/nvidia-smi"
+    fi
+    if [ -z "$_smi" ]; then echo "$_base/cpu"; return; fi
+    # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P)
+    _cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \
+        | sed -n 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \
+        | head -1)
+    if [ -z "$_cuda_ver" ]; then
+        echo "[WARN] Could not determine CUDA version from nvidia-smi, defaulting to cu126" >&2
+        echo "$_base/cu126"; return
+    fi
+    _major=${_cuda_ver%%.*}
+    _minor=${_cuda_ver#*.}
+    if [ "$_major" -ge 13 ]; then echo "$_base/cu130"
+    elif [ "$_major" -eq 12 ] && [ "$_minor" -ge 8 ]; then echo "$_base/cu128"
+    elif [ "$_major" -eq 12 ] && [ "$_minor" -ge 6 ]; then echo "$_base/cu126"
+    elif [ "$_major" -ge 12 ]; then echo "$_base/cu124"
+    elif [ "$_major" -ge 11 ]; then echo "$_base/cu118"
+    else echo "$_base/cpu"; fi
+}
+TORCH_INDEX_URL=$(get_torch_index_url)
+
 # ── Install unsloth directly into the venv (no activation needed) ──
-echo "==> Installing unsloth (this may take a few minutes)..."
-uv pip install --python "$VENV_NAME/bin/python" unsloth --torch-backend=auto
+_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..."
+    uv pip install --python "$_VENV_PY" \
+        --reinstall-package unsloth --reinstall-package unsloth-zoo \
+        "unsloth>=2026.3.11" unsloth-zoo
+    if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
+        echo "==> Overlaying local repo (editable)..."
+        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
+    echo "==> Installing PyTorch ($TORCH_INDEX_URL)..."
+    uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \
+        --index-url "$TORCH_INDEX_URL"
+    # Fresh: Step 2 - install unsloth, preserving pre-installed torch
+    echo "==> Installing unsloth (this may take a few minutes)..."
+    if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
+        uv pip install --python "$_VENV_PY" \
+            --upgrade-package unsloth "unsloth>=2026.3.11" unsloth-zoo
+        echo "==> Overlaying local repo (editable)..."
+        uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
+    else
+        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)..."
+    if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
+        uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.11" --torch-backend=auto
+        echo "==> Overlaying local repo (editable)..."
+        uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
+    else
+        uv pip install --python "$_VENV_PY" "$PACKAGE_NAME" --torch-backend=auto
+    fi
+fi
 
 # ── Run studio setup ──
-# Ensure the venv's Python is on PATH for setup.sh's Python discovery.
-# On macOS the system Python may be outside the 3.11-3.13 range that
-# setup.sh requires, but uv already installed a compatible interpreter
-# inside the venv.
-VENV_ABS_BIN="$(cd "$VENV_NAME/bin" && pwd)"
+# When --local, use the repo's own setup.sh directly.
+# Otherwise, find it inside the installed package.
+SETUP_SH=""
+if [ "$STUDIO_LOCAL_INSTALL" = true ] && [ -f "$_REPO_ROOT/studio/setup.sh" ]; then
+    SETUP_SH="$_REPO_ROOT/studio/setup.sh"
+fi
+
+if [ -z "$SETUP_SH" ] || [ ! -f "$SETUP_SH" ]; then
+    SETUP_SH=$("$VENV_DIR/bin/python" -c "
+import importlib.resources
+print(importlib.resources.files('studio') / 'setup.sh')
+" 2>/dev/null || echo "")
+fi
+
+# Fallback: search site-packages
+if [ -z "$SETUP_SH" ] || [ ! -f "$SETUP_SH" ]; then
+    SETUP_SH=$(find "$VENV_DIR" -path "*/studio/setup.sh" -print -quit 2>/dev/null || echo "")
+fi
+
+if [ -z "$SETUP_SH" ] || [ ! -f "$SETUP_SH" ]; then
+    echo "❌ ERROR: Could not find studio/setup.sh in the installed package."
+    exit 1
+fi
+
+# Ensure the venv's Python is on PATH so setup.sh can find it.
+VENV_ABS_BIN="$(cd "$VENV_DIR/bin" && pwd)"
 if [ -n "$VENV_ABS_BIN" ]; then
     export PATH="$VENV_ABS_BIN:$PATH"
 fi
 
-echo "==> Running unsloth studio setup..."
-REQUESTED_PYTHON_VERSION="$(cd "$VENV_NAME/bin" && pwd)/python" \
-"$VENV_NAME/bin/unsloth" studio setup </dev/null
+echo "==> Running unsloth setup..."
+if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
+    SKIP_STUDIO_BASE=1 \
+    STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
+    STUDIO_LOCAL_INSTALL=1 \
+    STUDIO_LOCAL_REPO="$_REPO_ROOT" \
+    bash "$SETUP_SH" </dev/null
+else
+    SKIP_STUDIO_BASE=1 \
+    STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
+    bash "$SETUP_SH" </dev/null
+fi
+
+# ── Make 'unsloth' available globally via ~/.local/bin ──
+mkdir -p "$HOME/.local/bin"
+ln -sf "$VENV_DIR/bin/unsloth" "$HOME/.local/bin/unsloth"
+
+_LOCAL_BIN="$HOME/.local/bin"
+case ":$PATH:" in
+    *":$_LOCAL_BIN:"*) ;;  # already on PATH
+    *)
+        _SHELL_PROFILE=""
+        if [ -n "${ZSH_VERSION:-}" ] || [ "$(basename "${SHELL:-}")" = "zsh" ]; then
+            _SHELL_PROFILE="$HOME/.zshrc"
+        elif [ -f "$HOME/.bashrc" ]; then
+            _SHELL_PROFILE="$HOME/.bashrc"
+        elif [ -f "$HOME/.profile" ]; then
+            _SHELL_PROFILE="$HOME/.profile"
+        fi
+
+        if [ -n "$_SHELL_PROFILE" ]; then
+            if ! grep -q '\.local/bin' "$_SHELL_PROFILE" 2>/dev/null; then
+                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"
+            fi
+        fi
+        export PATH="$_LOCAL_BIN:$PATH"
+        ;;
+esac
+
+create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS"
 
 echo ""
 echo "========================================="
@@ -257,8 +881,32 @@ echo "   Unsloth Studio installed!"
 echo "========================================="
 echo ""
 
-echo "  To launch, run:"
-echo ""
-echo "    source ${VENV_NAME}/bin/activate"
-echo "    unsloth studio -H 0.0.0.0 -p 8888"
-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 ""
+    "$VENV_DIR/bin/unsloth" studio -H 0.0.0.0 -p 8888
+    _LAUNCH_EXIT=$?
+    if [ "$_LAUNCH_EXIT" -ne 0 ] && [ "$_MIGRATED" = true ]; then
+        echo ""
+        echo "⚠️  Unsloth Studio failed to start after migration."
+        echo "   Your migrated environment may be incompatible."
+        echo "   To fix, remove the environment and reinstall:"
+        echo ""
+        echo "   rm -rf $VENV_DIR"
+        echo "   curl -fsSL https://unsloth.ai/install.sh | sh"
+        echo ""
+    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"
+    echo ""
+fi

pyproject.toml

@@ -69,12 +69,9 @@ studio = [
     "*.ps1",
     "*.bat",
     "frontend/dist/**/*",
-    "frontend/public/**/*",
-    "frontend/src/**/*",
     "frontend/*.json",
     "frontend/*.ts",
     "frontend/*.js",
-    "frontend/*.lock",
     "frontend/*.html",
     "frontend/*.yaml",
     "frontend/.git*",

studio/backend/colab.py

@@ -26,7 +26,7 @@ def _bootstrap_studio_venv() -> None:
     site-packages so that packages like structlog, fastapi, etc. are
     importable from notebook cells and take priority over system copies.
     """
-    venv_lib = Path.home() / ".unsloth" / "studio" / ".venv" / "lib"
+    venv_lib = Path.home() / ".unsloth" / "studio" / "unsloth_studio" / "lib"
     if not venv_lib.exists():
         import warnings
 

studio/backend/core/data_recipe/oxc-validator/package.json

@@ -4,7 +4,7 @@
   "version": "0.0.1",
   "type": "module",
   "dependencies": {
-    "oxc-parser": "^0.116.0",
+    "oxc-parser": "^0.121.0",
     "oxlint": "^1.51.0"
   }
 }

studio/backend/core/inference/llama_cpp.py

@@ -857,6 +857,9 @@ class LlamaCppBackend:
 
             if use_fit:
                 cmd.extend(["--fit", "on"])
+            elif gpu_indices is not None:
+                # Model fits on selected GPU(s) -- offload all layers
+                cmd.extend(["-ngl", "-1"])
 
             if n_threads is not None:
                 cmd.extend(["--threads", str(n_threads)])
@@ -966,6 +969,46 @@ class LlamaCppBackend:
 
                 lib_dirs = [binary_dir]
                 _arch = platform.machine()  # x86_64, aarch64, etc.
+
+                # Pip-installed nvidia CUDA runtime libs (e.g. torch's
+                # bundled cuda-bindings).  The prebuilt llama.cpp binary
+                # links against libcudart.so.13 / libcublas.so.13 which
+                # live here, not in /usr/local/cuda.
+                import glob as _glob
+
+                for _nv_pattern in [
+                    os.path.join(
+                        sys.prefix,
+                        "lib",
+                        "python*",
+                        "site-packages",
+                        "nvidia",
+                        "cu*",
+                        "lib",
+                    ),
+                    os.path.join(
+                        sys.prefix,
+                        "lib",
+                        "python*",
+                        "site-packages",
+                        "nvidia",
+                        "cudnn",
+                        "lib",
+                    ),
+                    os.path.join(
+                        sys.prefix,
+                        "lib",
+                        "python*",
+                        "site-packages",
+                        "nvidia",
+                        "nvjitlink",
+                        "lib",
+                    ),
+                ]:
+                    for _nv_dir in _glob.glob(_nv_pattern):
+                        if os.path.isdir(_nv_dir):
+                            lib_dirs.append(_nv_dir)
+
                 for cuda_lib in [
                     "/usr/local/cuda/lib64",
                     f"/usr/local/cuda/targets/{_arch}-linux/lib",

studio/backend/core/training/trainer.py

@@ -81,8 +81,8 @@ class TrainingProgress:
     epoch: float = 0
     step: int = 0
     total_steps: int = 0
-    loss: float = 0.0
-    learning_rate: float = 0.0
+    loss: Optional[float] = None
+    learning_rate: Optional[float] = None
     is_training: bool = False
     is_completed: bool = False
     error: Optional[str] = None
@@ -244,7 +244,7 @@ class UnslothTrainer:
             def on_log(self, args, state, control, logs = None, **kwargs):
                 if not logs:
                     return
-                loss_value = logs.get("loss", logs.get("train_loss", 0.0))
+                loss_value = logs.get("loss", logs.get("train_loss", None))
                 current_step = state.global_step
                 grad_norm = logs.get("grad_norm", None)
 
@@ -268,7 +268,7 @@ class UnslothTrainer:
                     step = current_step,
                     epoch = round(state.epoch, 2) if state.epoch else 0,
                     loss = loss_value,
-                    learning_rate = logs.get("learning_rate", 0.0),
+                    learning_rate = logs.get("learning_rate", None),
                     elapsed_seconds = elapsed_seconds,
                     eta_seconds = eta_seconds,
                     grad_norm = grad_norm,

studio/backend/core/training/training.py

@@ -14,12 +14,14 @@ worker's mp.Queue, and exposes the same API surface to routes/training.py.
 Pattern follows core/data_recipe/jobs/manager.py.
 """
 
+import json as _json
 import math
 import multiprocessing as mp
 import queue
 import threading
 import time
 import structlog
+from datetime import datetime, timezone
 from loggers import get_logger
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -44,8 +46,8 @@ class TrainingProgress:
     epoch: float = 0
     step: int = 0
     total_steps: int = 0
-    loss: float = 0.0
-    learning_rate: float = 0.0
+    loss: Optional[float] = None
+    learning_rate: Optional[float] = None
     is_training: bool = False
     is_completed: bool = False
     error: Optional[str] = None
@@ -63,6 +65,8 @@ class TrainingBackend:
     Launches a fresh subprocess per training job, communicates via mp.Queue.
     """
 
+    FLUSH_THRESHOLD: int = 10
+
     def __init__(self):
         # Subprocess state
         self._proc: Optional[mp.Process] = None
@@ -91,13 +95,21 @@ class TrainingBackend:
         self.current_job_id: Optional[str] = None
         self._output_dir: Optional[str] = None
 
+        # DB persistence
+        self._metric_buffer: list[dict] = []
+        self._run_finalized: bool = False
+        self._db_run_created: bool = False
+        self._db_total_steps_set: bool = False
+        self._db_config: Optional[dict] = None
+        self._db_started_at: Optional[str] = None
+
         logger.info("TrainingBackend initialized (subprocess mode)")
 
     # ------------------------------------------------------------------
     # Public API (called by routes/training.py)
     # ------------------------------------------------------------------
 
-    def start_training(self, **kwargs) -> bool:
+    def start_training(self, job_id: str, **kwargs) -> bool:
         """Spawn a subprocess to run the full training pipeline.
 
         All kwargs are serialized into a config dict and sent to the worker.
@@ -108,30 +120,16 @@ class TrainingBackend:
                 logger.warning("Training subprocess already running")
                 return False
 
-        # Join prior pump thread to prevent it from consuming events
-        # from the new job's queue (it reads self._event_queue dynamically).
+        # Join prior pump thread — refuse to start if it won't die
         if self._pump_thread is not None and self._pump_thread.is_alive():
             self._pump_thread.join(timeout = 5.0)
             if self._pump_thread.is_alive():
-                logger.warning("Previous pump thread did not exit within 5s")
+                logger.warning(
+                    "Previous pump thread did not exit within 5s — refusing to start"
+                )
+                return False
         self._pump_thread = None
 
-        # Reset state
-        self._should_stop = False
-        self._cancel_requested = False
-        self._progress = TrainingProgress(
-            is_training = True, status_message = "Initializing training..."
-        )
-        self.loss_history.clear()
-        self.lr_history.clear()
-        self.step_history.clear()
-        self.grad_norm_history.clear()
-        self.grad_norm_step_history.clear()
-        self.eval_loss_history.clear()
-        self.eval_step_history.clear()
-        self.eval_enabled = False
-        self._output_dir = None
-
         # Build config dict for the subprocess
         config = {
             "model_name": kwargs["model_name"],
@@ -193,23 +191,62 @@ class TrainingBackend:
         if config["training_type"] != "LoRA/QLoRA":
             config["load_in_4bit"] = False
 
-        # Spawn subprocess
+        # Spawn subprocess — use locals so state is untouched on failure
         from .worker import run_training_process
 
-        self._event_queue = _CTX.Queue()
-        self._stop_queue = _CTX.Queue()
+        event_queue = _CTX.Queue()
+        stop_queue = _CTX.Queue()
 
-        self._proc = _CTX.Process(
+        proc = _CTX.Process(
             target = run_training_process,
             kwargs = {
-                "event_queue": self._event_queue,
-                "stop_queue": self._stop_queue,
+                "event_queue": event_queue,
+                "stop_queue": stop_queue,
                 "config": config,
             },
             daemon = True,
         )
-        self._proc.start()
-        logger.info("Training subprocess started (pid=%s)", self._proc.pid)
+        try:
+            proc.start()
+        except Exception:
+            logger.error("Failed to start training subprocess", exc_info = True)
+            return False
+
+        logger.info("Training subprocess started (pid=%s)", proc.pid)
+
+        # Reset state — safe because old pump thread is confirmed dead
+        # and proc.start() succeeded
+        self.current_job_id = job_id
+        self._should_stop = False
+        self._cancel_requested = False
+        self._progress = TrainingProgress(
+            is_training = True, status_message = "Initializing training..."
+        )
+        self.loss_history.clear()
+        self.lr_history.clear()
+        self.step_history.clear()
+        self.grad_norm_history.clear()
+        self.grad_norm_step_history.clear()
+        self.eval_loss_history.clear()
+        self.eval_step_history.clear()
+        self.eval_enabled = False
+        self._output_dir = None
+        self._metric_buffer.clear()
+        self._run_finalized = False
+        self._db_run_created = False
+        self._db_total_steps_set = False
+        self._db_config = {
+            k: v for k, v in config.items() if k not in {"hf_token", "wandb_token"}
+        }
+        self._db_started_at = datetime.now(timezone.utc).isoformat()
+
+        # Assign subprocess handles after state reset
+        self._event_queue = event_queue
+        self._stop_queue = stop_queue
+        self._proc = proc
+
+        # Eagerly create DB run row so the run appears in history during model loading
+        self._ensure_db_run_created()
 
         # Start event pump thread
         self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
@@ -252,6 +289,11 @@ class TrainingBackend:
                 proc.kill()
                 proc.join(timeout = 2.0)
 
+        # Wait for pump thread to finish DB finalization before returning
+        # (8s covers SQLite's default 5s lock timeout plus execution overhead)
+        if self._pump_thread is not None and self._pump_thread.is_alive():
+            self._pump_thread.join(timeout = 8.0)
+
     def is_training_active(self) -> bool:
         """Check if training is currently active."""
         with self._lock:
@@ -389,20 +431,54 @@ class TrainingBackend:
                             self._progress.error
                             or "Training process exited unexpectedly"
                         )
+
+            self._ensure_db_run_created()
+            self._finalize_run_in_db(
+                status = "stopped" if self._should_stop else "error",
+                error_message = None
+                if self._should_stop
+                else "Training process terminated unexpectedly",
+            )
             return
 
     def _handle_event(self, event: dict) -> None:
-        """Apply a subprocess event to local state."""
+        """Apply a subprocess event to local state.
+
+        State updates happen inside self._lock; DB I/O happens after
+        releasing it so status-polling API endpoints are never blocked
+        by slow SQLite writes.
+        """
         etype = event.get("type")
+        db_action: Optional[str] = None
+        db_action_kwargs: dict = {}
 
         with self._lock:
             if etype == "progress":
                 self._progress.step = event.get("step", self._progress.step)
                 self._progress.epoch = event.get("epoch", self._progress.epoch)
-                self._progress.loss = event.get("loss", self._progress.loss)
-                self._progress.learning_rate = event.get(
-                    "learning_rate", self._progress.learning_rate
-                )
+                # loss/lr are sanitized below; update progress after coercion
+                _raw_loss = event.get("loss")
+                _raw_lr = event.get("learning_rate")
+                try:
+                    _safe_loss = float(_raw_loss) if _raw_loss is not None else None
+                except (TypeError, ValueError):
+                    logger.debug("Could not convert loss to float: %s", _raw_loss)
+                    _safe_loss = None
+                if _safe_loss is not None and not math.isfinite(_safe_loss):
+                    _safe_loss = None
+                try:
+                    _safe_lr = float(_raw_lr) if _raw_lr is not None else None
+                except (TypeError, ValueError):
+                    logger.debug(
+                        "Could not convert learning_rate to float: %s", _raw_lr
+                    )
+                    _safe_lr = None
+                if _safe_lr is not None and not math.isfinite(_safe_lr):
+                    _safe_lr = None
+                if _safe_loss is not None:
+                    self._progress.loss = _safe_loss
+                if _safe_lr is not None:
+                    self._progress.learning_rate = _safe_lr
                 self._progress.total_steps = event.get(
                     "total_steps", self._progress.total_steps
                 )
@@ -416,30 +492,85 @@ class TrainingBackend:
                 if status:
                     self._progress.status_message = status
 
-                # Update metric histories
+                # Update metric histories — reuse sanitized values from above
                 step = event.get("step", 0)
-                loss = event.get("loss", 0.0)
-                lr = event.get("learning_rate", 0.0)
-                if step >= 0 and loss > 0:
+                loss = _safe_loss
+                lr = _safe_lr
+                if step > 0 and loss is not None:
                     self.loss_history.append(loss)
-                    self.lr_history.append(lr)
+                    self.lr_history.append(lr if lr is not None else 0.0)
                     self.step_history.append(step)
 
                 grad_norm = event.get("grad_norm")
+                gn = None
                 if grad_norm is not None:
                     try:
                         gn = float(grad_norm)
                     except (TypeError, ValueError):
                         gn = None
-                    if gn is not None and math.isfinite(gn):
+                    if step > 0 and gn is not None and math.isfinite(gn):
                         self.grad_norm_history.append(gn)
                         self.grad_norm_step_history.append(step)
+                    else:
+                        gn = None
 
                 eval_loss = event.get("eval_loss")
                 if eval_loss is not None:
-                    self.eval_loss_history.append(eval_loss)
-                    self.eval_step_history.append(step)
-                    self.eval_enabled = True
+                    try:
+                        eval_loss = float(eval_loss)
+                    except (TypeError, ValueError):
+                        logger.debug(
+                            "Could not convert eval_loss to float: %s", eval_loss
+                        )
+                        eval_loss = None
+                    if step > 0 and eval_loss is not None and math.isfinite(eval_loss):
+                        self.eval_loss_history.append(eval_loss)
+                        self.eval_step_history.append(step)
+                        self.eval_enabled = True
+                    else:
+                        eval_loss = None
+
+                # Buffer metric for DB flush (loss/lr already sanitized above)
+                self._metric_buffer.append(
+                    {
+                        "step": step,
+                        "loss": loss,
+                        "learning_rate": lr,
+                        "grad_norm": gn,
+                        "eval_loss": eval_loss,
+                        "epoch": event.get("epoch"),
+                        "num_tokens": event.get("num_tokens"),
+                        "elapsed_seconds": event.get("elapsed_seconds"),
+                    }
+                )
+
+                # Decide which DB action to take after releasing the lock
+                if not self._db_run_created and self.current_job_id and self._db_config:
+                    db_action = "create_run"
+                    db_action_kwargs = {
+                        "job_id": self.current_job_id,
+                        "model_name": self._db_config["model_name"],
+                        "dataset_name": self._db_config.get("hf_dataset")
+                        or next(
+                            iter(self._db_config.get("local_datasets") or []), "unknown"
+                        ),
+                        "config_json": _json.dumps(self._db_config),
+                        "started_at": self._db_started_at
+                        or datetime.now(timezone.utc).isoformat(),
+                        "total_steps": event.get("total_steps"),
+                    }
+                elif (
+                    event.get("total_steps")
+                    and self._db_run_created
+                    and not self._db_total_steps_set
+                ):
+                    db_action = "update_total_steps"
+                    db_action_kwargs = {
+                        "job_id": self.current_job_id,
+                        "total_steps": event["total_steps"],
+                    }
+                elif len(self._metric_buffer) >= self.FLUSH_THRESHOLD:
+                    db_action = "flush"
 
             elif etype == "eval_configured":
                 self.eval_enabled = True
@@ -454,6 +585,14 @@ class TrainingBackend:
                 self._output_dir = event.get("output_dir")
                 msg = event.get("status_message", "Training completed")
                 self._progress.status_message = msg
+                if not self._db_run_created and self.current_job_id and self._db_config:
+                    db_action = "create_and_finalize"
+                else:
+                    db_action = "finalize"
+                db_action_kwargs = {
+                    "status": "stopped" if self._should_stop else "completed",
+                    "output_dir": self._output_dir,
+                }
 
             elif etype == "error":
                 self._progress.is_training = False
@@ -462,6 +601,149 @@ class TrainingBackend:
                 stack = event.get("stack", "")
                 if stack:
                     logger.error("Stack trace:\n%s", stack)
+                if not self._db_run_created and self.current_job_id and self._db_config:
+                    db_action = "create_and_finalize"
+                else:
+                    db_action = "finalize"
+                db_action_kwargs = {
+                    "status": "stopped" if self._should_stop else "error",
+                    "error_message": event.get("error", "Unknown error"),
+                }
+
+        # --- DB I/O outside the lock ---
+        if db_action == "create_run":
+            try:
+                from storage.studio_db import create_run
+
+                create_run(
+                    id = db_action_kwargs["job_id"],
+                    model_name = db_action_kwargs["model_name"],
+                    dataset_name = db_action_kwargs["dataset_name"],
+                    config_json = db_action_kwargs["config_json"],
+                    started_at = db_action_kwargs["started_at"],
+                    total_steps = db_action_kwargs["total_steps"],
+                )
+                self._db_run_created = True
+                if db_action_kwargs["total_steps"]:
+                    self._db_total_steps_set = True
+            except Exception:
+                logger.warning("Failed to create DB run record", exc_info = True)
+        elif db_action == "create_and_finalize":
+            self._ensure_db_run_created()
+            self._finalize_run_in_db(**db_action_kwargs)
+        elif db_action == "update_total_steps":
+            try:
+                from storage.studio_db import update_run_total_steps
+
+                update_run_total_steps(
+                    db_action_kwargs["job_id"], db_action_kwargs["total_steps"]
+                )
+                self._db_total_steps_set = True
+            except Exception:
+                logger.warning("Failed to update total_steps in DB", exc_info = True)
+        elif db_action == "flush":
+            self._flush_metrics_to_db()
+        elif db_action == "finalize":
+            self._finalize_run_in_db(**db_action_kwargs)
+
+    def _ensure_db_run_created(self) -> None:
+        """Create the DB row if it doesn't exist yet. Called outside the lock."""
+        if self._db_run_created or not self.current_job_id or not self._db_config:
+            return
+        try:
+            from storage.studio_db import create_run
+
+            dataset_name = self._db_config.get("hf_dataset") or next(
+                iter(self._db_config.get("local_datasets") or []), "unknown"
+            )
+            create_run(
+                id = self.current_job_id,
+                model_name = self._db_config["model_name"],
+                dataset_name = dataset_name,
+                config_json = _json.dumps(self._db_config),
+                started_at = self._db_started_at
+                or datetime.now(timezone.utc).isoformat(),
+                total_steps = self._progress.total_steps or None,
+            )
+            self._db_run_created = True
+        except Exception:
+            logger.warning(
+                "Failed to create DB run record for early failure", exc_info = True
+            )
+
+    def _finalize_run_in_db(
+        self,
+        status: str,
+        error_message: Optional[str] = None,
+        output_dir: Optional[str] = None,
+    ) -> None:
+        """Flush remaining metrics and mark a run as finished in the DB."""
+        if not self.current_job_id or not self._db_run_created or self._run_finalized:
+            return
+        self._flush_metrics_to_db()
+        try:
+            from storage.studio_db import finish_run
+            from utils.downsample import downsample
+
+            sparkline = downsample(self.loss_history, 50)
+            finish_run(
+                id = self.current_job_id,
+                status = status,
+                ended_at = datetime.now(timezone.utc).isoformat(),
+                final_step = self._progress.step,
+                final_loss = self._progress.loss
+                if (
+                    self._progress.loss is not None
+                    and math.isfinite(self._progress.loss)
+                )
+                else None,
+                duration_seconds = self._progress.elapsed_seconds,
+                loss_sparkline = _json.dumps(sparkline),
+                output_dir = output_dir,
+                error_message = error_message,
+            )
+            self._run_finalized = True
+        except Exception:
+            logger.warning(
+                "Failed to finalize run in DB (status=%s)", status, exc_info = True
+            )
+
+    def _flush_metrics_to_db(self) -> None:
+        """Flush buffered metrics to the database and update live progress."""
+        if (
+            not self._metric_buffer
+            or not self.current_job_id
+            or not self._db_run_created
+        ):
+            return
+        # Cap buffer to prevent unbounded memory growth
+        if len(self._metric_buffer) > 500:
+            logger.warning(
+                "Metric buffer exceeded 500 entries (%d) — trimming oldest",
+                len(self._metric_buffer),
+            )
+            self._metric_buffer = self._metric_buffer[-500:]
+        # Snapshot before insert so metrics arriving during the write are preserved
+        batch = list(self._metric_buffer)
+        try:
+            from storage.studio_db import insert_metrics_batch, update_run_progress
+
+            insert_metrics_batch(self.current_job_id, batch)
+            del self._metric_buffer[: len(batch)]
+            update_run_progress(
+                id = self.current_job_id,
+                step = self._progress.step,
+                loss = self._progress.loss
+                if (
+                    self._progress.loss is not None
+                    and math.isfinite(self._progress.loss)
+                )
+                else None,
+                duration_seconds = self._progress.elapsed_seconds,
+            )
+        except Exception:
+            # Leave buffer intact for retry on next flush
+            logger.warning("Failed to flush metrics to DB", exc_info = True)
 
     @staticmethod
     def _read_queue(q: Any, timeout_sec: float) -> Optional[dict]:
@@ -561,11 +843,13 @@ class TrainingBackend:
             if progress.error:
                 title = f"Error: {progress.error}"
             elif progress.is_completed:
-                title = f"Training completed! Final loss: {progress.loss:.4f}"
+                loss_str = f"{progress.loss:.4f}" if progress.loss is not None else "--"
+                title = f"Training completed! Final loss: {loss_str}"
             elif progress.status_message:
                 title = progress.status_message
             elif progress.step > 0:
-                title = f"Epoch: {progress.epoch} | Step: {progress.step}/{progress.total_steps} | Loss: {progress.loss:.4f}"
+                loss_str = f"{progress.loss:.4f}" if progress.loss is not None else "--"
+                title = f"Epoch: {progress.epoch} | Step: {progress.step}/{progress.total_steps} | Loss: {loss_str}"
             else:
                 title = "Training Loss"
 

studio/backend/core/training/worker.py

@@ -16,15 +16,294 @@ from __future__ import annotations
 import structlog
 from loggers import get_logger
 import os
+import platform
+import shutil
 import sys
 import time
 import traceback
+import json
+import subprocess as _sp
 from pathlib import Path
 from typing import Any
+import urllib.error
+import urllib.request
 
 logger = get_logger(__name__)
 
 
+_CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4"
+_CAUSAL_CONV1D_PACKAGE_VERSION = "1.6.1"
+_MAMBA_SSM_RELEASE_TAG = "v2.3.1"
+_MAMBA_SSM_PACKAGE_VERSION = "2.3.1"
+
+
+def _model_wants_causal_conv1d(model_name: str) -> bool:
+    name = model_name.lower()
+    return any(
+        key in name
+        for key in (
+            "qwen3.5",
+            "qwen3_5",
+            "qwen3-next",
+            "qwen3_next",
+            "nemotron_h",
+            "nemotron-h",
+            "nemotron-3-nano",
+            "falcon_h1",
+            "falcon-h1",
+            "granite-4.0-h",
+            "granitemoehybrid",
+            "lfm2",
+        )
+    )
+
+
+def _causal_conv1d_platform_tag() -> str | None:
+    machine = platform.machine().lower()
+    if sys.platform.startswith("linux"):
+        if machine in {"x86_64", "amd64"}:
+            return "linux_x86_64"
+        if machine in {"aarch64", "arm64"}:
+            return "linux_aarch64"
+        return None
+    # No prebuilt wheels published for macOS or Windows
+    return None
+
+
+def _probe_causal_conv1d_env() -> dict[str, str] | None:
+    try:
+        probe = _sp.run(
+            [
+                sys.executable,
+                "-c",
+                (
+                    "import json, sys, re, torch; "
+                    "parts = torch.__version__.split('+', 1)[0].split('.')[:2]; "
+                    "minor = re.sub(r'[^0-9].*', '', parts[1]) if len(parts) > 1 else '0'; "
+                    "torch_mm = parts[0] + '.' + minor; "
+                    "print(json.dumps({"
+                    "'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', "
+                    "'torch_mm': torch_mm, "
+                    "'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', "
+                    "'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()"
+                    "}))"
+                ),
+            ],
+            stdout = _sp.PIPE,
+            stderr = _sp.PIPE,
+            text = True,
+            timeout = 30,
+        )
+    except _sp.TimeoutExpired:
+        logger.warning("Torch environment probe timed out after 30s")
+        return None
+    if probe.returncode != 0:
+        logger.warning(
+            "Failed to probe torch environment for causal-conv1d wheel:\n%s",
+            probe.stdout,
+        )
+        return None
+
+    try:
+        return json.loads(probe.stdout.strip())
+    except json.JSONDecodeError:
+        logger.warning(
+            "Failed to parse torch environment probe output: %s", probe.stdout
+        )
+        return None
+
+
+def _direct_wheel_url(
+    *,
+    filename_prefix: str,
+    package_version: str,
+    release_tag: str,
+    release_base_url: str,
+    env: dict[str, str] | None = None,
+) -> str | None:
+    env = env or _probe_causal_conv1d_env()
+    platform_tag = _causal_conv1d_platform_tag()
+    if env is None or platform_tag is None or not env.get("cuda_major"):
+        return None
+
+    filename = (
+        f"{filename_prefix}-{package_version}"
+        f"+cu{env['cuda_major']}torch{env['torch_mm']}"
+        f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}-{platform_tag}.whl"
+    )
+    return f"{release_base_url}/{release_tag}/{filename}"
+
+
+def _url_exists(url: str) -> bool:
+    try:
+        request = urllib.request.Request(url, method = "HEAD")
+        with urllib.request.urlopen(request, timeout = 10):
+            return True
+    except urllib.error.HTTPError as exc:
+        if exc.code == 404:
+            return False
+        logger.warning("Unexpected HTTP error while probing %s: %s", url, exc)
+        return False
+    except Exception as exc:
+        logger.warning("Failed to probe %s: %s", url, exc)
+        return False
+
+
+def _install_package_wheel_first(
+    *,
+    event_queue: Any,
+    import_name: str,
+    display_name: str,
+    pypi_name: str,
+    pypi_version: str,
+    filename_prefix: str,
+    release_tag: str,
+    release_base_url: str,
+) -> None:
+    try:
+        __import__(import_name)
+        logger.info("%s already installed", display_name)
+        return
+    except ImportError:
+        pass
+
+    env = _probe_causal_conv1d_env()
+    wheel_url = _direct_wheel_url(
+        filename_prefix = filename_prefix,
+        package_version = pypi_version,
+        release_tag = release_tag,
+        release_base_url = release_base_url,
+        env = env,
+    )
+
+    if wheel_url is None:
+        logger.info("No compatible %s wheel candidate", display_name)
+    else:
+        if _url_exists(wheel_url):
+            _send_status(event_queue, f"Installing prebuilt {display_name} wheel...")
+            installed = False
+            # Try uv first if available, then fall back to pip
+            if shutil.which("uv"):
+                uv_cmd = [
+                    "uv",
+                    "pip",
+                    "install",
+                    "--python",
+                    sys.executable,
+                    "--no-deps",
+                    wheel_url,
+                ]
+                result = _sp.run(
+                    uv_cmd,
+                    stdout = _sp.PIPE,
+                    stderr = _sp.STDOUT,
+                    text = True,
+                )
+                if result.returncode == 0:
+                    installed = True
+                else:
+                    logger.warning(
+                        "uv failed to install %s wheel:\n%s",
+                        display_name,
+                        result.stdout,
+                    )
+            if not installed:
+                pip_cmd = [
+                    sys.executable,
+                    "-m",
+                    "pip",
+                    "install",
+                    "--no-deps",
+                    wheel_url,
+                ]
+                result = _sp.run(
+                    pip_cmd,
+                    stdout = _sp.PIPE,
+                    stderr = _sp.STDOUT,
+                    text = True,
+                )
+                if result.returncode == 0:
+                    installed = True
+                else:
+                    logger.warning(
+                        "pip failed to install %s wheel:\n%s",
+                        display_name,
+                        result.stdout,
+                    )
+            if installed:
+                logger.info("Installed prebuilt %s wheel successfully", display_name)
+                return
+        else:
+            logger.info("No published %s wheel found: %s", display_name, wheel_url)
+
+    _send_status(event_queue, f"Installing {display_name} from PyPI...")
+    pypi_cmd = [
+        sys.executable,
+        "-m",
+        "pip",
+        "install",
+        "--no-build-isolation",
+        "--no-deps",
+        "--no-cache-dir",
+        f"{pypi_name}=={pypi_version}",
+    ]
+    result = _sp.run(
+        pypi_cmd,
+        stdout = _sp.PIPE,
+        stderr = _sp.STDOUT,
+        text = True,
+    )
+    if result.returncode != 0:
+        logger.error("Failed to install %s from PyPI:\n%s", display_name, result.stdout)
+        return
+
+    logger.info("Installed %s from PyPI", display_name)
+
+
+def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None:
+    if not _model_wants_causal_conv1d(model_name):
+        return
+
+    _install_package_wheel_first(
+        event_queue = event_queue,
+        import_name = "causal_conv1d",
+        display_name = "causal-conv1d",
+        pypi_name = "causal-conv1d",
+        pypi_version = _CAUSAL_CONV1D_PACKAGE_VERSION,
+        filename_prefix = "causal_conv1d",
+        release_tag = _CAUSAL_CONV1D_RELEASE_TAG,
+        release_base_url = "https://github.com/Dao-AILab/causal-conv1d/releases/download",
+    )
+
+
+_SSM_MODEL_SUBSTRINGS = (
+    "nemotron_h",
+    "nemotron-h",
+    "nemotron-3-nano",
+    "falcon_h1",
+    "falcon-h1",
+    "granite-4.0-h",
+    "granitemoehybrid",
+)
+
+
+def _ensure_mamba_ssm(event_queue: Any, model_name: str) -> None:
+    if not any(sub in model_name.lower() for sub in _SSM_MODEL_SUBSTRINGS):
+        return
+
+    logger.info("SSM model detected; setting up mamba-ssm after causal-conv1d")
+    _install_package_wheel_first(
+        event_queue = event_queue,
+        import_name = "mamba_ssm",
+        display_name = "mamba-ssm",
+        pypi_name = "mamba-ssm",
+        pypi_version = _MAMBA_SSM_PACKAGE_VERSION,
+        filename_prefix = "mamba_ssm",
+        release_tag = _MAMBA_SSM_RELEASE_TAG,
+        release_base_url = "https://github.com/state-spaces/mamba/releases/download",
+    )
+
+
 def _activate_transformers_version(model_name: str) -> None:
     """Activate the correct transformers version BEFORE any ML imports.
 
@@ -121,45 +400,24 @@ def run_training_process(
             model_name,
         )
 
-    # ── 1b. Auto-install mamba-ssm for SSM/hybrid models (NemotronH, Falcon-H1) ──
-    _SSM_MODEL_SUBSTRINGS = ("nemotron_h", "nemotron-3-nano", "falcon_h1", "falcon-h1")
-    if any(sub in model_name.lower() for sub in _SSM_MODEL_SUBSTRINGS):
-        try:
-            import mamba_ssm  # noqa: F401
-
-            logger.info("mamba-ssm already installed")
-        except ImportError:
-            logger.info(
-                "SSM model detected — installing mamba-ssm and causal-conv1d (this may take several minutes)..."
-            )
-            _send_status(
-                event_queue, "Installing mamba-ssm (first time only, ~7 min)..."
-            )
-            import subprocess as _sp
-
-            # --no-build-isolation: compile against current torch (no version conflicts)
-            # --no-deps: don't pull in torch/transformers/triton (already installed)
-            for _pkg in ["causal_conv1d", "mamba_ssm"]:
-                _r = _sp.run(
-                    [
-                        sys.executable,
-                        "-m",
-                        "pip",
-                        "install",
-                        "--no-build-isolation",
-                        "--no-deps",
-                        "--no-cache-dir",
-                        _pkg,
-                    ],
-                    stdout = _sp.PIPE,
-                    stderr = _sp.STDOUT,
-                    text = True,
-                )
-                if _r.returncode != 0:
-                    logger.error("Failed to install %s:\n%s", _pkg, _r.stdout)
-                else:
-                    logger.info("Installed %s successfully", _pkg)
-            logger.info("mamba-ssm installation complete")
+    # ── 1b. Set up causal-conv1d first, then install mamba-ssm if needed ──
+    try:
+        _ensure_causal_conv1d_fast_path(event_queue, model_name)
+        _ensure_mamba_ssm(event_queue, model_name)
+    except Exception as exc:
+        event_queue.put(
+            {
+                "type": "error",
+                "error": (
+                    f"Please choose another model to train, since "
+                    f"causal-conv1d / mamba-ssm failed to install "
+                    f"with error: {exc}"
+                ),
+                "stack": traceback.format_exc(limit = 20),
+                "ts": time.time(),
+            }
+        )
+        return
 
     # ── 1c. Set fork start method so dataset.map() can multiprocess ──
     # The parent launched us via spawn (clean process), but the compiled
@@ -242,7 +500,7 @@ def run_training_process(
 
     # Wire up progress callback → event_queue
     def _on_progress(progress: TrainingProgress):
-        has_train_loss = progress.step >= 0 and progress.loss > 0
+        has_train_loss = progress.step > 0 and progress.loss is not None
         has_eval_loss = progress.eval_loss is not None
         if has_train_loss or has_eval_loss:
             event_queue.put(
@@ -918,7 +1176,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
         def on_log(self, args, state, control, logs = None, **kwargs):
             if not logs:
                 return
-            loss_value = logs.get("loss", logs.get("train_loss", 0.0))
+            loss_value = logs.get("loss", logs.get("train_loss", None))
             current_step = state.global_step
 
             elapsed = time.time() - training_start_time
@@ -934,7 +1192,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
                     "step": current_step,
                     "epoch": round(state.epoch, 2) if state.epoch else 0,
                     "loss": loss_value,
-                    "learning_rate": logs.get("learning_rate", 0.0),
+                    "learning_rate": logs.get("learning_rate", None),
                     "total_steps": total_steps,
                     "elapsed_seconds": elapsed,
                     "eta_seconds": eta,

studio/backend/main.py

@@ -49,6 +49,7 @@ from routes import (
     export_router,
     inference_router,
     models_router,
+    training_history_router,
     training_router,
 )
 from auth import storage
@@ -73,6 +74,17 @@ async def lifespan(app: FastAPI):
     # Detect hardware first — sets DEVICE global used everywhere
     detect_hardware()
 
+    from storage.studio_db import cleanup_orphaned_runs
+
+    try:
+        cleanup_orphaned_runs()
+    except Exception as exc:
+        import structlog
+
+        structlog.get_logger(__name__).warning(
+            "cleanup_orphaned_runs failed at startup: %s", exc
+        )
+
     # Pre-cache the helper GGUF model for LLM-assisted dataset detection.
     # Runs in a background thread so it doesn't block server startup.
     import threading
@@ -149,6 +161,9 @@ app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
 app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
 app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
 app.include_router(export_router, prefix = "/api/export", tags = ["export"])
+app.include_router(
+    training_history_router, prefix = "/api/train", tags = ["training-history"]
+)
 
 
 # ============ Health and System Endpoints ============

studio/backend/models/__init__.py

@@ -10,6 +10,11 @@ from .training import (
     TrainingJobResponse,
     TrainingStatus,
     TrainingProgress,
+    TrainingRunSummary,
+    TrainingRunListResponse,
+    TrainingRunMetrics,
+    TrainingRunDetailResponse,
+    TrainingRunDeleteResponse,
 )
 from .models import (
     CheckpointInfo,
@@ -71,6 +76,11 @@ __all__ = [
     "TrainingJobResponse",
     "TrainingStatus",
     "TrainingProgress",
+    "TrainingRunSummary",
+    "TrainingRunListResponse",
+    "TrainingRunMetrics",
+    "TrainingRunDetailResponse",
+    "TrainingRunDeleteResponse",
     # Model management schemas
     "ModelDetails",
     "LocalModelInfo",

studio/backend/models/training.py

@@ -177,8 +177,8 @@ class TrainingProgress(BaseModel):
     job_id: str = Field(..., description = "Training job identifier")
     step: int = Field(..., description = "Current training step")
     total_steps: int = Field(..., description = "Total training steps")
-    loss: float = Field(..., description = "Current loss value")
-    learning_rate: float = Field(..., description = "Current learning rate")
+    loss: Optional[float] = Field(None, description = "Current loss value")
+    learning_rate: Optional[float] = Field(None, description = "Current learning rate")
     progress_percent: float = Field(
         ..., description = "Progress percentage (0.0 to 100.0)"
     )
@@ -196,3 +196,59 @@ class TrainingProgress(BaseModel):
     eval_loss: Optional[float] = Field(
         None, description = "Eval loss from the most recent evaluation step"
     )
+
+
+class TrainingRunSummary(BaseModel):
+    """Summary of a training run for list views."""
+
+    id: str
+    status: Literal["running", "completed", "stopped", "error"]
+    model_name: str
+    dataset_name: str
+    started_at: str
+    ended_at: Optional[str] = None
+    total_steps: Optional[int] = None
+    final_step: Optional[int] = None
+    final_loss: Optional[float] = None
+    output_dir: Optional[str] = None
+    duration_seconds: Optional[float] = None
+    error_message: Optional[str] = None
+    loss_sparkline: Optional[List[float]] = None
+
+
+class TrainingRunListResponse(BaseModel):
+    """Response for listing training runs."""
+
+    runs: List[TrainingRunSummary]
+    total: int
+
+
+class TrainingRunMetrics(BaseModel):
+    """Metrics arrays for a training run, using paired step arrays per metric."""
+
+    step_history: List[int] = Field(default_factory = list)
+    loss_history: List[float] = Field(default_factory = list)
+    loss_step_history: List[int] = Field(default_factory = list)
+    lr_history: List[float] = Field(default_factory = list)
+    lr_step_history: List[int] = Field(default_factory = list)
+    grad_norm_history: List[float] = Field(default_factory = list)
+    grad_norm_step_history: List[int] = Field(default_factory = list)
+    eval_loss_history: List[float] = Field(default_factory = list)
+    eval_step_history: List[int] = Field(default_factory = list)
+    final_epoch: Optional[float] = None
+    final_num_tokens: Optional[int] = None
+
+
+class TrainingRunDetailResponse(BaseModel):
+    """Response for a single training run with config and metrics."""
+
+    run: TrainingRunSummary
+    config: dict
+    metrics: TrainingRunMetrics
+
+
+class TrainingRunDeleteResponse(BaseModel):
+    """Response for deleting a training run."""
+
+    status: str
+    message: str

studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml

@@ -11,7 +11,7 @@ version = "0.1.0"
 description = "Local Data Designer unstructured seed reader plugin"
 requires-python = ">=3.11"
 dependencies = [
-  "data-designer-engine>=0.5.1,<0.6",
+  "data-designer-engine>=0.5.4,<0.6",
   "pandas>=2,<3",
   "pymupdf>=1.24.0",
   "pymupdf4llm>=0.0.17",

studio/backend/requirements/single-env/data-designer-deps.txt

@@ -1,8 +1,10 @@
 # Data Designer runtime deps installed explicitly (single-env mode).
-# DuckDB 1.5 removed Relation.record_batch(); keep <1.5 until upstream ships the fix.
+# Synced with data-designer-engine==0.5.4 requirements.
 anyascii<1,>=0.3.3
-duckdb<1.5,>=1.1.3
+chardet<6,>=3.0.2
+duckdb<2,>=1.5.0
 faker<21,>=20.1.0
+fsspec<2026,>=2025.3.0
 httpx<1,>=0.27.2
 httpx-retries<1,>=0.4.2
 json-repair<1,>=0.48.0
@@ -10,12 +12,13 @@ jsonpath-rust-bindings<2,>=1.0
 jsonschema<5,>=4.0.0
 lxml<7,>=6.0.2
 marko<3,>=2.1.2
+mcp<2,>=1.26.0
 networkx<4,>=3.0
-python-json-logger<4,>=3
 ruff<1,>=0.14.10
 scipy<2,>=1.11.0
 sqlfluff<4,>=3.2.0
 tiktoken<1,>=0.8.0
+# Unstructured-seed plugin deps (plugin installed with --no-deps)
 pymupdf>=1.24.0
 pymupdf4llm>=0.0.17
 mammoth>=1.8.0

studio/backend/requirements/single-env/data-designer.txt

@@ -1,5 +1,5 @@
 # Install Data Designer in same env as Unsloth.
-data-designer==0.5.2
-data-designer-config==0.5.2
-data-designer-engine==0.5.2
+data-designer==0.5.4
+data-designer-config==0.5.4
+data-designer-engine==0.5.4
 prompt-toolkit>=3,<4

studio/backend/routes/__init__.py

@@ -12,6 +12,7 @@ from routes.datasets import router as datasets_router
 from routes.auth import router as auth_router
 from routes.data_recipe import router as data_recipe_router
 from routes.export import router as export_router
+from routes.training_history import router as training_history_router
 
 __all__ = [
     "training_router",
@@ -21,4 +22,5 @@ __all__ = [
     "auth_router",
     "data_recipe_router",
     "export_router",
+    "training_history_router",
 ]

studio/backend/routes/training.py

@@ -14,6 +14,7 @@ import structlog
 from loggers import get_logger
 import asyncio
 from datetime import datetime
+import uuid as _uuid
 
 # Add backend directory to path
 # The backend code should be in the same directory structure
@@ -115,15 +116,11 @@ async def start_training(
 
         backend = get_training_backend()
 
-        # Generate job ID and attach to backend for later status/progress calls
-        job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
-        backend.current_job_id = job_id
-
-        # Check if training is already active
+        # Check if training is already active (before mutating any state)
         if backend.is_training_active():
             existing_job_id: Optional[str] = getattr(backend, "current_job_id", "")
             return TrainingJobResponse(
-                job_id = existing_job_id or job_id,
+                job_id = existing_job_id or "",
                 status = "error",
                 message = (
                     "Training is already in progress. "
@@ -132,6 +129,12 @@ async def start_training(
                 error = "Training already active",
             )
 
+        # Generate job ID — passed into start_training() which sets it on the
+        # backend only after confirming the old pump thread is dead.
+        job_id = (
+            f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{_uuid.uuid4().hex[:8]}"
+        )
+
         # Validate dataset paths if provided
         if request.local_datasets:
             request.local_datasets = _validate_local_dataset_paths(
@@ -248,12 +251,12 @@ async def start_training(
             logger.warning("Could not shut down export subprocess: %s", e)
 
         # start_training now spawns a subprocess (non-blocking)
-        success = backend.start_training(**training_kwargs)
+        success = backend.start_training(job_id = job_id, **training_kwargs)
 
         if not success:
             progress_error = backend.trainer.training_progress.error
             return TrainingJobResponse(
-                job_id = job_id,
+                job_id = backend.current_job_id or "",
                 status = "error",
                 message = progress_error or "Failed to start training subprocess",
                 error = progress_error or "subprocess_start_failed",
@@ -345,7 +348,7 @@ async def reset_training(
             error = None,
             status_message = "Ready to train",
             step = 0,
-            loss = 0.0,
+            loss = None,
             epoch = 0,
             total_steps = 0,
         )
@@ -419,8 +422,8 @@ async def get_training_status(
                 "epoch": getattr(progress, "epoch", 0),
                 "step": getattr(progress, "step", 0),
                 "total_steps": getattr(progress, "total_steps", 0),
-                "loss": getattr(progress, "loss", 0.0),
-                "learning_rate": getattr(progress, "learning_rate", 0.0),
+                "loss": getattr(progress, "loss", None),
+                "learning_rate": getattr(progress, "learning_rate", None),
             }
 
         # Build metric history for chart recovery after SSE reconnection
@@ -526,8 +529,8 @@ async def stream_training_progress(
         # ── Helpers ──────────────────────────────────────────────
         def build_progress(
             step: int,
-            loss: float,
-            learning_rate: float,
+            loss: Optional[float],
+            learning_rate: Optional[float],
             total_steps: int,
             epoch: Optional[float] = None,
             progress: Optional[Any] = None,
@@ -604,10 +607,10 @@ async def stream_training_progress(
                     loss_val = (
                         backend.loss_history[i]
                         if i < len(backend.loss_history)
-                        else 0.0
+                        else None
                     )
                     lr_val = (
-                        backend.lr_history[i] if i < len(backend.lr_history) else 0.0
+                        backend.lr_history[i] if i < len(backend.lr_history) else None
                     )
                     tp_replay = getattr(
                         getattr(backend, "trainer", None), "training_progress", None
@@ -645,8 +648,8 @@ async def stream_training_progress(
 
             initial_progress = build_progress(
                 step = 0,
-                loss = 0.0,
-                learning_rate = 0.0,
+                loss = None,
+                learning_rate = None,
                 total_steps = initial_total_steps,
                 epoch = initial_epoch,
                 progress = tp,
@@ -660,9 +663,9 @@ async def stream_training_progress(
                 if backend.step_history:
                     final_step = backend.step_history[-1]
                     final_loss = (
-                        backend.loss_history[-1] if backend.loss_history else 0.0
+                        backend.loss_history[-1] if backend.loss_history else None
                     )
-                    final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
+                    final_lr = backend.lr_history[-1] if backend.lr_history else None
                     final_total_steps = (
                         getattr(tp, "total_steps", final_step) if tp else final_step
                     )
@@ -680,7 +683,9 @@ async def stream_training_progress(
                     )
                 else:
                     yield format_sse(
-                        build_progress(-1, 0.0, 0.0, 0, progress = tp).model_dump_json(),
+                        build_progress(
+                            -1, None, None, 0, progress = tp
+                        ).model_dump_json(),
                         event = "complete",
                         event_id = 0,
                     )
@@ -698,9 +703,9 @@ async def stream_training_progress(
                 if backend.step_history:
                     current_step = backend.step_history[-1]
                     current_loss = (
-                        backend.loss_history[-1] if backend.loss_history else 0.0
+                        backend.loss_history[-1] if backend.loss_history else None
                     )
-                    current_lr = backend.lr_history[-1] if backend.lr_history else 0.0
+                    current_lr = backend.lr_history[-1] if backend.lr_history else None
                     tp_inner = getattr(
                         getattr(backend, "trainer", None), "training_progress", None
                     )
@@ -763,8 +768,8 @@ async def stream_training_progress(
                         )
                         preparing_payload = build_progress(
                             0,
-                            0.0,
-                            0.0,
+                            None,
+                            None,
                             prep_total,
                             progress = tp_prep,
                         )
@@ -781,7 +786,7 @@ async def stream_training_progress(
                         getattr(backend, "trainer", None), "training_progress", None
                     )
                     timeout_payload = build_progress(
-                        last_step, 0.0, 0.0, 0, progress = tp_timeout
+                        last_step, None, None, 0, progress = tp_timeout
                     )
                     yield format_sse(
                         timeout_payload.model_dump_json(),
@@ -797,7 +802,7 @@ async def stream_training_progress(
                 tp_error = getattr(
                     getattr(backend, "trainer", None), "training_progress", None
                 )
-                error_payload = build_progress(0, 0.0, 0.0, 0, progress = tp_error)
+                error_payload = build_progress(0, None, None, 0, progress = tp_error)
                 yield format_sse(
                     error_payload.model_dump_json(),
                     event = "error",
@@ -807,8 +812,8 @@ async def stream_training_progress(
 
         # ── Final "complete" event ───────────────────────────────
         final_step = backend.step_history[-1] if backend.step_history else last_step
-        final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
-        final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
+        final_loss = backend.loss_history[-1] if backend.loss_history else None
+        final_lr = backend.lr_history[-1] if backend.lr_history else None
         final_tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
         final_total_steps = (
             getattr(final_tp, "total_steps", final_step) if final_tp else final_step

studio/backend/routes/training_history.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""
+Training history API routes — browse, view, and delete past training runs.
+"""
+
+import json
+
+from fastapi import APIRouter, Depends, HTTPException, Query
+from loggers import get_logger
+
+from auth.authentication import get_current_subject
+from models import (
+    TrainingRunDeleteResponse,
+    TrainingRunDetailResponse,
+    TrainingRunListResponse,
+    TrainingRunMetrics,
+    TrainingRunSummary,
+)
+from storage.studio_db import delete_run, get_run, get_run_metrics, list_runs
+
+logger = get_logger(__name__)
+
+router = APIRouter()
+
+
+@router.get("/runs", response_model = TrainingRunListResponse)
+async def list_training_runs(
+    limit: int = Query(50, ge = 1, le = 200),
+    offset: int = Query(0, ge = 0),
+    current_subject: str = Depends(get_current_subject),
+):
+    """List training runs, newest first."""
+    result = list_runs(limit = limit, offset = offset)
+    return TrainingRunListResponse(
+        runs = [TrainingRunSummary(**r) for r in result["runs"]],
+        total = result["total"],
+    )
+
+
+@router.get("/runs/{run_id}", response_model = TrainingRunDetailResponse)
+async def get_training_run_detail(
+    run_id: str,
+    current_subject: str = Depends(get_current_subject),
+):
+    """Get a single training run with full config and metrics."""
+    run = get_run(run_id)
+    if run is None:
+        raise HTTPException(status_code = 404, detail = f"Run {run_id} not found")
+
+    try:
+        config = json.loads(run.get("config_json", "{}"))
+    except (json.JSONDecodeError, TypeError):
+        logger.debug("Failed to parse config_json for run %s", run_id)
+        config = {}
+
+    metrics_data = get_run_metrics(run_id)
+
+    return TrainingRunDetailResponse(
+        run = TrainingRunSummary(**{k: v for k, v in run.items() if k != "config_json"}),
+        config = config,
+        metrics = TrainingRunMetrics(**metrics_data),
+    )
+
+
+@router.delete("/runs/{run_id}", response_model = TrainingRunDeleteResponse)
+async def delete_training_run(
+    run_id: str,
+    current_subject: str = Depends(get_current_subject),
+):
+    """Delete a training run and its metrics (CASCADE)."""
+    run = get_run(run_id)
+    if run is None:
+        raise HTTPException(status_code = 404, detail = f"Run {run_id} not found")
+    if run["status"] == "running":
+        raise HTTPException(
+            status_code = 409, detail = "Cannot delete a running training run"
+        )
+    logger.info("Deleting training run %s", run_id)
+    delete_run(run_id)
+    return TrainingRunDeleteResponse(
+        status = "deleted",
+        message = f"Run {run_id} deleted",
+    )

studio/backend/storage/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

studio/backend/storage/studio_db.py

@@ -0,0 +1,362 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""
+SQLite storage for training run history and metrics.
+
+Follows the same pattern as auth/storage.py — module-level functions,
+raw sqlite3, per-function connections. Enhancements over auth:
+  - WAL mode for concurrent read/write access
+  - PRAGMA foreign_keys = ON for CASCADE deletes
+"""
+
+import json
+import logging
+import sqlite3
+import threading
+
+logger = logging.getLogger(__name__)
+from typing import Optional
+
+from utils.paths import studio_db_path, ensure_dir
+
+_schema_lock = threading.Lock()
+_schema_ready = False
+
+
+def _ensure_schema(conn: sqlite3.Connection) -> None:
+    """Create tables and indexes if they don't exist. Called once per process."""
+    conn.execute("PRAGMA journal_mode=WAL")
+    conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS training_runs (
+            id TEXT NOT NULL PRIMARY KEY,
+            status TEXT NOT NULL DEFAULT 'running',
+            model_name TEXT NOT NULL,
+            dataset_name TEXT NOT NULL,
+            config_json TEXT NOT NULL,
+            started_at TEXT NOT NULL,
+            ended_at TEXT,
+            total_steps INTEGER,
+            final_step INTEGER,
+            final_loss REAL,
+            output_dir TEXT,
+            error_message TEXT,
+            duration_seconds REAL,
+            loss_sparkline TEXT
+        )
+        """
+    )
+    conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS training_metrics (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            run_id TEXT NOT NULL REFERENCES training_runs(id) ON DELETE CASCADE,
+            step INTEGER NOT NULL,
+            loss REAL,
+            learning_rate REAL,
+            grad_norm REAL,
+            eval_loss REAL,
+            epoch REAL,
+            num_tokens INTEGER,
+            elapsed_seconds REAL,
+            UNIQUE(run_id, step)
+        )
+        """
+    )
+    conn.execute(
+        "CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)"
+    )
+
+
+def get_connection() -> sqlite3.Connection:
+    """Open studio.db with WAL mode, create tables once per process, enable foreign keys."""
+    global _schema_ready
+    db_path = studio_db_path()
+    ensure_dir(db_path.parent)
+    conn = sqlite3.connect(str(db_path))
+    conn.row_factory = sqlite3.Row
+    # foreign_keys is session-scoped, must be set per connection
+    conn.execute("PRAGMA foreign_keys=ON")
+    if not _schema_ready:
+        with _schema_lock:
+            if not _schema_ready:
+                try:
+                    _ensure_schema(conn)
+                    _schema_ready = True
+                except Exception:
+                    conn.close()
+                    raise
+    return conn
+
+
+def create_run(
+    id: str,
+    model_name: str,
+    dataset_name: str,
+    config_json: str,
+    started_at: str,
+    total_steps: Optional[int],
+) -> None:
+    conn = get_connection()
+    try:
+        conn.execute(
+            """
+            INSERT INTO training_runs (id, model_name, dataset_name, config_json, started_at, total_steps)
+            VALUES (?, ?, ?, ?, ?, ?)
+            """,
+            (id, model_name, dataset_name, config_json, started_at, total_steps),
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def update_run_total_steps(id: str, total_steps: int) -> None:
+    conn = get_connection()
+    try:
+        conn.execute(
+            "UPDATE training_runs SET total_steps = ? WHERE id = ?",
+            (total_steps, id),
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def update_run_progress(
+    id: str, step: int, loss: Optional[float], duration_seconds: Optional[float]
+) -> None:
+    """Update current progress on a running training run (called on each metric flush)."""
+    conn = get_connection()
+    try:
+        conn.execute(
+            "UPDATE training_runs SET final_step = ?, final_loss = ?, duration_seconds = ? WHERE id = ?",
+            (step, loss, duration_seconds, id),
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def finish_run(
+    id: str,
+    status: str,
+    ended_at: str,
+    final_step: Optional[int],
+    final_loss: Optional[float],
+    duration_seconds: Optional[float],
+    loss_sparkline: Optional[str] = None,
+    output_dir: Optional[str] = None,
+    error_message: Optional[str] = None,
+) -> None:
+    conn = get_connection()
+    try:
+        conn.execute(
+            """
+            UPDATE training_runs
+            SET status = ?, ended_at = ?, final_step = ?, final_loss = ?,
+                duration_seconds = ?, loss_sparkline = ?, output_dir = ?,
+                error_message = ?
+            WHERE id = ?
+            """,
+            (
+                status,
+                ended_at,
+                final_step,
+                final_loss,
+                duration_seconds,
+                loss_sparkline,
+                output_dir,
+                error_message,
+                id,
+            ),
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def insert_metrics_batch(run_id: str, metrics: list[dict]) -> None:
+    if not metrics:
+        return
+    conn = get_connection()
+    try:
+        conn.executemany(
+            """
+            INSERT INTO training_metrics
+                (run_id, step, loss, learning_rate, grad_norm, eval_loss, epoch, num_tokens, elapsed_seconds)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+            ON CONFLICT(run_id, step) DO UPDATE SET
+                loss = COALESCE(excluded.loss, loss),
+                learning_rate = COALESCE(excluded.learning_rate, learning_rate),
+                grad_norm = COALESCE(excluded.grad_norm, grad_norm),
+                eval_loss = COALESCE(excluded.eval_loss, eval_loss),
+                epoch = COALESCE(excluded.epoch, epoch),
+                num_tokens = COALESCE(excluded.num_tokens, num_tokens),
+                elapsed_seconds = COALESCE(excluded.elapsed_seconds, elapsed_seconds)
+            """,
+            [
+                (
+                    run_id,
+                    m.get("step"),
+                    m.get("loss"),
+                    m.get("learning_rate"),
+                    m.get("grad_norm"),
+                    m.get("eval_loss"),
+                    m.get("epoch"),
+                    m.get("num_tokens"),
+                    m.get("elapsed_seconds"),
+                )
+                for m in metrics
+            ],
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def list_runs(limit: int = 50, offset: int = 0) -> dict:
+    conn = get_connection()
+    try:
+        total = conn.execute("SELECT COUNT(*) FROM training_runs").fetchone()[0]
+        rows = conn.execute(
+            """
+            SELECT id, status, model_name, dataset_name, started_at, ended_at,
+                   total_steps, final_step, final_loss, output_dir,
+                   duration_seconds, error_message, loss_sparkline
+            FROM training_runs
+            ORDER BY started_at DESC
+            LIMIT ? OFFSET ?
+            """,
+            (limit, offset),
+        ).fetchall()
+        runs = []
+        for row in rows:
+            run = dict(row)
+            sparkline = run.get("loss_sparkline")
+            if sparkline:
+                try:
+                    run["loss_sparkline"] = json.loads(sparkline)
+                except (json.JSONDecodeError, TypeError):
+                    logger.debug(
+                        "Failed to parse loss_sparkline for run %s", run.get("id")
+                    )
+                    run["loss_sparkline"] = None
+            runs.append(run)
+        return {"runs": runs, "total": total}
+    finally:
+        conn.close()
+
+
+def get_run(id: str) -> Optional[dict]:
+    conn = get_connection()
+    try:
+        row = conn.execute("SELECT * FROM training_runs WHERE id = ?", (id,)).fetchone()
+        if row is None:
+            return None
+        run = dict(row)
+        sparkline = run.get("loss_sparkline")
+        if sparkline:
+            try:
+                run["loss_sparkline"] = json.loads(sparkline)
+            except (json.JSONDecodeError, TypeError):
+                logger.debug("Failed to parse loss_sparkline for run %s", id)
+                run["loss_sparkline"] = None
+        return run
+    finally:
+        conn.close()
+
+
+def get_run_metrics(id: str) -> dict:
+    """Return metric arrays for a run, using paired step arrays per metric."""
+    conn = get_connection()
+    try:
+        rows = conn.execute(
+            """
+            SELECT step, loss, learning_rate, grad_norm, eval_loss, epoch,
+                   num_tokens, elapsed_seconds
+            FROM training_metrics
+            WHERE run_id = ?
+            ORDER BY step
+            """,
+            (id,),
+        ).fetchall()
+
+        step_history: list[int] = []
+        loss_history: list[float] = []
+        loss_step_history: list[int] = []
+        lr_history: list[float] = []
+        lr_step_history: list[int] = []
+        grad_norm_history: list[float] = []
+        grad_norm_step_history: list[int] = []
+        eval_loss_history: list[float] = []
+        eval_step_history: list[int] = []
+        final_epoch: float | None = None
+        final_num_tokens: int | None = None
+
+        for row in rows:
+            step = row["step"]
+            step_history.append(step)
+            if step > 0 and row["loss"] is not None:
+                loss_history.append(row["loss"])
+                loss_step_history.append(step)
+            if step > 0 and row["learning_rate"] is not None:
+                lr_history.append(row["learning_rate"])
+                lr_step_history.append(step)
+            if step > 0 and row["grad_norm"] is not None:
+                grad_norm_history.append(row["grad_norm"])
+                grad_norm_step_history.append(step)
+            if step > 0 and row["eval_loss"] is not None:
+                eval_loss_history.append(row["eval_loss"])
+                eval_step_history.append(step)
+            if row["epoch"] is not None:
+                final_epoch = row["epoch"]
+            if row["num_tokens"] is not None:
+                final_num_tokens = row["num_tokens"]
+
+        return {
+            "step_history": step_history,
+            "loss_history": loss_history,
+            "loss_step_history": loss_step_history,
+            "lr_history": lr_history,
+            "lr_step_history": lr_step_history,
+            "grad_norm_history": grad_norm_history,
+            "grad_norm_step_history": grad_norm_step_history,
+            "eval_loss_history": eval_loss_history,
+            "eval_step_history": eval_step_history,
+            "final_epoch": final_epoch,
+            "final_num_tokens": final_num_tokens,
+        }
+    finally:
+        conn.close()
+
+
+def delete_run(id: str) -> None:
+    conn = get_connection()
+    try:
+        conn.execute("DELETE FROM training_runs WHERE id = ?", (id,))
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def cleanup_orphaned_runs() -> None:
+    """Mark any 'running' rows as errored on startup (server restarted mid-training)."""
+    from datetime import datetime, timezone
+
+    conn = get_connection()
+    try:
+        conn.execute(
+            """
+            UPDATE training_runs
+            SET status = 'error',
+                error_message = 'Server restarted during training',
+                ended_at = ?
+            WHERE status = 'running'
+            """,
+            (datetime.now(timezone.utc).isoformat(),),
+        )
+        conn.commit()
+    finally:
+        conn.close()

studio/backend/utils/downsample.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Generic numeric downsampling utility."""
+
+
+def downsample(values: list[float], target_count: int) -> list[float]:
+    """Reduce a list to target_count points via evenly-spaced index sampling."""
+    if len(values) <= target_count:
+        return list(values)
+    if target_count <= 0:
+        return []
+    if target_count == 1:
+        return [values[-1]]
+    indices = [
+        round(i * (len(values) - 1) / (target_count - 1)) for i in range(target_count)
+    ]
+    return [values[i] for i in indices]

studio/backend/utils/paths/__init__.py

@@ -16,6 +16,7 @@ from .storage_roots import (
     exports_root,
     auth_root,
     auth_db_path,
+    studio_db_path,
     tmp_root,
     seed_uploads_root,
     unstructured_seed_cache_root,
@@ -45,6 +46,7 @@ __all__ = [
     "exports_root",
     "auth_root",
     "auth_db_path",
+    "studio_db_path",
     "tmp_root",
     "seed_uploads_root",
     "unstructured_seed_cache_root",

studio/backend/utils/paths/storage_roots.py

@@ -49,6 +49,10 @@ def auth_db_path() -> Path:
     return auth_root() / "auth.db"
 
 
+def studio_db_path() -> Path:
+    return studio_root() / "studio.db"
+
+
 def tmp_root() -> Path:
     return Path(tempfile.gettempdir()) / "unsloth-studio"
 

studio/frontend/.gitignore

@@ -11,6 +11,7 @@ pnpm-debug.log*
 lerna-debug.log*
 
 node_modules
+bun.lock
 dist
 dist-ssr
 test/

studio/frontend/package.json

@@ -3,6 +3,9 @@
   "private": true,
   "version": "0.0.0",
   "type": "module",
+  "engines": {
+    "node": "^20.19.0 || >=22.12.0"
+  },
   "scripts": {
     "dev": "vite",
     "build": "tsc -b && vite build",
@@ -35,7 +38,7 @@
     "@streamdown/code": "1.0.2",
     "@streamdown/math": "1.0.2",
     "@streamdown/mermaid": "1.0.2",
-    "@tailwindcss/vite": "^4.1.18",
+    "@tailwindcss/vite": "^4.2.2",
     "@tanstack/react-router": "^1.159.10",
     "@tanstack/react-table": "^8.21.3",
     "@toolwind/corner-shape": "^0.0.8-3",
@@ -48,7 +51,6 @@
     "cmdk": "^1.1.1",
     "date-fns": "^4.1.0",
     "dexie": "^4.3.0",
-    "framer-motion": "^11.18.2",
     "js-yaml": "^4.1.1",
     "katex": "^0.16.28",
     "lucide-react": "^0.577.0",
@@ -80,13 +82,13 @@
     "@types/node": "^24.10.1",
     "@types/react": "^19.2.5",
     "@types/react-dom": "^19.2.3",
-    "@vitejs/plugin-react": "^5.1.1",
+    "@vitejs/plugin-react": "^6.0.1",
     "eslint": "^9.39.1",
     "eslint-plugin-react-hooks": "^7.0.1",
     "eslint-plugin-react-refresh": "^0.4.26",
     "globals": "^16.5.0",
     "typescript": "~5.9.3",
     "typescript-eslint": "^8.55.0",
-    "vite": "^7.3.1"
+    "vite": "^8.0.1"
   }
 }

studio/frontend/src/components/assistant-ui/markdown-text.tsx

@@ -4,19 +4,39 @@
 "use client";
 
 import { copyToClipboard } from "@/lib/copy-to-clipboard";
+import { preprocessLaTeX } from "@/lib/latex";
 import { INTERNAL, useMessagePartText } from "@assistant-ui/react";
 import { Copy02Icon, Tick02Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { code } from "@streamdown/code";
-import { math } from "@streamdown/math";
+import { createMathPlugin } from "@streamdown/math";
 import { mermaid } from "@streamdown/mermaid";
 import { DownloadIcon, Maximize2Icon, Minimize2Icon } from "lucide-react";
-import { useEffect, useRef, useState } from "react";
+import { useEffect, useMemo, useRef, useState } from "react";
 import { Block, type BlockProps, Streamdown } from "streamdown";
 import "katex/dist/katex.min.css";
 import { AudioPlayer } from "./audio-player";
 
+const math = createMathPlugin({ singleDollarTextMath: true });
 const { withSmoothContextProvider } = INTERNAL;
+
+const STREAMDOWN_COMPONENTS = {
+  a: ({
+    href,
+    children,
+    ...props
+  }: React.ComponentProps<"a">) => (
+    <a
+      href={href}
+      target="_blank"
+      rel="noopener noreferrer"
+      className="text-primary underline underline-offset-2 decoration-primary/40 hover:decoration-primary transition-colors"
+      {...props}
+    >
+      {children}
+    </a>
+  ),
+};
 const COPY_RESET_MS = 2000;
 const MERMAID_SOURCE_RE = /```mermaid\s*([\s\S]*?)```/i;
 const CODE_FENCE_RE = /^```([^\r\n`]*)\r?\n([\s\S]*?)\r?\n?```$/;
@@ -375,6 +395,7 @@ const AUDIO_PLAYER_RE = /<audio-player\s+src="([^"]+)"\s*\/>/;
 
 const MarkdownTextImpl = () => {
   const { text, status } = useMessagePartText();
+  const processedText = useMemo(() => preprocessLaTeX(text), [text]);
 
   const audioMatch = text.match(AUDIO_PLAYER_RE);
   if (audioMatch) {
@@ -387,6 +408,7 @@ const MarkdownTextImpl = () => {
         mode="streaming"
         isAnimating={status.type === "running"}
         plugins={{ code, math, mermaid }}
+        components={STREAMDOWN_COMPONENTS}
         controls={{
           code: false,
           mermaid: {
@@ -399,7 +421,7 @@ const MarkdownTextImpl = () => {
         shikiTheme={["github-light", "github-dark"]}
         BlockComponent={StreamdownBlock}
       >
-        {text}
+        {processedText}
       </Streamdown>
     </div>
   );

studio/frontend/src/components/assistant-ui/reasoning.tsx

@@ -17,7 +17,6 @@ import {
   type ReasoningGroupComponent,
   type ReasoningMessagePartComponent,
   useAuiState,
-  useScrollLock,
 } from "@assistant-ui/react";
 import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import { Idea01Icon } from "@hugeicons/core-free-icons";
@@ -34,6 +33,7 @@ import {
   useState,
 } from "react";
 const ANIMATION_DURATION = 200;
+const AUTO_SCROLL_THRESHOLD_PX = 24;
 
 export const reasoningVariants = cva("aui-reasoning-root mb-4 w-full", {
   variants: {
@@ -68,8 +68,49 @@ function ReasoningRoot({
   ...props
 }: ReasoningRootProps) {
   const collapsibleRef = useRef<HTMLDivElement>(null);
+  const lockCleanupRef = useRef<(() => void) | null>(null);
   const [uncontrolledOpen, setUncontrolledOpen] = useState(defaultOpen);
-  const lockScroll = useScrollLock(collapsibleRef, ANIMATION_DURATION);
+
+  useEffect(() => {
+    return () => {
+      lockCleanupRef.current?.();
+    };
+  }, []);
+
+  const lockScroll = useCallback(() => {
+    lockCleanupRef.current?.();
+
+    const animatedElement = collapsibleRef.current;
+    if (!animatedElement) return;
+
+    let scrollContainer: HTMLElement | null = animatedElement;
+    while (scrollContainer) {
+      const { overflowY } = getComputedStyle(scrollContainer);
+      if (overflowY === "scroll" || overflowY === "auto") {
+        break;
+      }
+      scrollContainer = scrollContainer.parentElement;
+    }
+    if (!scrollContainer) return;
+
+    const scrollPosition = scrollContainer.scrollTop;
+    const resetPosition = () => {
+      scrollContainer.scrollTop = scrollPosition;
+    };
+
+    scrollContainer.addEventListener("scroll", resetPosition);
+    let timeoutId: ReturnType<typeof setTimeout> | null = null;
+    const cleanup = () => {
+      if (timeoutId !== null) {
+        clearTimeout(timeoutId);
+        timeoutId = null;
+      }
+      scrollContainer.removeEventListener("scroll", resetPosition);
+      lockCleanupRef.current = null;
+    };
+    timeoutId = setTimeout(cleanup, ANIMATION_DURATION);
+    lockCleanupRef.current = cleanup;
+  }, []);
 
   const isControlled = controlledOpen !== undefined;
   const isOpen = isControlled ? controlledOpen : uncontrolledOpen;
@@ -220,6 +261,8 @@ function ReasoningText({
 }: ComponentProps<"div"> & { streaming?: boolean }) {
   const scrollRef = useRef<HTMLDivElement>(null);
   const shouldAutoScrollRef = useRef(true);
+  const detachedFromBottomRef = useRef(false);
+  const lastScrollTopRef = useRef(0);
 
   useEffect(() => {
     if (!(streaming && scrollRef.current)) {
@@ -227,8 +270,25 @@ function ReasoningText({
     }
     const el = scrollRef.current;
     const updateAutoScroll = () => {
+      const currentScrollTop = el.scrollTop;
+      if (currentScrollTop < lastScrollTopRef.current) {
+        detachedFromBottomRef.current = true;
+      }
       const distanceFromBottom = el.scrollHeight - el.scrollTop - el.clientHeight;
-      shouldAutoScrollRef.current = distanceFromBottom <= 24;
+      if (
+        detachedFromBottomRef.current &&
+        distanceFromBottom <= AUTO_SCROLL_THRESHOLD_PX
+      ) {
+        detachedFromBottomRef.current = false;
+      }
+      shouldAutoScrollRef.current = !detachedFromBottomRef.current;
+      lastScrollTopRef.current = currentScrollTop;
+    };
+    const handleWheel = (event: WheelEvent) => {
+      if (event.deltaY < 0) {
+        detachedFromBottomRef.current = true;
+        shouldAutoScrollRef.current = false;
+      }
     };
     const observer = new MutationObserver(() => {
       if (shouldAutoScrollRef.current) {
@@ -236,16 +296,19 @@ function ReasoningText({
       }
     });
     el.addEventListener("scroll", updateAutoScroll);
+    el.addEventListener("wheel", handleWheel, { passive: true });
     observer.observe(el, {
       childList: true,
       subtree: true,
       characterData: true,
     });
-    shouldAutoScrollRef.current = true;
-    el.scrollTop = el.scrollHeight;
+    lastScrollTopRef.current = el.scrollTop;
+    detachedFromBottomRef.current = false;
+    updateAutoScroll();
     return () => {
       observer.disconnect();
       el.removeEventListener("scroll", updateAutoScroll);
+      el.removeEventListener("wheel", handleWheel);
     };
   }, [streaming]);
 

studio/frontend/src/components/assistant-ui/sources.tsx

@@ -43,7 +43,8 @@ function SourceIcon({
 }: ComponentProps<"span"> & { url: string; size?: number }) {
   const [hasError, setHasError] = useState(false);
   const domain = extractDomain(url);
-  const sizeClass = `size-${size}`;
+  const SIZE_CLASSES: Record<number, string> = { 3: "size-3", 4: "size-4", 5: "size-5" };
+  const sizeClass = SIZE_CLASSES[size] ?? "size-3";
 
   if (hasError) {
     return (

studio/frontend/src/components/assistant-ui/thread.tsx

@@ -16,6 +16,7 @@ import { WebSearchToolUI } from "@/components/assistant-ui/tool-ui-web-search";
 import { PythonToolUI } from "@/components/assistant-ui/tool-ui-python";
 import { TerminalToolUI } from "@/components/assistant-ui/tool-ui-terminal";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import { AnimatedShinyText } from "@/components/ui/animated-shiny-text";
 import { Button } from "@/components/ui/button";
 import { sentAudioNames } from "@/features/chat/api/chat-adapter";
 import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils";
@@ -35,7 +36,7 @@ import {
   useAuiEvent,
   useAuiState,
 } from "@assistant-ui/react";
-import { motion } from "framer-motion";
+import { motion } from "motion/react";
 import {
   ArrowDownIcon,
   ArrowUpIcon,
@@ -88,9 +89,8 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
           }}
         />
 
-        <ThreadPrimitive.ViewportFooter className="aui-thread-viewport-footer sticky bottom-0 mt-auto flex w-full flex-col gap-4 overflow-visible bg-background pb-4 md:pb-4">
+        <ThreadPrimitive.ViewportFooter className="aui-thread-viewport-footer sticky bottom-0 z-20 mt-auto flex w-full flex-col gap-4 overflow-visible bg-background pb-4 md:pb-4">
           <ThreadScrollToBottom />
-          <GeneratingSpinner />
           <AuiIf condition={({ thread }) => !thread.isEmpty}>
             {!hideComposer && <ComposerAnimated />}
           </AuiIf>
@@ -541,6 +541,17 @@ const MessageError: FC = () => {
   );
 };
 
+const GeneratingIndicator: FC = () => {
+  const show = useAuiState(
+    ({ message }) =>
+      message.content.length === 0 && message.status?.type === "running",
+  );
+  if (!show) return null;
+  return (
+    <AnimatedShinyText className="text-sm">Generating...</AnimatedShinyText>
+  );
+};
+
 const AssistantMessage: FC = () => {
   return (
     <MessagePrimitive.Root
@@ -548,6 +559,7 @@ const AssistantMessage: FC = () => {
       data-role="assistant"
     >
       <div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
+        <GeneratingIndicator />
         <MessagePrimitive.Parts
           components={{
             Text: MarkdownText,

studio/frontend/src/components/assistant-ui/tool-fallback.tsx

@@ -77,7 +77,7 @@ function ToolFallbackRoot({
       open={isOpen}
       onOpenChange={handleOpenChange}
       className={cn(
-        "aui-tool-fallback-root group/tool-fallback-root w-full corner-squircle rounded-lg border py-3",
+        "aui-tool-fallback-root group/tool-fallback-root w-full py-1",
         className,
       )}
       style={
@@ -124,7 +124,7 @@ function ToolFallbackTrigger({
     <CollapsibleTrigger
       data-slot="tool-fallback-trigger"
       className={cn(
-        "aui-tool-fallback-trigger group/trigger flex w-full items-center gap-2 px-4 text-sm transition-colors",
+        "aui-tool-fallback-trigger group/trigger flex w-full items-center gap-2 py-1.5 text-sm transition-colors",
         className,
       )}
       {...props}
@@ -207,7 +207,7 @@ function ToolFallbackContent({
       )}
       {...props}
     >
-      <div className="mt-3 flex flex-col gap-2 border-t pt-2">{children}</div>
+      <div className="mt-1 flex flex-col gap-2 pl-5">{children}</div>
     </CollapsibleContent>
   );
 }
@@ -226,7 +226,7 @@ function ToolFallbackArgs({
   return (
     <div
       data-slot="tool-fallback-args"
-      className={cn("aui-tool-fallback-args px-4", className)}
+      className={cn("aui-tool-fallback-args", className)}
       {...props}
     >
       <pre className="aui-tool-fallback-args-value whitespace-pre-wrap">
@@ -251,7 +251,7 @@ function ToolFallbackResult({
     <div
       data-slot="tool-fallback-result"
       className={cn(
-        "aui-tool-fallback-result border-t border-dashed px-4 pt-2",
+        "aui-tool-fallback-result pt-2",
         className,
       )}
       {...props}
@@ -292,7 +292,7 @@ function ToolFallbackError({
   return (
     <div
       data-slot="tool-fallback-error"
-      className={cn("aui-tool-fallback-error px-4", className)}
+      className={cn("aui-tool-fallback-error", className)}
       {...props}
     >
       <p className="aui-tool-fallback-error-header font-semibold text-muted-foreground">
@@ -316,7 +316,7 @@ const ToolFallbackImpl: ToolCallMessagePartComponent = ({
 
   return (
     <ToolFallbackRoot
-      className={cn(isCancelled && "border-muted-foreground/30 bg-muted/30")}
+      className={cn(isCancelled && "bg-muted/30")}
     >
       <ToolFallbackTrigger toolName={toolName} status={status} />
       <ToolFallbackContent>

studio/frontend/src/components/assistant-ui/tool-group.tsx

@@ -26,11 +26,11 @@ const toolGroupVariants = cva("aui-tool-group-root group/tool-group w-full", {
   variants: {
     variant: {
       outline: "corner-squircle rounded-lg border py-3",
-      ghost: "",
+      ghost: "rounded-lg bg-muted/10 py-2",
       muted: "corner-squircle rounded-lg border border-muted-foreground/30 bg-muted/30 py-3",
     },
   },
-  defaultVariants: { variant: "outline" },
+  defaultVariants: { variant: "ghost" },
 });
 
 export type ToolGroupRootProps = Omit<
@@ -76,7 +76,7 @@ function ToolGroupRoot({
     <Collapsible
       ref={collapsibleRef}
       data-slot="tool-group-root"
-      data-variant={variant ?? "outline"}
+      data-variant={variant ?? "ghost"}
       open={isOpen}
       onOpenChange={handleOpenChange}
       className={cn(
@@ -111,9 +111,10 @@ function ToolGroupTrigger({
     <CollapsibleTrigger
       data-slot="tool-group-trigger"
       className={cn(
-        "aui-tool-group-trigger group/trigger flex items-center gap-2 text-sm transition-colors",
-        "group-data-[variant=outline]/tool-group-root:w-full group-data-[variant=outline]/tool-group-root:px-4",
-        "group-data-[variant=muted]/tool-group-root:w-full group-data-[variant=muted]/tool-group-root:px-4",
+        "aui-tool-group-trigger group/trigger flex w-full items-center gap-2 text-sm transition-colors",
+        "group-data-[variant=outline]/tool-group-root:px-4",
+        "group-data-[variant=muted]/tool-group-root:px-4",
+        "group-data-[variant=ghost]/tool-group-root:px-0",
         className,
       )}
       {...props}
@@ -134,9 +135,7 @@ function ToolGroupTrigger({
       <span
         data-slot="tool-group-trigger-label"
         className={cn(
-          "aui-tool-group-trigger-label-wrapper relative inline-block text-left font-medium leading-none",
-          "group-data-[variant=outline]/tool-group-root:grow",
-          "group-data-[variant=muted]/tool-group-root:grow",
+          "aui-tool-group-trigger-label-wrapper relative inline-block grow text-left font-medium leading-none",
         )}
       >
         <span>{label}</span>
@@ -189,6 +188,7 @@ function ToolGroupContent({
           "mt-2 flex flex-col gap-2",
           "group-data-[variant=outline]/tool-group-root:mt-3 group-data-[variant=outline]/tool-group-root:border-t group-data-[variant=outline]/tool-group-root:px-4 group-data-[variant=outline]/tool-group-root:pt-3",
           "group-data-[variant=muted]/tool-group-root:mt-3 group-data-[variant=muted]/tool-group-root:border-t group-data-[variant=muted]/tool-group-root:px-4 group-data-[variant=muted]/tool-group-root:pt-3",
+          "group-data-[variant=ghost]/tool-group-root:mt-1 group-data-[variant=ghost]/tool-group-root:gap-1",
         )}
       >
         {children}

studio/frontend/src/components/assistant-ui/tool-ui-python.tsx

@@ -7,7 +7,7 @@ import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
 import { code as codePlugin } from "@streamdown/code";
 import { CheckIcon, CodeIcon, CopyIcon, LoaderIcon } from "lucide-react";
-import { memo, useCallback, useMemo, useRef, useState } from "react";
+import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { Streamdown } from "streamdown";
 import {
   ToolFallbackContent,
@@ -28,6 +28,15 @@ function truncate(text: string): string {
 function CopyBtn({ text }: { text: string }) {
   const [copied, setCopied] = useState(false);
   const timer = useRef<ReturnType<typeof setTimeout> | null>(null);
+
+  useEffect(() => {
+    return () => {
+      if (timer.current) {
+        clearTimeout(timer.current);
+      }
+    };
+  }, []);
+
   const copy = useCallback(() => {
     if (copyToClipboard(text)) {
       setCopied(true);
@@ -98,14 +107,14 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
         icon={CodeIcon}
       />
       <ToolFallbackContent>
-        <div className="flex flex-col px-4">
+        <div className="border-l-2 border-muted-foreground/20 pl-2">
           {/* Code + copy */}
           {code && (
             <div className="flex justify-end">
               <CopyBtn text={code} />
             </div>
           )}
-          <HighlightedCode code={code} language="python" />
+          {code && <HighlightedCode code={code} language="python" />}
 
           {/* Output */}
           {isRunning ? (

studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx

@@ -6,7 +6,7 @@
 import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
 import { CheckIcon, CopyIcon, LoaderIcon, TerminalIcon } from "lucide-react";
-import { memo, useCallback, useRef, useState } from "react";
+import { memo, useCallback, useEffect, useRef, useState } from "react";
 import {
   ToolFallbackContent,
   ToolFallbackRoot,
@@ -25,6 +25,15 @@ function truncate(text: string): string {
 function CopyBtn({ text }: { text: string }) {
   const [copied, setCopied] = useState(false);
   const timer = useRef<ReturnType<typeof setTimeout> | null>(null);
+
+  useEffect(() => {
+    return () => {
+      if (timer.current) {
+        clearTimeout(timer.current);
+      }
+    };
+  }, []);
+
   const copy = useCallback(() => {
     if (copyToClipboard(text)) {
       setCopied(true);
@@ -74,7 +83,7 @@ const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
         icon={TerminalIcon}
       />
       <ToolFallbackContent>
-        <div className="flex flex-col px-4">
+        <div className="border-l-2 border-muted-foreground/20 pl-2">
           {isRunning ? (
             <div className="flex items-center gap-2 text-sm text-muted-foreground">
               <LoaderIcon className="size-3.5 animate-spin" />

studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx

@@ -81,29 +81,27 @@ const WebSearchToolUIImpl: ToolCallMessagePartComponent = ({
       />
       <ToolFallbackContent>
         {isRunning ? (
-          <div className="flex items-center gap-2 px-4 text-sm text-muted-foreground">
+          <div className="flex items-center gap-2 text-sm text-muted-foreground">
             <LoaderIcon className="size-3.5 animate-spin" />
             <span>Searching for &ldquo;{query}&rdquo;&hellip;</span>
           </div>
         ) : sources.length > 0 ? (
-          <div className="flex flex-col gap-1.5 px-4">
-            {sources.map((source) => (
+          <div className="flex flex-wrap gap-1.5">
+            {sources.map((source, i) => (
               <Source
-                key={source.url}
+                key={`${source.url}-${i}`}
                 href={source.url}
                 variant="outline"
-                size="default"
-                className="flex w-full max-w-full items-center gap-2 py-1.5"
+                size="sm"
+                className="inline-flex items-center gap-1.5"
               >
-                <SourceIcon url={source.url} className="size-3.5" />
-                <SourceTitle className="max-w-none flex-1 truncate">
-                  {source.title}
-                </SourceTitle>
+                <SourceIcon url={source.url} size={3} />
+                <SourceTitle>{source.title}</SourceTitle>
               </Source>
             ))}
           </div>
         ) : result ? (
-          <div className="px-4">
+          <div>
             <pre className="max-h-40 overflow-auto whitespace-pre-wrap break-words rounded bg-muted/50 p-2 text-xs">
               {typeof result === "string"
                 ? result

studio/frontend/src/components/ui/collapsible.tsx

@@ -2,13 +2,18 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import { cn } from "@/lib/utils";
+import * as React from "react";
 import { Collapsible as CollapsiblePrimitive } from "radix-ui";
 
-function Collapsible({
-  ...props
-}: React.ComponentProps<typeof CollapsiblePrimitive.Root>) {
-  return <CollapsiblePrimitive.Root data-slot="collapsible" {...props} />;
-}
+const Collapsible = React.forwardRef<
+  React.ElementRef<typeof CollapsiblePrimitive.Root>,
+  React.ComponentPropsWithoutRef<typeof CollapsiblePrimitive.Root>
+>(({ ...props }, ref) => {
+  return (
+    <CollapsiblePrimitive.Root ref={ref} data-slot="collapsible" {...props} />
+  );
+});
+Collapsible.displayName = CollapsiblePrimitive.Root.displayName;
 
 function CollapsibleTrigger({
   ...props

studio/frontend/src/features/chat/chat-settings-sheet.tsx

@@ -28,6 +28,8 @@ import {
   PencilEdit01Icon,
   Settings02Icon,
   SlidersHorizontalIcon,
+  UserSettings01Icon,
+  Wrench01Icon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { AnimatePresence, motion } from "motion/react";
@@ -164,6 +166,39 @@ function ParamSlider({
   );
 }
 
+const COLLAPSIBLE_STATE_KEY = "unsloth_chat_collapsible_state";
+
+function loadCollapsibleState(): Record<string, boolean> {
+  if (!canUseStorage()) return {};
+  try {
+    const raw = localStorage.getItem(COLLAPSIBLE_STATE_KEY);
+    if (!raw) return {};
+    const parsed = JSON.parse(raw);
+    if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) {
+      return {};
+    }
+    return Object.fromEntries(
+      Object.entries(parsed).filter(
+        (entry): entry is [string, boolean] => typeof entry[1] === "boolean",
+      ),
+    );
+  } catch (error) {
+    console.warn("Failed to load collapsible state from localStorage:", error);
+    return {};
+  }
+}
+
+function saveCollapsibleOpen(label: string, open: boolean) {
+  if (!canUseStorage()) return;
+  try {
+    const state = loadCollapsibleState();
+    state[label] = open;
+    localStorage.setItem(COLLAPSIBLE_STATE_KEY, JSON.stringify(state));
+  } catch {
+    // ignore
+  }
+}
+
 function CollapsibleSection({
   icon,
   label,
@@ -175,13 +210,20 @@ function CollapsibleSection({
   children?: ReactNode;
   defaultOpen?: boolean;
 }) {
-  const [open, setOpen] = useState(defaultOpen);
+  const [open, setOpen] = useState(() => {
+    const saved = loadCollapsibleState();
+    return Object.hasOwn(saved, label) ? saved[label] : defaultOpen;
+  });
 
   return (
     <div>
       <button
         type="button"
-        onClick={() => setOpen(!open)}
+        onClick={() => {
+          const next = !open;
+          setOpen(next);
+          saveCollapsibleOpen(label, next);
+        }}
         className="flex w-full items-center corner-squircle gap-2.5 rounded-md px-2 py-2 text-sm transition-colors hover:bg-accent"
       >
         <HugeiconsIcon icon={icon} className="size-4 text-muted-foreground" />
@@ -421,6 +463,69 @@ export function ChatSettingsPanel({
             />
           </div>
 
+          <CollapsibleSection icon={Settings02Icon} label="Model" defaultOpen={true}>
+            <div className="flex flex-col gap-3 py-1">
+              {isGguf && (
+                <>
+                  <div className="flex items-center justify-between gap-3">
+                    <div className="min-w-0">
+                      <div className="text-xs font-medium">Context Length</div>
+                      <div className="text-[11px] text-muted-foreground">
+                        Reported by the loaded GGUF model.
+                      </div>
+                    </div>
+                    <Input
+                      value={ggufContextLength ?? ""}
+                      placeholder="Loading..."
+                      disabled={true}
+                      className="h-7 w-[90px] text-xs"
+                    />
+                  </div>
+                  <div className="flex items-center justify-between gap-3">
+                    <div className="min-w-0">
+                      <div className="text-xs font-medium">KV Cache Dtype</div>
+                      <div className="text-[11px] text-muted-foreground">
+                        Quantize KV cache to reduce VRAM. Reload to apply.
+                      </div>
+                    </div>
+                    <Select
+                      value={kvCacheDtype ?? "f16"}
+                      onValueChange={(v) => {
+                        setKvCacheDtype(v === "f16" ? null : v);
+                        onReloadModel?.();
+                      }}
+                    >
+                      <SelectTrigger className="h-7 w-[90px] text-xs">
+                        <SelectValue />
+                      </SelectTrigger>
+                      <SelectContent>
+                        <SelectItem value="f16">f16</SelectItem>
+                        <SelectItem value="bf16">bf16</SelectItem>
+                        <SelectItem value="q8_0">q8_0</SelectItem>
+                        <SelectItem value="q5_1">q5_1</SelectItem>
+                        <SelectItem value="q4_1">q4_1</SelectItem>
+                      </SelectContent>
+                    </Select>
+                  </div>
+                </>
+              )}
+              {!isGguf && (
+                <div className="flex items-center justify-between gap-3">
+                  <div className="min-w-0">
+                    <div className="text-xs font-medium">Trust remote code</div>
+                    <div className="text-[11px] text-muted-foreground">
+                      Allow models with custom code (e.g. Nemotron). Only enable for repos you trust.
+                    </div>
+                  </div>
+                  <Switch
+                    checked={params.trustRemoteCode ?? false}
+                    onCheckedChange={set("trustRemoteCode")}
+                  />
+                </div>
+              )}
+            </div>
+          </CollapsibleSection>
+
           <CollapsibleSection
             icon={SlidersHorizontalIcon}
             label="Sampling"
@@ -505,7 +610,15 @@ export function ChatSettingsPanel({
             </div>
           </CollapsibleSection>
 
-          <CollapsibleSection icon={Settings02Icon} label="Settings" defaultOpen={true}>
+          <CollapsibleSection icon={Wrench01Icon} label="Tools">
+            <div className="flex flex-col gap-3 py-1">
+              <AutoHealToolCallsToggle />
+              <MaxToolCallsSlider />
+              <ToolCallTimeoutSlider />
+            </div>
+          </CollapsibleSection>
+
+          <CollapsibleSection icon={UserSettings01Icon} label="Preferences" defaultOpen={true}>
             <div className="flex flex-col gap-3 py-1">
               <div className="flex items-center justify-between gap-3">
                 <div className="min-w-0">
@@ -519,49 +632,6 @@ export function ChatSettingsPanel({
                   onCheckedChange={onAutoTitleChange}
                 />
               </div>
-              <div className="flex items-center justify-between gap-3">
-                <div className="min-w-0">
-                  <div className="text-xs font-medium">Trust remote code</div>
-                  <div className="text-[11px] text-muted-foreground">
-                    Allow models with custom code (e.g. Nemotron). Only enable for repos you trust.
-                  </div>
-                </div>
-                <Switch
-                  checked={params.trustRemoteCode ?? false}
-                  onCheckedChange={set("trustRemoteCode")}
-                />
-              </div>
-              {isGguf && (
-                <div className="flex items-center justify-between gap-3">
-                  <div className="min-w-0">
-                    <div className="text-xs font-medium">KV Cache Dtype</div>
-                    <div className="text-[11px] text-muted-foreground">
-                      Quantize KV cache to reduce VRAM. Reload to apply.
-                    </div>
-                  </div>
-                  <Select
-                    value={kvCacheDtype ?? "f16"}
-                    onValueChange={(v) => {
-                      setKvCacheDtype(v === "f16" ? null : v);
-                      onReloadModel?.();
-                    }}
-                  >
-                    <SelectTrigger className="h-7 w-[90px] text-xs">
-                      <SelectValue />
-                    </SelectTrigger>
-                    <SelectContent>
-                      <SelectItem value="f16">f16</SelectItem>
-                      <SelectItem value="bf16">bf16</SelectItem>
-                      <SelectItem value="q8_0">q8_0</SelectItem>
-                      <SelectItem value="q5_1">q5_1</SelectItem>
-                      <SelectItem value="q4_1">q4_1</SelectItem>
-                    </SelectContent>
-                  </Select>
-                </div>
-              )}
-              <AutoHealToolCallsToggle />
-              <MaxToolCallsSlider />
-              <ToolCallTimeoutSlider />
             </div>
           </CollapsibleSection>
 

studio/frontend/src/features/studio/historical-training-view.tsx

@@ -0,0 +1,168 @@
+// 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 type { TrainingViewData } from "@/features/training";
+import { getTrainingRun } from "@/features/training";
+import type { TrainingRunDetailResponse } from "@/features/training";
+import { type ReactElement, useEffect, useState } from "react";
+import { ChartsSection } from "./sections/charts-section";
+import { ProgressSection } from "./sections/progress-section";
+
+interface HistoricalTrainingViewProps {
+  runId: string;
+}
+
+function normalizeTrainingMethod(config: Record<string, unknown>): string {
+  const type = config?.training_type as string | undefined;
+  if (!type || type === "Full Finetuning") return "full";
+  if (type === "LoRA/QLoRA") {
+    return config?.load_in_4bit ? "qlora" : "lora";
+  }
+  return "full";
+}
+
+function mapToViewData(detail: TrainingRunDetailResponse): TrainingViewData {
+  const { run, metrics } = detail;
+
+  const lossHistory = metrics.loss_step_history
+    .map((step, i) => ({ step, value: metrics.loss_history[i] }))
+    .filter((p): p is { step: number; value: number } => p.value != null);
+
+  const lrHistory = metrics.lr_step_history
+    .map((step, i) => ({ step, value: metrics.lr_history[i] }))
+    .filter((p): p is { step: number; value: number } => p.value != null);
+
+  const gradNormHistory = metrics.grad_norm_step_history
+    .map((step, i) => ({ step, value: metrics.grad_norm_history[i] }))
+    .filter((p): p is { step: number; value: number } => p.value != null);
+
+  const evalLossHistory = metrics.eval_step_history
+    .map((step, i) => ({ step, value: metrics.eval_loss_history[i] }))
+    .filter((p): p is { step: number; value: number } => p.value != null);
+
+  const phase =
+    run.status === "completed"
+      ? "completed"
+      : run.status === "stopped"
+        ? "stopped"
+        : run.status === "error"
+          ? "error"
+          : run.status === "running"
+            ? "training"
+            : "idle";
+
+  return {
+    phase,
+    currentStep: run.final_step ?? 0,
+    totalSteps: run.total_steps ?? 0,
+    currentLoss: run.final_loss,
+    currentLearningRate: metrics.lr_history.at(-1) ?? null,
+    currentGradNorm: metrics.grad_norm_history.at(-1) ?? null,
+    currentEpoch: metrics.final_epoch,
+    currentNumTokens: metrics.final_num_tokens ?? null,
+    progressPercent:
+      run.total_steps && run.final_step
+        ? (run.final_step / run.total_steps) * 100
+        : 0,
+    elapsedSeconds: run.duration_seconds,
+    etaSeconds: null,
+    evalEnabled: evalLossHistory.length > 0,
+    message:
+      run.status === "completed"
+        ? "Training completed"
+        : run.status === "stopped"
+          ? "Training stopped"
+          : run.status === "running"
+            ? "Training in progress"
+            : run.error_message ?? "Training errored",
+    error: run.status === "error" ? run.error_message : null,
+    isTrainingRunning: false,
+    modelName: run.model_name,
+    trainingMethod: normalizeTrainingMethod(detail.config),
+    lossHistory,
+    lrHistory,
+    gradNormHistory,
+    evalLossHistory,
+  };
+}
+
+export function HistoricalTrainingView({
+  runId,
+}: HistoricalTrainingViewProps): ReactElement {
+  const [detail, setDetail] = useState<TrainingRunDetailResponse | null>(null);
+  const [error, setError] = useState<string | null>(null);
+
+  // Derive loading from detail/error -- no separate state needed
+  const loading = detail === null && error === null;
+
+  useEffect(() => {
+    const controller = new AbortController();
+    getTrainingRun(runId, controller.signal)
+      .then((result) => {
+        setDetail(result);
+      })
+      .catch((err) => {
+        if (err instanceof DOMException && err.name === "AbortError") return;
+        setError(err instanceof Error ? err.message : "Failed to load run");
+      });
+    return () => {
+      controller.abort();
+      // Reset on runId change so loading derives correctly for the next fetch
+      setDetail(null);
+      setError(null);
+    };
+  }, [runId]);
+
+  if (loading) {
+    return (
+      <div className="rounded-xl border bg-card p-8 text-sm text-muted-foreground">
+        Loading training run...
+      </div>
+    );
+  }
+
+  if (error || !detail) {
+    return (
+      <div className="rounded-xl border border-destructive/30 bg-destructive/5 p-8 text-sm text-red-500">
+        {error ?? "Run not found"}
+      </div>
+    );
+  }
+
+  const viewData = mapToViewData(detail);
+  const configOverride = detail.config
+    ? {
+        epochs: detail.config.num_epochs as number | undefined,
+        batchSize: detail.config.batch_size as number | undefined,
+        learningRate: detail.config.learning_rate as string | undefined,
+        maxSteps: detail.config.max_steps as number | undefined,
+        contextLength: detail.config.max_seq_length as number | undefined,
+        warmupSteps: detail.config.warmup_steps as number | undefined,
+        optimizerType: detail.config.optim as string | undefined,
+        loraRank: detail.config.lora_r as number | undefined,
+        loraAlpha: detail.config.lora_alpha as number | undefined,
+        loraDropout: detail.config.lora_dropout as number | undefined,
+        loraVariant: detail.config.use_rslora ? "rsLoRA" : undefined,
+      }
+    : undefined;
+
+  return (
+    <div className="flex flex-col gap-6">
+      <ProgressSection
+        data={viewData}
+        isHistorical
+        configOverride={configOverride}
+      />
+      <ChartsSection
+        currentStep={viewData.currentStep}
+        totalSteps={viewData.totalSteps}
+        isTraining={false}
+        evalEnabled={viewData.evalEnabled}
+        lossHistory={viewData.lossHistory}
+        lrHistory={viewData.lrHistory}
+        gradNormHistory={viewData.gradNormHistory}
+        evalLossHistory={viewData.evalLossHistory}
+      />
+    </div>
+  );
+}

studio/frontend/src/features/studio/history-card-grid.tsx

@@ -0,0 +1,401 @@
+// 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 {
+  AlertDialog,
+  AlertDialogAction,
+  AlertDialogCancel,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogFooter,
+  AlertDialogHeader,
+  AlertDialogTitle,
+} from "@/components/ui/alert-dialog";
+import { Button } from "@/components/ui/button";
+import type { TrainingRunSummary } from "@/features/training";
+import { deleteTrainingRun, listTrainingRuns } from "@/features/training";
+import { cn } from "@/lib/utils";
+import { Delete02Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { type ReactElement, useCallback, useEffect, useRef, useState } from "react";
+import { Spinner } from "@/components/ui/spinner";
+
+const PAGE_SIZE = 12;
+const RUNNING_POLL_INTERVAL_MS = 5000;
+
+const statusBadge: Record<
+  string,
+  { label: string; className: string }
+> = {
+  completed: {
+    label: "Completed",
+    className:
+      "bg-emerald-100 text-emerald-700 dark:bg-emerald-950 dark:text-emerald-400",
+  },
+  stopped: {
+    label: "Stopped",
+    className:
+      "bg-amber-100 text-amber-700 dark:bg-amber-950 dark:text-amber-400",
+  },
+  error: {
+    label: "Error",
+    className: "bg-red-100 text-red-700 dark:bg-red-950 dark:text-red-400",
+  },
+  running: {
+    label: "Running",
+    className:
+      "bg-blue-100 text-blue-700 dark:bg-blue-950 dark:text-blue-400",
+  },
+};
+
+function catmullRomPath(points: { x: number; y: number }[]): string {
+  if (points.length < 2) return "";
+  const d = [`M${points[0]!.x.toFixed(1)},${points[0]!.y.toFixed(1)}`];
+  for (let i = 0; i < points.length - 1; i++) {
+    const p0 = points[Math.max(i - 1, 0)]!;
+    const p1 = points[i]!;
+    const p2 = points[i + 1]!;
+    const p3 = points[Math.min(i + 2, points.length - 1)]!;
+    const cp1x = p1.x + (p2.x - p0.x) / 6;
+    const cp1y = p1.y + (p2.y - p0.y) / 6;
+    const cp2x = p2.x - (p3.x - p1.x) / 6;
+    const cp2y = p2.y - (p3.y - p1.y) / 6;
+    d.push(
+      `C${cp1x.toFixed(1)},${cp1y.toFixed(1)} ${cp2x.toFixed(1)},${cp2y.toFixed(1)} ${p2.x.toFixed(1)},${p2.y.toFixed(1)}`,
+    );
+  }
+  return d.join(" ");
+}
+
+function Sparkline({ values, id }: { values: number[]; id: string }): ReactElement | null {
+  if (!values || values.length < 2) return null;
+  let min = values[0]!;
+  let max = values[0]!;
+  for (let i = 1; i < values.length; i++) {
+    if (values[i]! < min) min = values[i]!;
+    if (values[i]! > max) max = values[i]!;
+  }
+  const range = max - min || 1;
+  const pad = 1.5; // half stroke-width so peaks aren't clipped
+  const h = 32;
+  const w = 120;
+  const gradientId = `sparkFill-${id}`;
+
+  // Build points with vertical padding so the stroke isn't clipped
+  const pts = values.map((v, i) => ({
+    x: (i / (values.length - 1)) * w,
+    y: pad + (1 - (v - min) / range) * (h - pad * 2),
+  }));
+
+  const linePath = catmullRomPath(pts);
+  const last = pts[pts.length - 1]!;
+  const first = pts[0]!;
+  const fillPath = `${linePath} L${last.x.toFixed(1)},${h} L${first.x.toFixed(1)},${h} Z`;
+
+  return (
+    <svg viewBox={`0 0 ${w} ${h}`} className="h-8 w-full" preserveAspectRatio="none" role="img" aria-label="Loss trend sparkline">
+      <defs>
+        <linearGradient id={gradientId} x1="0" y1="0" x2="0" y2="1">
+          <stop offset="0%" stopColor="currentColor" stopOpacity="0.12" />
+          <stop offset="100%" stopColor="currentColor" stopOpacity="0" />
+        </linearGradient>
+      </defs>
+      <path
+        d={fillPath}
+        fill={`url(#${gradientId})`}
+        className="text-emerald-500"
+      />
+      <path
+        d={linePath}
+        fill="none"
+        stroke="currentColor"
+        strokeWidth="1.5"
+        strokeLinecap="round"
+        strokeLinejoin="round"
+        className="text-emerald-500"
+      />
+    </svg>
+  );
+}
+
+function formatRelativeTime(isoDate: string): string {
+  const diff = Date.now() - new Date(isoDate).getTime();
+  const mins = Math.floor(diff / 60000);
+  if (mins < 1) return "just now";
+  if (mins < 60) return `${mins}m ago`;
+  const hrs = Math.floor(mins / 60);
+  if (hrs < 24) return `${hrs}h ago`;
+  const days = Math.floor(hrs / 24);
+  return `${days}d ago`;
+}
+
+function formatDuration(seconds: number | null): string {
+  if (seconds == null) return "--";
+  const total = Math.floor(seconds);
+  if (total < 60) return `${total}s`;
+  const min = Math.floor(total / 60);
+  const sec = total % 60;
+  if (min < 60) return `${min}m ${sec}s`;
+  const hrs = Math.floor(min / 60);
+  return `${hrs}h ${min % 60}m`;
+}
+
+interface HistoryCardGridProps {
+  onSelectRun: (runId: string) => void;
+}
+
+export function HistoryCardGrid({
+  onSelectRun,
+}: HistoryCardGridProps): ReactElement {
+  const [runs, setRuns] = useState<TrainingRunSummary[]>([]);
+  const [total, setTotal] = useState(0);
+  const [loading, setLoading] = useState(true);
+  const [error, setError] = useState<string | null>(null);
+  const [deleteError, setDeleteError] = useState<string | null>(null);
+  const [deleteTarget, setDeleteTarget] = useState<string | null>(null);
+  const [manualFetchInFlight, setManualFetchInFlight] = useState(false);
+
+  const userControllerRef = useRef<AbortController | null>(null);
+  const pollControllerRef = useRef<AbortController | null>(null);
+  const fetchIdRef = useRef(0);
+  const pollIdRef = useRef(0);
+
+  const fetchRuns = useCallback(async (offset = 0, append = false, limit = PAGE_SIZE) => {
+    // Cancel any in-flight poll so its stale response can't clobber this fresher fetch
+    pollControllerRef.current?.abort();
+    userControllerRef.current?.abort();
+    const controller = new AbortController();
+    userControllerRef.current = controller;
+    const id = ++fetchIdRef.current;
+
+    setManualFetchInFlight(true);
+    setLoading(true);
+    setError(null);
+    try {
+      const result = await listTrainingRuns(limit, offset, controller.signal);
+      if (fetchIdRef.current !== id) return;
+      setRuns((prev) => (append ? [...prev, ...result.runs] : result.runs));
+      setTotal(result.total);
+    } catch (err) {
+      if (err instanceof DOMException && err.name === "AbortError") return;
+      if (fetchIdRef.current !== id) return;
+      if (!append) setError("Failed to load training runs");
+    } finally {
+      if (fetchIdRef.current === id) {
+        setLoading(false);
+        setManualFetchInFlight(false);
+      }
+    }
+  }, []);
+
+  useEffect(() => {
+    void fetchRuns(0);
+    return () => {
+      userControllerRef.current?.abort();
+    };
+  }, [fetchRuns]);
+
+  // Poll while any run is still "running" so the card shows live progress
+  const hasRunningRun = runs.some((r) => r.status === "running");
+  const visibleCount = runs.length;
+  useEffect(() => {
+    if (!hasRunningRun) return;
+    const timer = setInterval(async () => {
+      if (manualFetchInFlight) return;
+      pollControllerRef.current?.abort();
+      const controller = new AbortController();
+      pollControllerRef.current = controller;
+      const pid = ++pollIdRef.current;
+      try {
+        const limit = Math.max(PAGE_SIZE, visibleCount);
+        const result = await listTrainingRuns(limit, 0, controller.signal);
+        if (pollIdRef.current !== pid) return; // stale poll — discard
+        setRuns(result.runs);
+        setTotal(result.total);
+      } catch {
+        // silently handle — poll will retry
+      }
+    }, RUNNING_POLL_INTERVAL_MS);
+    return () => {
+      clearInterval(timer);
+      pollControllerRef.current?.abort();
+    };
+  }, [hasRunningRun, visibleCount, manualFetchInFlight]);
+
+  const handleDelete = async () => {
+    if (!deleteTarget) return;
+    setDeleteError(null);
+    try {
+      await deleteTrainingRun(deleteTarget);
+      // Optimistically remove the card so it disappears immediately
+      setRuns((prev) => prev.filter((r) => r.id !== deleteTarget));
+      setTotal((prev) => Math.max(0, prev - 1));
+      // Re-fetch preserving visible count so offsets stay consistent for "Load more"
+      const currentCount = runs.length - 1;
+      const limit = Math.max(PAGE_SIZE, currentCount);
+      fetchRuns(0, false, limit).catch(() => {
+        // Refresh failed — card is already removed, no stale display
+      });
+    } catch {
+      setDeleteError("Failed to delete training run. Please try again.");
+    }
+    setDeleteTarget(null);
+  };
+
+  if (!loading && error && runs.length === 0) {
+    return (
+      <div className="flex flex-col items-center gap-2 py-16 text-center">
+        <p className="text-sm text-destructive">{error}</p>
+        <Button variant="outline" size="sm" onClick={() => void fetchRuns(0)}>
+          Retry
+        </Button>
+      </div>
+    );
+  }
+
+  if (!loading && runs.length === 0) {
+    return (
+      <div className="flex flex-col items-center gap-2 py-16 text-center">
+        <p className="text-sm text-muted-foreground">
+          No training runs yet. Start your first training run in the Configure
+          tab.
+        </p>
+      </div>
+    );
+  }
+
+  return (
+    <>
+      {deleteError && (
+        <div className="mb-4 rounded-lg border border-destructive/50 bg-destructive/10 px-4 py-2 text-sm text-destructive">
+          {deleteError}
+        </div>
+      )}
+      <div className="grid grid-cols-1 gap-4 sm:grid-cols-2 xl:grid-cols-3">
+        {runs.map((run) => {
+          const badge = statusBadge[run.status] ?? statusBadge.error;
+          const isRunning = run.status === "running";
+          return (
+            <div
+              role="button"
+              tabIndex={0}
+              key={run.id}
+              className={cn(
+                "group relative flex cursor-pointer flex-col gap-3 rounded-xl border bg-card p-4 text-left transition-colors hover:border-border hover:bg-accent/30",
+                isRunning
+                  ? "border-blue-400/50 dark:border-blue-500/30"
+                  : "border-border/60",
+              )}
+              onClick={() => onSelectRun(run.id)}
+              onKeyDown={(e) => {
+                if (e.key === "Enter" || e.key === " ") {
+                  e.preventDefault();
+                  onSelectRun(run.id);
+                }
+              }}
+            >
+              <div className="flex items-center justify-between pr-6">
+                <span
+                  className={cn(
+                    "inline-flex items-center gap-1.5 rounded-full px-2 py-0.5 text-[10px] font-semibold",
+                    badge.className,
+                  )}
+                >
+                  {isRunning && <Spinner className="size-2.5" />}
+                  {badge.label}
+                </span>
+                <span className="text-[10px] text-muted-foreground">
+                  {formatRelativeTime(run.started_at)}
+                </span>
+              </div>
+              <div className="min-w-0">
+                <p
+                  className="truncate text-sm font-medium"
+                  title={run.model_name}
+                >
+                  {run.model_name}
+                </p>
+                <p className="truncate text-xs text-muted-foreground">
+                  {run.dataset_name}
+                </p>
+              </div>
+              {run.loss_sparkline && run.loss_sparkline.length >= 2 && (
+                <Sparkline values={run.loss_sparkline} id={run.id} />
+              )}
+              <div className="flex flex-wrap gap-x-4 gap-y-1 text-[11px] text-muted-foreground">
+                <span>
+                  Loss:{" "}
+                  {run.final_loss != null ? run.final_loss.toFixed(4) : "--"}
+                </span>
+                <span>
+                  Steps: {run.final_step ?? 0}/{run.total_steps ?? "--"}
+                </span>
+                <span>{formatDuration(run.duration_seconds)}</span>
+              </div>
+              {!isRunning && (
+                <button
+                  type="button"
+                  className="absolute right-3 top-3 rounded-md p-1 text-muted-foreground/50 opacity-0 transition-opacity hover:bg-destructive/10 hover:text-destructive group-hover:opacity-100 focus-visible:opacity-100"
+                  aria-label="Delete run"
+                  onClick={(e) => {
+                    e.stopPropagation();
+                    setDeleteTarget(run.id);
+                  }}
+                >
+                  <HugeiconsIcon icon={Delete02Icon} className="size-3.5" />
+                </button>
+              )}
+            </div>
+          );
+        })}
+      </div>
+      {runs.length < total && (
+        <div className="mt-4 flex justify-center">
+          <Button
+            variant="outline"
+            size="sm"
+            onClick={() => void fetchRuns(runs.length, true)}
+            disabled={loading}
+          >
+            {loading ? "Loading..." : "Load more"}
+          </Button>
+        </div>
+      )}
+      {loading && runs.length === 0 && (
+        <div className="grid grid-cols-1 gap-4 sm:grid-cols-2 xl:grid-cols-3">
+          {Array.from({ length: 3 }).map((_, i) => (
+            <div
+              key={`skeleton-${i}`}
+              className="h-40 animate-pulse rounded-xl border bg-muted/30"
+            />
+          ))}
+        </div>
+      )}
+      <AlertDialog
+        open={deleteTarget !== null}
+        onOpenChange={(open) => {
+          if (!open) setDeleteTarget(null);
+        }}
+      >
+        <AlertDialogContent>
+          <AlertDialogHeader>
+            <AlertDialogTitle>Delete training run?</AlertDialogTitle>
+            <AlertDialogDescription>
+              This will permanently delete this training run and all its metrics.
+              This action cannot be undone.
+            </AlertDialogDescription>
+          </AlertDialogHeader>
+          <AlertDialogFooter>
+            <AlertDialogCancel>Cancel</AlertDialogCancel>
+            <AlertDialogAction
+              variant="destructive"
+              onClick={() => void handleDelete()}
+            >
+              Delete
+            </AlertDialogAction>
+          </AlertDialogFooter>
+        </AlertDialogContent>
+      </AlertDialog>
+    </>
+  );
+}

studio/frontend/src/features/studio/live-training-view.tsx

@@ -0,0 +1,118 @@
+// 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 { cn } from "@/lib/utils";
+import {
+  useTrainingConfigStore,
+  useTrainingRuntimeStore,
+} from "@/features/training";
+import type { TrainingViewData } from "@/features/training";
+import type { ReactElement } from "react";
+import { useShallow } from "zustand/react/shallow";
+import { ChartsSection } from "./sections/charts-section";
+import { ProgressSection } from "./sections/progress-section";
+import { TrainingStartOverlay } from "./training-start-overlay";
+
+export function LiveTrainingView(): ReactElement {
+  const runtime = useTrainingRuntimeStore(
+    useShallow((state) => ({
+      jobId: state.jobId,
+      phase: state.phase,
+      message: state.message,
+      error: state.error,
+      currentStep: state.currentStep,
+      totalSteps: state.totalSteps,
+      currentEpoch: state.currentEpoch,
+      currentLoss: state.currentLoss,
+      currentLearningRate: state.currentLearningRate,
+      currentGradNorm: state.currentGradNorm,
+      currentNumTokens: state.currentNumTokens,
+      progressPercent: state.progressPercent,
+      elapsedSeconds: state.elapsedSeconds,
+      etaSeconds: state.etaSeconds,
+      evalEnabled: state.evalEnabled,
+      isTrainingRunning: state.isTrainingRunning,
+      lossHistory: state.lossHistory,
+      lrHistory: state.lrHistory,
+      gradNormHistory: state.gradNormHistory,
+      evalLossHistory: state.evalLossHistory,
+      firstStepReceived: state.firstStepReceived,
+      isStarting: state.isStarting,
+    })),
+  );
+
+  const config = useTrainingConfigStore(
+    useShallow((state) => ({
+      selectedModel: state.selectedModel,
+      trainingMethod: state.trainingMethod,
+    })),
+  );
+
+  const viewData: TrainingViewData = {
+    phase: runtime.phase,
+    currentStep: runtime.currentStep,
+    totalSteps: runtime.totalSteps,
+    currentLoss: runtime.currentLoss,
+    currentLearningRate: runtime.currentLearningRate,
+    currentGradNorm: runtime.currentGradNorm,
+    currentEpoch: runtime.currentEpoch,
+    currentNumTokens: runtime.currentNumTokens,
+    progressPercent: runtime.progressPercent,
+    elapsedSeconds: runtime.elapsedSeconds,
+    etaSeconds: runtime.etaSeconds,
+    evalEnabled: runtime.evalEnabled,
+    message: runtime.message,
+    error: runtime.error,
+    isTrainingRunning: runtime.isTrainingRunning,
+    modelName: config.selectedModel ?? "",
+    trainingMethod: config.trainingMethod ?? "",
+    lossHistory: runtime.lossHistory,
+    lrHistory: runtime.lrHistory,
+    gradNormHistory: runtime.gradNormHistory,
+    evalLossHistory: runtime.evalLossHistory,
+  };
+
+  const isPreparingPhase =
+    runtime.phase === "downloading_model" ||
+    runtime.phase === "downloading_dataset" ||
+    runtime.phase === "loading_model" ||
+    runtime.phase === "loading_dataset" ||
+    runtime.phase === "configuring";
+  const isWaitingForFirstStep =
+    runtime.phase === "training" && !runtime.firstStepReceived;
+  const showOverlay =
+    runtime.isStarting ||
+    isPreparingPhase ||
+    (isWaitingForFirstStep && runtime.currentStep <= 0);
+
+  return (
+    <div className={cn("relative", showOverlay && "min-h-[72vh]")}>
+      <div
+        className={cn(
+          "relative z-10 flex flex-col gap-6 transition-[filter]",
+          showOverlay && "blur",
+        )}
+      >
+        <div data-tour="studio-training-progress">
+          <ProgressSection key={runtime.jobId ?? "no-job"} data={viewData} />
+        </div>
+        <ChartsSection
+          currentStep={viewData.currentStep}
+          totalSteps={viewData.totalSteps}
+          isTraining={viewData.isTrainingRunning}
+          evalEnabled={viewData.evalEnabled}
+          lossHistory={viewData.lossHistory}
+          lrHistory={viewData.lrHistory}
+          gradNormHistory={viewData.gradNormHistory}
+          evalLossHistory={viewData.evalLossHistory}
+        />
+      </div>
+      {showOverlay ? (
+        <TrainingStartOverlay
+          message={runtime.message}
+          currentStep={runtime.currentStep}
+        />
+      ) : null}
+    </div>
+  );
+}

studio/frontend/src/features/studio/sections/charts-section.tsx

@@ -1,7 +1,7 @@
 // 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 { useTrainingRuntimeStore } from "@/features/training";
+import type { TrainingSeriesPoint } from "@/features/training";
 import { type ReactElement, Suspense, lazy, useMemo } from "react";
 
 const ChartsContent = lazy(() =>
@@ -16,42 +16,49 @@ const SKELETON_KEYS = [
   "chart-skeleton-4",
 ];
 
-export function ChartsSection(): ReactElement | null {
-  const currentStep = useTrainingRuntimeStore((state) => state.currentStep);
-  const totalSteps = useTrainingRuntimeStore((state) => state.totalSteps);
-  const isTraining = useTrainingRuntimeStore((state) => state.isTrainingRunning);
-  const evalEnabled = useTrainingRuntimeStore((state) => state.evalEnabled);
-  const lossHistoryRaw = useTrainingRuntimeStore((state) => state.lossHistory);
-  const lrHistoryRaw = useTrainingRuntimeStore((state) => state.lrHistory);
-  const gradNormHistoryRaw = useTrainingRuntimeStore(
-    (state) => state.gradNormHistory,
-  );
-  const evalLossHistoryRaw = useTrainingRuntimeStore(
-    (state) => state.evalLossHistory,
-  );
+interface ChartsSectionProps {
+  currentStep: number;
+  totalSteps: number;
+  isTraining: boolean;
+  evalEnabled: boolean;
+  lossHistory: TrainingSeriesPoint[];
+  lrHistory: TrainingSeriesPoint[];
+  gradNormHistory: TrainingSeriesPoint[];
+  evalLossHistory: TrainingSeriesPoint[];
+}
 
+export function ChartsSection({
+  currentStep,
+  totalSteps,
+  isTraining,
+  evalEnabled,
+  lossHistory,
+  lrHistory,
+  gradNormHistory,
+  evalLossHistory,
+}: ChartsSectionProps): ReactElement | null {
   const series = useMemo(
     () => ({
       currentStep,
       totalSteps,
-      lossHistory: lossHistoryRaw.map((point) => ({
+      lossHistory: lossHistory.map((point) => ({
         step: point.step,
         loss: point.value,
       })),
-      lrHistory: lrHistoryRaw.map((point) => ({
+      lrHistory: lrHistory.map((point) => ({
         step: point.step,
         lr: point.value,
       })),
-      gradNormHistory: gradNormHistoryRaw.map((point) => ({
+      gradNormHistory: gradNormHistory.map((point) => ({
         step: point.step,
         gradNorm: point.value,
       })),
-      evalLossHistory: evalLossHistoryRaw.map((point) => ({
+      evalLossHistory: evalLossHistory.map((point) => ({
         step: point.step,
         loss: point.value,
       })),
     }),
-    [currentStep, evalLossHistoryRaw, gradNormHistoryRaw, lossHistoryRaw, lrHistoryRaw, totalSteps],
+    [currentStep, evalLossHistory, gradNormHistory, lossHistory, lrHistory, totalSteps],
   );
 
   if (

studio/frontend/src/features/studio/sections/progress-section.tsx

@@ -26,6 +26,7 @@ import {
   useTrainingConfigStore,
   useTrainingRuntimeStore,
 } from "@/features/training";
+import type { TrainingViewData } from "@/features/training";
 import { useGpuUtilization } from "@/hooks";
 import { cn } from "@/lib/utils";
 import {
@@ -39,7 +40,7 @@ import {
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { Link, useNavigate } from "@tanstack/react-router";
-import { type ReactElement, type ReactNode, useEffect, useState } from "react";
+import { type ReactElement, type ReactNode, useState } from "react";
 import { useShallow } from "zustand/react/shallow";
 import { ChartSettingsSheet } from "./charts/chart-settings-sheet";
 import {
@@ -61,34 +62,33 @@ function configRow(
   return [label, value];
 }
 
-export function ProgressSection(): ReactElement {
+interface ProgressSectionProps {
+  data: TrainingViewData;
+  isHistorical?: boolean;
+  configOverride?: {
+    epochs?: number;
+    batchSize?: number;
+    learningRate?: string;
+    maxSteps?: number;
+    contextLength?: number;
+    warmupSteps?: number;
+    optimizerType?: string;
+    loraRank?: number;
+    loraAlpha?: number;
+    loraDropout?: number;
+    loraVariant?: string;
+  };
+}
+
+export function ProgressSection({
+  data,
+  isHistorical = false,
+  configOverride,
+}: ProgressSectionProps): ReactElement {
   const navigate = useNavigate();
-  const runtime = useTrainingRuntimeStore(
-    useShallow((state) => ({
-      phase: state.phase,
-      message: state.message,
-      error: state.error,
-      currentStep: state.currentStep,
-      totalSteps: state.totalSteps,
-      currentEpoch: state.currentEpoch,
-      currentLoss: state.currentLoss,
-      currentLearningRate: state.currentLearningRate,
-      currentGradNorm: state.currentGradNorm,
-      progressPercent: state.progressPercent,
-      elapsedSeconds: state.elapsedSeconds,
-      etaSeconds: state.etaSeconds,
-      currentNumTokens: state.currentNumTokens,
-      isTrainingRunning: state.isTrainingRunning,
-      lossHistory: state.lossHistory,
-      lrHistory: state.lrHistory,
-      gradNormHistory: state.gradNormHistory,
-    })),
-  );
 
   const config = useTrainingConfigStore(
     useShallow((state) => ({
-      selectedModel: state.selectedModel,
-      trainingMethod: state.trainingMethod,
       epochs: state.epochs,
       batchSize: state.batchSize,
       learningRate: state.learningRate,
@@ -103,98 +103,92 @@ export function ProgressSection(): ReactElement {
     })),
   );
 
-  const { stopTrainingRun } = useTrainingActions();
-  const gpu = useGpuUtilization(runtime.isTrainingRunning);
   const [stopDialogOpen, setStopDialogOpen] = useState(false);
-  const [stopRequested, setStopRequested] = useState(false);
+  const [stopRequestedLocal, setStopRequestedLocal] = useState(false);
 
-  useEffect(() => {
-    if (!runtime.isTrainingRunning) {
-      setStopRequested(false);
-    }
-  }, [runtime.isTrainingRunning]);
+  // Auto-reset when training stops -- no useEffect needed
+  const stopRequested = data.isTrainingRunning && stopRequestedLocal;
 
   const pct =
-    runtime.totalSteps > 0
+    data.totalSteps > 0
       ? Math.min(
           100,
           Math.max(
             0,
-            Math.round((runtime.currentStep / runtime.totalSteps) * 100),
+            Math.round((data.currentStep / data.totalSteps) * 100),
           ),
         )
-      : Math.round(runtime.progressPercent);
+      : Math.round(data.progressPercent);
 
-  const elapsed = runtime.elapsedSeconds;
+  const elapsed = data.elapsedSeconds;
   const derivedEta =
     elapsed != null && pct > 0
       ? Math.round((elapsed * (100 - pct)) / Math.max(pct, 1))
       : null;
-  const eta = runtime.etaSeconds ?? derivedEta;
+  const eta = data.etaSeconds ?? derivedEta;
 
   const stepsPerSecond =
-    elapsed != null && elapsed > 0 ? runtime.currentStep / elapsed : null;
+    elapsed != null && elapsed > 0 ? data.currentStep / elapsed : null;
   const showHalfwayHint =
-    runtime.phase === "training" && pct >= 50 && pct < 100;
-  const showCompletedHint = runtime.phase === "completed";
+    data.phase === "training" && pct >= 50 && pct < 100;
+  const showCompletedHint = data.phase === "completed";
   const handleCompareInChat = async () => {
-    setTrainingCompareHandoff(config.selectedModel);
+    setTrainingCompareHandoff(data.modelName);
     await navigate({ to: "/chat" });
   };
-  const requestStop = async (saveCheckpoint: boolean) => {
-    setStopRequested(true);
-    setStopDialogOpen(false);
-    useTrainingRuntimeStore.getState().setStopRequested(true);
-    try {
-      const ok = await stopTrainingRun(saveCheckpoint);
-      if (!ok) {
-        setStopRequested(false);
-      }
-    } catch {
-      setStopRequested(false);
-    }
-  };
 
   const stoppedLoss = getDisplayMetric(
-    runtime.isTrainingRunning,
-    runtime.currentLoss,
-    runtime.lossHistory,
+    data.isTrainingRunning,
+    data.currentLoss,
+    data.lossHistory,
   );
   const stoppedLr = getDisplayMetric(
-    runtime.isTrainingRunning,
-    runtime.currentLearningRate,
-    runtime.lrHistory,
+    data.isTrainingRunning,
+    data.currentLearningRate,
+    data.lrHistory,
   );
-  const stoppedGradNorm = runtime.isTrainingRunning
-    ? runtime.currentGradNorm
-    : (lastNonZeroValue(runtime.gradNormHistory) ?? runtime.currentGradNorm);
+  const stoppedGradNorm = data.isTrainingRunning
+    ? data.currentGradNorm
+    : (lastValue(data.gradNormHistory) ?? data.currentGradNorm);
+
+  const cfgEpochs = isHistorical ? configOverride?.epochs : config.epochs;
+  const cfgBatchSize = isHistorical ? configOverride?.batchSize : config.batchSize;
+  const cfgLearningRate = isHistorical ? configOverride?.learningRate : config.learningRate;
+  const cfgMaxSteps = isHistorical ? configOverride?.maxSteps : config.maxSteps;
+  const cfgContextLength = isHistorical ? configOverride?.contextLength : config.contextLength;
+  const cfgWarmupSteps = isHistorical ? configOverride?.warmupSteps : config.warmupSteps;
+  const cfgOptimizerType = isHistorical ? configOverride?.optimizerType : config.optimizerType;
+  const cfgLoraRank = isHistorical ? configOverride?.loraRank : config.loraRank;
+  const cfgLoraAlpha = isHistorical ? configOverride?.loraAlpha : config.loraAlpha;
+  const cfgLoraDropout = isHistorical ? configOverride?.loraDropout : config.loraDropout;
+  const cfgLoraVariant = isHistorical ? configOverride?.loraVariant : config.loraVariant;
 
   const optimizerLabel =
-    OPTIMIZER_OPTIONS.find((o) => o.value === config.optimizerType)?.label ??
-    config.optimizerType;
+    OPTIMIZER_OPTIONS.find((o) => o.value === cfgOptimizerType)?.label ??
+    cfgOptimizerType;
 
   const configItems: ConfigGroup[] = [
     {
       section: "Hyperparams",
       rows: [
-        configRow("Epochs", config.epochs),
-        configRow("Batch size", config.batchSize),
-        configRow("Learning rate", config.learningRate),
+        configRow("Epochs", cfgEpochs),
+        configRow("Batch size", cfgBatchSize),
+        configRow("Learning rate", cfgLearningRate),
         configRow("Optimizer", optimizerLabel),
-        configRow("Max steps", config.maxSteps),
-        configRow("Context length", config.contextLength),
-        configRow("Warmup steps", config.warmupSteps),
+        configRow("Max steps", cfgMaxSteps),
+        configRow("Context length", cfgContextLength),
+        configRow("Warmup steps", cfgWarmupSteps),
       ],
     },
-    ...(config.trainingMethod !== "full"
+    ...(data.trainingMethod !== "full"
       ? [
           {
             section: "LoRA",
             rows: [
-              configRow("Rank", config.loraRank),
-              configRow("Alpha", config.loraAlpha),
-              configRow("Dropout", config.loraDropout),
-              configRow("Variant", config.loraVariant),
+              configRow("Rank", cfgLoraRank),
+              configRow("Alpha", cfgLoraAlpha),
+              configRow("Dropout", cfgLoraDropout),
+              configRow("Variant", cfgLoraVariant),
             ],
           },
         ]
@@ -205,30 +199,34 @@ export function ProgressSection(): ReactElement {
     <SectionCard
       icon={<HugeiconsIcon icon={ChartAverageIcon} className="size-5" />}
       title="Training Progress"
-      description={runtime.message || "Live training metrics"}
+      description={data.message || "Live training metrics"}
       accent="emerald"
       className="shadow-border border border-border/60 bg-card/90 ring-0 backdrop-blur-sm"
       headerAction={
-        <TrainingHeaderActions
-          configItems={configItems}
-          isTrainingRunning={runtime.isTrainingRunning}
-          onOpenStopDialog={setStopDialogOpen}
-          onRequestStop={requestStop}
-          stopDialogOpen={stopDialogOpen}
-          stopRequested={stopRequested}
-        />
+        isHistorical ? (
+          <ConfigPopoverButton configItems={configItems} />
+        ) : (
+          <LiveTrainingHeaderActions
+            configItems={configItems}
+            isTrainingRunning={data.isTrainingRunning}
+            onOpenStopDialog={setStopDialogOpen}
+            stopDialogOpen={stopDialogOpen}
+            stopRequested={stopRequested}
+            onSetStopRequested={setStopRequestedLocal}
+          />
+        )
       }
     >
       <div className="grid grid-cols-1 gap-5 lg:grid-cols-[minmax(0,1.2fr)_minmax(18rem,0.8fr)]">
         <div className="flex flex-col gap-4">
           <div className="flex flex-wrap items-center gap-2">
             <span
-              className={`rounded-full px-2.5 py-1 text-[10px] font-semibold ${phaseColors[runtime.phase]}`}
+              className={`rounded-full px-2.5 py-1 text-[10px] font-semibold ${phaseColors[data.phase]}`}
             >
-              {phaseLabel[runtime.phase]}
+              {phaseLabel[data.phase]}
             </span>
             <span className="text-[10px] tabular-nums text-muted-foreground">
-              Epoch {runtime.currentEpoch.toFixed(2)}
+              Epoch {formatNumber(data.currentEpoch, 2)}
             </span>
             <span className="rounded-full border border-border/60 px-2.5 py-1 text-[10px] font-medium tabular-nums text-muted-foreground">
               {pct}% complete
@@ -238,22 +236,24 @@ export function ProgressSection(): ReactElement {
           <div className="flex flex-col gap-2">
             <div className="flex justify-between text-xs text-muted-foreground">
               <span>
-                Step {runtime.currentStep} / {runtime.totalSteps || "--"}
+                Step {data.currentStep} / {data.totalSteps || "--"}
               </span>
               <span>{pct}%</span>
             </div>
             <Progress value={pct} className="h-2 bg-foreground/[0.05]" />
           </div>
 
-          <MilestoneCallout
-            showCompletedHint={showCompletedHint}
-            showHalfwayHint={showHalfwayHint}
-            onCompareInChat={handleCompareInChat}
-          />
+          {!isHistorical && (
+            <MilestoneCallout
+              showCompletedHint={showCompletedHint}
+              showHalfwayHint={showHalfwayHint}
+              onCompareInChat={handleCompareInChat}
+            />
+          )}
 
-          {runtime.error && (
+          {data.error && (
             <p className="rounded-2xl border border-destructive/30 bg-destructive/5 px-3 py-2 text-xs text-red-500 leading-relaxed">
-              {runtime.error}
+              {data.error}
             </p>
           )}
 
@@ -262,97 +262,196 @@ export function ProgressSection(): ReactElement {
               label="Loss"
               valueClassName="text-2xl font-bold tracking-tight"
             >
-              {stoppedLoss.toFixed(4)}
+              {stoppedLoss != null ? stoppedLoss.toFixed(4) : "--"}
             </MetricStat>
-            <MetricStat label="LR">{stoppedLr.toExponential(2)}</MetricStat>
+            <MetricStat label="LR">{stoppedLr != null ? stoppedLr.toExponential(2) : "--"}</MetricStat>
             <MetricStat label="Grad Norm">
               {formatNumber(stoppedGradNorm, 3)}
             </MetricStat>
             <MetricStat label="Model" valueClassName="truncate">
-              {config.selectedModel ?? "--"}
+              {data.modelName || "--"}
             </MetricStat>
             <MetricStat label="Method">
-              {config.trainingMethod === "qlora" ? "QLoRA" : config.trainingMethod === "lora" ? "LoRA" : "Full"}
+              {data.trainingMethod === "qlora" ? "QLoRA" : data.trainingMethod === "lora" ? "LoRA" : "Full"}
             </MetricStat>
           </div>
 
           <div className="flex flex-wrap gap-x-4 gap-y-1 text-xs text-muted-foreground">
             <span>Elapsed: {formatDuration(elapsed)}</span>
-            <span>ETA: {formatDuration(eta)}</span>
+            {!isHistorical && <span>ETA: {formatDuration(eta)}</span>}
             <span>
               {stepsPerSecond == null
                 ? "-- steps/s"
                 : `${stepsPerSecond.toFixed(2)} steps/s`}
             </span>
-            {runtime.currentNumTokens != null && (
-              <span>Tokens: {runtime.currentNumTokens}</span>
+            {data.currentNumTokens != null && (
+              <span>Tokens: {data.currentNumTokens}</span>
             )}
           </div>
         </div>
 
-        <div className="flex flex-col gap-3">
-          <div className="flex items-center justify-between">
-            <p className="text-xs font-medium text-muted-foreground">
-              GPU Monitor
-            </p>
-            <span className="text-[11px] text-muted-foreground">Live</span>
-          </div>
-          <div className="grid grid-cols-2 gap-2.5">
-            <GpuStat
-              label="Utilization"
-              icon={
-                <HugeiconsIcon
-                  icon={DashboardSpeed01Icon}
-                  className="size-3.5"
-                />
-              }
-              value={
-                gpu.gpu_utilization_pct != null
-                  ? `${gpu.gpu_utilization_pct}%`
-                  : "--"
-              }
-              pct={gpu.gpu_utilization_pct ?? 0}
-            />
-            <GpuStat
-              label="Temperature"
-              icon={
-                <HugeiconsIcon icon={TemperatureIcon} className="size-3.5" />
-              }
-              value={
-                gpu.temperature_c != null ? `${gpu.temperature_c}°C` : "--"
-              }
-              pct={gpu.temperature_c ?? 0}
-              max={100}
-            />
-            <GpuStat
-              label="VRAM"
-              icon={<HugeiconsIcon icon={RamMemoryIcon} className="size-3.5" />}
-              value={
-                gpu.vram_used_gb != null && gpu.vram_total_gb != null
-                  ? `${gpu.vram_used_gb} / ${gpu.vram_total_gb} GB`
-                  : "--"
-              }
-              pct={gpu.vram_utilization_pct ?? 0}
-            />
-            <GpuStat
-              label="Power"
-              icon={<HugeiconsIcon icon={ZapIcon} className="size-3.5" />}
-              value={
-                gpu.power_draw_w != null
-                  ? gpu.power_limit_w != null
-                    ? `${gpu.power_draw_w} / ${gpu.power_limit_w} W`
-                    : `${gpu.power_draw_w} W`
-                  : "--"
-              }
-              pct={gpu.power_utilization_pct ?? 0}
-            />
-          </div>
-        </div>
+        {!isHistorical && (
+          <LiveGpuPanel isTrainingRunning={data.isTrainingRunning} />
+        )}
       </div>
     </SectionCard>
   );
 }
 
+function LiveGpuPanel({
+  isTrainingRunning,
+}: {
+  isTrainingRunning: boolean;
+}): ReactElement {
+  const gpu = useGpuUtilization(isTrainingRunning);
+
+  return (
+    <div className="flex flex-col gap-3">
+      <div className="flex items-center justify-between">
+        <p className="text-xs font-medium text-muted-foreground">
+          GPU Monitor
+        </p>
+        <span className="text-[11px] text-muted-foreground">Live</span>
+      </div>
+      <div className="grid grid-cols-2 gap-2.5">
+        <GpuStat
+          label="Utilization"
+          icon={
+            <HugeiconsIcon
+              icon={DashboardSpeed01Icon}
+              className="size-3.5"
+            />
+          }
+          value={
+            gpu.gpu_utilization_pct != null
+              ? `${gpu.gpu_utilization_pct}%`
+              : "--"
+          }
+          pct={gpu.gpu_utilization_pct ?? 0}
+        />
+        <GpuStat
+          label="Temperature"
+          icon={
+            <HugeiconsIcon icon={TemperatureIcon} className="size-3.5" />
+          }
+          value={
+            gpu.temperature_c != null ? `${gpu.temperature_c}°C` : "--"
+          }
+          pct={gpu.temperature_c ?? 0}
+          max={100}
+        />
+        <GpuStat
+          label="VRAM"
+          icon={<HugeiconsIcon icon={RamMemoryIcon} className="size-3.5" />}
+          value={
+            gpu.vram_used_gb != null && gpu.vram_total_gb != null
+              ? `${gpu.vram_used_gb} / ${gpu.vram_total_gb} GB`
+              : "--"
+          }
+          pct={gpu.vram_utilization_pct ?? 0}
+        />
+        <GpuStat
+          label="Power"
+          icon={<HugeiconsIcon icon={ZapIcon} className="size-3.5" />}
+          value={
+            gpu.power_draw_w != null
+              ? gpu.power_limit_w != null
+                ? `${gpu.power_draw_w} / ${gpu.power_limit_w} W`
+                : `${gpu.power_draw_w} W`
+              : "--"
+          }
+          pct={gpu.power_utilization_pct ?? 0}
+        />
+      </div>
+    </div>
+  );
+}
+
+function LiveTrainingHeaderActions({
+  configItems,
+  isTrainingRunning,
+  onOpenStopDialog,
+  stopDialogOpen,
+  stopRequested,
+  onSetStopRequested,
+}: {
+  configItems: ConfigGroup[];
+  isTrainingRunning: boolean;
+  onOpenStopDialog: (open: boolean) => void;
+  stopDialogOpen: boolean;
+  stopRequested: boolean;
+  onSetStopRequested: (v: boolean) => void;
+}): ReactElement {
+  const { stopTrainingRun } = useTrainingActions();
+
+  const requestStop = async (saveCheckpoint: boolean) => {
+    onSetStopRequested(true);
+    onOpenStopDialog(false);
+    useTrainingRuntimeStore.getState().setStopRequested(true);
+    try {
+      const ok = await stopTrainingRun(saveCheckpoint);
+      if (!ok) {
+        onSetStopRequested(false);
+      }
+    } catch {
+      onSetStopRequested(false);
+    }
+  };
+
+  return (
+    <TrainingHeaderActions
+      configItems={configItems}
+      isTrainingRunning={isTrainingRunning}
+      onOpenStopDialog={onOpenStopDialog}
+      onRequestStop={requestStop}
+      stopDialogOpen={stopDialogOpen}
+      stopRequested={stopRequested}
+    />
+  );
+}
+
+function ConfigPopoverButton({
+  configItems,
+}: {
+  configItems: ConfigGroup[];
+}): ReactElement {
+  return (
+    <Popover>
+      <PopoverTrigger asChild={true}>
+        <Button
+          type="button"
+          variant="ghost"
+          size="icon-sm"
+          className="rounded-full text-muted-foreground hover:bg-muted hover:text-foreground"
+          aria-label="Open training config"
+        >
+          <HugeiconsIcon icon={Notebook01Icon} className="size-4" />
+        </Button>
+      </PopoverTrigger>
+      <PopoverContent className="w-72" align="end">
+        <div className="flex flex-col gap-3">
+          <p className="text-xs font-semibold">Training Config</p>
+          {configItems.map((group) => (
+            <div key={group.section} className="flex flex-col gap-1">
+              <p className="text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
+                {group.section}
+              </p>
+              {group.rows.map(([label, value]) => (
+                <div key={label} className="flex justify-between text-xs">
+                  <span className="text-muted-foreground">{label}</span>
+                  <span className="font-medium tabular-nums">
+                    {value == null || value === "" ? "--" : String(value)}
+                  </span>
+                </div>
+              ))}
+            </div>
+          ))}
+        </div>
+      </PopoverContent>
+    </Popover>
+  );
+}
+
 function TrainingHeaderActions({
   configItems,
   isTrainingRunning,
@@ -370,39 +469,7 @@ function TrainingHeaderActions({
 }): ReactElement {
   return (
     <div className="flex items-center gap-2">
-      <Popover>
-        <PopoverTrigger asChild={true}>
-          <Button
-            type="button"
-            variant="ghost"
-            size="icon-sm"
-            className="rounded-full text-muted-foreground hover:bg-muted hover:text-foreground"
-            aria-label="Open training config"
-          >
-            <HugeiconsIcon icon={Notebook01Icon} className="size-4" />
-          </Button>
-        </PopoverTrigger>
-        <PopoverContent className="w-72" align="end">
-          <div className="flex flex-col gap-3">
-            <p className="text-xs font-semibold">Training Config</p>
-            {configItems.map((group) => (
-              <div key={group.section} className="flex flex-col gap-1">
-                <p className="text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
-                  {group.section}
-                </p>
-                {group.rows.map(([label, value]) => (
-                  <div key={label} className="flex justify-between text-xs">
-                    <span className="text-muted-foreground">{label}</span>
-                    <span className="font-medium tabular-nums">
-                      {String(value)}
-                    </span>
-                  </div>
-                ))}
-              </div>
-            ))}
-          </div>
-        </PopoverContent>
-      </Popover>
+      <ConfigPopoverButton configItems={configItems} />
       <ChartSettingsSheet />
       <AlertDialog open={stopDialogOpen} onOpenChange={onOpenStopDialog}>
         <Button
@@ -518,25 +585,21 @@ function MetricStat({
   );
 }
 
-function lastNonZeroValue(points: { value: number }[]): number | null {
-  for (let i = points.length - 1; i >= 0; i -= 1) {
-    const value = points[i]?.value;
-    if (Number.isFinite(value) && value !== 0) {
-      return value;
-    }
-  }
-  return null;
+function lastValue(points: { value: number }[]): number | null {
+  if (points.length === 0) return null;
+  const v = points[points.length - 1]?.value;
+  return v != null && Number.isFinite(v) ? v : null;
 }
 
 function getDisplayMetric(
   isTrainingRunning: boolean,
-  currentValue: number,
+  currentValue: number | null,
   history: { value: number }[],
-): number {
+): number | null {
   if (isTrainingRunning) {
-    return currentValue;
+    return currentValue != null ? currentValue : null;
   }
-  return lastNonZeroValue(history) ?? currentValue;
+  return lastValue(history) ?? (currentValue != null ? currentValue : null);
 }
 
 function GpuStat({

studio/frontend/src/features/studio/studio-page.tsx

@@ -1,26 +1,28 @@
 // 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 { Button } from "@/components/ui/button";
 import {
   shouldShowTrainingView,
   useDatasetPreviewDialogStore,
-  useTrainingActions,
   useTrainingConfigStore,
   useTrainingRuntimeLifecycle,
   useTrainingRuntimeStore,
 } from "@/features/training";
 import { GuidedTour, useGuidedTourController } from "@/features/tour";
 import { studioTourSteps, studioTrainingTourSteps } from "./tour";
+import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
+import { Button } from "@/components/ui/button";
 import { ArrowLeft01Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import { type ReactElement, useEffect } from "react";
+import { type ReactElement, useEffect, useState } from "react";
 import { DatasetPreviewDialog } from "./sections/dataset-preview-dialog";
 import { DatasetSection } from "./sections/dataset-section";
 import { ModelSection } from "./sections/model-section";
 import { ParamsSection } from "./sections/params-section";
 import { TrainingSection } from "./sections/training-section";
-import { TrainingView } from "./training-view";
+import { LiveTrainingView } from "./live-training-view";
+import { HistoricalTrainingView } from "./historical-training-view";
+import { HistoryCardGrid } from "./history-card-grid";
 
 const STUDIO_TOUR_KEY = "tour:studio:v1";
 
@@ -28,11 +30,10 @@ export function StudioPage(): ReactElement {
   useTrainingRuntimeLifecycle();
   const showTrainingView = useTrainingRuntimeStore(shouldShowTrainingView);
   const isTrainingRunning = useTrainingRuntimeStore((state) => state.isTrainingRunning);
+  const currentJobId = useTrainingRuntimeStore((state) => state.jobId);
   const runtimeMessage = useTrainingRuntimeStore((state) => state.message);
-  const runtimePhase = useTrainingRuntimeStore((state) => state.phase);
   const isHydratingRuntime = useTrainingRuntimeStore((state) => state.isHydrating);
   const hasHydratedRuntime = useTrainingRuntimeStore((state) => state.hasHydrated);
-  const { dismissTrainingRun } = useTrainingActions();
 
   const config = useTrainingConfigStore();
   const selectedModel = useTrainingConfigStore((s) => s.selectedModel);
@@ -47,19 +48,22 @@ export function StudioPage(): ReactElement {
   const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData);
   const closeDialog = useDatasetPreviewDialogStore((s) => s.close);
 
-  const stopRequested = useTrainingRuntimeStore((state) => state.stopRequested);
-  const canGoBack =
-    showTrainingView &&
-    !isHydratingRuntime &&
-    (stopRequested ||
-      (!isTrainingRunning &&
-        (runtimePhase === "stopped" ||
-          runtimePhase === "error" ||
-          runtimePhase === "completed" ||
-          runtimePhase === "idle")));
+  const [requestedTab, setRequestedTab] = useState("configure");
+  const [selectedHistoryRunId, setSelectedHistoryRunId] = useState<string | null>(null);
+
+  // Derive activeTab: auto-switch to "current-run" only while training is
+  // genuinely running.  Once training ends, honour whatever tab the user clicks.
+  // If requestedTab is "current-run" but there's nothing to show, fall back to "configure".
+  const activeTab =
+    isTrainingRunning && requestedTab !== "history"
+      ? "current-run"
+      : requestedTab === "current-run" && !showTrainingView
+        ? "configure"
+        : requestedTab;
+
   const tourEnabled = hasHydratedRuntime && !isHydratingRuntime;
-  const isConfigTour = !showTrainingView;
-  const tourSteps = showTrainingView ? studioTrainingTourSteps : studioTourSteps;
+  const isConfigTour = activeTab === "configure";
+  const tourSteps = activeTab === "current-run" ? studioTrainingTourSteps : studioTourSteps;
   const tour = useGuidedTourController({
     id: "studio",
     steps: tourSteps,
@@ -71,13 +75,36 @@ export function StudioPage(): ReactElement {
   const setTourOpen = tour.setOpen;
   useEffect(() => {
     setTourOpen(false);
-  }, [showTrainingView, setTourOpen]);
+  }, [activeTab, setTourOpen]);
+
+  // When training auto-switches us to "current-run", persist that in
+  // requestedTab so the user stays on results after training ends.
+  useEffect(() => {
+    if (isTrainingRunning && requestedTab !== "history" && requestedTab !== "current-run") {
+      setRequestedTab("current-run");
+      setSelectedHistoryRunId(null);
+    }
+  }, [isTrainingRunning, requestedTab]);
 
   useEffect(() => {
     ensureModelDefaultsLoaded();
     ensureDatasetChecked();
   }, [selectedModel, ensureModelDefaultsLoaded, ensureDatasetChecked]);
 
+  function handleTabChange(value: string) {
+    setRequestedTab(value);
+    if (value !== "history") {
+      setSelectedHistoryRunId(null);
+    }
+  }
+
+  const subtitle = (() => {
+    if (activeTab === "current-run") return runtimeMessage || "Training in progress";
+    if (activeTab === "history")
+      return selectedHistoryRunId ? "Viewing past run" : "View past training runs";
+    return "Configure and start training";
+  })();
+
   return (
     <div className="relative min-h-screen overflow-hidden bg-background">
       <main className="relative z-10 mx-auto max-w-7xl px-4 py-4 sm:px-6">
@@ -100,42 +127,69 @@ export function StudioPage(): ReactElement {
           isVlm={config.isVisionModel && config.isDatasetImage === true}
         />
 
-        {canGoBack && (
-          <Button
-            variant="ghost"
-            size="sm"
-            className="mb-2 cursor-pointer gap-1.5 text-muted-foreground"
-            onClick={() => void dismissTrainingRun()}
-          >
-            <HugeiconsIcon icon={ArrowLeft01Icon} className="size-4" />
-            Back to configuration
-          </Button>
-        )}
-
         <div className="mb-6 flex flex-col gap-0.5 sm:mb-8">
           <h1 className="text-2xl font-semibold tracking-tight">
             Fine-tuning Studio
           </h1>
-          <p className="text-sm text-muted-foreground">
-            {showTrainingView
-              ? runtimeMessage || "Training in progress"
-              : "Configure and start training"}
-          </p>
+          <p className="text-sm text-muted-foreground">{subtitle}</p>
         </div>
 
         {!hasHydratedRuntime && isHydratingRuntime ? (
           <div className="rounded-xl border bg-card p-8 text-sm text-muted-foreground">
             Loading training runtime...
           </div>
-        ) : showTrainingView ? (
-          <TrainingView />
         ) : (
-          <div className="grid grid-cols-1 items-start gap-4 md:grid-cols-2 md:gap-6 xl:grid-cols-12">
-            <ModelSection />
-            <DatasetSection />
-            <ParamsSection />
-            <TrainingSection />
-          </div>
+          <Tabs value={activeTab} onValueChange={handleTabChange}>
+            <div className="flex items-center gap-3">
+              {selectedHistoryRunId && activeTab === "history" && (
+                <Button
+                  variant="ghost"
+                  size="icon-sm"
+                  className="rounded-full text-muted-foreground"
+                  onClick={() => setSelectedHistoryRunId(null)}
+                  aria-label="Back to history"
+                >
+                  <HugeiconsIcon icon={ArrowLeft01Icon} className="size-4" />
+                </Button>
+              )}
+              <TabsList variant="line">
+                <TabsTrigger value="configure" disabled={isTrainingRunning}>
+                  Configure
+                </TabsTrigger>
+                <TabsTrigger value="current-run" disabled={!showTrainingView}>
+                  Current Run
+                </TabsTrigger>
+                <TabsTrigger value="history">History</TabsTrigger>
+              </TabsList>
+            </div>
+
+            <TabsContent value="configure">
+              <div className="grid grid-cols-1 items-start gap-4 md:grid-cols-2 md:gap-6 xl:grid-cols-12">
+                <ModelSection />
+                <DatasetSection />
+                <ParamsSection />
+                <TrainingSection />
+              </div>
+            </TabsContent>
+
+            <TabsContent value="current-run">
+              <LiveTrainingView />
+            </TabsContent>
+
+            <TabsContent value="history">
+              {selectedHistoryRunId ? (
+                <HistoricalTrainingView runId={selectedHistoryRunId} />
+              ) : (
+                <HistoryCardGrid onSelectRun={(runId) => {
+                  if (runId === currentJobId && isTrainingRunning) {
+                    handleTabChange("current-run");
+                  } else {
+                    setSelectedHistoryRunId(runId);
+                  }
+                }} />
+              )}
+            </TabsContent>
+          </Tabs>
         )}
       </main>
     </div>

studio/frontend/src/features/studio/training-view.tsx

@@ -1,57 +0,0 @@
-// 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 { cn } from "@/lib/utils";
-import { useTrainingRuntimeStore } from "@/features/training";
-import type { ReactElement } from "react";
-import { useShallow } from "zustand/react/shallow";
-import { ChartsSection } from "./sections/charts-section";
-import { ProgressSection } from "./sections/progress-section";
-import { TrainingStartOverlay } from "./training-start-overlay";
-
-export function TrainingView(): ReactElement {
-  const runtime = useTrainingRuntimeStore(
-    useShallow((state) => ({
-      phase: state.phase,
-      message: state.message,
-      currentStep: state.currentStep,
-      firstStepReceived: state.firstStepReceived,
-      isStarting: state.isStarting,
-    })),
-  );
-
-  const isPreparingPhase =
-    runtime.phase === "downloading_model" ||
-    runtime.phase === "downloading_dataset" ||
-    runtime.phase === "loading_model" ||
-    runtime.phase === "loading_dataset" ||
-    runtime.phase === "configuring";
-  const isWaitingForFirstStep =
-    runtime.phase === "training" && !runtime.firstStepReceived;
-  const showOverlay =
-    runtime.isStarting ||
-    isPreparingPhase ||
-    (isWaitingForFirstStep && runtime.currentStep <= 0);
-
-  return (
-    <div className={cn("relative", showOverlay && "min-h-[72vh]")}>
-      <div
-        className={cn(
-          "relative z-10 flex flex-col gap-6 transition-[filter]",
-          showOverlay && "blur",
-        )}
-      >
-        <div data-tour="studio-training-progress">
-          <ProgressSection />
-        </div>
-        <ChartsSection />
-      </div>
-      {showOverlay ? (
-        <TrainingStartOverlay
-          message={runtime.message}
-          currentStep={runtime.currentStep}
-        />
-      ) : null}
-    </div>
-  );
-}

studio/frontend/src/features/training/api/history-api.ts

@@ -0,0 +1,59 @@
+// 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 { authFetch } from "@/features/auth";
+import type {
+  TrainingRunDeleteResponse,
+  TrainingRunDetailResponse,
+  TrainingRunListResponse,
+} from "../types/history";
+
+async function readError(response: Response): Promise<string> {
+  try {
+    const payload = (await response.json()) as { detail?: string; message?: string };
+    return payload.detail || payload.message || `Request failed (${response.status})`;
+  } catch {
+    return `Request failed (${response.status})`;
+  }
+}
+
+async function parseJson<T>(response: Response): Promise<T> {
+  if (!response.ok) {
+    throw new Error(await readError(response));
+  }
+  return (await response.json()) as T;
+}
+
+export async function listTrainingRuns(
+  limit = 50,
+  offset = 0,
+  signal?: AbortSignal,
+): Promise<TrainingRunListResponse> {
+  const response = await authFetch(
+    `/api/train/runs?limit=${limit}&offset=${offset}`,
+    { signal },
+  );
+  return parseJson<TrainingRunListResponse>(response);
+}
+
+export async function getTrainingRun(
+  runId: string,
+  signal?: AbortSignal,
+): Promise<TrainingRunDetailResponse> {
+  const response = await authFetch(
+    `/api/train/runs/${encodeURIComponent(runId)}`,
+    { signal },
+  );
+  return parseJson<TrainingRunDetailResponse>(response);
+}
+
+export async function deleteTrainingRun(
+  runId: string,
+  signal?: AbortSignal,
+): Promise<TrainingRunDeleteResponse> {
+  const response = await authFetch(
+    `/api/train/runs/${encodeURIComponent(runId)}`,
+    { method: "DELETE", signal },
+  );
+  return parseJson<TrainingRunDeleteResponse>(response);
+}

studio/frontend/src/features/training/index.ts

@@ -14,6 +14,14 @@ export { useDatasetPreviewDialogStore } from "./stores/dataset-preview-dialog-st
 export { uploadTrainingDataset } from "./api/datasets-api";
 export { listLocalModels } from "./api/models-api";
 export type { LocalModelInfo } from "./api/models-api";
-export type { TrainingPhase } from "./types/runtime";
+export type { TrainingPhase, TrainingViewData, TrainingSeriesPoint } from "./types/runtime";
+export type {
+  TrainingRunSummary,
+  TrainingRunListResponse,
+  TrainingRunMetrics,
+  TrainingRunDetailResponse,
+  TrainingRunDeleteResponse,
+} from "./types/history";
+export { listTrainingRuns, getTrainingRun, deleteTrainingRun } from "./api/history-api";
 export { parseYamlConfig, serializeConfigToYaml } from "./lib/yaml-config";
 export { validateTrainingConfig } from "./lib/validation";

studio/frontend/src/features/training/types/history.ts

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+export interface TrainingRunSummary {
+  id: string;
+  status: "running" | "completed" | "stopped" | "error";
+  model_name: string;
+  dataset_name: string;
+  started_at: string;
+  ended_at: string | null;
+  total_steps: number | null;
+  final_step: number | null;
+  final_loss: number | null;
+  output_dir: string | null;
+  duration_seconds: number | null;
+  error_message: string | null;
+  loss_sparkline: number[] | null;
+}
+
+export interface TrainingRunListResponse {
+  runs: TrainingRunSummary[];
+  total: number;
+}
+
+export interface TrainingRunMetrics {
+  step_history: number[];
+  loss_history: number[];
+  loss_step_history: number[];
+  lr_history: number[];
+  lr_step_history: number[];
+  grad_norm_history: number[];
+  grad_norm_step_history: number[];
+  eval_loss_history: number[];
+  eval_step_history: number[];
+  final_epoch: number | null;
+  final_num_tokens: number | null;
+}
+
+export interface TrainingRunDetailResponse {
+  run: TrainingRunSummary;
+  config: Record<string, unknown>;
+  metrics: TrainingRunMetrics;
+}
+
+export interface TrainingRunDeleteResponse {
+  status: string;
+  message: string;
+}

studio/frontend/src/features/training/types/runtime.ts

@@ -53,8 +53,8 @@ export interface TrainingProgressPayload {
   job_id: string;
   step: number;
   total_steps: number;
-  loss: number;
-  learning_rate: number;
+  loss: number | null;
+  learning_rate: number | null;
   progress_percent: number;
   epoch: number | null;
   elapsed_seconds: number | null;
@@ -118,3 +118,32 @@ export interface TrainingRuntimeActions {
 }
 
 export type TrainingRuntimeStore = TrainingRuntimeState & TrainingRuntimeActions;
+
+export interface TrainingViewData {
+  // Current metrics (for ProgressSection)
+  phase: TrainingPhase;
+  currentStep: number;
+  totalSteps: number;
+  currentLoss: number | null;
+  currentLearningRate: number | null;
+  currentGradNorm: number | null;
+  currentEpoch: number | null;
+  currentNumTokens: number | null;
+  progressPercent: number;
+  elapsedSeconds: number | null;
+  etaSeconds: number | null;
+  evalEnabled: boolean;
+  message: string;
+  error: string | null;
+  isTrainingRunning: boolean;
+
+  // Config summary
+  modelName: string;
+  trainingMethod: string;
+
+  // Time-series (for ChartsSection)
+  lossHistory: TrainingSeriesPoint[];
+  lrHistory: TrainingSeriesPoint[];
+  gradNormHistory: TrainingSeriesPoint[];
+  evalLossHistory: TrainingSeriesPoint[];
+}

studio/frontend/src/lib/latex.ts

@@ -0,0 +1,97 @@
+// Adapted from LibreChat's latex.ts
+// https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts
+//
+// Escapes currency dollar signs so they are not misinterpreted as LaTeX math
+// delimiters when singleDollarTextMath is enabled.
+
+/**
+ * Matches a single $ followed by a number pattern (currency), e.g.:
+ *   $5, $1,000, $5.99, $100K, $3.5M
+ *
+ * Does NOT match:
+ *   $$ (display math), \$ (already escaped), $\alpha (LaTeX command)
+ */
+const CURRENCY_REGEX =
+  /(?<![\\$])\$(?!\$)(?=\d+(?:,\d{3})*(?:\.\d+)?[KMBkmb]?(?:\s|$|[^a-zA-Z\d]))/g;
+
+/**
+ * Find regions inside code blocks (``` ... ``` and ` ... `) so we can skip them.
+ * Returns sorted array of [start, end] index pairs.
+ */
+function findCodeBlockRegions(content: string): Array<[number, number]> {
+  const regions: Array<[number, number]> = [];
+
+  // Fenced code blocks: ```...```
+  const fencedRe = /```[\s\S]*?```/g;
+  let match: RegExpExecArray | null;
+  while ((match = fencedRe.exec(content)) !== null) {
+    regions.push([match.index, match.index + match[0].length]);
+  }
+
+  // Inline code: `...` (but not inside fenced blocks -- we filter below)
+  const inlineRe = /`[^`\n]+`/g;
+  while ((match = inlineRe.exec(content)) !== null) {
+    const start = match.index;
+    const end = start + match[0].length;
+    // Skip if this backtick span falls inside a fenced block
+    let inside = false;
+    for (const [rs, re] of regions) {
+      if (start >= rs && end <= re) {
+        inside = true;
+        break;
+      }
+    }
+    if (!inside) {
+      regions.push([start, end]);
+    }
+  }
+
+  // Sort by start position for binary search
+  regions.sort((a, b) => a[0] - b[0]);
+  return regions;
+}
+
+/**
+ * Binary search to check if a position falls inside any code region.
+ */
+function isInCodeBlock(
+  position: number,
+  regions: Array<[number, number]>,
+): boolean {
+  let lo = 0;
+  let hi = regions.length - 1;
+  while (lo <= hi) {
+    const mid = (lo + hi) >>> 1;
+    const [start, end] = regions[mid];
+    if (position < start) {
+      hi = mid - 1;
+    } else if (position >= end) {
+      lo = mid + 1;
+    } else {
+      return true;
+    }
+  }
+  return false;
+}
+
+/**
+ * Preprocess a markdown string to escape currency dollar signs so they are not
+ * parsed as LaTeX math delimiters.
+ *
+ * - `$5` alone becomes `\$5` (currency, not math)
+ * - `$\alpha$` is untouched (real LaTeX)
+ * - `$$E = mc^2$$` is untouched (display math)
+ * - Currency inside code blocks/spans is untouched
+ */
+export function preprocessLaTeX(content: string): string {
+  if (!content.includes("$")) return content;
+
+  const codeRegions = findCodeBlockRegions(content);
+
+  return content.replace(CURRENCY_REGEX, (match, offset) => {
+    if (isInCodeBlock(offset, codeRegions)) {
+      return match;
+    }
+    return "\\" + match;
+  });
+}

studio/install_python_stack.py

@@ -22,20 +22,20 @@ from pathlib import Path
 
 IS_WINDOWS = sys.platform == "win32"
 
-# ── Verbosity control ──────────────────────────────────────────────────────────
+# -- Verbosity control ----------------------------------------------------------
 # By default the installer shows a minimal progress bar (one line, in-place).
 # Set UNSLOTH_VERBOSE=1 in the environment to restore full per-step output:
 #   Linux/Mac:  UNSLOTH_VERBOSE=1 ./studio/setup.sh
 #   Windows:    $env:UNSLOTH_VERBOSE="1" ; .\studio\setup.ps1
 VERBOSE: bool = os.environ.get("UNSLOTH_VERBOSE", "0") == "1"
 
-# Progress bar state — updated by _progress() as each install step runs.
+# Progress bar state -- updated by _progress() as each install step runs.
 # _TOTAL counts: pip-upgrade + 7 shared steps + triton (non-Windows) + local-plugin + finalize
 # Update _TOTAL here if you add or remove install steps in install_python_stack().
 _STEP: int = 0
 _TOTAL: int = 0  # set at runtime in install_python_stack() based on platform
 
-# ── Paths ──────────────────────────────────────────────────────────────
+# -- Paths --------------------------------------------------------------
 SCRIPT_DIR = Path(__file__).resolve().parent
 REQ_ROOT = SCRIPT_DIR / "backend" / "requirements"
 SINGLE_ENV = REQ_ROOT / "single-env"
@@ -44,7 +44,39 @@ LOCAL_DD_UNSTRUCTURED_PLUGIN = (
     SCRIPT_DIR / "backend" / "plugins" / "data-designer-unstructured-seed"
 )
 
-# ── Color support ──────────────────────────────────────────────────────
+# -- Unicode-safe printing ---------------------------------------------
+# On Windows the default console encoding can be a legacy code page
+# (e.g. CP1252) that cannot represent Unicode glyphs such as ✅ or ❌.
+# _safe_print() gracefully degrades to ASCII equivalents so the
+# installer never crashes just because of a status glyph.
+
+_UNICODE_TO_ASCII: dict[str, str] = {
+    "\u2705": "[OK]",  # ✅
+    "\u274c": "[FAIL]",  # ❌
+    "\u26a0\ufe0f": "[!]",  # ⚠️  (warning + variation selector)
+    "\u26a0": "[!]",  # ⚠  (warning without variation selector)
+}
+
+
+def _safe_print(*args: object, **kwargs: object) -> None:
+    """Drop-in print() replacement that survives non-UTF-8 consoles."""
+    try:
+        print(*args, **kwargs)
+    except UnicodeEncodeError:
+        # Stringify, then swap emoji for ASCII equivalents
+        text = " ".join(str(a) for a in args)
+        for uni, ascii_alt in _UNICODE_TO_ASCII.items():
+            text = text.replace(uni, ascii_alt)
+        # Final fallback: replace any remaining unencodable chars
+        print(
+            text.encode(sys.stdout.encoding or "ascii", errors = "replace").decode(
+                sys.stdout.encoding or "ascii", errors = "replace"
+            ),
+            **kwargs,
+        )
+
+
+# -- Color support ------------------------------------------------------
 
 
 def _enable_colors() -> bool:
@@ -72,7 +104,7 @@ def _enable_colors() -> bool:
     return True  # Unix terminals support ANSI by default
 
 
-# Colors disabled — Colab and most CI runners render ANSI fine, but plain output
+# Colors disabled -- Colab and most CI runners render ANSI fine, but plain output
 # is cleaner in the notebook cell. Re-enable by setting _HAS_COLOR = _enable_colors()
 _HAS_COLOR = False
 
@@ -92,7 +124,7 @@ def _red(msg: str) -> str:
 def _progress(label: str) -> None:
     """Print an in-place progress bar for the current install step.
 
-    Uses only stdlib (sys.stdout) — no extra packages required.
+    Uses only stdlib (sys.stdout) -- no extra packages required.
     In VERBOSE mode this is a no-op; per-step labels are printed by run() instead.
     """
     global _STEP
@@ -119,7 +151,7 @@ def run(
         stderr = subprocess.STDOUT if quiet else None,
     )
     if result.returncode != 0:
-        print(_red(f"❌ {label} failed (exit code {result.returncode}):"))
+        _safe_print(_red(f"❌ {label} failed (exit code {result.returncode}):"))
         if result.stdout:
             print(result.stdout.decode(errors = "replace"))
         sys.exit(result.returncode)
@@ -129,7 +161,7 @@ def run(
 # Packages to skip on Windows (require special build steps)
 WINDOWS_SKIP_PACKAGES = {"open_spiel", "triton_kernels"}
 
-# ── uv bootstrap ──────────────────────────────────────────────────────
+# -- uv bootstrap ------------------------------------------------------
 
 USE_UV = False  # Set by _bootstrap_uv() at the start of install_python_stack()
 UV_NEEDS_SYSTEM = False  # Set by _bootstrap_uv() via probe
@@ -193,9 +225,20 @@ def _translate_pip_args_for_uv(args: tuple[str, ...]) -> list[str]:
 
 
 def _build_pip_cmd(args: tuple[str, ...]) -> list[str]:
-    """Build a standard pip install command."""
+    """Build a standard pip install command.
+
+    Strips uv-only flags like --upgrade-package that pip doesn't understand.
+    """
     cmd = [sys.executable, "-m", "pip", "install"]
-    cmd.extend(args)
+    skip_next = False
+    for arg in args:
+        if skip_next:
+            skip_next = False
+            continue
+        if arg == "--upgrade-package":
+            skip_next = True  # skip the flag and its value
+            continue
+        cmd.append(arg)
     return cmd
 
 
@@ -209,7 +252,12 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]:
     # the system Python (observed on Colab and similar environments).
     cmd.extend(["--python", sys.executable])
     cmd.extend(_translate_pip_args_for_uv(args))
-    cmd.append("--torch-backend=auto")
+    # Torch is pre-installed by install.sh/setup.ps1.  Do not add
+    # --torch-backend by default -- it can cause solver dead-ends on
+    # CPU-only machines.  Callers that need it can set UV_TORCH_BACKEND.
+    _tb = os.environ.get("UV_TORCH_BACKEND", "")
+    if _tb:
+        cmd.append(f"--torch-backend={_tb}")
     return cmd
 
 
@@ -267,7 +315,9 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None:
         text = True,
     )
     if result.returncode != 0:
-        print(_red(f"   ⚠️  Could not find package {package_name}, skipping patch"))
+        _safe_print(
+            _red(f"   ⚠️  Could not find package {package_name}, skipping patch")
+        )
         return
 
     location = None
@@ -277,7 +327,7 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None:
             break
 
     if not location:
-        print(_red(f"   ⚠️  Could not determine location of {package_name}"))
+        _safe_print(_red(f"   ⚠️  Could not determine location of {package_name}"))
         return
 
     dest = Path(location) / relative_path
@@ -285,28 +335,96 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None:
     download_file(url, dest)
 
 
-# ── Main install sequence ─────────────────────────────────────────────
+# -- Main install sequence ---------------------------------------------
 
 
 def install_python_stack() -> int:
     global USE_UV, _STEP, _TOTAL
     _STEP = 0
-    _TOTAL = 10 if IS_WINDOWS else 11
 
-    # 1. Upgrade pip (needed even with uv as fallback and for bootstrapping)
-    _progress("pip upgrade")
-    run("Upgrading pip", [sys.executable, "-m", "pip", "install", "--upgrade", "pip"])
+    # When called from install.sh (which already installed unsloth into the venv),
+    # SKIP_STUDIO_BASE=1 is set to avoid redundant reinstallation of base packages.
+    # When called from "unsloth studio update", it is NOT set so base packages
+    # (unsloth + unsloth-zoo) are always reinstalled to pick up new versions.
+    skip_base = os.environ.get("SKIP_STUDIO_BASE", "0") == "1"
+    # When --package is used, install a different package name (e.g. roland-sloth for testing)
+    package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth")
+    # When --local is used, overlay a local repo checkout after updating deps
+    local_repo = os.environ.get("STUDIO_LOCAL_REPO", "")
+    base_total = 10 if IS_WINDOWS else 11
+    _TOTAL = (base_total - 1) if skip_base else base_total
 
-    # Try to use uv for faster installs
+    # 1. Try to use uv for faster installs (must happen before pip upgrade
+    #    because uv venvs don't include pip by default)
     USE_UV = _bootstrap_uv()
 
-    # 2. Core packages: unsloth-zoo + unsloth
-    _progress("base packages")
-    pip_install(
-        "Installing base packages",
-        "--no-cache-dir",
-        req = REQ_ROOT / "base.txt",
-    )
+    # 2. Ensure pip is available (uv venvs created by install.sh don't include pip)
+    _progress("pip bootstrap")
+    if USE_UV:
+        run(
+            "Bootstrapping pip via uv",
+            [
+                "uv",
+                "pip",
+                "install",
+                "--python",
+                sys.executable,
+                "pip",
+            ],
+        )
+    else:
+        run(
+            "Upgrading pip",
+            [sys.executable, "-m", "pip", "install", "--upgrade", "pip"],
+        )
+
+    # 3. Core packages: unsloth-zoo + unsloth (or custom package name)
+    if skip_base:
+        print(_green(f"✅ {package_name} already installed — skipping base packages"))
+    elif local_repo:
+        # Local dev install: update deps from base.txt, then overlay the
+        # local checkout as an editable install (--no-deps so torch is
+        # never re-resolved).
+        _progress("base packages")
+        pip_install(
+            "Updating base packages",
+            "--no-cache-dir",
+            "--upgrade-package",
+            "unsloth",
+            "--upgrade-package",
+            "unsloth-zoo",
+            req = REQ_ROOT / "base.txt",
+        )
+        pip_install(
+            "Overlaying local repo (editable)",
+            "--no-cache-dir",
+            "--no-deps",
+            "-e",
+            local_repo,
+            constrain = False,
+        )
+    elif package_name != "unsloth":
+        # Custom package name (e.g. roland-sloth for testing) — install directly
+        _progress("base packages")
+        pip_install(
+            f"Installing {package_name}",
+            "--no-cache-dir",
+            package_name,
+        )
+    else:
+        # Update path: upgrade only unsloth + unsloth-zoo while preserving
+        # existing torch/CUDA installations.  Torch is pre-installed by
+        # install.sh / setup.ps1; --upgrade-package targets only base pkgs.
+        _progress("base packages")
+        pip_install(
+            "Updating base packages",
+            "--no-cache-dir",
+            "--upgrade-package",
+            "unsloth",
+            "--upgrade-package",
+            "unsloth-zoo",
+            req = REQ_ROOT / "base.txt",
+        )
 
     # 3. Extra dependencies
     _progress("unsloth extras")
@@ -316,7 +434,7 @@ def install_python_stack() -> int:
         req = REQ_ROOT / "extras.txt",
     )
 
-    # 3b. Extra dependencies (no-deps) — audio model support etc.
+    # 3b. Extra dependencies (no-deps) -- audio model support etc.
     _progress("extra codecs")
     pip_install(
         "Installing extras (no-deps)",
@@ -325,7 +443,7 @@ def install_python_stack() -> int:
         req = REQ_ROOT / "extras-no-deps.txt",
     )
 
-    # 4. Overrides (torchao, transformers) — force-reinstall
+    # 4. Overrides (torchao, transformers) -- force-reinstall
     _progress("dependency overrides")
     pip_install(
         "Installing dependency overrides",
@@ -393,7 +511,7 @@ def install_python_stack() -> int:
 
     # 11. Local Data Designer seed plugin
     if not LOCAL_DD_UNSTRUCTURED_PLUGIN.is_dir():
-        print(
+        _safe_print(
             _red(
                 f"❌ Missing local plugin directory: {LOCAL_DD_UNSTRUCTURED_PLUGIN}",
             ),
@@ -422,7 +540,7 @@ def install_python_stack() -> int:
         stderr = subprocess.DEVNULL,
     )
 
-    print(_green("✅ Python dependencies installed"))
+    _safe_print(_green("✅ Python dependencies installed"))
     return 0
 
 

studio/setup.ps1

@@ -250,9 +250,15 @@ function Find-VsBuildTools {
 # ─────────────────────────────────────────────
 # Banner
 # ─────────────────────────────────────────────
-Write-Host "+==============================================+" -ForegroundColor Green
-Write-Host "|       Unsloth Studio Setup (Windows)         |" -ForegroundColor Green
-Write-Host "+==============================================+" -ForegroundColor Green
+if ($env:SKIP_STUDIO_BASE -eq "1") {
+    Write-Host "+==============================================+" -ForegroundColor Green
+    Write-Host "|       Unsloth Studio Setup (Windows)         |" -ForegroundColor Green
+    Write-Host "+==============================================+" -ForegroundColor Green
+} else {
+    Write-Host "+==============================================+" -ForegroundColor Green
+    Write-Host "|      Unsloth Studio Update (Windows)         |" -ForegroundColor Green
+    Write-Host "+==============================================+" -ForegroundColor Green
+}
 
 # ==========================================================================
 #  PHASE 1: System-level prerequisites (winget installs, env vars)
@@ -497,7 +503,6 @@ if ($DriverMaxCuda) {
             $isCompat = ($tkMaj -lt $drMajorCuda) -or ($tkMaj -eq $drMajorCuda -and $tkMin -le $drMinorCuda)
             if ($isCompat) {
                 # Also verify the toolkit supports our GPU architecture
-                Write-Host "   [DEBUG] Checking CUDA compatibility: toolkit=$tkMaj.$tkMin arch=sm_$CudaArch" -ForegroundColor Magenta
                 $archOk = $true
                 if ($CudaArch) {
                     $archOk = Test-NvccArchSupport -NvccExe $candidateNvcc -Arch $CudaArch
@@ -728,16 +733,22 @@ if ($IsPipInstall) {
     Write-Host "[OK] Running from pip install - frontend already bundled, skipping Node/npm check" -ForegroundColor Green
 } else {
     # setup.sh installs Node LTS (v22) via nvm. We enforce the same range here:
-    # Node >= 20, npm >= 11.
+    # Vite 8 requires Node ^20.19.0 || >=22.12.0, npm >= 11.
     $NeedNode = $true
     try {
         $NodeVersion = (node -v 2>$null)
         $NpmVersion = (npm -v 2>$null)
         if ($NodeVersion -and $NpmVersion) {
-            $NodeMajor = [int]($NodeVersion -replace 'v','').Split('.')[0]
+            $NodeParts = ($NodeVersion -replace 'v','').Split('.')
+            $NodeMajor = [int]$NodeParts[0]
+            $NodeMinor = [int]$NodeParts[1]
             $NpmMajor = [int]$NpmVersion.Split('.')[0]
 
-            if ($NodeMajor -ge 20 -and $NpmMajor -ge 11) {
+            # Vite 8: ^20.19.0 || >=22.12.0
+            $NodeOk = ($NodeMajor -eq 20 -and $NodeMinor -ge 19) -or
+                      ($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
                 $NeedNode = $false
             } else {
@@ -761,6 +772,24 @@ if ($IsPipInstall) {
     }
 
     Write-Host "[OK] Node $(node -v) | npm $(npm -v)" -ForegroundColor Green
+
+    # ── 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
+        $prevEAP_bun = $ErrorActionPreference
+        $ErrorActionPreference = "Continue"
+        npm install -g bun 2>&1 | Out-Null
+        $ErrorActionPreference = $prevEAP_bun
+        Refresh-Environment
+        if (Get-Command bun -ErrorAction SilentlyContinue) {
+            Write-Host "[OK] bun installed ($(bun --version))" -ForegroundColor Green
+        } else {
+            Write-Host "[OK] bun install skipped (npm will be used instead)" -ForegroundColor DarkGray
+        }
+    } else {
+        Write-Host "[OK] bun already installed ($(bun --version))" -ForegroundColor Green
+    }
 }
 
 # ============================================
@@ -844,10 +873,10 @@ if ($IsPipInstall) {
             if ($NewerFile) { break }
         }
     }
-    # Also check all top-level files (package.json, bun.lock, vite.config.ts, index.html, etc.)
+    # Also check all top-level files (package.json, vite.config.ts, index.html, etc.)
     if (-not $NewerFile) {
         $NewerFile = Get-ChildItem -Path $FrontendDir -File -ErrorAction SilentlyContinue |
-            Where-Object { $_.LastWriteTime -gt $DistTime } |
+            Where-Object { $_.Name -ne "bun.lock" -and $_.LastWriteTime -gt $DistTime } |
             Select-Object -First 1
     }
     if (-not $NewerFile) {
@@ -882,26 +911,73 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) {
         $WalkDir = Split-Path $WalkDir -Parent
     }
 
-    # npm writes warnings to stderr; lower ErrorActionPreference so PS doesn't
-    # treat them as terminating errors (same pattern as the pip section below).
+    # Use bun if available (faster install), fall back to npm.
+    # Bun is used only as package manager; Node runs the actual build (Vite 8).
     $prevEAP_npm = $ErrorActionPreference
     $ErrorActionPreference = "Continue"
     Push-Location $FrontendDir
-    npm install 2>&1 | Out-Null
-    if ($LASTEXITCODE -ne 0) {
-        Pop-Location
-        $ErrorActionPreference = $prevEAP_npm
-        foreach ($gi in $HiddenGitignores) { Rename-Item -Path "$gi._twbuild" -NewName (Split-Path $gi -Leaf) -Force -ErrorAction SilentlyContinue }
-        Write-Host "[ERROR] npm install failed (exit code $LASTEXITCODE)" -ForegroundColor Red
-        Write-Host "   Try running 'npm install' manually in frontend/ to see errors" -ForegroundColor Yellow
-        exit 1
+
+    $UseBun = $null -ne (Get-Command bun -ErrorAction SilentlyContinue)
+
+    # bun's package cache can become corrupt -- packages get stored with only
+    # metadata but no actual content (bin/, lib/). When this happens bun install
+    # exits 0 but leaves binaries missing. We validate after install and clear
+    # the cache + retry once before falling back to npm.
+    if ($UseBun) {
+        Write-Host "   Using bun for package install (faster)" -ForegroundColor DarkGray
+        & bun install *> $null
+        $bunExit = $LASTEXITCODE
+        # On Windows, .bin/ entries can be tsc, tsc.cmd, or tsc.ps1
+        $hasTsc = (Test-Path "node_modules\.bin\tsc") -or (Test-Path "node_modules\.bin\tsc.cmd")
+        $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd")
+        if ($bunExit -eq 0 -and $hasTsc -and $hasVite) {
+            # bun install succeeded and critical binaries are present
+        } elseif ($bunExit -eq 0) {
+            Write-Host "   bun install exited 0 but critical binaries are missing, clearing cache and retrying..." -ForegroundColor Yellow
+            if (Test-Path "node_modules") {
+                Remove-Item "node_modules" -Recurse -Force -ErrorAction SilentlyContinue
+            }
+            & bun pm cache rm *> $null
+            & bun install *> $null
+            $bunExit = $LASTEXITCODE
+            $hasTsc = (Test-Path "node_modules\.bin\tsc") -or (Test-Path "node_modules\.bin\tsc.cmd")
+            $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd")
+            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") {
+                    Remove-Item "node_modules" -Recurse -Force -ErrorAction SilentlyContinue
+                }
+                $UseBun = $false
+            }
+        } else {
+            Write-Host "   [WARN] bun install failed (exit $bunExit), falling back to npm" -ForegroundColor Yellow
+            if (Test-Path "node_modules") {
+                Remove-Item "node_modules" -Recurse -Force -ErrorAction SilentlyContinue
+            }
+            $UseBun = $false
+        }
     }
-    npm run build 2>&1 | Out-Null
-    if ($LASTEXITCODE -ne 0) {
+    if (-not $UseBun) {
+        & npm install *> $null
+        $npmExit = $LASTEXITCODE
+        if ($npmExit -ne 0) {
+            Pop-Location
+            $ErrorActionPreference = $prevEAP_npm
+            foreach ($gi in $HiddenGitignores) { Rename-Item -Path "$gi._twbuild" -NewName (Split-Path $gi -Leaf) -Force -ErrorAction SilentlyContinue }
+            Write-Host "[ERROR] npm install failed (exit code $npmExit)" -ForegroundColor Red
+            Write-Host "   Try running 'npm install' manually in frontend/ to see errors" -ForegroundColor Yellow
+            exit 1
+        }
+    }
+
+    # Always use npm to run the build (Node runtime — avoids bun Windows runtime issues)
+    & npm run build *> $null
+    $buildExit = $LASTEXITCODE
+    if ($buildExit -ne 0) {
         Pop-Location
         $ErrorActionPreference = $prevEAP_npm
         foreach ($gi in $HiddenGitignores) { Rename-Item -Path "$gi._twbuild" -NewName (Split-Path $gi -Leaf) -Force -ErrorAction SilentlyContinue }
-        Write-Host "[ERROR] npm run build failed (exit code $LASTEXITCODE)" -ForegroundColor Red
+        Write-Host "[ERROR] npm run build failed (exit code $buildExit)" -ForegroundColor Red
         exit 1
     }
     Pop-Location
@@ -1030,9 +1106,9 @@ if (-not $PythonCmd) {
 
 Write-Host "[OK] Using $PythonCmd ($(& $PythonCmd --version 2>&1))" -ForegroundColor Green
 
-# Always create a .venv for isolation -- even for pip installs.
-# Created in the repo root (parent of studio/).
-$VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv"
+# The venv must already exist (created by install.ps1).
+# This script (setup.ps1 / "unsloth studio update") only updates packages.
+$VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio"
 
 # Stale-venv detection: if the venv exists but its torch flavor no longer
 # matches the current machine, wipe it so we get a clean install.
@@ -1095,8 +1171,10 @@ if (Test-Path $VenvDir -PathType Container) {
 }
 
 if (-not (Test-Path $VenvDir)) {
-    Write-Host "   Creating virtual environment at $VenvDir..." -ForegroundColor Cyan
-    & $PythonCmd -m venv $VenvDir
+    Write-Host "[ERROR] Virtual environment not found at $VenvDir" -ForegroundColor Red
+    Write-Host "        Run install.ps1 first to create the environment:" -ForegroundColor Yellow
+    Write-Host "        irm https://unsloth.ai/install.ps1 | iex" -ForegroundColor Yellow
+    exit 1
 } else {
     Write-Host "   Reusing existing virtual environment at $VenvDir" -ForegroundColor Green
 }
@@ -1243,6 +1321,96 @@ if ($LASTEXITCODE -ne 0) {
 $ErrorActionPreference = $prevEAP_t5
 Write-Host "[OK] Transformers 5.x pre-installed to .venv_t5/" -ForegroundColor Green
 
+# ==========================================================================
+#  PHASE 3.4: Prefer prebuilt llama.cpp bundles before source build
+# ==========================================================================
+$UnslothHome = Join-Path $env:USERPROFILE ".unsloth"
+if (-not (Test-Path $UnslothHome)) { New-Item -ItemType Directory -Force $UnslothHome | Out-Null }
+$LlamaCppDir = Join-Path $UnslothHome "llama.cpp"
+$NeedLlamaSourceBuild = $false
+$SkipPrebuiltInstall = $false
+$RequestedLlamaTag = if ($env:UNSLOTH_LLAMA_TAG) { $env:UNSLOTH_LLAMA_TAG } else { "latest" }
+$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO) { $env:UNSLOTH_LLAMA_RELEASE_REPO } else { "unslothai/llama.cpp" }
+$resolveOutput = & python "$PSScriptRoot\install_llama_prebuilt.py" --resolve-install-tag $RequestedLlamaTag --published-repo $HelperReleaseRepo 2>&1
+$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 $_ }
+    }
+    # 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.
+    $fallbackOutput = & python "$PSScriptRoot\install_llama_prebuilt.py" --resolve-llama-tag $RequestedLlamaTag --published-repo $HelperReleaseRepo 2>$null
+    $fallbackExit = $LASTEXITCODE
+    $ResolvedLlamaTag = if ($fallbackExit -eq 0 -and $fallbackOutput) {
+        ($fallbackOutput | Select-Object -Last 1).ToString().Trim()
+    } elseif ($RequestedLlamaTag -eq "latest") {
+        # Try Unsloth release repo first, then fall back to ggml-org upstream
+        $resolvedLatest = $null
+        try {
+            $latestRelease = Invoke-RestMethod -Uri "https://api.github.com/repos/$HelperReleaseRepo/releases/latest" -ErrorAction Stop
+            $resolvedLatest = $latestRelease.tag_name
+        } catch {}
+        if (-not $resolvedLatest) {
+            try {
+                $latestRelease = Invoke-RestMethod -Uri "https://api.github.com/repos/ggml-org/llama.cpp/releases/latest" -ErrorAction Stop
+                $resolvedLatest = $latestRelease.tag_name
+            } catch {}
+        }
+        if ($resolvedLatest) { $resolvedLatest } else { $RequestedLlamaTag }
+    } else {
+        $RequestedLlamaTag
+    }
+    $NeedLlamaSourceBuild = $true
+    $SkipPrebuiltInstall = $true
+}
+
+Write-Host ""
+Write-Host "Resolved llama.cpp release tag: $ResolvedLlamaTag" -ForegroundColor Gray
+
+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
+    $NeedLlamaSourceBuild = $true
+} else {
+    Write-Host ""
+    Write-Host "Installing prebuilt llama.cpp bundle (preferred path)..." -ForegroundColor Cyan
+    if (Test-Path $LlamaCppDir) {
+        Write-Host "Existing llama.cpp install detected -- validating staged prebuilt update before replacement" -ForegroundColor Gray
+    }
+    if ($SkipPrebuiltInstall) {
+        Write-Host "[WARN] Skipping prebuilt install because prebuilt tag resolution failed -- falling back to source build" -ForegroundColor Yellow
+    } else {
+        $prebuiltArgs = @(
+            "$PSScriptRoot\install_llama_prebuilt.py",
+            "--install-dir", $LlamaCppDir,
+            "--llama-tag", $ResolvedLlamaTag,
+            "--published-repo", $HelperReleaseRepo
+        )
+        if ($env:UNSLOTH_LLAMA_RELEASE_TAG) {
+            $prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG)
+        }
+        $prevEAPPrebuilt = $ErrorActionPreference
+        $ErrorActionPreference = "Continue"
+        & python @prebuiltArgs
+        $prebuiltExit = $LASTEXITCODE
+        $ErrorActionPreference = $prevEAPPrebuilt
+
+        if ($prebuiltExit -eq 0) {
+            Write-Host "[OK] Prebuilt llama.cpp installed and validated" -ForegroundColor Green
+        } else {
+            if (Test-Path $LlamaCppDir) {
+                Write-Host "[WARN] Prebuilt update failed; existing install was restored or cleaned before source build fallback" -ForegroundColor Yellow
+            }
+            Write-Host "[WARN] Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build" -ForegroundColor Yellow
+            $NeedLlamaSourceBuild = $true
+        }
+    }
+}
+
 # ==========================================================================
 #  PHASE 3.5: Install OpenSSL dev (for HTTPS support in llama-server)
 # ==========================================================================
@@ -1250,42 +1418,46 @@ Write-Host "[OK] Transformers 5.x pre-installed to .venv_t5/" -ForegroundColor G
 # ShiningLight.OpenSSL.Dev includes headers + libs that cmake can find.
 $OpenSslAvailable = $false
 
-# Check if OpenSSL dev is already installed (look for include dir)
-$OpenSslRoots = @(
-    'C:\Program Files\OpenSSL-Win64',
-    'C:\Program Files\OpenSSL',
-    'C:\OpenSSL-Win64'
-)
-$OpenSslRoot = $null
-foreach ($root in $OpenSslRoots) {
-    if (Test-Path (Join-Path $root 'include\openssl\ssl.h')) {
-        $OpenSslRoot = $root
-        break
-    }
-}
-
-if ($OpenSslRoot) {
-    $OpenSslAvailable = $true
-    Write-Host "[OK] OpenSSL dev found at $OpenSslRoot" -ForegroundColor Green
-} else {
-    Write-Host "" 
-    Write-Host "Installing OpenSSL dev (for HTTPS in llama-server)..." -ForegroundColor Cyan
-    $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue)
-    if ($HasWinget) {
-        winget install -e --id ShiningLight.OpenSSL.Dev --accept-package-agreements --accept-source-agreements
-        # Re-check after install
-        foreach ($root in $OpenSslRoots) {
-            if (Test-Path (Join-Path $root 'include\openssl\ssl.h')) {
-                $OpenSslRoot = $root
-                $OpenSslAvailable = $true
-                Write-Host "[OK] OpenSSL dev installed at $OpenSslRoot" -ForegroundColor Green
-                break
-            }
+if ($NeedLlamaSourceBuild) {
+    # Check if OpenSSL dev is already installed (look for include dir)
+    $OpenSslRoots = @(
+        'C:\Program Files\OpenSSL-Win64',
+        'C:\Program Files\OpenSSL',
+        'C:\OpenSSL-Win64'
+    )
+    $OpenSslRoot = $null
+    foreach ($root in $OpenSslRoots) {
+        if (Test-Path (Join-Path $root 'include\openssl\ssl.h')) {
+            $OpenSslRoot = $root
+            break
         }
     }
-    if (-not $OpenSslAvailable) {
-        Write-Host "[WARN] OpenSSL dev not available -- llama-server will be built without HTTPS" -ForegroundColor Yellow
+
+    if ($OpenSslRoot) {
+        $OpenSslAvailable = $true
+        Write-Host "[OK] OpenSSL dev found at $OpenSslRoot" -ForegroundColor Green
+    } else {
+        Write-Host "" 
+        Write-Host "Installing OpenSSL dev (for HTTPS in llama-server)..." -ForegroundColor Cyan
+        $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue)
+        if ($HasWinget) {
+            winget install -e --id ShiningLight.OpenSSL.Dev --accept-package-agreements --accept-source-agreements
+            # Re-check after install
+            foreach ($root in $OpenSslRoots) {
+                if (Test-Path (Join-Path $root 'include\openssl\ssl.h')) {
+                    $OpenSslRoot = $root
+                    $OpenSslAvailable = $true
+                    Write-Host "[OK] OpenSSL dev installed at $OpenSslRoot" -ForegroundColor Green
+                    break
+                }
+            }
+        }
+        if (-not $OpenSslAvailable) {
+            Write-Host "[WARN] OpenSSL dev not available -- llama-server will be built without HTTPS" -ForegroundColor Yellow
+        }
     }
+} else {
+    Write-Host "[SKIP] OpenSSL dev install -- prebuilt llama.cpp already validated" -ForegroundColor Yellow
 }
 
 # ==========================================================================
@@ -1298,9 +1470,7 @@ if ($OpenSslRoot) {
 #   - llama-server:   for GGUF model inference (with HTTPS if OpenSSL available)
 #   - llama-quantize: for GGUF export quantization
 # Prerequisites (git, cmake, VS Build Tools, CUDA Toolkit) already installed in Phase 1.
-$UnslothHome = Join-Path $env:USERPROFILE ".unsloth"
-if (-not (Test-Path $UnslothHome)) { New-Item -ItemType Directory -Force $UnslothHome | Out-Null }
-$LlamaCppDir = Join-Path $UnslothHome "llama.cpp"
+$OriginalLlamaCppDir = $LlamaCppDir
 $BuildDir = Join-Path $LlamaCppDir "build"
 $LlamaServerBin = Join-Path $BuildDir "bin\Release\llama-server.exe"
 
@@ -1323,7 +1493,10 @@ if (Test-Path $LlamaServerBin) {
     }
 }
 
-if ((Test-Path $LlamaServerBin) -and -not $NeedRebuild) {
+if (-not $NeedLlamaSourceBuild) {
+    Write-Host ""
+    Write-Host "[OK] Using validated prebuilt llama.cpp install at $LlamaCppDir" -ForegroundColor Green
+} elseif ((Test-Path $LlamaServerBin) -and -not $NeedRebuild) {
     Write-Host ""
     Write-Host "[OK] llama-server already exists at $LlamaServerBin" -ForegroundColor Green
 } elseif (-not $HasCmakeForBuild) {
@@ -1379,29 +1552,49 @@ if ((Test-Path $LlamaServerBin) -and -not $NeedRebuild) {
 
     # -- Step A: Clone or pull llama.cpp --
 
+    $UseConcreteRef = ($ResolvedLlamaTag -ne "latest" -and -not [string]::IsNullOrWhiteSpace($ResolvedLlamaTag))
+
     if (Test-Path (Join-Path $LlamaCppDir ".git")) {
-        Write-Host "   llama.cpp repo already cloned, pulling latest..." -ForegroundColor Gray
-        git -C $LlamaCppDir pull 2>&1 | Out-Null
+        Write-Host "   Syncing llama.cpp to $ResolvedLlamaTag..." -ForegroundColor Gray
+        if ($UseConcreteRef) {
+            git -C $LlamaCppDir fetch --depth 1 origin $ResolvedLlamaTag 2>&1 | Out-Null
+        } else {
+            git -C $LlamaCppDir fetch --depth 1 origin 2>&1 | Out-Null
+        }
         if ($LASTEXITCODE -ne 0) {
-            Write-Host "   [WARN] git pull failed -- using existing source" -ForegroundColor Yellow
+            Write-Host "   [WARN] git fetch failed -- using existing source" -ForegroundColor Yellow
+        } else {
+            git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD 2>&1 | Out-Null
+            if ($LASTEXITCODE -ne 0) {
+                $BuildOk = $false
+                $FailedStep = "git checkout"
+            } else {
+                git -C $LlamaCppDir clean -fdx 2>&1 | Out-Null
+            }
         }
     } else {
-        Write-Host "   Cloning llama.cpp..." -ForegroundColor Gray
-        if (Test-Path $LlamaCppDir) { Remove-Item -Recurse -Force $LlamaCppDir }
-        git clone --depth 1 https://github.com/ggml-org/llama.cpp.git $LlamaCppDir 2>&1 | Out-Null
+        Write-Host "   Cloning llama.cpp @ $ResolvedLlamaTag..." -ForegroundColor Gray
+        $buildTmp = "$LlamaCppDir.build.$PID"
+        if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+        $cloneArgs = @("clone", "--depth", "1")
+        if ($UseConcreteRef) {
+            $cloneArgs += @("--branch", $ResolvedLlamaTag)
+        }
+        $cloneArgs += @("https://github.com/ggml-org/llama.cpp.git", $buildTmp)
+        git @cloneArgs 2>&1 | Out-Null
         if ($LASTEXITCODE -ne 0) {
             $BuildOk = $false
             $FailedStep = "git clone"
+            if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+        }
+        # Use temp dir for build; swap into $LlamaCppDir only after build succeeds
+        if ($BuildOk) {
+            $LlamaCppDir = $buildTmp
+            $BuildDir = Join-Path $LlamaCppDir "build"
         }
     }
 
     # -- Step B: cmake configure --
-    # Clean stale CMake cache to prevent previous CUDA settings from leaking
-    # into a CPU-only rebuild (or vice versa).
-    $CmakeCacheFile = Join-Path $BuildDir "CMakeCache.txt"
-    if (Test-Path $CmakeCacheFile) {
-        Remove-Item -Recurse -Force $BuildDir
-    }
 
     if ($BuildOk) {
         Write-Host ""
@@ -1502,6 +1695,21 @@ if ((Test-Path $LlamaServerBin) -and -not $NeedRebuild) {
         }
     }
 
+    # Swap temp build dir into final location (only if we built in a temp dir)
+    if ($BuildOk -and $LlamaCppDir -ne $OriginalLlamaCppDir) {
+        if (Test-Path $OriginalLlamaCppDir) { Remove-Item -Recurse -Force $OriginalLlamaCppDir }
+        Move-Item $LlamaCppDir $OriginalLlamaCppDir
+        $LlamaCppDir = $OriginalLlamaCppDir
+        $BuildDir = Join-Path $LlamaCppDir "build"
+        $LlamaServerBin = Join-Path $BuildDir "bin\Release\llama-server.exe"
+    } elseif (-not $BuildOk -and $LlamaCppDir -ne $OriginalLlamaCppDir) {
+        # Build failed -- clean up temp dir, preserve existing install
+        if (Test-Path $LlamaCppDir) { Remove-Item -Recurse -Force $LlamaCppDir }
+        $LlamaCppDir = $OriginalLlamaCppDir
+        $BuildDir = Join-Path $LlamaCppDir "build"
+        $LlamaServerBin = Join-Path $BuildDir "bin\Release\llama-server.exe"
+    }
+
     # Restore ErrorActionPreference
     $ErrorActionPreference = $prevEAP
 
@@ -1537,8 +1745,9 @@ if ((Test-Path $LlamaServerBin) -and -not $NeedRebuild) {
 # Done
 # ============================================
 Write-Host ""
+$doneLine = if ($env:SKIP_STUDIO_BASE -eq "1") { "Setup Complete!" } else { "Update Complete!" }
 Write-Host "+===============================================+" -ForegroundColor Green
-Write-Host "|           Setup Complete!                     |" -ForegroundColor Green
+Write-Host "|           $doneLine                    |" -ForegroundColor Green
 Write-Host "|                                               |" -ForegroundColor Green
 Write-Host "|  Launch with:                                 |" -ForegroundColor Green
 Write-Host "|    unsloth studio -H 0.0.0.0 -p 8888          |" -ForegroundColor Green

studio/setup.sh

@@ -44,9 +44,15 @@ run_quiet_no_exit() {
     _run_quiet return "$@"
 }
 
-echo "╔══════════════════════════════════════╗"
-echo "║     Unsloth Studio Setup Script      ║"
-echo "╚══════════════════════════════════════╝"
+if [ "${SKIP_STUDIO_BASE:-0}" = "1" ]; then
+    echo "╔══════════════════════════════════════╗"
+    echo "║     Unsloth Studio Setup Script      ║"
+    echo "╚══════════════════════════════════════╝"
+else
+    echo "╔══════════════════════════════════════╗"
+    echo "║     Unsloth Studio Update Script     ║"
+    echo "╚══════════════════════════════════════╝"
+fi
 
 # ── Clean up stale Unsloth compiled caches ──
 rm -rf "$REPO_ROOT/unsloth_compiled_cache"
@@ -69,6 +75,7 @@ _NEED_FRONTEND_BUILD=true
 if [ -d "$SCRIPT_DIR/frontend/dist" ]; then
     # Check all top-level files (package.json, bun.lock, vite.config.ts, index.html, etc.)
     _changed=$(find "$SCRIPT_DIR/frontend" -maxdepth 1 -type f \
+        ! -name 'bun.lock' \
         -newer "$SCRIPT_DIR/frontend/dist" -print -quit 2>/dev/null)
     # Check src/ and public/ recursively (|| true guards against set -e when dirs are missing)
     if [ -z "$_changed" ]; then
@@ -85,12 +92,18 @@ else
 NEED_NODE=true
 if command -v node &>/dev/null && command -v npm &>/dev/null; then
     NODE_MAJOR=$(node -v | sed 's/v//' | cut -d. -f1)
+    NODE_MINOR=$(node -v | sed 's/v//' | cut -d. -f2)
     NPM_MAJOR=$(npm -v | cut -d. -f1)
-    if [ "$NODE_MAJOR" -ge 20 ] && [ "$NPM_MAJOR" -ge 11 ]; then
+    # Vite 8 requires Node ^20.19.0 || >=22.12.0
+    NODE_OK=false
+    if [ "$NODE_MAJOR" -eq 20 ] && [ "$NODE_MINOR" -ge 19 ]; then NODE_OK=true; fi
+    if [ "$NODE_MAJOR" -eq 22 ] && [ "$NODE_MINOR" -ge 12 ]; then NODE_OK=true; fi
+    if [ "$NODE_MAJOR" -ge 23 ]; then NODE_OK=true; fi
+    if [ "$NODE_OK" = true ] && [ "$NPM_MAJOR" -ge 11 ]; then
         echo "✅ Node $(node -v) and npm $(npm -v) already meet requirements. Skipping nvm install."
         NEED_NODE=false
     else
-        if [ "$IS_COLAB" = true ]; then
+        if [ "$IS_COLAB" = true ] && [ "$NODE_OK" = true ]; then
             echo "✅ Node $(node -v) and npm $(npm -v) detected in Colab."
             # In Colab, just upgrade npm directly - nvm doesn't work well
             if [ "$NPM_MAJOR" -lt 11 ]; then
@@ -150,6 +163,20 @@ fi
 
 echo "✅ Node $(node -v) | npm $(npm -v)"
 
+# ── 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
+    echo "   Installing bun (faster frontend package installs)..."
+    if npm install -g bun > /dev/null 2>&1 && command -v bun &>/dev/null; then
+        echo "   bun installed ($(bun --version))"
+    else
+        echo "   bun install skipped (npm will be used instead)"
+    fi
+else
+    echo "   bun already installed ($(bun --version))"
+fi
+
 # ── 5. Build frontend ──
 cd "$SCRIPT_DIR/frontend"
 
@@ -174,7 +201,57 @@ _restore_gitignores() {
 }
 trap _restore_gitignores EXIT
 
-run_quiet "npm install" npm install
+# Use bun for install if available (faster), fall back to npm.
+# Build always uses npm (Node runtime -- avoids bun runtime issues on some platforms).
+# NOTE: We intentionally avoid run_quiet for the bun install attempt because
+# run_quiet calls exit on failure, which would kill the script before the npm
+# fallback can run. Instead we capture output manually and only show it on failure.
+#
+# IMPORTANT: bun's package cache can become corrupt -- packages get stored
+# with only metadata (package.json, README) but no actual content (bin/,
+# lib/). When this happens bun install exits 0 but leaves binaries missing.
+# We verify critical binaries after install. If missing, we clear the cache
+# and retry once before falling back to npm.
+_try_bun_install() {
+    local _log _exit_code=0
+    _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
+        rm -f "$_log"
+        return 0
+    fi
+
+    # Either bun install failed or it exited 0 but left packages missing
+    if [ "$_exit_code" -ne 0 ]; then
+        echo "   bun install failed (exit code $_exit_code):"
+    else
+        echo "   bun install exited 0 but critical binaries are missing:"
+    fi
+    sed 's/^/   | /' "$_log" >&2
+    rm -f "$_log"
+    rm -rf node_modules
+    return 1
+}
+
+_bun_install_ok=false
+if command -v bun &>/dev/null; then
+    echo "   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
+        if _try_bun_install; then
+            _bun_install_ok=true
+        fi
+    fi
+fi
+if [ "$_bun_install_ok" = false ]; then
+    run_quiet "npm install" npm install
+fi
 run_quiet "npm run build" npm run build
 
 _restore_gitignores
@@ -203,114 +280,31 @@ fi
 
 # ── 6. Python venv + deps ──
 
-# ── 6a. Discover best Python >= 3.11 and < 3.14 (i.e. 3.11.x, 3.12.x, or 3.13.x) ──
-MIN_PY_MINOR=11   # minimum minor version (>= 3.11)
-MAX_PY_MINOR=13   # maximum minor version (< 3.14)
-BEST_PY=""
-BEST_MINOR=0
-
-# If the caller (e.g. install.sh) already chose a Python, use it directly.
-if [ -n "${REQUESTED_PYTHON_VERSION:-}" ] && [ -x "$REQUESTED_PYTHON_VERSION" ]; then
-    _req_ver=$("$REQUESTED_PYTHON_VERSION" --version 2>&1 | awk '{print $2}')
-    _req_major=$(echo "$_req_ver" | cut -d. -f1)
-    _req_minor=$(echo "$_req_ver" | cut -d. -f2)
-    if [ "$_req_major" -eq 3 ] 2>/dev/null && \
-       [ "$_req_minor" -ge "$MIN_PY_MINOR" ] 2>/dev/null && \
-       [ "$_req_minor" -le "$MAX_PY_MINOR" ] 2>/dev/null; then
-        BEST_PY="$REQUESTED_PYTHON_VERSION"
-        echo "Using requested Python version: $BEST_PY"
-    else
-        echo "Ignoring requested Python $REQUESTED_PYTHON_VERSION ($_req_ver) -- outside supported range"
-    fi
-fi
-
-if [ -z "$BEST_PY" ]; then
-# Collect candidate python3 binaries (python3, python3.9, python3.10, …)
-for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)?$' | sort -u); do
-    if ! command -v "$candidate" &>/dev/null; then
-        continue
-    fi
-    # Get version string, e.g. "Python 3.12.5"
-    ver_str=$("$candidate" --version 2>&1) || continue
-    ver_str=$(echo "$ver_str" | awk '{print $2}')
-    py_major=$(echo "$ver_str" | cut -d. -f1)
-    py_minor=$(echo "$ver_str" | cut -d. -f2)
-
-    # Skip anything that isn't Python 3
-    if [ "$py_major" -ne 3 ] 2>/dev/null; then
-        continue
-    fi
-
-    # Skip versions below 3.11
-    if [ "$py_minor" -lt "$MIN_PY_MINOR" ] 2>/dev/null; then
-        continue
-    fi
-
-    # Skip versions above 3.13 (require < 3.14)
-    if [ "$py_minor" -gt "$MAX_PY_MINOR" ] 2>/dev/null; then
-        continue
-    fi
-
-    # Keep the highest qualifying version
-    if [ "$py_minor" -gt "$BEST_MINOR" ]; then
-        BEST_PY="$candidate"
-        BEST_MINOR="$py_minor"
-    fi
-done
-fi
-
-if [ -z "$BEST_PY" ]; then
-    echo "❌ ERROR: No Python version between 3.${MIN_PY_MINOR} and 3.${MAX_PY_MINOR} found on this system."
-    echo "   Detected Python 3 installations:"
-    for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)?$' | sort -u); do
-        if command -v "$candidate" &>/dev/null; then
-            echo "     - $candidate ($($candidate --version 2>&1))"
-        fi
-    done
-    echo ""
-    echo "   Please install Python 3.${MIN_PY_MINOR} or 3.${MAX_PY_MINOR}."
-    echo "   For example:  sudo apt install python3.12 python3.12-venv"
-    exit 1
-fi
-
-BEST_VER=$("$BEST_PY" --version 2>&1 | awk '{print $2}')
-echo "✅ Using $BEST_PY ($BEST_VER) — compatible (3.${MIN_PY_MINOR}.x – 3.${MAX_PY_MINOR}.x)"
-
-REQ_ROOT="$SCRIPT_DIR/backend/requirements"
-SINGLE_ENV_CONSTRAINTS="$REQ_ROOT/single-env/constraints.txt"
-SINGLE_ENV_DATA_DESIGNER="$REQ_ROOT/single-env/data-designer.txt"
-SINGLE_ENV_DATA_DESIGNER_DEPS="$REQ_ROOT/single-env/data-designer-deps.txt"
-SINGLE_ENV_PATCH="$REQ_ROOT/single-env/patch_metadata.py"
-
-install_python_stack() {
-    python "$SCRIPT_DIR/install_python_stack.py"
-}
-
-# Create venv under ~/.unsloth/studio/ (shared location, not in repo).
-# All platforms (including Colab) use the same isolated venv so that
-# studio dependencies are never installed into the system Python.
+# The venv must already exist (created by install.sh).
+# This script (setup.sh / "unsloth studio update") only updates packages.
 STUDIO_HOME="$HOME/.unsloth/studio"
-VENV_DIR="$STUDIO_HOME/.venv"
+VENV_DIR="$STUDIO_HOME/unsloth_studio"
 VENV_T5_DIR="$STUDIO_HOME/.venv_t5"
-mkdir -p "$STUDIO_HOME"
 
 # Clean up legacy in-repo venvs if they exist
 [ -d "$REPO_ROOT/.venv" ] && rm -rf "$REPO_ROOT/.venv"
 [ -d "$REPO_ROOT/.venv_overlay" ] && rm -rf "$REPO_ROOT/.venv_overlay"
 [ -d "$REPO_ROOT/.venv_t5" ] && rm -rf "$REPO_ROOT/.venv_t5"
+# Note: do NOT delete $STUDIO_HOME/.venv here — install.sh handles migration
 
-rm -rf "$VENV_DIR"
-rm -rf "$VENV_T5_DIR"
-# Try creating venv with pip; fall back to --without-pip + bootstrap
-# (some environments like Colab have broken ensurepip)
-if ! "$BEST_PY" -m venv "$VENV_DIR" 2>/dev/null; then
-    "$BEST_PY" -m venv --without-pip "$VENV_DIR"
-    source "$VENV_DIR/bin/activate"
-    curl -sS https://bootstrap.pypa.io/get-pip.py | python > /dev/null
-else
-    source "$VENV_DIR/bin/activate"
+if [ ! -x "$VENV_DIR/bin/python" ]; then
+    echo "❌ ERROR: Virtual environment not found at $VENV_DIR"
+    echo "   Run install.sh first to create the environment:"
+    echo "   curl -fsSL https://unsloth.ai/install.sh | sh"
+    exit 1
 fi
 
+source "$VENV_DIR/bin/activate"
+
+install_python_stack() {
+    python "$SCRIPT_DIR/install_python_stack.py"
+}
+
 # ── Ensure uv is available (much faster than pip) ──
 USE_UV=false
 if command -v uv &>/dev/null; then
@@ -329,27 +323,149 @@ fast_install() {
 }
 
 cd "$SCRIPT_DIR"
-install_python_stack
 
-# ── 6b. Pre-install transformers 5.x into .venv_t5/ ──
-# Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing
-# at runtime (slow, ~10-15s), we pre-install into a separate directory.
-# The training subprocess just prepends .venv_t5/ to sys.path -- instant switch.
+# ── Check if Python deps need updating ──
+# Compare installed package version against PyPI latest.
+# Skip all Python dependency work if versions match (fast update path).
+_PKG_NAME="${STUDIO_PACKAGE_NAME:-unsloth}"
+_SKIP_PYTHON_DEPS=false
+if [ "${SKIP_STUDIO_BASE:-0}" != "1" ] && [ "${STUDIO_LOCAL_INSTALL:-0}" != "1" ]; then
+    # Only check when NOT called from install.sh (which just installed the package)
+    INSTALLED_VER=$("$VENV_DIR/bin/python" -c "
+from importlib.metadata import version
+print(version('$_PKG_NAME'))
+" 2>/dev/null || echo "")
+
+    LATEST_VER=$(curl -fsSL --max-time 5 "https://pypi.org/pypi/$_PKG_NAME/json" 2>/dev/null \
+        | "$VENV_DIR/bin/python" -c "import sys,json; print(json.load(sys.stdin)['info']['version'])" 2>/dev/null \
+        || echo "")
+
+    if [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ] && [ "$INSTALLED_VER" = "$LATEST_VER" ]; then
+        echo "✅ $_PKG_NAME $INSTALLED_VER is up to date (matches PyPI latest)"
+        _SKIP_PYTHON_DEPS=true
+    elif [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ]; then
+        echo "⬆️  $_PKG_NAME $INSTALLED_VER → $LATEST_VER available, updating dependencies..."
+    elif [ -z "$LATEST_VER" ]; then
+        echo "⚠️  Could not reach PyPI, updating dependencies to be safe..."
+    fi
+fi
+
+if [ "$_SKIP_PYTHON_DEPS" = false ]; then
+    install_python_stack
+
+    # ── 6b. Pre-install transformers 5.x into .venv_t5/ ──
+    # Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing
+    # at runtime (slow, ~10-15s), we pre-install into a separate directory.
+    # The training subprocess just prepends .venv_t5/ to sys.path -- instant switch.
+    echo ""
+    echo "   Pre-installing transformers 5.x for newer model support..."
+    mkdir -p "$VENV_T5_DIR"
+    run_quiet "install transformers 5.x" fast_install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0"
+    run_quiet "install huggingface_hub for t5" fast_install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.7.1"
+    run_quiet "install hf_xet for t5" fast_install --target "$VENV_T5_DIR" --no-deps "hf_xet==1.4.2"
+    # tiktoken is needed by Qwen-family tokenizers. Install with deps since
+    # regex/requests may be missing on Windows.
+    run_quiet "install tiktoken for t5" fast_install --target "$VENV_T5_DIR" "tiktoken"
+    echo "✅ Transformers 5.x pre-installed to $VENV_T5_DIR/"
+else
+    echo "✅ Python dependencies up to date — skipping"
+fi
+
+# ── 7. Prefer prebuilt llama.cpp bundles before any source build path ──
+UNSLOTH_HOME="$HOME/.unsloth"
+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_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}"
+_REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-latest}"
+_HELPER_RELEASE_REPO="${UNSLOTH_LLAMA_RELEASE_REPO:-unslothai/llama.cpp}"
+_RESOLVE_LLAMA_LOG="$(mktemp)"
+set +e
+python "$SCRIPT_DIR/install_llama_prebuilt.py" \
+    --resolve-install-tag "$_REQUESTED_LLAMA_TAG" \
+    --published-repo "$_HELPER_RELEASE_REPO" >"$_RESOLVE_LLAMA_LOG" 2>&1
+_RESOLVE_LLAMA_STATUS=$?
+set -e
+if [ "$_RESOLVE_LLAMA_STATUS" -eq 0 ]; then
+    _RESOLVED_LLAMA_TAG="$(tail -n 1 "$_RESOLVE_LLAMA_LOG" | tr -d '\r')"
+else
+    _RESOLVED_LLAMA_TAG=""
+fi
+if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
+    echo ""
+    echo "⚠️  Failed to resolve an installable prebuilt llama.cpp tag via $_HELPER_RELEASE_REPO"
+    cat "$_RESOLVE_LLAMA_LOG" >&2 || true
+    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
+    # bleeding-edge tag (e.g. b8514) from ggml-org/llama.cpp.
+    _RESOLVED_LLAMA_TAG="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" --resolve-llama-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO" 2>/dev/null)"
+    _RESOLVE_UPSTREAM_STATUS=$?
+    set -e
+    if [ "$_RESOLVE_UPSTREAM_STATUS" -ne 0 ] || [ -z "$_RESOLVED_LLAMA_TAG" ]; then
+        if [ "$_REQUESTED_LLAMA_TAG" = "latest" ]; then
+            # Try Unsloth release repo first, then fall back to ggml-org upstream
+            _RESOLVED_LLAMA_TAG="$(curl -fsSL "https://api.github.com/repos/${_HELPER_RELEASE_REPO}/releases/latest" 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG=""
+            if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
+                _RESOLVED_LLAMA_TAG="$(curl -fsSL https://api.github.com/repos/ggml-org/llama.cpp/releases/latest 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG=""
+            fi
+        fi
+        if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
+            _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
+        fi
+    fi
+    _NEED_LLAMA_SOURCE_BUILD=true
+    _SKIP_PREBUILT_INSTALL=true
+fi
+rm -f "$_RESOLVE_LLAMA_LOG"
+
 echo ""
-echo "   Pre-installing transformers 5.x for newer model support..."
-mkdir -p "$VENV_T5_DIR"
-run_quiet "install transformers 5.x" fast_install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0"
-run_quiet "install huggingface_hub for t5" fast_install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.7.1"
-run_quiet "install hf_xet for t5" fast_install --target "$VENV_T5_DIR" --no-deps "hf_xet==1.4.2"
-# tiktoken is needed by Qwen-family tokenizers. Install with deps since
-# regex/requests may be missing on Windows.
-run_quiet "install tiktoken for t5" fast_install --target "$VENV_T5_DIR" "tiktoken"
-echo "✅ Transformers 5.x pre-installed to $VENV_T5_DIR/"
+echo "Resolved llama.cpp release tag: $_RESOLVED_LLAMA_TAG"
 
-# ── 7. WSL: pre-install GGUF build dependencies ──
+if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then
+    echo ""
+    echo "⚠️  UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install"
+    _NEED_LLAMA_SOURCE_BUILD=true
+else
+    echo ""
+    echo "Installing prebuilt llama.cpp bundle (preferred path)..."
+    if [ -d "$LLAMA_CPP_DIR" ]; then
+        echo "Existing llama.cpp install detected -- validating staged prebuilt update before replacement"
+    fi
+    if [ "${_SKIP_PREBUILT_INSTALL:-false}" = true ]; then
+        echo "⚠️  Skipping prebuilt install because prebuilt tag resolution failed -- falling back to source build"
+    else
+        _PREBUILT_CMD=(
+            python "$SCRIPT_DIR/install_llama_prebuilt.py"
+            --install-dir "$LLAMA_CPP_DIR"
+            --llama-tag "$_RESOLVED_LLAMA_TAG"
+            --published-repo "$_HELPER_RELEASE_REPO"
+        )
+        if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
+            _PREBUILT_CMD+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
+        fi
+        set +e
+        "${_PREBUILT_CMD[@]}"
+        _PREBUILT_STATUS=$?
+        set -e
+
+        if [ "$_PREBUILT_STATUS" -eq 0 ]; then
+            echo "✅ Prebuilt llama.cpp installed and validated"
+        else
+            if [ -d "$LLAMA_CPP_DIR" ]; then
+                echo "⚠️  Prebuilt update failed; existing install was restored or cleaned before source build fallback"
+            fi
+            echo "⚠️  Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build"
+            _NEED_LLAMA_SOURCE_BUILD=true
+        fi
+    fi
+fi
+
+# ── 8. WSL: pre-install GGUF build dependencies for fallback source builds ──
 # On WSL, sudo requires a password and can't be entered during GGUF export
 # (runs in a non-interactive subprocess). Install build deps here instead.
-if grep -qi microsoft /proc/version 2>/dev/null; then
+if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] && grep -qi microsoft /proc/version 2>/dev/null; then
     echo ""
     echo "⚠️  WSL detected -- installing build dependencies for GGUF export..."
     _GGUF_DEPS="pciutils build-essential cmake curl git libcurl4-openssl-dev"
@@ -407,22 +523,19 @@ if grep -qi microsoft /proc/version 2>/dev/null; then
     fi
 fi
 
-# ── 8. Build llama.cpp binaries for GGUF inference + export ──
+# ── 9. Build llama.cpp binaries for GGUF inference + export when prebuilt install fails ──
 # Builds at ~/.unsloth/llama.cpp — a single shared location under the user's
 # home directory. This is used by both the inference server and the GGUF
 # export pipeline (unsloth-zoo).
 #   - llama-server: for GGUF model inference
 #   - llama-quantize: for GGUF export quantization (symlinked to root for check_llama_cpp())
-UNSLOTH_HOME="$HOME/.unsloth"
-mkdir -p "$UNSLOTH_HOME"
-LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp"
-LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server"
-if [ "${_SKIP_GGUF_BUILD:-}" = true ]; then
+if [ "$_NEED_LLAMA_SOURCE_BUILD" = false ]; then
+    :
+elif [ "${_SKIP_GGUF_BUILD:-}" = true ]; then
     echo ""
     echo "Skipping llama-server build (missing dependencies)"
     echo "   Install the missing packages and re-run setup to enable GGUF inference."
 else
-rm -rf "$LLAMA_CPP_DIR"
 {
     # Check prerequisites
     if ! command -v cmake &>/dev/null; then
@@ -437,7 +550,13 @@ rm -rf "$LLAMA_CPP_DIR"
         echo "Building llama-server for GGUF inference..."
 
         BUILD_OK=true
-        run_quiet_no_exit "clone llama.cpp" git clone --depth 1 https://github.com/ggml-org/llama.cpp.git "$LLAMA_CPP_DIR" || BUILD_OK=false
+        _CLONE_BRANCH_ARGS=()
+        if [ "$_RESOLVED_LLAMA_TAG" != "latest" ] && [ -n "$_RESOLVED_LLAMA_TAG" ]; then
+            _CLONE_BRANCH_ARGS=(--branch "$_RESOLVED_LLAMA_TAG")
+        fi
+        _BUILD_TMP="${LLAMA_CPP_DIR}.build.$$"
+        rm -rf "$_BUILD_TMP"
+        run_quiet_no_exit "clone llama.cpp" git clone --depth 1 "${_CLONE_BRANCH_ARGS[@]}" https://github.com/ggml-org/llama.cpp.git "$_BUILD_TMP" || BUILD_OK=false
 
         if [ "$BUILD_OK" = true ]; then
             # Skip tests/examples we don't need (faster build)
@@ -449,17 +568,40 @@ rm -rf "$LLAMA_CPP_DIR"
                 echo "   Using ccache for faster compilation"
             fi
 
-            # Detect CUDA: check nvcc on PATH, then common install locations
+            # Detect GPU backend: CUDA (NVIDIA) or ROCm (AMD)
+            GPU_BACKEND=""
+
+            # Check for CUDA: check nvcc on PATH, then common install locations
             NVCC_PATH=""
             if command -v nvcc &>/dev/null; then
                 NVCC_PATH="$(command -v nvcc)"
+                GPU_BACKEND="cuda"
             elif [ -x /usr/local/cuda/bin/nvcc ]; then
                 NVCC_PATH="/usr/local/cuda/bin/nvcc"
                 export PATH="/usr/local/cuda/bin:$PATH"
+                GPU_BACKEND="cuda"
             elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
                 # Pick the newest cuda-XX.X directory
                 NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
                 export PATH="$(dirname "$NVCC_PATH"):$PATH"
+                GPU_BACKEND="cuda"
+            fi
+
+            # Check for ROCm (AMD) only if CUDA was not already selected
+            ROCM_HIPCC=""
+            if [ -z "$GPU_BACKEND" ]; then
+                if command -v hipcc &>/dev/null; then
+                    ROCM_HIPCC="$(command -v hipcc)"
+                    GPU_BACKEND="rocm"
+                elif [ -x /opt/rocm/bin/hipcc ]; then
+                    ROCM_HIPCC="/opt/rocm/bin/hipcc"
+                    export PATH="/opt/rocm/bin:$PATH"
+                    GPU_BACKEND="rocm"
+                elif ls /opt/rocm-*/bin/hipcc &>/dev/null 2>&1; then
+                    ROCM_HIPCC="$(ls -d /opt/rocm-*/bin/hipcc 2>/dev/null | sort -V | tail -1)"
+                    export PATH="$(dirname "$ROCM_HIPCC"):$PATH"
+                    GPU_BACKEND="rocm"
+                fi
             fi
 
             if [ -n "$NVCC_PATH" ]; then
@@ -494,9 +636,53 @@ rm -rf "$LLAMA_CPP_DIR"
 
                 # Multi-threaded nvcc compilation (uses all CPU cores per .cu file)
                 CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_FLAGS=--threads=0"
+            elif [ "$GPU_BACKEND" = "rocm" ]; then
+                # Resolve hipcc symlinks to find the real ROCm root
+                _HIPCC_REAL="$(readlink -f "$ROCM_HIPCC" 2>/dev/null || printf '%s' "$ROCM_HIPCC")"
+                ROCM_ROOT=""
+                if command -v hipconfig &>/dev/null; then
+                    ROCM_ROOT="$(hipconfig -R 2>/dev/null || true)"
+                fi
+                if [ -z "$ROCM_ROOT" ]; then
+                    ROCM_ROOT="$(cd "$(dirname "$_HIPCC_REAL")/.." 2>/dev/null && pwd)"
+                fi
+
+                echo "   Building with ROCm support (AMD GPU, hipcc: $_HIPCC_REAL)..."
+                CMAKE_ARGS="$CMAKE_ARGS -DGGML_HIP=ON"
+                export ROCM_PATH="$ROCM_ROOT"
+                export HIP_PATH="$ROCM_ROOT"
+
+                # Use upstream-recommended HIP compiler (not legacy hipcc-as-CXX)
+                if command -v hipconfig &>/dev/null; then
+                    _HIP_CLANG_DIR="$(hipconfig -l 2>/dev/null || true)"
+                    [ -n "$_HIP_CLANG_DIR" ] && export HIPCXX="$_HIP_CLANG_DIR/clang"
+                fi
+
+                # Detect AMD GPU architecture (gfx target)
+                GPU_TARGETS=""
+                if command -v rocminfo &>/dev/null; then
+                    _gfx_list=$(rocminfo 2>/dev/null | grep -oE 'gfx[0-9]{2,4}[a-z]?' | sort -u || true)
+                    _valid_gfx=""
+                    for _gfx in $_gfx_list; do
+                        if [[ "$_gfx" =~ ^gfx[0-9]{2,4}[a-z]?$ ]]; then
+                            _valid_gfx="${_valid_gfx}${_valid_gfx:+;}$_gfx"
+                        fi
+                    done
+                    [ -n "$_valid_gfx" ] && GPU_TARGETS="$_valid_gfx"
+                fi
+
+                if [ -n "$GPU_TARGETS" ]; then
+                    echo "   AMD GPU architectures: ${GPU_TARGETS//;/, } -- limiting build to detected targets"
+                    CMAKE_ARGS="$CMAKE_ARGS -DGPU_TARGETS=${GPU_TARGETS}"
+                else
+                    echo "   Could not detect AMD GPU arch -- building for default targets (cmake will auto-detect)"
+                fi
             elif [ -d /usr/local/cuda ] || nvidia-smi &>/dev/null; then
                 echo "   CUDA driver detected but nvcc not found — building CPU-only"
                 echo "   To enable GPU: install cuda-toolkit or add nvcc to PATH"
+            elif [ -d /opt/rocm ] || command -v rocm-smi &>/dev/null; then
+                echo "   ROCm driver detected but hipcc not found — building CPU-only"
+                echo "   To enable GPU: install rocm-dev or add hipcc to PATH"
             else
                 echo "   Building CPU-only (no CUDA detected)..."
             fi
@@ -509,21 +695,29 @@ rm -rf "$LLAMA_CPP_DIR"
                 CMAKE_GENERATOR_ARGS="-G Ninja"
             fi
 
-            run_quiet_no_exit "cmake llama.cpp" cmake $CMAKE_GENERATOR_ARGS -S "$LLAMA_CPP_DIR" -B "$LLAMA_CPP_DIR/build" $CMAKE_ARGS || BUILD_OK=false
+            run_quiet_no_exit "cmake llama.cpp" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS || BUILD_OK=false
         fi
 
         if [ "$BUILD_OK" = true ]; then
-            run_quiet_no_exit "build llama-server" cmake --build "$LLAMA_CPP_DIR/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
+            run_quiet_no_exit "build llama-server" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
         fi
 
         # Also build llama-quantize (needed by unsloth-zoo's GGUF export pipeline)
         if [ "$BUILD_OK" = true ]; then
-            run_quiet_no_exit "build llama-quantize" cmake --build "$LLAMA_CPP_DIR/build" --config Release --target llama-quantize -j"$NCPU" || true
-            # Symlink to llama.cpp root — check_llama_cpp() looks for the binary there
+            run_quiet_no_exit "build llama-quantize" cmake --build "$_BUILD_TMP/build" --config Release --target llama-quantize -j"$NCPU" || true
+        fi
+
+        # Swap only after build succeeds -- preserves existing install on failure
+        if [ "$BUILD_OK" = true ]; then
+            rm -rf "$LLAMA_CPP_DIR"
+            mv "$_BUILD_TMP" "$LLAMA_CPP_DIR"
+            # Symlink to llama.cpp root -- check_llama_cpp() looks for the binary there
             QUANTIZE_BIN="$LLAMA_CPP_DIR/build/bin/llama-quantize"
             if [ -f "$QUANTIZE_BIN" ]; then
                 ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize"
             fi
+        else
+            rm -rf "$_BUILD_TMP"
         fi
 
         if [ "$BUILD_OK" = true ]; then
@@ -543,9 +737,15 @@ rm -rf "$LLAMA_CPP_DIR"
 fi  # end _SKIP_GGUF_BUILD check
 
 echo ""
+if [ "${SKIP_STUDIO_BASE:-0}" = "1" ]; then
+    _DONE_LINE="║          Setup Complete!             ║"
+else
+    _DONE_LINE="║          Update Complete!            ║"
+fi
+
 if [ "$IS_COLAB" = true ]; then
     echo "╔══════════════════════════════════════╗"
-    echo "║           Setup Complete!            ║"
+    echo "$_DONE_LINE"
     echo "╠══════════════════════════════════════╣"
     echo "║ Unsloth Studio is ready to start     ║"
     echo "║ in your Colab notebook!              ║"
@@ -555,7 +755,7 @@ if [ "$IS_COLAB" = true ]; then
     echo "╚══════════════════════════════════════╝"
 else
     echo "╔══════════════════════════════════════╗"
-    echo "║           Setup Complete!            ║"
+    echo "$_DONE_LINE"
     echo "╠══════════════════════════════════════╣"
     echo "║ Launch with:                         ║"
     echo "║                                      ║"

tests/python/test_cross_platform_parity.py

@@ -0,0 +1,137 @@
+"""Cross-platform parity tests between install.sh and install.ps1."""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+INSTALL_SH = REPO_ROOT / "install.sh"
+INSTALL_PS1 = REPO_ROOT / "install.ps1"
+
+
+class TestNoTorchBackendAutoInInstallSh:
+    """install.sh primary install paths must not use --torch-backend=auto.
+
+    The fallback else-branch (when TORCH_INDEX_URL is empty) is allowed to
+    use --torch-backend=auto since that is the last-resort recovery path.
+    """
+
+    def test_no_torch_backend_auto_outside_fallback(self):
+        lines = INSTALL_SH.read_text().splitlines()
+        # Find the fallback block: starts with the "else" after the
+        # TORCH_INDEX_URL check and ends at the next "fi".
+        fallback_start = None
+        fallback_end = None
+        for i, line in enumerate(lines):
+            if fallback_start is None and "GPU detection failed" in line:
+                fallback_start = i
+            elif (
+                fallback_start is not None
+                and fallback_end is None
+                and line.strip() == "fi"
+            ):
+                fallback_end = i
+                break
+        fallback_range = (
+            range(fallback_start or 0, (fallback_end or 0) + 1)
+            if fallback_start
+            else range(0)
+        )
+
+        matches = [
+            (i + 1, line)
+            for i, line in enumerate(lines)
+            if "--torch-backend=auto" in line
+            and not line.lstrip().startswith("#")
+            and i not in fallback_range
+        ]
+        assert matches == [], (
+            f"install.sh contains --torch-backend=auto outside the fallback block at lines: "
+            f"{[m[0] for m in matches]}"
+        )
+
+    def test_fallback_uses_torch_backend_auto(self):
+        """The fallback branch should use --torch-backend=auto as recovery."""
+        text = INSTALL_SH.read_text()
+        assert (
+            "GPU detection failed" in text
+        ), "install.sh should have a fallback branch for when GPU detection fails"
+
+
+class TestInstallShHasGpuDetection:
+    """install.sh must contain the get_torch_index_url function."""
+
+    def test_function_exists(self):
+        text = INSTALL_SH.read_text()
+        assert (
+            "get_torch_index_url()" in text
+        ), "install.sh is missing the get_torch_index_url() function"
+
+    def test_torch_index_url_assigned(self):
+        text = INSTALL_SH.read_text()
+        assert (
+            "TORCH_INDEX_URL=$(get_torch_index_url)" in text
+        ), "install.sh should assign TORCH_INDEX_URL from get_torch_index_url()"
+
+
+class TestCudaMappingParity:
+    """CUDA version thresholds must match between install.sh and install.ps1."""
+
+    @staticmethod
+    def _extract_cuda_thresholds_sh(text: str) -> list[str]:
+        """Extract cu* suffixes from the major/minor comparison chain in install.sh."""
+        # Only match lines in the if/elif chain that compare _major/_minor
+        in_func = False
+        results = []
+        for line in text.splitlines():
+            if "get_torch_index_url()" in line:
+                in_func = True
+                continue
+            if in_func and line.startswith("}"):
+                break
+            if in_func and ("_major" in line or "_minor" in line):
+                m = re.search(r"/(cu\d+|cpu)", line)
+                if m:
+                    results.append(m.group(1))
+        return results
+
+    @staticmethod
+    def _extract_cuda_thresholds_ps1(text: str) -> list[str]:
+        """Extract cu* suffixes from the major/minor comparison chain in install.ps1."""
+        in_func = False
+        depth = 0
+        results = []
+        for line in text.splitlines():
+            if "function Get-TorchIndexUrl" in line:
+                in_func = True
+                depth = 1
+                continue
+            if in_func:
+                depth += line.count("{") - line.count("}")
+                if depth <= 0:
+                    break
+                # Only match the if-chain lines that compare $major/$minor
+                if "$major" in line or "$minor" in line:
+                    m = re.search(r"/(cu\d+|cpu)", line)
+                    if m:
+                        results.append(m.group(1))
+        return results
+
+    def test_same_cuda_suffixes(self):
+        """Both scripts should produce the same ordered list of CUDA index suffixes."""
+        sh_text = INSTALL_SH.read_text()
+        ps1_text = INSTALL_PS1.read_text()
+
+        sh_thresholds = self._extract_cuda_thresholds_sh(sh_text)
+        ps1_thresholds = self._extract_cuda_thresholds_ps1(ps1_text)
+
+        assert len(sh_thresholds) > 0, "Could not extract thresholds from install.sh"
+        assert len(ps1_thresholds) > 0, "Could not extract thresholds from install.ps1"
+        assert sh_thresholds == ps1_thresholds, (
+            f"CUDA mapping mismatch:\n"
+            f"  install.sh:  {sh_thresholds}\n"
+            f"  install.ps1: {ps1_thresholds}"
+        )

tests/python/test_install_python_stack.py

@@ -0,0 +1,56 @@
+"""Tests for install_python_stack._build_uv_cmd torch-backend handling."""
+
+from __future__ import annotations
+
+import importlib
+import os
+import sys
+from pathlib import Path
+from unittest import mock
+
+import pytest
+
+# Add the studio directory so we can import install_python_stack
+STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
+sys.path.insert(0, str(STUDIO_DIR))
+
+# _build_uv_cmd lives at module level; import after path setup.
+# We need to mock parts of the module that do work at import time.
+import install_python_stack as ips
+
+
+class TestBuildUvCmdTorchBackend:
+    """Verify _build_uv_cmd only adds --torch-backend when UV_TORCH_BACKEND is set."""
+
+    def _call(self, args: tuple[str, ...] = ()) -> list[str]:
+        return ips._build_uv_cmd(args)
+
+    def test_default_no_torch_backend(self):
+        """Without UV_TORCH_BACKEND env var, no --torch-backend flag."""
+        env = os.environ.copy()
+        env.pop("UV_TORCH_BACKEND", None)
+        with mock.patch.dict(os.environ, env, clear = True):
+            cmd = self._call(("somepackage",))
+        assert not any(
+            a.startswith("--torch-backend") for a in cmd
+        ), f"--torch-backend should not appear by default, got: {cmd}"
+
+    def test_uv_torch_backend_auto(self):
+        """UV_TORCH_BACKEND=auto adds --torch-backend=auto."""
+        with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "auto"}):
+            cmd = self._call(("somepackage",))
+        assert "--torch-backend=auto" in cmd
+
+    def test_uv_torch_backend_cpu(self):
+        """UV_TORCH_BACKEND=cpu adds --torch-backend=cpu."""
+        with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
+            cmd = self._call(("somepackage",))
+        assert "--torch-backend=cpu" in cmd
+
+    def test_uv_torch_backend_empty(self):
+        """UV_TORCH_BACKEND="" (empty string) should NOT add --torch-backend."""
+        with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": ""}):
+            cmd = self._call(("somepackage",))
+        assert not any(
+            a.startswith("--torch-backend") for a in cmd
+        ), f"Empty UV_TORCH_BACKEND should not add flag, got: {cmd}"

tests/run_all.sh

@@ -0,0 +1,16 @@
+#!/bin/sh
+# Run all installer tests.
+set -e
+
+TESTS_DIR="$(cd "$(dirname "$0")" && pwd)"
+
+echo "=== Bash tests ==="
+sh "$TESTS_DIR/sh/test_get_torch_index_url.sh"
+
+echo ""
+echo "=== Python tests ==="
+python -m pytest "$TESTS_DIR/python/test_install_python_stack.py" -v
+python -m pytest "$TESTS_DIR/python/test_cross_platform_parity.py" -v
+
+echo ""
+echo "All tests passed."

tests/sh/test_get_torch_index_url.sh

@@ -0,0 +1,128 @@
+#!/bin/bash
+# Unit tests for get_torch_index_url() from install.sh
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+INSTALL_SH="$SCRIPT_DIR/../../install.sh"
+PASS=0
+FAIL=0
+
+# Extract only the get_torch_index_url function from install.sh
+# Also replace the hardcoded /usr/bin/nvidia-smi fallback with a
+# controllable path so we can test the "no GPU" scenario on GPU machines.
+_FUNC_FILE=$(mktemp)
+_FAKE_SMI_DIR=$(mktemp -d)
+sed -n '/^get_torch_index_url()/,/^}/p' "$INSTALL_SH" \
+    | sed "s|/usr/bin/nvidia-smi|$_FAKE_SMI_DIR/nvidia-smi-absent|g" \
+    > "$_FUNC_FILE"
+
+# Save system PATH so we always have basic tools (uname, grep, head, etc.)
+_SYS_PATH="/usr/local/bin:/usr/bin:/bin"
+
+assert_eq() {
+    _label="$1"; _expected="$2"; _actual="$3"
+    if [ "$_actual" = "$_expected" ]; then
+        echo "  PASS: $_label"
+        PASS=$((PASS + 1))
+    else
+        echo "  FAIL: $_label (expected '$_expected', got '$_actual')"
+        FAIL=$((FAIL + 1))
+    fi
+}
+
+# Helper: create a mock nvidia-smi that prints a given CUDA version string
+make_mock_smi() {
+    _dir=$(mktemp -d)
+    cat > "$_dir/nvidia-smi" <<MOCK
+#!/bin/sh
+cat <<'SMI_OUT'
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 550.54.15              Driver Version: 550.54.15      CUDA Version: $1     |
++-----------------------------------------------------------------------------------------+
+SMI_OUT
+MOCK
+    chmod +x "$_dir/nvidia-smi"
+    echo "$_dir"
+}
+
+# Build a minimal tools directory with symlinks to essential commands
+# (uname, grep, head, etc.) but WITHOUT nvidia-smi.
+_TOOLS_DIR=$(mktemp -d)
+for _cmd in uname grep sed head sh bash cat; do
+    _real=$(command -v "$_cmd" 2>/dev/null || true)
+    [ -n "$_real" ] && ln -sf "$_real" "$_TOOLS_DIR/$_cmd"
+done
+
+# Helper: run get_torch_index_url with a custom PATH
+# $1 = directory with mock nvidia-smi (prepended to PATH), or "none" for no-GPU test
+run_func() {
+    _mock_dir="$1"
+    if [ "$_mock_dir" = "none" ]; then
+        # Minimal PATH with only basic tools, no nvidia-smi anywhere
+        PATH="$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
+    else
+        # Put mock nvidia-smi dir first, then basic tools
+        PATH="$_mock_dir:$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
+    fi
+}
+
+echo "=== test_get_torch_index_url ==="
+
+# 1) No nvidia-smi available -> cpu
+_result=$(run_func "none")
+assert_eq "no nvidia-smi -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
+
+# 2) CUDA 12.6 -> cu126
+_dir=$(make_mock_smi "12.6")
+_result=$(run_func "$_dir")
+assert_eq "CUDA 12.6 -> cu126" "https://download.pytorch.org/whl/cu126" "$_result"
+rm -rf "$_dir"
+
+# 3) CUDA 12.8 -> cu128
+_dir=$(make_mock_smi "12.8")
+_result=$(run_func "$_dir")
+assert_eq "CUDA 12.8 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
+rm -rf "$_dir"
+
+# 4) CUDA 13.0 -> cu130
+_dir=$(make_mock_smi "13.0")
+_result=$(run_func "$_dir")
+assert_eq "CUDA 13.0 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
+rm -rf "$_dir"
+
+# 5) CUDA 12.4 -> cu124
+_dir=$(make_mock_smi "12.4")
+_result=$(run_func "$_dir")
+assert_eq "CUDA 12.4 -> cu124" "https://download.pytorch.org/whl/cu124" "$_result"
+rm -rf "$_dir"
+
+# 6) CUDA 11.8 -> cu118
+_dir=$(make_mock_smi "11.8")
+_result=$(run_func "$_dir")
+assert_eq "CUDA 11.8 -> cu118" "https://download.pytorch.org/whl/cu118" "$_result"
+rm -rf "$_dir"
+
+# 7) CUDA 10.2 (too old) -> cpu
+_dir=$(make_mock_smi "10.2")
+_result=$(run_func "$_dir")
+assert_eq "CUDA 10.2 -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
+rm -rf "$_dir"
+
+# 8) Unparseable nvidia-smi output -> cu126 default
+_dir=$(mktemp -d)
+cat > "$_dir/nvidia-smi" <<'MOCK'
+#!/bin/sh
+echo "something completely unexpected"
+MOCK
+chmod +x "$_dir/nvidia-smi"
+_result=$(run_func "$_dir")
+assert_eq "unparseable -> cu126" "https://download.pytorch.org/whl/cu126" "$_result"
+rm -rf "$_dir"
+
+rm -f "$_FUNC_FILE"
+rm -rf "$_FAKE_SMI_DIR"
+rm -rf "$_TOOLS_DIR"
+
+echo ""
+echo "Results: $PASS passed, $FAIL failed"
+[ "$FAIL" -eq 0 ] || exit 1

tests/studio/install/smoke_test_llama_prebuilt.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import shutil
+import sys
+import tempfile
+import time
+from pathlib import Path
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+INSTALLER_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
+
+
+def load_installer_module():
+    spec = importlib.util.spec_from_file_location(
+        "studio_install_llama_prebuilt", INSTALLER_PATH
+    )
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"unable to load installer module from {INSTALLER_PATH}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+installer = load_installer_module()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description = (
+            "Run a real end-to-end prebuilt llama.cpp install into an isolated temporary "
+            "directory on the current machine."
+        )
+    )
+    parser.add_argument(
+        "--llama-tag",
+        default = "latest",
+        help = "llama.cpp tag to resolve. Defaults to the approved prebuilt tag for this host.",
+    )
+    parser.add_argument(
+        "--published-repo",
+        default = installer.DEFAULT_PUBLISHED_REPO,
+        help = "Published bundle repository used for Linux CUDA selection.",
+    )
+    parser.add_argument(
+        "--published-release-tag",
+        default = installer.DEFAULT_PUBLISHED_TAG or "",
+        help = "Optional published GitHub release tag to pin.",
+    )
+    parser.add_argument(
+        "--work-dir",
+        default = "",
+        help = (
+            "Optional directory under which the smoke install temp dir will be created. "
+            "If omitted, defaults to ./.tmp/llama-prebuilt-smoke under the current directory."
+        ),
+    )
+    parser.add_argument(
+        "--keep-temp",
+        action = "store_true",
+        help = "Keep the temporary smoke install directory after success.",
+    )
+    return parser.parse_args()
+
+
+def smoke_root_base(work_dir: str) -> Path:
+    if work_dir:
+        return Path(work_dir).expanduser().resolve()
+    return (Path.cwd() / ".tmp" / "llama-prebuilt-smoke").resolve()
+
+
+def make_smoke_root(base_dir: Path) -> Path:
+    base_dir.mkdir(parents = True, exist_ok = True)
+    timestamp = time.strftime("%Y%m%d%H%M%S", time.gmtime())
+    return Path(tempfile.mkdtemp(prefix = f"run-{timestamp}-", dir = base_dir))
+
+
+def main() -> int:
+    args = parse_args()
+    host = installer.detect_host()
+    smoke_base = smoke_root_base(args.work_dir)
+    smoke_root = make_smoke_root(smoke_base)
+    install_dir = smoke_root / "install" / "llama.cpp"
+    choice = None
+
+    print(f"[smoke] host={host.system} machine={host.machine}")
+    print(f"[smoke] temp_root={smoke_root}")
+
+    try:
+        requested_tag, resolved_tag, attempts, _approved_checksums = (
+            installer.resolve_install_attempts(
+                args.llama_tag,
+                host,
+                args.published_repo,
+                args.published_release_tag,
+            )
+        )
+        choice = attempts[0]
+        print(f"[smoke] requested_tag={requested_tag}")
+        print(f"[smoke] resolved_tag={resolved_tag}")
+        print(f"[smoke] selected_asset={choice.name}")
+        print(f"[smoke] selected_source={choice.source_label}")
+        print(f"[smoke] install_dir={install_dir}")
+        installer.install_prebuilt(
+            install_dir = install_dir,
+            llama_tag = args.llama_tag,
+            published_repo = args.published_repo,
+            published_release_tag = args.published_release_tag,
+        )
+        print(f"[smoke] PASS install_dir={install_dir}")
+        print(
+            "[smoke] note=This was a real prebuilt install into an isolated temp directory."
+        )
+        return installer.EXIT_SUCCESS
+    except SystemExit as exc:
+        code = int(exc.code) if isinstance(exc.code, int) else installer.EXIT_ERROR
+        if code == installer.EXIT_FALLBACK:
+            print(f"[smoke] FALLBACK install_dir={install_dir}")
+            print(
+                "[smoke] note=Prebuilt path failed and would fall back to source build in setup."
+            )
+            print(installer.collect_system_report(host, choice, install_dir))
+        else:
+            print(f"[smoke] ERROR exit_code={code} install_dir={install_dir}")
+        return code
+    except Exception as exc:
+        print(f"[smoke] ERROR {exc}")
+        print(installer.collect_system_report(host, choice, install_dir))
+        return installer.EXIT_ERROR
+    finally:
+        if args.keep_temp:
+            print(f"[smoke] keeping_temp_root={smoke_root}")
+        elif smoke_root.exists():
+            shutil.rmtree(smoke_root, ignore_errors = True)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/studio/install/test_install_llama_prebuilt_logic.py

@@ -0,0 +1,630 @@
+import importlib.util
+import io
+import json
+import os
+import sys
+import tarfile
+import zipfile
+from pathlib import Path
+
+import pytest
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
+SPEC = importlib.util.spec_from_file_location(
+    "studio_install_llama_prebuilt", MODULE_PATH
+)
+assert SPEC is not None and SPEC.loader is not None
+INSTALL_LLAMA_PREBUILT = importlib.util.module_from_spec(SPEC)
+sys.modules[SPEC.name] = INSTALL_LLAMA_PREBUILT
+SPEC.loader.exec_module(INSTALL_LLAMA_PREBUILT)
+
+PrebuiltFallback = INSTALL_LLAMA_PREBUILT.PrebuiltFallback
+extract_archive = INSTALL_LLAMA_PREBUILT.extract_archive
+binary_env = INSTALL_LLAMA_PREBUILT.binary_env
+HostInfo = INSTALL_LLAMA_PREBUILT.HostInfo
+AssetChoice = INSTALL_LLAMA_PREBUILT.AssetChoice
+ApprovedArtifactHash = INSTALL_LLAMA_PREBUILT.ApprovedArtifactHash
+ApprovedReleaseChecksums = INSTALL_LLAMA_PREBUILT.ApprovedReleaseChecksums
+hydrate_source_tree = INSTALL_LLAMA_PREBUILT.hydrate_source_tree
+validate_prebuilt_choice = INSTALL_LLAMA_PREBUILT.validate_prebuilt_choice
+activate_install_tree = INSTALL_LLAMA_PREBUILT.activate_install_tree
+create_install_staging_dir = INSTALL_LLAMA_PREBUILT.create_install_staging_dir
+sha256_file = INSTALL_LLAMA_PREBUILT.sha256_file
+source_archive_logical_name = INSTALL_LLAMA_PREBUILT.source_archive_logical_name
+
+
+def approved_checksums_for(
+    upstream_tag: str, *, source_archive: Path, bundle_archive: Path, bundle_name: str
+) -> ApprovedReleaseChecksums:
+    return ApprovedReleaseChecksums(
+        repo = "local",
+        release_tag = upstream_tag,
+        upstream_tag = upstream_tag,
+        source_commit = None,
+        artifacts = {
+            source_archive_logical_name(upstream_tag): ApprovedArtifactHash(
+                asset_name = source_archive_logical_name(upstream_tag),
+                sha256 = sha256_file(source_archive),
+                repo = "ggml-org/llama.cpp",
+                kind = "upstream-source",
+            ),
+            bundle_name: ApprovedArtifactHash(
+                asset_name = bundle_name,
+                sha256 = sha256_file(bundle_archive),
+                repo = "local",
+                kind = "local-test-bundle",
+            ),
+        },
+    )
+
+
+def test_extract_archive_allows_safe_tar_symlink_chain(tmp_path: Path):
+    archive_path = tmp_path / "bundle.tar.gz"
+    payload = b"shared-object"
+
+    with tarfile.open(archive_path, "w:gz") as archive:
+        versioned = tarfile.TarInfo("libllama.so.0.0.1")
+        versioned.size = len(payload)
+        archive.addfile(versioned, io_bytes(payload))
+
+        soname = tarfile.TarInfo("libllama.so.0")
+        soname.type = tarfile.SYMTYPE
+        soname.linkname = "libllama.so.0.0.1"
+        archive.addfile(soname)
+
+        linker_name = tarfile.TarInfo("libllama.so")
+        linker_name.type = tarfile.SYMTYPE
+        linker_name.linkname = "libllama.so.0"
+        archive.addfile(linker_name)
+
+    destination = tmp_path / "extract"
+    extract_archive(archive_path, destination)
+
+    assert (destination / "libllama.so.0.0.1").read_bytes() == payload
+    assert (destination / "libllama.so.0").is_symlink()
+    assert (destination / "libllama.so").is_symlink()
+    assert (destination / "libllama.so").resolve().read_bytes() == payload
+
+
+def test_extract_archive_allows_safe_tar_hardlink(tmp_path: Path):
+    archive_path = tmp_path / "bundle.tar.gz"
+    payload = b"quantize"
+
+    with tarfile.open(archive_path, "w:gz") as archive:
+        target = tarfile.TarInfo("llama-quantize")
+        target.size = len(payload)
+        archive.addfile(target, io_bytes(payload))
+
+        hardlink = tarfile.TarInfo("llama-quantize-copy")
+        hardlink.type = tarfile.LNKTYPE
+        hardlink.linkname = "llama-quantize"
+        archive.addfile(hardlink)
+
+    destination = tmp_path / "extract"
+    extract_archive(archive_path, destination)
+
+    assert (destination / "llama-quantize-copy").read_bytes() == payload
+    assert not (destination / "llama-quantize-copy").is_symlink()
+
+
+def test_extract_archive_rejects_absolute_tar_symlink_target(tmp_path: Path):
+    archive_path = tmp_path / "bundle.tar.gz"
+
+    with tarfile.open(archive_path, "w:gz") as archive:
+        entry = tarfile.TarInfo("libllama.so")
+        entry.type = tarfile.SYMTYPE
+        entry.linkname = "/tmp/libllama.so.0"
+        archive.addfile(entry)
+
+    with pytest.raises(PrebuiltFallback, match = "archive link used an absolute target"):
+        extract_archive(archive_path, tmp_path / "extract")
+
+
+def test_extract_archive_rejects_escaping_tar_symlink_target(tmp_path: Path):
+    archive_path = tmp_path / "bundle.tar.gz"
+
+    with tarfile.open(archive_path, "w:gz") as archive:
+        entry = tarfile.TarInfo("libllama.so")
+        entry.type = tarfile.SYMTYPE
+        entry.linkname = "../outside/libllama.so.0"
+        archive.addfile(entry)
+
+    with pytest.raises(PrebuiltFallback, match = "archive link escaped destination"):
+        extract_archive(archive_path, tmp_path / "extract")
+
+
+def test_extract_archive_rejects_unresolved_tar_symlink_target(tmp_path: Path):
+    archive_path = tmp_path / "bundle.tar.gz"
+
+    with tarfile.open(archive_path, "w:gz") as archive:
+        entry = tarfile.TarInfo("libllama.so")
+        entry.type = tarfile.SYMTYPE
+        entry.linkname = "libllama.so.0"
+        archive.addfile(entry)
+
+    with pytest.raises(PrebuiltFallback, match = "unresolved link entries"):
+        extract_archive(archive_path, tmp_path / "extract")
+
+
+def test_extract_archive_rejects_zip_symlink_entry(tmp_path: Path):
+    archive_path = tmp_path / "bundle.zip"
+
+    with zipfile.ZipFile(archive_path, "w") as archive:
+        info = zipfile.ZipInfo("libllama.so")
+        info.create_system = 3
+        info.external_attr = 0o120777 << 16
+        archive.writestr(info, "libllama.so.0")
+
+    with pytest.raises(PrebuiltFallback, match = "zip archive contained a symlink entry"):
+        extract_archive(archive_path, tmp_path / "extract")
+
+
+def test_hydrate_source_tree_extracts_upstream_archive_contents(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    upstream_tag = "b9999"
+    archive_path = tmp_path / "llama.cpp-source.tar.gz"
+    with tarfile.open(archive_path, "w:gz") as archive:
+        add_bytes_to_tar(
+            archive,
+            f"llama.cpp-{upstream_tag}/CMakeLists.txt",
+            b"cmake_minimum_required(VERSION 3.14)\n",
+        )
+        add_bytes_to_tar(
+            archive,
+            f"llama.cpp-{upstream_tag}/convert_hf_to_gguf.py",
+            b"#!/usr/bin/env python3\nimport gguf\n",
+        )
+        add_bytes_to_tar(
+            archive,
+            f"llama.cpp-{upstream_tag}/gguf-py/gguf/__init__.py",
+            b"__all__ = []\n",
+        )
+
+    source_urls = set(INSTALL_LLAMA_PREBUILT.upstream_source_archive_urls(upstream_tag))
+
+    def fake_download_file(url: str, destination: Path) -> None:
+        assert url in source_urls
+        destination.write_bytes(archive_path.read_bytes())
+
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
+
+    install_dir = tmp_path / "install"
+    work_dir = tmp_path / "work"
+    work_dir.mkdir()
+    hydrate_source_tree(
+        upstream_tag, install_dir, work_dir, expected_sha256 = sha256_file(archive_path)
+    )
+
+    assert (install_dir / "CMakeLists.txt").exists()
+    assert (install_dir / "convert_hf_to_gguf.py").exists()
+    assert (install_dir / "gguf-py" / "gguf" / "__init__.py").exists()
+    assert not (install_dir / f"llama.cpp-{upstream_tag}").exists()
+
+
+def test_validate_prebuilt_choice_creates_repo_shaped_linux_install(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    upstream_tag = "b9998"
+    bundle_name = "app-b9998-linux-x64-cuda13-newer.tar.gz"
+    source_archive = tmp_path / "source.tar.gz"
+    bundle_archive = tmp_path / "bundle.tar.gz"
+    with tarfile.open(source_archive, "w:gz") as archive:
+        add_bytes_to_tar(
+            archive,
+            f"llama.cpp-{upstream_tag}/CMakeLists.txt",
+            b"cmake_minimum_required(VERSION 3.14)\n",
+        )
+        add_bytes_to_tar(
+            archive,
+            f"llama.cpp-{upstream_tag}/convert_hf_to_gguf.py",
+            b"#!/usr/bin/env python3\nimport gguf\n",
+        )
+        add_bytes_to_tar(
+            archive,
+            f"llama.cpp-{upstream_tag}/gguf-py/gguf/__init__.py",
+            b"__all__ = []\n",
+        )
+    with tarfile.open(bundle_archive, "w:gz") as archive:
+        add_bytes_to_tar(archive, "llama-server", b"#!/bin/sh\nexit 0\n", mode = 0o755)
+        add_bytes_to_tar(archive, "llama-quantize", b"#!/bin/sh\nexit 0\n", mode = 0o755)
+        add_bytes_to_tar(archive, "libllama.so.0.0.1", b"libllama")
+        add_symlink_to_tar(archive, "libllama.so.0", "libllama.so.0.0.1")
+        add_symlink_to_tar(archive, "libllama.so", "libllama.so.0")
+        add_bytes_to_tar(archive, "libggml.so.0.9.8", b"libggml")
+        add_symlink_to_tar(archive, "libggml.so.0", "libggml.so.0.9.8")
+        add_symlink_to_tar(archive, "libggml.so", "libggml.so.0")
+        add_bytes_to_tar(archive, "libggml-base.so.0.9.8", b"libggml-base")
+        add_symlink_to_tar(archive, "libggml-base.so.0", "libggml-base.so.0.9.8")
+        add_symlink_to_tar(archive, "libggml-base.so", "libggml-base.so.0")
+        add_bytes_to_tar(archive, "libggml-cpu-x64.so.0.9.8", b"libggml-cpu")
+        add_symlink_to_tar(archive, "libggml-cpu-x64.so.0", "libggml-cpu-x64.so.0.9.8")
+        add_symlink_to_tar(archive, "libggml-cpu-x64.so", "libggml-cpu-x64.so.0")
+        add_bytes_to_tar(archive, "libmtmd.so.0.0.1", b"libmtmd")
+        add_symlink_to_tar(archive, "libmtmd.so.0", "libmtmd.so.0.0.1")
+        add_symlink_to_tar(archive, "libmtmd.so", "libmtmd.so.0")
+        add_bytes_to_tar(archive, "BUILD_INFO.txt", b"bundle metadata\n")
+        add_bytes_to_tar(archive, "THIRD_PARTY_LICENSES.txt", b"licenses\n")
+
+    source_urls = set(INSTALL_LLAMA_PREBUILT.upstream_source_archive_urls(upstream_tag))
+
+    def fake_download_file(url: str, destination: Path) -> None:
+        if url in source_urls:
+            destination.write_bytes(source_archive.read_bytes())
+            return
+        if url == "file://bundle":
+            destination.write_bytes(bundle_archive.read_bytes())
+            return
+        raise AssertionError(f"unexpected download url: {url}")
+
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT,
+        "download_bytes",
+        lambda url, **_: b"#!/usr/bin/env python3\nimport gguf\n",
+    )
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT,
+        "preflight_linux_installed_binaries",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None
+    )
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None
+    )
+
+    host = HostInfo(
+        system = "Linux",
+        machine = "x86_64",
+        is_windows = False,
+        is_linux = True,
+        is_macos = False,
+        is_x86_64 = True,
+        is_arm64 = False,
+        nvidia_smi = None,
+        driver_cuda_version = None,
+        compute_caps = [],
+        visible_cuda_devices = None,
+        has_physical_nvidia = False,
+        has_usable_nvidia = False,
+    )
+    choice = AssetChoice(
+        repo = "local",
+        tag = upstream_tag,
+        name = bundle_name,
+        url = "file://bundle",
+        source_label = "local",
+        is_ready_bundle = True,
+        install_kind = "linux-cuda",
+        bundle_profile = "cuda13-newer",
+        runtime_line = "cuda13",
+        expected_sha256 = sha256_file(bundle_archive),
+    )
+
+    install_dir = tmp_path / "install"
+    work_dir = tmp_path / "work"
+    work_dir.mkdir()
+    probe_path = tmp_path / "stories260K.gguf"
+    quantized_path = tmp_path / "stories260K-q4.gguf"
+    validate_prebuilt_choice(
+        choice,
+        host,
+        install_dir,
+        work_dir,
+        probe_path,
+        requested_tag = upstream_tag,
+        llama_tag = upstream_tag,
+        approved_checksums = approved_checksums_for(
+            upstream_tag,
+            source_archive = source_archive,
+            bundle_archive = bundle_archive,
+            bundle_name = bundle_name,
+        ),
+        prebuilt_fallback_used = False,
+        quantized_path = quantized_path,
+    )
+
+    assert (install_dir / "gguf-py" / "gguf" / "__init__.py").exists()
+    assert (install_dir / "convert_hf_to_gguf.py").exists()
+    assert (install_dir / "build" / "bin" / "llama-server").exists()
+    assert (install_dir / "build" / "bin" / "llama-quantize").exists()
+    assert (install_dir / "build" / "bin" / "libllama.so").exists()
+    assert (install_dir / "llama-server").exists()
+    assert (install_dir / "llama-quantize").exists()
+    assert (install_dir / "UNSLOTH_PREBUILT_INFO.json").exists()
+    assert (install_dir / "BUILD_INFO.txt").exists()
+
+
+def test_validate_prebuilt_choice_creates_repo_shaped_windows_install(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    upstream_tag = "b9997"
+    bundle_name = "app-b9997-windows-x64-cpu.zip"
+    source_archive = tmp_path / "source.tar.gz"
+    bundle_archive = tmp_path / "bundle.zip"
+    with tarfile.open(source_archive, "w:gz") as archive:
+        add_bytes_to_tar(
+            archive,
+            f"llama.cpp-{upstream_tag}/CMakeLists.txt",
+            b"cmake_minimum_required(VERSION 3.14)\n",
+        )
+        add_bytes_to_tar(
+            archive,
+            f"llama.cpp-{upstream_tag}/convert_hf_to_gguf.py",
+            b"#!/usr/bin/env python3\nimport gguf\n",
+        )
+        add_bytes_to_tar(
+            archive,
+            f"llama.cpp-{upstream_tag}/gguf-py/gguf/__init__.py",
+            b"__all__ = []\n",
+        )
+    with zipfile.ZipFile(bundle_archive, "w") as archive:
+        archive.writestr("llama-server.exe", b"MZ")
+        archive.writestr("llama-quantize.exe", b"MZ")
+        archive.writestr("llama.dll", b"DLL")
+        archive.writestr("BUILD_INFO.txt", b"bundle metadata\n")
+
+    source_urls = set(INSTALL_LLAMA_PREBUILT.upstream_source_archive_urls(upstream_tag))
+
+    def fake_download_file(url: str, destination: Path) -> None:
+        if url in source_urls:
+            destination.write_bytes(source_archive.read_bytes())
+            return
+        if url == "file://bundle.zip":
+            destination.write_bytes(bundle_archive.read_bytes())
+            return
+        raise AssertionError(f"unexpected download url: {url}")
+
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT,
+        "download_bytes",
+        lambda url, **_: b"#!/usr/bin/env python3\nimport gguf\n",
+    )
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT,
+        "preflight_linux_installed_binaries",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None
+    )
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None
+    )
+
+    host = HostInfo(
+        system = "Windows",
+        machine = "AMD64",
+        is_windows = True,
+        is_linux = False,
+        is_macos = False,
+        is_x86_64 = True,
+        is_arm64 = False,
+        nvidia_smi = None,
+        driver_cuda_version = None,
+        compute_caps = [],
+        visible_cuda_devices = None,
+        has_physical_nvidia = False,
+        has_usable_nvidia = False,
+    )
+    choice = AssetChoice(
+        repo = "local",
+        tag = upstream_tag,
+        name = bundle_name,
+        url = "file://bundle.zip",
+        source_label = "local",
+        is_ready_bundle = True,
+        install_kind = "windows-cpu",
+        expected_sha256 = sha256_file(bundle_archive),
+    )
+
+    install_dir = tmp_path / "install"
+    work_dir = tmp_path / "work"
+    work_dir.mkdir()
+    probe_path = tmp_path / "stories260K.gguf"
+    quantized_path = tmp_path / "stories260K-q4.gguf"
+    validate_prebuilt_choice(
+        choice,
+        host,
+        install_dir,
+        work_dir,
+        probe_path,
+        requested_tag = upstream_tag,
+        llama_tag = upstream_tag,
+        approved_checksums = approved_checksums_for(
+            upstream_tag,
+            source_archive = source_archive,
+            bundle_archive = bundle_archive,
+            bundle_name = bundle_name,
+        ),
+        prebuilt_fallback_used = False,
+        quantized_path = quantized_path,
+    )
+
+    assert (install_dir / "gguf-py" / "gguf" / "__init__.py").exists()
+    assert (install_dir / "convert_hf_to_gguf.py").exists()
+    assert (install_dir / "build" / "bin" / "Release" / "llama-server.exe").exists()
+    assert (install_dir / "build" / "bin" / "Release" / "llama-quantize.exe").exists()
+    assert (install_dir / "build" / "bin" / "Release" / "llama.dll").exists()
+    assert not (install_dir / "llama-server.exe").exists()
+    assert (install_dir / "UNSLOTH_PREBUILT_INFO.json").exists()
+    assert (install_dir / "BUILD_INFO.txt").exists()
+
+
+def test_activate_install_tree_restores_existing_install_after_activation_failure(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    capsys: pytest.CaptureFixture[str],
+):
+    install_dir = tmp_path / "llama.cpp"
+    install_dir.mkdir()
+    (install_dir / "old.txt").write_text("old install\n")
+
+    staging_dir = create_install_staging_dir(install_dir)
+    (staging_dir / "new.txt").write_text("new install\n")
+
+    host = HostInfo(
+        system = "Linux",
+        machine = "x86_64",
+        is_windows = False,
+        is_linux = True,
+        is_macos = False,
+        is_x86_64 = True,
+        is_arm64 = False,
+        nvidia_smi = None,
+        driver_cuda_version = None,
+        compute_caps = [],
+        visible_cuda_devices = None,
+        has_physical_nvidia = False,
+        has_usable_nvidia = False,
+    )
+
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT,
+        "confirm_install_tree",
+        lambda *_args, **_kwargs: (_ for _ in ()).throw(
+            RuntimeError("activation confirm failed")
+        ),
+    )
+
+    with pytest.raises(
+        PrebuiltFallback,
+        match = "activation failed; restored previous install",
+    ):
+        activate_install_tree(staging_dir, install_dir, host)
+
+    assert (install_dir / "old.txt").read_text() == "old install\n"
+    assert not (install_dir / "new.txt").exists()
+    assert not staging_dir.exists()
+    assert not (tmp_path / ".staging").exists()
+
+    output = capsys.readouterr().out
+    assert "moving existing install to rollback path" in output
+    assert "restored previous install from rollback path" in output
+
+
+def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    capsys: pytest.CaptureFixture[str],
+):
+    install_dir = tmp_path / "llama.cpp"
+    install_dir.mkdir()
+    (install_dir / "old.txt").write_text("old install\n")
+
+    staging_dir = create_install_staging_dir(install_dir)
+    (staging_dir / "new.txt").write_text("new install\n")
+
+    host = HostInfo(
+        system = "Linux",
+        machine = "x86_64",
+        is_windows = False,
+        is_linux = True,
+        is_macos = False,
+        is_x86_64 = True,
+        is_arm64 = False,
+        nvidia_smi = None,
+        driver_cuda_version = None,
+        compute_caps = [],
+        visible_cuda_devices = None,
+        has_physical_nvidia = False,
+        has_usable_nvidia = False,
+    )
+
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT,
+        "confirm_install_tree",
+        lambda *_args, **_kwargs: (_ for _ in ()).throw(
+            RuntimeError("activation confirm failed")
+        ),
+    )
+
+    original_replace = INSTALL_LLAMA_PREBUILT.os.replace
+
+    def flaky_replace(src, dst):
+        src_path = Path(src)
+        dst_path = Path(dst)
+        if "rollback-" in src_path.name and dst_path == install_dir:
+            raise OSError("restore failed")
+        return original_replace(src, dst)
+
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT.os, "replace", flaky_replace)
+
+    with pytest.raises(
+        PrebuiltFallback,
+        match = "activation and rollback failed; cleaned install state for fresh source build",
+    ):
+        activate_install_tree(staging_dir, install_dir, host)
+
+    assert not install_dir.exists()
+    assert not staging_dir.exists()
+    assert not (tmp_path / ".staging").exists()
+
+    output = capsys.readouterr().out
+    assert "rollback after failed activation also failed: restore failed" in output
+    assert (
+        "cleaning staging, install, and rollback paths before source build fallback"
+        in output
+    )
+    assert "removing failed install path" in output
+    assert "removing rollback path" in output
+
+
+def test_binary_env_linux_includes_binary_parent_in_ld_library_path(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    install_dir = tmp_path / "llama.cpp"
+    bin_dir = install_dir / "build" / "bin"
+    bin_dir.mkdir(parents = True)
+    binary_path = bin_dir / "llama-server"
+    binary_path.write_bytes(b"fake")
+
+    host = HostInfo(
+        system = "Linux",
+        machine = "x86_64",
+        is_windows = False,
+        is_linux = True,
+        is_macos = False,
+        is_x86_64 = True,
+        is_arm64 = False,
+        nvidia_smi = None,
+        driver_cuda_version = None,
+        compute_caps = [],
+        visible_cuda_devices = None,
+        has_physical_nvidia = False,
+        has_usable_nvidia = False,
+    )
+
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "linux_runtime_dirs", lambda _bp: [])
+
+    env = binary_env(binary_path, install_dir, host)
+    ld_dirs = env["LD_LIBRARY_PATH"].split(os.pathsep)
+    assert (
+        str(bin_dir) in ld_dirs
+    ), f"binary_path.parent ({bin_dir}) must be in LD_LIBRARY_PATH, got: {ld_dirs}"
+    assert str(install_dir) in ld_dirs
+
+
+def io_bytes(data: bytes):
+    return io.BytesIO(data)
+
+
+def add_bytes_to_tar(
+    archive: tarfile.TarFile, name: str, data: bytes, *, mode: int = 0o644
+) -> None:
+    info = tarfile.TarInfo(name)
+    info.size = len(data)
+    info.mode = mode
+    archive.addfile(info, io_bytes(data))
+
+
+def add_symlink_to_tar(archive: tarfile.TarFile, name: str, target: str) -> None:
+    info = tarfile.TarInfo(name)
+    info.type = tarfile.SYMTYPE
+    info.linkname = target
+    archive.addfile(info)

tests/studio/install/test_pr4562_bugfixes.py

@@ -0,0 +1,687 @@
+"""
+Comprehensive tests for PR #4562 bug fixes.
+
+Tests cover:
+  - Bug 1: PS1 detached HEAD on re-run (fetch + checkout -B pattern)
+  - Bug 2: Source-build fallback ignores pinned tag (both .sh and .ps1)
+  - Bug 3: Unix fallback deletes install before checking prerequisites
+  - Bug 4: Linux LD_LIBRARY_PATH missing build/bin
+  - "latest" tag resolution fallback chain (Unsloth -> ggml-org -> raw)
+  - Cross-platform binary_env (Linux, macOS, Windows)
+  - Edge cases: malformed JSON, empty responses, env overrides
+
+Run: pytest tests/studio/install/test_pr4562_bugfixes.py -v
+"""
+
+import importlib.util
+import json
+import os
+import subprocess
+import sys
+import textwrap
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Load the module under test (same pattern as existing test files)
+# ---------------------------------------------------------------------------
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
+SPEC = importlib.util.spec_from_file_location(
+    "studio_install_llama_prebuilt", MODULE_PATH
+)
+assert SPEC is not None and SPEC.loader is not None
+MOD = importlib.util.module_from_spec(SPEC)
+sys.modules[SPEC.name] = MOD
+SPEC.loader.exec_module(MOD)
+
+binary_env = MOD.binary_env
+HostInfo = MOD.HostInfo
+resolve_requested_llama_tag = MOD.resolve_requested_llama_tag
+
+SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
+SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+def make_host(*, system: str) -> HostInfo:
+    """Create a HostInfo for the given OS."""
+    return HostInfo(
+        system = system,
+        machine = "x86_64" if system != "Darwin" else "arm64",
+        is_windows = (system == "Windows"),
+        is_linux = (system == "Linux"),
+        is_macos = (system == "Darwin"),
+        is_x86_64 = (system != "Darwin"),
+        is_arm64 = (system == "Darwin"),
+        nvidia_smi = None,
+        driver_cuda_version = None,
+        compute_caps = [],
+        visible_cuda_devices = None,
+        has_physical_nvidia = False,
+        has_usable_nvidia = False,
+    )
+
+
+BASH = "/bin/bash"
+
+
+def run_bash(script: str, *, timeout: int = 10, env: dict | None = None) -> str:
+    """Run a bash script fragment and return its stdout."""
+    run_env = os.environ.copy()
+    if env:
+        run_env.update(env)
+    result = subprocess.run(
+        [BASH, "-c", script],
+        capture_output = True,
+        text = True,
+        timeout = timeout,
+        env = run_env,
+    )
+    return result.stdout.strip()
+
+
+# =========================================================================
+# TEST GROUP A: binary_env across all platforms (Bug 4 + cross-platform)
+# =========================================================================
+class TestBinaryEnvCrossPlatform:
+    """Test that binary_env returns correct library paths for all OSes."""
+
+    def test_linux_includes_binary_parent_in_ld_library_path(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ):
+        install_dir = tmp_path / "llama.cpp"
+        bin_dir = install_dir / "build" / "bin"
+        bin_dir.mkdir(parents = True)
+        binary_path = bin_dir / "llama-server"
+        binary_path.write_bytes(b"fake")
+
+        host = make_host(system = "Linux")
+        monkeypatch.setattr(MOD, "linux_runtime_dirs", lambda _bp: [])
+
+        env = binary_env(binary_path, install_dir, host)
+        ld_dirs = env["LD_LIBRARY_PATH"].split(os.pathsep)
+        assert str(bin_dir) in ld_dirs, f"build/bin not in LD_LIBRARY_PATH: {ld_dirs}"
+        assert (
+            str(install_dir) in ld_dirs
+        ), f"install_dir not in LD_LIBRARY_PATH: {ld_dirs}"
+
+    def test_linux_binary_parent_comes_before_install_dir(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ):
+        """build/bin should be searched before install_dir for .so files."""
+        install_dir = tmp_path / "llama.cpp"
+        bin_dir = install_dir / "build" / "bin"
+        bin_dir.mkdir(parents = True)
+        binary_path = bin_dir / "llama-server"
+        binary_path.write_bytes(b"fake")
+
+        host = make_host(system = "Linux")
+        monkeypatch.setattr(MOD, "linux_runtime_dirs", lambda _bp: [])
+
+        env = binary_env(binary_path, install_dir, host)
+        ld_dirs = env["LD_LIBRARY_PATH"].split(os.pathsep)
+        bin_idx = ld_dirs.index(str(bin_dir))
+        install_idx = ld_dirs.index(str(install_dir))
+        assert (
+            bin_idx < install_idx
+        ), "binary_path.parent should come before install_dir"
+
+    def test_linux_deduplicates_when_binary_parent_equals_install_dir(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ):
+        """When binary is directly in install_dir, no duplicate entries."""
+        install_dir = tmp_path / "llama.cpp"
+        install_dir.mkdir(parents = True)
+        binary_path = install_dir / "llama-server"
+        binary_path.write_bytes(b"fake")
+
+        host = make_host(system = "Linux")
+        monkeypatch.setattr(MOD, "linux_runtime_dirs", lambda _bp: [])
+
+        env = binary_env(binary_path, install_dir, host)
+        ld_dirs = [d for d in env["LD_LIBRARY_PATH"].split(os.pathsep) if d]
+        count = ld_dirs.count(str(install_dir))
+        assert count == 1, f"install_dir appears {count} times in LD_LIBRARY_PATH"
+
+    def test_linux_preserves_existing_ld_library_path(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ):
+        install_dir = tmp_path / "llama.cpp"
+        bin_dir = install_dir / "build" / "bin"
+        bin_dir.mkdir(parents = True)
+        binary_path = bin_dir / "llama-server"
+        binary_path.write_bytes(b"fake")
+
+        # Create real directories so dedupe_existing_dirs keeps them
+        custom_lib = tmp_path / "custom_lib"
+        other_lib = tmp_path / "other_lib"
+        custom_lib.mkdir()
+        other_lib.mkdir()
+
+        host = make_host(system = "Linux")
+        monkeypatch.setattr(MOD, "linux_runtime_dirs", lambda _bp: [])
+        original = os.environ.get("LD_LIBRARY_PATH", "")
+        os.environ["LD_LIBRARY_PATH"] = f"{custom_lib}:{other_lib}"
+        try:
+            env = binary_env(binary_path, install_dir, host)
+        finally:
+            if original:
+                os.environ["LD_LIBRARY_PATH"] = original
+            else:
+                os.environ.pop("LD_LIBRARY_PATH", None)
+        ld_dirs = env["LD_LIBRARY_PATH"].split(os.pathsep)
+        assert str(custom_lib.resolve()) in ld_dirs
+        assert str(other_lib.resolve()) in ld_dirs
+
+    def test_windows_includes_binary_parent_in_path(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ):
+        install_dir = tmp_path / "llama.cpp"
+        bin_dir = install_dir / "build" / "bin" / "Release"
+        bin_dir.mkdir(parents = True)
+        binary_path = bin_dir / "llama-server.exe"
+        binary_path.write_bytes(b"MZ")
+
+        host = make_host(system = "Windows")
+        monkeypatch.setattr(
+            MOD, "windows_runtime_dirs_for_runtime_line", lambda _rt: []
+        )
+
+        env = binary_env(binary_path, install_dir, host)
+        path_dirs = env["PATH"].split(os.pathsep)
+        assert str(bin_dir) in path_dirs, f"build/bin/Release not in PATH: {path_dirs}"
+
+    def test_macos_sets_dyld_library_path(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ):
+        install_dir = tmp_path / "llama.cpp"
+        install_dir.mkdir(parents = True)
+        bin_dir = install_dir / "build" / "bin"
+        binary_path = bin_dir / "llama-server"
+        binary_path.parent.mkdir(parents = True)
+        binary_path.write_bytes(b"fake")
+
+        host = make_host(system = "Darwin")
+        monkeypatch.delenv("DYLD_LIBRARY_PATH", raising = False)
+
+        env = binary_env(binary_path, install_dir, host)
+        dyld_parts = [p for p in env["DYLD_LIBRARY_PATH"].split(os.pathsep) if p]
+        assert (
+            str(bin_dir) in dyld_parts
+        ), f"build/bin not in DYLD_LIBRARY_PATH: {dyld_parts}"
+        assert (
+            str(install_dir) in dyld_parts
+        ), f"install_dir not in DYLD_LIBRARY_PATH: {dyld_parts}"
+        # binary_path.parent (build/bin) should come before install_dir
+        assert dyld_parts.index(str(bin_dir)) < dyld_parts.index(str(install_dir))
+
+
+# =========================================================================
+# TEST GROUP B: resolve_requested_llama_tag (Python function)
+# =========================================================================
+class TestResolveRequestedLlamaTag:
+    def test_concrete_tag_passes_through(self):
+        assert resolve_requested_llama_tag("b8508") == "b8508"
+
+    def test_none_resolves_to_latest(self, monkeypatch: pytest.MonkeyPatch):
+        monkeypatch.setattr(MOD, "latest_upstream_release_tag", lambda: "b9999")
+        assert resolve_requested_llama_tag(None) == "b9999"
+
+    def test_latest_resolves_to_upstream(self, monkeypatch: pytest.MonkeyPatch):
+        monkeypatch.setattr(MOD, "latest_upstream_release_tag", lambda: "b1234")
+        assert resolve_requested_llama_tag("latest") == "b1234"
+
+    def test_empty_string_resolves_to_latest(self, monkeypatch: pytest.MonkeyPatch):
+        monkeypatch.setattr(MOD, "latest_upstream_release_tag", lambda: "b5555")
+        assert resolve_requested_llama_tag("") == "b5555"
+
+
+# =========================================================================
+# TEST GROUP C: setup.sh logic (bash subprocess tests)
+# =========================================================================
+class TestSetupShLogic:
+    """Test setup.sh fragments via bash subprocess with controlled PATH."""
+
+    def test_cmake_missing_preserves_install(self, tmp_path: Path):
+        """Bug 3: When cmake is missing, rm -rf should NOT run."""
+        llama_dir = tmp_path / "llama.cpp"
+        llama_dir.mkdir()
+        marker = llama_dir / "marker.txt"
+        marker.write_text("existing")
+
+        mock_bin = tmp_path / "mock_bin"
+        mock_bin.mkdir()
+        # Create mock git but NOT cmake
+        (mock_bin / "git").write_text("#!/bin/bash\nexit 0\n")
+        (mock_bin / "git").chmod(0o755)
+
+        # Build PATH: mock_bin first, then system dirs WITHOUT cmake
+        safe_dirs = [str(mock_bin)]
+        for d in os.environ.get("PATH", "").split(":"):
+            if d and not os.path.isfile(os.path.join(d, "cmake")):
+                safe_dirs.append(d)
+
+        script = textwrap.dedent(f"""\
+            export LLAMA_CPP_DIR="{llama_dir}"
+            if ! command -v cmake &>/dev/null; then
+                echo "cmake_missing"
+            elif ! command -v git &>/dev/null; then
+                echo "git_missing"
+            else
+                rm -rf "$LLAMA_CPP_DIR"
+                echo "would_clone"
+            fi
+        """)
+        output = run_bash(script, env = {"PATH": ":".join(safe_dirs)})
+        assert "cmake_missing" in output
+        assert marker.exists(), "Install dir was deleted despite cmake missing!"
+
+    def test_git_missing_preserves_install(self, tmp_path: Path):
+        """Bug 3: When git is missing, rm -rf should NOT run."""
+        llama_dir = tmp_path / "llama.cpp"
+        llama_dir.mkdir()
+        marker = llama_dir / "marker.txt"
+        marker.write_text("existing")
+
+        mock_bin = tmp_path / "mock_bin"
+        mock_bin.mkdir()
+        # Create mock cmake but NOT git
+        (mock_bin / "cmake").write_text("#!/bin/bash\nexit 0\n")
+        (mock_bin / "cmake").chmod(0o755)
+
+        # Build PATH: mock_bin first, then system dirs WITHOUT git
+        safe_dirs = [str(mock_bin)]
+        for d in os.environ.get("PATH", "").split(":"):
+            if d and not os.path.isfile(os.path.join(d, "git")):
+                safe_dirs.append(d)
+
+        script = textwrap.dedent(f"""\
+            export LLAMA_CPP_DIR="{llama_dir}"
+            if ! command -v cmake &>/dev/null; then
+                echo "cmake_missing"
+            elif ! command -v git &>/dev/null; then
+                echo "git_missing"
+            else
+                rm -rf "$LLAMA_CPP_DIR"
+                echo "would_clone"
+            fi
+        """)
+        output = run_bash(script, env = {"PATH": ":".join(safe_dirs)})
+        assert "git_missing" in output
+        assert marker.exists(), "Install dir was deleted despite git missing!"
+
+    def test_both_present_runs_rm_and_clone(self, tmp_path: Path):
+        """Bug 3: When both present, rm -rf runs before clone."""
+        llama_dir = tmp_path / "llama.cpp"
+        llama_dir.mkdir()
+        marker = llama_dir / "marker.txt"
+        marker.write_text("existing")
+
+        mock_bin = tmp_path / "mock_bin"
+        mock_bin.mkdir()
+        (mock_bin / "cmake").write_text("#!/bin/bash\nexit 0\n")
+        (mock_bin / "cmake").chmod(0o755)
+        (mock_bin / "git").write_text("#!/bin/bash\nexit 0\n")
+        (mock_bin / "git").chmod(0o755)
+
+        script = textwrap.dedent(f"""\
+            export PATH="{mock_bin}:$PATH"
+            export LLAMA_CPP_DIR="{llama_dir}"
+            if ! command -v cmake &>/dev/null; then
+                echo "cmake_missing"
+            elif ! command -v git &>/dev/null; then
+                echo "git_missing"
+            else
+                rm -rf "$LLAMA_CPP_DIR"
+                echo "would_clone"
+            fi
+        """)
+        output = run_bash(script)
+        assert "would_clone" in output
+        assert not marker.exists(), "Install dir should have been deleted"
+
+    def test_clone_uses_pinned_tag(self, tmp_path: Path):
+        """Bug 2: git clone should use --branch with the resolved tag."""
+        mock_bin = tmp_path / "mock_bin"
+        mock_bin.mkdir()
+        log_file = tmp_path / "git_calls.log"
+        (mock_bin / "git").write_text(f'#!/bin/bash\necho "$*" >> {log_file}\nexit 0\n')
+        (mock_bin / "git").chmod(0o755)
+
+        script = textwrap.dedent(f"""\
+            export PATH="{mock_bin}:$PATH"
+            git clone --depth 1 --branch "b8508" https://github.com/ggml-org/llama.cpp.git /tmp/llama_test
+        """)
+        run_bash(script)
+        log = log_file.read_text()
+        assert "--branch b8508" in log, f"Expected --branch b8508 in: {log}"
+
+    def test_fetch_checkout_b_pattern(self, tmp_path: Path):
+        """Bug 1: Re-run should use fetch + checkout -B, not pull + checkout FETCH_HEAD."""
+        mock_bin = tmp_path / "mock_bin"
+        mock_bin.mkdir()
+        log_file = tmp_path / "git_calls.log"
+        (mock_bin / "git").write_text(f'#!/bin/bash\necho "$*" >> {log_file}\nexit 0\n')
+        (mock_bin / "git").chmod(0o755)
+
+        llama_dir = tmp_path / "llama.cpp"
+        llama_dir.mkdir()
+        (llama_dir / ".git").mkdir()
+
+        script = textwrap.dedent(f"""\
+            export PATH="{mock_bin}:$PATH"
+            LlamaCppDir="{llama_dir}"
+            ResolvedLlamaTag="b8508"
+            if [ -d "$LlamaCppDir/.git" ]; then
+                git -C "$LlamaCppDir" fetch --depth 1 origin "$ResolvedLlamaTag"
+                if [ $? -ne 0 ]; then
+                    echo "WARN: fetch failed"
+                else
+                    git -C "$LlamaCppDir" checkout -B unsloth-llama-build FETCH_HEAD
+                fi
+            fi
+        """)
+        run_bash(script)
+        log = log_file.read_text()
+        assert "fetch --depth 1 origin b8508" in log
+        assert "checkout -B unsloth-llama-build FETCH_HEAD" in log
+        assert "pull" not in log, "Should use fetch, not pull"
+
+    def test_fetch_failure_warns_not_aborts(self, tmp_path: Path):
+        """Bug 1: fetch failure should warn and continue, not set BuildOk=false."""
+        mock_bin = tmp_path / "mock_bin"
+        mock_bin.mkdir()
+        (mock_bin / "git").write_text(
+            '#!/bin/bash\nif echo "$*" | grep -q fetch; then exit 1; fi\nexit 0\n'
+        )
+        (mock_bin / "git").chmod(0o755)
+
+        llama_dir = tmp_path / "llama.cpp"
+        llama_dir.mkdir()
+        (llama_dir / ".git").mkdir()
+
+        script = textwrap.dedent(f"""\
+            export PATH="{mock_bin}:$PATH"
+            LlamaCppDir="{llama_dir}"
+            ResolvedLlamaTag="b8508"
+            BuildOk=true
+            if [ -d "$LlamaCppDir/.git" ]; then
+                git -C "$LlamaCppDir" fetch --depth 1 origin "$ResolvedLlamaTag"
+                if [ $? -ne 0 ]; then
+                    echo "WARN: fetch failed -- using existing source"
+                else
+                    git -C "$LlamaCppDir" checkout -B unsloth-llama-build FETCH_HEAD
+                fi
+            fi
+            echo "BuildOk=$BuildOk"
+        """)
+        output = run_bash(script)
+        assert "WARN: fetch failed" in output
+        assert "BuildOk=true" in output
+
+
+# =========================================================================
+# TEST GROUP D: "latest" tag resolution (bash subprocess)
+# =========================================================================
+class TestLatestTagResolution:
+    """Test the fallback chain: Unsloth API -> ggml-org API -> raw."""
+
+    RESOLVE_TEMPLATE = textwrap.dedent("""\
+        export PATH="{mock_bin}:$PATH"
+        _REQUESTED_LLAMA_TAG="{requested_tag}"
+        _RESOLVED_LLAMA_TAG=""
+        _RESOLVE_UPSTREAM_STATUS=1
+        _HELPER_RELEASE_REPO="unslothai/llama.cpp"
+        if [ "$_RESOLVE_UPSTREAM_STATUS" -ne 0 ] || [ -z "$_RESOLVED_LLAMA_TAG" ]; then
+            if [ "$_REQUESTED_LLAMA_TAG" = "latest" ]; then
+                _RESOLVED_LLAMA_TAG="$(curl -fsSL "https://api.github.com/repos/${{_HELPER_RELEASE_REPO}}/releases/latest" 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG=""
+                if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
+                    _RESOLVED_LLAMA_TAG="$(curl -fsSL https://api.github.com/repos/ggml-org/llama.cpp/releases/latest 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG=""
+                fi
+            fi
+            if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
+                _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
+            fi
+        fi
+        echo "$_RESOLVED_LLAMA_TAG"
+    """)
+
+    @staticmethod
+    def _make_curl_mock(
+        mock_bin: Path, unsloth_response: str | None, ggml_response: str | None
+    ):
+        """Create a curl mock that returns different responses per repo."""
+        lines = ["#!/bin/bash"]
+        if unsloth_response is not None:
+            lines.append(
+                f'if echo "$*" | grep -q "unslothai/llama.cpp"; then echo \'{unsloth_response}\'; exit 0; fi'
+            )
+        else:
+            lines.append(
+                'if echo "$*" | grep -q "unslothai/llama.cpp"; then exit 1; fi'
+            )
+        if ggml_response is not None:
+            lines.append(
+                f'if echo "$*" | grep -q "ggml-org/llama.cpp"; then echo \'{ggml_response}\'; exit 0; fi'
+            )
+        else:
+            lines.append('if echo "$*" | grep -q "ggml-org/llama.cpp"; then exit 1; fi')
+        lines.append("exit 1")
+        curl_path = mock_bin / "curl"
+        curl_path.write_text("\n".join(lines) + "\n")
+        curl_path.chmod(0o755)
+
+    def _run_resolve(
+        self,
+        tmp_path: Path,
+        requested_tag: str,
+        unsloth_resp: str | None,
+        ggml_resp: str | None,
+    ) -> str:
+        mock_bin = tmp_path / "mock_bin"
+        mock_bin.mkdir(exist_ok = True)
+        self._make_curl_mock(mock_bin, unsloth_resp, ggml_resp)
+        script = self.RESOLVE_TEMPLATE.format(
+            mock_bin = mock_bin, requested_tag = requested_tag
+        )
+        return run_bash(script)
+
+    def test_unsloth_succeeds(self, tmp_path: Path):
+        output = self._run_resolve(
+            tmp_path,
+            "latest",
+            unsloth_resp = '{"tag_name":"b8508"}',
+            ggml_resp = '{"tag_name":"b9000"}',
+        )
+        assert output == "b8508"
+
+    def test_unsloth_fails_ggml_succeeds(self, tmp_path: Path):
+        output = self._run_resolve(
+            tmp_path,
+            "latest",
+            unsloth_resp = None,
+            ggml_resp = '{"tag_name":"b9000"}',
+        )
+        assert output == "b9000"
+
+    def test_both_fail_raw_fallback(self, tmp_path: Path):
+        output = self._run_resolve(
+            tmp_path,
+            "latest",
+            unsloth_resp = None,
+            ggml_resp = None,
+        )
+        assert output == "latest"
+
+    def test_concrete_tag_passes_through(self, tmp_path: Path):
+        output = self._run_resolve(
+            tmp_path,
+            "b7777",
+            unsloth_resp = '{"tag_name":"b8508"}',
+            ggml_resp = '{"tag_name":"b9000"}',
+        )
+        assert output == "b7777"
+
+    def test_unsloth_malformed_json_falls_through(self, tmp_path: Path):
+        output = self._run_resolve(
+            tmp_path,
+            "latest",
+            unsloth_resp = '{"bad_key":"no_tag"}',
+            ggml_resp = '{"tag_name":"b9001"}',
+        )
+        assert output == "b9001"
+
+    def test_both_malformed_json_raw_fallback(self, tmp_path: Path):
+        output = self._run_resolve(
+            tmp_path,
+            "latest",
+            unsloth_resp = '{"bad":"data"}',
+            ggml_resp = '{"also":"bad"}',
+        )
+        assert output == "latest"
+
+    def test_unsloth_empty_body_falls_through(self, tmp_path: Path):
+        output = self._run_resolve(
+            tmp_path,
+            "latest",
+            unsloth_resp = "",
+            ggml_resp = '{"tag_name":"b7000"}',
+        )
+        assert output == "b7000"
+
+    def test_unsloth_empty_tag_name_falls_through(self, tmp_path: Path):
+        output = self._run_resolve(
+            tmp_path,
+            "latest",
+            unsloth_resp = '{"tag_name":""}',
+            ggml_resp = '{"tag_name":"b6000"}',
+        )
+        assert output == "b6000"
+
+    def test_env_override_unsloth_llama_tag(self):
+        output = run_bash(
+            'echo "${UNSLOTH_LLAMA_TAG:-latest}"',
+            env = {"UNSLOTH_LLAMA_TAG": "b1234"},
+        )
+        assert output == "b1234"
+
+    def test_env_unset_defaults_to_latest(self):
+        env = os.environ.copy()
+        env.pop("UNSLOTH_LLAMA_TAG", None)
+        output = run_bash('echo "${UNSLOTH_LLAMA_TAG:-latest}"', env = env)
+        assert output == "latest"
+
+    def test_env_empty_defaults_to_latest(self):
+        output = run_bash(
+            'echo "${UNSLOTH_LLAMA_TAG:-latest}"',
+            env = {"UNSLOTH_LLAMA_TAG": ""},
+        )
+        assert output == "latest"
+
+
+# =========================================================================
+# TEST GROUP E: Source file verification
+# =========================================================================
+class TestSourceCodePatterns:
+    """Verify the actual source files contain the expected fix patterns."""
+
+    def test_setup_sh_no_rm_before_prereq_check(self):
+        """rm -rf must appear AFTER cmake/git checks, not before."""
+        content = SETUP_SH.read_text()
+        # Find the source-build block
+        idx_else = content.find("# Check prerequisites")
+        assert idx_else != -1
+        block = content[idx_else:]
+        # rm -rf should appear after the cmake/git checks
+        idx_cmake = block.find("command -v cmake")
+        idx_git = block.find("command -v git")
+        idx_rm = block.find("rm -rf")
+        assert idx_rm > idx_cmake, "rm -rf should come after cmake check"
+        assert idx_rm > idx_git, "rm -rf should come after git check"
+
+    def test_setup_sh_clone_uses_branch_tag(self):
+        """git clone in source-build should use --branch via _CLONE_BRANCH_ARGS."""
+        content = SETUP_SH.read_text()
+        # The clone line should use _CLONE_BRANCH_ARGS (which conditionally includes --branch)
+        assert (
+            "_CLONE_BRANCH_ARGS" in content
+        ), "Clone should use _CLONE_BRANCH_ARGS array"
+        assert (
+            '--branch "$_RESOLVED_LLAMA_TAG"' in content
+        ), "_CLONE_BRANCH_ARGS should be set to --branch $_RESOLVED_LLAMA_TAG"
+        # Verify the guard: --branch is only used when tag is not "latest"
+        assert (
+            '_RESOLVED_LLAMA_TAG" != "latest"' in content
+        ), "Should guard against literal 'latest' tag"
+
+    def test_setup_sh_latest_resolution_queries_unsloth_first(self):
+        """The Unsloth repo should be queried before ggml-org."""
+        content = SETUP_SH.read_text()
+        idx_unsloth = content.find("_HELPER_RELEASE_REPO}/releases/latest")
+        idx_ggml = content.find("ggml-org/llama.cpp/releases/latest")
+        assert idx_unsloth != -1, "Unsloth API query not found"
+        assert idx_ggml != -1, "ggml-org API query not found"
+        assert idx_unsloth < idx_ggml, "Unsloth should be queried before ggml-org"
+
+    def test_setup_ps1_uses_checkout_b(self):
+        """PS1 should use checkout -B, not checkout --force FETCH_HEAD."""
+        content = SETUP_PS1.read_text()
+        assert "checkout -B unsloth-llama-build" in content
+        assert "checkout --force FETCH_HEAD" not in content
+
+    def test_setup_ps1_clone_uses_branch_tag(self):
+        """PS1 clone should use --branch with the resolved tag."""
+        content = SETUP_PS1.read_text()
+        assert "--branch" in content and "$ResolvedLlamaTag" in content
+        # The old commented-out line should be gone
+        assert "# git clone --depth 1 --branch" not in content
+
+    def test_setup_ps1_no_git_pull(self):
+        """PS1 should use fetch, not pull (which fails in detached HEAD)."""
+        content = SETUP_PS1.read_text()
+        # In the source-build section, there should be no "git pull"
+        # (git pull is only valid on a branch)
+        lines = content.splitlines()
+        for i, line in enumerate(lines):
+            stripped = line.strip()
+            if "git pull" in stripped and not stripped.startswith("#"):
+                # Check context -- should not be in the llama.cpp build section
+                # Allow git pull in other contexts
+                context = "\n".join(lines[max(0, i - 5) : i + 5])
+                if "LlamaCppDir" in context:
+                    pytest.fail(
+                        f"Found 'git pull' in llama.cpp build section at line {i+1}"
+                    )
+
+    def test_setup_ps1_latest_resolution_queries_unsloth_first(self):
+        """PS1 should query Unsloth repo before ggml-org."""
+        content = SETUP_PS1.read_text()
+        idx_unsloth = content.find("$HelperReleaseRepo/releases/latest")
+        idx_ggml = content.find("ggml-org/llama.cpp/releases/latest")
+        assert idx_unsloth != -1, "Unsloth API query not found in PS1"
+        assert idx_ggml != -1, "ggml-org API query not found in PS1"
+        assert idx_unsloth < idx_ggml, "Unsloth should be queried before ggml-org"
+
+    def test_binary_env_linux_has_binary_parent(self):
+        """The Linux branch of binary_env should include binary_path.parent."""
+        content = MODULE_PATH.read_text()
+        # Find the binary_env function
+        in_func = False
+        in_linux = False
+        found = False
+        for line in content.splitlines():
+            if "def binary_env(" in line:
+                in_func = True
+            elif in_func and line and not line[0].isspace() and "def " in line:
+                break
+            if in_func and "host.is_linux" in line:
+                in_linux = True
+            if in_linux and "binary_path.parent" in line:
+                found = True
+                break
+        assert found, "binary_path.parent not found in Linux branch of binary_env"

tests/studio/install/test_selection_logic.py

@@ -0,0 +1,903 @@
+"""Tests for binary selection logic in install_llama_prebuilt.py.
+
+Covers: normalize_compute_cap, normalize_compute_caps, parse_cuda_visible_devices,
+supports_explicit_visible_device_matching, select_visible_gpu_rows,
+compatible_linux_runtime_lines, pick_windows_cuda_runtime,
+compatible_windows_runtime_lines, runtime_line_from_cuda_version,
+apply_approved_hashes, linux_cuda_choice_from_release, windows_cuda_attempts,
+resolve_upstream_asset_choice.
+
+No GPU, no network, no torch required -- all I/O is monkeypatched.
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+
+import pytest
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
+SPEC = importlib.util.spec_from_file_location(
+    "studio_install_llama_prebuilt", MODULE_PATH
+)
+assert SPEC is not None and SPEC.loader is not None
+INSTALL_LLAMA_PREBUILT = importlib.util.module_from_spec(SPEC)
+sys.modules[SPEC.name] = INSTALL_LLAMA_PREBUILT
+SPEC.loader.exec_module(INSTALL_LLAMA_PREBUILT)
+
+HostInfo = INSTALL_LLAMA_PREBUILT.HostInfo
+AssetChoice = INSTALL_LLAMA_PREBUILT.AssetChoice
+PublishedLlamaArtifact = INSTALL_LLAMA_PREBUILT.PublishedLlamaArtifact
+PublishedReleaseBundle = INSTALL_LLAMA_PREBUILT.PublishedReleaseBundle
+ApprovedArtifactHash = INSTALL_LLAMA_PREBUILT.ApprovedArtifactHash
+ApprovedReleaseChecksums = INSTALL_LLAMA_PREBUILT.ApprovedReleaseChecksums
+PrebuiltFallback = INSTALL_LLAMA_PREBUILT.PrebuiltFallback
+LinuxCudaSelection = INSTALL_LLAMA_PREBUILT.LinuxCudaSelection
+UPSTREAM_REPO = INSTALL_LLAMA_PREBUILT.UPSTREAM_REPO
+
+normalize_compute_cap = INSTALL_LLAMA_PREBUILT.normalize_compute_cap
+normalize_compute_caps = INSTALL_LLAMA_PREBUILT.normalize_compute_caps
+parse_cuda_visible_devices = INSTALL_LLAMA_PREBUILT.parse_cuda_visible_devices
+supports_explicit_visible_device_matching = (
+    INSTALL_LLAMA_PREBUILT.supports_explicit_visible_device_matching
+)
+select_visible_gpu_rows = INSTALL_LLAMA_PREBUILT.select_visible_gpu_rows
+compatible_linux_runtime_lines = INSTALL_LLAMA_PREBUILT.compatible_linux_runtime_lines
+pick_windows_cuda_runtime = INSTALL_LLAMA_PREBUILT.pick_windows_cuda_runtime
+compatible_windows_runtime_lines = (
+    INSTALL_LLAMA_PREBUILT.compatible_windows_runtime_lines
+)
+runtime_line_from_cuda_version = INSTALL_LLAMA_PREBUILT.runtime_line_from_cuda_version
+apply_approved_hashes = INSTALL_LLAMA_PREBUILT.apply_approved_hashes
+linux_cuda_choice_from_release = INSTALL_LLAMA_PREBUILT.linux_cuda_choice_from_release
+windows_cuda_attempts = INSTALL_LLAMA_PREBUILT.windows_cuda_attempts
+resolve_upstream_asset_choice = INSTALL_LLAMA_PREBUILT.resolve_upstream_asset_choice
+
+
+# ---------------------------------------------------------------------------
+# Helper factories
+# ---------------------------------------------------------------------------
+
+
+def make_host(**overrides):
+    system = overrides.pop("system", "Linux")
+    machine = overrides.pop("machine", "x86_64")
+    defaults = dict(
+        system = system,
+        machine = machine,
+        is_linux = system == "Linux",
+        is_windows = system == "Windows",
+        is_macos = system == "Darwin",
+        is_x86_64 = machine.lower() in {"x86_64", "amd64"},
+        is_arm64 = machine.lower() in {"arm64", "aarch64"},
+        nvidia_smi = "/usr/bin/nvidia-smi",
+        driver_cuda_version = (12, 8),
+        compute_caps = ["86"],
+        visible_cuda_devices = None,
+        has_physical_nvidia = True,
+        has_usable_nvidia = True,
+    )
+    defaults.update(overrides)
+    return HostInfo(**defaults)
+
+
+def make_artifact(asset_name, **overrides):
+    defaults = dict(
+        asset_name = asset_name,
+        install_kind = "linux-cuda",
+        runtime_line = "cuda12",
+        coverage_class = "targeted",
+        supported_sms = ["75", "80", "86", "89", "90"],
+        min_sm = 75,
+        max_sm = 90,
+        bundle_profile = "cuda12-newer",
+        rank = 100,
+    )
+    defaults.update(overrides)
+    return PublishedLlamaArtifact(**defaults)
+
+
+def make_release(artifacts, **overrides):
+    defaults = dict(
+        repo = "unslothai/llama.cpp",
+        release_tag = "v1.0",
+        upstream_tag = "b8508",
+        assets = {a.asset_name: f"https://example.com/{a.asset_name}" for a in artifacts},
+        manifest_asset_name = "llama-prebuilt-manifest.json",
+        artifacts = artifacts,
+        selection_log = [],
+    )
+    defaults.update(overrides)
+    return PublishedReleaseBundle(**defaults)
+
+
+def make_checksums(asset_names):
+    return ApprovedReleaseChecksums(
+        repo = "unslothai/llama.cpp",
+        release_tag = "v1.0",
+        upstream_tag = "b8508",
+        source_commit = None,
+        artifacts = {
+            name: ApprovedArtifactHash(
+                asset_name = name,
+                sha256 = "a" * 64,
+                repo = "unslothai/llama.cpp",
+                kind = "prebuilt",
+            )
+            for name in asset_names
+        },
+    )
+
+
+def mock_linux_runtime(monkeypatch, lines):
+    dirs = {line: ["/usr/lib/stub"] for line in lines}
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT,
+        "detected_linux_runtime_lines",
+        lambda: (list(lines), dict(dirs)),
+    )
+
+
+def mock_windows_runtime(monkeypatch, lines):
+    dirs = {line: ["C:\\Windows\\System32"] for line in lines}
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT,
+        "detected_windows_runtime_lines",
+        lambda: (list(lines), dict(dirs)),
+    )
+
+
+# ===========================================================================
+# A. normalize_compute_cap
+# ===========================================================================
+
+
+class TestNormalizeComputeCap:
+    def test_dotted_86(self):
+        assert normalize_compute_cap("8.6") == "86"
+
+    def test_dotted_leading_zero(self):
+        assert normalize_compute_cap("07.05") == "75"
+
+    def test_already_normalized(self):
+        assert normalize_compute_cap("75") == "75"
+
+    def test_int_input(self):
+        assert normalize_compute_cap(86) == "86"
+
+    def test_empty_string(self):
+        assert normalize_compute_cap("") is None
+
+    def test_whitespace(self):
+        assert normalize_compute_cap("  ") is None
+
+    def test_non_numeric(self):
+        assert normalize_compute_cap("x.y") is None
+
+    def test_triple_part(self):
+        assert normalize_compute_cap("8.6.0") is None
+
+    def test_zero_minor(self):
+        assert normalize_compute_cap("9.0") == "90"
+
+
+# ===========================================================================
+# B. normalize_compute_caps
+# ===========================================================================
+
+
+class TestNormalizeComputeCaps:
+    def test_deduplication(self):
+        assert normalize_compute_caps(["8.6", "86", "8.6"]) == ["86"]
+
+    def test_numeric_sort(self):
+        assert normalize_compute_caps(["9.0", "7.5", "8.6"]) == ["75", "86", "90"]
+
+    def test_drops_invalid(self):
+        assert normalize_compute_caps(["8.6", "bad", "", "7.5"]) == ["75", "86"]
+
+    def test_empty_input(self):
+        assert normalize_compute_caps([]) == []
+
+
+# ===========================================================================
+# C. parse_cuda_visible_devices
+# ===========================================================================
+
+
+class TestParseCudaVisibleDevices:
+    def test_none(self):
+        assert parse_cuda_visible_devices(None) is None
+
+    def test_empty(self):
+        assert parse_cuda_visible_devices("") == []
+
+    def test_minus_one(self):
+        assert parse_cuda_visible_devices("-1") == []
+
+    def test_single(self):
+        assert parse_cuda_visible_devices("0") == ["0"]
+
+    def test_multi(self):
+        assert parse_cuda_visible_devices("0,1,2") == ["0", "1", "2"]
+
+    def test_whitespace_stripped(self):
+        assert parse_cuda_visible_devices(" 0 , 1 ") == ["0", "1"]
+
+
+# ===========================================================================
+# D. supports_explicit_visible_device_matching
+# ===========================================================================
+
+
+class TestSupportsExplicitVisibleDeviceMatching:
+    def test_all_digits(self):
+        assert supports_explicit_visible_device_matching(["0", "1", "2"]) is True
+
+    def test_gpu_prefix(self):
+        assert supports_explicit_visible_device_matching(["GPU-abc123"]) is True
+
+    def test_none(self):
+        assert supports_explicit_visible_device_matching(None) is False
+
+    def test_empty(self):
+        assert supports_explicit_visible_device_matching([]) is False
+
+    def test_mixed_invalid(self):
+        assert supports_explicit_visible_device_matching(["0", "MIG-device"]) is False
+
+
+# ===========================================================================
+# E. select_visible_gpu_rows
+# ===========================================================================
+
+
+class TestSelectVisibleGpuRows:
+    ROWS = [
+        ("0", "GPU-aaa", "8.6"),
+        ("1", "GPU-bbb", "7.5"),
+        ("2", "GPU-ccc", "8.9"),
+    ]
+
+    def test_none_returns_all(self):
+        assert select_visible_gpu_rows(self.ROWS, None) == list(self.ROWS)
+
+    def test_empty_returns_empty(self):
+        assert select_visible_gpu_rows(self.ROWS, []) == []
+
+    def test_filter_by_index(self):
+        result = select_visible_gpu_rows(self.ROWS, ["0", "2"])
+        assert result == [("0", "GPU-aaa", "8.6"), ("2", "GPU-ccc", "8.9")]
+
+    def test_filter_by_uuid_case_insensitive(self):
+        result = select_visible_gpu_rows(self.ROWS, ["gpu-bbb"])
+        assert result == [("1", "GPU-bbb", "7.5")]
+
+    def test_dedup_same_device(self):
+        result = select_visible_gpu_rows(self.ROWS, ["0", "0"])
+        assert result == [("0", "GPU-aaa", "8.6")]
+
+    def test_missing_token(self):
+        result = select_visible_gpu_rows(self.ROWS, ["99"])
+        assert result == []
+
+
+# ===========================================================================
+# F. compatible_linux_runtime_lines
+# ===========================================================================
+
+
+class TestCompatibleLinuxRuntimeLines:
+    def test_no_driver(self):
+        host = make_host(driver_cuda_version = None)
+        assert compatible_linux_runtime_lines(host) == []
+
+    def test_driver_11_8(self):
+        host = make_host(driver_cuda_version = (11, 8))
+        assert compatible_linux_runtime_lines(host) == []
+
+    def test_driver_12_4(self):
+        host = make_host(driver_cuda_version = (12, 4))
+        assert compatible_linux_runtime_lines(host) == ["cuda12"]
+
+    def test_driver_13_0(self):
+        host = make_host(driver_cuda_version = (13, 0))
+        assert compatible_linux_runtime_lines(host) == ["cuda13", "cuda12"]
+
+
+# ===========================================================================
+# G. pick_windows_cuda_runtime + compatible_windows_runtime_lines
+# ===========================================================================
+
+
+class TestPickWindowsCudaRuntime:
+    def test_no_driver(self):
+        host = make_host(driver_cuda_version = None)
+        assert pick_windows_cuda_runtime(host) is None
+
+    def test_below_threshold(self):
+        host = make_host(driver_cuda_version = (12, 3))
+        assert pick_windows_cuda_runtime(host) is None
+
+    def test_driver_12_4(self):
+        host = make_host(driver_cuda_version = (12, 4))
+        assert pick_windows_cuda_runtime(host) == "12.4"
+
+    def test_driver_13_1(self):
+        host = make_host(driver_cuda_version = (13, 1))
+        assert pick_windows_cuda_runtime(host) == "13.1"
+
+
+class TestCompatibleWindowsRuntimeLines:
+    def test_no_driver(self):
+        host = make_host(driver_cuda_version = None)
+        assert compatible_windows_runtime_lines(host) == []
+
+    def test_driver_12_4(self):
+        host = make_host(driver_cuda_version = (12, 4))
+        assert compatible_windows_runtime_lines(host) == ["cuda12"]
+
+    def test_driver_13_1(self):
+        host = make_host(driver_cuda_version = (13, 1))
+        assert compatible_windows_runtime_lines(host) == ["cuda13", "cuda12"]
+
+
+# ===========================================================================
+# H. runtime_line_from_cuda_version
+# ===========================================================================
+
+
+class TestRuntimeLineFromCudaVersion:
+    def test_cuda_12(self):
+        assert runtime_line_from_cuda_version("12.6") == "cuda12"
+
+    def test_cuda_13(self):
+        assert runtime_line_from_cuda_version("13.0") == "cuda13"
+
+    def test_cuda_11(self):
+        assert runtime_line_from_cuda_version("11.8") is None
+
+    def test_none(self):
+        assert runtime_line_from_cuda_version(None) is None
+
+    def test_empty(self):
+        assert runtime_line_from_cuda_version("") is None
+
+
+# ===========================================================================
+# I. apply_approved_hashes
+# ===========================================================================
+
+
+class TestApplyApprovedHashes:
+    def _choice(self, name):
+        return AssetChoice(
+            repo = "test",
+            tag = "v1",
+            name = name,
+            url = f"https://x/{name}",
+            source_label = "test",
+        )
+
+    def test_both_approved(self):
+        c1, c2 = self._choice("a.tar.gz"), self._choice("b.tar.gz")
+        checksums = make_checksums(["a.tar.gz", "b.tar.gz"])
+        result = apply_approved_hashes([c1, c2], checksums)
+        assert len(result) == 2
+        assert all(c.expected_sha256 == "a" * 64 for c in result)
+
+    def test_one_approved(self):
+        c1, c2 = self._choice("a.tar.gz"), self._choice("missing.tar.gz")
+        checksums = make_checksums(["a.tar.gz"])
+        result = apply_approved_hashes([c1, c2], checksums)
+        assert len(result) == 1
+        assert result[0].name == "a.tar.gz"
+
+    def test_none_approved(self):
+        c1 = self._choice("missing.tar.gz")
+        checksums = make_checksums(["other.tar.gz"])
+        with pytest.raises(PrebuiltFallback, match = "approved checksum"):
+            apply_approved_hashes([c1], checksums)
+
+    def test_empty_input(self):
+        checksums = make_checksums(["a.tar.gz"])
+        with pytest.raises(PrebuiltFallback, match = "approved checksum"):
+            apply_approved_hashes([], checksums)
+
+
+# ===========================================================================
+# J. linux_cuda_choice_from_release -- core selection
+# ===========================================================================
+
+
+class TestLinuxCudaChoiceFromRelease:
+    # --- Runtime line resolution ---
+
+    def test_no_runtime_lines_detected(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, [])
+        host = make_host(driver_cuda_version = (12, 8))
+        art = make_artifact("bundle-cuda12.tar.gz")
+        release = make_release([art])
+        assert linux_cuda_choice_from_release(host, release) is None
+
+    def test_detected_lines_incompatible_with_driver(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda13"])
+        host = make_host(driver_cuda_version = (12, 4))
+        art = make_artifact("bundle-cuda13.tar.gz", runtime_line = "cuda13")
+        release = make_release([art])
+        assert linux_cuda_choice_from_release(host, release) is None
+
+    def test_driver_13_only_cuda12_detected(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(driver_cuda_version = (13, 0))
+        art = make_artifact("bundle-cuda12.tar.gz", runtime_line = "cuda12")
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is not None
+        assert result.primary.runtime_line == "cuda12"
+
+    def test_preferred_runtime_line_reorders(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
+        host = make_host(driver_cuda_version = (13, 0))
+        art12 = make_artifact("bundle-cuda12.tar.gz", runtime_line = "cuda12")
+        art13 = make_artifact("bundle-cuda13.tar.gz", runtime_line = "cuda13")
+        release = make_release([art12, art13])
+        result = linux_cuda_choice_from_release(
+            host, release, preferred_runtime_line = "cuda12"
+        )
+        assert result is not None
+        assert result.primary.runtime_line == "cuda12"
+
+    def test_preferred_runtime_line_unavailable(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(driver_cuda_version = (12, 8))
+        art = make_artifact("bundle-cuda12.tar.gz", runtime_line = "cuda12")
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(
+            host, release, preferred_runtime_line = "cuda13"
+        )
+        assert result is not None
+        assert result.primary.runtime_line == "cuda12"
+        log_entries = result.selection_log
+        assert any("unavailable_on_host" in entry for entry in log_entries)
+
+    # --- SM matching ---
+
+    def test_exact_sm_match(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        art = make_artifact(
+            "bundle.tar.gz", supported_sms = ["75", "86", "89"], min_sm = 75, max_sm = 89
+        )
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is not None
+        assert result.primary.name == "bundle.tar.gz"
+
+    def test_sm_not_in_supported_sms(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        art = make_artifact(
+            "bundle.tar.gz", supported_sms = ["75", "80", "89"], min_sm = 75, max_sm = 89
+        )
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    def test_sm_outside_min_range(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["50"])
+        art = make_artifact(
+            "bundle.tar.gz", supported_sms = ["50", "75", "86"], min_sm = 75, max_sm = 90
+        )
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    def test_sm_outside_max_range(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["100"])
+        art = make_artifact(
+            "bundle.tar.gz", supported_sms = ["100", "75", "86"], min_sm = 75, max_sm = 90
+        )
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    def test_very_old_sm(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["50"])
+        art = make_artifact("bundle.tar.gz", min_sm = 75, max_sm = 90)
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    def test_very_new_sm(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["100"])
+        art = make_artifact("bundle.tar.gz", min_sm = 75, max_sm = 90)
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    # --- Unknown compute caps (empty list) ---
+
+    def test_unknown_caps_only_portable(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = [])
+        targeted = make_artifact("targeted.tar.gz", coverage_class = "targeted")
+        portable = make_artifact("portable.tar.gz", coverage_class = "portable")
+        release = make_release([targeted, portable])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is not None
+        assert result.primary.name == "portable.tar.gz"
+
+    def test_unknown_caps_no_portable(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = [])
+        targeted = make_artifact("targeted.tar.gz", coverage_class = "targeted")
+        release = make_release([targeted])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    # --- Multi-GPU ---
+
+    def test_multi_gpu_all_covered(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["75", "89"])
+        art = make_artifact(
+            "bundle.tar.gz",
+            supported_sms = ["75", "80", "86", "89", "90"],
+            min_sm = 75,
+            max_sm = 90,
+        )
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is not None
+
+    def test_multi_gpu_not_all_covered(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["50", "89"])
+        art = make_artifact(
+            "bundle.tar.gz", supported_sms = ["75", "89"], min_sm = 75, max_sm = 89
+        )
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    # --- Artifact selection priority ---
+
+    def test_narrowest_sm_range_wins(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        wide = make_artifact(
+            "wide.tar.gz",
+            supported_sms = ["75", "86", "90"],
+            min_sm = 75,
+            max_sm = 90,
+            rank = 100,
+        )
+        narrow = make_artifact(
+            "narrow.tar.gz",
+            supported_sms = ["80", "86", "89"],
+            min_sm = 80,
+            max_sm = 89,
+            rank = 100,
+        )
+        release = make_release([wide, narrow])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is not None
+        assert result.primary.name == "narrow.tar.gz"
+
+    def test_range_tie_lower_rank_wins(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        high = make_artifact(
+            "high.tar.gz",
+            supported_sms = ["75", "86", "90"],
+            min_sm = 75,
+            max_sm = 90,
+            rank = 200,
+        )
+        low = make_artifact(
+            "low.tar.gz",
+            supported_sms = ["75", "86", "90"],
+            min_sm = 75,
+            max_sm = 90,
+            rank = 50,
+        )
+        release = make_release([high, low])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is not None
+        assert result.primary.name == "low.tar.gz"
+
+    def test_targeted_preferred_portable_fallback(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        targeted = make_artifact("targeted.tar.gz", coverage_class = "targeted", rank = 100)
+        portable = make_artifact("portable.tar.gz", coverage_class = "portable", rank = 100)
+        release = make_release([targeted, portable])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is not None
+        assert result.primary.name == "targeted.tar.gz"
+        assert len(result.attempts) == 2
+        assert result.attempts[1].name == "portable.tar.gz"
+
+    # --- Edge cases ---
+
+    def test_asset_missing_from_release_assets(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        art = make_artifact("bundle.tar.gz")
+        release = make_release([art], assets = {})
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    def test_artifact_empty_supported_sms(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        art = make_artifact("bundle.tar.gz", supported_sms = [])
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    def test_artifact_missing_min_sm(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        art = make_artifact("bundle.tar.gz", min_sm = None, max_sm = 90)
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    def test_artifact_missing_max_sm(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        art = make_artifact("bundle.tar.gz", min_sm = 75, max_sm = None)
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    def test_no_linux_cuda_artifacts(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        art = make_artifact("bundle.tar.gz", install_kind = "windows-cuda")
+        release = make_release([art])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+    def test_empty_artifacts_list(self, monkeypatch):
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(compute_caps = ["86"])
+        release = make_release([])
+        result = linux_cuda_choice_from_release(host, release)
+        assert result is None
+
+
+# ===========================================================================
+# K. windows_cuda_attempts
+# ===========================================================================
+
+
+class TestWindowsCudaAttempts:
+    TAG = "b8508"
+
+    def _upstream(self, *runtime_versions):
+        assets = {}
+        for rv in runtime_versions:
+            name = f"llama-{self.TAG}-bin-win-cuda-{rv}-x64.zip"
+            assets[name] = f"https://example.com/{name}"
+        return assets
+
+    def test_driver_12_4_no_dlls_fallback(self, monkeypatch):
+        mock_windows_runtime(monkeypatch, [])
+        host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (12, 4))
+        assets = self._upstream("12.4")
+        result = windows_cuda_attempts(host, self.TAG, assets, None)
+        assert len(result) == 1
+        assert result[0].runtime_line == "cuda12"
+
+    def test_driver_13_1_both_dlls(self, monkeypatch):
+        mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
+        host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1))
+        assets = self._upstream("13.1", "12.4")
+        result = windows_cuda_attempts(host, self.TAG, assets, None)
+        assert len(result) == 2
+        assert result[0].runtime_line == "cuda13"
+        assert result[1].runtime_line == "cuda12"
+
+    def test_preferred_reorders(self, monkeypatch):
+        mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
+        host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1))
+        assets = self._upstream("13.1", "12.4")
+        result = windows_cuda_attempts(host, self.TAG, assets, "cuda12")
+        assert len(result) == 2
+        assert result[0].runtime_line == "cuda12"
+
+    def test_preferred_unavailable(self, monkeypatch):
+        mock_windows_runtime(monkeypatch, ["cuda12"])
+        host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (12, 4))
+        assets = self._upstream("12.4")
+        result = windows_cuda_attempts(host, self.TAG, assets, "cuda13")
+        assert len(result) == 1
+        assert result[0].runtime_line == "cuda12"
+
+    def test_detected_incompatible_with_driver(self, monkeypatch):
+        mock_windows_runtime(monkeypatch, ["cuda13"])
+        host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (12, 4))
+        assets = self._upstream("12.4")
+        result = windows_cuda_attempts(host, self.TAG, assets, None)
+        assert len(result) == 1
+        assert result[0].runtime_line == "cuda12"
+
+    def test_driver_too_old(self, monkeypatch):
+        mock_windows_runtime(monkeypatch, [])
+        host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (11, 8))
+        assets = self._upstream("12.4")
+        result = windows_cuda_attempts(host, self.TAG, assets, None)
+        assert result == []
+
+    def test_asset_missing_from_upstream(self, monkeypatch):
+        mock_windows_runtime(monkeypatch, ["cuda12"])
+        host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (12, 4))
+        result = windows_cuda_attempts(host, self.TAG, {}, None)
+        assert result == []
+
+    def test_both_assets_present(self, monkeypatch):
+        mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
+        host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1))
+        assets = self._upstream("13.1", "12.4")
+        result = windows_cuda_attempts(host, self.TAG, assets, None)
+        assert len(result) == 2
+
+
+# ===========================================================================
+# L. resolve_upstream_asset_choice -- platform routing
+# ===========================================================================
+
+
+class TestResolveUpstreamAssetChoice:
+    TAG = "b8508"
+
+    def _mock_github_assets(self, monkeypatch, assets):
+        monkeypatch.setattr(
+            INSTALL_LLAMA_PREBUILT,
+            "github_release_assets",
+            lambda repo, tag: assets,
+        )
+
+    def test_linux_x86_64_cpu(self, monkeypatch):
+        name = f"llama-{self.TAG}-bin-ubuntu-x64.tar.gz"
+        self._mock_github_assets(monkeypatch, {name: f"https://x/{name}"})
+        host = make_host(
+            has_usable_nvidia = False, nvidia_smi = None, has_physical_nvidia = False
+        )
+        result = resolve_upstream_asset_choice(host, self.TAG)
+        assert result.install_kind == "linux-cpu"
+        assert result.name == name
+
+    def test_linux_cpu_missing(self, monkeypatch):
+        self._mock_github_assets(monkeypatch, {})
+        host = make_host(
+            has_usable_nvidia = False, nvidia_smi = None, has_physical_nvidia = False
+        )
+        with pytest.raises(PrebuiltFallback, match = "Linux CPU"):
+            resolve_upstream_asset_choice(host, self.TAG)
+
+    def test_windows_x86_64_cpu(self, monkeypatch):
+        name = f"llama-{self.TAG}-bin-win-cpu-x64.zip"
+        self._mock_github_assets(monkeypatch, {name: f"https://x/{name}"})
+        host = make_host(
+            system = "Windows",
+            machine = "AMD64",
+            has_usable_nvidia = False,
+            nvidia_smi = None,
+            has_physical_nvidia = False,
+        )
+        result = resolve_upstream_asset_choice(host, self.TAG)
+        assert result.install_kind == "windows-cpu"
+        assert result.name == name
+
+    def test_windows_cpu_missing(self, monkeypatch):
+        self._mock_github_assets(monkeypatch, {})
+        host = make_host(
+            system = "Windows",
+            machine = "AMD64",
+            has_usable_nvidia = False,
+            nvidia_smi = None,
+            has_physical_nvidia = False,
+        )
+        with pytest.raises(PrebuiltFallback, match = "Windows CPU"):
+            resolve_upstream_asset_choice(host, self.TAG)
+
+    def test_macos_arm64(self, monkeypatch):
+        name = f"llama-{self.TAG}-bin-macos-arm64.tar.gz"
+        self._mock_github_assets(monkeypatch, {name: f"https://x/{name}"})
+        host = make_host(
+            system = "Darwin",
+            machine = "arm64",
+            nvidia_smi = None,
+            driver_cuda_version = None,
+            compute_caps = [],
+            has_physical_nvidia = False,
+            has_usable_nvidia = False,
+        )
+        result = resolve_upstream_asset_choice(host, self.TAG)
+        assert result.install_kind == "macos-arm64"
+        assert result.name == name
+
+    def test_macos_arm64_missing(self, monkeypatch):
+        self._mock_github_assets(monkeypatch, {})
+        host = make_host(
+            system = "Darwin",
+            machine = "arm64",
+            nvidia_smi = None,
+            driver_cuda_version = None,
+            compute_caps = [],
+            has_physical_nvidia = False,
+            has_usable_nvidia = False,
+        )
+        with pytest.raises(PrebuiltFallback, match = "macOS arm64"):
+            resolve_upstream_asset_choice(host, self.TAG)
+
+    def test_macos_x86_64(self, monkeypatch):
+        name = f"llama-{self.TAG}-bin-macos-x64.tar.gz"
+        self._mock_github_assets(monkeypatch, {name: f"https://x/{name}"})
+        host = make_host(
+            system = "Darwin",
+            machine = "x86_64",
+            nvidia_smi = None,
+            driver_cuda_version = None,
+            compute_caps = [],
+            has_physical_nvidia = False,
+            has_usable_nvidia = False,
+        )
+        result = resolve_upstream_asset_choice(host, self.TAG)
+        assert result.install_kind == "macos-x64"
+        assert result.name == name
+
+    def test_linux_aarch64(self, monkeypatch):
+        self._mock_github_assets(monkeypatch, {})
+        host = make_host(
+            system = "Linux",
+            machine = "aarch64",
+            nvidia_smi = None,
+            driver_cuda_version = None,
+            compute_caps = [],
+            has_physical_nvidia = False,
+            has_usable_nvidia = False,
+        )
+        with pytest.raises(
+            PrebuiltFallback, match = "no prebuilt policy exists for Linux aarch64"
+        ):
+            resolve_upstream_asset_choice(host, self.TAG)
+
+    def test_windows_usable_nvidia_delegates(self, monkeypatch):
+        cuda_name = f"llama-{self.TAG}-bin-win-cuda-12.4-x64.zip"
+        self._mock_github_assets(monkeypatch, {cuda_name: f"https://x/{cuda_name}"})
+        mock_windows_runtime(monkeypatch, ["cuda12"])
+        monkeypatch.setattr(
+            INSTALL_LLAMA_PREBUILT,
+            "resolve_windows_cuda_choices",
+            lambda host, tag, assets: [
+                AssetChoice(
+                    repo = UPSTREAM_REPO,
+                    tag = tag,
+                    name = cuda_name,
+                    url = f"https://x/{cuda_name}",
+                    source_label = "upstream",
+                    install_kind = "windows-cuda",
+                    runtime_line = "cuda12",
+                )
+            ],
+        )
+        host = make_host(
+            system = "Windows",
+            machine = "AMD64",
+            driver_cuda_version = (12, 4),
+            has_usable_nvidia = True,
+        )
+        result = resolve_upstream_asset_choice(host, self.TAG)
+        assert result.install_kind == "windows-cuda"
+        assert result.name == cuda_name

tests/utils/test_q_galore.py

@@ -0,0 +1,528 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Tests for Q-GaLore integration (unsloth/optimizers/).
+
+import pytest
+import sys
+import os
+import torch
+import torch.nn as nn
+
+# Import the optimizers module directly to avoid triggering unsloth.__init__
+# which requires unsloth_zoo and other heavy dependencies.
+_repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+_optimizers_dir = os.path.join(_repo_root, "unsloth", "optimizers")
+if _repo_root not in sys.path:
+    sys.path.insert(0, _repo_root)
+
+# Direct import of the actual modules (avoids unsloth/__init__.py)
+import importlib.util
+
+
+def _load_module(name, filepath):
+    spec = importlib.util.spec_from_file_location(name, filepath)
+    mod = importlib.util.module_from_spec(spec)
+    sys.modules[name] = mod
+    spec.loader.exec_module(mod)
+    return mod
+
+
+# Load projector module first (no dependencies on unsloth)
+_projector_mod = _load_module(
+    "unsloth.optimizers.q_galore_projector",
+    os.path.join(_optimizers_dir, "q_galore_projector.py"),
+)
+GaLoreProjector = _projector_mod.GaLoreProjector
+_quantize = _projector_mod._quantize
+_dequantize = _projector_mod._dequantize
+_quantize_stochastic = _projector_mod._quantize_stochastic
+
+# Load adamw module (depends on projector, may skip bitsandbytes)
+_adamw_mod = _load_module(
+    "unsloth.optimizers.q_galore_adamw",
+    os.path.join(_optimizers_dir, "q_galore_adamw.py"),
+)
+make_q_galore_param_groups = _adamw_mod.make_q_galore_param_groups
+
+# ======================================================================
+# Projector tests
+# ======================================================================
+
+
+class TestGaLoreProjector:
+    """Tests for the GaLore low-rank gradient projector."""
+
+    def test_project_and_back_tall(self):
+        """Project → project_back preserves shape for tall matrices."""
+        proj = GaLoreProjector(rank = 4, update_proj_gap = 1)
+        grad = torch.randn(16, 8)  # tall
+        low = proj.project(grad, step = 0)
+        assert low.shape == (16, 4)
+
+        full = proj.project_back(low)
+        assert full.shape == grad.shape
+
+    def test_project_and_back_wide(self):
+        """Project → project_back preserves shape for wide matrices."""
+        proj = GaLoreProjector(rank = 4, update_proj_gap = 1)
+        grad = torch.randn(8, 16)  # wide
+        low = proj.project(grad, step = 0)
+        assert low.shape == (4, 16)
+
+        full = proj.project_back(low)
+        assert full.shape == grad.shape
+
+    def test_project_reuses_cached_svd(self):
+        """SVD is not recomputed when step is not a multiple of update_proj_gap."""
+        proj = GaLoreProjector(rank = 4, update_proj_gap = 100)
+        grad = torch.randn(16, 8)
+        proj.project(grad, step = 0)
+        assert proj.svd_count == 1
+
+        proj.project(grad, step = 1)
+        assert proj.svd_count == 1  # No recomputation
+
+        proj.project(grad, step = 100)
+        assert proj.svd_count == 2  # Recomputed
+
+    def test_quantized_projection(self):
+        """Quantized projection matrix stores and restores with bounded error."""
+        proj = GaLoreProjector(rank = 4, update_proj_gap = 1, quant = True, n_bit = 8)
+        grad = torch.randn(16, 8)
+        low = proj.project(grad, step = 0)
+        assert low.shape == (16, 4)
+
+        # The projection matrix should be stored as uint8
+        assert proj.ortho_matrix.dtype == torch.uint8
+
+    def test_quantized_projection_int4(self):
+        """INT4 quantized projection stores correctly."""
+        proj = GaLoreProjector(rank = 4, update_proj_gap = 1, quant = True, n_bit = 4)
+        grad = torch.randn(16, 8)
+        proj.project(grad, step = 0)
+        assert proj.ortho_matrix.dtype == torch.uint8
+        # INT4 values should be in range [0, 15]
+        assert proj.ortho_matrix.max() <= 15
+
+    def test_adaptive_scheduling(self):
+        """update_proj_gap increases when cosine similarity exceeds threshold."""
+        proj = GaLoreProjector(
+            rank = 4,
+            update_proj_gap = 10,
+            cos_threshold = 0.0,  # Very low threshold → always triggers
+            gamma_proj = 2.0,
+            queue_size = 2,
+        )
+        # Use very similar gradients so cosine similarity is high
+        base_grad = torch.randn(16, 8)
+        for i in range(5):
+            grad = base_grad + torch.randn_like(base_grad) * 0.001
+            proj.project(grad, step = i * 10)
+
+        # After several similar SVDs, update_proj_gap should have increased
+        assert proj.update_proj_gap > 10
+
+    def test_scale_applied(self):
+        """project_back applies the scale factor."""
+        proj = GaLoreProjector(rank = 4, update_proj_gap = 1, scale = 0.5)
+        grad = torch.randn(16, 8)
+        low = proj.project(grad, step = 0)
+
+        proj2 = GaLoreProjector(rank = 4, update_proj_gap = 1, scale = 1.0)
+        low2 = proj2.project(grad, step = 0)
+
+        full_half = proj.project_back(low)
+        full_one = proj2.project_back(low2)
+
+        # The ratio should be exactly 0.5 (SVD is deterministic on same input)
+        ratio = full_half.norm() / full_one.norm()
+        assert abs(ratio - 0.5) < 1e-5, f"Expected ratio ~0.5, got {ratio:.8f}"
+
+
+# ======================================================================
+# Quantization utility tests
+# ======================================================================
+
+
+class TestQuantizationUtils:
+    """Tests for _quantize, _dequantize, _quantize_stochastic."""
+
+    def test_quantize_dequantize_roundtrip(self):
+        """Quantize → dequantize has bounded error."""
+        w = torch.randn(32, 64)
+        q, scales, zeros, shape = _quantize(w, n_bit = 8)
+        w_hat = _dequantize(q, scales, zeros, shape)
+
+        # Error should be bounded by the quantization step size
+        error = (w - w_hat).abs().max()
+        assert error < 0.1, f"Max error {error} exceeds threshold"
+
+    def test_quantize_group_roundtrip(self):
+        """Grouped quantization → dequantization has bounded error."""
+        w = torch.randn(32, 64)
+        q, scales, zeros, shape = _quantize(w, q_group_size = 32, n_bit = 8)
+        w_hat = _dequantize(q, scales, zeros, shape)
+        error = (w - w_hat).abs().max()
+        assert error < 0.1
+
+    def test_quantize_dtype(self):
+        """Quantized output should be uint8."""
+        w = torch.randn(16, 16)
+        q, _, _, _ = _quantize(w, n_bit = 8)
+        assert q.dtype == torch.uint8
+
+    def test_quantize_int4_range(self):
+        """INT4 values should be in [0, 15]."""
+        w = torch.randn(16, 16)
+        q, _, _, _ = _quantize(w, n_bit = 4)
+        assert q.max() <= 15
+        assert q.min() >= 0
+
+    def test_stochastic_rounding_unbiased(self):
+        """Stochastic rounding should be approximately unbiased."""
+        torch.manual_seed(42)
+        w = torch.randn(64, 64)
+        errors = []
+        for _ in range(50):
+            q, scales, zeros, shape = _quantize_stochastic(w, n_bit = 8)
+            w_hat = _dequantize(q, scales, zeros, shape)
+            errors.append((w - w_hat).mean().item())
+
+        mean_error = sum(errors) / len(errors)
+        assert (
+            abs(mean_error) < 0.01
+        ), f"Mean error {mean_error} suggests biased rounding"
+
+
+# ======================================================================
+# Param group helper tests
+# ======================================================================
+
+
+class TestParamGroupHelper:
+    """Tests for make_q_galore_param_groups."""
+
+    def test_param_group_separation(self):
+        """GaLore vs non-GaLore params are correctly separated."""
+
+        # Create a mini-transformer-like model
+        model = nn.Module()
+        model.q_proj = nn.Linear(64, 64, bias = False)
+        model.k_proj = nn.Linear(64, 64, bias = False)
+        model.embed = nn.Embedding(100, 64)
+        model.norm = nn.LayerNorm(64)
+
+        groups = make_q_galore_param_groups(model, rank = 8, weight_quant = False)
+
+        # Should have 2 groups: galore and non-galore
+        assert len(groups) == 2
+
+        galore_group = [g for g in groups if "rank" in g][0]
+        non_galore_group = [g for g in groups if "rank" not in g][0]
+
+        # q_proj and k_proj should be in galore group (2 params)
+        assert len(galore_group["params"]) == 2
+        # embed and norm should be in non-galore group
+        assert (
+            len(non_galore_group["params"]) == 3
+        )  # embed weight + norm weight + norm bias
+
+    def test_custom_target_modules(self):
+        """Custom target_modules narrows GaLore scope."""
+
+        model = nn.Module()
+        model.q_proj = nn.Linear(64, 64, bias = False)
+        model.k_proj = nn.Linear(64, 64, bias = False)
+        model.v_proj = nn.Linear(64, 64, bias = False)
+        model.embed = nn.Embedding(100, 64)
+
+        groups = make_q_galore_param_groups(
+            model,
+            rank = 8,
+            target_modules = ["q_proj"],
+            weight_quant = False,
+        )
+
+        galore_group = [g for g in groups if "rank" in g][0]
+        assert len(galore_group["params"]) == 1  # Only q_proj
+
+    def test_bias_excluded_from_galore(self):
+        """1D bias params matching target names must NOT be in the GaLore group.
+
+        GaLoreProjector.project requires 2-D gradients, so bias vectors
+        (e.g. q_proj.bias) that match a target name must be excluded.
+        """
+        model = nn.Module()
+        model.q_proj = nn.Linear(64, 64, bias = True)  # has .weight AND .bias
+        model.embed = nn.Embedding(100, 64)
+
+        groups = make_q_galore_param_groups(model, rank = 8, weight_quant = False)
+
+        galore_group = [g for g in groups if "rank" in g][0]
+        non_galore_group = [g for g in groups if "rank" not in g][0]
+
+        # Only the 2-D q_proj.weight should be in the GaLore group
+        assert len(galore_group["params"]) == 1
+        assert galore_group["params"][0].dim() == 2
+
+        # q_proj.bias (1-D) + embed.weight should be in non-GaLore
+        assert any(p.dim() == 1 for p in non_galore_group["params"])
+
+    def test_empty_target_modules_no_galore(self):
+        """target_modules=[] should result in no GaLore params."""
+        model = nn.Module()
+        model.q_proj = nn.Linear(64, 64, bias = False)
+
+        # Pass empty list, should NOT fall back to defaults
+        groups = make_q_galore_param_groups(
+            model,
+            rank = 8,
+            target_modules = [],
+            weight_quant = False,
+        )
+
+        galore_groups = [g for g in groups if "rank" in g]
+        assert (
+            len(galore_groups) == 0
+        ), "Expected no GaLore groups when target_modules=[]"
+
+
+# ======================================================================
+# Optimizer tests (CPU-only, no bitsandbytes dependency)
+# ======================================================================
+
+
+class TestQGaLoreIntegration:
+    """Integration tests that work without bitsandbytes on CPU."""
+
+    def test_projector_training_loop(self):
+        """A simple training loop using manual GaLore projection converges."""
+        torch.manual_seed(42)
+
+        # Tiny model: single linear layer
+        model = nn.Linear(32, 16, bias = False)
+        target = torch.randn(4, 16)
+        x = torch.randn(4, 32)
+
+        proj = GaLoreProjector(rank = 8, update_proj_gap = 1, scale = 1.0)
+        optimizer = torch.optim.AdamW(model.parameters(), lr = 0.01)
+
+        losses = []
+        for step in range(20):
+            optimizer.zero_grad()
+            out = model(x)
+            loss = nn.functional.mse_loss(out, target)
+            loss.backward()
+            losses.append(loss.item())
+
+            # Manual GaLore projection
+            for p in model.parameters():
+                if p.grad is not None and p.grad.dim() == 2:
+                    low = proj.project(p.grad, step)
+                    p._saved = p.data.clone()
+                    update = torch.zeros_like(low)
+                    update.add_(low)  # Simplified update
+                    full_update = proj.project_back(update)
+                    p.grad.copy_(full_update)
+
+            optimizer.step()
+
+        # Loss should decrease
+        assert (
+            losses[-1] < losses[0]
+        ), f"Loss did not decrease: {losses[0]:.4f} → {losses[-1]:.4f}"
+
+    def test_full_projector_roundtrip_quality(self):
+        """project → project_back captures the dominant gradient directions."""
+        torch.manual_seed(42)
+        # Create a gradient with clear low-rank structure
+        u = torch.randn(32, 4)
+        v = torch.randn(4, 16)
+        grad = u @ v  # rank-4 gradient
+
+        proj = GaLoreProjector(rank = 4, update_proj_gap = 1, scale = 1.0)
+        low = proj.project(grad, step = 0)
+        reconstructed = proj.project_back(low)
+
+        # For a rank-4 gradient with rank-4 projection, reconstruction
+        # should be very close to original
+        relative_error = (grad - reconstructed).norm() / grad.norm()
+        assert (
+            relative_error < 0.05
+        ), f"Reconstruction error too high: {relative_error:.4f}"
+
+    def test_weight_quant_activates_on_first_step(self):
+        """_has_weight_quant returns True even when _q_scales is None (first step)."""
+        _adamw_mod_local = sys.modules["unsloth.optimizers.q_galore_adamw"]
+        QGaLoreAdamW8bit = _adamw_mod_local.QGaLoreAdamW8bit
+
+        p = torch.nn.Parameter(torch.randn(16, 16))
+        # Simulate init_weight_quantization tagging
+        p._q_scales = None
+        p._q_zeros = None
+        p._q_shape = p.data.shape
+
+        group = {"weight_quant": True}
+
+        # _has_weight_quant must return True even on first step (_q_scales=None)
+        assert QGaLoreAdamW8bit._has_weight_quant(p, group) is True
+
+        # Without the tag, it should return False
+        p2 = torch.nn.Parameter(torch.randn(16, 16))
+        assert QGaLoreAdamW8bit._has_weight_quant(p2, group) is False
+
+    def test_embedding_lr_param_group_split(self):
+        """Embedding params can be split into a separate group with custom LR."""
+        # This tests the logic that make_q_galore_param_groups produces groups
+        # that can be further split by the trainer for embedding LR.
+        model = nn.Module()
+        model.q_proj = nn.Linear(64, 64, bias = False)
+        model.embed = nn.Embedding(100, 64)
+
+        groups = make_q_galore_param_groups(model, rank = 8, weight_quant = False)
+
+        # Simulate splitting non-GaLore group for embedding LR
+        embed_lr = 5e-5
+        new_groups = []
+        for group in groups:
+            if "rank" in group:
+                new_groups.append(group)
+                continue
+            embed_params = []
+            other_params = []
+            for p in group["params"]:
+                # In real usage, we'd check the name; here just split by shape
+                if p.shape[0] == 100:  # embedding
+                    embed_params.append(p)
+                else:
+                    other_params.append(p)
+            if other_params:
+                g = dict(group)
+                g["params"] = other_params
+                new_groups.append(g)
+            if embed_params:
+                g = dict(group)
+                g["params"] = embed_params
+                g["lr"] = embed_lr
+                new_groups.append(g)
+
+        # Should have 3 groups: galore, non-galore non-embed, embed
+        embed_groups = [g for g in new_groups if g.get("lr") == embed_lr]
+        assert len(embed_groups) == 1
+        assert embed_groups[0]["lr"] == embed_lr
+
+    def test_optimizer_hyperparams_forwarded(self):
+        """QGaLoreAdamW8bit accepts betas and eps keyword arguments."""
+        # Verify the constructor signature accepts these params.
+        # Without bitsandbytes we can't instantiate, but we can check the
+        # function signature.
+        import inspect
+
+        _adamw_mod_local = sys.modules["unsloth.optimizers.q_galore_adamw"]
+        QGaLoreAdamW8bit = _adamw_mod_local.QGaLoreAdamW8bit
+
+        sig = inspect.signature(QGaLoreAdamW8bit.__init__)
+        param_names = list(sig.parameters.keys())
+        assert "betas" in param_names, "betas not in QGaLoreAdamW8bit.__init__ params"
+        assert "eps" in param_names, "eps not in QGaLoreAdamW8bit.__init__ params"
+
+    def test_weight_decay_uses_saved_data(self):
+        """Weight decay should apply standard decoupled AdamW decay on current weights."""
+        _adamw_mod_local = sys.modules["unsloth.optimizers.q_galore_adamw"]
+
+        # Create a mock parameter and group
+        p = torch.nn.Parameter(torch.ones(4, 4))
+        p._saved_data = torch.ones(4, 4) * 2.0  # Pre-update weights
+        # Simulate project-back: p.data = p._saved_data + projected update
+        p.data = p._saved_data.add_(torch.ones(4, 4) * 1.0)  # p.data is now 3.0
+
+        group = {"weight_decay": 0.1, "lr": 1.0, "_wd_saved": 0.1}
+
+        # Replicate the fixed decoupled weight decay logic (uses p.data, not p._saved_data)
+        p.data.add_(
+            p.data,
+            alpha = -group["lr"] * group["_wd_saved"],
+        )
+
+        del p._saved_data  # Clean up after all uses, matching fixed code
+
+        # Decoupled weight decay: 3.0 - (1.0 * 0.1 * 3.0) = 2.7
+        assert torch.allclose(
+            p.data, torch.tensor(2.7)
+        ), "Weight decay didn't use p.data for decoupled decay!"
+
+    def test_params_float_after_weight_quant_step(self):
+        """After a step with weight_quant=True, parameters must remain floating point."""
+        _adamw_mod_local = sys.modules["unsloth.optimizers.q_galore_adamw"]
+        _projector_mod_local = sys.modules["unsloth.optimizers.q_galore_projector"]
+
+        _quantize = _projector_mod_local._quantize
+
+        p = torch.nn.Parameter(torch.randn(16, 16))
+        group = {
+            "weight_quant": True,
+            "stochastic_round": False,
+            "weight_group_size": 16,
+        }
+
+        # Replicate the re-quantize logic at the end of optimizer step
+        float_data = p.data.clone()
+        q, scales, zeros, shape = _quantize(
+            float_data, q_group_size = group["weight_group_size"]
+        )
+
+        # The key assertion: p.data stays float, _q_data holds uint8
+        p._q_data = q.to(p.data.device)
+        p._q_scales = scales
+        p._q_zeros = zeros
+        p._q_shape = shape
+
+        assert p.data.is_floating_point(), "p.data was converted to uint8!"
+        assert p._q_data.dtype == torch.uint8, "_q_data should be uint8!"
+
+    def test_weight_quant_hook_restores_float(self):
+        """Forward pre-hook should dequantize INT8 weights before forward pass."""
+        _adamw_mod_local = sys.modules["unsloth.optimizers.q_galore_adamw"]
+        _projector_mod_local = sys.modules["unsloth.optimizers.q_galore_projector"]
+        install_hook = _adamw_mod_local.install_weight_quant_hooks
+
+        linear = nn.Linear(16, 8, bias = False)
+        original = linear.weight.data.clone()
+
+        # Quantize the weight and replace with placeholder (simulates post-step)
+        q, scales, zeros, shape = _projector_mod_local._quantize(
+            linear.weight.data.clone(), q_group_size = 16
+        )
+        linear.weight._q_data = q
+        linear.weight._q_scales = scales
+        linear.weight._q_zeros = zeros
+        linear.weight._q_shape = shape
+        linear.weight.data = torch.zeros(1, dtype = linear.weight.dtype)
+        assert linear.weight.data.numel() == 1, "placeholder should be 1 element"
+
+        # Install hook and run forward -- should restore float weights
+        handles = install_hook(linear)
+        x = torch.randn(2, 16)
+        out = linear(x)  # triggers pre-hook
+
+        assert linear.weight.data.shape == (8, 16), "weight shape not restored"
+        assert linear.weight.data.is_floating_point(), "weight not float after hook"
+        # Check values are close to original (quantization introduces small error)
+        assert torch.allclose(
+            linear.weight.data, original, atol = 0.15
+        ), "dequantized weight too far from original"
+
+        for h in handles:
+            h.remove()

unsloth/models/_utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "2026.3.11"
+__version__ = "2026.3.12"
 
 __all__ = [
     "SUPPORTS_BFLOAT16",

unsloth/models/rl.py

@@ -93,6 +93,58 @@ def vLLMSamplingParams(**kwargs):
     return sampling_params
 
 
+def _maybe_prepare_vllm_for_resume(trainer):
+    if not torch.cuda.is_available():
+        return
+
+    llm = getattr(trainer, "llm", None)
+    if llm is None:
+        llm = getattr(getattr(trainer, "model", None), "vllm_engine", None)
+    if llm is None:
+        return
+
+    model_config = getattr(
+        getattr(getattr(llm, "llm_engine", None), "vllm_config", None),
+        "model_config",
+        None,
+    )
+    if not getattr(model_config, "enable_sleep_mode", False):
+        return
+
+    try:
+        llm.sleep(1)
+    except Exception:
+        pass
+
+    import gc
+
+    for _ in range(3):
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+def _patch_resume_from_checkpoint_memory(trainer_class):
+    original_train = getattr(trainer_class, "train", None)
+    if original_train is None:
+        return
+    if getattr(original_train, "_unsloth_resume_guard", False):
+        return
+
+    def _unsloth_train_with_resume_guard(self, *args, **kwargs):
+        resume_from_checkpoint = kwargs.get("resume_from_checkpoint", None)
+        if resume_from_checkpoint is None:
+            resume_from_checkpoint = kwargs.get("model_path", None)
+        if resume_from_checkpoint is None and len(args) != 0:
+            resume_from_checkpoint = args[0]
+
+        if resume_from_checkpoint:
+            _maybe_prepare_vllm_for_resume(self)
+        return original_train(self, *args, **kwargs)
+
+    _unsloth_train_with_resume_guard._unsloth_resume_guard = True
+    trainer_class.train = _unsloth_train_with_resume_guard
+
+
 def PatchRL(FastLanguageModel):
     try:
         from trl.models.utils import unwrap_model_for_generation
@@ -305,7 +357,6 @@ from transformers.training_args import ParallelMode
 from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize
 
 # Wrap trainer with padding to right and enable training mode
-# Also patches W&B since multiple runs must use wandb.finish()
 import functools
 from types import MethodType
 try:
@@ -315,6 +366,23 @@ except:
 def prepare_for_training_mode(f):
     @functools.wraps(f)
     def wrapper(self, *args, **kwargs):
+        # Finish the previous W&B run if this is a subsequent train() call.
+        # We do this at the START of train() (not the end) so that
+        # evaluate() / log() still work after train() completes.
+        # HF's WandbCallback.setup() will call wandb.init() for the new run.
+        # See: https://github.com/unslothai/unsloth/issues/3954
+        if getattr(self, '_unsloth_training_completed', False):
+            try:
+                import wandb
+                if wandb.run is not None:
+                    wandb.finish()
+                    # Reset HF's WandbCallback so it calls wandb.init() for the new run
+                    for cb in self.callback_handler.callbacks:
+                        if type(cb).__name__ == 'WandbCallback':
+                            cb._initialized = False
+                            break
+            except:
+                pass
         # Enable training mode
         _was_training = None
         # Get gradient checkpointing setting from training arguments
@@ -335,12 +403,9 @@ def prepare_for_training_mode(f):
             reset_unsloth_gradient_checkpointing_buffers()
         except:
             pass
-        # Patch W&B to enable logging on future runs, otherwise it'll overwrite the first run
-        try:
-            import wandb
-            wandb.finish()
-        except:
-            pass
+        # Mark that training completed so the next train() call can
+        # finish this W&B run before starting a new one
+        self._unsloth_training_completed = True
         return output
     return wrapper
 pass
@@ -686,8 +751,8 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
             else:
                 continue
             call_args.append(f"{k} = {k}")
-        arguments = f"\n{' '*8}" + f",\n{' '*8}".join(arguments)
-        call_args = f"\n{' '*12}" + f",\n{' '*12}".join(call_args)
+        arguments = f"\n{' ' * 8}" + f",\n{' ' * 8}".join(arguments)
+        call_args = f"\n{' ' * 12}" + f",\n{' ' * 12}".join(call_args)
         processed.append(
             (
                 arguments,
@@ -701,7 +766,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
 
     # Add tokenizer if not seen
     if "tokenizer" not in parameters and "processing_class" in parameters:
-        arguments += f",\n{' '*8}tokenizer = None"
+        arguments += f",\n{' ' * 8}tokenizer = None"
         call_args = call_args.replace(
             "processing_class = processing_class",
             "processing_class = tokenizer if tokenizer is not None else processing_class",
@@ -1490,6 +1555,9 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         imports,
         overwrite = False,
     )
+    patched_trainer = getattr(created_module, f"Unsloth{RLTrainer_name}")
+    if trainer_file == "grpo_trainer":
+        _patch_resume_from_checkpoint_memory(patched_trainer)
 
     # Patch Trainer
     exec(
@@ -1706,8 +1774,8 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
                 sampling_params = re.sub(r"[\,][\s]{0,}\,", ",", sampling_params)
 
                 new_vllm_part = (
-                    f"\n{' '*8}if {args}.use_vllm:\n{sampling_params}"
-                    f"\n{' '*8}else:\n"
+                    f"\n{' ' * 8}if {args}.use_vllm:\n{sampling_params}"
+                    f"\n{' ' * 8}else:\n"
                 )
 
         if trl_version >= Version("0.18.0"):

unsloth/optimizers/__init__.py

@@ -0,0 +1,21 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .q_galore_projector import GaLoreProjector
+from .q_galore_adamw import QGaLoreAdamW8bit
+
+__all__ = [
+    "GaLoreProjector",
+    "QGaLoreAdamW8bit",
+]

unsloth/optimizers/q_galore_adamw.py

@@ -0,0 +1,424 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Adapted from Q-GaLore (https://github.com/VITA-Group/Q-GaLore)
+# Original paper: "Q-GaLore: Quantized GaLore with INT4 Projection and
+# Layer-Adaptive Low-Rank Gradients" (arXiv:2407.08296)
+
+import torch
+from typing import Optional, List
+
+from .q_galore_projector import (
+    GaLoreProjector,
+    _quantize,
+    _quantize_stochastic,
+    _dequantize,
+)
+
+__all__ = ["QGaLoreAdamW8bit", "install_weight_quant_hooks"]
+
+try:
+    import bitsandbytes.functional as bnb_F
+    from bitsandbytes.optim.optimizer import Optimizer2State
+
+    _HAS_BNB = True
+except ImportError:
+    _HAS_BNB = False
+    # Provide a fallback base so the module can at least be imported.
+    Optimizer2State = torch.optim.Optimizer
+
+
+def _require_bnb():
+    if not _HAS_BNB:
+        raise ImportError(
+            "Unsloth: Q-GaLore requires bitsandbytes. "
+            "Install it with: pip install bitsandbytes"
+        )
+
+
+class QGaLoreAdamW8bit(Optimizer2State):
+    """AdamW optimizer with 8-bit states, GaLore low-rank gradient projection,
+    and optional INT8 weight quantization.
+
+    This optimizer combines three memory-saving techniques:
+
+    1. **8-bit optimizer states** (via bitsandbytes) — Adam's first and second
+       moments are stored in 8-bit, reducing optimizer state memory by ~4×.
+
+    2. **GaLore low-rank gradient projection** — gradients are projected into a
+       low-rank subspace before the optimizer step, then projected back.  The
+       projection matrix itself can be quantized to INT4.
+
+    3. **INT8 weight quantization** — model weights are stored in INT8 during
+       training with stochastic rounding, reducing weight memory by ~2× for
+       eligible layers.
+
+    Param group keys consumed by GaLore projection:
+        ``rank``, ``update_proj_gap``, ``scale``, ``proj_type``,
+        ``quant`` (projection quantization), ``quant_group_size``,
+        ``quant_n_bit``, ``cos_threshold``, ``gamma_proj``, ``queue_size``
+
+    Param group keys for weight quantization:
+        ``weight_quant``, ``stochastic_round``, ``weight_group_size``
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-3,
+        betas: tuple = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 1e-2,
+        min_8bit_size: int = 4096,
+        percentile_clipping: int = 100,
+        block_wise: bool = True,
+        is_paged: bool = False,
+    ):
+        _require_bnb()
+        super().__init__(
+            "adam",
+            params,
+            lr,
+            betas,
+            eps,
+            weight_decay,
+            8,  # optim_bits
+            None,  # args
+            min_8bit_size,
+            percentile_clipping,
+            block_wise,
+            is_paged = is_paged,
+        )
+
+    # ------------------------------------------------------------------
+    # Core step
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def step(self, closure = None):
+        """Perform a single optimization step.
+
+        For each parameter that has a ``rank`` key in its param group, the
+        following sequence is executed:
+
+        1. If ``weight_quant`` is set, dequantize the INT8 weight to float.
+        2. Project the gradient to low-rank via the cached ``GaLoreProjector``.
+        3. Perform the 8-bit Adam update in the low-rank space.
+        4. Project the update back to full rank and add to saved weight.
+        5. If ``weight_quant`` is set, re-quantize the weight to INT8.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        if not self.initialized:
+            self.check_overrides()
+            self.to_gpu()
+            self.initialized = True
+
+        for gindex, group in enumerate(self.param_groups):
+            for pindex, p in enumerate(group["params"]):
+                if p.grad is None:
+                    continue
+
+                state = self.state[p]
+                if "step" not in state:
+                    state["step"] = 0
+
+                has_weight_quant = self._has_weight_quant(p, group)
+
+                # --- Dequantize weight if INT8 ---
+                if has_weight_quant:
+                    if p._q_scales is not None:
+                        float_weight = _dequantize(
+                            p._q_data,
+                            p._q_scales,
+                            p._q_zeros,
+                            p._q_shape,
+                        )
+                        p.data = float_weight
+                    # else: first step, weights are still float — skip dequantize
+
+                # --- GaLore projection ---
+                if "rank" in group:
+                    if "projector" not in state:
+                        state["projector"] = GaLoreProjector(
+                            rank = group["rank"],
+                            update_proj_gap = group.get("update_proj_gap", 200),
+                            scale = group.get("scale", 0.25),
+                            proj_type = group.get("proj_type", "std"),
+                            quant = group.get("quant", False),
+                            group_size = group.get("quant_group_size", -1),
+                            n_bit = group.get("quant_n_bit", 4),
+                            cos_threshold = group.get("cos_threshold", 0.4),
+                            gamma_proj = group.get("gamma_proj", 2.0),
+                            queue_size = group.get("queue_size", 5),
+                        )
+
+                    # Temporarily disable weight decay for GaLore params
+                    # (we apply it manually after project-back)
+                    if "weight_decay" in group and group["weight_decay"] > 0:
+                        group["_wd_saved"] = group["weight_decay"]
+                        group["weight_decay"] = 0
+
+                    grad = state["projector"].project(p.grad, state["step"])
+
+                    # Save current weight; replace p.data with zeros so
+                    # the 8-bit update writes the pure weight delta.
+                    p._saved_data = p.data.clone()
+                    p.data = torch.zeros_like(
+                        grad, dtype = p.data.dtype, device = p.data.device
+                    )
+                    p.grad = grad
+
+                # --- 8-bit Adam update ---
+                if "state1" not in state:
+                    self.init_state(group, p, gindex, pindex)
+
+                self.prefetch_state(p)
+                self.update_step(group, p, gindex, pindex)
+
+                # --- GaLore project-back ---
+                if "rank" in group:
+                    # p.data now holds the weight update in low-rank space
+                    p.data = p._saved_data.add_(state["projector"].project_back(p.data))
+
+                    # Re-apply decoupled weight decay using pre-update weights
+                    if "_wd_saved" in group:
+                        p.data.add_(
+                            p.data,
+                            alpha = -group["lr"] * group["_wd_saved"],
+                        )
+                        group["weight_decay"] = group["_wd_saved"]
+                        del group["_wd_saved"]
+
+                    del p._saved_data
+
+                # --- Re-quantize weight to INT8 ---
+                if has_weight_quant:
+                    float_data = p.data
+                    stochastic = group.get("stochastic_round", True)
+                    gsize = group.get("weight_group_size", 128)
+                    quant_fn = _quantize_stochastic if stochastic else _quantize
+                    q, scales, zeros, shape = quant_fn(float_data, q_group_size = gsize)
+                    p._q_data = q.to(p.data.device)
+                    p._q_scales = scales
+                    p._q_zeros = zeros
+                    p._q_shape = shape
+                    # Replace p.data with a scalar placeholder to free float memory.
+                    # A forward pre-hook (install_weight_quant_hooks) will
+                    # dequantize back to float before the next forward pass.
+                    p.data = torch.empty(1, dtype = p.data.dtype, device = p.data.device)
+
+                state["step"] += 1
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+        return loss
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _has_weight_quant(p: torch.Tensor, group: dict) -> bool:
+        """Check if this parameter uses INT8 weight quantization."""
+        return (
+            group.get("weight_quant", False)
+            and hasattr(p, "_q_scales")  # tag set by init_weight_quantization()
+        )
+
+    @staticmethod
+    def init_weight_quantization(
+        model: torch.nn.Module,
+        param_groups: list,
+        group_size: int = 128,
+        stochastic: bool = True,
+    ) -> None:
+        """Tag parameters for INT8 weight quantization.
+
+        This marks eligible weights with quantization metadata so that
+        the optimizer knows to quantize/dequantize them during ``step()``.
+        **Weights are NOT converted to uint8 here** — they remain in float
+        so that the first forward/backward pass runs correctly.  The actual
+        quantization happens at the end of the first ``step()`` call.
+        """
+        weight_quant_params = set()
+        for group in param_groups:
+            if group.get("weight_quant", False):
+                for p in group["params"]:
+                    weight_quant_params.add(id(p))
+
+        for name, p in model.named_parameters():
+            if id(p) in weight_quant_params:
+                # Store quantization metadata WITHOUT converting weights to
+                # uint8.  The first optimizer.step() will quantize after the
+                # update.  We store dummy scales/zeros so _has_weight_quant()
+                # returns True on the first step.
+                p._q_scales = None
+                p._q_zeros = None
+                p._q_shape = p.data.shape
+                p._stochastic_round = stochastic
+                p._weight_group_size = group_size
+
+
+def _weight_quant_pre_hook(module, args):
+    """Forward pre-hook: dequantize INT8 weights to float before forward."""
+    for p in module.parameters(recurse = False):
+        if hasattr(p, "_q_scales") and p._q_scales is not None:
+            float_weight = _dequantize(
+                p._q_data,
+                p._q_scales,
+                p._q_zeros,
+                p._q_shape,
+            )
+            p.data = float_weight.to(p.data.device)
+
+
+def install_weight_quant_hooks(model: torch.nn.Module) -> list:
+    """Register forward pre-hooks on modules whose weights are INT8-quantized.
+
+    Returns a list of hook handles so the caller can remove them if needed.
+    """
+    handles = []
+    for module in model.modules():
+        has_quant_param = any(
+            hasattr(p, "_q_scales") for p in module.parameters(recurse = False)
+        )
+        if has_quant_param:
+            h = module.register_forward_pre_hook(_weight_quant_pre_hook)
+            handles.append(h)
+    return handles
+
+
+# ======================================================================
+# Param-group construction helper
+# ======================================================================
+
+# Default linear layer names in transformer blocks that should use GaLore.
+_DEFAULT_GALORE_TARGETS = {
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "down_proj",
+}
+
+
+def make_q_galore_param_groups(
+    model: torch.nn.Module,
+    lr: float = 1e-3,
+    weight_decay: float = 0.0,
+    rank: int = 256,
+    update_proj_gap: int = 200,
+    scale: float = 0.25,
+    proj_quant: bool = True,
+    proj_quant_group_size: int = -1,
+    proj_quant_n_bit: int = 4,
+    weight_quant: bool = False,
+    stochastic_round: bool = True,
+    weight_group_size: int = 128,
+    cos_threshold: float = 0.4,
+    gamma_proj: float = 2.0,
+    queue_size: int = 5,
+    target_modules: Optional[List[str]] = None,
+) -> list:
+    """Build param groups suitable for :class:`QGaLoreAdamW8bit`.
+
+    Parameters matching ``target_modules`` (or the default set of attention
+    and MLP projection names) are placed in the GaLore group.  All other
+    trainable parameters go into the non-GaLore group.
+
+    Args:
+        model: The model whose parameters to partition.
+        lr: Learning rate for all parameter groups.
+        weight_decay: Weight decay coefficient.
+        rank: GaLore projection rank.
+        update_proj_gap: Steps between SVD recomputations.
+        scale: Scaling factor for project-back.
+        proj_quant: Quantize projection matrices.
+        proj_quant_group_size: Group size for projection quantization.
+        proj_quant_n_bit: Bit-width for projection quantization.
+        weight_quant: Enable INT8 weight quantization for GaLore params.
+        stochastic_round: Use stochastic rounding for weight quantization.
+        weight_group_size: Group size for weight quantization.
+        cos_threshold: Cosine similarity threshold for adaptive scheduling.
+        gamma_proj: Multiplier for update_proj_gap when subspace is stable.
+        queue_size: Rolling window size for stability tracking.
+        target_modules: Module name substrings to match for GaLore.  If None,
+            uses the default set of attention/MLP projection names.
+
+    Returns:
+        List of two param group dicts: ``[galore_group, non_galore_group]``.
+    """
+    targets = (
+        set(target_modules) if target_modules is not None else _DEFAULT_GALORE_TARGETS
+    )
+
+    galore_params = []
+    non_galore_params = []
+
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+
+        # Check if any target module name appears as a component in the param name.
+        # Exclude 1-D parameters (biases, norms) because GaLoreProjector.project
+        # requires 2-D gradients.
+        name_parts = name.split(".")
+        is_galore = param.dim() >= 2 and any(t in name_parts for t in targets)
+
+        if is_galore:
+            galore_params.append(param)
+        else:
+            non_galore_params.append(param)
+
+    groups = []
+
+    if galore_params:
+        groups.append(
+            {
+                "params": galore_params,
+                "lr": lr,
+                "weight_decay": weight_decay,
+                "rank": rank,
+                "update_proj_gap": update_proj_gap,
+                "scale": scale,
+                "proj_type": "std",
+                "quant": proj_quant,
+                "quant_group_size": proj_quant_group_size,
+                "quant_n_bit": proj_quant_n_bit,
+                "weight_quant": weight_quant,
+                "stochastic_round": stochastic_round,
+                "weight_group_size": weight_group_size,
+                "cos_threshold": cos_threshold,
+                "gamma_proj": gamma_proj,
+                "queue_size": queue_size,
+            }
+        )
+
+    if non_galore_params:
+        groups.append(
+            {
+                "params": non_galore_params,
+                "lr": lr,
+                "weight_decay": weight_decay,
+            }
+        )
+
+    return groups

unsloth/optimizers/q_galore_projector.py

@@ -0,0 +1,385 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Adapted from Q-GaLore (https://github.com/VITA-Group/Q-GaLore)
+# Original paper: "Q-GaLore: Quantized GaLore with INT4 Projection and
+# Layer-Adaptive Low-Rank Gradients" (arXiv:2407.08296)
+
+from collections import deque
+
+import torch
+
+__all__ = ["GaLoreProjector"]
+
+
+class GaLoreProjector:
+    """Low-rank gradient projector with optional INT4/INT8 quantized projection
+    matrices and layer-adaptive subspace update scheduling.
+
+    The projector computes an SVD of the gradient to obtain an orthogonal basis
+    for the top-``rank`` subspace.  Gradients are projected into this subspace
+    for the optimizer step, then projected back to full rank for the weight
+    update.
+
+    Two key Q-GaLore innovations are implemented:
+
+    1. **Quantized projection matrices** — when ``quant=True``, the orthogonal
+       matrix is stored in INT4/INT8, reducing the memory cost of keeping the
+       projector state.
+
+    2. **Layer-adaptive update scheduling** — a rolling queue of cosine
+       similarities between consecutive orthogonal vectors is maintained.  When
+       the average exceeds ``cos_threshold``, ``update_proj_gap`` is multiplied
+       by ``gamma_proj``, effectively reducing the frequency of expensive SVD
+       recomputations for layers whose subspace has stabilized.
+
+    Args:
+        rank: Target rank for the low-rank projection.
+        update_proj_gap: Number of steps between SVD recomputations.
+        scale: Scaling factor applied when projecting back to full rank.
+        proj_type: Projection type.  Only ``'std'`` is supported.
+        quant: Whether to quantize the projection matrix.
+        group_size: Group size for projection matrix quantization.
+        n_bit: Bit-width for projection matrix quantization (4 or 8).
+        cos_threshold: Cosine similarity threshold for adaptive scheduling.
+        gamma_proj: Multiplier for ``update_proj_gap`` on stability detection.
+        queue_size: Number of recent cosine similarities to average.
+    """
+
+    __slots__ = (
+        "rank",
+        "update_proj_gap",
+        "scale",
+        "proj_type",
+        "quant",
+        "quant_group_size",
+        "quant_n_bit",
+        "cos_threshold",
+        "gamma_proj",
+        "queue_size",
+        "ortho_matrix",
+        "ortho_matrix_scales",
+        "ortho_matrix_zeros",
+        "ortho_matrix_shape",
+        "past_ortho_vector",
+        "queue",
+        "svd_count",
+        "_ortho_float_cache",
+    )
+
+    def __init__(
+        self,
+        rank: int,
+        update_proj_gap: int = 200,
+        scale: float = 1.0,
+        proj_type: str = "std",
+        quant: bool = False,
+        group_size: int = -1,
+        n_bit: int = 4,
+        cos_threshold: float = 0.4,
+        gamma_proj: float = 2.0,
+        queue_size: int = 5,
+    ):
+        self.rank = rank
+        self.update_proj_gap = update_proj_gap
+        self.scale = scale
+        self.proj_type = proj_type
+
+        # Quantization settings for the projection matrix
+        self.quant = quant
+        self.quant_group_size = group_size
+        self.quant_n_bit = n_bit
+
+        # Adaptive update scheduling state
+        self.cos_threshold = cos_threshold
+        self.gamma_proj = gamma_proj
+        self.queue_size = queue_size
+        self.past_ortho_vector = None
+        self.queue = deque(maxlen = queue_size)
+        self.svd_count = 0
+        self._ortho_float_cache = None
+
+        # Projection matrix state
+        self.ortho_matrix = None
+        self.ortho_matrix_scales = None
+        self.ortho_matrix_zeros = None
+        self.ortho_matrix_shape = None
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def project(self, full_rank_grad: torch.Tensor, step: int) -> torch.Tensor:
+        """Project a full-rank gradient into the low-rank subspace.
+
+        The SVD is recomputed every ``update_proj_gap`` steps (subject to
+        adaptive scheduling).  Between recomputations the cached orthogonal
+        matrix is reused.
+
+        Args:
+            full_rank_grad: The full-rank gradient tensor (2-D).
+            step: The current optimizer step (0-indexed).
+
+        Returns:
+            The low-rank gradient tensor.
+        """
+        assert self.proj_type == "std", "Only proj_type='std' is supported."
+
+        if full_rank_grad.shape[0] >= full_rank_grad.shape[1]:
+            # "tall" matrix → right projection  (grad @ Q^T)
+            if self.ortho_matrix is None or step % self.update_proj_gap == 0:
+                float_ortho = self._compute_orthogonal(
+                    full_rank_grad,
+                    self.rank,
+                    side = "right",
+                )
+                self._update_adaptive_schedule(float_ortho, side = "right")
+                self._store_ortho(float_ortho)
+
+            self._ortho_float_cache = self._load_ortho()
+            low_rank_grad = torch.matmul(full_rank_grad, self._ortho_float_cache.t())
+        else:
+            # "wide" matrix → left projection  (Q^T @ grad)
+            if self.ortho_matrix is None or step % self.update_proj_gap == 0:
+                float_ortho = self._compute_orthogonal(
+                    full_rank_grad,
+                    self.rank,
+                    side = "left",
+                )
+                self._update_adaptive_schedule(float_ortho, side = "left")
+                self._store_ortho(float_ortho)
+
+            self._ortho_float_cache = self._load_ortho()
+            low_rank_grad = torch.matmul(self._ortho_float_cache.t(), full_rank_grad)
+
+        return low_rank_grad
+
+    def project_back(self, low_rank_grad: torch.Tensor) -> torch.Tensor:
+        """Project a low-rank update back to full rank.
+
+        Args:
+            low_rank_grad: The low-rank gradient/update tensor.
+
+        Returns:
+            The full-rank update scaled by ``self.scale``.
+        """
+        float_ortho = self._ortho_float_cache
+        self._ortho_float_cache = None
+        if float_ortho is None:
+            float_ortho = self._load_ortho()
+
+        if low_rank_grad.shape[0] >= low_rank_grad.shape[1]:
+            full_rank_grad = torch.matmul(low_rank_grad, float_ortho)
+        else:
+            full_rank_grad = torch.matmul(float_ortho, low_rank_grad)
+
+        return full_rank_grad * self.scale
+
+    # ------------------------------------------------------------------
+    # SVD
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _compute_orthogonal(
+        weights: torch.Tensor,
+        rank: int,
+        side: str,
+    ) -> torch.Tensor:
+        """Compute the top-``rank`` orthogonal matrix via truncated SVD.
+
+        Args:
+            weights: 2-D tensor (typically the gradient).
+            rank: Number of singular vectors to keep.
+            side: ``'left'`` returns U[:, :rank], ``'right'`` returns Vh[:rank, :].
+
+        Returns:
+            Orthogonal matrix of shape ``(rank, N)`` (right) or ``(M, rank)`` (left).
+        """
+        original_dtype = weights.dtype
+        original_device = weights.device
+
+        matrix = weights.float() if original_dtype != torch.float32 else weights
+
+        if side not in ("right", "left"):
+            raise ValueError(f"side must be 'left' or 'right', got '{side}'")
+
+        m, n = matrix.shape
+        if min(m, n) <= rank * 2:
+            U, s, Vh = torch.linalg.svd(matrix, full_matrices = False)
+            result = Vh[:rank, :] if side == "right" else U[:, :rank]
+        else:
+            # Oversampling p=10 per Halko et al. 2009 (arXiv:0909.4061)
+            # recommendation of p=5..10 for large low-rank matrices.
+            q = min(rank + 10, min(m, n))
+            U, s, V = torch.svd_lowrank(matrix, q = q, niter = 2)
+            result = V[:, :rank].t() if side == "right" else U[:, :rank]
+
+        if original_dtype != torch.float32:
+            result = result.to(device = original_device, dtype = original_dtype)
+        return result
+
+    # ------------------------------------------------------------------
+    # Adaptive scheduling
+    # ------------------------------------------------------------------
+
+    def _update_adaptive_schedule(
+        self,
+        float_ortho: torch.Tensor,
+        side: str,
+    ) -> None:
+        """Track subspace stability and increase ``update_proj_gap`` if stable."""
+        self.svd_count += 1
+
+        if side == "right":
+            current_vector = float_ortho[:1, :].flatten()
+        else:
+            current_vector = float_ortho[:, :1].flatten()
+
+        if self.past_ortho_vector is not None:
+            cos_sim = torch.dot(self.past_ortho_vector, current_vector).item()
+
+            self.queue.append(cos_sim)
+
+            if (
+                len(self.queue) == self.queue.maxlen
+                and sum(self.queue) / len(self.queue) >= self.cos_threshold
+            ):
+                self.update_proj_gap = int(self.update_proj_gap * self.gamma_proj)
+
+        self.past_ortho_vector = current_vector.clone()
+
+    # ------------------------------------------------------------------
+    # Quantized projection matrix storage
+    # ------------------------------------------------------------------
+
+    def _store_ortho(self, float_ortho: torch.Tensor) -> None:
+        """Store the orthogonal matrix, optionally quantized."""
+        if self.quant:
+            q, scales, zeros, shape = _quantize(
+                float_ortho,
+                q_group_size = self.quant_group_size,
+                n_bit = self.quant_n_bit,
+            )
+            self.ortho_matrix = q
+            self.ortho_matrix_scales = scales
+            self.ortho_matrix_zeros = zeros
+            self.ortho_matrix_shape = shape
+        else:
+            self.ortho_matrix = float_ortho
+
+    def _load_ortho(self) -> torch.Tensor:
+        """Load the orthogonal matrix, dequantizing if necessary."""
+        if self.quant:
+            return _dequantize(
+                self.ortho_matrix,
+                self.ortho_matrix_scales,
+                self.ortho_matrix_zeros,
+                self.ortho_matrix_shape,
+            )
+        return self.ortho_matrix
+
+
+# ======================================================================
+# Quantization utilities (shared with the optimizer)
+# ======================================================================
+
+
+@torch.no_grad()
+def _quantize(
+    w: torch.Tensor,
+    q_group_size: int = -1,
+    n_bit: int = 8,
+) -> tuple:
+    """Asymmetric min-max quantization to unsigned int.
+
+    Returns:
+        ``(quantized_uint8, scales, zeros, original_shape)``
+    """
+    org_shape = w.shape
+    if q_group_size > 0:
+        assert (
+            w.nelement() % q_group_size == 0
+        ), f"Tensor size {w.nelement()} not divisible by group_size {q_group_size}"
+        w = w.reshape(-1, q_group_size)
+    assert w.dim() == 2
+
+    max_val = w.amax(dim = 1, keepdim = True)
+    min_val = w.amin(dim = 1, keepdim = True)
+    max_int = 2**n_bit - 1
+    min_int = 0
+    scales = (max_val - min_val).clamp(min = 1e-5) / max_int
+    zeros = (-torch.round(min_val / scales)).clamp_(min_int, max_int)
+
+    w = torch.clamp(torch.round(w / scales) + zeros, min_int, max_int)
+    w = w.reshape(org_shape).to(torch.uint8)
+
+    return w, scales, zeros, org_shape
+
+
+@torch.no_grad()
+def _dequantize(
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    zeros: torch.Tensor,
+    original_shape: tuple,
+) -> torch.Tensor:
+    """Dequantize from uint8 back to float."""
+    # Infer group size: scales has shape (n_groups, 1), so n_groups = scales.shape[0]
+    total = w.numel()
+    n_groups = scales.shape[0] if scales.dim() > 1 else scales.numel()
+    group_size = total // n_groups if n_groups > 0 else total
+
+    float_w = w.to(scales.dtype).reshape(-1, group_size)
+    float_w = (float_w - zeros) * scales
+    return float_w.reshape(original_shape)
+
+
+@torch.no_grad()
+def _quantize_stochastic(
+    w: torch.Tensor,
+    q_group_size: int = -1,
+    n_bit: int = 8,
+) -> tuple:
+    """Asymmetric min-max quantization with stochastic rounding.
+
+    Instead of deterministic ``round()``, the rounding direction is chosen
+    probabilistically proportional to the fractional part. This gives an
+    unbiased estimator of the original value in expectation.
+
+    Returns:
+        ``(quantized_uint8, scales, zeros, original_shape)``
+    """
+    org_shape = w.shape
+    if q_group_size > 0:
+        assert w.nelement() % q_group_size == 0
+        w = w.reshape(-1, q_group_size)
+    assert w.dim() == 2
+
+    max_val = w.amax(dim = 1, keepdim = True)
+    min_val = w.amin(dim = 1, keepdim = True)
+    max_int = 2**n_bit - 1
+    min_int = 0
+    scales = (max_val - min_val).clamp(min = 1e-5) / max_int
+    zeros = (-torch.round(min_val / scales)).clamp_(min_int, max_int)
+
+    w_scaled = w / scales
+    up = torch.ceil(w_scaled)
+    down = torch.floor(w_scaled)
+    prob = w_scaled - down
+    rng = torch.rand_like(prob)
+    w = torch.where(rng < prob, up, down)
+    w = torch.clamp(w + zeros, min_int, max_int)
+    w = w.reshape(org_shape).to(torch.uint8)
+
+    return w, scales, zeros, org_shape

unsloth/tokenizer_utils.py

@@ -42,6 +42,7 @@ __all__ = [
     "check_tokenizer",
     "add_new_tokens",
     "fix_sentencepiece_gguf",
+    "get_tokenizer_info",
 ]
 
 
@@ -896,6 +897,55 @@ def check_tokenizer(
     return convert_to_fast_tokenizer(tokenizer)
 
 
+def get_tokenizer_info(tokenizer) -> dict:
+    """Return a concise diagnostic summary of a tokenizer instance.
+
+    Collects key properties into a plain dict suitable for logging, debugging,
+    or displaying in the Unsloth Studio UI. All fields are safe to access —
+    missing attributes fall back to ``None`` rather than raising.
+
+    Example output::
+
+        {
+            "name_or_path": "unsloth/Llama-3.2-1B-Instruct",
+            "tokenizer_class": "PreTrainedTokenizerFast",
+            "is_fast": True,
+            "vocab_size": 128000,
+            "added_tokens_count": 256,
+            "model_max_length": 131072,
+            "padding_side": "right",
+            "bos_token": "<|begin_of_text|>",
+            "eos_token": "<|eot_id|>",
+            "pad_token": "<|finetune_right_pad_id|>",
+            "unk_token": None,
+            "has_chat_template": True,
+            "special_tokens_count": 3,
+        }
+
+    Args:
+        tokenizer: Any HuggingFace ``PreTrainedTokenizer`` or
+                   ``PreTrainedTokenizerFast`` instance.
+
+    Returns:
+        A ``dict`` of tokenizer properties. Safe to serialize to JSON.
+    """
+    return {
+        "name_or_path": getattr(tokenizer, "name_or_path", None),
+        "tokenizer_class": type(tokenizer).__name__,
+        "is_fast": getattr(tokenizer, "is_fast", False),
+        "vocab_size": getattr(tokenizer, "vocab_size", None),
+        "added_tokens_count": len(getattr(tokenizer, "added_tokens_decoder", {})),
+        "model_max_length": getattr(tokenizer, "model_max_length", None),
+        "padding_side": getattr(tokenizer, "padding_side", None),
+        "bos_token": getattr(tokenizer, "bos_token", None),
+        "eos_token": getattr(tokenizer, "eos_token", None),
+        "pad_token": getattr(tokenizer, "pad_token", None),
+        "unk_token": getattr(tokenizer, "unk_token", None),
+        "has_chat_template": getattr(tokenizer, "chat_template", None) is not None,
+        "special_tokens_count": len(getattr(tokenizer, "all_special_tokens", [])),
+    }
+
+
 import inspect
 from inspect import getsource
 import trl

unsloth/trainer.py

@@ -17,7 +17,7 @@ import os
 import psutil
 import warnings
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Optional, List
 from functools import wraps
 
 import trl
@@ -46,6 +46,7 @@ __all__ = [
     "unsloth_train",
     "_patch_trl_trainer",
     "UnslothVisionDataCollator",
+    "QGaloreConfig",
 ]
 
 logger = logging.getLogger(__name__)
@@ -130,8 +131,39 @@ except:
     from transformers import TrainingArguments
 
 
+@dataclass
+class QGaloreConfig:
+    """Configuration for Q-GaLore optimizer integration.
+
+    Pass an instance of this class to ``UnslothTrainingArguments`` (via
+    ``q_galore_config``) to enable Q-GaLore training.
+    """
+
+    rank: int = 256
+    update_proj_gap: int = 200
+    scale: float = 0.25
+    proj_quant: bool = True
+    proj_quant_group_size: int = -1
+    proj_quant_n_bit: int = 4
+    weight_quant: bool = False
+    stochastic_round: bool = True
+    weight_group_size: int = 128
+    cos_threshold: float = 0.4
+    gamma_proj: float = 2.0
+    queue_size: int = 5
+    target_modules: Optional[List[str]] = None
+
+
 class UnslothTrainingArguments(TrainingArguments):
-    def __init__(self, embedding_learning_rate: float = None, *args, **kwargs):
+    def __init__(
+        self,
+        embedding_learning_rate: float = None,
+        q_galore_config: Optional[QGaloreConfig] = None,
+        *args,
+        **kwargs,
+    ):
+        self.q_galore_config = q_galore_config
+        self.embedding_learning_rate = embedding_learning_rate
         super().__init__(*args, **kwargs)
         self.embedding_learning_rate = embedding_learning_rate
 
@@ -181,6 +213,13 @@ def _create_unsloth_optimizer(
 
 class UnslothTrainer(SFTTrainer):
     def create_optimizer(self):
+        # --- Q-GaLore optimizer ---
+        q_galore_config = getattr(self.args, "q_galore_config", None)
+        if q_galore_config is not None and self.optimizer is None:
+            embedding_lr = getattr(self.args, "embedding_learning_rate", None)
+            return self._create_q_galore_optimizer(q_galore_config, embedding_lr)
+
+        # --- Embedding-LR optimizer ---
         embedding_learning_rate = getattr(self.args, "embedding_learning_rate", None)
         if embedding_learning_rate is None:
             return super().create_optimizer()
@@ -197,6 +236,105 @@ class UnslothTrainer(SFTTrainer):
             )
         return self.optimizer
 
+    def _create_q_galore_optimizer(self, config: "QGaloreConfig", embedding_lr = None):
+        """Build the Q-GaLore optimizer from a QGaloreConfig."""
+        from unsloth.optimizers.q_galore_adamw import (
+            QGaLoreAdamW8bit,
+            make_q_galore_param_groups,
+            install_weight_quant_hooks,
+        )
+
+        lr = self.args.learning_rate
+        weight_decay = self.args.weight_decay
+
+        param_groups = make_q_galore_param_groups(
+            self.model,
+            lr = lr,
+            weight_decay = weight_decay,
+            rank = config.rank,
+            update_proj_gap = config.update_proj_gap,
+            scale = config.scale,
+            proj_quant = config.proj_quant,
+            proj_quant_group_size = config.proj_quant_group_size,
+            proj_quant_n_bit = config.proj_quant_n_bit,
+            weight_quant = config.weight_quant,
+            stochastic_round = config.stochastic_round,
+            weight_group_size = config.weight_group_size,
+            cos_threshold = config.cos_threshold,
+            gamma_proj = config.gamma_proj,
+            queue_size = config.queue_size,
+            target_modules = config.target_modules,
+        )
+
+        # --- Split embedding params with custom LR (Fix #2) ---
+        if embedding_lr is not None:
+            # Build a fast param->name lookup (O(N) instead of O(N*M))
+            param_to_name = {id(p): name for name, p in self.model.named_parameters()}
+
+            new_groups = []
+            for group in param_groups:
+                if "rank" in group:
+                    # GaLore group — keep as-is (embeddings are never in here)
+                    new_groups.append(group)
+                    continue
+                # Non-GaLore group: split out embedding params
+                embed_params = []
+                other_params = []
+                for p in group["params"]:
+                    # Check if this param belongs to a modules_to_save embedding
+                    name = param_to_name.get(id(p))
+                    if name and name.endswith("modules_to_save.default.weight"):
+                        partial_name = name[: -len(".modules_to_save.default.weight")]
+                        partial_name = partial_name[partial_name.rfind(".") + 1 :]
+                        print(
+                            f"Unsloth: Setting lr = {embedding_lr:.2e} instead of {lr:.2e} for {partial_name}."
+                        )
+                        embed_params.append(p)
+                    else:
+                        other_params.append(p)
+                if other_params:
+                    other_group = dict(group)
+                    other_group["params"] = other_params
+                    new_groups.append(other_group)
+                if embed_params:
+                    embed_group = dict(group)
+                    embed_group["params"] = embed_params
+                    embed_group["lr"] = embedding_lr
+                    new_groups.append(embed_group)
+            param_groups = new_groups
+
+        # --- Forward optimizer hyperparameters (Fix #3) ---
+        self.optimizer = QGaLoreAdamW8bit(
+            param_groups,
+            lr = lr,
+            weight_decay = weight_decay,
+            betas = (self.args.adam_beta1, self.args.adam_beta2),
+            eps = self.args.adam_epsilon,
+        )
+
+        # Initialize INT8 weight quantization if enabled
+        if config.weight_quant:
+            QGaLoreAdamW8bit.init_weight_quantization(
+                self.model,
+                param_groups,
+                group_size = config.weight_group_size,
+                stochastic = config.stochastic_round,
+            )
+            # Forward pre-hooks dequantize INT8 weights to float before each
+            # forward pass, allowing the optimizer to free float weight memory
+            # between steps.
+            install_weight_quant_hooks(self.model)
+
+        n_galore = sum(len(g["params"]) for g in param_groups if "rank" in g)
+        n_other = sum(len(g["params"]) for g in param_groups if "rank" not in g)
+        print(
+            f"🦥 Unsloth: Q-GaLore enabled — "
+            f"{n_galore} GaLore params (rank={config.rank}), "
+            f"{n_other} standard params."
+        )
+
+        return self.optimizer
+
 
 # From `trl>=0.13.0`, they changed how to pass several params to the trainer
 # We need to patch to make the transition smooth

unsloth_cli/__init__.py

@@ -6,7 +6,6 @@ import typer
 from unsloth_cli.commands.train import train
 from unsloth_cli.commands.inference import inference
 from unsloth_cli.commands.export import export, list_checkpoints
-from unsloth_cli.commands.ui import ui
 from unsloth_cli.commands.studio import studio_app
 
 app = typer.Typer(
@@ -18,5 +17,4 @@ app.command()(train)
 app.command()(inference)
 app.command()(export)
 app.command("list-checkpoints")(list_checkpoints)
-app.command()(ui)
 app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.")

unsloth_cli/commands/studio.py

@@ -22,9 +22,9 @@ _PACKAGE_ROOT = Path(__file__).resolve().parent.parent.parent
 def _studio_venv_python() -> Optional[Path]:
     """Return the studio venv Python binary, or None if not set up."""
     if platform.system() == "Windows":
-        p = STUDIO_HOME / ".venv" / "Scripts" / "python.exe"
+        p = STUDIO_HOME / "unsloth_studio" / "Scripts" / "python.exe"
     else:
-        p = STUDIO_HOME / ".venv" / "bin" / "python"
+        p = STUDIO_HOME / "unsloth_studio" / "bin" / "python"
     return p if p.is_file() else None
 
 
@@ -44,7 +44,7 @@ def _find_run_py() -> Optional[Path]:
         "lib/python*/site-packages/studio/backend/run.py",
         "Lib/site-packages/studio/backend/run.py",
     ):
-        for match in (STUDIO_HOME / ".venv").glob(pattern):
+        for match in (STUDIO_HOME / "unsloth_studio").glob(pattern):
             return match
     return None
 
@@ -64,7 +64,7 @@ def _find_setup_script() -> Optional[Path]:
         f"lib/python*/site-packages/studio/{name}",
         f"Lib/site-packages/studio/{name}",
     ):
-        for match in (STUDIO_HOME / ".venv").glob(pattern):
+        for match in (STUDIO_HOME / "unsloth_studio").glob(pattern):
             return match
     return None
 
@@ -85,7 +85,7 @@ def studio_default(
         return
 
     # Always use the studio venv if it exists and we're not already in it
-    studio_venv_dir = STUDIO_HOME / ".venv"
+    studio_venv_dir = STUDIO_HOME / "unsloth_studio"
     in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
 
     if not in_studio_venv:
@@ -132,7 +132,7 @@ def studio_default(
             else:
                 os.execvp(str(studio_python), args)
         else:
-            typer.echo("Studio not set up. Run 'unsloth studio setup' first.")
+            typer.echo("Studio not set up. Run install.sh first.")
             raise typer.Exit(1)
 
     from studio.backend.run import run_server
@@ -166,12 +166,11 @@ def studio_default(
         typer.echo("\nShutting down...")
 
 
-# ── unsloth studio setup ─────────────────────────────────────────────
+# ── unsloth studio setup / update ─────────────────────────────────────
 
 
-@studio_app.command()
-def setup():
-    """Run one-time Studio environment setup."""
+def _run_setup_script() -> None:
+    """Find and run the studio setup/update script."""
     script = _find_setup_script()
     if not script:
         typer.echo("Error: Could not find setup script (setup.sh / setup.ps1).")
@@ -188,6 +187,35 @@ def setup():
         raise typer.Exit(result.returncode)
 
 
+@studio_app.command(hidden = True)
+def setup():
+    """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_setup_script()
+
+
+@studio_app.command()
+def update(
+    local: bool = typer.Option(
+        False, "--local", help = "Install from local repo instead of PyPI"
+    ),
+    package: str = typer.Option(
+        "unsloth", "--package", help = "Package name to install/update (for testing)"
+    ),
+):
+    """Update Unsloth Studio dependencies and rebuild."""
+    os.environ["STUDIO_LOCAL_INSTALL"] = "1" if local else "0"
+    os.environ["STUDIO_PACKAGE_NAME"] = package
+    if local:
+        # 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)
+    _run_setup_script()
+
+
 # ── unsloth studio reset-password ────────────────────────────────────
 
 

unsloth_cli/commands/ui.py

@@ -1,103 +0,0 @@
-# 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 os
-import sys
-import time
-from pathlib import Path
-from typing import Optional
-
-import typer
-
-
-def ui(
-    port: int = typer.Option(
-        8888, "--port", "-p", help = "Port to run the UI server on."
-    ),
-    host: str = typer.Option(
-        "0.0.0.0", "--host", "-H", help = "Host address to bind to."
-    ),
-    frontend: Optional[Path] = typer.Option(
-        None, "--frontend", "-f", help = "Path to frontend build directory."
-    ),
-    silent: bool = typer.Option(
-        False, "--silent", "-q", help = "Suppress startup messages."
-    ),
-):
-    """Launch the Unsloth web UI backend server (alias for 'unsloth studio')."""
-    from unsloth_cli.commands.studio import (
-        _studio_venv_python,
-        _find_run_py,
-        STUDIO_HOME,
-    )
-
-    # Re-execute in studio venv if available and not already inside it
-    studio_venv_dir = STUDIO_HOME / ".venv"
-    in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
-
-    if not in_studio_venv:
-        studio_python = _studio_venv_python()
-        run_py = _find_run_py()
-        if studio_python and run_py:
-            if not silent:
-                typer.echo("Launching Unsloth Studio... Please wait...")
-            args = [
-                str(studio_python),
-                str(run_py),
-                "--host",
-                host,
-                "--port",
-                str(port),
-            ]
-            if frontend:
-                args.extend(["--frontend", str(frontend)])
-            if silent:
-                args.append("--silent")
-            # On Windows, os.execvp() spawns a child but the parent lingers,
-            # so Ctrl+C only kills the parent leaving the child orphaned.
-            # Use subprocess.run() on Windows so the parent waits for the child.
-            if sys.platform == "win32":
-                import subprocess as _sp
-
-                proc = _sp.Popen(args)
-                try:
-                    rc = proc.wait()
-                except KeyboardInterrupt:
-                    # Child has its own signal handler — let it finish
-                    rc = proc.wait()
-                raise typer.Exit(rc)
-            else:
-                os.execvp(str(studio_python), args)
-        else:
-            typer.echo("Studio not set up. Run 'unsloth studio setup' first.")
-            raise typer.Exit(1)
-
-    from studio.backend.run import run_server
-
-    if not silent:
-        from studio.backend.run import _resolve_external_ip
-
-        display_host = _resolve_external_ip() if host == "0.0.0.0" else host
-        typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}")
-
-    run_kwargs = dict(host = host, port = port, silent = silent)
-    if frontend is not None:
-        run_kwargs["frontend_path"] = frontend
-    run_server(**run_kwargs)
-
-    from studio.backend.run import _shutdown_event
-
-    try:
-        if _shutdown_event is not None:
-            # NOTE: Event.wait() without a timeout blocks at the C level
-            # on Linux, preventing Python from delivering SIGINT (Ctrl+C).
-            while not _shutdown_event.is_set():
-                _shutdown_event.wait(timeout = 1)
-        else:
-            while True:
-                time.sleep(1)
-    except KeyboardInterrupt:
-        from studio.backend.run import _graceful_shutdown, _server
-
-        _graceful_shutdown(_server)
-        typer.echo("\nShutting down...")