Is it feasible to train Mamba+MoE using torchtitan?

[imhuim982]

Torchtitan is a very exciting project, and its support for native torch models is more researcher-friendly. I am currently working on a mamba+moe+sparse attention architecture, and I want to train a model with 15B parameters and about 1.5B activations. However, I am not sure if expert parallelism is well supported, and whether architectures like Mamba are fully compatible. I would greatly appreciate it if someone could give some suggestions.