nanogpt 0.1.0

data/data/shakespeare_char/prepare.rb

@@ -0,0 +1,73 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+# Prepare the Shakespeare dataset for character-level language modeling.
+# Downloads the tiny shakespeare dataset and creates train.bin, val.bin, and meta.json
+
+require "net/http"
+require "openssl"
+require "numo/narray"
+require "json"
+
+DATA_URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
+SCRIPT_DIR = File.dirname(__FILE__)
+
+def download_file(url)
+  uri = URI(url)
+  http = Net::HTTP.new(uri.host, uri.port)
+  http.use_ssl = true
+  http.verify_mode = OpenSSL::SSL::VERIFY_NONE  # GitHub CDN has CRL issues
+  response = http.get(uri.request_uri)
+  response.body
+end
+
+# Download data if not exists
+input_path = File.join(SCRIPT_DIR, "input.txt")
+unless File.exist?(input_path)
+  puts "Downloading tiny shakespeare..."
+  data = download_file(DATA_URL)
+  File.write(input_path, data)
+end
+
+data = File.read(input_path)
+puts "Length of dataset in characters: #{data.length}"
+
+# Build vocabulary from all unique characters
+chars = data.chars.uniq.sort
+vocab_size = chars.size
+puts "All unique characters: #{chars.join.inspect}"
+puts "Vocab size: #{vocab_size}"
+
+# Create mappings
+stoi = chars.each_with_index.to_h
+itos = chars.each_with_index.map { |c, i| [i, c] }.to_h
+
+# Encode function
+encode = ->(s) { s.chars.map { |c| stoi[c] } }
+
+# Train/val split (90/10)
+n = data.length
+train_data = data[0...(n * 0.9).to_i]
+val_data = data[(n * 0.9).to_i..]
+
+# Encode to integers
+train_ids = encode.call(train_data)
+val_ids = encode.call(val_data)
+puts "Train has #{train_ids.length} tokens"
+puts "Val has #{val_ids.length} tokens"
+
+# Export to binary files (uint16)
+train_arr = Numo::UInt16.cast(train_ids)
+val_arr = Numo::UInt16.cast(val_ids)
+File.binwrite(File.join(SCRIPT_DIR, "train.bin"), train_arr.to_binary)
+File.binwrite(File.join(SCRIPT_DIR, "val.bin"), val_arr.to_binary)
+
+# Save meta information as JSON
+meta = {
+  "vocab_size" => vocab_size,
+  "itos" => itos.transform_keys(&:to_s),
+  "stoi" => stoi
+}
+File.write(File.join(SCRIPT_DIR, "meta.json"), JSON.pretty_generate(meta))
+
+puts "Done! Created train.bin, val.bin, and meta.json"

data/exe/nanogpt

