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

Merge branch 'main' into feature/chat-api

Browse source
This commit is contained in:
Lee Jackson 2026-05-06 11:35:33 +01:00 • committed by GitHub
commit 01099f729d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
31 changed files with 1487 additions and 133 deletions
Download patch file Download diff file Expand all files Collapse all files

4
.gitignore vendored
View file

@ -24,8 +24,8 @@ dist/
downloads/
eggs/
.eggs/
lib/
lib64/
/lib/
/lib64/
parts/
sdist/
var/

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

@ -62,6 +62,7 @@ from datasets import Dataset, load_dataset
from utils.models import is_vision_model, detect_audio_type
from utils.datasets import format_and_template_dataset
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER
from utils.datasets.raw_text import prepare_raw_text_dataset
from utils.paths import (
ensure_dir,
resolve_dataset_path,
@ -125,6 +126,7 @@ class UnslothTrainer:
self.load_in_4bit = True # Track quantization mode for metadata
# Model state tracking
self.is_cpt = False # Set to True for Continued Pretraining
self.is_vlm = False
self.is_audio = False
self.is_audio_vlm = (
@ -925,6 +927,7 @@ class UnslothTrainer:
use_gradient_checkpointing: str = "unsloth",
use_rslora: bool = False,
use_loftq: bool = False,
modules_to_save: list = None,
) -> bool:
"""
Prepare model for training (with optional LoRA).
@ -1121,11 +1124,14 @@ class UnslothTrainer:
loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
if use_loftq
else None,
modules_to_save = modules_to_save,
)
else:
# Text model LoRA
logger.info(f"Text model LoRA configuration:")
logger.info(f" - Target modules: {target_modules}\n")
if modules_to_save:
logger.info(f" - Modules to save: {modules_to_save}\n")
self.model = FastLanguageModel.get_peft_model(
self.model,
@ -1140,6 +1146,7 @@ class UnslothTrainer:
loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
if use_loftq
else None,
modules_to_save = modules_to_save,
)
# Check if stopped during LoRA preparation
@ -2342,6 +2349,7 @@ class UnslothTrainer:
eval_steps: float = 0.00,
dataset_slice_start: int = None,
dataset_slice_end: int = None,
is_cpt: bool = False,
) -> Optional[tuple]:
"""
Load and prepare dataset for training.
@ -2360,6 +2368,35 @@ class UnslothTrainer:
False # True if eval comes from a separate HF split
)
eval_enabled = eval_steps is not None and eval_steps > 0
raw_text_mode = is_cpt or format_type == "raw"
def _raw_mode_label() -> str:
return "CPT" if is_cpt else "raw text"
def _apply_raw_text_prep(ds: Dataset, split_name: str) -> Dataset:
try:
result = prepare_raw_text_dataset(
ds,
mode_label = _raw_mode_label(),
split_name = split_name,
eos_token = getattr(self.tokenizer, "eos_token", None),
append_eos = True,
)
except ValueError as exc:
error_msg = str(exc)
logger.error(error_msg)
self._update_progress(error = error_msg)
raise
for notice in result.notices:
if notice.level == "warning":
logger.warning(notice.message)
if notice.update_status:
self._update_progress(status_message = notice.message)
else:
logger.info(f"{notice.message}\n")
return result.dataset
if local_datasets:
# Load local datasets using load_dataset() so the result is
@ -2534,6 +2571,48 @@ class UnslothTrainer:
processed = self._preprocess_dac_dataset(dataset, custom_format_mapping)
return ({"dataset": processed, "final_format": "audio_dac"}, None)
# ========== RAW TEXT BYPASS ==========
if raw_text_mode:
logger.info(
f"{_raw_mode_label().capitalize()} mode: bypassing chat template, "
"using raw text\n"
)
dataset = _apply_raw_text_prep(dataset, "train")
if has_separate_eval_source and eval_dataset is not None:
eval_dataset = _apply_raw_text_prep(eval_dataset, "eval")
dataset_info = {
"dataset": dataset,
"detected_format": "raw_text",
"final_format": "raw_text",
"success": True,
}
if has_separate_eval_source and eval_dataset is not None:
logger.info(
f"{_raw_mode_label().capitalize()}: eval dataset "
f"({len(eval_dataset)} rows) kept as raw text\n"
)
elif eval_enabled and not has_separate_eval_source:
split_result = self._resolve_eval_split_from_dataset(dataset)
if split_result is not None:
train_portion, eval_dataset = split_result
dataset_info["dataset"] = train_portion
train_dataset = dataset_info["dataset"]
n = len(train_dataset) if hasattr(train_dataset, "__len__") else None
n_display = f"{n:,}" if isinstance(n, int) else "streaming"
self._update_progress(
status_message = f"Dataset ready ({n_display} samples, raw text)"
)
logger.info(f"Raw-text dataset ready ({n_display} samples)\n")
if "text" not in train_dataset.column_names:
raise ValueError(
f"Raw-text dataset missing 'text' column: {train_dataset.column_names}"
)
return (dataset_info, eval_dataset)
elif self.is_audio_vlm:
formatted = self._format_audio_vlm_dataset(
dataset, custom_format_mapping
@ -2676,6 +2755,7 @@ class UnslothTrainer:
output_dir: str | None = None,
num_epochs: int = 3,
learning_rate: float = 2e-4,
embedding_learning_rate: float | None = None,
batch_size: int = 2,
gradient_accumulation_steps: int = 4,
warmup_steps: int = None,
@ -2728,6 +2808,7 @@ class UnslothTrainer:
"output_dir": output_dir,
"num_epochs": num_epochs,
"learning_rate": learning_rate,
"embedding_learning_rate": embedding_learning_rate,
"batch_size": batch_size,
"gradient_accumulation_steps": gradient_accumulation_steps,
"warmup_steps": warmup_steps,
@ -2945,6 +3026,13 @@ class UnslothTrainer:
logger.info("Configuring data collator...\n")
dataset_final_format = (
str(dataset.get("final_format", "")).lower()
if isinstance(dataset, dict)
else ""
)
raw_text_mode = dataset_final_format == "raw_text"
data_collator = None # Default to built-in data collator
if is_deepseek_ocr:
# Special DeepSeek OCR collator - auto-install if needed
@ -2984,7 +3072,7 @@ class UnslothTrainer:
self._update_progress(error = error_msg, is_training = False)
return
elif self.is_audio_vlm:
elif self.is_audio_vlm and not raw_text_mode:
# Audio VLM collator (e.g. Gemma 3N with audio data)
# Mirrors the collate_fn from Gemma3N_(4B)-Audio notebook
logger.info("Configuring audio VLM data collator...\n")
@ -3026,7 +3114,7 @@ class UnslothTrainer:
data_collator = audio_vlm_collate_fn
logger.info("Audio VLM data collator configured\n")
elif self.is_vlm:
elif self.is_vlm and not raw_text_mode:
# Standard VLM collator (images)
logger.info("Using UnslothVisionDataCollator for vision model\n")
from unsloth.trainer import UnslothVisionDataCollator
@ -3137,8 +3225,9 @@ class UnslothTrainer:
optim_value = training_args.get("optim", "adamw_8bit")
lr_scheduler_type_value = training_args.get("lr_scheduler_type", "linear")
if self.is_vlm or self.is_audio_vlm:
if (self.is_vlm or self.is_audio_vlm) and not raw_text_mode:
# Vision / audio VLM config (both need skip_prepare_dataset + remove_unused_columns)
# Raw-text runs on VLM-capable models are routed to the text path below.
label = "audio VLM" if self.is_audio_vlm else "vision"
logger.info(f"Configuring {label} model training parameters\n")
# Use provided values or defaults for vision models
@ -3160,7 +3249,14 @@ class UnslothTrainer:
}
)
else:
logger.info("Configuring text model training parameters\n")
is_cpt = training_args.get("is_cpt", False)
self.is_cpt = is_cpt
if is_cpt:
logger.info("Configuring Continued Pretraining (CPT) parameters\n")
elif raw_text_mode:
logger.info("Configuring raw-text training parameters\n")
else:
logger.info("Configuring text model training parameters\n")
config_args.update(
{
"optim": optim_value,
@ -3189,9 +3285,10 @@ class UnslothTrainer:
logger.info("Training configuration prepared\n")
# ========== TRAINER INITIALIZATION ==========
if self.is_audio_vlm:
if self.is_audio_vlm and not raw_text_mode:
# Audio VLM (e.g. Gemma 3N + audio): raw Dataset from _format_audio_vlm_dataset
# Notebook uses processing_class=processor.tokenizer (text tokenizer only)
# Raw-text runs are routed to the text path below.
train_dataset = (
dataset if isinstance(dataset, Dataset) else dataset["dataset"]
)
@ -3210,8 +3307,9 @@ class UnslothTrainer:
if eval_dataset is not None:
trainer_kwargs["eval_dataset"] = eval_dataset
self.trainer = SFTTrainer(**trainer_kwargs)
elif self.is_vlm:
elif self.is_vlm and not raw_text_mode:
# Image VLM: dataset is dict wrapper from format_and_template_dataset
# Raw-text runs are routed to the text path below.
train_dataset = (
dataset["dataset"] if isinstance(dataset, dict) else dataset
)
@ -3242,16 +3340,48 @@ class UnslothTrainer:
)
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)
if is_cpt:
try:
from unsloth import (
UnslothTrainer as _UnslothCPTTrainer,
UnslothTrainingArguments as _UnslothTrainingArguments,
)
except ImportError as exc:
raise RuntimeError(
"CPT requires a newer Unsloth install that exports "
"`UnslothTrainer` and `UnslothTrainingArguments` "
"(for embedding_learning_rate support). "
"Upgrade with: `pip install -U unsloth unsloth_zoo`."
) from exc
embedding_lr = training_args.get("embedding_learning_rate")
logger.info(
f"CPT: using UnslothTrainer with embedding_learning_rate={embedding_lr}\n"
)
trainer_kwargs = {
"model": self.model,
"tokenizer": sft_tokenizer,
"train_dataset": dataset["dataset"],
"data_collator": data_collator,
"args": _UnslothTrainingArguments(
embedding_learning_rate = embedding_lr,
**config_args,
),
}
if eval_dataset is not None:
trainer_kwargs["eval_dataset"] = eval_dataset
self.trainer = _UnslothCPTTrainer(**trainer_kwargs)
else:
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)
# Restore the full processor as processing_class so checkpoint
# saves include preprocessor_config.json (needed for GGUF export).
if sft_tokenizer is not self.tokenizer:
@ -3260,19 +3390,32 @@ class UnslothTrainer:
# ========== TRAIN ON RESPONSES ONLY ==========
# Determine if we should train on responses only
# Raw-text datasets always train on all tokens.
instruction_part = None
response_part = None
train_on_responses_enabled = training_args.get(
"train_on_completions", False
is_cpt = training_args.get("is_cpt", False)
train_on_responses_enabled = (
False
if (is_cpt or raw_text_mode)
else training_args.get("train_on_completions", False)
)
if is_cpt:
logger.info(
"CPT mode: skipping train_on_responses_only — training on all tokens\n"
)
elif raw_text_mode:
logger.info(
"Raw-text mode: skipping train_on_responses_only — training on all tokens\n"
)
# DeepSeek OCR handles this internally in its collator, so skip
# Audio VLM handles label masking in its collator, so skip
if (
train_on_responses_enabled
and not self.is_audio_vlm
and not self.is_audio
and not (is_deepseek_ocr or dataset["final_format"].lower() == "alpaca")
and not (is_deepseek_ocr or dataset_final_format == "alpaca")
):
try:
logger.info("Configuring train on responses only...\n")
@ -3318,7 +3461,7 @@ class UnslothTrainer:
and response_part
and not self.is_audio_vlm
and not self.is_audio
and not (is_deepseek_ocr or dataset["final_format"].lower() == "alpaca")
and not (is_deepseek_ocr or dataset_final_format == "alpaca")
):
try:
from unsloth.chat_templates import train_on_responses_only
@ -3451,7 +3594,9 @@ class UnslothTrainer:
config = json.load(f)
# Determine the training method
if self.load_in_4bit:
if self.is_cpt:
method = "CPT"
elif self.load_in_4bit:
method = "qlora"
else:
method = "lora"

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

