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