From 0d0a6d2b9603ae75532f2e44cffcf972e4db8882 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 13:39:34 +0000 Subject: [PATCH] Drop piecewise_constant from the trainable lr_scheduler allow-list piecewise_constant is the only diffusers scheduler that needs a step_rules string, and neither diffusion trainer passes one (get_scheduler is called with only warmup/training steps, and there is no config field for it). Accepting it let /diffusion/start pass normalized(), free the resident GPU workloads, spawn the trainer, and only then crash in the subprocess (get_piecewise_constant_schedule does step_rules.split(",") on None) -- the exact evict-then-fail the up-front validation exists to prevent. Reject it now with a clear 400. The remaining six schedulers all run with only warmup/training steps. --- .../core/training/diffusion_train_common.py | 9 +++++-- .../tests/test_diffusion_lora_trainer.py | 26 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index aabfbf42a6..aefe2de3e1 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -35,7 +35,13 @@ from core.inference.diffusion_families import ( # set overrides this in its own defaults; kept here so DiffusionLoraConfig has a sane fallback. DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0") -# diffusers' SchedulerType names (diffusers.optimization.get_scheduler). +# diffusers' SchedulerType names (diffusers.optimization.get_scheduler). piecewise_constant is +# intentionally excluded: it is the only scheduler that needs a `step_rules` string, which the +# trainers never pass (get_scheduler is called with only warmup/training steps, and there is no +# config field for it). Accepting it would pass normalized(), free the resident GPU workloads, +# then crash in the trainer subprocess (get_piecewise_constant_schedule does step_rules.split(",") +# on None) -- the exact evict-then-fail the up-front validation exists to prevent. The remaining +# six all run with only warmup/training steps. _LR_SCHEDULERS: frozenset[str] = frozenset( { "linear", @@ -44,7 +50,6 @@ _LR_SCHEDULERS: frozenset[str] = frozenset( "polynomial", "constant", "constant_with_warmup", - "piecewise_constant", } ) diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index bd91eae9c8..d89a472bd3 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -114,6 +114,32 @@ def test_config_normalized_validation(kw): DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw).normalized() +def test_normalized_rejects_piecewise_constant(): + # piecewise_constant needs a step_rules string the trainers never supply, so get_scheduler() + # would crash in the trainer subprocess AFTER the resident GPU workloads are freed. It must be + # rejected up front (a clean ValueError -> 400), not accepted like the other schedulers. + with pytest.raises(ValueError, match = "lr_scheduler"): + DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = "piecewise_constant" + ).normalized() + + +def test_normalized_accepts_supported_schedulers(): + # Every scheduler in the allow-list runs with only warmup/training steps (no extra required arg). + for sched in ( + "linear", + "cosine", + "cosine_with_restarts", + "polynomial", + "constant", + "constant_with_warmup", + ): + cfg = DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = sched + ).normalized() + assert cfg.lr_scheduler == sched + + def test_compute_sdxl_add_time_ids(): assert compute_sdxl_add_time_ids(1024) == (1024, 1024, 0, 0, 1024, 1024)