MicroGPT in Rust

The most atomic way to train and inference a GPT in pure, dependency-free Rust.

Note: This file is the complete algorithm. Everything else is just efficiency.

Overview

This is a minimal, dependency-free implementation for understanding Generative Pre-trained Transformers (GPT).

GitHub Repository: https://github.com/aarambh-darshan/microgpt-rust

The implementation demonstrates the core concepts of modern large language models (LLMs) including:

• Token embeddings and positional encodings
• Multi-head self-attention mechanism
• Feed-forward neural networks (MLP)
• Layer normalization (RMSNorm)
• Autoregressive training with cross-entropy loss
• Adam optimizer with gradient clipping

Features

• Zero Dependencies: Pure Rust implementation with only the standard library
• Complete Training Pipeline: From random initialization to trained model
• Text Generation: Generate new text samples after training
• Educational: Clean, readable code that demonstrates core concepts
• Fast: Optimized release builds for reasonable training times

Architecture

The model follows the GPT-2 architecture with minor modifications:

• Embedding dimension: 16
• Number of attention heads: 4
• Number of layers: 1
• Context length (block size): 16
• Vocabulary size: 27 (A-Z characters + BOS token)
• Total parameters: ~4,192

Modifications from standard GPT-2:

• LayerNorm → RMSNorm
• GeLU → ReLU activation
• No bias terms in linear layers

Quick Start

Prerequisites

• Rust toolchain (1.70+)
• curl (for downloading the dataset)

Running

# Clone the repository
git clone https://github.com/aarambh-darshan/microgpt-rust.git
cd microgpt-rust

# Run in release mode for optimal performance
cargo run --release

The program will:

1. Download the names dataset (if not present)
2. Train for 1000 steps (time varies by hardware, use --release for best performance)
3. Generate 20 new sample names

Expected Output

num docs: 32033
vocab size: 27
num params: 4192
step    1 / 1000 | loss 3.2889
step    2 / 1000 | loss 3.3999
step    3 / 1000 | loss 3.3450
...
step 1000 / 1000 | loss 2.6440

--- inference (new, hallucinated names) ---
sample  1: ama
sample  2: adana
sample  3: sada
sample  4: salan
sample  5: caanan
...

How It Works

Autograd System

The implementation includes a minimal automatic differentiation engine:

// Values track their computation graph for backprop
let x = Value::new(2.0);
let y = x.mul(&x);  // y = x²
y.backward();        // dy/dx = 4.0

GPT Forward Pass

fn forward(&self, token_id: usize, pos_id: usize, ...) -> Vec<Rc<Value>> {
    // Token + positional embedding
    let tok_emb = &self.wte[token_id];
    let pos_emb = &self.wpe[pos_id];
    let mut x: Vec<Rc<Value>> = tok_emb.iter()
        .zip(pos_emb.iter())
        .map(|(t, p)| t.add(p))
        .collect();
    
    // Multi-head attention + MLP
    // ... (see src/main.rs for full implementation)
    
    linear(&x, &self.lm_head)
}

Training Loop

The training follows standard supervised learning:

1. Sample a document from the dataset
2. Tokenize and create input-target pairs
3. Forward pass through the model
4. Compute cross-entropy loss
5. Backward pass to compute gradients
6. Update parameters with Adam optimizer

Key Implementation Details

Reference-Counted Values

The implementation uses Rc<Value> (reference counting) to manage the computation graph:

struct Value {
    data: f64,
    grad: RefCell<f64>,
    children: Vec<Rc<Value>>,  // References to parent nodes
    local_grads: Vec<f64>,     // Local derivatives
}

Gradient Clipping

To ensure stable training, gradients are clipped to a maximum norm:

let max_grad_norm = 1.0;
let grad_norm = compute_gradient_norm(&grads);
let clip_scale = if grad_norm > max_grad_norm {
    max_grad_norm / grad_norm
} else { 1.0 };

RMSNorm

Layer normalization without centering:

fn rmsnorm(x: &[Rc<Value>]) -> Vec<Rc<Value>> {
    let ms = mean_square(x);
    let scale = (ms + 1e-5).pow(-0.5);
    x.iter().map(|xi| xi.mul(&scale)).collect()
}

Performance

Training performance varies by hardware. On modern CPUs with --release:

• Typical training: ~45 steps/second
• Memory usage: <50MB

Note: Always run with --release flag for optimal performance:

cargo run --release

Limitations

This is an educational implementation, not production-ready:

• Single-threaded (no parallelization)
• No GPU support
• Fixed hyperparameters
• Character-level tokenization only
• No checkpointing/resuming

Educational Resources

To understand the concepts behind this implementation:

1. Andrej Karpathy's Neural Networks: Zero to Hero

   • YouTube playlist covering backpropagation, MLPs, RNNs, Transformers
   • https://karpathy.ai/zero-to-hero.html

2. The Illustrated Transformer by Jay Alammar

   • Visual explanation of attention mechanisms
   • https://jalammar.github.io/illustrated-transformer/

3. Attention Is All You Need (Original Transformer Paper)

   • Vaswani et al., 2017
   • https://arxiv.org/abs/1706.03762

4. GPT-2 Paper

   • Radford et al., 2019
   • https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf

License

This project is released under the MIT License.

Acknowledgments

• Darshan Vichhi (@aarambh-darshan) - Implementation
• Andrej Karpathy (@karpathy) for inspiration and excellent educational content
• The Rust community for providing a zero-cost abstraction language perfect for systems programming

Author

Darshan Vichhi - GitHub

Contributing

This is primarily an educational project. Feel free to fork and experiment with:

• Different model architectures (deeper networks, different attention heads)
• Alternative optimizers (SGD with momentum, AdaGrad, etc.)
• Different datasets (text generation, code completion, etc.)
• Performance optimizations (parallelization, SIMD, etc.)

---

Built with ❤️ for the love of understanding how transformers work.