nanogpt 0.1.0

data/lib/nano_gpt/trainer.rb

@@ -0,0 +1,221 @@
+# frozen_string_literal: true
+
+require "fileutils"
+
+module NanoGPT
+  # Training loop for GPT models
+  class Trainer
+    attr_reader :model, :optimizer, :config, :iter_num, :best_val_loss
+
+    def initialize(model:, data_loader:, config: {})
+      @model = model
+      @data_loader = data_loader
+      @config = default_config.merge(config)
+
+      @iter_num = 0
+      @best_val_loss = Float::INFINITY
+
+      setup_optimizer
+      setup_lr_scheduler
+    end
+
+    def default_config
+      {
+        out_dir: "out",
+        eval_interval: 250,
+        log_interval: 10,
+        eval_iters: 200,
+        eval_only: false,
+        always_save_checkpoint: false,
+
+        # Optimizer
+        learning_rate: 1e-3,
+        weight_decay: 1e-1,
+        beta1: 0.9,
+        beta2: 0.99,
+        grad_clip: 1.0,
+
+        # LR scheduler
+        decay_lr: true,
+        warmup_iters: 100,
+        lr_decay_iters: 5000,
+        min_lr: 1e-4,
+
+        # Training
+        max_iters: 5000,
+        gradient_accumulation_steps: 1,
+
+        device: "cpu"
+      }
+    end
+
+    def train
+      puts "Starting training..."
+      puts "Tokens per iteration: #{tokens_per_iter}"
+
+      @model.train
+      x, y = @data_loader.get_batch(:train)
+      t0 = Time.now
+
+      while @iter_num <= @config[:max_iters]
+        # Set learning rate for this iteration
+        lr = @config[:decay_lr] ? @lr_scheduler.step(@optimizer, @iter_num) : @config[:learning_rate]
+
+        # Evaluate and checkpoint
+        if @iter_num % @config[:eval_interval] == 0
+          losses = estimate_loss
+          puts "step #{@iter_num}: train loss #{losses[:train].round(4)}, val loss #{losses[:val].round(4)}"
+
+          if losses[:val] < @best_val_loss || @config[:always_save_checkpoint]
+            @best_val_loss = [losses[:val], @best_val_loss].min
+            save_checkpoint if @iter_num > 0
+          end
+        end
+
+        break if @iter_num == 0 && @config[:eval_only]
+
+        # Forward/backward with gradient accumulation
+        @optimizer.zero_grad
+
+        accumulated_loss = 0.0
+        @config[:gradient_accumulation_steps].times do |micro_step|
+          logits, loss = @model.call(x, targets: y)
+          loss = loss / @config[:gradient_accumulation_steps]
+          accumulated_loss += loss.item
+          loss.backward
+
+          # Prefetch next batch
+          x, y = @data_loader.get_batch(:train)
+        end
+
+        # Gradient clipping (manual implementation since torch.rb lacks clip_grad_norm_)
+        if @config[:grad_clip] > 0.0
+          clip_grad_norm(@model.parameters, @config[:grad_clip])
+        end
+
+        # Optimizer step
+        @optimizer.step
+
+        # Logging
+        t1 = Time.now
+        dt = t1 - t0
+        t0 = t1
+
+        if @iter_num % @config[:log_interval] == 0
+          puts "iter #{@iter_num}: loss #{accumulated_loss.round(4)}, time #{(dt * 1000).round(2)}ms, lr #{lr.round(6)}"
+        end
+
+        @iter_num += 1
+      end
+
+      puts "Training complete!"
+    end
+
+    def estimate_loss
+      @model.eval
+      out = {}
+
+      [:train, :val].each do |split|
+        losses = []
+        @config[:eval_iters].times do
+          x, y = @data_loader.get_batch(split)
+          Torch.no_grad do
+            _logits, loss = @model.call(x, targets: y)
+            losses << loss.item
+          end
+        end
+        out[split] = losses.sum / losses.size
+      end
+
+      @model.train
+      out
+    end
+
+    def save_checkpoint
+      FileUtils.mkdir_p(@config[:out_dir])
+      path = File.join(@config[:out_dir], "ckpt.pt")
+
+      # Note: torch.rb doesn't support optimizer.state_dict yet
+      # We save model state and training metadata
+      # Convert symbol keys to strings for Torch.save compatibility
+      checkpoint = {
+        "model" => @model.state_dict,
+        "model_args" => stringify_keys(@model.config.to_h),
+        "iter_num" => @iter_num,
+        "best_val_loss" => @best_val_loss,
+        "config" => stringify_keys(@config)
+      }
+
+      Torch.save(checkpoint, path)
+      puts "Saved checkpoint to #{path}"
+    end
+
+    def load_checkpoint(path)
+      checkpoint = Torch.load(path)
+
+      @model.load_state_dict(checkpoint["model"])
+      @iter_num = checkpoint["iter_num"]
+      @best_val_loss = checkpoint["best_val_loss"]
+
+      # Reinitialize optimizer (since we can't restore optimizer state in torch.rb)
+      setup_optimizer
+
+      puts "Loaded checkpoint from #{path} (iter #{@iter_num})"
+      checkpoint
+    end
+
+    private
+
+    # Convert symbol keys to strings recursively (for Torch.save)
+    def stringify_keys(hash)
+      hash.transform_keys(&:to_s).transform_values do |v|
+        v.is_a?(Hash) ? stringify_keys(v) : v
+      end
+    end
+
+    # Manual gradient clipping (torch.rb doesn't have clip_grad_norm_)
+    def clip_grad_norm(parameters, max_norm)
+      total_norm = 0.0
+      parameters.each do |p|
+        next unless p.grad
+
+        param_norm = p.grad.data.norm(2).item
+        total_norm += param_norm ** 2
+      end
+      total_norm = Math.sqrt(total_norm)
+
+      clip_coef = max_norm / (total_norm + 1e-6)
+      if clip_coef < 1
+        parameters.each do |p|
+          next unless p.grad
+
+          p.grad.data.mul!(clip_coef)
+        end
+      end
+
+      total_norm
+    end
+
+    def setup_optimizer
+      @optimizer = @model.configure_optimizers(
+        weight_decay: @config[:weight_decay],
+        learning_rate: @config[:learning_rate],
+        betas: [@config[:beta1], @config[:beta2]],
+        device_type: @config[:device]
+      )
+    end
+
+    def setup_lr_scheduler
+      @lr_scheduler = LRScheduler.new(
+        learning_rate: @config[:learning_rate],
+        min_lr: @config[:min_lr],
+        warmup_iters: @config[:warmup_iters],
+        lr_decay_iters: @config[:lr_decay_iters]
+      )
+    end
+
+    def tokens_per_iter
+      @config[:gradient_accumulation_steps] * @data_loader.batch_size * @data_loader.block_size
+    end
+  end
+end

