Wrong learning rate scheduler training step count for examples with multi-gpu.

[eliphatfs]

Describe the bug

For cosine learning rate schedules, the learning rate should approach zero when the training ends. However, due to accelerator prepare logic a learning rate scheduler will be stepped N times for each optimization step when there are N GPUs. Thus, we need to multiply the steps by num_processes when constructing schedulers.

Sample learning rate curve running dreambooth example on 4 GPUs:

Reproduction

Run the example scripts passing cosine learning rate schedules.

Logs

No response

System Info

• diffusers version: 0.17.1
• Platform: Linux-5.4.0-150-generic-x86_64-with-debian-buster-sid
• Python version: 3.7.13
• PyTorch version (GPU?): 1.12.1 (True)
• Huggingface_hub version: 0.15.1
• Transformers version: 4.30.2
• Accelerate version: 0.20.3
• xFormers version: 0.0.20+1dc3d7a.d20230628
• Using GPU in script?: Y
• Using distributed or parallel set-up in script?: DDP

Who can help?

@williamberman, @sayakpaul, @yiyixuxu

[sayakpaul]

Thanks fore reporting. Would you maybe like to open a PR to fix this?

[eliphatfs]

Do you think this should be fixed inside accelerator, as stepping xGPU times is counter-intuitive to me and there doesn't seem to be evidence that stepping xGPU times is equivalent to step fewer times with single GPU; or should we just change the number of steps in examples?

[sayakpaul]

I think modifying the number of steps is a good idea when there are multiple GPUs being used. But I would like to also check with @muellerzr regarding:

Do you think this should be fixed inside accelerator,

[muellerzr]

Yep, steps should be modified, this is how Accelerate schedulers work

[eliphatfs]

In case anyone else does unnecessary and duplicate work, I have forked the repository and will start a pull request soon.

[eliphatfs]

I have finished the changes and proposed a PR. Please review. Thanks!

[eliphatfs]

Another problem. Seems that we should not multiply the steps by gradient accumulation steps.
https://github.com/huggingface/accelerate/blob/main/src/accelerate/scheduler.py#L61

[sayakpaul]

This is something we follow across all our examples. Could you expand on why?

[eliphatfs]

Inside the accelerate scheduler (I attached a pointer in the previous comment), the code is already skipping steps when gradient sync is inactive (accumulation steps), so we shall not multiply with accumulation for total training steps.
I happen to run an experiment for my research with accumulation and cosine annealing. The schedule is long-streched. I may attach the full curve after the run finishes.

Cosine schedule with 30k steps (completed, red) with peak 1e-4 and 20k steps (running, orange) with peak 1e-5. It may be somewhat hard to visually compare, but you can see that it is decreasing much slower than expected.

[sayakpaul]

Hmm makes sense. Thanks for flagging this.

I would also like to take advice from @muellerzr here a bit.

Once we have a consensus, let's jam #3983.

Also cc: @williamberman and @patil-suraj so that they are aware.

[muellerzr]

Yes, so as noted you should multiply by the number of processes, but not by the gradient accumulation steps. Accelerate handles that already internally.

[eliphatfs]

I have updated my PR with removal of multiplications by the accumulation steps.