| --- |
| license: apache-2.0 |
| language: |
| - en |
| tags: |
| - n-gram |
| - olmo-core |
| - language-modeling |
| datasets: |
| - IanMagnusson/count-ngram-tables |
| - allenai/dolma3_mix-6T-1025-7B |
| --- |
| |
| # COUNT checkpoints |
|
|
| Pretrained COUNT models from **The COUNT: Counting N-grams Unlocks Scaling on Host Memory and |
| Training-Free Adaptation** ([paper](https://arxiv.org/abs/2610.20456), [code](https://github.com/IanMagnusson/COUNT)). |
|
|
| COUNT adds an n-gram embedding to a standard transformer without adding learned parameters. For each |
| position it looks up the most frequent next tokens of the longest matching n-grams in a table of corpus |
| counts held in CPU RAM, averages those tokens' rows of the model's own output head (weighted by their |
| count scores), and adds the result to the token embedding. |
|
|
| For the full configuration of each model, see the paper's appendix, the |
| `config.json` and `release_args.json` in each step directory, and the ladder configuration in the |
| [code](https://github.com/IanMagnusson/COUNT). |
|
|
| ## Checkpoints |
|
|
| Seven backbone sizes on the OLMo-core model ladder, each trained on Dolma 3 |
| (`allenai/dolma3_mix-6T-1025-7B`) and saved at three training budgets: 10, 20 and 40 tokens per |
| backbone parameter (D/N). Every checkpoint is a full OLMo-core distributed checkpoint with model |
| weights, optimizer state and trainer state. |
|
|
| | Size | Table used in training | D/N = 10 | D/N = 20 | D/N = 40 | Global batch (tokens) | |
| |---|---|---|---|---|---:| |
| | 60M | `general-17gb` | step 2190 | step 4381 | step 8762 | 262,144 | |
| | 100M | `general-17gb` | step 3881 | step 7762 | step 15524 | 262,144 | |
| | 190M | `general-17gb` | step 3631 | step 7261 | step 14523 | 524,288 | |
| | 370M | `general-797gb` | step 4721 | step 9442 | step 18883 | 786,432 | |
| | 600M | `general-797gb` | step 5574 | step 11148 | step 22297 | 983,040 | |
| | 760M | `general-797gb` | step 5785 | step 11570 | step 23139 | 1,310,720 | |
| | 1B | `general-797gb` | step 8134 | step 16268 | step 32537 | 1,572,864 | |
|
|
| Paths follow the code's ladder layout, `model-ladders/count-<size>/<SIZE>/step<N>/`, so pointing |
| `OLMO_ROOT_DIR` at a download lets the code find them. Each step directory has a `release_args.json` |
| with the ladder settings it was trained with. |
|
|
| The tables are in [count-ngram-tables](https://huggingface.co/datasets/IanMagnusson/count-ngram-tables). |
| A checkpoint can be evaluated with its training table or with any other table there, including the |
| smaller general tables and the math, biomedical and law tables. |
|
|
| **190M note.** The 190M checkpoints come from a rerun with the same recipe, because the original run's |
| D/N = 10 and 20 checkpoints were lost. They are a reproduction of the paper's 190M COUNT model, not the |
| exact weights behind its reported numbers. |
|
|
| ## Usage |
|
|
| Evaluate a checkpoint with the [COUNT code](https://github.com/IanMagnusson/COUNT): |
|
|
| ```bash |
| hf download IanMagnusson/count-checkpoints-count --include "model-ladders/count-190m/*" \ |
| --local-dir $OLMO_ROOT_DIR |
| torchrun --standalone --nproc-per-node=1 olmo-core/src/scripts/train/ladder/count_ladder.py \ |
| eval-checkpoint --name count-190m --cluster local/h200 --size 190M --sequence-length 8192 \ |
| --chinchilla-multiple 2.0 --max-gpus 1 --minimum-world-size 1 --fixed-global-batch-size 524288 \ |
| --rank-mbz 8 --data-root /path/to/tokenized \ |
| --checkpoint-path $OLMO_ROOT_DIR/model-ladders/count-190m/190M/step14523 \ |
| --eval-output-dir /path/to/eval-out --no-load-trainer-state \ |
| --count-table-dir /dev/shm/count-tables/tables/general-17gb \ |
| --count-k 8 --count-max-order 6 --count-num-orders 6 |
| ``` |
|
|
| Evaluation loads model weights only. To resume training on your own copy of the tokenized data, |
| pass `--ignore-data-fingerprints`: the checkpoints record a fingerprint of the data paths they were |
| trained on, and without the flag the code refuses to resume from different paths. |
|
|
| ## Related |
|
|
| - [N-gram tables](https://huggingface.co/datasets/IanMagnusson/count-ngram-tables) |
| - [Learned n-gram embedding checkpoints](https://huggingface.co/IanMagnusson/count-checkpoints-learned-ngram) |
| - [Standard transformer checkpoints](https://huggingface.co/IanMagnusson/count-checkpoints-standard-transformer) |
|
|
| ## Citation |
|
|
| ```bibtex |
| @article{magnusson2026count, |
| title = {The {COUNT}: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation}, |
| author = {Magnusson, Ian and Smith, Noah A. and Koh, Pang Wei}, |
| journal = {arXiv preprint arXiv:2610.20456}, |
| year = {2026} |
| } |
| ``` |
|
|
