Skip to content

[train_t2i_adapter_sdxl.py] Fix the LR scheduler when num_train_epochs is passed in a distributed training env - #14540

Open
adi-IL wants to merge 1 commit into
huggingface:mainfrom
adi-IL:fix/sdxl-t2i-adapter-lr-scheduler-distributed
Open

[train_t2i_adapter_sdxl.py] Fix the LR scheduler when num_train_epochs is passed in a distributed training env#14540
adi-IL wants to merge 1 commit into
huggingface:mainfrom
adi-IL:fix/sdxl-t2i-adapter-lr-scheduler-distributed

Conversation

@adi-IL

@adi-IL adi-IL commented Aug 20, 2026

Copy link
Copy Markdown

Fixes #8384

Follow-up to #14527 and #14528, still one script. This PR updates examples/t2i_adapter/train_t2i_adapter_sdxl.py.

What was wrong

When --num_train_epochs is used, train_t2i_adapter_sdxl.py built the LR schedule from the unsharded dataloader and ignored accelerator.num_processes, resulting in mismatched warmup and total training steps across distributed environments.

What changed

Applies the #8312 pattern already established across the DreamBooth, ControlNet, and Text-to-Image trainers:

  • Warmup and training steps passed to get_scheduler are scaled by accelerator.num_processes.
  • Step counts used to build the schedule account for the sharded dataloader length (math.ceil(len(train_dataloader) / accelerator.num_processes)).
  • After accelerator.prepare, a warning is logged if the prepared dataloader length deviates from that expectation.
  • Added test_t2i_adapter_sdxl_num_train_epochs in examples/t2i_adapter/test_t2i_adapter.py to cover the --num_train_epochs path.

Coordination

Covered by the ongoing community tracker #8384. cc @sayakpaul @geniuspatrick

Minimal training command using num_train_epochs

export MODEL_DIR="stabilityai/stable-diffusion-xl-base-1.0"
export OUTPUT_DIR="sdxl-t2i-adapter-model"

accelerate launch train_t2i_adapter_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_DIR \
  --output_dir=$OUTPUT_DIR \
  --dataset_name=fusing/fill50k \
  --mixed_precision="fp16" \
  --resolution=1024 \
  --learning_rate=1e-5 \
  --num_train_epochs=2 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --seed=42

Tests I ran

pytest examples/t2i_adapter/test_t2i_adapter.py
python3 -m py_compile examples/t2i_adapter/train_t2i_adapter_sdxl.py examples/t2i_adapter/test_t2i_adapter.py
ruff check examples/t2i_adapter/train_t2i_adapter_sdxl.py examples/t2i_adapter/test_t2i_adapter.py
ruff format --check examples/t2i_adapter/train_t2i_adapter_sdxl.py examples/t2i_adapter/test_t2i_adapter.py

All passed (2 passed in 84.45s).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Community] Help us fix the LR schedulers when num_train_epochs is passed in a distributed training env

1 participant