IanMagnusson/count-checkpoints-learned-ngram

Learned n-gram embedding checkpoints

The learned n-gram embedding baseline from The COUNT: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation (paper, code).

These models add learned n-gram embeddings to a standard transformer, following Over-Encoding as implemented in LongCat: for orders 2 and 3, four hashed embedding tables per order are looked up for each position, projected to model width, and combined with the token embedding. Following LongCat's suggested 1:1 ratio of embedding to backbone parameters, each model's embedding parameters (token embeddings plus n-gram tables) equal all its other learned parameters, including the output head and projections. 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	D/N = 10	D/N = 20	D/N = 40	Target rows per hashed table (m)	Global batch (tokens)
60M	step 2190	step 4381	step 8762	49,563	262,144
100M	step 3881	step 7762	step 15524	98,859	262,144
190M	step 3631	step 7261	step 14523	148,267	524,288
370M	step 4721	step 9442	step 18883	263,227	786,432
600M	step 5574	step 11148	step 22297	329,019	983,040
760M	step 5785	step 11570	step 23139	394,811	1,310,720
1B	step 8134	step 16268	step 32537	526,395	1,572,864

Paths follow the code's ladder layout, model-ladders/learned-<size>/<SIZE>/step<N>/, and each step directory has a release_args.json with its ladder settings.

Usage

Evaluate a checkpoint with the COUNT code, passing the size's target rows per table m from the table above:

hf download IanMagnusson/count-checkpoints-learned-ngram --include "model-ladders/learned-190m/*" \
  --local-dir $OLMO_ROOT_DIR
torchrun --standalone --nproc-per-node=1 olmo-core/src/scripts/train/ladder/count_ladder.py \
  eval-checkpoint --name learned-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/learned-190m/190M/step14523 \
  --eval-output-dir /path/to/eval-out --no-load-trainer-state \
  --learned-ngram-rows-per-table 148267

Evaluation loads model weights only.

Loading: the hashed tables have unequal sizes

All eight hashed tables use the same hash function and differ only in the modulus, their row count. To make the tables hash differently, table r gets m + 2r + 1 rows (m+1, m+3, …, m+15) around the target m in the table above. Pass that m to the code's ladder as --learned-ngram-rows-per-table, and it sizes the tables the same way. Each step directory's release_args.json also lists the exact sizes.

Optimizer state

These models were trained with sparse row-wise AdamW updates to the hashed tables, and the checkpoints keep that optimizer state exactly: per-row step counts and moments for the tables, plus the sharding buffer learned_ngram.tables.shard_meta. The released code trains these tables with a dense optimizer, which cannot load the per-row state, so it cannot resume training from these checkpoints; asking it to load their optimizer state stops with an error that says so. Evaluation loads model weights only and is unaffected; see Usage above.

Related

• COUNT checkpoints
• Standard transformer checkpoints
• N-gram tables

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