[Community] Help us fix the LR schedulers when num_train_epochs is passed in a distributed training env

@sayakpaul

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_epochs CLI arg.
Enjoy!

@sayakpaul

@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?

I think we can still have this opened as it's good chance for the community for contributions.

@sayakpaul

@sayakpaul is this issue still looking for contribution?

@kghamilton89

@kghamilton89 I just updated the list of the scripts that have been updated. So, yes, very much open for contributions.

@sayakpaul

@sayakpaul, the above PRs close out the DreamBooth trainers.

@sayakpaul

I'd like to take examples/text_to_image/train_text_to_image_sdxl.py for this, one script as requested.

That trainer still computes max_train_steps from the unsharded dataloader and does not scale warmup or training steps by accelerator.num_processes. I'll match the #8312 / DreamBooth pattern and open a PR that only touches that file.

@sayakpaul @geniuspatrick

@sayakpaul

I'd like to take examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py — adding the len_train_dataloader_after_sharding / num_warmup_steps_for_scheduler computation following the pattern in #11239, with the minimal training command included. cc @sayakpaul @geniuspatrick

Sure.

@sayakpaul

Hi @sayakpaul @geniuspatrick — I’d like to take only examples/unconditional_image_generation/train_unconditional.py for #8384. The script creates its scheduler before accelerator.prepare from 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_epochs command. This will be AI-assisted; I’ll self-review and share the final notes. Is this scope approved?

Correction: this target script exposes --num_epochs, not --num_train_epochs; I would use its documented --num_epochs as the equivalent minimal command, subject to your confirmation.

@sayakpaul

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 at 0.0003333333333333333; with the #8312 process-aware calculation it reaches 0.0 as expected.

Validation completed:

full examples/unconditional_image_generation/test_unconditional.py: 3 passed
two-process Multi-CPU/Gloo run: 2 optimizer steps, final scheduler LR 0.0
make style, make fix-copies, python -m py_compile, and git diff --check: passed
final self-review: no blocking, non-blocking, or dead-code findings

The 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?

@sayakpaul

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?

@sayakpaul

Hi @sayakpaul @geniuspatrick — I would like to take only examples/research_projects/multi_subject_dreambooth/train_multi_subject_dreambooth.py for #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_epochs validation command, run the required style/fix-copies checks and self-review, and disclose that the contribution is AI-assisted. Is this scope approved?

@sayakpaul

Hi @sayakpaul @geniuspatrick — I'd like to take only examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py for this tracker. The script still derives scheduler steps from the unsharded dataloader before accelerator.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?

@sayakpaul

Hi @sayakpaul @geniuspatrick - I'd like to take [wuerstchen] text_to_image/train_text_to_image_prior.py.

@sayakpaul

Hi @sayakpaul @geniuspatrick — I’d like to take only examples/research_projects/wuerstchen/text_to_image/train_text_to_image_lora_prior.py for #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 --check pass.

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_epochs command.

@sayakpaul

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.

@sayakpaul

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.

We don't need more contributions on this one. Thanks, all, for your support!

Navigation Menu

Repository navigation

[Community] Help us fix the LR schedulers when `num_train_epochs` is passed in a distributed training env #8384

Context

Steps to follow when opening PRs

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[Community] Help us fix the LR schedulers when num_train_epochs is passed in a distributed training env #8384

Description

Context

Steps to follow when opening PRs

Activity

a-r-r-o-w commented on Nov 20, 2024

sayakpaul commented on Nov 20, 2024

kghamilton89 commented on Apr 4, 2025

sayakpaul commented on Apr 8, 2025

kghamilton89 commented on Apr 9, 2025

adi-IL commented on Aug 18, 2026

shoutoutuoadi325 commented on Aug 20, 2026

sayakpaul commented on Aug 20, 2026

primorLee commented on Aug 21, 2026

primorLee commented on Aug 21, 2026

primorLee commented on Aug 24, 2026

gitedmond commented on Aug 28, 2026

Alphaxiaoteng commented on Aug 30, 2026

yushuosun commented on Aug 31, 2026

ViktoriiaRomanova commented on Sep 1, 2026

wunianze666-netizen commented on Sep 8, 2026

wunianze666-netizen commented on Sep 27, 2026

825pranav commented on Sep 27, 2026

sayakpaul commented on Sep 28, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

[Community] Help us fix the LR schedulers when `num_train_epochs` is passed in a distributed training env #8384