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

feat: add eval split auto-detection, eval_steps hyperparam, and eval_loss chart integration

Browse source
This commit is contained in:
Roland Tannous 2026-02-16 13:38:54 +00:00
commit 37452d56cf
16 changed files with 347 additions and 82 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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):

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

@ -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,
)

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

@ -1,8 +1,8 @@
"""
Dataset-related Pydantic models for API requests and responses.
"""
from pydantic import BaseModel
from typing import Optional, Dict, List
from pydantic import BaseModel, model_validator
from typing import Any, Optional, Dict, List
class CheckFormatRequest(BaseModel):
@ -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):

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

@ -1,8 +1,8 @@
"""
Pydantic schemas for Training API
"""
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Literal
from pydantic import BaseModel, Field, model_validator
from typing import Any, Optional, List, Dict, Literal
class TrainingStartRequest(BaseModel):
@ -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")

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

@ -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:

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

@ -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(

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

@ -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,

120
studio/frontend/src/features/studio/sections/charts-content.tsx
View file

@ -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({
<Card data-tour="studio-eval-loss" size="sm">
<CardHeader>
<CardTitle className="text-sm text-muted-foreground pl-2">
<CardTitle className={`text-sm pl-2${reducedEvalLossData.length > 0 ? "" : " text-muted-foreground"}`}>
Eval Loss
</CardTitle>
</CardHeader>
<CardContent>
<div className="relative">
{reducedEvalLossData.length > 0 ? (
<ChartContainer
config={evalLossConfig}
className="-ml-3 h-[220px] w-full blur"
className="-ml-3 h-[220px] w-full"
>
<LineChart
data={placeholderEvalData}
data={reducedEvalLossData}
accessibilityLayer={true}
margin={{ left: 0, right: 8 }}
>
@ -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"
/>
<YAxis
domain={evalLossDomain}
allowDataOverflow={true}
tickLine={false}
axisLine={false}
tickMargin={4}
fontSize={10}
width={40}
width={52}
tickFormatter={(value) => Number(value).toFixed(2)}
/>
<ChartTooltip
content={
<ChartTooltipContent
labelFormatter={(_value, payload) =>
`Step ${payload?.[0]?.payload?.step ?? ""}`
}
/>
}
/>
<Line
type="monotone"
dataKey="loss"
stroke="var(--color-loss)"
strokeWidth={2}
dot={false}
dot={{ r: 3, strokeWidth: 0, fill: "#ef4444" }}
activeDot={{ r: 4, strokeWidth: 0 }}
connectNulls={true}
isAnimationActive={false}
/>
<ChartLegend content={<ChartLegendContent />} />
</LineChart>
</ChartContainer>
<div className="absolute inset-0 flex flex-col items-center justify-center gap-1">
<HugeiconsIcon
icon={ChartAverageIcon}
className="size-5 text-muted-foreground/50"
/>
<p className="text-sm font-medium text-muted-foreground">
Evaluation not configured
</p>
<p className="text-xs text-muted-foreground/60">
Set eval dataset & eval_steps to track eval loss
</p>
) : (
<div className="relative">
<ChartContainer
config={evalLossConfig}
className="-ml-3 h-[220px] w-full blur"
>
<LineChart
data={placeholderEvalData}
accessibilityLayer={true}
margin={{ left: 0, right: 8 }}
>
<CartesianGrid vertical={false} strokeDasharray="3 3" />
<XAxis
dataKey="step"
type="number"
domain={["dataMin", "dataMax"]}
tickLine={false}
axisLine={false}
tickMargin={8}
fontSize={10}
interval="preserveStartEnd"
/>
<YAxis
tickLine={false}
axisLine={false}
tickMargin={4}
fontSize={10}
width={40}
/>
<Line
type="monotone"
dataKey="loss"
stroke="var(--color-loss)"
strokeWidth={2}
dot={false}
isAnimationActive={false}
/>
</LineChart>
</ChartContainer>
<div className="absolute inset-0 flex flex-col items-center justify-center gap-1">
<HugeiconsIcon
icon={ChartAverageIcon}
className="size-5 text-muted-foreground/50"
/>
<p className="text-sm font-medium text-muted-foreground">
Evaluation not configured
</p>
<p className="text-xs text-muted-foreground/60">
Set eval dataset & eval_steps to track eval loss
</p>
</div>
</div>
</div>
)}
</CardContent>
</Card>
</div>

9
studio/frontend/src/features/studio/sections/charts-section.tsx
View file

@ -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 (

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

@ -648,6 +648,20 @@ export function ParamsSection(): ReactElement {
className="w-28 font-mono"
/>
</Row>
<Row
label="Eval Steps"
tooltip="Fraction of total training steps between evaluations. E.g. 0.01 = evaluate every 1% of steps."
>
<Input
type="number"
step="0.01"
min="0.001"
max="1"
value={store.evalSteps}
onChange={(e) => store.setEvalSteps(Number(e.target.value))}
className="w-28 font-mono"
/>
</Row>
<Row label="Seed" tooltip="Random seed for reproducibility.">
<Input
type="number"

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

@ -36,6 +36,7 @@ export function buildTrainingStartPayload(
warmup_ratio: null,
max_steps: config.maxSteps,
save_steps: config.saveSteps,
eval_steps: config.evalSteps,
weight_decay: config.weightDecay,
random_seed: config.randomSeed,
packing: config.packing,

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

@ -94,6 +94,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
setWarmupSteps: (warmupSteps) => set({ warmupSteps }),
setMaxSteps: (maxSteps) => set({ maxSteps }),
setSaveSteps: (saveSteps) => set({ saveSteps }),
setEvalSteps: (evalSteps) => set({ evalSteps }),
setPacking: (packing) => set({ packing }),
setTrainOnCompletions: (trainOnCompletions) =>
set({ trainOnCompletions }),

16
studio/frontend/src/features/training/stores/training-runtime-store.ts
View file

@ -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<TrainingRuntimeStore>()((set) => ({
@ -101,6 +107,7 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
lossHistory: [],
lrHistory: [],
gradNormHistory: [],
evalLossHistory: [],
resetGeneration: state.resetGeneration + 1,
})),
@ -154,6 +161,7 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((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<TrainingRuntimeStore>()((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,
};
}),
}));

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

@ -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;

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

@ -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;

4
studio/frontend/src/features/training/types/runtime.ts
View file

@ -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;
}