trainers-rb 0.1.0

data/lib/trainers/lora/lora_linear.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+module Trainers
+  # Wraps a frozen Torch::NN::Linear with low-rank A/B adapter matrices.
+  #
+  # Forward: y = Wx + b + (x @ A^T @ B^T) * scaling
+  #
+  # Only lora_A and lora_B are trainable. The original weight and bias
+  # are frozen, keeping ~99% of parameters fixed during fine-tuning.
+  class LoraLinear < Torch::NN::Module
+    attr_reader :in_features, :out_features, :r, :scaling
+
+    def initialize(original_linear, r:, lora_alpha:, lora_dropout: 0.0)
+      super()
+
+      @in_features  = original_linear.instance_variable_get(:@in_features) ||
+                      original_linear.weight.size(1)
+      @out_features = original_linear.instance_variable_get(:@out_features) ||
+                      original_linear.weight.size(0)
+      @r            = r
+      @scaling      = lora_alpha.to_f / r
+
+      # Freeze original parameters
+      @weight = original_linear.weight
+      @weight.requires_grad = false
+
+      @bias = original_linear.bias
+      @bias.requires_grad = false if @bias
+
+      # LoRA low-rank matrices (these are the only trainable parameters)
+      # A: (r, in_features)  — initialized with Kaiming uniform
+      # B: (out_features, r) — initialized to zero so LoRA starts as identity
+      @lora_A = Torch::NN::Parameter.new(Torch.empty(@r, @in_features))
+      Torch::NN::Init.kaiming_uniform!(@lora_A, a: Math.sqrt(5))
+
+      @lora_B = Torch::NN::Parameter.new(Torch.zeros(@out_features, @r))
+
+      @lora_dropout = lora_dropout > 0 ? Torch::NN::Dropout.new(p: lora_dropout) : nil
+    end
+
+    def forward(x)
+      # Original linear
+      base_output = Torch::NN::F.linear(x, @weight, @bias)
+
+      # LoRA path
+      lora_input = @lora_dropout ? @lora_dropout.call(x) : x
+      lora_output = lora_input.matmul(@lora_A.t).matmul(@lora_B.t) * @scaling
+
+      base_output + lora_output
+    end
+
+    # Merge LoRA weights into the base weight matrix (for inference)
+    def merge!
+      Torch.no_grad do
+        @weight.add!(@lora_B.matmul(@lora_A) * @scaling)
+      end
+      self
+    end
+
+    # Extract LoRA adapter state for saving
+    def lora_state_dict
+      { "lora_A" => @lora_A.data, "lora_B" => @lora_B.data }
+    end
+
+    # Load LoRA weights from saved state
+    def load_lora_weights(lora_a_tensor, lora_b_tensor)
+      Torch.no_grad do
+        @lora_A.copy!(lora_a_tensor)
+        @lora_B.copy!(lora_b_tensor)
+      end
+    end
+
+    def extra_repr
+      "in_features=#{@in_features}, out_features=#{@out_features}, " \
+      "r=#{@r}, scaling=#{@scaling}"
+    end
+  end
+end

data/lib/trainers/lora/lora_model.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+module Trainers
+  class LoraModel
+    # Apply LoRA adapters to a model.
+    #
+    # 1. Finds all Linear layers matching config.target_modules
+    # 2. Replaces each with a LoraLinear that freezes the original weight
+    # 3. Freezes all base model parameters
+    # 4. Only LoRA A/B matrices remain trainable
+    #
+    # Returns the modified model (in-place).
+    def self.apply(model, config)
+      config = LoraConfig.new(**config) if config.is_a?(Hash)
+
+      # Freeze all base parameters first
+      model.parameters.each { |p| p.requires_grad = false }
+
+      # Find and replace target modules
+      targets = LoraUtils.find_target_modules(model, config.target_modules)
+
+      if targets.empty?
+        raise "No Linear modules found matching target_modules: #{config.target_modules.inspect}. " \
+              "Available modules: #{model.named_modules.map(&:first).join(', ')}"
+      end
+
+      targets.each do |name, linear|
+        lora_linear = LoraLinear.new(
+          linear,
+          r:            config.r,
+          lora_alpha:   config.lora_alpha,
+          lora_dropout: config.lora_dropout
+        )
+
+        LoraUtils.replace_module(model, name, lora_linear)
+      end
+
+      # Handle bias training based on config
+      case config.bias
+      when :all
+        model.named_parameters.each do |name, param|
+          param.requires_grad = true if name.include?("bias")
+        end
+      when :lora_only
+        model.named_modules.each do |_, mod|
+          if mod.is_a?(LoraLinear) && mod.instance_variable_get(:@bias)
+            mod.instance_variable_get(:@bias).requires_grad = true
+          end
+        end
+      end
+      # :none — biases stay frozen (default)
+
+      puts "LoRA applied to #{targets.size} modules: #{targets.keys.join(', ')}"
+      LoraUtils.print_trainable_parameters(model)
+
+      model
+    end
+
+    # Merge all LoRA weights back into base weights (for inference)
+    def self.merge(model)
+      count = 0
+      model.named_modules.each do |_, mod|
+        if mod.is_a?(LoraLinear)
+          mod.merge!
+          count += 1
+        end
+      end
+      puts "Merged #{count} LoRA adapters into base model"
+      model
+    end
+
+    # Save only the LoRA adapter weights
+    def self.save_adapters(model, output_dir, config: nil)
+      SaveUtils.save_lora_adapters(model, output_dir)
+
+      if config
+        config_path = File.join(output_dir, "lora_config.json")
+        File.write(config_path, JSON.pretty_generate(config.to_h))
+      end
+    end
+
+    # Load LoRA adapter weights into a model that already has LoRA applied
+    def self.load_adapters(model, input_dir)
+      SaveUtils.load_lora_adapters(model, input_dir)
+    end
+  end
+end

