Add files using upload-large-folder tool
Browse files- .gitattributes +8 -0
- DRMoET-1.7B/iter_0032000/.metadata +3 -0
- DRMoET-1.7B/iter_0032000/__0_0.distcp +3 -0
- DRMoET-1.7B/iter_0032000/common.pt +3 -0
- DRMoET-1.7B/iter_0032000/metadata.json +1 -0
- DRMoET-1.7B/latest_checkpointed_iteration.txt +1 -0
- DRMoET-290M-32E/iter_0016000/.metadata +3 -0
- DRMoET-290M-32E/iter_0016000/__0_0.distcp +3 -0
- DRMoET-290M-32E/iter_0016000/common.pt +3 -0
- DRMoET-290M-32E/iter_0016000/metadata.json +1 -0
- DRMoET-290M-32E/latest_checkpointed_iteration.txt +1 -0
- FLAME-MoE-1.7B/iter_0032000/.metadata +3 -0
- FLAME-MoE-1.7B/iter_0032000/__0_0.distcp +3 -0
- FLAME-MoE-1.7B/iter_0032000/common.pt +3 -0
- FLAME-MoE-1.7B/iter_0032000/metadata.json +1 -0
- FLAME-MoE-1.7B/latest_checkpointed_iteration.txt +1 -0
- FLAME-MoE-290M-32E/iter_0016000/.metadata +3 -0
- FLAME-MoE-290M-32E/iter_0016000/__0_0.distcp +3 -0
- FLAME-MoE-290M-32E/iter_0016000/common.pt +3 -0
- FLAME-MoE-290M-32E/iter_0016000/metadata.json +1 -0
- FLAME-MoE-290M-32E/latest_checkpointed_iteration.txt +1 -0
- README.md +68 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,11 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
DRMoET-1.7B/iter_0032000/.metadata filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
DRMoET-290M-32E/iter_0016000/.metadata filter=lfs diff=lfs merge=lfs -text
|
| 38 |
+
FLAME-MoE-1.7B/iter_0032000/.metadata filter=lfs diff=lfs merge=lfs -text
|
| 39 |
+
FLAME-MoE-290M-32E/iter_0016000/.metadata filter=lfs diff=lfs merge=lfs -text
|
| 40 |
+
FLAME-MoE-290M-32E/iter_0016000/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
|
| 41 |
+
DRMoET-290M-32E/iter_0016000/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
|
| 42 |
+
FLAME-MoE-1.7B/iter_0032000/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
|
| 43 |
+
DRMoET-1.7B/iter_0032000/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
|
DRMoET-1.7B/iter_0032000/.metadata
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:d4ebd59326a09a000300b47177bf4ecfe64615053acd3f0f7c56ee98c2e8a7a2
|
| 3 |
+
size 394290
|
DRMoET-1.7B/iter_0032000/__0_0.distcp
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:3afc3926ff6289ceacf912bf42875329fb37986d7ab0fb177f7b535085119734
|
| 3 |
+
size 20571786340
|
DRMoET-1.7B/iter_0032000/common.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:9fdd1d523ba75d8e4bb4a1154c2d737d64dc837225a01c5166b4ae57841d0fb2
|
| 3 |
+
size 328348
|
DRMoET-1.7B/iter_0032000/metadata.json
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{"sharded_backend": "torch_dist", "sharded_backend_version": 1, "common_backend": "torch", "common_backend_version": 1}
|
DRMoET-1.7B/latest_checkpointed_iteration.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
32000
|
DRMoET-290M-32E/iter_0016000/.metadata
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e4f12a6fddb7fb780b8df50c5978ba6b2a9b1bc6a923bd6037bbe93dd22806bf
|
| 3 |
+
size 108668
|
DRMoET-290M-32E/iter_0016000/__0_0.distcp
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:0b0611d3e66a9f432593a9fcaa4f1d00c4954521ad5a8de611e6372fbec0d186
|
| 3 |
+
size 1493356136
|
DRMoET-290M-32E/iter_0016000/common.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:01d007cda840f2d16c9dbf3e50a00cb600b025039e50e04dba85c7ab95707be2
|
| 3 |
+
size 54876
|
DRMoET-290M-32E/iter_0016000/metadata.json
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{"sharded_backend": "torch_dist", "sharded_backend_version": 1, "common_backend": "torch", "common_backend_version": 1}
|
DRMoET-290M-32E/latest_checkpointed_iteration.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
16000
|
FLAME-MoE-1.7B/iter_0032000/.metadata
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:4a7236eea57ec40d314285273f9346345296210750aea2b5adb799865972d340
|
| 3 |
+
size 394293
|
FLAME-MoE-1.7B/iter_0032000/__0_0.distcp
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:37fc21be2a5d1acf4078b186bc545f7d237aedc55b176e6bc08d31e802c50aad
|
| 3 |
+
size 20571786340
|
FLAME-MoE-1.7B/iter_0032000/common.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:bf25307ed9dc0dd4b61abfb382ce248bfc5bc6c56c0de39aebee846ec17ee195
|
| 3 |
+
size 327900
|
FLAME-MoE-1.7B/iter_0032000/metadata.json
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{"sharded_backend": "torch_dist", "sharded_backend_version": 1, "common_backend": "torch", "common_backend_version": 1}
|
FLAME-MoE-1.7B/latest_checkpointed_iteration.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
32000
|
FLAME-MoE-290M-32E/iter_0016000/.metadata
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:2b2f309990c2eebb67ada31adc51f971597075847488d56e575674047acff2df
|
| 3 |
+
size 108671
|
FLAME-MoE-290M-32E/iter_0016000/__0_0.distcp
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b34648069e08f9efda16d3e07108e835aacf447279f779769a96b57e0fdff197
|
| 3 |
+
size 1493356136
|
FLAME-MoE-290M-32E/iter_0016000/common.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a32cd49ec181fc2e1aea1d36db035bebce7d72e9f664dd60c474c132f361fe64
|
| 3 |
+
size 54428
|
FLAME-MoE-290M-32E/iter_0016000/metadata.json
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{"sharded_backend": "torch_dist", "sharded_backend_version": 1, "common_backend": "torch", "common_backend_version": 1}
|
FLAME-MoE-290M-32E/latest_checkpointed_iteration.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
16000
|
README.md
CHANGED
|
@@ -1,3 +1,71 @@
|
|
| 1 |
---
|
| 2 |
license: mit
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
---
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
license: mit
|
| 3 |
+
library_name: megatron-lm
|
| 4 |
+
tags:
|
| 5 |
+
- mixture-of-experts
|
| 6 |
+
- moe
|
| 7 |
+
- distributionally-robust-optimization
|
| 8 |
+
- megatron
|
| 9 |
+
datasets:
|
| 10 |
+
- mlfoundations/dclm-baseline-1.0
|
| 11 |
---
|
| 12 |
+
|
| 13 |
+
# DRMoET checkpoints
|
| 14 |
+
|
| 15 |
+
<!-- TODO: paper title, arXiv link, authors -->
|
| 16 |
+
|
| 17 |
+
Pretrained Mixture-of-Experts checkpoints for **DRMoET**, which applies distributionally robust optimization over experts during MoE pretraining, together with matched [FLAME-MoE](https://arxiv.org/abs/2505.20225) baselines trained with the same data, seed, and schedule.
|
| 18 |
+
|
| 19 |
+
Code: <https://github.com/MAPS-research/DRMoET>
|
| 20 |
+
|
| 21 |
+
## Checkpoints
|
| 22 |
+
|
| 23 |
+
| Folder | Method | Experts (top-k) | Layers / hidden | Train iters | DRO (β, η, α, credit) |
|
| 24 |
+
|---|---|---|---|---|---|
|
| 25 |
+
| `DRMoET-1.7B` | DRMoET, activation credit | 64 (top-6) + shared | 18 / 2048 | 32,000 | 0.999, 0.1, 1.0, L2 activation norm |
|
| 26 |
+
| `FLAME-MoE-1.7B` | FLAME-MoE baseline | 64 (top-6) + shared | 18 / 2048 | 32,000 | — |
|
| 27 |
+
| `DRMoET-290M-32E` | DRMoET, activation credit | 32 (top-6) + shared | 9 / 1024 | 16,000 | 0.999, 0.001, 1.0, L2 activation norm |
|
| 28 |
+
| `FLAME-MoE-290M-32E` | FLAME-MoE baseline | 32 (top-6) + shared | 9 / 1024 | 16,000 | — |
|
| 29 |
+
|
| 30 |
+
All models use sequence length 2048, global batch size 1024, learning rate 3e-4, auxiliary load-balancing loss coefficient 0.01, and the `EleutherAI/pythia-12b` tokenizer, trained on DCLM. Each DRMoET checkpoint and its baseline share the same seed (1,234 for 1.7B and 3,407 for 290M-32E).
|
| 31 |
+
|
| 32 |
+
## Results (1.7B, 0-shot)
|
| 33 |
+
|
| 34 |
+
| Model | ARC-C | ARC-E | HellaSwag | PIQA | WinoGrande | SciQ | ReCoRD | Avg |
|
| 35 |
+
|---|---|---|---|---|---|---|---|---|
|
| 36 |
+
| FLAME-MoE-1.7B | 34.81 | 69.70 | 49.42 | 75.63 | 59.04 | **90.50** | 84.67 | 66.25 |
|
| 37 |
+
| **DRMoET-1.7B** | **38.31** | **71.84** | **49.90** | **76.66** | **62.51** | 90.40 | **84.79** | **67.77** |
|
| 38 |
+
|
| 39 |
+
Accuracy (`acc`) from lm-evaluation-harness, final checkpoint (iteration 32,000).
|
| 40 |
+
|
| 41 |
+
## Results (290M, 32 experts, 0-shot)
|
| 42 |
+
|
| 43 |
+
| Model | ARC-C | ARC-E | HellaSwag | PIQA | WinoGrande | SciQ | ReCoRD | Avg |
|
| 44 |
+
|---|---|---|---|---|---|---|---|---|
|
| 45 |
+
| FLAME-MoE-290M-32E | 23.72 | 55.72 | 34.73 | 69.26 | 51.14 | 81.70 | 71.47 | 55.39 |
|
| 46 |
+
| DRMoET-290M-32E | 23.46 | 57.03 | 34.90 | 68.50 | 50.99 | 82.10 | 71.45 | 55.49 |
|
| 47 |
+
|
| 48 |
+
Accuracy (`acc`) from lm-evaluation-harness, final checkpoint (iteration 16,000), evaluated from the released (optimizer-stripped) checkpoints.
|
| 49 |
+
|
| 50 |
+
## Format
|
| 51 |
+
|
| 52 |
+
Each folder is a Megatron-LM distributed (`torch_dist`) checkpoint containing **model weights only**; optimizer state has been removed. Load with `--no-load-optim` (and `--no-load-rng`), as for fine-tuning or evaluation.
|
| 53 |
+
|
| 54 |
+
```
|
| 55 |
+
DRMoET-1.7B/
|
| 56 |
+
├── latest_checkpointed_iteration.txt
|
| 57 |
+
└── iter_0032000/
|
| 58 |
+
```
|
| 59 |
+
|
| 60 |
+
## Usage
|
| 61 |
+
|
| 62 |
+
```bash
|
| 63 |
+
git clone --recursive https://github.com/MAPS-research/DRMoET.git && cd DRMoET
|
| 64 |
+
huggingface-cli download TonyTeng/DRMoET --include "DRMoET-1.7B/*" --local-dir weights
|
| 65 |
+
```
|
| 66 |
+
|
| 67 |
+
Then pass `--load weights/DRMoET-1.7B --no-load-optim --no-load-rng` with the matching model arguments from `configs/model/flame-moe.sh` and the corresponding recipe in `scripts/release/` (e.g. `flame-moe-1.7b-dro-activation.sh`). See the repository README for evaluation with lm-evaluation-harness.
|
| 68 |
+
|
| 69 |
+
## Citation
|
| 70 |
+
|
| 71 |
+
<!-- TODO: BibTeX -->
|