@ -159,6 +159,7 @@ class TrainingBackend:
"is_embedding": kwargs.get("is_embedding", False),
"num_epochs": kwargs.get("num_epochs", 3),
"learning_rate": kwargs.get("learning_rate", "2e-4"),
"embedding_learning_rate": kwargs.get("embedding_learning_rate"),
"batch_size": kwargs.get("batch_size", 2),
"gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4),
"warmup_steps": kwargs.get("warmup_steps"),
@ -195,8 +196,9 @@ class TrainingBackend:
"gpu_ids": kwargs.get("gpu_ids"),
}
# Derive load_in_4bit from training_type
if config["training_type"] != "LoRA/QLoRA":
# Full finetuning always runs in 16-bit. LoRA/QLoRA and CPT preserve the
# explicit request so 4-bit adapter/raw-text runs remain possible.
if config["training_type"] == "Full Finetuning":
config["load_in_4bit"] = False
# Spawn subprocess — use locals so state is untouched on failure

74
studio/backend/core/training/worker.py
View file

@ -15,6 +15,7 @@ from __future__ import annotations
import structlog
from loggers import get_logger
import math
import os
import shutil
import sys
@ -1208,6 +1209,8 @@ def run_training_process(
# ── 4b. Load and format dataset (LLM helper may use VRAM briefly) ──
_send_status(event_queue, "Loading and formatting dataset...")
hf_dataset = config.get("hf_dataset", "")
training_type = config.get("training_type", "LoRA/QLoRA")
_is_cpt_for_dataset = training_type == "Continued Pretraining"
dataset_result = trainer.load_and_format_dataset(
dataset_source = hf_dataset if hf_dataset and hf_dataset.strip() else None,
format_type = config.get("format_type", ""),
@ -1220,6 +1223,7 @@ def run_training_process(
eval_steps = config.get("eval_steps", 0.00),
dataset_slice_start = config.get("dataset_slice_start"),
dataset_slice_end = config.get("dataset_slice_end"),
is_cpt = _is_cpt_for_dataset,
)
if isinstance(dataset_result, tuple):
@ -1305,7 +1309,9 @@ def run_training_process(
_tqdm_thread.start()
training_type = config.get("training_type", "LoRA/QLoRA")
use_lora = training_type == "LoRA/QLoRA"
is_cpt = training_type == "Continued Pretraining"
use_lora = training_type in ("LoRA/QLoRA", "Continued Pretraining")
cpt_trains_embeddings = False
# ── 4c. Load training model (uses VRAM — dataset already formatted) ──
_send_status(event_queue, "Loading model...")
@ -1337,8 +1343,41 @@ def run_training_process(
)
return
# ── 4d. Prepare model (LoRA or full finetuning) ──
if use_lora:
# ── 4d. Prepare model (LoRA, full finetuning, or CPT) ──
if is_cpt:
_send_status(event_queue, "Configuring LoRA for continued pretraining...")
# embed_tokens (if the user included it) goes to modules_to_save —
# trained full-precision at embedding_learning_rate. lm_head stays as
# a LoRA target for merge compatibility (see unsloth PR #4106).
_user_modules = config.get("target_modules") or []
wants_embed = "embed_tokens" in _user_modules
cpt_trains_embeddings = wants_embed
cpt_target_modules = [m for m in _user_modules if m != "embed_tokens"]
if not cpt_target_modules:
cpt_target_modules = [
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
"lm_head",
]
success = trainer.prepare_model_for_training(
use_lora = True,
target_modules = cpt_target_modules,
modules_to_save = ["embed_tokens"] if wants_embed else None,
lora_r = config.get("lora_r", 128),
lora_alpha = config.get("lora_alpha", 32),
lora_dropout = config.get("lora_dropout", 0.0),
use_gradient_checkpointing = config.get(
"gradient_checkpointing", "unsloth"
),
use_rslora = config.get("use_rslora", False),
use_loftq = config.get("use_loftq", False),
)
elif use_lora:
_send_status(event_queue, "Configuring LoRA adapters...")
success = trainer.prepare_model_for_training(
use_lora = True,
@ -1379,9 +1418,9 @@ def run_training_process(
)
return
# Convert learning rate
lr_default = "5e-5" if is_cpt else "2e-4"
try:
lr_value = float(config.get("learning_rate", "2e-4"))
lr_value = float(config.get("learning_rate", lr_default))
except ValueError:
event_queue.put(
{
@ -1393,6 +1432,25 @@ def run_training_process(
)
return
# embedding_learning_rate is validated by the Pydantic model (Optional[float],
# gt=0, lt=1.0); if present it is already a finite float in range.
embedding_lr_value = config.get("embedding_learning_rate")
if is_cpt:
if cpt_trains_embeddings:
if embedding_lr_value is None:
# Default embedding_learning_rate = lr/10 per Unsloth's CPT notebook.
embedding_lr_value = lr_value / 10.0
logger.info(
f"CPT: using default embedding_learning_rate={embedding_lr_value:.1e} "
f"(lr/10). Set explicitly to override.\n"
)
elif embedding_lr_value is not None:
logger.warning(
"CPT: embedding_learning_rate was provided but embed_tokens is "
"not being trained; ignoring the override.\n"
)
embedding_lr_value = None
# Generate output dir
resume_from_checkpoint = config.get("resume_from_checkpoint")
output_dir = config.get("output_dir") or _output_dir_from_resume_checkpoint(
@ -1425,6 +1483,7 @@ def run_training_process(
output_dir = output_dir,
num_epochs = config.get("num_epochs", 3),
learning_rate = lr_value,
embedding_learning_rate = embedding_lr_value,
batch_size = config.get("batch_size", 2),
gradient_accumulation_steps = config.get("gradient_accumulation_steps", 4),
warmup_steps = config.get("warmup_steps"),
@ -1434,7 +1493,9 @@ def run_training_process(
weight_decay = config.get("weight_decay", 0.001),
random_seed = config.get("random_seed", 3407),
packing = config.get("packing", False),
train_on_completions = config.get("train_on_completions", False),
train_on_completions = False
if is_cpt
else config.get("train_on_completions", False),
enable_wandb = config.get("enable_wandb", False),
wandb_project = config.get("wandb_project", "unsloth-training"),
wandb_token = config.get("wandb_token"),
@ -1445,6 +1506,7 @@ def run_training_process(
max_seq_length = config.get("max_seq_length", 2048),
optim = config.get("optim", "adamw_8bit"),
lr_scheduler_type = config.get("lr_scheduler_type", "linear"),
is_cpt = is_cpt,
resume_from_checkpoint = resume_from_checkpoint,
)

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

@ -16,8 +16,11 @@ class TrainingStartRequest(BaseModel):
model_name: str = Field(
..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')"
)
training_type: str = Field(
..., description = "Training type: 'LoRA/QLoRA' or 'Full Finetuning'"
training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = (
Field(
...,
description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'",
)
)
hf_token: Optional[str] = Field(None, description = "HuggingFace token")
load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization")
@ -86,6 +89,13 @@ class TrainingStartRequest(BaseModel):
packing: bool = Field(False, description = "Enable sequence packing")
optim: str = Field("adamw_8bit", description = "Optimizer")
lr_scheduler_type: str = Field("linear", description = "Learning rate scheduler type")
embedding_learning_rate: Optional[float] = Field(
None,
gt = 0,
lt = 1.0,
description = "Separate learning rate for embedding matrices (CPT). "
"Must be in (0, 1). Should be 2-10x smaller than the main learning rate.",
)
# LoRA parameters
use_lora: bool = Field(True, description = "Use LoRA (derived from training_type)")

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

@ -207,6 +207,7 @@ async def start_training(
"custom_format_mapping": request.custom_format_mapping,
"num_epochs": request.num_epochs,
"learning_rate": request.learning_rate,
"embedding_learning_rate": request.embedding_learning_rate,
"batch_size": request.batch_size,
"gradient_accumulation_steps": request.gradient_accumulation_steps,
"warmup_steps": request.warmup_steps,

183
studio/backend/tests/test_training_raw_support.py Normal file
View file

@ -0,0 +1,183 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import asyncio
import importlib.util
import unittest
from pathlib import Path
from unittest.mock import patch
from datasets import Dataset
from core.training.training import TrainingBackend
from models.training import TrainingStartRequest
from utils.datasets import format_dataset, format_and_template_dataset
from utils.datasets.raw_text import prepare_raw_text_dataset
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
def _load_route_module(name: str, relative_path: str):
spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
class TestTrainingRawSupport(unittest.TestCase):
def test_training_backend_preserves_cpt_4bit_and_embedding_lr(self):
backend = TrainingBackend()
class DummyProcess:
pid = 12345
def start(self):
return None
class DummyThread:
def start(self):
return None
dummy_queue = object()
with (
patch(
"core.training.training.prepare_gpu_selection",
return_value = ([0], {"selection_mode": "auto"}),
),
patch(
"core.training.training._CTX.Queue",
side_effect = [dummy_queue, dummy_queue],
),
patch(
"core.training.training._CTX.Process", return_value = DummyProcess()
) as mock_process,
patch(
"core.training.training.threading.Thread",
return_value = DummyThread(),
),
):
backend.start_training(
job_id = "test-cpt-raw",
model_name = "unsloth/test-bnb-4bit",
training_type = "Continued Pretraining",
format_type = "raw",
load_in_4bit = True,
embedding_learning_rate = 1e-5,
)
config = mock_process.call_args.kwargs["kwargs"]["config"]
self.assertTrue(config["load_in_4bit"])
self.assertEqual(config["embedding_learning_rate"], 1e-5)
def test_training_route_forwards_embedding_learning_rate(self):
training_route = _load_route_module(
"training_route_module_raw_support",
"routes/training.py",
)
captured: dict = {}
class DummyBackend:
current_job_id = None
def is_training_active(self):
return False
def start_training(self, **kwargs):
captured.update(kwargs)
return True
request = TrainingStartRequest(
model_name = "unsloth/test-bnb-4bit",
training_type = "Continued Pretraining",
format_type = "raw",
load_in_4bit = True,
embedding_learning_rate = 1e-5,
)
with (
patch.object(
training_route,
"get_training_backend",
return_value = DummyBackend(),
),
patch.object(training_route, "load_model_defaults", return_value = {}),
patch(
"core.inference.get_inference_backend",
return_value = type(
"InferenceBackend",
(),
{"active_model_name": None},
)(),
),
patch(
"core.export.get_export_backend",
return_value = type(
"ExportBackend",
(),
{"current_checkpoint": None},
)(),
),
):
response = asyncio.run(
training_route.start_training(request, current_subject = "test-user")
)
self.assertEqual(response.status, "queued")
self.assertEqual(captured["embedding_learning_rate"], 1e-5)
self.assertTrue(captured["load_in_4bit"])
def test_format_dataset_supports_raw_text(self):
dataset = Dataset.from_dict(
{
"body": ["hello", "world"],
"title": ["a", "b"],
"id": [1, 2],
}
)
result = format_dataset(dataset, format_type = "raw")
self.assertEqual(result["final_format"], "raw_text")
self.assertIn("text", result["dataset"].column_names)
self.assertEqual(result["dataset"][0]["text"], "hello")
self.assertFalse(result["requires_manual_mapping"])
def test_format_and_template_dataset_supports_raw_text_without_template(self):
dataset = Dataset.from_dict({"body": ["hello raw world"]})
result = format_and_template_dataset(
dataset,
model_name = "unsloth/test",
tokenizer = None,
format_type = "raw",
)
self.assertTrue(result["success"])
self.assertEqual(result["final_format"], "raw_text")
self.assertEqual(result["dataset"][0]["text"], "hello raw world")
def test_prepare_raw_text_dataset_drops_null_rows_before_appending_eos(self):
dataset = Dataset.from_dict({"text": ["hello", None, "world"]})
result = prepare_raw_text_dataset(
dataset,
mode_label = "CPT",
split_name = "train",
eos_token = "<eos>",
append_eos = True,
)
self.assertEqual(len(result.dataset), 2)
self.assertEqual(result.dataset[0]["text"], "hello<eos>")
self.assertEqual(result.dataset[1]["text"], "world<eos>")
self.assertTrue(
any(
"null or non-string 'text' values" in notice.message
for notice in result.notices
)
)
if __name__ == "__main__":
unittest.main()

30
studio/backend/utils/datasets/dataset_utils.py
View file

@ -41,6 +41,7 @@ from .chat_templates import (
get_tokenizer_chat_template,
DEFAULT_ALPACA_TEMPLATE,
)
from .raw_text import prepare_raw_text_dataset
from .vlm_processing import generate_smart_vlm_instruction
from .data_collators import DeepSeekOCRDataCollator, VLMDataCollator
from .model_mappings import TEMPLATE_TO_MODEL_MAPPER
@ -437,6 +438,20 @@ def format_dataset(
# Detect multimodal first (needed for all flows)
multimodal_info = detect_multimodal_dataset(dataset)
if format_type == "raw":
raw_result = prepare_raw_text_dataset(dataset)
return {
"dataset": raw_result.dataset,
"detected_format": "raw_text",
"final_format": "raw_text",
"chat_column": "text",
"is_standardized": True,
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": [notice.message for notice in raw_result.notices],
}
# If user provided explicit mapping, skip detection and apply in the requested format
if custom_format_mapping:
try:
@ -1105,6 +1120,21 @@ def format_and_template_dataset(
num_proc = num_proc,
)
if dataset_info["final_format"] == "raw_text":
summary = get_dataset_info_summary(dataset_info)
return {
"dataset": dataset_info["dataset"],
"detected_format": dataset_info["detected_format"],
"final_format": dataset_info["final_format"],
"chat_column": dataset_info.get("chat_column"),
"is_vlm": False,
"success": True,
"requires_manual_mapping": False,
"warnings": dataset_info.get("warnings", []),
"errors": [],
"summary": summary,
}
# Step 2: Apply chat template
detected = dataset_info.get("detected_format", "unknown")
if progress_callback and n_rows:

142
studio/backend/utils/datasets/raw_text.py Normal file
View file

@ -0,0 +1,142 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Shared helpers for raw-text dataset preparation.
"""
from dataclasses import dataclass
from typing import Literal
from datasets import Dataset
@dataclass(frozen = True)
class RawTextNotice:
message: str
level: Literal["info", "warning"]
update_status: bool = False
@dataclass(frozen = True)
class RawTextPreparationResult:
dataset: Dataset
notices: list[RawTextNotice]
def _string_columns(dataset: Dataset) -> list[str]:
feature_map = getattr(dataset, "features", {}) or {}
string_cols: list[str] = []
for col in dataset.column_names:
feature = feature_map.get(col)
dtype = str(getattr(feature, "dtype", ""))
if dtype in {"string", "large_string"}:
string_cols.append(col)
return string_cols
def _split_scope(split_name: str | None) -> str:
return f"the {split_name} split" if split_name else "this dataset"
def _drop_invalid_text_rows(
dataset: Dataset,
*,
mode_title: str,
split_scope: str,
) -> tuple[Dataset, list[RawTextNotice]]:
filtered_dataset = dataset.filter(lambda ex: isinstance(ex["text"], str))
dropped_rows = len(dataset) - len(filtered_dataset)
if not dropped_rows:
return filtered_dataset, []
if len(filtered_dataset) == 0:
raise ValueError(
f"{mode_title} training requires at least one string 'text' value "
f"in {split_scope}; all {dropped_rows} rows were null or non-string."
)
return filtered_dataset, [
RawTextNotice(
message = (
f"{mode_title}: dropped {dropped_rows:,} row(s) with null or "
f"non-string 'text' values from {split_scope}"
),
level = "warning",
update_status = True,
)
]
def prepare_raw_text_dataset(
dataset: Dataset,
*,
mode_label: str = "raw text",
split_name: str | None = None,
eos_token: str | None = None,
append_eos: bool = False,
) -> RawTextPreparationResult:
notices: list[RawTextNotice] = []
mode_title = mode_label.capitalize()
split_scope = _split_scope(split_name)
if "text" not in dataset.column_names:
string_cols = _string_columns(dataset)
if not string_cols:
raise ValueError(
f"{mode_title} training requires a string 'text' column but none "
f"was found in {split_scope} (columns: {dataset.column_names})."
)
renamed_col = string_cols[0]
if len(string_cols) > 1:
notices.append(
RawTextNotice(
message = (
f"{mode_title}: dataset has {len(string_cols)} string "
f"columns ({string_cols}); auto-selecting '{renamed_col}' "
"as the training text. Rename the intended column to "
"'text' to override."
),
level = "warning",
update_status = True,
)
)
notices.append(
RawTextNotice(
message = (
f"{mode_title}: renaming column '{renamed_col}' -> 'text' "
f"for {split_scope}"
),
level = "info",
)
)
dataset = dataset.rename_column(renamed_col, "text")
dataset, invalid_row_notices = _drop_invalid_text_rows(
dataset,
mode_title = mode_title,
split_scope = split_scope,
)
notices.extend(invalid_row_notices)
if append_eos:
if not eos_token:
notices.append(
RawTextNotice(
message = (
f"{mode_title}: tokenizer has no eos_token; skipping EOS "
"append. Model will not learn document boundaries."
),
level = "warning",
)
)
else:
def _append_eos(ex, _eos = eos_token):
text = ex["text"]
return {"text": text if text.endswith(_eos) else text + _eos}
dataset = dataset.map(_append_eos)
return RawTextPreparationResult(dataset = dataset, notices = notices)

10
studio/frontend/src/config/training.ts
View file

@ -76,6 +76,13 @@ export const TARGET_MODULES = [
"down_proj",
];
/** CPT requires embed_tokens and lm_head in addition to standard LoRA modules. */
export const CPT_TARGET_MODULES = [
...TARGET_MODULES,
"embed_tokens",
"lm_head",
];
export const OPTIMIZER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [
{ value: "adamw_8bit", label: "AdamW 8-bit" },
{ value: "paged_adamw_8bit", label: "Paged AdamW 8-bit" },
@ -96,11 +103,14 @@ export const LR_SCHEDULER_OPTIONS: ReadonlyArray<{ value: string; label: string
*/
export const LR_DEFAULT_LORA = 2e-4;
export const LR_DEFAULT_FULL = 2e-5;
export const LR_DEFAULT_CPT = 5e-5;
export const DEFAULT_HYPERPARAMS = {
epochs: 3,
contextLength: 2048,
learningRate: LR_DEFAULT_LORA,
// null = let backend auto-compute (lr/10 per Unsloth CPT recipe). Only used by CPT.
embeddingLearningRate: null as number | null,
optimizerType: "adamw_8bit",
lrSchedulerType: "linear",
loraRank: 16,

1
studio/frontend/src/features/export/constants.ts
View file

@ -74,6 +74,7 @@ export const METHOD_LABELS: Record<TrainingMethod, string> = {
qlora: "QLoRA",
lora: "LoRA",
full: "Full Fine-tune",
cpt: "Continued Pretraining",
};
export const GUIDE_STEPS = [

1
studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx
View file

@ -63,6 +63,7 @@ const FORMAT_OPTIONS: { value: DatasetFormat; label: string }[] = [
{ value: "alpaca", label: "Alpaca" },
{ value: "chatml", label: "ChatML" },
{ value: "sharegpt", label: "ShareGPT" },
{ value: "raw", label: "Raw Text" },
];
export function DatasetStep() {

1
studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx
View file

@ -366,6 +366,7 @@ export function ModelSelectionStep() {
<SelectItem value="qlora">QLoRA (4-bit)</SelectItem>
<SelectItem value="lora">LoRA (16-bit)</SelectItem>
<SelectItem value="full">Full Fine-tune</SelectItem>
<SelectItem value="cpt">Continued Pretraining</SelectItem>
</SelectContent>
</Select>
</div>

6
studio/frontend/src/features/onboarding/components/steps/summary-step.tsx
View file

@ -5,6 +5,7 @@ import { Badge } from "@/components/ui/badge";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Separator } from "@/components/ui/separator";
import { useTrainingConfigStore } from "@/features/training";
import { getTrainingMethodLabel } from "@/features/training/lib/training-methods";
import { useHardwareInfo } from "@/hooks";
import { isAdapterMethod } from "@/types/training";
import { ChipIcon, Database02Icon, GpuIcon, Settings04Icon } from "@hugeicons/core-free-icons";
@ -102,6 +103,7 @@ export function SummaryStep() {
const showLoraParams = isAdapterMethod(trainingMethod);
const datasetName = datasetSource === "upload" ? uploadedFile : dataset;
const trainingMethodLabel = getTrainingMethodLabel(trainingMethod);
return (
<div className="grid grid-cols-2 gap-3">
@ -150,7 +152,7 @@ export function SummaryStep() {
<Separator className="my-2" />
<div className="space-y-1 text-sm">
<Row label="Type" value={modelType} capitalize />
<Row label="Method" value={trainingMethod === "qlora" ? "QLoRA" : trainingMethod === "lora" ? "LoRA" : "Full"} />
<Row label="Method" value={trainingMethodLabel} />
</div>
</CardContent>
</Card>
@ -199,7 +201,7 @@ export function SummaryStep() {
<div className="flex flex-1 flex-col">
<span className="text-xs text-muted-foreground">Training</span>
<span className="text-sm font-medium">
{trainingMethod === "qlora" ? "QLoRA" : trainingMethod === "lora" ? "LoRA" : "Full"}
{trainingMethodLabel}
</span>
</div>
</div>

21
studio/frontend/src/features/studio/historical-training-view.tsx
View file

@ -4,6 +4,7 @@
import type { TrainingViewData } from "@/features/training";
import { getTrainingRun } from "@/features/training";
import type { TrainingRunDetailResponse } from "@/features/training";
import { parseBackendTrainingMethod } from "@/features/training/lib/training-methods";
import { type ReactElement, useEffect, useState } from "react";
import { ChartsSection } from "./sections/charts-section";
import { ProgressSection } from "./sections/progress-section";
@ -12,15 +13,6 @@ interface HistoricalTrainingViewProps {
runId: string;
}
function normalizeTrainingMethod(config: Record<string, unknown>): string {
const type = config?.training_type as string | undefined;
if (!type || type === "Full Finetuning") return "full";
if (type === "LoRA/QLoRA") {
return config?.load_in_4bit ? "qlora" : "lora";
}
return "full";
}
function mapToViewData(detail: TrainingRunDetailResponse): TrainingViewData {
const { run, metrics } = detail;
@ -79,7 +71,10 @@ function mapToViewData(detail: TrainingRunDetailResponse): TrainingViewData {
error: run.status === "error" ? run.error_message : null,
isTrainingRunning: false,
modelName: run.model_name,
trainingMethod: normalizeTrainingMethod(detail.config),
trainingMethod: parseBackendTrainingMethod(
detail.config?.training_type,
detail.config?.load_in_4bit,
),
lossHistory,
lrHistory,
gradNormHistory,
@ -143,7 +138,11 @@ export function HistoricalTrainingView({
loraRank: detail.config.lora_r as number | undefined,
loraAlpha: detail.config.lora_alpha as number | undefined,
loraDropout: detail.config.lora_dropout as number | undefined,
loraVariant: detail.config.use_rslora ? "rsLoRA" : undefined,
loraVariant: detail.config.use_rslora
? "rslora"
: detail.config.use_loftq
? "loftq"
: "lora",
}
: undefined;

10
studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
View file

@ -15,6 +15,7 @@ import { Badge } from "@/components/ui/badge";
import { Spinner } from "@/components/ui/spinner";
import { useTrainingActions, useTrainingConfigStore } from "@/features/training";
import { checkDatasetFormat } from "@/features/training/api/datasets-api";
import { isRawTextDatasetFormat } from "@/features/training/lib/training-methods";
import type { CheckFormatResponse } from "@/features/training/types/datasets";
import { Database02Icon, AlertCircleIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
@ -90,10 +91,11 @@ export function DatasetPreviewDialog({
const effectiveIsAudio = !!data?.is_audio;
const effectiveIsVlm = isVlm || !!data?.is_image;
const isRawFormat = isRawTextDatasetFormat(datasetFormat);
const hasHeuristicMapping = !data?.requires_manual_mapping && !!data?.suggested_mapping;
const mappingEnabled = !!data?.requires_manual_mapping || hasHeuristicMapping;
const mappingEnabled = !isRawFormat && (!!data?.requires_manual_mapping || hasHeuristicMapping);
const showMappingFooter = mode === "mapping" && mappingEnabled;
const mappingOk = isMappingComplete(manualMapping, effectiveIsVlm, datasetFormat, effectiveIsAudio);
const mappingOk = isRawFormat || isMappingComplete(manualMapping, effectiveIsVlm, datasetFormat, effectiveIsAudio);
const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat, effectiveIsAudio);
const isHfDataset = datasetSource === "huggingface";
@ -413,7 +415,7 @@ export function DatasetPreviewDialog({
<MetaRow label="Source" value={sourceLabel} />
<MetaRow
label="Format"
value={data.detected_format || "--"}
value={isRawFormat ? "Raw Text" : (data.detected_format || "--")}
/>
<MetaRow
label="Total Rows"
@ -441,7 +443,7 @@ export function DatasetPreviewDialog({
/>
</div>
{data.warning && (
{data.warning && !isRawFormat && (
<div className="rounded-lg border border-amber-200 bg-amber-50 px-4 py-3 text-xs text-amber-700 dark:border-amber-800 dark:bg-amber-950 dark:text-amber-400 mb-4 flex items-start gap-2.5">
<HugeiconsIcon icon={AlertCircleIcon} className="size-4 shrink-0 mt-0.5" />
<span>{data.warning}</span>

1
studio/frontend/src/features/studio/sections/dataset-section.tsx
View file

@ -913,6 +913,7 @@ export function DatasetSection() {
<SelectItem value="alpaca">Alpaca</SelectItem>
<SelectItem value="chatml">ChatML</SelectItem>
<SelectItem value="sharegpt">ShareGPT</SelectItem>
<SelectItem value="raw">Raw Text</SelectItem>
</SelectContent>
</Select>
</div>

13
studio/frontend/src/features/studio/sections/model-section.tsx
View file

@ -67,6 +67,7 @@ const METHOD_DOTS: Record<string, string> = {
qlora: "bg-emerald-400",
lora: "bg-blue-400",
full: "bg-amber-400",
cpt: "bg-purple-400",
};
const DARK_TRIGGER =
@ -570,7 +571,9 @@ export function ModelSection() {
</TooltipTrigger>
<TooltipContent className="max-w-xs">
QLoRA uses 4-bit quantization for lowest VRAM. LoRA uses
16-bit. Full updates all weights.{" "}
16-bit. Full updates all weights. CPT (Continued Pretraining)
trains on raw text to adapt the model to a new domain without
chat formatting.{" "}
<a
href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
target="_blank"
@ -617,6 +620,14 @@ export function ModelSection() {
Full Fine-tune
</span>
</SelectItem>
<SelectItem value="cpt">
<span className="flex items-center gap-2">
<span
className={`size-2 shrink-0 rounded-full ${METHOD_DOTS.cpt}`}
/>
Continued Pretraining
</span>
</SelectItem>
</SelectContent>
</Select>
</div>

68
studio/frontend/src/features/studio/sections/params-section.tsx
View file

@ -34,11 +34,14 @@ import {
} from "@/components/ui/tooltip";
import {
CONTEXT_LENGTHS,
CPT_TARGET_MODULES,
LR_SCHEDULER_OPTIONS,
OPTIMIZER_OPTIONS,
TARGET_MODULES,
} from "@/config/training";
import { useMaxStepsEpochsToggle, useTrainingConfigStore } from "@/features/training";
import { isRawTextDatasetFormat } from "@/features/training/lib/training-methods";
import { isAdapterMethod } from "@/types/training";
import type { GradientCheckpointing } from "@/types/training";
import {
ArrowDown01Icon,
@ -126,10 +129,13 @@ function SliderRow({
export function ParamsSection(): ReactElement {
const store = useTrainingConfigStore();
const platformDeviceType = usePlatformStore((s) => s.deviceType);
const isLora = store.trainingMethod !== "full";
const isLora = isAdapterMethod(store.trainingMethod);
const isCpt = store.trainingMethod === "cpt";
const isRawText = isRawTextDatasetFormat(store.datasetFormat);
const showVisionLora = store.isVisionModel && store.isDatasetImage === true;
const [loraOpen, setLoraOpen] = useState(false);
const [hyperOpen, setHyperOpen] = useState(false);
const needsExpandedHeight = isCpt || (isLora && loraOpen) || hyperOpen;
const [ctxInput, setCtxInput] = useState(String(store.contextLength));
const ctxAnchorRef = useRef<HTMLDivElement>(null);
const ctxItems = CONTEXT_LENGTHS.map(String);
@ -168,7 +174,7 @@ export function ParamsSection(): ReactElement {
title="Parameters"
description="Configure training hyperparameters"
accent="orange"
className={`${(isLora && loraOpen) || hyperOpen
className={`${needsExpandedHeight
? "min-h-studio-config-column"
: "h-studio-config-column"} duration-150`}
>
@ -378,10 +384,62 @@ export function ParamsSection(): ReactElement {
className="w-full font-mono"
/>
<p className="text-[10px] text-muted-foreground">
Recommended: 2e-4 for LoRA, 2e-5 for full fine-tune
Recommended: 2e-4 for LoRA, 5e-5 for CPT, 2e-5 for full fine-tune
</p>
</div>
{/* Embedding Learning Rate (CPT only) */}
{isCpt && (
<div className="flex flex-col gap-2">
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
Embedding Learning Rate
<Tooltip>
<TooltipTrigger asChild={true}>
<button
type="button"
className="text-foreground/70 hover:text-foreground"
>
<HugeiconsIcon
icon={InformationCircleIcon}
className="size-3"
/>
</button>
</TooltipTrigger>
<TooltipContent>
Only used when CPT is training <code>embed_tokens</code>.
Embeddings are easier to destabilize than LoRA weights, so
they usually need a smaller LR. Leave blank to use
<code>lr/10</code>; typical working range is 2x-10x smaller
than the main LR. Increase it only if vocabulary or
domain-token adaptation is too slow.
</TooltipContent>
</Tooltip>
</span>
<Input
type="number"
step="0.00001"
min="0"
max="1"
placeholder={`auto (${(store.learningRate / 10).toExponential(1)})`}
value={store.embeddingLearningRate ?? ""}
onChange={(e) => {
const raw = e.target.value;
if (raw === "") {
store.setEmbeddingLearningRate(null);
return;
}
const n = Number(raw);
store.setEmbeddingLearningRate(Number.isFinite(n) ? n : null);
}}
className="w-full font-mono"
/>
<p className="text-[10px] text-muted-foreground">
Leave blank to use lr/10 (recommended). Typical range is
2x-10x smaller than the main learning rate.
</p>
</div>
)}
{/* LoRA Settings */}
{isLora && (
<Collapsible open={loraOpen} onOpenChange={setLoraOpen}>
@ -516,7 +574,7 @@ export function ParamsSection(): ReactElement {
Target Modules
</span>
<div className="flex flex-wrap gap-1.5">
{TARGET_MODULES.map((mod) => {
{(isCpt ? CPT_TARGET_MODULES : TARGET_MODULES).map((mod) => {
const active = store.targetModules.includes(mod);
return (
<button
@ -908,7 +966,7 @@ export function ParamsSection(): ReactElement {
</label>
</div>
)}
{!store.isEmbeddingModel && (
{!store.isEmbeddingModel && !isCpt && !isRawText && (
<div className="flex items-center gap-2">
<Checkbox
id="trainOnCompletions"

4
studio/frontend/src/features/studio/sections/progress-section.tsx
View file

@ -26,6 +26,7 @@ import {
useTrainingConfigStore,
useTrainingRuntimeStore,
} from "@/features/training";
import { getTrainingMethodLabel } from "@/features/training/lib/training-methods";
import type { TrainingViewData } from "@/features/training";
import { useGpuUtilization } from "@/hooks";
import { cn } from "@/lib/utils";
@ -86,6 +87,7 @@ export function ProgressSection({
configOverride,
}: ProgressSectionProps): ReactElement {
const navigate = useNavigate();
const trainingMethodLabel = getTrainingMethodLabel(data.trainingMethod);
const config = useTrainingConfigStore(
useShallow((state) => ({
@ -272,7 +274,7 @@ export function ProgressSection({
{data.modelName || "--"}
</MetricStat>
<MetricStat label="Method">
{data.trainingMethod === "qlora" ? "QLoRA" : data.trainingMethod === "lora" ? "LoRA" : "Full"}
{trainingMethodLabel}
</MetricStat>
</div>

23
studio/frontend/src/features/training/api/mappers.ts
View file

@ -3,9 +3,10 @@
import type { TrainingConfigState } from "../types/config";
import type { TrainingStartRequest } from "../types/api";
const BACKEND_LORA_TYPE = "LoRA/QLoRA";
const BACKEND_FULL_TYPE = "Full Finetuning";
import {
isRawTextDatasetFormat,
toBackendTrainingType,
} from "../lib/training-methods";
function parseSliceValue(value: string | null): number | null {
if (value == null) return null;
@ -16,16 +17,15 @@ function parseSliceValue(value: string | null): number | null {
return num;
}
export function toBackendTrainingType(trainingMethod: string): string {
return trainingMethod === "full" ? BACKEND_FULL_TYPE : BACKEND_LORA_TYPE;
}
export function buildTrainingStartPayload(
config: TrainingConfigState,
): TrainingStartRequest {
const isCpt = config.trainingMethod === "cpt";
const adapterMethod = config.trainingMethod !== "full";
const isQloraMethod = config.trainingMethod === "qlora";
const isFourBitModel = (config.selectedModel ?? "").toLowerCase().includes("4bit");
const isEmbedding = config.isEmbeddingModel;
const isRawText = isRawTextDatasetFormat(config.datasetFormat);
const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null;
const localDatasets =
config.datasetSource === "upload" && config.uploadedFile
@ -53,7 +53,7 @@ export function buildTrainingStartPayload(
model_name: config.selectedModel ?? "",
training_type: toBackendTrainingType(config.trainingMethod),
hf_token: config.hfToken.trim() || null,
load_in_4bit: adapterMethod ? isQloraMethod : false,
load_in_4bit: (adapterMethod && isQloraMethod) || (isCpt && isFourBitModel),
max_seq_length: config.contextLength,
trust_remote_code: config.trustRemoteCode ?? false,
hf_dataset: hfDataset,
@ -71,6 +71,10 @@ export function buildTrainingStartPayload(
custom_format_mapping: customFormatMapping,
num_epochs: config.epochs,
learning_rate: String(config.learningRate),
embedding_learning_rate:
isCpt && config.embeddingLearningRate != null
? config.embeddingLearningRate
: null,
batch_size: config.batchSize,
gradient_accumulation_steps: config.gradientAccumulation,
warmup_steps: isEmbedding ? null : config.warmupSteps,
@ -91,7 +95,8 @@ export function buildTrainingStartPayload(
gradient_checkpointing: config.gradientCheckpointing,
use_rslora: config.loraVariant === "rslora",
use_loftq: config.loraVariant === "loftq",
train_on_completions: isEmbedding ? false : config.trainOnCompletions,
// CPT always trains on full sequences (no chat format masking)
train_on_completions: (isEmbedding || isCpt || isRawText) ? false : config.trainOnCompletions,
finetune_vision_layers: config.finetuneVisionLayers,
finetune_language_layers: config.finetuneLanguageLayers,
finetune_attention_modules: config.finetuneAttentionModules,

6
studio/frontend/src/features/training/hooks/use-training-actions.ts
View file

@ -8,6 +8,7 @@ import { checkDatasetFormat } from "../api/datasets-api";
import { getTrainingRun } from "../api/history-api";
import { buildTrainingStartPayload } from "../api/mappers";
import { resetTraining, startTraining, stopTraining } from "../api/train-api";
import { isRawTextDatasetFormat } from "../lib/training-methods";
import { syncTrainingRuntimeFromBackend } from "../lib/sync-runtime";
import { validateTrainingConfig } from "../lib/validation";
import { useDatasetPreviewDialogStore } from "../stores/dataset-preview-dialog-store";
@ -88,7 +89,10 @@ export function useTrainingActions() {
});
}
const needsReview = check.requires_manual_mapping || check.detected_format === "custom_heuristic";
const isRawFormat = isRawTextDatasetFormat(config.datasetFormat);
const needsReview =
!isRawFormat &&
(check.requires_manual_mapping || check.detected_format === "custom_heuristic");
if (needsReview && !hasManualMapping(config, isVlm, isAudio)) {
// Pre-fill from suggested_mapping or VLM detected columns
const hint: Record<string, string> = {};

48
studio/frontend/src/features/training/lib/training-methods.ts Normal file
View file

@ -0,0 +1,48 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import type { DatasetFormat, TrainingMethod } from "@/types/training";
const BACKEND_TRAINING_TYPE: Record<TrainingMethod, string> = {
qlora: "LoRA/QLoRA",
lora: "LoRA/QLoRA",
full: "Full Finetuning",
cpt: "Continued Pretraining",
};
const TRAINING_METHOD_LABELS: Record<TrainingMethod, string> = {
qlora: "QLoRA",
lora: "LoRA",
full: "Full",
cpt: "CPT",
};
export function toBackendTrainingType(trainingMethod: TrainingMethod): string {
return BACKEND_TRAINING_TYPE[trainingMethod];
}
export function getTrainingMethodLabel(
trainingMethod: TrainingMethod | string,
): string {
if (Object.prototype.hasOwnProperty.call(TRAINING_METHOD_LABELS, trainingMethod)) {
return TRAINING_METHOD_LABELS[trainingMethod as TrainingMethod];
}
return TRAINING_METHOD_LABELS.full;
}
export function parseBackendTrainingMethod(
trainingType: unknown,
loadIn4Bit: unknown,
): TrainingMethod {
if (trainingType === "Continued Pretraining") return "cpt";
if (trainingType === "LoRA/QLoRA") {
return loadIn4Bit ? "qlora" : "lora";
}
return "full";
}
export function isRawTextDatasetFormat(
datasetFormat: DatasetFormat,
): boolean {
return datasetFormat === "raw";
}

207
studio/frontend/src/features/training/stores/training-config-store.ts
View file

@ -1,15 +1,17 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { DEFAULT_HYPERPARAMS, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS } from "@/config/training";
import { CPT_TARGET_MODULES, DEFAULT_HYPERPARAMS, LR_DEFAULT_CPT, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS, TARGET_MODULES } from "@/config/training";
import { authFetch } from "@/features/auth";
import { isAdapterMethod } from "@/types/training";
import type { DatasetFormat } from "@/types/training";
import type { ModelType, StepNumber, TrainingMethod } from "@/types/training";
import { create } from "zustand";
import { persist } from "zustand/middleware";
import { checkDatasetFormat } from "../api/datasets-api";
import { checkVisionModel, getModelConfig } from "../api/models-api";
import { mapBackendModelConfigToTrainingPatch } from "../lib/model-defaults";
import { isRawTextDatasetFormat } from "../lib/training-methods";
import type { BackendModelConfig } from "../api/models-api";
import type { TrainingConfigState, TrainingConfigStore } from "../types/config";
@ -108,6 +110,11 @@ let _learningRateManuallySet = false;
// setTrainingMethod can restore it when switching back from full to adapter.
let _yamlLearningRate: number | undefined = undefined;
// Track whether entering CPT auto-forced datasetFormat="raw" so that
// leaving CPT can restore the prior user-visible format.
let _datasetFormatBeforeCpt: DatasetFormat | null = null;
let _datasetFormatAutoForcedByCpt = false;
const NON_PERSISTED_STATE_KEYS: ReadonlySet<keyof TrainingConfigState> = new Set([
"modelType",
"isCheckingVision",
@ -156,6 +163,123 @@ function canProceedForStep(state: TrainingConfigState): boolean {
}
}
type TrainingMethodStatePatch = Partial<
Pick<
TrainingConfigState,
| "trainingMethod"
| "learningRate"
| "loraRank"
| "loraAlpha"
| "loraVariant"
| "targetModules"
| "datasetFormat"
| "trainOnCompletions"
>
>;
function getCptTrainingPatch(): TrainingMethodStatePatch {
return {
loraRank: 128,
loraAlpha: 32,
loraVariant: "rslora",
targetModules: CPT_TARGET_MODULES,
datasetFormat: "raw",
trainOnCompletions: false,
};
}
function getCptModelDefaultsPatch(): TrainingMethodStatePatch {
return {
...getCptTrainingPatch(),
learningRate: LR_DEFAULT_CPT,
};
}
function getRestoreFromCptPatch(): TrainingMethodStatePatch {
return {
loraRank: DEFAULT_HYPERPARAMS.loraRank,
loraAlpha: DEFAULT_HYPERPARAMS.loraAlpha,
loraVariant: DEFAULT_HYPERPARAMS.loraVariant,
targetModules: TARGET_MODULES,
};
}
function clearCptDatasetFormatTracking(): void {
_datasetFormatBeforeCpt = null;
_datasetFormatAutoForcedByCpt = false;
}
function recordCptDatasetFormatOverride(currentDatasetFormat: DatasetFormat): void {
if (isRawTextDatasetFormat(currentDatasetFormat)) {
clearCptDatasetFormatTracking();
return;
}
_datasetFormatBeforeCpt = currentDatasetFormat;
_datasetFormatAutoForcedByCpt = true;
}
function getRestoreDatasetFormatFromCptPatch(): TrainingMethodStatePatch {
if (!_datasetFormatAutoForcedByCpt || _datasetFormatBeforeCpt == null) {
clearCptDatasetFormatTracking();
return {};
}
const previousDatasetFormat = _datasetFormatBeforeCpt;
clearCptDatasetFormatTracking();
return { datasetFormat: previousDatasetFormat };
}
function resolveTrainingMethodLearningRate(
prevMethod: TrainingMethod,
nextMethod: TrainingMethod,
): number | undefined {
if (_learningRateManuallySet) {
return undefined;
}
const wasCpt = prevMethod === "cpt";
const wasAdapter = isAdapterMethod(prevMethod);
const nowAdapter = isAdapterMethod(nextMethod);
if (nextMethod === "cpt") {
return LR_DEFAULT_CPT;
}
if (wasCpt && nowAdapter) {
return _yamlLearningRate ?? LR_DEFAULT_LORA;
}
if (wasAdapter && nowAdapter) {
return undefined;
}
return nowAdapter ? _yamlLearningRate ?? LR_DEFAULT_LORA : LR_DEFAULT_FULL;
}
function buildTrainingMethodPatch(
prevMethod: TrainingMethod,
nextMethod: TrainingMethod,
currentDatasetFormat: DatasetFormat,
): TrainingMethodStatePatch {
const patch: TrainingMethodStatePatch = { trainingMethod: nextMethod };
if (prevMethod !== "cpt" && nextMethod === "cpt") {
recordCptDatasetFormatOverride(currentDatasetFormat);
Object.assign(patch, getCptTrainingPatch());
}
if (prevMethod === "cpt" && nextMethod !== "cpt") {
Object.assign(
patch,
getRestoreFromCptPatch(),
getRestoreDatasetFormatFromCptPatch(),
);
}
const learningRate = resolveTrainingMethodLearningRate(prevMethod, nextMethod);
if (learningRate !== undefined) {
patch.learningRate = learningRate;
}
return patch;
}
export const useTrainingConfigStore = create<TrainingConfigStore>()(
persist(
(set, get) => {
@ -216,11 +340,14 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
// Auto-select training method based on model size vs GPU memory.
// If model_size * 1.5 * context_scale fits in free VRAM, use LoRA 16-bit.
// Otherwise use QLoRA 4-bit.
// Auto-select LoRA vs QLoRA based on GPU memory.
// Skip if user has manually chosen CPT -- don't override it.
const modelSizeBytes = modelDetails.model_size_bytes;
if (modelSizeBytes && modelSizeBytes > 0) {
if (modelSizeBytes && modelSizeBytes > 0 && get().trainingMethod !== "cpt") {
void autoSelectTrainingMethod(modelSizeBytes, patch.contextLength ?? get().contextLength)
.then((method) => {
if (get().selectedModel !== modelName) return;
if (get().trainingMethod === "cpt") return;
if (method) {
const lrPatch = !_learningRateManuallySet && !modelConfigHasLR
? { learningRate: method === "full" ? LR_DEFAULT_FULL : LR_DEFAULT_LORA }
@ -230,8 +357,16 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
});
}
// Preserve CPT hyperparams: YAML adapter defaults (r/alpha/targets/LR)
// are tuned for standard LoRA and would otherwise clobber CPT settings.
const cptOverrides =
get().trainingMethod === "cpt"
? getCptModelDefaultsPatch()
: {};
set({
...patch,
...cptOverrides,
modelType: inferredModelType,
isVisionModel: modelDetails.is_vision,
isEmbeddingModel: isEmbedding,
@ -396,29 +531,14 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
void loadAndApplyModelDefaults(state.selectedModel);
},
setTrainingMethod: (trainingMethod) => {
if (_learningRateManuallySet) {
set({ trainingMethod });
return;
}
const prev = get().trainingMethod;
const wasAdapter = isAdapterMethod(prev);
const nowAdapter = isAdapterMethod(trainingMethod);
// qlora <-> lora: same LR range, don't touch learning rate
if (wasAdapter && nowAdapter) {
set({ trainingMethod });
return;
}
// Category changed (adapter <-> full)
if (nowAdapter) {
// Switching TO adapter: restore YAML LR if available
set({ trainingMethod, learningRate: _yamlLearningRate ?? LR_DEFAULT_LORA });
} else {
// Switching TO full: no YAML full-LR exists, use constant
set({ trainingMethod, learningRate: LR_DEFAULT_FULL });
}
const state = get();
set(
buildTrainingMethodPatch(
state.trainingMethod,
trainingMethod,
state.datasetFormat,
),
);
},
setHfToken: (hfToken) =>
set({ hfToken: hfToken.trim().replace(/^["']+|["']+$/g, "") }),
@ -448,7 +568,26 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
runDatasetCheck(uploadedFile, "train");
}
},
setDatasetFormat: (datasetFormat) => set({ datasetFormat }),
setDatasetFormat: (datasetFormat) =>
set((state) => {
if (state.trainingMethod === "cpt") {
if (isRawTextDatasetFormat(datasetFormat)) {
clearCptDatasetFormatTracking();
}
return {
datasetFormat: "raw",
trainOnCompletions: false,
};
}
return {
datasetFormat,
trainOnCompletions:
isRawTextDatasetFormat(datasetFormat)
? false
: state.trainOnCompletions,
};
}),
setDataset: (dataset) => {
_datasetCheckController?.abort();
_datasetCheckController = null;
@ -566,6 +705,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
_learningRateManuallySet = true;
set({ learningRate });
},
setEmbeddingLearningRate: (embeddingLearningRate) =>
set({ embeddingLearningRate }),
setOptimizerType: (optimizerType) => set({ optimizerType }),
setLrSchedulerType: (lrSchedulerType) => set({ lrSchedulerType }),
setLoraRank: (loraRank) => set({ loraRank }),
@ -608,6 +749,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
_trainOnCompletionsManuallySet = false;
_learningRateManuallySet = false;
_yamlLearningRate = undefined;
clearCptDatasetFormatTracking();
set(initialState);
},
resetToModelDefaults: () => {
@ -629,7 +771,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
},
{
name: "unsloth_training_config_v1",
version: 9,
version: 10,
migrate: (persisted, version) => {
const s = persisted as Record<string, unknown>;
if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
@ -665,6 +807,17 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
s.weightDecay = DEFAULT_HYPERPARAMS.weightDecay;
}
}
if (version < 10 && s.trainingMethod === "cpt") {
// Backfill CPT defaults for state persisted before they existed.
s.loraRank = 128;
s.loraAlpha = 32;
s.loraVariant = "rslora";
s.targetModules = CPT_TARGET_MODULES;
s.datasetFormat = "raw";
if (s.learningRate == null || s.learningRate === LR_DEFAULT_LORA) {
s.learningRate = LR_DEFAULT_CPT;
}
}
return s as unknown as TrainingConfigStore;
},
partialize: partializePersistedState,

2
studio/frontend/src/features/training/types/api.ts
View file

@ -21,6 +21,8 @@ export interface TrainingStartRequest {
custom_format_mapping?: Record<string, unknown> | null;
num_epochs: number;
learning_rate: string;
/** Optional CPT embedding LR. If omitted, backend uses lr/10; typical range is 2x-10x smaller than main LR. */
embedding_learning_rate?: number | null;
batch_size: number;
gradient_accumulation_steps: number;
warmup_steps: number | null;

2
studio/frontend/src/features/training/types/config.ts
View file

@ -41,6 +41,7 @@ export interface TrainingConfigState {
epochs: number;
contextLength: number;
learningRate: number;
embeddingLearningRate: number | null;
optimizerType: string;
lrSchedulerType: string;
loraRank: number;
@ -115,6 +116,7 @@ export interface TrainingConfigActions {
setEpochs: (epochs: number) => void;
setContextLength: (length: number) => void;
setLearningRate: (rate: number) => void;
setEmbeddingLearningRate: (rate: number | null) => void;
setOptimizerType: (value: string) => void;
setLrSchedulerType: (value: string) => void;
setLoraRank: (rank: number) => void;

16
studio/frontend/src/lib/vram.ts
View file

@ -57,7 +57,12 @@ export type VramFitStatus = "fits" | "tight" | "exceeds";
*/
export const FP16_LOADING_BYTES = 2.0;
export type TrainingMethod = "qlora" | "lora" | "full";
export type TrainingMethod = "qlora" | "lora" | "full" | "cpt";
function usesQuantizedLoading(method: TrainingMethod, modelId?: string): boolean {
if (method === "qlora") return true;
return method === "cpt" && (modelId ?? "").toLowerCase().includes("4bit");
}
/**
* Estimate VRAM (GB) needed to load a model with Unsloth.
@ -66,15 +71,18 @@ export type TrainingMethod = "qlora" | "lora" | "full";
* - QLoRA : 4-bit quantized via bnb -> 0.90 bytes/param (calibrated)
* - LoRA : fp16 -> 2.0 bytes/param (theoretical)
* - Full : fp16 -> 2.0 bytes/param (theoretical)
* - CPT : fp16 LoRA (16-bit base) -> 2.0 bytes/param (theoretical)
*
* Formula: totalParams * bytesPerParam + 1.4 GB overhead
*/
export function estimateLoadingVram(
totalParams: number,
method: TrainingMethod = "qlora",
modelId?: string,
): number {
const bytesPerParam =
method === "qlora" ? BNB_4BIT_LOADING_BYTES : FP16_LOADING_BYTES;
const bytesPerParam = usesQuantizedLoading(method, modelId)
? BNB_4BIT_LOADING_BYTES
: FP16_LOADING_BYTES;
const gb = (totalParams / 1e9) * bytesPerParam + LOADING_OVERHEAD_GB;
return Math.round(gb * 10) / 10;
}
@ -119,7 +127,7 @@ export function buildModelVramMap(
continue;
}
const est = estimateLoadingVram(model.totalParams, method);
const est = estimateLoadingVram(model.totalParams, method, model.id);
const status = gpu.available ? checkVramFit(est, gpu.memoryTotalGb) : null;
map.set(model.id, { est, status });
}

6
studio/frontend/src/types/training.ts
View file

@ -2,14 +2,14 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export type ModelType = "vision" | "audio" | "embeddings" | "text";
export type TrainingMethod = "qlora" | "lora" | "full";
export type TrainingMethod = "qlora" | "lora" | "full" | "cpt";
export function isAdapterMethod(method: TrainingMethod): boolean {
return method === "lora" || method === "qlora";
return method === "lora" || method === "qlora" || method === "cpt";
}
export type StepNumber = 1 | 2 | 3 | 4 | 5;
export type DatasetSource = "huggingface" | "upload";
export type DatasetFormat = "auto" | "alpaca" | "chatml" | "sharegpt";
export type DatasetFormat = "auto" | "alpaca" | "chatml" | "sharegpt" | "raw";
export type GradientCheckpointing = "none" | "true" | "unsloth" | "mlx";
export interface WizardState {

3
tests/studio/install/test_install_llama_prebuilt_logic.py
View file

@ -754,6 +754,9 @@ def write_linux_install_shape(install_dir: Path) -> None:
(install_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8")
(runtime_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8")
(runtime_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8")
# Mirror the runtime payload health groups in install_llama_prebuilt.py:
# libllama-common.so* was added by PR #5135 and is required.
(runtime_dir / "libllama-common.so.0").write_bytes(b"DLL")
(runtime_dir / "libllama.so.0").write_bytes(b"DLL")
(runtime_dir / "libggml.so.0").write_bytes(b"DLL")
(runtime_dir / "libggml-base.so.0").write_bytes(b"DLL")

149
tests/studio/install/test_rocm_support.py
View file

@ -53,6 +53,32 @@ _has_usable_nvidia_gpu = stack_mod._has_usable_nvidia_gpu
_ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX
def _extract_sh_function_body(source: str, name: str) -> str:
"""Return the body of a shell function from `source` by brace matching.
Used by structural tests that need to assert ordering of helper
calls inside a specific function rather than across the whole
install.sh file.
"""
needle = f"{name}() {{"
start = source.find(needle)
if start < 0:
return ""
depth = 0
i = start + len(needle) - 1 # land on the opening brace
n = len(source)
while i < n:
ch = source[i]
if ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
return source[start : i + 1]
i += 1
return source[start:]
# ── Helper: build HostInfo for different scenarios ──────────────────────────
@ -561,12 +587,13 @@ class TestEnsureRocmTorch:
_ensure_rocm_torch()
mock_pip.assert_not_called()
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_cpu_torch_gets_rocm_reinstall(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""CPU-only torch on ROCm host should trigger reinstall."""
mock_probe = MagicMock()
@ -575,12 +602,11 @@ class TestEnsureRocmTorch:
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
# Should call pip_install twice: once for torch, once for bitsandbytes
assert mock_pip.call_count == 2
torch_call = mock_pip.call_args_list[0]
assert "rocm7.1" in str(torch_call)
bnb_call = mock_pip.call_args_list[1]
assert "bitsandbytes" in str(bnb_call)
# Should install torch via pip_install and bitsandbytes via pip_install_try.
assert mock_pip.call_count == 1
assert "rocm7.1" in str(mock_pip.call_args_list[0])
assert mock_pip_try.call_count >= 1
assert "bitsandbytes" in str(mock_pip_try.call_args_list[0])
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@ -642,12 +668,13 @@ class TestEnsureRocmTorch:
torch_call = mock_pip.call_args_list[0]
assert "rocm7.1" in str(torch_call)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_probe_timeout_triggers_reinstall(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""Probe subprocess timeout should not crash; should proceed to reinstall."""
with patch("os.path.isdir", return_value = True):
@ -656,8 +683,10 @@ class TestEnsureRocmTorch:
):
_ensure_rocm_torch()
# If probe times out, the function should treat torch as unusable and reinstall
assert mock_pip.call_count == 2
# both torch (via pip_install) and bitsandbytes (via pip_install_try).
assert mock_pip.call_count == 1
assert "rocm7.1" in str(mock_pip.call_args_list[0])
assert mock_pip_try.call_count >= 1
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@ -857,15 +886,33 @@ class TestInstallShStructure:
assert "rocm" in source.lower()
def test_cuda_precedence(self):
"""ROCm detection should only run when nvidia-smi is absent."""
"""ROCm detection should only run when nvidia-smi is absent.
install.sh defines _has_amd_rocm_gpu and _has_usable_nvidia_gpu
helpers near each other (file-position order has no semantic
meaning), so check the runtime ordering inside
get_torch_index_url instead: NVIDIA branch runs first and the
AMD/ROCm branch only fires inside the `if [ -z "$_smi" ]`
block.
"""
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text()
# The ROCm block should be inside the "if [ -z "$_smi" ]" branch
smi_block_start = source.find('if [ -z "$_smi" ]')
rocm_block_start = source.find("amd-smi")
body = _extract_sh_function_body(source, "get_torch_index_url")
nvidia_call = body.find("_has_usable_nvidia_gpu")
no_nvidia_branch = body.find('if [ -z "$_smi" ]')
rocm_call = body.find("_has_amd_rocm_gpu")
assert (
smi_block_start < rocm_block_start
), "ROCm detection should be inside the 'no nvidia-smi' branch"
nvidia_call >= 0
), "get_torch_index_url should call _has_usable_nvidia_gpu"
assert (
no_nvidia_branch >= 0
), "get_torch_index_url should gate ROCm on no-nvidia-smi"
assert (
rocm_call > no_nvidia_branch
), "ROCm detection should sit inside the 'no nvidia-smi' branch"
assert (
nvidia_call < no_nvidia_branch
), "NVIDIA detection should run before the no-nvidia-smi branch"
def test_bitsandbytes_amd_install(self):
"""install.sh should install bitsandbytes for AMD when ROCm detected."""
@ -963,16 +1010,32 @@ class TestLiveRegression:
if not shutil.which("nvidia-smi"):
pytest.skip("No nvidia-smi available")
sh_path = PACKAGE_ROOT / "install.sh"
# Extract just the function (don't source the whole installer)
result = subprocess.run(
# Skip if nvidia-smi exists but does not actually list a GPU on this
# host (containers occasionally ship the binary without a driver).
check = subprocess.run(
[
"bash",
"-c",
f"eval \"$(sed -n '/^get_torch_index_url()/,/^}}/p' '{sh_path}')\"; "
"get_torch_index_url",
"nvidia-smi -L 2>/dev/null | "
"awk '/^GPU[[:space:]]+[0-9]+:/{f=1} END{exit !f}'",
],
capture_output = True,
)
if check.returncode != 0:
pytest.skip("nvidia-smi is on PATH but no GPU is listed")
sh_path = PACKAGE_ROOT / "install.sh"
# get_torch_index_url calls _has_usable_nvidia_gpu and
# _has_amd_rocm_gpu, so all three function definitions must be
# in scope when we eval the extract.
extract_cmd = (
f"sed -n '/^_has_amd_rocm_gpu()/,/^}}$/p; "
f"/^_has_usable_nvidia_gpu()/,/^}}$/p; "
f"/^get_torch_index_url()/,/^}}$/p' '{sh_path}'"
)
result = subprocess.run(
["bash", "-c", f'eval "$({extract_cmd})"; get_torch_index_url'],
capture_output = True,
text = True,
timeout = 30,
)
@ -988,19 +1051,23 @@ class TestLiveRegression:
# Load worker.py module
_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py"
# The wheel-probe subprocess was hoisted out of worker.py into wheel_utils
# during the wheel-resolver refactor; the probe script literal lives there.
_WHEEL_UTILS_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "wheel_utils.py"
class TestWorkerRocmMambaSsm:
"""Verify worker.py Mamba/SSM install logic on ROCm."""
def test_probe_returns_hip_version_field(self):
"""_probe_causal_conv1d_env probe script should include hip_version."""
source = _WORKER_PATH.read_text()
assert "hip_version" in source
"""The wheel probe should include hip_version, and worker.py should
consume it."""
assert "hip_version" in _WHEEL_UTILS_PATH.read_text()
assert "hip_version" in _WORKER_PATH.read_text()
def test_probe_script_has_getattr_hip(self):
"""Probe script should use getattr for torch.version.hip (safe on CUDA)."""
source = _WORKER_PATH.read_text()
source = _WHEEL_UTILS_PATH.read_text()
assert "getattr(torch.version, 'hip', None)" in source
def test_direct_wheel_url_returns_none_without_cuda_major(self):
@ -1216,27 +1283,45 @@ class TestHardwareAmdBranching:
assert "from . import amd" in source
def test_hardware_branches_on_is_rocm_for_utilization(self):
"""get_gpu_utilization should check IS_ROCM before choosing backend."""
"""get_gpu_utilization should dispatch to amd.py via _smi_query
when IS_ROCM, and the dispatcher itself must check IS_ROCM and
import the amd backend."""
hw_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
)
source = hw_path.read_text()
# Find the get_gpu_utilization function
func_start = source.find("def get_gpu_utilization")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "IS_ROCM" in func_body
assert "amd.get_primary_gpu_utilization" in func_body
assert '_smi_query("get_primary_gpu_utilization"' in func_body
smi = source[
source.find("def _smi_query") : source.find(
"\ndef ", source.find("def _smi_query") + 1
)
]
assert "IS_ROCM" in smi
assert "from . import amd" in smi
def test_hardware_branches_on_is_rocm_for_visible(self):
"""get_visible_gpu_utilization should check IS_ROCM."""
"""get_visible_gpu_utilization should dispatch to amd.py via
_smi_query when IS_ROCM."""
hw_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
)
source = hw_path.read_text()
func_start = source.find("def get_visible_gpu_utilization")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "IS_ROCM" in func_body
assert "amd.get_visible_gpu_utilization" in func_body
# The dispatcher call may wrap onto multiple lines; allow whitespace
# between the open paren and the literal func name argument.
import re as _re
assert _re.search(r'_smi_query\(\s*"get_visible_gpu_utilization"', func_body)
smi = source[
source.find("def _smi_query") : source.find(
"\ndef ", source.find("def _smi_query") + 1
)
]
assert "IS_ROCM" in smi
assert "from . import amd" in smi
def test_hardware_branches_on_is_rocm_for_physical_count(self):
"""get_physical_gpu_count should try amd.py when IS_ROCM."""
@ -1247,7 +1332,7 @@ class TestHardwareAmdBranching:
func_start = source.find("def get_physical_gpu_count")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert "IS_ROCM" in func_body
assert "amd.get_physical_gpu_count" in func_body
assert "from . import amd" in func_body
# =============================================================================

381
tests/studio/test_hardware_dispatch_matrix.py Normal file
View file

@ -0,0 +1,381 @@
# SPDX-License-Identifier: AGPL-3.0-only
"""
Comprehensive hardware dispatch matrix for Studio.
Drives every supported hardware profile from a single test host by
spoofing platform / torch.cuda / torch.xpu / sys.modules['mlx'] so we
can exercise the CUDA, ROCm, XPU, MLX, and CPU dispatch paths
deterministically without real hardware.
Profiles checked:
nvidia_cuda Linux x86_64 + torch.cuda.is_available()=True,
torch.version.hip=None
amd_rocm Linux x86_64 + torch.cuda.is_available()=True,
torch.version.hip="6.1" (PyTorch ROCm aliases
torch.cuda.* over HIP)
intel_xpu Linux x86_64 + torch.cuda off, torch.xpu.is_available()=True
apple_silicon_mlx Darwin arm64 + cuda off + xpu off + mlx importable
apple_silicon_no_mlx Darwin arm64 + everything off (no mlx pkg)
linux_arm64_with_mlx Linux arm64 + mlx importable -- gate must NOT activate
(canary against accidental Linux-arm64 hijack)
cpu_only Linux x86_64 + nothing -- pure CPU fallback
For each profile we assert three contracts:
1. ``unsloth._IS_MLX`` (re-evaluated under the spoof).
2. ``utils.hardware.detect_hardware()`` ``DeviceType`` and ``IS_ROCM``.
3. ``utils.hardware.is_apple_silicon()``.
Add a row to ``PROFILES`` to extend coverage; tests parametrize over it
automatically. No real hardware required.
"""
from __future__ import annotations
import importlib
import importlib.machinery
import importlib.util
import sys
import types
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
# ---------------------------------------------------------------------------
# Profile definition
# ---------------------------------------------------------------------------
@dataclass
class HardwareProfile:
name: str
system: str # platform.system() value
machine: str # platform.machine() value
cuda_available: bool # torch.cuda.is_available() value
hip_version: Optional[
str
] # torch.version.hip; None for NVIDIA, "6.1" etc. for ROCm
xpu_available: bool # torch.xpu.is_available() value
has_mlx: bool # whether to inject a fake mlx into sys.modules
mps_available: bool # torch.backends.mps.is_available() value
expect_is_mlx: bool # unsloth._IS_MLX
expect_device_type: (
str # Studio DeviceType (uppercased name: "CUDA"/"XPU"/"MLX"/"CPU")
)
expect_is_rocm: bool # Studio IS_ROCM
expect_apple_silicon: bool # Studio is_apple_silicon()
extra_notes: str = ""
PROFILES = [
HardwareProfile(
name = "nvidia_cuda",
system = "Linux",
machine = "x86_64",
cuda_available = True,
hip_version = None,
xpu_available = False,
has_mlx = False,
mps_available = False,
expect_is_mlx = False,
expect_device_type = "CUDA",
expect_is_rocm = False,
expect_apple_silicon = False,
),
HardwareProfile(
name = "amd_rocm",
system = "Linux",
machine = "x86_64",
cuda_available = True,
hip_version = "6.1",
xpu_available = False,
has_mlx = False,
mps_available = False,
expect_is_mlx = False,
expect_device_type = "CUDA",
expect_is_rocm = True,
expect_apple_silicon = False,
extra_notes = "PyTorch ROCm reuses torch.cuda.* over HIP; "
"Studio still uses DeviceType.CUDA but flips IS_ROCM=True.",
),
HardwareProfile(
name = "intel_xpu",
system = "Linux",
machine = "x86_64",
cuda_available = False,
hip_version = None,
xpu_available = True,
has_mlx = False,
mps_available = False,
expect_is_mlx = False,
expect_device_type = "XPU",
expect_is_rocm = False,
expect_apple_silicon = False,
),
HardwareProfile(
name = "apple_silicon_mlx",
system = "Darwin",
machine = "arm64",
cuda_available = False,
hip_version = None,
xpu_available = False,
has_mlx = True,
mps_available = True,
expect_is_mlx = True,
expect_device_type = "MLX",
expect_is_rocm = False,
expect_apple_silicon = True,
),
HardwareProfile(
name = "apple_silicon_no_mlx",
system = "Darwin",
machine = "arm64",
cuda_available = False,
hip_version = None,
xpu_available = False,
has_mlx = False,
mps_available = True,
expect_is_mlx = False,
expect_device_type = "CPU",
expect_is_rocm = False,
expect_apple_silicon = True,
extra_notes = "Mac without mlx falls through to CPU (chat-only).",
),
HardwareProfile(
name = "linux_arm64_with_mlx",
system = "Linux",
machine = "arm64",
cuda_available = False,
hip_version = None,
xpu_available = False,
has_mlx = True,
mps_available = False,
expect_is_mlx = False,
expect_device_type = "CPU",
expect_is_rocm = False,
expect_apple_silicon = False,
extra_notes = "Canary: Linux ARM64 with mlx package installed must NOT "
"trigger MLX dispatch; the system check is what guards it.",
),
HardwareProfile(
name = "cpu_only",
system = "Linux",
machine = "x86_64",
cuda_available = False,
hip_version = None,
xpu_available = False,
has_mlx = False,
mps_available = False,
expect_is_mlx = False,
expect_device_type = "CPU",
expect_is_rocm = False,
expect_apple_silicon = False,
),
]
PROFILE_IDS = [p.name for p in PROFILES]
# ---------------------------------------------------------------------------
# Spoofing helpers
# ---------------------------------------------------------------------------
@pytest.fixture
def spoof_hardware(monkeypatch):
"""Return a function that applies a HardwareProfile to the live process.
Idempotent: each call re-applies the profile. Cleanup happens
automatically when the test exits via monkeypatch.
"""
def _apply(profile: HardwareProfile) -> None:
import platform
import torch
# platform spoof (used by both the unsloth gate and Studio's helpers)
monkeypatch.setattr(platform, "system", lambda: profile.system)
monkeypatch.setattr(platform, "machine", lambda: profile.machine)
# torch.cuda.is_available
monkeypatch.setattr(torch.cuda, "is_available", lambda: profile.cuda_available)
# torch.version.hip — None on NVIDIA, "6.1" etc. on ROCm
torch_version = torch.version
monkeypatch.setattr(torch_version, "hip", profile.hip_version, raising = False)
# torch.xpu.is_available + get_device_name -- detect_hardware reads both.
# Real torch.xpu.get_device_name requires the XPU-compiled torch build,
# so always stub it under the spoof to keep tests hardware-agnostic.
if hasattr(torch, "xpu"):
monkeypatch.setattr(
torch.xpu, "is_available", lambda: profile.xpu_available
)
monkeypatch.setattr(
torch.xpu,
"get_device_name",
lambda i = 0: "Intel XPU (stub)",
raising = False,
)
elif profile.xpu_available:
xpu_stub = types.SimpleNamespace(
is_available = lambda: True,
get_device_name = lambda i = 0: "Intel XPU (stub)",
)
monkeypatch.setattr(torch, "xpu", xpu_stub, raising = False)
# torch.backends.mps.is_available
if hasattr(torch.backends, "mps"):
monkeypatch.setattr(
torch.backends.mps, "is_available", lambda: profile.mps_available
)
# mlx + mlx.core in sys.modules
if profile.has_mlx:
fake_mlx = types.ModuleType("mlx")
fake_mlx.__spec__ = importlib.machinery.ModuleSpec("mlx", loader = None)
fake_mlx.__path__ = []
fake_mlx_core = types.ModuleType("mlx.core")
fake_mlx.core = fake_mlx_core
monkeypatch.setitem(sys.modules, "mlx", fake_mlx)
monkeypatch.setitem(sys.modules, "mlx.core", fake_mlx_core)
else:
monkeypatch.delitem(sys.modules, "mlx", raising = False)
monkeypatch.delitem(sys.modules, "mlx.core", raising = False)
real_find_spec = importlib.util.find_spec
def _no_mlx(name, *args, **kwargs):
if name == "mlx":
return None
return real_find_spec(name, *args, **kwargs)
monkeypatch.setattr(importlib.util, "find_spec", _no_mlx)
return _apply
def _evaluate_unsloth_is_mlx_gate() -> bool:
"""Re-evaluate the exact expression from unsloth/__init__.py:20-24."""
import importlib.util
import platform
return (
platform.system() == "Darwin"
and platform.machine() == "arm64"
and importlib.util.find_spec("mlx") is not None
)
def _import_studio_hardware_module():
"""Lazy-load Studio's hardware module under the bare-imports layout."""
if str(STUDIO_BACKEND) not in sys.path:
sys.path.insert(0, str(STUDIO_BACKEND))
# Force a fresh import so detect_hardware re-runs under the current spoofs.
sys.modules.pop("utils.hardware.hardware", None)
sys.modules.pop("utils.hardware", None)
from utils.hardware import hardware as hw # type: ignore
return hw
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS)
def test_unsloth_is_mlx_gate_matches_profile(profile, spoof_hardware):
"""The _IS_MLX expression in unsloth/__init__.py flips correctly per profile."""
spoof_hardware(profile)
actual = _evaluate_unsloth_is_mlx_gate()
assert actual is profile.expect_is_mlx, (
f"profile {profile.name}: expected _IS_MLX={profile.expect_is_mlx}, "
f"got {actual}. {profile.extra_notes}"
)
@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS)
def test_studio_detect_hardware_matches_profile(profile, spoof_hardware):
"""Studio's detect_hardware() routes to the right DeviceType per profile."""
spoof_hardware(profile)
hw = _import_studio_hardware_module()
detected = hw.detect_hardware()
expected = getattr(hw.DeviceType, profile.expect_device_type)
assert detected == expected, (
f"profile {profile.name}: expected {profile.expect_device_type}, "
f"got {detected!r}. {profile.extra_notes}"
)
assert hw.IS_ROCM is profile.expect_is_rocm, (
f"profile {profile.name}: expected IS_ROCM={profile.expect_is_rocm}, "
f"got {hw.IS_ROCM}"
)
@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS)
def test_studio_is_apple_silicon_matches_profile(profile, spoof_hardware):
"""Studio's is_apple_silicon() helper agrees with platform spoof."""
spoof_hardware(profile)
hw = _import_studio_hardware_module()
assert hw.is_apple_silicon() is profile.expect_apple_silicon, (
f"profile {profile.name}: expected is_apple_silicon={profile.expect_apple_silicon}, "
f"got {hw.is_apple_silicon()}"
)
# ---------------------------------------------------------------------------
# Negative-space tests: catch regressions where the dispatch order changes.
# ---------------------------------------------------------------------------
def test_cuda_takes_priority_over_mlx_when_both_available(spoof_hardware):
"""If both CUDA and MLX are available, Studio MUST pick CUDA. This is the
canary that protects every existing GPU user from being silently routed
to MLX after future refactors.
"""
profile = HardwareProfile(
name = "cuda_plus_mlx",
system = "Darwin",
machine = "arm64",
cuda_available = True,
hip_version = None,
xpu_available = False,
has_mlx = True,
mps_available = True,
expect_is_mlx = True,
expect_device_type = "CUDA",
expect_is_rocm = False,
expect_apple_silicon = True,
)
spoof_hardware(profile)
hw = _import_studio_hardware_module()
assert hw.detect_hardware() == hw.DeviceType.CUDA
def test_xpu_takes_priority_over_mlx_when_both_available(spoof_hardware):
"""XPU is selected over MLX in the dispatch order."""
profile = HardwareProfile(
name = "xpu_plus_mlx",
system = "Darwin",
machine = "arm64",
cuda_available = False,
hip_version = None,
xpu_available = True,
has_mlx = True,
mps_available = True,
expect_is_mlx = True,
expect_device_type = "XPU",
expect_is_rocm = False,
expect_apple_silicon = True,
)
spoof_hardware(profile)
hw = _import_studio_hardware_module()
assert hw.detect_hardware() == hw.DeviceType.XPU