{{ message }}
Repository navigation
[Community] Help us fix the LR schedulers when num_train_epochs is passed in a distributed training env #8384
Description
Context
Refer to #8312 for the full context. The changes introduced in the PR should be propagated to the following scripts, too:
-
advanced_diffusion_training
- train_dreambooth_lora_sd15_advanced.py
- train_dreambooth_lora_sdxl_advanced.py
-
consistency_distillation
- train_lcm_distill_lora_sdxl.py
-
controlnet
- train_controlnet.py
- train_controlnet_sdxl.py
-
custom_diffusion
- train_custom_diffusion.py
-
dreambooth
- train_dreambooth.py
- train_dreambooth_lora.py
- train_dreambooth_lora_sdxl.py
-
instruct_pix2pix
- train_instruct_pix2pix.py
- rain_instruct_pix2pix_sdxl.py
-
kandinsky2_2/text_to_image
- train_text_to_image_decoder.py
- train_text_to_image_prior.py
- train_text_to_image_lora_decoder.py
- train_text_to_image_lora_prior.py
-
t2i_adapter
- train_t2i_adapter_sdxl.py
-
text_to_image
- train_text_to_image.py
- train_text_to_image_sdxl.py
- train_text_to_image_lora.py
- train_text_to_image_lora_sdxl.py
-
textual_inversion
- textual_inversion.py
- textual_inversion_sdxl.py
-
unconditional_image_generation
- train_unconditional.py
-
wuerstchen
- text_to_image/train_text_to_image_prior.py
- text_to_image/train_text_to_image_lora_prior.py
-
research_projects (low-priority)
- consistency_training/train_cm_ct_unconditional.py
- diffusion_dpo/train_diffusion_dpo.py
- diffusion_dpo/train_diffusion_dpo_sdxl.py
- diffusion_orpo/train_diffusion_orpo_sdxl_lora.py
- dreambooth_inpaint/train_dreambooth_inpaint.py
- dreambooth_inpaint/train_dreambooth_inpaint_lora.py
- instructpix2pix_lora/train_instruct_pix2pix_lora.py
- intel_opts/textual_inversion/textual_inversion_bf16.py
- intel_opts/textual_inversion_dfq/textual_inversion.py
- lora/train_text_to_image_lora.py
- multi_subject_dreambooth/train_multi_subject_dreambooth.py
- multi_token_textual_inversion/textual_inversion.py
- onnxruntime/text_to_image/train_text_to_image.py
- onnxruntime/textual_inversion/textual_inversion.py
- onnxruntime/unconditional_image_generation/train_unconditional.py
- realfill/train_realfill.py
- scheduled_huber_loss_training/dreambooth/train_dreambooth.py
- scheduled_huber_loss_training/dreambooth/train_dreambooth_lora.py
- scheduled_huber_loss_training/dreambooth/train_dreambooth_lora_sdxl.py
- scheduled_huber_loss_training/text_to_image/train_text_to_image.py
- scheduled_huber_loss_training/text_to_image/train_text_to_image_sdxl.py
- scheduled_huber_loss_training/text_to_image/train_text_to_image_lora.py
- scheduled_huber_loss_training/text_to_image/train_text_to_image_lora_sdxl.py
The following scripts do not have the argument --num_train_epochs:
- amused
- train_amused.py
- research_projects
- multi_subject_dreambooth_inpainting/train_multi_subject_dreambooth_inpainting.py
So, they don't need to be updated.
Then we have the following scripts that don't use accelerator to prepare the datasets:
Distributed dataset sharding is done by WebDataset, not accelerator. So, we can skip them for now.
- consistency_distillation
- train_lcm_distill_sd_wds.py
- train_lcm_distill_sdxl_wds.py
- train_lcm_distill_lora_sd_wds.py
- train_lcm_distill_lora_sdxl_wds.py
- research_projects
- controlnet/train_controlnet_webdataset.py
- diffusion_orpo/train_diffusion_orpo_sdxl_lora_wds.py
Steps to follow when opening PRs
- Target one AND only one training script in a single PR.
- When you open a PR, please mention this issue.
- Mention @sayakpaul and @geniuspatrick for a review.
- Accompany your PR with a minimal training command using the
num_train_epochsCLI arg. - Enjoy!
Activity
a-r-r-o-w commentedon Nov 20, 2024
@sayakpaul Are all of the above resolved? Maybe for the ones that are remaining, we can complete it ourselves - in which case, do you know which ones are remaining?
sayakpaul commentedon Nov 20, 2024
I think we can still have this opened as it's good chance for the community for contributions.
kghamilton89 commentedon Apr 4, 2025
@sayakpaul is this issue still looking for contribution?
sayakpaul commentedon Apr 8, 2025
@kghamilton89 I just updated the list of the scripts that have been updated. So, yes, very much open for contributions.
kghamilton89 commentedon Apr 9, 2025
@sayakpaul, the above PRs close out the DreamBooth trainers.
adi-IL commentedon Aug 18, 2026
I'd like to take
examples/text_to_image/train_text_to_image_sdxl.pyfor this, one script as requested.That trainer still computes
max_train_stepsfrom the unsharded dataloader and does not scale warmup or training steps byaccelerator.num_processes. I'll match the #8312 / DreamBooth pattern and open a PR that only touches that file.@sayakpaul @geniuspatrick
shoutoutuoadi325 commentedon Aug 20, 2026
I'd like to take
examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py— adding thelen_train_dataloader_after_sharding/num_warmup_steps_for_schedulercomputation following the pattern in #11239, with the minimal training command included. cc @sayakpaul @geniuspatricksayakpaul commentedon Aug 20, 2026
Sure.
primorLee commentedon Aug 21, 2026
Hi @sayakpaul @geniuspatrick — I’d like to take only
examples/unconditional_image_generation/train_unconditional.pyfor #8384. The script creates its scheduler beforeaccelerator.preparefrom the unsharded dataloader length (len(train_dataloader) * args.num_train_epochs), while later training uses post-sharding update-step math. I plan to apply the established #8312 process-aware calculation, add focused CPU regression coverage using the existing dummy-model example test, and include the required minimal--num_train_epochscommand. This will be AI-assisted; I’ll self-review and share the final notes. Is this scope approved?primorLee commentedon Aug 21, 2026
Correction: this target script exposes
--num_epochs, not--num_train_epochs; I would use its documented--num_epochsas the equivalent minimal command, subject to your confirmation.primorLee commentedon Aug 24, 2026
The focused patch and regression are ready locally for
examples/unconditional_image_generation/train_unconditional.py; I am still waiting for explicit approval before opening the PR.The regression now proves the bug rather than only exercising the branch: with the six-image dummy dataset, batch size 2, one epoch, and gradient accumulation 2, training has two optimizer steps. On current
main, the saved linear scheduler stops at0.0003333333333333333; with the #8312 process-aware calculation it reaches0.0as expected.Validation completed:
examples/unconditional_image_generation/test_unconditional.py: 3 passedmake style,make fix-copies,python -m py_compile, andgit diff --check: passedThe diff remains limited to this one training script and its existing example test (
+25/-5). @sayakpaul @geniuspatrick, may I open the PR for this scope?gitedmond commentedon Aug 28, 2026
Hi @sayakpaul @geniuspatrick , I’d like to take examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py as a single-script contribution for #8384. I’ll follow the established #8312 pattern, include a minimal --num_train_epochs validation command, run the required self-review, and disclose AI assistance. I found no current claim or PR for this exact file. Is this scope approved?
Alphaxiaoteng commentedon Aug 30, 2026
Hi @sayakpaul @geniuspatrick — I would like to take only
examples/research_projects/multi_subject_dreambooth/train_multi_subject_dreambooth.pyfor #8384.The script still derives scheduler steps from the unsharded dataloader before
accelerator.prepare(). In a focused two-process CPU reproduction with three pre-sharding batches, the current calculation creates four scheduler steps although training performs two; applying the established #8312 process-aware calculation gives two scheduler steps. I found no open PR or current claim for this exact file.I plan to keep the change to this one research training script, include a minimal
--num_train_epochsvalidation command, run the required style/fix-copies checks and self-review, and disclose that the contribution is AI-assisted. Is this scope approved?yushuosun commentedon Aug 31, 2026
Hi @sayakpaul @geniuspatrick — I'd like to take only
examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.pyfor this tracker. The script still derives scheduler steps from the unsharded dataloader beforeaccelerator.prepare()and scales them by gradient accumulation rather than process count. I found no open PR or current claim for this exact file.I plan to apply the established #8312 process-aware calculation, keep the change to this one script, run a focused two-process CPU validation if the example dependencies allow it, and complete the required style, fix-copies, and self-review checks. Is this scope okay?
ViktoriiaRomanova commentedon Sep 1, 2026
Hi @sayakpaul @geniuspatrick - I'd like to take [wuerstchen] text_to_image/train_text_to_image_prior.py.
wunianze666-netizen commentedon Sep 8, 2026
Hi @sayakpaul @geniuspatrick — I’d like to take only
examples/research_projects/wuerstchen/text_to_image/train_text_to_image_lora_prior.pyfor #8384.The script still creates its scheduler from the unsharded dataloader length and scales warmup/training steps by gradient accumulation rather than process count. I found no open PR or current claim for this exact file.
I prepared a single-script patch matching the established #8312 pattern. In a focused two-process CPU/Gloo reproduction with five pre-sharding batches and gradient accumulation 2, the current calculation creates 6 scheduler steps for 2 optimizer updates and stops at LR 0.333333; the corrected calculation creates 4 distributed scheduler steps and reaches LR 0.0. Ruff check, Ruff format check, Python compilation, and
git diff --checkpass.This contribution is AI-assisted. I completed the repository self-review rubric with no blocking, non-blocking, or dead-code findings (verdict: READY). Is this scope approved? If so, I’ll open the PR with the requested minimal
--num_train_epochscommand.wunianze666-netizen commentedon Sep 27, 2026
Hi @sayakpaul @geniuspatrick — I would like to take only examples/research_projects/diffusion_dpo/train_diffusion_dpo.py for #8384. Its scheduler is sized from the unsharded loader, then max_train_steps is recalculated after accelerator.prepare: with five batches, two processes and accumulation two, the old path schedules six steps for two optimizer updates and ends at LR 0.333333. I prepared a one-file, AI-assisted fix at https://github.com/wunianze666-netizen/diffusers/tree/fix/dpo-scheduler-epochs. A deterministic execution of the trainer scheduler block fails on main (6 vs 4 steps) and passes with the patch; a two-process CPU/Gloo probe ends at LR 0 on both ranks. Explicit max_train_steps behavior is unchanged. Ruff check/format, py_compile and git diff --check pass; self-review found no blocking or dead-code issues. Please confirm this single-script scope before I open a PR, per the AI-assisted contribution policy. I will include the requested minimal --num_train_epochs command in the PR.
825pranav commentedon Sep 27, 2026
Hi @sayakpaul @geniuspatrick, I'd like to take
examples/research_projects/realfill/train_realfill.py. It still sizes the LR scheduler from the unsharded dataloader, so with 2 processes the LR never fully decays. I have a fix ready and tested it with a 2-process CPU run.sayakpaul commentedon Sep 28, 2026
We don't need more contributions on this one. Thanks, all, for your support!