TokSuite Training Code (Meta Lingua)

This repository contains the training and evaluation code for the TokSuite paper. It is built on top of Meta Lingua — see README_lingua.md for framework setup, SLURM job launching, debugging, and checkpoint management. Meta Lingua is a minimal and fast LLM training and inference library for research.

Set-up

Python environment

Set up the Python environment, we recommend uv but you can use the original conda setting as well:

bash setup/create_env.sh
# or, on a SLURM cluster
sbatch setup/create_env.sh

# then activate the environment created by the script
source .venv/bin/activate

TokSuite setup

In the paper we train multiple models from the same initialization. To do this we compute a unified "super-vocabulary" $\mathcal S$ and construct a shared embedding table of size $|\mathcal S| \times d_model$. For each tokenizer we then produce a corresponding checkpoint so that all non-embedding parameters are identically initialized across models, while embeddings/unembeddings for tokens present in multiple tokenizers are initialized to the same values. You can find the initialization for the fourteen models in the paper at toksuite/initializations.

TokSuite Training

Training settings from the paper

Practical note: the repository includes utilities to generate the per-tokenizer configs and to save/load checkpoints that use the shared super-vocabulary; follow the apps/main/configs/toksuite/create_toksuite_configs.py workflow and point init_ckpt_path to the downloaded initialization to reproduce the shared initialization setup.

Before running the config generator, fill in these variables in apps/main/configs/toksuite/create_toksuite_configs.py:

• dump_dir: where generated configs and experiment outputs should be written
• root_dir: the local path to the downloaded TokSuite pretraining data
• init_ckpt_path: the local path to the downloaded initialization checkpoint

Download the required assets first:

• TokSuite training data: huggingface-cli download toksuite/toksuite_pretraining_data --repo-type dataset --local-dir /path/to/toksuite_pretraining_data
• Initializations: download one of the model folders from toksuite/initializations into /path/to/init_ckpt

Create the TokSuite configs by editing LLAMA_1B_OPTIONS in apps/main/configs/toksuite/create_toksuite_configs.py and then running:

python apps/main/configs/toksuite/create_toksuite_configs.py

The script writes generated YAMLs next to itself in apps/main/configs/toksuite/ and will stop with a clear message if any of dump_dir, root_dir, or init_ckpt_path are still empty.

Reproducing the TokSuite models (training)

The steps below reproduce the TokSuite training runs. Follow them in order; for full-scale runs use SLURM and multiple GPU nodes.

1. Prepare the environment, see above
2. Download data and initial checkpoint

huggingface-cli login
huggingface-cli download toksuite/toksuite_pretraining_data --repo-type dataset --local-dir /path/to/toksuite_pretraining_data
huggingface-cli download <INIT_MODEL_FROM_COLLECTION> --local-dir /path/to/init_ckpt

3. Edit the generator

Open apps/main/configs/toksuite/create_toksuite_configs.py and set:

• dump_dir → where outputs and generated configs will be saved
• root_dir → path to /path/to/toksuite_pretraining_data
• init_ckpt_path → path to /path/to/init_ckpt

4. Generate configs

python apps/main/configs/toksuite/create_toksuite_configs.py

This writes _toksuite_llama_1B.yaml (shared base config) and toksuite_<tokenizer>.yaml files for each tokenizer into apps/main/configs/toksuite/.

5. Representative training recipe (paper)

• Steps: 100000
• Effective batch: bs (4) x grad_accum (8) x n_gpus (8) x seq_len (4096) => 1,048,576 tokens (set data.batch_size and distribute across GPUs accordingly, in the paper we use 8 gpus)
• Optimizer: AdamW, weight decay 0.1
• Peak LR: 1e-3 with cosine decay and 2000 warmup steps
• Model: Llama-3.2-style 1B (untied embeddings)

6. Launch a single run

python -m lingua.stool \
  script=apps.main.train \
  config=[apps/main/configs/toksuite/_base.yaml,apps/main/configs/toksuite/_toksuite_llama_1B.yaml,apps/main/configs/toksuite/toksuite_meta-llama-Llama-3.2-1B.yaml] \
  nodes=1 launcher=sbatch 

7. Launch the full suite (all tokenizers)

for cfg in apps/main/configs/toksuite/toksuite_*.yaml; do
  python -m lingua.stool \
    script=apps.main.train \
    config=[apps/main/configs/toksuite/_base.yaml,apps/main/configs/toksuite/_toksuite_llama_1B.yaml,$cfg] \
    nodes=1 launcher=sbatch &
done
wait

8. Checkpoints & resuming

• Checkpoints are saved under <dump_dir>/checkpoints by default (configured in generated YAMLs).
• Resume by setting checkpoint.path or passing a CLI override pointing to the checkpoint directory.

9. Reproducibility tips

• Use the same init_ckpt_path for all tokenizers to ensure identical initialization.
• Use explicit seeds (see apps/main/configs/toksuite/ablations/ for examples).
• Keep _toksuite_llama_1B.yaml as the shared base for all runs to avoid path-dependent differences.

Training data format

TokSuite training data should be provided as jsonl files and chunked using the same layout expected by Lingua, i.e. max. 1 chunk per device, with each line having the text key.

Each chunk should follow the Lingua data conventions used by the training code, so the data loader can read it without extra conversion steps.

Lingua docs

For the full project documentation, setup instructions, and experiment details, see README_lingua.md.

Citation

If you use TokSuite or Meta Lingua in your work please cite the papers/projects below. BibTeX entries are provided for convenience.

@inproceedings{altintas2026toksuite,
  author       = {G{"u}l Sena Altınta\c{s} and Malikeh Ehghaghi and Brian Lester and Fengyuan Liu and Wanru Zhao and Marco Ciccone and Colin Raffel},
  title        = {{TokSuite}: Measuring the Impact of Tokenizer Choice on Language Model Behavior},
  booktitle    = {Proceedings of the 43rd International Conference on Machine Learning (ICML)},
  year         = {2026},
  eprint       = {2512.20757},
  archivePrefix= {arXiv},
  url          = {https://arxiv.org/abs/2512.20757}
}

@misc{meta_lingua,
  author = {Mathurin Videau and Badr Youbi Idrissi and Daniel Haziza and Luca Wehrstedt and Jade Copet and Olivier Teytaud and David Lopez-Paz},
  title  = {{Meta Lingua}: A minimal {PyTorch LLM} training library},
  year   = {2024},
  url    = {https://github.com/facebookresearch/lingua}
}