[realfill] Fix the LR scheduler when num_train_epochs is passed in a distributed training env

[825pranav]

What does this PR do?

Part of #8384, for examples/research_projects/realfill/train_realfill.py (claimed here).

The LR scheduler was sized from the dataloader length before accelerator.prepare(), so with multiple processes it was built for the unsharded length and the LR never finished decaying. This computes the step count from the per-process length, same approach as #8312.

One difference from the other scripts: realfill doesn't pass lr_scheduler to accelerator.prepare(), so it steps once per batch on each process. I kept that and scale by gradient_accumulation_steps instead of num_processes.

Checked with a toy model using the same loop (unprepared scheduler, accelerator.accumulate, 2 epochs, GA=2, 24 samples), run with 2 CPU processes (gloo):

	scheduler total	steps taken	final lr
before, 1 proc	48	48	0.0000
after, 1 proc	48	48	0.0000
before, 2 procs	48	24	0.5455
after, 2 procs	24	24	0.0000

Minimal command:

accelerate launch --num_processes 2 train_realfill.py --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-inpainting --train_data_dir=data/flowerwoman --output_dir=out --resolution=512 --train_batch_size=4 --gradient_accumulation_steps=1 --unet_learning_rate=2e-4 --text_encoder_learning_rate=4e-5 --lr_scheduler=constant --lr_warmup_steps=100 --num_train_epochs=2

ruff check and ruff format --check pass.

Before submitting

• Did you read the contributor guideline?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.

Who can review?

@sayakpaul @geniuspatrick

[github-actions]

Hi @825pranav, 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.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[825pranav]

This covers one script from the umbrella tracker #8384, so it deliberately says "Part of #8384" instead of a closing keyword (a Fixes would close the whole tracker). Could a maintainer add no-issue-needed? Thanks!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[825pranav]

Thanks @sayakpaul for reviewing and merging! Happy to take another script from #8384 if that's useful.