[examples] flux-control: use num_training_steps_for_scheduler (#11662)

Markus-Pobitzer · sayakpaul · web-flow · commit 745199a86975 · 2025-06-05T14:56:25.000+05:30
[examples] flux-control: use num_training_steps_for_scheduler in get_scheduler instead of args.max_train_steps * accelerator.num_processes

Co-authored-by: Sayak Paul &lt;spsayakpaul@gmail.com&gt;
diff --git a/examples/flux-control/train_control_flux.py b/examples/flux-control/train_control_flux.py
@@ -915,7 +915,7 @@ def load_model_hook(models, input_dir):
         args.lr_scheduler,
         optimizer=optimizer,
         num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
-        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_training_steps=num_training_steps_for_scheduler,
         num_cycles=args.lr_num_cycles,
         power=args.lr_power,
     )
diff --git a/examples/flux-control/train_control_lora_flux.py b/examples/flux-control/train_control_lora_flux.py
@@ -1060,7 +1060,7 @@ def load_model_hook(models, input_dir):
         args.lr_scheduler,
         optimizer=optimizer,
         num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
-        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_training_steps=num_training_steps_for_scheduler,
         num_cycles=args.lr_num_cycles,
         power=args.lr_power,
     )

Original file line number	Diff line number	Diff line change
`@@ -915,7 +915,7 @@ def load_model_hook(models, input_dir):`
`915`	`915`	`args.lr_scheduler,`
`916`	`916`	`optimizer=optimizer,`
`917`	`917`	`num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,`
`918`		`- num_training_steps=args.max_train_steps * accelerator.num_processes,`
	`918`	`+ num_training_steps=num_training_steps_for_scheduler,`
`919`	`919`	`num_cycles=args.lr_num_cycles,`
`920`	`920`	`power=args.lr_power,`
`921`	`921`	`)`
Original file line number	Diff line number	Diff line change
`@@ -1060,7 +1060,7 @@ def load_model_hook(models, input_dir):`
`1060`	`1060`	`args.lr_scheduler,`
`1061`	`1061`	`optimizer=optimizer,`
`1062`	`1062`	`num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,`
`1063`		`- num_training_steps=args.max_train_steps * accelerator.num_processes,`
	`1063`	`+ num_training_steps=num_training_steps_for_scheduler,`
`1064`	`1064`	`num_cycles=args.lr_num_cycles,`
`1065`	`1065`	`power=args.lr_power,`
`1066`	`1066`	`)`