Fix diffusion training validation and honor lr_scheduler and batch size in the DiT trainer

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,

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:

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)

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(

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"})