Horovod: adjust base LR used by schedulers to scale with the number of workers

[tgaddair]

In #2574 it was observed that the learning rate used to initialize learning rate schedulers is in conflict with the way we scale the learning rate with the number of Horovod workers. Because the learning rate schedulers are initialized before the scaling, the scaling would be overridden. This PR fixes this so that optimizers and LR schedulers scale up correctly.

Follow-ups to consider would be:

• Make scaling work for other backends like DDP.
• Add an option to configure the LR scaling in the Trainer.

[codecov]

Codecov Report

Merging #2626 into master will increase coverage by 0%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #2626    +/-   ##
=======================================
  Coverage      91%     92%            
=======================================
  Files          72      74     +2     
  Lines        6131    6321   +190     
=======================================
+ Hits         5599    5791   +192     
+ Misses        532     530     -2     

[mergify]

This pull request is now in conflict... :(

[tgaddair]

@williamFalcon @Borda failing tests seem unrelated. Could you take a look at the PR and see if everything makes sense?

As a follow-up, I want to add a param to Trainer so we can override the LR scale when using Horovod/DDP.

[Borda]

@williamFalcon @Borda failing tests seem unrelated. Could you take a look at the PR and see if everything makes sense?

yes, no connection to this PR

[amorehead]

Hello. I was curious if there have been any other PRs that have added the ability for users to specify which learning rate scaling strategy should be used with e.g., DDP/Horovod. As a frequent user of DDP with Lightning, I would love to see the ability to have my optimizer's learning rate automatically scaled according to the effective batch size across nodes and their GPUs (i.e., the total world size).