ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'nightly' into feature/canvas-lab

Browse source
This commit is contained in:
Wasim Yousef Said 2026-02-20 01:23:44 -08:00 • committed by GitHub
commit a5ad6ff12f
210 changed files with 9634 additions and 3788 deletions
Download patch file Download diff file Expand all files Collapse all files

8
.gitignore vendored
View file

@ -12,12 +12,15 @@ __pycache__/
.venv/
venv/
env/
environment.yaml
# Unsloth cache
unsloth_compiled_cache/
# ML artifacts (large files)
outputs/
exports/
unsloth_training_checkpoints/
*.gguf
*.safetensors
@ -28,6 +31,9 @@ outputs/
*.swp
*.swo
# oh-my-codex
.omx/
# OS
.DS_Store
Thumbs.db
@ -47,3 +53,5 @@ docs/canvas-lab-architecture.md
studio/frontend/test/
studio/tests/
studio/backend/tests/
log_rtx.txt
log.txt

4
README.md
View file

@ -45,12 +45,14 @@ This script will:
### Launch the studio
```bash
# After setup, open a new terminal (or source ~/.bashrc), then:
# After setup, open a new terminal (or source ~/.bashrc), then inside your working directory:
unsloth-ui -H 0.0.0.0 -p 8000
```
On **first launch**, a one-time setup token is printed to the console. Use it in the browser to create your admin account.
As this repo is in continuous development, please make sure to run the setup.sh file everytime you pull new changes from the repo.
## API Reference
All endpoints require a valid JWT `Authorization: Bearer <token>` header (except `/api/auth/*` and `/api/health`).

99
Unsloth_Studio_Colab.ipynb Normal file
View file

@ -0,0 +1,99 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "447c1156",
"metadata": {},
"outputs": [],
"source": [
"# ===========================================\n",
"# ⚠️ GPU Check - Run This First!\n",
"# ===========================================\n",
"import torch\n",
"\n",
"print(\"🔍 Checking for GPU...\")\n",
"if not torch.cuda.is_available():\n",
" print(\"❌ ERROR: No GPU detected!\")\n",
" print(\"\\n📋 To enable GPU:\")\n",
" print(\" 1. Go to: Runtime → Change runtime type\")\n",
" print(\" 2. Select: Hardware accelerator → GPU (T4 is free)\")\n",
" print(\" 3. Click: Save\")\n",
" print(\" 4. Restart and re-run all cells\")\n",
" raise RuntimeError(\"⛔ GPU required for Unsloth Studio\")\n",
"else:\n",
" gpu_name = torch.cuda.get_device_name(0)\n",
" print(f\"✅ GPU detected: {gpu_name}\")\n",
" print(\" Ready to proceed!\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f04a9b46",
"metadata": {},
"outputs": [],
"source": [
"# ===========================================\n",
"# GitHub Authentication (Private Repo)\n",
"# ===========================================\n",
"from getpass import getpass\n",
"import os\n",
"\n",
"print(\"🔐 GitHub Token Required\")\n",
"print(\"Get token: https://github.com/settings/tokens\")\n",
"print(\"Scope needed: 'repo'\")\n",
"print(\"-\" * 50)\n",
"\n",
"github_token = getpass(\"Enter GitHub Token: \")\n",
"os.environ['GITHUB_TOKEN'] = github_token\n",
"print(\"✅ Token stored\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "27e68f91",
"metadata": {},
"outputs": [],
"source": [
"# ===========================================\n",
"# Setup: Clone repo and run setup\n",
"# ===========================================\n",
"\n",
"import os\n",
"github_token = os.environ['GITHUB_TOKEN']\n",
"!git clone -b feature/colab-notebook https://{github_token}@github.com/unslothai/new-ui-prototype.git\n",
"%cd /content/new-ui-prototype\n",
"\n",
"# Run setup script\n",
"!chmod +x setup.sh\n",
"!./setup.sh"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "277e431e",
"metadata": {},
"outputs": [],
"source": [
"# ===========================================\n",
"# Start Unsloth Studio\n",
"# ===========================================\n",
"import sys\n",
"sys.path.insert(0, '/content/new-ui-prototype/studio/backend')\n",
"\n",
"from colab import start\n",
"start()"
]
}
],
"metadata": {
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

33
cli/commands/studio.py
View file

@ -1,19 +1,34 @@
import time
from pathlib import Path
from typing import Optional
import typer
def studio(
port: int = typer.Option(8000, "--port", "-p", help="Port to run the UI server on."),
host: str = typer.Option("0.0.0.0", "--host", "-H", help="Host address to bind to."),
share: bool = typer.Option(True, "--share", "-s", help="Create a public Gradio share link."),
frontend: Optional[Path] = typer.Option(None, "--frontend", "-f", help="Path to frontend build directory."),
silent: bool = typer.Option(False, "--silent", "-q", help="Suppress startup messages."),
):
"""Launch the Unsloth web UI for training, inference, and export."""
from app import demo, script_dir
"""Launch the Unsloth web UI backend server."""
from studio.backend.run import run_server
typer.echo(f"Starting Unsloth UI on http://{host}:{port}")
if not silent:
from studio.backend.run import _resolve_external_ip
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}")
demo.launch(
share=share,
server_port=port,
server_name=host,
favicon_path=f"{script_dir}/assets/favicon-32x32.png",
run_server(
host=host,
port=port,
frontend_path=frontend,
silent=silent,
)
# Keep running until interrupted
try:
while True:
time.sleep(1)
except KeyboardInterrupt:
typer.echo("\nShutting down...")

204
setup.sh
View file

@ -24,6 +24,13 @@ echo "╔═══════════════════════
echo "║ Unsloth Studio Setup Script ║"
echo "╚══════════════════════════════════════╝"
# ── Detect Colab (like unsloth does) ──
IS_COLAB=false
keynames=$'\n'$(printenv | cut -d= -f1)
if [[ "$keynames" == *$'\nCOLAB_'* ]]; then
IS_COLAB=true
fi
# ── 1. Check existing Node/npm versions ──
NEED_NODE=true
if command -v node &>/dev/null && command -v npm &>/dev/null; then
@ -33,7 +40,17 @@ if command -v node &>/dev/null && command -v npm &>/dev/null; then
echo "✅ Node $(node -v) and npm $(npm -v) already meet requirements. Skipping nvm install."
NEED_NODE=false
else
echo "⚠️ Node $(node -v) / npm $(npm -v) too old. Installing via nvm..."
if [ "$IS_COLAB" = true ]; then
echo "✅ Node $(node -v) and npm $(npm -v) detected in Colab."
# In Colab, just upgrade npm directly - nvm doesn't work well
if [ "$NPM_MAJOR" -lt 11 ]; then
echo " Upgrading npm to latest..."
npm install -g npm@latest > /dev/null 2>&1
fi
NEED_NODE=false
else
echo "⚠️ Node $(node -v) / npm $(npm -v) too old. Installing via nvm..."
fi
fi
else
echo "⚠️ Node/npm not found. Installing via nvm..."
@ -81,38 +98,179 @@ echo "✅ Frontend built to studio/frontend/dist"
# ── 6. Python venv + deps ──
echo ""
echo "Setting up Python environment..."
python3 -m venv .venv
source .venv/bin/activate
run_quiet "pip upgrade" pip install --upgrade pip
echo " Installing unsloth-zoo + unsloth..."
run_quiet "pip install unsloth" pip install unsloth-zoo unsloth
echo " Installing studio dependencies..."
run_quiet "pip install extras" pip install typer fastapi uvicorn pydantic matplotlib pandas nest_asyncio "datasets==4.3.0" pyjwt easydict addict
echo "✅ Python dependencies installed"
# ── 7. Add shell alias ──
# ── 6a. Discover best Python <= 3.12.x ──
BEST_PY=""
BEST_MAJOR=0
BEST_MINOR=0
# Collect candidate python3 binaries (python3, python3.9, python3.10, …)
for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)?$' | sort -u); do
if ! command -v "$candidate" &>/dev/null; then
continue
fi
# Get version string, e.g. "Python 3.11.5"
ver_str=$("$candidate" --version 2>&1 | awk '{print $2}')
py_major=$(echo "$ver_str" | cut -d. -f1)
py_minor=$(echo "$ver_str" | cut -d. -f2)
# Skip anything that isn't Python 3
if [ "$py_major" -ne 3 ] 2>/dev/null; then
continue
fi
# Skip versions above 3.12
if [ "$py_minor" -gt 12 ] 2>/dev/null; then
continue
fi
# Keep the highest qualifying version
if [ "$py_minor" -gt "$BEST_MINOR" ]; then
BEST_PY="$candidate"
BEST_MAJOR="$py_major"
BEST_MINOR="$py_minor"
fi
done
if [ -z "$BEST_PY" ]; then
echo "❌ ERROR: No Python version <= 3.12.x found on this system."
echo " Detected Python 3 installations:"
for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)?$' | sort -u); do
if command -v "$candidate" &>/dev/null; then
echo " - $candidate ($($candidate --version 2>&1))"
fi
done
echo ""
echo " Please install Python <= 3.12.x for maximum compatibility."
echo " For example: sudo apt install python3.12 python3.12-venv"
exit 1
fi
BEST_VER=$("$BEST_PY" --version 2>&1 | awk '{print $2}')
echo "✅ Using $BEST_PY ($BEST_VER) — compatible (≤ 3.12.x)"
if [ "$IS_COLAB" = true ]; then
# Colab: install packages directly without venv
run_quiet "pip upgrade" pip install --upgrade pip
echo " Installing unsloth-zoo + unsloth..."
run_quiet "pip install unsloth" pip install -r "$SCRIPT_DIR/studio/backend/requirements/base.txt"
echo " Installing additional unsloth dependencies..."
run_quiet "pip install extras" pip install --no-cache-dir -r "$SCRIPT_DIR/studio/backend/requirements/extras.txt"
run_quiet "pip install extras" pip install --no-deps --no-cache-dir -r "$SCRIPT_DIR/studio/backend/requirements/extras-no-deps.txt"
run_quiet "pip install torchao+transformers" pip install --force-reinstall --no-cache-dir -r "$SCRIPT_DIR/studio/backend/requirements/overrides.txt"
run_quiet "pip install triton_kernels" pip install --no-deps -r "$SCRIPT_DIR/studio/backend/requirements/triton-kernels.txt"
# Patch: override llama_cpp.py with fix from unsloth-zoo branch
LLAMA_CPP_DST="$(pip show unsloth-zoo | grep -i '^Location:' | awk '{print $2}')/unsloth_zoo/llama_cpp.py"
curl -sSL "https://raw.githubusercontent.com/unslothai/unsloth-zoo/refs/heads/main/unsloth_zoo/llama_cpp.py" \
-o "$LLAMA_CPP_DST"
echo " Installing studio dependencies..."
run_quiet "pip install studio" pip install -r "$SCRIPT_DIR/studio/backend/requirements/studio.txt"
echo "✅ Python dependencies installed"
else
# Local: create venv (always start fresh to preserve correct install order)
rm -rf .venv
"$BEST_PY" -m venv .venv
source .venv/bin/activate
run_quiet "pip upgrade" pip install --upgrade pip
echo " Installing unsloth-zoo + unsloth..."
run_quiet "pip install unsloth" pip install -r "$SCRIPT_DIR/studio/backend/requirements/base.txt"
echo " Installing additional unsloth dependencies..."
run_quiet "pip install extras" pip install --no-cache-dir -r "$SCRIPT_DIR/studio/backend/requirements/extras.txt"
run_quiet "pip install extras" pip install --no-deps --no-cache-dir -r "$SCRIPT_DIR/studio/backend/requirements/extras-no-deps.txt"
run_quiet "pip install torchao+transformers" pip install --force-reinstall --no-cache-dir -r "$SCRIPT_DIR/studio/backend/requirements/overrides.txt"
run_quiet "pip install triton_kernels" pip install --no-deps -r "$SCRIPT_DIR/studio/backend/requirements/triton-kernels.txt"
# Patch: override llama_cpp.py with fix from unsloth-zoo branch
LLAMA_CPP_DST="$(pip show unsloth-zoo | grep -i '^Location:' | awk '{print $2}')/unsloth_zoo/llama_cpp.py"
curl -sSL "https://raw.githubusercontent.com/unslothai/unsloth-zoo/refs/heads/main/unsloth_zoo/llama_cpp.py" \
-o "$LLAMA_CPP_DST"
echo " Installing studio dependencies..."
run_quiet "pip install studio" pip install -r "$SCRIPT_DIR/studio/backend/requirements/studio.txt"
echo "✅ Python dependencies installed"
# ── 7. WSL: pre-install GGUF build dependencies ──
# On WSL, sudo requires a password and can't be entered during GGUF export
# (runs in a non-interactive subprocess). Install build deps here instead.
if grep -qi microsoft /proc/version 2>/dev/null; then
echo ""
echo "⚠️ WSL detected — installing build dependencies for GGUF export..."
echo " You may be prompted for your password."
sudo apt-get update -y
sudo apt-get install -y build-essential cmake curl git libcurl4-openssl-dev
echo "✅ GGUF build dependencies installed"
fi
fi
# ── 8. Add shell alias (skip in Colab) ──
# Note: venv activation does NOT persist across terminal sessions.
# This alias hardcodes the venv python path so users don't need to activate.
if [ "$IS_COLAB" = false ]; then
echo ""
REPO_DIR="$SCRIPT_DIR"
if ! grep -qF "unsloth-ui" ~/.bashrc 2>/dev/null; then
cat >> ~/.bashrc <<UNSLOTH_EOF
# Detect the user's default shell and pick the right rc file
USER_SHELL="$(basename "${SHELL:-/bin/bash}")"
case "$USER_SHELL" in
zsh)
SHELL_RC="$HOME/.zshrc"
ALIAS_BLOCK="alias unsloth-studio='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'
alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'"
;;
fish)
SHELL_RC="$HOME/.config/fish/config.fish"
# fish uses 'abbr' or 'function'; a simple alias works via 'alias' in config.fish
ALIAS_BLOCK="alias unsloth-studio '${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'
alias unsloth-ui '${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'"
;;
ksh)
SHELL_RC="$HOME/.kshrc"
ALIAS_BLOCK="alias unsloth-studio='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'
alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'"
;;
*)
# Default to bash for bash and any other POSIX-compatible shell
SHELL_RC="$HOME/.bashrc"
ALIAS_BLOCK="alias unsloth-studio='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'
alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'"
;;
esac
echo " Detected shell: $USER_SHELL → $SHELL_RC"
ALIAS_ADDED=false
if ! grep -qF "unsloth-studio" "$SHELL_RC" 2>/dev/null; then
mkdir -p "$(dirname "$SHELL_RC")" # needed for fish's nested config path
cat >> "$SHELL_RC" <<UNSLOTH_EOF
# Unsloth Studio launcher
alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py ui -f ${REPO_DIR}/studio/frontend/dist'
$ALIAS_BLOCK
UNSLOTH_EOF
echo "✅ Alias 'unsloth-ui' added to ~/.bashrc"
echo "✅ Aliases 'unsloth-studio' and 'unsloth-ui' added to $SHELL_RC"
ALIAS_ADDED=true
else
echo "✅ Alias 'unsloth-ui' already exists in ~/.bashrc"
echo "✅ Aliases 'unsloth-studio' and 'unsloth-ui' already exist in $SHELL_RC"
fi
fi # End of "if not Colab" for shell alias setup
echo ""
echo "╔══════════════════════════════════════╗"
echo "║ Setup Complete! ║"
echo "╠══════════════════════════════════════╣"
echo "║ Run 'source ~/.bashrc' or open a ║"
echo "║ new terminal, then launch with: ║"
echo "║ ║"
echo "║ unsloth-ui -H 0.0.0.0 -p 8000 ║"
echo "╚══════════════════════════════════════╝"
if [ "$IS_COLAB" = true ]; then
echo "╔══════════════════════════════════════╗"
echo "║ Setup Complete! ║"
echo "╠══════════════════════════════════════╣"
echo "║ Unsloth Studio is ready to start ║"
echo "║ in your Colab notebook! ║"
echo "╚══════════════════════════════════════╝"
else
echo "╔══════════════════════════════════════╗"
echo "║ Setup Complete! ║"
echo "╠══════════════════════════════════════╣"
if [ "$ALIAS_ADDED" = true ]; then
echo "║ Run 'source $SHELL_RC'"
echo "║ or open a new terminal, then: ║"
else
echo "║ Launch with: ║"
fi
echo "║ ║"
echo "║ unsloth-studio -H 0.0.0.0 -p 8000 ║"
echo "╚══════════════════════════════════════╝"
fi