data/lib/trainers/lora/lora_utils.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module Trainers
+  module LoraUtils
+    # Find all Linear modules in a model matching the target pattern
+    def self.find_target_modules(model, target_modules)
+      targets = {}
+
+      model.named_modules.each do |name, mod|
+        next unless mod.is_a?(Torch::NN::Linear)
+
+        if target_modules == :all_linear
+          targets[name] = mod
+        elsif target_modules.is_a?(Array)
+          if target_modules.any? { |pattern| name.include?(pattern) }
+            targets[name] = mod
+          end
+        end
+      end
+
+      targets
+    end
+
+    # Replace a child module on a parent by setting the instance variable.
+    # torch-rb's named_children discovers modules via instance variable scan,
+    # so setting the ivar is the correct replacement mechanism.
+    def self.replace_module(model, target_name, new_module)
+      parts = target_name.split(".")
+      parent = model
+
+      # Navigate to the parent module
+      parts[0...-1].each do |part|
+        if part =~ /\A\d+\z/
+          # Numeric index — likely a ModuleList element
+          parent = parent.instance_variable_get(:@modules)[part.to_i]
+        else
+          parent = parent.instance_variable_get(:"@#{part}")
+        end
+
+        raise "Could not find module part '#{part}' in #{target_name}" unless parent
+      end
+
+      child_name = parts.last
+      if child_name =~ /\A\d+\z/
+        parent.instance_variable_get(:@modules)[child_name.to_i] = new_module
+      else
+        parent.instance_variable_set(:"@#{child_name}", new_module)
+      end
+    end
+
+    # Count total and trainable parameters
+    def self.print_trainable_parameters(model)
+      total    = 0
+      trainable = 0
+
+      model.parameters.each do |p|
+        total += p.numel
+        trainable += p.numel if p.requires_grad
+      end
+
+      pct = total > 0 ? (trainable.to_f / total * 100) : 0.0
+      puts "trainable params: #{format_number(trainable)} || " \
+           "all params: #{format_number(total)} || " \
+           "trainable%: #{format('%.4f', pct)}%"
+
+      { total: total, trainable: trainable, percentage: pct }
+    end
+
+    def self.format_number(n)
+      n.to_s.reverse.gsub(/(\d{3})(?=\d)/, '\\1,').reverse
+    end
+  end
+end

data/lib/trainers/optimization/optimizer.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module Trainers
+  module Optimization
+    # Builds AdamW with two param groups:
+    #   - Parameters with weight decay (Linear weights, Embedding weights)
+    #   - Parameters without weight decay (biases, LayerNorm/layer_norm params)
+    #
+    # This split is critical for transformer fine-tuning — regularizing biases
+    # and normalization weights hurts convergence.
+    def self.create_optimizer(model, args)
+      decay_params    = []
+      no_decay_params = []
+
+      no_decay_patterns = ["bias", "LayerNorm", "layer_norm", "layernorm"]
+
+      model.named_parameters.each do |name, param|
+        next unless param.requires_grad
+
+        if no_decay_patterns.any? { |pattern| name.include?(pattern) }
+          no_decay_params << param
+        else
+          decay_params << param
+        end
+      end
+
+      param_groups = []
+
+      if decay_params.any?
+        param_groups << { params: decay_params, weight_decay: args.weight_decay }
+      end
+
+      if no_decay_params.any?
+        param_groups << { params: no_decay_params, weight_decay: 0.0 }
+      end
+
+      if param_groups.empty?
+        raise "No trainable parameters found. Did you forget to unfreeze the model or apply LoRA?"
+      end
+
+      Torch::Optim::AdamW.new(
+        param_groups,
+        lr:    args.learning_rate,
+        betas: [args.adam_beta1, args.adam_beta2],
+        eps:   args.adam_epsilon
+      )
+    end
+  end
+end

