From f58c3ddb07654743875b8ae09c67ca5a709d155a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 05:46:48 +0000 Subject: [PATCH] Count LR scheduler warmup/decay in optimizer steps, not micro-steps lr_sched.step() runs once per outer optimizer step (after the gradient accumulation inner loop), for train_steps total. The scheduler was configured with num_warmup_steps and num_training_steps multiplied by gradient_accumulation_steps, so with accumulation > 1 a warmup or non-constant schedule stretched past the run and never reached the intended decay. Count both in optimizer steps. --- studio/backend/core/training/diffusion_lora_trainer.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 70f7c8aedf..0e20b6ede5 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -391,11 +391,15 @@ def run_diffusion_lora_training( lora_params = [p for p in unet.parameters() if p.requires_grad] optimizer = torch.optim.AdamW(lora_params, lr = cfg.learning_rate) + # The scheduler advances once per optimizer update: lr_sched.step() runs a single + # time per outer opt_step (after the accumulation inner loop), for cfg.train_steps + # total. Count warmup/decay in those optimizer steps -- multiplying by the + # accumulation factor would stretch warmup past the run and never reach the decay. lr_sched = get_scheduler( cfg.lr_scheduler, optimizer = optimizer, - num_warmup_steps = cfg.lr_warmup_steps * cfg.gradient_accumulation_steps, - num_training_steps = cfg.train_steps * cfg.gradient_accumulation_steps, + num_warmup_steps = cfg.lr_warmup_steps, + num_training_steps = cfg.train_steps, ) vae_scale = vae.config.scaling_factor