Fix Wuerstchen LoRA scheduler steps for distributed epochs

[wunianze666-netizen]

Part of #8384. This PR addresses one unchecked script and must not close the umbrella tracker.

Coordination and reproduction evidence: #8384 (comment)

This updates only examples/research_projects/wuerstchen/text_to_image/train_text_to_image_lora_prior.py to size the learning-rate scheduler using the post-sharding update count and Accelerate's process-scaled scheduler contract.

Root cause

The scheduler was created before accelerator.prepare() from the unsharded dataloader length, while warmup and training steps were multiplied by gradient accumulation instead of the number of processes. In distributed epoch-based training, the optimizer therefore consumed only part of the intended schedule.

The change mirrors the established #8312 pattern:

• estimate the per-process dataloader length before scheduler creation;
• scale scheduler warmup/training steps by accelerator.num_processes;
• recompute the actual update count after accelerator.prepare();
• warn if the prepared dataloader length differs from the estimate.

Focused reproduction

With five pre-sharding batches, two CPU/Gloo processes, gradient accumulation 2, and one epoch:

• before: 2 optimizer updates, 6 scheduler steps, final linear-schedule LR 0.333333;
• after: 2 optimizer updates, 4 process-scaled scheduler steps, final LR 0.000000 on both ranks.

Additional uneven-sharding arithmetic cases (7 batches / 3 processes and 11 batches / 2 processes) also match the actual process-scaled scheduler call count after the change.

Minimal training command

accelerate launch examples/research_projects/wuerstchen/text_to_image/train_text_to_image_lora_prior.py \
  --dataset_name="lambdalabs/naruto-blip-captions" \
  --caption_column="text" \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=2 \
  --num_train_epochs=1 \
  --learning_rate=1e-4 \
  --lr_scheduler="linear" \
  --lr_warmup_steps=0 \
  --rank=4 \
  --output_dir="wuerstchen-prior-lora-test"

Validation

• ruff check on the changed script
• ruff format --check on the changed script
• python -m py_compile on the changed script
• git diff --check upstream/main...HEAD
• focused arithmetic regression: corrected totals 4/6/8 versus old under-consumption of 33.3%/25%/66.7%
• real two-process CPU/Gloo reproduction for both old and corrected contracts

AI assistance and self-review

This contribution was developed with Codex assistance. I read the repository's AI-agent guidance and ran the repository self-review rubric after syncing with current main.

• Verdict: READY
• Blocking findings: none
• Non-blocking findings: none
• Dead-code findings: none
• Final diff: one production script, +18/-7

@sayakpaul @geniuspatrick

[github-actions]

Hi @wunianze666-netizen, 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.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[wunianze666-netizen]

Thanks for the reminder. This PR is one item in the multi-script umbrella tracker #8384, which still has many unchecked entries, so adding Fixes #8384 would close the tracker prematurely when this one-script PR merges. The body already links both the tracker and the exact coordination comment. Could a maintainer apply no-issue-needed (or the project's preferred non-closing equivalent) for this umbrella-issue case?