autoresearch-rust

Pure Rust port of karpathy/autoresearch — an autonomous pretraining research swarm.

Overview

Trains a GPT transformer on climbmix-400b-shuffle parquet shards with:

• Muon optimizer — 3-term Newton-Schulz matrix orthogonalization
• Grouped-query attention (GQA) with configurable n_head / n_kv_head
• PAR-16 RoPE rotary embeddings
• Parquet dataloader — Arrow 55 + Parquet 55, no Python needed
• Pure Rust BPE tokenizer via rustbpe
• .safetensors checkpointing via tch 0.23
• MFU tracking (Model FLOPs Utilization) against H100 BF16 baseline
• Time-budgeted training — stops automatically after wall-clock budget
• Obsidian sync — trains, logs, and pushes observations to a local vault

Prerequisites

• Rust toolchain (1.80+)
• PyTorch 2.5+ (via conda/mamba)
• LIBTORCH_USE_PYTORCH=1 environment variable

Build

export LD_LIBRARY_PATH=/home/mctouch/anaconda3/lib/python3.13/site-packages/torch/lib
export LIBTORCH_USE_PYTORCH=1
cargo build

Expected output: Finished dev profile [unoptimized + debuginfo] target(s) in Xs

Run tests

cargo run --bin test_synthetic
# All 8 tests pass: model creation, forward/backward, checkpoint I/O, memory estimation

Run training

cargo run --bin train -- --n-layer 6 --n-head 6 --n-kv-head 2 --seq-len 512 --batch-size 32 --num-shards 10

CLI reference

Flag	Default	Description
-n, --num-shards	0 (all)	Training shards to download
-r, --lr	3e-4	Learning rate
-b, --batch-size	32	Batch size
-s, --seq-len	512	Sequence length
-l, --n-layer	6	Transformer layers
-a, --n-head	6	Attention heads
--n-kv-head	2	KV heads (GQA)
--embd-pct	1.0	Embedding size fraction
--max-steps	0 (time-limited)	Max steps (0 = time-budgeted)
--time-budget	300s	Wall-clock training budget
--eval-interval	500	Validation interval
--eval-iters	20	Validation iterations
--log-interval	1	Logging interval
--device	auto	Device ("cuda:0", "cpu")
--checkpoint-dir	./checkpoints	Checkpoint directory
--obsidian-vault	~/.obsidian/vaults/llm	Obsidian vault path
--warmup-iters	100	Warmup iterations
--weight-decay	0.1	Weight decay
--beta2	0.95	AdamW beta2

Architecture

train.rs (entrypoint)
  ├── Dataloader (parquet Arrow reader)
  │     ├── Tokenizer (BPE via rustbpe)
  │     └── Shard fetcher (HTTP → parquet)
  ├── GPT model (tch::nn)
  │     ├── GPT blocks (attention + MLP)
  │     ├── PAR-16 RoPE embeddings
  │     └── LM head
  ├── Muon optimizer (3-term Newton-Schulz)
  ├── EMA tracking
  └── sync_to_obsidian (markdown logger)

Key crates

Crate	Version	Purpose
tch	0.23	PyTorch Rust bindings
parquet	55	Parquet I/O
arrow	55	Columnar format
rustbpe	0.3	BPE tokenizer
clap	4	CLI parsing
serde	1	Serialization
rand	0.8	RNG

Environment variables

Variable	Required	Purpose
LIBTORCH_USE_PYTORCH	Yes	Set to 1 to use conda PyTorch
LD_LIBRARY_PATH	Yes	Path to torch/lib directory
RUST_LOG	No	Log level

Memory estimation

Total params = 6 * (embd_dim * embd_dim * 4) + (embd_dim * 3 * 8192) + (vocab_size * embd_dim)
Memory (fp32) ≈ params * 12 (AdamW state)

Default config (6 layers, 6 heads, embd-pct=1.0): ~555M params, ~6.4 GB fp32 memory.

License

MIT