[QUESTION] How to enable ZeRO 2/3 stages ?

[polisettyvarma]

How to enable ZeRO 2/3 stages ?
similar to #589

[lmcafee-nvidia]

I responded to this on #589.

[polisettyvarma]

please convert this issue to feature request for ZeRO 2/3
Thank you.

[carolove]

i think article,https://www.deepspeed.ai/tutorials/megatron/, is useful.
deepspeed ZeRO 1/2 works with Megatron-lm latest code.

[polisettyvarma]

@carolove Thanks for the inputs, i am familiar with deepspeed framework to enable all ZeRO stages. here query is regarding enabling ZeRO in this repo natively.
can you please share commits which added ZeRO 2 support in latest code of this repo.
Thank you.

[carolove]

I also look for such example~.

[SeunghyunSEO]

megatron-lm now has its own zero-1 (it is called distributed optimizer in this project), but if u are more familiar with deepspeed, then how bout using deepspeed-megatron, @polisettyvarma ?
And to my best knowledge, zero-3 is not compatible with model parallelism (TP or PP) of megatron-lm.
zero-3 reduce vram memory and improve throughput by partitioning and broadcasting model parameters but TP or PP partition its own way and rather communicate activations (all-reduce activations for backward and forward).
So TP or PP has no room for communicating model parameters.

[polisettyvarma]

Thank you @SeunghyunSEO for your inputs. Yes Megatron-DeepSpeed repo can be used but it's not up to date with Megatron-LM. I agree on Zero > 1 is not compatible with PP.
My request here is some similar feature of ZeRO on Megatron-LM.

[deepakn94]

We should have PyTorch FSDP support compatible with TP in the next couple of weeks.

[polisettyvarma]

Thank you @deepakn94 for sharing this information.