MoE kernel

[ByronHsu]

🚀 The feature, motivation and pitch

Currently the most popular library might be https://github.com/databricks/megablocks. Would be interesting if we can implement it in triton and make it HF compatible

Alternatives

No response

Additional context

No response

[S1ro1]

Will do a research on this more, if anyone has any insights on what could/should be implemented, resp. details on to how, cc me.

[S1ro1]

Maybe a preliminary would be to support for example mixtral/nllb_moe from huggingface, to have the integration ready when the layers are done?

[yundai424]

@S1ro1 one straightforward idea is to parallelize expert forward (just like what megablock impl does). Right now in HF model code the MoE block is performed sequentiallyexpert-by-expert. Not sure if it's worth implementing load balancing loss too, haven't seen an actual profiling trace of MoE model training

[S1ro1]

@yundai424 Haven't seen one either, gonna try patching either Mixtral or Nllb with our kernels and profile it, will decide on what to do after that I guess.
Implementing dMoE (dropless MoE) could also be interesting. Will try to send the profiler benchmarks tomorrow and could discuss more in depth. Also I suppose Mixtral > Nllb.

Edit: to address your comment, parallelizing the experts is certainly a low hanging fruit