randyGPT - Mini GPT in Rust

A GPT-style language model implemented from scratch in Rust. Training runs on Metal GPU via Candle autograd on Apple Silicon, with a CPU fallback (Rayon + Accelerate BLAS) on other systems.

Features

• Transformer Architecture: Multi-head causal attention, feed-forward layers, residual connections, RMSNorm
• Metal GPU Training: Candle autograd on Apple M-series — forward + backward + AdamW fully on GPU
• CPU / Linux Fallback: Rayon multi-core + Accelerate BLAS when Metal unavailable; builds natively on Linux
• Full Training Loop: AdamW, gradient clipping, dropout, ReduceLROnPlateau, cosine decay
• GPT-2 Style Init: Scaled output projections for stable deep training
• Checkpoint Save/Resume: RGPT0003 format (weights + optimizer moments); resumes LR schedule exactly
• Character-Level + BPE Tokenization: char vocab or learned BPE with --bpe [N]
• True KV Cache: Single-token decode — prefill prompt once, O(1) projections per new token
• Top-P + Temperature Sampling: Nucleus sampling for text generation
• HTTP Inference Server: --serve flag; no output token cap; returns completion only (prompt stripped)

Model Presets (select at build time)

Preset	Dim	Heads	Layers	Params	Build flag
model-xs	116	4	3	~746K	--features model-xs
model-s	128	4	8	~1.6M	--features model-s
model-ds (deep-s)	128	4	12	~2.78M	--features model-ds
model-m	192	6	6	~2.7M	--features model-m
model-l (default)	256	8	6	~4.82M	(none)
model-deep	192	6	16	~7.5M	--features model-deep
model-xl	384	8	8	~10.8M	--features model-xl

Default hyperparameters (v0.9.5):

Hyperparameter	Value
Block size	256 tokens
Batch size	64
Learning rate	1e-4 → 1e-5 (cosine)
Weight decay	0.1
Dropout	0.1
Early stopping	ReduceLROnPlateau (patience=30, ×3 reductions)

Usage

Build

cargo build --release

Train from scratch

# model-l (default), char-level
./target/release/randygpt --iters 5000

# model-s with BPE-500 vocab (recommended for novels dataset)
cargo build --release --features model-s
./target/release/randygpt --iters 25000 --bpe 500

Resume training from a checkpoint

# Resume from last checkpoint (auto path)
./target/release/randygpt --iters 2000 --resume

# Resume from the best-loss checkpoint
./target/release/randygpt --iters 2000 --resume checkpoint_best.bin

# Resume from a specific file
./target/release/randygpt --iters 5000 --resume my_run.bin

Run inference server

# Serve on default address (0.0.0.0:8080) — auto-loads checkpoint_best.bin
./target/release/randygpt --bpe --serve

# Custom address
./target/release/randygpt --bpe --serve 127.0.0.1:9000

# With bearer-token authentication
./target/release/randygpt --bpe --serve --api-key mysecret

The server accepts POST / with a JSON body and returns generated text:

curl -s http://localhost:8080/ \
  -H 'Content-Type: application/json' \
  -d '{"prompt":"Once upon a time","max_tokens":200,"temperature":0.8}'
# → {"text":"…","model":"randygpt-6L-8H-256D","usage":{"prompt_tokens":4,"completion_tokens":196}}

Field	Default	Description
prompt	required	Text prompt
max_tokens	256 (BLOCK_SIZE)	Max new tokens; capped at context window
temperature	0.7	Sampling temperature

If --api-key is set, every request must include Authorization: Bearer <key>.

Checkpoint files

Weights are serialized to memory every EVAL_INTERVAL iterations and flushed to disk when training finishes or you press Ctrl-C — zero disk I/O in the hot loop.

File	Contents
checkpoint.bin	Latest periodic checkpoint
checkpoint_best.bin	Best-loss checkpoint

Both files use the RGPT0003 format and include model weights + full AdamW optimizer state, so the LR schedule resumes correctly from exactly where you left off.

