diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index da5b11c60d..231566cc22 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -2,6 +2,7 @@ """ Export backend - handles model exporting in various formats """ +import json import logging import os from pathlib import Path @@ -200,6 +201,18 @@ class ExportBackend: logger.error(traceback.format_exc()) return False, f"Failed to load checkpoint: {str(e)}" + def _write_export_metadata(self, save_directory: str): + """Write export_metadata.json with base model info for Chat page discovery.""" + try: + base_model = get_base_model_from_lora(self.current_checkpoint) if self.current_checkpoint else None + metadata = {"base_model": base_model} + metadata_path = os.path.join(save_directory, "export_metadata.json") + with open(metadata_path, "w") as f: + json.dump(metadata, f, indent=2) + logger.info(f"Wrote export metadata to {metadata_path}") + except Exception as e: + logger.warning(f"Could not write export metadata: {e}") + def export_merged_model(self, save_directory: str, format_type: str = "16-bit (FP16)", @@ -244,6 +257,9 @@ class ExportBackend: self.current_tokenizer, save_method=save_method ) + + # Write export metadata so the Chat page can identify the base model + self._write_export_metadata(save_directory) logger.info(f"Model saved successfully to {save_directory}") # Push to hub if requested @@ -297,6 +313,9 @@ class ExportBackend: self.current_model.save_pretrained(save_directory) self.current_tokenizer.save_pretrained(save_directory) + + # Write export metadata so the Chat page can identify the base model + self._write_export_metadata(save_directory) logger.info(f"Model saved successfully to {save_directory}") # Push to hub if requested diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index e90e6c0c2a..1147c281b7 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -6,8 +6,10 @@ from unsloth.chat_templates import get_chat_template from transformers import TextStreamer from peft import PeftModel, PeftModelForCausalLM +import json import sys import torch +from pathlib import Path from typing import Optional, Union, Generator, Tuple from utils.models import ModelConfig, get_base_model_from_lora from utils.paths import is_model_cached @@ -112,7 +114,18 @@ class InferenceBackend: # In that case, load the real processor from the base model. from transformers import ProcessorMixin if not (isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")): + # For LoRA adapters, use the base model. For local merged exports, + # read export_metadata.json to find the original base model. processor_source = config.base_model if config.is_lora else config.identifier + if not config.is_lora and config.is_local: + _meta_path = Path(config.path) / "export_metadata.json" + try: + if _meta_path.exists(): + _meta = json.loads(_meta_path.read_text()) + if _meta.get("base_model"): + processor_source = _meta["base_model"] + except Exception: + pass logger.warning( f"FastVisionModel returned {type(processor).__name__} (no image_processor) " f"for '{model_name}' — loading proper processor from '{processor_source}'" diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 804baa33c0..43cb5ff1a1 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -343,7 +343,8 @@ class UnslothTrainer: custom_format_mapping: dict = None, subset: str = None, train_split: str = "train", - eval_split: str = None) -> Optional[tuple]: + eval_split: str = None, + eval_steps: float = 0.00) -> Optional[tuple]: """ Load and prepare dataset for training. @@ -358,6 +359,7 @@ class UnslothTrainer: dataset = None eval_dataset = None has_separate_eval_source = False # True if eval comes from a separate HF split + eval_enabled = eval_steps is not None and eval_steps > 0 if local_datasets: # Load local datasets @@ -410,23 +412,26 @@ class UnslothTrainer: 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: + if eval_enabled: + 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 + else: + print("Eval disabled (eval_steps <= 0), skipping eval split detection\n") if dataset is None: raise ValueError("No dataset provided") @@ -472,7 +477,7 @@ class UnslothTrainer: ) eval_dataset = eval_info["dataset"] print(f"Eval dataset formatted successfully\n") - elif not has_separate_eval_source: + elif eval_enabled and 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) @@ -543,7 +548,7 @@ class UnslothTrainer: def start_training(self, dataset: Dataset, eval_dataset: Dataset = None, - eval_steps: float = 0.01, + eval_steps: float = 0.00, output_dir: str = "./outputs", num_epochs: int = 3, learning_rate: float = 5e-5, @@ -743,12 +748,16 @@ class UnslothTrainer: # ========== EVAL CONFIGURATION ========== eval_dataset = training_args.get('eval_dataset', None) - eval_steps_val = training_args.get('eval_steps', 0.01) + eval_steps_val = training_args.get('eval_steps', 0.00) 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") + if eval_steps_val > 0: + 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(f"⚠️ Eval dataset provided but eval_steps={eval_steps_val} (disabled)\n") + print("To enable evaluation, set eval_steps > 0.0\n") else: print("No eval dataset — evaluation disabled\n") diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 6fe08c2b9e..9123d36b39 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -115,7 +115,7 @@ class TrainingBackend: subset: str = None, train_split: str = "train", eval_split: str = None, - eval_steps: float = 0.01, + eval_steps: float = 0.00, is_dataset_multimodal: bool = False) -> bool: """ Start training. @@ -223,6 +223,7 @@ class TrainingBackend: subset=subset, train_split=train_split, eval_split=eval_split, + eval_steps=eval_steps, ) # Unpack: load_and_format_dataset returns (dataset, eval_dataset) @@ -232,10 +233,6 @@ class TrainingBackend: 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 diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 5542db76d5..10d87d7825 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -57,10 +57,12 @@ class ModelDetails(BaseModel): class LoRAInfo(BaseModel): - """LoRA adapter information""" + """LoRA adapter or exported model information""" display_name: str = Field(..., description="Display name for the LoRA") - adapter_path: str = Field(..., description="Path to the LoRA adapter") + adapter_path: str = Field(..., description="Path to the LoRA adapter or exported model") base_model: Optional[str] = Field(None, description="Base model identifier") + source: Optional[str] = Field(None, description="'training' or 'exported'") + export_type: Optional[str] = Field(None, description="'lora' or 'merged' (for exports)") class LoRAScanResponse(BaseModel): diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 2b989e6a82..54de974100 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -21,7 +21,7 @@ class TrainingStartRequest(BaseModel): 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)") + eval_steps: float = Field(0.00, description="Fraction of total steps between evals (0-1)") @model_validator(mode="before") @classmethod diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 761c04d3e7..8c96ee5f28 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -18,6 +18,7 @@ from auth.authentication import get_current_subject try: from utils.models import ( scan_trained_loras, + scan_exported_models, load_model_defaults, get_base_model_from_lora, is_vision_model, @@ -32,6 +33,7 @@ except ImportError: sys.path.insert(0, str(parent_backend)) from utils.models import ( scan_trained_loras, + scan_exported_models, load_model_defaults, get_base_model_from_lora, is_vision_model, @@ -289,35 +291,45 @@ async def get_model_config( @router.get("/loras") async def scan_loras( outputs_dir: str = Query(default="./outputs", description="Directory to scan for LoRA adapters"), + exports_dir: str = Query(default="./exports", description="Directory to scan for exported models"), current_subject: str = Depends(get_current_subject), ): """ - Scan for trained LoRA adapters in the outputs directory. - - This endpoint wraps the backend scan_trained_loras function. + Scan for trained LoRA adapters and exported models. + + Returns both training outputs (from outputs_dir) and exported models + (from exports_dir) in a single list, distinguished by source field. """ try: - # Call backend scan function - trained_loras = scan_trained_loras(outputs_dir=outputs_dir) - - # Convert to LoRAInfo objects lora_list = [] + + # Scan training outputs + trained_loras = scan_trained_loras(outputs_dir=outputs_dir) for display_name, adapter_path in trained_loras: - # Get base model if available base_model = get_base_model_from_lora(adapter_path) - - lora_info = LoRAInfo( + lora_list.append(LoRAInfo( display_name=display_name, adapter_path=adapter_path, - base_model=base_model - ) - lora_list.append(lora_info) - + base_model=base_model, + source="training", + )) + + # Scan exported models (merged, LoRA, base — skips GGUF) + exported = scan_exported_models(exports_dir=exports_dir) + for display_name, model_path, export_type, base_model in exported: + lora_list.append(LoRAInfo( + display_name=display_name, + adapter_path=model_path, + base_model=base_model, + source="exported", + export_type=export_type, + )) + return LoRAScanResponse( loras=lora_list, outputs_dir=outputs_dir ) - + except Exception as e: logger.error(f"Error scanning LoRAs: {e}", exc_info=True) raise HTTPException( diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 505fd35edd..006deb99c0 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -5,6 +5,7 @@ from .model_config import ( ModelConfig, is_vision_model, scan_trained_loras, + scan_exported_models, load_model_defaults, get_base_model_from_lora, load_model_config, @@ -17,6 +18,7 @@ __all__ = [ 'ModelConfig', 'is_vision_model', 'scan_trained_loras', + 'scan_exported_models', 'load_model_defaults', 'get_base_model_from_lora', 'load_model_config', diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index fdf89fce39..f7b95fad11 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -465,6 +465,90 @@ def scan_trained_loras(outputs_dir: str = "./outputs") -> List[Tuple[str, str]]: logger.error(f"Error scanning outputs folder: {e}") return [] +def scan_exported_models(exports_dir: str = "./exports") -> List[Tuple[str, str, str, Optional[str]]]: + """ + Scan exports folder for exported models (merged, LoRA, base). + Skips GGUF-only exports (not loadable by Unsloth inference backend). + + The exports directory is two levels deep: {run}/{checkpoint}/ + + Returns: + List of tuples: [(display_name, model_path, export_type, base_model), ...] + export_type: "lora" | "merged" + """ + results = [] + exports_path = Path(exports_dir) + + if not exports_path.exists(): + return results + + try: + for run_dir in exports_path.iterdir(): + if not run_dir.is_dir(): + continue + for checkpoint_dir in run_dir.iterdir(): + if not checkpoint_dir.is_dir(): + continue + + adapter_config = checkpoint_dir / "adapter_config.json" + config_file = checkpoint_dir / "config.json" + has_weights = ( + any(checkpoint_dir.glob("*.safetensors")) + or any(checkpoint_dir.glob("*.bin")) + ) + has_gguf = any(checkpoint_dir.glob("*.gguf")) + + base_model = None + export_type = None + + if adapter_config.exists(): + export_type = "lora" + try: + cfg = json.loads(adapter_config.read_text()) + base_model = cfg.get("base_model_name_or_path") + except Exception: + pass + elif config_file.exists() and has_weights: + export_type = "merged" + # Read base model from export_metadata.json (written at export time) + export_meta = checkpoint_dir / "export_metadata.json" + try: + if export_meta.exists(): + meta = json.loads(export_meta.read_text()) + base_model = meta.get("base_model") + except Exception: + pass + elif has_gguf: + # GGUF-only — not loadable by current inference backend + continue + else: + continue + + # Fallback: read base model from the original training run's + # adapter_config.json in ./outputs/{run_name}/ + if not base_model: + outputs_adapter_cfg = Path("./outputs") / run_dir.name / "adapter_config.json" + try: + if outputs_adapter_cfg.exists(): + cfg = json.loads(outputs_adapter_cfg.read_text()) + base_model = cfg.get("base_model_name_or_path") + except Exception: + pass + + display_name = f"{run_dir.name} / {checkpoint_dir.name}" + model_path = str(checkpoint_dir) + results.append((display_name, model_path, export_type, base_model)) + logger.debug(f"Found exported model: {display_name} ({export_type})") + + results.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True) + logger.info(f"Found {len(results)} exported models in {exports_dir}") + return results + + except Exception as e: + logger.error(f"Error scanning exports folder: {e}") + return [] + + def get_base_model_from_lora(lora_path: str) -> Optional[str]: """ Read the base model name from a LoRA adapter's config. diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index bf5f219558..d470034832 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -62,7 +62,7 @@ function ModelSelectorTrigger({ className={cn( "flex items-center gap-2 transition-colors", variant === "outline" && - "rounded-full border border-border/60 hover:bg-accent", + "rounded-full border border-border/60 hover:bg-accent", variant === "ghost" && "rounded-md hover:bg-accent", variant === "muted" && "rounded-md bg-muted hover:bg-muted/80", size === "sm" && "h-8 px-3 text-xs", @@ -183,9 +183,20 @@ export function ModelSelector({ all.set(model.id, model); } for (const lora of loraModels) { + // Strip "/ suffix" from display name (e.g. "foo_123/foo" → "foo_123") + const displayName = lora.name.includes("/") + ? lora.name.split("/")[0].trim() + : lora.name; + // Show type tag instead of base model name + const isExported = lora.source === "exported"; + const isMerged = lora.exportType === "merged"; + const tag = isExported + ? isMerged ? "Merged · Exported" : "LoRA" + : "LoRA"; all.set(lora.id, { ...lora, - description: lora.baseModel || lora.description, + name: displayName, + description: tag, }); } return all; diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 80e9dd0b6b..441fdf0f02 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -365,15 +365,26 @@ export function LoraModelPicker({
{index > 0 ?
: null} {baseModel} - {adapters.map((adapter) => ( - onSelect(adapter.id, { source: "lora", isLora: true })} - /> - ))} + {adapters.map((adapter) => { + const isExported = adapter.source === "exported"; + const isMerged = adapter.exportType === "merged"; + const tag = isExported + ? isMerged ? "Merged" : "LoRA" + : "LoRA"; + const meta = isExported ? `${tag} · Exported` : tag; + return ( + onSelect(adapter.id, { + source: isExported ? "exported" : "lora", + isLora: !isMerged, + })} + /> + ); + })}
)) )} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts index dcf110bfb7..b8df0f6c6c 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts @@ -10,10 +10,12 @@ export interface ModelOption { export interface LoraModelOption extends ModelOption { baseModel?: string; updatedAt?: number; + source?: "training" | "exported"; + exportType?: "lora" | "merged"; } export interface ModelSelectorChangeMeta { - source: "hub" | "lora"; + source: "hub" | "lora" | "exported"; isLora: boolean; } diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index da60328d40..33249a044a 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -103,7 +103,7 @@ export const DEFAULT_HYPERPARAMS = { warmupSteps: 5, maxSteps: 0, saveSteps: 0, - evalSteps: 0.01, + evalSteps: 0.00, packing: false, trainOnCompletions: false, gradientCheckpointing: "unsloth" as const, diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index c363704d61..6e37468e71 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -410,6 +410,8 @@ export function ChatPage(): ReactElement { name: lora.name, baseModel: lora.baseModel, updatedAt: lora.updatedAt, + source: lora.source, + exportType: lora.exportType, })), [lorasFromStore], ); diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 3dca94d7f5..1c8922e929 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -68,6 +68,8 @@ function toLoraSummary(lora: { display_name: string; adapter_path: string; base_model?: string | null; + source?: "training" | "exported" | null; + export_type?: "lora" | "merged" | null; }): ChatLoraSummary { const idTail = lora.adapter_path.split("/").filter(Boolean).at(-1) ?? ""; const updatedAt = @@ -78,6 +80,8 @@ function toLoraSummary(lora: { name: stripTrailingEpoch(lora.display_name), baseModel: lora.base_model || "Unknown base model", updatedAt, + source: lora.source ?? undefined, + exportType: lora.export_type ?? undefined, }; } diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index cadcad152d..003c3b3629 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -14,6 +14,8 @@ export interface BackendLoraInfo { display_name: string; adapter_path: string; base_model?: string | null; + source?: "training" | "exported" | null; + export_type?: "lora" | "merged" | null; } export interface ListLorasResponse { diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts index 47478e5aae..87eb4c565c 100644 --- a/studio/frontend/src/features/chat/types/runtime.ts +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -33,4 +33,6 @@ export interface ChatLoraSummary { name: string; baseModel: string; updatedAt?: number; + source?: "training" | "exported"; + exportType?: "lora" | "merged"; } diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index de82fb0fb5..144a1f34d7 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -130,11 +130,62 @@ export function ParamsSection(): ReactElement { className="md:min-h-[450px]" >
- {/* Max Steps */} -
-
+ {/* Max Steps */} +
+
+ + Max Steps + + + + + + Override total steps. Set 0 to use epochs instead.{" "} + + Read more + + + + + store.setMaxSteps(Number(e.target.value))} + min={0} + max={maxStepsSliderMax} + step={1} + className="w-16 text-right font-mono text-xs font-medium bg-muted/50 border border-border rounded-lg px-1.5 py-0.5 focus:outline-none focus:ring-1 focus:ring-primary/30 [&::-webkit-inner-spin-button]:appearance-none" + /> +
+ store.setMaxSteps(v)} + min={0} + max={maxStepsSliderMax} + step={1} + /> +

+ Total optimizer steps. Use 0 to run by epochs. +

+
+ + {/* Context length */} +
- Max Steps + Context Length
- store.setMaxSteps(v)} - min={0} - max={maxStepsSliderMax} - step={1} - /> -

