very cool work! we did something similar in the context of the Swiss AI initiative (https://www.swiss-ai.org/ https://www.swiss-ai.org/) here: https://github.co

very cool work! we did something similar in the context of the Swiss AI initiative (https://www.swiss-ai.org/ https://www.swiss-ai.org/) here: https://github.com/swiss-ai/MoE https://github.com/swiss-ai/MoE. The implementation is as simple and fast as nanoGPT and works with our modular llm-baselines codebase (https://github.com/epfml/llm-baselines https://github.com/epfml/llm-baselines) for experimenting with transformers and different datasets :)