Modularize components of megablocks layer builder

[dakinggg]

Cleans up the sections of the builder functions for megablocks to be more modular, and modularizes the device mesh creation in the megablocks args.

TODO:

• Manual test that training still works

[mvpatel2000]

LGTM but @vchiley should look as well

[vchiley]

Manual test that training still works

Don't we have tests for MoEs (MoE testing config here)?
(I'm not against doing more testing, just asking)

[dakinggg]

@vchiley there are some simple MoE tests, I just wanted to double check e2e still works.

[dakinggg]

manual test done, merging