Merge main into feat/chat-ui-enhancements and resolve db.ts conflict
Combined the extended type imports (FolderRecord, MemoryRecord, PromptRecord) from the PR branch with the useRef addition from main.
This commit is contained in:
Daniel Han
commit
9d61147d65
126 changed files with 21263 additions and 4542 deletions
|
|
@ -1,6 +1,6 @@
|
|||
repos:
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.15.7
|
||||
rev: v0.15.8
|
||||
hooks:
|
||||
- id: ruff
|
||||
args:
|
||||
|
|
|
|||
149
README.md
149
README.md
|
|
@ -57,19 +57,20 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**.
|
|||
```bash
|
||||
curl -fsSL https://unsloth.ai/install.sh | sh
|
||||
```
|
||||
If you don't have `curl`, use `wget`. Launch after setup via:
|
||||
```bash
|
||||
source unsloth_studio/bin/activate
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
|
||||
#### Windows:
|
||||
```powershell
|
||||
irm https://unsloth.ai/install.ps1 | iex
|
||||
```
|
||||
Launch after setup via:
|
||||
```powershell
|
||||
& .\unsloth_studio\Scripts\unsloth.exe studio -H 0.0.0.0 -p 8888
|
||||
|
||||
#### Launch
|
||||
```bash
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
|
||||
#### Update
|
||||
To update, use the same install commands as above. Or run (does not work on Windows):
|
||||
```bash
|
||||
unsloth studio update
|
||||
```
|
||||
|
||||
#### Docker
|
||||
|
|
@ -82,64 +83,8 @@ docker run -d -e JUPYTER_PASSWORD="mypassword" \
|
|||
unsloth/unsloth
|
||||
```
|
||||
|
||||
#### macOS, Linux, WSL developer installs:
|
||||
```bash
|
||||
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
uv venv unsloth_studio --python 3.13
|
||||
source unsloth_studio/bin/activate
|
||||
uv pip install unsloth --torch-backend=auto
|
||||
unsloth studio setup
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
|
||||
#### Windows PowerShell developer installs:
|
||||
```powershell
|
||||
winget install -e --id Python.Python.3.13
|
||||
winget install --id=astral-sh.uv -e
|
||||
uv venv unsloth_studio --python 3.13
|
||||
.\unsloth_studio\Scripts\activate
|
||||
uv pip install unsloth --torch-backend=auto
|
||||
unsloth studio setup
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
|
||||
#### Nightly - MacOS, Linux, WSL:
|
||||
```bash
|
||||
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
git clone --filter=blob:none https://github.com/unslothai/unsloth.git unsloth_studio
|
||||
cd unsloth_studio
|
||||
uv venv --python 3.13
|
||||
source .venv/bin/activate
|
||||
uv pip install -e . --torch-backend=auto
|
||||
unsloth studio setup
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
Then to launch every time:
|
||||
```bash
|
||||
cd unsloth_studio
|
||||
source .venv/bin/activate
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
|
||||
#### Nightly - Windows:
|
||||
Run in Windows Powershell:
|
||||
```bash
|
||||
winget install -e --id Python.Python.3.13
|
||||
winget install --id=astral-sh.uv -e
|
||||
git clone --filter=blob:none https://github.com/unslothai/unsloth.git unsloth_studio
|
||||
cd unsloth_studio
|
||||
uv venv --python 3.13
|
||||
.\.venv\Scripts\activate
|
||||
uv pip install -e . --torch-backend=auto
|
||||
unsloth studio setup
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
Then to launch every time:
|
||||
```bash
|
||||
cd unsloth_studio
|
||||
.\.venv\Scripts\activate
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
#### Developer, Nightly, Uninstall
|
||||
To see developer, nightly and uninstallation etc. instructions, see [advanced installation](#-advanced-installation).
|
||||
|
||||
### Unsloth Core (code-based)
|
||||
#### Linux, WSL:
|
||||
|
|
@ -197,6 +142,76 @@ Train for free with our notebooks. Read our [guide](https://unsloth.ai/docs/get-
|
|||
- **FP8 & Vision RL**: You can now do FP8 & VLM GRPO on consumer GPUs. [FP8 Blog](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) • [Vision RL](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/vision-reinforcement-learning-vlm-rl)
|
||||
- **gpt-oss** by OpenAI: Read our [RL blog](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/gpt-oss-reinforcement-learning), [Flex Attention](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/long-context-gpt-oss-training) blog and [Guide](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune).
|
||||
|
||||
## 📥 Advanced Installation
|
||||
The below advanced instructions are for Unsloth Studio. For Unsloth Core advanced installation, [view our docs](https://unsloth.ai/docs/get-started/install/pip-install#advanced-pip-installation).
|
||||
#### Developer installs: macOS, Linux, WSL:
|
||||
```bash
|
||||
git clone https://github.com/unslothai/unsloth
|
||||
cd unsloth
|
||||
./install.sh --local
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
Then to update :
|
||||
```bash
|
||||
unsloth studio update
|
||||
```
|
||||
|
||||
#### Developer installs: Windows PowerShell:
|
||||
```powershell
|
||||
git clone https://github.com/unslothai/unsloth.git
|
||||
cd unsloth
|
||||
Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass
|
||||
.\install.ps1 --local
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
Then to update :
|
||||
```bash
|
||||
unsloth studio update
|
||||
```
|
||||
|
||||
#### Nightly: MacOS, Linux, WSL:
|
||||
```bash
|
||||
git clone https://github.com/unslothai/unsloth
|
||||
cd unsloth
|
||||
git checkout nightly
|
||||
./install.sh --local
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
Then to launch every time:
|
||||
```bash
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
|
||||
#### Nightly: Windows:
|
||||
Run in Windows Powershell:
|
||||
```bash
|
||||
git clone https://github.com/unslothai/unsloth.git
|
||||
cd unsloth
|
||||
git checkout nightly
|
||||
Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass
|
||||
.\install.ps1 --local
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
Then to launch every time:
|
||||
```bash
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
|
||||
#### Uninstall
|
||||
You can uninstall Unsloth Studio by deleting its install folder usually located under `$HOME/.unsloth/studio` on Mac/Linux/WSL and `%USERPROFILE%\.unsloth\studio` on Windows. Using the `rm -rf` commands will **delete everything**, including your history, cache:
|
||||
|
||||
* **MacOS, WSL, Linux:** `rm -rf ~/.unsloth/studio`
|
||||
* **Windows (PowerShell):** `Remove-Item -Recurse -Force "$HOME\.unsloth\studio"`
|
||||
|
||||
For more info, [see our docs](https://unsloth.ai/docs/new/studio/install#uninstall).
|
||||
|
||||
#### Deleting model files
|
||||
|
||||
You can delete old model files either from the bin icon in model search or by removing the relevant cached model folder from the default Hugging Face cache directory. By default, HF uses:
|
||||
|
||||
* **MacOS, Linux, WSL:** `~/.cache/huggingface/hub/`
|
||||
* **Windows:** `%USERPROFILE%\.cache\huggingface\hub\`
|
||||
|
||||
## 💚 Community and Links
|
||||
| Type | Links |
|
||||
| ----------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
|
||||
|
|
|
|||
17
build.sh
17
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
|
||||
|
|
|
|||
515
install.ps1
515
install.ps1
|
|
@ -1,17 +1,102 @@
|
|||
# Unsloth Studio Installer for Windows PowerShell
|
||||
# Usage: irm https://raw.githubusercontent.com/unslothai/unsloth/main/install.ps1 | iex
|
||||
# Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1
|
||||
# Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 --local
|
||||
# NoTorch: .\install.ps1 --no-torch (skip PyTorch, GGUF-only mode)
|
||||
# Test: .\install.ps1 --package roland-sloth
|
||||
|
||||
function Install-UnslothStudio {
|
||||
$ErrorActionPreference = "Stop"
|
||||
$script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq "1")
|
||||
|
||||
# ── Parse flags ──
|
||||
$StudioLocalInstall = $false
|
||||
$PackageName = "unsloth"
|
||||
$RepoRoot = ""
|
||||
$SkipTorch = $false
|
||||
$argList = $args
|
||||
for ($i = 0; $i -lt $argList.Count; $i++) {
|
||||
switch ($argList[$i]) {
|
||||
"--local" { $StudioLocalInstall = $true }
|
||||
"--no-torch" { $SkipTorch = $true }
|
||||
"--verbose" { $script:UnslothVerbose = $true }
|
||||
"-v" { $script:UnslothVerbose = $true }
|
||||
"--package" {
|
||||
$i++
|
||||
if ($i -ge $argList.Count) {
|
||||
Write-Host "[ERROR] --package requires an argument." -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
$PackageName = $argList[$i]
|
||||
}
|
||||
}
|
||||
}
|
||||
# Propagate to child processes so they also respect verbose mode.
|
||||
# Process-scoped -- does not persist.
|
||||
if ($script:UnslothVerbose) {
|
||||
$env:UNSLOTH_VERBOSE = '1'
|
||||
}
|
||||
|
||||
if ($StudioLocalInstall) {
|
||||
$RepoRoot = (Resolve-Path (Split-Path -Parent $PSCommandPath)).Path
|
||||
if (-not (Test-Path (Join-Path $RepoRoot "pyproject.toml"))) {
|
||||
Write-Host "[ERROR] --local must be run from the unsloth repo root (pyproject.toml not found at $RepoRoot)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
$VenvName = "unsloth_studio"
|
||||
$PythonVersion = "3.13"
|
||||
$StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
|
||||
$VenvDir = Join-Path $StudioHome "unsloth_studio"
|
||||
|
||||
$Rule = [string]::new([char]0x2500, 52)
|
||||
$Sloth = [char]::ConvertFromUtf32(0x1F9A5)
|
||||
|
||||
function Enable-StudioVirtualTerminal {
|
||||
if ($env:NO_COLOR) { return $false }
|
||||
try {
|
||||
if (-not ("StudioVT.Native" -as [type])) {
|
||||
Add-Type -Namespace StudioVT -Name Native -MemberDefinition @'
|
||||
[DllImport("kernel32.dll")] public static extern IntPtr GetStdHandle(int nStdHandle);
|
||||
[DllImport("kernel32.dll")] public static extern bool GetConsoleMode(IntPtr h, out uint m);
|
||||
[DllImport("kernel32.dll")] public static extern bool SetConsoleMode(IntPtr h, uint m);
|
||||
'@ -ErrorAction Stop
|
||||
}
|
||||
$h = [StudioVT.Native]::GetStdHandle(-11)
|
||||
[uint32]$mode = 0
|
||||
if (-not [StudioVT.Native]::GetConsoleMode($h, [ref]$mode)) { return $false }
|
||||
$mode = $mode -bor 0x0004
|
||||
return [StudioVT.Native]::SetConsoleMode($h, $mode)
|
||||
} catch {
|
||||
return $false
|
||||
}
|
||||
}
|
||||
$script:StudioVtOk = Enable-StudioVirtualTerminal
|
||||
|
||||
function Get-StudioAnsi {
|
||||
param(
|
||||
[Parameter(Mandatory = $true)]
|
||||
[ValidateSet('Title', 'Dim', 'Ok', 'Warn', 'Err', 'Reset')]
|
||||
[string]$Kind
|
||||
)
|
||||
$e = [char]27
|
||||
switch ($Kind) {
|
||||
'Title' { return "${e}[38;5;150m" }
|
||||
'Dim' { return "${e}[38;5;245m" }
|
||||
'Ok' { return "${e}[38;5;108m" }
|
||||
'Warn' { return "${e}[38;5;136m" }
|
||||
'Err' { return "${e}[91m" }
|
||||
'Reset' { return "${e}[0m" }
|
||||
}
|
||||
}
|
||||
|
||||
Write-Host ""
|
||||
Write-Host "========================================="
|
||||
Write-Host " Unsloth Studio Installer (Windows)"
|
||||
Write-Host "========================================="
|
||||
if ($script:StudioVtOk -and -not $env:NO_COLOR) {
|
||||
Write-Host (" " + (Get-StudioAnsi Title) + $Sloth + " Unsloth Studio Installer (Windows)" + (Get-StudioAnsi Reset))
|
||||
Write-Host (" {0}{1}{2}" -f (Get-StudioAnsi Dim), $Rule, (Get-StudioAnsi Reset))
|
||||
} else {
|
||||
Write-Host (" {0} Unsloth Studio Installer (Windows)" -f $Sloth) -ForegroundColor DarkGreen
|
||||
Write-Host " $Rule" -ForegroundColor DarkGray
|
||||
}
|
||||
Write-Host ""
|
||||
|
||||
# ── Helper: refresh PATH from registry (deduplicating entries) ──
|
||||
|
|
@ -31,13 +116,96 @@ function Install-UnslothStudio {
|
|||
$env:Path = $unique -join ";"
|
||||
}
|
||||
|
||||
function step {
|
||||
param(
|
||||
[Parameter(Mandatory = $true)][string]$Label,
|
||||
[Parameter(Mandatory = $true)][string]$Value,
|
||||
[string]$Color = "Green"
|
||||
)
|
||||
if ($script:StudioVtOk -and -not $env:NO_COLOR) {
|
||||
$dim = Get-StudioAnsi Dim
|
||||
$rst = Get-StudioAnsi Reset
|
||||
$val = switch ($Color) {
|
||||
'Green' { Get-StudioAnsi Ok }
|
||||
'Yellow' { Get-StudioAnsi Warn }
|
||||
'Red' { Get-StudioAnsi Err }
|
||||
'DarkGray' { Get-StudioAnsi Dim }
|
||||
default { Get-StudioAnsi Ok }
|
||||
}
|
||||
$padded = if ($Label.Length -ge 15) { $Label.Substring(0, 15) } else { $Label.PadRight(15) }
|
||||
Write-Host (" {0}{1}{2}{3}{4}{2}" -f $dim, $padded, $rst, $val, $Value)
|
||||
} else {
|
||||
$padded = if ($Label.Length -ge 15) { $Label.Substring(0, 15) } else { $Label.PadRight(15) }
|
||||
Write-Host (" {0}" -f $padded) -NoNewline -ForegroundColor DarkGray
|
||||
$fc = switch ($Color) {
|
||||
'Green' { 'DarkGreen' }
|
||||
'Yellow' { 'Yellow' }
|
||||
'Red' { 'Red' }
|
||||
'DarkGray' { 'DarkGray' }
|
||||
default { 'DarkGreen' }
|
||||
}
|
||||
Write-Host $Value -ForegroundColor $fc
|
||||
}
|
||||
}
|
||||
|
||||
function substep {
|
||||
param(
|
||||
[Parameter(Mandatory = $true)][string]$Message,
|
||||
[string]$Color = "DarkGray"
|
||||
)
|
||||
if ($script:StudioVtOk -and -not $env:NO_COLOR) {
|
||||
$msgCol = switch ($Color) {
|
||||
'Yellow' { (Get-StudioAnsi Warn) }
|
||||
'Red' { (Get-StudioAnsi Err) }
|
||||
default { (Get-StudioAnsi Dim) }
|
||||
}
|
||||
$pad = "".PadRight(15)
|
||||
Write-Host (" {0}{1}{2}{3}" -f $msgCol, $pad, $Message, (Get-StudioAnsi Reset))
|
||||
} else {
|
||||
$fc = switch ($Color) {
|
||||
'Yellow' { 'Yellow' }
|
||||
'Red' { 'Red' }
|
||||
default { 'DarkGray' }
|
||||
}
|
||||
Write-Host (" {0,-15}{1}" -f "", $Message) -ForegroundColor $fc
|
||||
}
|
||||
}
|
||||
|
||||
# Run native commands quietly by default to match install.sh behavior.
|
||||
# Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1.
|
||||
function Invoke-InstallCommand {
|
||||
param(
|
||||
[Parameter(Mandatory = $true)][ScriptBlock]$Command
|
||||
)
|
||||
$prevEap = $ErrorActionPreference
|
||||
$ErrorActionPreference = "Continue"
|
||||
try {
|
||||
# Reset to avoid stale values from prior native commands.
|
||||
$global:LASTEXITCODE = 0
|
||||
if ($script:UnslothVerbose) {
|
||||
# Merge stderr into stdout so progress/warning output stays visible
|
||||
# without flipping $? on successful native commands (PS 5.1 treats
|
||||
# stderr records as errors that set $? = $false even on exit code 0).
|
||||
& $Command 2>&1 | Out-Host
|
||||
} else {
|
||||
$output = & $Command 2>&1 | Out-String
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host $output -ForegroundColor Red
|
||||
}
|
||||
}
|
||||
return [int]$LASTEXITCODE
|
||||
} finally {
|
||||
$ErrorActionPreference = $prevEap
|
||||
}
|
||||
}
|
||||
|
||||
function New-StudioShortcuts {
|
||||
param(
|
||||
[Parameter(Mandatory = $true)][string]$UnslothExePath
|
||||
)
|
||||
|
||||
if (-not (Test-Path $UnslothExePath)) {
|
||||
Write-Host "[WARN] Cannot create shortcuts: unsloth.exe not found at $UnslothExePath" -ForegroundColor Yellow
|
||||
substep "cannot create shortcuts, unsloth.exe not found at $UnslothExePath" "Yellow"
|
||||
return
|
||||
}
|
||||
try {
|
||||
|
|
@ -50,7 +218,7 @@ function Install-UnslothStudio {
|
|||
|
||||
$localAppDataDir = $env:LOCALAPPDATA
|
||||
if (-not $localAppDataDir -or [string]::IsNullOrWhiteSpace($localAppDataDir)) {
|
||||
Write-Host "[WARN] LOCALAPPDATA path unavailable; skipped shortcut creation" -ForegroundColor Yellow
|
||||
substep "LOCALAPPDATA path unavailable; skipped shortcut creation" "Yellow"
|
||||
return
|
||||
}
|
||||
$appDir = Join-Path $localAppDataDir "Unsloth Studio"
|
||||
|
|
@ -73,10 +241,10 @@ function Install-UnslothStudio {
|
|||
$null
|
||||
}
|
||||
if (-not $desktopLink) {
|
||||
Write-Host "[WARN] Desktop path unavailable; skipped desktop shortcut creation" -ForegroundColor Yellow
|
||||
substep "Desktop path unavailable; skipped desktop shortcut creation" "Yellow"
|
||||
}
|
||||
if (-not $startMenuLink) {
|
||||
Write-Host "[WARN] APPDATA/Start Menu path unavailable; skipped Start menu shortcut creation" -ForegroundColor Yellow
|
||||
substep "APPDATA/Start Menu path unavailable; skipped Start menu shortcut creation" "Yellow"
|
||||
}
|
||||
$iconPath = Join-Path $appDir "unsloth.ico"
|
||||
$bundledIcon = $null
|
||||
|
|
@ -135,6 +303,44 @@ function Find-HealthyStudioPort {
|
|||
return `$null
|
||||
}
|
||||
|
||||
function Test-PortBusy {
|
||||
param([Parameter(Mandatory = `$true)][int]`$Port)
|
||||
`$listener = `$null
|
||||
try {
|
||||
`$listener = [System.Net.Sockets.TcpListener]::new([System.Net.IPAddress]::Any, `$Port)
|
||||
`$listener.Start()
|
||||
return `$false
|
||||
} catch {
|
||||
return `$true
|
||||
} finally {
|
||||
if (`$listener) { try { `$listener.Stop() } catch {} }
|
||||
}
|
||||
}
|
||||
|
||||
function Find-FreeLaunchPort {
|
||||
`$maxPort = `$basePort + `$maxPortOffset
|
||||
try {
|
||||
`$listening = Get-NetTCPConnection -State Listen -ErrorAction Stop |
|
||||
Where-Object { `$_.LocalPort -ge `$basePort -and `$_.LocalPort -le `$maxPort } |
|
||||
Select-Object -ExpandProperty LocalPort
|
||||
for (`$offset = 0; `$offset -le `$maxPortOffset; `$offset++) {
|
||||
`$candidate = `$basePort + `$offset
|
||||
if (`$candidate -notin `$listening) {
|
||||
return `$candidate
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
# Get-NetTCPConnection unavailable or restricted; probe ports directly
|
||||
for (`$offset = 0; `$offset -le `$maxPortOffset; `$offset++) {
|
||||
`$candidate = `$basePort + `$offset
|
||||
if (-not (Test-PortBusy -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) {
|
||||
|
|
@ -163,7 +369,16 @@ try {
|
|||
|
||||
`$powershellExe = Join-Path `$env:SystemRoot 'System32\WindowsPowerShell\v1.0\powershell.exe'
|
||||
`$studioExe = '$SingleQuotedExePath'
|
||||
`$studioCommand = '& "' + `$studioExe + '" studio -H 0.0.0.0 -p ' + `$basePort
|
||||
`$launchPort = Find-FreeLaunchPort
|
||||
if (-not `$launchPort) {
|
||||
`$msg = "No free port found in range `$basePort-`$(`$basePort + `$maxPortOffset)"
|
||||
try {
|
||||
Add-Type -AssemblyName System.Windows.Forms -ErrorAction Stop
|
||||
[System.Windows.Forms.MessageBox]::Show(`$msg, 'Unsloth Studio') | Out-Null
|
||||
} catch {}
|
||||
exit 1
|
||||
}
|
||||
`$studioCommand = '& "' + `$studioExe + '" studio -H 0.0.0.0 -p ' + `$launchPort
|
||||
`$launchArgs = @(
|
||||
'-NoExit',
|
||||
'-NoProfile',
|
||||
|
|
@ -284,27 +499,27 @@ shell.Run cmd, 0, False
|
|||
$shortcut.Save()
|
||||
$createdShortcutCount++
|
||||
} catch {
|
||||
Write-Host "[WARN] Could not create shortcut at ${linkPath}: $($_.Exception.Message)" -ForegroundColor Yellow
|
||||
substep "could not create shortcut at ${linkPath}: $($_.Exception.Message)" "Yellow"
|
||||
}
|
||||
}
|
||||
if ($createdShortcutCount -gt 0) {
|
||||
Write-Host "[OK] Created Unsloth Studio shortcut(s): $createdShortcutCount" -ForegroundColor Green
|
||||
substep "Created Unsloth Studio shortcut"
|
||||
} else {
|
||||
Write-Host "[WARN] No Unsloth Studio shortcuts were created" -ForegroundColor Yellow
|
||||
substep "no Unsloth Studio shortcuts were created" "Yellow"
|
||||
}
|
||||
} catch {
|
||||
Write-Host "[WARN] Shortcut creation unavailable: $($_.Exception.Message)" -ForegroundColor Yellow
|
||||
substep "shortcut creation unavailable: $($_.Exception.Message)" "Yellow"
|
||||
}
|
||||
} catch {
|
||||
Write-Host "[WARN] Shortcut setup failed; skipping shortcuts: $($_.Exception.Message)" -ForegroundColor Yellow
|
||||
substep "shortcut setup failed; skipping shortcuts: $($_.Exception.Message)" "Yellow"
|
||||
}
|
||||
}
|
||||
|
||||
# ── Check winget ──
|
||||
if (-not (Get-Command winget -ErrorAction SilentlyContinue)) {
|
||||
Write-Host "Error: winget is not available." -ForegroundColor Red
|
||||
Write-Host " Install it from https://aka.ms/getwinget" -ForegroundColor Yellow
|
||||
Write-Host " or install Python $PythonVersion and uv manually, then re-run." -ForegroundColor Yellow
|
||||
step "winget" "not available" "Red"
|
||||
substep "Install it from https://aka.ms/getwinget" "Yellow"
|
||||
substep "or install Python $PythonVersion and uv manually, then re-run." "Yellow"
|
||||
return
|
||||
}
|
||||
|
||||
|
|
@ -382,10 +597,10 @@ shell.Run cmd, 0, False
|
|||
# Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null.
|
||||
$DetectedPython = Find-CompatiblePython
|
||||
if ($DetectedPython) {
|
||||
Write-Host "==> Python already installed: Python $($DetectedPython.Version)"
|
||||
step "python" "Python $($DetectedPython.Version) already installed"
|
||||
}
|
||||
if (-not $DetectedPython) {
|
||||
Write-Host "==> Installing Python ${PythonVersion}..."
|
||||
substep "installing Python ${PythonVersion}..."
|
||||
$pythonPackageId = "Python.Python.$PythonVersion"
|
||||
# Temporarily lower ErrorActionPreference so that winget stderr
|
||||
# (progress bars, warnings) does not become a terminating error
|
||||
|
|
@ -407,7 +622,7 @@ shell.Run cmd, 0, False
|
|||
# This handles both real failures AND "already installed" codes where
|
||||
# winget thinks Python is present but it's not actually on PATH
|
||||
# (e.g. user partially uninstalled, or installed via a different method).
|
||||
Write-Host " Python not found on PATH after winget. Retrying with --force..."
|
||||
substep "Python not found on PATH after winget. Retrying with --force..." "Yellow"
|
||||
$ErrorActionPreference = "Continue"
|
||||
try {
|
||||
winget install -e --id $pythonPackageId --accept-package-agreements --accept-source-agreements --force
|
||||
|
|
@ -429,7 +644,7 @@ shell.Run cmd, 0, False
|
|||
|
||||
# ── Install uv if not present ──
|
||||
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
|
||||
Write-Host "==> Installing uv package manager..."
|
||||
substep "installing uv package manager..."
|
||||
$prevEAP = $ErrorActionPreference
|
||||
$ErrorActionPreference = "Continue"
|
||||
try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {}
|
||||
|
|
@ -437,32 +652,73 @@ shell.Run cmd, 0, False
|
|||
Refresh-SessionPath
|
||||
# Fallback: if winget didn't put uv on PATH, try the PowerShell installer
|
||||
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
|
||||
Write-Host " Trying alternative uv installer..."
|
||||
substep "trying alternative uv installer..." "Yellow"
|
||||
powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex"
|
||||
Refresh-SessionPath
|
||||
}
|
||||
}
|
||||
|
||||
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
|
||||
Write-Host "Error: uv could not be installed." -ForegroundColor Red
|
||||
Write-Host " Install it from https://docs.astral.sh/uv/" -ForegroundColor Yellow
|
||||
step "uv" "could not be installed" "Red"
|
||||
substep "Install it from https://docs.astral.sh/uv/" "Yellow"
|
||||
return
|
||||
}
|
||||
|
||||
# ── 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
|
||||
substep "removing existing environment for fresh install..."
|
||||
Remove-Item -Recurse -Force $VenvDir
|
||||
} elseif (Test-Path (Join-Path $StudioHome ".venv\Scripts\python.exe")) {
|
||||
# Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating
|
||||
$OldVenv = Join-Path $StudioHome ".venv"
|
||||
$OldPy = Join-Path $OldVenv "Scripts\python.exe"
|
||||
substep "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) {
|
||||
substep "legacy environment is healthy -- migrating..."
|
||||
Move-Item -Path $OldVenv -Destination $VenvDir -Force
|
||||
substep "moved .venv -> unsloth_studio"
|
||||
$_Migrated = $true
|
||||
} else {
|
||||
substep "legacy environment failed validation -- creating fresh environment" "Yellow"
|
||||
Remove-Item -Recurse -Force $OldVenv -ErrorAction SilentlyContinue
|
||||
}
|
||||
} elseif (Test-Path (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) {
|
||||
# CWD-relative venv from old install.ps1 -- migrate to absolute path
|
||||
$CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio"
|
||||
substep "found CWD-relative Studio environment, migrating to $VenvDir..."
|
||||
Move-Item -Path $CwdVenv -Destination $VenvDir -Force
|
||||
substep "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)"
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host "[ERROR] Failed to create virtual environment (exit code $LASTEXITCODE)" -ForegroundColor Red
|
||||
step "venv" "creating Python $($DetectedPython.Version) virtual environment"
|
||||
substep "$VenvDir"
|
||||
$venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" }
|
||||
if ($venvExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
} else {
|
||||
Write-Host "==> Virtual environment ${VenvName} already exists, skipping creation."
|
||||
step "venv" "using migrated environment"
|
||||
substep "$VenvDir"
|
||||
}
|
||||
|
||||
# ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ──
|
||||
|
|
@ -471,7 +727,7 @@ shell.Run cmd, 0, False
|
|||
try {
|
||||
$nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue
|
||||
if ($nvSmiCmd) {
|
||||
& $nvSmiCmd.Source 2>&1 | Out-Null
|
||||
& $nvSmiCmd.Source *> $null
|
||||
if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source }
|
||||
}
|
||||
} catch {}
|
||||
|
|
@ -482,18 +738,18 @@ shell.Run cmd, 0, False
|
|||
)) {
|
||||
if (Test-Path $p) {
|
||||
try {
|
||||
& $p 2>&1 | Out-Null
|
||||
& $p *> $null
|
||||
if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break }
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
}
|
||||
if ($HasNvidiaSmi) {
|
||||
Write-Host "[OK] NVIDIA GPU detected" -ForegroundColor Green
|
||||
step "gpu" "NVIDIA GPU detected"
|
||||
} else {
|
||||
Write-Host "[WARN] No NVIDIA GPU detected. Studio will run in chat-only (GGUF) mode." -ForegroundColor Yellow
|
||||
Write-Host " Training and GPU inference require an NVIDIA GPU with drivers installed." -ForegroundColor Yellow
|
||||
Write-Host " https://www.nvidia.com/Download/index.aspx" -ForegroundColor Yellow
|
||||
step "gpu" "none (chat-only / GGUF)" "Yellow"
|
||||
substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow"
|
||||
substep "https://www.nvidia.com/Download/index.aspx" "Yellow"
|
||||
}
|
||||
|
||||
# ── Choose the correct PyTorch index URL based on driver CUDA version ──
|
||||
|
|
@ -513,11 +769,21 @@ shell.Run cmd, 0, False
|
|||
return "$baseUrl/cpu"
|
||||
}
|
||||
} catch {}
|
||||
Write-Host "[WARN] Could not determine CUDA version from nvidia-smi, defaulting to cu126" -ForegroundColor Yellow
|
||||
substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow"
|
||||
return "$baseUrl/cu126"
|
||||
}
|
||||
$TorchIndexUrl = Get-TorchIndexUrl
|
||||
|
||||
# ── Print CPU-only hint when no GPU detected ──
|
||||
if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") {
|
||||
Write-Host ""
|
||||
Write-Host " NOTE: No NVIDIA GPU detected." -ForegroundColor Yellow
|
||||
Write-Host " Installing CPU-only PyTorch. If you only need GGUF chat/inference,"
|
||||
Write-Host " re-run with --no-torch for a faster, lighter install:"
|
||||
Write-Host " .\install.ps1 --no-torch"
|
||||
Write-Host ""
|
||||
}
|
||||
|
||||
# ── Install PyTorch first, then unsloth separately ──
|
||||
#
|
||||
# Why two steps?
|
||||
|
|
@ -536,25 +802,116 @@ shell.Run cmd, 0, False
|
|||
# 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
|
||||
# ── Helper: find no-torch-runtime.txt ──
|
||||
function Find-NoTorchRuntimeFile {
|
||||
if ($StudioLocalInstall -and (Test-Path (Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt"))) {
|
||||
return Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt"
|
||||
}
|
||||
$installed = Get-ChildItem -Path $VenvDir -Recurse -Filter "no-torch-runtime.txt" -ErrorAction SilentlyContinue |
|
||||
Where-Object { $_.FullName -like "*studio*backend*requirements*no-torch-runtime.txt" } |
|
||||
Select-Object -ExpandProperty FullName -First 1
|
||||
return $installed
|
||||
}
|
||||
|
||||
Write-Host "==> Installing unsloth (this may take a few minutes)..."
|
||||
uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11"
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install unsloth (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
|
||||
substep "upgrading unsloth in migrated environment..."
|
||||
if ($SkipTorch) {
|
||||
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
|
||||
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo }
|
||||
if ($baseInstallExit -eq 0) {
|
||||
$NoTorchReq = Find-NoTorchRuntimeFile
|
||||
if ($NoTorchReq) {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
|
||||
}
|
||||
}
|
||||
} else {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo }
|
||||
}
|
||||
if ($baseInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
if ($StudioLocalInstall) {
|
||||
substep "overlaying local repo (editable)..."
|
||||
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
if ($overlayExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
}
|
||||
} elseif ($TorchIndexUrl) {
|
||||
if ($SkipTorch) {
|
||||
substep "skipping PyTorch (--no-torch flag set)." "Yellow"
|
||||
} else {
|
||||
substep "installing PyTorch ($TorchIndexUrl)..."
|
||||
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
|
||||
if ($torchInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
substep "installing unsloth (this may take a few minutes)..."
|
||||
if ($SkipTorch) {
|
||||
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
|
||||
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo }
|
||||
if ($baseInstallExit -eq 0) {
|
||||
$NoTorchReq = Find-NoTorchRuntimeFile
|
||||
if ($NoTorchReq) {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
|
||||
}
|
||||
}
|
||||
} elseif ($StudioLocalInstall) {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo }
|
||||
} else {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" }
|
||||
}
|
||||
if ($baseInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
|
||||
if ($StudioLocalInstall) {
|
||||
substep "overlaying local repo (editable)..."
|
||||
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
if ($overlayExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
}
|
||||
} else {
|
||||
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
|
||||
substep "installing unsloth (this may take a few minutes)..."
|
||||
if ($StudioLocalInstall) {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto }
|
||||
if ($baseInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
substep "overlaying local repo (editable)..."
|
||||
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
if ($overlayExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
} else {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "$PackageName" --torch-backend=auto }
|
||||
if ($baseInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# ── Run studio setup ──
|
||||
# setup.ps1 will handle installing Git, CMake, Visual Studio Build Tools,
|
||||
# CUDA Toolkit, Node.js, and other dependencies automatically via winget.
|
||||
Write-Host "==> Running unsloth studio setup..."
|
||||
$UnslothExe = Join-Path $VenvName "Scripts\unsloth.exe"
|
||||
step "setup" "running unsloth studio setup..."
|
||||
$UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe"
|
||||
if (-not (Test-Path $UnslothExe)) {
|
||||
Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red
|
||||
Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow
|
||||
|
|
@ -562,35 +919,57 @@ shell.Run cmd, 0, False
|
|||
Write-Host " Try re-running the installer or see: https://github.com/unslothai/unsloth?tab=readme-ov-file#-quickstart" -ForegroundColor Yellow
|
||||
return
|
||||
}
|
||||
& $UnslothExe studio setup
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host "[ERROR] unsloth studio setup failed (exit code $LASTEXITCODE)" -ForegroundColor Red
|
||||
# Tell setup.ps1 to skip base package installation (install.ps1 already did it)
|
||||
$env:SKIP_STUDIO_BASE = "1"
|
||||
$env:STUDIO_PACKAGE_NAME = $PackageName
|
||||
$env:UNSLOTH_NO_TORCH = if ($SkipTorch) { "true" } else { "false" }
|
||||
# Always set STUDIO_LOCAL_INSTALL explicitly to avoid stale values from
|
||||
# a previous --local run in the same PowerShell session.
|
||||
if ($StudioLocalInstall) {
|
||||
$env:STUDIO_LOCAL_INSTALL = "1"
|
||||
$env:STUDIO_LOCAL_REPO = $RepoRoot
|
||||
} else {
|
||||
$env:STUDIO_LOCAL_INSTALL = "0"
|
||||
Remove-Item Env:STUDIO_LOCAL_REPO -ErrorAction SilentlyContinue
|
||||
}
|
||||
# Use 'studio setup' (not 'studio update') because 'update' pops
|
||||
# SKIP_STUDIO_BASE, which would cause redundant package reinstallation
|
||||
# and bypass the fast-path version check from PR #4667.
|
||||
$studioArgs = @('studio', 'setup')
|
||||
if ($script:UnslothVerbose) { $studioArgs += '--verbose' }
|
||||
& $UnslothExe @studioArgs
|
||||
$setupExit = $LASTEXITCODE
|
||||
if ($setupExit -ne 0) {
|
||||
Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
|
||||
New-StudioShortcuts -UnslothExePath $UnslothExe
|
||||
|
||||
Write-Host ""
|
||||
Write-Host "========================================="
|
||||
Write-Host " Unsloth Studio installed!"
|
||||
Write-Host "========================================="
|
||||
Write-Host ""
|
||||
# ── Add venv Scripts dir to User PATH so `unsloth studio` works from any terminal ──
|
||||
$ScriptsDir = Join-Path $VenvDir "Scripts"
|
||||
$UserPath = [System.Environment]::GetEnvironmentVariable("Path", "User")
|
||||
if (-not $UserPath -or $UserPath -notlike "*$ScriptsDir*") {
|
||||
if ($UserPath) {
|
||||
[System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir;$UserPath", "User")
|
||||
} else {
|
||||
[System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir", "User")
|
||||
}
|
||||
Refresh-SessionPath
|
||||
step "path" "added unsloth to PATH"
|
||||
}
|
||||
|
||||
# Launch studio automatically in interactive terminals;
|
||||
# in non-interactive environments (CI, Docker) just print instructions.
|
||||
$IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected)
|
||||
if ($IsInteractive) {
|
||||
Write-Host "==> Launching Unsloth Studio..."
|
||||
Write-Host ""
|
||||
$UnslothExe = 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 " unsloth studio -H 0.0.0.0 -p 8888"
|
||||
step "launch" "manual commands:"
|
||||
substep "& `"$VenvDir\Scripts\Activate.ps1`""
|
||||
substep "unsloth studio -H 0.0.0.0 -p 8888"
|
||||
Write-Host ""
|
||||
}
|
||||
}
|
||||
|
||||
Install-UnslothStudio
|
||||
Install-UnslothStudio @args
|
||||
|
|
|
|||
702
install.sh
702
install.sh
|
|
@ -1,11 +1,111 @@
|
|||
#!/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 (no-torch): ./install.sh --no-torch (skip PyTorch, GGUF-only mode)
|
||||
# Usage (test): ./install.sh --package roland-sloth (install a different package name)
|
||||
# Usage (py): ./install.sh --python 3.12 (override auto-detected Python version)
|
||||
set -e
|
||||
|
||||
VENV_NAME="unsloth_studio"
|
||||
PYTHON_VERSION="3.13"
|
||||
# ── Output style (aligned with studio/setup.sh) ──
|
||||
RULE=""
|
||||
_rule_i=0
|
||||
while [ "$_rule_i" -lt 52 ]; do
|
||||
RULE="${RULE}─"
|
||||
_rule_i=$((_rule_i + 1))
|
||||
done
|
||||
if [ -n "${NO_COLOR:-}" ]; then
|
||||
C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
|
||||
elif [ -t 1 ] || [ -n "${FORCE_COLOR:-}" ]; then
|
||||
_ESC="$(printf '\033')"
|
||||
C_TITLE="${_ESC}[38;5;150m"
|
||||
C_DIM="${_ESC}[38;5;245m"
|
||||
C_OK="${_ESC}[38;5;108m"
|
||||
C_WARN="${_ESC}[38;5;136m"
|
||||
C_ERR="${_ESC}[91m"
|
||||
C_RST="${_ESC}[0m"
|
||||
else
|
||||
C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
|
||||
fi
|
||||
|
||||
step() { printf " ${C_DIM}%-15.15s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; }
|
||||
substep() { printf " ${C_DIM}%-15s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; }
|
||||
|
||||
# ── Parse flags ──
|
||||
STUDIO_LOCAL_INSTALL=false
|
||||
PACKAGE_NAME="unsloth"
|
||||
_USER_PYTHON=""
|
||||
_NO_TORCH_FLAG=false
|
||||
_VERBOSE=false
|
||||
_next_is_package=false
|
||||
_next_is_python=false
|
||||
for arg in "$@"; do
|
||||
if [ "$_next_is_package" = true ]; then
|
||||
PACKAGE_NAME="$arg"
|
||||
_next_is_package=false
|
||||
continue
|
||||
fi
|
||||
if [ "$_next_is_python" = true ]; then
|
||||
_USER_PYTHON="$arg"
|
||||
_next_is_python=false
|
||||
continue
|
||||
fi
|
||||
case "$arg" in
|
||||
--local) STUDIO_LOCAL_INSTALL=true ;;
|
||||
--package) _next_is_package=true ;;
|
||||
--python) _next_is_python=true ;;
|
||||
--no-torch) _NO_TORCH_FLAG=true ;;
|
||||
--verbose|-v) _VERBOSE=true ;;
|
||||
esac
|
||||
done
|
||||
|
||||
if [ "$_VERBOSE" = true ]; then
|
||||
export UNSLOTH_VERBOSE=1
|
||||
fi
|
||||
|
||||
_is_verbose() {
|
||||
[ "${UNSLOTH_VERBOSE:-0}" = "1" ]
|
||||
}
|
||||
|
||||
run_maybe_quiet() {
|
||||
if _is_verbose; then
|
||||
"$@"
|
||||
else
|
||||
"$@" > /dev/null 2>&1
|
||||
fi
|
||||
}
|
||||
|
||||
run_install_cmd() {
|
||||
_label="$1"
|
||||
shift
|
||||
if _is_verbose; then
|
||||
"$@" && return 0
|
||||
_rc=$?
|
||||
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
|
||||
return "$_rc"
|
||||
fi
|
||||
_log=$(mktemp)
|
||||
"$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; }
|
||||
_rc=$?
|
||||
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
|
||||
cat "$_log" >&2
|
||||
rm -f "$_log"
|
||||
return $_rc
|
||||
}
|
||||
|
||||
if [ "$_next_is_package" = true ]; then
|
||||
echo "❌ ERROR: --package requires an argument." >&2
|
||||
exit 1
|
||||
fi
|
||||
if [ "$_next_is_python" = true ]; then
|
||||
echo "❌ ERROR: --python requires a version argument (e.g. --python 3.12)." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
PYTHON_VERSION="" # resolved after platform detection
|
||||
STUDIO_HOME="$HOME/.unsloth/studio"
|
||||
VENV_DIR="$STUDIO_HOME/unsloth_studio"
|
||||
|
||||
# ── Helper: download a URL to a file (supports curl and wget) ──
|
||||
download() {
|
||||
|
|
@ -92,17 +192,12 @@ _smart_apt_install() {
|
|||
# ── 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).
|
||||
# plus platform-specific shortcuts (Linux .desktop / macOS .app bundle /
|
||||
# WSL Windows Desktop+Start Menu .lnk).
|
||||
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"
|
||||
|
|
@ -231,6 +326,17 @@ _open_browser() {
|
|||
_url="$1"
|
||||
if [ "$(uname)" = "Darwin" ] && command -v open >/dev/null 2>&1; then
|
||||
open "$_url"
|
||||
elif grep -qi microsoft /proc/version 2>/dev/null; then
|
||||
# WSL: xdg-open is unreliable; use Windows browser via PowerShell or cmd
|
||||
if command -v powershell.exe >/dev/null 2>&1; then
|
||||
powershell.exe -NoProfile -Command "Start-Process '$_url'" >/dev/null 2>&1 &
|
||||
elif command -v cmd.exe >/dev/null 2>&1; then
|
||||
cmd.exe /c start "" "$_url" >/dev/null 2>&1 &
|
||||
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
|
||||
elif command -v xdg-open >/dev/null 2>&1; then
|
||||
xdg-open "$_url" >/dev/null 2>&1 &
|
||||
else
|
||||
|
|
@ -294,6 +400,8 @@ _acquire_lock() {
|
|||
}
|
||||
|
||||
_release_lock() {
|
||||
[ -d "$LOCK_DIR" ] || return 0
|
||||
[ "$(cat "$LOCK_DIR/pid" 2>/dev/null)" = "$$" ] || return 0
|
||||
rm -rf "$LOCK_DIR"
|
||||
}
|
||||
|
||||
|
|
@ -319,24 +427,48 @@ _launch_port=$(_find_launch_port) || {
|
|||
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"
|
||||
if [ -t 1 ]; then
|
||||
# ── Foreground mode (TTY available) ──
|
||||
# Background subshell: wait for studio to become healthy, release the
|
||||
# single-instance lock, then open the browser. The lock stays held until
|
||||
# health is confirmed so a second launcher cannot race during startup.
|
||||
(
|
||||
_obwr_deadline=$(($(date +%s) + TIMEOUT_SEC))
|
||||
while [ "$(date +%s)" -lt "$_obwr_deadline" ]; do
|
||||
if _check_health "$_launch_port"; then
|
||||
_release_lock
|
||||
_open_browser "http://localhost:$_launch_port"
|
||||
exit 0
|
||||
fi
|
||||
sleep "$POLL_INTERVAL_SEC"
|
||||
done
|
||||
# Timed out -- release the lock anyway so future launches are not blocked
|
||||
_release_lock
|
||||
) &
|
||||
# Clear traps so exec does not trigger _release_lock (the subshell owns it)
|
||||
trap - EXIT INT TERM
|
||||
exec "$UNSLOTH_EXE" studio -H 0.0.0.0 -p "$_launch_port"
|
||||
else
|
||||
# ── Background mode (no TTY) ──
|
||||
# Used by macOS .app and headless invocations.
|
||||
_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
|
||||
# Poll for health on the specific port we launched on
|
||||
_deadline=$(($(date +%s) + TIMEOUT_SEC))
|
||||
while [ "$(date +%s)" -lt "$_deadline" ]; do
|
||||
if _check_health "$_launch_port"; then
|
||||
_open_browser "http://localhost:$_launch_port"
|
||||
exit 0
|
||||
fi
|
||||
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
|
||||
echo "Unsloth Studio did not become healthy within ${TIMEOUT_SEC}s." >&2
|
||||
echo "Check logs at: $LOG_FILE" >&2
|
||||
exit 1
|
||||
fi
|
||||
LAUNCHER_EOF
|
||||
|
||||
chmod +x "$_css_launcher"
|
||||
|
|
@ -348,44 +480,33 @@ LAUNCHER_EOF
|
|||
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)
|
||||
# rounded-512.png used for both Linux and macOS icons
|
||||
_css_script_dir=""
|
||||
if [ -n "${0:-}" ] && [ -f "$0" ]; then
|
||||
_css_script_dir=$(cd "$(dirname "$0")" 2>/dev/null && pwd) || true
|
||||
fi
|
||||
|
||||
# Try to find favicon.png from installed package (site-packages) or local repo
|
||||
_css_found_favicon=""
|
||||
_css_found_gem=""
|
||||
# Try to find rounded-512.png from installed package (site-packages) or local repo
|
||||
_css_found_icon=""
|
||||
_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"
|
||||
if [ -f "$_sp/rounded-512.png" ]; then
|
||||
_css_found_icon="$_sp/rounded-512.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"
|
||||
if [ -z "$_css_found_icon" ] && [ -n "$_css_script_dir" ] && [ -f "$_css_script_dir/studio/frontend/public/rounded-512.png" ]; then
|
||||
_css_found_icon="$_css_script_dir/studio/frontend/public/rounded-512.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
|
||||
# Copy or download rounded-512.png (used for both Linux icon and macOS icns)
|
||||
if [ -n "$_css_found_icon" ]; then
|
||||
cp "$_css_found_icon" "$_css_icon_png" 2>/dev/null || true
|
||||
cp "$_css_found_icon" "$_css_gem_png" 2>/dev/null || true
|
||||
else
|
||||
download "https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/rounded-512.png" "$_css_icon_png" 2>/dev/null || true
|
||||
cp "$_css_icon_png" "$_css_gem_png" 2>/dev/null || true
|
||||
fi
|
||||
|
||||
# Validate PNG header (first 4 bytes: \x89PNG)
|
||||
|
|
@ -421,7 +542,7 @@ Name=Unsloth Studio
|
|||
Comment=Launch Unsloth Studio
|
||||
Exec="$_css_exec_escaped"
|
||||
Icon=$_css_icon_escaped
|
||||
Terminal=false
|
||||
Terminal=true
|
||||
StartupNotify=true
|
||||
Categories=Development;Science;
|
||||
DESKTOP_EOF
|
||||
|
|
@ -517,17 +638,78 @@ STUB_EOF
|
|||
ln -sf "$_css_app" "$HOME/Desktop/Unsloth Studio" 2>/dev/null || true
|
||||
fi
|
||||
_css_created=1
|
||||
|
||||
elif [ "$_css_os" = "wsl" ]; then
|
||||
# ── WSL: create Windows Desktop and Start Menu shortcuts ──
|
||||
# Detect current WSL distro for targeted shortcut
|
||||
_css_distro="${WSL_DISTRO_NAME:-}"
|
||||
|
||||
# Build the wsl.exe arguments.
|
||||
# Double-quote distro name and launcher path for Windows command line
|
||||
# parsing so values with spaces (e.g. "Ubuntu Preview") are kept as
|
||||
# single arguments.
|
||||
_css_wsl_args=""
|
||||
if [ -n "$_css_distro" ]; then
|
||||
_css_wsl_args="-d \"$_css_distro\" "
|
||||
fi
|
||||
_css_wsl_args="${_css_wsl_args}-- bash -l -c \"exec \\\"$_css_launcher\\\"\""
|
||||
|
||||
# Detect whether Windows Terminal (wt.exe) is available (better UX)
|
||||
_css_use_wt=false
|
||||
if command -v wt.exe >/dev/null 2>&1; then
|
||||
_css_use_wt=true
|
||||
fi
|
||||
|
||||
if [ "$_css_use_wt" = true ]; then
|
||||
_css_sc_target='wt.exe'
|
||||
_css_sc_args="wsl.exe $_css_wsl_args"
|
||||
else
|
||||
_css_sc_target='wsl.exe'
|
||||
_css_sc_args="$_css_wsl_args"
|
||||
fi
|
||||
|
||||
# Escape single quotes for PowerShell single-quoted string embedding
|
||||
_css_sc_args_ps=$(printf '%s' "$_css_sc_args" | sed "s/'/''/g")
|
||||
|
||||
# Create shortcuts via a temp PowerShell script to avoid escaping issues
|
||||
_css_ps1_tmp=$(mktemp /tmp/unsloth-shortcut-XXXXXX.ps1 2>/dev/null) || true
|
||||
if [ -n "$_css_ps1_tmp" ]; then
|
||||
cat > "$_css_ps1_tmp" << WSLPS1_EOF
|
||||
\$WshShell = New-Object -ComObject WScript.Shell
|
||||
\$targetExe = (Get-Command '$_css_sc_target' -ErrorAction SilentlyContinue).Source
|
||||
if (-not \$targetExe) { exit 1 }
|
||||
\$locations = @(
|
||||
[Environment]::GetFolderPath('Desktop'),
|
||||
(Join-Path \$env:APPDATA 'Microsoft\Windows\Start Menu\Programs')
|
||||
)
|
||||
foreach (\$dir in \$locations) {
|
||||
if (-not \$dir -or -not (Test-Path \$dir)) { continue }
|
||||
\$linkPath = Join-Path \$dir 'Unsloth Studio.lnk'
|
||||
\$shortcut = \$WshShell.CreateShortcut(\$linkPath)
|
||||
\$shortcut.TargetPath = \$targetExe
|
||||
\$shortcut.Arguments = '$_css_sc_args_ps'
|
||||
\$shortcut.Description = 'Launch Unsloth Studio'
|
||||
\$shortcut.Save()
|
||||
}
|
||||
WSLPS1_EOF
|
||||
|
||||
# Convert WSL path to Windows path for powershell.exe
|
||||
_css_ps1_win=$(wslpath -w "$_css_ps1_tmp" 2>/dev/null)
|
||||
if [ -n "$_css_ps1_win" ]; then
|
||||
powershell.exe -NoProfile -ExecutionPolicy Bypass -File "$_css_ps1_win" >/dev/null 2>&1 && _css_created=1
|
||||
fi
|
||||
rm -f "$_css_ps1_tmp"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "$_css_created" -eq 1 ]; then
|
||||
echo "[OK] Created Unsloth Studio shortcut(s)"
|
||||
substep "Created Unsloth Studio shortcut"
|
||||
fi
|
||||
}
|
||||
|
||||
echo ""
|
||||
echo "========================================="
|
||||
echo " Unsloth Studio Installer"
|
||||
echo "========================================="
|
||||
printf " ${C_TITLE}%s${C_RST}\n" "🦥 Unsloth Studio Installer"
|
||||
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
|
||||
echo ""
|
||||
|
||||
# ── Detect platform ──
|
||||
|
|
@ -537,7 +719,48 @@ if [ "$(uname)" = "Darwin" ]; then
|
|||
elif grep -qi microsoft /proc/version 2>/dev/null; then
|
||||
OS="wsl"
|
||||
fi
|
||||
echo "==> Platform: $OS"
|
||||
step "platform" "$OS"
|
||||
|
||||
# ── Architecture detection & Python version ──
|
||||
_ARCH=$(uname -m)
|
||||
MAC_INTEL=false
|
||||
if [ "$OS" = "macos" ] && [ "$_ARCH" = "x86_64" ]; then
|
||||
# Guard against Apple Silicon running under Rosetta (reports x86_64).
|
||||
# sysctl hw.optional.arm64 returns "1" on Apple Silicon even in Rosetta.
|
||||
if [ "$(sysctl -in hw.optional.arm64 2>/dev/null || echo 0)" = "1" ]; then
|
||||
echo ""
|
||||
echo " WARNING: Apple Silicon detected, but this shell is running under Rosetta (x86_64)."
|
||||
echo " Re-run install.sh from a native arm64 terminal for full PyTorch support."
|
||||
echo " Continuing in GGUF-only mode for now."
|
||||
echo ""
|
||||
fi
|
||||
MAC_INTEL=true
|
||||
fi
|
||||
|
||||
if [ -n "$_USER_PYTHON" ]; then
|
||||
PYTHON_VERSION="$_USER_PYTHON"
|
||||
echo " Using user-specified Python $PYTHON_VERSION (--python override)"
|
||||
elif [ "$MAC_INTEL" = true ]; then
|
||||
PYTHON_VERSION="3.12"
|
||||
else
|
||||
PYTHON_VERSION="3.13"
|
||||
fi
|
||||
|
||||
if [ "$MAC_INTEL" = true ]; then
|
||||
echo ""
|
||||
echo " NOTE: Intel Mac (x86_64) detected."
|
||||
echo " PyTorch is unavailable for this platform (dropped Jan 2024)."
|
||||
echo " Studio will install in GGUF-only mode."
|
||||
echo " Chat, inference via GGUF, and data recipes will work."
|
||||
echo " Training requires Apple Silicon or Linux with GPU."
|
||||
echo ""
|
||||
fi
|
||||
|
||||
# ── Unified SKIP_TORCH: --no-torch flag OR Intel Mac auto-detection ──
|
||||
SKIP_TORCH=false
|
||||
if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then
|
||||
SKIP_TORCH=true
|
||||
fi
|
||||
|
||||
# ── Check system dependencies ──
|
||||
# cmake and git are needed by unsloth studio setup to build the GGUF inference
|
||||
|
|
@ -576,8 +799,8 @@ MISSING=$(echo "$MISSING" | sed 's/^ *//')
|
|||
|
||||
if [ -n "$MISSING" ]; then
|
||||
echo ""
|
||||
echo "==> Unsloth Studio needs these packages: $MISSING"
|
||||
echo " These are needed to build the GGUF inference engine."
|
||||
step "deps" "missing: $MISSING" "$C_WARN"
|
||||
substep "These are needed to build the GGUF inference engine."
|
||||
|
||||
case "$OS" in
|
||||
macos)
|
||||
|
|
@ -602,7 +825,7 @@ if [ -n "$MISSING" ]; then
|
|||
esac
|
||||
echo ""
|
||||
else
|
||||
echo "==> All system dependencies found."
|
||||
step "deps" "all system dependencies found"
|
||||
fi
|
||||
|
||||
# ── Install uv ──
|
||||
|
|
@ -648,10 +871,10 @@ _uv_version_ok() {
|
|||
}
|
||||
|
||||
if ! command -v uv >/dev/null 2>&1 || ! _uv_version_ok uv; then
|
||||
echo "==> Installing uv package manager..."
|
||||
substep "installing uv package manager..."
|
||||
_uv_tmp=$(mktemp)
|
||||
download "https://astral.sh/uv/install.sh" "$_uv_tmp"
|
||||
sh "$_uv_tmp" </dev/null
|
||||
run_maybe_quiet sh "$_uv_tmp" </dev/null
|
||||
rm -f "$_uv_tmp"
|
||||
if [ -f "$HOME/.local/bin/env" ]; then
|
||||
. "$HOME/.local/bin/env"
|
||||
|
|
@ -659,43 +882,346 @@ 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
|
||||
substep "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 an Intel Mac has a stale 3.13 venv from a previous failed install, recreate
|
||||
# (skip when the user explicitly chose a version via --python)
|
||||
if [ "$SKIP_TORCH" = true ] && [ "$MAC_INTEL" = true ] && [ -z "$_USER_PYTHON" ] && [ -x "$VENV_DIR/bin/python" ]; then
|
||||
_PY_MM=$("$VENV_DIR/bin/python" -c \
|
||||
"import sys; print('{}.{}'.format(*sys.version_info[:2]))" 2>/dev/null || echo "")
|
||||
if [ "$_PY_MM" != "3.12" ]; then
|
||||
echo " Recreating Intel Mac environment with Python 3.12 (was $_PY_MM)..."
|
||||
rm -rf "$VENV_DIR"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ ! -x "$VENV_DIR/bin/python" ]; then
|
||||
step "venv" "creating Python ${PYTHON_VERSION} virtual environment"
|
||||
substep "$VENV_DIR"
|
||||
run_install_cmd "create venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
|
||||
fi
|
||||
|
||||
# Guard against Python 3.13.8 torch import bug on Apple Silicon
|
||||
# (skip when the user explicitly chose a version via --python)
|
||||
if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
|
||||
_PY_VER=$("$VENV_DIR/bin/python" -c \
|
||||
"import sys; print('{}.{}.{}'.format(*sys.version_info[:3]))" 2>/dev/null || echo "")
|
||||
if [ "$_PY_VER" = "3.13.8" ]; then
|
||||
echo " WARNING: Python 3.13.8 has a known torch import bug."
|
||||
echo " Recreating venv with Python 3.12..."
|
||||
rm -rf "$VENV_DIR"
|
||||
PYTHON_VERSION="3.12"
|
||||
run_install_cmd "recreate venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -x "$VENV_DIR/bin/python" ]; then
|
||||
step "venv" "using environment"
|
||||
substep "${VENV_DIR}"
|
||||
fi
|
||||
|
||||
# ── Resolve repo root (for --local installs) ──
|
||||
_REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)"
|
||||
|
||||
# ── Helper: find no-torch-runtime.txt (local repo or site-packages) ──
|
||||
_find_no_torch_runtime() {
|
||||
# Check local repo first (for --local installs)
|
||||
if [ -f "$_REPO_ROOT/studio/backend/requirements/no-torch-runtime.txt" ]; then
|
||||
echo "$_REPO_ROOT/studio/backend/requirements/no-torch-runtime.txt"
|
||||
return
|
||||
fi
|
||||
# Check inside installed package
|
||||
_rt=$(find "$VENV_DIR" -path "*/studio/backend/requirements/no-torch-runtime.txt" -print -quit 2>/dev/null || echo "")
|
||||
if [ -n "$_rt" ]; then
|
||||
echo "$_rt"
|
||||
return
|
||||
fi
|
||||
}
|
||||
|
||||
# ── 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)
|
||||
|
||||
# ── Print CPU-only hint when no GPU detected ──
|
||||
case "$TORCH_INDEX_URL" in
|
||||
*/cpu)
|
||||
if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then
|
||||
echo ""
|
||||
echo " NOTE: No NVIDIA GPU detected (nvidia-smi not found)."
|
||||
echo " Installing CPU-only PyTorch. If you only need GGUF chat/inference,"
|
||||
echo " re-run with --no-torch for a faster, lighter install:"
|
||||
echo " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch"
|
||||
echo ""
|
||||
fi
|
||||
;;
|
||||
esac
|
||||
|
||||
# ── 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>=2026.3.11" --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
|
||||
substep "upgrading unsloth in migrated environment..."
|
||||
if [ "$SKIP_TORCH" = true ]; then
|
||||
# No-torch: install unsloth + unsloth-zoo with --no-deps (current
|
||||
# PyPI metadata still declares torch as a hard dep), then install
|
||||
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps
|
||||
# to prevent transitive torch resolution.
|
||||
run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
|
||||
--reinstall-package unsloth --reinstall-package unsloth-zoo \
|
||||
"unsloth>=2026.3.16" unsloth-zoo
|
||||
_NO_TORCH_RT="$(_find_no_torch_runtime)"
|
||||
if [ -n "$_NO_TORCH_RT" ]; then
|
||||
run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
|
||||
fi
|
||||
else
|
||||
run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
|
||||
--reinstall-package unsloth --reinstall-package unsloth-zoo \
|
||||
"unsloth>=2026.3.16" unsloth-zoo
|
||||
fi
|
||||
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
||||
substep "overlaying local repo (editable)..."
|
||||
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
|
||||
fi
|
||||
elif [ -n "$TORCH_INDEX_URL" ]; then
|
||||
# Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac)
|
||||
if [ "$SKIP_TORCH" = true ]; then
|
||||
substep "skipping PyTorch (--no-torch or Intel Mac x86_64)." "$C_WARN"
|
||||
else
|
||||
substep "installing PyTorch ($TORCH_INDEX_URL)..."
|
||||
run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \
|
||||
--index-url "$TORCH_INDEX_URL"
|
||||
fi
|
||||
# Fresh: Step 2 - install unsloth, preserving pre-installed torch
|
||||
substep "installing unsloth (this may take a few minutes)..."
|
||||
if [ "$SKIP_TORCH" = true ]; then
|
||||
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
|
||||
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
|
||||
run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
|
||||
--upgrade-package unsloth --upgrade-package unsloth-zoo \
|
||||
"unsloth>=2026.3.16" unsloth-zoo
|
||||
_NO_TORCH_RT="$(_find_no_torch_runtime)"
|
||||
if [ -n "$_NO_TORCH_RT" ]; then
|
||||
run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
|
||||
fi
|
||||
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
||||
substep "overlaying local repo (editable)..."
|
||||
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
|
||||
fi
|
||||
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
||||
run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \
|
||||
--upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo
|
||||
substep "overlaying local repo (editable)..."
|
||||
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
|
||||
else
|
||||
run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \
|
||||
--upgrade-package unsloth "$PACKAGE_NAME"
|
||||
fi
|
||||
else
|
||||
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
|
||||
substep "installing unsloth (this may take a few minutes)..."
|
||||
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
||||
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto
|
||||
substep "overlaying local repo (editable)..."
|
||||
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
|
||||
else
|
||||
run_install_cmd "install unsloth (auto torch backend)" 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
|
||||
if ! command -v bash >/dev/null 2>&1; then
|
||||
step "setup" "bash is required to run studio setup" "$C_ERR"
|
||||
substep "Please install bash and re-run install.sh"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
step "setup" "running unsloth studio update..."
|
||||
# install.sh already installs base packages (unsloth + unsloth-zoo) and
|
||||
# no-torch-runtime.txt above, so tell install_python_stack.py to skip
|
||||
# the base step to avoid redundant reinstallation.
|
||||
_SKIP_BASE=1
|
||||
# Run setup.sh outside set -e so that a llama.cpp build failure (exit 1)
|
||||
# does not skip PATH setup, shortcuts, and launch below. We capture the
|
||||
# exit code and propagate it after post-install steps finish.
|
||||
_SETUP_EXIT=0
|
||||
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
||||
SKIP_STUDIO_BASE="$_SKIP_BASE" \
|
||||
STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
|
||||
STUDIO_LOCAL_INSTALL=1 \
|
||||
STUDIO_LOCAL_REPO="$_REPO_ROOT" \
|
||||
UNSLOTH_NO_TORCH="$SKIP_TORCH" \
|
||||
bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
|
||||
else
|
||||
SKIP_STUDIO_BASE="$_SKIP_BASE" \
|
||||
STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
|
||||
UNSLOTH_NO_TORCH="$SKIP_TORCH" \
|
||||
bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
|
||||
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"
|
||||
step "path" "added ~/.local/bin to PATH in $_SHELL_PROFILE"
|
||||
fi
|
||||
fi
|
||||
export PATH="$_LOCAL_BIN:$PATH"
|
||||
;;
|
||||
esac
|
||||
|
||||
create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS"
|
||||
|
||||
# If setup.sh failed, report and exit now.
|
||||
# PATH and shortcuts are already set up so the user can fix and retry.
|
||||
if [ "$_SETUP_EXIT" -ne 0 ]; then
|
||||
echo ""
|
||||
step "error" "studio setup failed (exit code $_SETUP_EXIT)" "$C_ERR"
|
||||
substep "Check the output above for details, then re-run:"
|
||||
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
||||
substep " unsloth studio update --local"
|
||||
else
|
||||
substep " unsloth studio update"
|
||||
fi
|
||||
echo ""
|
||||
exit "$_SETUP_EXIT"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "========================================="
|
||||
echo " Unsloth Studio installed!"
|
||||
echo "========================================="
|
||||
printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!"
|
||||
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
|
||||
echo ""
|
||||
|
||||
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
|
||||
step "launch" "starting Unsloth Studio..."
|
||||
"$VENV_DIR/bin/unsloth" studio -H 0.0.0.0 -p 8888
|
||||
_LAUNCH_EXIT=$?
|
||||
if [ "$_LAUNCH_EXIT" -ne 0 ] && [ "$_MIGRATED" = true ]; then
|
||||
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
|
||||
step "launch" "manual commands:"
|
||||
substep "unsloth studio -H 0.0.0.0 -p 8888"
|
||||
substep "or activate env first:"
|
||||
substep "source ${VENV_DIR}/bin/activate"
|
||||
substep "unsloth studio -H 0.0.0.0 -p 8888"
|
||||
echo ""
|
||||
fi
|
||||
|
|
|
|||
|
|
@ -58,7 +58,8 @@ studio = [
|
|||
]
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
exclude = ["images*", "tests*", "kernels/moe*"]
|
||||
include = ["unsloth*", "unsloth_cli*", "studio", "studio.backend*"]
|
||||
exclude = ["images*", "tests*", "*.node_modules", "*.node_modules.*"]
|
||||
|
||||
[project.optional-dependencies]
|
||||
triton = [
|
||||
|
|
@ -87,7 +88,7 @@ huggingfacenotorch = [
|
|||
]
|
||||
huggingface = [
|
||||
"unsloth[huggingfacenotorch]",
|
||||
"unsloth_zoo>=2026.3.5",
|
||||
"unsloth_zoo>=2026.3.6",
|
||||
"torchvision",
|
||||
"unsloth[triton]",
|
||||
]
|
||||
|
|
@ -577,7 +578,7 @@ colab-ampere-torch220 = [
|
|||
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
|
||||
]
|
||||
colab-new = [
|
||||
"unsloth_zoo>=2026.3.5",
|
||||
"unsloth_zoo>=2026.3.6",
|
||||
"packaging",
|
||||
"tyro",
|
||||
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.3.0",
|
||||
|
|
|
|||
|
|
@ -1,157 +1,153 @@
|
|||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "view-in-github",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"<a href=\"https://colab.research.google.com/github/unslothai/unsloth/blob/main/studio/Unsloth_Studio_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6b87de59",
|
||||
"metadata": {
|
||||
"id": "6b87de59"
|
||||
},
|
||||
"source": [
|
||||
"To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n",
|
||||
"<div class=\"align-center\">\n",
|
||||
"<a href=\"https://unsloth.ai/\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
|
||||
"<a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord button.png\" width=\"145\"></a>\n",
|
||||
"<a href=\"https://unsloth.ai/docs/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a> Join Discord if you need help + ⭐ <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> ⭐\n",
|
||||
"</div>\n",
|
||||
"\n",
|
||||
"To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n",
|
||||
"\n",
|
||||
"### Unsloth Studio\n",
|
||||
"\n",
|
||||
"Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). Currently, installation may take 30+ mins so use a newer GPU.\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n",
|
||||
"\n",
|
||||
"[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Studio Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e4206349",
|
||||
"metadata": {
|
||||
"id": "e4206349"
|
||||
},
|
||||
"source": [
|
||||
"<p align=\"left\"><img src=\"https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/studio%20github%20landscape%20colab%20display.png\" width=\"600\"></p>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "27da2957",
|
||||
"metadata": {
|
||||
"id": "27da2957"
|
||||
},
|
||||
"source": [
|
||||
"### Setup: Clone repo and run setup"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "27e68f91",
|
||||
"metadata": {
|
||||
"id": "27e68f91"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n",
|
||||
"%cd /content/unsloth\n",
|
||||
"!chmod +x studio/setup.sh && ./studio/setup.sh"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3e1771a9",
|
||||
"metadata": {
|
||||
"id": "3e1771a9"
|
||||
},
|
||||
"source": [
|
||||
"### Start Unsloth Studio"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "277e431e",
|
||||
"metadata": {
|
||||
"id": "277e431e"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import sys, time\n",
|
||||
"sys.path.insert(0, \"/content/unsloth/studio/backend\")\n",
|
||||
"from colab import start\n",
|
||||
"start()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from google.colab import output\n",
|
||||
"output.serve_kernel_port_as_iframe(8888, height = 1200, width = \"100%\")\n",
|
||||
"for _ in range(10000): time.sleep(300), print(\"=\", end = \"\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "wb9UELh--XzX"
|
||||
},
|
||||
"id": "wb9UELh--XzX",
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f2b0c6a1",
|
||||
"metadata": {
|
||||
"id": "f2b0c6a1"
|
||||
},
|
||||
"source": [
|
||||
"And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n",
|
||||
"\n",
|
||||
"Some other resources:\n",
|
||||
"1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n",
|
||||
"2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n",
|
||||
"3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n",
|
||||
"4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n",
|
||||
"5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n",
|
||||
"\n",
|
||||
"<div class=\"align-center\">\n",
|
||||
" <a href=\"https://unsloth.ai\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
|
||||
" <a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord.png\" width=\"145\"></a>\n",
|
||||
" <a href=\"https://unsloth.ai/docs/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a>\n",
|
||||
"\n",
|
||||
" Join Discord if you need help + ⭐️ <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> ⭐️\n",
|
||||
"\n",
|
||||
" <b>This notebook is licensed <a href=\"https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0\">AGPL-3.0</a></b>\n",
|
||||
"</div>"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"provenance": [],
|
||||
"include_colab_link": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "view-in-github",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"<a href=\"https://colab.research.google.com/github/unslothai/unsloth/blob/main/studio/Unsloth_Studio_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
||||
]
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6b87de59",
|
||||
"metadata": {
|
||||
"id": "6b87de59"
|
||||
},
|
||||
"source": [
|
||||
"To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n",
|
||||
"<div class=\"align-center\">\n",
|
||||
"<a href=\"https://unsloth.ai/\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
|
||||
"<a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord button.png\" width=\"145\"></a>\n",
|
||||
"<a href=\"https://unsloth.ai/docs/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a> Join Discord if you need help + ⭐ <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> ⭐\n",
|
||||
"</div>\n",
|
||||
"\n",
|
||||
"To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n",
|
||||
"\n",
|
||||
"### Unsloth Studio\n",
|
||||
"\n",
|
||||
"Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n",
|
||||
"\n",
|
||||
"[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Studio Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e4206349",
|
||||
"metadata": {
|
||||
"id": "e4206349"
|
||||
},
|
||||
"source": [
|
||||
"<p align=\"left\"><img src=\"https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/studio%20github%20landscape%20colab%20display.png\" width=\"600\"></p>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "27da2957",
|
||||
"metadata": {
|
||||
"id": "27da2957"
|
||||
},
|
||||
"source": [
|
||||
"### Setup: Clone repo and run setup"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "27e68f91",
|
||||
"metadata": {
|
||||
"id": "27e68f91"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh"
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3e1771a9",
|
||||
"metadata": {
|
||||
"id": "3e1771a9"
|
||||
},
|
||||
"source": [
|
||||
"### Start Unsloth Studio"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "277e431e",
|
||||
"metadata": {
|
||||
"id": "277e431e"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import sys, time\n",
|
||||
"sys.path.insert(0, \"/content/unsloth/studio/backend\")\n",
|
||||
"from colab import start\n",
|
||||
"start()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from google.colab import output\n",
|
||||
"output.serve_kernel_port_as_iframe(8888, height = 1200, width = \"100%\")\n",
|
||||
"for _ in range(10000): time.sleep(300), print(\"=\", end = \"\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "wb9UELh--XzX"
|
||||
},
|
||||
"id": "wb9UELh--XzX",
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f2b0c6a1",
|
||||
"metadata": {
|
||||
"id": "f2b0c6a1"
|
||||
},
|
||||
"source": [
|
||||
"And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n",
|
||||
"\n",
|
||||
"Some other resources:\n",
|
||||
"1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n",
|
||||
"2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n",
|
||||
"3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n",
|
||||
"4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n",
|
||||
"5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n",
|
||||
"\n",
|
||||
"<div class=\"align-center\">\n",
|
||||
" <a href=\"https://unsloth.ai\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png\" width=\"115\"></a>\n",
|
||||
" <a href=\"https://discord.gg/unsloth\"><img src=\"https://github.com/unslothai/unsloth/raw/main/images/Discord.png\" width=\"145\"></a>\n",
|
||||
" <a href=\"https://unsloth.ai/docs/\"><img src=\"https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true\" width=\"125\"></a>\n",
|
||||
"\n",
|
||||
" Join Discord if you need help + ⭐️ <i>Star us on <a href=\"https://github.com/unslothai/unsloth\">Github</a> </i> ⭐️\n",
|
||||
"\n",
|
||||
" <b>This notebook is licensed <a href=\"https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0\">AGPL-3.0</a></b>\n",
|
||||
"</div>"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"provenance": [],
|
||||
"include_colab_link": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,13 +10,13 @@ training:
|
|||
load_in_4bit: false
|
||||
output_dir: outputs
|
||||
num_epochs: 1
|
||||
learning_rate: 0.0002
|
||||
learning_rate: 2e-5
|
||||
batch_size: 1
|
||||
gradient_accumulation_steps: 4
|
||||
warmup_steps: 5
|
||||
max_steps: 0
|
||||
save_steps: 0
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: false
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ training:
|
|||
warmup_steps: 5
|
||||
max_steps: 0
|
||||
save_steps: 0
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: false
|
||||
|
|
|
|||
|
|
@ -6,13 +6,13 @@ training:
|
|||
max_seq_length: 2048
|
||||
# num_epochs: 4
|
||||
num_epochs: 0
|
||||
learning_rate: 5e-5
|
||||
learning_rate: 2e-4
|
||||
batch_size: 2
|
||||
gradient_accumulation_steps: 4
|
||||
warmup_ratio: 0.1
|
||||
max_steps: 30
|
||||
save_steps: 30
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: true
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ training:
|
|||
warmup_ratio: 0.03
|
||||
max_steps: 30
|
||||
save_steps: 30
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: false
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ training:
|
|||
warmup_ratio: 0.03
|
||||
max_steps: 30
|
||||
save_steps: 30
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: false
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ training:
|
|||
warmup_ratio: 0.03
|
||||
max_steps: 30
|
||||
save_steps: 30
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: false
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ training:
|
|||
warmup_ratio: 0.03
|
||||
max_steps: 30
|
||||
save_steps: 30
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: false
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ training:
|
|||
warmup_ratio: 0.03
|
||||
max_steps: 30
|
||||
save_steps: 30
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: false
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ training:
|
|||
warmup_steps: 5
|
||||
max_steps: 30
|
||||
save_steps: 30
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: true
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ training:
|
|||
warmup_steps: 5
|
||||
max_steps: 30
|
||||
save_steps: 30
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: true
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ training:
|
|||
warmup_steps: 0
|
||||
max_steps: 30
|
||||
save_steps: 30
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: true
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ training:
|
|||
warmup_steps: 5
|
||||
max_steps: 0
|
||||
save_steps: 0
|
||||
weight_decay: 0.01
|
||||
weight_decay: 0.001
|
||||
random_seed: 3407
|
||||
packing: false
|
||||
train_on_completions: false
|
||||
|
|
|
|||
|
|
@ -18,31 +18,6 @@ if _backend_dir not in sys.path:
|
|||
import _platform_compat # noqa: F401
|
||||
|
||||
|
||||
def _bootstrap_studio_venv() -> None:
|
||||
"""Expose the Studio venv's site-packages to the current interpreter.
|
||||
|
||||
On Colab, notebook cells run outside the venv subshell. Instead of
|
||||
installing the full stack into system Python, we prepend the venv's
|
||||
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"
|
||||
if not venv_lib.exists():
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
f"Studio venv not found at {venv_lib.parent} -- run 'unsloth studio setup' first",
|
||||
stacklevel = 2,
|
||||
)
|
||||
return
|
||||
for sp in venv_lib.glob("python*/site-packages"):
|
||||
sp_str = str(sp)
|
||||
if sp_str not in sys.path:
|
||||
sys.path.insert(0, sp_str)
|
||||
|
||||
|
||||
_bootstrap_studio_venv()
|
||||
|
||||
from loggers import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
|
|
|||
|
|
@ -217,6 +217,7 @@ class ExportOrchestrator:
|
|||
max_seq_length: int = 2048,
|
||||
load_in_4bit: bool = True,
|
||||
trust_remote_code: bool = False,
|
||||
hf_token: Optional[str] = None,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Load a checkpoint for export.
|
||||
|
||||
|
|
@ -227,6 +228,7 @@ class ExportOrchestrator:
|
|||
"max_seq_length": max_seq_length,
|
||||
"load_in_4bit": load_in_4bit,
|
||||
"trust_remote_code": trust_remote_code,
|
||||
"hf_token": hf_token,
|
||||
}
|
||||
|
||||
# Always kill existing subprocess and spawn fresh.
|
||||
|
|
|
|||
|
|
@ -18,7 +18,14 @@ from typing import Optional, Union, Generator, Tuple
|
|||
from utils.models import ModelConfig, get_base_model_from_lora
|
||||
from utils.paths import is_model_cached
|
||||
from utils.utils import format_error_message
|
||||
from utils.hardware import get_device, clear_gpu_cache, log_gpu_memory
|
||||
from utils.hardware import (
|
||||
get_device,
|
||||
clear_gpu_cache,
|
||||
log_gpu_memory,
|
||||
get_device_map,
|
||||
raise_if_offloaded,
|
||||
get_visible_gpu_count,
|
||||
)
|
||||
from core.inference.audio_codecs import AudioCodecManager
|
||||
from io import StringIO
|
||||
import structlog
|
||||
|
|
@ -241,6 +248,7 @@ class InferenceBackend:
|
|||
load_in_4bit: bool = True,
|
||||
hf_token: Optional[str] = None,
|
||||
trust_remote_code: bool = False,
|
||||
gpu_ids: Optional[list[int]] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Load any model: base, LoRA adapter, text, or vision.
|
||||
|
|
@ -260,6 +268,10 @@ class InferenceBackend:
|
|||
return False
|
||||
|
||||
self.loading_models.add(model_name)
|
||||
device_map = get_device_map(gpu_ids)
|
||||
logger.info(
|
||||
f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)"
|
||||
)
|
||||
|
||||
self.models[model_name] = {
|
||||
"is_vision": config.is_vision,
|
||||
|
|
@ -290,6 +302,7 @@ class InferenceBackend:
|
|||
config.path,
|
||||
auto_model = CsmForConditionalGeneration,
|
||||
load_in_4bit = False,
|
||||
device_map = device_map,
|
||||
token = hf_token if hf_token and hf_token.strip() else None,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -325,6 +338,7 @@ class InferenceBackend:
|
|||
config.path,
|
||||
dtype = torch.float32,
|
||||
load_in_4bit = False,
|
||||
device_map = device_map,
|
||||
token = hf_token if hf_token and hf_token.strip() else None,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -345,6 +359,7 @@ class InferenceBackend:
|
|||
llm_path,
|
||||
dtype = torch.float32,
|
||||
load_in_4bit = False,
|
||||
device_map = device_map,
|
||||
token = hf_token if hf_token and hf_token.strip() else None,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -361,6 +376,7 @@ class InferenceBackend:
|
|||
config.path,
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = False,
|
||||
device_map = device_map,
|
||||
token = hf_token if hf_token and hf_token.strip() else None,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -378,6 +394,7 @@ class InferenceBackend:
|
|||
whisper_language = "English",
|
||||
whisper_task = "transcribe",
|
||||
load_in_4bit = False,
|
||||
device_map = device_map,
|
||||
token = hf_token if hf_token and hf_token.strip() else None,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -405,6 +422,7 @@ class InferenceBackend:
|
|||
model_name = config.path,
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = False,
|
||||
device_map = device_map,
|
||||
token = hf_token if hf_token and hf_token.strip() else None,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -420,6 +438,11 @@ class InferenceBackend:
|
|||
audio_type, self.device, model_repo_path = model_repo_path
|
||||
)
|
||||
|
||||
# Reject CPU/disk offload for audio models too
|
||||
raise_if_offloaded(
|
||||
self.models[model_name]["model"], device_map, "Inference"
|
||||
)
|
||||
|
||||
self.active_model_name = model_name
|
||||
self.loading_models.discard(model_name)
|
||||
logger.info(f"Successfully loaded audio model: {model_name}")
|
||||
|
|
@ -441,6 +464,7 @@ class InferenceBackend:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = dtype,
|
||||
load_in_4bit = load_in_4bit,
|
||||
device_map = device_map,
|
||||
token = hf_token if hf_token and hf_token.strip() else None,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -497,6 +521,7 @@ class InferenceBackend:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = dtype,
|
||||
load_in_4bit = load_in_4bit,
|
||||
device_map = device_map,
|
||||
token = hf_token if hf_token and hf_token.strip() else None,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -507,6 +532,10 @@ class InferenceBackend:
|
|||
self.models[model_name]["model"] = model
|
||||
self.models[model_name]["tokenizer"] = tokenizer
|
||||
|
||||
raise_if_offloaded(
|
||||
self.models[model_name]["model"], device_map, "Inference"
|
||||
)
|
||||
|
||||
# Load chat template info
|
||||
self._load_chat_template_info(model_name)
|
||||
|
||||
|
|
@ -615,6 +644,7 @@ class InferenceBackend:
|
|||
dtype = None,
|
||||
load_in_4bit: bool = True,
|
||||
hf_token: Optional[str] = None,
|
||||
gpu_ids: Optional[list[int]] = None,
|
||||
) -> Tuple[bool, Optional[str], Optional[str]]:
|
||||
"""
|
||||
Final Corrected Version:
|
||||
|
|
@ -639,7 +669,12 @@ class InferenceBackend:
|
|||
base_model_name, None, is_lora = False
|
||||
)
|
||||
if not self.load_model(
|
||||
base_config, max_seq_length, dtype, load_in_4bit, hf_token
|
||||
base_config,
|
||||
max_seq_length,
|
||||
dtype,
|
||||
load_in_4bit,
|
||||
hf_token,
|
||||
gpu_ids = gpu_ids,
|
||||
):
|
||||
return False, None, None
|
||||
|
||||
|
|
@ -1037,12 +1072,12 @@ class InferenceBackend:
|
|||
input_text,
|
||||
add_special_tokens = False,
|
||||
return_tensors = "pt",
|
||||
).to(self.device)
|
||||
).to(model.device)
|
||||
else:
|
||||
# Text-only for vision model
|
||||
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
|
||||
inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(
|
||||
self.device
|
||||
model.device
|
||||
)
|
||||
|
||||
# Stream with TextIteratorStreamer + background thread
|
||||
|
|
@ -1182,7 +1217,7 @@ class InferenceBackend:
|
|||
return_dict = True,
|
||||
return_tensors = "pt",
|
||||
truncation = False,
|
||||
).to(self.device)
|
||||
).to(model.device)
|
||||
|
||||
try:
|
||||
from transformers import TextIteratorStreamer
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -17,6 +17,7 @@ Pattern follows core/training/training.py.
|
|||
|
||||
import atexit
|
||||
import base64
|
||||
import os
|
||||
import structlog
|
||||
from loggers import get_logger
|
||||
import multiprocessing as mp
|
||||
|
|
@ -27,11 +28,17 @@ import uuid
|
|||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any, Generator, Optional, Tuple, Union
|
||||
from utils.hardware import prepare_gpu_selection
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_CTX = mp.get_context("spawn")
|
||||
|
||||
|
||||
class DownloadStallError(RuntimeError):
|
||||
"""Raised when the worker reports no download progress for too long."""
|
||||
|
||||
|
||||
# Dispatcher timeout constants (seconds)
|
||||
_DISPATCH_READ_TIMEOUT = 30.0
|
||||
_DISPATCH_POLL_INTERVAL = 0.5
|
||||
|
|
@ -262,12 +269,17 @@ class InferenceOrchestrator:
|
|||
except (EOFError, OSError, ValueError):
|
||||
return None
|
||||
|
||||
def _wait_response(self, expected_type: str, timeout: float = 120.0) -> dict:
|
||||
def _wait_response(self, expected_type: str, timeout: float = 300.0) -> dict:
|
||||
"""Block until a response of the expected type arrives.
|
||||
|
||||
Also handles 'status' and 'error' events during the wait.
|
||||
Returns the matching response dict.
|
||||
Raises RuntimeError on timeout or subprocess crash.
|
||||
|
||||
The *timeout* is an **inactivity** timeout: it resets whenever the
|
||||
subprocess sends a status message, so long-running operations (large
|
||||
downloads, slow model loads) won't be killed as long as the subprocess
|
||||
keeps reporting progress.
|
||||
"""
|
||||
deadline = time.monotonic() + timeout
|
||||
|
||||
|
|
@ -292,8 +304,15 @@ class InferenceOrchestrator:
|
|||
|
||||
if rtype == "status":
|
||||
logger.info("Subprocess status: %s", resp.get("message", ""))
|
||||
# Reset deadline — subprocess is still alive and working
|
||||
deadline = time.monotonic() + timeout
|
||||
continue
|
||||
|
||||
if rtype == "stall":
|
||||
msg = resp.get("message", "Download stalled")
|
||||
logger.warning("Subprocess reported stall: %s", msg)
|
||||
raise DownloadStallError(msg)
|
||||
|
||||
# Other response types during wait — skip
|
||||
logger.debug(
|
||||
"Skipping response type '%s' while waiting for '%s'",
|
||||
|
|
@ -302,7 +321,8 @@ class InferenceOrchestrator:
|
|||
)
|
||||
|
||||
raise RuntimeError(
|
||||
f"Timeout waiting for '{expected_type}' response after {timeout}s"
|
||||
f"Timeout waiting for '{expected_type}' response "
|
||||
f"(no activity for {timeout}s)"
|
||||
)
|
||||
|
||||
def _drain_queue(self) -> list:
|
||||
|
|
@ -571,6 +591,7 @@ class InferenceOrchestrator:
|
|||
load_in_4bit: bool = True,
|
||||
hf_token: Optional[str] = None,
|
||||
trust_remote_code: bool = False,
|
||||
gpu_ids: Optional[list[int]] = None,
|
||||
) -> bool:
|
||||
"""Load a model for inference.
|
||||
|
||||
|
|
@ -594,7 +615,16 @@ class InferenceOrchestrator:
|
|||
"hf_token": hf_token or "",
|
||||
"gguf_variant": getattr(config, "gguf_variant", None),
|
||||
"trust_remote_code": trust_remote_code,
|
||||
"gpu_ids": gpu_ids,
|
||||
}
|
||||
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
|
||||
gpu_ids,
|
||||
model_name = model_name,
|
||||
hf_token = hf_token,
|
||||
load_in_4bit = load_in_4bit,
|
||||
)
|
||||
sub_config["resolved_gpu_ids"] = resolved_gpu_ids
|
||||
sub_config["gpu_selection"] = gpu_selection
|
||||
|
||||
# Always kill existing subprocess and spawn fresh.
|
||||
# Reusing a subprocess after unsloth patches torch internals
|
||||
|
|
@ -608,36 +638,66 @@ class InferenceOrchestrator:
|
|||
# Dead subprocess — clean up
|
||||
self._shutdown_subprocess(timeout = 2)
|
||||
|
||||
logger.info(
|
||||
"Spawning fresh inference subprocess for '%s' (transformers %s.x)",
|
||||
model_name,
|
||||
needed_major,
|
||||
disable_xet = sub_config.get("disable_xet", False) or (
|
||||
os.environ.get("HF_HUB_DISABLE_XET") == "1"
|
||||
)
|
||||
self._spawn_subprocess(sub_config)
|
||||
resp = self._wait_response("loaded", timeout = 180)
|
||||
|
||||
# Update local state from response
|
||||
if resp.get("success"):
|
||||
self._current_transformers_major = needed_major
|
||||
model_info = resp.get("model_info", {})
|
||||
self.active_model_name = model_info.get("identifier", model_name)
|
||||
self.models[self.active_model_name] = {
|
||||
"is_vision": model_info.get("is_vision", False),
|
||||
"is_lora": model_info.get("is_lora", False),
|
||||
"display_name": model_info.get("display_name", model_name),
|
||||
"is_audio": model_info.get("is_audio", False),
|
||||
"audio_type": model_info.get("audio_type"),
|
||||
"has_audio_input": model_info.get("has_audio_input", False),
|
||||
}
|
||||
self.loading_models.discard(model_name)
|
||||
logger.info("Model '%s' loaded successfully in subprocess", model_name)
|
||||
return True
|
||||
else:
|
||||
error = resp.get("error", "Failed to load model")
|
||||
self.loading_models.discard(model_name)
|
||||
self.active_model_name = None
|
||||
self.models.clear()
|
||||
raise Exception(error)
|
||||
for attempt in range(2):
|
||||
logger.info(
|
||||
"Spawning fresh inference subprocess for '%s' "
|
||||
"(transformers %s.x, attempt %d/2%s)",
|
||||
model_name,
|
||||
needed_major,
|
||||
attempt + 1,
|
||||
", xet disabled" if disable_xet else "",
|
||||
)
|
||||
sub_config["disable_xet"] = disable_xet
|
||||
self._spawn_subprocess(sub_config)
|
||||
|
||||
try:
|
||||
resp = self._wait_response("loaded")
|
||||
except DownloadStallError:
|
||||
# First stall and Xet was enabled -> retry with Xet disabled
|
||||
if attempt == 0 and not disable_xet:
|
||||
logger.warning(
|
||||
"Download stalled for '%s' -- retrying with "
|
||||
"HF_HUB_DISABLE_XET=1",
|
||||
model_name,
|
||||
)
|
||||
self._shutdown_subprocess(timeout = 5)
|
||||
disable_xet = True
|
||||
continue
|
||||
# Second stall (or already had xet disabled) -> give up
|
||||
self._shutdown_subprocess(timeout = 5)
|
||||
raise RuntimeError(
|
||||
f"Download stalled for '{model_name}' even with "
|
||||
f"HF_HUB_DISABLE_XET=1 -- check your network connection"
|
||||
)
|
||||
|
||||
# Got a response — check success
|
||||
if resp.get("success"):
|
||||
self._current_transformers_major = needed_major
|
||||
model_info = resp.get("model_info", {})
|
||||
self.active_model_name = model_info.get("identifier", model_name)
|
||||
self.models[self.active_model_name] = {
|
||||
"is_vision": model_info.get("is_vision", False),
|
||||
"is_lora": model_info.get("is_lora", False),
|
||||
"display_name": model_info.get("display_name", model_name),
|
||||
"is_audio": model_info.get("is_audio", False),
|
||||
"audio_type": model_info.get("audio_type"),
|
||||
"has_audio_input": model_info.get("has_audio_input", False),
|
||||
}
|
||||
self.loading_models.discard(model_name)
|
||||
logger.info(
|
||||
"Model '%s' loaded successfully in subprocess", model_name
|
||||
)
|
||||
return True
|
||||
else:
|
||||
error = resp.get("error", "Failed to load model")
|
||||
self.loading_models.discard(model_name)
|
||||
self.active_model_name = None
|
||||
self.models.clear()
|
||||
raise Exception(error)
|
||||
|
||||
except Exception:
|
||||
self.loading_models.discard(model_name)
|
||||
|
|
@ -661,7 +721,7 @@ class InferenceOrchestrator:
|
|||
"model_name": model_name,
|
||||
}
|
||||
)
|
||||
resp = self._wait_response("unloaded", timeout = 30)
|
||||
resp = self._wait_response("unloaded")
|
||||
|
||||
# Update local state
|
||||
self.models.pop(model_name, None)
|
||||
|
|
|
|||
|
|
@ -57,16 +57,23 @@ WEB_SEARCH_TOOL = {
|
|||
"type": "function",
|
||||
"function": {
|
||||
"name": "web_search",
|
||||
"description": "Search the web for current information, recent events, or facts you are uncertain about.",
|
||||
"description": (
|
||||
"Search the web and fetch page content. Returns snippets for all results. "
|
||||
"Use the url parameter to fetch full page text from a specific URL."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "The search query",
|
||||
}
|
||||
},
|
||||
"url": {
|
||||
"type": "string",
|
||||
"description": "A URL to fetch full page content from (instead of searching). Use this to read a page found in search results.",
|
||||
},
|
||||
},
|
||||
"required": ["query"],
|
||||
"required": [],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
|
@ -131,7 +138,11 @@ def execute_tool(
|
|||
)
|
||||
effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
|
||||
if name == "web_search":
|
||||
return _web_search(arguments.get("query", ""), timeout = effective_timeout)
|
||||
return _web_search(
|
||||
arguments.get("query", ""),
|
||||
url = arguments.get("url"),
|
||||
timeout = effective_timeout,
|
||||
)
|
||||
if name == "python":
|
||||
return _python_exec(
|
||||
arguments.get("code", ""), cancel_event, effective_timeout, session_id
|
||||
|
|
@ -143,9 +154,180 @@ def execute_tool(
|
|||
return f"Unknown tool: {name}"
|
||||
|
||||
|
||||
def _web_search(query: str, max_results: int = 5, timeout: int = _EXEC_TIMEOUT) -> str:
|
||||
"""Search the web using DuckDuckGo and return formatted results."""
|
||||
if not query.strip():
|
||||
_MAX_PAGE_CHARS = 16000 # limit fetched page text
|
||||
_MAX_FETCH_BYTES = _MAX_PAGE_CHARS * 4 + 1 # cap raw download size
|
||||
|
||||
|
||||
def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str]:
|
||||
"""Resolve *hostname*, reject non-public IPs, return a pinned IP string.
|
||||
|
||||
Returns ``(ok, reason_or_empty, resolved_ip)``. The caller should
|
||||
connect to *resolved_ip* (with a ``Host`` header) to prevent DNS
|
||||
rebinding between validation and the actual fetch.
|
||||
"""
|
||||
import ipaddress
|
||||
import socket
|
||||
|
||||
try:
|
||||
infos = socket.getaddrinfo(hostname, port, type = socket.SOCK_STREAM)
|
||||
except OSError as e:
|
||||
return False, f"Failed to resolve host: {e}", ""
|
||||
|
||||
if not infos:
|
||||
return False, f"Failed to resolve host: no addresses for {hostname!r}", ""
|
||||
|
||||
for *_, sockaddr in infos:
|
||||
ip = ipaddress.ip_address(sockaddr[0])
|
||||
if (
|
||||
ip.is_private
|
||||
or ip.is_loopback
|
||||
or ip.is_link_local
|
||||
or ip.is_multicast
|
||||
or ip.is_reserved
|
||||
or ip.is_unspecified
|
||||
):
|
||||
return False, f"Blocked: refusing to fetch non-public address {ip}.", ""
|
||||
|
||||
# Return the first resolved address for pinning
|
||||
first_ip = infos[0][4][0]
|
||||
return True, "", first_ip
|
||||
|
||||
|
||||
def _fetch_page_text(
|
||||
url: str, max_chars: int = _MAX_PAGE_CHARS, timeout: int = 30
|
||||
) -> str:
|
||||
"""Fetch a URL and return plain text content (HTML tags stripped).
|
||||
|
||||
Blocks private/loopback/link-local targets (SSRF protection) and caps
|
||||
the download size to avoid unbounded memory usage.
|
||||
"""
|
||||
import re as _re
|
||||
from urllib.parse import urlparse
|
||||
|
||||
parsed = urlparse(url)
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
return f"Blocked: only http/https URLs are allowed (got {parsed.scheme!r})."
|
||||
if not parsed.hostname:
|
||||
return "Blocked: URL is missing a hostname."
|
||||
|
||||
port = parsed.port or (443 if parsed.scheme == "https" else 80)
|
||||
ok, reason, pinned_ip = _validate_and_resolve_host(parsed.hostname, port)
|
||||
if not ok:
|
||||
return reason
|
||||
|
||||
try:
|
||||
import urllib.request
|
||||
from urllib.error import HTTPError as _HTTPError
|
||||
from urllib.parse import urljoin, urlunparse
|
||||
|
||||
# Disable auto-redirect so we can validate each hop for SSRF.
|
||||
# urllib raises HTTPError for 3xx when the handler returns None,
|
||||
# so we catch that and extract the Location header manually.
|
||||
class _NoRedirect(urllib.request.HTTPRedirectHandler):
|
||||
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
||||
return None
|
||||
|
||||
opener = urllib.request.build_opener(_NoRedirect)
|
||||
max_bytes = max_chars * 4 + 1
|
||||
current_url = url
|
||||
current_host = parsed.hostname
|
||||
|
||||
for _hop in range(5):
|
||||
# Pin to the validated IP to prevent DNS rebinding.
|
||||
# Rewrite the URL to use the IP and set the Host header.
|
||||
cp = urlparse(current_url)
|
||||
ip_netloc = f"{pinned_ip}:{cp.port}" if cp.port else pinned_ip
|
||||
pinned_url = urlunparse(cp._replace(netloc = ip_netloc))
|
||||
|
||||
req = urllib.request.Request(
|
||||
pinned_url,
|
||||
headers = {
|
||||
"User-Agent": "UnslothStudio/1.0",
|
||||
"Host": current_host,
|
||||
},
|
||||
)
|
||||
try:
|
||||
resp = opener.open(req, timeout = timeout)
|
||||
except _HTTPError as e:
|
||||
if e.code not in (301, 302, 303, 307, 308):
|
||||
return (
|
||||
f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}"
|
||||
)
|
||||
location = e.headers.get("Location")
|
||||
if not location:
|
||||
return "Failed to fetch URL: redirect missing Location header."
|
||||
current_url = urljoin(current_url, location)
|
||||
rp = urlparse(current_url)
|
||||
if rp.scheme not in ("http", "https") or not rp.hostname:
|
||||
return "Blocked: redirect target is not a valid http/https URL."
|
||||
rp_port = rp.port or (443 if rp.scheme == "https" else 80)
|
||||
ok2, reason2, pinned_ip = _validate_and_resolve_host(
|
||||
rp.hostname,
|
||||
rp_port,
|
||||
)
|
||||
if not ok2:
|
||||
return reason2
|
||||
current_host = rp.hostname
|
||||
continue
|
||||
# Success -- read capped body
|
||||
raw_bytes = resp.read(max_bytes)
|
||||
break
|
||||
else:
|
||||
return "Failed to fetch URL: too many redirects."
|
||||
|
||||
charset = resp.headers.get_content_charset() or "utf-8"
|
||||
raw_html = raw_bytes.decode(charset, errors = "replace")
|
||||
except _HTTPError as e:
|
||||
return f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}"
|
||||
except Exception as e:
|
||||
return f"Failed to fetch URL: {e}"
|
||||
|
||||
# Convert HTML to text -- prefer html2text for clean markdown output
|
||||
try:
|
||||
import html2text as _h2t
|
||||
|
||||
converter = _h2t.HTML2Text()
|
||||
converter.ignore_links = False
|
||||
converter.ignore_images = True
|
||||
converter.body_width = 0 # no wrapping
|
||||
text = converter.handle(raw_html).strip()
|
||||
except ImportError:
|
||||
# Fallback: regex-based stripping
|
||||
text = _re.sub(
|
||||
r"<script[^>]*>.*?</script[^>]*>",
|
||||
"",
|
||||
raw_html,
|
||||
flags = _re.DOTALL | _re.IGNORECASE,
|
||||
)
|
||||
text = _re.sub(
|
||||
r"<style[^>]*>.*?</style[^>]*>", "", text, flags = _re.DOTALL | _re.IGNORECASE
|
||||
)
|
||||
text = _re.sub(r"<[^>]+>", " ", text)
|
||||
text = _re.sub(r"\s+", " ", text).strip()
|
||||
|
||||
if not text:
|
||||
return "(page returned no readable text)"
|
||||
if len(text) > max_chars:
|
||||
text = text[:max_chars] + f"\n\n... (truncated, {len(text)} chars total)"
|
||||
return text
|
||||
|
||||
|
||||
def _web_search(
|
||||
query: str,
|
||||
max_results: int = 5,
|
||||
timeout: int = _EXEC_TIMEOUT,
|
||||
url: str | None = None,
|
||||
) -> str:
|
||||
"""Search the web using DuckDuckGo and return formatted results.
|
||||
|
||||
If ``url`` is provided, fetches that page directly instead of searching.
|
||||
"""
|
||||
# Direct URL fetch mode
|
||||
if url and url.strip():
|
||||
fetch_timeout = 60 if timeout is None else min(timeout, 60)
|
||||
return _fetch_page_text(url.strip(), timeout = fetch_timeout)
|
||||
|
||||
if not query or not query.strip():
|
||||
return "No query provided."
|
||||
try:
|
||||
from ddgs import DDGS
|
||||
|
|
@ -160,7 +342,13 @@ def _web_search(query: str, max_results: int = 5, timeout: int = _EXEC_TIMEOUT)
|
|||
f"URL: {r.get('href', '')}\n"
|
||||
f"Snippet: {r.get('body', '')}"
|
||||
)
|
||||
return "\n\n---\n\n".join(parts)
|
||||
text = "\n\n---\n\n".join(parts)
|
||||
text += (
|
||||
"\n\n---\n\nIMPORTANT: These are only short snippets. "
|
||||
"To get the full page content, call web_search with "
|
||||
'the url parameter (e.g. {"url": "<URL>"}).'
|
||||
)
|
||||
return text
|
||||
except Exception as e:
|
||||
return f"Search failed: {e}"
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from loggers import get_logger
|
|||
import os
|
||||
import queue as _queue
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from io import BytesIO
|
||||
|
|
@ -29,6 +30,7 @@ from pathlib import Path
|
|||
from typing import Any
|
||||
|
||||
logger = get_logger(__name__)
|
||||
from utils.hardware import apply_gpu_ids
|
||||
|
||||
|
||||
def _activate_transformers_version(model_name: str) -> None:
|
||||
|
|
@ -113,6 +115,154 @@ def _build_model_config(config: dict):
|
|||
return mc
|
||||
|
||||
|
||||
def _get_hf_download_state(
|
||||
model_names: list[str] | None = None,
|
||||
) -> tuple[int, bool] | None:
|
||||
"""Return (total_bytes, has_incomplete) for the HF Hub cache, or None on error.
|
||||
|
||||
When *model_names* is provided, only those models' ``blobs/``
|
||||
directories are checked instead of scanning every cached model --
|
||||
much faster on systems with many models. Accepts multiple names so
|
||||
that LoRA loads can watch both the adapter repo and the base model
|
||||
repo simultaneously.
|
||||
|
||||
*has_incomplete* is True when any ``*.incomplete`` files exist in the
|
||||
watched blobs directories, indicating that ``huggingface_hub`` is
|
||||
actively downloading.
|
||||
|
||||
Returns None if the state cannot be determined (import error,
|
||||
permission error, etc.) so callers can skip stall logic.
|
||||
"""
|
||||
try:
|
||||
from huggingface_hub.constants import HF_HUB_CACHE
|
||||
|
||||
cache = Path(HF_HUB_CACHE)
|
||||
if not cache.exists():
|
||||
return (0, False)
|
||||
|
||||
total = 0
|
||||
has_incomplete = False
|
||||
blobs_dirs: list[Path] = []
|
||||
|
||||
if model_names:
|
||||
for name in model_names:
|
||||
if not name:
|
||||
continue
|
||||
# Skip local filesystem paths -- HF model IDs use forward
|
||||
# slashes (org/model) but never start with / . ~ or contain
|
||||
# backslashes. This distinguishes them from absolute paths,
|
||||
# relative paths, and Windows paths.
|
||||
if name.startswith(("/", ".", "~")) or "\\" in name:
|
||||
continue
|
||||
# HF cache dir format: models--org--name (slashes -> --)
|
||||
cache_dir_name = "models--" + name.replace("/", "--")
|
||||
blobs_dir = cache / cache_dir_name / "blobs"
|
||||
if blobs_dir.exists():
|
||||
blobs_dirs.append(blobs_dir)
|
||||
else:
|
||||
blobs_dirs = list(cache.glob("models--*/blobs"))
|
||||
|
||||
for bdir in blobs_dirs:
|
||||
for f in bdir.iterdir():
|
||||
try:
|
||||
if f.is_file():
|
||||
total += f.stat().st_size
|
||||
if f.name.endswith(".incomplete"):
|
||||
has_incomplete = True
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
return (total, has_incomplete)
|
||||
except Exception as e:
|
||||
logger.debug("Failed to determine HF download state: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def _start_heartbeat(
|
||||
resp_queue: Any,
|
||||
interval: float = 30.0,
|
||||
stall_timeout: float = 180.0,
|
||||
xet_disabled: bool = False,
|
||||
model_names: list[str] | None = None,
|
||||
) -> threading.Event:
|
||||
"""Start a daemon thread that sends periodic status heartbeats.
|
||||
|
||||
Monitors the HF Hub cache directory for download activity. A stall
|
||||
is only reported when ``*.incomplete`` files are present (indicating
|
||||
``huggingface_hub`` is actively downloading) **and** the total cache
|
||||
size has not changed for *stall_timeout* seconds.
|
||||
|
||||
Once the download finishes (no more ``.incomplete`` files), the stall
|
||||
timer resets, so post-download initialization (quantization, GPU
|
||||
weight loading) is never misclassified as a stalled download.
|
||||
|
||||
Returns a stop event -- set it to terminate the heartbeat thread.
|
||||
"""
|
||||
stop = threading.Event()
|
||||
transport = "https" if xet_disabled else "xet"
|
||||
|
||||
def _beat():
|
||||
state = _get_hf_download_state(model_names)
|
||||
last_size = state[0] if state is not None else 0
|
||||
last_change = time.monotonic()
|
||||
|
||||
while not stop.wait(interval):
|
||||
state = _get_hf_download_state(model_names)
|
||||
now = time.monotonic()
|
||||
|
||||
# Skip stall logic if we cannot measure the cache
|
||||
if state is None:
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{
|
||||
"type": "status",
|
||||
"message": f"Loading model ({transport} transport)...",
|
||||
"ts": time.time(),
|
||||
},
|
||||
)
|
||||
continue
|
||||
|
||||
current_size, has_incomplete = state
|
||||
|
||||
if current_size != last_size:
|
||||
last_size = current_size
|
||||
last_change = now
|
||||
|
||||
# Only fire stall when .incomplete files are present,
|
||||
# confirming a download is actively in progress.
|
||||
# Once downloads finish (no .incomplete), reset the timer
|
||||
# so model init time is not counted as a stall.
|
||||
if not has_incomplete:
|
||||
last_change = now
|
||||
elif now - last_change >= stall_timeout:
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{
|
||||
"type": "stall",
|
||||
"message": (
|
||||
f"Download appears stalled ({transport} transport) "
|
||||
f"-- no progress for {int(now - last_change)}s"
|
||||
),
|
||||
"ts": time.time(),
|
||||
},
|
||||
)
|
||||
# Only fire once -- the orchestrator will kill us
|
||||
return
|
||||
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{
|
||||
"type": "status",
|
||||
"message": f"Loading model ({transport} transport)...",
|
||||
"ts": time.time(),
|
||||
},
|
||||
)
|
||||
|
||||
t = threading.Thread(target = _beat, daemon = True)
|
||||
t.start()
|
||||
return stop
|
||||
|
||||
|
||||
def _handle_load(backend, config: dict, resp_queue: Any) -> None:
|
||||
"""Handle a load command: load a model into the backend."""
|
||||
try:
|
||||
|
|
@ -156,13 +306,50 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
|
|||
except Exception as e:
|
||||
logger.warning("Could not read adapter_config.json: %s", e)
|
||||
|
||||
success = backend.load_model(
|
||||
config = mc,
|
||||
max_seq_length = config.get("max_seq_length", 2048),
|
||||
load_in_4bit = load_in_4bit,
|
||||
hf_token = hf_token,
|
||||
trust_remote_code = config.get("trust_remote_code", False),
|
||||
# Auto-enable trust_remote_code for unsloth/* transformers 5.x models
|
||||
# (matches the training worker logic in core/training/worker.py)
|
||||
trust_remote_code = config.get("trust_remote_code", False)
|
||||
if not trust_remote_code:
|
||||
from utils.transformers_version import needs_transformers_5
|
||||
|
||||
model_name = config["model_name"]
|
||||
if needs_transformers_5(model_name) and model_name.lower().startswith(
|
||||
"unsloth/"
|
||||
):
|
||||
trust_remote_code = True
|
||||
logger.info(
|
||||
"Auto-enabled trust_remote_code for unsloth/* transformers 5.x model: %s",
|
||||
model_name,
|
||||
)
|
||||
|
||||
# Send heartbeats every 30s so the orchestrator knows we're still alive
|
||||
# (download / weight loading can take a long time on slow connections)
|
||||
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1"
|
||||
|
||||
# Watch both the model repo and base model repo (for LoRA loads
|
||||
# where the base model download is the actual bottleneck)
|
||||
watch_repos = [mc.identifier]
|
||||
base = getattr(mc, "base_model", None)
|
||||
if base and str(base) != mc.identifier:
|
||||
watch_repos.append(str(base))
|
||||
|
||||
heartbeat_stop = _start_heartbeat(
|
||||
resp_queue,
|
||||
interval = 30.0,
|
||||
xet_disabled = xet_disabled,
|
||||
model_names = watch_repos,
|
||||
)
|
||||
try:
|
||||
success = backend.load_model(
|
||||
config = mc,
|
||||
max_seq_length = config.get("max_seq_length", 2048),
|
||||
load_in_4bit = load_in_4bit,
|
||||
hf_token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
gpu_ids = config.get("resolved_gpu_ids"),
|
||||
)
|
||||
finally:
|
||||
heartbeat_stop.set()
|
||||
|
||||
if success:
|
||||
# Build model_info for the parent to mirror
|
||||
|
|
@ -474,6 +661,10 @@ def run_inference_process(
|
|||
"ignore" # Suppress warnings at C-level before imports
|
||||
)
|
||||
|
||||
if config.get("disable_xet"):
|
||||
os.environ["HF_HUB_DISABLE_XET"] = "1"
|
||||
logger.info("Xet transport disabled (HF_HUB_DISABLE_XET=1)")
|
||||
|
||||
import warnings
|
||||
from loggers.config import LogConfig
|
||||
|
||||
|
|
@ -485,6 +676,8 @@ def run_inference_process(
|
|||
env = os.getenv("ENVIRONMENT_TYPE", "production"),
|
||||
)
|
||||
|
||||
apply_gpu_ids(config.get("resolved_gpu_ids"))
|
||||
|
||||
model_name = config["model_name"]
|
||||
|
||||
# ── 1. Activate correct transformers version BEFORE any ML imports ──
|
||||
|
|
|
|||
|
|
@ -33,7 +33,14 @@ if sys.platform in ("win32", "darwin"):
|
|||
sys.path.insert(0, _compile_cache)
|
||||
|
||||
import torch
|
||||
from utils.hardware import clear_gpu_cache, safe_num_proc, dataset_map_num_proc
|
||||
from utils.hardware import (
|
||||
clear_gpu_cache,
|
||||
safe_num_proc,
|
||||
dataset_map_num_proc,
|
||||
get_device_map,
|
||||
raise_if_offloaded,
|
||||
get_visible_gpu_count,
|
||||
)
|
||||
|
||||
torch._dynamo.config.recompile_limit = 64
|
||||
from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
|
||||
|
|
@ -487,6 +494,7 @@ class UnslothTrainer:
|
|||
is_dataset_audio: bool = False,
|
||||
trust_remote_code: bool = False,
|
||||
full_finetuning: bool = False,
|
||||
gpu_ids: Optional[list[int]] = None,
|
||||
) -> bool:
|
||||
"""Load model for training (supports both text and vision models)"""
|
||||
self.load_in_4bit = load_in_4bit # Store for training_meta.json
|
||||
|
|
@ -624,6 +632,11 @@ class UnslothTrainer:
|
|||
self._update_progress(error = friendly, is_training = False)
|
||||
return False
|
||||
|
||||
device_map = get_device_map(gpu_ids)
|
||||
logger.info(
|
||||
f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)"
|
||||
)
|
||||
|
||||
# Branch based on model type
|
||||
if self._audio_type == "csm":
|
||||
# CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False
|
||||
|
|
@ -636,6 +649,7 @@ class UnslothTrainer:
|
|||
dtype = None,
|
||||
auto_model = CsmForConditionalGeneration,
|
||||
load_in_4bit = False,
|
||||
device_map = device_map,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
|
|
@ -651,6 +665,7 @@ class UnslothTrainer:
|
|||
model_name = model_name,
|
||||
dtype = None,
|
||||
load_in_4bit = False,
|
||||
device_map = device_map,
|
||||
full_finetuning = full_finetuning,
|
||||
auto_model = WhisperForConditionalGeneration,
|
||||
whisper_language = "English",
|
||||
|
|
@ -672,6 +687,7 @@ class UnslothTrainer:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = None,
|
||||
load_in_4bit = load_in_4bit,
|
||||
device_map = device_map,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
|
|
@ -711,6 +727,7 @@ class UnslothTrainer:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = torch.float32, # Spark-TTS requires float32
|
||||
load_in_4bit = False,
|
||||
device_map = device_map,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
|
|
@ -725,6 +742,7 @@ class UnslothTrainer:
|
|||
model_name,
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = False,
|
||||
device_map = device_map,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
|
|
@ -741,6 +759,7 @@ class UnslothTrainer:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = None,
|
||||
load_in_4bit = load_in_4bit,
|
||||
device_map = device_map,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
|
|
@ -754,6 +773,7 @@ class UnslothTrainer:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = None, # Auto-detect
|
||||
load_in_4bit = load_in_4bit,
|
||||
device_map = device_map,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
|
|
@ -786,12 +806,15 @@ class UnslothTrainer:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = None, # Auto-detect
|
||||
load_in_4bit = load_in_4bit,
|
||||
device_map = device_map,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
logger.info("Loaded text model")
|
||||
|
||||
raise_if_offloaded(self.model, device_map, "Studio training")
|
||||
|
||||
if self.should_stop:
|
||||
return False
|
||||
|
||||
|
|
@ -824,6 +847,7 @@ class UnslothTrainer:
|
|||
is_dataset_audio = is_dataset_audio,
|
||||
trust_remote_code = trust_remote_code,
|
||||
full_finetuning = full_finetuning,
|
||||
gpu_ids = gpu_ids,
|
||||
)
|
||||
error_msg = str(e)
|
||||
error_lower = error_msg.lower()
|
||||
|
|
@ -2634,14 +2658,14 @@ class UnslothTrainer:
|
|||
eval_steps: float = 0.00,
|
||||
output_dir: str | None = None,
|
||||
num_epochs: int = 3,
|
||||
learning_rate: float = 5e-5,
|
||||
learning_rate: float = 2e-4,
|
||||
batch_size: int = 2,
|
||||
gradient_accumulation_steps: int = 4,
|
||||
warmup_steps: int = None,
|
||||
warmup_ratio: float = None,
|
||||
max_steps: int = 0,
|
||||
save_steps: int = 0,
|
||||
weight_decay: float = 0.01,
|
||||
weight_decay: float = 0.001,
|
||||
random_seed: int = 3407,
|
||||
packing: bool = False,
|
||||
train_on_completions: bool = False,
|
||||
|
|
@ -3010,7 +3034,7 @@ class UnslothTrainer:
|
|||
"fp16": not is_bfloat16_supported(),
|
||||
"bf16": is_bfloat16_supported(),
|
||||
"logging_steps": 1,
|
||||
"weight_decay": training_args.get("weight_decay", 0.01),
|
||||
"weight_decay": training_args.get("weight_decay", 0.001),
|
||||
"seed": training_args.get("random_seed", 3407),
|
||||
"output_dir": output_dir,
|
||||
"report_to": _build_report_targets(training_args),
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from pathlib import Path
|
|||
from typing import Optional, Tuple, Any
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
from utils.hardware import prepare_gpu_selection
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -159,7 +160,7 @@ class TrainingBackend:
|
|||
"warmup_ratio": kwargs.get("warmup_ratio"),
|
||||
"max_steps": kwargs.get("max_steps", 0),
|
||||
"save_steps": kwargs.get("save_steps", 0),
|
||||
"weight_decay": kwargs.get("weight_decay", 0.01),
|
||||
"weight_decay": kwargs.get("weight_decay", 0.001),
|
||||
"random_seed": kwargs.get("random_seed", 3407),
|
||||
"packing": kwargs.get("packing", False),
|
||||
"optim": kwargs.get("optim", "adamw_8bit"),
|
||||
|
|
@ -185,6 +186,7 @@ class TrainingBackend:
|
|||
"enable_tensorboard": kwargs.get("enable_tensorboard", False),
|
||||
"tensorboard_dir": kwargs.get("tensorboard_dir", "runs"),
|
||||
"trust_remote_code": kwargs.get("trust_remote_code", False),
|
||||
"gpu_ids": kwargs.get("gpu_ids"),
|
||||
}
|
||||
|
||||
# Derive load_in_4bit from training_type
|
||||
|
|
@ -192,6 +194,22 @@ class TrainingBackend:
|
|||
config["load_in_4bit"] = False
|
||||
|
||||
# Spawn subprocess — use locals so state is untouched on failure
|
||||
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
|
||||
kwargs.get("gpu_ids"),
|
||||
model_name = config["model_name"],
|
||||
hf_token = config["hf_token"] or None,
|
||||
training_type = config["training_type"],
|
||||
load_in_4bit = config["load_in_4bit"],
|
||||
batch_size = config.get("batch_size", 4),
|
||||
max_seq_length = config.get("max_seq_length", 2048),
|
||||
lora_rank = config.get("lora_r", 16),
|
||||
target_modules = config.get("target_modules"),
|
||||
gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
|
||||
optimizer = config.get("optim", "adamw_8bit"),
|
||||
)
|
||||
config["resolved_gpu_ids"] = resolved_gpu_ids
|
||||
config["gpu_selection"] = gpu_selection
|
||||
|
||||
from .worker import run_training_process
|
||||
|
||||
event_queue = _CTX.Queue()
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ import urllib.error
|
|||
import urllib.request
|
||||
|
||||
logger = get_logger(__name__)
|
||||
from utils.hardware import apply_gpu_ids
|
||||
|
||||
|
||||
_CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4"
|
||||
|
|
@ -367,6 +368,8 @@ def run_training_process(
|
|||
env = os.getenv("ENVIRONMENT_TYPE", "production"),
|
||||
)
|
||||
|
||||
apply_gpu_ids(config.get("resolved_gpu_ids"))
|
||||
|
||||
model_name = config["model_name"]
|
||||
|
||||
# ── 1. Activate correct transformers version BEFORE any ML imports ──
|
||||
|
|
@ -682,6 +685,7 @@ def run_training_process(
|
|||
is_dataset_image = config.get("is_dataset_image", False),
|
||||
is_dataset_audio = config.get("is_dataset_audio", False),
|
||||
trust_remote_code = config.get("trust_remote_code", False),
|
||||
gpu_ids = config.get("resolved_gpu_ids"),
|
||||
)
|
||||
if not success or trainer.should_stop:
|
||||
if trainer.should_stop:
|
||||
|
|
@ -791,7 +795,7 @@ def run_training_process(
|
|||
warmup_ratio = config.get("warmup_ratio"),
|
||||
max_steps = max_steps if max_steps and max_steps > 0 else 0,
|
||||
save_steps = save_steps if save_steps and save_steps > 0 else 0,
|
||||
weight_decay = config.get("weight_decay", 0.01),
|
||||
weight_decay = config.get("weight_decay", 0.001),
|
||||
random_seed = config.get("random_seed", 3407),
|
||||
packing = config.get("packing", False),
|
||||
train_on_completions = config.get("train_on_completions", False),
|
||||
|
|
@ -1137,7 +1141,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
|
|||
"lr_scheduler_type": config.get("lr_scheduler_type", "linear"),
|
||||
"batch_sampler": BatchSamplers.NO_DUPLICATES,
|
||||
"optim": config.get("optim", "adamw_8bit"),
|
||||
"weight_decay": config.get("weight_decay", 0.01),
|
||||
"weight_decay": config.get("weight_decay", 0.001),
|
||||
"seed": config.get("random_seed", 3407),
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -23,10 +23,23 @@ if _backend_dir not in sys.path:
|
|||
# See: https://github.com/python/cpython/issues/102396
|
||||
import _platform_compat # noqa: F401
|
||||
|
||||
import mimetypes
|
||||
import shutil
|
||||
import warnings
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
# Fix broken Windows registry MIME types. Some Windows installs map .js to
|
||||
# "text/plain" in the registry (HKCR\.js\Content Type). Python's mimetypes
|
||||
# module reads from the registry, and FastAPI/Starlette's StaticFiles uses
|
||||
# mimetypes.guess_type() to set Content-Type headers. Browsers enforce strict
|
||||
# MIME checking for ES module scripts (<script type="module">) and will refuse
|
||||
# to execute .js files served as text/plain — resulting in a blank page.
|
||||
# Calling add_type() *before* StaticFiles is instantiated ensures the correct
|
||||
# types are used regardless of the OS registry.
|
||||
if sys.platform == "win32":
|
||||
mimetypes.add_type("application/javascript", ".js")
|
||||
mimetypes.add_type("text/css", ".css")
|
||||
|
||||
# Suppress annoying dependency warnings in production
|
||||
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
|
||||
warnings.filterwarnings("ignore")
|
||||
|
|
@ -34,7 +47,7 @@ if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
|
|||
# warnings.filterwarnings("ignore", category=DeprecationWarning)
|
||||
# warnings.filterwarnings("ignore", module="triton.*")
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi import Depends, FastAPI, Request
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.responses import FileResponse, HTMLResponse, Response
|
||||
|
|
@ -53,7 +66,13 @@ from routes import (
|
|||
training_router,
|
||||
)
|
||||
from auth import storage
|
||||
from utils.hardware import detect_hardware, get_device, DeviceType
|
||||
from auth.authentication import get_current_subject
|
||||
from utils.hardware import (
|
||||
detect_hardware,
|
||||
get_device,
|
||||
DeviceType,
|
||||
get_backend_visible_gpu_info,
|
||||
)
|
||||
import utils.hardware.hardware as _hw_module
|
||||
|
||||
from utils.cache_cleanup import clear_unsloth_compiled_cache
|
||||
|
|
@ -184,73 +203,46 @@ async def health_check():
|
|||
}
|
||||
|
||||
|
||||
@app.post("/api/shutdown")
|
||||
async def shutdown_server(
|
||||
request: Request,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Gracefully shut down the Unsloth Studio server.
|
||||
|
||||
Called by the frontend quit dialog so users can stop the server from the UI
|
||||
without needing to use the CLI or kill the process manually.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
async def _delayed_shutdown():
|
||||
await asyncio.sleep(0.2) # Let the HTTP response return first
|
||||
trigger = getattr(request.app.state, "trigger_shutdown", None)
|
||||
if trigger is not None:
|
||||
trigger()
|
||||
else:
|
||||
# Fallback when not launched via run_server() (e.g. direct uvicorn)
|
||||
import signal
|
||||
import os
|
||||
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
|
||||
request.app.state._shutdown_task = asyncio.create_task(_delayed_shutdown())
|
||||
return {"status": "shutting_down"}
|
||||
|
||||
|
||||
@app.get("/api/system")
|
||||
async def get_system_info():
|
||||
"""Get system information"""
|
||||
import platform
|
||||
import subprocess
|
||||
import psutil
|
||||
from utils.hardware import get_device, get_gpu_memory_info, DeviceType
|
||||
from utils.hardware import get_device
|
||||
|
||||
# GPU Info — query nvidia-smi for physical GPUs, filtered by
|
||||
# CUDA_VISIBLE_DEVICES when set (the frontend uses this for GGUF
|
||||
# fit estimation and llama-server respects CVD too).
|
||||
import os
|
||||
|
||||
gpu_info: dict = {"available": False, "devices": []}
|
||||
|
||||
device = get_device()
|
||||
if device == DeviceType.CUDA:
|
||||
# Parse CUDA_VISIBLE_DEVICES allowlist
|
||||
allowed_indices = None
|
||||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||||
if cvd is not None and cvd.strip():
|
||||
try:
|
||||
allowed_indices = set(int(x.strip()) for x in cvd.split(","))
|
||||
except ValueError:
|
||||
pass # Non-numeric (e.g. GPU-uuid), show all
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"nvidia-smi",
|
||||
"--query-gpu=index,name,memory.total",
|
||||
"--format=csv,noheader,nounits",
|
||||
],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 10,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
for line in result.stdout.strip().splitlines():
|
||||
parts = [p.strip() for p in line.split(",")]
|
||||
if len(parts) == 3:
|
||||
idx = int(parts[0])
|
||||
if allowed_indices is not None and idx not in allowed_indices:
|
||||
continue
|
||||
gpu_info["devices"].append(
|
||||
{
|
||||
"index": idx,
|
||||
"name": parts[1],
|
||||
"memory_total_gb": round(int(parts[2]) / 1024, 2),
|
||||
}
|
||||
)
|
||||
gpu_info["available"] = len(gpu_info["devices"]) > 0
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Fallback to torch-based single-GPU detection
|
||||
if not gpu_info["available"]:
|
||||
mem_info = get_gpu_memory_info()
|
||||
if mem_info.get("available"):
|
||||
gpu_info["available"] = True
|
||||
gpu_info["devices"].append(
|
||||
{
|
||||
"index": mem_info.get("device", 0),
|
||||
"name": mem_info.get("device_name", "Unknown"),
|
||||
"memory_total_gb": round(mem_info.get("total_gb", 0), 2),
|
||||
}
|
||||
)
|
||||
visibility_info = get_backend_visible_gpu_info()
|
||||
gpu_info = {
|
||||
"available": visibility_info["available"],
|
||||
"devices": visibility_info["devices"],
|
||||
}
|
||||
|
||||
# CPU & Memory
|
||||
memory = psutil.virtual_memory()
|
||||
|
|
@ -269,6 +261,13 @@ async def get_system_info():
|
|||
}
|
||||
|
||||
|
||||
@app.get("/api/system/gpu-visibility")
|
||||
async def get_gpu_visibility(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
return get_backend_visible_gpu_info()
|
||||
|
||||
|
||||
@app.get("/api/system/hardware")
|
||||
async def get_hardware_info():
|
||||
"""Return GPU name, total VRAM, and key ML package versions."""
|
||||
|
|
|
|||
|
|
@ -22,7 +22,10 @@ class LoadRequest(BaseModel):
|
|||
None, description = "HuggingFace token for gated models"
|
||||
)
|
||||
max_seq_length: int = Field(
|
||||
4096, ge = 128, le = 32768, description = "Maximum sequence length"
|
||||
0,
|
||||
ge = 0,
|
||||
le = 1048576,
|
||||
description = "Maximum sequence length (0 = model default for GGUF)",
|
||||
)
|
||||
load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization")
|
||||
is_lora: bool = Field(False, description = "Whether this is a LoRA adapter")
|
||||
|
|
@ -41,6 +44,10 @@ class LoadRequest(BaseModel):
|
|||
None,
|
||||
description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')",
|
||||
)
|
||||
gpu_ids: Optional[List[int]] = Field(
|
||||
None,
|
||||
description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.",
|
||||
)
|
||||
|
||||
|
||||
class UnloadRequest(BaseModel):
|
||||
|
|
@ -129,10 +136,17 @@ class LoadResponse(BaseModel):
|
|||
context_length: Optional[int] = Field(
|
||||
None, description = "Model's native context length (from GGUF metadata)"
|
||||
)
|
||||
max_context_length: Optional[int] = Field(
|
||||
None, description = "Maximum context length currently available on this hardware"
|
||||
)
|
||||
supports_reasoning: bool = Field(
|
||||
False,
|
||||
description = "Whether model supports thinking/reasoning mode (enable_thinking)",
|
||||
)
|
||||
reasoning_always_on: bool = Field(
|
||||
False,
|
||||
description = "Whether reasoning is always on (hardcoded <think> tags, not toggleable)",
|
||||
)
|
||||
supports_tools: bool = Field(
|
||||
False,
|
||||
description = "Whether model supports tool calling (web search, etc.)",
|
||||
|
|
@ -190,12 +204,19 @@ class InferenceStatusResponse(BaseModel):
|
|||
supports_reasoning: bool = Field(
|
||||
False, description = "Whether the active model supports reasoning/thinking mode"
|
||||
)
|
||||
reasoning_always_on: bool = Field(
|
||||
False, description = "Whether reasoning is always on (not toggleable)"
|
||||
)
|
||||
supports_tools: bool = Field(
|
||||
False, description = "Whether the active model supports tool calling"
|
||||
)
|
||||
context_length: Optional[int] = Field(
|
||||
None, description = "Context length of the active model"
|
||||
)
|
||||
max_context_length: Optional[int] = Field(
|
||||
None,
|
||||
description = "Maximum context length currently available for the active model",
|
||||
)
|
||||
|
||||
|
||||
# =====================================================================
|
||||
|
|
@ -288,7 +309,7 @@ class ChatCompletionRequest(BaseModel):
|
|||
0.01, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold"
|
||||
)
|
||||
repetition_penalty: float = Field(
|
||||
1.1, ge = 1.0, le = 2.0, description = "[x-unsloth] Repetition penalty"
|
||||
1.0, ge = 1.0, le = 2.0, description = "[x-unsloth] Repetition penalty"
|
||||
)
|
||||
image_base64: Optional[str] = Field(
|
||||
None, description = "[x-unsloth] Base64-encoded image for vision models"
|
||||
|
|
@ -323,7 +344,7 @@ class ChatCompletionRequest(BaseModel):
|
|||
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
|
||||
)
|
||||
max_tool_calls_per_message: Optional[int] = Field(
|
||||
10,
|
||||
25,
|
||||
ge = 0,
|
||||
description = "[x-unsloth] Maximum number of tool call iterations per message (0 = disabled, 9999 = unlimited).",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -165,7 +165,7 @@ class LocalModelInfo(BaseModel):
|
|||
id: str = Field(..., description = "Identifier to use for loading/training")
|
||||
display_name: str = Field(..., description = "Display label")
|
||||
path: str = Field(..., description = "Local path where model data was discovered")
|
||||
source: Literal["models_dir", "hf_cache"] = Field(
|
||||
source: Literal["models_dir", "hf_cache", "lmstudio"] = Field(
|
||||
...,
|
||||
description = "Discovery source",
|
||||
)
|
||||
|
|
@ -189,6 +189,10 @@ class LocalModelListResponse(BaseModel):
|
|||
None,
|
||||
description = "HF cache root that was scanned",
|
||||
)
|
||||
lmstudio_dirs: List[str] = Field(
|
||||
default_factory = list,
|
||||
description = "LM Studio model directories that were scanned",
|
||||
)
|
||||
models: List[LocalModelInfo] = Field(
|
||||
default_factory = list,
|
||||
description = "Discovered local/cached models",
|
||||
|
|
|
|||
|
|
@ -81,7 +81,7 @@ class TrainingStartRequest(BaseModel):
|
|||
warmup_ratio: Optional[float] = Field(None, description = "Warmup ratio")
|
||||
max_steps: Optional[int] = Field(None, description = "Maximum training steps")
|
||||
save_steps: int = Field(100, description = "Steps between checkpoints")
|
||||
weight_decay: float = Field(0.01, description = "Weight decay")
|
||||
weight_decay: float = Field(0.001, description = "Weight decay")
|
||||
random_seed: int = Field(42, description = "Random seed")
|
||||
packing: bool = Field(False, description = "Enable sequence packing")
|
||||
optim: str = Field("adamw_8bit", description = "Optimizer")
|
||||
|
|
@ -128,6 +128,12 @@ class TrainingStartRequest(BaseModel):
|
|||
enable_tensorboard: bool = Field(False, description = "Enable TensorBoard logging")
|
||||
tensorboard_dir: Optional[str] = Field(None, description = "TensorBoard directory")
|
||||
|
||||
# GPU selection
|
||||
gpu_ids: Optional[List[int]] = Field(
|
||||
None,
|
||||
description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.",
|
||||
)
|
||||
|
||||
|
||||
class TrainingJobResponse(BaseModel):
|
||||
"""Immediate response when training is initiated"""
|
||||
|
|
|
|||
|
|
@ -12,3 +12,5 @@ git+https://github.com/meta-pytorch/OpenEnv.git
|
|||
torch-c-dlpack-ext
|
||||
sentence_transformers==5.2.0
|
||||
transformers==4.57.6
|
||||
pytorch_tokenizers
|
||||
kernels
|
||||
|
|
|
|||
35
studio/backend/requirements/no-torch-runtime.txt
Normal file
35
studio/backend/requirements/no-torch-runtime.txt
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
# Runtime dependencies for no-torch (GGUF-only) mode.
|
||||
# Installed with --no-deps to prevent transitive torch resolution
|
||||
# from packages like accelerate, peft, trl, sentence-transformers.
|
||||
#
|
||||
# Includes unsloth's own direct deps (typer, pydantic, pyyaml,
|
||||
# nest-asyncio) since unsloth is also installed with --no-deps
|
||||
# (current PyPI metadata still declares torch as a hard dep).
|
||||
|
||||
# unsloth direct deps (from pyproject.toml [project].dependencies)
|
||||
typer
|
||||
pydantic
|
||||
pyyaml
|
||||
nest-asyncio
|
||||
|
||||
# HF ecosystem (from [huggingfacenotorch] extras in pyproject.toml)
|
||||
wheel>=0.42.0
|
||||
packaging
|
||||
numpy
|
||||
tqdm
|
||||
psutil
|
||||
tyro
|
||||
protobuf
|
||||
sentencepiece>=0.2.0
|
||||
safetensors>=0.4.3
|
||||
datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0
|
||||
accelerate>=0.34.1
|
||||
peft>=0.18.0,!=0.11.0
|
||||
huggingface_hub>=0.34.0
|
||||
hf_transfer
|
||||
diffusers
|
||||
transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.3.0
|
||||
trl>=0.18.2,!=0.19.0,<=0.24.0
|
||||
sentence-transformers
|
||||
cut_cross_entropy
|
||||
pillow
|
||||
|
|
@ -1,6 +1,2 @@
|
|||
# Torch AO overrides (installed with --force-reinstall --no-cache-dir)
|
||||
torchao==0.14.0
|
||||
pytorch_tokenizers
|
||||
|
||||
# Kernel packages
|
||||
kernels
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ lxml<7,>=6.0.2
|
|||
marko<3,>=2.1.2
|
||||
mcp<2,>=1.26.0
|
||||
networkx<4,>=3.0
|
||||
python-json-logger>=3,<4
|
||||
ruff<1,>=0.14.10
|
||||
scipy<2,>=1.11.0
|
||||
sqlfluff<4,>=3.2.0
|
||||
|
|
|
|||
|
|
@ -388,7 +388,7 @@ def _extract_text_from_file(file_path: Path, ext: str) -> str:
|
|||
import pymupdf4llm
|
||||
|
||||
raw = pymupdf4llm.to_markdown(
|
||||
str(file_path), write_images = False, show_progress = False
|
||||
str(file_path), write_images = False, show_progress = False, use_ocr = False
|
||||
)
|
||||
elif ext == ".docx":
|
||||
import mammoth
|
||||
|
|
|
|||
|
|
@ -19,6 +19,27 @@ import asyncio
|
|||
import threading
|
||||
|
||||
|
||||
import re as _re
|
||||
|
||||
|
||||
def _friendly_error(exc: Exception) -> str:
|
||||
"""Extract a user-friendly message from known llama-server errors."""
|
||||
msg = str(exc)
|
||||
m = _re.search(
|
||||
r"request \((\d+) tokens?\) exceeds the available context size \((\d+) tokens?\)",
|
||||
msg,
|
||||
)
|
||||
if m:
|
||||
return (
|
||||
f"Message too long: {m.group(1)} tokens exceeds the {m.group(2)}-token "
|
||||
f"context window. Try increasing the Context Length in Model settings, "
|
||||
f"or shorten the conversation."
|
||||
)
|
||||
if "Lost connection to llama-server" in msg:
|
||||
return "Lost connection to the model server. It may have crashed -- try reloading the model."
|
||||
return "An internal error occurred"
|
||||
|
||||
|
||||
# Add backend directory to path
|
||||
backend_path = Path(__file__).parent.parent.parent
|
||||
if str(backend_path) not in sys.path:
|
||||
|
|
@ -65,8 +86,15 @@ import io
|
|||
import wave
|
||||
import base64
|
||||
import numpy as np
|
||||
from datetime import date as _date
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# Regex for stripping leaked tool-call XML from assistant messages/stream
|
||||
_TOOL_XML_RE = _re.compile(
|
||||
r"<tool_call>.*?</tool_call>|<function=\w+>.*?</function>",
|
||||
_re.DOTALL,
|
||||
)
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
|
|
@ -134,7 +162,9 @@ async def load_model(
|
|||
else False,
|
||||
inference = inference_config,
|
||||
context_length = llama_backend.context_length,
|
||||
max_context_length = llama_backend.max_context_length,
|
||||
supports_reasoning = llama_backend.supports_reasoning,
|
||||
reasoning_always_on = llama_backend.reasoning_always_on,
|
||||
chat_template = llama_backend.chat_template,
|
||||
)
|
||||
else:
|
||||
|
|
@ -183,8 +213,17 @@ async def load_model(
|
|||
detail = f"Invalid model identifier: {request.model_path}",
|
||||
)
|
||||
|
||||
# Normalize gpu_ids: empty list means auto-selection, same as None
|
||||
effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
|
||||
|
||||
# ── GGUF path: load via llama-server ──────────────────────
|
||||
if config.is_gguf:
|
||||
if effective_gpu_ids is not None:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "gpu_ids is not supported for GGUF models yet.",
|
||||
)
|
||||
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
unsloth_backend = get_inference_backend()
|
||||
|
||||
|
|
@ -258,7 +297,9 @@ async def load_model(
|
|||
has_audio_input = is_audio_input_type(_gguf_audio),
|
||||
inference = inference_config,
|
||||
context_length = llama_backend.context_length,
|
||||
max_context_length = llama_backend.max_context_length,
|
||||
supports_reasoning = llama_backend.supports_reasoning,
|
||||
reasoning_always_on = llama_backend.reasoning_always_on,
|
||||
supports_tools = llama_backend.supports_tools,
|
||||
cache_type_kv = llama_backend.cache_type_kv,
|
||||
chat_template = llama_backend.chat_template,
|
||||
|
|
@ -344,6 +385,7 @@ async def load_model(
|
|||
load_in_4bit = load_in_4bit,
|
||||
hf_token = request.hf_token,
|
||||
trust_remote_code = request.trust_remote_code,
|
||||
gpu_ids = effective_gpu_ids,
|
||||
)
|
||||
|
||||
if not success:
|
||||
|
|
@ -395,6 +437,9 @@ async def load_model(
|
|||
|
||||
except HTTPException:
|
||||
raise
|
||||
except ValueError as e:
|
||||
logger.warning("Rejected inference GPU selection: %s", e)
|
||||
raise HTTPException(status_code = 400, detail = str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading model: {e}", exc_info = True)
|
||||
msg = str(e)
|
||||
|
|
@ -550,7 +595,7 @@ async def generate_stream(
|
|||
except Exception as e:
|
||||
backend.reset_generation_state()
|
||||
logger.error(f"Error during generation: {e}", exc_info = True)
|
||||
yield f"data: {json.dumps({'error': 'An internal error occurred'})}\n\n"
|
||||
yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n"
|
||||
|
||||
return StreamingResponse(
|
||||
stream(),
|
||||
|
|
@ -588,8 +633,10 @@ async def get_status(
|
|||
loaded = [_model_id],
|
||||
inference = _inference_cfg,
|
||||
supports_reasoning = llama_backend.supports_reasoning,
|
||||
reasoning_always_on = llama_backend.reasoning_always_on,
|
||||
supports_tools = llama_backend.supports_tools,
|
||||
context_length = llama_backend.context_length,
|
||||
max_context_length = llama_backend.max_context_length,
|
||||
)
|
||||
|
||||
# Otherwise, report Unsloth backend status
|
||||
|
|
@ -944,7 +991,7 @@ async def openai_chat_completions(
|
|||
logger.error(
|
||||
f"Error during audio input streaming: {e}", exc_info = True
|
||||
)
|
||||
yield f"data: {json.dumps({'error': {'message': 'An internal error occurred', 'type': 'server_error'}})}\n\n"
|
||||
yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
|
||||
|
||||
return StreamingResponse(
|
||||
audio_input_stream(),
|
||||
|
|
@ -1038,6 +1085,68 @@ async def openai_chat_completions(
|
|||
else:
|
||||
tools_to_use = ALL_TOOLS
|
||||
|
||||
# ── Tool-use system prompt nudge ──────────────────────
|
||||
_tool_names = {t["function"]["name"] for t in tools_to_use}
|
||||
_has_web = "web_search" in _tool_names
|
||||
_has_code = "python" in _tool_names or "terminal" in _tool_names
|
||||
|
||||
_date_line = f"The current date is {_date.today().isoformat()}."
|
||||
|
||||
_web_tips = (
|
||||
"When you search and find a relevant URL in the results, "
|
||||
"fetch its full content by calling web_search with the url parameter. "
|
||||
"Do not repeat the same search query. If a search returns "
|
||||
"no useful results, try rephrasing or fetching a result URL directly."
|
||||
)
|
||||
_code_tips = (
|
||||
"Use code execution for math, calculations, data processing, "
|
||||
"or to parse and analyze information from tool results."
|
||||
)
|
||||
|
||||
if _has_web and _has_code:
|
||||
_nudge = (
|
||||
_date_line + " "
|
||||
"You have access to tools. When appropriate, prefer using "
|
||||
"tools rather than answering from memory. "
|
||||
+ _web_tips
|
||||
+ " "
|
||||
+ _code_tips
|
||||
)
|
||||
elif _has_code:
|
||||
_nudge = (
|
||||
_date_line + " "
|
||||
"You have access to tools. When appropriate, prefer using "
|
||||
"code execution rather than answering from memory. " + _code_tips
|
||||
)
|
||||
elif _has_web:
|
||||
_nudge = (
|
||||
_date_line + " "
|
||||
"You have access to tools. When appropriate, prefer using "
|
||||
"web search for up-to-date or uncertain factual "
|
||||
"information rather than answering from memory. " + _web_tips
|
||||
)
|
||||
else:
|
||||
_nudge = ""
|
||||
|
||||
if _nudge:
|
||||
# Append nudge to system prompt (preserve user's prompt)
|
||||
if system_prompt:
|
||||
system_prompt = system_prompt.rstrip() + "\n\n" + _nudge
|
||||
else:
|
||||
system_prompt = _nudge
|
||||
# Rebuild gguf_messages with updated system prompt
|
||||
gguf_messages = []
|
||||
if system_prompt:
|
||||
gguf_messages.append({"role": "system", "content": system_prompt})
|
||||
gguf_messages.extend(chat_messages)
|
||||
|
||||
# ── Strip stale tool-call XML from conversation history ─
|
||||
for _msg in gguf_messages:
|
||||
if _msg.get("role") == "assistant" and isinstance(
|
||||
_msg.get("content"), str
|
||||
):
|
||||
_msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip()
|
||||
|
||||
def gguf_generate_with_tools():
|
||||
return llama_backend.generate_chat_completion_with_tools(
|
||||
messages = gguf_messages,
|
||||
|
|
@ -1056,7 +1165,7 @@ async def openai_chat_completions(
|
|||
else True,
|
||||
max_tool_iterations = payload.max_tool_calls_per_message
|
||||
if payload.max_tool_calls_per_message is not None
|
||||
else 10,
|
||||
else 25,
|
||||
tool_call_timeout = payload.tool_call_timeout
|
||||
if payload.tool_call_timeout is not None
|
||||
else 300,
|
||||
|
|
@ -1107,6 +1216,8 @@ async def openai_chat_completions(
|
|||
continue
|
||||
|
||||
if event["type"] in ("tool_start", "tool_end"):
|
||||
if event["type"] == "tool_start":
|
||||
prev_text = ""
|
||||
yield f"data: {json.dumps(event)}\n\n"
|
||||
continue
|
||||
|
||||
|
|
@ -1116,9 +1227,13 @@ async def openai_chat_completions(
|
|||
continue
|
||||
|
||||
# "content" type -- cumulative text
|
||||
cumulative = event.get("text", "")
|
||||
new_text = cumulative[len(prev_text) :]
|
||||
prev_text = cumulative
|
||||
# Sanitize the full cumulative then diff against
|
||||
# the last sanitized snapshot so cross-chunk XML
|
||||
# tags are handled correctly.
|
||||
raw_cumulative = event.get("text", "")
|
||||
clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative)
|
||||
new_text = clean_cumulative[len(prev_text) :]
|
||||
prev_text = clean_cumulative
|
||||
if not new_text:
|
||||
continue
|
||||
chunk = ChatCompletionChunk(
|
||||
|
|
@ -1176,7 +1291,7 @@ async def openai_chat_completions(
|
|||
logger.error(f"Error during GGUF tool streaming: {e}\n{tb}")
|
||||
error_chunk = {
|
||||
"error": {
|
||||
"message": "An internal error occurred",
|
||||
"message": _friendly_error(e),
|
||||
"type": "server_error",
|
||||
},
|
||||
}
|
||||
|
|
@ -1314,7 +1429,7 @@ async def openai_chat_completions(
|
|||
logger.error(f"Error during GGUF streaming: {e}", exc_info = True)
|
||||
error_chunk = {
|
||||
"error": {
|
||||
"message": "An internal error occurred",
|
||||
"message": _friendly_error(e),
|
||||
"type": "server_error",
|
||||
},
|
||||
}
|
||||
|
|
@ -1495,7 +1610,7 @@ async def openai_chat_completions(
|
|||
logger.error(f"Error during OpenAI streaming: {e}", exc_info = True)
|
||||
error_chunk = {
|
||||
"error": {
|
||||
"message": "An internal error occurred",
|
||||
"message": _friendly_error(e),
|
||||
"type": "server_error",
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -210,6 +210,77 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
|
|||
return found
|
||||
|
||||
|
||||
def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]:
|
||||
"""Scan an LM Studio models directory for model files.
|
||||
|
||||
LM Studio uses a ``publisher/model-name`` folder structure containing
|
||||
GGUF files, or standalone GGUF files at the top level.
|
||||
"""
|
||||
if not lm_dir.exists() or not lm_dir.is_dir():
|
||||
return []
|
||||
|
||||
found: List[LocalModelInfo] = []
|
||||
for child in lm_dir.iterdir():
|
||||
if not child.is_dir():
|
||||
if child.suffix == ".gguf" and child.is_file():
|
||||
try:
|
||||
updated_at = child.stat().st_mtime
|
||||
except OSError:
|
||||
updated_at = None
|
||||
found.append(
|
||||
LocalModelInfo(
|
||||
id = str(child),
|
||||
display_name = child.stem,
|
||||
path = str(child),
|
||||
source = "lmstudio",
|
||||
updated_at = updated_at,
|
||||
),
|
||||
)
|
||||
continue
|
||||
|
||||
# child is a publisher directory — scan its sub-directories
|
||||
for model_dir in child.iterdir():
|
||||
if model_dir.is_dir():
|
||||
has_model = (
|
||||
any(model_dir.glob("*.gguf"))
|
||||
or (model_dir / "config.json").exists()
|
||||
or any(model_dir.glob("*.safetensors"))
|
||||
)
|
||||
if not has_model:
|
||||
continue
|
||||
model_id = f"{child.name}/{model_dir.name}"
|
||||
try:
|
||||
updated_at = model_dir.stat().st_mtime
|
||||
except OSError:
|
||||
updated_at = None
|
||||
found.append(
|
||||
LocalModelInfo(
|
||||
id = str(model_dir),
|
||||
model_id = model_id,
|
||||
display_name = model_dir.name,
|
||||
path = str(model_dir),
|
||||
source = "lmstudio",
|
||||
updated_at = updated_at,
|
||||
),
|
||||
)
|
||||
elif model_dir.suffix == ".gguf" and model_dir.is_file():
|
||||
try:
|
||||
updated_at = model_dir.stat().st_mtime
|
||||
except OSError:
|
||||
updated_at = None
|
||||
found.append(
|
||||
LocalModelInfo(
|
||||
id = str(model_dir),
|
||||
model_id = f"{child.name}/{model_dir.stem}",
|
||||
display_name = model_dir.stem,
|
||||
path = str(model_dir),
|
||||
source = "lmstudio",
|
||||
updated_at = updated_at,
|
||||
),
|
||||
)
|
||||
return found
|
||||
|
||||
|
||||
@router.get("/local", response_model = LocalModelListResponse)
|
||||
async def list_local_models(
|
||||
models_dir: str = Query(
|
||||
|
|
@ -218,13 +289,29 @@ async def list_local_models(
|
|||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
List local model candidates from custom models dir and HF cache.
|
||||
List local model candidates from custom models dir, HF cache,
|
||||
legacy Unsloth HF cache, and LM Studio directories.
|
||||
"""
|
||||
from utils.paths import (
|
||||
legacy_hf_cache_dir,
|
||||
hf_default_cache_dir,
|
||||
lmstudio_model_dirs,
|
||||
)
|
||||
|
||||
# Resolve all scan directories up front.
|
||||
hf_cache_dir = _resolve_hf_cache_dir()
|
||||
legacy_hf = legacy_hf_cache_dir()
|
||||
hf_default = hf_default_cache_dir()
|
||||
lm_dirs = lmstudio_model_dirs()
|
||||
|
||||
# Validate models_dir against an allowlist of trusted directories.
|
||||
# Only the trusted Path objects are used for filesystem access -- the
|
||||
# user-supplied string is only used for matching, never for path construction.
|
||||
hf_cache_dir = _resolve_hf_cache_dir()
|
||||
allowed_roots = [Path("./models").resolve(), hf_cache_dir]
|
||||
allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir]
|
||||
if legacy_hf.is_dir():
|
||||
allowed_roots.append(legacy_hf)
|
||||
if hf_default.is_dir():
|
||||
allowed_roots.append(hf_default)
|
||||
try:
|
||||
from utils.paths import studio_root, outputs_root
|
||||
|
||||
|
|
@ -248,6 +335,22 @@ async def list_local_models(
|
|||
try:
|
||||
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
|
||||
|
||||
# Scan legacy Unsloth HF cache for backward compatibility
|
||||
if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve():
|
||||
local_models += _scan_hf_cache(legacy_hf)
|
||||
|
||||
# Scan HF system default cache (may differ when env vars are overridden)
|
||||
if (
|
||||
hf_default.is_dir()
|
||||
and hf_default.resolve() != hf_cache_dir.resolve()
|
||||
and hf_default.resolve() != legacy_hf.resolve()
|
||||
):
|
||||
local_models += _scan_hf_cache(hf_default)
|
||||
|
||||
# Scan LM Studio directories
|
||||
for lm_dir in lm_dirs:
|
||||
local_models += _scan_lmstudio_dir(lm_dir)
|
||||
|
||||
deduped: dict[str, LocalModelInfo] = {}
|
||||
for model in local_models:
|
||||
if model.id not in deduped:
|
||||
|
|
@ -262,6 +365,7 @@ async def list_local_models(
|
|||
return LocalModelListResponse(
|
||||
models_dir = str(models_root),
|
||||
hf_cache_dir = str(hf_cache_dir),
|
||||
lmstudio_dirs = [str(d) for d in lm_dirs],
|
||||
models = models,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -622,13 +726,40 @@ async def get_gguf_variants(
|
|||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
List available GGUF quantization variants for a HuggingFace repo.
|
||||
List available GGUF quantization variants for a HuggingFace repo
|
||||
or a local directory (e.g. LM Studio model folder).
|
||||
|
||||
Returns all available quantization variants (Q4_K_M, Q8_0, BF16, etc.)
|
||||
with file sizes, whether the model supports vision, and the recommended
|
||||
default variant.
|
||||
"""
|
||||
try:
|
||||
from utils.models.model_config import is_local_path, list_local_gguf_variants
|
||||
|
||||
# Local directory path (e.g. LM Studio models) — scan filesystem
|
||||
if is_local_path(repo_id):
|
||||
variants, has_vision = list_local_gguf_variants(repo_id)
|
||||
|
||||
filenames = [v.filename for v in variants]
|
||||
best = _pick_best_gguf(filenames)
|
||||
default_variant = _extract_quant_label(best) if best else None
|
||||
|
||||
return GgufVariantsResponse(
|
||||
repo_id = repo_id,
|
||||
variants = [
|
||||
GgufVariantDetail(
|
||||
filename = v.filename,
|
||||
quant = v.quant,
|
||||
size_bytes = v.size_bytes,
|
||||
downloaded = True, # all local variants are downloaded
|
||||
)
|
||||
for v in variants
|
||||
],
|
||||
has_vision = has_vision,
|
||||
default_variant = default_variant,
|
||||
)
|
||||
|
||||
# Remote HuggingFace repo — query HF API
|
||||
variants, has_vision = list_gguf_variants(repo_id, hf_token = hf_token)
|
||||
|
||||
# Determine default variant
|
||||
|
|
@ -846,46 +977,65 @@ def _get_repo_size_cached(repo_id: str) -> int:
|
|||
return 0
|
||||
|
||||
|
||||
def _all_hf_cache_scans():
|
||||
"""Return scan_cache_dir results for the active, legacy, and default HF caches."""
|
||||
from huggingface_hub import scan_cache_dir
|
||||
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir
|
||||
|
||||
scans = [scan_cache_dir()]
|
||||
seen: set[str] = set()
|
||||
try:
|
||||
# Resolve the active cache dir so we can dedup
|
||||
from huggingface_hub.constants import HF_HUB_CACHE
|
||||
|
||||
seen.add(str(Path(HF_HUB_CACHE).resolve()))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir):
|
||||
extra = extra_fn()
|
||||
if extra.is_dir() and str(extra.resolve()) not in seen:
|
||||
seen.add(str(extra.resolve()))
|
||||
try:
|
||||
scans.append(scan_cache_dir(cache_dir = str(extra)))
|
||||
except Exception as exc:
|
||||
logger.warning("Could not scan HF cache %s: %s", extra, exc)
|
||||
return scans
|
||||
|
||||
|
||||
@router.get("/cached-gguf")
|
||||
async def list_cached_gguf(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""List GGUF repos that have already been downloaded to the HF cache.
|
||||
|
||||
Uses scan_cache_dir() for proper repo IDs, then deduplicates by
|
||||
lowercased key (HF cache dirs are lowercased but the canonical repo
|
||||
ID preserves casing).
|
||||
"""
|
||||
"""List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache."""
|
||||
try:
|
||||
from huggingface_hub import scan_cache_dir
|
||||
cache_scans = _all_hf_cache_scans()
|
||||
|
||||
hf_cache = scan_cache_dir()
|
||||
seen_lower: dict[str, dict] = {}
|
||||
for repo_info in hf_cache.repos:
|
||||
if repo_info.repo_type != "model":
|
||||
continue
|
||||
repo_id = repo_info.repo_id
|
||||
if not repo_id.upper().endswith("-GGUF"):
|
||||
continue
|
||||
# Check for actual .gguf files and sum sizes
|
||||
total_size = 0
|
||||
has_gguf = False
|
||||
for revision in repo_info.revisions:
|
||||
for f in revision.files:
|
||||
if f.file_name.endswith(".gguf"):
|
||||
has_gguf = True
|
||||
total_size += f.size_on_disk
|
||||
if not has_gguf:
|
||||
continue
|
||||
# Deduplicate: keep the entry with the most data
|
||||
key = repo_id.lower()
|
||||
existing = seen_lower.get(key)
|
||||
if existing is None or total_size > existing["size_bytes"]:
|
||||
seen_lower[key] = {
|
||||
"repo_id": repo_id,
|
||||
"size_bytes": total_size,
|
||||
"cache_path": str(repo_info.repo_path),
|
||||
}
|
||||
for hf_cache in cache_scans:
|
||||
for repo_info in hf_cache.repos:
|
||||
if repo_info.repo_type != "model":
|
||||
continue
|
||||
repo_id = repo_info.repo_id
|
||||
if not repo_id.upper().endswith("-GGUF"):
|
||||
continue
|
||||
total_size = 0
|
||||
has_gguf = False
|
||||
for revision in repo_info.revisions:
|
||||
for f in revision.files:
|
||||
if f.file_name.endswith(".gguf"):
|
||||
has_gguf = True
|
||||
total_size += f.size_on_disk
|
||||
if not has_gguf:
|
||||
continue
|
||||
key = repo_id.lower()
|
||||
existing = seen_lower.get(key)
|
||||
if existing is None or total_size > existing["size_bytes"]:
|
||||
seen_lower[key] = {
|
||||
"repo_id": repo_id,
|
||||
"size_bytes": total_size,
|
||||
"cache_path": str(repo_info.repo_path),
|
||||
}
|
||||
cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"])
|
||||
return {"cached": cached}
|
||||
except Exception as e:
|
||||
|
|
@ -897,44 +1047,39 @@ async def list_cached_gguf(
|
|||
async def list_cached_models(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""List non-GGUF model repos that have been downloaded to the HF cache.
|
||||
|
||||
Only includes repos that actually contain model weight files
|
||||
(.safetensors, .bin), not repos with only config/metadata.
|
||||
"""
|
||||
"""List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache."""
|
||||
_WEIGHT_EXTENSIONS = (".safetensors", ".bin")
|
||||
|
||||
try:
|
||||
from huggingface_hub import scan_cache_dir
|
||||
cache_scans = _all_hf_cache_scans()
|
||||
|
||||
hf_cache = scan_cache_dir()
|
||||
seen_lower: dict[str, dict] = {}
|
||||
for repo_info in hf_cache.repos:
|
||||
if repo_info.repo_type != "model":
|
||||
continue
|
||||
repo_id = repo_info.repo_id
|
||||
if repo_id.upper().endswith("-GGUF"):
|
||||
continue
|
||||
total_size = sum(
|
||||
f.size_on_disk for rev in repo_info.revisions for f in rev.files
|
||||
)
|
||||
if total_size == 0:
|
||||
continue
|
||||
# Skip repos that only have config/metadata files (no weights)
|
||||
has_weights = any(
|
||||
f.file_name.endswith(_WEIGHT_EXTENSIONS)
|
||||
for rev in repo_info.revisions
|
||||
for f in rev.files
|
||||
)
|
||||
if not has_weights:
|
||||
continue
|
||||
key = repo_id.lower()
|
||||
existing = seen_lower.get(key)
|
||||
if existing is None or total_size > existing["size_bytes"]:
|
||||
seen_lower[key] = {
|
||||
"repo_id": repo_id,
|
||||
"size_bytes": total_size,
|
||||
}
|
||||
for hf_cache in cache_scans:
|
||||
for repo_info in hf_cache.repos:
|
||||
if repo_info.repo_type != "model":
|
||||
continue
|
||||
repo_id = repo_info.repo_id
|
||||
if repo_id.upper().endswith("-GGUF"):
|
||||
continue
|
||||
total_size = sum(
|
||||
f.size_on_disk for rev in repo_info.revisions for f in rev.files
|
||||
)
|
||||
if total_size == 0:
|
||||
continue
|
||||
has_weights = any(
|
||||
f.file_name.endswith(_WEIGHT_EXTENSIONS)
|
||||
for rev in repo_info.revisions
|
||||
for f in rev.files
|
||||
)
|
||||
if not has_weights:
|
||||
continue
|
||||
key = repo_id.lower()
|
||||
existing = seen_lower.get(key)
|
||||
if existing is None or total_size > existing["size_bytes"]:
|
||||
seen_lower[key] = {
|
||||
"repo_id": repo_id,
|
||||
"size_bytes": total_size,
|
||||
}
|
||||
cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"])
|
||||
return {"cached": cached}
|
||||
except Exception as e:
|
||||
|
|
@ -989,15 +1134,17 @@ async def delete_cached_model(
|
|||
pass
|
||||
|
||||
try:
|
||||
from huggingface_hub import scan_cache_dir
|
||||
cache_scans = _all_hf_cache_scans()
|
||||
|
||||
hf_cache = scan_cache_dir()
|
||||
target_repo = None
|
||||
for repo_info in hf_cache.repos:
|
||||
if repo_info.repo_type != "model":
|
||||
continue
|
||||
if repo_info.repo_id.lower() == repo_id.lower():
|
||||
target_repo = repo_info
|
||||
for hf_cache in cache_scans:
|
||||
for repo_info in hf_cache.repos:
|
||||
if repo_info.repo_type != "model":
|
||||
continue
|
||||
if repo_info.repo_id.lower() == repo_id.lower():
|
||||
target_repo = repo_info
|
||||
break
|
||||
if target_repo is not None:
|
||||
break
|
||||
|
||||
if target_repo is None:
|
||||
|
|
|
|||
|
|
@ -88,14 +88,22 @@ async def get_hardware_utilization(
|
|||
Get a live snapshot of GPU hardware utilization.
|
||||
|
||||
Designed to be polled by the frontend during training.
|
||||
Returns GPU utilization %, temperature, VRAM usage, and power draw
|
||||
via nvidia-smi for maximum accuracy.
|
||||
Returns live GPU memory usage information for the active backend.
|
||||
"""
|
||||
from utils.hardware import get_gpu_utilization
|
||||
|
||||
return get_gpu_utilization()
|
||||
|
||||
|
||||
@router.get("/hardware/visible")
|
||||
async def get_visible_hardware_utilization(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
from utils.hardware import get_visible_gpu_utilization
|
||||
|
||||
return get_visible_gpu_utilization()
|
||||
|
||||
|
||||
@router.post("/start")
|
||||
async def start_training(
|
||||
request: TrainingStartRequest,
|
||||
|
|
@ -202,6 +210,7 @@ async def start_training(
|
|||
"enable_tensorboard": request.enable_tensorboard,
|
||||
"tensorboard_dir": request.tensorboard_dir or "",
|
||||
"trust_remote_code": request.trust_remote_code,
|
||||
"gpu_ids": request.gpu_ids,
|
||||
}
|
||||
|
||||
# Training page has no trust_remote_code toggle — the value comes from
|
||||
|
|
@ -269,6 +278,9 @@ async def start_training(
|
|||
error = None,
|
||||
)
|
||||
|
||||
except ValueError as e:
|
||||
logger.warning("Rejected training GPU selection: %s", e)
|
||||
raise HTTPException(status_code = 400, detail = str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Error starting training: {e}", exc_info = True)
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ if str(backend_dir) not in sys.path:
|
|||
import _platform_compat # noqa: F401
|
||||
|
||||
from loggers import get_logger
|
||||
from startup_banner import print_studio_access_banner
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -158,6 +159,29 @@ def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int:
|
|||
)
|
||||
|
||||
|
||||
_PID_FILE = Path.home() / ".unsloth" / "studio" / "studio.pid"
|
||||
|
||||
|
||||
def _write_pid_file():
|
||||
"""Write the current process PID to the studio PID file."""
|
||||
try:
|
||||
_PID_FILE.parent.mkdir(parents = True, exist_ok = True)
|
||||
_PID_FILE.write_text(str(os.getpid()))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _remove_pid_file():
|
||||
"""Remove the PID file if it belongs to this process."""
|
||||
try:
|
||||
if _PID_FILE.is_file():
|
||||
stored = _PID_FILE.read_text().strip()
|
||||
if stored == str(os.getpid()):
|
||||
_PID_FILE.unlink(missing_ok = True)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _graceful_shutdown(server = None):
|
||||
"""Explicitly shut down all subprocess backends and the uvicorn server.
|
||||
|
||||
|
|
@ -165,6 +189,7 @@ def _graceful_shutdown(server = None):
|
|||
before the parent exits. This is critical on Windows where atexit
|
||||
handlers are unreliable after Ctrl+C.
|
||||
"""
|
||||
_remove_pid_file()
|
||||
logger.info("Graceful shutdown initiated — cleaning up subprocesses...")
|
||||
|
||||
# 1. Shut down uvicorn server (releases the listening socket)
|
||||
|
|
@ -287,10 +312,10 @@ def run_server(
|
|||
if frontend_path:
|
||||
if setup_frontend(app, frontend_path):
|
||||
if not silent:
|
||||
print(f"✅ Frontend loaded from {frontend_path}")
|
||||
print(f"[OK] Frontend loaded from {frontend_path}")
|
||||
else:
|
||||
if not silent:
|
||||
print(f"⚠️ Frontend not found at {frontend_path}")
|
||||
print(f"[WARNING] Frontend not found at {frontend_path}")
|
||||
|
||||
# Create the uvicorn server and expose it for signal handlers
|
||||
config = uvicorn.Config(
|
||||
|
|
@ -307,21 +332,27 @@ def run_server(
|
|||
thread.start()
|
||||
time.sleep(3)
|
||||
|
||||
_write_pid_file()
|
||||
import atexit
|
||||
|
||||
atexit.register(_remove_pid_file)
|
||||
|
||||
# Expose a shutdown callable via app.state so the /api/shutdown endpoint
|
||||
# can trigger graceful shutdown without circular imports.
|
||||
def _trigger_shutdown():
|
||||
_graceful_shutdown(_server)
|
||||
if _shutdown_event is not None:
|
||||
_shutdown_event.set()
|
||||
|
||||
app.state.trigger_shutdown = _trigger_shutdown
|
||||
|
||||
if not silent:
|
||||
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
||||
|
||||
print("")
|
||||
print("=" * 50)
|
||||
print(f"🦥 Open your web browser, and enter http://localhost:{port}")
|
||||
print("=" * 50)
|
||||
print("")
|
||||
print("=" * 50)
|
||||
print(f"🦥 Unsloth Studio is running on port {port}")
|
||||
print(f" Local Access: http://localhost:{port}")
|
||||
print(f" Worldwide Web Address: http://{display_host}:{port}")
|
||||
print(f" API: http://{display_host}:{port}/api")
|
||||
print(f" Health: http://{display_host}:{port}/api/health")
|
||||
print("=" * 50)
|
||||
print_studio_access_banner(
|
||||
port = port,
|
||||
bind_host = host,
|
||||
display_host = display_host,
|
||||
)
|
||||
|
||||
return app
|
||||
|
||||
|
|
|
|||
123
studio/backend/startup_banner.py
Normal file
123
studio/backend/startup_banner.py
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Terminal banner for Studio startup.
|
||||
|
||||
Stdlib only — safe to import without the rest of the backend (no structlog/uvicorn).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
def stdout_supports_color() -> bool:
|
||||
"""True if we should emit ANSI colors."""
|
||||
if os.environ.get("NO_COLOR", "").strip():
|
||||
return False
|
||||
if os.environ.get("FORCE_COLOR", "").strip():
|
||||
return True
|
||||
try:
|
||||
return sys.stdout.isatty()
|
||||
except (AttributeError, OSError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def print_port_in_use_notice(original_port: int, new_port: int) -> None:
|
||||
"""Message when the requested port is taken and another is chosen."""
|
||||
msg = f"Port {original_port} is in use, using port {new_port} instead."
|
||||
if stdout_supports_color():
|
||||
print(f"\033[38;5;245m{msg}\033[0m")
|
||||
else:
|
||||
print(msg)
|
||||
|
||||
|
||||
def print_studio_access_banner(
|
||||
*,
|
||||
port: int,
|
||||
bind_host: str,
|
||||
display_host: str,
|
||||
) -> None:
|
||||
"""Pretty-print URLs after the server is listening (beginner-friendly)."""
|
||||
use_color = stdout_supports_color()
|
||||
dim = "\033[38;5;245m"
|
||||
title = "\033[38;5;150m"
|
||||
local_url_style = "\033[38;5;108;1m"
|
||||
secondary = "\033[38;5;109m"
|
||||
reset = "\033[0m"
|
||||
|
||||
def style(text: str, code: str) -> str:
|
||||
return f"{code}{text}{reset}" if use_color else text
|
||||
|
||||
ipv6_bind = bind_host in ("::", "::1")
|
||||
if ipv6_bind:
|
||||
loopback_url = f"http://[::1]:{port}"
|
||||
alt_local = f"http://localhost:{port}"
|
||||
else:
|
||||
loopback_url = f"http://127.0.0.1:{port}"
|
||||
alt_local = f"http://localhost:{port}"
|
||||
if ":" in display_host:
|
||||
external_url = f"http://[{display_host}]:{port}"
|
||||
else:
|
||||
external_url = f"http://{display_host}:{port}"
|
||||
|
||||
listen_all = bind_host in ("0.0.0.0", "::")
|
||||
loopback_bind = bind_host in ("127.0.0.1", "localhost", "::1")
|
||||
|
||||
# Use loopback URL only when the server is reachable on loopback;
|
||||
# otherwise show the actual bound address.
|
||||
primary_url = loopback_url if listen_all or loopback_bind else external_url
|
||||
tip_url = alt_local if listen_all or loopback_bind else external_url
|
||||
api_base = primary_url
|
||||
|
||||
lines: list[str] = [
|
||||
"",
|
||||
style("🦥 Unsloth Studio is running", title),
|
||||
style("─" * 52, dim),
|
||||
style(" On this machine -- open this in your browser:", dim),
|
||||
style(f" {primary_url}", local_url_style),
|
||||
]
|
||||
|
||||
if (listen_all or loopback_bind) and primary_url != alt_local:
|
||||
lines.append(style(f" (same as {alt_local})", dim))
|
||||
|
||||
if listen_all and display_host not in (
|
||||
"127.0.0.1",
|
||||
"localhost",
|
||||
"::1",
|
||||
"0.0.0.0",
|
||||
"::",
|
||||
):
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
style(" From another device on your network / to share:", dim),
|
||||
style(f" {external_url}", secondary),
|
||||
]
|
||||
)
|
||||
elif not listen_all and not loopback_bind and external_url != primary_url:
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
style(" Bound address:", dim),
|
||||
style(f" {external_url}", secondary),
|
||||
]
|
||||
)
|
||||
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
style(" API & health:", dim),
|
||||
style(f" {api_base}/api", secondary),
|
||||
style(f" {api_base}/api/health", secondary),
|
||||
style("─" * 52, dim),
|
||||
style(
|
||||
f" Tip: if you are on this computer, open {tip_url}/ in your browser.",
|
||||
dim,
|
||||
),
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
print("\n".join(lines))
|
||||
1103
studio/backend/tests/test_gpu_selection.py
Normal file
1103
studio/backend/tests/test_gpu_selection.py
Normal file
File diff suppressed because it is too large
Load diff
544
studio/backend/tests/test_gpu_selection_sandbox.py
Normal file
544
studio/backend/tests/test_gpu_selection_sandbox.py
Normal file
|
|
@ -0,0 +1,544 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Sandbox test for multi-GPU selection logic.
|
||||
|
||||
Tests the core GPU selection, memory estimation, and device_map logic
|
||||
in an isolated environment. Can be run on Linux, macOS, and Windows
|
||||
without requiring actual GPUs -- all hardware calls are mocked.
|
||||
|
||||
Usage:
|
||||
python -m pytest studio/backend/tests/test_gpu_selection_sandbox.py -v
|
||||
# or directly:
|
||||
python studio/backend/tests/test_gpu_selection_sandbox.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
# Ensure backend is on sys.path
|
||||
_backend_root = Path(__file__).resolve().parent.parent
|
||||
if str(_backend_root) not in sys.path:
|
||||
sys.path.insert(0, str(_backend_root))
|
||||
|
||||
|
||||
def _make_fake_config(
|
||||
vocab_size = 32000,
|
||||
hidden_size = 4096,
|
||||
intermediate_size = 11008,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
tie_word_embeddings = False,
|
||||
):
|
||||
"""Create a fake HF config-like object for estimation tests."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
return SimpleNamespace(
|
||||
vocab_size = vocab_size,
|
||||
hidden_size = hidden_size,
|
||||
intermediate_size = intermediate_size,
|
||||
num_hidden_layers = num_hidden_layers,
|
||||
num_attention_heads = num_attention_heads,
|
||||
num_key_value_heads = num_key_value_heads,
|
||||
tie_word_embeddings = tie_word_embeddings,
|
||||
)
|
||||
|
||||
|
||||
class TestEstimateFP16ModelSizeFromConfig(unittest.TestCase):
|
||||
"""Test the config-based model size estimation."""
|
||||
|
||||
def test_llama_8b_size_reasonable(self):
|
||||
from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config
|
||||
|
||||
config = _make_fake_config(
|
||||
vocab_size = 128256,
|
||||
hidden_size = 4096,
|
||||
intermediate_size = 14336,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
tie_word_embeddings = False,
|
||||
)
|
||||
size = _estimate_fp16_model_size_bytes_from_config(config)
|
||||
self.assertIsNotNone(size)
|
||||
size_gb = size / (1024**3)
|
||||
# Llama 3.1 8B should be ~15GB in fp16
|
||||
self.assertGreater(size_gb, 12)
|
||||
self.assertLess(size_gb, 20)
|
||||
|
||||
def test_small_model(self):
|
||||
from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config
|
||||
|
||||
config = _make_fake_config(
|
||||
vocab_size = 32000,
|
||||
hidden_size = 2048,
|
||||
intermediate_size = 5504,
|
||||
num_hidden_layers = 22,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 4,
|
||||
)
|
||||
size = _estimate_fp16_model_size_bytes_from_config(config)
|
||||
self.assertIsNotNone(size)
|
||||
size_gb = size / (1024**3)
|
||||
# ~1B model should be ~2GB in fp16
|
||||
self.assertGreater(size_gb, 1)
|
||||
self.assertLess(size_gb, 5)
|
||||
|
||||
def test_returns_none_for_incomplete_config(self):
|
||||
from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config
|
||||
from types import SimpleNamespace
|
||||
|
||||
config = SimpleNamespace(vocab_size = 32000) # Missing most fields
|
||||
size = _estimate_fp16_model_size_bytes_from_config(config)
|
||||
self.assertIsNone(size)
|
||||
|
||||
def test_moe_model(self):
|
||||
from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config
|
||||
from types import SimpleNamespace
|
||||
|
||||
config = SimpleNamespace(
|
||||
vocab_size = 152064,
|
||||
hidden_size = 3584,
|
||||
intermediate_size = 18944,
|
||||
num_hidden_layers = 28,
|
||||
num_attention_heads = 28,
|
||||
num_key_value_heads = 4,
|
||||
tie_word_embeddings = False,
|
||||
num_local_experts = 64,
|
||||
moe_intermediate_size = 2560,
|
||||
)
|
||||
size = _estimate_fp16_model_size_bytes_from_config(config)
|
||||
self.assertIsNotNone(size)
|
||||
size_gb = size / (1024**3)
|
||||
# MoE model with 64 experts should be large
|
||||
self.assertGreater(size_gb, 50)
|
||||
|
||||
|
||||
class TestEstimateRequiredModelMemory(unittest.TestCase):
|
||||
"""Test memory requirement estimation."""
|
||||
|
||||
def test_inference_fp16_uses_1_3x(self):
|
||||
from utils.hardware.hardware import estimate_required_model_memory_gb
|
||||
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (10 * (1024**3), "config"), # 10GB model
|
||||
):
|
||||
required, meta = estimate_required_model_memory_gb(
|
||||
"test/model",
|
||||
training_type = None, # inference
|
||||
load_in_4bit = False,
|
||||
)
|
||||
self.assertIsNotNone(required)
|
||||
self.assertAlmostEqual(required, 13.0, places = 0)
|
||||
self.assertEqual(meta["mode"], "inference")
|
||||
|
||||
def test_inference_4bit_uses_reduced_estimate(self):
|
||||
from utils.hardware.hardware import estimate_required_model_memory_gb
|
||||
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (30 * (1024**3), "config"), # 30GB fp16 model
|
||||
):
|
||||
required, meta = estimate_required_model_memory_gb(
|
||||
"test/model",
|
||||
training_type = None, # inference
|
||||
load_in_4bit = True,
|
||||
)
|
||||
self.assertIsNotNone(required)
|
||||
# 4bit base = 30/3.2 = 9.375GB, required = 9.375 + max(9.375*0.3, 2) = 12.19GB
|
||||
self.assertAlmostEqual(required, 12.2, places = 0)
|
||||
|
||||
def test_4bit_training_reduces_base(self):
|
||||
from utils.hardware.hardware import estimate_required_model_memory_gb
|
||||
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (30 * (1024**3), "config"), # 30GB fp16 model
|
||||
):
|
||||
required, meta = estimate_required_model_memory_gb(
|
||||
"test/model",
|
||||
training_type = "LoRA/QLoRA",
|
||||
load_in_4bit = True,
|
||||
)
|
||||
self.assertIsNotNone(required)
|
||||
# fallback: base=30/3.2=9.375, lora=30*0.04=1.2, act=30*0.15=4.5, cuda=1.4
|
||||
self.assertAlmostEqual(required, 16.5, places = 0)
|
||||
|
||||
def test_full_finetune_uses_3_5x(self):
|
||||
from utils.hardware.hardware import estimate_required_model_memory_gb
|
||||
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (10 * (1024**3), "config"), # 10GB model
|
||||
):
|
||||
required, meta = estimate_required_model_memory_gb(
|
||||
"test/model",
|
||||
training_type = "Full Finetuning",
|
||||
)
|
||||
self.assertIsNotNone(required)
|
||||
# fallback: 10 * 3.5 + 1.4 cuda overhead = 36.4
|
||||
self.assertAlmostEqual(required, 36.4, places = 0)
|
||||
|
||||
def test_returns_none_when_unavailable(self):
|
||||
from utils.hardware.hardware import estimate_required_model_memory_gb
|
||||
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (None, "unavailable"),
|
||||
):
|
||||
required, meta = estimate_required_model_memory_gb("test/model")
|
||||
self.assertIsNone(required)
|
||||
|
||||
|
||||
class TestAutoSelectGpuIds(unittest.TestCase):
|
||||
"""Test automatic GPU selection based on model size and free memory."""
|
||||
|
||||
def _make_utilization(self, devices):
|
||||
"""Create a fake utilization response."""
|
||||
return {
|
||||
"available": True,
|
||||
"devices": [
|
||||
{
|
||||
"index": idx,
|
||||
"vram_total_gb": total,
|
||||
"vram_used_gb": total - free,
|
||||
}
|
||||
for idx, total, free in devices
|
||||
],
|
||||
}
|
||||
|
||||
def test_single_gpu_sufficient(self):
|
||||
from utils.hardware.hardware import auto_select_gpu_ids
|
||||
import utils.hardware.hardware as hw
|
||||
|
||||
with (
|
||||
patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
|
||||
patch.object(
|
||||
hw,
|
||||
"estimate_required_model_memory_gb",
|
||||
return_value = (
|
||||
10.0,
|
||||
{
|
||||
"mode": "inference",
|
||||
"required_gb": 10.0,
|
||||
"model_size_source": "config",
|
||||
"model_size_gb": 7.7,
|
||||
},
|
||||
),
|
||||
),
|
||||
patch.object(
|
||||
hw,
|
||||
"_get_parent_visible_gpu_spec",
|
||||
return_value = {
|
||||
"raw": "0,1,2,3",
|
||||
"numeric_ids": [0, 1, 2, 3],
|
||||
"supports_explicit_gpu_ids": True,
|
||||
},
|
||||
),
|
||||
patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1, 2, 3]),
|
||||
patch.object(
|
||||
hw,
|
||||
"get_visible_gpu_utilization",
|
||||
return_value = self._make_utilization(
|
||||
[
|
||||
(0, 80.0, 75.0),
|
||||
(1, 80.0, 78.0),
|
||||
(2, 80.0, 70.0),
|
||||
(3, 80.0, 72.0),
|
||||
]
|
||||
),
|
||||
),
|
||||
):
|
||||
selected, meta = auto_select_gpu_ids("test/model")
|
||||
# Should pick GPU 1 (most free memory: 78GB) -- enough for 10GB
|
||||
self.assertEqual(len(selected), 1)
|
||||
self.assertEqual(selected[0], 1)
|
||||
|
||||
def test_two_gpus_needed(self):
|
||||
from utils.hardware.hardware import auto_select_gpu_ids
|
||||
import utils.hardware.hardware as hw
|
||||
|
||||
with (
|
||||
patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
|
||||
patch.object(
|
||||
hw,
|
||||
"estimate_required_model_memory_gb",
|
||||
return_value = (
|
||||
50.0,
|
||||
{
|
||||
"mode": "inference",
|
||||
"required_gb": 50.0,
|
||||
"model_size_source": "config",
|
||||
"model_size_gb": 38.0,
|
||||
},
|
||||
),
|
||||
),
|
||||
patch.object(
|
||||
hw,
|
||||
"_get_parent_visible_gpu_spec",
|
||||
return_value = {
|
||||
"raw": "0,1",
|
||||
"numeric_ids": [0, 1],
|
||||
"supports_explicit_gpu_ids": True,
|
||||
},
|
||||
),
|
||||
patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]),
|
||||
patch.object(
|
||||
hw,
|
||||
"get_visible_gpu_utilization",
|
||||
return_value = self._make_utilization(
|
||||
[
|
||||
(0, 40.0, 30.0), # 30GB free
|
||||
(1, 40.0, 35.0), # 35GB free
|
||||
]
|
||||
),
|
||||
),
|
||||
):
|
||||
selected, meta = auto_select_gpu_ids("test/model")
|
||||
# 35GB (first) + 30*0.85 (second) = 60.5GB > 50GB
|
||||
self.assertEqual(len(selected), 2)
|
||||
|
||||
def test_non_cuda_returns_none(self):
|
||||
from utils.hardware.hardware import auto_select_gpu_ids
|
||||
import utils.hardware.hardware as hw
|
||||
|
||||
with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU):
|
||||
selected, meta = auto_select_gpu_ids("test/model")
|
||||
self.assertIsNone(selected)
|
||||
self.assertEqual(meta["selection_mode"], "non_cuda")
|
||||
|
||||
|
||||
class TestGetDeviceMap(unittest.TestCase):
|
||||
"""Test device_map string generation."""
|
||||
|
||||
def test_single_gpu_returns_sequential(self):
|
||||
from utils.hardware.hardware import get_device_map
|
||||
import utils.hardware.hardware as hw
|
||||
|
||||
with (
|
||||
patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
|
||||
patch.object(
|
||||
hw,
|
||||
"_get_parent_visible_gpu_spec",
|
||||
return_value = {
|
||||
"raw": "0",
|
||||
"numeric_ids": [0],
|
||||
"supports_explicit_gpu_ids": True,
|
||||
},
|
||||
),
|
||||
patch.object(hw, "get_visible_gpu_count", return_value = 1),
|
||||
):
|
||||
dm = get_device_map(gpu_ids = [0])
|
||||
self.assertEqual(dm, "sequential")
|
||||
|
||||
def test_multi_gpu_returns_balanced(self):
|
||||
from utils.hardware.hardware import get_device_map
|
||||
import utils.hardware.hardware as hw
|
||||
|
||||
with patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA):
|
||||
dm = get_device_map(gpu_ids = [0, 1])
|
||||
self.assertEqual(dm, "balanced")
|
||||
|
||||
def test_cpu_returns_sequential(self):
|
||||
from utils.hardware.hardware import get_device_map
|
||||
import utils.hardware.hardware as hw
|
||||
|
||||
with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU):
|
||||
dm = get_device_map(gpu_ids = None)
|
||||
self.assertEqual(dm, "sequential")
|
||||
|
||||
|
||||
class TestResolveRequestedGpuIds(unittest.TestCase):
|
||||
"""Test GPU ID validation."""
|
||||
|
||||
def test_none_returns_parent_visible(self):
|
||||
from utils.hardware.hardware import resolve_requested_gpu_ids
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False),
|
||||
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
||||
):
|
||||
result = resolve_requested_gpu_ids(None)
|
||||
self.assertEqual(result, [2, 3])
|
||||
|
||||
def test_empty_list_returns_parent_visible(self):
|
||||
from utils.hardware.hardware import resolve_requested_gpu_ids
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False),
|
||||
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
||||
):
|
||||
result = resolve_requested_gpu_ids([])
|
||||
self.assertEqual(result, [2, 3])
|
||||
|
||||
def test_duplicates_rejected(self):
|
||||
from utils.hardware.hardware import resolve_requested_gpu_ids
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1,2"}, clear = False),
|
||||
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
||||
):
|
||||
with self.assertRaises(ValueError):
|
||||
resolve_requested_gpu_ids([1, 1])
|
||||
|
||||
def test_out_of_range_rejected(self):
|
||||
from utils.hardware.hardware import resolve_requested_gpu_ids
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"}, clear = False),
|
||||
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 4),
|
||||
):
|
||||
with self.assertRaises(ValueError):
|
||||
resolve_requested_gpu_ids([5])
|
||||
|
||||
def test_uuid_env_var_rejects_explicit_ids(self):
|
||||
from utils.hardware.hardware import resolve_requested_gpu_ids
|
||||
|
||||
with (
|
||||
patch.dict(
|
||||
os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-abc,GPU-def"}, clear = False
|
||||
),
|
||||
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
||||
):
|
||||
with self.assertRaises(ValueError):
|
||||
resolve_requested_gpu_ids([0])
|
||||
|
||||
|
||||
class TestApplyGpuIds(unittest.TestCase):
|
||||
"""Test CUDA_VISIBLE_DEVICES environment variable setting."""
|
||||
|
||||
def test_apply_list(self):
|
||||
from utils.hardware.hardware import apply_gpu_ids
|
||||
|
||||
with patch.dict(os.environ, {}, clear = False):
|
||||
apply_gpu_ids([3, 5])
|
||||
self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), "3,5")
|
||||
|
||||
def test_apply_none_does_nothing(self):
|
||||
from utils.hardware.hardware import apply_gpu_ids
|
||||
|
||||
original = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||||
apply_gpu_ids(None)
|
||||
self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), original)
|
||||
|
||||
|
||||
class TestMultiGpuOverheadAccounting(unittest.TestCase):
|
||||
"""Test that multi-GPU overhead is applied correctly.
|
||||
|
||||
The first GPU should keep its full free memory, and only
|
||||
additional GPUs should have the overhead factor applied.
|
||||
"""
|
||||
|
||||
def _make_utilization(self, devices):
|
||||
return {
|
||||
"available": True,
|
||||
"devices": [
|
||||
{
|
||||
"index": idx,
|
||||
"vram_total_gb": total,
|
||||
"vram_used_gb": total - free,
|
||||
}
|
||||
for idx, total, free in devices
|
||||
],
|
||||
}
|
||||
|
||||
def test_first_gpu_not_penalized(self):
|
||||
"""A model that just fits on 1 GPU should not require 2 GPUs."""
|
||||
from utils.hardware.hardware import auto_select_gpu_ids
|
||||
import utils.hardware.hardware as hw
|
||||
|
||||
# Model requires 79GB, GPU has 80GB free
|
||||
with (
|
||||
patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
|
||||
patch.object(
|
||||
hw,
|
||||
"estimate_required_model_memory_gb",
|
||||
return_value = (
|
||||
79.0,
|
||||
{
|
||||
"mode": "inference",
|
||||
"required_gb": 79.0,
|
||||
"model_size_source": "config",
|
||||
"model_size_gb": 60.0,
|
||||
},
|
||||
),
|
||||
),
|
||||
patch.object(
|
||||
hw,
|
||||
"_get_parent_visible_gpu_spec",
|
||||
return_value = {
|
||||
"raw": "0,1",
|
||||
"numeric_ids": [0, 1],
|
||||
"supports_explicit_gpu_ids": True,
|
||||
},
|
||||
),
|
||||
patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]),
|
||||
patch.object(
|
||||
hw,
|
||||
"get_visible_gpu_utilization",
|
||||
return_value = self._make_utilization(
|
||||
[
|
||||
(0, 80.0, 80.0),
|
||||
(1, 80.0, 80.0),
|
||||
]
|
||||
),
|
||||
),
|
||||
):
|
||||
selected, meta = auto_select_gpu_ids("test/model")
|
||||
# Should fit on 1 GPU (80GB >= 79GB)
|
||||
self.assertEqual(len(selected), 1)
|
||||
|
||||
def test_second_gpu_has_overhead(self):
|
||||
"""When 2 GPUs are needed, the second one's contribution is reduced."""
|
||||
from utils.hardware.hardware import auto_select_gpu_ids
|
||||
import utils.hardware.hardware as hw
|
||||
|
||||
# Model requires 110GB. First GPU has 80GB, second has 40GB.
|
||||
# With overhead: 80 + 40*0.85 = 114GB -- just enough
|
||||
with (
|
||||
patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
|
||||
patch.object(
|
||||
hw,
|
||||
"estimate_required_model_memory_gb",
|
||||
return_value = (
|
||||
110.0,
|
||||
{
|
||||
"mode": "inference",
|
||||
"required_gb": 110.0,
|
||||
"model_size_source": "config",
|
||||
"model_size_gb": 85.0,
|
||||
},
|
||||
),
|
||||
),
|
||||
patch.object(
|
||||
hw,
|
||||
"_get_parent_visible_gpu_spec",
|
||||
return_value = {
|
||||
"raw": "0,1",
|
||||
"numeric_ids": [0, 1],
|
||||
"supports_explicit_gpu_ids": True,
|
||||
},
|
||||
),
|
||||
patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]),
|
||||
patch.object(
|
||||
hw,
|
||||
"get_visible_gpu_utilization",
|
||||
return_value = self._make_utilization(
|
||||
[
|
||||
(0, 80.0, 80.0),
|
||||
(1, 80.0, 40.0),
|
||||
]
|
||||
),
|
||||
),
|
||||
):
|
||||
selected, meta = auto_select_gpu_ids("test/model")
|
||||
# Should use both GPUs
|
||||
self.assertEqual(len(selected), 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -285,7 +285,7 @@ class TestLogGpuMemory:
|
|||
def test_does_not_raise(self):
|
||||
log_gpu_memory("test")
|
||||
|
||||
def test_logs_gpu_info_when_available(self, caplog):
|
||||
def test_logs_gpu_info_when_available(self, capfd):
|
||||
fake_info = {
|
||||
"available": True,
|
||||
"backend": "cuda",
|
||||
|
|
@ -295,35 +295,27 @@ class TestLogGpuMemory:
|
|||
"utilization_pct": 12.5,
|
||||
"free_gb": 14.0,
|
||||
}
|
||||
import structlog
|
||||
from loggers import get_logger
|
||||
|
||||
with (
|
||||
patch(
|
||||
"utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
|
||||
),
|
||||
caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"),
|
||||
with patch(
|
||||
"utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
|
||||
):
|
||||
log_gpu_memory("unit-test")
|
||||
|
||||
assert "unit-test" in caplog.text
|
||||
assert "CUDA" in caplog.text
|
||||
assert "FakeGPU" in caplog.text
|
||||
captured = capfd.readouterr()
|
||||
assert "unit-test" in captured.out
|
||||
assert "CUDA" in captured.out
|
||||
assert "FakeGPU" in captured.out
|
||||
|
||||
def test_logs_cpu_fallback_when_no_gpu(self, caplog):
|
||||
def test_logs_cpu_fallback_when_no_gpu(self, capfd):
|
||||
fake_info = {"available": False, "backend": "cpu"}
|
||||
import structlog
|
||||
from loggers import get_logger
|
||||
|
||||
with (
|
||||
patch(
|
||||
"utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
|
||||
),
|
||||
caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"),
|
||||
with patch(
|
||||
"utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
|
||||
):
|
||||
log_gpu_memory("cpu-test")
|
||||
|
||||
assert "No GPU available" in caplog.text
|
||||
captured = capfd.readouterr()
|
||||
assert "No GPU available" in captured.out
|
||||
|
||||
|
||||
# ========== format_error_message() ==========
|
||||
|
|
|
|||
695
studio/backend/tests/test_vram_estimation.py
Normal file
695
studio/backend/tests/test_vram_estimation.py
Normal file
|
|
@ -0,0 +1,695 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from utils.hardware.vram_estimation import (
|
||||
ModelArchConfig,
|
||||
TrainingVramConfig,
|
||||
extract_arch_config,
|
||||
compute_model_weights_bytes,
|
||||
compute_total_params,
|
||||
compute_lora_params,
|
||||
compute_lora_adapter_bytes,
|
||||
compute_optimizer_bytes,
|
||||
compute_gradient_bytes,
|
||||
compute_activation_bytes,
|
||||
estimate_training_vram,
|
||||
DEFAULT_TARGET_MODULES,
|
||||
)
|
||||
|
||||
|
||||
def _gb(b: int) -> float:
|
||||
return b / (1024**3)
|
||||
|
||||
|
||||
LLAMA_8B = ModelArchConfig(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 128256,
|
||||
tie_word_embeddings = False,
|
||||
)
|
||||
|
||||
QWEN_05B = ModelArchConfig(
|
||||
hidden_size = 896,
|
||||
num_hidden_layers = 24,
|
||||
num_attention_heads = 14,
|
||||
num_key_value_heads = 2,
|
||||
intermediate_size = 4864,
|
||||
vocab_size = 151936,
|
||||
tie_word_embeddings = True,
|
||||
)
|
||||
|
||||
MOE_CONFIG = ModelArchConfig(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 32000,
|
||||
tie_word_embeddings = False,
|
||||
num_experts = 8,
|
||||
)
|
||||
|
||||
DEEPSEEK_V3 = ModelArchConfig(
|
||||
hidden_size = 7168,
|
||||
num_hidden_layers = 61,
|
||||
num_attention_heads = 128,
|
||||
num_key_value_heads = 128,
|
||||
intermediate_size = 18432,
|
||||
vocab_size = 129280,
|
||||
tie_word_embeddings = False,
|
||||
num_experts = 256,
|
||||
moe_intermediate_size = 2048,
|
||||
n_shared_experts = 1,
|
||||
num_dense_layers = 3,
|
||||
q_lora_rank = 1536,
|
||||
kv_lora_rank = 512,
|
||||
qk_nope_head_dim = 128,
|
||||
qk_rope_head_dim = 64,
|
||||
v_head_dim = 128,
|
||||
)
|
||||
|
||||
QWEN3_MOE_30B = ModelArchConfig(
|
||||
hidden_size = 2048,
|
||||
num_hidden_layers = 48,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 4,
|
||||
intermediate_size = 8192,
|
||||
vocab_size = 151936,
|
||||
tie_word_embeddings = True,
|
||||
num_experts = 128,
|
||||
moe_intermediate_size = 768,
|
||||
n_shared_experts = 0,
|
||||
num_dense_layers = 0,
|
||||
)
|
||||
|
||||
GLM4_MOE = ModelArchConfig(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 46,
|
||||
num_attention_heads = 96,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 10944,
|
||||
vocab_size = 151552,
|
||||
tie_word_embeddings = False,
|
||||
num_experts = 128,
|
||||
moe_intermediate_size = 1408,
|
||||
n_shared_experts = 1,
|
||||
num_dense_layers = 1,
|
||||
)
|
||||
|
||||
GPT_OSS = ModelArchConfig(
|
||||
hidden_size = 6144,
|
||||
num_hidden_layers = 64,
|
||||
num_attention_heads = 64,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 2880,
|
||||
vocab_size = 200064,
|
||||
tie_word_embeddings = False,
|
||||
num_experts = 128,
|
||||
moe_intermediate_size = None,
|
||||
n_shared_experts = 0,
|
||||
num_dense_layers = 0,
|
||||
)
|
||||
|
||||
|
||||
class TestExtractArchConfig(unittest.TestCase):
|
||||
def test_basic_config(self):
|
||||
hf_config = SimpleNamespace(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 128256,
|
||||
tie_word_embeddings = False,
|
||||
)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertIsNotNone(arch)
|
||||
self.assertEqual(arch.hidden_size, 4096)
|
||||
self.assertEqual(arch.num_hidden_layers, 32)
|
||||
self.assertEqual(arch.num_key_value_heads, 8)
|
||||
self.assertIsNone(arch.num_experts)
|
||||
|
||||
def test_vlm_text_config(self):
|
||||
text_cfg = SimpleNamespace(
|
||||
hidden_size = 2048,
|
||||
num_hidden_layers = 24,
|
||||
num_attention_heads = 16,
|
||||
num_key_value_heads = 4,
|
||||
intermediate_size = 8192,
|
||||
vocab_size = 32000,
|
||||
tie_word_embeddings = True,
|
||||
)
|
||||
hf_config = SimpleNamespace(text_config = text_cfg)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertIsNotNone(arch)
|
||||
self.assertEqual(arch.hidden_size, 2048)
|
||||
|
||||
def test_moe_detection(self):
|
||||
hf_config = SimpleNamespace(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 32000,
|
||||
tie_word_embeddings = False,
|
||||
num_local_experts = 8,
|
||||
)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertEqual(arch.num_experts, 8)
|
||||
|
||||
def test_missing_fields_returns_none(self):
|
||||
hf_config = SimpleNamespace(hidden_size = 4096)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertIsNone(arch)
|
||||
|
||||
def test_intermediate_size_list(self):
|
||||
hf_config = SimpleNamespace(
|
||||
hidden_size = 2048,
|
||||
num_hidden_layers = 24,
|
||||
num_attention_heads = 16,
|
||||
num_key_value_heads = 4,
|
||||
intermediate_size = [8192, 8192],
|
||||
vocab_size = 32000,
|
||||
tie_word_embeddings = True,
|
||||
)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertEqual(arch.intermediate_size, 8192)
|
||||
|
||||
|
||||
class TestModelWeightsBytes(unittest.TestCase):
|
||||
def test_llama_8b_fp16(self):
|
||||
weight_bytes = compute_model_weights_bytes(LLAMA_8B, "full", False)
|
||||
weight_gb = _gb(weight_bytes)
|
||||
self.assertGreater(weight_gb, 14.0)
|
||||
self.assertLess(weight_gb, 18.0)
|
||||
|
||||
def test_llama_8b_qlora_4bit(self):
|
||||
weight_bytes = compute_model_weights_bytes(LLAMA_8B, "qlora", True)
|
||||
weight_gb = _gb(weight_bytes)
|
||||
self.assertGreater(weight_gb, 4.0)
|
||||
self.assertLess(weight_gb, 7.0)
|
||||
|
||||
def test_4bit_smaller_than_fp16(self):
|
||||
fp16 = compute_model_weights_bytes(LLAMA_8B, "full", False)
|
||||
q4 = compute_model_weights_bytes(LLAMA_8B, "qlora", True)
|
||||
self.assertLess(q4, fp16)
|
||||
ratio = fp16 / q4
|
||||
self.assertGreater(ratio, 2.0)
|
||||
self.assertLess(ratio, 4.0)
|
||||
|
||||
def test_moe_larger_than_dense(self):
|
||||
dense = compute_model_weights_bytes(LLAMA_8B, "full", False)
|
||||
moe = compute_model_weights_bytes(MOE_CONFIG, "full", False)
|
||||
self.assertGreater(moe, dense * 3)
|
||||
|
||||
|
||||
class TestLoraParams(unittest.TestCase):
|
||||
def test_llama_8b_default_modules_rank16(self):
|
||||
lora_p = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES)
|
||||
total_p = compute_total_params(LLAMA_8B)
|
||||
ratio = lora_p / total_p
|
||||
self.assertGreater(ratio, 0.005)
|
||||
self.assertLess(ratio, 0.05)
|
||||
|
||||
def test_higher_rank_more_params(self):
|
||||
r16 = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES)
|
||||
r64 = compute_lora_params(LLAMA_8B, 64, DEFAULT_TARGET_MODULES)
|
||||
self.assertAlmostEqual(r64 / r16, 4.0, places = 1)
|
||||
|
||||
def test_fewer_modules_fewer_params(self):
|
||||
all_mods = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES)
|
||||
qv_only = compute_lora_params(LLAMA_8B, 16, ["q_proj", "v_proj"])
|
||||
self.assertLess(qv_only, all_mods)
|
||||
|
||||
def test_moe_mlp_modules_scale_with_experts(self):
|
||||
dense_lora = compute_lora_params(
|
||||
LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"]
|
||||
)
|
||||
moe_lora = compute_lora_params(
|
||||
MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"]
|
||||
)
|
||||
ratio = moe_lora / dense_lora
|
||||
self.assertAlmostEqual(ratio, 8.0, delta = 0.5)
|
||||
|
||||
def test_attention_modules_same_for_moe(self):
|
||||
dense_attn = compute_lora_params(
|
||||
LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]
|
||||
)
|
||||
moe_attn = compute_lora_params(
|
||||
MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]
|
||||
)
|
||||
self.assertEqual(dense_attn, moe_attn)
|
||||
|
||||
|
||||
class TestOptimizerBytes(unittest.TestCase):
|
||||
def test_adamw_8bit(self):
|
||||
self.assertEqual(compute_optimizer_bytes(1_000_000, "adamw_8bit"), 4_000_000)
|
||||
|
||||
def test_adamw_torch(self):
|
||||
self.assertEqual(compute_optimizer_bytes(1_000_000, "adamw_torch"), 6_000_000)
|
||||
|
||||
def test_sgd(self):
|
||||
self.assertEqual(compute_optimizer_bytes(1_000_000, "sgd"), 4_000_000)
|
||||
|
||||
def test_unknown_defaults_to_4(self):
|
||||
self.assertEqual(compute_optimizer_bytes(1_000_000, "some_new_opt"), 4_000_000)
|
||||
|
||||
|
||||
class TestGradientBytes(unittest.TestCase):
|
||||
def test_fp16_gradients(self):
|
||||
self.assertEqual(compute_gradient_bytes(1_000_000), 2_000_000)
|
||||
|
||||
|
||||
class TestActivationBytes(unittest.TestCase):
|
||||
def test_no_gc_scales_with_layers(self):
|
||||
act_none = compute_activation_bytes(LLAMA_8B, 2, 2048, "none")
|
||||
act_gc = compute_activation_bytes(LLAMA_8B, 2, 2048, "true")
|
||||
self.assertGreater(act_none, act_gc * 10)
|
||||
|
||||
def test_unsloth_gc_smaller_than_standard(self):
|
||||
act_true = compute_activation_bytes(LLAMA_8B, 2, 2048, "true")
|
||||
act_unsloth = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth")
|
||||
self.assertLess(act_unsloth, act_true)
|
||||
|
||||
def test_lora_activations_smaller_than_full_ft(self):
|
||||
full_ft = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth", is_lora = False)
|
||||
lora = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth", is_lora = True)
|
||||
self.assertLess(lora, full_ft)
|
||||
|
||||
def test_scales_with_batch_size(self):
|
||||
act_bsz2 = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth")
|
||||
act_bsz4 = compute_activation_bytes(LLAMA_8B, 4, 2048, "unsloth")
|
||||
self.assertAlmostEqual(act_bsz4 / act_bsz2, 2.0, delta = 0.1)
|
||||
|
||||
def test_scales_with_seq_len(self):
|
||||
act_2k = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth")
|
||||
act_4k = compute_activation_bytes(LLAMA_8B, 2, 4096, "unsloth")
|
||||
self.assertAlmostEqual(act_4k / act_2k, 2.0, delta = 0.1)
|
||||
|
||||
|
||||
class TestEstimateTrainingVram(unittest.TestCase):
|
||||
def test_llama_8b_qlora_reasonable_total(self):
|
||||
config = TrainingVramConfig(
|
||||
training_method = "qlora",
|
||||
batch_size = 2,
|
||||
max_seq_length = 2048,
|
||||
lora_rank = 16,
|
||||
gradient_checkpointing = "unsloth",
|
||||
optimizer = "adamw_8bit",
|
||||
load_in_4bit = True,
|
||||
)
|
||||
breakdown = estimate_training_vram(LLAMA_8B, config)
|
||||
total_gb = _gb(breakdown.total)
|
||||
self.assertGreater(total_gb, 5.0)
|
||||
self.assertLess(total_gb, 12.0)
|
||||
|
||||
def test_llama_8b_full_ft_reasonable_total(self):
|
||||
config = TrainingVramConfig(
|
||||
training_method = "full",
|
||||
batch_size = 2,
|
||||
max_seq_length = 2048,
|
||||
gradient_checkpointing = "unsloth",
|
||||
optimizer = "adamw_8bit",
|
||||
load_in_4bit = False,
|
||||
)
|
||||
breakdown = estimate_training_vram(LLAMA_8B, config)
|
||||
total_gb = _gb(breakdown.total)
|
||||
self.assertGreater(total_gb, 50.0)
|
||||
self.assertLess(total_gb, 75.0)
|
||||
|
||||
def test_qlora_much_less_than_full_ft(self):
|
||||
qlora_config = TrainingVramConfig(
|
||||
training_method = "qlora",
|
||||
load_in_4bit = True,
|
||||
batch_size = 2,
|
||||
max_seq_length = 2048,
|
||||
)
|
||||
full_config = TrainingVramConfig(
|
||||
training_method = "full",
|
||||
load_in_4bit = False,
|
||||
batch_size = 2,
|
||||
max_seq_length = 2048,
|
||||
)
|
||||
qlora = estimate_training_vram(LLAMA_8B, qlora_config)
|
||||
full = estimate_training_vram(LLAMA_8B, full_config)
|
||||
self.assertLess(qlora.total, full.total / 3)
|
||||
|
||||
def test_qwen_05b_qlora_fits_in_4gb(self):
|
||||
config = TrainingVramConfig(
|
||||
training_method = "qlora",
|
||||
batch_size = 2,
|
||||
max_seq_length = 2048,
|
||||
lora_rank = 16,
|
||||
gradient_checkpointing = "unsloth",
|
||||
optimizer = "adamw_8bit",
|
||||
load_in_4bit = True,
|
||||
)
|
||||
breakdown = estimate_training_vram(QWEN_05B, config)
|
||||
total_gb = _gb(breakdown.total)
|
||||
self.assertLess(total_gb, 5.0)
|
||||
|
||||
def test_breakdown_components_positive(self):
|
||||
config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True)
|
||||
breakdown = estimate_training_vram(LLAMA_8B, config)
|
||||
self.assertGreater(breakdown.model_weights, 0)
|
||||
self.assertGreater(breakdown.lora_adapters, 0)
|
||||
self.assertGreater(breakdown.optimizer_states, 0)
|
||||
self.assertGreater(breakdown.gradients, 0)
|
||||
self.assertGreater(breakdown.activations, 0)
|
||||
self.assertGreater(breakdown.cuda_overhead, 0)
|
||||
|
||||
def test_full_ft_no_lora_adapters(self):
|
||||
config = TrainingVramConfig(training_method = "full", load_in_4bit = False)
|
||||
breakdown = estimate_training_vram(LLAMA_8B, config)
|
||||
self.assertEqual(breakdown.lora_adapters, 0)
|
||||
|
||||
def test_to_gb_dict_keys(self):
|
||||
config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True)
|
||||
breakdown = estimate_training_vram(LLAMA_8B, config)
|
||||
gb_dict = breakdown.to_gb_dict()
|
||||
expected_keys = {
|
||||
"model_weights_gb",
|
||||
"lora_adapters_gb",
|
||||
"optimizer_states_gb",
|
||||
"gradients_gb",
|
||||
"activations_gb",
|
||||
"cuda_overhead_gb",
|
||||
"total_gb",
|
||||
}
|
||||
self.assertEqual(set(gb_dict.keys()), expected_keys)
|
||||
|
||||
def test_total_equals_sum_of_parts(self):
|
||||
config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True)
|
||||
breakdown = estimate_training_vram(LLAMA_8B, config)
|
||||
parts_sum = (
|
||||
breakdown.model_weights
|
||||
+ breakdown.lora_adapters
|
||||
+ breakdown.optimizer_states
|
||||
+ breakdown.gradients
|
||||
+ breakdown.activations
|
||||
+ breakdown.cuda_overhead
|
||||
)
|
||||
self.assertEqual(breakdown.total, parts_sum)
|
||||
|
||||
def test_larger_batch_increases_total(self):
|
||||
small = TrainingVramConfig(
|
||||
training_method = "qlora",
|
||||
load_in_4bit = True,
|
||||
batch_size = 1,
|
||||
)
|
||||
large = TrainingVramConfig(
|
||||
training_method = "qlora",
|
||||
load_in_4bit = True,
|
||||
batch_size = 8,
|
||||
)
|
||||
small_v = estimate_training_vram(LLAMA_8B, small)
|
||||
large_v = estimate_training_vram(LLAMA_8B, large)
|
||||
self.assertGreater(large_v.total, small_v.total)
|
||||
|
||||
def test_adamw_fp32_uses_more_optimizer_memory(self):
|
||||
opt8 = TrainingVramConfig(
|
||||
training_method = "full",
|
||||
load_in_4bit = False,
|
||||
optimizer = "adamw_8bit",
|
||||
)
|
||||
opt32 = TrainingVramConfig(
|
||||
training_method = "full",
|
||||
load_in_4bit = False,
|
||||
optimizer = "adamw_torch",
|
||||
)
|
||||
v8 = estimate_training_vram(LLAMA_8B, opt8)
|
||||
v32 = estimate_training_vram(LLAMA_8B, opt32)
|
||||
self.assertAlmostEqual(
|
||||
v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1
|
||||
)
|
||||
|
||||
|
||||
class TestExtractArchConfigMoE(unittest.TestCase):
|
||||
def test_deepseek_v3_shared_experts(self):
|
||||
hf_config = SimpleNamespace(
|
||||
hidden_size = 7168,
|
||||
num_hidden_layers = 61,
|
||||
num_attention_heads = 128,
|
||||
num_key_value_heads = 128,
|
||||
intermediate_size = 18432,
|
||||
vocab_size = 129280,
|
||||
tie_word_embeddings = False,
|
||||
n_routed_experts = 256,
|
||||
moe_intermediate_size = 2048,
|
||||
n_shared_experts = 1,
|
||||
first_k_dense_replace = 3,
|
||||
q_lora_rank = 1536,
|
||||
kv_lora_rank = 512,
|
||||
qk_nope_head_dim = 128,
|
||||
qk_rope_head_dim = 64,
|
||||
v_head_dim = 128,
|
||||
)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertEqual(arch.num_experts, 256)
|
||||
self.assertEqual(arch.n_shared_experts, 1)
|
||||
self.assertEqual(arch.num_dense_layers, 3)
|
||||
self.assertEqual(arch.q_lora_rank, 1536)
|
||||
self.assertEqual(arch.kv_lora_rank, 512)
|
||||
|
||||
def test_qwen3_moe_decoder_sparse_step(self):
|
||||
hf_config = SimpleNamespace(
|
||||
hidden_size = 2048,
|
||||
num_hidden_layers = 48,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 4,
|
||||
intermediate_size = 8192,
|
||||
vocab_size = 151936,
|
||||
tie_word_embeddings = True,
|
||||
num_local_experts = 128,
|
||||
moe_intermediate_size = 768,
|
||||
decoder_sparse_step = 1,
|
||||
mlp_only_layers = [],
|
||||
)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertEqual(arch.num_experts, 128)
|
||||
self.assertEqual(arch.num_dense_layers, 0)
|
||||
self.assertIsNone(arch.q_lora_rank)
|
||||
|
||||
def test_qwen3_moe_with_mlp_only_layers(self):
|
||||
hf_config = SimpleNamespace(
|
||||
hidden_size = 2048,
|
||||
num_hidden_layers = 24,
|
||||
num_attention_heads = 16,
|
||||
num_key_value_heads = 4,
|
||||
intermediate_size = 8192,
|
||||
vocab_size = 151936,
|
||||
tie_word_embeddings = True,
|
||||
num_local_experts = 60,
|
||||
moe_intermediate_size = 1408,
|
||||
decoder_sparse_step = 1,
|
||||
mlp_only_layers = [0, 1, 2, 3],
|
||||
)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertEqual(arch.num_dense_layers, 4)
|
||||
|
||||
def test_glm4_moe_first_k_dense(self):
|
||||
hf_config = SimpleNamespace(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 46,
|
||||
num_attention_heads = 96,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 10944,
|
||||
vocab_size = 151552,
|
||||
tie_word_embeddings = False,
|
||||
n_routed_experts = 128,
|
||||
moe_intermediate_size = 1408,
|
||||
n_shared_experts = 1,
|
||||
first_k_dense_replace = 1,
|
||||
)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertEqual(arch.num_dense_layers, 1)
|
||||
self.assertEqual(arch.n_shared_experts, 1)
|
||||
|
||||
def test_gpt_oss_no_moe_intermediate(self):
|
||||
hf_config = SimpleNamespace(
|
||||
hidden_size = 6144,
|
||||
num_hidden_layers = 64,
|
||||
num_attention_heads = 64,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 2880,
|
||||
vocab_size = 200064,
|
||||
tie_word_embeddings = False,
|
||||
num_local_experts = 128,
|
||||
)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertEqual(arch.num_experts, 128)
|
||||
self.assertIsNone(arch.moe_intermediate_size)
|
||||
self.assertEqual(arch.num_dense_layers, 0)
|
||||
|
||||
def test_backward_compat_no_new_fields(self):
|
||||
hf_config = SimpleNamespace(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 128256,
|
||||
tie_word_embeddings = False,
|
||||
)
|
||||
arch = extract_arch_config(hf_config)
|
||||
self.assertEqual(arch.n_shared_experts, 0)
|
||||
self.assertEqual(arch.num_dense_layers, 0)
|
||||
self.assertIsNone(arch.q_lora_rank)
|
||||
|
||||
|
||||
class TestSharedExperts(unittest.TestCase):
|
||||
def test_shared_experts_increase_weight_bytes(self):
|
||||
no_shared = ModelArchConfig(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 32000,
|
||||
tie_word_embeddings = False,
|
||||
num_experts = 64,
|
||||
moe_intermediate_size = 1407,
|
||||
n_shared_experts = 0,
|
||||
)
|
||||
with_shared = ModelArchConfig(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 32000,
|
||||
tie_word_embeddings = False,
|
||||
num_experts = 64,
|
||||
moe_intermediate_size = 1407,
|
||||
n_shared_experts = 2,
|
||||
)
|
||||
w_no = compute_model_weights_bytes(no_shared, "full", False)
|
||||
w_yes = compute_model_weights_bytes(with_shared, "full", False)
|
||||
self.assertGreater(w_yes, w_no)
|
||||
delta_per_layer = 4096 * 1407 * 3 * 2
|
||||
expected_delta = delta_per_layer * 32 * 2
|
||||
actual_delta = w_yes - w_no
|
||||
self.assertAlmostEqual(
|
||||
actual_delta, expected_delta, delta = expected_delta * 0.01
|
||||
)
|
||||
|
||||
def test_deepseek_v3_params_in_range(self):
|
||||
total = compute_total_params(DEEPSEEK_V3)
|
||||
total_b = total / 1e9
|
||||
self.assertGreater(total_b, 600)
|
||||
self.assertLess(total_b, 750)
|
||||
|
||||
|
||||
class TestMLA(unittest.TestCase):
|
||||
def test_mla_different_from_standard(self):
|
||||
from utils.hardware.vram_estimation import _compute_attn_elements
|
||||
|
||||
mla_arch = DEEPSEEK_V3
|
||||
std_arch = ModelArchConfig(
|
||||
hidden_size = 7168,
|
||||
num_hidden_layers = 61,
|
||||
num_attention_heads = 128,
|
||||
num_key_value_heads = 128,
|
||||
intermediate_size = 18432,
|
||||
vocab_size = 129280,
|
||||
)
|
||||
mla_attn = _compute_attn_elements(mla_arch)
|
||||
std_attn = _compute_attn_elements(std_arch)
|
||||
self.assertNotEqual(mla_attn, std_attn)
|
||||
|
||||
def test_mla_lora_produces_values(self):
|
||||
lora_p = compute_lora_params(DEEPSEEK_V3, 16, ["q_proj", "v_proj", "o_proj"])
|
||||
self.assertGreater(lora_p, 0)
|
||||
|
||||
|
||||
class TestDenseMoEMix(unittest.TestCase):
|
||||
def test_dense_layers_change_total(self):
|
||||
all_moe = ModelArchConfig(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 46,
|
||||
num_attention_heads = 96,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 10944,
|
||||
vocab_size = 151552,
|
||||
tie_word_embeddings = False,
|
||||
num_experts = 128,
|
||||
moe_intermediate_size = 1408,
|
||||
n_shared_experts = 1,
|
||||
num_dense_layers = 0,
|
||||
)
|
||||
mixed = ModelArchConfig(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 46,
|
||||
num_attention_heads = 96,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 10944,
|
||||
vocab_size = 151552,
|
||||
tie_word_embeddings = False,
|
||||
num_experts = 128,
|
||||
moe_intermediate_size = 1408,
|
||||
n_shared_experts = 1,
|
||||
num_dense_layers = 1,
|
||||
)
|
||||
w_all = compute_model_weights_bytes(all_moe, "full", False)
|
||||
w_mixed = compute_model_weights_bytes(mixed, "full", False)
|
||||
self.assertNotEqual(w_all, w_mixed)
|
||||
|
||||
def test_glm4_moe_params_reasonable(self):
|
||||
total = compute_total_params(GLM4_MOE)
|
||||
total_b = total / 1e9
|
||||
self.assertGreater(total_b, 80)
|
||||
self.assertLess(total_b, 120)
|
||||
|
||||
def test_qwen3_moe_30b_params_reasonable(self):
|
||||
total = compute_total_params(QWEN3_MOE_30B)
|
||||
total_b = total / 1e9
|
||||
self.assertGreater(total_b, 20)
|
||||
self.assertLess(total_b, 50)
|
||||
|
||||
def test_gpt_oss_uses_intermediate_size(self):
|
||||
total = compute_total_params(GPT_OSS)
|
||||
total_b = total / 1e9
|
||||
self.assertGreater(total_b, 350)
|
||||
self.assertLess(total_b, 500)
|
||||
|
||||
def test_lora_dense_vs_moe_layers_differ(self):
|
||||
all_moe = ModelArchConfig(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 10,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 32000,
|
||||
tie_word_embeddings = False,
|
||||
num_experts = 8,
|
||||
moe_intermediate_size = 1024,
|
||||
num_dense_layers = 0,
|
||||
)
|
||||
mixed = ModelArchConfig(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 10,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 32000,
|
||||
tie_word_embeddings = False,
|
||||
num_experts = 8,
|
||||
moe_intermediate_size = 1024,
|
||||
num_dense_layers = 5,
|
||||
)
|
||||
lora_all = compute_lora_params(
|
||||
all_moe, 16, ["gate_proj", "up_proj", "down_proj"]
|
||||
)
|
||||
lora_mix = compute_lora_params(mixed, 16, ["gate_proj", "up_proj", "down_proj"])
|
||||
self.assertNotEqual(lora_all, lora_mix)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
62
studio/backend/tests/tool_calling_benchmark_results.md
Normal file
62
studio/backend/tests/tool_calling_benchmark_results.md
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
# GGUF Tool Calling Benchmark Results
|
||||
|
||||
Prompt: "List and categorize all the songs that charted #3 on the Billboard Hot 100 in 2015."
|
||||
10 runs per configuration, web search + code execution + thinking enabled.
|
||||
GPU: NVIDIA B200, CUDA_VISIBLE_DEVICES=2.
|
||||
|
||||
Ground truth: 4 songs peaked at #3 in 2015 -- "Love Me like You Do" (Ellie Goulding), "Earned It" (The Weeknd), "Watch Me" (Silento), "Drag Me Down" (One Direction).
|
||||
|
||||
## Cartesian Grid: Model x Quant x KV Cache
|
||||
|
||||
| Model | Quant | KV Cache | OK/10 | Avg Time | Avg Tools | XML Leaks | URL Fetch | Peak3 Avg | All 4/4 | Best Songs |
|
||||
|-------|-------|----------|-------|----------|-----------|-----------|-----------|-----------|---------|------------|
|
||||
| 4B | UD-Q4_K_XL | f16 | 10/10 | 9.8s | 3.5 | 0/10 | 4/10 | 0.8/4 | 2/10 | 9 |
|
||||
| 4B | UD-Q4_K_XL | bf16 | 10/10 | 10.6s | 4.5 | 0/10 | 4/10 | 0.4/4 | 1/10 | 5 |
|
||||
| 4B | Q8_0 | f16 | 10/10 | 4.9s | 2.4 | 0/10 | 8/10 | 0.4/4 | 1/10 | 5 |
|
||||
| 4B | Q8_0 | bf16 | 10/10 | 8.0s | 3.0 | 0/10 | 5/10 | 0.0/4 | 0/10 | 0 |
|
||||
| 9B | UD-Q4_K_XL | f16 | 10/10 | 6.7s | 2.0 | 0/10 | 5/10 | 0.0/4 | 0/10 | 3 |
|
||||
| 9B | UD-Q4_K_XL | bf16 | 9/10 | 49.5s | 2.4 | 0/10 | 5/10 | 0.0/4 | 0/10 | 1 |
|
||||
| 9B | Q8_0 | f16 | 10/10 | 7.4s | 2.5 | 0/10 | 5/10 | 0.0/4 | 0/10 | 2 |
|
||||
| 9B | Q8_0 | bf16 | 10/10 | 10.4s | 2.7 | 0/10 | 6/10 | 1.0/4 | 2/10 | 15 |
|
||||
| **27B** | **UD-Q4_K_XL** | **bf16** | **9/10** | **131.1s** | **13.8** | **0/10** | **7/10** | **2.7/4** | **6/10** | **27** |
|
||||
| 27B | UD-Q4_K_XL | f16 | 7/10 | 201.6s | 14.1 | 0/10 | 8/10 | 2.0/4 | 5/10 | 26 |
|
||||
| 27B | Q8_0 | f16 | 4/10 | 312.5s | 16.0 | 1/10 | 10/10 | 2.4/4 | 6/10 | 28 |
|
||||
| 27B | Q8_0 | bf16 | 5/10 | 258.4s | 16.5 | 2/10 | 10/10 | 0.9/4 | 1/10 | 27 |
|
||||
| 35B-A3B | UD-Q4_K_XL | f16 | 3/10 | 353.6s | 14.7 | 1/10 | 6/10 | 1.2/4 | 3/10 | 27 |
|
||||
| 35B-A3B | UD-Q4_K_XL | bf16 | 3/10 | 356.2s | 17.2 | 1/10 | 8/10 | 1.6/4 | 4/10 | 27 |
|
||||
| 35B-A3B | Q8_0 | f16 | 2/10 | 372.1s | 17.6 | 1/10 | 7/10 | 1.2/4 | 3/10 | 26 |
|
||||
| 35B-A3B | Q8_0 | bf16 | 6/10 | 267.7s | 17.5 | 1/10 | 8/10 | 2.4/4 | 6/10 | 27 |
|
||||
|
||||
**Column definitions:**
|
||||
- **Peak3 Avg**: Average number of correct peak-#3 songs found per run (out of 4)
|
||||
- **All 4/4**: Runs where all 4 correct songs were identified
|
||||
- **Best Songs**: Maximum number of Billboard 2015 songs mentioned in any single run (out of 31 tracked)
|
||||
- **URL Fetch**: Runs where the model used web_search with `url` parameter to fetch full page content
|
||||
|
||||
## Key Findings
|
||||
|
||||
1. **27B UD-Q4_K_XL + bf16 KV is the sweet spot.** 6/10 runs found all 4 correct songs, 0 XML leaks, 131s average. Best balance of accuracy, speed, and reliability.
|
||||
|
||||
2. **Larger models use tools more effectively.** 27B and 35B-A3B models used 13-17 tool calls per query (vs 2-4 for 4B/9B), performing multiple searches and URL fetches to find the answer.
|
||||
|
||||
3. **27B Q8_0 had the highest raw accuracy (6/10 all-4/4) but lower reliability** -- only 4/10 OK runs due to timeouts on long agentic chains. The UD-Q4_K_XL quant is more practical.
|
||||
|
||||
4. **4B models were fastest (5-10s) but least accurate.** They occasionally found all 4 songs (2/10 best case) when they happened to fetch the right Wikipedia page.
|
||||
|
||||
5. **9B was surprisingly weaker than 4B on this task.** It used fewer tool calls and rarely extracted song data from fetched pages. The 9B model may need higher temperature or different prompting for this specific task type.
|
||||
|
||||
6. **35B-A3B had reliability issues.** Most runs timed out or errored due to slow per-token generation with many tool iterations. When it completed (2-6/10 OK), accuracy was comparable to 27B.
|
||||
|
||||
7. **bf16 KV cache had mixed effects.** For 27B it improved both speed (131s vs 202s) and accuracy (6/10 vs 5/10 all-4/4). For smaller models it had no consistent benefit.
|
||||
|
||||
8. **XML leaks are nearly eliminated.** 0/10 for all 4B and 9B configs, and only 1-2/10 for the largest models (which generate much more text in complex agentic loops).
|
||||
|
||||
## Before vs After (4B UD-Q4_K_XL, f16 KV)
|
||||
|
||||
| Metric | Before Changes | After Changes |
|
||||
|--------|---------------|---------------|
|
||||
| XML leaks | 10/10 | 0/10 |
|
||||
| URL fetches | 0/10 | 4/10 |
|
||||
| Peak3 accuracy | 0.0/4 | 0.8/4 |
|
||||
| Runs with all 4 songs | 0/10 | 2/10 |
|
||||
| Avg time | 12.3s | 9.8s |
|
||||
|
|
@ -8,8 +8,6 @@ This module contains functions for applying chat templates to datasets
|
|||
and generating dataset info summaries.
|
||||
"""
|
||||
|
||||
from torch.utils.data import IterableDataset
|
||||
|
||||
from .format_detection import detect_dataset_format, detect_multimodal_dataset, detect_custom_format_heuristic
|
||||
from .model_mappings import MODEL_TO_TEMPLATE_MAPPER
|
||||
from loggers import get_logger
|
||||
|
|
@ -290,7 +288,13 @@ def apply_chat_template_to_dataset(
|
|||
'batch_size': batch_size,
|
||||
}
|
||||
|
||||
if not isinstance(dataset, IterableDataset):
|
||||
try:
|
||||
from torch.utils.data import IterableDataset
|
||||
_is_torch_iterable = isinstance(dataset, IterableDataset)
|
||||
except ImportError:
|
||||
_is_torch_iterable = False
|
||||
|
||||
if not _is_torch_iterable:
|
||||
from utils.hardware import dataset_map_num_proc
|
||||
if num_proc is None or type(num_proc) is not int:
|
||||
num_proc = dataset_map_num_proc()
|
||||
|
|
@ -351,12 +355,18 @@ def apply_chat_template_to_dataset(
|
|||
return {"text": texts}
|
||||
|
||||
try:
|
||||
try:
|
||||
from torch.utils.data import IterableDataset
|
||||
_is_torch_iterable = isinstance(dataset, IterableDataset)
|
||||
except ImportError:
|
||||
_is_torch_iterable = False
|
||||
|
||||
dataset_map_kwargs = {
|
||||
'batched': True,
|
||||
'batch_size': batch_size,
|
||||
}
|
||||
|
||||
if not isinstance(dataset, IterableDataset):
|
||||
if not _is_torch_iterable:
|
||||
from utils.hardware import dataset_map_num_proc
|
||||
if num_proc is None or type(num_proc) is not int:
|
||||
num_proc = dataset_map_num_proc()
|
||||
|
|
@ -367,7 +377,7 @@ def apply_chat_template_to_dataset(
|
|||
|
||||
# Monitor tqdm progress from dataset.map() and relay to callback
|
||||
_tqdm_monitor_stop = None
|
||||
if progress_callback and not isinstance(dataset, IterableDataset):
|
||||
if progress_callback and not _is_torch_iterable:
|
||||
import threading
|
||||
from tqdm.auto import tqdm as _tqdm_cls
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ This module contains custom data collators for training,
|
|||
particularly for VLM/OCR processing.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, List, Optional, Union
|
||||
from loggers import get_logger
|
||||
|
|
|
|||
|
|
@ -149,7 +149,12 @@ def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None):
|
|||
- "messages" or "conversations" column
|
||||
- "role"/"content" (standard) or "from"/"value" (ShareGPT)
|
||||
"""
|
||||
from torch.utils.data import IterableDataset
|
||||
try:
|
||||
from torch.utils.data import IterableDataset
|
||||
|
||||
_is_torch_iterable = isinstance(dataset, IterableDataset)
|
||||
except ImportError:
|
||||
_is_torch_iterable = False
|
||||
|
||||
def _convert(examples):
|
||||
# Auto-detect which column name is used
|
||||
|
|
@ -196,7 +201,7 @@ def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None):
|
|||
"batch_size": batch_size,
|
||||
}
|
||||
|
||||
if not isinstance(dataset, IterableDataset):
|
||||
if not _is_torch_iterable:
|
||||
from utils.hardware import dataset_map_num_proc
|
||||
|
||||
if num_proc is None or type(num_proc) is not int:
|
||||
|
|
@ -216,7 +221,12 @@ def convert_alpaca_to_chatml(dataset, batch_size = 1000, num_proc = None):
|
|||
|
||||
Output format: Uses 'conversations' column with standard 'role'/'content' structure.
|
||||
"""
|
||||
from torch.utils.data import IterableDataset
|
||||
try:
|
||||
from torch.utils.data import IterableDataset
|
||||
|
||||
_is_torch_iterable = isinstance(dataset, IterableDataset)
|
||||
except ImportError:
|
||||
_is_torch_iterable = False
|
||||
|
||||
def _convert(examples):
|
||||
conversations = []
|
||||
|
|
@ -246,7 +256,7 @@ def convert_alpaca_to_chatml(dataset, batch_size = 1000, num_proc = None):
|
|||
"batch_size": batch_size,
|
||||
}
|
||||
|
||||
if not isinstance(dataset, IterableDataset):
|
||||
if not _is_torch_iterable:
|
||||
from utils.hardware import dataset_map_num_proc
|
||||
|
||||
if num_proc is None or type(num_proc) is not int:
|
||||
|
|
@ -543,13 +553,9 @@ def convert_to_vlm_format(
|
|||
batch_results[idx] = future.result()
|
||||
except Exception as e:
|
||||
failed_count += 1
|
||||
if failed_count == 1:
|
||||
print(
|
||||
f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
|
||||
)
|
||||
if failed_count == 1:
|
||||
logger.info(
|
||||
f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
|
||||
f"First VLM conversion failure: {type(e).__name__}: {e}"
|
||||
)
|
||||
|
||||
converted_list.extend(r for r in batch_results if r is not None)
|
||||
|
|
@ -573,13 +579,10 @@ def convert_to_vlm_format(
|
|||
converted_list.append(_convert_single_sample(sample))
|
||||
except Exception as e:
|
||||
failed_count += 1
|
||||
if failed_count == 1:
|
||||
# Log the first failure to aid debugging
|
||||
print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
|
||||
if failed_count == 1:
|
||||
# Log the first failure to aid debugging
|
||||
logger.info(
|
||||
f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
|
||||
f"First VLM conversion failure: {type(e).__name__}: {e}"
|
||||
)
|
||||
pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False)
|
||||
pbar.close()
|
||||
|
|
|
|||
161
studio/backend/utils/hardware/VRAM_ESTIMATION.md
Normal file
161
studio/backend/utils/hardware/VRAM_ESTIMATION.md
Normal file
|
|
@ -0,0 +1,161 @@
|
|||
# VRAM Estimation for Training
|
||||
|
||||
```
|
||||
Total VRAM = Weights + LoRA Adapters + Optimizer + Gradients + Activations + CUDA Overhead
|
||||
```
|
||||
|
||||
| Symbol | Meaning |
|
||||
|--------|---------|
|
||||
| `H` | `hidden_size` |
|
||||
| `L` | `num_hidden_layers` |
|
||||
| `V` | `vocab_size` |
|
||||
| `K` | `(H / num_attention_heads) * num_key_value_heads` |
|
||||
| `M` | `intermediate_size` (or `moe_intermediate_size`) |
|
||||
| `E` | `num_experts` (1 for dense) |
|
||||
| `r` | LoRA rank |
|
||||
| `B` | `per_device_train_batch_size` |
|
||||
| `S` | `max_seq_length` |
|
||||
|
||||
---
|
||||
|
||||
## 1. Model Weights
|
||||
|
||||
```
|
||||
QKVO = (H + K + K + H) * H
|
||||
MLP = H * M * 3 * E + (E * H if E > 1 else 0)
|
||||
|
||||
Quantizable = (QKVO + MLP) * L
|
||||
Non-quantizable = 2*H*L + V*H + (V*H if not tie_embeddings else 0)
|
||||
```
|
||||
|
||||
| Mode | Bytes |
|
||||
|------|-------|
|
||||
| QLoRA 4-bit | `Quantizable * 2 / 3.2 + Non-quantizable * 2` |
|
||||
| LoRA / Full fp16 | `(Quantizable + Non-quantizable) * 2` |
|
||||
|
||||
The 3.2 factor (`16/5`) accounts for BNB NF4 blockwise scales.
|
||||
|
||||
## 2. LoRA Adapters
|
||||
|
||||
| Module | A | B |
|
||||
|--------|---|---|
|
||||
| q_proj | `H×r` | `r×H` |
|
||||
| k_proj | `H×r` | `r×K` |
|
||||
| v_proj | `H×r` | `r×K` |
|
||||
| o_proj | `H×r` | `r×H` |
|
||||
| gate_proj | `H×r` | `r×M` |
|
||||
| up_proj | `H×r` | `r×M` |
|
||||
| down_proj | `M×r` | `r×H` |
|
||||
|
||||
MLP modules multiply by `E` for MoE.
|
||||
|
||||
```
|
||||
LoRA_bytes = sum(A + B per selected module) * L * 2
|
||||
```
|
||||
|
||||
## 3. Optimizer States (calibrated)
|
||||
|
||||
| Optimizer | Bytes/param | Notes |
|
||||
|-----------|------------|-------|
|
||||
| `adamw_8bit` | 4 | BNB upcasts to fp32 during step |
|
||||
| `adamw_torch` | 6 | Fused, no master copy |
|
||||
| `paged_adamw_32bit` | 8 | Full fp32 states |
|
||||
| `sgd` | 4 | |
|
||||
|
||||
Trainable params = all params (Full FT) or LoRA params only.
|
||||
|
||||
## 4. Gradients
|
||||
|
||||
```
|
||||
Gradient_bytes = trainable_params * 2 (fp16, accumulated in-place)
|
||||
```
|
||||
|
||||
## 5. Activations
|
||||
|
||||
Per-layer (from `unsloth_zoo/vllm_utils.py`):
|
||||
```
|
||||
Per_layer = (S*B*(H+K+K) + S*B*2 + S*B*(M+M)) * 2 * 1.25
|
||||
```
|
||||
|
||||
| GC Mode | Full FT | LoRA/QLoRA |
|
||||
|---------|---------|------------|
|
||||
| none | `L` layers | `L` layers |
|
||||
| true (HF) | 2.0 | 1.0 |
|
||||
| unsloth | 1.5 | 1.0 |
|
||||
|
||||
## 6. Floors
|
||||
|
||||
Gradients and activations have minimum floors at **15% of model weight memory** to account for autograd overhead, attention score matrices, NCCL buffers, mixed-precision scaling, and PyTorch fragmentation.
|
||||
|
||||
```
|
||||
gradient_bytes = max(computed, weights * 0.15)
|
||||
activation_bytes = max(computed, weights * 0.15 * B/2)
|
||||
```
|
||||
|
||||
## 7. CUDA Overhead
|
||||
|
||||
**1.4 GB** fixed — CUDA driver + PyTorch runtime, calibrated on RTX 5070 Ti.
|
||||
|
||||
## 8. Multi-GPU Overhead
|
||||
|
||||
When sharding across multiple GPUs, each additional GPU (beyond the first) contributes only **85%** of its free VRAM to the usable pool. The 15% discount accounts for NCCL all-reduce buffers, PCIe/NVLink transfer overhead, synchronization barriers, and memory fragmentation from non-uniform shard sizes. Calibrated empirically on 2-8 GPU setups with NVLink and PCIe topologies.
|
||||
|
||||
```
|
||||
usable_gb = free[gpu_0] + sum(free[gpu_i] * 0.85 for i in 1..N)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Reference Table (bsz=2, seq=2048, rank=16, GC=unsloth, adamw_8bit)
|
||||
|
||||
| Model | Weights | LoRA | Optim | Grad | Act | CUDA | Total |
|
||||
|-------|---------|------|-------|------|-----|------|-------|
|
||||
| 0.5B QLoRA | 0.5 | 0.0 | 0.0 | 0.1 | 0.1 | 1.4 | **2.1** |
|
||||
| 1B QLoRA | 1.1 | 0.0 | 0.0 | 0.2 | 0.2 | 1.4 | **2.9** |
|
||||
| 3B QLoRA | 2.4 | 0.0 | 0.1 | 0.5 | 0.5 | 1.4 | **4.9** |
|
||||
| 8B QLoRA | 6.0 | 0.1 | 0.2 | 1.2 | 1.2 | 1.4 | **10.1** |
|
||||
| 8B LoRA fp16 | 15.0 | 0.1 | 0.2 | 3.0 | 3.0 | 1.4 | **22.6** |
|
||||
| 8B Full FT | 15.0 | — | 29.9 | 15.0 | 3.0 | 1.4 | **64.2** |
|
||||
| 32B LoRA fp16 | 61.0 | 0.2 | 0.5 | 12.2 | 12.2 | 1.4 | **87.6** |
|
||||
| 72B QLoRA | 45.5 | 0.4 | 0.8 | 9.1 | 9.1 | 1.4 | **66.3** |
|
||||
|
||||
## E2E Validation (Llama-3.2-1B, B200 emulating 24GB)
|
||||
|
||||
| Config | Estimated | Actual (nvsmi) | Error |
|
||||
|--------|----------|----------------|-------|
|
||||
| QLoRA bsz=2 seq=512 | 2.55 GB | 2.65 GB | -3.7% |
|
||||
| QLoRA bsz=2 seq=2048 | 2.60 GB | 2.65 GB | -1.8% |
|
||||
| QLoRA bsz=4 seq=2048 | 2.65 GB | 2.65 GB | +0.0% |
|
||||
| LoRA fp16 bsz=2 | 3.84 GB | 3.88 GB | -1.0% |
|
||||
| Full FT adamw_8bit | 10.89 GB | 10.80 GB | +0.8% |
|
||||
| Full FT adamw_torch | 13.19 GB | 12.93 GB | +2.0% |
|
||||
|
||||
*Note: e2e numbers predate the 15% floors, which add safety margin on top.*
|
||||
|
||||
---
|
||||
|
||||
## Parameter Flow
|
||||
|
||||
```
|
||||
Frontend -> routes/{training,inference}.py
|
||||
-> prepare_gpu_selection(gpu_ids, model_name, ...)
|
||||
|
|
||||
+-- gpu_ids is explicit (e.g. [5,6,7])
|
||||
| -> resolve_requested_gpu_ids: validate against parent-visible set
|
||||
| -> return all requested GPUs (model sharded across all of them)
|
||||
|
|
||||
+-- gpu_ids is None or []
|
||||
-> auto_select_gpu_ids: estimate VRAM, pick minimum GPUs needed
|
||||
-> estimate_required_model_memory_gb -> estimate_training_vram
|
||||
-> greedy selection: rank GPUs by free VRAM, add until model fits
|
||||
|
||||
-> get_device_map(resolved_gpu_ids)
|
||||
-> "balanced" if >1 GPU, "sequential" otherwise
|
||||
|
||||
-> worker subprocess: apply_gpu_ids(resolved_gpu_ids)
|
||||
-> sets CUDA_VISIBLE_DEVICES before torch/CUDA init
|
||||
```
|
||||
|
||||
Threaded params: `batch_size`, `max_seq_length`, `lora_r`, `target_modules`, `gradient_checkpointing`, `optim`.
|
||||
|
||||
Source: `studio/backend/utils/hardware/vram_estimation.py`
|
||||
|
|
@ -18,11 +18,31 @@ from .hardware import (
|
|||
get_gpu_summary,
|
||||
get_package_versions,
|
||||
get_gpu_utilization,
|
||||
get_visible_gpu_utilization,
|
||||
get_backend_visible_gpu_info,
|
||||
get_physical_gpu_count,
|
||||
get_visible_gpu_count,
|
||||
get_parent_visible_gpu_ids,
|
||||
resolve_requested_gpu_ids,
|
||||
estimate_fp16_model_size_bytes,
|
||||
estimate_required_model_memory_gb,
|
||||
auto_select_gpu_ids,
|
||||
prepare_gpu_selection,
|
||||
safe_num_proc,
|
||||
safe_thread_num_proc,
|
||||
dataset_map_num_proc,
|
||||
get_device_map,
|
||||
get_offloaded_device_map_entries,
|
||||
raise_if_offloaded,
|
||||
apply_gpu_ids,
|
||||
)
|
||||
|
||||
from .vram_estimation import (
|
||||
ModelArchConfig,
|
||||
TrainingVramConfig,
|
||||
VramBreakdown,
|
||||
extract_arch_config,
|
||||
estimate_training_vram,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
|
|
@ -38,9 +58,26 @@ __all__ = [
|
|||
"get_gpu_summary",
|
||||
"get_package_versions",
|
||||
"get_gpu_utilization",
|
||||
"get_visible_gpu_utilization",
|
||||
"get_backend_visible_gpu_info",
|
||||
"get_physical_gpu_count",
|
||||
"get_visible_gpu_count",
|
||||
"get_parent_visible_gpu_ids",
|
||||
"resolve_requested_gpu_ids",
|
||||
"estimate_fp16_model_size_bytes",
|
||||
"estimate_required_model_memory_gb",
|
||||
"auto_select_gpu_ids",
|
||||
"prepare_gpu_selection",
|
||||
"safe_num_proc",
|
||||
"safe_thread_num_proc",
|
||||
"dataset_map_num_proc",
|
||||
"get_device_map",
|
||||
"get_offloaded_device_map_entries",
|
||||
"raise_if_offloaded",
|
||||
"apply_gpu_ids",
|
||||
"ModelArchConfig",
|
||||
"TrainingVramConfig",
|
||||
"VramBreakdown",
|
||||
"extract_arch_config",
|
||||
"estimate_training_vram",
|
||||
]
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
279
studio/backend/utils/hardware/nvidia.py
Normal file
279
studio/backend/utils/hardware/nvidia.py
Normal file
|
|
@ -0,0 +1,279 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import subprocess
|
||||
from typing import Any, Optional
|
||||
|
||||
from loggers import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _parse_smi_value(raw: str):
|
||||
raw = raw.strip()
|
||||
if not raw or raw == "[N/A]":
|
||||
return None
|
||||
try:
|
||||
return float(raw)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _build_gpu_metrics(
|
||||
vram_used_mb,
|
||||
vram_total_mb,
|
||||
power_draw,
|
||||
power_limit,
|
||||
**extra,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
**extra,
|
||||
"vram_used_gb": round(vram_used_mb / 1024, 2)
|
||||
if vram_used_mb is not None
|
||||
else None,
|
||||
"vram_total_gb": round(vram_total_mb / 1024, 2)
|
||||
if vram_total_mb is not None
|
||||
else None,
|
||||
"vram_utilization_pct": round((vram_used_mb / vram_total_mb) * 100, 1)
|
||||
if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0
|
||||
else None,
|
||||
"power_draw_w": power_draw,
|
||||
"power_limit_w": power_limit,
|
||||
"power_utilization_pct": round((power_draw / power_limit) * 100, 1)
|
||||
if power_draw is not None and power_limit and power_limit > 0
|
||||
else None,
|
||||
}
|
||||
|
||||
|
||||
def _visible_ordinal_map(
|
||||
parent_visible_ids: Optional[list[int]],
|
||||
) -> Optional[dict[int, int]]:
|
||||
if parent_visible_ids is None:
|
||||
return None
|
||||
return {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)}
|
||||
|
||||
|
||||
def get_physical_gpu_count() -> Optional[int]:
|
||||
"""Return physical GPU count via nvidia-smi, or None on failure."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["nvidia-smi", "-L"],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 5,
|
||||
)
|
||||
if result.returncode == 0 and result.stdout.strip():
|
||||
return len(result.stdout.strip().splitlines())
|
||||
logger.warning(
|
||||
"nvidia-smi -L returned code %d; caller should fall back to torch",
|
||||
result.returncode,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("nvidia-smi -L failed: %s; caller should fall back to torch", e)
|
||||
return None
|
||||
|
||||
|
||||
def get_primary_gpu_utilization() -> dict[str, Any]:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"nvidia-smi",
|
||||
"--query-gpu=utilization.gpu,temperature.gpu,"
|
||||
"memory.used,memory.total,power.draw,power.limit",
|
||||
"--format=csv,noheader,nounits",
|
||||
],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 5,
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
logger.warning("nvidia-smi query failed in get_primary_gpu_utilization: %s", e)
|
||||
return {"available": False}
|
||||
if result.returncode != 0 or not result.stdout.strip():
|
||||
return {"available": False}
|
||||
|
||||
first_line = result.stdout.strip().splitlines()[0]
|
||||
parts = [p.strip() for p in first_line.split(",")]
|
||||
if len(parts) < 6:
|
||||
return {"available": False}
|
||||
|
||||
return _build_gpu_metrics(
|
||||
vram_used_mb = _parse_smi_value(parts[2]),
|
||||
vram_total_mb = _parse_smi_value(parts[3]),
|
||||
power_draw = _parse_smi_value(parts[4]),
|
||||
power_limit = _parse_smi_value(parts[5]),
|
||||
available = True,
|
||||
gpu_utilization_pct = _parse_smi_value(parts[0]),
|
||||
temperature_c = _parse_smi_value(parts[1]),
|
||||
)
|
||||
|
||||
|
||||
def get_visible_gpu_utilization(
|
||||
parent_visible_ids: Optional[list[int]],
|
||||
parent_cuda_visible_devices: Optional[str] = None,
|
||||
) -> dict[str, Any]:
|
||||
# When parent_visible_ids is None (UUID/MIG mask), we cannot safely
|
||||
# map nvidia-smi rows to the process's visible devices. Return empty
|
||||
# instead of exposing all physical GPUs.
|
||||
if parent_visible_ids is None:
|
||||
return {
|
||||
"available": False,
|
||||
"backend_cuda_visible_devices": parent_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": [],
|
||||
"devices": [],
|
||||
"index_kind": "unresolved",
|
||||
}
|
||||
visible_ordinals = _visible_ordinal_map(parent_visible_ids)
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"nvidia-smi",
|
||||
"--query-gpu=index,utilization.gpu,temperature.gpu,"
|
||||
"memory.used,memory.total,power.draw,power.limit",
|
||||
"--format=csv,noheader,nounits",
|
||||
],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 5,
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
logger.warning("nvidia-smi query failed in get_visible_gpu_utilization: %s", e)
|
||||
return {
|
||||
"available": False,
|
||||
"backend_cuda_visible_devices": parent_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": parent_visible_ids or [],
|
||||
"devices": [],
|
||||
"index_kind": "physical",
|
||||
}
|
||||
if result.returncode != 0 or not result.stdout.strip():
|
||||
return {
|
||||
"available": False,
|
||||
"backend_cuda_visible_devices": parent_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": parent_visible_ids or [],
|
||||
"devices": [],
|
||||
"index_kind": "physical",
|
||||
}
|
||||
|
||||
devices = []
|
||||
for line in result.stdout.strip().splitlines():
|
||||
parts = [p.strip() for p in line.split(",")]
|
||||
if len(parts) < 7:
|
||||
continue
|
||||
|
||||
try:
|
||||
idx = int(parts[0])
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
|
||||
if visible_ordinals is not None and idx not in visible_ordinals:
|
||||
continue
|
||||
|
||||
devices.append(
|
||||
_build_gpu_metrics(
|
||||
vram_used_mb = _parse_smi_value(parts[3]),
|
||||
vram_total_mb = _parse_smi_value(parts[4]),
|
||||
power_draw = _parse_smi_value(parts[5]),
|
||||
power_limit = _parse_smi_value(parts[6]),
|
||||
index = idx,
|
||||
index_kind = "physical",
|
||||
visible_ordinal = (
|
||||
visible_ordinals[idx]
|
||||
if visible_ordinals is not None
|
||||
else len(devices)
|
||||
),
|
||||
gpu_utilization_pct = _parse_smi_value(parts[1]),
|
||||
temperature_c = _parse_smi_value(parts[2]),
|
||||
)
|
||||
)
|
||||
|
||||
return {
|
||||
"available": len(devices) > 0,
|
||||
"backend_cuda_visible_devices": parent_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": parent_visible_ids or [],
|
||||
"devices": devices,
|
||||
"index_kind": "physical",
|
||||
}
|
||||
|
||||
|
||||
def get_backend_visible_gpu_info(
|
||||
parent_visible_ids: Optional[list[int]],
|
||||
backend_cuda_visible_devices: Optional[str],
|
||||
) -> dict[str, Any]:
|
||||
# When parent_visible_ids is None (UUID/MIG mask), we cannot safely
|
||||
# map nvidia-smi rows to the process's visible devices.
|
||||
if parent_visible_ids is None:
|
||||
return {
|
||||
"available": False,
|
||||
"backend_cuda_visible_devices": backend_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": [],
|
||||
"devices": [],
|
||||
"index_kind": "unresolved",
|
||||
}
|
||||
visible_ordinals = _visible_ordinal_map(parent_visible_ids)
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"nvidia-smi",
|
||||
"--query-gpu=index,name,memory.total",
|
||||
"--format=csv,noheader,nounits",
|
||||
],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 10,
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
logger.warning("nvidia-smi query failed in get_backend_visible_gpu_info: %s", e)
|
||||
return {
|
||||
"available": False,
|
||||
"backend_cuda_visible_devices": backend_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": parent_visible_ids or [],
|
||||
"devices": [],
|
||||
"index_kind": "physical",
|
||||
}
|
||||
if result.returncode != 0:
|
||||
return {
|
||||
"available": False,
|
||||
"backend_cuda_visible_devices": backend_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": parent_visible_ids or [],
|
||||
"devices": [],
|
||||
"index_kind": "physical",
|
||||
}
|
||||
|
||||
devices = []
|
||||
for line in result.stdout.strip().splitlines():
|
||||
parts = [p.strip() for p in line.split(",")]
|
||||
if len(parts) < 3:
|
||||
continue
|
||||
try:
|
||||
idx = int(parts[0])
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
if visible_ordinals is not None and idx not in visible_ordinals:
|
||||
continue
|
||||
# Use split with limit to handle GPU names containing commas
|
||||
name = parts[1] if len(parts) == 3 else ", ".join(parts[1:-1])
|
||||
try:
|
||||
mem_total_mb = int(parts[-1])
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
devices.append(
|
||||
{
|
||||
"index": idx,
|
||||
"index_kind": "physical",
|
||||
"visible_ordinal": (
|
||||
visible_ordinals[idx]
|
||||
if visible_ordinals is not None
|
||||
else len(devices)
|
||||
),
|
||||
"name": name,
|
||||
"memory_total_gb": round(mem_total_mb / 1024, 2),
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"available": len(devices) > 0,
|
||||
"backend_cuda_visible_devices": backend_cuda_visible_devices,
|
||||
"parent_visible_gpu_ids": parent_visible_ids or [],
|
||||
"devices": devices,
|
||||
"index_kind": "physical",
|
||||
}
|
||||
501
studio/backend/utils/hardware/vram_estimation.py
Normal file
501
studio/backend/utils/hardware/vram_estimation.py
Normal file
|
|
@ -0,0 +1,501 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""
|
||||
Training VRAM estimation.
|
||||
|
||||
Total VRAM = weights + LoRA adapters + optimizer states + gradients
|
||||
+ activations + CUDA overhead.
|
||||
Activation formula from unsloth_zoo/vllm_utils.py.
|
||||
All constants empirically calibrated against Llama-3.2-1B on B200.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, Optional
|
||||
|
||||
QUANT_4BIT_FACTOR = 16 / 5
|
||||
CUDA_OVERHEAD_BYTES = int(1.4 * 1024**3) # calibrated on RTX 5070 Ti
|
||||
|
||||
DEFAULT_TARGET_MODULES = [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
"down_proj",
|
||||
]
|
||||
|
||||
# Empirically calibrated bytes/param — see VRAM_ESTIMATION.md for rationale.
|
||||
OPTIMIZER_BYTES_PER_PARAM: Dict[str, int] = {
|
||||
"adamw_8bit": 4, # BNB upcasts to fp32 during step
|
||||
"paged_adamw_8bit": 4,
|
||||
"adamw_bnb_8bit": 4,
|
||||
"paged_adamw_32bit": 8,
|
||||
"adamw_torch": 6, # fused, no master copy
|
||||
"adamw_torch_fused": 6,
|
||||
"sgd": 4,
|
||||
}
|
||||
|
||||
# (full_ft_multiplier, lora_multiplier) — fraction of num_layers.
|
||||
# LoRA: frozen base layers skip activation storage, but you always need
|
||||
# at least ~1 layer in flight during backprop recomputation.
|
||||
GC_LAYER_MULTIPLIERS = {
|
||||
"none": (None, None),
|
||||
"true": (2.0, 1.0),
|
||||
"unsloth": (1.5, 1.0),
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArchConfig:
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
intermediate_size: int
|
||||
vocab_size: int
|
||||
tie_word_embeddings: bool = True
|
||||
num_experts: Optional[int] = None
|
||||
moe_intermediate_size: Optional[int] = None
|
||||
n_shared_experts: int = 0
|
||||
num_dense_layers: int = 0
|
||||
q_lora_rank: Optional[int] = None
|
||||
kv_lora_rank: Optional[int] = None
|
||||
qk_nope_head_dim: Optional[int] = None
|
||||
qk_rope_head_dim: Optional[int] = None
|
||||
v_head_dim: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainingVramConfig:
|
||||
training_method: str = "qlora"
|
||||
batch_size: int = 4
|
||||
max_seq_length: int = 2048
|
||||
lora_rank: int = 16
|
||||
target_modules: list = field(default_factory = lambda: list(DEFAULT_TARGET_MODULES))
|
||||
gradient_checkpointing: str = "unsloth"
|
||||
optimizer: str = "adamw_8bit"
|
||||
load_in_4bit: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class VramBreakdown:
|
||||
model_weights: int
|
||||
lora_adapters: int
|
||||
optimizer_states: int
|
||||
gradients: int
|
||||
activations: int
|
||||
cuda_overhead: int
|
||||
# The computed (formula-based) activation cost before floors.
|
||||
# This is the true per-layer cost that doesn't shard across GPUs.
|
||||
activations_computed: int = 0
|
||||
|
||||
@property
|
||||
def total(self) -> int:
|
||||
return (
|
||||
self.model_weights
|
||||
+ self.lora_adapters
|
||||
+ self.optimizer_states
|
||||
+ self.gradients
|
||||
+ self.activations
|
||||
+ self.cuda_overhead
|
||||
)
|
||||
|
||||
def min_gpu_vram(self, n_gpus: int) -> int:
|
||||
"""Minimum VRAM a single GPU needs: its shard + non-shardable costs.
|
||||
|
||||
Weights/LoRA/optimizer/gradients shard across GPUs.
|
||||
The computed activation cost does NOT shard (one GPU runs the layer).
|
||||
The floor portion (activations - computed) is overhead that shards.
|
||||
"""
|
||||
shardable = (
|
||||
self.model_weights
|
||||
+ self.lora_adapters
|
||||
+ self.optimizer_states
|
||||
+ self.gradients
|
||||
+ (self.activations - self.activations_computed) # floor overhead shards
|
||||
)
|
||||
per_gpu_fixed = self.activations_computed + self.cuda_overhead
|
||||
return shardable // max(n_gpus, 1) + per_gpu_fixed
|
||||
|
||||
def to_gb_dict(self) -> Dict[str, float]:
|
||||
return {
|
||||
"model_weights_gb": round(self.model_weights / (1024**3), 3),
|
||||
"lora_adapters_gb": round(self.lora_adapters / (1024**3), 3),
|
||||
"optimizer_states_gb": round(self.optimizer_states / (1024**3), 3),
|
||||
"gradients_gb": round(self.gradients / (1024**3), 3),
|
||||
"activations_gb": round(self.activations / (1024**3), 3),
|
||||
"cuda_overhead_gb": round(self.cuda_overhead / (1024**3), 3),
|
||||
"total_gb": round(self.total / (1024**3), 3),
|
||||
}
|
||||
|
||||
|
||||
def _compute_num_dense_layers(text_config, total_layers: int) -> int:
|
||||
"""Count how many layers use dense MLP instead of MoE."""
|
||||
first_k = getattr(text_config, "first_k_dense_replace", None)
|
||||
if first_k is not None:
|
||||
return min(int(first_k), total_layers)
|
||||
|
||||
sparse_step = getattr(text_config, "decoder_sparse_step", None)
|
||||
mlp_only = getattr(text_config, "mlp_only_layers", None) or []
|
||||
if sparse_step is not None and sparse_step > 0:
|
||||
mlp_only_set = set(mlp_only)
|
||||
moe_count = sum(
|
||||
1
|
||||
for i in range(total_layers)
|
||||
if i not in mlp_only_set and (i + 1) % sparse_step == 0
|
||||
)
|
||||
return total_layers - moe_count
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
|
||||
text_config = getattr(hf_config, "text_config", None) or hf_config
|
||||
|
||||
hidden_size = getattr(text_config, "hidden_size", None)
|
||||
num_layers = getattr(text_config, "num_hidden_layers", None)
|
||||
num_heads = getattr(text_config, "num_attention_heads", None)
|
||||
intermediate_size = getattr(text_config, "intermediate_size", None)
|
||||
vocab_size = getattr(text_config, "vocab_size", None)
|
||||
|
||||
if isinstance(intermediate_size, (list, tuple)):
|
||||
intermediate_size = intermediate_size[0] if intermediate_size else None
|
||||
if intermediate_size is None and hidden_size is not None:
|
||||
intermediate_size = hidden_size * 4
|
||||
|
||||
if not all(
|
||||
v is not None
|
||||
for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size)
|
||||
):
|
||||
return None
|
||||
if num_heads <= 0:
|
||||
return None
|
||||
|
||||
num_kv_heads = getattr(text_config, "num_key_value_heads", num_heads)
|
||||
|
||||
num_experts = None
|
||||
for attr in ("num_local_experts", "num_experts", "n_routed_experts"):
|
||||
num_experts = getattr(text_config, attr, None)
|
||||
if num_experts is not None:
|
||||
break
|
||||
|
||||
moe_intermediate = getattr(text_config, "moe_intermediate_size", None)
|
||||
n_shared_experts = getattr(text_config, "n_shared_experts", None) or 0
|
||||
|
||||
num_dense_layers = 0
|
||||
if num_experts is not None and num_experts > 1:
|
||||
num_dense_layers = _compute_num_dense_layers(text_config, num_layers)
|
||||
|
||||
q_lora_rank = getattr(text_config, "q_lora_rank", None)
|
||||
kv_lora_rank = getattr(text_config, "kv_lora_rank", None)
|
||||
qk_nope_head_dim = getattr(text_config, "qk_nope_head_dim", None)
|
||||
qk_rope_head_dim = getattr(text_config, "qk_rope_head_dim", None)
|
||||
v_head_dim = getattr(text_config, "v_head_dim", None)
|
||||
|
||||
return ModelArchConfig(
|
||||
hidden_size = hidden_size,
|
||||
num_hidden_layers = num_layers,
|
||||
num_attention_heads = num_heads,
|
||||
num_key_value_heads = num_kv_heads,
|
||||
intermediate_size = intermediate_size,
|
||||
vocab_size = vocab_size,
|
||||
tie_word_embeddings = getattr(text_config, "tie_word_embeddings", True),
|
||||
num_experts = num_experts,
|
||||
moe_intermediate_size = moe_intermediate,
|
||||
n_shared_experts = n_shared_experts,
|
||||
num_dense_layers = num_dense_layers,
|
||||
q_lora_rank = q_lora_rank,
|
||||
kv_lora_rank = kv_lora_rank,
|
||||
qk_nope_head_dim = qk_nope_head_dim,
|
||||
qk_rope_head_dim = qk_rope_head_dim,
|
||||
v_head_dim = v_head_dim,
|
||||
)
|
||||
|
||||
|
||||
def _get_kv_size(arch: ModelArchConfig) -> int:
|
||||
return (arch.hidden_size // arch.num_attention_heads) * arch.num_key_value_heads
|
||||
|
||||
|
||||
def _get_mlp_size(arch: ModelArchConfig) -> int:
|
||||
if arch.moe_intermediate_size is not None:
|
||||
return arch.moe_intermediate_size
|
||||
return arch.intermediate_size
|
||||
|
||||
|
||||
def _get_num_experts(arch: ModelArchConfig) -> int:
|
||||
return arch.num_experts if arch.num_experts and arch.num_experts > 1 else 1
|
||||
|
||||
|
||||
def _compute_attn_elements(arch: ModelArchConfig) -> int:
|
||||
"""Attention weight elements per layer."""
|
||||
hd = arch.hidden_size
|
||||
if arch.q_lora_rank is not None:
|
||||
nh = arch.num_attention_heads
|
||||
qk_head = arch.qk_nope_head_dim + arch.qk_rope_head_dim
|
||||
q_a = hd * arch.q_lora_rank
|
||||
q_b = arch.q_lora_rank * (nh * qk_head)
|
||||
kv_a = hd * (arch.kv_lora_rank + arch.qk_rope_head_dim)
|
||||
kv_b = arch.kv_lora_rank * (nh * (arch.qk_nope_head_dim + arch.v_head_dim))
|
||||
o = (nh * arch.v_head_dim) * hd
|
||||
norms = arch.q_lora_rank + arch.kv_lora_rank
|
||||
return q_a + q_b + kv_a + kv_b + o + norms
|
||||
kv_size = _get_kv_size(arch)
|
||||
return (hd + kv_size + kv_size + hd) * hd
|
||||
|
||||
|
||||
def _compute_dense_mlp_elements(arch: ModelArchConfig) -> int:
|
||||
return arch.hidden_size * arch.intermediate_size * 3
|
||||
|
||||
|
||||
def _compute_moe_mlp_elements(arch: ModelArchConfig) -> int:
|
||||
hd = arch.hidden_size
|
||||
mlp_size = _get_mlp_size(arch)
|
||||
n_experts = _get_num_experts(arch)
|
||||
return hd * mlp_size * 3 * (n_experts + arch.n_shared_experts) + n_experts * hd
|
||||
|
||||
|
||||
def _compute_layer_elements(arch: ModelArchConfig):
|
||||
"""Return (total_quantizable, layernorms_per_layer, embed, lm_head) element counts.
|
||||
|
||||
total_quantizable is summed across ALL layers (not per-layer).
|
||||
"""
|
||||
hd = arch.hidden_size
|
||||
n_layers = arch.num_hidden_layers
|
||||
n_experts = _get_num_experts(arch)
|
||||
|
||||
attn_total = _compute_attn_elements(arch) * n_layers
|
||||
|
||||
if n_experts > 1:
|
||||
n_dense = arch.num_dense_layers
|
||||
n_moe = n_layers - n_dense
|
||||
mlp_total = (
|
||||
_compute_moe_mlp_elements(arch) * n_moe
|
||||
+ _compute_dense_mlp_elements(arch) * n_dense
|
||||
)
|
||||
else:
|
||||
mlp_total = _compute_dense_mlp_elements(arch) * n_layers
|
||||
|
||||
layernorms = 2 * hd
|
||||
embed_tokens = arch.vocab_size * hd
|
||||
lm_head = 0 if arch.tie_word_embeddings else arch.vocab_size * hd
|
||||
return attn_total + mlp_total, layernorms, embed_tokens, lm_head
|
||||
|
||||
|
||||
def compute_model_weights_bytes(
|
||||
arch: ModelArchConfig,
|
||||
training_method: str,
|
||||
load_in_4bit: bool,
|
||||
) -> int:
|
||||
total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch)
|
||||
n_layers = arch.num_hidden_layers
|
||||
non_quantizable = layernorms * n_layers + embed_tokens + lm_head
|
||||
|
||||
if training_method == "qlora" and load_in_4bit:
|
||||
return int(total_quantizable * 2 / QUANT_4BIT_FACTOR + non_quantizable * 2)
|
||||
|
||||
return int((total_quantizable + non_quantizable) * 2)
|
||||
|
||||
|
||||
def compute_total_params(arch: ModelArchConfig) -> int:
|
||||
total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch)
|
||||
n_layers = arch.num_hidden_layers
|
||||
return total_quantizable + layernorms * n_layers + embed_tokens + lm_head
|
||||
|
||||
|
||||
def _lora_attn_elements(
|
||||
arch: ModelArchConfig,
|
||||
r: int,
|
||||
target_modules: list,
|
||||
) -> int:
|
||||
hd = arch.hidden_size
|
||||
if arch.q_lora_rank is not None:
|
||||
# MLA: q_proj->q_b, k_proj->kv_a, v_proj->kv_b, o_proj->o
|
||||
nh = arch.num_attention_heads
|
||||
qk_head = arch.qk_nope_head_dim + arch.qk_rope_head_dim
|
||||
kv_out = nh * (arch.qk_nope_head_dim + arch.v_head_dim)
|
||||
o_in = nh * arch.v_head_dim
|
||||
dims = {
|
||||
"q_proj": (arch.q_lora_rank, nh * qk_head),
|
||||
"k_proj": (hd, arch.kv_lora_rank + arch.qk_rope_head_dim),
|
||||
"v_proj": (arch.kv_lora_rank, kv_out),
|
||||
"o_proj": (o_in, hd),
|
||||
}
|
||||
else:
|
||||
kv_size = _get_kv_size(arch)
|
||||
dims = {
|
||||
"q_proj": (hd, hd),
|
||||
"k_proj": (hd, kv_size),
|
||||
"v_proj": (hd, kv_size),
|
||||
"o_proj": (hd, hd),
|
||||
}
|
||||
total = 0
|
||||
for name, (in_dim, out_dim) in dims.items():
|
||||
if name in target_modules:
|
||||
total += in_dim * r + r * out_dim
|
||||
return total
|
||||
|
||||
|
||||
def _lora_mlp_elements(
|
||||
hd: int,
|
||||
mlp_size: int,
|
||||
r: int,
|
||||
target_modules: list,
|
||||
expert_mult: int,
|
||||
) -> int:
|
||||
module_ab = {
|
||||
"gate_proj": (hd * r, r * mlp_size),
|
||||
"up_proj": (hd * r, r * mlp_size),
|
||||
"down_proj": (mlp_size * r, r * hd),
|
||||
}
|
||||
total = 0
|
||||
for name, (a, b) in module_ab.items():
|
||||
if name in target_modules:
|
||||
total += (a + b) * expert_mult
|
||||
return total
|
||||
|
||||
|
||||
def compute_lora_params(
|
||||
arch: ModelArchConfig,
|
||||
lora_rank: int,
|
||||
target_modules: list,
|
||||
) -> int:
|
||||
hd = arch.hidden_size
|
||||
r = lora_rank
|
||||
n_layers = arch.num_hidden_layers
|
||||
n_experts = _get_num_experts(arch)
|
||||
|
||||
attn_total = _lora_attn_elements(arch, r, target_modules) * n_layers
|
||||
|
||||
if n_experts > 1:
|
||||
n_dense = arch.num_dense_layers
|
||||
n_moe = n_layers - n_dense
|
||||
# Include shared experts alongside routed experts
|
||||
moe_expert_mult = n_experts + arch.n_shared_experts
|
||||
moe_mlp = _lora_mlp_elements(
|
||||
hd,
|
||||
_get_mlp_size(arch),
|
||||
r,
|
||||
target_modules,
|
||||
moe_expert_mult,
|
||||
)
|
||||
dense_mlp = _lora_mlp_elements(
|
||||
hd,
|
||||
arch.intermediate_size,
|
||||
r,
|
||||
target_modules,
|
||||
1,
|
||||
)
|
||||
mlp_total = moe_mlp * n_moe + dense_mlp * n_dense
|
||||
else:
|
||||
mlp_total = (
|
||||
_lora_mlp_elements(
|
||||
hd,
|
||||
arch.intermediate_size,
|
||||
r,
|
||||
target_modules,
|
||||
1,
|
||||
)
|
||||
* n_layers
|
||||
)
|
||||
|
||||
return attn_total + mlp_total
|
||||
|
||||
|
||||
def compute_lora_adapter_bytes(lora_params: int) -> int:
|
||||
return lora_params * 2
|
||||
|
||||
|
||||
def compute_optimizer_bytes(trainable_params: int, optimizer: str) -> int:
|
||||
optimizer_key = optimizer.lower().replace("-", "_")
|
||||
bytes_per_param = OPTIMIZER_BYTES_PER_PARAM.get(optimizer_key, 4)
|
||||
return trainable_params * bytes_per_param
|
||||
|
||||
|
||||
def compute_gradient_bytes(trainable_params: int) -> int:
|
||||
return trainable_params * 2
|
||||
|
||||
|
||||
def compute_activation_bytes(
|
||||
arch: ModelArchConfig,
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
gradient_checkpointing: str,
|
||||
is_lora: bool = False,
|
||||
) -> int:
|
||||
hd = arch.hidden_size
|
||||
kv_size = _get_kv_size(arch)
|
||||
mlp_size = _get_mlp_size(arch)
|
||||
bsz = batch_size
|
||||
n_layers = arch.num_hidden_layers
|
||||
|
||||
activation_qkv = seq_len * bsz * (hd + kv_size + kv_size)
|
||||
residual_memory = (seq_len * bsz) * 2
|
||||
activation_mlp = seq_len * bsz * (mlp_size + mlp_size)
|
||||
|
||||
per_layer_bytes = (activation_qkv + residual_memory + activation_mlp) * 2
|
||||
per_layer_bytes = int(per_layer_bytes * 1.25)
|
||||
|
||||
gc_key = gradient_checkpointing.lower()
|
||||
gc_entry = GC_LAYER_MULTIPLIERS.get(gc_key, (None, None))
|
||||
full_ft_mult, lora_mult = gc_entry
|
||||
gc_multiplier = lora_mult if is_lora else full_ft_mult
|
||||
|
||||
if gc_multiplier is None:
|
||||
effective_layers = n_layers
|
||||
else:
|
||||
effective_layers = gc_multiplier
|
||||
|
||||
return int(per_layer_bytes * effective_layers)
|
||||
|
||||
|
||||
def estimate_training_vram(
|
||||
arch: ModelArchConfig,
|
||||
config: TrainingVramConfig,
|
||||
) -> VramBreakdown:
|
||||
method = config.training_method.lower()
|
||||
is_lora = method in ("qlora", "lora")
|
||||
load_in_4bit = config.load_in_4bit or method == "qlora"
|
||||
|
||||
model_weights = compute_model_weights_bytes(arch, method, load_in_4bit)
|
||||
|
||||
lora_params = 0
|
||||
lora_adapter_bytes = 0
|
||||
if is_lora:
|
||||
lora_params = compute_lora_params(
|
||||
arch,
|
||||
config.lora_rank,
|
||||
config.target_modules,
|
||||
)
|
||||
lora_adapter_bytes = compute_lora_adapter_bytes(lora_params)
|
||||
|
||||
trainable_params = lora_params if is_lora else compute_total_params(arch)
|
||||
optimizer_bytes = compute_optimizer_bytes(trainable_params, config.optimizer)
|
||||
gradient_bytes = max(
|
||||
compute_gradient_bytes(trainable_params),
|
||||
int(model_weights * 0.15),
|
||||
)
|
||||
activations_computed = compute_activation_bytes(
|
||||
arch,
|
||||
config.batch_size,
|
||||
config.max_seq_length,
|
||||
config.gradient_checkpointing,
|
||||
is_lora = is_lora,
|
||||
)
|
||||
activation_bytes = max(
|
||||
activations_computed,
|
||||
int(model_weights * 0.15 * (config.batch_size / 2)),
|
||||
)
|
||||
|
||||
return VramBreakdown(
|
||||
model_weights = model_weights,
|
||||
lora_adapters = lora_adapter_bytes,
|
||||
optimizer_states = optimizer_bytes,
|
||||
gradients = gradient_bytes,
|
||||
activations = activation_bytes,
|
||||
cuda_overhead = CUDA_OVERHEAD_BYTES,
|
||||
activations_computed = activations_computed,
|
||||
)
|
||||
|
|
@ -973,6 +973,73 @@ def list_gguf_variants(
|
|||
return variants, has_vision
|
||||
|
||||
|
||||
def list_local_gguf_variants(
|
||||
directory: str,
|
||||
) -> tuple[list[GgufVariantInfo], bool]:
|
||||
"""List GGUF quantization variants in a local directory.
|
||||
|
||||
Mirrors :func:`list_gguf_variants` but reads from the filesystem
|
||||
instead of the HuggingFace API. Aggregates shard sizes by quant
|
||||
label so that split GGUFs appear as a single variant.
|
||||
|
||||
Returns:
|
||||
(variants, has_vision): list of non-mmproj GGUF variants + vision flag.
|
||||
"""
|
||||
p = Path(directory)
|
||||
if not p.is_dir():
|
||||
return [], False
|
||||
|
||||
quant_totals: dict[str, int] = {}
|
||||
quant_first_file: dict[str, str] = {}
|
||||
has_vision = False
|
||||
|
||||
for f in sorted(p.glob("*.gguf")):
|
||||
if _is_mmproj(f.name):
|
||||
has_vision = True
|
||||
continue
|
||||
try:
|
||||
size = f.stat().st_size
|
||||
except OSError:
|
||||
size = 0
|
||||
quant = _extract_quant_label(f.name)
|
||||
quant_totals[quant] = quant_totals.get(quant, 0) + size
|
||||
if quant not in quant_first_file:
|
||||
quant_first_file[quant] = f.name
|
||||
|
||||
variants = [
|
||||
GgufVariantInfo(
|
||||
filename = quant_first_file[q],
|
||||
quant = q,
|
||||
size_bytes = s,
|
||||
)
|
||||
for q, s in quant_totals.items()
|
||||
]
|
||||
variants.sort(key = lambda v: -v.size_bytes)
|
||||
return variants, has_vision
|
||||
|
||||
|
||||
def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
|
||||
"""Find the GGUF file in *directory* matching a quantization *variant*.
|
||||
|
||||
For sharded GGUFs (multiple files with the same quant label), returns
|
||||
the first shard (sorted by name) which is what ``llama-server -m`` expects.
|
||||
|
||||
Returns the resolved absolute path, or ``None`` if no match.
|
||||
"""
|
||||
p = Path(directory)
|
||||
if not p.is_dir():
|
||||
return None
|
||||
|
||||
matches = sorted(
|
||||
f
|
||||
for f in p.glob("*.gguf")
|
||||
if not _is_mmproj(f.name) and _extract_quant_label(f.name) == variant
|
||||
)
|
||||
if matches:
|
||||
return str(matches[0].resolve())
|
||||
return None
|
||||
|
||||
|
||||
def detect_gguf_model_remote(
|
||||
repo_id: str,
|
||||
hf_token: Optional[str] = None,
|
||||
|
|
@ -1530,7 +1597,10 @@ class ModelConfig:
|
|||
|
||||
# Auto-detect GGUF models (check before LoRA/vision detection)
|
||||
if is_local:
|
||||
gguf_file = detect_gguf_model(path)
|
||||
if gguf_variant:
|
||||
gguf_file = _find_local_gguf_by_variant(path, gguf_variant)
|
||||
else:
|
||||
gguf_file = detect_gguf_model(path)
|
||||
if gguf_file:
|
||||
display_name = Path(gguf_file).stem
|
||||
logger.info(f"Detected local GGUF model: {gguf_file}")
|
||||
|
|
|
|||
|
|
@ -23,6 +23,9 @@ from .storage_roots import (
|
|||
unstructured_uploads_root,
|
||||
oxc_validator_tmp_root,
|
||||
tensorboard_root,
|
||||
legacy_hf_cache_dir,
|
||||
hf_default_cache_dir,
|
||||
lmstudio_model_dirs,
|
||||
ensure_dir,
|
||||
ensure_studio_directories,
|
||||
resolve_under_root,
|
||||
|
|
@ -53,6 +56,9 @@ __all__ = [
|
|||
"unstructured_uploads_root",
|
||||
"oxc_validator_tmp_root",
|
||||
"tensorboard_root",
|
||||
"legacy_hf_cache_dir",
|
||||
"hf_default_cache_dir",
|
||||
"lmstudio_model_dirs",
|
||||
"ensure_dir",
|
||||
"ensure_studio_directories",
|
||||
"resolve_under_root",
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ Path utilities for model and dataset handling
|
|||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
import structlog
|
||||
|
|
@ -14,12 +15,33 @@ from loggers import get_logger
|
|||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _is_wsl() -> bool:
|
||||
"""Detect if we are running inside WSL (Windows Subsystem for Linux)."""
|
||||
if sys.platform == "win32":
|
||||
return False
|
||||
try:
|
||||
with open("/proc/version", "r") as f:
|
||||
return "microsoft" in f.read().lower()
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
_IS_WSL: bool = _is_wsl()
|
||||
|
||||
|
||||
def normalize_path(path: str) -> str:
|
||||
"""
|
||||
Convert Windows paths to WSL format if needed.
|
||||
Normalize filesystem paths for cross-platform use.
|
||||
|
||||
Examples:
|
||||
On WSL, converts Windows drive-letter paths to ``/mnt/<drive>/...``.
|
||||
On native Windows, keeps the drive letter and normalizes separators.
|
||||
On Linux/macOS (non-WSL), paths are returned with forward slashes.
|
||||
|
||||
Examples (WSL):
|
||||
C:\\Users\\... -> /mnt/c/Users/...
|
||||
Examples (native Windows):
|
||||
C:\\Users\\... -> C:/Users/...
|
||||
Examples (Linux/macOS):
|
||||
/home/user/... -> /home/user/... (unchanged)
|
||||
"""
|
||||
if not path:
|
||||
|
|
@ -27,9 +49,13 @@ def normalize_path(path: str) -> str:
|
|||
|
||||
# Handle Windows drive letters (C:\\ or c:\\)
|
||||
if len(path) >= 3 and path[1] == ":" and path[2] in ("\\", "/"):
|
||||
drive = path[0].lower()
|
||||
rest = path[3:].replace("\\", "/")
|
||||
return f"/mnt/{drive}/{rest}"
|
||||
# Only map to /mnt/<drive>/ when running under WSL;
|
||||
# on native Windows the drive letter must be preserved.
|
||||
if _IS_WSL:
|
||||
drive = path[0].lower()
|
||||
rest = path[3:].replace("\\", "/")
|
||||
return f"/mnt/{drive}/{rest}"
|
||||
return path.replace("\\", "/")
|
||||
|
||||
# Already Unix-style or relative
|
||||
return path.replace("\\", "/")
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
|
|
@ -82,19 +83,75 @@ def ensure_dir(path: Path) -> Path:
|
|||
return path
|
||||
|
||||
|
||||
def legacy_hf_cache_dir() -> Path:
|
||||
"""Old Unsloth-specific HF hub cache, kept for backward-compat scanning."""
|
||||
return cache_root() / "huggingface" / "hub"
|
||||
|
||||
|
||||
def hf_default_cache_dir() -> Path:
|
||||
"""Return the platform default HuggingFace hub cache (ignoring env overrides).
|
||||
|
||||
This is the location HF uses when no ``HF_HUB_CACHE`` / ``HF_HOME``
|
||||
env var is set. We scan it so that models a user downloaded *before*
|
||||
installing Unsloth Studio are still discovered.
|
||||
"""
|
||||
return Path.home() / ".cache" / "huggingface" / "hub"
|
||||
|
||||
|
||||
def lmstudio_model_dirs() -> list[Path]:
|
||||
"""Return LM Studio model directories that exist on disk."""
|
||||
dirs: list[Path] = []
|
||||
seen: set[Path] = set()
|
||||
|
||||
def _add(p: Path) -> None:
|
||||
resolved = p.resolve()
|
||||
if resolved not in seen and p.is_dir():
|
||||
seen.add(resolved)
|
||||
dirs.append(p)
|
||||
|
||||
# 1. Check LM Studio settings.json for custom downloads folder
|
||||
settings_path = Path.home() / ".lmstudio" / "settings.json"
|
||||
if settings_path.is_file():
|
||||
try:
|
||||
with open(settings_path) as f:
|
||||
settings = json.load(f)
|
||||
downloads = settings.get("downloadsFolder", "")
|
||||
if downloads:
|
||||
_add(Path(downloads).expanduser())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 2. LM Studio current default models directory (all platforms)
|
||||
_add(Path.home() / ".lmstudio" / "models")
|
||||
|
||||
# 3. Legacy LM Studio cache location
|
||||
_add(Path.home() / ".cache" / "lm-studio" / "models")
|
||||
|
||||
return dirs
|
||||
|
||||
|
||||
def _setup_cache_env() -> None:
|
||||
"""Set cache environment variables for HuggingFace, uv, and vLLM.
|
||||
|
||||
Respects the standard HF cache resolution chain: explicit ``HF_HOME``
|
||||
/ ``HF_HUB_CACHE`` env vars take priority, then ``XDG_CACHE_HOME``,
|
||||
then the platform default (``~/.cache/huggingface``). The legacy
|
||||
Unsloth cache is still *scanned* for models but is never set as the
|
||||
active download target.
|
||||
|
||||
Only sets variables that are not already set by the user, so
|
||||
explicit overrides (e.g. HF_HOME=/data/hf) are respected.
|
||||
Works on Linux, macOS, and Windows.
|
||||
"""
|
||||
root = cache_root()
|
||||
hf_dir = root / "huggingface"
|
||||
defaults = {
|
||||
"HF_HOME": str(hf_dir),
|
||||
"HF_HUB_CACHE": str(hf_dir / "hub"),
|
||||
"HF_XET_CACHE": str(hf_dir / "xet"),
|
||||
xdg_cache = Path(
|
||||
os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")
|
||||
).expanduser()
|
||||
hf_default = xdg_cache / "huggingface"
|
||||
defaults: dict[str, str] = {
|
||||
"HF_HOME": str(hf_default),
|
||||
"HF_HUB_CACHE": str(hf_default / "hub"),
|
||||
"HF_XET_CACHE": str(hf_default / "xet"),
|
||||
"UV_CACHE_DIR": str(root / "uv"),
|
||||
"VLLM_CACHE_ROOT": str(root / "vllm"),
|
||||
}
|
||||
|
|
|
|||
1
studio/frontend/.gitignore
vendored
1
studio/frontend/.gitignore
vendored
|
|
@ -11,6 +11,7 @@ pnpm-debug.log*
|
|||
lerna-debug.log*
|
||||
|
||||
node_modules
|
||||
bun.lock
|
||||
dist
|
||||
dist-ssr
|
||||
test/
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
|
|
|
|||
BIN
studio/frontend/public/rounded-512.png
Normal file
BIN
studio/frontend/public/rounded-512.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 50 KiB |
BIN
studio/frontend/public/rounded.png
Normal file
BIN
studio/frontend/public/rounded.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 84 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 157 KiB After Width: | Height: | Size: 361 KiB |
|
|
@ -18,8 +18,8 @@ import {
|
|||
TooltipContent,
|
||||
TooltipTrigger,
|
||||
} from "@/components/ui/tooltip";
|
||||
import { deleteCachedModel, listCachedGguf, listCachedModels, listGgufVariants } from "@/features/chat/api/chat-api";
|
||||
import type { CachedGgufRepo, CachedModelRepo } from "@/features/chat/api/chat-api";
|
||||
import { deleteCachedModel, listCachedGguf, listCachedModels, listGgufVariants, listLocalModels } from "@/features/chat/api/chat-api";
|
||||
import type { CachedGgufRepo, CachedModelRepo, LocalModelInfo } from "@/features/chat/api/chat-api";
|
||||
import type { GgufVariantDetail } from "@/features/chat/types/api";
|
||||
import { usePlatformStore } from "@/config/env";
|
||||
import {
|
||||
|
|
@ -47,6 +47,11 @@ function dedupe(values: string[]): string[] {
|
|||
return [...new Set(values.filter(Boolean))];
|
||||
}
|
||||
|
||||
/** Normalize a string for fuzzy search: lowercase, strip separators. */
|
||||
function normalizeForSearch(s: string): string {
|
||||
return s.toLowerCase().replace(/[\s\-_\.]/g, "");
|
||||
}
|
||||
|
||||
function ListLabel({ children }: { children: ReactNode }) {
|
||||
return (
|
||||
<div className="px-2.5 py-1.5 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
|
||||
|
|
@ -198,17 +203,20 @@ function GgufVariantExpander({
|
|||
};
|
||||
}, [repoId]);
|
||||
|
||||
// Covers Unix absolute (/), Windows drive (C:\, D:/), UNC (\\server), relative (./, ../), tilde (~/)
|
||||
const isLocalPath = /^(\/|\.{1,2}[\\\/]|~[\\\/]|[A-Za-z]:[\\\/]|\\\\)/.test(repoId);
|
||||
|
||||
const handleVariantClick = useCallback(
|
||||
(quant: string, downloaded?: boolean, sizeBytes?: number) => {
|
||||
onSelect(repoId, {
|
||||
source: "hub",
|
||||
source: isLocalPath ? "local" : "hub",
|
||||
isLora: false,
|
||||
ggufVariant: quant,
|
||||
isDownloaded: downloaded,
|
||||
isDownloaded: isLocalPath ? true : downloaded,
|
||||
expectedBytes: sizeBytes,
|
||||
});
|
||||
},
|
||||
[repoId, onSelect],
|
||||
[repoId, isLocalPath, onSelect],
|
||||
);
|
||||
|
||||
// GGUF fit classification matching llama-server's _select_gpus logic:
|
||||
|
|
@ -375,6 +383,17 @@ function extractParamLabel(id: string): string | undefined {
|
|||
// Module-level caches so re-mounting the popover shows results instantly
|
||||
let _cachedGgufCache: CachedGgufRepo[] = [];
|
||||
let _cachedModelsCache: CachedModelRepo[] = [];
|
||||
let _lmStudioCache: LocalModelInfo[] = [];
|
||||
|
||||
/** Sort LM Studio models with unsloth publisher first. */
|
||||
function sortLmStudio(models: LocalModelInfo[]): LocalModelInfo[] {
|
||||
return [...models].sort((a, b) => {
|
||||
const aUnsloth = (a.model_id ?? "").startsWith("unsloth/") ? 0 : 1;
|
||||
const bUnsloth = (b.model_id ?? "").startsWith("unsloth/") ? 0 : 1;
|
||||
if (aUnsloth !== bUnsloth) return aUnsloth - bUnsloth;
|
||||
return (a.model_id ?? a.display_name).localeCompare(b.model_id ?? b.display_name);
|
||||
});
|
||||
}
|
||||
|
||||
// ── Hub Model Picker ──────────────────────────────────────────
|
||||
|
||||
|
|
@ -408,12 +427,28 @@ export function HubModelPicker({
|
|||
const alreadyCached = _cachedGgufCache.length > 0 || _cachedModelsCache.length > 0;
|
||||
const [cachedReady, setCachedReady] = useState(alreadyCached);
|
||||
|
||||
// LM Studio local models -- module-level cache so re-mounting the
|
||||
// popover does not flash an empty section (same pattern as GGUF/models).
|
||||
const [lmStudioModels, setLmStudioModels] = useState<LocalModelInfo[]>(_lmStudioCache);
|
||||
|
||||
const refreshCachedLists = useCallback(() => {
|
||||
listCachedGguf().then((v) => { _cachedGgufCache = v; setCachedGguf(v); }).catch(() => {});
|
||||
listCachedModels().then((v) => { _cachedModelsCache = v; setCachedModels(v); }).catch(() => {});
|
||||
listLocalModels().then((res) => {
|
||||
const next = sortLmStudio(res.models.filter((m) => m.source === "lmstudio"));
|
||||
_lmStudioCache = next;
|
||||
setLmStudioModels(next);
|
||||
}).catch(() => {});
|
||||
}, []);
|
||||
|
||||
useEffect(() => {
|
||||
// Always refresh LM Studio models (not gated by alreadyCached)
|
||||
listLocalModels().then((res) => {
|
||||
const next = sortLmStudio(res.models.filter((m) => m.source === "lmstudio"));
|
||||
_lmStudioCache = next;
|
||||
setLmStudioModels(next);
|
||||
}).catch(() => {});
|
||||
|
||||
if (alreadyCached) return;
|
||||
let done = 0;
|
||||
const check = () => { if (++done >= 2) setCachedReady(true); };
|
||||
|
|
@ -454,7 +489,8 @@ export function HubModelPicker({
|
|||
const recommendedIds = useMemo(() => {
|
||||
const all = dedupe([...models.map((model) => model.id), value ?? ""])
|
||||
.filter((id) => !downloadedSet.has(id.toLowerCase()))
|
||||
.filter((id) => !chatOnly || isGgufRepo(id));
|
||||
.filter((id) => !chatOnly || isGgufRepo(id))
|
||||
.filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id));
|
||||
// Sort: GGUFs first, then hub models
|
||||
const gguf: string[] = [];
|
||||
const hub: string[] = [];
|
||||
|
|
@ -485,20 +521,40 @@ export function HubModelPicker({
|
|||
|
||||
const hasMoreRecommended = visibleRecommendedIds.length < recommendedIds.length;
|
||||
|
||||
// Fetch VRAM info for the full pool once (recommendedIds is stable across
|
||||
// page increments) so we don't re-fetch on every scroll.
|
||||
const { paramCountById: recommendedParamCountById } =
|
||||
useRecommendedModelVram(recommendedIds);
|
||||
|
||||
const showHfSection = debouncedQuery.trim().length > 0;
|
||||
const recommendedSet = useMemo(() => new Set(visibleRecommendedIds), [visibleRecommendedIds]);
|
||||
|
||||
// Recommended models that match the current search query
|
||||
const filteredRecommendedIds = useMemo(() => {
|
||||
if (!showHfSection) return [];
|
||||
const q = normalizeForSearch(debouncedQuery.trim());
|
||||
return recommendedIds.filter((id) => normalizeForSearch(id).includes(q));
|
||||
}, [showHfSection, debouncedQuery, recommendedIds]);
|
||||
|
||||
// Fetch VRAM info for visible models, plus any models surfaced by a search
|
||||
// query so that filtered recommended models also show VRAM badges.
|
||||
// Skip GGUF repos: they have no safetensors metadata and the render layer
|
||||
// already shows a static "GGUF" badge instead of VRAM data.
|
||||
const idsForVram = useMemo(() => {
|
||||
const ids = showHfSection
|
||||
? [...new Set([...visibleRecommendedIds, ...filteredRecommendedIds])]
|
||||
: visibleRecommendedIds;
|
||||
return ids.filter((id) => !isGgufRepo(id));
|
||||
}, [visibleRecommendedIds, showHfSection, filteredRecommendedIds]);
|
||||
const { paramCountById: recommendedParamCountById } =
|
||||
useRecommendedModelVram(idsForVram);
|
||||
|
||||
const recommendedSet = useMemo(
|
||||
() => new Set(showHfSection ? filteredRecommendedIds : visibleRecommendedIds),
|
||||
[showHfSection, filteredRecommendedIds, visibleRecommendedIds],
|
||||
);
|
||||
|
||||
const hfIds = useMemo(() => {
|
||||
if (!showHfSection) return [];
|
||||
return results
|
||||
.map((result) => result.id)
|
||||
.filter((id) => !recommendedSet.has(id))
|
||||
.filter((id) => !chatOnly || isGgufRepo(id));
|
||||
.filter((id) => !chatOnly || isGgufRepo(id))
|
||||
.filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id));
|
||||
}, [recommendedSet, results, showHfSection, chatOnly]);
|
||||
|
||||
const metricsById = useMemo(
|
||||
|
|
@ -541,7 +597,8 @@ export function HubModelPicker({
|
|||
string,
|
||||
{ est: number; status: VramFitStatus | null; detail: string | null }
|
||||
>();
|
||||
for (const id of visibleRecommendedIds) {
|
||||
const ids = showHfSection ? filteredRecommendedIds : visibleRecommendedIds;
|
||||
for (const id of ids) {
|
||||
const totalParams = recommendedParamCountById.get(id);
|
||||
if (totalParams) {
|
||||
const est = estimateLoadingVram(totalParams, "qlora");
|
||||
|
|
@ -553,7 +610,7 @@ export function HubModelPicker({
|
|||
}
|
||||
}
|
||||
return map;
|
||||
}, [visibleRecommendedIds, recommendedParamCountById, gpu]);
|
||||
}, [showHfSection, filteredRecommendedIds, visibleRecommendedIds, recommendedParamCountById, gpu]);
|
||||
|
||||
const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length);
|
||||
|
||||
|
|
@ -667,6 +724,40 @@ export function HubModelPicker({
|
|||
</>
|
||||
) : null}
|
||||
|
||||
{!showHfSection && chatOnly && lmStudioModels.length > 0 ? (
|
||||
<>
|
||||
<ListLabel>LM Studio</ListLabel>
|
||||
{lmStudioModels.map((m) => {
|
||||
const isGguf = isGgufRepo(m.id) || isGgufRepo(m.display_name);
|
||||
return (
|
||||
<div key={m.id}>
|
||||
<ModelRow
|
||||
label={m.model_id ?? m.display_name}
|
||||
meta={isGguf || m.path.endsWith(".gguf") ? "GGUF" : "Local"}
|
||||
selected={value === m.id}
|
||||
onClick={() => {
|
||||
if (isGguf) {
|
||||
setExpandedGguf((prev) => (prev === m.id ? null : m.id));
|
||||
} else {
|
||||
onSelect(m.id, { source: "local", isLora: false, isDownloaded: true });
|
||||
}
|
||||
}}
|
||||
vramStatus={null}
|
||||
/>
|
||||
{expandedGguf === m.id && (
|
||||
<GgufVariantExpander
|
||||
repoId={m.id}
|
||||
onSelect={onSelect}
|
||||
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
|
||||
systemRamGb={gpu.available ? gpu.systemRamAvailableGb : undefined}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</>
|
||||
) : null}
|
||||
|
||||
{!showHfSection && cachedReady ? (
|
||||
<>
|
||||
<ListLabel>{"\uD83E\uDDA5"} Recommended</ListLabel>
|
||||
|
|
@ -710,13 +801,44 @@ export function HubModelPicker({
|
|||
</>
|
||||
) : null}
|
||||
|
||||
{showHfSection && filteredRecommendedIds.length > 0 ? (
|
||||
<>
|
||||
<ListLabel>{"\uD83E\uDDA5"} Recommended</ListLabel>
|
||||
{filteredRecommendedIds.map((id) => {
|
||||
const vram = recommendedVramMap.get(id);
|
||||
return (
|
||||
<div key={id}>
|
||||
<ModelRow
|
||||
label={id}
|
||||
meta={
|
||||
isGgufRepo(id)
|
||||
? "GGUF"
|
||||
: vram?.detail ?? extractParamLabel(id)
|
||||
}
|
||||
selected={value === id}
|
||||
onClick={() => handleModelClick(id)}
|
||||
vramStatus={isGgufRepo(id) ? null : vram?.status ?? null}
|
||||
vramEst={isGgufRepo(id) ? undefined : vram?.est}
|
||||
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
|
||||
/>
|
||||
{expandedGguf === id && (
|
||||
<GgufVariantExpander repoId={id} onSelect={onSelect} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={gpu.available ? gpu.systemRamAvailableGb : undefined} />
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</>
|
||||
) : null}
|
||||
|
||||
{showHfSection ? (
|
||||
<>
|
||||
<ListLabel>Hugging Face</ListLabel>
|
||||
{(hfIds.length > 0 || isLoading) && <ListLabel>Hugging Face</ListLabel>}
|
||||
{hfIds.length === 0 && !isLoading ? (
|
||||
<div className="px-2.5 py-2 text-xs text-muted-foreground">
|
||||
No matching models.
|
||||
</div>
|
||||
filteredRecommendedIds.length === 0 ? (
|
||||
<div className="px-2.5 py-2 text-xs text-muted-foreground">
|
||||
No matching models.
|
||||
</div>
|
||||
) : null
|
||||
) : (
|
||||
hfIds.map((id) => {
|
||||
const vram = vramMap.get(id);
|
||||
|
|
@ -787,6 +909,8 @@ export function LoraModelPicker({
|
|||
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
|
||||
}) {
|
||||
const [query, setQuery] = useState("");
|
||||
const [expandedGguf, setExpandedGguf] = useState<string | null>(null);
|
||||
const gpu = useGpuInfo();
|
||||
|
||||
const normalized = useMemo(
|
||||
() =>
|
||||
|
|
@ -796,22 +920,28 @@ export function LoraModelPicker({
|
|||
baseModel: model.baseModel || model.description || "Unknown base model",
|
||||
}))
|
||||
.sort((a, b) => {
|
||||
const baseCmp = a.baseModel.localeCompare(b.baseModel);
|
||||
if (baseCmp !== 0) return baseCmp;
|
||||
// Prioritize unsloth publisher within LM Studio group
|
||||
if (a.baseModel === "LM Studio" && b.baseModel === "LM Studio") {
|
||||
const aUnsloth = a.name.startsWith("unsloth/") ? 0 : 1;
|
||||
const bUnsloth = b.name.startsWith("unsloth/") ? 0 : 1;
|
||||
if (aUnsloth !== bUnsloth) return aUnsloth - bUnsloth;
|
||||
}
|
||||
const aTime = a.updatedAt ?? -1;
|
||||
const bTime = b.updatedAt ?? -1;
|
||||
if (aTime !== bTime) return bTime - aTime;
|
||||
const baseCmp = a.baseModel.localeCompare(b.baseModel);
|
||||
if (baseCmp !== 0) return baseCmp;
|
||||
return a.name.localeCompare(b.name);
|
||||
}),
|
||||
[loraModels],
|
||||
);
|
||||
|
||||
const grouped = useMemo(() => {
|
||||
const needle = query.trim().toLowerCase();
|
||||
const needle = normalizeForSearch(query.trim());
|
||||
const out = new Map<string, LoraModelOption[]>();
|
||||
|
||||
for (const model of normalized) {
|
||||
const searchText = `${model.name} ${model.baseModel} ${model.id}`.toLowerCase();
|
||||
const searchText = normalizeForSearch(`${model.name} ${model.baseModel} ${model.id}`);
|
||||
if (needle && !searchText.includes(needle)) continue;
|
||||
|
||||
const key = model.baseModel || "Unknown base model";
|
||||
|
|
@ -855,34 +985,54 @@ export function LoraModelPicker({
|
|||
{index > 0 ? <div className="my-1" /> : null}
|
||||
<ListLabel>{baseModel}</ListLabel>
|
||||
{adapters.map((adapter) => {
|
||||
const isLocal = adapter.source === "local";
|
||||
const isExported = adapter.source === "exported";
|
||||
const isMerged = adapter.exportType === "merged";
|
||||
const isGguf = adapter.exportType === "gguf";
|
||||
const tag = isGguf
|
||||
? "GGUF"
|
||||
: isExported
|
||||
? isMerged ? "Merged" : "LoRA"
|
||||
: "LoRA";
|
||||
const meta = isExported ? `${tag} · Exported` : tag;
|
||||
const isLocalGgufDir = isLocal && (isGgufRepo(adapter.id) || isGgufRepo(adapter.name));
|
||||
const tag = isLocal
|
||||
? isLocalGgufDir ? "GGUF" : "Local"
|
||||
: isGguf
|
||||
? "GGUF"
|
||||
: isExported
|
||||
? isMerged ? "Merged" : "LoRA"
|
||||
: "LoRA";
|
||||
const meta = isLocal ? (isLocalGgufDir ? "GGUF" : "Local") : isExported ? `${tag} · Exported` : tag;
|
||||
return (
|
||||
<ModelRow
|
||||
key={adapter.id}
|
||||
label={adapter.name}
|
||||
meta={meta}
|
||||
selected={value === adapter.id}
|
||||
onClick={() => onSelect(adapter.id, {
|
||||
source: isExported ? "exported" : "lora",
|
||||
isLora: !isMerged && !isGguf,
|
||||
})}
|
||||
tooltipText={
|
||||
<>
|
||||
<span className="block break-words">{adapter.name}</span>
|
||||
<span className="block mt-1 text-[10px] text-muted-foreground break-all">
|
||||
{adapter.id}
|
||||
</span>
|
||||
</>
|
||||
}
|
||||
/>
|
||||
<div key={adapter.id}>
|
||||
<ModelRow
|
||||
label={adapter.name}
|
||||
meta={meta}
|
||||
selected={value === adapter.id}
|
||||
onClick={() => {
|
||||
if (isLocalGgufDir) {
|
||||
setExpandedGguf((prev) => (prev === adapter.id ? null : adapter.id));
|
||||
} else {
|
||||
onSelect(adapter.id, {
|
||||
source: isLocal ? "local" : isExported ? "exported" : "lora",
|
||||
isLora: !isLocal && !isMerged && !isGguf,
|
||||
isDownloaded: true,
|
||||
});
|
||||
}
|
||||
}}
|
||||
tooltipText={
|
||||
<>
|
||||
<span className="block break-words">{adapter.name}</span>
|
||||
<span className="block mt-1 text-[10px] text-muted-foreground break-all">
|
||||
{adapter.id}
|
||||
</span>
|
||||
</>
|
||||
}
|
||||
/>
|
||||
{expandedGguf === adapter.id && (
|
||||
<GgufVariantExpander
|
||||
repoId={adapter.id}
|
||||
onSelect={onSelect}
|
||||
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
|
||||
systemRamGb={gpu.available ? gpu.systemRamAvailableGb : undefined}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -13,12 +13,12 @@ export interface ModelOption {
|
|||
export interface LoraModelOption extends ModelOption {
|
||||
baseModel?: string;
|
||||
updatedAt?: number;
|
||||
source?: "training" | "exported";
|
||||
source?: "training" | "exported" | "local";
|
||||
exportType?: "lora" | "merged" | "gguf";
|
||||
}
|
||||
|
||||
export interface ModelSelectorChangeMeta {
|
||||
source: "hub" | "lora" | "exported";
|
||||
source: "hub" | "lora" | "exported" | "local";
|
||||
isLora: boolean;
|
||||
ggufVariant?: string;
|
||||
isDownloaded?: boolean;
|
||||
|
|
|
|||
|
|
@ -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]);
|
||||
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ import {
|
|||
useAuiEvent,
|
||||
useAuiState,
|
||||
} from "@assistant-ui/react";
|
||||
import { motion } from "framer-motion";
|
||||
import { motion } from "motion/react";
|
||||
import {
|
||||
ArrowDownIcon,
|
||||
ArrowUpIcon,
|
||||
|
|
@ -92,7 +92,7 @@ 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 />
|
||||
<AuiIf condition={({ thread }) => !thread.isEmpty}>
|
||||
{!hideComposer && <ComposerAnimated />}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,11 @@ import {
|
|||
HoverCardContent,
|
||||
HoverCardTrigger,
|
||||
} from "@/components/ui/hover-card";
|
||||
import {
|
||||
Collapsible,
|
||||
CollapsibleContent,
|
||||
CollapsibleTrigger,
|
||||
} from "@/components/ui/collapsible";
|
||||
import { AnimatedThemeToggler } from "@/components/ui/animated-theme-toggler";
|
||||
import {
|
||||
Sheet,
|
||||
|
|
@ -16,21 +21,28 @@ import {
|
|||
} from "@/components/ui/sheet";
|
||||
import { cn } from "@/lib/utils";
|
||||
import {
|
||||
ArrowReloadHorizontalIcon,
|
||||
ArrowRight01Icon,
|
||||
Cancel01Icon,
|
||||
Book03Icon,
|
||||
BubbleChatIcon,
|
||||
ChefHatIcon,
|
||||
Copy01Icon,
|
||||
CursorInfo02Icon,
|
||||
PackageIcon,
|
||||
Tick02Icon,
|
||||
ZapIcon,
|
||||
} from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { useTrainingRuntimeStore } from "@/features/training";
|
||||
import { usePlatformStore } from "@/config/env";
|
||||
import { Link, useRouterState } from "@tanstack/react-router";
|
||||
import { motion } from "motion/react";
|
||||
import { useState } from "react";
|
||||
import { AnimatePresence, motion, useReducedMotion } from "motion/react";
|
||||
import type { ReactElement } from "react";
|
||||
import { useEffect, useRef, useState } from "react";
|
||||
import { TOUR_OPEN_EVENT } from "@/features/tour";
|
||||
import { ShutdownDialog } from "@/components/shutdown-dialog";
|
||||
|
||||
const NAV_ITEMS = [
|
||||
{ label: "Studio", href: "/studio", icon: ZapIcon, enabled: true },
|
||||
|
|
@ -39,6 +51,185 @@ const NAV_ITEMS = [
|
|||
{ label: "Chat", href: "/chat", icon: BubbleChatIcon, enabled: true },
|
||||
];
|
||||
|
||||
const STUDIO_UPDATE_CMD = "unsloth studio update";
|
||||
const STUDIO_UPDATE_FALLBACK_UNIX_CMD =
|
||||
"curl -fsSL https://unsloth.ai/install.sh | sh";
|
||||
const STUDIO_UPDATE_FALLBACK_WINDOWS_CMD =
|
||||
"irm https://unsloth.ai/install.ps1 | iex";
|
||||
|
||||
type UpdateShell = "windows" | "unix";
|
||||
|
||||
function getDefaultUpdateShell(deviceType: string): UpdateShell {
|
||||
return deviceType === "windows" ? "windows" : "unix";
|
||||
}
|
||||
|
||||
function getStudioUpdateInstructionLine(shell: UpdateShell): string {
|
||||
return shell === "windows" ? "Open PowerShell and run:" : "Open Terminal and run:";
|
||||
}
|
||||
|
||||
function CopyableCommand({
|
||||
command,
|
||||
copyLabel,
|
||||
}: {
|
||||
command: string;
|
||||
copyLabel: string;
|
||||
}): ReactElement {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const timerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
if (timerRef.current) {
|
||||
clearTimeout(timerRef.current);
|
||||
}
|
||||
};
|
||||
}, []);
|
||||
|
||||
const handleCopy = () => {
|
||||
if (!copyToClipboard(command)) {
|
||||
return;
|
||||
}
|
||||
setCopied(true);
|
||||
if (timerRef.current) {
|
||||
clearTimeout(timerRef.current);
|
||||
}
|
||||
timerRef.current = setTimeout(() => setCopied(false), 2000);
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="flex min-w-0 items-stretch overflow-hidden rounded-md border border-border bg-muted/40">
|
||||
<input
|
||||
type="text"
|
||||
readOnly
|
||||
value={command}
|
||||
className="min-w-0 flex-1 bg-transparent px-2 py-1.5 font-mono text-[11px] text-foreground outline-none"
|
||||
title={command}
|
||||
aria-label={`${copyLabel} text`}
|
||||
/>
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleCopy}
|
||||
className="flex shrink-0 items-center justify-center border-l border-border px-2 text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
title={copied ? "Copied" : "Copy command"}
|
||||
aria-label={copied ? `${copyLabel} copied` : `Copy ${copyLabel}`}
|
||||
>
|
||||
{copied ? (
|
||||
<HugeiconsIcon icon={Tick02Icon} className="size-4 text-emerald-600" />
|
||||
) : (
|
||||
<HugeiconsIcon icon={Copy01Icon} className="size-4" />
|
||||
)}
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function UpdateStudioInstructions({
|
||||
className,
|
||||
defaultShell,
|
||||
showTitle = true,
|
||||
}: {
|
||||
className?: string;
|
||||
defaultShell: UpdateShell;
|
||||
showTitle?: boolean;
|
||||
}): ReactElement {
|
||||
const [shell, setShell] = useState<UpdateShell>(defaultShell);
|
||||
const prefersReducedMotion = useReducedMotion();
|
||||
const windows = shell === "windows";
|
||||
const fadeTransition = prefersReducedMotion
|
||||
? { duration: 0 }
|
||||
: { duration: 0.16, ease: [0.165, 0.84, 0.44, 1] as const };
|
||||
const fadeInitial = prefersReducedMotion ? { opacity: 1 } : { opacity: 0, y: 2 };
|
||||
const fadeAnimate = { opacity: 1, y: 0 };
|
||||
const fadeExit = prefersReducedMotion ? { opacity: 1 } : { opacity: 0, y: -2 };
|
||||
|
||||
useEffect(() => {
|
||||
setShell(defaultShell);
|
||||
}, [defaultShell]);
|
||||
|
||||
return (
|
||||
<div className={cn("flex flex-col gap-3", className)}>
|
||||
<div
|
||||
className={cn(
|
||||
"flex items-center gap-3",
|
||||
showTitle ? "justify-between" : "justify-start",
|
||||
)}
|
||||
>
|
||||
{showTitle ? (
|
||||
<p className="shrink-0 whitespace-nowrap text-sm font-semibold font-heading">
|
||||
Update Unsloth Studio
|
||||
</p>
|
||||
) : null}
|
||||
<div className="flex shrink-0 items-center gap-0.5 text-[11px]">
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setShell("windows")}
|
||||
className={cn(
|
||||
"px-0.5 py-0.5 font-medium transition-colors",
|
||||
windows
|
||||
? "text-foreground"
|
||||
: "text-muted-foreground hover:text-emerald-600",
|
||||
)}
|
||||
aria-pressed={windows}
|
||||
>
|
||||
Windows
|
||||
</button>
|
||||
<span className="text-border">/</span>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setShell("unix")}
|
||||
className={cn(
|
||||
"px-0.5 py-0.5 font-medium transition-colors",
|
||||
!windows
|
||||
? "text-foreground"
|
||||
: "text-muted-foreground hover:text-emerald-600",
|
||||
)}
|
||||
aria-pressed={!windows}
|
||||
>
|
||||
macOS/Linux
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
<AnimatePresence mode="wait" initial={false}>
|
||||
<motion.p
|
||||
key={`instruction-${shell}`}
|
||||
initial={fadeInitial}
|
||||
animate={fadeAnimate}
|
||||
exit={fadeExit}
|
||||
transition={fadeTransition}
|
||||
className="text-xs text-muted-foreground leading-relaxed"
|
||||
>
|
||||
{getStudioUpdateInstructionLine(shell)}
|
||||
</motion.p>
|
||||
</AnimatePresence>
|
||||
<CopyableCommand command={STUDIO_UPDATE_CMD} copyLabel="update command" />
|
||||
<p className="text-xs text-muted-foreground leading-relaxed">
|
||||
If that fails or unsloth studio update is unavailable, run:
|
||||
</p>
|
||||
<AnimatePresence mode="wait" initial={false}>
|
||||
<motion.div
|
||||
key={`fallback-${shell}`}
|
||||
initial={fadeInitial}
|
||||
animate={fadeAnimate}
|
||||
exit={fadeExit}
|
||||
transition={fadeTransition}
|
||||
>
|
||||
<CopyableCommand
|
||||
command={
|
||||
windows
|
||||
? STUDIO_UPDATE_FALLBACK_WINDOWS_CMD
|
||||
: STUDIO_UPDATE_FALLBACK_UNIX_CMD
|
||||
}
|
||||
copyLabel="fallback command"
|
||||
/>
|
||||
</motion.div>
|
||||
</AnimatePresence>
|
||||
<p className="text-xs text-muted-foreground leading-relaxed">
|
||||
Restart Studio after updating for changes to take effect.
|
||||
</p>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function getTourId(pathname: string): "studio" | "chat" | "export" | null {
|
||||
if (pathname === "/studio") return "studio";
|
||||
if (pathname === "/chat") return "chat";
|
||||
|
|
@ -50,8 +241,37 @@ export function Navbar() {
|
|||
const pathname = useRouterState({ select: (s) => s.location.pathname });
|
||||
const isTrainingRunning = useTrainingRuntimeStore((s) => s.isTrainingRunning);
|
||||
const [mobileOpen, setMobileOpen] = useState(false);
|
||||
const [mobileUpdateOpen, setMobileUpdateOpen] = useState(false);
|
||||
const [shutdownOpen, setShutdownOpen] = useState(false);
|
||||
|
||||
const deviceType = usePlatformStore((s) => s.deviceType);
|
||||
const chatOnly = usePlatformStore((s) => s.isChatOnly());
|
||||
const defaultUpdateShell = getDefaultUpdateShell(deviceType);
|
||||
|
||||
// Warn before closing the tab only when training is running (data loss risk).
|
||||
// We store the handler in a ref so removeUnloadHandler() can clean it up
|
||||
// before the "Server stopped" page renders.
|
||||
const unloadHandlerRef = useRef<((e: BeforeUnloadEvent) => void) | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
const handler = (e: BeforeUnloadEvent) => {
|
||||
if (!useTrainingRuntimeStore.getState().isTrainingRunning) return;
|
||||
e.preventDefault();
|
||||
e.returnValue = "";
|
||||
};
|
||||
unloadHandlerRef.current = handler;
|
||||
window.addEventListener("beforeunload", handler);
|
||||
return () => {
|
||||
window.removeEventListener("beforeunload", handler);
|
||||
};
|
||||
}, []);
|
||||
|
||||
const removeUnloadHandler = () => {
|
||||
if (unloadHandlerRef.current) {
|
||||
window.removeEventListener("beforeunload", unloadHandlerRef.current);
|
||||
unloadHandlerRef.current = null;
|
||||
}
|
||||
};
|
||||
|
||||
const tourId = getTourId(pathname);
|
||||
|
||||
|
|
@ -63,6 +283,7 @@ export function Navbar() {
|
|||
};
|
||||
|
||||
return (
|
||||
<>
|
||||
<header className="relative top-0 z-40 h-16 w-full">
|
||||
<div className="mx-auto grid h-full max-w-7xl grid-cols-[1fr_auto_1fr] items-center px-4 sm:px-6">
|
||||
{/* Left: logo */}
|
||||
|
|
@ -206,6 +427,35 @@ export function Navbar() {
|
|||
<HugeiconsIcon icon={CursorInfo02Icon} className="size-4" />
|
||||
<span className="text-sm font-medium">Tour</span>
|
||||
</button>
|
||||
|
||||
<HoverCard openDelay={200} closeDelay={100}>
|
||||
<HoverCardTrigger asChild={true}>
|
||||
<button
|
||||
type="button"
|
||||
className="flex h-9 items-center gap-1.5 rounded-md px-3 text-sm font-medium text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
|
||||
aria-label="How to update Unsloth Studio"
|
||||
>
|
||||
<HugeiconsIcon icon={ArrowReloadHorizontalIcon} className="size-4" />
|
||||
Update
|
||||
</button>
|
||||
</HoverCardTrigger>
|
||||
<HoverCardContent align="end" className="w-[22.5rem] p-0">
|
||||
<UpdateStudioInstructions
|
||||
className="p-4"
|
||||
defaultShell={defaultUpdateShell}
|
||||
/>
|
||||
</HoverCardContent>
|
||||
</HoverCard>
|
||||
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setShutdownOpen(true)}
|
||||
className="-mr-1.5 flex h-9 w-9 items-center justify-center rounded-md text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
|
||||
title="Shut down Unsloth Studio server"
|
||||
aria-label="Shut down Unsloth Studio server"
|
||||
>
|
||||
<HugeiconsIcon icon={Cancel01Icon} className="size-5" />
|
||||
</button>
|
||||
</div>
|
||||
|
||||
{/* Right: mobile */}
|
||||
|
|
@ -220,7 +470,13 @@ export function Navbar() {
|
|||
<HugeiconsIcon icon={CursorInfo02Icon} className="size-4" />
|
||||
</button>
|
||||
) : null}
|
||||
<Sheet open={mobileOpen} onOpenChange={setMobileOpen}>
|
||||
<Sheet
|
||||
open={mobileOpen}
|
||||
onOpenChange={(open) => {
|
||||
setMobileOpen(open);
|
||||
if (!open) setMobileUpdateOpen(false);
|
||||
}}
|
||||
>
|
||||
<SheetTrigger asChild={true}>
|
||||
<button
|
||||
type="button"
|
||||
|
|
@ -234,7 +490,7 @@ export function Navbar() {
|
|||
<SheetHeader>
|
||||
<SheetTitle>Navigate</SheetTitle>
|
||||
</SheetHeader>
|
||||
<div className="mt-6 flex flex-col gap-2">
|
||||
<div className="mt-6 flex max-h-[calc(100dvh-8rem)] flex-col gap-2 overflow-y-auto pr-1">
|
||||
{NAV_ITEMS.filter((item) => item.enabled).map((item) => {
|
||||
const active = pathname === item.href;
|
||||
const disabledByTraining =
|
||||
|
|
@ -273,7 +529,7 @@ export function Navbar() {
|
|||
href="https://unsloth.ai/docs"
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="mt-2 flex items-center gap-2 rounded-md border border-border px-3 py-2 text-sm font-medium text-foreground hover:bg-accent"
|
||||
className="mt-3 flex items-center gap-2 rounded-md border border-border px-3 py-2 text-sm font-medium text-foreground hover:bg-accent"
|
||||
onClick={() => setMobileOpen(false)}
|
||||
>
|
||||
<HugeiconsIcon icon={Book03Icon} className="size-4" />
|
||||
|
|
@ -292,6 +548,48 @@ export function Navbar() {
|
|||
Start tour
|
||||
</button>
|
||||
) : null}
|
||||
<Collapsible
|
||||
open={mobileUpdateOpen}
|
||||
onOpenChange={setMobileUpdateOpen}
|
||||
className="rounded-md border border-border"
|
||||
>
|
||||
<CollapsibleTrigger asChild={true}>
|
||||
<button
|
||||
type="button"
|
||||
className="flex w-full items-center justify-between rounded-md px-3 py-2 text-left text-sm font-medium text-foreground transition-colors hover:bg-accent"
|
||||
aria-label="Toggle update instructions"
|
||||
>
|
||||
<span className="flex items-center gap-2">
|
||||
<HugeiconsIcon icon={ArrowReloadHorizontalIcon} className="size-4" />
|
||||
Update Unsloth Studio
|
||||
</span>
|
||||
<HugeiconsIcon
|
||||
icon={ArrowRight01Icon}
|
||||
className={cn(
|
||||
"size-4 text-muted-foreground transition-transform",
|
||||
mobileUpdateOpen && "rotate-90",
|
||||
)}
|
||||
/>
|
||||
</button>
|
||||
</CollapsibleTrigger>
|
||||
<CollapsibleContent className="border-t border-border p-3 pt-2">
|
||||
<UpdateStudioInstructions
|
||||
defaultShell={defaultUpdateShell}
|
||||
showTitle={false}
|
||||
/>
|
||||
</CollapsibleContent>
|
||||
</Collapsible>
|
||||
<button
|
||||
type="button"
|
||||
className="mt-3 flex items-center gap-2 rounded-md border border-border px-3 py-2 text-left text-sm font-medium text-foreground hover:bg-accent"
|
||||
onClick={() => {
|
||||
setMobileOpen(false);
|
||||
setShutdownOpen(true);
|
||||
}}
|
||||
>
|
||||
<HugeiconsIcon icon={Cancel01Icon} className="size-5" />
|
||||
Quit Unsloth Studio
|
||||
</button>
|
||||
<div className="mt-2 flex items-center justify-between rounded-md border border-border px-3 py-2">
|
||||
<span className="text-sm font-medium text-foreground">Theme</span>
|
||||
<AnimatedThemeToggler
|
||||
|
|
@ -306,5 +604,12 @@ export function Navbar() {
|
|||
</div>
|
||||
</div>
|
||||
</header>
|
||||
|
||||
<ShutdownDialog
|
||||
open={shutdownOpen}
|
||||
onOpenChange={setShutdownOpen}
|
||||
onBeforeShutdown={removeUnloadHandler}
|
||||
/>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
84
studio/frontend/src/components/shutdown-dialog.tsx
Normal file
84
studio/frontend/src/components/shutdown-dialog.tsx
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
// 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 { toastError } from "@/shared/toast";
|
||||
import { useState } from "react";
|
||||
import {
|
||||
AlertDialog,
|
||||
AlertDialogAction,
|
||||
AlertDialogCancel,
|
||||
AlertDialogContent,
|
||||
AlertDialogDescription,
|
||||
AlertDialogFooter,
|
||||
AlertDialogHeader,
|
||||
AlertDialogTitle,
|
||||
} from "@/components/ui/alert-dialog";
|
||||
|
||||
interface ShutdownDialogProps {
|
||||
open: boolean;
|
||||
onOpenChange: (open: boolean) => void;
|
||||
/** Called right before the shutdown API request so callers can remove the
|
||||
* beforeunload listener — otherwise the "Server stopped" page would still
|
||||
* trigger a "Leave site?" prompt when the user tries to close it. */
|
||||
onBeforeShutdown?: () => void;
|
||||
}
|
||||
|
||||
export function ShutdownDialog({
|
||||
open,
|
||||
onOpenChange,
|
||||
onBeforeShutdown,
|
||||
}: ShutdownDialogProps) {
|
||||
const [stopping, setStopping] = useState(false);
|
||||
|
||||
const handleStop = async () => {
|
||||
setStopping(true);
|
||||
let accepted = false;
|
||||
try {
|
||||
const res = await authFetch("/api/shutdown", { method: "POST" });
|
||||
accepted = res.ok;
|
||||
if (!accepted) {
|
||||
toastError("Failed to shut down server");
|
||||
setStopping(false);
|
||||
return;
|
||||
}
|
||||
} catch {
|
||||
// Network error — shutdown request never reached the server
|
||||
toastError("Could not reach server");
|
||||
setStopping(false);
|
||||
return;
|
||||
}
|
||||
|
||||
onBeforeShutdown?.();
|
||||
document.body.innerHTML = `
|
||||
<div style="display:flex;flex-direction:column;align-items:center;justify-content:center;height:100vh;font-family:sans-serif;gap:12px">
|
||||
<p style="font-size:1.1rem;font-weight:600;margin:0">Unsloth Studio has stopped.</p>
|
||||
<p style="font-size:0.9rem;color:#888;margin:0">You can now close this tab.</p>
|
||||
</div>`;
|
||||
};
|
||||
|
||||
return (
|
||||
<AlertDialog open={open} onOpenChange={onOpenChange}>
|
||||
<AlertDialogContent>
|
||||
<AlertDialogHeader>
|
||||
<AlertDialogTitle>Stop Unsloth Studio?</AlertDialogTitle>
|
||||
<AlertDialogDescription>
|
||||
This will shut down the server. Any active training or inference
|
||||
jobs will be terminated. You can restart it any time from the
|
||||
desktop shortcut.
|
||||
</AlertDialogDescription>
|
||||
</AlertDialogHeader>
|
||||
<AlertDialogFooter>
|
||||
<AlertDialogCancel>Cancel</AlertDialogCancel>
|
||||
<AlertDialogAction
|
||||
onClick={handleStop}
|
||||
disabled={stopping}
|
||||
variant="destructive"
|
||||
>
|
||||
{stopping ? "Stopping…" : "Stop server"}
|
||||
</AlertDialogAction>
|
||||
</AlertDialogFooter>
|
||||
</AlertDialogContent>
|
||||
</AlertDialog>
|
||||
);
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -90,10 +90,17 @@ export const LR_SCHEDULER_OPTIONS: ReadonlyArray<{ value: string; label: string
|
|||
{ value: "cosine", label: "Cosine" },
|
||||
];
|
||||
|
||||
/**
|
||||
* Method-aware learning rate defaults.
|
||||
* Backend mirrors these in the YAML configs under studio/backend/assets/configs/.
|
||||
*/
|
||||
export const LR_DEFAULT_LORA = 2e-4;
|
||||
export const LR_DEFAULT_FULL = 2e-5;
|
||||
|
||||
export const DEFAULT_HYPERPARAMS = {
|
||||
epochs: 3,
|
||||
contextLength: 2048,
|
||||
learningRate: 2e-4,
|
||||
learningRate: LR_DEFAULT_LORA,
|
||||
optimizerType: "adamw_8bit",
|
||||
lrSchedulerType: "linear",
|
||||
loraRank: 16,
|
||||
|
|
@ -102,7 +109,7 @@ export const DEFAULT_HYPERPARAMS = {
|
|||
loraVariant: "lora" as const,
|
||||
batchSize: 4,
|
||||
gradientAccumulation: 8,
|
||||
weightDecay: 0.01,
|
||||
weightDecay: 0.001,
|
||||
warmupSteps: 5,
|
||||
maxSteps: 60,
|
||||
saveSteps: 0,
|
||||
|
|
|
|||
|
|
@ -253,6 +253,7 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise<void> {
|
|||
* falls back to smallest cached safetensors model.
|
||||
*/
|
||||
async function autoLoadSmallestModel(): Promise<boolean> {
|
||||
const hfToken = useChatRuntimeStore.getState().hfToken || null;
|
||||
const toastId = toast("Loading a model…", {
|
||||
description: "Auto-selecting the smallest downloaded model.",
|
||||
duration: 5000,
|
||||
|
|
@ -278,8 +279,8 @@ async function autoLoadSmallestModel(): Promise<boolean> {
|
|||
const variant = downloaded[0];
|
||||
const loadResp = await loadModel({
|
||||
model_path: repo.repo_id,
|
||||
hf_token: null,
|
||||
max_seq_length: 4096,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: 0,
|
||||
load_in_4bit: true,
|
||||
is_lora: false,
|
||||
gguf_variant: variant.quant,
|
||||
|
|
@ -305,11 +306,15 @@ async function autoLoadSmallestModel(): Promise<boolean> {
|
|||
}
|
||||
useChatRuntimeStore.setState({
|
||||
ggufContextLength: loadResp.context_length ?? 131072,
|
||||
ggufMaxContextLength: loadResp.max_context_length ?? loadResp.context_length ?? 131072,
|
||||
supportsReasoning: loadResp.supports_reasoning ?? false,
|
||||
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
|
||||
reasoningEnabled: loadResp.supports_reasoning ?? false,
|
||||
supportsTools: loadResp.supports_tools ?? false,
|
||||
toolsEnabled: false,
|
||||
codeToolsEnabled: false,
|
||||
toolsEnabled: loadResp.supports_tools ?? false,
|
||||
codeToolsEnabled: loadResp.supports_tools ?? false,
|
||||
kvCacheDtype: loadResp.cache_type_kv ?? null,
|
||||
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
|
||||
defaultChatTemplate: loadResp.chat_template ?? null,
|
||||
chatTemplateOverride: null,
|
||||
});
|
||||
|
|
@ -329,7 +334,7 @@ async function autoLoadSmallestModel(): Promise<boolean> {
|
|||
try {
|
||||
const sfLoadResp = await loadModel({
|
||||
model_path: repo.repo_id,
|
||||
hf_token: null,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: 4096,
|
||||
load_in_4bit: true,
|
||||
is_lora: false,
|
||||
|
|
@ -366,8 +371,8 @@ async function autoLoadSmallestModel(): Promise<boolean> {
|
|||
try {
|
||||
const loadResp = await loadModel({
|
||||
model_path: "unsloth/Qwen3.5-4B-GGUF",
|
||||
hf_token: null,
|
||||
max_seq_length: 4096,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: 0,
|
||||
load_in_4bit: true,
|
||||
is_lora: false,
|
||||
gguf_variant: "UD-Q4_K_XL",
|
||||
|
|
@ -388,10 +393,15 @@ async function autoLoadSmallestModel(): Promise<boolean> {
|
|||
}
|
||||
useChatRuntimeStore.setState({
|
||||
ggufContextLength: loadResp.context_length ?? 131072,
|
||||
ggufMaxContextLength: loadResp.max_context_length ?? loadResp.context_length ?? 131072,
|
||||
supportsReasoning: loadResp.supports_reasoning ?? false,
|
||||
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
|
||||
reasoningEnabled: loadResp.supports_reasoning ?? false,
|
||||
supportsTools: loadResp.supports_tools ?? false,
|
||||
toolsEnabled: false,
|
||||
toolsEnabled: loadResp.supports_tools ?? false,
|
||||
codeToolsEnabled: loadResp.supports_tools ?? false,
|
||||
kvCacheDtype: loadResp.cache_type_kv ?? null,
|
||||
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
|
||||
defaultChatTemplate: loadResp.chat_template ?? null,
|
||||
chatTemplateOverride: null,
|
||||
});
|
||||
|
|
@ -410,8 +420,7 @@ async function autoLoadSmallestModel(): Promise<boolean> {
|
|||
export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
||||
return {
|
||||
async *run({ messages, abortSignal, unstable_threadId }) {
|
||||
const runtime = useChatRuntimeStore.getState();
|
||||
const { params } = runtime;
|
||||
let runtime = useChatRuntimeStore.getState();
|
||||
|
||||
// Wait for in-progress model load to finish before inferring
|
||||
if (runtime.modelLoading) {
|
||||
|
|
@ -430,6 +439,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
}
|
||||
}
|
||||
|
||||
// Re-read store after potential auto-load / model ready wait
|
||||
runtime = useChatRuntimeStore.getState();
|
||||
const { params } = runtime;
|
||||
const {
|
||||
supportsTools,
|
||||
toolsEnabled,
|
||||
|
|
|
|||
|
|
@ -125,6 +125,27 @@ export async function getDownloadProgress(
|
|||
return parseJsonOrThrow(response);
|
||||
}
|
||||
|
||||
export interface LocalModelInfo {
|
||||
id: string;
|
||||
display_name: string;
|
||||
path: string;
|
||||
source: "models_dir" | "hf_cache" | "lmstudio";
|
||||
model_id?: string | null;
|
||||
updated_at?: number | null;
|
||||
}
|
||||
|
||||
interface LocalModelListResponse {
|
||||
models_dir: string;
|
||||
hf_cache_dir?: string | null;
|
||||
lmstudio_dirs: string[];
|
||||
models: LocalModelInfo[];
|
||||
}
|
||||
|
||||
export async function listLocalModels(): Promise<LocalModelListResponse> {
|
||||
const response = await authFetch("/api/models/local");
|
||||
return parseJsonOrThrow<LocalModelListResponse>(response);
|
||||
}
|
||||
|
||||
export async function listCachedGguf(): Promise<CachedGgufRepo[]> {
|
||||
const response = await authFetch("/api/models/cached-gguf");
|
||||
const data = await parseJsonOrThrow<{ cached: CachedGgufRepo[] }>(response);
|
||||
|
|
|
|||
|
|
@ -61,6 +61,7 @@ import { ArtifactPanel } from "./components/artifact-panel";
|
|||
import { useArtifactStore } from "./stores/artifact-store";
|
||||
import { PromptLibrarySheet } from "./components/prompt-library-sheet";
|
||||
import { GuidedTour, useGuidedTourController } from "@/features/tour";
|
||||
import { listLocalModels } from "./api/chat-api";
|
||||
import { ChatSettingsPanel } from "./chat-settings-sheet";
|
||||
import { ContextUsageBar } from "./components/context-usage-bar";
|
||||
import { ModelLoadInlineStatus } from "./components/model-load-status";
|
||||
|
|
@ -606,22 +607,36 @@ export function ChatPage(): ReactElement {
|
|||
[modelsFromStore],
|
||||
);
|
||||
|
||||
const loraModels = useMemo<LoraModelOption[]>(
|
||||
() =>
|
||||
lorasFromStore.map((lora) => ({
|
||||
id: lora.id,
|
||||
name: lora.name,
|
||||
baseModel: lora.baseModel,
|
||||
updatedAt: lora.updatedAt,
|
||||
source: lora.source,
|
||||
exportType: lora.exportType,
|
||||
})),
|
||||
[lorasFromStore],
|
||||
);
|
||||
const [localModels, setLocalModels] = useState<LoraModelOption[]>([]);
|
||||
|
||||
const loraModels = useMemo<LoraModelOption[]>(() => {
|
||||
const fromLoras = lorasFromStore.map((lora) => ({
|
||||
id: lora.id,
|
||||
name: lora.name,
|
||||
baseModel: lora.baseModel,
|
||||
updatedAt: lora.updatedAt,
|
||||
source: lora.source,
|
||||
exportType: lora.exportType,
|
||||
}));
|
||||
return [...fromLoras, ...localModels];
|
||||
}, [lorasFromStore, localModels]);
|
||||
|
||||
useEffect(() => {
|
||||
if (getTrainingCompareHandoff()) return;
|
||||
void refresh();
|
||||
void listLocalModels().then((res) => {
|
||||
setLocalModels(
|
||||
res.models
|
||||
.filter((m) => m.source === "lmstudio" || m.source === "models_dir")
|
||||
.map((m) => ({
|
||||
id: m.id,
|
||||
name: m.source === "lmstudio" && m.model_id ? m.model_id : m.display_name,
|
||||
baseModel: m.source === "lmstudio" ? "LM Studio" : "Local models",
|
||||
updatedAt: m.updated_at ?? undefined,
|
||||
source: "local" as const,
|
||||
})),
|
||||
);
|
||||
}).catch(() => {});
|
||||
}, [refresh]);
|
||||
|
||||
useEffect(() => {
|
||||
|
|
|
|||
|
|
@ -278,8 +278,18 @@ export function ChatSettingsPanel({
|
|||
const isMobile = useIsMobile();
|
||||
const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null;
|
||||
const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
|
||||
const ggufMaxContextLength = useChatRuntimeStore((s) => s.ggufMaxContextLength);
|
||||
const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype);
|
||||
const setKvCacheDtype = useChatRuntimeStore((s) => s.setKvCacheDtype);
|
||||
const loadedKvCacheDtype = useChatRuntimeStore((s) => s.loadedKvCacheDtype);
|
||||
const customContextLength = useChatRuntimeStore((s) => s.customContextLength);
|
||||
const setCustomContextLength = useChatRuntimeStore((s) => s.setCustomContextLength);
|
||||
|
||||
const ctxDisplayValue = customContextLength ?? ggufContextLength ?? "";
|
||||
const ctxMaxValue = ggufMaxContextLength ?? ggufContextLength ?? null;
|
||||
const kvDirty = kvCacheDtype !== loadedKvCacheDtype;
|
||||
const ctxDirty = customContextLength !== null;
|
||||
const modelSettingsDirty = kvDirty || ctxDirty;
|
||||
const [customPresets, setCustomPresets] = useState<Preset[]>(() =>
|
||||
loadSavedCustomPresets(),
|
||||
);
|
||||
|
|
@ -468,32 +478,53 @@ export function ChatSettingsPanel({
|
|||
<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 className="space-y-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium">Context Length</span>
|
||||
<Input
|
||||
type="number"
|
||||
value={typeof ctxDisplayValue === "number" ? ctxDisplayValue : (ggufContextLength ?? "")}
|
||||
placeholder="..."
|
||||
min={128}
|
||||
max={ctxMaxValue ?? undefined}
|
||||
step={1024}
|
||||
className="h-6 w-[100px] text-right text-xs tabular-nums"
|
||||
onChange={(e) => {
|
||||
const raw = e.target.value;
|
||||
if (raw === "") {
|
||||
setCustomContextLength(null);
|
||||
return;
|
||||
}
|
||||
const v = parseInt(raw, 10);
|
||||
if (!Number.isNaN(v) && v >= 0) {
|
||||
const maxCtx = ctxMaxValue ?? Infinity;
|
||||
const clamped = Math.min(v, maxCtx);
|
||||
setCustomContextLength(clamped === (ggufContextLength ?? 0) ? null : clamped);
|
||||
}
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
<Input
|
||||
value={ggufContextLength ?? ""}
|
||||
placeholder="Loading..."
|
||||
disabled={true}
|
||||
className="h-7 w-[90px] text-xs"
|
||||
<Slider
|
||||
min={1024}
|
||||
max={ctxMaxValue ?? 4096}
|
||||
step={1024}
|
||||
value={[Math.min(typeof ctxDisplayValue === "number" ? ctxDisplayValue : (ggufContextLength ?? 4096), ctxMaxValue ?? 4096)]}
|
||||
onValueChange={([v]) => {
|
||||
setCustomContextLength(v === (ggufContextLength ?? 0) ? null : v);
|
||||
}}
|
||||
/>
|
||||
</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.
|
||||
Quantize KV cache to reduce VRAM.
|
||||
</div>
|
||||
</div>
|
||||
<Select
|
||||
value={kvCacheDtype ?? "f16"}
|
||||
onValueChange={(v) => {
|
||||
setKvCacheDtype(v === "f16" ? null : v);
|
||||
onReloadModel?.();
|
||||
}}
|
||||
>
|
||||
<SelectTrigger className="h-7 w-[90px] text-xs">
|
||||
|
|
@ -508,14 +539,35 @@ export function ChatSettingsPanel({
|
|||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
{modelSettingsDirty && (
|
||||
<div className="flex flex-wrap gap-1.5 pt-1">
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => onReloadModel?.()}
|
||||
className="rounded-md bg-primary px-2.5 py-1 text-[11px] font-medium text-primary-foreground transition-colors hover:bg-primary/90"
|
||||
>
|
||||
Apply
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setCustomContextLength(null);
|
||||
setKvCacheDtype(loadedKvCacheDtype);
|
||||
}}
|
||||
className="rounded-md border px-2.5 py-1 text-[11px] font-medium text-muted-foreground transition-colors hover:bg-accent"
|
||||
>
|
||||
Reset
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
{!isGguf && (
|
||||
{!isGguf && params.checkpoint && (
|
||||
<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-xs font-medium">Enable custom code</div>
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
Allow models with custom code (e.g. Nemotron). Only enable for repos you trust.
|
||||
Allow models with custom code (e.g. Nemotron). Only enable if sure.
|
||||
</div>
|
||||
</div>
|
||||
<Switch
|
||||
|
|
@ -633,6 +685,7 @@ export function ChatSettingsPanel({
|
|||
onCheckedChange={onAutoTitleChange}
|
||||
/>
|
||||
</div>
|
||||
<HfTokenField />
|
||||
</div>
|
||||
</CollapsibleSection>
|
||||
|
||||
|
|
@ -780,6 +833,29 @@ function AutoHealToolCallsToggle() {
|
|||
);
|
||||
}
|
||||
|
||||
function HfTokenField() {
|
||||
const hfToken = useChatRuntimeStore((s) => s.hfToken);
|
||||
const setHfToken = useChatRuntimeStore((s) => s.setHfToken);
|
||||
|
||||
return (
|
||||
<div className="flex flex-col gap-1.5">
|
||||
<div className="min-w-0">
|
||||
<div className="text-xs font-medium">Hugging Face Token</div>
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
For downloading gated or private models.
|
||||
</div>
|
||||
</div>
|
||||
<Input
|
||||
type="password"
|
||||
value={hfToken}
|
||||
placeholder="hf_..."
|
||||
className="h-7 text-xs font-mono"
|
||||
onChange={(e) => setHfToken(e.target.value)}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function ChatTemplateSection({
|
||||
onReloadModel,
|
||||
}: {
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import Dexie, { type EntityTable, liveQuery } from "dexie";
|
||||
import { useEffect, useState } from "react";
|
||||
import { useEffect, useRef, useState } from "react";
|
||||
import type {
|
||||
FolderRecord,
|
||||
MemoryRecord,
|
||||
|
|
@ -84,18 +84,29 @@ db.version(4)
|
|||
|
||||
export { db };
|
||||
|
||||
/**
|
||||
* Wraps Dexie liveQuery for React state updates.
|
||||
*
|
||||
* Important: include every semantic query input in `deps` (filters, sort keys,
|
||||
* IDs, etc). `querier` identity is intentionally ignored to avoid re-subscribing
|
||||
* on every render when callers pass inline functions.
|
||||
*/
|
||||
export function useLiveQuery<T>(
|
||||
querier: () => Promise<T>,
|
||||
deps: unknown[] = [],
|
||||
): T | undefined {
|
||||
const [value, setValue] = useState<T>();
|
||||
const querierRef = useRef(querier);
|
||||
querierRef.current = querier;
|
||||
|
||||
useEffect(() => {
|
||||
const sub = liveQuery(querier).subscribe({
|
||||
const sub = liveQuery(() => querierRef.current()).subscribe({
|
||||
next: setValue,
|
||||
error: (err) => console.error("useLiveQuery:", err),
|
||||
});
|
||||
return () => sub.unsubscribe();
|
||||
// Intentionally omit `querier` from deps: inline functions would re-subscribe every render.
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [querier, ...deps]);
|
||||
}, deps);
|
||||
return value;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -238,11 +238,20 @@ export function useChatModelRuntime() {
|
|||
|
||||
// Restore reasoning/tools support flags and context length
|
||||
const supportsReasoning = statusRes.supports_reasoning ?? false;
|
||||
const reasoningAlwaysOn = statusRes.reasoning_always_on ?? false;
|
||||
const supportsTools = statusRes.supports_tools ?? false;
|
||||
const currentGgufContextLength = statusRes.is_gguf
|
||||
? (statusRes.context_length ?? null)
|
||||
: null;
|
||||
const ggufMaxContextLength = statusRes.is_gguf
|
||||
? (statusRes.max_context_length ?? null)
|
||||
: null;
|
||||
useChatRuntimeStore.setState({
|
||||
supportsReasoning,
|
||||
reasoningAlwaysOn,
|
||||
supportsTools,
|
||||
ggufContextLength: statusRes.is_gguf ? (statusRes.context_length ?? null) : null,
|
||||
ggufContextLength: currentGgufContextLength,
|
||||
ggufMaxContextLength,
|
||||
});
|
||||
|
||||
// Set reasoning default for Qwen3.5 small models
|
||||
|
|
@ -354,12 +363,13 @@ export function useChatModelRuntime() {
|
|||
useChatRuntimeStore.getState().params.checkpoint;
|
||||
const paramsBeforeLoad = useChatRuntimeStore.getState().params;
|
||||
const maxSeqLength = paramsBeforeLoad.maxSeqLength;
|
||||
const hfToken = useChatRuntimeStore.getState().hfToken || null;
|
||||
try {
|
||||
// Lightweight pre-flight validation: avoid unloading a working model
|
||||
// if the new identifier is clearly invalid (e.g. bad HF id / path).
|
||||
await validateModel({
|
||||
model_path: modelId,
|
||||
hf_token: null,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: maxSeqLength,
|
||||
load_in_4bit: true,
|
||||
is_lora: isLora,
|
||||
|
|
@ -371,11 +381,16 @@ export function useChatModelRuntime() {
|
|||
previousWasUnloaded = true;
|
||||
}
|
||||
|
||||
const { chatTemplateOverride, kvCacheDtype } = useChatRuntimeStore.getState();
|
||||
const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength } = useChatRuntimeStore.getState();
|
||||
// GGUF: use custom context length, or 0 = model's native context
|
||||
// Non-GGUF: use the Max Seq Length slider value
|
||||
const effectiveMaxSeqLength = customContextLength != null
|
||||
? customContextLength
|
||||
: ggufVariant != null ? (ggufContextLength ?? 0) : maxSeqLength;
|
||||
const loadResponse = await loadModel({
|
||||
model_path: modelId,
|
||||
hf_token: null,
|
||||
max_seq_length: maxSeqLength,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: effectiveMaxSeqLength,
|
||||
load_in_4bit: true,
|
||||
is_lora: isLora,
|
||||
gguf_variant: ggufVariant ?? null,
|
||||
|
|
@ -403,15 +418,30 @@ export function useChatModelRuntime() {
|
|||
}
|
||||
}
|
||||
}
|
||||
const loadedKv = loadResponse.cache_type_kv ?? null;
|
||||
const nativeCtx = loadResponse.is_gguf
|
||||
? (loadResponse.context_length ?? 131072)
|
||||
: null;
|
||||
const reportedMaxCtx = loadResponse.is_gguf
|
||||
? (loadResponse.max_context_length ?? null)
|
||||
: null;
|
||||
// A successful reload has applied settings, so clear pending custom
|
||||
// context state and display the backend-reported effective context.
|
||||
const keepCustomCtx = null;
|
||||
const reasoningAlwaysOn = loadResponse.reasoning_always_on ?? false;
|
||||
const ggufMaxContextLength = reportedMaxCtx;
|
||||
useChatRuntimeStore.setState({
|
||||
ggufContextLength: loadResponse.is_gguf
|
||||
? (loadResponse.context_length ?? 131072)
|
||||
: null,
|
||||
ggufContextLength: nativeCtx,
|
||||
ggufMaxContextLength,
|
||||
supportsReasoning: loadResponse.supports_reasoning ?? false,
|
||||
reasoningEnabled: reasoningDefault,
|
||||
reasoningAlwaysOn,
|
||||
reasoningEnabled: reasoningAlwaysOn ? true : reasoningDefault,
|
||||
supportsTools: loadResponse.supports_tools ?? false,
|
||||
toolsEnabled: false,
|
||||
kvCacheDtype: loadResponse.cache_type_kv ?? null,
|
||||
toolsEnabled: loadResponse.supports_tools ?? false,
|
||||
codeToolsEnabled: loadResponse.supports_tools ?? false,
|
||||
kvCacheDtype: loadedKv,
|
||||
loadedKvCacheDtype: loadedKv,
|
||||
customContextLength: keepCustomCtx,
|
||||
defaultChatTemplate: loadResponse.chat_template ?? null,
|
||||
chatTemplateOverride: null,
|
||||
});
|
||||
|
|
@ -432,7 +462,7 @@ export function useChatModelRuntime() {
|
|||
try {
|
||||
await loadModel({
|
||||
model_path: previousCheckpoint,
|
||||
hf_token: null,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: maxSeqLength,
|
||||
load_in_4bit: true,
|
||||
is_lora: previousIsLora,
|
||||
|
|
|
|||
|
|
@ -241,6 +241,7 @@ export function SharedComposer({
|
|||
(s) => !!s.params.checkpoint && !s.modelLoading,
|
||||
);
|
||||
const supportsReasoning = useChatRuntimeStore((s) => s.supportsReasoning);
|
||||
const reasoningAlwaysOn = useChatRuntimeStore((s) => s.reasoningAlwaysOn);
|
||||
const reasoningEnabled = useChatRuntimeStore((s) => s.reasoningEnabled);
|
||||
const setReasoningEnabled = useChatRuntimeStore((s) => s.setReasoningEnabled);
|
||||
const supportsTools = useChatRuntimeStore((s) => s.supportsTools);
|
||||
|
|
@ -337,7 +338,7 @@ export function SharedComposer({
|
|||
async function ensureModelLoaded(sel: CompareModelSelection): Promise<string> {
|
||||
const resp = await loadModel({
|
||||
model_path: sel.id,
|
||||
hf_token: null,
|
||||
hf_token: useChatRuntimeStore.getState().hfToken || null,
|
||||
max_seq_length: maxSeqLength,
|
||||
load_in_4bit: true,
|
||||
is_lora: sel.isLora,
|
||||
|
|
@ -528,6 +529,7 @@ export function SharedComposer({
|
|||
type="button"
|
||||
disabled={reasoningDisabled}
|
||||
onClick={() => {
|
||||
if (reasoningAlwaysOn) return;
|
||||
const next = !reasoningEnabled;
|
||||
setReasoningEnabled(next);
|
||||
// Qwen3/3.5: adjust params for thinking on/off
|
||||
|
|
@ -544,13 +546,13 @@ export function SharedComposer({
|
|||
"flex items-center gap-0.5 rounded-full px-2 py-0.5 text-xs font-medium transition-colors",
|
||||
reasoningDisabled
|
||||
? "cursor-not-allowed opacity-40"
|
||||
: reasoningEnabled
|
||||
: (reasoningEnabled || reasoningAlwaysOn)
|
||||
? "bg-primary/10 text-primary hover:bg-primary/20"
|
||||
: "bg-muted text-muted-foreground hover:bg-muted-foreground/15",
|
||||
)}
|
||||
aria-label={reasoningEnabled ? "Disable thinking" : "Enable thinking"}
|
||||
>
|
||||
{reasoningEnabled && !reasoningDisabled ? (
|
||||
{(reasoningEnabled || reasoningAlwaysOn) && !reasoningDisabled ? (
|
||||
<LightbulbIcon className="size-3" />
|
||||
) : (
|
||||
<LightbulbOffIcon className="size-3" />
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ const AUTO_TITLE_KEY = "unsloth_chat_auto_title";
|
|||
const AUTO_HEAL_TOOL_CALLS_KEY = "unsloth_auto_heal_tool_calls";
|
||||
const MAX_TOOL_CALLS_KEY = "unsloth_max_tool_calls_per_message";
|
||||
const TOOL_CALL_TIMEOUT_KEY = "unsloth_tool_call_timeout";
|
||||
const HF_TOKEN_KEY = "unsloth_hf_token";
|
||||
const INFERENCE_PARAMS_KEY = "unsloth_chat_inference_params";
|
||||
let hasShownInferencePersistenceWarning = false;
|
||||
|
||||
|
|
@ -62,6 +63,24 @@ function saveInt(key: string, value: number): void {
|
|||
}
|
||||
}
|
||||
|
||||
function loadString(key: string, fallback: string): string {
|
||||
if (!canUseStorage()) return fallback;
|
||||
try {
|
||||
return localStorage.getItem(key) ?? fallback;
|
||||
} catch {
|
||||
return fallback;
|
||||
}
|
||||
}
|
||||
|
||||
function saveString(key: string, value: string): void {
|
||||
if (!canUseStorage()) return;
|
||||
try {
|
||||
localStorage.setItem(key, value);
|
||||
} catch {
|
||||
// ignore
|
||||
}
|
||||
}
|
||||
|
||||
function asFiniteNumber(value: unknown, fallback: number): number {
|
||||
return typeof value === "number" && Number.isFinite(value) ? value : fallback;
|
||||
}
|
||||
|
|
@ -127,10 +146,13 @@ type ChatRuntimeStore = {
|
|||
loras: ChatLoraSummary[];
|
||||
runningByThreadId: Record<string, boolean>;
|
||||
autoTitle: boolean;
|
||||
hfToken: string;
|
||||
modelsError: string | null;
|
||||
activeGgufVariant: string | null;
|
||||
ggufContextLength: number | null;
|
||||
ggufMaxContextLength: number | null;
|
||||
supportsReasoning: boolean;
|
||||
reasoningAlwaysOn: boolean;
|
||||
reasoningEnabled: boolean;
|
||||
supportsTools: boolean;
|
||||
toolsEnabled: boolean;
|
||||
|
|
@ -141,6 +163,8 @@ type ChatRuntimeStore = {
|
|||
maxToolCallsPerMessage: number;
|
||||
toolCallTimeout: number;
|
||||
kvCacheDtype: string | null;
|
||||
loadedKvCacheDtype: string | null;
|
||||
customContextLength: number | null;
|
||||
defaultChatTemplate: string | null;
|
||||
chatTemplateOverride: string | null;
|
||||
activeThreadId: string | null;
|
||||
|
|
@ -159,6 +183,7 @@ type ChatRuntimeStore = {
|
|||
setLoras: (loras: ChatLoraSummary[]) => void;
|
||||
setThreadRunning: (threadId: string, running: boolean) => void;
|
||||
setAutoTitle: (enabled: boolean) => void;
|
||||
setHfToken: (token: string) => void;
|
||||
setModelsError: (error: string | null) => void;
|
||||
setCheckpoint: (modelId: string, ggufVariant?: string | null) => void;
|
||||
setActiveThreadId: (threadId: string | null) => void;
|
||||
|
|
@ -172,6 +197,7 @@ type ChatRuntimeStore = {
|
|||
setMaxToolCallsPerMessage: (value: number) => void;
|
||||
setToolCallTimeout: (value: number) => void;
|
||||
setKvCacheDtype: (dtype: string | null) => void;
|
||||
setCustomContextLength: (v: number | null) => void;
|
||||
setChatTemplateOverride: (template: string | null) => void;
|
||||
setPendingAudio: (base64: string, name: string) => void;
|
||||
clearPendingAudio: () => void;
|
||||
|
|
@ -184,10 +210,13 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
loras: [],
|
||||
runningByThreadId: {},
|
||||
autoTitle: loadBool(AUTO_TITLE_KEY, false),
|
||||
hfToken: loadString(HF_TOKEN_KEY, ""),
|
||||
modelsError: null,
|
||||
activeGgufVariant: null,
|
||||
ggufContextLength: null,
|
||||
ggufMaxContextLength: null,
|
||||
supportsReasoning: false,
|
||||
reasoningAlwaysOn: false,
|
||||
reasoningEnabled: true,
|
||||
supportsTools: false,
|
||||
toolsEnabled: false,
|
||||
|
|
@ -195,9 +224,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
toolStatus: null,
|
||||
generatingStatus: null,
|
||||
autoHealToolCalls: loadBool(AUTO_HEAL_TOOL_CALLS_KEY, true),
|
||||
maxToolCallsPerMessage: loadInt(MAX_TOOL_CALLS_KEY, 10),
|
||||
maxToolCallsPerMessage: loadInt(MAX_TOOL_CALLS_KEY, 25),
|
||||
toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5),
|
||||
kvCacheDtype: null,
|
||||
loadedKvCacheDtype: null,
|
||||
customContextLength: null,
|
||||
defaultChatTemplate: null,
|
||||
chatTemplateOverride: null,
|
||||
activeThreadId: null,
|
||||
|
|
@ -235,6 +266,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
saveBool(AUTO_TITLE_KEY, autoTitle);
|
||||
return { autoTitle };
|
||||
}),
|
||||
setHfToken: (hfToken) =>
|
||||
set(() => {
|
||||
saveString(HF_TOKEN_KEY, hfToken);
|
||||
return { hfToken };
|
||||
}),
|
||||
setModelsError: (modelsError) => set({ modelsError }),
|
||||
setCheckpoint: (modelId, ggufVariant) =>
|
||||
set((state) => ({
|
||||
|
|
@ -253,6 +289,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
},
|
||||
activeGgufVariant: null,
|
||||
ggufContextLength: null,
|
||||
ggufMaxContextLength: null,
|
||||
contextUsage: null,
|
||||
supportsReasoning: false,
|
||||
reasoningEnabled: true,
|
||||
|
|
@ -261,6 +298,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
codeToolsEnabled: false,
|
||||
toolStatus: null,
|
||||
kvCacheDtype: null,
|
||||
loadedKvCacheDtype: null,
|
||||
customContextLength: null,
|
||||
defaultChatTemplate: null,
|
||||
chatTemplateOverride: null,
|
||||
})),
|
||||
|
|
@ -285,6 +324,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
return { toolCallTimeout };
|
||||
}),
|
||||
setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }),
|
||||
setCustomContextLength: (customContextLength) => set({ customContextLength }),
|
||||
setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }),
|
||||
setPendingAudio: (base64, name) =>
|
||||
set({ pendingAudioBase64: base64, pendingAudioName: name }),
|
||||
|
|
|
|||
|
|
@ -455,7 +455,7 @@ export function ThreadSidebar({
|
|||
<span>Learn more in docs</span>
|
||||
</a>
|
||||
<a
|
||||
href="https://unsloth.ai/blog"
|
||||
href="https://unsloth.ai/docs/new/changelog"
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="flex items-center gap-2 rounded-md px-2 py-1.5 text-xs text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
|
||||
|
|
|
|||
|
|
@ -86,7 +86,9 @@ export interface LoadModelResponse {
|
|||
trust_remote_code?: boolean;
|
||||
};
|
||||
context_length?: number | null;
|
||||
max_context_length?: number | null;
|
||||
supports_reasoning?: boolean;
|
||||
reasoning_always_on?: boolean;
|
||||
supports_tools?: boolean;
|
||||
cache_type_kv?: string | null;
|
||||
chat_template?: string | null;
|
||||
|
|
@ -115,8 +117,10 @@ export interface InferenceStatusResponse {
|
|||
trust_remote_code?: boolean;
|
||||
};
|
||||
supports_reasoning?: boolean;
|
||||
reasoning_always_on?: boolean;
|
||||
supports_tools?: boolean;
|
||||
context_length?: number | null;
|
||||
max_context_length?: number | null;
|
||||
}
|
||||
|
||||
export interface AudioGenerationResponse {
|
||||
|
|
|
|||
|
|
@ -3,6 +3,19 @@
|
|||
|
||||
import { SectionCard } from "@/components/section-card";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
Combobox,
|
||||
ComboboxContent,
|
||||
ComboboxEmpty,
|
||||
ComboboxInput,
|
||||
ComboboxItem,
|
||||
ComboboxList,
|
||||
} from "@/components/ui/combobox";
|
||||
import {
|
||||
InputGroup,
|
||||
InputGroupAddon,
|
||||
InputGroupInput,
|
||||
} from "@/components/ui/input-group";
|
||||
import {
|
||||
Select,
|
||||
SelectContent,
|
||||
|
|
@ -11,17 +24,34 @@ import {
|
|||
SelectValue,
|
||||
} from "@/components/ui/select";
|
||||
import { Separator } from "@/components/ui/separator";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import {
|
||||
Tooltip,
|
||||
TooltipContent,
|
||||
TooltipTrigger,
|
||||
} from "@/components/ui/tooltip";
|
||||
import { useTrainingConfigStore } from "@/features/training";
|
||||
import { AlertCircleIcon, InformationCircleIcon, PackageIcon } from "@hugeicons/core-free-icons";
|
||||
import {
|
||||
listLocalModels,
|
||||
type LocalModelInfo,
|
||||
useTrainingConfigStore,
|
||||
} from "@/features/training";
|
||||
import {
|
||||
useDebouncedValue,
|
||||
useHfModelSearch,
|
||||
useHfTokenValidation,
|
||||
} from "@/hooks";
|
||||
import {
|
||||
AlertCircleIcon,
|
||||
FolderSearchIcon,
|
||||
InformationCircleIcon,
|
||||
Key01Icon,
|
||||
PackageIcon,
|
||||
Search01Icon,
|
||||
} from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { AnimatePresence, motion } from "motion/react";
|
||||
import { useCallback, useEffect, useMemo, useState } from "react";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { useShallow } from "zustand/react/shallow";
|
||||
import { collapseAnim } from "./anim";
|
||||
import type { ModelCheckpoints } from "./api/export-api";
|
||||
|
|
@ -60,6 +90,22 @@ export function ExportPage() {
|
|||
|
||||
const [selectedModelIdx, setSelectedModelIdx] = useState<string | null>(null);
|
||||
const [checkpoint, setCheckpoint] = useState<string | null>(null);
|
||||
const [sourceMode, setSourceMode] = useState<"checkpoint" | "model">(
|
||||
"checkpoint",
|
||||
);
|
||||
const [modelSource, setModelSource] = useState<"hf" | "local">("hf");
|
||||
const [hfExportTrustRemoteCode, setHfExportTrustRemoteCode] =
|
||||
useState(true);
|
||||
const [modelInput, setModelInput] = useState("");
|
||||
const [selectedSourceModel, setSelectedSourceModel] = useState<string | null>(
|
||||
null,
|
||||
);
|
||||
const [localModelInput, setLocalModelInput] = useState("");
|
||||
const [localModels, setLocalModels] = useState<LocalModelInfo[]>([]);
|
||||
const [isLoadingLocalModels, setIsLoadingLocalModels] = useState(true);
|
||||
const [localModelsError, setLocalModelsError] = useState<string | null>(null);
|
||||
const debouncedModelQuery = useDebouncedValue(modelInput);
|
||||
const debouncedHfToken = useDebouncedValue(hfToken, 500);
|
||||
|
||||
const [exportMethod, setExportMethod] = useState<ExportMethod | null>(null);
|
||||
const [quantLevels, setQuantLevels] = useState<string[]>([]);
|
||||
|
|
@ -74,6 +120,9 @@ export function ExportPage() {
|
|||
const [exportError, setExportError] = useState<string | null>(null);
|
||||
const [exportSuccess, setExportSuccess] = useState(false);
|
||||
|
||||
const hfComboboxAnchorRef = useRef<HTMLDivElement>(null);
|
||||
const localComboboxAnchorRef = useRef<HTMLDivElement>(null);
|
||||
|
||||
const tour = useGuidedTourController({
|
||||
id: "export",
|
||||
steps: exportTourSteps,
|
||||
|
|
@ -105,6 +154,27 @@ export function ExportPage() {
|
|||
};
|
||||
}, []);
|
||||
|
||||
// ---- Fetch local models for direct export ----
|
||||
useEffect(() => {
|
||||
const controller = new AbortController();
|
||||
void listLocalModels(controller.signal)
|
||||
.then((models) => {
|
||||
if (controller.signal.aborted) return;
|
||||
setLocalModels(models);
|
||||
})
|
||||
.catch((error) => {
|
||||
if (controller.signal.aborted) return;
|
||||
setLocalModelsError(
|
||||
error instanceof Error ? error.message : "Failed to load local models",
|
||||
);
|
||||
})
|
||||
.finally(() => {
|
||||
if (controller.signal.aborted) return;
|
||||
setIsLoadingLocalModels(false);
|
||||
});
|
||||
return () => controller.abort();
|
||||
}, []);
|
||||
|
||||
// ---- Derived state ----
|
||||
const selectedModelData = useMemo(
|
||||
() =>
|
||||
|
|
@ -127,6 +197,83 @@ export function ExportPage() {
|
|||
const trainingMethodLabel = selectedModelData?.peft_type
|
||||
? "LoRA / QLoRA"
|
||||
: "Full Fine-tune";
|
||||
const sourceBaseModelName = sourceMode === "model"
|
||||
? selectedSourceModel ?? "—"
|
||||
: baseModelName;
|
||||
|
||||
const {
|
||||
results: hfResults,
|
||||
isLoading: isLoadingHfModels,
|
||||
error: hfSearchError,
|
||||
} = useHfModelSearch(debouncedModelQuery, {
|
||||
accessToken: debouncedHfToken || undefined,
|
||||
excludeGguf: true,
|
||||
});
|
||||
const { error: tokenValidationError, isChecking: isCheckingToken } =
|
||||
useHfTokenValidation(hfToken);
|
||||
|
||||
const hfResultIds = useMemo(() => {
|
||||
const ids = hfResults.map((r) => r.id);
|
||||
if (
|
||||
selectedSourceModel &&
|
||||
modelSource === "hf" &&
|
||||
!ids.includes(selectedSourceModel)
|
||||
) {
|
||||
ids.push(selectedSourceModel);
|
||||
}
|
||||
return ids;
|
||||
}, [hfResults, modelSource, selectedSourceModel]);
|
||||
|
||||
const exportableLocalModels = useMemo(
|
||||
() =>
|
||||
localModels.filter((m) => {
|
||||
if (m.path.endsWith(".gguf")) return false;
|
||||
if (m.id.toLowerCase().includes("-gguf")) return false;
|
||||
return true;
|
||||
}),
|
||||
[localModels],
|
||||
);
|
||||
|
||||
const localMetaById = useMemo(() => {
|
||||
const map = new Map<string, LocalModelInfo>();
|
||||
for (const model of exportableLocalModels) map.set(model.id, model);
|
||||
return map;
|
||||
}, [exportableLocalModels]);
|
||||
|
||||
const localResultIds = useMemo(() => {
|
||||
const ids = exportableLocalModels.map((model) => model.id);
|
||||
const manual = localModelInput.trim();
|
||||
if (manual && !ids.includes(manual)) {
|
||||
ids.unshift(manual);
|
||||
}
|
||||
return ids;
|
||||
}, [exportableLocalModels, localModelInput]);
|
||||
|
||||
const localFilteredIds = useMemo(() => {
|
||||
const q = localModelInput.trim().toLowerCase();
|
||||
if (!q) return localResultIds;
|
||||
return localResultIds.filter((id) => {
|
||||
const meta = localMetaById.get(id);
|
||||
if (id.toLowerCase().includes(q)) return true;
|
||||
if (meta?.display_name.toLowerCase().includes(q)) return true;
|
||||
if (meta?.path.toLowerCase().includes(q)) return true;
|
||||
return false;
|
||||
});
|
||||
}, [localMetaById, localModelInput, localResultIds]);
|
||||
|
||||
const exportGuideSteps = useMemo(
|
||||
() =>
|
||||
sourceMode === "model"
|
||||
? [
|
||||
"Select a Hugging Face or local model to export from",
|
||||
"GGUF is used for non-finetuned model exports",
|
||||
"Pick one or more GGUF quantization levels",
|
||||
"Click Export and choose your destination",
|
||||
"Test your model and compare outputs in Chat",
|
||||
]
|
||||
: GUIDE_STEPS,
|
||||
[sourceMode],
|
||||
);
|
||||
|
||||
// Reset checkpoint when the selected model changes
|
||||
useEffect(() => {
|
||||
|
|
@ -144,6 +291,25 @@ export function ExportPage() {
|
|||
}
|
||||
}, [isAdapter, isQuantized, exportMethod]);
|
||||
|
||||
const handleSourceModeSwitch = useCallback(
|
||||
(next: "checkpoint" | "model") => {
|
||||
setSourceMode(next);
|
||||
if (next === "model") {
|
||||
setExportMethod("gguf");
|
||||
}
|
||||
setSelectedSourceModel(null);
|
||||
setLocalModelInput("");
|
||||
setModelInput("");
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
setSelectedSourceModel(null);
|
||||
setLocalModelInput("");
|
||||
setModelInput("");
|
||||
}, [modelSource]);
|
||||
|
||||
const handleMethodChange = (method: ExportMethod) => {
|
||||
setExportMethod(method);
|
||||
if (method !== "gguf") {
|
||||
|
|
@ -152,19 +318,24 @@ export function ExportPage() {
|
|||
};
|
||||
|
||||
const estimatedSize = getEstimatedSize(exportMethod, quantLevels);
|
||||
const canExport =
|
||||
checkpoint &&
|
||||
const selectedExportSource =
|
||||
sourceMode === "checkpoint" ? checkpoint : selectedSourceModel;
|
||||
const canExport = !!(
|
||||
selectedExportSource &&
|
||||
exportMethod &&
|
||||
(exportMethod !== "gguf" || quantLevels.length > 0);
|
||||
(exportMethod !== "gguf" || quantLevels.length > 0)
|
||||
);
|
||||
|
||||
// ---- Export handler ----
|
||||
const handleExport = useCallback(async () => {
|
||||
if (!checkpoint) return;
|
||||
const source = sourceMode === "checkpoint" ? checkpoint : selectedSourceModel;
|
||||
if (!source) return;
|
||||
|
||||
const selectedCp = checkpointsForModel.find(
|
||||
(cp) => cp.display_name === checkpoint,
|
||||
);
|
||||
if (!selectedCp) return;
|
||||
const selectedCp = sourceMode === "checkpoint"
|
||||
? checkpointsForModel.find((cp) => cp.display_name === checkpoint)
|
||||
: null;
|
||||
if (sourceMode === "checkpoint" && !selectedCp) return;
|
||||
const checkpointPath = selectedCp?.path;
|
||||
|
||||
setExporting(true);
|
||||
setExportError(null);
|
||||
|
|
@ -174,7 +345,8 @@ export function ExportPage() {
|
|||
// For other formats, nest under training-run/checkpoint
|
||||
const saveDir =
|
||||
exportMethod === "gguf"
|
||||
? `${baseModelName.split("/").pop() ?? selectedModelIdx ?? "model"}-finetune-gguf`
|
||||
? `${(sourceBaseModelName.split("/").pop() ?? selectedModelIdx ?? "model")
|
||||
.replace(/[^a-zA-Z0-9._-]/g, "-")}-gguf`
|
||||
: `${selectedModelIdx ?? "model"}/${checkpoint}`;
|
||||
const pushToHub = destination === "hub";
|
||||
const repoId = pushToHub && hfUsername && modelName
|
||||
|
|
@ -183,8 +355,18 @@ export function ExportPage() {
|
|||
const token = pushToHub && hfToken ? hfToken : undefined;
|
||||
|
||||
try {
|
||||
// 1. Load checkpoint
|
||||
await loadCheckpoint({ checkpoint_path: selectedCp.path });
|
||||
// 1. Load model source
|
||||
if (sourceMode === "checkpoint") {
|
||||
if (!checkpointPath) return;
|
||||
await loadCheckpoint({ checkpoint_path: checkpointPath });
|
||||
} else {
|
||||
await loadCheckpoint({
|
||||
checkpoint_path: source,
|
||||
load_in_4bit: false,
|
||||
trust_remote_code:
|
||||
modelSource === "hf" ? hfExportTrustRemoteCode : true,
|
||||
});
|
||||
}
|
||||
|
||||
// 2. Run export based on method
|
||||
if (exportMethod === "merged") {
|
||||
|
|
@ -242,16 +424,21 @@ export function ExportPage() {
|
|||
}, [
|
||||
checkpoint,
|
||||
checkpointsForModel,
|
||||
sourceMode,
|
||||
selectedSourceModel,
|
||||
selectedModelIdx,
|
||||
selectedModelData,
|
||||
exportMethod,
|
||||
isAdapter,
|
||||
sourceBaseModelName,
|
||||
quantLevels,
|
||||
destination,
|
||||
hfUsername,
|
||||
modelName,
|
||||
hfToken,
|
||||
privateRepo,
|
||||
modelSource,
|
||||
hfExportTrustRemoteCode,
|
||||
]);
|
||||
|
||||
// ---- Render ----
|
||||
|
|
@ -265,14 +452,14 @@ export function ExportPage() {
|
|||
Export Model
|
||||
</h1>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
Export your fine-tuned model for deployment
|
||||
Export fine-tuned or base models for deployment
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<SectionCard
|
||||
icon={<HugeiconsIcon icon={PackageIcon} className="size-5" />}
|
||||
title="Export Configuration"
|
||||
description="Select checkpoint, method, and quantization"
|
||||
description="Select source, method, and quantization"
|
||||
accent="emerald"
|
||||
featured={true}
|
||||
className="shadow-border ring-1 ring-border"
|
||||
|
|
@ -296,11 +483,10 @@ export function ExportPage() {
|
|||
<>
|
||||
{/* Top row: Dropdowns + metadata | Guide */}
|
||||
<div className="grid grid-cols-1 gap-6 md:grid-cols-2 md:gap-8">
|
||||
<div className="flex flex-col gap-4">
|
||||
{/* Training run dropdown */}
|
||||
<div data-tour="export-training-run" className="flex flex-col gap-2">
|
||||
<div className="flex flex-col gap-2">
|
||||
<div className="flex items-end justify-between">
|
||||
<label className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
|
||||
Training Run
|
||||
{sourceMode === "checkpoint" ? "Training Run" : "Model Source"}
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<button
|
||||
|
|
@ -314,147 +500,415 @@ export function ExportPage() {
|
|||
</button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent>
|
||||
Select the training run that produced the checkpoints
|
||||
you want to export.
|
||||
{sourceMode === "checkpoint"
|
||||
? "Select the training run that produced the checkpoints you want to export."
|
||||
: "Select a Hugging Face model or local model path to export directly to GGUF."}
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
</label>
|
||||
<Select
|
||||
value={selectedModelIdx ?? ""}
|
||||
onValueChange={setSelectedModelIdx}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() =>
|
||||
handleSourceModeSwitch(
|
||||
sourceMode === "checkpoint" ? "model" : "checkpoint",
|
||||
)
|
||||
}
|
||||
className="text-xs text-primary underline cursor-pointer leading-none"
|
||||
>
|
||||
<SelectTrigger className="w-full">
|
||||
<SelectValue
|
||||
placeholder={
|
||||
models.length === 0
|
||||
? "No training runs found"
|
||||
: "Select a training run…"
|
||||
}
|
||||
/>
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{models.map((m) => {
|
||||
const tsMatch = m.name.match(/_(\d{10,})$/);
|
||||
const displayName = tsMatch ? m.name.slice(0, tsMatch.index) : m.name;
|
||||
const timeStr = tsMatch
|
||||
? new Date(Number(tsMatch[1]) * 1000).toLocaleString(undefined, {
|
||||
dateStyle: "medium",
|
||||
timeStyle: "short",
|
||||
})
|
||||
: null;
|
||||
return (
|
||||
<SelectItem key={m.name} value={m.name}>
|
||||
<span className="flex items-center gap-2">
|
||||
{displayName}
|
||||
<span className="text-muted-foreground text-xs">
|
||||
{m.checkpoints.length} checkpoint
|
||||
{m.checkpoints.length !== 1 ? "s" : ""}
|
||||
</span>
|
||||
{timeStr && (
|
||||
<span className="text-muted-foreground text-xs">
|
||||
· {timeStr}
|
||||
{sourceMode === "checkpoint"
|
||||
? "Use Hugging Face / Local Model"
|
||||
: "Use Training Checkpoints"}
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<AnimatePresence mode="wait" initial={false}>
|
||||
{sourceMode === "checkpoint" ? (
|
||||
<motion.div
|
||||
key="checkpoint"
|
||||
initial={{ height: 0, opacity: 0 }}
|
||||
animate={{ height: "auto", opacity: 1 }}
|
||||
exit={{ height: 0, opacity: 0 }}
|
||||
transition={{ duration: 0.25, ease: [0.25, 0.1, 0.25, 1] }}
|
||||
className="flex flex-col gap-2 overflow-visible"
|
||||
>
|
||||
<div data-tour="export-training-run" className="flex flex-col gap-2">
|
||||
<Select
|
||||
value={selectedModelIdx ?? ""}
|
||||
onValueChange={setSelectedModelIdx}
|
||||
>
|
||||
<SelectTrigger className="w-full">
|
||||
<SelectValue
|
||||
placeholder={
|
||||
models.length === 0
|
||||
? "No training runs found"
|
||||
: "Select a training run…"
|
||||
}
|
||||
/>
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{models.map((m) => {
|
||||
const tsMatch = m.name.match(/_(\d{10,})$/);
|
||||
const displayName = tsMatch
|
||||
? m.name.slice(0, tsMatch.index)
|
||||
: m.name;
|
||||
const timeStr = tsMatch
|
||||
? new Date(Number(tsMatch[1]) * 1000).toLocaleString(
|
||||
undefined,
|
||||
{
|
||||
dateStyle: "medium",
|
||||
timeStyle: "short",
|
||||
},
|
||||
)
|
||||
: null;
|
||||
return (
|
||||
<SelectItem key={m.name} value={m.name}>
|
||||
<span className="flex items-center gap-2">
|
||||
{displayName}
|
||||
<span className="text-muted-foreground text-xs">
|
||||
{m.checkpoints.length} checkpoint
|
||||
{m.checkpoints.length !== 1 ? "s" : ""}
|
||||
</span>
|
||||
{timeStr && (
|
||||
<span className="text-muted-foreground text-xs">
|
||||
· {timeStr}
|
||||
</span>
|
||||
)}
|
||||
</span>
|
||||
)}
|
||||
</span>
|
||||
</SelectItem>
|
||||
);
|
||||
})}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
</SelectItem>
|
||||
);
|
||||
})}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
|
||||
{/* Checkpoint dropdown */}
|
||||
<div data-tour="export-checkpoint" className="flex flex-col gap-2">
|
||||
<label className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
|
||||
Checkpoint
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<button
|
||||
type="button"
|
||||
className="text-foreground/70 hover:text-foreground"
|
||||
>
|
||||
<HugeiconsIcon
|
||||
icon={InformationCircleIcon}
|
||||
className="size-3"
|
||||
<div data-tour="export-checkpoint" className="flex flex-col gap-2">
|
||||
<label className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
|
||||
Checkpoint
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<button
|
||||
type="button"
|
||||
className="text-foreground/70 hover:text-foreground"
|
||||
>
|
||||
<HugeiconsIcon
|
||||
icon={InformationCircleIcon}
|
||||
className="size-3"
|
||||
/>
|
||||
</button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent>
|
||||
Choose a saved checkpoint to export. Lower loss
|
||||
generally means better quality.{" "}
|
||||
<a
|
||||
href="https://unsloth.ai/docs/basics/inference-and-deployment"
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="text-primary underline"
|
||||
>
|
||||
Read more
|
||||
</a>
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
</label>
|
||||
<Select
|
||||
value={checkpoint ?? ""}
|
||||
onValueChange={setCheckpoint}
|
||||
disabled={!selectedModelIdx}
|
||||
>
|
||||
<SelectTrigger className="w-full">
|
||||
<SelectValue
|
||||
placeholder={
|
||||
!selectedModelIdx
|
||||
? "Select a training run first"
|
||||
: checkpointsForModel.length === 0
|
||||
? "No checkpoints found"
|
||||
: "Select a checkpoint…"
|
||||
}
|
||||
/>
|
||||
</button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent>
|
||||
Choose a saved checkpoint to export. Lower loss
|
||||
generally means better quality.{" "}
|
||||
<a
|
||||
href="https://unsloth.ai/docs/basics/inference-and-deployment"
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="text-primary underline"
|
||||
>
|
||||
Read more
|
||||
</a>
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
</label>
|
||||
<Select
|
||||
value={checkpoint ?? ""}
|
||||
onValueChange={setCheckpoint}
|
||||
disabled={!selectedModelIdx}
|
||||
>
|
||||
<SelectTrigger className="w-full">
|
||||
<SelectValue
|
||||
placeholder={
|
||||
!selectedModelIdx
|
||||
? "Select a training run first"
|
||||
: checkpointsForModel.length === 0
|
||||
? "No checkpoints found"
|
||||
: "Select a checkpoint…"
|
||||
}
|
||||
/>
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{checkpointsForModel.map((cp) => (
|
||||
<SelectItem key={cp.path} value={cp.display_name}>
|
||||
<span className="flex items-center gap-2">
|
||||
{cp.display_name}
|
||||
{cp.loss != null && (
|
||||
<span className="text-muted-foreground text-xs">
|
||||
loss: {cp.loss.toFixed(4)}
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{checkpointsForModel.map((cp) => (
|
||||
<SelectItem key={cp.path} value={cp.display_name}>
|
||||
<span className="flex items-center gap-2">
|
||||
{cp.display_name}
|
||||
{cp.loss != null && (
|
||||
<span className="text-muted-foreground text-xs">
|
||||
loss: {cp.loss.toFixed(4)}
|
||||
</span>
|
||||
)}
|
||||
</span>
|
||||
)}
|
||||
</span>
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
</motion.div>
|
||||
) : (
|
||||
<motion.div
|
||||
key="model"
|
||||
initial={{ height: 0, opacity: 0 }}
|
||||
animate={{ height: "auto", opacity: 1 }}
|
||||
exit={{ height: 0, opacity: 0 }}
|
||||
transition={{ duration: 0.25, ease: [0.25, 0.1, 0.25, 1] }}
|
||||
className="flex flex-col gap-2 overflow-visible"
|
||||
>
|
||||
<div className="flex gap-2">
|
||||
<Button
|
||||
variant={modelSource === "hf" ? "dark" : "outline"}
|
||||
className="flex-1"
|
||||
onClick={() => setModelSource("hf")}
|
||||
>
|
||||
Hugging Face
|
||||
</Button>
|
||||
<Button
|
||||
variant={modelSource === "local" ? "dark" : "outline"}
|
||||
className="flex-1"
|
||||
onClick={() => setModelSource("local")}
|
||||
>
|
||||
Local Model
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
<div className="rounded-xl bg-muted/50 p-3 flex flex-col gap-2">
|
||||
<span className="text-[11px] font-medium text-muted-foreground uppercase tracking-wider">
|
||||
Training Info
|
||||
</span>
|
||||
<div className="grid grid-cols-1 gap-x-6 gap-y-1.5 text-xs sm:grid-cols-2">
|
||||
<div className="flex justify-between">
|
||||
<span className="text-muted-foreground">Base Model</span>
|
||||
<span className="font-medium">{baseModelName}</span>
|
||||
</div>
|
||||
<div className="flex justify-between">
|
||||
<span className="text-muted-foreground">Method</span>
|
||||
<span className="font-medium">
|
||||
{trainingMethodLabel}
|
||||
</span>
|
||||
</div>
|
||||
<div className="flex justify-between">
|
||||
<span className="text-muted-foreground">Checkpoints</span>
|
||||
<span className="font-medium">
|
||||
{checkpointsForModel.length}
|
||||
</span>
|
||||
</div>
|
||||
{isAdapter && (
|
||||
<div className="flex justify-between">
|
||||
<span className="text-muted-foreground">LoRA Rank</span>
|
||||
<span className="font-medium">{loraRank}</span>
|
||||
{modelSource === "hf" ? (
|
||||
<>
|
||||
<div className="flex flex-col gap-2">
|
||||
<label className="text-xs font-medium text-muted-foreground">
|
||||
Hugging Face Model
|
||||
</label>
|
||||
<div ref={hfComboboxAnchorRef}>
|
||||
<Combobox
|
||||
items={hfResultIds}
|
||||
filteredItems={hfResultIds}
|
||||
filter={null}
|
||||
value={selectedSourceModel}
|
||||
onValueChange={setSelectedSourceModel}
|
||||
onInputValueChange={(val) => {
|
||||
setModelInput(val);
|
||||
setSelectedSourceModel(null);
|
||||
}}
|
||||
itemToStringValue={(id) => id}
|
||||
autoHighlight={true}
|
||||
>
|
||||
<ComboboxInput placeholder="Search models..." className="w-full">
|
||||
<InputGroupAddon>
|
||||
<HugeiconsIcon icon={Search01Icon} className="size-4" />
|
||||
</InputGroupAddon>
|
||||
</ComboboxInput>
|
||||
<ComboboxContent anchor={hfComboboxAnchorRef}>
|
||||
{isLoadingHfModels ? (
|
||||
<div className="flex items-center justify-center py-4 gap-2 text-xs text-muted-foreground">
|
||||
<Spinner className="size-4" /> Searching…
|
||||
</div>
|
||||
) : (
|
||||
<ComboboxEmpty>No models found</ComboboxEmpty>
|
||||
)}
|
||||
<ComboboxList className="p-1 !max-h-none !overflow-visible">
|
||||
{(id: string) => (
|
||||
<ComboboxItem key={id} value={id} className="gap-2">
|
||||
<span className="block min-w-0 flex-1 truncate">
|
||||
{id}
|
||||
</span>
|
||||
</ComboboxItem>
|
||||
)}
|
||||
</ComboboxList>
|
||||
</ComboboxContent>
|
||||
</Combobox>
|
||||
</div>
|
||||
{(tokenValidationError ?? hfSearchError) && (
|
||||
<p className="text-xs text-destructive">
|
||||
{tokenValidationError ?? hfSearchError}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
<div className="flex items-center gap-2">
|
||||
<Switch
|
||||
id="hf-export-trust-remote-code"
|
||||
size="sm"
|
||||
checked={hfExportTrustRemoteCode}
|
||||
onCheckedChange={setHfExportTrustRemoteCode}
|
||||
disabled={exporting}
|
||||
/>
|
||||
<label
|
||||
htmlFor="hf-export-trust-remote-code"
|
||||
className="cursor-pointer text-xs font-medium text-muted-foreground hover:text-foreground"
|
||||
>
|
||||
Trust remote code
|
||||
</label>
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<button
|
||||
type="button"
|
||||
className="text-muted-foreground hover:text-foreground -m-1 inline-flex rounded p-1"
|
||||
aria-label="About trust remote code"
|
||||
>
|
||||
<HugeiconsIcon
|
||||
icon={InformationCircleIcon}
|
||||
className="size-3.5"
|
||||
/>
|
||||
</button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent
|
||||
side="top"
|
||||
className="max-w-[260px] text-xs"
|
||||
>
|
||||
Loads custom Python from the repo if the model
|
||||
needs it. Turn off if you do not trust the
|
||||
source.
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
</div>
|
||||
<div className="flex flex-col gap-1.5">
|
||||
<label className="text-xs font-medium text-muted-foreground">
|
||||
Hugging Face Token (Optional)
|
||||
</label>
|
||||
<InputGroup>
|
||||
<InputGroupAddon>
|
||||
<HugeiconsIcon icon={Key01Icon} className="size-4" />
|
||||
</InputGroupAddon>
|
||||
<InputGroupInput
|
||||
type="password"
|
||||
autoComplete="new-password"
|
||||
name="hf-token-export-source"
|
||||
placeholder="hf_..."
|
||||
value={hfToken}
|
||||
onChange={(e) => setHfToken(e.target.value)}
|
||||
/>
|
||||
</InputGroup>
|
||||
{isCheckingToken && (
|
||||
<p className="text-xs text-muted-foreground">Checking token…</p>
|
||||
)}
|
||||
</div>
|
||||
</>
|
||||
) : (
|
||||
<div className="flex flex-col gap-2">
|
||||
<label className="text-xs font-medium text-muted-foreground">
|
||||
Local Model Path
|
||||
</label>
|
||||
<div ref={localComboboxAnchorRef}>
|
||||
<Combobox
|
||||
items={localResultIds}
|
||||
filteredItems={localFilteredIds}
|
||||
filter={null}
|
||||
value={localModelInput || null}
|
||||
onValueChange={(id) => {
|
||||
const next = id ?? "";
|
||||
setLocalModelInput(next);
|
||||
setSelectedSourceModel(next || null);
|
||||
}}
|
||||
onInputValueChange={setLocalModelInput}
|
||||
itemToStringValue={(id) => id}
|
||||
autoHighlight={true}
|
||||
>
|
||||
<ComboboxInput
|
||||
placeholder={
|
||||
isLoadingLocalModels
|
||||
? "Scanning local and cached models..."
|
||||
: "./models/my-model"
|
||||
}
|
||||
className="w-full"
|
||||
onBlur={() =>
|
||||
setSelectedSourceModel(localModelInput.trim() || null)
|
||||
}
|
||||
onKeyDown={(event) => {
|
||||
if (event.key !== "Enter") return;
|
||||
event.preventDefault();
|
||||
setSelectedSourceModel(localModelInput.trim() || null);
|
||||
}}
|
||||
>
|
||||
<InputGroupAddon>
|
||||
<HugeiconsIcon icon={FolderSearchIcon} className="size-4" />
|
||||
</InputGroupAddon>
|
||||
</ComboboxInput>
|
||||
<ComboboxContent anchor={localComboboxAnchorRef}>
|
||||
{isLoadingLocalModels ? (
|
||||
<div className="flex items-center justify-center gap-2 py-4 text-xs text-muted-foreground">
|
||||
<Spinner className="size-4" /> Scanning...
|
||||
</div>
|
||||
) : localModelsError ? (
|
||||
<div className="px-3 py-2 text-xs text-red-500">
|
||||
{localModelsError}
|
||||
</div>
|
||||
) : (
|
||||
<ComboboxEmpty>No local models found</ComboboxEmpty>
|
||||
)}
|
||||
<ComboboxList className="p-1 !max-h-none !overflow-visible">
|
||||
{(id: string) => {
|
||||
const model = localMetaById.get(id);
|
||||
const source =
|
||||
model?.source === "hf_cache"
|
||||
? "HF cache"
|
||||
: "Local dir";
|
||||
return (
|
||||
<ComboboxItem key={id} value={id} className="gap-2">
|
||||
<span className="block min-w-0 flex-1 truncate">
|
||||
{model?.display_name ?? id}
|
||||
</span>
|
||||
<span className="ml-auto shrink-0 text-[10px] text-muted-foreground">
|
||||
{source}
|
||||
</span>
|
||||
</ComboboxItem>
|
||||
);
|
||||
}}
|
||||
</ComboboxList>
|
||||
</ComboboxContent>
|
||||
</Combobox>
|
||||
</div>
|
||||
{isLoadingLocalModels ? (
|
||||
<p className="text-[10px] text-muted-foreground">
|
||||
Scanning local models...
|
||||
</p>
|
||||
) : localModelsError ? (
|
||||
<p className="text-[10px] text-red-500">{localModelsError}</p>
|
||||
) : (
|
||||
<p className="text-[10px] text-muted-foreground">
|
||||
{exportableLocalModels.length > 0
|
||||
? `${exportableLocalModels.length} local/cached models found`
|
||||
: "No local models found. Enter path manually."}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="rounded-xl bg-muted/50 p-3">
|
||||
<p className="text-[11px] text-muted-foreground">
|
||||
Direct model exports currently support GGUF only.
|
||||
</p>
|
||||
</div>
|
||||
</motion.div>
|
||||
)}
|
||||
</AnimatePresence>
|
||||
|
||||
{sourceMode === "checkpoint" && (
|
||||
<div className="rounded-xl bg-muted/50 p-3 flex flex-col gap-2">
|
||||
<span className="text-[11px] font-medium text-muted-foreground uppercase tracking-wider">
|
||||
Training Info
|
||||
</span>
|
||||
<div className="grid grid-cols-1 gap-x-6 gap-y-1.5 text-xs sm:grid-cols-2">
|
||||
<div className="flex justify-between">
|
||||
<span className="text-muted-foreground">Base Model</span>
|
||||
<span className="font-medium">{baseModelName}</span>
|
||||
</div>
|
||||
<div className="flex justify-between">
|
||||
<span className="text-muted-foreground">Method</span>
|
||||
<span className="font-medium">
|
||||
{trainingMethodLabel}
|
||||
</span>
|
||||
</div>
|
||||
<div className="flex justify-between">
|
||||
<span className="text-muted-foreground">Checkpoints</span>
|
||||
<span className="font-medium">
|
||||
{checkpointsForModel.length}
|
||||
</span>
|
||||
</div>
|
||||
{isAdapter && (
|
||||
<div className="flex justify-between">
|
||||
<span className="text-muted-foreground">LoRA Rank</span>
|
||||
<span className="font-medium">{loraRank}</span>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="flex flex-col gap-2.5">
|
||||
|
|
@ -462,7 +916,7 @@ export function ExportPage() {
|
|||
Quick Guide
|
||||
</span>
|
||||
<ol className="flex flex-col gap-3">
|
||||
{GUIDE_STEPS.map((step, i) => (
|
||||
{exportGuideSteps.map((step, i) => (
|
||||
<li
|
||||
key={step}
|
||||
className="flex items-start gap-2 text-xs text-muted-foreground"
|
||||
|
|
@ -483,22 +937,24 @@ export function ExportPage() {
|
|||
disabledMethods={
|
||||
!isAdapter && isQuantized
|
||||
? ["merged", "lora", "gguf"]
|
||||
: !isAdapter
|
||||
: !isAdapter || sourceMode === "model"
|
||||
? ["merged", "lora"]
|
||||
: []
|
||||
}
|
||||
disabledReason={
|
||||
!isAdapter && isQuantized
|
||||
? "Pre-quantized (BNB 4-bit) models cannot be exported without LoRA adapters"
|
||||
: !isAdapter
|
||||
? "Not available for full fine-tune checkpoints (no LoRA adapters)"
|
||||
: undefined
|
||||
: sourceMode === "model"
|
||||
? "Only GGUF export is available for direct model export"
|
||||
: !isAdapter
|
||||
? "Not available for full fine-tune checkpoints (no LoRA adapters)"
|
||||
: undefined
|
||||
}
|
||||
/>
|
||||
|
||||
<AnimatePresence>
|
||||
{exportMethod === "gguf" && (
|
||||
<motion.div {...collapseAnim} className="overflow-hidden">
|
||||
<motion.div {...collapseAnim} className="overflow-visible">
|
||||
<QuantPicker value={quantLevels} onChange={setQuantLevels} />
|
||||
</motion.div>
|
||||
)}
|
||||
|
|
@ -530,12 +986,12 @@ export function ExportPage() {
|
|||
<ExportDialog
|
||||
open={dialogOpen}
|
||||
onOpenChange={setDialogOpen}
|
||||
checkpoint={checkpoint}
|
||||
exportMethod={exportMethod}
|
||||
quantLevels={quantLevels}
|
||||
estimatedSize={estimatedSize}
|
||||
baseModelName={baseModelName}
|
||||
isAdapter={isAdapter}
|
||||
checkpoint={selectedExportSource}
|
||||
baseModelName={sourceBaseModelName}
|
||||
isAdapter={sourceMode === "checkpoint" && isAdapter}
|
||||
destination={destination}
|
||||
onDestinationChange={setDestination}
|
||||
hfUsername={hfUsername}
|
||||
|
|
|
|||
|
|
@ -58,6 +58,13 @@ import { HugeiconsIcon } from "@hugeicons/react";
|
|||
import { useEffect, useMemo, useRef, useState } from "react";
|
||||
import { useShallow } from "zustand/react/shallow";
|
||||
|
||||
/** Extract param count label from model name (e.g. "Qwen3-0.6B" -> "0.6B"). */
|
||||
function extractParamLabel(id: string): string | null {
|
||||
const name = id.split("/").pop() ?? id;
|
||||
const match = name.match(/(?:^|[-_])(\d+(?:\.\d+)?)[Bb](?:[-_]|$)/);
|
||||
return match ? `${match[1]}B` : null;
|
||||
}
|
||||
|
||||
export function ModelSelectionStep() {
|
||||
const gpu = useGpuInfo();
|
||||
const {
|
||||
|
|
@ -85,6 +92,7 @@ export function ModelSelectionStep() {
|
|||
const [inputValue, setInputValue] = useState("");
|
||||
const selectingRef = useRef(false);
|
||||
const debouncedQuery = useDebouncedValue(inputValue);
|
||||
const debouncedHfToken = useDebouncedValue(hfToken, 500);
|
||||
const task = modelType ? MODEL_TYPE_TO_HF_TASK[modelType] : undefined;
|
||||
const {
|
||||
results: hfResults,
|
||||
|
|
@ -94,7 +102,7 @@ export function ModelSelectionStep() {
|
|||
error: hfSearchError,
|
||||
} = useHfModelSearch(debouncedQuery, {
|
||||
task,
|
||||
accessToken: hfToken || undefined,
|
||||
accessToken: debouncedHfToken || undefined,
|
||||
excludeGguf: true,
|
||||
priorityIds: PRIORITY_TRAINING_MODELS,
|
||||
});
|
||||
|
|
@ -119,7 +127,7 @@ export function ModelSelectionStep() {
|
|||
const fit = fitMap.get(r.id);
|
||||
map.set(r.id, {
|
||||
status: fit?.status ?? null,
|
||||
detail: r.totalParams ? formatCompact(r.totalParams) : null,
|
||||
detail: r.totalParams ? formatCompact(r.totalParams) : extractParamLabel(r.id),
|
||||
});
|
||||
}
|
||||
return map;
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import {
|
|||
import { Button } from "@/components/ui/button";
|
||||
import type { TrainingRunSummary } from "@/features/training";
|
||||
import { deleteTrainingRun, listTrainingRuns } from "@/features/training";
|
||||
import { formatDuration } from "@/features/studio/sections/progress-section-lib";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { Delete02Icon } from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
|
|
@ -129,16 +130,6 @@ function formatRelativeTime(isoDate: string): string {
|
|||
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;
|
||||
|
|
|
|||
|
|
@ -410,16 +410,20 @@ export function DatasetSection() {
|
|||
}, [navigate]);
|
||||
|
||||
return (
|
||||
<div data-tour="studio-dataset" className="col-span-1 xl:col-span-4">
|
||||
<div data-tour="studio-dataset" className="min-w-0">
|
||||
<SectionCard
|
||||
icon={<HugeiconsIcon icon={Database02Icon} className="size-5" />}
|
||||
title="Dataset"
|
||||
description="Select or upload training data"
|
||||
accent="indigo"
|
||||
className="dark:shadow-border"
|
||||
className={`dark:shadow-border ${
|
||||
advancedOpen || (datasetSource === "upload" && uploadedFile)
|
||||
? "min-h-studio-config-column"
|
||||
: "h-studio-config-column"
|
||||
}`}
|
||||
>
|
||||
<div className="flex flex-col gap-4">
|
||||
<div className="flex flex-col gap-2">
|
||||
<div className="flex min-w-0 flex-col gap-4">
|
||||
<div className="flex min-w-0 flex-col gap-2">
|
||||
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
|
||||
Choose dataset
|
||||
<span className="rounded-full border border-border/70 bg-muted/40 px-2 py-0.5 text-[10px] font-medium text-foreground/80">
|
||||
|
|
@ -453,6 +457,7 @@ export function DatasetSection() {
|
|||
</span>
|
||||
<div
|
||||
ref={comboboxAnchorRef}
|
||||
className="min-w-0"
|
||||
onKeyDown={(event) => {
|
||||
if (event.key !== "Enter") return;
|
||||
if (!(event.target instanceof HTMLInputElement)) return;
|
||||
|
|
@ -521,7 +526,7 @@ export function DatasetSection() {
|
|||
? "Search Hugging Face datasets..."
|
||||
: "Search local datasets..."
|
||||
}
|
||||
className="w-full"
|
||||
className="w-full min-w-0 overflow-hidden leading-5"
|
||||
showClear={true}
|
||||
>
|
||||
<InputGroupAddon>
|
||||
|
|
|
|||
|
|
@ -72,6 +72,13 @@ const DARK_CONTENT =
|
|||
const DARK_COMBOBOX_CONTENT =
|
||||
"bg-foreground text-background shadow-xl border-background/10 dark:[--accent:rgba(2,6,23,0.08)] dark:[--accent-foreground:rgb(2,6,23)] dark:[&_[data-slot=combobox-item]]:text-slate-900 dark:[&_.text-muted-foreground]:text-slate-500";
|
||||
|
||||
/** Extract param count label from model name (e.g. "Qwen3-0.6B" -> "0.6B"). */
|
||||
function extractParamLabel(id: string): string | null {
|
||||
const name = id.split("/").pop() ?? id;
|
||||
const match = name.match(/(?:^|[-_])(\d+(?:\.\d+)?)[Bb](?:[-_]|$)/);
|
||||
return match ? `${match[1]}B` : null;
|
||||
}
|
||||
|
||||
export function ModelSection() {
|
||||
const gpu = useGpuInfo();
|
||||
|
||||
|
|
@ -112,6 +119,7 @@ export function ModelSection() {
|
|||
const [localModelsError, setLocalModelsError] = useState<string | null>(null);
|
||||
const selectingRef = useRef(false);
|
||||
const debouncedQuery = useDebouncedValue(inputValue);
|
||||
const debouncedHfToken = useDebouncedValue(hfToken, 500);
|
||||
|
||||
function handleModelSelect(id: string | null) {
|
||||
selectingRef.current = true;
|
||||
|
|
@ -160,7 +168,7 @@ export function ModelSection() {
|
|||
error: hfSearchError,
|
||||
} = useHfModelSearch(debouncedQuery, {
|
||||
task,
|
||||
accessToken: hfToken || undefined,
|
||||
accessToken: debouncedHfToken || undefined,
|
||||
excludeGguf: true,
|
||||
priorityIds: PRIORITY_TRAINING_MODELS,
|
||||
});
|
||||
|
|
@ -181,6 +189,7 @@ export function ModelSection() {
|
|||
const trainableLocalModels = useMemo(
|
||||
() =>
|
||||
localModels.filter((m) => {
|
||||
if (m.source === "lmstudio") return false;
|
||||
if (m.path.endsWith(".gguf")) return false;
|
||||
if (m.id.toLowerCase().includes("-gguf")) return false;
|
||||
return true;
|
||||
|
|
@ -232,7 +241,7 @@ export function ModelSection() {
|
|||
{ est: number; status: VramFitStatus | null; detail: string | null }
|
||||
>();
|
||||
for (const r of hfResults) {
|
||||
const detail = r.totalParams ? formatCompact(r.totalParams) : null;
|
||||
const detail = r.totalParams ? formatCompact(r.totalParams) : extractParamLabel(r.id);
|
||||
const fit = fitMap.get(r.id);
|
||||
map.set(r.id, {
|
||||
est: fit?.est ?? 0,
|
||||
|
|
@ -251,7 +260,7 @@ export function ModelSection() {
|
|||
);
|
||||
|
||||
return (
|
||||
<div data-tour="studio-model" className="col-span-1 md:col-span-2 xl:col-span-12">
|
||||
<div data-tour="studio-model" className="w-full min-w-0">
|
||||
<SectionCard
|
||||
icon={<HugeiconsIcon icon={ChipIcon} className="size-5" />}
|
||||
title="Model"
|
||||
|
|
@ -259,10 +268,10 @@ export function ModelSection() {
|
|||
accent="emerald"
|
||||
featured={true}
|
||||
badge="2x Faster Training"
|
||||
className="shadow-border ring-1 ring-border"
|
||||
className="shadow-border ring-border"
|
||||
>
|
||||
<div className="grid gap-4 md:grid-cols-2 xl:grid-cols-4">
|
||||
<div data-tour="studio-local-model" className="flex flex-col gap-2">
|
||||
<div className="grid min-w-0 gap-4 md:grid-cols-2 xl:grid-cols-4">
|
||||
<div data-tour="studio-local-model" className="flex min-w-0 flex-col gap-2">
|
||||
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
|
||||
Local Model
|
||||
<Tooltip>
|
||||
|
|
@ -282,7 +291,7 @@ export function ModelSection() {
|
|||
</TooltipContent>
|
||||
</Tooltip>
|
||||
</span>
|
||||
<div ref={localComboboxAnchorRef}>
|
||||
<div ref={localComboboxAnchorRef} className="min-w-0">
|
||||
<Combobox
|
||||
items={localResultIds}
|
||||
filteredItems={localFilteredIds}
|
||||
|
|
@ -334,7 +343,11 @@ export function ModelSection() {
|
|||
{(id: string) => {
|
||||
const model = localMetaById.get(id);
|
||||
const source =
|
||||
model?.source === "hf_cache" ? "HF cache" : "Local dir";
|
||||
model?.source === "hf_cache"
|
||||
? "HF cache"
|
||||
: model?.source === "lmstudio"
|
||||
? "LM Studio"
|
||||
: "Local dir";
|
||||
return (
|
||||
<ComboboxItem key={id} value={id} className="gap-2">
|
||||
<Tooltip>
|
||||
|
|
@ -370,7 +383,7 @@ export function ModelSection() {
|
|||
)}
|
||||
</div>
|
||||
|
||||
<div data-tour="studio-base-model" className="flex flex-col gap-2">
|
||||
<div data-tour="studio-base-model" className="flex min-w-0 flex-col gap-2">
|
||||
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
|
||||
Hugging Face Model
|
||||
<Tooltip>
|
||||
|
|
@ -400,6 +413,7 @@ export function ModelSection() {
|
|||
</span>
|
||||
<div
|
||||
ref={comboboxAnchorRef}
|
||||
className="min-w-0"
|
||||
onKeyDown={(event) => {
|
||||
if (event.key !== "Enter") return;
|
||||
if (!(event.target instanceof HTMLInputElement)) return;
|
||||
|
|
@ -422,7 +436,10 @@ export function ModelSection() {
|
|||
itemToStringValue={(id) => id}
|
||||
autoHighlight={true}
|
||||
>
|
||||
<ComboboxInput placeholder="Search models..." className="w-full">
|
||||
<ComboboxInput
|
||||
placeholder="Search models..."
|
||||
className="w-full leading-5"
|
||||
>
|
||||
<InputGroupAddon>
|
||||
<HugeiconsIcon icon={Search01Icon} className="size-4" />
|
||||
</InputGroupAddon>
|
||||
|
|
@ -508,7 +525,7 @@ export function ModelSection() {
|
|||
</div>
|
||||
</div>
|
||||
|
||||
<div data-tour="studio-method" className="flex flex-col gap-2">
|
||||
<div data-tour="studio-method" className="flex min-w-0 flex-col gap-2">
|
||||
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
|
||||
Method
|
||||
<Tooltip>
|
||||
|
|
@ -576,7 +593,7 @@ export function ModelSection() {
|
|||
</Select>
|
||||
</div>
|
||||
|
||||
<div className="flex flex-col gap-2">
|
||||
<div className="flex min-w-0 flex-col gap-2">
|
||||
<span className="text-xs font-medium text-muted-foreground">
|
||||
Hugging Face Token (Optional)
|
||||
</span>
|
||||
|
|
|
|||
|
|
@ -160,13 +160,15 @@ export function ParamsSection(): ReactElement {
|
|||
const epochsSliderMax = Math.max(20, store.epochs, 1);
|
||||
|
||||
return (
|
||||
<div data-tour="studio-params" className="col-span-1 xl:col-span-4">
|
||||
<div data-tour="studio-params" className="min-w-0">
|
||||
<SectionCard
|
||||
icon={<HugeiconsIcon icon={Settings04Icon} className="size-5" />}
|
||||
title="Parameters"
|
||||
description="Configure training hyperparameters"
|
||||
accent="orange"
|
||||
className="md:min-h-[470px]"
|
||||
className={`${(isLora && loraOpen) || hyperOpen
|
||||
? "min-h-studio-config-column"
|
||||
: "h-studio-config-column"} duration-150`}
|
||||
>
|
||||
<div className="flex flex-col gap-4">
|
||||
{/* Max Steps / Epochs */}
|
||||
|
|
@ -380,21 +382,16 @@ export function ParamsSection(): ReactElement {
|
|||
|
||||
{/* LoRA Settings */}
|
||||
{isLora && (
|
||||
<div>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setLoraOpen(!loraOpen)}
|
||||
className="flex w-full cursor-pointer items-center gap-1.5 text-xs text-muted-foreground"
|
||||
>
|
||||
<Collapsible open={loraOpen} onOpenChange={setLoraOpen}>
|
||||
<CollapsibleTrigger className="flex w-full cursor-pointer items-center gap-1.5 text-xs text-muted-foreground">
|
||||
<HugeiconsIcon
|
||||
icon={ArrowDown01Icon}
|
||||
className={`size-3.5 transition-transform ${loraOpen ? "rotate-180" : ""}`}
|
||||
/>
|
||||
LoRA Settings
|
||||
</button>
|
||||
<div
|
||||
className={`${loraOpen ? "" : "hidden"} pt-1.5 mt-4 flex flex-col gap-4`}
|
||||
>
|
||||
</CollapsibleTrigger>
|
||||
<CollapsibleContent className="mt-3 data-[state=open]:overflow-visible">
|
||||
<div className="pt-1.5 flex flex-col gap-4">
|
||||
<SliderRow
|
||||
label="Rank"
|
||||
tooltip={
|
||||
|
|
@ -576,8 +573,9 @@ export function ParamsSection(): ReactElement {
|
|||
</button>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</CollapsibleContent>
|
||||
</Collapsible>
|
||||
)}
|
||||
|
||||
{/* Training Hyperparams */}
|
||||
|
|
|
|||
|
|
@ -35,12 +35,17 @@ export const phaseColors: Record<TrainingPhase, string> = {
|
|||
stopped: "bg-muted text-muted-foreground",
|
||||
};
|
||||
|
||||
export function formatDuration(seconds: number | null): string {
|
||||
if (seconds == null || seconds < 0) return "--";
|
||||
export function formatDuration(seconds: number | null | undefined): string {
|
||||
if (seconds == null || !Number.isFinite(seconds) || seconds < 0) return "--";
|
||||
const total = Math.floor(seconds);
|
||||
const min = Math.floor(total / 60);
|
||||
const sec = total % 60;
|
||||
return `${min}m ${sec}s`;
|
||||
const d = Math.floor(total / 86400);
|
||||
const h = Math.floor((total % 86400) / 3600);
|
||||
const m = Math.floor((total % 3600) / 60);
|
||||
const s = total % 60;
|
||||
if (d > 0) return `${d}d ${h}h ${m}m`;
|
||||
if (h > 0) return `${h}h ${m}m ${s}s`;
|
||||
if (m > 0) return `${m}m ${s}s`;
|
||||
return `${s}s`;
|
||||
}
|
||||
|
||||
export function formatNumber(value: number | null | undefined, digits: number): string {
|
||||
|
|
|
|||
|
|
@ -45,10 +45,14 @@ const placeholderData = [
|
|||
export function TrainingSection() {
|
||||
const store = useTrainingConfigStore();
|
||||
const { isStarting, startError, startTrainingRun } = useTrainingActions();
|
||||
const isLoadingModel = store.isLoadingModelDefaults || store.isCheckingVision;
|
||||
const isModelCapabilitiesSettled = !!store.selectedModel && !isLoadingModel;
|
||||
const isIncompatible =
|
||||
(!store.isVisionModel && store.isDatasetImage === true) ||
|
||||
(!store.isAudioModel && store.isDatasetAudio === true);
|
||||
isModelCapabilitiesSettled &&
|
||||
((!store.isVisionModel && store.isDatasetImage === true) ||
|
||||
(!store.isAudioModel && store.isDatasetAudio === true));
|
||||
const configValidation = validateTrainingConfig(store);
|
||||
const hasMessage = !!(startError || isIncompatible || (!configValidation.ok && configValidation.message));
|
||||
const fileInputRef = useRef<HTMLInputElement>(null);
|
||||
|
||||
const handleFileUpload = (e: React.ChangeEvent<HTMLInputElement>) => {
|
||||
|
|
@ -98,13 +102,13 @@ export function TrainingSection() {
|
|||
};
|
||||
|
||||
return (
|
||||
<div data-tour="studio-training" className="col-span-1 xl:col-span-4">
|
||||
<div data-tour="studio-training" className="min-w-0">
|
||||
<SectionCard
|
||||
icon={<HugeiconsIcon icon={ChartAverageIcon} className="size-5" />}
|
||||
title="Training"
|
||||
description="Monitor and control training"
|
||||
accent="blue"
|
||||
className="md:min-h-[470px]"
|
||||
className={hasMessage ? "min-h-studio-config-column" : "h-studio-config-column"}
|
||||
>
|
||||
<div className="flex flex-col gap-4">
|
||||
{/* Loss chart */}
|
||||
|
|
@ -156,17 +160,19 @@ export function TrainingSection() {
|
|||
data-tour="studio-start"
|
||||
className="w-full cursor-pointer bg-gradient-to-r from-emerald-500 to-teal-500 text-white hover:from-emerald-600 hover:to-teal-600"
|
||||
onClick={() => void startTrainingRun()}
|
||||
disabled={isStarting || isIncompatible || store.isCheckingDataset || !configValidation.ok}
|
||||
disabled={isStarting || isIncompatible || store.isCheckingDataset || isLoadingModel || !configValidation.ok}
|
||||
>
|
||||
<HugeiconsIcon icon={Rocket01Icon} className="size-4" />
|
||||
{isStarting ? "Starting..." : store.isCheckingDataset ? "Checking dataset..." : "Start Training"}
|
||||
{isStarting ? "Starting..." : isLoadingModel ? "Loading model..." : store.isCheckingDataset ? "Checking dataset..." : "Start Training"}
|
||||
</Button>
|
||||
{startError && (
|
||||
<p className="text-xs text-red-500 leading-relaxed">{startError}</p>
|
||||
)}
|
||||
{isIncompatible && (
|
||||
<p className="text-xs text-red-500 leading-relaxed">
|
||||
Text model is not compatible with a multimodal dataset. Switch to a vision model or choose a text-only dataset.
|
||||
{!store.isAudioModel && store.isDatasetAudio === true
|
||||
? "This model does not support audio. Switch to an audio-capable model or choose a non-audio dataset."
|
||||
: "Text model is not compatible with a multimodal dataset. Switch to a vision model or choose a text-only dataset."}
|
||||
</p>
|
||||
)}
|
||||
{!configValidation.ok && configValidation.message && !isIncompatible && (
|
||||
|
|
|
|||
|
|
@ -164,11 +164,13 @@ export function StudioPage(): ReactElement {
|
|||
</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">
|
||||
<div className="flex min-w-0 flex-col gap-4 md:gap-6">
|
||||
<ModelSection />
|
||||
<DatasetSection />
|
||||
<ParamsSection />
|
||||
<TrainingSection />
|
||||
<div className="grid min-w-0 grid-cols-1 items-start gap-4 md:grid-cols-2 md:gap-6 xl:grid-cols-3 xl:gap-6">
|
||||
<DatasetSection />
|
||||
<ParamsSection />
|
||||
<TrainingSection />
|
||||
</div>
|
||||
</div>
|
||||
</TabsContent>
|
||||
|
||||
|
|
|
|||
|
|
@ -1,10 +1,15 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
const EXTERNAL_URL_RE = /^https?:\/\//;
|
||||
|
||||
export function ReadMore({ href = "#" }: { href?: string }) {
|
||||
const isExternal = EXTERNAL_URL_RE.test(href);
|
||||
return (
|
||||
<a
|
||||
href={href}
|
||||
target={isExternal ? "_blank" : undefined}
|
||||
rel={isExternal ? "noopener noreferrer" : undefined}
|
||||
onClick={(e) => {
|
||||
if (href === "#") e.preventDefault();
|
||||
}}
|
||||
|
|
|
|||
|
|
@ -79,7 +79,7 @@ export interface LocalModelInfo {
|
|||
id: string;
|
||||
display_name: string;
|
||||
path: string;
|
||||
source: "models_dir" | "hf_cache";
|
||||
source: "models_dir" | "hf_cache" | "lmstudio";
|
||||
model_id?: string | null;
|
||||
updated_at?: number | null;
|
||||
}
|
||||
|
|
@ -87,6 +87,7 @@ export interface LocalModelInfo {
|
|||
interface LocalModelListResponse {
|
||||
models_dir: string;
|
||||
hf_cache_dir?: string | null;
|
||||
lmstudio_dirs: string[];
|
||||
models: LocalModelInfo[];
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -91,7 +91,7 @@ export function HfDatasetSubsetSplitSelectors({
|
|||
<>
|
||||
{showPlaceholderDropdowns && (
|
||||
<>
|
||||
<div className="grid gap-3 sm:grid-cols-2">
|
||||
<div className="grid min-w-0 gap-3 sm:grid-cols-2">
|
||||
<SelectorDropdown
|
||||
variant={variant}
|
||||
label="Subset"
|
||||
|
|
@ -132,7 +132,7 @@ export function HfDatasetSubsetSplitSelectors({
|
|||
className={
|
||||
variant === "wizard"
|
||||
? "flex items-center gap-2 text-xs text-muted-foreground py-1"
|
||||
: "flex items-center gap-2 rounded-lg border bg-muted/20 px-3.5 py-3 text-xs text-muted-foreground"
|
||||
: "flex min-w-0 items-center gap-2 rounded-lg border bg-muted/20 px-3.5 py-3 text-xs text-muted-foreground"
|
||||
}
|
||||
>
|
||||
<Spinner className="size-3.5" />
|
||||
|
|
@ -145,7 +145,7 @@ export function HfDatasetSubsetSplitSelectors({
|
|||
className={
|
||||
variant === "wizard"
|
||||
? "rounded-lg border border-amber-200 bg-amber-50 px-3 py-2 text-xs text-amber-700 dark:border-amber-800 dark:bg-amber-950 dark:text-amber-400"
|
||||
: "rounded-lg border border-amber-200 bg-amber-50 px-3.5 py-2.5 text-xs text-amber-700 dark:border-amber-800 dark:bg-amber-950 dark:text-amber-400"
|
||||
: "min-w-0 rounded-lg border border-amber-200 bg-amber-50 px-3.5 py-2.5 text-xs text-amber-700 dark:border-amber-800 dark:bg-amber-950 dark:text-amber-400"
|
||||
}
|
||||
>
|
||||
{error}
|
||||
|
|
@ -155,7 +155,7 @@ export function HfDatasetSubsetSplitSelectors({
|
|||
{showDropdowns && (
|
||||
<>
|
||||
{variant === "studio" ? (
|
||||
<div className="grid gap-3 sm:grid-cols-2">
|
||||
<div className="grid min-w-0 gap-3 sm:grid-cols-2">
|
||||
<SelectorDropdown
|
||||
variant={variant}
|
||||
label="Subset"
|
||||
|
|
@ -283,7 +283,7 @@ function SelectorDropdown({
|
|||
}
|
||||
|
||||
return (
|
||||
<div className="flex flex-col gap-1.5">
|
||||
<div className="flex min-w-0 flex-col gap-1.5">
|
||||
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
|
||||
{label}
|
||||
<Tooltip>
|
||||
|
|
@ -308,7 +308,7 @@ function SelectorDropdown({
|
|||
onValueChange={(v) => onChange(v === "_none" ? null : v)}
|
||||
disabled={disabled}
|
||||
>
|
||||
<SelectTrigger className="w-full">
|
||||
<SelectTrigger className="w-full min-w-0">
|
||||
<SelectValue placeholder={placeholder} />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
// 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 { DEFAULT_HYPERPARAMS, STEPS } from "@/config/training";
|
||||
import { DEFAULT_HYPERPARAMS, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS } from "@/config/training";
|
||||
import { authFetch } from "@/features/auth";
|
||||
import { isAdapterMethod } from "@/types/training";
|
||||
import type { ModelType, StepNumber, TrainingMethod } from "@/types/training";
|
||||
import { create } from "zustand";
|
||||
import { persist } from "zustand/middleware";
|
||||
|
|
@ -98,6 +99,15 @@ let _modelConfigController: AbortController | null = null;
|
|||
// since the last auto-set (model load or dataset change).
|
||||
let _trainOnCompletionsManuallySet = false;
|
||||
|
||||
// Track whether the user has manually edited the learning rate
|
||||
// since the last model load. When false, switching training method
|
||||
// auto-sets LR to 2e-4 (LoRA/QLoRA) or 2e-5 (full fine-tune).
|
||||
let _learningRateManuallySet = false;
|
||||
|
||||
// Stash the model-config-provided (YAML) learning rate so that
|
||||
// setTrainingMethod can restore it when switching back from full to adapter.
|
||||
let _yamlLearningRate: number | undefined = undefined;
|
||||
|
||||
const NON_PERSISTED_STATE_KEYS: ReadonlySet<keyof TrainingConfigState> = new Set([
|
||||
"modelType",
|
||||
"isCheckingVision",
|
||||
|
|
@ -165,8 +175,22 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
if (get().selectedModel !== modelName) return;
|
||||
|
||||
_trainOnCompletionsManuallySet = false;
|
||||
_learningRateManuallySet = false;
|
||||
_yamlLearningRate = undefined;
|
||||
const patch = mapBackendModelConfigToTrainingPatch(modelDetails.config);
|
||||
|
||||
// If the model config provides a specific learning rate, treat
|
||||
// it as authoritative so the async auto-select does not overwrite it.
|
||||
const modelConfigHasLR = patch.learningRate !== undefined;
|
||||
_yamlLearningRate = patch.learningRate;
|
||||
|
||||
// YAML learning rates are tuned for adapter methods (LoRA/QLoRA).
|
||||
// If the user is currently on full fine-tune, override with the
|
||||
// full-finetune default instead of applying the YAML adapter LR.
|
||||
if (modelConfigHasLR && !isAdapterMethod(get().trainingMethod)) {
|
||||
patch.learningRate = LR_DEFAULT_FULL;
|
||||
}
|
||||
|
||||
// If vision model + image dataset already known, override
|
||||
// trainOnCompletions to false regardless of backend default.
|
||||
if (modelDetails.is_vision && get().isDatasetImage === true) {
|
||||
|
|
@ -174,11 +198,11 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
}
|
||||
|
||||
const isAudio = !!modelDetails.is_audio;
|
||||
// Pure audio model → always uncheck trainOnCompletions.
|
||||
// Pure audio model -> always uncheck trainOnCompletions.
|
||||
if (isAudio && !modelDetails.is_vision) {
|
||||
patch.trainOnCompletions = false;
|
||||
}
|
||||
// Audio-capable vision model (e.g. gemma3n) + audio dataset → uncheck.
|
||||
// Audio-capable vision model (e.g. gemma3n) + audio dataset -> uncheck.
|
||||
if (isAudio && modelDetails.is_vision && get().isDatasetAudio) {
|
||||
patch.trainOnCompletions = false;
|
||||
}
|
||||
|
|
@ -197,7 +221,12 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
void autoSelectTrainingMethod(modelSizeBytes, patch.contextLength ?? get().contextLength)
|
||||
.then((method) => {
|
||||
if (get().selectedModel !== modelName) return;
|
||||
if (method) set({ trainingMethod: method });
|
||||
if (method) {
|
||||
const lrPatch = !_learningRateManuallySet && !modelConfigHasLR
|
||||
? { learningRate: method === "full" ? LR_DEFAULT_FULL : LR_DEFAULT_LORA }
|
||||
: {};
|
||||
set({ trainingMethod: method, ...lrPatch });
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -366,7 +395,31 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
if (state.modelDefaultsAppliedFor === state.selectedModel) return;
|
||||
void loadAndApplyModelDefaults(state.selectedModel);
|
||||
},
|
||||
setTrainingMethod: (trainingMethod) => set({ trainingMethod }),
|
||||
setTrainingMethod: (trainingMethod) => {
|
||||
if (_learningRateManuallySet) {
|
||||
set({ trainingMethod });
|
||||
return;
|
||||
}
|
||||
|
||||
const prev = get().trainingMethod;
|
||||
const wasAdapter = isAdapterMethod(prev);
|
||||
const nowAdapter = isAdapterMethod(trainingMethod);
|
||||
|
||||
// qlora <-> lora: same LR range, don't touch learning rate
|
||||
if (wasAdapter && nowAdapter) {
|
||||
set({ trainingMethod });
|
||||
return;
|
||||
}
|
||||
|
||||
// Category changed (adapter <-> full)
|
||||
if (nowAdapter) {
|
||||
// Switching TO adapter: restore YAML LR if available
|
||||
set({ trainingMethod, learningRate: _yamlLearningRate ?? LR_DEFAULT_LORA });
|
||||
} else {
|
||||
// Switching TO full: no YAML full-LR exists, use constant
|
||||
set({ trainingMethod, learningRate: LR_DEFAULT_FULL });
|
||||
}
|
||||
},
|
||||
setHfToken: (hfToken) =>
|
||||
set({ hfToken: hfToken.trim().replace(/^["']+|["']+$/g, "") }),
|
||||
setDatasetSource: (datasetSource) => set({ datasetSource }),
|
||||
|
|
@ -509,7 +562,10 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
}),
|
||||
setEpochs: (epochs) => set({ epochs }),
|
||||
setContextLength: (contextLength) => set({ contextLength }),
|
||||
setLearningRate: (learningRate) => set({ learningRate }),
|
||||
setLearningRate: (learningRate) => {
|
||||
_learningRateManuallySet = true;
|
||||
set({ learningRate });
|
||||
},
|
||||
setOptimizerType: (optimizerType) => set({ optimizerType }),
|
||||
setLrSchedulerType: (lrSchedulerType) => set({ lrSchedulerType }),
|
||||
setLoraRank: (loraRank) => set({ loraRank }),
|
||||
|
|
@ -548,7 +604,12 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
set({ finetuneMLPModules }),
|
||||
setTargetModules: (targetModules) => set({ targetModules }),
|
||||
canProceed: () => canProceedForStep(get()),
|
||||
reset: () => set(initialState),
|
||||
reset: () => {
|
||||
_trainOnCompletionsManuallySet = false;
|
||||
_learningRateManuallySet = false;
|
||||
_yamlLearningRate = undefined;
|
||||
set(initialState);
|
||||
},
|
||||
resetToModelDefaults: () => {
|
||||
const { selectedModel } = get();
|
||||
if (!selectedModel) return;
|
||||
|
|
@ -557,13 +618,18 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
},
|
||||
applyConfigPatch: (config: BackendModelConfig) => {
|
||||
const patch = mapBackendModelConfigToTrainingPatch(config);
|
||||
// Only clear the manual-edit flag when the config provides a LR,
|
||||
// so unrelated config patches don't silently disarm the guard.
|
||||
if (patch.learningRate !== undefined) {
|
||||
_learningRateManuallySet = false;
|
||||
}
|
||||
set(patch);
|
||||
},
|
||||
};
|
||||
},
|
||||
{
|
||||
name: "unsloth_training_config_v1",
|
||||
version: 8,
|
||||
version: 9,
|
||||
migrate: (persisted, version) => {
|
||||
const s = persisted as Record<string, unknown>;
|
||||
if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
|
||||
|
|
@ -593,6 +659,12 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
s.datasetLabelMapping ??= {};
|
||||
s.datasetAdvisorNotification ??= null;
|
||||
}
|
||||
if (version < 9) {
|
||||
// weight_decay default changed from 0.01 to 0.001.
|
||||
if (s.weightDecay === 0.01) {
|
||||
s.weightDecay = DEFAULT_HYPERPARAMS.weightDecay;
|
||||
}
|
||||
}
|
||||
return s as unknown as TrainingConfigStore;
|
||||
},
|
||||
partialize: partializePersistedState,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,8 @@
|
|||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import type { PipelineType } from "@huggingface/hub";
|
||||
import { listModels, modelInfo } from "@huggingface/hub";
|
||||
import { listModels } from "@huggingface/hub";
|
||||
import { type CachedResult, cachedModelInfo, primeCacheFromListing } from "@/lib/hf-cache";
|
||||
import { useCallback, useMemo } from "react";
|
||||
import { useHfPaginatedSearch } from "./use-hf-paginated-search";
|
||||
|
||||
|
|
@ -104,6 +105,24 @@ function makeMapModel(excludeGguf: boolean) {
|
|||
/** Number of unsloth results to pull up-front before yielding general results. */
|
||||
const UNSLOTH_PREFETCH = 20;
|
||||
|
||||
/**
|
||||
* Prime the hf-cache from a listModels result. For public (non-gated,
|
||||
* non-private) models, also prime the anonymous slot so the VRAM hook
|
||||
* gets cache hits without re-fetching. Gated/private models are only
|
||||
* cached under the caller's token to avoid auth leakage.
|
||||
*/
|
||||
function primeFromListing(
|
||||
name: string,
|
||||
accessToken: string | undefined,
|
||||
model: unknown,
|
||||
): void {
|
||||
const data = model as CachedResult;
|
||||
primeCacheFromListing(name, accessToken, data);
|
||||
if (accessToken && !data.private && !data.gated) {
|
||||
primeCacheFromListing(name, undefined, data);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a merged async generator that yields unsloth-owned models first,
|
||||
* then general results (with deduplication).
|
||||
|
|
@ -134,7 +153,10 @@ async function* mergedModelIterator(
|
|||
let count = 0;
|
||||
for await (const model of unslothIter) {
|
||||
const m = model as { name?: string };
|
||||
if (m.name) seen.add(m.name);
|
||||
if (m.name) {
|
||||
seen.add(m.name);
|
||||
primeFromListing(m.name, accessToken, model);
|
||||
}
|
||||
yield model;
|
||||
count++;
|
||||
if (count >= UNSLOTH_PREFETCH) break;
|
||||
|
|
@ -144,6 +166,9 @@ async function* mergedModelIterator(
|
|||
for await (const model of generalIter) {
|
||||
const m = model as { name?: string };
|
||||
if (m.name && seen.has(m.name)) continue;
|
||||
if (m.name) {
|
||||
primeFromListing(m.name, accessToken, model);
|
||||
}
|
||||
yield model;
|
||||
}
|
||||
}
|
||||
|
|
@ -167,7 +192,7 @@ async function* priorityThenListingIterator(
|
|||
const seen = new Set<string>();
|
||||
const settled = await Promise.allSettled(
|
||||
priorityIds.map((id) =>
|
||||
modelInfo({
|
||||
cachedModelInfo({
|
||||
name: id,
|
||||
additionalFields: ["safetensors", "tags"],
|
||||
...(accessToken ? { credentials: { accessToken } } : {}),
|
||||
|
|
@ -192,6 +217,9 @@ async function* priorityThenListingIterator(
|
|||
for await (const model of generalIter) {
|
||||
const m = model as { name?: string };
|
||||
if (m.name && seen.has(m.name)) continue;
|
||||
if (m.name) {
|
||||
primeFromListing(m.name, accessToken, model);
|
||||
}
|
||||
yield model;
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 { modelInfo } from "@huggingface/hub";
|
||||
import { cachedModelInfo } from "@/lib/hf-cache";
|
||||
import { useEffect, useState } from "react";
|
||||
|
||||
/**
|
||||
|
|
@ -10,9 +10,9 @@ import { useEffect, useState } from "react";
|
|||
* models in the chat model dropdown.
|
||||
*/
|
||||
export function useRecommendedModelVram(ids: string[]) {
|
||||
const [paramCountById, setParamCountById] = useState<
|
||||
Map<string, number>
|
||||
>(new Map());
|
||||
const [paramCountById, setParamCountById] = useState<Map<string, number>>(
|
||||
new Map(),
|
||||
);
|
||||
const [isLoading, setIsLoading] = useState(false);
|
||||
|
||||
const stableKey = [...ids].filter(Boolean).sort().join(",");
|
||||
|
|
@ -30,14 +30,15 @@ export function useRecommendedModelVram(ids: string[]) {
|
|||
const next = new Map<string, number>();
|
||||
await Promise.all(
|
||||
stableIds.map(async (id) => {
|
||||
if (canceled) return;
|
||||
if (canceled) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const info = await modelInfo({
|
||||
const info = await cachedModelInfo({
|
||||
name: id,
|
||||
additionalFields: ["safetensors"],
|
||||
});
|
||||
const raw = info as { safetensors?: { total?: number } };
|
||||
const total = raw.safetensors?.total;
|
||||
const total = info.safetensors?.total;
|
||||
if (typeof total === "number" && total > 0) {
|
||||
next.set(id, total);
|
||||
}
|
||||
|
|
@ -47,7 +48,9 @@ export function useRecommendedModelVram(ids: string[]) {
|
|||
}),
|
||||
);
|
||||
if (!canceled) {
|
||||
setParamCountById(next);
|
||||
// Merge with previous state so that VRAM badges for already-visible
|
||||
// models are preserved while newly-visible models are still loading.
|
||||
setParamCountById((prev) => new Map([...prev, ...next]));
|
||||
setIsLoading(false);
|
||||
}
|
||||
})();
|
||||
|
|
|
|||
|
|
@ -307,6 +307,14 @@
|
|||
--tw-shadow: 0 4px 16px rgba(0, 0, 0, 0.3);
|
||||
}
|
||||
|
||||
/* Fine-tuning Studio: equal default height, expandable when needed (md+) */
|
||||
.min-h-studio-config-column {
|
||||
@apply md:min-h-[470px];
|
||||
}
|
||||
.h-studio-config-column {
|
||||
@apply md:h-[470px];
|
||||
}
|
||||
|
||||
[data-streamdown="unordered-list"] {
|
||||
list-style-type: disc;
|
||||
list-style-position: outside;
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue