IanMagnusson/count-checkpoints-count

COUNT checkpoints

Pretrained COUNT models from The COUNT: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation (paper, code).

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.

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. 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:

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
• Learned n-gram embedding checkpoints
• Standard transformer checkpoints

Citation

@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}
}