[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.

[adi-IL]

Hi @sayakpaul, gentle check-in on this fix for the LR scheduler in train_text_to_image_sdxl.py under distributed training. Ready whenever you have a chance to review.

[adi-IL]

Closing in favor of the umbrella PR that packs the same LR-scheduler num_train_epochs fix across all of these example scripts: #14743

That PR supersedes #14527#14528#14540#14546#14597 so review can happen once. Tip contents were ported onto current main there.