data/lib/trainers/optimization/scheduler.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module Trainers
+  module Optimization
+    # Linear warmup then linear decay to 0
+    def self.get_linear_schedule_with_warmup(optimizer, num_warmup_steps:, num_training_steps:)
+      lr_lambda = ->(current_step) {
+        if current_step < num_warmup_steps
+          current_step.to_f / [1, num_warmup_steps].max
+        else
+          remaining   = num_training_steps - current_step
+          total_decay = num_training_steps - num_warmup_steps
+          [0.0, remaining.to_f / [1.0, total_decay].max].max
+        end
+      }
+
+      Torch::Optim::LRScheduler::LambdaLR.new(optimizer, lr_lambda)
+    end
+
+    # Linear warmup then cosine decay to 0
+    def self.get_cosine_schedule_with_warmup(optimizer, num_warmup_steps:, num_training_steps:, num_cycles: 0.5)
+      lr_lambda = ->(current_step) {
+        if current_step < num_warmup_steps
+          current_step.to_f / [1, num_warmup_steps].max
+        else
+          progress = (current_step - num_warmup_steps).to_f /
+                     [1, num_training_steps - num_warmup_steps].max
+          [0.0, 0.5 * (1.0 + Math.cos(Math::PI * num_cycles * 2.0 * progress))].max
+        end
+      }
+
+      Torch::Optim::LRScheduler::LambdaLR.new(optimizer, lr_lambda)
+    end
+
+    # Linear warmup then constant LR
+    def self.get_constant_schedule_with_warmup(optimizer, num_warmup_steps:)
+      lr_lambda = ->(current_step) {
+        if current_step < num_warmup_steps
+          current_step.to_f / [1, num_warmup_steps].max
+        else
+          1.0
+        end
+      }
+
+      Torch::Optim::LRScheduler::LambdaLR.new(optimizer, lr_lambda)
+    end
+
+    # Dispatcher: pick scheduler by type symbol
+    def self.create_scheduler(type, optimizer, num_warmup_steps:, num_training_steps:)
+      case type
+      when :linear
+        get_linear_schedule_with_warmup(optimizer,
+          num_warmup_steps: num_warmup_steps,
+          num_training_steps: num_training_steps)
+      when :cosine
+        get_cosine_schedule_with_warmup(optimizer,
+          num_warmup_steps: num_warmup_steps,
+          num_training_steps: num_training_steps)
+      when :constant
+        get_constant_schedule_with_warmup(optimizer,
+          num_warmup_steps: num_warmup_steps)
+      else
+        raise ArgumentError, "Unknown scheduler type: #{type}. Use :linear, :cosine, or :constant"
+      end
+    end
+  end
+end

data/lib/trainers/save_utils.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+require "json"
+require "fileutils"
+
+module Trainers
+  module SaveUtils
+    def self.save_pretrained(model, tokenizer, output_dir, training_args: nil)
+      FileUtils.mkdir_p(output_dir)
+
+      # Save model weights via safetensors
+      state_dict = model.state_dict
+      tensors = {}
+      state_dict.each do |name, param|
+        tensors[name] = param.is_a?(Torch::NN::Parameter) ? param.data : param
+      end
+
+      model_path = File.join(output_dir, "model.safetensors")
+      require "safetensors"
+      Safetensors::Torch.save_file(tensors, model_path)
+      puts "Model saved to #{model_path}"
+
+      # Save tokenizer
+      if tokenizer
+        if tokenizer.respond_to?(:save)
+          tokenizer.save(File.join(output_dir, "tokenizer.json"))
+        end
+      end
+
+      # Save training args
+      if training_args
+        args_path = File.join(output_dir, "training_args.json")
+        File.write(args_path, JSON.pretty_generate(training_args.to_h))
+      end
+
+      output_dir
+    end
+
+    def self.save_lora_adapters(model, output_dir)
+      FileUtils.mkdir_p(output_dir)
+
+      lora_state = {}
+      model.named_modules.each do |name, mod|
+        if mod.is_a?(Trainers::LoraLinear)
+          mod.lora_state_dict.each do |param_name, tensor|
+            lora_state["#{name}.#{param_name}"] = tensor
+          end
+        end
+      end
+
+      if lora_state.empty?
+        puts "No LoRA adapters found in model"
+        return
+      end
+
+      adapter_path = File.join(output_dir, "adapter_model.safetensors")
+      require "safetensors"
+      Safetensors::Torch.save_file(lora_state, adapter_path)
+      puts "LoRA adapters saved to #{adapter_path} (#{lora_state.size} tensors)"
+
+      output_dir
+    end
+
+    def self.load_lora_adapters(model, input_dir)
+      adapter_path = File.join(input_dir, "adapter_model.safetensors")
+      raise "Adapter file not found: #{adapter_path}" unless File.exist?(adapter_path)
+
+      require "safetensors"
+      lora_state = Safetensors::Torch.load_file(adapter_path)
+
+      model.named_modules.each do |name, mod|
+        if mod.is_a?(Trainers::LoraLinear)
+          a_key = "#{name}.lora_A"
+          b_key = "#{name}.lora_B"
+          if lora_state[a_key] && lora_state[b_key]
+            mod.load_lora_weights(lora_state[a_key], lora_state[b_key])
+          end
+        end
+      end
+
+      puts "LoRA adapters loaded from #{adapter_path}"
+    end
+  end
+end