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

[adi-IL]

Fixes#8384

Follow-up to #14527, still one script. This PR only updates examples/text_to_image/train_text_to_image_lora_sdxl.py.

What was wrong

When --num_train_epochs is used, this trainer still built the LR schedule from the unsharded dataloader and ignored accelerator.num_processes. The official README already documents --num_train_epochs=2 for this script, so the broken path is the documented one.

What changed

Same #8312 math already in train_text_to_image_lora.py and the DreamBooth trainers:

• 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 test_text_to_image_lora_sdxl_num_train_epochs so the --num_train_epochs path is covered.

Coordination

Covered by the existing #8384 thread. No new issue 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_lora_sdxl.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--pretrained_vae_model_name_or_path=$VAE_NAME \
--dataset_name=$DATASET_NAME --caption_column="text" \
--resolution=1024 --random_flip \
--train_batch_size=1 \
--num_train_epochs=2 --checkpointing_steps=500 \
--learning_rate=1e-04 --lr_scheduler="constant" --lr_warmup_steps=0 \
--mixed_precision="fp16" \
--seed=42 \
--output_dir="sd-naruto-model-lora-sdxl"

Tests I ran

python3 -m py_compile examples/text_to_image/train_text_to_image_lora_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. CI should run the new test_text_to_image_lora_sdxl_num_train_epochs plus the existing LoRA SDXL smoke tests.

Self-review

• Scope is one official trainer, as requested on the issue.
• The new math is copied from train_text_to_image_lora.py, not reinvented.
• No public API change.
• I left the remaining [Community] Help us fix the LR schedulers when num_train_epochs is passed in a distributed training env #8384 scripts alone.

[sayakpaul]

What is the difference from #14527?

[adi-IL]

#14527 updates the full fine-tuning script (examples/text_to_image/train_text_to_image_sdxl.py), while this PR (#14528) updates the LoRA fine-tuning script (examples/text_to_image/train_text_to_image_lora_sdxl.py). Both apply the same LR scheduler fix for #8384 across multi-process environments, separated to keep changes isolated to one script per PR.

[adi-IL]

Hi @sayakpaul, following up on the explanation above. Let me know if this clarifies the difference between the full SDXL trainer and LoRA trainer fixes, or if anything else is needed.

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