oxidizr

Layers mix within one model. mamba_layers lists the SSM layer indices, and every other layer uses attention.

Mamba2 and Mamba3 require hidden_size * mamba2_expand == mamba2_num_heads * mamba2_head_dim. For example, 384 * 2 == 48 * 16.

Installation

oxidizr needs a Rust toolchain. It fetches boostr from crates.io. Clone and build:

git clone https://github.com/ml-rust/oxidizr
cd oxidizr
cargo build --release

Add --features cuda to the build for GPU support. The CUDA build uses the toolkit installed on the system.

Quick Start

Build the project, download the educational dataset, and start training:

git clone https://github.com/ml-rust/oxidizr
cd oxidizr
cargo build --release

pip install huggingface_hub   # only to download this dataset
hf download fs90/nano-start-data-bin --local-dir data/nano-start/tokenized --repo-type dataset

cargo run --release -- train -f models/nano-start/nano-start.yaml -d data/nano-start/tokenized/combined.bin

The default build trains on CPU. Checkpoints for this config go to ./checkpoints/nano-start.

To train on a GPU, build and run with --features cuda:

cargo run --release --features cuda -- train -f models/nano-start/nano-start.yaml -d data/nano-start/tokenized/combined.bin

After training, package the model. See Model Distribution.

Bring Your Own Config and Data

Write a Config

A config has two top-level keys: model and trainer. This dense GQA example is a complete starting point:

# my-model.yaml
model:
  name: my-model
  hidden_size: 512
  num_layers: 8
  num_heads: 8
  kv_heads: 4              # Fewer KV heads than query heads (GQA)
  intermediate_size: 2048
  vocab_size: 100315
  tokenizer_vocab: cl100k_base   # Must match the tokenizer that produced the data
  max_seq_len: 512
  rope_theta: 10000.0

trainer:
  target_device: gpu
  learning_rate: 0.0003
  batch_size: 2
  gradient_accumulation: 1
  max_steps: 5000
  num_epochs: 2
  checkpoint_dir: "./checkpoints/my-model"
  log_interval: 10
  save_interval: 500
  max_checkpoints: 3
  load_balance_alpha: 0.0  # MoE auxiliary loss weight. Required. 0.0 disables it.

Run it:

cargo run --release -- train -f my-model.yaml -d data/my-dataset.bin

tokenizer_vocab is required. Set it to cl100k_base, o200k_base, llama3, deepseek_v3, or custom.

Key trainer fields:

The full key reference is in docs/configuration.md.

Behavior changes that alter what an existing config trains:

• MLA (kv_latent_dim) applies rope_scaling. It used to ignore it. A checkpoint trained before this change with kv_latent_dim and rope_scaling resumes with different rotary frequencies.
• MLA refuses an explicit head_dim other than hidden_size / num_heads.
• A fresh run with no weight_decay_rule uses skip_1d. A resume keeps the checkpoint's rule.

See Behavior Changes.

Sample Configs

Start from a file in models/. models/README.md describes each one. models/nano-start/nano-start.yaml is the educational hybrid model (Mamba2 + MLA + MoE).

Prepare the Data

oxidizr memory-maps a binary file of little-endian u32 tokens. The tokenizer must match model.tokenizer_vocab.

oxidizr tokenize writes that file from .txt and .jsonl inputs:

cargo run --release --features vocab-llama3 -- tokenize \
  -i data/raw -o data/my-dataset.bin --tokenizer llama3

A vocabulary name needs its vocab-* feature. A tokenizer.json path needs none. See Tokenization.

Pass the file with -d. Without -d, train generates 1,000,000 random tokens for a smoke test.

The oxidizr::data::create_dummy_data(vocab_size, total_tokens) function returns the same random tokens as a Vec<u32>.

CLI Reference

oxidizr <SUBCOMMAND> [OPTIONS]

Common train flags:

Every run writes run.json to its run directory when it ends, in both output modes. It records status (completed or failed), the error text of a failed run, the final step, the last loss, and start and end times. train exits with code 1 when the run failed. steps.json holds one JSON line per log interval and one for the first step of each session.

Without --headless, train shows a progress bar. Run oxidizr <SUBCOMMAND> --help for every flag. docs/cli.md documents each subcommand.

Model Distribution

Package a checkpoint, then upload it:

oxidizr pack --checkpoint latest      # Or: final, or a step number. Omit to pick interactively
cargo build --release --features huggingface
pip install huggingface_hub
oxidizr push --model hf/<username>/<model-name>

pack writes model.safetensors, config.json, and a generated README.md to hf/<username>/<model>/. It also copies training_config.json when the checkpoint has a run.json.

push runs the hf CLI, so hf must be on PATH. It accepts --private, -t, --token, and -y, --yes.

Put these values in a .env file, copied from .env.example:

HF_USERNAME=your-username
HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxx

Get a token from the HuggingFace token page. See hf/README.md and docs/cli.md for the full flag list.

Multi-GPU Training

A build with --features cuda,nccl trains one model on several GPUs of one machine. This is data parallelism: every GPU holds a full replica.

cargo build --release --features cuda,nccl
./target/release/oxidizr train -f models/nano.yaml --gpus 0,1

• --gpus 0,1 starts one rank per GPU, in one process. Rank i uses the i-th id. An id named twice is an error.
• batch_size is the batch of one GPU. One optimizer step reads batch_size x GPUs x gradient_accumulation rows. Two GPUs at batch_size: 4 match one GPU at batch_size: 8 up to float rounding.
• Each step sums the gradients of every GPU with one NCCL all-reduce. Every replica stays bit-equal.
• Rank 0 alone prints, evaluates, and writes run.json, steps.json, evals.json and the checkpoints. Its early-stop decision reaches every rank.
• A failure or panic on one GPU aborts the others. The run fails with the error of the rank that failed first.
• The checkpoint meta records world_size. A resume needs the same GPU count. A checkpoint of another count is an error that names both counts.
• --target-device cpu with more than one id is an error. A build without nccl refuses more than one id with the feature to add.

Not supported: FSDP, ZeRO sharding, tensor parallelism, pipeline parallelism, and more than one machine.

The CPU test suite runs every data-parallel path on in-process CPU ranks. The NCCL communicator, the NCCL abort, and the multi-GPU device setup are untested on real GPUs.

Tracking a Run

Each run keeps its own history in its run directory:

A build with --features tensorboard also writes TensorBoard event files of the training and eval scalars. --headless prints JSON metrics to stdout. --resume continues a run from its weights, optimizer state and data position.

Features

Bundled vocabulary names: cl100k, o200k, gpt-oss, phi4, olmo2, llama3, llama2, deepseek, qwen, glm, gemma2, gemma3, gemma4, kimi, mistral, modernbert.

Example: cargo build --release --features cuda,huggingface.

Development

See CONTRIBUTING.md for build, test, and lint commands, and the project layout.

Documentation: docs/README.md indexes the guides. docs/configuration.md is the config reference. docs/cli.md is the CLI reference. docs/architecture/overview.md describes the design.

Citation

@software{oxidizr,
  author = {Farhan Syah},
  title = {oxidizr: A Rust-based LLM training framework},
  year = {2025},
  url = {https://github.com/ml-rust/oxidizr}
}

Contributing

Contributions are welcome. See CONTRIBUTING.md for conventions and checks.

Questions and design discussion happen on Discord.

License

Licensed under Apache-2.0. See LICENSE.