Implementation of LAMB optimizer

[Atharva-Phatak]

🚀 The feature

#6323 Mentions that torchvision is looking to implement LAMB optimizer.
@datumbox I would very much like to take this issue and create a PR.

Motivation, pitch

LAMB optimizer was created because LARS optimizer performed poorly on model with attention mechanism (mainly). LAMB has shown to achieve very good performance gains across various tasks and I believe that it should be implemented in torchvision.

Alternatives

No response

Additional context

No response

[Dsantra92]

Are you still working on this issue?

[Atharva-Phatak]

@Dsantra92 Yes

[ashikshafi08]

Is anyone currently working on this, or can I plan to work on it?

[Atharva-Phatak]

@ashikshafi08 Go ahead and implement it as I don't have the bandwidth to do so.

[ashikshafi08]

I am not sure where would I push this code. Like under what directory should the optimizers go in?

[datumbox]

@ashikshafi08 See this PR which is in progress for adding LARS. pytorch/pytorch#88106

[janeyx99]

Hey! To set some expectations on this from the core side: we currently do not have much review bandwidth for optimizers and so it may take a while for us to get around to verifying whether LAMB should live in core and then reviewing a LAMB implementation.