9
studio/backend/assets/configs/model_defaults/default.yaml
View file

@ -3,17 +3,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 5e-5
batch_size: 2
gradient_accumulation_steps: 4
warmup_ratio: 0.1
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.01
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

7
studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml
View file

@ -4,13 +4,14 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 4
gradient_accumulation_steps: 2
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false

9
studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 2
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 8
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.01
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: false
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 4096
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

7
studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml
View file

@ -5,13 +5,14 @@
training:
max_seq_length: 4096
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 4
gradient_accumulation_steps: 2
warmup_steps: 10
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false

9
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml
View file

@ -3,17 +3,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.01
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 5e-5
batch_size: 4
gradient_accumulation_steps: 1
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 2
# num_epochs: 2
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 4
warmup_ratio: 0.03
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: true
optim: "adamw_torch_fused"
lr_scheduler_type: "cosine"

9
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 1024
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 2
# num_epochs: 2
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 4
warmup_ratio: 0.03
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: true
optim: "adamw_torch_fused"
lr_scheduler_type: "cosine"

9
studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 4096
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 4
gradient_accumulation_steps: 1
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

7
studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml
View file

@ -5,13 +5,14 @@
training:
max_seq_length: 1024
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false

7
studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml
View file

@ -5,13 +5,14 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false

7
studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml
View file

@ -5,13 +5,14 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false

9
studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 5
# num_epochs: 5
num_epochs: 0
learning_rate: 2e-5
batch_size: 1
gradient_accumulation_steps: 8
warmup_steps: 0
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.01
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: true
optim: "adamw_torch"
lr_scheduler_type: "cosine"

9
studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 8192
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 5e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 2
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 4
gradient_accumulation_steps: 2
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "paged_adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml
View file

@ -3,17 +3,18 @@
# Also applies to: "unsloth/mistral-7b-v0.3", "mistralai/Mistral-7B-v0.3",
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

7
studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml
View file

@ -5,13 +5,14 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false

7
studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml
View file

@ -5,13 +5,14 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 4
gradient_accumulation_steps: 2
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false

9
studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 5e-5
batch_size: 4
gradient_accumulation_steps: 2
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml
View file

@ -3,17 +3,18 @@
training:
max_seq_length: 2048
num_epochs: 1
# num_epochs: 1
num_epochs: 0
learning_rate: 5e-5
batch_size: 32
gradient_accumulation_steps: 1
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 4096
num_epochs: 1
# num_epochs: 1
num_epochs: 0
learning_rate: 2e-5
batch_size: 2
gradient_accumulation_steps: 4
warmup_ratio: 0.1
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.1
random_seed: 3407
packing: true
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 448
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 1e-4
batch_size: 1
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4-14B-Instruct.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 4096
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 1e-5
batch_size: 1
gradient_accumulation_steps: 64
warmup_ratio: 0.1
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 42
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "paged_adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml
View file

@ -4,17 +4,18 @@
training:
max_seq_length: 32768
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 1024
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 5e-5
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

7
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml
View file

@ -5,13 +5,14 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false

9
studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml
View file

@ -5,17 +5,18 @@
training:
max_seq_length: 2048
num_epochs: 4
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 0
save_steps: 0
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"

90
studio/backend/colab.py Normal file
View file

