[train_text_to_image_decoder.py] Fix the LR scheduler when num_train_epochs is passed in a distributed training env

[shoutoutuoadi325]

What does this PR do?

Part of #8384. Propagates the LR scheduler fix from #8312 to examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py.

When --num_train_epochs is used with more than one process, the scheduler's num_warmup_steps/num_training_steps were computed from the unsharded dataloader length and scaled by gradient_accumulation_steps, so the LR schedule was wrong. This follows the same pattern already applied to the other scripts in #8384 (e.g. #11239, #11557), including the mismatch warning when the dataloader length changes after accelerator.prepare. Per #8312, the scheduler is stepped num_processes times per optimizer step, so the gradient_accumulation_steps multiplier is removed. In single-process runs with gradient_accumulation_steps > 1, --lr_warmup_steps and --max_train_steps are now counted in optimizer steps, as documented.

Coordination: claimed here and acknowledged by @sayakpaul here.

Test commands and results

Minimal 2-process run exercising the num_train_epochs path this fix changes:

accelerate launch --multi_gpu --num_processes=2 \
  examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py \
  --dataset_name="lambdalabs/naruto-blip-captions" \
  --max_train_samples=8 --resolution=64 \
  --train_batch_size=1 --gradient_accumulation_steps=1 \
  --num_train_epochs=1 \
  --learning_rate=1e-05 --lr_scheduler="linear" --lr_warmup_steps=2 \
  --report_to="tensorboard" \
  --output_dir="kandi22-decoder-minimal"

Result: 8 samples over 2 processes = 4 optimizer steps per process, ran to completion, models saved:

08/20/2026 12:20:19 - INFO - __main__ - [RANK 0] ***** Running training *****
08/20/2026 12:20:19 - INFO - __main__ - [RANK 0]   Num Epochs = 1
08/20/2026 12:20:19 - INFO - __main__ - [RANK 0]   Instantaneous batch size per device = 1
08/20/2026 12:20:19 - INFO - __main__ - [RANK 0]   Total train batch size (w. parallel, distributed & accumulation) = 2
08/20/2026 12:20:19 - INFO - __main__ - [RANK 0]   Gradient Accumulation steps = 1
08/20/2026 12:20:19 - INFO - __main__ - [RANK 0]   Total optimization steps = 4
Steps: 100%|██████████| 4/4 [00:04<00:00, 1.19it/s, lr=0, step_loss=0.0803]
Model weights saved in kandi22-decoder-minimal/unet/diffusion_pytorch_model.safetensors
Model weights saved in kandi22-decoder-minimal/movq/diffusion_pytorch_model.safetensors

The lr goes 5e-6 -> 1e-5 -> 5e-6 -> 0, which matches num_warmup_steps_for_scheduler = 2 * 2 processes = 4 and num_training_steps_for_scheduler = 1 epoch * 2 processes * 4 steps = 8. No post-prepare mismatch warning was emitted.

Self-review notes

Final self-review round against .ai/review-rules.md (diff is a single examples/ training script; no src/, tests/, or .ai/ changes):

• Blocking issues: none.
• Non-blocking issues: none.
• Dead code: none. The overrode_max_train_steps flag and the branch it guarded are both gone, and nothing else references it. All three new variables are used.
• Intentionally not done: no README change (the fix doesn't alter any documented command) and no new test (examples/kandinsky2_2/ has no test module; the issue asks for a minimal training command instead of unit tests).
• ruff check, ruff format --check, and python -m py_compile all pass on the touched file.
• Verdict: READY. Leave for the actual review: the post-prepare warning condition uses num_training_steps_for_scheduler != args.max_train_steps * accelerator.num_processes (same as examples/text_to_image/train_text_to_image.py:892), not the older bare != args.max_train_steps check some dreambooth scripts still use.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
  (no updates needed: no documented behavior changed)
• Did you write any new necessary tests? (none needed — see self-review notes)
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

cc @sayakpaul @geniuspatrick

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[github-actions]

Hi @shoutoutuoadi325, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

[shoutoutuoadi325]

Keeping Part of #8384 instead of a closing keyword on purpose — the issue is a checklist covering several scripts, and this PR fixes only one of them, so it shouldn't auto-close it.

[github-actions]

Hi @shoutoutuoadi325, a follow-up on the reminder above: this PR still does not link an issue it fixes.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice. Adding a closing keyword (e.g. Fixes #1234) to the PR description, or a maintainer adding the no-issue-needed label, will prevent that.