checkpoint_best.bin is written to disk immediately whenever validation loss improves.

Checkpoint size: ~55 MB for the current 4.77M-param config.

Training data

Place your training text in train.txt:

# Shakespeare (classic benchmark)
curl -o train.txt https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt

# Or any plain text
cp my_text.txt train.txt

If train.txt is absent, the model falls back to a tiny built-in sample.

Architecture Details

Forward Pass

1. Token + Position Embedding — learned lookup tables [vocab, D] and [T, D]
2. Transformer Layers (×6):
   • RMSNorm → Multi-head causal self-attention → Residual
   • RMSNorm → MLP (squared ReLU, 4× expansion) → Residual
   • Dropout after attention projection and MLP output (training only)
3. LM Head — linear projection to vocab logits

Training (Metal GPU path)

• Forward + Backward: Candle autograd on Metal — fully batched [B, T, D] tensor ops
• Loss: candle_nn::loss::cross_entropy over all sequence positions
• Optimizer: Full GPU AdamW via GpuAdamState — moments live as Metal Vars, zero CPU transfers in hot loop
• Gradient clipping: L∞ norm clipped at 1.0
• LR schedule: constant at max_lr → cosine decay to min_lr over final 40%; ReduceLROnPlateau halves max_lr on stagnation (up to 3×)

Training (CPU fallback)

• Rayon: Each batch item processed on a separate thread; gradients summed after
• BLAS: Accelerate cblas_sgemv for d_x, cblas_sgemm for batched d_W (one call per weight matrix instead of T rank-1 updates)

Performance (Apple M-series, Shakespeare 1.1M tokens)

Version	ms/iter	1000 iter	Val ppl @1000	Notes
v0.7.1 CPU	~964ms	~96s	—	12 cores, batch=128, SGEMM backward
v0.8.0 GPU	~488ms	~49s	10.8	Metal autograd, batch=128, 60.9% GPU

RSS memory: ~400 MB real; Activity Monitor shows ~3 GB (Metal unified memory mapped into virtual space, not CPU-resident).

Parameter Scaling Reference

Config	Layers	Dim	Heads	Params
Tiny	2	64	4	~80K
Small	6	128	8	~1.2M
Current	6	256	8	~4.77M
Large	12	512	16	~40M

Completed Improvements

• Backward pass for training (v0.2)
• Adam optimizer (v0.2)
• Multi-core training via Rayon (v0.3)
• Gradient clipping, LR scheduling, dropout (v0.4)
• GPT-2 style initialization (v0.4)
• Model checkpointing (v0.5)
• Metal GPU inference via Candle (v0.5)
• Full attention gradient computation (v0.6)
• Validation split and perplexity metrics (v0.6)
• Accelerate BLAS for CPU matmuls (v0.7.0)
• SGEMM batched backward, 256-dim, batch=128 (v0.7.1)
• Metal GPU training via Candle autograd (v0.8.0)
• BPE tokenization with --bpe [N] flag (v0.9.2)
• Fast CPU-only --generate mode (v0.9.3)
• HTTP inference server via --serve (v0.9.4)
• True KV-cache single-token decode (v0.9.5)
• Generation returns completion only — prompt stripped (v0.9.5)
• ReduceLROnPlateau — halves LR on stagnation, resumes; no premature early stop (v0.9.5)
• LR warmup removed — redundant with gradient clipping + Adam bias correction (v0.9.5)
• Linux / cross-platform builds — Metal conditional on macOS (v0.9.5)
• Systemd deployment to oss.xenon.fun.local (v0.9.5)
• Gutenberg corpus cleaning + deduplication (scripts/clean_gutenberg.py) (v0.9.6)
• Document-boundary-aware batch sampling — no batches crossing <|eos|> separators (v0.9.6)
• Best checkpoint flushed to disk immediately on val loss improvement (v0.9.6)
• HuggingFace export — safetensors + PyTorch modeling + tokenizer scripts (v0.9.6)
• HuggingFace Space — Gradio UI loading model from Hub (v0.9.6)
• model-ds preset — 128-dim, 4-head, 12-layer (~2.78M params, deep-s) (v0.9.6)
• Per-model HF Spaces via MODEL_REPO env var (v0.9.6)
• PyTorch training on NVIDIA GPU — Colab T4/V100/A100, RGPT0003-compatible checkpoints (v0.9.7)
• Gutenberg corpus v2 cleaning — Ulysses excluded, pipes/emails/repeated chars stripped (v0.9.7)
• 1500-token BPE vocab rebuilt from cleaned v2 corpus (v0.9.7)

