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?

[wunianze666-netizen]

Synced this branch with current main (c419dac01) in merge commit be6f6ad3e. The six upstream commits do not touch the Wuerstchen LoRA prior trainer, and the PR comparison remains exactly one production script (+18/-7).

Post-sync validation passed: Ruff check/format, Python compilation, git diff --check, the three-case scheduler arithmetic regression, and the real two-process CPU/Gloo probe. The old contract again ended after two optimizer updates with 6 process-scaled scheduler steps and LR 0.333333; the corrected contract used 4 steps and reached LR 0.000000 on both ranks. The dependency environment was updated to huggingface-hub 1.31.0 to match the newly merged upstream httpx migration before running the probe.

[github-actions]

This PR has been automatically closed because it does not link an issue and the reminder above was not addressed within 10 days. If this PR is still relevant, please link the issue it fixes (e.g. Fixes #1234) or ask a maintainer to add the no-issue-needed label, and it can be reopened.

We are experimenting with this process to keep the review queue manageable, and it will sometimes get it wrong. If you think this PR should stay open, please just say so here and we will reopen it — no need to justify it at length.

Thanks again for contributing, and sorry for the noise if we closed this by mistake!

[wunianze666-netizen]

This fix is still relevant and remains an unfinished subtask of the umbrella tracker #8384. Could a maintainer please reopen the PR and apply the appropriate non-closing/no-issue-needed label for this umbrella-issue case?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

This PR has been automatically closed because it does not link an issue and the reminder above was not addressed within 10 days. If this PR is still relevant, please link the issue it fixes (e.g. Fixes #1234) or ask a maintainer to add the no-issue-needed label, and it can be reopened.

We are experimenting with this process to keep the review queue manageable, and it will sometimes get it wrong. If you think this PR should stay open, please just say so here and we will reopen it — no need to justify it at length.

Thanks again for contributing, and sorry for the noise if we closed this by mistake!

[sayakpaul]

Failing tests are unrelated.