nanogpt 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: b10e62747a63be10b519fdd8eb1eaabfc63214b01ab553d4230af865a0a8f327
+  data.tar.gz: 69db126e3d02c897045e543981223d1d29d72f7625bd381e3ba4badee88b6d20
+SHA512:
+  metadata.gz: 97f5846373224b889f22b80302a05271c7ab0276037e519c19c58cc0954aa982a4d3ac9992766aa73f6d1eec1dfe2b96b8a750414dd02818529e37789671cd87
+  data.tar.gz: accd2112564f004cc763dbaac418f35ce24426d4d86e53935a32bdf6411659463ed6cb8e9415146136ee78fa79748e39a1e506b5c44848a7b6e503dfee86e310

data/Gemfile

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+source "https://rubygems.org"
+
+gemspec
+
+gem "parquet", "~> 0.5", require: false # For OpenWebText data prep

data/Gemfile.lock

@@ -0,0 +1,42 @@
+PATH
+  remote: .
+  specs:
+    nanogpt (0.1.0)
+      numo-narray (~> 0.9)
+      tiktoken_ruby (~> 0.0)
+      torch-rb (~> 0.14)
+
+GEM
+  remote: https://rubygems.org/
+  specs:
+    diff-lcs (1.6.2)
+    numo-narray (0.9.2.1)
+    parquet (0.7.3-arm64-darwin)
+    rice (4.7.1)
+    rspec (3.13.2)
+      rspec-core (~> 3.13.0)
+      rspec-expectations (~> 3.13.0)
+      rspec-mocks (~> 3.13.0)
+    rspec-core (3.13.6)
+      rspec-support (~> 3.13.0)
+    rspec-expectations (3.13.5)
+      diff-lcs (>= 1.2.0, < 2.0)
+      rspec-support (~> 3.13.0)
+    rspec-mocks (3.13.7)
+      diff-lcs (>= 1.2.0, < 2.0)
+      rspec-support (~> 3.13.0)
+    rspec-support (3.13.6)
+    tiktoken_ruby (0.0.13-arm64-darwin)
+    torch-rb (0.22.2)
+      rice (>= 4.7)
+
+PLATFORMS
+  arm64-darwin-24
+
+DEPENDENCIES
+  nanogpt!
+  parquet (~> 0.5)
+  rspec (~> 3.12)
+
+BUNDLED WITH
+   2.7.2

data/README.md