HuggingFace Export & Deployment

Export a trained checkpoint to HuggingFace format and publish to the Hub in one step:

./scripts/publish_hf.sh                        # model-s, checkpoint_best.bin (defaults)
./scripts/publish_hf.sh s checkpoint_best.bin  # explicit

This script:

1. Runs scripts/export_hf.py → produces hf_export/ with model.safetensors, config.json, tokenizer.json, etc.
2. Uploads model weights to MonumentalSystems/randygpt-{size} on the Hub
3. Uploads the Gradio Space to MonumentalSystems/randygpt-space

One-time Space setup

hf repo create MonumentalSystems/randygpt-space --repo-type space --space-sdk gradio

Export scripts

Script	Purpose
scripts/export_hf.py	Export RGPT0003 checkpoint → safetensors + HF config files
scripts/modeling_randygpt.py	PyTorch PreTrainedModel matching the Rust forward pass exactly
scripts/tokenizer_randygpt.py	Python BPE tokenizer matching the Rust encoder exactly
scripts/train_torch.py	PyTorch training loop for NVIDIA GPU (Colab T4/V100/A100)
scripts/write_rgpt_checkpoint.py	Write RGPT0003 binary from PyTorch model + optimizer state
scripts/publish_hf.sh	One-shot export + upload to Hub + upload Space
spaces/app.py	Gradio Space — loads model from Hub, HF hosts the compute

Colab GPU Training

Train on Google Colab NVIDIA GPU (~6-10x faster than Metal for larger models). Produces RGPT0003-compatible checkpoints loadable by the Rust inference server.

See colab/RUNBOOK.md for the full step-by-step guide including rclone setup, pre-generating the token cache locally, and resuming after session timeout.

Quick start

# 1. Generate token cache locally (do once, upload to Drive — saves ~15 min per Colab session)
python3 -c "
import sys, numpy as np, time
sys.path.insert(0, 'scripts')
from tokenizer_randygpt import RandyGPTTokenizer
tok = RandyGPTTokenizer.from_file('vocab.json')
ids = tok.encode(open('train.txt', encoding='utf-8').read())
np.array(ids, dtype=np.uint32).tofile('train.txt.tokens.bin')
"

# 2. Upload to Drive
rclone copy train.txt.tokens.bin gdrive:randyGPT/ --progress
rclone sync scripts/ gdrive:randyGPT/scripts/ --include "*.py" --exclude "__pycache__/**" --progress

# 3. After training, pull checkpoint
rclone copy gdrive:randyGPT/checkpoint_best.bin . --progress

Open colab/randygpt_train.ipynb in Colab, set MODEL_SIZE, ITERS, DTYPE, run cells 1-5.

Recommended settings by GPU

GPU	dtype	Model	Batch	ms/iter
T4	fp16	s	256	~500ms
V100	fp16	ds	256	~380ms
A100	bf16	l	512	~220ms

Manual export

python3 scripts/export_hf.py \
    --checkpoint checkpoint_best.bin \
    --vocab vocab.json \
    --output hf_export \
    --model-size s \
    --repo MonumentalSystems/randygpt-s

hf upload MonumentalSystems/randygpt-s ./hf_export . --repo-type model

Planned

• Gnostr model distribution — publish/fetch checkpoints as nostr events via gnostr CLI

Credits

Based on mini_gpt.rs by @RandyMcMillan Enhanced by Claude Sonnet 4.5 / Monumental Systems

License

MIT