[feature request] Support for MoE models

[Cameron7195]

I believe MoE models are not yet supported with ArcticTrainer? According to this thread, they are not yet optimized, and have trouble if one does not implement the leaf module patch
#250 (comment)

[sfc-gh-sbekman]

In theory it should just take telling Deepspeed to have MoE modules as leaf modules, in practice it only works in some situations. I found that setting zero's overlap_comm: false + reduce_bucket_size big enough that it doesn't trigger grad reduction during the backward (only at the end) helps, but it still breaks in some situations on the Deepspeed side.

@sfc-gh-reyazda and I have just started working on adding MoE support in AT here #272 - but it's still very early. we want to redesign DeepspeedMoE to bring it to up to modern techniques, while simplifying it on the way and using AT as its day-0 integration framework.

[Cameron7195]

Got it, will try these thing. Thanks!