olmo-cpp

A C++17 / LibTorch implementation of an OLMo-2 style transformer for training and inference. CUDA is used when available, CPU otherwise.

Build

Requires CMake 3.18+, a C++17 compiler, zlib, and LibTorch (the pip install torch install works). For GPU builds, install a matching CUDA toolkit (>= 12.8 for RTX 50-series / Blackwell, >= 12.1 for Ampere/Ada/Hopper).

./build.sh

The script auto-detects LibTorch and CUDA. Overrides:

LIBTORCH_DIR=/opt/libtorch ./build.sh
CUDA_HOME=/usr/local/cuda-12.8 ./build.sh
CMAKE_CUDA_ARCHITECTURES="120" ./build.sh   # 5060 Ti only

Train

A small tokenized dataset (data/tokens.npy) ships with the repo.

./build/olmo_train --train \
  --config configs/olmo2_125M.json \
  --data-path data/tokens.npy \
  --device auto \
  --steps 5000 \
  --lr 3e-4

Available configs: olmo2_100M.json, olmo2_125M.json, olmo2_125M_tinystories.json, olmo2_1B.json, olmo2_3B.json, olmo2_7B.json.

To tokenize your own data:

./build/prepare_data --download-hf roneneldan/TinyStories \
  --output data/tokens.npy \
  --vocab-file data/gpt2/vocab.json \
  --merges-file data/gpt2/merges.txt

Inference

./build/chat --checkpoint path/to/checkpoint.pt --config configs/olmo2_125M.json

Layout

• include/olmo_cpp/, src/ — library sources
• tools/ — training entry point and CLI tools (chat, prepare_data, ...)
• configs/ — model configs
• third_party/cnpy/ — vendored numpy reader (compiled into the library)