miniMoE

miniMoE is a sparse GPT-style language model I trained from scratch on 10B FineWeb-Edu tokens across 8x A100 GPUs with DDP. This repo holds the trained checkpoint. The code, training pipeline, and full write-up are on GitHub: https://github.com/maokner/miniMoE

What's in here

minimoe_step_0019073.pt is the final checkpoint, taken at step 19,073 after ~10B tokens. It is a full training checkpoint, so alongside the model weights it keeps the optimizer and RNG state. That means you can either sample from it or resume training where it left off.

Model

Item Value
Context length 1,024
Vocabulary 50,304 GPT-2 BPE tokens
Transformer blocks 6
Hidden size 768
Experts per block 8
Active experts per token 2
Total parameters 280.4M
Active parameters per token 110.4M

Every dense feed-forward block is replaced with an 8-expert top-2 MoE layer. Attention, embeddings, layer norms, the router, and the tied output head are shared across tokens.

Results

These come from the completed training log.

Metric miniMoE
Train loss 3.0869
Validation loss 3.0409
Test loss 3.0725

On the full HellaSwag benchmark (10,042 examples, completion-style scoring), this checkpoint edges out GPT-2 124M while activating fewer parameters per token.

Model Active params per token HellaSwag
miniMoE (this checkpoint) 110.4M 31.2%
GPT-2 124M 124M 29.6%

The instruction-tuned variant holds HellaSwag steady at 30.6%; see the GitHub repo for the full evaluation and routing ablations.

Usage

Clone the GitHub repo, drop the checkpoint next to the code, and sample from it:

pip install -r requirements.txt
python sample.py -c minimoe_step_0019073.pt -p "Sparse expert language models"

To pull just the checkpoint from here:

from huggingface_hub import hf_hub_download

path = hf_hub_download("mokner123/miniMoE", "minimoe_step_0019073.pt")
Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train mokner123/miniMoE