Skip to content

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

Open
adi-IL wants to merge 1 commit into
huggingface:mainfrom
adi-IL:fix/textual-inversion-lr-scheduler-epochs
Open

[textual_inversion.py] Fix the LR scheduler when num_train_epochs is passed in a distributed training env#14546
adi-IL wants to merge 1 commit into
huggingface:mainfrom
adi-IL:fix/textual-inversion-lr-scheduler-epochs

Conversation

@adi-IL

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

Copy link
Copy Markdown

Fixes #8384

Fixes the LR scheduler step miscalculation in examples/textual_inversion/textual_inversion.py when --num_train_epochs is used with distributed training.

What was wrong
When --max_train_steps is None, the script computed num_update_steps_per_epoch from the unsharded dataloader length and set max_train_steps = num_train_epochs * num_update_steps_per_epoch before accelerator.prepare. After sharding, the effective dataloader length is ceil(len(dataloader)/num_processes), so the scheduler was created with training_steps = max_train_steps * num_processes while the actual optimization steps were smaller. The mismatch caused the scheduler to step incorrectly and max_train_steps to be recomputed without the sharding guard.

What changed
Mirrors the fix already applied to textual_inversion_sdxl.py in #11557 and the pattern from #8312. Now computes len_train_dataloader_after_sharding, derives num_update_steps_per_epoch from the sharded length, and sets num_training_steps_for_scheduler = num_train_epochs * num_update_steps_per_epoch * num_processes (or max_train_steps * num_processes when provided). Introduces num_warmup_steps_for_scheduler and defers max_train_steps assignment until after accelerator.prepare with the warning guard when lengths do not match. One file, one script as requested in #8384.

Coordination
Tracks #8384. The issue lists this script as [ ] textual_inversion.py and says to target one script per PR and mention @sayakpaul and @geniuspatrick for review. Standing permission for this checklist item, same as the already merged siblings #14527, #14528, #14540.

Test plan

  • /home/adi-IL/Desktop/.gh-pr-routine/.venv/bin/python -m py_compile examples/textual_inversion/textual_inversion.py -> py_compile OK
  • python -c "import ast; ast.parse(open('examples/textual_inversion/textual_inversion.py').read())" -> ast parse ok
  • Manual arithmetic check with dummy lengths (10/100 samples, 2/4 processes, grad_accum 1/2) confirms new num_training_steps_for_scheduler equals max_train_steps * num_processes after sharding, while the old code produced double or mismatched steps.

Self-review notes
Checked against .ai/references/review-rules.md: no new dependencies, no model or pipeline changes, preserves existing docstrings and comments, matches the sibling fix in textual_inversion_sdxl.py:802-853, no dead code, only the scheduler math block changed. The warning uses the same message as the sibling.

cc @sayakpaul @geniuspatrick

Scale warmup and training steps by accelerator.num_processes so

--num_train_epochs builds the same schedule as the other official

trainers after huggingface#8312.

Tracks huggingface#8384.
@github-actions github-actions Bot added examples size/S PR with diff < 50 LOC labels Aug 20, 2026
@github-actions

Copy link
Copy Markdown
Contributor

Hi @adi-IL, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

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

Labels

examples size/S PR with diff < 50 LOC

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