add dense and moe checkpoints

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,164 @@
+---
+license: mit
+tags:
+  - mixture-of-experts
+  - gpt2
+  - research
+  - expert-specialization
+language:
+  - en
+datasets:
+  - codeparrot/codeparrot-clean
+  - allenai/ai2_arc
+  - allenai/c4
+base_model:
+  - openai-community/gpt2
+library_name: transformers
+pipeline_tag: text-generation
+---
+
+# MoE Emergence
+
+Checkpoints from a research project studying expert specialization in Mixture-of-Experts models. I fine-tuned GPT-2 small on three domains -- code, math, and prose -- to see whether experts naturally specialize by domain when given the right routing incentives.
+
+Short answer: they do. MoE beats the dense baseline by 3.6% overall and 14% on math, with zero expert collapse across 10,000 training steps.
+
+---
+
+## 1. Results
+
+| Metric | Dense Baseline | MoE (8 experts, top-1) | Delta |
+|---|---|---|---|
+| eval/loss | 2.157 | 2.080 | -3.6% |
+| loss_code | 1.554 | 1.521 | -2.1% |
+| loss_math | 2.023 | 1.740 | -14.0% |
+| loss_prose | 3.485 | 3.541 | +1.6% |
+| perplexity | 8.64 | 7.91 | -8.4% |
+
+Math benefits the most from expert routing. Prose is the one domain where dense wins; diverse web text doesn't lend itself to clean expert specialization. The MoE model crossed the dense baseline at step ~3,600 (36% of training).
+
+---
+
+## 2. Files
+
+```
+dense-baseline/
+├── final-model.safetensors    # 622 MB -- dense GPT-2, 124M params
+├── final-model.json           # metadata sidecar (config, metrics)
+├── ckpt-step-4999.pt          # 1.4 GB -- full resume checkpoint (optimizer, scheduler, RNG)
+└── metrics.jsonl              # per-step training + eval metrics
+
+moe-main/
+├── final-model.safetensors    # 1.1 GB -- MoE GPT-2, 257M params (8 experts × 4 layers)
+├── final-model.json           # metadata sidecar (config, metrics)
+├── ckpt-step-9999.pt          # 2.9 GB -- full resume checkpoint
+└── metrics.jsonl              # per-step training + eval metrics
+```
+
+The `.safetensors` files are the trained model weights. The `.pt` files contain the full training state for resuming runs (optimizer, LR scheduler, RNG states). The `.json` sidecars store architecture config and final eval metrics.
+
+---
+
+## 3. Usage
+
+Clone the [source repo](https://github.com/sumitdotml/moe-emergence) and install dependencies:
+
+```bash
+git clone https://github.com/sumitdotml/moe-emergence.git
+cd moe-emergence
+uv sync
+```
+
+Run inference with a trained checkpoint:
+
+```bash
+# MoE model
+uv run python moe_emergence/gpt2_inference.py \
+  --checkpoint checkpoints/moe-main/final-model \
+  --prompt "def fibonacci(n):" \
+  --sample --temperature 0.8
+
+# Dense baseline
+uv run python moe_emergence/gpt2_inference.py \
+  --checkpoint checkpoints/dense-baseline/final-model \
+  --prompt "The meaning of life is"
+```
+
+The inference script reads the `.json` sidecar to detect mode (dense vs MoE) and architecture config automatically.
+
+To resume training from a checkpoint:
+
+```bash
+uv run python -m moe_emergence.train \
+  --preset moe-main --run-name moe-main \
+  --device cuda \
+  --resume checkpoints/moe-main/ckpt-step-9999.pt
+```
+
+---
+
+## 4. Architecture
+
+The dense baseline is standard GPT-2 small (124M parameters, 12 transformer layers).
+
+The MoE model takes GPT-2 small and replaces layers 8-11 with MoE layers. Each MoE layer has 8 experts -- deep copies of the original GPT-2 MLP, warm-started from pretrained weights -- and a learned router with top-1 routing. Total: 257M parameters.
+
+Routing uses the Straight-Through Estimator. Forward pass routes to one expert with weight 1.0, backward pass flows gradients through the soft probability from the router.
+
+| Component | Detail |
+|---|---|
+| Base model | GPT-2 small (124M) |
+| MoE layers | 8, 9, 10, 11 |
+| Experts per layer | 8 |
+| Routing | Top-1, STE |
+| Expert init | `deepcopy(original_mlp)` + tiny noise |
+| Load balance loss | `0.01 × n_experts × Σ(f_i × P_i)` |
+| Z-loss | `0.001 × mean(logsumexp(logits)²)` |
+
+---
+
+## 5. Training
+
+Both models trained on ~6.6M tokens across three domains, balanced to equal token counts:
+
+| Domain | Source | Size |
+|---|---|---|
+| Code | CodeParrot-clean (Python) | 10 MB |
+| Math | MathQA (allenai) | 10 MB |
+| Prose | C4 English (allenai) | 10 MB |
+
+Training config:
+
+| Parameter | Dense | MoE |
+|---|---|---|
+| Max steps | 5,000 | 10,000 |
+| Batch size | 2 × 4 grad accum = 8 | 2 × 4 grad accum = 8 |
+| Block size | 512 | 512 |
+| Learning rate | 5e-5 | 5e-5 |
+| Warmup | 10% | 10% |
+| Schedule | Cosine | Cosine |
+| Hardware | 1× RTX 4090 24GB | 1× RTX 4090 24GB |
+| Wall time | ~30 min | ~85 min |
+| Throughput | ~25.7k tok/s | ~14.2k tok/s |
+
+---
+
+## 6. W&B
+
+Training curves are on Weights & Biases:
+
+- [Dense baseline](https://wandb.ai/sumit-ml/moe-emergence/runs/fqhfblfv)
+- [MoE main run](https://wandb.ai/sumit-ml/moe-emergence/runs/j08s2d1m)
+
+---
+
+## 7. Links
+
+- Code: [github.com/sumitdotml/moe-emergence](https://github.com/sumitdotml/moe-emergence)
+- Experiment docs: [run-004 (dense)](https://github.com/sumitdotml/moe-emergence/blob/main/docs/experiments/run-004-dense-baseline.md), [run-005 (MoE)](https://github.com/sumitdotml/moe-emergence/blob/main/docs/experiments/run-005-moe-main.md)
+
+---
+
+## License
+
+MIT. See the [source repo](https://github.com/sumitdotml/moe-emergence/blob/main/LICENSE) for details. Third-party dataset licenses are documented in [THIRD-PARTY-NOTICES.md](https://github.com/sumitdotml/moe-emergence/blob/main/THIRD-PARTY-NOTICES.md).

dense-baseline/ckpt-step-4999.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89b689d49b1524a6c4e48a8035ad51ddc2979d2086576995d3feb8a541636c65
+size 1493476483

dense-baseline/final-model.json

@@ -0,0 +1,36 @@
+{
+  "format_version": 1,
+  "step": 4999,
+  "preset": "dense",
+  "mode": "dense",
+  "config": {
+    "preset": "dense",
+    "mode": "dense",
+    "run_name": "dense-baseline",
+    "seed": 42,
+    "max_steps": 5000,
+    "batch_size": 2,
+    "grad_accum_steps": 4,
+    "effective_batch_size": 8,
+    "block_size": 512,
+    "learning_rate": 5e-05,
+    "weight_decay": 0.01,
+    "warmup_fraction": 0.1,
+    "max_grad_norm": 1.0,
+    "lb_coef": 0.0,
+    "z_coef": 0.0,
+    "n_experts": 8,
+    "topk": 1,
+    "noise_std": 0.0,
+    "moe_layers": [],
+    "size_mb": 10.0,
+    "balance_tokens": true,
+    "eval_every": 200,
+    "save_every": 500,
+    "collapse_early_stop": false
+  },
+  "metrics_summary": {
+    "eval_loss": 2.1567,
+    "eval_perplexity": 8.6424
+  }
+}

dense-baseline/final-model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e90989f75165f95fd08c679d6ef3798405470b42ac170f477022b4c2a20c8909
+size 652163888

moe-main/ckpt-step-9999.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df9d01fbef22c9a74cc7d11a05a743f3aa551f80c5201a02ba1ff0d897c46c0d
+size 3080659837

moe-main/final-model.json

@@ -0,0 +1,41 @@
+{
+  "format_version": 1,
+  "step": 9999,
+  "preset": "moe-main",
+  "mode": "moe",
+  "config": {
+    "preset": "moe-main",
+    "mode": "moe",
+    "run_name": "moe-main",
+    "seed": 42,
+    "max_steps": 10000,
+    "batch_size": 2,
+    "grad_accum_steps": 4,
+    "effective_batch_size": 8,
+    "block_size": 512,
+    "learning_rate": 5e-05,
+    "weight_decay": 0.01,
+    "warmup_fraction": 0.1,
+    "max_grad_norm": 1.0,
+    "lb_coef": 0.01,
+    "z_coef": 0.001,
+    "n_experts": 8,
+    "topk": 1,
+    "noise_std": 0.1,
+    "moe_layers": [
+      8,
+      9,
+      10,
+      11
+    ],
+    "size_mb": 10.0,
+    "balance_tokens": true,
+    "eval_every": 200,
+    "save_every": 500,
+    "collapse_early_stop": false
+  },
+  "metrics_summary": {
+    "eval_loss": 2.0798,
+    "eval_perplexity": 7.9147
+  }
+}

moe-main/final-model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d62d8897155f0a5670ad95033b2e5a4726013823c2632c35a45fab2df2d4d2d
+size 1181188800