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:

• Run mutransfer transformer example to gain familiarity with how their codebase behaves.
• Make necessary changes in gpt-neox/megatron/model/ to use mup. Probably mostly in transformer.py
• Add muSGD and muAdam to gpt-neox/megatron/optimizers.py
• Update gpt-neox/megatron/training.py to allow previous features to be selected during training
• 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 #704 :)

[StellaAthena]

Closed as completed by #704