Add support for mutransfer

[Quentin-Anthony]

We should add support for mutransfer: https://github.com/microsoft/mup

Appears non-trivial, but not as difficult as MoE. We'd have to modify the model itself. https://github.com/microsoft/mup/blob/main/examples/Transformer/model.py appears especially relevant. A good workflow would be:

• [x] Run mutransfer transformer example to gain familiarity with how their codebase behaves.
• [x] Make necessary changes in gpt-neox/megatron/model/ to use mup. Probably mostly in transformer.py
• [x] Add muSGD and muAdam to gpt-neox/megatron/optimizers.py
• [x] Update gpt-neox/megatron/training.py to allow previous features to be selected during training
• [x] Add neox_args to choose the options in step 4
• [ ] Test and add example config
• [ ] Add args for saving coord check plot

[nsarka]

nsarka/mup-support has my changes for this so far. I haven't tested it.

There's one more thing to add to this list. Mup can generate a plot that's helpful for checking the correctness of the implementation. https://github.com/microsoft/mup#checking-correctness-of-parametrization

• [ ] Add args for saving coord check plot

[StellaAthena]

nsarka/mup-support has my changes for this so far. I haven't tested it.

There's one more thing to add to this list. Mup can generate a plot that's helpful for checking the correctness of the implementation. https://github.com/microsoft/mup#checking-correctness-of-parametrization

• [ ] Add args for saving coord check plot

Great work! Thank you for this contribution ^_^

Don’t forget to add yourself as a library contributor in the readme as well 😉

[nsarka]

Thanks Stella! I added myself as a contributor in the draft PR here https://github.com/EleutherAI/gpt-neox/pull/704 :)

[StellaAthena]

Closed as completed by #704