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

[adi-IL]

Fixes #8384

This PR updates only examples/text_to_image/train_text_to_image_sdxl.py, matching the #8312 pattern already applied to train_text_to_image.py and the DreamBooth trainers.

What was wrong

When --num_train_epochs is used (so max_train_steps is derived), the SDXL trainer still built the LR schedule from the unsharded dataloader length and ignored accelerator.num_processes. In a multi-process run that makes the scheduler finish too early.

What changed

• Warmup and training steps passed to get_scheduler are scaled by accelerator.num_processes.
• Step counts used to build the schedule assume the dataloader will be sharded.
• After accelerator.prepare, we warn if the prepared dataloader length does not match that assumption.
• Added a smoke test that runs the script with --num_train_epochs 1.

Coordination

Claimed on #8384: #8384 (comment)

@sayakpaul @geniuspatrick

Minimal training command using num_train_epochs

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fix"
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch train_text_to_image_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --pretrained_vae_model_name_or_path=$VAE_NAME \
  --dataset_name=$DATASET_NAME \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 --gradient_checkpointing \
  --num_train_epochs=1 \
  --learning_rate=1e-06 --lr_scheduler="constant" --lr_warmup_steps=0 \
  --mixed_precision="fp16" \
  --output_dir="sdxl-naruto-model"

Tests I ran

python3 -m py_compile examples/text_to_image/train_text_to_image_sdxl.py

That passed. I did not run the example pytest here because this machine does not have the package or its test extras installed. The new test_text_to_image_sdxl_num_train_epochs plus the existing test_text_to_image_sdxl should cover both --max_train_steps and --num_train_epochs in CI.

Self-review

• Scope is one official trainer, as requested on the issue.
• The new math is copied from train_text_to_image.py, not reinvented.
• No public API change.
• I left train_text_to_image_lora_sdxl.py and the other remaining scripts alone so this stays one script per PR.