ZeRO optimizer LAMB compatibility

[msdejong]

My use case for this library is mostly for BERT models, as opposed to Megatron+ sized LMs. ZeRO in that context is mainly useful for fitting larger batch sizes and increasing throughput. For that reason, I'm wondering if/when you are planning on adding a ZeRO compatible LAMB optimizer.

[jeffra]

This feature is definitely on our radar, I don’t have a concrete timeline though. But we’ll be sure to update this thread when we have more info.

[samyam]

@Michiel29 as Jeff mentioned, this is something we definitely want, but currently its not on the top of our priority list at the moment. We be more than happy to accept contributions on this from the community. If you would like to contribute, we can definitely share our discussions on what would be needed to support Lamb with ZeRO. Please let us know.

[LiweiPeng]

@jeffra I'd like to use lamb with zero too. Can you share updates on using lamb + zero?

[e-sha]

@Michiel29 as Jeff mentioned, this is something we definitely want, but currently its not on the top of our priority list at the moment. We be more than happy to accept contributions on this from the community. If you would like to contribute, we can definitely share our discussions on what would be needed to support Lamb with ZeRO. Please let us know.

Why Lamb is not in a list of supported optimizers?
Please share your ideas of what would be needed to support Lamb with ZeRO.

[e-sha]

I realized that we have to add two updates to use LAMB with ZeRO:

1. Split model parameters preserving whole tensors on a single node (here)
2. Preserve splitting of a parameter vector into tensors at optimizer initialization stage (here)

Am I right?

[gongjingcs]

I realized that we have to add two updates to use LAMB with ZeRO:

1. Split model parameters preserving whole tensors on a single node (here)
2. Preserve splitting of a parameter vector into tensors at optimizer initialization stage (here)

Am I right?

Hi，I am not very familiar with lamb, can you explain in detail your thoughts on lamb combined with zero ？thanks alot

[gongjingcs]

@Michiel29 as Jeff mentioned, this is something we definitely want, but currently its not on the top of our priority list at the moment. We be more than happy to accept contributions on this from the community. If you would like to contribute, we can definitely share our discussions on what would be needed to support Lamb with ZeRO. Please let us know.

Hi, I'd like to use lamb with zero stage1. could you share your ideas of what would be needed to support Lamb with ZeRO.
Do we need to modify the implementation of fused lamb or modify the implementation of zero ？
If I know the detailed idea, I can contribute the code. Looking forward to your reply