nanochat-rs

nanochat-rs is a Rust CLI for training and evaluating tiny language models from a text corpus through a configurable, reproducible workflow. Users pass corpus paths and tune the full training/sampling surface from the command line: execution mode (train, sample, ablate), engine (scalar pure-Rust or tensor), model kind (bigram or mini-gpt), optimizer (sgd / adamw), learning-rate schedule, seed, step budget, batch/sequence settings, checkpoint cadence, sampler settings, and output locations. A run produces persisted artifacts (checkpoints and outputs), generated text for qualitative inspection, structured metrics and logs for drift/quality tracking, and ablation summaries for controlled comparisons; this enables direct experimentation over speed, quality, and behavior under equivalent settings. The intent is to keep runs deterministic and comparable by design, with CLI-defined controls for each experiment and machine-readable artifacts aligned to benchmark and ablation workflows. Hardware support is explicitly surfaced by mode: CPU runs work through the built-in tensor path, while GPU use depends on the optional tch backend and local PyTorch/CUDA availability. Reproducibility is enforced with explicit CLI parameters and stable artifact naming/output layout, which makes it straightforward to rerun the same experiment, compare checkpoints across commits, and audit benchmark deltas over time.

nanochat-rs is a Rust-first continuation of AntigmaLabs/nanochat-rs, benchmarked against karpathy/nanochat.

Repository Layout

• src/: Rust implementation (CLI, training, sampling, tensor/scalar runtimes, benchmarks, and tests).
• scripts/: Benchmark and utility scripts.
• docs/: Design notes, experimental notes, and reference docs.
• references/: External reference scripts and baselines.
• input.txt: Input corpus file used by example flows.
• Cargo.toml / Cargo.lock: Rust manifest and lockfile.

Goal

Build a Rust repo that can match or beat karpathy/nanochat on:

1. Quality (comparable eval metrics at matched budget)
2. Speed (throughput and time-to-target)
3. Reliability (stability of eval/checkpoint behavior)
4. Reproducibility (pinned refs + machine-readable artifacts)

Working Today

1. CLI flows: train, sample, ablate with full argument parsing.
2. Scalar mode: --model-kind bigram (default) and --model-kind mini-gpt (1-layer GPT).
3. Tensor mode: CPU-native path + optional tch backend with CUDA (--features tch-backend). Supports both bigram and mini-gpt model kinds.
4. Optimizers: SGD (default) and AdamW (--optimizer adamw) for both scalar and tensor paths.
5. Ablation sweep: 4 variants × (tie_lm_head, input_rmsnorm), writes artifacts to results/.
6. Guards: prompt-length guard, memory-drift guard, checkpoint-before-eval persistence guard.
7. Checkpointing: atomic persist-then-eval across tensor train loops with configurable interval.
8. Training utilities: shared LR scheduling (linear warmdown), train/val splitting, transition count building.
9. Benchmark harness against nanochat (scripts/benchmark_karpathy.py) with pinned baseline ref support, timeout-safe runs, and GPU profiling scripts.
10. 62 tests covering loss convergence, round-trip encoding, guard rejection, optimizer comparison, checkpoint ordering, and Python parity traces.

Known Gaps

1. No demonstrated end-to-end benchmark win vs karpathy/nanochat yet.
2. Full training/eval parity with upstream nanochat is not complete.
3. Mini-gpt is minimal parity scaffolding, not a full production training stack.
4. No distributed training.
5. Char-level tokenizer only (no BPE/subword).
6. Ablation sweep is scalar-bigram only.
7. Local CPU baseline runs of upstream nanochat can be slow; meaningful comparisons should use GPU runs.

Current benchmark status (as-of this environment)

• CPU scalar path (default): cargo run --release -- train --mode scalar --model-kind bigram --steps 200 --optimizer adamw --seed 42 took ~0.40s.
• CPU tensor-native path: cargo run --release -- train --mode tensor --model-kind bigram --steps 200 --optimizer sgd --seed 42 took ~0.26s.
• tensor --optimizer adamw is currently unsupported in CPU-native tensor mode in this codebase.
• The tch backend could not run in this environment yet because the local libtorch/PyTorch setup is incomplete (version mismatch and missing runtime libtorch_cpu.dylib lookup path).
• GeForce 4060 support is expected through the tch backend when a compatible libtorch/PyTorch/CUDA runtime is available.

Lineage and Baseline

• Evolving from: AntigmaLabs/nanochat-rs
• Benchmark source of truth: karpathy/nanochat
• Current pinned baseline ref used by harness defaults: 2f096867244e3d00a50284d1be05fa3f5dcfb84b (observed 2026-02-12)

Quick Start

cargo check
cargo test

# scalar bigram (default)
cargo run --release -- train --mode scalar --steps 500
cargo run --release -- sample --mode scalar --temperature 0.8 --max-new-tokens 120

# scalar mini-gpt (opt-in)
cargo run --release -- train --mode scalar --model-kind mini-gpt --style classic --steps 200
cargo run --release -- sample --mode scalar --model-kind mini-gpt --style classic --temperature 0.8 --max-new-tokens 120

# tensor (CPU fallback)
cargo run --release -- train --mode tensor --steps 500

# tensor with tch backend (uses CUDA when available)
LIBTORCH_USE_PYTORCH=1 cargo run --release --features tch-backend -- train --mode tensor --steps 500

Run Matrix

Use these presets as a starting point:

• CPU scalar baseline:
  • cargo run --release -- train --mode scalar --steps 500
• CPU tensor baseline:
  • cargo run --release -- train --mode tensor --steps 500
• GPU tensor baseline (when tch backend is available):
  • LIBTORCH_USE_PYTORCH=1 cargo run --release --features tch-backend -- train --mode tensor --steps 500
• Deterministic sampling from a checkpoint:
  • cargo run --release -- sample --mode tensor --checkpoint-path <path> --seed 42 --temperature 0.8 --max-new-tokens 120
• Ablation sweep:
  • cargo run --release -- ablate --model-kind bigram --base-seed 42 --steps 120 --checkpoint-interval 40

Benchmark Commands

# GPU-first benchmark
bash scripts/colab_gpu_benchmark.sh

# explicit profile
PROFILE=full BASELINE=nanochat bash scripts/colab_gpu_benchmark.sh

# local CPU smoke benchmark (expected to be limited)
python3 scripts/benchmark_karpathy.py \
  --baseline nanochat \
  --no-require-gpu \
  --install-deps \
  --ours-cargo-features "" \
  --nanochat-device-type cpu \
  --nanochat-disable-compile \
  --nanochat-target-flops -1 \
  --nanochat-target-param-data-ratio -1 \
  --nanochat-num-iterations 1 \
  --nanochat-eval-every -1 \
  --nanochat-max-chars 50000 \
  --nanochat-depth 1 \
  --nanochat-head-dim 16 \
  --nanochat-max-seq-len 64 \
  --nanochat-device-batch-size 1 \
  --nanochat-total-batch-size 8 \
  --nanochat-train-timeout-sec 120

Artifacts are written to results/.

Sources

• https://github.com/AntigmaLabs/nanochat-rs
• https://github.com/karpathy/nanochat

License

This project is licensed under Apache-2.0.

Originally developed by Ricardo Frantz. See LICENSE and NOTICE for license terms and attribution notices.

Disclaimer

This software is provided "as is", without warranty of any kind. To the extent permitted by law, the authors and contributors are not liable for any damage, loss or claim arising from its use or misuse. You are responsible for how you use it and for following the laws and rules that apply to you. The full terms are in LICENSE.

Models trained with this code can produce wrong or offensive text. Review generated text before you use or publish it.