Implement Mixture of Depth and Experts (MoDE)

[casper-hansen]

Given that MegaBlocks is highly optimized for sparse MoE models like Mixtral, I am requesting support for a variant recently termed as MoDE by Google DeepMind. Benefits include much faster training and inference due to increased sparsity.

Paper: https://arxiv.org/abs/2404.02258

I found two implementations:

• https://github.com/epfml/llm-baselines/blob/mixture_of_depth/src/models/mod.py
• https://github.com/kyegomez/Mixture-of-Depths/blob/main/mixture_of_depths/main.py

[ehartford]

Very interested in this

[mvpatel2000]

We'd love community PRs for this! Happy to help review and design. It's not currently on our roadmap, but we are evaluating it.