- Total optimizer steps. Use 0 to run by epochs. -

-
- - {/* Context length */} -
- - Context Length - - - - - - Maximum number of tokens per training sample.{" "} - - Read more - - - - - -

- Max sequence length for training samples -

-
- - {/* Learning Rate */} -
- - Learning Rate - - - - - - Step size for weight updates. Lower values train slower but more - stably.{" "} - - Read more - - - - - store.setLearningRate(Number(e.target.value))} - className="w-full font-mono" - /> -

- Recommended: 2e-4 for LoRA, 2e-5 for full fine-tune -

-
- - {/* LoRA Settings */} - {isLora && ( -
-
- )} - {/* Training Hyperparams */} - - - - Training Hyperparameters - - - - - - Optimization - - - Schedule - - - Memory - - - - - - Optimization algorithm. 8-bit variants reduce memory usage. - Fused is recommended for vision models.{" "} - - Read more - - - } - > - - - - How the learning rate changes over training. Linear decays - steadily; cosine decays in a curve.{" "} - - Read more - - - } - > - - - - Samples processed per step. Higher uses more VRAM.{" "} - - Read more - - - } - value={store.batchSize} - onChange={store.setBatchSize} - min={1} - max={32} - step={1} - /> - - Simulates larger batch sizes without extra VRAM.{" "} - - Read more - - - } - value={store.gradientAccumulation} - onChange={store.setGradientAccumulation} - min={1} - max={64} - step={1} - /> - - L2 regularization to prevent overfitting.{" "} - - Read more - - - } - > - - store.setWeightDecay(Number(e.target.value)) - } - className="w-28 font-mono" - /> - - - - - - Gradually increase LR at training start for stability.{" "} - - Read more - - - } - value={store.warmupSteps} - onChange={store.setWarmupSteps} - min={0} - max={100} - step={1} - /> - - Number of full passes over the dataset. Set 0 to run by - max steps.{" "} - - Read more - - - } - value={store.epochs} - onChange={store.setEpochs} - min={0} - max={epochsSliderMax} - step={1} - /> - - Save a checkpoint every N steps. 0 to disable.{" "} - - Read more - - - } - > - store.setSaveSteps(Number(e.target.value))} - className="w-28 font-mono" - /> - - - store.setEvalSteps(Number(e.target.value))} - className="w-28 font-mono" - /> - - - - store.setRandomSeed(Number(e.target.value)) - } - className="w-28 font-mono" - /> - - - - - - Trade compute for memory by recomputing activations.{" "} - - Read more - - - } - > - - - {!showVisionLora && ( -
- store.setPacking(!!v)} + - + + + + Step size for weight updates. Lower values train slower but more + stably.{" "} + + Read more + + + + + store.setLearningRate(Number(e.target.value))} + className="w-full font-mono" + /> +

