diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index d29a85a140..54ff06d2ce 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -592,6 +592,16 @@ def run_dit_lora_training( lora_params = [p for p in transformer.parameters() if p.requires_grad] optimizer = _make_optimizer(lora_params, cfg.learning_rate) + # One lr_sched.step() per optimizer update (cfg.train_steps total), matching the + # SDXL trainer: counting micro-steps instead would stretch warmup past the run. + from diffusers.optimization import get_scheduler + + lr_sched = get_scheduler( + cfg.lr_scheduler, + optimizer = optimizer, + num_warmup_steps = cfg.lr_warmup_steps, + num_training_steps = cfg.train_steps, + ) scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( cfg.base_model, subfolder = "scheduler", token = cfg.hf_token ) @@ -604,10 +614,15 @@ def run_dit_lora_training( peak_gb = 0.0 t_start = time.time() done = 0 + # Honor train_batch_size by folding it into the micro-step count: averaging the + # gradient over batch * accum single-image passes is mathematically identical to + # true batching with a mean loss, and keeps the QLoRA memory profile flat (one + # image's activations at a time). Previously batch_size > 1 silently trained at 1. + micro_steps = cfg.gradient_accumulation_steps * cfg.train_batch_size for opt_step in range(cfg.train_steps): optimizer.zero_grad(set_to_none = True) step_loss = 0.0 - for _ in range(cfg.gradient_accumulation_steps): + for _ in range(micro_steps): i = rng.randrange(len(image_paths)) px = ( _load_pixel_tensor( @@ -645,8 +660,8 @@ def run_dit_lora_training( ) target = noise - latents loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean") - (loss / cfg.gradient_accumulation_steps).backward() - step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps + (loss / micro_steps).backward() + step_loss += float(loss.detach()) / micro_steps grad_norm: Optional[float] = None if cfg.max_grad_norm and cfg.max_grad_norm > 0: @@ -654,6 +669,7 @@ def run_dit_lora_training( # chart wants (spikes stay visible even when clipping flattens the update). grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)) optimizer.step() + lr_sched.step() running_loss += step_loss done = opt_step + 1 @@ -672,7 +688,7 @@ def run_dit_lora_training( total_steps = cfg.train_steps, loss = round(step_loss, 5), avg_loss = round(running_loss / done, 5), - learning_rate = cfg.learning_rate, + learning_rate = lr_sched.get_last_lr()[0], grad_norm = round(grad_norm, 5) if grad_norm is not None else None, samples_per_second = sps, peak_memory_gb = peak_gb or None, diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 62a16eef8c..d70f9bce00 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -259,6 +259,10 @@ class DiffusionLoraConfig: raise ValueError("resolution must be a multiple of 8 and >= 64") if self.mixed_precision not in ("bf16", "fp16", "no"): raise ValueError("mixed_precision must be one of bf16 / fp16 / no") + # A zero/negative gamma would zero out (or invert) the min-SNR weight and + # silently train on a degenerate loss; None is the documented disable. + if self.snr_gamma is not None and float(self.snr_gamma) <= 0: + raise ValueError("snr_gamma must be > 0, or null to disable min-SNR weighting") # learning_rate can arrive as a string ("1e-4") from the Studio config path, which # preserves it as a string after validation; coerce so AdamW receives a float. try: diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 965a854042..3256dfa8e7 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -704,10 +704,14 @@ class DiffusionTrainingStartRequest(BaseModel): default_factory = lambda: ["to_k", "to_q", "to_v", "to_out.0"], description = "U-Net modules to attach LoRA to", ) - max_grad_norm: float = Field(1.0, gt = 0, description = "Gradient clipping max-norm") + max_grad_norm: float = Field( + 1.0, ge = 0, description = "Gradient clipping max-norm; 0 disables clipping" + ) seed: int = Field(42) mixed_precision: Literal["bf16", "fp16", "no"] = Field("bf16") - snr_gamma: Optional[float] = Field(5.0, description = "Min-SNR loss weighting; null disables") + snr_gamma: Optional[float] = Field( + 5.0, gt = 0, description = "Min-SNR loss weighting; null disables" + ) gradient_checkpointing: bool = Field(True) lr_scheduler: str = Field("constant") lr_warmup_steps: int = Field(0, ge = 0) diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index 75046165d0..388a2379e8 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -138,6 +138,16 @@ def test_config_rejects_zero_lora_alpha(): DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", lora_alpha = 0).normalized() +def test_config_rejects_nonpositive_snr_gamma(): + # gamma <= 0 zeroes/inverts the min-SNR weight; None is the documented disable. + with pytest.raises(ValueError, match = "snr_gamma"): + DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", snr_gamma = 0).normalized() + cfg = DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", snr_gamma = None + ).normalized() + assert cfg.snr_gamma is None + + def test_config_coerces_string_learning_rate(): # The Studio config path preserves learning_rate as a string; normalize to float. cfg = DiffusionLoraConfig( diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index a7ef2240f4..1f85b08ccb 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -303,6 +303,20 @@ 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_accepts_zero_max_grad_norm(client): + # 0 is the documented "disable clipping" value (the trainer skips clip_grad_norm_); + # the request model must not reject it. + r = client.post("/api/train/diffusion/start", json = {**_BODY, "max_grad_norm": 0.0}) + assert r.status_code == 200, r.text + assert client._fake.started_with["max_grad_norm"] == 0.0 + + +def test_route_start_rejects_nonpositive_snr_gamma(client): + # gamma <= 0 zeroes/inverts the min-SNR loss weight; null is the disable value. + r = client.post("/api/train/diffusion/start", json = {**_BODY, "snr_gamma": 0}) + assert r.status_code == 422 + + 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"})