From df1ecee819173618fa33a7e583a2e76372af8cc6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 13:25:16 +0000 Subject: [PATCH] Train tab: epochs run length + stop dialog wrap fix - num_epochs on the diffusion train request and config: > 0 overrides train_steps with epochs x ceil(N / (batch x grad_accum)) optimizer steps, resolved against the dataset size in both the DiT and SDXL trainers - Train settings: run length control with a Steps / Epochs unit select - Stop dialog: flex-wrap footer so Stop and save wraps instead of clipping out of frame at narrow window widths --- .../core/training/diffusion_dit_trainer.py | 7 ++- .../core/training/diffusion_lora_trainer.py | 6 ++ .../core/training/diffusion_train_common.py | 22 +++++++ studio/backend/models/training.py | 9 +++ .../tests/test_diffusion_lora_trainer.py | 55 ++++++++++++++++ .../backend/tests/test_diffusion_training.py | 23 +++++++ studio/frontend/src/features/images/api.ts | 3 + .../images/train/diffusion-train-panel.tsx | 62 +++++++++++++++++-- 8 files changed, 180 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 5776f87652..44597ce561 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -31,7 +31,7 @@ import os import random import time from contextlib import nullcontext -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import Any, Callable, Optional @@ -48,6 +48,7 @@ from core.training.diffusion_train_common import ( _restore_perf_flags, discover_image_caption_pairs, repo_is_prequantized, + resolve_train_steps, ) # Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks @@ -948,6 +949,10 @@ def run_dit_lora_training( pairs = discover_image_caption_pairs( cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column ) + # Resolve num_epochs -> a concrete train_steps now that the dataset size is known, and + # rebind cfg so every downstream read (scheduler length, the loop range, progress + # total_steps, steps_run) sees the same resolved value. + cfg = replace(cfg, train_steps = resolve_train_steps(cfg, len(pairs)), num_epochs = 0) _emit(on_event, "model_load_started", num_images = len(pairs)) if _check_stop(): out_dir = Path(cfg.output_dir).expanduser() diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 84bbb6bf74..09d1c5618b 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -36,6 +36,7 @@ import gc import os import random import time +from dataclasses import replace from pathlib import Path from typing import Any, Optional @@ -58,6 +59,7 @@ from core.training.diffusion_train_common import ( # noqa: F401 _restore_perf_flags, discover_image_caption_pairs, get_trainer, + resolve_train_steps, ) @@ -301,6 +303,10 @@ def run_diffusion_lora_training( pairs = discover_image_caption_pairs( cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column ) + # Resolve num_epochs -> a concrete train_steps now that the dataset size is known, and + # rebind cfg so every downstream read (scheduler length, the loop range, progress + # total_steps, steps_run) sees the same resolved value. + cfg = replace(cfg, train_steps = resolve_train_steps(cfg, len(pairs)), num_epochs = 0) _emit(on_event, "model_load_started", num_images = len(pairs)) # Honour a stop requested before the (potentially large / slow) base model loads, the diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 56b1020cd5..05faaa5f3b 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -16,6 +16,7 @@ actual training loop; this module only routes a request to the right one. from __future__ import annotations import json +import math import os import random import re @@ -245,6 +246,9 @@ class DiffusionLoraConfig: instance_prompt: Optional[str] = None resolution: int = 1024 train_steps: int = 500 + # 0 = disabled (train for train_steps). > 0 overrides train_steps with a run length of + # num_epochs full passes over the dataset, in optimizer steps (see resolve_train_steps). + num_epochs: int = 0 learning_rate: float = 1e-4 train_batch_size: int = 1 gradient_accumulation_steps: int = 1 @@ -298,6 +302,8 @@ class DiffusionLoraConfig: resolved_family = resolve_trainable_family(self.base_model, self.model_family) if self.train_steps < 1: raise ValueError("train_steps must be >= 1") + if not 0 <= int(self.num_epochs) <= 1000: + raise ValueError("num_epochs must be between 0 and 1000 (0 uses train_steps)") if self.train_batch_size < 1: raise ValueError("train_batch_size must be >= 1") if self.gradient_accumulation_steps < 1: @@ -359,6 +365,18 @@ class DiffusionLoraConfig: ) +def resolve_train_steps(cfg: "DiffusionLoraConfig", n_images: int) -> int: + """The effective optimizer-step count for a run. When ``cfg.num_epochs`` is set (> 0), + one epoch is one full pass over the dataset in optimizer steps -- ceil(N / (batch x + grad_accum)) steps -- so the run is ``num_epochs`` such passes, capped at 100000. With + ``num_epochs == 0`` the explicit ``cfg.train_steps`` is used unchanged.""" + if cfg.num_epochs > 0: + per_step = max(1, cfg.train_batch_size * cfg.gradient_accumulation_steps) + steps_per_epoch = max(1, math.ceil(n_images / per_step)) + return min(100000, cfg.num_epochs * steps_per_epoch) + return cfg.train_steps + + def discover_image_caption_pairs( data_dir: str | os.PathLike[str], *, @@ -596,6 +614,10 @@ def _write_lora_sidecar(sidecar_path: Path, cfg: DiffusionLoraConfig) -> None: _CONFIG_ALIASES = { "model_name": "base_model", "max_steps": "train_steps", + # The generic payload's num_epochs already matches the diffusion field name, but list it + # so the epochs override is threaded through the shared-payload path as explicitly as + # max_steps -> train_steps is. + "num_epochs": "num_epochs", "batch_size": "train_batch_size", "lora_r": "lora_rank", "lr_scheduler_type": "lr_scheduler", diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index d7703f7141..aa7b1f29fe 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -691,6 +691,15 @@ class DiffusionTrainingStartRequest(BaseModel): 1024, ge = 64, le = 2048, description = "Square training resolution (multiple of 8)" ) train_steps: int = Field(500, ge = 1, le = 100000) + num_epochs: int = Field( + 0, + ge = 0, + le = 1000, + description = ( + "0 = use train_steps; > 0 overrides train_steps with epochs x " + "ceil(N / (batch x grad_accum)) optimizer steps over the N-image dataset" + ), + ) learning_rate: float = Field(1e-4, gt = 0) train_batch_size: int = Field(1, ge = 1, le = 64) gradient_accumulation_steps: int = Field(1, ge = 1, le = 256) diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index ff8f3ec89e..3465408a0e 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -20,6 +20,7 @@ from core.training.diffusion_lora_trainer import ( _config_from_dict, compute_sdxl_add_time_ids, discover_image_caption_pairs, + resolve_train_steps, ) @@ -102,6 +103,60 @@ def test_config_normalized_validation(kw): DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw).normalized() +def _cfg(**kw): + return DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw) + + +def test_resolve_train_steps_uses_train_steps_when_epochs_disabled(): + # num_epochs == 0 leaves the explicit train_steps untouched, whatever the image count. + cfg = _cfg(train_steps = 300, num_epochs = 0) + assert resolve_train_steps(cfg, 20) == 300 + assert resolve_train_steps(cfg, 1) == 300 + + +def test_resolve_train_steps_epochs_ceil_over_batch_and_grad_accum(): + # One epoch = ceil(N / (batch x grad_accum)) optimizer steps; num_epochs multiplies it. + # 10 images, batch 4, grad_accum 1 -> ceil(10/4)=3 steps/epoch. + assert resolve_train_steps(_cfg(num_epochs = 1, train_batch_size = 4), 10) == 3 + assert resolve_train_steps(_cfg(num_epochs = 5, train_batch_size = 4), 10) == 15 + # grad_accum widens the effective batch: 100 images, batch 2, grad_accum 3 -> per_step=6, + # ceil(100/6)=17 steps/epoch, 2 epochs -> 34. + cfg = _cfg(num_epochs = 2, train_batch_size = 2, gradient_accumulation_steps = 3) + assert resolve_train_steps(cfg, 100) == 34 + # An exact multiple does not round up: 8 images / batch 4 -> 2 steps/epoch. + assert resolve_train_steps(_cfg(num_epochs = 3, train_batch_size = 4), 8) == 6 + + +def test_resolve_train_steps_single_image_dataset(): + # A one-image dataset is one optimizer step per epoch, so num_epochs == steps. + assert resolve_train_steps(_cfg(num_epochs = 7, train_batch_size = 4), 1) == 7 + + +def test_resolve_train_steps_caps_at_100000(): + # The run length is capped at 100000 even for absurd epoch counts (matches the request + # model's train_steps ceiling), so a huge epochs x dataset never overflows the loop. + cfg = _cfg(num_epochs = 1000, train_batch_size = 1) + assert resolve_train_steps(cfg, 10_000) == 100000 + + +def test_config_normalized_num_epochs_bounds(): + # 0 (disabled) and the 1..1000 range normalise; out-of-range is rejected. + assert _cfg(num_epochs = 0).normalized().num_epochs == 0 + assert _cfg(num_epochs = 1000).normalized().num_epochs == 1000 + with pytest.raises(ValueError, match = "num_epochs"): + _cfg(num_epochs = -1).normalized() + with pytest.raises(ValueError, match = "num_epochs"): + _cfg(num_epochs = 1001).normalized() + + +def test_config_from_dict_threads_num_epochs(): + # num_epochs flows through the shared-payload adapter onto the diffusion field. + cfg = _config_from_dict( + {"base_model": "b", "data_dir": "d", "output_dir": "o", "num_epochs": 12} + ) + assert cfg.num_epochs == 12 + + def test_compute_sdxl_add_time_ids(): assert compute_sdxl_add_time_ids(1024) == (1024, 1024, 0, 0, 1024, 1024) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index c9bb62335a..6320b7fc42 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -316,6 +316,29 @@ def test_route_start_forwards_extra_training_knobs(client): assert client._fake.started_with["lora_target_modules"] == ["to_q", "to_v"] +def test_route_start_forwards_num_epochs(client): + # Epochs mode: the frontend omits train_steps and sends num_epochs; it must reach the + # service so the trainer can resolve it against the dataset size. + body = {k: v for k, v in _BODY.items() if k != "train_steps"} + r = client.post("/api/train/diffusion/start", json = {**body, "num_epochs": 8}) + assert r.status_code == 200, r.text + assert client._fake.started_with["num_epochs"] == 8 + + +def test_request_model_num_epochs_bounds(): + # The request schema mirrors DiffusionLoraConfig's 0..1000 num_epochs range. + from pydantic import ValidationError + + from models.training import DiffusionTrainingStartRequest + + base = {"base_model": "b", "data_dir": "d", "output_dir": "o"} + assert DiffusionTrainingStartRequest(**base).num_epochs == 0 # default = use train_steps + assert DiffusionTrainingStartRequest(**base, num_epochs = 1000).num_epochs == 1000 + for bad in (-1, 1001): + with pytest.raises(ValidationError): + DiffusionTrainingStartRequest(**base, num_epochs = bad) + + def test_route_start_rejects_uncontained_paths(client): # An absolute path outside the Studio dataset roots is a 400, not silently accepted. r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "/etc"}) diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 61726a02ff..eb8ecd785e 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -270,6 +270,9 @@ export interface DiffusionTrainingStartRequest { instance_prompt?: string | null; resolution?: number; train_steps?: number; + // 0 or omitted uses train_steps. > 0 overrides train_steps with that many epochs + // (full passes over the dataset, in optimizer steps). + num_epochs?: number; learning_rate?: number; train_batch_size?: number; gradient_accumulation_steps?: number; diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 0a9f44b839..a76e8e5209 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -210,6 +210,10 @@ export function DiffusionTrainPanel({ const [instancePrompt, setInstancePrompt] = useState(""); const [steps, setSteps] = useState(500); + // Run length is set in either steps or epochs; the trainer resolves epochs -> steps once + // the dataset size is known (num_epochs overrides train_steps on the backend). + const [durationUnit, setDurationUnit] = useState<"steps" | "epochs">("steps"); + const [epochs, setEpochs] = useState(10); const [learningRate, setLearningRate] = useState(family?.defaults.lr ?? 0.0001); const [rank, setRank] = useState(family?.defaults.rank ?? 16); const [resolution, setResolution] = useState(family?.defaults.resolution ?? 768); @@ -553,7 +557,11 @@ export function DiffusionTrainPanel({ ); return; } - if (steps < 1) return toast.error("Steps must be at least 1."); + if (durationUnit === "epochs") { + if (epochs < 1) return toast.error("Epochs must be at least 1."); + } else if (steps < 1) { + return toast.error("Steps must be at least 1."); + } if (rank < 1) return toast.error("LoRA rank must be at least 1."); if (resolution < 64 || resolution % 8 !== 0) { return toast.error("Resolution must be a multiple of 8 and at least 64."); @@ -577,7 +585,10 @@ export function DiffusionTrainPanel({ output_dir: outputDir.trim(), instance_prompt: instancePrompt.trim() || undefined, resolution, - train_steps: steps, + // Epochs mode overrides train_steps on the backend, so send num_epochs and omit + // train_steps (the backend default is unused when num_epochs > 0). + train_steps: durationUnit === "epochs" ? undefined : steps, + num_epochs: durationUnit === "epochs" ? epochs : undefined, learning_rate: learningRate, train_batch_size: batchSize, gradient_accumulation_steps: gradAccum, @@ -610,6 +621,8 @@ export function DiffusionTrainPanel({ instancePrompt, resolution, steps, + durationUnit, + epochs, learningRate, batchSize, gradAccum, @@ -690,6 +703,40 @@ export function DiffusionTrainPanel({ ); + // Run length: a number paired with a compact unit select (Steps / Epochs). Epochs mode + // trains for that many full passes over the dataset; the backend resolves it to steps. + const durationField = ( +
+ +
+ { + settingsDirty.current = true; + const n = Number(e.target.value) || 1; + if (durationUnit === "epochs") setEpochs(n); + else setSteps(n); + }} + className="h-8 min-w-0 flex-1 text-xs" + /> + +
+
+ ); + const precisionLabel = (m: "nf4" | "bf16" | "int8" | "fp8" | "auto"): string => { if (m === "auto") return "Auto (recommended)"; if (m === "nf4") return "nf4 (4-bit QLoRA, lowest VRAM)"; @@ -704,7 +751,7 @@ export function DiffusionTrainPanel({ const trainingSettings = (
- {numberField("Steps", steps, setSteps, 1)} + {durationField} {numberField("LoRA rank", rank, setRank, 1)} {numberField("Resolution", resolution, setResolution, 512, { min: 64, step: 64 })} {numberField("Batch", batchSize, setBatchSize, 1)} @@ -1292,9 +1339,12 @@ export function DiffusionTrainPanel({ finishes first. - {/* items-center + a real label on the destructive action: a bare "Stop" rendered - as a stubby pill between two wide ones and read as misaligned. */} - + {/* flex-wrap keeps all three buttons visible when the sm:flex-row row is wider than + the dialog at narrow widths (down to ~480px); it wraps instead of clipping the + last button past the right edge. items-center + a real label on the destructive + action: a bare "Stop" rendered as a stubby pill between two wide ones and read + as misaligned. */} + Continue training void onStop(false)}> Stop without saving