Nova-MoE is a clean, modular, from-scratch PyTorch implementation of a Mixture-of-Experts (MoE) Causal Language Model designed for controlled sparsity ablations against an active-parameter-matched dense baseline on the TinyStories dataset.
- Hugging Face Model Hub: sarimahsan101/Nova-MoE-122M
# 1. Clone the repository
git clone https://github.com/sarimahsan/nova-moe.git
cd nova-moe
# 2. Install requirements
pip install -r requirements.txt
# 3. Generate story pulling model weights from Hugging Face
python generate.py --hf_repo sarimahsan101/Nova-MoE-122M --prompt "Once upon a time, there was a little girl named Lily"Nova-MoE consists of two matched model variants:
-
Dense Baseline: 8 transformer layers with dense SwiGLU Feed-Forward Networks (
$d_{ffn} = 2048$ ). -
MoE Model: 4 dense SwiGLU layers (layers 1, 3, 5, 7) + 4 MoE layers (layers 2, 4, 6, 8) with 8 experts using Expert-Choice Routing (
$d_{\text{expert ffn}} = 2048$ , capacity factor = 1.25).
graph TD
subgraph Input_Layer ["Input & Embedding Layer"]
Tokens["Input Token IDs (Batch, SeqLen)"] --> Embed["Token Embedding (Vocab: 8,000, d_model: 512)"]
end
subgraph Transformer_Block ["Nova-MoE Transformer Layer (Pre-Norm)"]
Embed --> Norm1["RMSNorm (d_model=512)"]
subgraph Attention_Sublayer ["Grouped-Query Attention (GQA 8:2) + RoPE"]
Norm1 --> Q_proj["Q Projection (8 heads, dim 64)"]
Norm1 --> K_proj["K Projection (2 heads, dim 64)"]
Norm1 --> V_proj["V Projection (2 heads, dim 64)"]
Q_proj --> RoPE["Apply RoPE (Rotary Position Emb)"]
K_proj --> RoPE
RoPE --> SDPA["PyTorch SDPA (Causal Masking)"]
V_proj --> SDPA
SDPA --> O_proj["Output Projection (o_proj)"]
end
O_proj --> Add1["Residual Add (+ x)"]
Embed -. Residual Connection .-> Add1
Add1 --> Norm2["RMSNorm"]
subgraph MoE_Sublayer ["Expert-Choice Router & SwiGLU Experts (Layers 1, 3, 5, 7)"]
Norm2 --> Router["Router Linear (W_router: 512 -> 8)"]
Router --> ColSoftmax["Column-wise Softmax over Tokens (dim=0)"]
ColSoftmax --> TopK["Expert Top-C Token Selection (C = ceil(1.25 * T / 8))"]
TopK --> Exp1["SwiGLU Expert 1"]
TopK --> Exp2["SwiGLU Expert 2"]
TopK --> ExpDots["... Experts 3-7 ..."]
TopK --> Exp8["SwiGLU Expert 8"]
Exp1 --> IndexAdd["Index Add & Recombination (Weighted Sum)"]
Exp2 --> IndexAdd
ExpDots --> IndexAdd
Exp8 --> IndexAdd
end
IndexAdd --> Add2["Residual Add (+ x_attn)"]
Add1 -. Residual Connection .-> Add2
end
subgraph Output_Layer ["Final Norm & Projection"]
Add2 --> FinalNorm["Final RMSNorm"]
FinalNorm --> LMHead["LM Output Head (Tied Weight E^T)"]
LMHead --> Logits["Logits (Batch, SeqLen, 8000)"]
end
Input: x (Batch, SeqLen, d_model)
│
┌─────────────┴─────────────┐
│ │
▼ │
RMSNorm(x) │
│ │
GroupedQueryAttention │ (8 Q-heads, 2 KV-heads,
+ RoPE + SDPA │ head_dim=64)
│ │
▼ │
attn_output │
│ │
├───────────────────────────┘ (Residual Connection)
│
▼
x_attn = x + attn_output
│
┌─────────────┴─────────────┐
│ │
▼ │
RMSNorm(x_attn) │
│ │
┌───────────┴───────────┐ │
│ Config-Switchable │ │
│ FFN Layer │ │
└───────────┬───────────┘ │
│ │
Dense SwiGLU OR MoE FFN │ (MoE in layers 2, 4, 6, 8)
│ │
▼ │
ffn_output │
│ │
├───────────────────────────┘ (Residual Connection)
│
▼
Output: x_final
Unlike token-choice routing where each token chooses top-$k$ experts (requiring auxiliary load-balancing losses to prevent expert collapse), Expert-Choice Routing flips the choice: each expert selects its top-$C$ tokens.
Token Representations X (T × d_model)
│
▼
Router Linear (W_router)
│
▼
Softmax over Tokens (dim=0)
Scores S (T × n_experts)
│
┌─────────────────────────┼─────────────────────────┐
▼ ▼ ▼
Expert 1 Expert 2 Expert 8
Top-C Tokens Top-C Tokens Top-C Tokens
│ │ │
▼ ▼ ▼
SwiGLU Expert 1 SwiGLU Expert 2 SwiGLU Expert 8
│ │ │
└─────────────────────────┼─────────────────────────┘
│
▼
Weighted Recombination
(index_add_ to out_flat)
│
Unselected / Dropped Tokens = 0
(Pass through unchanged via residual: x + 0)
Key Formulas:
-
Expert Capacity:
$C = \lceil \text{capacity factor} \times \frac{T}{E} \rceil$ where$T = \text{Batch} \times \text{SeqLen}$ and$E = \text{num experts} = 8$ . -
Self-Balancing: Every expert processes exactly
$C$ tokens — perfect load balancing without auxiliary loss. -
Residual Safety: Tokens not selected by any expert return
$0$ from the MoE layer, passing through unchanged via the block's residual connection ($x + 0$ ).
By construction, active parameter counts are strictly matched between the Dense Baseline and the MoE variant to isolate sparsity benefits.
| Model | Active Parameters | Total Parameters | Parameter Gap |
|---|---|---|---|
|
Dense Baseline (8 dense layers, |
34,513,408 | 34,513,408 | Baseline |
|
MoE Model (4 dense + 4 MoE layers, |
34,529,792 | 122,610,688 | 0.047% (< 0.1%) |
Note: Since 1 expert is activated per token in each MoE layer, setting $d_{\text{expert ffn}} = 2048$ makes the active compute of 1 expert equal to 1 dense layer.
├── components/
│ ├── norms.py # RMSNorm (Pre-norm, learnable weight)
│ ├── rope.py # Rotary Positional Embeddings (RoPE)
│ ├── attention.py # Grouped-Query Attention (GQA 8:2, RoPE + SDPA)
│ ├── router.py # Expert-Choice Router (top-C tokens per expert, drop metrics)
│ ├── moe.py # MoE FFN layer (experts + router + residual safety)
│ ├── feedforward.py # SwiGLU FFN (dense layers)
│ └── block.py # TransformerBlock with config-switchable dense/MoE FFN
├── models/
│ └── transformer.py # CausalLM with tied embeddings & active/total param counting
├── trainer/
│ ├── trainer.py # Modern PyTorch AMP (fp16) trainer with cosine LR schedule
│ └── callbacks.py # CheckpointCallback & ExpertUtilizationCallback
├── utils/
│ ├── data.py # TextDataset & create_dataloader (DistributedSampler for DDP)
│ ├── tokenizer.py # Custom BPE Tokenizer (vocab size 8,000 via HF tokenizers)
│ ├── seed.py # Deterministic seed utility
│ └── config.py # ModelConfig & TrainingConfig dataclasses with YAML parsing
├── analysis/
│ └── expert_utilization.py # Expert distribution & capacity drop rate analysis tools
├── configs/
│ ├── dense_baseline.yaml # Configuration for Dense Baseline model
│ └── moe.yaml # Configuration for MoE variant
├── tests/
│ ├── test_environment.py # T4 sanity, fp16 autocast, SDPA backend, DDP checks
│ └── test_model.py # Unit test suite (components, router edge cases, param match, grad check)
└── train.py # Main entrypoint supporting single GPU / CPU and torchrun DDP
Run the full pytest suite locally before training:
python -m pytest tests/ -vtest_rmsnorm: Verifies unit RMS scaling and finite gradient flow.test_rope: Verifies position 0 rotation identity and position-dependent rotary encoding.test_gqa: Verifies GQA (8:2 head ratio) shape transformation and causal masking.test_swiglu_ffn: Verifies SwiGLU gate/up/down matrix operations.test_dense_transformer_forward_and_params: Verifies causal LM forward pass, tied embedding memory sharing, and 34.5M param count.test_expert_choice_router_edge_cases_and_gradients: Verifies ceiling rounding when token count is odd, capacity drop calculation, and gradient flow toW_router.test_moe_ffn_forward_and_metrics: Verifies MoE forward output shape and drop metrics recording.test_param_matching_dense_vs_moe: Asserts active parameter gap between Dense and MoE is < 0.05%.test_moe_gradient_check_all_experts_receive_gradients: Verifies every expert across all 4 MoE layers receives non-zero gradients over multiple steps.
Training uses Data Parallelism (DDP) across both T4 GPUs (each GPU holds a full model replica; no cross-GPU expert sharding over PCIe).
# Train Dense Baseline
torchrun --nproc_per_node=2 train.py --config configs/dense_baseline.yaml
# Train MoE Variant
torchrun --nproc_per_node=2 train.py --config configs/moe.yaml- Checkpoints are saved every 30 minutes in
checkpoints/(resilient to Kaggle's 12-hour session cap). - Expert token assignments and capacity drop rates are logged to
logs/expert_utilization.json.
After training the MoE model, run the analysis helper:
from analysis.expert_utilization import print_expert_summary
print_expert_summary("logs/expert_utilization.json")This outputs average, peak, and final capacity drop rates across training steps.