+ Recommended: 2e-4 for LoRA, 2e-5 for full fine-tune +

+
+ + {/* LoRA Settings */} + {isLora && ( +
+ +
+ + Dimension of the low-rank matrices. Higher = more capacity.{" "} + + Read more + + + } + value={store.loraRank} + onChange={store.setLoraRank} + min={4} + max={128} + step={4} + /> + + Scaling factor for LoRA updates. Usually 2x rank.{" "} + + Read more + + + } + value={store.loraAlpha} + onChange={store.setLoraAlpha} + min={4} + max={256} + step={4} + /> + + Dropout probability for LoRA layers to reduce overfitting.{" "} + + Read more + + + } + value={store.loraDropout} + onChange={store.setLoraDropout} + min={0} + max={0.5} + step={0.01} + format={(v) => v.toFixed(2)} + /> + + {/* Vision checkboxes */} + {showVisionLora && ( +
+ {( + [ + [ + "finetuneVisionLayers", + "Vision layers", + store.finetuneVisionLayers, + store.setFinetuneVisionLayers, + ], + [ + "finetuneLanguageLayers", + "Language layers", + store.finetuneLanguageLayers, + store.setFinetuneLanguageLayers, + ], + [ + "finetuneAttentionModules", + "Attention modules", + store.finetuneAttentionModules, + store.setFinetuneAttentionModules, + ], + [ + "finetuneMLPModules", + "MLP modules", + store.finetuneMLPModules, + store.setFinetuneMLPModules, + ], + ] as const + ).map(([key, label, value, setter]) => ( +
+ + (setter as (v: boolean) => void)(!!v) + } + /> + +
+ ))}
)} -
- store.setTrainOnCompletions(!!v)} - /> - + + {/* Text target modules */} + {!showVisionLora && ( +
+ + Target Modules + +
+ {TARGET_MODULES.map((mod) => { + const active = store.targetModules.includes(mod); + return ( + + ); + })} +
+
+ )} + + {/* LoRA variant */} +
+ {( + [ + { + value: "lora", + label: "Enable LoRA", + desc: "Train with LoRA", + }, + { value: "rslora", label: "RS-LoRA", desc: "Stable Rank" }, + { + value: "loftq", + label: "LoftQ", + desc: "Memory Efficient", + }, + ] as const + ).map((opt) => ( + + ))}
- - - - +
+
+ )} + + {/* Training Hyperparams */} + + + + Training Hyperparameters + + + + + + Optimization + + + Schedule + + + Memory + + + + + + Optimization algorithm. 8-bit variants reduce memory usage. + Fused is recommended for vision models.{" "} + + Read more + + + } + > + + + + How the learning rate changes over training. Linear decays + steadily; cosine decays in a curve.{" "} + + Read more + + + } + > + + + + Samples processed per step. Higher uses more VRAM.{" "} + + Read more + + + } + value={store.batchSize} + onChange={store.setBatchSize} + min={1} + max={32} + step={1} + /> + + Simulates larger batch sizes without extra VRAM.{" "} + + Read more + + + } + value={store.gradientAccumulation} + onChange={store.setGradientAccumulation} + min={1} + max={64} + step={1} + /> + + L2 regularization to prevent overfitting.{" "} + + Read more + + + } + > + + store.setWeightDecay(Number(e.target.value)) + } + className="w-28 font-mono" + /> + + + + + + Gradually increase LR at training start for stability.{" "} + + Read more + + + } + value={store.warmupSteps} + onChange={store.setWarmupSteps} + min={0} + max={100} + step={1} + /> + + Number of full passes over the dataset. Set 0 to run by + max steps.{" "} + + Read more + + + } + value={store.epochs} + onChange={store.setEpochs} + min={0} + max={epochsSliderMax} + step={1} + /> + + Save a checkpoint every N steps. 0 to disable.{" "} + + Read more + + + } + > + store.setSaveSteps(Number(e.target.value))} + className="w-28 font-mono" + /> + + + store.setEvalSteps(Number(e.target.value))} + className="w-28 font-mono" + /> + + + + store.setRandomSeed(Number(e.target.value)) + } + className="w-28 font-mono" + /> + + + + + + Trade compute for memory by recomputing activations.{" "} + + Read more + + + } + > + + + {!showVisionLora && ( +
+ store.setPacking(!!v)} + /> + +
+ )} +
+ store.setTrainOnCompletions(!!v)} + /> + +
+ + + +
diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index a04273f568..10a547b82d 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -237,6 +237,7 @@ export function ProgressSection(): ReactElement { onClick={() => { setStopRequested(true); setStopDialogOpen(false); + useTrainingRuntimeStore.getState().setStopRequested(true); void stopTrainingRun(false).then((ok) => { if (!ok) setStopRequested(false); }); @@ -248,6 +249,7 @@ export function ProgressSection(): ReactElement { onClick={() => { setStopRequested(true); setStopDialogOpen(false); + useTrainingRuntimeStore.getState().setStopRequested(true); void stopTrainingRun(true).then((ok) => { if (!ok) setStopRequested(false); }); diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index 26353b9b7e..69abc1969e 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/studio/frontend/src/features/studio/studio-page.tsx @@ -45,11 +45,16 @@ export function StudioPage(): ReactElement { const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData); const closeDialog = useDatasetPreviewDialogStore((s) => s.close); + const stopRequested = useTrainingRuntimeStore((state) => state.stopRequested); const canGoBack = showTrainingView && - !isTrainingRunning && !isHydratingRuntime && - (runtimePhase === "stopped" || runtimePhase === "error" || runtimePhase === "completed" || runtimePhase === "idle"); + (stopRequested || + (!isTrainingRunning && + (runtimePhase === "stopped" || + runtimePhase === "error" || + runtimePhase === "completed" || + runtimePhase === "idle"))); const tourEnabled = hasHydratedRuntime && !isHydratingRuntime; const isConfigTour = !showTrainingView; const tourSteps = showTrainingView ? studioTrainingTourSteps : studioTourSteps; diff --git a/studio/frontend/src/features/studio/training-start-overlay.tsx b/studio/frontend/src/features/studio/training-start-overlay.tsx index b6c413b400..03078302f6 100644 --- a/studio/frontend/src/features/studio/training-start-overlay.tsx +++ b/studio/frontend/src/features/studio/training-start-overlay.tsx @@ -1,9 +1,23 @@ +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { Button } from "@/components/ui/button"; import { AnimatedSpan, Terminal, TypingAnimation, -} from "@/components/ui/terminal" -import type { ReactElement } from "react" +} from "@/components/ui/terminal"; +import { useTrainingActions, useTrainingRuntimeStore } from "@/features/training"; +import { Cancel01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useEffect, useState, type ReactElement } from "react"; type TrainingStartOverlayProps = { message: string @@ -14,18 +28,65 @@ export function TrainingStartOverlay({ message, currentStep, }: TrainingStartOverlayProps): ReactElement { + const { stopTrainingRun } = useTrainingActions(); + const isStarting = useTrainingRuntimeStore((s) => s.isStarting); + const [cancelDialogOpen, setCancelDialogOpen] = useState(false); + const [cancelRequested, setCancelRequested] = useState(false); + + useEffect(() => { + if (!isStarting) { + setCancelRequested(false); + } + }, [isStarting]); + return (
-
+
Unsloth mascot - +
+ + + + + Cancel Training + + Do you want to cancel the current training run? + + + + Continue Training + { + setCancelRequested(true); + setCancelDialogOpen(false); + useTrainingRuntimeStore.getState().setStopRequested(true); + void stopTrainingRun(false).then((ok) => { + if (!ok) setCancelRequested(false); + }); + }} + > + Cancel Training + + + + + {`> ${message || "starting training..."} | waiting for first step... (${currentStep})`} - + +
) diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts index a80425188d..7d5c86549b 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -37,6 +37,7 @@ const initialState: TrainingRuntimeState = { gradNormHistory: [], evalLossHistory: [], resetGeneration: 0, + stopRequested: false, }; function sortSeries(points: TrainingSeriesPoint[]): TrainingSeriesPoint[] { @@ -110,6 +111,7 @@ function applyMetricHistoryFromStatus(payload: TrainingStatusResponse): { export const useTrainingRuntimeStore = create()((set) => ({ ...initialState, + setStopRequested: (value) => set({ stopRequested: value }), setHydrating: (value) => set({ isHydrating: value }), setHasHydrated: (value) => set({ hasHydrated: value }), setStarting: (value) => set({ isStarting: value }), @@ -173,12 +175,15 @@ export const useTrainingRuntimeStore = create()((set) => ( const detailLoss = payload.details?.loss; const detailLr = payload.details?.learning_rate; const detailEpoch = payload.details?.epoch; + const stopRequested = + payload.is_training_running ? state.stopRequested : false; return { ...state, jobId: payload.job_id || state.jobId, phase: payload.phase, isTrainingRunning: payload.is_training_running, + stopRequested, evalEnabled: payload.eval_enabled ?? state.evalEnabled, message: payload.message, error: payload.error, diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index 7ebf09518d..b6e417c963 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -95,9 +95,11 @@ export interface TrainingRuntimeState { gradNormHistory: TrainingSeriesPoint[]; evalLossHistory: TrainingSeriesPoint[]; resetGeneration: number; + stopRequested: boolean; } export interface TrainingRuntimeActions { + setStopRequested: (value: boolean) => void; setHydrating: (value: boolean) => void; setHasHydrated: (value: boolean) => void; setStarting: (value: boolean) => void;