diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 8a164193ea..8ada3d30f4 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -51,6 +51,7 @@ class TrainingProgress: eta_seconds: Optional[float] = None grad_norm: Optional[float] = None num_tokens: Optional[int] = None + eval_loss: Optional[float] = None class UnslothTrainer: """ @@ -323,12 +324,22 @@ class UnslothTrainer: local_datasets: list = None, custom_format_mapping: dict = None, subset: str = None, - split: str = "train") -> Optional[Dataset]: + train_split: str = "train", + eval_split: str = None) -> Optional[tuple]: """ - Load and prepare dataset for training + Load and prepare dataset for training. + + Strategy: format first, then split — ensures both train and eval + portions are properly formatted and templated. + + Returns: + Tuple of (dataset_info, eval_dataset) or None on error. + eval_dataset may be None if no eval split is available. """ try: dataset = None + eval_dataset = None + has_separate_eval_source = False # True if eval comes from a separate HF split if local_datasets: # Load local datasets @@ -367,7 +378,7 @@ class UnslothTrainer: elif dataset_source: # Load from Hugging Face - load_kwargs = {"path": dataset_source, "split": split or "train"} + load_kwargs = {"path": dataset_source, "split": train_split or "train"} if subset: load_kwargs["name"] = subset dataset = load_dataset(**load_kwargs) @@ -380,6 +391,25 @@ class UnslothTrainer: self._update_progress(status_message=f"Loaded dataset from HuggingFace: {dataset_source}") print(f"Loaded dataset from Hugging Face: {dataset_source}\n") + # Resolve eval split from a separate HF split (explicit or auto-detected) + if eval_split: + # Explicit eval split provided - load it directly + print(f"Loading explicit eval split: '{eval_split}'\n") + eval_load_kwargs = {"path": dataset_source, "split": eval_split} + if subset: + eval_load_kwargs["name"] = subset + eval_dataset = load_dataset(**eval_load_kwargs) + has_separate_eval_source = True + print(f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n") + else: + # Auto-detect eval split from HF (returns a separate dataset, or None) + eval_dataset = self._auto_detect_eval_split_from_hf( + dataset_source=dataset_source, + subset=subset, + ) + if eval_dataset is not None: + has_separate_eval_source = True + if dataset is None: raise ValueError("No dataset provided") @@ -388,16 +418,15 @@ class UnslothTrainer: print("Stopped before applying chat template\n") return None - # NEW: Use unified format_and_template_dataset + # ========== FORMAT FIRST ========== print(f"Formatting dataset with format_type='{format_type}'...\n") - #breakpoint() dataset_info = format_and_template_dataset( dataset, model_name=self.model_name, - tokenizer=self.tokenizer, # Works for both text and vision models + tokenizer=self.tokenizer, is_vlm=self.is_vlm, - format_type=format_type, # "auto", "alpaca", "chatml", "sharegpt" + format_type=format_type, dataset_name=dataset_source, custom_format_mapping=custom_format_mapping, ) @@ -409,15 +438,94 @@ class UnslothTrainer: self._update_progress(status_message=f"Dataset formatted and ready for training") print(f"Dataset formatted successfully\n") - return dataset_info + + # ========== THEN SPLIT ========== + if has_separate_eval_source and eval_dataset is not None: + # Eval came from a separate HF split — format it too + print(f"Formatting eval dataset ({len(eval_dataset)} rows)...\n") + eval_info = format_and_template_dataset( + eval_dataset, + model_name=self.model_name, + tokenizer=self.tokenizer, + is_vlm=self.is_vlm, + format_type=format_type, + dataset_name=dataset_source, + custom_format_mapping=custom_format_mapping, + ) + eval_dataset = eval_info["dataset"] + print(f"Eval dataset formatted successfully\n") + elif not has_separate_eval_source: + # No separate eval source — split the already-formatted dataset + formatted_dataset = dataset_info["dataset"] + split_result = self._resolve_eval_split_from_dataset(formatted_dataset) + if split_result is not None: + train_portion, eval_dataset = split_result + dataset_info["dataset"] = train_portion + + return (dataset_info, eval_dataset) except Exception as e: logger.error(f"Error loading dataset: {e}") self._update_progress(error=str(e)) return None + def _auto_detect_eval_split_from_hf(self, dataset_source: str, + subset: str) -> Optional[Dataset]: + """Auto-detect an eval split from HF dataset (separate named split only).""" + try: + from datasets import get_dataset_split_names + load_kwargs = {"path": dataset_source} + if subset: + load_kwargs["name"] = subset + available_splits = get_dataset_split_names(**load_kwargs) + print(f"Available splits: {available_splits}\n") + + # Check for common eval split names + for candidate in ["eval", "validation", "valid", "val", "test"]: + if candidate in available_splits: + eval_load_kwargs = {"path": dataset_source, "split": candidate} + if subset: + eval_load_kwargs["name"] = subset + candidate_ds = load_dataset(**eval_load_kwargs) + if len(candidate_ds) >= 16: + print(f"Auto-detected eval split '{candidate}' with {len(candidate_ds)} rows\n") + return candidate_ds + else: + print(f"Found eval split '{candidate}' but only {len(candidate_ds)} rows (< 16), skipping\n") + + except Exception as e: + logger.warning(f"Could not check dataset splits: {e}") + + # No separate HF eval split found — caller will handle programmatic splitting + return None + + def _resolve_eval_split_from_dataset(self, dataset) -> Optional[tuple]: + """Split a dataset into train and eval portions. + + Returns: + Tuple of (train_dataset, eval_dataset), or None if dataset too small. + """ + MIN_EVAL_ROWS = 16 + MIN_TOTAL_ROWS = 32 # Need at least 16 train + 16 eval + + n = len(dataset) + if n < MIN_TOTAL_ROWS: + print(f"Dataset too small ({n} rows) for eval split, skipping eval\n") + return None + + eval_size = max(MIN_EVAL_ROWS, min(128, int(0.05 * n))) + # Ensure we don't take more than half the dataset + eval_size = min(eval_size, n // 2) + + print(f"Auto-splitting: {eval_size} rows for eval from {n} total\n") + split_result = dataset.train_test_split(test_size=eval_size, seed=3407) + print(f"Split complete: {len(split_result['train'])} train, {len(split_result['test'])} eval\n") + return (split_result['train'], split_result['test']) + def start_training(self, dataset: Dataset, + eval_dataset: Dataset = None, + eval_steps: float = 0.01, output_dir: str = "./outputs", num_epochs: int = 3, learning_rate: float = 5e-5, @@ -471,6 +579,8 @@ class UnslothTrainer: 'wandb_token': wandb_token, 'enable_tensorboard': enable_tensorboard, 'tensorboard_dir': tensorboard_dir, + 'eval_dataset': eval_dataset, + 'eval_steps': eval_steps, **kwargs } ) @@ -606,6 +716,17 @@ class UnslothTrainer: else: print(f"Training for {config_args['num_train_epochs']} epochs\n") + # ========== EVAL CONFIGURATION ========== + eval_dataset = training_args.get('eval_dataset', None) + eval_steps_val = training_args.get('eval_steps', 0.01) + if eval_dataset is not None: + config_args["eval_strategy"] = "steps" + config_args["eval_steps"] = eval_steps_val + print(f"Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n") + print(f"Eval dataset: {len(eval_dataset)} rows\n") + else: + print("No eval dataset — evaluation disabled\n") + # Add model-specific parameters # Use optim and lr_scheduler_type from training_args if provided, otherwise use defaults optim_value = training_args.get('optim', "adamw_8bit") @@ -647,21 +768,27 @@ class UnslothTrainer: print("Training configuration prepared\n") # ========== TRAINER INITIALIZATION ========== if self.is_vlm: - self.trainer = SFTTrainer( - model=self.model, - train_dataset=dataset['dataset'], - processing_class = self.tokenizer.tokenizer, - data_collator=data_collator, - args=SFTConfig(**config_args), - ) + trainer_kwargs = { + "model": self.model, + "train_dataset": dataset['dataset'], + "processing_class": self.tokenizer.tokenizer, + "data_collator": data_collator, + "args": SFTConfig(**config_args), + } + if eval_dataset is not None: + trainer_kwargs["eval_dataset"] = eval_dataset + self.trainer = SFTTrainer(**trainer_kwargs) else: - self.trainer = SFTTrainer( - model=self.model, - tokenizer=self.tokenizer, - train_dataset=dataset['dataset'], - data_collator=data_collator, - args=SFTConfig(**config_args), - ) + trainer_kwargs = { + "model": self.model, + "tokenizer": self.tokenizer, + "train_dataset": dataset['dataset'], + "data_collator": data_collator, + "args": SFTConfig(**config_args), + } + if eval_dataset is not None: + trainer_kwargs["eval_dataset"] = eval_dataset + self.trainer = SFTTrainer(**trainer_kwargs) print("Trainer initialized\n") # ========== TRAIN ON RESPONSES ONLY ========== @@ -761,14 +888,15 @@ class UnslothTrainer: self.trainer_instance._update_progress( step=current_step, - epoch=round(state.epoch, 2) if state.epoch else 0, # Round epoch to 2 decimals + epoch=round(state.epoch, 2) if state.epoch else 0, loss=loss_value, learning_rate=logs.get('learning_rate', 0.0), elapsed_seconds=elapsed_seconds, eta_seconds=eta_seconds, grad_norm=grad_norm, num_tokens=num_tokens, - status_message="" # Clear status message so metrics show + eval_loss=logs.get('eval_loss', None), + status_message="" ) def on_epoch_end(self, args, state, control, **kwargs): diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 251b246f35..51d411be78 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -28,6 +28,9 @@ class TrainingBackend: self.loss_history = [] self.lr_history = [] self.step_history = [] + self.eval_loss_history = [] + self.eval_step_history = [] + self.eval_enabled = False self.current_theme = "light" self.trainer.add_progress_callback(self._on_progress_update) @@ -40,6 +43,9 @@ class TrainingBackend: self.loss_history.append(progress.loss) self.lr_history.append(progress.learning_rate) self.step_history.append(progress.step) + if progress.eval_loss is not None: + self.eval_loss_history.append(progress.eval_loss) + self.eval_step_history.append(progress.step) def start_training(self, # Model parameters @@ -96,7 +102,9 @@ class TrainingBackend: # Optional parameters custom_format_mapping: dict = None, subset: str = None, - split: str = "train", + train_split: str = "train", + eval_split: str = None, + eval_steps: float = 0.01, is_dataset_multimodal: bool = False) -> bool: """ Start training. @@ -136,6 +144,9 @@ class TrainingBackend: self.loss_history = [] self.lr_history = [] self.step_history = [] + self.eval_loss_history = [] + self.eval_step_history = [] + self.eval_enabled = False import time output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}" @@ -191,15 +202,30 @@ class TrainingBackend: # ========== LOAD DATASET ========== logger.info("Loading dataset...") #breakpoint() - dataset = self.trainer.load_and_format_dataset( + dataset_result = self.trainer.load_and_format_dataset( dataset_source=hf_dataset if hf_dataset.strip() else None, format_type=format_type, local_datasets=local_datasets if local_datasets else None, custom_format_mapping=custom_format_mapping, subset=subset, - split=split, + train_split=train_split, + eval_split=eval_split, ) + # Unpack: load_and_format_dataset returns (dataset, eval_dataset) + if isinstance(dataset_result, tuple): + dataset, eval_dataset = dataset_result + else: + dataset = dataset_result + eval_dataset = None + + # If user set eval_steps to 0, disable evaluation entirely + if eval_steps is not None and float(eval_steps) <= 0: + eval_dataset = None + + # Track whether eval is enabled for status reporting + self.eval_enabled = eval_dataset is not None + if dataset is None or self.trainer.should_stop: logger.error("Failed to load dataset or stopped by user") return False @@ -219,7 +245,8 @@ class TrainingBackend: logger.info("Starting training worker thread...") success = self.trainer.start_training( dataset=dataset, - #output_dir=f"./outputs/{model_name.replace('/', '_')}_{int(__import__('time').time())}", + eval_dataset=eval_dataset, + eval_steps=eval_steps, output_dir=output_dir, num_epochs=num_epochs, learning_rate=lr_value, @@ -238,7 +265,7 @@ class TrainingBackend: wandb_token=wandb_token if wandb_token.strip() else None, enable_tensorboard=enable_tensorboard, tensorboard_dir=tensorboard_dir, - max_seq_length=max_seq_length, # Pass through for config + max_seq_length=max_seq_length, optim=optim, lr_scheduler_type=lr_scheduler_type, ) diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index 4afe087481..81adef7577 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -1,8 +1,8 @@ """ Dataset-related Pydantic models for API requests and responses. """ -from pydantic import BaseModel -from typing import Optional, Dict, List +from pydantic import BaseModel, model_validator +from typing import Any, Optional, Dict, List class CheckFormatRequest(BaseModel): @@ -11,7 +11,15 @@ class CheckFormatRequest(BaseModel): is_vlm: bool = False hf_token: Optional[str] = None subset: Optional[str] = None - split: Optional[str] = "train" + train_split: Optional[str] = "train" + + @model_validator(mode="before") + @classmethod + def _compat_split(cls, values: Any) -> Any: + """Accept legacy 'split' field as alias for 'train_split'.""" + if isinstance(values, dict) and "split" in values: + values.setdefault("train_split", values.pop("split")) + return values class CheckFormatResponse(BaseModel): diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index dcc6aff0ae..fd04baf3a0 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -1,8 +1,8 @@ """ Pydantic schemas for Training API """ -from pydantic import BaseModel, Field -from typing import Optional, List, Dict, Literal +from pydantic import BaseModel, Field, model_validator +from typing import Any, Optional, List, Dict, Literal class TrainingStartRequest(BaseModel): @@ -19,7 +19,17 @@ class TrainingStartRequest(BaseModel): local_datasets: List[str] = Field(default_factory=list, description="List of local dataset paths") format_type: str = Field(..., description="Dataset format type") subset: Optional[str] = None - split: Optional[str] = "train" + train_split: Optional[str] = Field("train", description="Training split name") + eval_split: Optional[str] = Field(None, description="Eval split name. None = auto-detect") + eval_steps: float = Field(0.01, description="Fraction of total steps between evals (0-1)") + + @model_validator(mode="before") + @classmethod + def _compat_split(cls, values: Any) -> Any: + """Accept legacy 'split' field as alias for 'train_split'.""" + if isinstance(values, dict) and "split" in values: + values.setdefault("train_split", values.pop("split")) + return values custom_format_mapping: Optional[Dict[str, str]] = Field( None, description="User-provided column-to-role mapping, e.g. {'image': 'image', 'caption': 'text'} for VLM or {'instruction': 'user', 'output': 'assistant'} for LLM" @@ -87,6 +97,7 @@ class TrainingStatus(BaseModel): "stopped" ] = Field(..., description="Current phase of training pipeline") is_training_running: bool = Field(..., description="True if training loop is actively running") + eval_enabled: bool = Field(False, description="True if evaluation dataset is configured for this training run") message: str = Field(..., description="Human-readable status message") error: Optional[str] = Field(None, description="Error details if phase is 'error'") details: Optional[dict] = Field(None, description="Phase-specific info, e.g. {'model_size': '8B'}") @@ -110,4 +121,5 @@ class TrainingProgress(BaseModel): eta_seconds: Optional[float] = Field(None, description="Estimated time remaining") grad_norm: Optional[float] = Field(None, description="L2 norm of gradients, computed before gradient clipping") num_tokens: Optional[int] = Field(None, description="Total number of tokens processed so far") + eval_loss: Optional[float] = Field(None, description="Eval loss from the most recent evaluation step") diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index fc6b93c90e..223f701d8c 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -93,11 +93,11 @@ async def check_format(request: CheckFormatRequest): if dataset_path.exists(): # Local dataset if dataset_path.suffix in ['.json', '.jsonl']: - dataset = load_dataset('json', data_files=str(dataset_path), split=request.split) + dataset = load_dataset('json', data_files=str(dataset_path), split=request.train_split) elif dataset_path.suffix == '.csv': - dataset = load_dataset('csv', data_files=str(dataset_path), split=request.split) + dataset = load_dataset('csv', data_files=str(dataset_path), split=request.train_split) elif dataset_path.suffix == '.parquet': - dataset = load_dataset('parquet', data_files=str(dataset_path), split=request.split) + dataset = load_dataset('parquet', data_files=str(dataset_path), split=request.train_split) else: raise HTTPException( status_code=400, @@ -105,7 +105,7 @@ async def check_format(request: CheckFormatRequest): ) else: # HuggingFace dataset - load_kwargs = {"path": request.dataset_name, "split": request.split} + load_kwargs = {"path": request.dataset_name, "split": request.train_split} if request.subset: load_kwargs["name"] = request.subset if request.hf_token: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 13f7135ba4..c09b814fa1 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -146,7 +146,9 @@ async def start_training( "local_datasets": request.local_datasets, "format_type": request.format_type, "subset": request.subset, - "split": request.split, + "train_split": request.train_split, + "eval_split": request.eval_split, + "eval_steps": request.eval_steps, "custom_format_mapping": request.custom_format_mapping, "num_epochs": request.num_epochs, "learning_rate": request.learning_rate, @@ -406,12 +408,15 @@ async def get_training_status( "steps": list(backend.step_history), "loss": list(backend.loss_history), "lr": list(backend.lr_history), + "eval_loss": list(backend.eval_loss_history), + "eval_steps": list(backend.eval_step_history), } return TrainingStatus( job_id=job_id, phase=phase, is_training_running=is_active, + eval_enabled=backend.eval_enabled, message=status_message, error=error_message, details=details, @@ -514,6 +519,7 @@ async def stream_training_progress( eta_seconds = getattr(progress, 'eta_seconds', None) if progress else None grad_norm = getattr(progress, 'grad_norm', None) if progress else None num_tokens = getattr(progress, 'num_tokens', None) if progress else None + eval_loss = getattr(progress, 'eval_loss', None) if progress else None return TrainingProgress( job_id=job_id, @@ -527,6 +533,7 @@ async def stream_training_progress( eta_seconds=eta_seconds, grad_norm=grad_norm, num_tokens=num_tokens, + eval_loss=eval_loss, ) def format_sse( diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index b870f81dbc..8840b9363c 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -39,27 +39,27 @@ export const MODEL_TYPES: ReadonlyArray<{ label: string; description: string; }> = [ - { - value: "vision", - label: "Vision", - description: "Image understanding models", - }, - { - value: "tts", - label: "TTS", - description: "Text-to-speech models", - }, - { - value: "embeddings", - label: "Embeddings", - description: "Text embedding models", - }, - { - value: "text", - label: "Text", - description: "Language models", - }, -]; + { + value: "vision", + label: "Vision", + description: "Image understanding models", + }, + { + value: "tts", + label: "TTS", + description: "Text-to-speech models", + }, + { + value: "embeddings", + label: "Embeddings", + description: "Text embedding models", + }, + { + value: "text", + label: "Text", + description: "Language models", + }, + ]; export const CONTEXT_LENGTHS = [512, 1024, 2048, 4096, 8192, 16384, 32768]; @@ -87,6 +87,7 @@ export const DEFAULT_HYPERPARAMS = { warmupSteps: 5, maxSteps: 0, saveSteps: 0, + evalSteps: 0.01, packing: false, trainOnCompletions: false, gradientCheckpointing: "unsloth" as const, diff --git a/studio/frontend/src/features/studio/sections/charts-content.tsx b/studio/frontend/src/features/studio/sections/charts-content.tsx index a6c5035551..61cbb0a3ea 100644 --- a/studio/frontend/src/features/studio/sections/charts-content.tsx +++ b/studio/frontend/src/features/studio/sections/charts-content.tsx @@ -67,6 +67,7 @@ interface TrainingChartSeries { lossHistory: LossHistoryItem[]; lrHistory: { step: number; lr: number }[]; gradNormHistory: { step: number; gradNorm: number }[]; + evalLossHistory: { step: number; loss: number }[]; } const CHART_SYNC_ID = "train-metrics-sync"; @@ -147,7 +148,9 @@ function ema(data: LossHistoryItem[], alpha: number): SmoothedLossItem[] { export function ChartsContent({ metrics, -}: { metrics: TrainingChartSeries }): ReactElement { + isTraining, + evalEnabled, +}: { metrics: TrainingChartSeries; isTraining: boolean; evalEnabled: boolean }): ReactElement { const [smoothing, setSmoothing] = useState(0.75); const [showRaw, setShowRaw] = useState(true); const [showSmoothed, setShowSmoothed] = useState(true); @@ -174,6 +177,11 @@ export function ChartsContent({ [metrics.lrHistory], ); + const reducedEvalLossData = useMemo( + () => compressSeries(metrics.evalLossHistory, MAX_RENDER_POINTS), + [metrics.evalLossHistory], + ); + const visibleStepDomain = useMemo<[number, number]>(() => { const allSteps = [ ...reducedLossData.map((point) => point.step), @@ -254,12 +262,24 @@ export function ChartsContent({ const gradDomain = useMemo(() => buildYDomain(visibleGradValues), [visibleGradValues]); const lrDomain = useMemo(() => buildYDomain(visibleLrValues), [visibleLrValues]); + const evalLossDomain = useMemo(() => { + const vals = reducedEvalLossData.map((p) => p.loss); + return buildYDomain(vals); + }, [reducedEvalLossData]); + + const evalLossStepTicks = useMemo(() => { + if (reducedEvalLossData.length < 2) return undefined; + const min = reducedEvalLossData[0].step; + const max = reducedEvalLossData[reducedEvalLossData.length - 1].step; + return buildStepTicks(min, max); + }, [reducedEvalLossData]); + const avg = metrics.lossHistory.length > 0 ? +( - metrics.lossHistory.reduce((a, b) => a + b.loss, 0) / - metrics.lossHistory.length - ).toFixed(4) + metrics.lossHistory.reduce((a, b) => a + b.loss, 0) / + metrics.lossHistory.length + ).toFixed(4) : 0; return ( @@ -548,18 +568,18 @@ export function ChartsContent({ - + 0 ? "" : " text-muted-foreground"}`}> Eval Loss -
+ {reducedEvalLossData.length > 0 ? ( @@ -568,42 +588,102 @@ export function ChartsContent({ dataKey="step" type="number" domain={["dataMin", "dataMax"]} + ticks={evalLossStepTicks} + allowDataOverflow={true} + allowDecimals={false} + minTickGap={28} tickLine={false} axisLine={false} tickMargin={8} fontSize={10} + tickFormatter={(value) => formatStepTick(Number(value))} interval="preserveStartEnd" /> Number(value).toFixed(2)} + /> + + `Step ${payload?.[0]?.payload?.step ?? ""}` + } + /> + } /> + } /> -
- -

- Evaluation not configured -

-

- Set eval dataset & eval_steps to track eval loss -

+ ) : ( +
+ + + + + + + + +
+ +

+ {isTraining && evalEnabled ? "Waiting for first evaluation step…" : "Evaluation not configured"} +

+

+ {isTraining && evalEnabled ? "Chart will appear once eval_steps is reached" : "Set eval dataset & eval_steps to track eval loss"} +

+
-
+ )}
diff --git a/studio/frontend/src/features/studio/sections/charts-section.tsx b/studio/frontend/src/features/studio/sections/charts-section.tsx index 868df740d8..7c2c248608 100644 --- a/studio/frontend/src/features/studio/sections/charts-section.tsx +++ b/studio/frontend/src/features/studio/sections/charts-section.tsx @@ -16,11 +16,16 @@ const SKELETON_KEYS = [ export function ChartsSection(): ReactElement | null { const currentStep = useTrainingRuntimeStore((state) => state.currentStep); const totalSteps = useTrainingRuntimeStore((state) => state.totalSteps); + const isTraining = useTrainingRuntimeStore((state) => state.isTrainingRunning); + const evalEnabled = useTrainingRuntimeStore((state) => state.evalEnabled); const lossHistoryRaw = useTrainingRuntimeStore((state) => state.lossHistory); const lrHistoryRaw = useTrainingRuntimeStore((state) => state.lrHistory); const gradNormHistoryRaw = useTrainingRuntimeStore( (state) => state.gradNormHistory, ); + const evalLossHistoryRaw = useTrainingRuntimeStore( + (state) => state.evalLossHistory, + ); const series = useMemo( () => ({ @@ -38,8 +43,12 @@ export function ChartsSection(): ReactElement | null { step: point.step, gradNorm: point.value, })), + evalLossHistory: evalLossHistoryRaw.map((point) => ({ + step: point.step, + loss: point.value, + })), }), - [currentStep, gradNormHistoryRaw, lossHistoryRaw, lrHistoryRaw, totalSteps], + [currentStep, evalLossHistoryRaw, gradNormHistoryRaw, lossHistoryRaw, lrHistoryRaw, totalSteps], ); if ( @@ -63,7 +72,7 @@ export function ChartsSection(): ReactElement | null { } > - + ); } diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 9a3495120d..141a38286d 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -648,6 +648,20 @@ export function ParamsSection(): ReactElement { className="w-28 font-mono" /> + + store.setEvalSteps(Number(e.target.value))} + className="w-28 font-mono" + /> + ()( setWarmupSteps: (warmupSteps) => set({ warmupSteps }), setMaxSteps: (maxSteps) => set({ maxSteps }), setSaveSteps: (saveSteps) => set({ saveSteps }), + setEvalSteps: (evalSteps) => set({ evalSteps }), setPacking: (packing) => set({ packing }), setTrainOnCompletions: (trainOnCompletions) => set({ trainOnCompletions }), 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 94db0f28f5..fd0d3d5733 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -12,6 +12,7 @@ const initialState: TrainingRuntimeState = { jobId: null, phase: "idle", isTrainingRunning: false, + evalEnabled: false, message: "Ready to train", error: null, isHydrating: false, @@ -34,6 +35,7 @@ const initialState: TrainingRuntimeState = { lossHistory: [], lrHistory: [], gradNormHistory: [], + evalLossHistory: [], resetGeneration: 0, }; @@ -72,17 +74,22 @@ function upsertPoint( function applyMetricHistoryFromStatus(payload: TrainingStatusResponse): { lossHistory: TrainingSeriesPoint[] | null; lrHistory: TrainingSeriesPoint[] | null; + evalLossHistory: TrainingSeriesPoint[] | null; } { const history = payload.metric_history; if (!history || !history.steps?.length) { - return { lossHistory: null, lrHistory: null }; + return { lossHistory: null, lrHistory: null, evalLossHistory: null }; } const steps = history.steps; const lossHistory = history.loss ? toSeries(steps, history.loss) : null; const lrHistory = history.lr ? toSeries(steps, history.lr) : null; + const evalLossHistory = + history.eval_loss && history.eval_steps + ? toSeries(history.eval_steps, history.eval_loss) + : null; - return { lossHistory, lrHistory }; + return { lossHistory, lrHistory, evalLossHistory }; } export const useTrainingRuntimeStore = create()((set) => ({ @@ -101,6 +108,7 @@ export const useTrainingRuntimeStore = create()((set) => ( lossHistory: [], lrHistory: [], gradNormHistory: [], + evalLossHistory: [], resetGeneration: state.resetGeneration + 1, })), @@ -137,6 +145,7 @@ export const useTrainingRuntimeStore = create()((set) => ( jobId: payload.job_id || state.jobId, phase: payload.phase, isTrainingRunning: payload.is_training_running, + evalEnabled: payload.eval_enabled ?? state.evalEnabled, message: payload.message, error: payload.error, startError: null, @@ -154,6 +163,7 @@ export const useTrainingRuntimeStore = create()((set) => ( typeof detailEpoch === "number" ? detailEpoch : state.currentEpoch, lossHistory: metricHistory.lossHistory ?? state.lossHistory, lrHistory: metricHistory.lrHistory ?? state.lrHistory, + evalLossHistory: metricHistory.evalLossHistory ?? state.evalLossHistory, }; }), @@ -216,6 +226,10 @@ export const useTrainingRuntimeStore = create()((set) => ( step > 0 && typeof payload.grad_norm === "number" ? upsertPoint(state.gradNormHistory, step, payload.grad_norm) : state.gradNormHistory, + evalLossHistory: + step > 0 && typeof payload.eval_loss === "number" + ? upsertPoint(state.evalLossHistory, step, payload.eval_loss) + : state.evalLossHistory, }; }), })); diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index 95de92cfdb..fce7e18b0b 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -18,6 +18,7 @@ export interface TrainingStartRequest { warmup_ratio: number | null; max_steps: number | null; save_steps: number; + eval_steps: number; weight_decay: number; random_seed: number; packing: boolean; diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index 551d947617..ad883bc357 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -40,6 +40,7 @@ export interface TrainingConfigState { warmupSteps: number; maxSteps: number; saveSteps: number; + evalSteps: number; packing: boolean; trainOnCompletions: boolean; gradientCheckpointing: GradientCheckpointing; @@ -89,6 +90,7 @@ export interface TrainingConfigActions { setWarmupSteps: (value: number) => void; setMaxSteps: (value: number) => void; setSaveSteps: (value: number) => void; + setEvalSteps: (value: number) => void; setPacking: (value: boolean) => void; setTrainOnCompletions: (value: boolean) => void; setGradientCheckpointing: (value: GradientCheckpointing) => void; diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index 389418680a..8e0b43bd52 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -12,6 +12,7 @@ export interface TrainingStatusResponse { job_id: string; phase: TrainingPhase; is_training_running: boolean; + eval_enabled: boolean; message: string; error: string | null; details?: { @@ -25,6 +26,8 @@ export interface TrainingStatusResponse { steps?: number[]; loss?: number[]; lr?: number[]; + eval_loss?: number[]; + eval_steps?: number[]; } | null; } @@ -49,6 +52,7 @@ export interface TrainingProgressPayload { eta_seconds: number | null; grad_norm: number | null; num_tokens: number | null; + eval_loss: number | null; } export interface TrainingSeriesPoint { @@ -60,6 +64,7 @@ export interface TrainingRuntimeState { jobId: string | null; phase: TrainingPhase; isTrainingRunning: boolean; + evalEnabled: boolean; message: string; error: string | null; isHydrating: boolean; @@ -82,6 +87,7 @@ export interface TrainingRuntimeState { lossHistory: TrainingSeriesPoint[]; lrHistory: TrainingSeriesPoint[]; gradNormHistory: TrainingSeriesPoint[]; + evalLossHistory: TrainingSeriesPoint[]; resetGeneration: number; }