data/lib/nano_gpt/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module NanoGPT
+  VERSION = "0.1.0"
+end

data/lib/nano_gpt.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+require "torch"
+require "numo/narray"
+
+require_relative "nano_gpt/version"
+require_relative "nano_gpt/device"
+require_relative "nano_gpt/config"
+require_relative "nano_gpt/layers/layer_norm"
+require_relative "nano_gpt/layers/mlp"
+require_relative "nano_gpt/layers/causal_self_attention"
+require_relative "nano_gpt/layers/block"
+require_relative "nano_gpt/model"
+require_relative "nano_gpt/tokenizer"
+require_relative "nano_gpt/data_loader"
+require_relative "nano_gpt/lr_scheduler"
+require_relative "nano_gpt/trainer"
+require_relative "nano_gpt/train_config"

data/nanogpt.gemspec

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require_relative "lib/nano_gpt/version"
+
+Gem::Specification.new do |spec|
+  spec.name = "nanogpt"
+  spec.version = NanoGPT::VERSION
+  spec.authors = ["Chris Hasiński"]
+  spec.email = ["krzysztof.hasinski@gmail.com"]
+
+  spec.summary = "A Ruby port of Karpathy's nanoGPT"
+  spec.description = "Train and run GPT models in Ruby using torch.rb. " \
+                     "A minimal, educational implementation of GPT-2 style models."
+  spec.homepage = "https://github.com/khasinski/nanogpt-rb"
+  spec.license = "MIT"
+  spec.required_ruby_version = ">= 3.1.0"
+
+  spec.metadata["homepage_uri"] = spec.homepage
+  spec.metadata["source_code_uri"] = spec.homepage
+  spec.metadata["changelog_uri"] = "#{spec.homepage}/blob/main/CHANGELOG.md"
+
+  spec.files = Dir.chdir(__dir__) do
+    `git ls-files -z`.split("\x0").reject do |f|
+      (File.expand_path(f) == __FILE__) ||
+        f.start_with?("spec/", "test/", ".git", ".github", "vendor/", "python/")
+    end
+  end
+  spec.bindir = "exe"
+  spec.executables = ["nanogpt"]
+  spec.require_paths = ["lib"]
+
+  spec.add_dependency "torch-rb", "~> 0.14"
+  spec.add_dependency "numo-narray", "~> 0.9"
+  spec.add_dependency "tiktoken_ruby", "~> 0.0"
+
+  spec.add_development_dependency "rspec", "~> 3.12"
+end