@@ -0,0 +1,102 @@
+# nanoGPT
+
+A Ruby port of Karpathy's [nanoGPT](https://github.com/karpathy/nanoGPT). Train GPT-2 style language models from scratch using [torch.rb](https://github.com/ankane/torch.rb).
+
+Built for Ruby developers who want to understand how LLMs work by building one.
+
+## Quick Start
+
+```bash
+gem install nanogpt
+
+# Prepare Shakespeare dataset with character-level tokenizer
+nanogpt prepare shakespeare_char
+
+# Train (use MPS on Apple Silicon for 17x speedup)
+nanogpt train --dataset=shakespeare_char --device=mps
+
+# Generate text
+nanogpt sample --dataset=shakespeare_char
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/khasinski/nanogpt-rb
+cd nanogpt-rb
+bundle install
+
+# Prepare data
+bundle exec ruby data/shakespeare_char/prepare.rb
+
+# Train
+bundle exec exe/nanogpt train --dataset=shakespeare_char --device=mps
+
+# Sample
+bundle exec exe/nanogpt sample --dataset=shakespeare_char
+```
+
+## Performance (M1 Max)
+
+Training the default 10.65M parameter model on Shakespeare:
+
+| Device | Time/iter | Notes |
+|--------|-----------|-------|
+| MPS    | ~500ms    | Recommended for Apple Silicon |
+| CPU    | ~8,500ms  | 17x slower |
+
+After ~2000 iterations (~20 min on MPS), the model generates coherent Shakespeare-like text.
+
+## Commands
+
+```bash
+nanogpt train [options]    # Train a model
+nanogpt sample [options]   # Generate text from trained model
+nanogpt bench [options]    # Run performance benchmarks
+```
+
+### Training Options
+
+```bash
+--dataset=NAME        # Dataset to use (default: shakespeare_char)
+--device=DEVICE       # cpu or mps, cuda might work too 🤞(default: auto)
+--max_iters=N         # Training iterations (default: 5000)
+--batch_size=N        # Batch size (default: 64)
+--block_size=N        # Context length (default: 256)
+--n_layer=N           # Transformer layers (default: 6)
+--n_head=N            # Attention heads (default: 6)
+--n_embd=N            # Embedding dimension (default: 384)
+--learning_rate=F     # Learning rate (default: 1e-3)
+--config=FILE         # Load settings from JSON file
+```
+
+### Sampling Options
+
+```bash
+--dataset=NAME        # Dataset (for tokenizer)
+--out_dir=DIR         # Checkpoint directory
+--num_samples=N       # Number of samples to generate
+--max_new_tokens=N    # Tokens per sample (default: 500)
+--temperature=F       # Sampling temperature (default: 0.8)
+--top_k=N            # Top-k sampling (default: 200)
+```
+
+## Features
+
+- Full GPT-2 architecture (attention, MLP, layer norm, embeddings)
+- MPS (Metal) and CUDA GPU acceleration via torch.rb
+- Flash attention when dropout=0 (5x faster attention)
+- Cosine learning rate schedule with warmup
+- Gradient accumulation for larger effective batch sizes
+- Checkpointing and resumption
+- Character-level and GPT-2 BPE tokenizers
+
+## Requirements
+
+- Ruby >= 3.1
+- LibTorch (installed automatically with torch-rb)
+- For MPS: macOS 12.3+ with Apple Silicon
+
+## License
+
+MIT

data/bin/bench