@ -0,0 +1,90 @@
"""
Colab-specific helpers for running Unsloth Studio.
Uses Colab's built-in proxy - no external tunneling needed!
"""
from pathlib import Path
def get_colab_url(port: int = 8000) -> str:
"""
Get the actual Colab proxy URL for a port.
"""
try:
from google.colab.output import eval_js
# Use Colab's proxy mechanism
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec=5)
return url if url else f"http://localhost:{port}"
except Exception as e:
print(f"Note: Could not get Colab URL ({e})")
return f"http://localhost:{port}"
def show_link(port: int = 8000):
"""Display a styled clickable link to the UI."""
from IPython.display import display, HTML
# Get real Colab proxy URL
url = get_colab_url(port)
html = f"""
<div style="padding: 20px; background: linear-gradient(135deg, #22c55e 0%, #16a34a 100%);
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
<h2 style="color: white; margin: 0 0 12px 0; font-size: 24px;">
🦥 Unsloth Studio is Ready!
</h2>
<a href="{url}" target="_blank"
style="display: inline-block; padding: 14px 28px; background: white; color: #16a34a;
text-decoration: none; border-radius: 8px; font-weight: 600; font-size: 16px;
box-shadow: 0 4px 6px rgba(0,0,0,0.1);">
🚀 Open Unsloth Studio
</a>
<p style="color: rgba(255,255,255,0.9); margin: 16px 0 0 0; font-size: 13px;
word-break: break-all; font-family: monospace;">
{url}
</p>
</div>
"""
display(HTML(html))
def start(port: int = 8000):
"""
Start Unsloth Studio server in Colab and display the URL.
Usage:
from colab import start
start()
"""
import sys
print("🦥 Starting Unsloth Studio...")
# Add backend to path
backend_path = str(Path(__file__).parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
print(" Loading backend...")
from run import run_server
# Auto-detect frontend path
repo_root = Path(__file__).parent.parent
frontend_path = repo_root / "frontend" / "dist"
if not frontend_path.exists():
print("❌ Frontend not built! Please run the setup cell first.")
return
print(" Starting server...")
# Start server silently
run_server(host="0.0.0.0", port=port, frontend_path=frontend_path, silent=True)
print(" Server started!")
# Show the clickable link with real URL
show_link(port)
if __name__ == "__main__":
start()

87
studio/backend/core/export/export.py
View file

@ -18,6 +18,45 @@ from core.inference import get_inference_backend
logger = logging.getLogger(__name__)
def _is_wsl():
"""Detect if running under Windows Subsystem for Linux."""
try:
return "microsoft" in open("/proc/version").read().lower()
except Exception:
return False
def _apply_wsl_sudo_patch():
"""On WSL, monkey-patch do_we_need_sudo() to return False.
WSL doesn't have passwordless sudo, and do_we_need_sudo() runs
`sudo apt-get update` which hangs waiting for a stdin password
inside a non-interactive subprocess. setup.sh pre-installs the
build dependencies on WSL, so sudo is not needed at runtime.
"""
if not _is_wsl():
return
try:
import unsloth_zoo.llama_cpp as llama_cpp_module
def _wsl_do_we_need_sudo(system_type="debian"):
logger.info(
"WSL detected — skipping sudo check "
"(build deps pre-installed by setup.sh)"
)
return False
llama_cpp_module.do_we_need_sudo = _wsl_do_we_need_sudo
logger.info(
"Applied WSL sudo patch to "
"unsloth_zoo.llama_cpp.do_we_need_sudo"
)
except Exception as e:
logger.warning(f"Could not apply WSL sudo patch: {e}")
# Model card template
MODEL_CARD = \
"""---
@ -80,43 +119,15 @@ class ExportBackend:
logger.error(f"Error during memory cleanup: {e}")
return False
def scan_checkpoints(self, outputs_dir: str = "./outputs") -> List[Tuple[str, str]]:
def scan_checkpoints(self, outputs_dir: str = "./outputs") -> List[Tuple[str, List[Tuple[str, str]]]]:
"""
Scan outputs folder for model checkpoints.
Scan outputs folder for training runs and their checkpoints.
Returns:
List of tuples: [(display_name, checkpoint_path), ...]
List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...]
"""
checkpoints = []
outputs_path = Path(outputs_dir)
if not outputs_path.exists():
logger.warning(f"Outputs directory not found: {outputs_dir}")
return checkpoints
try:
for item in outputs_path.iterdir():
if item.is_dir():
# Check if this directory contains a model
config_file = item / "config.json"
adapter_config = item / "adapter_config.json"
if config_file.exists() or adapter_config.exists():
# This is a valid checkpoint
display_name = item.name
checkpoint_path = str(item)
checkpoints.append((display_name, checkpoint_path))
logger.debug(f"Found checkpoint: {display_name}")
# Sort by modification time (newest first)
checkpoints.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True)
logger.info(f"Found {len(checkpoints)} checkpoints in {outputs_dir}")
return checkpoints
except Exception as e:
logger.error(f"Error scanning checkpoints: {e}")
return []
from utils.models.checkpoints import scan_checkpoints
return scan_checkpoints(outputs_dir=outputs_dir)
def load_checkpoint(self,
checkpoint_path: str,
@ -264,7 +275,8 @@ class ExportBackend:
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False) -> Tuple[bool, str]:
private: bool = False,
base_model_id: Optional[str] = None) -> Tuple[bool, str]:
"""
Export base model (for non-PEFT models).
@ -294,8 +306,8 @@ class ExportBackend:
logger.info(f"Pushing base model to Hub: {repo_id}")
# Get base model name
base_model = self.current_model.config._name_or_path
# Get base model name from request or model config
base_model = base_model_id or self.current_model.config._name_or_path or "unknown"
# Create repo
hf_api = HfApi(token=hf_token)
@ -380,6 +392,9 @@ class ExportBackend:
os.chdir(save_directory)
logger.info(f"Changed directory to: {save_directory}")
# On WSL, patch out sudo check before llama.cpp build
_apply_wsl_sudo_patch()
# Now save (will save in current directory)
self.current_model.save_pretrained_gguf(
"model", # Base filename

268
studio/backend/core/inference/inference.py
View file

@ -38,13 +38,22 @@ class InferenceBackend:
]
self.device = get_device().value
# Thread safety
# Thread safety — _generation_lock serializes model.generate() calls.
# Must be a regular Lock (NOT RLock) because in async FastAPI, multiple
# requests share the same event-loop thread, so RLock reentrancy lets
# concurrent compare-mode requests race on the GPU. The lock is
# acquired by the *background generation thread*, not the event-loop.
import threading
self._generation_lock = threading.RLock()
self._generation_lock = threading.Lock()
self._model_state_lock = threading.Lock()
logger.info(f"InferenceBackend initialized on {self.device}")
@staticmethod
def _normalize_top_k(top_k: int) -> int:
# API supports -1 as "disable top-k"; transformers expects 0 to disable.
return 0 if top_k < 0 else top_k
def load_model(self,
config: ModelConfig,
max_seq_length: int = 2048,
@ -443,9 +452,10 @@ class InferenceBackend:
"""
Apply adapter state before generation. Must be called under _generation_lock.
Uses revert_to_base_model() / activate_lora_adapter() which work correctly
for models loaded by Unsloth as complete PeftModels (via model.unload() /
model.load_adapter()), matching the proven pattern from the Gradio eval page.
Uses PEFT's disable_adapter_layers() / enable_adapter_layers() which toggle
a boolean flag on each LoRA layer. Unsloth's fast_linear_forward checks this
flag (proj.disable_adapters) and skips LoRA computation when True.
This is non-destructive — no model unloading/reloading needed.
Args:
use_adapter: None = no change, False = disable (base model),
@ -459,53 +469,56 @@ class InferenceBackend:
return
model_info = self.models[base]
model = model_info.get("model")
if model is None:
return
if use_adapter is False:
# Revert to pure base model by unloading adapter weights
logger.info(f"Compare mode: reverting '{base}' to base model for generation")
self.revert_to_base_model(base)
# Disable LoRA layers → base model output
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
logger.info(f"Compare mode: disabling adapters on '{base}' for base model generation")
model.base_model.disable_adapter_layers()
else:
logger.info(f"Compare mode: model '{base}' is not a PeftModel, already base")
elif use_adapter is True:
# Activate the LoRA adapter from the original model path
lora_path = model_info.get("model_path")
if lora_path and model_info.get("is_lora"):
logger.info(f"Compare mode: activating LoRA adapter from '{lora_path}' on '{base}'")
self.activate_lora_adapter(base, lora_path)
# Re-enable LoRA layers → adapter output
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
logger.info(f"Compare mode: enabling adapters on '{base}' for LoRA generation")
model.base_model.enable_adapter_layers()
else:
# Fallback for dynamically attached adapters
loaded = model_info.get("loaded_adapters", {})
if loaded:
adapter_name = list(loaded.keys())[-1]
logger.info(f"Compare mode: enabling adapter '{adapter_name}' on '{base}'")
self.set_active_adapter(base, adapter_name)
else:
logger.warning("use_adapter=true but no adapter path/adapters on model")
logger.warning("use_adapter=true but model is not a PeftModel")
elif isinstance(use_adapter, str):
# Activate a specific adapter by path
logger.info(f"Compare mode: activating specific adapter '{use_adapter}' on '{base}'")
self.activate_lora_adapter(base, use_adapter)
# Enable adapters and set the specific one active
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
logger.info(f"Compare mode: enabling adapter '{use_adapter}' on '{base}'")
model.base_model.enable_adapter_layers()
self.set_active_adapter(base, use_adapter)
else:
logger.warning(f"use_adapter='{use_adapter}' but model is not a PeftModel")
def generate_with_adapter_control(
self,
use_adapter: Optional[Union[bool, str]] = None,
cancel_event=None,
**gen_kwargs,
) -> Generator[str, None, None]:
"""
Thread-safe generation with optional adapter toggling.
Acquires the generation lock, applies adapter state, then generates.
This ensures adapter toggle + generation are atomic — critical for
compare mode where base and LoRA panes fire concurrently.
The adapter toggle + model.generate() are serialized by _generation_lock
inside the background generation thread — NOT in the event-loop thread.
This prevents the RLock-reentrant race that occurs when two async SSE
handlers share the same event-loop thread.
Args:
use_adapter: Adapter control (None/False/True/str). See _apply_adapter_state.
**gen_kwargs: Forwarded to generate_chat_response.
"""
with self._generation_lock:
self._apply_adapter_state(use_adapter)
# Delegate to the lock-free generation path
yield from self._generate_chat_response_inner(**gen_kwargs)
yield from self._generate_chat_response_inner(
cancel_event=cancel_event, _adapter_state=use_adapter, **gen_kwargs
)
def generate_chat_response(self,
messages: list,
@ -514,24 +527,26 @@ class InferenceBackend:
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.1) -> Generator[str, None, None]:
repetition_penalty: float = 1.1,
cancel_event=None) -> Generator[str, None, None]:
"""
Generate response for text or vision models.
Acquires the generation lock. For adapter-controlled generation,
use generate_with_adapter_control() instead.
The generation lock is acquired by the background generation thread.
"""
with self._generation_lock:
yield from self._generate_chat_response_inner(
messages=messages,
system_prompt=system_prompt,
image=image,
temperature=temperature,
top_p=top_p,
top_k=top_k,
max_new_tokens=max_new_tokens,
repetition_penalty=repetition_penalty,
)
yield from self._generate_chat_response_inner(
messages=messages,
system_prompt=system_prompt,
image=image,
temperature=temperature,
top_p=top_p,
top_k=top_k,
min_p=min_p,
max_new_tokens=max_new_tokens,
repetition_penalty=repetition_penalty,
cancel_event=cancel_event,
)
def _generate_chat_response_inner(self,
messages: list,
@ -540,11 +555,17 @@ class InferenceBackend:
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.1) -> Generator[str, None, None]:
repetition_penalty: float = 1.1,
cancel_event=None,
_adapter_state=None) -> Generator[str, None, None]:
"""
Inner generation logic (no lock). Called by both generate_chat_response
Inner generation logic. Called by both generate_chat_response
and generate_with_adapter_control.
_adapter_state is passed to generate_stream/vision so the background
thread can toggle adapters under the generation lock.
"""
if not self.active_model_name:
yield "Error: No active model"
@ -553,12 +574,14 @@ class InferenceBackend:
model_info = self.models[self.active_model_name]
is_vision = model_info.get("is_vision", False)
tokenizer = model_info.get("tokenizer") or model_info.get("processor")
top_k = self._normalize_top_k(top_k)
if is_vision:
# Vision model generation
yield from self._generate_vision_response(
messages, system_prompt, image,
temperature, top_p, top_k, max_new_tokens, repetition_penalty
temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty,
cancel_event=cancel_event,
)
else:
# Text model: Use training pipeline approach
@ -600,12 +623,14 @@ class InferenceBackend:
# Step 3: Generate
yield from self.generate_stream(
formatted_prompt, temperature, top_p, top_k, max_new_tokens, repetition_penalty
formatted_prompt, temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty,
cancel_event=cancel_event,
_adapter_state=_adapter_state,
)
def _generate_vision_response(self, messages, system_prompt, image,
temperature, top_p, top_k, max_new_tokens,
repetition_penalty) -> Generator[str, None, None]:
temperature, top_p, top_k, min_p, max_new_tokens,
repetition_penalty, cancel_event=None) -> Generator[str, None, None]:
"""Handle vision model generation with true token-by-token streaming."""
model_info = self.models[self.active_model_name]
model = model_info["model"]
@ -651,7 +676,10 @@ class InferenceBackend:
import threading
streamer = TextIteratorStreamer(
processor.tokenizer, skip_prompt=True, skip_special_tokens=True
processor.tokenizer,
skip_prompt=True,
skip_special_tokens=True,
timeout=0.2,
)
generation_kwargs = dict(
@ -659,28 +687,58 @@ class InferenceBackend:
streamer=streamer,
max_new_tokens=max_new_tokens,
use_cache=True,
do_sample=temperature > 0,
temperature=temperature,
top_p=top_p,
top_k=top_k,
min_p=min_p,
)
err: dict[str, str] = {}
def generate_fn():
try:
model.generate(**generation_kwargs)
except Exception as e:
logger.error(f"Vision generation error in thread: {e}")
with self._generation_lock:
try:
model.generate(**generation_kwargs)
except Exception as e:
err["msg"] = str(e)
logger.error(f"Vision generation error in thread: {e}")
finally:
try:
streamer.end()
except Exception:
pass
thread = threading.Thread(target=generate_fn)
thread.start()
output = ""
for new_token in streamer:
if new_token:
output += new_token
cleaned = self._clean_generated_text(output)
yield cleaned
from queue import Empty
try:
while True:
if cancel_event is not None and cancel_event.is_set():
break
try:
new_token = next(streamer)
except StopIteration:
break
except Empty:
if not thread.is_alive():
break
continue
if new_token:
output += new_token
cleaned = self._clean_generated_text(output)
yield cleaned
finally:
if cancel_event is not None:
cancel_event.set()
thread.join(timeout=10)
if thread.is_alive():
logger.warning("Vision generation thread did not exit after cancel/join timeout")
thread.join()
if err.get("msg"):
yield f"Error: {err['msg']}"
except Exception as e:
logger.error(f"Vision generation error: {e}")
@ -692,9 +750,16 @@ class InferenceBackend:
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.1) -> Generator[str, None, None]:
"""Generate streaming text response (text models only)."""
repetition_penalty: float = 1.1,
cancel_event=None,
_adapter_state=None) -> Generator[str, None, None]:
"""Generate streaming text response (text models only).
_adapter_state: if not None, the background thread toggles adapters
before model.generate(), all under _generation_lock.
"""
if not self.active_model_name:
yield "Error: No active model"
return
@ -709,7 +774,12 @@ class InferenceBackend:
from transformers import TextIteratorStreamer
import threading
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt=True,
skip_special_tokens=True,
timeout=0.2,
)
generation_kwargs = dict(
**inputs,
@ -718,29 +788,75 @@ class InferenceBackend:
temperature=temperature,
top_p=top_p,
top_k=top_k,
min_p=min_p,
repetition_penalty=repetition_penalty,
do_sample=True,
do_sample=temperature > 0,
eos_token_id=tokenizer.eos_token_id,
pad_token_id=tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id,
)
if cancel_event is not None:
from transformers.generation.stopping_criteria import (
StoppingCriteria,
StoppingCriteriaList,
)
class _CancelCriteria(StoppingCriteria):
def __init__(self, ev):
self.ev = ev
def __call__(self, input_ids, scores, **kwargs):
return self.ev.is_set()
generation_kwargs["stopping_criteria"] = StoppingCriteriaList(
[_CancelCriteria(cancel_event)]
)
def generate_fn():
try:
model.generate(**generation_kwargs)
except Exception as e:
logger.error(f"Generation error: {e}")
with self._generation_lock:
try:
if _adapter_state is not None:
self._apply_adapter_state(_adapter_state)
model.generate(**generation_kwargs)
except Exception as e:
err["msg"] = str(e)
logger.error(f"Generation error: {e}")
finally:
try:
streamer.end()
except Exception:
pass
err: dict[str, str] = {}
thread = threading.Thread(target=generate_fn)
thread.start()
output = ""
for new_token in streamer:
if new_token:
output += new_token
cleaned = self._clean_generated_text(output)
yield cleaned
from queue import Empty
try:
while True:
if cancel_event is not None and cancel_event.is_set():
break
try:
new_token = next(streamer)
except StopIteration:
break
except Empty:
if not thread.is_alive():
break
continue
if new_token:
output += new_token
cleaned = self._clean_generated_text(output)
yield cleaned
finally:
if cancel_event is not None:
cancel_event.set()
thread.join(timeout=10)
if thread.is_alive():
logger.warning("Generation thread did not exit after cancel/join timeout")
thread.join()
if err.get("msg"):
yield f"Error: {err['msg']}"
except Exception as e:
logger.error(f"Error during generation: {e}")

291
studio/backend/core/training/trainer.py
View file

@ -16,6 +16,7 @@ import json
import threading
import math
import logging
import time
from typing import Optional, Callable
from dataclasses import dataclass
import pandas as pd
@ -46,6 +47,11 @@ class TrainingProgress:
is_completed: bool = False
error: Optional[str] = None
status_message: str = "Ready to train" # Current stage message
elapsed_seconds: Optional[float] = None
eta_seconds: Optional[float] = None
grad_norm: Optional[float] = None
num_tokens: Optional[int] = None
eval_loss: Optional[float] = None
class UnslothTrainer:
"""
@ -67,6 +73,12 @@ class UnslothTrainer:
self.is_vlm = False
self.model_name = None
# Training metrics tracking
self.training_start_time: Optional[float] = None
self.batch_size: Optional[int] = None
self.max_seq_length: Optional[int] = None
self.gradient_accumulation_steps: Optional[int] = None
# Thread safety
self._lock = threading.Lock()
@ -99,17 +111,29 @@ class UnslothTrainer:
model_name: str,
max_seq_length: int = 2048,
load_in_4bit: bool = True,
hf_token: Optional[str] = None) -> bool:
hf_token: Optional[str] = None,
is_dataset_multimodal: bool = False) -> bool:
"""Load model for training (supports both text and vision models)"""
try:
if self.model is not None:
del self.model
if self.tokenizer is not None:
del self.tokenizer
if self.trainer is not None:
del self.trainer
print("\nClearing GPU memory before training...")
clear_gpu_cache()
# Detect if this is a vision model first
self.is_vlm = is_vision_model(model_name)
# Detect if this is a vision model AND dataset is multimodal
# A vision-capable model with a text-only dataset should use FastLanguageModel
self.is_vlm = is_vision_model(model_name) and is_dataset_multimodal
self.model_name = model_name
logger.info(f"Model type detected: {'Vision' if self.is_vlm else 'Text'}")
logger.info(f"Model architecture is vision: {is_vision_model(model_name)}")
logger.info(f"Dataset is multimodal: {is_dataset_multimodal}")
logger.info(f"Using VLM path: {self.is_vlm}")
# Reset training state for new run
self._update_progress(
@ -198,7 +222,13 @@ class UnslothTrainer:
return True
# LoRA/QLoRA mode - apply PEFT
if target_modules is None or (isinstance(target_modules, list) and len(target_modules) == 0):
# "all-linear" is a PEFT keyword that targets every linear layer
if isinstance(target_modules, list) and "all-linear" in target_modules:
if len(target_modules) == 1:
target_modules = "all-linear"
else:
target_modules = [m for m in target_modules if m != "all-linear"]
elif target_modules is None or (isinstance(target_modules, list) and len(target_modules) == 0):
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"]
@ -306,12 +336,24 @@ class UnslothTrainer:
dataset_source: str,
format_type: str = "auto",
local_datasets: list = None,
custom_format_mapping: dict = None) -> Optional[Dataset]:
custom_format_mapping: dict = None,
subset: str = None,
train_split: str = "train",
eval_split: str = None) -> Optional[tuple]:
"""
Load and prepare dataset for training
Load and prepare dataset for training.
Strategy: format first, then split — ensures both train and eval
portions are properly formatted and templated.
Returns:
Tuple of (dataset_info, eval_dataset) or None on error.
eval_dataset may be None if no eval split is available.
"""
try:
dataset = None
eval_dataset = None
has_separate_eval_source = False # True if eval comes from a separate HF split
if local_datasets:
# Load local datasets
@ -350,7 +392,10 @@ class UnslothTrainer:
elif dataset_source:
# Load from Hugging Face
dataset = load_dataset(dataset_source, split="train")
load_kwargs = {"path": dataset_source, "split": train_split or "train"}
if subset:
load_kwargs["name"] = subset
dataset = load_dataset(**load_kwargs)
# Check if stopped during dataset loading
if self.should_stop:
@ -360,6 +405,25 @@ class UnslothTrainer:
self._update_progress(status_message=f"Loaded dataset from HuggingFace: {dataset_source}")
print(f"Loaded dataset from Hugging Face: {dataset_source}\n")
# Resolve eval split from a separate HF split (explicit or auto-detected)
if eval_split:
# Explicit eval split provided - load it directly
print(f"Loading explicit eval split: '{eval_split}'\n")
eval_load_kwargs = {"path": dataset_source, "split": eval_split}
if subset:
eval_load_kwargs["name"] = subset
eval_dataset = load_dataset(**eval_load_kwargs)
has_separate_eval_source = True
print(f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n")
else:
# Auto-detect eval split from HF (returns a separate dataset, or None)
eval_dataset = self._auto_detect_eval_split_from_hf(
dataset_source=dataset_source,
subset=subset,
)
if eval_dataset is not None:
has_separate_eval_source = True
if dataset is None:
raise ValueError("No dataset provided")
@ -368,16 +432,15 @@ class UnslothTrainer:
print("Stopped before applying chat template\n")
return None
# NEW: Use unified format_and_template_dataset
# ========== FORMAT FIRST ==========
print(f"Formatting dataset with format_type='{format_type}'...\n")
#breakpoint()
dataset_info = format_and_template_dataset(
dataset,
model_name=self.model_name,
tokenizer=self.tokenizer, # Works for both text and vision models
tokenizer=self.tokenizer,
is_vlm=self.is_vlm,
format_type=format_type, # "auto", "alpaca", "chatml", "sharegpt"
format_type=format_type,
dataset_name=dataset_source,
custom_format_mapping=custom_format_mapping,
)
@ -389,15 +452,94 @@ class UnslothTrainer:
self._update_progress(status_message=f"Dataset formatted and ready for training")
print(f"Dataset formatted successfully\n")
return dataset_info
# ========== THEN SPLIT ==========
if has_separate_eval_source and eval_dataset is not None:
# Eval came from a separate HF split — format it too
print(f"Formatting eval dataset ({len(eval_dataset)} rows)...\n")
eval_info = format_and_template_dataset(
eval_dataset,
model_name=self.model_name,
tokenizer=self.tokenizer,
is_vlm=self.is_vlm,
format_type=format_type,
dataset_name=dataset_source,
custom_format_mapping=custom_format_mapping,
)
eval_dataset = eval_info["dataset"]
print(f"Eval dataset formatted successfully\n")
elif not has_separate_eval_source:
# No separate eval source — split the already-formatted dataset
formatted_dataset = dataset_info["dataset"]
split_result = self._resolve_eval_split_from_dataset(formatted_dataset)
if split_result is not None:
train_portion, eval_dataset = split_result
dataset_info["dataset"] = train_portion
return (dataset_info, eval_dataset)
except Exception as e:
logger.error(f"Error loading dataset: {e}")
self._update_progress(error=str(e))
return None
def _auto_detect_eval_split_from_hf(self, dataset_source: str,
subset: str) -> Optional[Dataset]:
"""Auto-detect an eval split from HF dataset (separate named split only)."""
try:
from datasets import get_dataset_split_names
load_kwargs = {"path": dataset_source}
if subset:
load_kwargs["name"] = subset
available_splits = get_dataset_split_names(**load_kwargs)
print(f"Available splits: {available_splits}\n")
# Check for common eval split names
for candidate in ["eval", "validation", "valid", "val", "test"]:
if candidate in available_splits:
eval_load_kwargs = {"path": dataset_source, "split": candidate}
if subset:
eval_load_kwargs["name"] = subset
candidate_ds = load_dataset(**eval_load_kwargs)
if len(candidate_ds) >= 16:
print(f"Auto-detected eval split '{candidate}' with {len(candidate_ds)} rows\n")
return candidate_ds
else:
print(f"Found eval split '{candidate}' but only {len(candidate_ds)} rows (< 16), skipping\n")
except Exception as e:
logger.warning(f"Could not check dataset splits: {e}")
# No separate HF eval split found — caller will handle programmatic splitting
return None
def _resolve_eval_split_from_dataset(self, dataset) -> Optional[tuple]:
"""Split a dataset into train and eval portions.
Returns:
Tuple of (train_dataset, eval_dataset), or None if dataset too small.
"""
MIN_EVAL_ROWS = 16
MIN_TOTAL_ROWS = 32 # Need at least 16 train + 16 eval
n = len(dataset)
if n < MIN_TOTAL_ROWS:
print(f"Dataset too small ({n} rows) for eval split, skipping eval\n")
return None
eval_size = max(MIN_EVAL_ROWS, min(128, int(0.05 * n)))
# Ensure we don't take more than half the dataset
eval_size = min(eval_size, n // 2)
print(f"Auto-splitting: {eval_size} rows for eval from {n} total\n")
split_result = dataset.train_test_split(test_size=eval_size, seed=3407)
print(f"Split complete: {len(split_result['train'])} train, {len(split_result['test'])} eval\n")
return (split_result['train'], split_result['test'])
def start_training(self,
dataset: Dataset,
eval_dataset: Dataset = None,
eval_steps: float = 0.01,
output_dir: str = "./outputs",
num_epochs: int = 3,
learning_rate: float = 5e-5,
@ -451,17 +593,33 @@ class UnslothTrainer:
'wandb_token': wandb_token,
'enable_tensorboard': enable_tensorboard,
'tensorboard_dir': tensorboard_dir,
'eval_dataset': eval_dataset,
'eval_steps': eval_steps,
**kwargs
}
)
self.should_stop = False
self.training_thread.start()
return True
self.is_training = True
try:
self.training_thread.start()
return True
except Exception as e:
self.is_training = False
logger.error(f"Failed to start training thread: {e}")
return False
def _train_worker(self, dataset: Dataset, **training_args):
"""Worker function for training (runs in separate thread)"""
try:
# Store training parameters for metrics calculation
self.batch_size = training_args.get('batch_size', 2)
self.max_seq_length = training_args.get('max_seq_length', 2048)
self.gradient_accumulation_steps = training_args.get('gradient_accumulation_steps', 4)
# Set training start time
self.training_start_time = time.time()
self._update_progress(is_training=True, error=None)
# Setup logging
@ -548,6 +706,8 @@ class UnslothTrainer:
"seed": training_args.get('random_seed', 3407),
"output_dir": output_dir,
"report_to": ["wandb"] if training_args.get('enable_wandb', False) else "none",
"include_num_input_tokens_seen": True, # Enable token counting
"dataset_num_proc": max(1, os.cpu_count() // 4),
}
# Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps
@ -577,6 +737,17 @@ class UnslothTrainer:
else:
print(f"Training for {config_args['num_train_epochs']} epochs\n")
# ========== EVAL CONFIGURATION ==========
eval_dataset = training_args.get('eval_dataset', None)
eval_steps_val = training_args.get('eval_steps', 0.01)
if eval_dataset is not None:
config_args["eval_strategy"] = "steps"
config_args["eval_steps"] = eval_steps_val
print(f"Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n")
print(f"Eval dataset: {len(eval_dataset)} rows\n")
else:
print("No eval dataset — evaluation disabled\n")
# Add model-specific parameters
# Use optim and lr_scheduler_type from training_args if provided, otherwise use defaults
optim_value = training_args.get('optim', "adamw_8bit")
@ -618,21 +789,38 @@ class UnslothTrainer:
print("Training configuration prepared\n")
# ========== TRAINER INITIALIZATION ==========
if self.is_vlm:
self.trainer = SFTTrainer(
model=self.model,
train_dataset=dataset['dataset'],
processing_class = self.tokenizer.tokenizer,
data_collator=data_collator,
args=SFTConfig(**config_args),
)
trainer_kwargs = {
"model": self.model,
"train_dataset": dataset['dataset'],
"processing_class": self.tokenizer.tokenizer,
"data_collator": data_collator,
"args": SFTConfig(**config_args),
}
if eval_dataset is not None:
trainer_kwargs["eval_dataset"] = eval_dataset
self.trainer = SFTTrainer(**trainer_kwargs)
else:
self.trainer = SFTTrainer(
model=self.model,
tokenizer=self.tokenizer,
train_dataset=dataset['dataset'],
data_collator=data_collator,
args=SFTConfig(**config_args),
)
# For text-only training, if the tokenizer is actually a Processor
# (e.g., Gemma-3 returns a ProcessorMixin even for text), we must
# unwrap to the raw tokenizer. Otherwise Unsloth's SFTTrainer detects
# ProcessorMixin → sets _is_vlm=True → skips _prepare_dataset entirely,
# and the 'text' column never gets tokenized to 'input_ids'.
from transformers import ProcessorMixin
sft_tokenizer = self.tokenizer
if isinstance(self.tokenizer, ProcessorMixin) and hasattr(self.tokenizer, 'tokenizer'):
print(f" ⚠️ Unwrapping Processor → raw tokenizer for text-only SFTTrainer")
sft_tokenizer = self.tokenizer.tokenizer
trainer_kwargs = {
"model": self.model,
"tokenizer": sft_tokenizer,
"train_dataset": dataset['dataset'],
"data_collator": data_collator,
"args": SFTConfig(**config_args),
}
if eval_dataset is not None:
trainer_kwargs["eval_dataset"] = eval_dataset
self.trainer = SFTTrainer(**trainer_kwargs)
print("Trainer initialized\n")
# ========== TRAIN ON RESPONSES ONLY ==========
@ -679,6 +867,7 @@ class UnslothTrainer:
self.trainer,
instruction_part=instruction_part,
response_part=response_part,
num_proc=config_args.get("dataset_num_proc", max(1, os.cpu_count() // 4)),
)
print("Train on responses only configured successfully\n")
except Exception as e:
@ -706,12 +895,41 @@ class UnslothTrainer:
if logs:
# Get loss from either 'loss' or 'train_loss' key
loss_value = logs.get('loss', logs.get('train_loss', 0.0))
current_step = state.global_step
# Extract grad_norm from logs (available when gradient clipping is enabled)
grad_norm = logs.get('grad_norm', None)
# Calculate elapsed_seconds
elapsed_seconds = None
if self.trainer_instance.training_start_time is not None:
elapsed_seconds = time.time() - self.trainer_instance.training_start_time
# Calculate eta_seconds
eta_seconds = None
if elapsed_seconds is not None and current_step > 0:
total_steps = self.trainer_instance.training_progress.total_steps
if total_steps > 0:
steps_remaining = total_steps - current_step
if steps_remaining > 0:
time_per_step = elapsed_seconds / current_step
eta_seconds = time_per_step * steps_remaining
# Extract num_tokens from TRL SFTTrainer state (real counter)
# Requires include_num_input_tokens_seen=True in SFTConfig
num_tokens = getattr(state, "num_input_tokens_seen", None)
self.trainer_instance._update_progress(
step=state.global_step,
epoch=round(state.epoch, 2) if state.epoch else 0, # Round epoch to 2 decimals
step=current_step,
epoch=round(state.epoch, 2) if state.epoch else 0,
loss=loss_value,
learning_rate=logs.get('learning_rate', 0.0),
status_message="" # Clear status message so metrics show
elapsed_seconds=elapsed_seconds,
eta_seconds=eta_seconds,
grad_norm=grad_norm,
num_tokens=num_tokens,
eval_loss=logs.get('eval_loss', None),
status_message=""
)
def on_epoch_end(self, args, state, control, **kwargs):
@ -805,9 +1023,12 @@ class UnslothTrainer:
print(f"\nStopping training (save={save})...")
self.should_stop = True
self.save_on_stop = save
self.is_training = False
# Clear the status message so timer doesn't show stale status
self._update_progress(is_training=False, status_message="")
stop_msg = (
"Stopping training and saving checkpoint..."
if save
else "Cancelling training..."
)
self._update_progress(status_message=stop_msg)
# If trainer exists, try to stop it gracefully
if self.trainer:

85
studio/backend/core/training/training.py
View file

@ -4,6 +4,7 @@ Training backend for FastAPI integration
import matplotlib.pyplot as plt
from typing import Any, Generator, Tuple
import logging
import math
from .trainer import get_trainer, TrainingProgress
from utils.hardware import clear_gpu_cache
@ -28,6 +29,11 @@ class TrainingBackend:
self.loss_history = []
self.lr_history = []
self.step_history = []
self.grad_norm_history = []
self.grad_norm_step_history = []
self.eval_loss_history = []
self.eval_step_history = []
self.eval_enabled = False
self.current_theme = "light"
self.trainer.add_progress_callback(self._on_progress_update)
@ -40,6 +46,17 @@ class TrainingBackend:
self.loss_history.append(progress.loss)
self.lr_history.append(progress.learning_rate)
self.step_history.append(progress.step)
if progress.step >= 0 and progress.grad_norm is not None:
try:
grad_norm = float(progress.grad_norm)
except (TypeError, ValueError):
grad_norm = None
if grad_norm is not None and math.isfinite(grad_norm):
self.grad_norm_history.append(grad_norm)
self.grad_norm_step_history.append(progress.step)
if progress.eval_loss is not None:
self.eval_loss_history.append(progress.eval_loss)
self.eval_step_history.append(progress.step)
def start_training(self,
# Model parameters
@ -93,8 +110,13 @@ class TrainingBackend:
enable_tensorboard: bool,
tensorboard_dir: str,
# Optional: user-provided column mapping
custom_format_mapping: dict = None) -> bool:
# Optional parameters
custom_format_mapping: dict = None,
subset: str = None,
train_split: str = "train",
eval_split: str = None,
eval_steps: float = 0.01,
is_dataset_multimodal: bool = False) -> bool:
"""
Start training.
@ -133,6 +155,11 @@ class TrainingBackend:
self.loss_history = []
self.lr_history = []
self.step_history = []
self.grad_norm_history = []
self.grad_norm_step_history = []
self.eval_loss_history = []
self.eval_step_history = []
self.eval_enabled = False
import time
output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}"
@ -148,7 +175,8 @@ class TrainingBackend:
model_name=model_name,
max_seq_length=max_seq_length,
load_in_4bit=load_in_4bit if use_lora_actual else False, # Only 4bit for LoRA
hf_token=hf_token if hf_token.strip() else None
hf_token=hf_token if hf_token.strip() else None,
is_dataset_multimodal=is_dataset_multimodal,
)
if not success or self.trainer.should_stop:
@ -187,13 +215,30 @@ class TrainingBackend:
# ========== LOAD DATASET ==========
logger.info("Loading dataset...")
#breakpoint()
dataset = self.trainer.load_and_format_dataset(
dataset_result = self.trainer.load_and_format_dataset(
dataset_source=hf_dataset if hf_dataset.strip() else None,
format_type=format_type,
local_datasets=local_datasets if local_datasets else None,
custom_format_mapping=custom_format_mapping,
subset=subset,
train_split=train_split,
eval_split=eval_split,
)
# Unpack: load_and_format_dataset returns (dataset, eval_dataset)
if isinstance(dataset_result, tuple):
dataset, eval_dataset = dataset_result
else:
dataset = dataset_result
eval_dataset = None
# If user set eval_steps to 0, disable evaluation entirely
if eval_steps is not None and float(eval_steps) <= 0:
eval_dataset = None
# Track whether eval is enabled for status reporting
self.eval_enabled = eval_dataset is not None
if dataset is None or self.trainer.should_stop:
logger.error("Failed to load dataset or stopped by user")
return False
@ -213,7 +258,8 @@ class TrainingBackend:
logger.info("Starting training worker thread...")
success = self.trainer.start_training(
dataset=dataset,
#output_dir=f"./outputs/{model_name.replace('/', '_')}_{int(__import__('time').time())}",
eval_dataset=eval_dataset,
eval_steps=eval_steps,
output_dir=output_dir,
num_epochs=num_epochs,
learning_rate=lr_value,
@ -232,7 +278,7 @@ class TrainingBackend:
wandb_token=wandb_token if wandb_token.strip() else None,
enable_tensorboard=enable_tensorboard,
tensorboard_dir=tensorboard_dir,
max_seq_length=max_seq_length, # Pass through for config
max_seq_length=max_seq_length,
optim=optim,
lr_scheduler_type=lr_scheduler_type,
)
@ -323,8 +369,13 @@ class TrainingBackend:
True if training is in progress, False otherwise
"""
try:
# If user requested stop, training is no longer considered active
if self.trainer.should_stop:
training_thread = getattr(self.trainer, "training_thread", None)
if training_thread and training_thread.is_alive():
return True
# Stop requested and worker already exited => inactive.
# This allows UI to show stopped state + "Back to configuration".
if getattr(self.trainer, "should_stop", False):
return False
progress = self.trainer.get_training_progress()
@ -335,7 +386,23 @@ class TrainingBackend:
# but haven't completed or errored yet
if not is_active and not progress.is_completed and not progress.error:
status = progress.status_message or ""
if any(keyword in status.lower() for keyword in ["loading", "preparing", "training"]):
status_lower = status.lower()
if any(
keyword in status_lower
for keyword in ["cancelled", "canceled", "stopped", "completed", "ready to train"]
):
return False
if any(
keyword in status_lower
for keyword in [
"loading",
"preparing",
"training",
"configuring",
"tokenizing",
"starting",
]
):
is_active = True
return is_active
except Exception as e:

26
studio/backend/main.py
View file

@ -1,6 +1,7 @@
"""
Main FastAPI application for Unsloth UI Backend
"""
import os
import secrets
import shutil
from contextlib import asynccontextmanager
@ -23,7 +24,7 @@ from routes import (
training_router,
)
from auth import storage
from utils.hardware import detect_hardware
from utils.hardware import detect_hardware, get_device, DeviceType
import utils.hardware.hardware as _hw_module
UNSLOTH_CACHE_DIR = Path(__file__).parent / "unsloth_compiled_cache"
@ -35,6 +36,18 @@ async def lifespan(app: FastAPI):
# Detect hardware first — sets DEVICE global used everywhere
detect_hardware()
# Disable flex attention on Blackwell+ GPUs (sm_120 and above)
if get_device() == DeviceType.CUDA:
import torch
props = torch.cuda.get_device_properties(0)
sm_version = props.major * 10 + props.minor
if sm_version >= 120:
os.environ["UNSLOTH_ENABLE_FLEX_ATTENTION"] = "0"
import logging
logging.getLogger(__name__).info(
f"GPU sm_{sm_version} detected — setting UNSLOTH_FLEX_ATTENTION=0"
)
if not storage.is_initialized():
setup_token = secrets.token_urlsafe(32)
storage.save_setup_token(setup_token)
@ -126,6 +139,17 @@ async def get_system_info():
}
@app.get("/api/system/hardware")
async def get_hardware_info():
"""Return GPU name, total VRAM, and key ML package versions."""
from utils.hardware import get_gpu_summary, get_package_versions
return {
"gpu": get_gpu_summary(),
"versions": get_package_versions(),
}
# ============ Serve Frontend (Optional) ============
def setup_frontend(app: FastAPI, build_path: Path):