@@ -0,0 +1,338 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+$stdout.sync = true
+
+require "nano_gpt"
+
+class NanoGPTCLI
+  COMMANDS = %w[train sample bench version help].freeze
+
+  def initialize(args)
+    @command = args.shift
+    @args = args
+  end
+
+  def run
+    case @command
+    when "train"
+      train
+    when "sample"
+      sample
+    when "bench"
+      bench
+    when "version", "-v", "--version"
+      version
+    when "help", "-h", "--help", nil
+      help
+    else
+      puts "Unknown command: #{@command}"
+      puts ""
+      help
+      exit 1
+    end
+  end
+
+  private
+
+  def train
+    config = NanoGPT::TrainConfig.load(@args)
+
+    if config[:device] == "auto"
+      config[:device] = NanoGPT::Device.auto
+      puts "Auto-detected device: #{config[:device]}"
+    end
+
+    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}"
+
+    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]
+    )
+
+    model = NanoGPT::GPT.new(model_config)
+
+    device = config[:device]
+    if device != "cpu"
+      model.to(device)
+      puts "Model moved to #{device}"
+    end
+
+    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"
+
+    trainer = NanoGPT::Trainer.new(
+      model: model,
+      data_loader: data_loader,
+      config: config.to_h
+    )
+
+    trainer.train
+
+    puts "\nTraining complete! Checkpoint saved to #{config[:out_dir]}/ckpt.pt"
+  end
+
+  def sample
+    config = NanoGPT::SampleConfig.load(@args)
+
+    if config[:device] == "auto"
+      config[:device] = NanoGPT::Device.auto
+      puts "Auto-detected device: #{config[:device]}"
+    end
+    device = config[:device]
+
+    Torch.manual_seed(config[:seed])
+
+    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: nanogpt train"
+      exit 1
+    end
+
+    puts "Loading checkpoint from #{ckpt_path}..."
+    checkpoint = Torch.load(ckpt_path)
+
+    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
+
+    dataset_dir = File.join("data", config[:dataset])
+    tokenizer = NanoGPT::Tokenizer.for_dataset(dataset_dir)
+    puts "number of parameters: #{model.num_params / 1e6}M"
+
+    start_text = config[:start]
+    if start_text.start_with?("FILE:")
+      start_text = File.read(start_text[5..])
+    end
+
+    start_ids = tokenizer.encode(start_text)
+    x = Torch.tensor([start_ids], dtype: :long, device: device)
+
+    puts "Generating #{config[:num_samples]} samples..."
+    puts "=" * 50
+
+    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
+
+  def bench
+    bench_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"
+    }
+
+    # Parse args
+    @args.each do |arg|
+      next unless arg.start_with?("--") && arg.include?("=")
+
+      key, val = arg[2..].split("=", 2)
+      key = key.to_sym
+
+      next unless bench_config.key?(key)
+
+      bench_config[key] = case bench_config[key]
+                          when Integer then val.to_i
+                          when Float then val.to_f
+                          when TrueClass, FalseClass then val.downcase == "true"
+                          else val
+                          end
+    end
+
+    puts "=" * 60
+    puts "NanoGPT Benchmark"
+    puts "=" * 60
+    puts ""
+    puts "Configuration:"
+    puts "  batch_size: #{bench_config[:batch_size]}"
+    puts "  block_size: #{bench_config[:block_size]}"
+    puts "  n_layer: #{bench_config[:n_layer]}"
+    puts "  n_head: #{bench_config[:n_head]}"
+    puts "  n_embd: #{bench_config[:n_embd]}"
+    puts "  real_data: #{bench_config[:real_data]}"
+    puts ""
+
+    if bench_config[:device] == "auto"
+      bench_config[:device] = NanoGPT::Device.auto
+    end
+    device = bench_config[:device]
+    puts "Device: #{device}"
+
+    Torch.manual_seed(bench_config[:seed])
+
+    if bench_config[:real_data]
+      data_dir = File.join("data", bench_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/#{bench_config[:dataset]}/prepare.rb"
+        puts ""
+        bench_config[:real_data] = false
+      end
+    end
+
+    if bench_config[:real_data]
+      bytes = File.binread(File.join("data", bench_config[:dataset], "train.bin"))
+      train_data = bytes.unpack("S<*")
+      puts "Loaded #{train_data.size} tokens from #{bench_config[:dataset]}"
+
+      get_batch = lambda do
+        max_start = train_data.size - bench_config[:block_size] - 1
+        indices = Array.new(bench_config[:batch_size]) { rand(0..max_start) }
+        x_arrays = indices.map { |i| train_data[i, bench_config[:block_size]] }
+        y_arrays = indices.map { |i| train_data[i + 1, bench_config[:block_size]] }
+        x = Torch.tensor(x_arrays, dtype: :long)
+        y = Torch.tensor(y_arrays, dtype: :long)
+        x = x.to(device) if device != "cpu"
+        y = y.to(device) if device != "cpu"
+        [x, y]
+      end
+    else
+      vocab_size = 50304
+      puts "Using random data (vocab_size=#{vocab_size})"
+
+      get_batch = lambda do
+        x = Torch.randint(vocab_size, [bench_config[:batch_size], bench_config[:block_size]], dtype: :long)
+        y = Torch.randint(vocab_size, [bench_config[:batch_size], bench_config[:block_size]], dtype: :long)
+        x = x.to(device) if device != "cpu"
+        y = y.to(device) if device != "cpu"
+        [x, y]
+      end
+    end
+
+    puts ""
+    puts "Initializing model..."
+    model_config = NanoGPT::GPTConfig.new(
+      block_size: bench_config[:block_size],
+      vocab_size: 50304,
+      n_layer: bench_config[:n_layer],
+      n_head: bench_config[:n_head],
+      n_embd: bench_config[:n_embd],
+      dropout: bench_config[:dropout],
+      bias: bench_config[:bias]
+    )
+
+    model = NanoGPT::GPT.new(model_config)
+    model.to(device) if device != "cpu"
+
+    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
+
+    [{ 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|
+        _logits, loss = model.forward(x, targets: y)
+        x, y = get_batch.call
+        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(bench_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
+
+  def version
+    puts "nanogpt #{NanoGPT::VERSION}"
+  end
+
+  def help
+    puts <<~HELP
+      nanogpt - A Ruby port of Karpathy's nanoGPT
+
+      Usage: nanogpt <command> [options]
+
+      Commands:
+        train     Train a GPT model
+        sample    Generate text from a trained model
+        bench     Run performance benchmarks
+        version   Show version
+        help      Show this help message
+
+      Examples:
+        nanogpt train --config=config/train_shakespeare_char.json
+        nanogpt train --dataset=shakespeare_char --max_iters=1000
+        nanogpt sample --out_dir=out-shakespeare-char --num_samples=3
+        nanogpt bench --batch_size=8 --block_size=512
+
+      For more information, visit: https://github.com/khasinski/nanogpt-rb
+    HELP
+  end
+end
+
+NanoGPTCLI.new(ARGV).run

data/lib/nano_gpt/config.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module NanoGPT
+  # Configuration for GPT model architecture
+  class GPTConfig
+    attr_accessor :block_size, :vocab_size, :n_layer, :n_head, :n_embd, :dropout, :bias
+
+    def initialize(
+      block_size: 1024,
+      vocab_size: 50304,
+      n_layer: 12,
+      n_head: 12,
+      n_embd: 768,
+      dropout: 0.0,
+      bias: true
+    )
+      @block_size = block_size
+      @vocab_size = vocab_size
+      @n_layer = n_layer
+      @n_head = n_head
+      @n_embd = n_embd
+      @dropout = dropout
+      @bias = bias
+    end
+
+    def to_h
+      {
+        block_size: @block_size,
+        vocab_size: @vocab_size,
+        n_layer: @n_layer,
+        n_head: @n_head,
+        n_embd: @n_embd,
+        dropout: @dropout,
+        bias: @bias
+      }
+    end
+
+    def head_size
+      @n_embd / @n_head
+    end
+  end
+end

data/lib/nano_gpt/data_loader.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module NanoGPT
+  # Loads batches from binary token files
+  # Memory-efficient: reads from file each batch (like Python's memmap recreation)
+  class DataLoader
+    attr_reader :block_size, :batch_size
+
+    BYTES_PER_TOKEN = 2  # uint16
+
+    def initialize(data_dir:, block_size:, batch_size:, device: "cpu")
+      @data_dir = data_dir
+      @block_size = block_size
+      @batch_size = batch_size
+      @device = device
+
+      # Store file paths and sizes (NOT the data itself)
+      @train_path = File.join(data_dir, "train.bin")
+      @val_path = File.join(data_dir, "val.bin")
+
+      @train_size = File.size(@train_path) / BYTES_PER_TOKEN
+      @val_size = File.size(@val_path) / BYTES_PER_TOKEN
+    end
+
+    def train_size
+      @train_size
+    end
+
+    def val_size
+      @val_size
+    end
+
+    # Get a batch of data
+    # Memory-efficient: recreates data view per batch to avoid memory leak
+    # (matches Python's memmap recreation pattern)
+    def get_batch(split)
+      path = split == :train ? @train_path : @val_path
+      data_size = split == :train ? @train_size : @val_size
+
+      # Random starting indices
+      max_start = data_size - @block_size - 1
+      indices = Array.new(@batch_size) { rand(0..max_start) }
+
+      # Read only the bytes we need from file (memory-efficient)
+      # This mimics Python's memmap recreation per batch
+      x_arrays = []
+      y_arrays = []
+
+      File.open(path, "rb") do |f|
+        indices.each do |i|
+          # Read x: tokens[i:i+block_size]
+          f.seek(i * BYTES_PER_TOKEN)
+          x_bytes = f.read((@block_size + 1) * BYTES_PER_TOKEN)
+          tokens = x_bytes.unpack("S<*")  # uint16 little-endian
+
+          x_arrays << tokens[0...@block_size]
+          y_arrays << tokens[1..@block_size]
+        end
+      end
+
+      # Create tensors directly from arrays (avoiding Numo intermediate)
+      x = Torch.tensor(x_arrays, dtype: :long)
+      y = Torch.tensor(y_arrays, dtype: :long)
+
+      # Move to device (CPU, CUDA, or MPS)
+      if @device != "cpu"
+        x = x.to(@device)
+        y = y.to(@device)
+      end
+
+      [x, y]
+    end
+  end
+end

data/lib/nano_gpt/device.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module NanoGPT
+  # Device detection and management
+  module Device
+    class << self
+      # Auto-detect the best available device
+      # Priority: CUDA > MPS > CPU
+      def auto
+        return "cuda" if cuda_available?
+        return "mps" if mps_available?
+
+        "cpu"
+      end
+
+      # Check if CUDA is available
+      def cuda_available?
+        Torch::CUDA.available?
+      rescue StandardError
+        false
+      end
+
+      # Check if MPS (Metal Performance Shaders) is available
+      # MPS is Apple Silicon GPU acceleration
+      def mps_available?
+        # Try to create a tensor on MPS device
+        Torch.tensor([1.0], device: "mps")
+        true
+      rescue StandardError
+        false
+      end
+
+      # Get device type string (for optimizer configuration, etc.)
+      def type(device)
+        case device.to_s
+        when /cuda/ then "cuda"
+        when /mps/ then "mps"
+        else "cpu"
+        end
+      end
+
+      # Check if device is GPU (CUDA or MPS)
+      def gpu?(device)
+        %w[cuda mps].include?(type(device))
+      end
+
+      # Print device info
+      def info
+        puts "Device detection:"
+        puts "  CUDA available: #{cuda_available?}"
+        puts "  MPS available: #{mps_available?}"
+        puts "  Auto-selected: #{auto}"
+      end
+    end
+  end
+end

data/lib/nano_gpt/layers/block.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module NanoGPT
+  module Layers
+    # Transformer block: LayerNorm -> Attention -> LayerNorm -> MLP
+    class Block < Torch::NN::Module
+      def initialize(config)
+        super()
+        @ln_1 = LayerNorm.new(config.n_embd, bias: config.bias)
+        @attn = CausalSelfAttention.new(config)
+        @ln_2 = LayerNorm.new(config.n_embd, bias: config.bias)
+        @mlp = MLP.new(config)
+      end
+
+      def forward(x)
+        x = x + @attn.call(@ln_1.call(x))
+        x = x + @mlp.call(@ln_2.call(x))
+        # Trigger GC to free intermediate tensors (critical for torch.rb memory management)
+        # Ruby's GC doesn't run frequently enough during forward pass, causing memory accumulation
+        GC.start(full_mark: false, immediate_sweep: true)
+        x
+      end
+    end
+  end
+end

data/lib/nano_gpt/layers/causal_self_attention.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module NanoGPT
+  module Layers
+    # Multi-head causal self-attention
+    class CausalSelfAttention < Torch::NN::Module
+      def initialize(config)
+        super()
+        raise ArgumentError, "n_embd must be divisible by n_head" unless (config.n_embd % config.n_head).zero?
+
+        @n_head = config.n_head
+        @n_embd = config.n_embd
+        @head_size = config.n_embd / config.n_head
+        @dropout_p = config.dropout
+
+        # Key, query, value projections for all heads, combined
+        @c_attn = Torch::NN::Linear.new(config.n_embd, 3 * config.n_embd, bias: config.bias)
+        # Output projection
+        @c_proj = Torch::NN::Linear.new(config.n_embd, config.n_embd, bias: config.bias)
+        # Regularization
+        @attn_dropout = Torch::NN::Dropout.new(p: config.dropout)
+        @resid_dropout = Torch::NN::Dropout.new(p: config.dropout)
+
+        # Use native scaled_dot_product_attention with is_causal=true when dropout is 0
+        # Native SDPA is ~5x faster but doesn't support dropout with is_causal mode
+        @flash = config.dropout == 0.0
+
+        # Causal mask for manual attention (only used when @flash is false)
+        unless @flash
+          mask = Torch.tril(Torch.ones(config.block_size, config.block_size))
+          register_buffer("mask", mask.view(1, 1, config.block_size, config.block_size))
+        end
+      end
+
+      def forward(x)
+        b, t, c = x.shape
+
+        # Calculate Q, K, V
+        qkv = @c_attn.call(x)
+        q, k, v = qkv.split(@n_embd, 2)
+
+        # Reshape: (B, T, C) -> (B, nh, T, hs)
+        q = q.view(b, t, @n_head, @head_size).transpose(1, 2)
+        k = k.view(b, t, @n_head, @head_size).transpose(1, 2)
+        v = v.view(b, t, @n_head, @head_size).transpose(1, 2)
+
+        y = if @flash
+              # Native scaled_dot_product_attention with is_causal=true
+              # Uses Flash Attention on CUDA, optimized kernel on MPS
+              Torch::NN.scaled_dot_product_attention(q, k, v, nil, 0.0, true)
+            else
+              # Manual attention implementation with causal mask
+              scale = 1.0 / Math.sqrt(@head_size)
+              att = q.matmul(k.transpose(-2, -1))
+              att.mul!(scale)
+
+              # Apply causal mask - slice mask to current sequence length
+              mask_slice = @mask.narrow(2, 0, t).narrow(3, 0, t)
+              att.masked_fill!(mask_slice.eq(0), -Float::INFINITY)
+              att = Torch::NN::Functional.softmax(att, dim: -1)
+              att = @attn_dropout.call(att)
+              att.matmul(v)
+            end
+
+        # Reassemble heads: (B, nh, T, hs) -> (B, T, C)
+        y = y.transpose(1, 2).contiguous.view(b, t, c)
+
+        # Output projection
+        @resid_dropout.call(@c_proj.call(y))
+      end
+    end
+  end
+end

data/lib/nano_gpt/layers/layer_norm.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module NanoGPT
+  module Layers
+    # LayerNorm with optional bias (PyTorch doesn't support bias=false directly)
+    class LayerNorm < Torch::NN::Module
+      attr_reader :weight, :bias
+
+      def initialize(ndim, bias: true)
+        super()
+        @ndim = ndim
+        @weight = Torch::NN::Parameter.new(Torch.ones(ndim))
+        @bias = bias ? Torch::NN::Parameter.new(Torch.zeros(ndim)) : nil
+      end
+
+      def forward(input)
+        Torch::NN::Functional.layer_norm(input, [@ndim], weight: @weight, bias: @bias, eps: 1e-5)
+      end
+    end
+  end
+end