Add support for MoE model architecture

[gnadathur]

### Tasks
- [ ] Add barebone MoE without expert parallelism
- [ ] Prototype expert parallel with MoE
- [ ] e2e integration and perf validation with torchtrain

[lessw2020]

Some additional comments on MoE:

1 - A recent Meta paper highlights that you can train the experts independently (as usual dense models) and then merge the FFN layers to form the experts and do a bit of final training to integrate. They claim this is more efficient than training the MoE from scratch, and might be an option to support given the progress on normal dense training while still building MoE's:
Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM
https://arxiv.org/abs/2403.07816

2 - for 3D /Pippy - this paper shows how to support Pipeline Parallel MoE while lowering communication:
https://arxiv.org/abs/2304.11414

[aonrdbit]

Great! is there any update to this feature?

[tianyu-l]

This is under planning and will be worked on soon. Any suggestions are welcome.