@@ -0,0 +1,210 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+$stdout.sync = true
+
+# Benchmark script for measuring training performance
+# A shorter version of train.rb focused on benchmarking
+#
+# Usage:
+#   bundle exec ruby bin/bench
+#   bundle exec ruby bin/bench --batch_size=8 --block_size=512
+#   bundle exec ruby bin/bench --real_data=false  # use random data
+
+require_relative "../lib/nano_gpt"
+
+# Default configuration (GPT-2 style model)
+CONFIG = {
+  batch_size: 12,
+  block_size: 1024,
+  n_layer: 12,
+  n_head: 12,
+  n_embd: 768,
+  dropout: 0.0,
+  bias: false,
+  real_data: true,
+  dataset: "openwebtext",
+  seed: 1337,
+  device: "auto"
+}.freeze
+
+def parse_args(args, config)
+  result = config.dup
+
+  args.each do |arg|
+    next unless arg.start_with?("--") && arg.include?("=")
+
+    key, val = arg[2..].split("=", 2)
+    key = key.to_sym
+
+    unless result.key?(key)
+      puts "Warning: Unknown config key: #{key}"
+      next
+    end
+
+    result[key] = case result[key]
+                  when Integer then val.to_i
+                  when Float then val.to_f
+                  when TrueClass, FalseClass then val.downcase == "true"
+                  else val
+                  end
+  end
+
+  result
+end
+
+def get_batch_from_data(data, batch_size, block_size, device)
+  max_start = data.size - block_size - 1
+  indices = Array.new(batch_size) { rand(0..max_start) }
+
+  x_arrays = []
+  y_arrays = []
+
+  indices.each do |i|
+    x_arrays << data[i, block_size]
+    y_arrays << data[i + 1, block_size]
+  end
+
+  x = Torch.tensor(x_arrays, dtype: :long)
+  y = Torch.tensor(y_arrays, dtype: :long)
+
+  if device != "cpu"
+    x = x.to(device)
+    y = y.to(device)
+  end
+
+  [x, y]
+end
+
+def main
+  config = parse_args(ARGV, CONFIG)
+
+  puts "=" * 60
+  puts "NanoGPT Benchmark"
+  puts "=" * 60
+  puts ""
+  puts "Configuration:"
+  puts "  batch_size: #{config[:batch_size]}"
+  puts "  block_size: #{config[:block_size]}"
+  puts "  n_layer: #{config[:n_layer]}"
+  puts "  n_head: #{config[:n_head]}"
+  puts "  n_embd: #{config[:n_embd]}"
+  puts "  real_data: #{config[:real_data]}"
+  puts ""
+
+  # Resolve device
+  if config[:device] == "auto"
+    config[:device] = NanoGPT::Device.auto
+  end
+  device = config[:device]
+  puts "Device: #{device}"
+
+  Torch.manual_seed(config[:seed])
+
+  # Data loading
+  if config[:real_data]
+    data_dir = File.join("data", config[:dataset])
+    train_bin = File.join(data_dir, "train.bin")
+
+    unless File.exist?(train_bin)
+      puts ""
+      puts "Warning: #{train_bin} not found, using random data instead."
+      puts "To use real data, run: bundle exec ruby data/#{config[:dataset]}/prepare.rb"
+      puts ""
+      config[:real_data] = false
+    end
+  end
+
+  if config[:real_data]
+    # Load training data
+    bytes = File.binread(File.join("data", config[:dataset], "train.bin"))
+    train_data = bytes.unpack("S<*")  # uint16 little-endian
+    puts "Loaded #{train_data.size} tokens from #{config[:dataset]}"
+
+    get_batch = lambda do
+      get_batch_from_data(train_data, config[:batch_size], config[:block_size], device)
+    end
+  else
+    # Use random data
+    vocab_size = 50304  # GPT-2 vocab size rounded up for efficiency
+    puts "Using random data (vocab_size=#{vocab_size})"
+
+    get_batch = lambda do
+      x = Torch.randint(vocab_size, [config[:batch_size], config[:block_size]], dtype: :long)
+      y = Torch.randint(vocab_size, [config[:batch_size], config[:block_size]], dtype: :long)
+      x = x.to(device) if device != "cpu"
+      y = y.to(device) if device != "cpu"
+      [x, y]
+    end
+  end
+
+  # Model init
+  puts ""
+  puts "Initializing model..."
+  model_config = NanoGPT::GPTConfig.new(
+    block_size: config[:block_size],
+    vocab_size: 50304,  # GPT-2 vocab rounded up
+    n_layer: config[:n_layer],
+    n_head: config[:n_head],
+    n_embd: config[:n_embd],
+    dropout: config[:dropout],
+    bias: config[:bias]
+  )
+
+  model = NanoGPT::GPT.new(model_config)
+  model.to(device) if device != "cpu"
+
+  # Optimizer
+  optimizer = model.configure_optimizers(
+    weight_decay: 1e-2,
+    learning_rate: 1e-4,
+    betas: [0.9, 0.95],
+    device_type: NanoGPT::Device.type(device)
+  )
+
+  puts ""
+  puts "Starting benchmark..."
+  puts "-" * 60
+
+  # Benchmark: burn-in phase then measurement phase
+  [{ name: "burn-in", steps: 10 }, { name: "benchmark", steps: 20 }].each do |phase|
+    puts ""
+    puts "Phase: #{phase[:name]} (#{phase[:steps]} steps)"
+
+    x, y = get_batch.call
+    t0 = Time.now
+
+    phase[:steps].times do |k|
+      # Forward pass
+      _logits, loss = model.forward(x, targets: y)
+
+      # Get next batch
+      x, y = get_batch.call
+
+      # Backward pass
+      optimizer.zero_grad
+      loss.backward
+      optimizer.step
+
+      loss_val = loss.item
+      puts "  #{k}/#{phase[:steps]} loss: #{format('%.4f', loss_val)}"
+    end
+
+    t1 = Time.now
+    dt = t1 - t0
+
+    if phase[:name] == "benchmark"
+      mfu = model.estimate_mfu(config[:batch_size] * phase[:steps], dt)
+      time_per_iter = dt / phase[:steps] * 1000
+
+      puts ""
+      puts "=" * 60
+      puts "Results:"
+      puts "  Time per iteration: #{format('%.2f', time_per_iter)}ms"
+      puts "  MFU: #{format('%.2f', mfu * 100)}%"
+      puts "=" * 60
+    end
+  end
+end
+
+main

