Add support for Expert Parallel

[a-r-r-o-w]

#371 adds support for HiDream, which is a text-to-image generation model based on an MoE architecture. This is a good time to add experimental expert parallel training support