ZeRO 3 example does not run

[StellaAthena]

The ZeRO 3 example does not run. The main problem appears to be that the InitContext function does not actually exist despite being called by pretrain_gpt2.py. I have tried to introduce some changes to get it to run (incl. changing the batch size, the initialization function, and some of the inputs to the initialization function) but gave up after it threw the error variable beta1 is referenced before assignment. I think that has to do with something wonky in the optimizer?

[ShadenSmith]

Hi @StellaAthena , thanks for catching this. I missed merging a PR before release. Please let us know if you have further problems with the example.

[StellaAthena]

@ShadenSmith Unfortunately it still does not run. The current error is that beta1 is referenced before assignment. I can get a stack trace in a couple minutes.

[thies1006]

Same issue. Is there a workaround to get Megatron with Zero3 running?

[StellaAthena]

Same issue. Is there a workaround to get Megatron with Zero3 running?

I was unable to ever get this implementation to run. However I do have ZeRO-3 implemented in my megatron variant. You can check out the ZeRO-3-Brrrrr branch of www.github.com/eleutherai/gpt-neox if you want to give it a try. Unfortunately it’s not currently compatible with the main branch but I hope to have that remedied this weekend or next week (it has nothing to do with ZeRO, just internal changes that need to be made for consistency with other commits).

[mouxiaofeng1981]

the same issue

[fumihwh]

beta1 referenced before assignment is related to pytorch issue pytorch/pytorch#53354
I've added --cpu-optimizer to avoid using pytorch AdamW and try use DeepSpeedCPUAdam.
But still not work, get SIGSEGV: 11