data/bin/sample

@@ -0,0 +1,76 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+# Generate text from a trained model
+# Usage:
+#   bundle exec ruby bin/sample [options]
+#   bundle exec ruby bin/sample --config=config/sample.json
+#   bundle exec ruby bin/sample --out_dir=out-shakespeare --num_samples=3
+
+require_relative "../lib/nano_gpt"
+
+def main
+  config = NanoGPT::SampleConfig.load(ARGV)
+
+  # Resolve device (auto-detect if "auto")
+  if config[:device] == "auto"
+    config[:device] = NanoGPT::Device.auto
+    puts "Auto-detected device: #{config[:device]}"
+  end
+  device = config[:device]
+
+  Torch.manual_seed(config[:seed])
+
+  # Load checkpoint
+  ckpt_path = File.join(config[:out_dir], "ckpt.pt")
+  unless File.exist?(ckpt_path)
+    puts "Error: Checkpoint not found at #{ckpt_path}"
+    puts "Train a model first with: bundle exec ruby bin/train"
+    exit 1
+  end
+
+  puts "Loading checkpoint from #{ckpt_path}..."
+  checkpoint = Torch.load(ckpt_path)
+
+  # Recreate model from checkpoint
+  model_args = checkpoint["model_args"].transform_keys(&:to_sym)
+  model_config = NanoGPT::GPTConfig.new(**model_args)
+  model = NanoGPT::GPT.new(model_config)
+  model.load_state_dict(checkpoint["model"])
+  model.to(device) if device != "cpu"
+  model.eval
+
+  # Load tokenizer (auto-detects character-level vs GPT-2 BPE)
+  dataset_dir = File.join("data", config[:dataset])
+  tokenizer = NanoGPT::Tokenizer.for_dataset(dataset_dir)
+  puts "number of parameters: #{model.num_params / 1e6}M"
+
+  # Handle start text
+  start_text = config[:start]
+  if start_text.start_with?("FILE:")
+    start_text = File.read(start_text[5..])
+  end
+
+  # Encode starting prompt
+  start_ids = tokenizer.encode(start_text)
+  x = Torch.tensor([start_ids], dtype: :long, device: device)
+
+  puts "Generating #{config[:num_samples]} samples..."
+  puts "=" * 50
+
+  # Generate samples
+  config[:num_samples].times do |k|
+    y = model.generate(
+      x,
+      config[:max_new_tokens],
+      temperature: config[:temperature],
+      top_k: config[:top_k]
+    )
+
+    output = tokenizer.decode(y[0].to_a)
+    puts output
+    puts "-" * 50
+  end
+end
+
+main

data/bin/train