metadata

@@ -0,0 +1,133 @@
+--- !ruby/object:Gem::Specification
+name: nanogpt
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Chris Hasiński
+autorequire:
+bindir: exe
+cert_chain: []
+date: 2025-12-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: torch-rb
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.14'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.14'
+- !ruby/object:Gem::Dependency
+  name: numo-narray
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.9'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.9'
+- !ruby/object:Gem::Dependency
+  name: tiktoken_ruby
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.0'
+- !ruby/object:Gem::Dependency
+  name: rspec
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.12'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.12'
+description: Train and run GPT models in Ruby using torch.rb. A minimal, educational
+  implementation of GPT-2 style models.
+email:
+- krzysztof.hasinski@gmail.com
+executables:
+- nanogpt
+extensions: []
+extra_rdoc_files: []
+files:
+- Gemfile
+- Gemfile.lock
+- README.md
+- bin/bench
+- bin/sample
+- bin/train
+- config/train_gpt2.json
+- config/train_shakespeare.json
+- config/train_shakespeare_char.json
+- data/openwebtext/prepare.rb
+- data/shakespeare/prepare.rb
+- data/shakespeare_char/input.txt
+- data/shakespeare_char/prepare.rb
+- exe/nanogpt
+- lib/nano_gpt.rb
+- lib/nano_gpt/config.rb
+- lib/nano_gpt/data_loader.rb
+- lib/nano_gpt/device.rb
+- lib/nano_gpt/layers/block.rb
+- lib/nano_gpt/layers/causal_self_attention.rb
+- lib/nano_gpt/layers/layer_norm.rb
+- lib/nano_gpt/layers/mlp.rb
+- lib/nano_gpt/lr_scheduler.rb
+- lib/nano_gpt/model.rb
+- lib/nano_gpt/tokenizer.rb
+- lib/nano_gpt/train_config.rb
+- lib/nano_gpt/trainer.rb
+- lib/nano_gpt/version.rb
+- nanogpt.gemspec
+homepage: https://github.com/khasinski/nanogpt-rb
+licenses:
+- MIT
+metadata:
+  homepage_uri: https://github.com/khasinski/nanogpt-rb
+  source_code_uri: https://github.com/khasinski/nanogpt-rb
+  changelog_uri: https://github.com/khasinski/nanogpt-rb/blob/main/CHANGELOG.md
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.1.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.1.6
+signing_key:
+specification_version: 4
+summary: A Ruby port of Karpathy's nanoGPT
+test_files: []