10
studio/backend/models/__init__.py
View file

@ -8,7 +8,12 @@ from .training import (
TrainingProgress,
)
from .models import (
CheckpointInfo,
ModelCheckpoints,
CheckpointListResponse,
ModelDetails,
LocalModelInfo,
LocalModelListResponse,
LoRAInfo,
LoRAScanResponse,
ModelListResponse,
@ -20,8 +25,6 @@ from .auth import (
AuthStatusResponse,
)
from .export import (
CheckpointInfo,
CheckpointListResponse,
LoadCheckpointRequest,
ExportStatusResponse,
ExportOperationResponse,
@ -65,6 +68,8 @@ __all__ = [
"TrainingProgress",
# Model management schemas
"ModelDetails",
"LocalModelInfo",
"LocalModelListResponse",
"LoRAInfo",
"LoRAScanResponse",
"ModelListResponse",
@ -75,6 +80,7 @@ __all__ = [
"AuthStatusResponse",
# Export schemas
"CheckpointInfo",
"ModelCheckpoints",
"CheckpointListResponse",
"LoadCheckpointRequest",
"ExportStatusResponse",

15
studio/backend/models/datasets.py
View file

@ -1,8 +1,8 @@
"""
Dataset-related Pydantic models for API requests and responses.
"""
from pydantic import BaseModel
from typing import Optional, Dict, List
from pydantic import BaseModel, model_validator
from typing import Any, Optional, Dict, List
class CheckFormatRequest(BaseModel):
@ -10,7 +10,16 @@ class CheckFormatRequest(BaseModel):
dataset_name: str # HuggingFace dataset name or local path
is_vlm: bool = False
hf_token: Optional[str] = None
split: Optional[str] = "train"
subset: Optional[str] = None
train_split: Optional[str] = "train"
@model_validator(mode="before")
@classmethod
def _compat_split(cls, values: Any) -> Any:
"""Accept legacy 'split' field as alias for 'train_split'."""
if isinstance(values, dict) and "split" in values:
values.setdefault("train_split", values.pop("split"))
return values
class CheckFormatResponse(BaseModel):

21
studio/backend/models/export.py
View file

@ -5,23 +5,6 @@ from pydantic import BaseModel, Field
from typing import List, Optional, Literal, Dict, Any
class CheckpointInfo(BaseModel):
"""Information about a discovered checkpoint directory."""
display_name: str = Field(..., description="User-friendly checkpoint name (folder name)")
path: str = Field(..., description="Full path to the checkpoint directory")
class CheckpointListResponse(BaseModel):
"""Response for listing available checkpoints in an outputs directory."""
outputs_dir: str = Field(..., description="Directory that was scanned")
checkpoints: List[CheckpointInfo] = Field(
default_factory=list,
description="List of discovered checkpoints",
)
class LoadCheckpointRequest(BaseModel):
"""Request for loading a checkpoint into the export backend."""
@ -89,6 +72,10 @@ class ExportCommonOptions(BaseModel):
False,
description="If True, create a private repository on the Hub (where applicable)",
)
base_model_id: Optional[str] = Field(
None,
description="HuggingFace model ID of the base model (for model card metadata)",
)
class ExportMergedModelRequest(ExportCommonOptions):

5
studio/backend/models/inference.py
View file

@ -30,7 +30,7 @@ class GenerateRequest(BaseModel):
system_prompt: str = Field("You are a helpful AI assistant.", description="System prompt")
temperature: float = Field(0.7, ge=0.0, le=2.0, description="Sampling temperature")
top_p: float = Field(0.9, ge=0.0, le=1.0, description="Top-p sampling")
top_k: int = Field(40, ge=1, le=100, description="Top-k sampling")
top_k: int = Field(40, ge=-1, le=100, description="Top-k sampling")
max_new_tokens: int = Field(512, ge=1, le=4096, description="Maximum tokens to generate")
repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="Repetition penalty")
image_base64: Optional[str] = Field(None, description="Base64 encoded image for vision models")
@ -128,7 +128,8 @@ class ChatCompletionRequest(BaseModel):
max_tokens: Optional[int] = Field(512, ge=1, le=4096, description="Maximum tokens to generate")
# ── Unsloth extensions (ignored by standard OpenAI clients) ──
top_k: int = Field(40, ge=1, le=100, description="[x-unsloth] Top-k sampling")
top_k: int = Field(40, ge=-1, le=100, description="[x-unsloth] Top-k sampling")
min_p: float = Field(0.0, 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")
image_base64: Optional[str] = Field(None, description="[x-unsloth] Base64-encoded image for vision models")
use_adapter: Optional[Union[bool, str]] = Field(

73
studio/backend/models/models.py
View file

@ -2,7 +2,47 @@
Pydantic schemas for Model Management API
"""
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
from typing import Optional, List, Dict, Any, Literal
class CheckpointInfo(BaseModel):
"""Information about a discovered checkpoint directory."""
display_name: str = Field(..., description="User-friendly checkpoint name (folder name)")
path: str = Field(..., description="Full path to the checkpoint directory")
loss: Optional[float] = Field(None, description="Training loss at this checkpoint")
class ModelCheckpoints(BaseModel):
"""A training run and its associated checkpoints."""
name: str = Field(..., description="Training run folder name")
checkpoints: List[CheckpointInfo] = Field(
default_factory=list,
description="List of checkpoints for this training run (final + intermediate)",
)
base_model: Optional[str] = Field(
None,
description="Base model name from adapter_config.json or config.json",
)
peft_type: Optional[str] = Field(
None,
description="PEFT type (e.g. LORA) if adapter training, None for full fine-tune",
)
lora_rank: Optional[int] = Field(
None,
description="LoRA rank (r) if applicable",
)
class CheckpointListResponse(BaseModel):
"""Response for listing available checkpoints in an outputs directory."""
outputs_dir: str = Field(..., description="Directory that was scanned")
models: List[ModelCheckpoints] = Field(
default_factory=list,
description="List of training runs with their checkpoints",
)
class ModelDetails(BaseModel):
@ -34,3 +74,34 @@ class ModelListResponse(BaseModel):
models: List[ModelDetails] = Field(default_factory=list, description="List of models")
default_models: List[str] = Field(default_factory=list, description="List of default model IDs")
class LocalModelInfo(BaseModel):
"""Discovered local model candidate."""
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(
...,
description="Discovery source",
)
model_id: Optional[str] = Field(
None,
description="HF repo id for cached models, e.g. org/model",
)
updated_at: Optional[float] = Field(
None,
description="Unix timestamp of latest observed update",
)
class LocalModelListResponse(BaseModel):
"""Response schema for listing local/cached models."""
models_dir: str = Field(..., description="Directory scanned for custom local models")
hf_cache_dir: Optional[str] = Field(
None,
description="HF cache root that was scanned",
)
models: List[LocalModelInfo] = Field(
default_factory=list,
description="Discovered local/cached models",
)

2
studio/backend/models/responses.py
View file

@ -19,6 +19,8 @@ class TrainingMetricsResponse(BaseModel):
loss_history: List[float] = Field(default_factory=list, description="Loss values per step")
lr_history: List[float] = Field(default_factory=list, description="Learning rate per step")
step_history: List[int] = Field(default_factory=list, description="Step numbers")
grad_norm_history: List[float] = Field(default_factory=list, description="Gradient norm values")
grad_norm_step_history: List[int] = Field(default_factory=list, description="Step numbers for gradient norm values")
current_loss: Optional[float] = Field(None, description="Most recent loss value")
current_lr: Optional[float] = Field(None, description="Most recent learning rate")
current_step: Optional[int] = Field(None, description="Most recent step number")

22
studio/backend/models/training.py
View file

@ -1,8 +1,8 @@
"""
Pydantic schemas for Training API
"""
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Literal
from pydantic import BaseModel, Field, model_validator
from typing import Any, Optional, List, Dict, Literal
class TrainingStartRequest(BaseModel):
@ -18,6 +18,18 @@ class TrainingStartRequest(BaseModel):
hf_dataset: Optional[str] = Field(None, description="HuggingFace dataset identifier")
local_datasets: List[str] = Field(default_factory=list, description="List of local dataset paths")
format_type: str = Field(..., description="Dataset format type")
subset: Optional[str] = None
train_split: Optional[str] = Field("train", description="Training split name")
eval_split: Optional[str] = Field(None, description="Eval split name. None = auto-detect")
eval_steps: float = Field(0.01, description="Fraction of total steps between evals (0-1)")
@model_validator(mode="before")
@classmethod
def _compat_split(cls, values: Any) -> Any:
"""Accept legacy 'split' field as alias for 'train_split'."""
if isinstance(values, dict) and "split" in values:
values.setdefault("train_split", values.pop("split"))
return values
custom_format_mapping: Optional[Dict[str, str]] = Field(
None,
description="User-provided column-to-role mapping, e.g. {'image': 'image', 'caption': 'text'} for VLM or {'instruction': 'user', 'output': 'assistant'} for LLM"
@ -53,6 +65,7 @@ class TrainingStartRequest(BaseModel):
finetune_language_layers: bool = Field(False, description="Finetune language layers")
finetune_attention_modules: bool = Field(False, description="Finetune attention modules")
finetune_mlp_modules: bool = Field(False, description="Finetune MLP modules")
is_dataset_multimodal: bool = Field(False, description="Whether the dataset contains multimodal (image) data")
# Logging parameters
enable_wandb: bool = Field(False, description="Enable Weights & Biases logging")
@ -84,13 +97,14 @@ class TrainingStatus(BaseModel):
"stopped"
] = Field(..., description="Current phase of training pipeline")
is_training_running: bool = Field(..., description="True if training loop is actively running")
eval_enabled: bool = Field(False, description="True if evaluation dataset is configured for this training run")
message: str = Field(..., description="Human-readable status message")
error: Optional[str] = Field(None, description="Error details if phase is 'error'")
details: Optional[dict] = Field(None, description="Phase-specific info, e.g. {'model_size': '8B'}")
metric_history: Optional[dict] = Field(
None,
description="Full metric history arrays for chart recovery after SSE reconnection. "
"Keys: 'steps', 'loss', 'lr' — each a list of numeric values.",
"Keys: 'steps', 'loss', 'lr', 'grad_norm', 'grad_norm_steps' — each a list of numeric values.",
)
@ -107,4 +121,4 @@ class TrainingProgress(BaseModel):
eta_seconds: Optional[float] = Field(None, description="Estimated time remaining")
grad_norm: Optional[float] = Field(None, description="L2 norm of gradients, computed before gradient clipping")
num_tokens: Optional[int] = Field(None, description="Total number of tokens processed so far")
eval_loss: Optional[float] = Field(None, description="Eval loss from the most recent evaluation step")

7
studio/backend/requirements.txt
View file

@ -1,7 +0,0 @@
fastapi>=0.100.0
uvicorn>=0.27.0
pydantic>=2.0
torch
psutil
nest-asyncio>=1.5.8

3
studio/backend/requirements/base.txt Normal file
View file

@ -0,0 +1,3 @@
# Core unsloth packages
unsloth-zoo
unsloth

13
studio/backend/requirements/extras-no-deps.txt Normal file
View file

@ -0,0 +1,13 @@
# Audio extras (installed with --no-deps --no-cache-dir)
descript-audio-codec
descript-audiotools
julius
torchcodec
snac
# TRL and related packages
trl==0.23.1
git+https://github.com/meta-pytorch/OpenEnv.git
executorch==1.0.1
torch-c-dlpack-ext
sentence_transformers==5.2.0

56
studio/backend/requirements/extras.txt Normal file
View file

@ -0,0 +1,56 @@
# OpenEnv dependencies
tomli
tomli-w
# ExecuTorch dependencies
ruamel.yaml
coremltools
expecttest
flatbuffers
hydra-core
hypothesis
kgb
parameterized
pytest<9.0
pytest-json-report
pytest-rerunfailures==15.1
pytest-xdist
# Also needed by sentence_transformers
scikit-learn==1.7.1
# Additional extras
pybind11
langid
jiwer
omegaconf
einx
pyloudnorm
openai-whisper
uroman
MeCab
loguru
flatten_dict
ffmpy
randomname
argbind
tiktoken
ftfy
importlib-resources
librosa
markdown2
matplotlib
pystoi
soundfile
tensorboard
torch-stoi
evaluate
timm
transformers-cfg
open_spiel
addict
easydict
einops
tabulate
fastmcp>=2.0.0
openai>=2.7.2
websockets>=13.0,<14

7
studio/backend/requirements/overrides.txt Normal file
View file

@ -0,0 +1,7 @@
# Torch AO overrides (installed with --force-reinstall --no-cache-dir)
torchao==0.14.0
transformers==4.57.1
pytorch_tokenizers
# Kernel packages
kernels

13
studio/backend/requirements/studio.txt Normal file
View file

@ -0,0 +1,13 @@
# Studio UI backend dependencies
typer
fastapi
uvicorn
pydantic
matplotlib
pandas
nest_asyncio
datasets==4.3.0
pyjwt
easydict
addict
gradio>=4.0.0

2
studio/backend/requirements/triton-kernels.txt Normal file
View file

@ -0,0 +1,2 @@
# Triton kernels (installed with --no-deps, from source)
triton_kernels @ git+https://github.com/triton-lang/triton.git@release/3.6.x#subdirectory=python/triton_kernels

131
studio/backend/routes/datasets.py
View file

@ -70,76 +70,139 @@ def _serialize_preview_rows(rows):
# --- Endpoints ---
# Recognized data-file extensions for the single-file fallback approach.
DATA_EXTS = (
'.parquet',
'.json', '.jsonl',
'.csv', '.tsv',
'.txt',
'.arrow',
'.tar', '.tar.gz', '.tgz',
'.gz', '.zst',
'.zip',
)
@router.post("/check-format", response_model=CheckFormatResponse)
async def check_format(request: CheckFormatRequest):
def check_format(request: CheckFormatRequest):
"""
Check if a dataset requires manual column mapping.
This is a lightweight check that loads only the first 10 rows,
runs format detection, and (if processable) returns processed
preview samples. The full dataset is re-processed at training time.
Strategy for HuggingFace datasets:
1. list_repo_files → pick the first data file → load_dataset(data_files=[…])
Avoids resolving thousands of files; typically ~2-4 s.
2. Full streaming load_dataset as a last-resort fallback.
Local files are loaded directly.
Using a plain `def` (not async) so FastAPI runs this in a thread-pool,
preventing any blocking IO from freezing the event loop.
"""
try:
from datasets import load_dataset
from itertools import islice
from datasets import Dataset, load_dataset
from utils.datasets import format_dataset
PREVIEW_SIZE = 10
logger.info(f"Checking format for dataset: {request.dataset_name}")
# Load dataset
dataset_path = Path(request.dataset_name)
total_rows = None
if dataset_path.exists():
# Local dataset
# ── Local file ──────────────────────────────────────────
if dataset_path.suffix in ['.json', '.jsonl']:
dataset = load_dataset('json', data_files=str(dataset_path), split=request.split)
dataset = load_dataset('json', data_files=str(dataset_path), split=request.train_split)
elif dataset_path.suffix == '.csv':
dataset = load_dataset('csv', data_files=str(dataset_path), split=request.split)
dataset = load_dataset('csv', data_files=str(dataset_path), split=request.train_split)
elif dataset_path.suffix == '.parquet':
dataset = load_dataset('parquet', data_files=str(dataset_path), split=request.split)
dataset = load_dataset('parquet', data_files=str(dataset_path), split=request.train_split)
else:
raise HTTPException(
status_code=400,
detail=f"Unsupported file format: {dataset_path.suffix}"
)
total_rows = len(dataset)
preview_slice = dataset.select(range(min(PREVIEW_SIZE, total_rows)))
else:
# HuggingFace dataset
load_kwargs = {"path": request.dataset_name, "split": request.split}
if request.hf_token:
load_kwargs["token"] = request.hf_token
dataset = load_dataset(**load_kwargs)
# Slice to top N rows — all detection and preview runs on this subset
total_rows = len(dataset)
preview_slice = dataset.select(range(min(PREVIEW_SIZE, total_rows)))
# ── HuggingFace dataset ─────────────────────────────────
# Tier 1: list_repo_files → load only the first data file
preview_slice = None
try:
from huggingface_hub import HfApi
api = HfApi()
repo_files = api.list_repo_files(
request.dataset_name,
repo_type="dataset",
token=request.hf_token or None,
)
data_files = [f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)]
if data_files:
first_file = data_files[0]
logger.info(f"Tier 1: loading single file {first_file}")
load_kwargs = {
"path": request.dataset_name,
"data_files": [first_file],
"split": "train",
"streaming": True,
}
if request.hf_token:
load_kwargs["token"] = request.hf_token
streamed_ds = load_dataset(**load_kwargs)
rows = list(islice(streamed_ds, PREVIEW_SIZE))
if rows:
preview_slice = Dataset.from_list(rows)
except Exception as e:
logger.warning(f"Tier 1 (single-file) failed: {e}")
if preview_slice is None:
# Tier 2: full streaming (resolves all files — slow for large repos)
logger.info("Tier 2: falling back to full streaming load_dataset")
load_kwargs = {"path": request.dataset_name, "split": request.train_split, "streaming": True}
if request.subset:
load_kwargs["name"] = request.subset
if request.hf_token:
load_kwargs["token"] = request.hf_token
streamed_ds = load_dataset(**load_kwargs)
rows = list(islice(streamed_ds, PREVIEW_SIZE))
if not rows:
raise HTTPException(
status_code=400,
detail="Dataset appears to be empty or could not be streamed"
)
preview_slice = Dataset.from_list(rows)
total_rows = None
# Run lightweight format check on the preview slice
result = check_dataset_format(preview_slice, is_vlm=request.is_vlm)
logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}")
logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_multimodal={result.get('is_multimodal', False)}")
# Generate preview samples
preview_samples = None
if not result["requires_manual_mapping"]:
# Format detected — return processed preview
try:
format_result = format_dataset(
preview_slice,
format_type="auto",
custom_format_mapping=result.get("suggested_mapping"),
num_proc=1, # Only 10 preview rows — no need for multiprocessing
)
processed = format_result["dataset"]
preview_samples = _serialize_preview_rows(processed)
except Exception as e:
logger.warning(f"Processed preview generation failed (non-fatal): {e}")
# Fall back to raw samples so frontend still has something
preview_samples = _serialize_preview_rows(preview_slice)
else:
# Format detection failed — return raw samples so user can
# see actual data and map columns in the frontend
preview_samples = _serialize_preview_rows(preview_slice)
return CheckFormatResponse(
requires_manual_mapping=result["requires_manual_mapping"],
detected_format=result["detected_format"],
@ -152,7 +215,7 @@ async def check_format(request: CheckFormatRequest):
preview_samples=preview_samples,
total_rows=total_rows,
)
except HTTPException:
raise
except Exception as e:

35
studio/backend/routes/export.py
View file

@ -26,8 +26,6 @@ except ImportError:
# Import Pydantic models
from models import (
CheckpointInfo,
CheckpointListResponse,
LoadCheckpointRequest,
ExportStatusResponse,
ExportOperationResponse,
@ -50,38 +48,6 @@ if not logger.handlers:
logger.setLevel(logging.INFO)
@router.get("/checkpoints", response_model=CheckpointListResponse)
async def list_checkpoints(
outputs_dir: str = Query(
default="./outputs",
description="Directory to scan for checkpoints",
),
current_subject: str = Depends(get_current_subject),
):
"""
List available checkpoints in the outputs directory.
Wraps ExportBackend.scan_checkpoints.
"""
try:
backend = get_export_backend()
raw_checkpoints = backend.scan_checkpoints(outputs_dir=outputs_dir)
checkpoints = [
CheckpointInfo(display_name=display_name, path=path)
for display_name, path in raw_checkpoints
]
return CheckpointListResponse(
outputs_dir=outputs_dir,
checkpoints=checkpoints,
)
except Exception as e:
logger.error(f"Error listing checkpoints: {e}", exc_info=True)
raise HTTPException(
status_code=500,
detail=f"Failed to list checkpoints: {str(e)}",
)
@router.post("/load-checkpoint", response_model=ExportOperationResponse)
@ -224,6 +190,7 @@ async def export_base_model(
repo_id=request.repo_id,
hf_token=request.hf_token,
private=request.private,
base_model_id=request.base_model_id,
)
if not success:

44
studio/backend/routes/inference.py
View file

@ -5,11 +5,13 @@ import sys
import time
import uuid
from pathlib import Path
from fastapi import APIRouter, HTTPException
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse, JSONResponse
from typing import Optional
import json
import logging
import asyncio
import threading
@ -304,7 +306,7 @@ def _extract_content_parts(
@router.post("/chat/completions")
async def openai_chat_completions(request: ChatCompletionRequest):
async def openai_chat_completions(payload: ChatCompletionRequest, request: Request):
"""
OpenAI-compatible chat completions endpoint.
@ -324,7 +326,7 @@ async def openai_chat_completions(request: ChatCompletionRequest):
# ── Parse messages (handles multimodal content parts) ─────
system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(
request.messages
payload.messages
)
# If no non-system messages were provided, error out
@ -336,7 +338,7 @@ async def openai_chat_completions(request: ChatCompletionRequest):
# ── Decode image (from content parts OR legacy field) ─────
# Content-part images take priority; fall back to legacy field
image_b64 = extracted_image_b64 or request.image_base64
image_b64 = extracted_image_b64 or payload.image_base64
image = None
if image_b64:
@ -366,31 +368,36 @@ async def openai_chat_completions(request: ChatCompletionRequest):
messages=chat_messages,
system_prompt=system_prompt,
image=image,
temperature=request.temperature,
top_p=request.top_p,
top_k=request.top_k,
max_new_tokens=request.max_tokens or 512,
repetition_penalty=request.repetition_penalty,
temperature=payload.temperature,
top_p=payload.top_p,
top_k=payload.top_k,
min_p=payload.min_p,
max_new_tokens=payload.max_tokens or 512,
repetition_penalty=payload.repetition_penalty,
)
# ── Choose generation path (adapter-controlled or standard) ──
if request.use_adapter is not None:
cancel_event = threading.Event()
if payload.use_adapter is not None:
# Compare mode: toggle adapter state atomically with generation
def generate():
return backend.generate_with_adapter_control(
use_adapter=request.use_adapter, **gen_kwargs
use_adapter=payload.use_adapter,
cancel_event=cancel_event,
**gen_kwargs,
)
else:
# Standard path: no adapter toggling
def generate():
return backend.generate_chat_response(**gen_kwargs)
return backend.generate_chat_response(cancel_event=cancel_event, **gen_kwargs)
model_name = backend.active_model_name or request.model
model_name = backend.active_model_name or payload.model
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
# ── Streaming response ────────────────────────────────────────
if request.stream:
if payload.stream:
async def stream_chunks():
try:
# First chunk: send the role
@ -409,6 +416,10 @@ async def openai_chat_completions(request: ChatCompletionRequest):
# text, so we diff to get incremental deltas.
prev_text = ""
for cumulative in generate():
if await request.is_disconnected():
cancel_event.set()
backend.reset_generation_state()
return
new_text = cumulative[len(prev_text):]
prev_text = cumulative
if not new_text:
@ -437,6 +448,10 @@ async def openai_chat_completions(request: ChatCompletionRequest):
yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
backend.reset_generation_state()
raise
except Exception as e:
backend.reset_generation_state()
logger.error(f"Error during OpenAI streaming: {e}", exc_info=True)
@ -477,4 +492,3 @@ async def openai_chat_completions(request: ChatCompletionRequest):
backend.reset_generation_state()
logger.error(f"Error during OpenAI completion: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))

161
studio/backend/routes/models.py
View file

@ -21,6 +21,7 @@ try:
load_model_defaults,
get_base_model_from_lora,
is_vision_model,
scan_checkpoints,
ModelConfig,
)
from core.inference import get_inference_backend
@ -34,11 +35,17 @@ except ImportError:
load_model_defaults,
get_base_model_from_lora,
is_vision_model,
scan_checkpoints,
ModelConfig,
)
from core.inference import get_inference_backend
from models import (
CheckpointInfo,
CheckpointListResponse,
LocalModelInfo,
LocalModelListResponse,
ModelCheckpoints,
ModelDetails,
LoRAScanResponse,
LoRAInfo,
@ -59,6 +66,116 @@ if not logger.handlers:
logger.setLevel(logging.INFO)
def _resolve_hf_cache_dir() -> Path:
"""Resolve local HF cache root used by hub downloads."""
try:
from huggingface_hub.constants import HF_HUB_CACHE
return Path(HF_HUB_CACHE)
except Exception:
return Path.home() / ".cache" / "huggingface" / "hub"
def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
if not models_dir.exists() or not models_dir.is_dir():
return []
found: List[LocalModelInfo] = []
for child in models_dir.iterdir():
if not child.is_dir():
continue
has_model_files = (
(child / "config.json").exists()
or (child / "adapter_config.json").exists()
or any(child.glob("*.safetensors"))
or any(child.glob("*.bin"))
)
if not has_model_files:
continue
try:
updated_at = child.stat().st_mtime
except OSError:
updated_at = None
found.append(
LocalModelInfo(
id=str(child),
display_name=child.name,
path=str(child),
source="models_dir",
updated_at=updated_at,
),
)
return found
def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
if not cache_dir.exists() or not cache_dir.is_dir():
return []
found: List[LocalModelInfo] = []
for repo_dir in cache_dir.glob("models--*"):
if not repo_dir.is_dir():
continue
repo_name = repo_dir.name[len("models--"):]
if not repo_name:
continue
model_id = repo_name.replace("--", "/")
try:
updated_at = repo_dir.stat().st_mtime
except OSError:
updated_at = None
found.append(
LocalModelInfo(
id=model_id,
model_id=model_id,
display_name=model_id.split("/")[-1],
path=str(repo_dir),
source="hf_cache",
updated_at=updated_at,
),
)
return found
@router.get("/local", response_model=LocalModelListResponse)
async def list_local_models(
models_dir: str = Query(default="./models", description="Directory to scan for local model folders"),
current_subject: str = Depends(get_current_subject),
):
"""
List local model candidates from custom models dir and HF cache.
"""
try:
models_root = Path(models_dir).expanduser().resolve()
hf_cache_dir = _resolve_hf_cache_dir()
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
if model.id not in deduped:
deduped[model.id] = model
models = sorted(
deduped.values(),
key=lambda item: (item.updated_at or 0),
reverse=True,
)
return LocalModelListResponse(
models_dir=str(models_root),
hf_cache_dir=str(hf_cache_dir),
models=models,
)
except Exception as e:
logger.error(f"Error listing local models: {e}", exc_info=True)
raise HTTPException(
status_code=500,
detail=f"Failed to list local models: {str(e)}",
)
@router.get("/list")
@ -131,6 +248,7 @@ async def get_model_config(
This endpoint wraps the backend load_model_defaults function.
"""
try:
logger.info(f"Getting model config for: {model_name}")
# Load model defaults from backend
config_dict = load_model_defaults(model_name)
@ -150,6 +268,7 @@ async def get_model_config(
# If ModelConfig creation fails, use defaults
pass
logger.info(f"Model config result for {model_name}: is_vision={is_vision}, is_lora={is_lora}, base_model={base_model}")
return ModelDetails(
id=model_name,
model_name=model_name,
@ -252,8 +371,10 @@ async def check_vision_model(
This endpoint wraps the backend is_vision_model function.
"""
try:
logger.info(f"Checking if vision model: {model_name}")
is_vision = is_vision_model(model_name)
logger.info(f"Vision check result for {model_name}: is_vision={is_vision}")
return VisionCheckResponse(
model_name=model_name,
is_vision=is_vision,
@ -266,3 +387,43 @@ async def check_vision_model(
detail=f"Failed to check vision model: {str(e)}"
)
@router.get("/checkpoints", response_model=CheckpointListResponse)
async def list_checkpoints(
outputs_dir: str = Query(
default="./outputs",
description="Directory to scan for checkpoints",
),
current_subject: str = Depends(get_current_subject),
):
"""
List available checkpoints in the outputs directory.
Scans the outputs folder for training runs and their checkpoints.
"""
try:
raw_models = scan_checkpoints(outputs_dir=outputs_dir)
models = [
ModelCheckpoints(
name=model_name,
checkpoints=[
CheckpointInfo(display_name=display_name, path=path, loss=loss)
for display_name, path, loss in checkpoints
],
base_model=metadata.get("base_model"),
peft_type=metadata.get("peft_type"),
lora_rank=metadata.get("lora_rank"),
)
for model_name, checkpoints, metadata in raw_models
]
return CheckpointListResponse(
outputs_dir=outputs_dir,
models=models,
)
except Exception as e:
logger.error(f"Error listing checkpoints: {e}", exc_info=True)
raise HTTPException(
status_code=500,
detail=f"Failed to list checkpoints: {str(e)}",
)

131
studio/backend/routes/training.py
View file

@ -5,7 +5,7 @@ import sys
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse
from typing import Dict, Optional
from typing import Dict, Optional, Any
import logging
import asyncio
from datetime import datetime
@ -56,6 +56,22 @@ if not logger.handlers:
logger.setLevel(logging.INFO)
@router.get("/hardware")
async def get_hardware_utilization(
current_subject: str = Depends(get_current_subject),
):
"""
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.
"""
from utils.hardware import get_gpu_utilization
return get_gpu_utilization()
@router.post("/start")
async def start_training(
request: TrainingStartRequest,
@ -129,6 +145,10 @@ async def start_training(
"hf_dataset": request.hf_dataset or "",
"local_datasets": request.local_datasets,
"format_type": request.format_type,
"subset": request.subset,
"train_split": request.train_split,
"eval_split": request.eval_split,
"eval_steps": request.eval_steps,
"custom_format_mapping": request.custom_format_mapping,
"num_epochs": request.num_epochs,
"learning_rate": request.learning_rate,
@ -158,6 +178,7 @@ async def start_training(
"finetune_language_layers": request.finetune_language_layers,
"finetune_attention_modules": request.finetune_attention_modules,
"finetune_mlp_modules": request.finetune_mlp_modules,
"is_dataset_multimodal": request.is_dataset_multimodal,
"enable_wandb": request.enable_wandb,
"wandb_token": request.wandb_token or "",
"wandb_project": request.wandb_project or "",
@ -267,21 +288,31 @@ async def stop_training(
"""
try:
backend = get_training_backend()
if not backend.is_training_active():
trainer_thread = getattr(getattr(backend, "trainer", None), "training_thread", None)
thread_alive = bool(trainer_thread and trainer_thread.is_alive())
is_active = backend.is_training_active()
logger.info(
"Stop requested: save=%s is_active=%s thread_alive=%s should_stop=%s",
body.save,
is_active,
thread_alive,
getattr(getattr(backend, "trainer", None), "should_stop", None),
)
if not is_active and not thread_alive:
return TrainingStopResponse(
status="idle",
message="No training job is currently running"
)
# Call backend stop method
backend.stop_training(save=body.save)
return TrainingStopResponse(
status="stopped",
message="Training job stopped successfully"
message="Stop requested. Training will stop at the next safe step."
)
except Exception as e:
logger.error(f"Error stopping training: {e}", exc_info=True)
raise HTTPException(
@ -299,12 +330,33 @@ async def reset_training(
"""
try:
backend = get_training_backend()
trainer_thread = getattr(getattr(backend, "trainer", None), "training_thread", None)
thread_alive = bool(trainer_thread and trainer_thread.is_alive())
is_active = backend.is_training_active()
if is_active or thread_alive:
logger.warning(
"Rejected reset while training active: is_active=%s thread_alive=%s should_stop=%s",
is_active,
thread_alive,
getattr(getattr(backend, "trainer", None), "should_stop", None),
)
raise HTTPException(
status_code=409,
detail="Training is still running. Stop training and wait for it to finish before resetting.",
)
logger.info("Reset training state: clearing runtime + metric history")
backend.trainer.should_stop = False
backend.trainer.training_progress = backend.trainer.training_progress.__class__()
backend.loss_history = []
backend.lr_history = []
backend.step_history = []
backend.grad_norm_history = []
backend.grad_norm_step_history = []
return {"status": "ok"}
except HTTPException:
raise
except Exception as e:
logger.error(f"Error resetting training: {e}", exc_info=True)
raise HTTPException(
@ -387,12 +439,17 @@ async def get_training_status(
"steps": list(backend.step_history),
"loss": list(backend.loss_history),
"lr": list(backend.lr_history),
"grad_norm": list(getattr(backend, "grad_norm_history", [])),
"grad_norm_steps": list(getattr(backend, "grad_norm_step_history", [])),
"eval_loss": list(backend.eval_loss_history),
"eval_steps": list(backend.eval_step_history),
}
return TrainingStatus(
job_id=job_id,
phase=phase,
is_training_running=is_active,
eval_enabled=backend.eval_enabled,
message=status_message,
error=error_message,
details=details,
@ -421,6 +478,8 @@ async def get_training_metrics(
loss_history = backend.loss_history
lr_history = backend.lr_history
step_history = backend.step_history
grad_norm_history = getattr(backend, "grad_norm_history", [])
grad_norm_step_history = getattr(backend, "grad_norm_step_history", [])
# Get current values
current_loss = loss_history[-1] if loss_history else None
@ -431,6 +490,8 @@ async def get_training_metrics(
loss_history=loss_history,
lr_history=lr_history,
step_history=step_history,
grad_norm_history=grad_norm_history,
grad_norm_step_history=grad_norm_step_history,
current_loss=current_loss,
current_lr=current_lr,
current_step=current_step,
@ -480,6 +541,9 @@ async def stream_training_progress(
learning_rate: float,
total_steps: int,
epoch: Optional[float] = None,
progress: Optional[Any] = None,
grad_norm_override: Optional[float] = None,
eval_loss_override: Optional[float] = None,
) -> TrainingProgress:
total = max(total_steps, 0)
if step < 0 or total == 0:
@ -489,6 +553,17 @@ async def stream_training_progress(
float(step) / float(total) * 100.0 if total > 0 else 0.0
)
# Get actual values from progress object if available
elapsed_seconds = getattr(progress, 'elapsed_seconds', None) if progress else None
eta_seconds = getattr(progress, 'eta_seconds', None) if progress else None
grad_norm = grad_norm_override
if grad_norm is None and progress:
grad_norm = getattr(progress, 'grad_norm', None)
num_tokens = getattr(progress, 'num_tokens', None) if progress else None
eval_loss = eval_loss_override
if eval_loss is None and progress:
eval_loss = getattr(progress, 'eval_loss', None)
return TrainingProgress(
job_id=job_id,
step=step,
@ -497,10 +572,11 @@ async def stream_training_progress(
learning_rate=learning_rate,
progress_percent=progress_percent,
epoch=epoch,
elapsed_seconds=None,
eta_seconds=None,
grad_norm=None,
num_tokens=None,
elapsed_seconds=elapsed_seconds,
eta_seconds=eta_seconds,
grad_norm=grad_norm,
num_tokens=num_tokens,
eval_loss=eval_loss,
)
def format_sse(
@ -525,6 +601,13 @@ async def stream_training_progress(
# ── Replay missed steps on reconnect ─────────────────────
if resume_from_step is not None and backend.step_history:
replayed = 0
grad_norm_by_step = {
step_val: grad_val
for step_val, grad_val in zip(
getattr(backend, "grad_norm_step_history", []),
getattr(backend, "grad_norm_history", []),
)
}
for i, step_val in enumerate(backend.step_history):
if step_val > resume_from_step:
loss_val = backend.loss_history[i] if i < len(backend.loss_history) else 0.0
@ -534,7 +617,15 @@ async def stream_training_progress(
)
total_replay = getattr(tp_replay, "total_steps", step_val) if tp_replay else step_val
epoch_replay = getattr(tp_replay, "epoch", None) if tp_replay else None
payload = build_progress(step_val, loss_val, lr_val, total_replay, epoch_replay)
payload = build_progress(
step_val,
loss_val,
lr_val,
total_replay,
epoch_replay,
progress=tp_replay,
grad_norm_override=grad_norm_by_step.get(step_val),
)
yield format_sse(payload.model_dump_json(), event="progress", event_id=step_val)
replayed += 1
if replayed:
@ -553,6 +644,7 @@ async def stream_training_progress(
learning_rate=0.0,
total_steps=initial_total_steps,
epoch=initial_epoch,
progress=tp,
)
yield format_sse(initial_progress.model_dump_json(), event="progress", event_id=0)
@ -566,11 +658,11 @@ async def stream_training_progress(
getattr(tp, "total_steps", final_step) if tp else final_step
)
final_epoch = getattr(tp, "epoch", None) if tp else None
payload = build_progress(final_step, final_loss, final_lr, final_total_steps, final_epoch)
payload = build_progress(final_step, final_loss, final_lr, final_total_steps, final_epoch, progress=tp)
yield format_sse(payload.model_dump_json(), event="complete", event_id=final_step)
else:
yield format_sse(
build_progress(-1, 0.0, 0.0, 0).model_dump_json(),
build_progress(-1, 0.0, 0.0, 0, progress=tp).model_dump_json(),
event="complete",
event_id=0,
)
@ -605,6 +697,7 @@ async def stream_training_progress(
current_lr,
current_total_steps,
current_epoch,
progress=tp_inner,
)
yield format_sse(
progress_payload.model_dump_json(),
@ -623,6 +716,7 @@ async def stream_training_progress(
current_lr,
current_total_steps,
current_epoch,
progress=tp_inner,
)
yield format_sse(
heartbeat_payload.model_dump_json(),
@ -644,7 +738,7 @@ async def stream_training_progress(
if tp_prep else 0
)
preparing_payload = build_progress(
0, 0.0, 0.0, prep_total,
0, 0.0, 0.0, prep_total, progress=tp_prep,
)
yield format_sse(
preparing_payload.model_dump_json(),
@ -655,7 +749,8 @@ async def stream_training_progress(
# Timeout check
if no_update_count > max_no_updates:
logger.warning("Progress stream timeout - no updates received")
timeout_payload = build_progress(last_step, 0.0, 0.0, 0)
tp_timeout = getattr(getattr(backend, "trainer", None), "training_progress", None)
timeout_payload = build_progress(last_step, 0.0, 0.0, 0, progress=tp_timeout)
yield format_sse(
timeout_payload.model_dump_json(),
event="error",
@ -667,7 +762,8 @@ async def stream_training_progress(
except Exception as e:
logger.error(f"Error in progress stream: {e}", exc_info=True)
error_payload = build_progress(0, 0.0, 0.0, 0)
tp_error = getattr(getattr(backend, "trainer", None), "training_progress", None)
error_payload = build_progress(0, 0.0, 0.0, 0, progress=tp_error)
yield format_sse(
error_payload.model_dump_json(),
event="error",
@ -692,6 +788,7 @@ async def stream_training_progress(
final_lr,
final_total_steps,
final_epoch,
progress=final_tp,
)
yield format_sse(
final_payload.model_dump_json(),

33
studio/backend/utils/datasets/chat_templates.py
View file

@ -60,7 +60,24 @@ def get_tokenizer_chat_template(tokenizer, model_name):
print(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}")
print(f" Falling back to tokenizer's default chat template")
else:
print(f"📝 Using tokenizer's default chat template (no Unsloth template match)")
# Check if tokenizer actually has a chat_template set
has_chat_template = (
hasattr(tokenizer, 'chat_template')
and tokenizer.chat_template is not None
)
if has_chat_template:
print(f"📝 Using tokenizer's own chat template (no Unsloth template match)")
else:
# Base model with no chat template — apply default ChatML
print(f"📝 No chat template found — applying default ChatML template (base model)")
try:
tokenizer = get_chat_template(
tokenizer,
chat_template="chatml",
)
except Exception as e:
print(f"⚠️ Failed to apply default ChatML template: {e}")
print(f" Falling back to tokenizer as-is")
return tokenizer
@ -227,6 +244,16 @@ def apply_chat_template_to_dataset(
# ALPACA FORMAT
if final_format == "alpaca":
# Set alpaca chat template on tokenizer for saving (if not already set)
# This ensures the template is saved with the model for inference
if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template):
try:
from unsloth.chat_templates import get_chat_template
tokenizer = get_chat_template(tokenizer, chat_template="alpaca")
print(f"📝 Set alpaca chat template on tokenizer for model saving")
except Exception as e:
print(f"⚠️ Could not set alpaca template on tokenizer: {e}")
# Use custom template if provided
def _format_alpaca_custom(examples):
texts = []
@ -258,7 +285,7 @@ def apply_chat_template_to_dataset(
if not isinstance(dataset, IterableDataset):
from multiprocessing import cpu_count
if num_proc is None or type(num_proc) is not int:
num_proc = cpu_count()
num_proc = max(1, cpu_count() // 3)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Applying template to Alpaca format"
@ -322,7 +349,7 @@ def apply_chat_template_to_dataset(
if not isinstance(dataset, IterableDataset):
from multiprocessing import cpu_count
if num_proc is None or type(num_proc) is not int:
num_proc = cpu_count()
num_proc = max(1, cpu_count() // 3)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = f"Applying chat template to {final_format}"

6
studio/backend/utils/datasets/format_conversion.py
View file

@ -110,7 +110,7 @@ def standardize_chat_format(
from multiprocessing import cpu_count
if num_proc is None or type(num_proc) is not int:
num_proc = cpu_count()
num_proc = max(1, cpu_count() // 3)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Standardizing chat format"
@ -176,7 +176,7 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
from multiprocessing import cpu_count
if num_proc is None or type(num_proc) is not int:
num_proc = cpu_count()
num_proc = max(1, cpu_count() // 3)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Converting ChatML to Alpaca format"
@ -224,7 +224,7 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
from multiprocessing import cpu_count
if num_proc is None or type(num_proc) is not int:
num_proc = cpu_count()
num_proc = max(1, cpu_count() // 3)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Converting Alpaca to ChatML format"

71
studio/backend/utils/datasets/format_detection.py
View file

@ -328,6 +328,11 @@ def detect_multimodal_dataset(dataset):
"""
Detects if dataset contains multimodal data (images/vision).
Two-pass approach:
1. Column-name heuristic (fast): checks for keywords like 'image', 'img', 'pixel'.
2. Value-type inspection (reliable): checks if actual values are PIL Images,
bytes with image headers, or HF Image-feature dicts.
Returns:
dict: {
"is_multimodal": bool,
@ -339,11 +344,16 @@ def detect_multimodal_dataset(dataset):
column_names = list(sample.keys())
# Keywords that indicate multimodal/image data
multimodal_keywords = ['image', 'img', 'pixel']
multimodal_keywords = [
'image', 'img', 'pixel',
'jpg', 'jpeg', 'png', 'webp', 'bmp', 'gif', 'tiff', 'svg',
'photo', 'pic', 'picture', 'visual',
]
multimodal_columns = []
modality_types = set()
# ── Pass 1: column-name heuristic ───────────────────────
for col_name in column_names:
col_lower = col_name.lower()
@ -353,6 +363,17 @@ def detect_multimodal_dataset(dataset):
modality_types.add(keyword)
break # Don't check other keywords for this column
# ── Pass 2: inspect actual values ───────────────────────
# Catches columns with non-obvious names (e.g. "jpg", "photo", "pic")
already_detected = set(multimodal_columns)
for col_name in column_names:
if col_name in already_detected:
continue
value = sample[col_name]
if _is_image_value(value):
multimodal_columns.append(col_name)
modality_types.add("image")
return {
"is_multimodal": len(multimodal_columns) > 0,
"multimodal_columns": multimodal_columns,
@ -360,6 +381,54 @@ def detect_multimodal_dataset(dataset):
}
def _is_image_value(value) -> bool:
"""Check if a single sample value looks like image data."""
if value is None:
return False
# PIL Image instance
try:
from PIL.Image import Image as PILImage
if isinstance(value, PILImage):
return True
except ImportError:
pass
# HF datasets Image feature stores decoded images as PIL or dicts with
# {"bytes": b"...", "path": "..."} when not yet decoded
if isinstance(value, dict):
if "bytes" in value and "path" in value:
return True
# Raw bytes with a known image magic header
if isinstance(value, (bytes, bytearray)):
return _has_image_header(value)
return False
def _has_image_header(data: bytes) -> bool:
"""Quick magic-byte check for common image formats."""
if len(data) < 4:
return False
# JPEG
if data[:2] == b'\xff\xd8':
return True
# PNG
if data[:4] == b'\x89PNG':
return True
# GIF
if data[:3] == b'GIF':
return True
# WebP
if data[:4] == b'RIFF' and len(data) >= 12 and data[8:12] == b'WEBP':
return True
# BMP
if data[:2] == b'BM':
return True
return False
def detect_vlm_dataset_structure(dataset):
"""
Detects if VLM dataset is:

6
studio/backend/utils/hardware/__init__.py
View file

@ -10,6 +10,9 @@ from .hardware import (
clear_gpu_cache,
get_gpu_memory_info,
log_gpu_memory,
get_gpu_summary,
get_package_versions,
get_gpu_utilization,
)
__all__ = [
@ -21,4 +24,7 @@ __all__ = [
'clear_gpu_cache',
'get_gpu_memory_info',
'log_gpu_memory',
'get_gpu_summary',
'get_package_versions',
'get_gpu_utilization',
]

179
studio/backend/utils/hardware/hardware.py
View file

@ -120,6 +120,7 @@ def clear_gpu_cache():
if device == DeviceType.CUDA:
import torch
torch.cuda.synchronize()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
elif device == DeviceType.MLX:
@ -206,3 +207,181 @@ def log_gpu_memory(context: str):
)
else:
logger.info(f"GPU Memory [{context}]: No GPU available (CPU-only)")
# ========== GPU Summary & Package Versions ==========
def get_gpu_summary() -> Dict[str, Any]:
"""
Return a compact summary of the primary GPU.
Returns dict with keys:
gpu_name – e.g. "NVIDIA L4" (or None)
vram_total_gb – e.g. 22.17 (or None)
"""
mem = get_gpu_memory_info()
if mem.get("available"):
return {
"gpu_name": mem.get("device_name"),
"vram_total_gb": round(mem.get("total_gb", 0), 2),
}
return {"gpu_name": None, "vram_total_gb": None}
def get_package_versions() -> Dict[str, Optional[str]]:
"""
Return the installed versions of key ML packages.
Uses importlib.metadata (stdlib) so no subprocess is needed.
CUDA version comes from torch.version.cuda.
Returns dict with keys: unsloth, torch, transformers, cuda.
Missing packages yield None.
"""
from importlib.metadata import version as pkg_version, PackageNotFoundError
packages = ("unsloth", "torch", "transformers")
versions: Dict[str, Optional[str]] = {}
for name in packages:
try:
versions[name] = pkg_version(name)
except PackageNotFoundError:
versions[name] = None
# CUDA toolkit version bundled with torch
try:
import torch
versions["cuda"] = getattr(torch.version, "cuda", None)
except Exception:
versions["cuda"] = None
return versions
# ========== Live GPU Utilization (nvidia-smi) ==========
def get_gpu_utilization() -> Dict[str, Any]:
"""
Return a live snapshot of GPU utilization via ``nvidia-smi``.
Designed to be polled by the frontend during training (not streaming).
Uses ``nvidia-smi --query-gpu`` which is the most accurate source for
utilization %, temperature, and power draw – stats that PyTorch does
not expose.
Returns dict with keys:
available – bool, whether stats could be retrieved
gpu_utilization_pct – GPU core utilization %
temperature_c – GPU temperature in °C
vram_used_gb – VRAM currently used (GiB)
vram_total_gb – VRAM total (GiB)
vram_utilization_pct – VRAM used / total * 100
power_draw_w – current power draw (W)
power_limit_w – power limit (W)
power_utilization_pct – power draw / limit * 100
"""
device = get_device()
if device != DeviceType.CUDA:
return {"available": False, "backend": device.value}
def _parse_smi_value(raw: str):
"""Parse a single nvidia-smi CSV value. Returns float or None for [N/A]."""
raw = raw.strip()
if not raw or raw == "[N/A]":
return None
try:
return float(raw)
except (ValueError, TypeError):
return None
# ── nvidia-smi (most complete source) ───────────────────────
smi_data = {}
try:
import subprocess
result = subprocess.run(
[
"nvidia-smi",
"--query-gpu=utilization.gpu,temperature.gpu,"
"memory.used,memory.total,power.draw,power.limit",
"--format=csv,noheader,nounits",
],
capture_output=True,
text=True,
timeout=5,
)
if result.returncode == 0 and result.stdout.strip():
# nvidia-smi outputs one line per GPU; take GPU 0
first_line = result.stdout.strip().splitlines()[0]
parts = [p.strip() for p in first_line.split(",")]
if len(parts) >= 6:
smi_data = {
"gpu_util": _parse_smi_value(parts[0]),
"temp": _parse_smi_value(parts[1]),
"vram_used_mb": _parse_smi_value(parts[2]),
"vram_total_mb": _parse_smi_value(parts[3]),
"power_draw": _parse_smi_value(parts[4]),
"power_limit": _parse_smi_value(parts[5]),
}
except FileNotFoundError:
logger.debug("nvidia-smi not found, falling back to torch.cuda")
except Exception as e:
logger.warning(f"nvidia-smi query failed: {e}")
# ── Backfill VRAM from torch.cuda if nvidia-smi returned [N/A] ──
vram_used_mb = smi_data.get("vram_used_mb")
vram_total_mb = smi_data.get("vram_total_mb")
if vram_used_mb is None or vram_total_mb is None:
try:
import torch
idx = torch.cuda.current_device()
props = torch.cuda.get_device_properties(idx)
if vram_total_mb is None:
vram_total_mb = props.total_memory / (1024**2) # bytes → MiB
if vram_used_mb is None:
vram_used_mb = torch.cuda.memory_allocated(idx) / (1024**2)
except Exception as e:
logger.debug(f"torch.cuda VRAM backfill failed: {e}")
# ── Build response ──────────────────────────────────────────
gpu_util = smi_data.get("gpu_util")
temp = smi_data.get("temp")
power_draw = smi_data.get("power_draw")
power_limit = smi_data.get("power_limit")
vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None
vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
vram_pct = (
round((vram_used_mb / vram_total_mb) * 100, 1)
if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0
else None
)
power_pct = (
round((power_draw / power_limit) * 100, 1)
if power_draw is not None and power_limit and power_limit > 0
else None
)
# If we got at least something useful, report available
has_any = any(v is not None for v in [gpu_util, temp, vram_used_gb, power_draw])
if not has_any:
return {"available": False, "backend": device.value}
return {
"available": True,
"backend": device.value,
"gpu_utilization_pct": gpu_util,
"temperature_c": temp,
"vram_used_gb": vram_used_gb,
"vram_total_gb": vram_total_gb,
"vram_utilization_pct": vram_pct,
"power_draw_w": power_draw,
"power_limit_w": power_limit,
"power_utilization_pct": power_pct,
}

2
studio/backend/utils/models/__init__.py
View file

@ -11,6 +11,7 @@ from .model_config import (
MODEL_NAME_MAPPING,
UI_STATUS_INDICATORS,
)
from .checkpoints import scan_checkpoints
__all__ = [
'ModelConfig',
@ -21,4 +22,5 @@ __all__ = [
'load_model_config',
'MODEL_NAME_MAPPING',
'UI_STATUS_INDICATORS',
'scan_checkpoints',
]

120
studio/backend/utils/models/checkpoints.py Normal file
View file

@ -0,0 +1,120 @@
"""
Checkpoint scanning utilities for discovering training runs and their checkpoints.
"""
import json
import logging
from pathlib import Path
from typing import List, Optional, Tuple
logger = logging.getLogger(__name__)
def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]:
"""
Read the training loss from a checkpoint's trainer_state.json.
Returns the loss from the last log_history entry, or None if unavailable.
"""
trainer_state = checkpoint_path / "trainer_state.json"
if not trainer_state.exists():
return None
try:
with open(trainer_state) as f:
state = json.load(f)
log_history = state.get("log_history", [])
if log_history:
return log_history[-1].get("loss")
except Exception as e:
logger.debug(f"Could not read loss from {trainer_state}: {e}")
return None
def scan_checkpoints(
outputs_dir: str = "./outputs",
) -> List[Tuple[str, List[Tuple[str, str, Optional[float]]], dict]]:
"""
Scan outputs folder for training runs and their checkpoints.
Returns:
List of tuples: [(model_name, [(display_name, checkpoint_path, loss), ...], metadata), ...]
metadata keys: base_model, peft_type, lora_rank (all optional)
The first entry in each checkpoint list is the main adapter; its loss is
set to the loss of the last (highest-step) intermediate checkpoint.
"""
models = []
outputs_path = Path(outputs_dir)
if not outputs_path.exists():
logger.warning(f"Outputs directory not found: {outputs_dir}")
return models
try:
for item in outputs_path.iterdir():
if not item.is_dir():
continue
config_file = item / "config.json"
adapter_config = item / "adapter_config.json"
if not (config_file.exists() or adapter_config.exists()):
continue
# Extract training metadata from adapter_config.json / config.json
metadata: dict = {}
try:
if adapter_config.exists():
cfg = json.loads(adapter_config.read_text())
metadata["base_model"] = cfg.get("base_model_name_or_path")
metadata["peft_type"] = cfg.get("peft_type")
metadata["lora_rank"] = cfg.get("r")
elif config_file.exists():
cfg = json.loads(config_file.read_text())
metadata["base_model"] = cfg.get("_name_or_path")
except Exception:
pass
# Fallback: extract base model name from folder name
# e.g. "unsloth_Llama-3.2-3B-Instruct_1771227800" → "unsloth/Llama-3.2-3B-Instruct"
if not metadata.get("base_model"):
parts = item.name.rsplit("_", 1)
if len(parts) == 2 and parts[1].isdigit():
name_part = parts[0]
idx = name_part.find("_")
if idx > 0:
metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1:]
else:
metadata["base_model"] = name_part
# This is a valid training run
checkpoints = []
# Placeholder for the main adapter — loss filled from last checkpoint below
checkpoints.append((item.name, str(item), None))
# Scan for intermediate checkpoints (checkpoint-N subdirs)
for sub in sorted(item.iterdir()):
if not sub.is_dir() or not sub.name.startswith("checkpoint-"):
continue
sub_config = sub / "config.json"
sub_adapter = sub / "adapter_config.json"
if sub_config.exists() or sub_adapter.exists():
loss = _read_checkpoint_loss(sub)
checkpoints.append((sub.name, str(sub), loss))
# Assign the last checkpoint's loss to the main adapter entry
if len(checkpoints) > 1:
last_checkpoint_loss = checkpoints[-1][2]
checkpoints[0] = (checkpoints[0][0], checkpoints[0][1], last_checkpoint_loss)
models.append((item.name, checkpoints, metadata))
logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)")
# Sort by modification time (newest first)
models.sort(key=lambda x: Path(x[1][0][1]).stat().st_mtime, reverse=True)
logger.info(f"Found {len(models)} training runs in {outputs_dir}")
return models
except Exception as e:
logger.error(f"Error scanning checkpoints: {e}")
return []

Some files were not shown because too many files have changed in this diff Show more