@@ -0,0 +1,82 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+$stdout.sync = true
+
+# Train a GPT model
+# Usage:
+#   bundle exec ruby bin/train [options]
+#   bundle exec ruby bin/train --config=config/train_shakespeare.json
+#   bundle exec ruby bin/train --config=config/train_gpt2.json --max_iters=1000
+
+require_relative "../lib/nano_gpt"
+
+def main
+  config = NanoGPT::TrainConfig.load(ARGV)
+
+  # Resolve device (auto-detect if "auto")
+  if config[:device] == "auto"
+    config[:device] = NanoGPT::Device.auto
+    puts "Auto-detected device: #{config[:device]}"
+  end
+
+  # Load tokenizer to get vocab_size (auto-detects char-level vs GPT-2 BPE)
+  data_dir = File.join("data", config[:dataset])
+  train_bin = File.join(data_dir, "train.bin")
+
+  unless File.exist?(train_bin)
+    puts "Error: #{train_bin} not found. Run the data preparation script first:"
+    puts "  bundle exec ruby data/#{config[:dataset]}/prepare.rb"
+    exit 1
+  end
+
+  tokenizer = NanoGPT::Tokenizer.for_dataset(data_dir)
+  tokenizer_type = tokenizer.is_a?(NanoGPT::GPT2Tokenizer) ? "GPT-2 BPE" : "character-level"
+  puts "Loaded #{tokenizer_type} tokenizer with vocab_size=#{tokenizer.vocab_size}"
+
+  # Create model config
+  model_config = NanoGPT::GPTConfig.new(
+    block_size: config[:block_size],
+    vocab_size: tokenizer.vocab_size,
+    n_layer: config[:n_layer],
+    n_head: config[:n_head],
+    n_embd: config[:n_embd],
+    dropout: config[:dropout],
+    bias: config[:bias]
+  )
+
+  # Create model
+  model = NanoGPT::GPT.new(model_config)
+
+  # Move model to device (MPS, CUDA, or CPU)
+  device = config[:device]
+  if device != "cpu"
+    model.to(device)
+    puts "Model moved to #{device}"
+  end
+
+  # Create data loader
+  data_loader = NanoGPT::DataLoader.new(
+    data_dir: data_dir,
+    block_size: config[:block_size],
+    batch_size: config[:batch_size],
+    device: config[:device]
+  )
+
+  puts "Train data: #{data_loader.train_size} tokens"
+  puts "Val data: #{data_loader.val_size} tokens"
+
+  # Create trainer
+  trainer = NanoGPT::Trainer.new(
+    model: model,
+    data_loader: data_loader,
+    config: config.to_h
+  )
+
+  # Train!
+  trainer.train
+
+  puts "\nTraining complete! Checkpoint saved to #{config[:out_dir]}/ckpt.pt"
+end
+
+main

data/config/train_gpt2.json

@@ -0,0 +1,19 @@
+{
+  "out_dir": "out-gpt2",
+  "dataset": "openwebtext",
+  "batch_size": 12,
+  "block_size": 1024,
+  "n_layer": 12,
+  "n_head": 12,
+  "n_embd": 768,
+  "dropout": 0.0,
+  "bias": true,
+  "learning_rate": 6e-4,
+  "max_iters": 600000,
+  "warmup_iters": 2000,
+  "lr_decay_iters": 600000,
+  "min_lr": 6e-5,
+  "eval_interval": 1000,
+  "eval_iters": 200,
+  "device": "auto"
+}

data/config/train_shakespeare.json

@@ -0,0 +1,14 @@
+{
+  "out_dir": "out-shakespeare",
+  "dataset": "shakespeare",
+  "batch_size": 64,
+  "block_size": 256,
+  "n_layer": 6,
+  "n_head": 6,
+  "n_embd": 384,
+  "dropout": 0.2,
+  "learning_rate": 1e-3,
+  "max_iters": 5000,
+  "eval_interval": 250,
+  "device": "auto"
+}

data/config/train_shakespeare_char.json

@@ -0,0 +1,14 @@
+{
+  "out_dir": "out-shakespeare-char",
+  "dataset": "shakespeare_char",
+  "batch_size": 64,
+  "block_size": 256,
+  "n_layer": 6,
+  "n_head": 6,
+  "n_embd": 384,
+  "dropout": 0.2,
+  "learning_rate": 1e-3,
+  "max_iters": 5000,
+  "eval_interval": 250,
+  "device": "auto"
+}