diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index b3116c5c04..652bff7888 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: """ @@ -319,12 +320,18 @@ 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. + + 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 if local_datasets: # Load local datasets @@ -361,9 +368,12 @@ class UnslothTrainer: self._update_progress(status_message=f"Loaded {len(all_data)} samples from local files") print(f"Loaded {len(all_data)} samples from local files\n") + # For local datasets, use train_test_split if dataset is large enough + eval_dataset = self._resolve_eval_split_from_dataset(dataset) + 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) @@ -376,6 +386,23 @@ 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 + 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) + print(f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n") + else: + # Auto-detect eval split + eval_dataset = self._auto_detect_eval_split( + dataset_source=dataset_source, + subset=subset, + train_dataset=dataset, + ) + if dataset is None: raise ValueError("No dataset provided") @@ -384,10 +411,14 @@ class UnslothTrainer: print("Stopped before applying chat template\n") return None + # If we auto-split, use only the train portion for formatting + if hasattr(self, '_split_train_dataset') and self._split_train_dataset is not None: + dataset = self._split_train_dataset + self._split_train_dataset = None # Clear after use + # NEW: Use unified format_and_template_dataset print(f"Formatting dataset with format_type='{format_type}'...\n") - #breakpoint() dataset_info = format_and_template_dataset( dataset, model_name=self.model_name, @@ -405,15 +436,67 @@ class UnslothTrainer: self._update_progress(status_message=f"Dataset formatted and ready for training") print(f"Dataset formatted successfully\n") - return dataset_info + 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(self, dataset_source: str, subset: str, + train_dataset: Dataset) -> Optional[Dataset]: + """Auto-detect an eval split from HF dataset, or split the train set.""" + 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}") + + # Fallback: split the train set + return self._resolve_eval_split_from_dataset(train_dataset) + + def _resolve_eval_split_from_dataset(self, dataset: Dataset) -> Optional[Dataset]: + """Split the training dataset to create an eval set.""" + MIN_EVAL_ROWS = 16 + MIN_TOTAL_ROWS = 32 # Need at least 16 train + 16 eval + + if len(dataset) < MIN_TOTAL_ROWS: + print(f"Dataset too small ({len(dataset)} rows) for eval split, skipping eval\n") + return None + + eval_size = max(MIN_EVAL_ROWS, min(128, int(0.05 * len(dataset)))) + # Ensure we don't take more than half the dataset + eval_size = min(eval_size, len(dataset) // 2) + + print(f"Auto-splitting: {eval_size} rows for eval from {len(dataset)} 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") + # Store the train portion so the caller can use it instead of the original full dataset + self._split_train_dataset = split_result['train'] + return 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, @@ -602,6 +685,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") @@ -643,21 +737,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 ========== @@ -757,14 +857,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 d9aaa8ca0e..8a9d9c883e 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -28,6 +28,8 @@ class TrainingBackend: self.loss_history = [] self.lr_history = [] self.step_history = [] + self.eval_loss_history = [] + self.eval_step_history = [] self.current_theme = "light" self.trainer.add_progress_callback(self._on_progress_update) @@ -40,6 +42,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 +101,9 @@ class TrainingBackend: # Optional parameters custom_format_mapping: dict = None, subset: str = None, - split: str = "train") -> bool: + train_split: str = "train", + eval_split: str = None, + eval_steps: float = 0.01) -> bool: """ Start training. @@ -135,6 +142,8 @@ class TrainingBackend: self.loss_history = [] self.lr_history = [] self.step_history = [] + self.eval_loss_history = [] + self.eval_step_history = [] import time output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}" @@ -189,15 +198,23 @@ 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 dataset is None or self.trainer.should_stop: logger.error("Failed to load dataset or stopped by user") return False @@ -217,7 +234,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, @@ -236,7 +254,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 e0839ae485..abf441fcd0 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" @@ -109,4 +119,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 5bf59d85d2..07f6468c6a 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, @@ -405,6 +407,8 @@ 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( @@ -513,6 +517,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, @@ -526,6 +531,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..c05679bce4 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"; @@ -174,6 +175,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 +260,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 +566,18 @@ export function ChartsContent({ - + 0 ? "" : " text-muted-foreground"}`}> Eval Loss -
+ {reducedEvalLossData.length > 0 ? ( @@ -568,42 +586,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 -

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

+ Evaluation not configured +

+

+ 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..4daf9c461c 100644 --- a/studio/frontend/src/features/studio/sections/charts-section.tsx +++ b/studio/frontend/src/features/studio/sections/charts-section.tsx @@ -21,6 +21,9 @@ export function ChartsSection(): ReactElement | null { const gradNormHistoryRaw = useTrainingRuntimeStore( (state) => state.gradNormHistory, ); + const evalLossHistoryRaw = useTrainingRuntimeStore( + (state) => state.evalLossHistory, + ); const series = useMemo( () => ({ @@ -38,8 +41,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 ( diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 415aca3fac..0052fae5b6 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..c67e8d8244 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -34,6 +34,7 @@ const initialState: TrainingRuntimeState = { lossHistory: [], lrHistory: [], gradNormHistory: [], + evalLossHistory: [], resetGeneration: 0, }; @@ -72,17 +73,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 +107,7 @@ export const useTrainingRuntimeStore = create()((set) => ( lossHistory: [], lrHistory: [], gradNormHistory: [], + evalLossHistory: [], resetGeneration: state.resetGeneration + 1, })), @@ -154,6 +161,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 +224,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 f6c43616d9..68a1f02a74 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 6a28500b7e..b811536085 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; @@ -85,6 +86,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..5c32ac1835 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -25,6 +25,8 @@ export interface TrainingStatusResponse { steps?: number[]; loss?: number[]; lr?: number[]; + eval_loss?: number[]; + eval_steps?: number[]; } | null; } @@ -49,6 +51,7 @@ export interface TrainingProgressPayload { eta_seconds: number | null; grad_norm: number | null; num_tokens: number | null; + eval_loss: number | null; } export interface TrainingSeriesPoint { @@ -82,6 +85,7 @@ export interface TrainingRuntimeState { lossHistory: TrainingSeriesPoint[]; lrHistory: TrainingSeriesPoint[]; gradNormHistory: TrainingSeriesPoint[]; + evalLossHistory: TrainingSeriesPoint[]; resetGeneration: number; }