trainers-rb 0.1.1 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d574e3490a0aed1a0286cb4ff5489cf0581d05f033a30a9b62afc748a9bf0b91
-  data.tar.gz: 4a4a764def0528487ee20adb5e1edce52d7a3d267f1c99e399f5f48a7a5d9e2b
+  metadata.gz: 8db3543c735b87abb79d2de2d8278f11c3999289f2c68929a6cc7c855568b6a1
+  data.tar.gz: ff4783fde9226bff71d59e00960997109a5426f407572ea41859657a120e907f
 SHA512:
-  metadata.gz: 5a1e967d0eb2358f4aef3ac6dbb5c7a182984673330be8bc3141badd99e60e43af3c325f5a02de635e7cb9a329b498f650305bb83c46d0aa54e6ba219f7fd84a
-  data.tar.gz: da2ec1f4fd3d69e76f3d6a812aec49454ff9d99b99006e1b53f0086e4a60d1958d69b6a1b7447e033fcbf927c412cc3fea50954d0c3cd72e5ce49adf4140d95f
+  metadata.gz: 492022d310c9dc8b59e823e550769bdabbbd6af100d503ca82421c4fc3cef3aaea041f667c21d96d223b19fd5441c5b85fbd515a44f4ccb0a6820b471adcd28c
+  data.tar.gz: 92678a59b82c31a055c3e28bfe40b60a6c830368c9707d00ccb2481c43213127e8210d7547d2eff32f75d636c48632b8ef44056e9c0692841cc983b983c5297b

data/CHANGELOG.md

@@ -1,5 +1,63 @@
 # Changelog
 
+## 0.2.0 (2026-10-31)
+
+**This release changes training results.** Runs that use `gradient_accumulation_steps > 1` now follow the intended learning-rate schedule. All runs now apply `seed:`, so data order is different from 0.1.x and reproducible. Expect different (better) final models than 0.1.x with the same arguments.
+
+### Changed
+
+- **A step is now one optimizer update** (#12). `global_step`, `max_steps`, `logging_steps`, `eval_steps`, `save_steps`, `warmup_steps` and the scheduler all count optimizer updates, not batches. The scheduler is sized as `ceil(batches_per_epoch / gradient_accumulation_steps) * num_train_epochs`, so it now decays to the end. Previously, with accumulation 4, a linear schedule ended near 75% of the peak learning rate.
+  - The last, shorter accumulation window of an epoch now makes an optimizer update, with its loss averaged over the batches it holds. Previously its gradients leaked into the next epoch.
+  - `on_step_begin` fires at the start of each accumulation window and `on_step_end` after each optimizer update. Nothing changes with `gradient_accumulation_steps: 1`.
+- **`seed:` is applied** (#13). `train` seeds Torch, Ruby's RNG and its own shuffling. `seed: nil` disables this. New `Trainers.set_seed(seed)` seeds model initialisation.
+- **Progress goes to `Trainers.logger`** on `$stderr`, not stdout (#15). Set `Trainers.logger = nil` for a silent library. The default callback is now `Trainers::LoggingCallback`; `PrinterCallback` still prints to stdout if you add it to `callbacks:`.
+- **Errors descend from `Trainers::Error`** (#16): `ConfigurationError`, `AdapterError` and `CheckpointError`. Code that rescued `RuntimeError` for these must rescue `Trainers::Error` instead.
+- `SaveUtils.save_lora_adapters` and `LoraModel.save_adapters` raise `AdapterError` when the model has no adapters, instead of printing and returning `nil` (#16).
+- **Requires Ruby 3.3 or newer.** Ruby 3.1 and 3.2 are end-of-life.
+- New runtime dependency: `logger`, which Ruby 3.5+ no longer ships as a default gem.
+
+### Fixed
+
+- `dataloader_drop_last: true` skips the final partial batch in training, evaluation and prediction (#14). `train` raises `ConfigurationError` if that leaves no batches.
+- `label_names:` is honoured (#14). Every listed key is passed to the model, and the first is used for the built-in cross-entropy loss and `evaluate`.
+- `LoraModel.apply` works on models whose layers sit in a `ModuleList`, which covers every HuggingFace transformer, for example DistilBERT's `transformer.layer.0.attention.q_lin`. It used to raise `Could not find module part '0'`, because `ModuleList` stores its children in a Hash keyed by `"0"`, not in an Array.
+- When an epoch's last optimizer step was already logged, the end-of-epoch log no longer repeats it.
+- Learning rates and other floats below 0.001 are logged in scientific notation (`2.000e-05`), not as `0.0000`.
+
+### Added
+
+- `TrainingArguments` validates `seed` (nil or a non-negative Integer) and `label_names` (an Array of Strings or Symbols).
+
+## 0.1.2 (2026-10-30)
+
+### Fixed
+
+- `LoraModel.merge` no longer applies adapters twice. Merged layers kept adding the LoRA path on top of the merged weights, so a merged model's outputs were wrong. Each `LoraLinear` is now replaced by a plain `Torch::NN::Linear` holding the merged weights, so a saved merged model has the base model's parameter names. `LoraLinear#merge!` is idempotent and adds `#merged?` and `#to_linear`.
+- `Trainers::VERSION` now matches the gem version. The gemspec reads it from `lib/trainers/version.rb`. (#2)
+
+### Changed
+
+- `TrainingArguments` validates its values at construction, and raises `ArgumentError` naming the argument, the value received and what was expected (#8):
+  - positive: `learning_rate`
+  - positive Integer: `num_train_epochs`, both batch sizes, `gradient_accumulation_steps`, `logging_steps`, `save_steps`
+  - non-negative: `weight_decay`, `max_grad_norm`; `warmup_steps` must also be an Integer
+  - `warmup_ratio` must be within 0..1, and `eval_steps` must be nil or a positive Integer
+  - `lr_scheduler_type` must be one of `:linear`, `:cosine`, `:constant`; `eval_strategy` and `save_strategy` one of `:no`, `:epoch`, `:steps`
+
+  Previously these were accepted and failed later, or were silently ignored. Defaults are unchanged.
+
+### Internal
+
+- CI on Ruby 3.3, 3.4 and 4.0 against torch-rb 0.17 (LibTorch 2.4.0) and 0.27 (LibTorch 2.15.0), using Appraisal gemfiles, plus macOS arm64; every job is required (#9)
+- Tests for LR schedulers and warmup (#6), callback order and arguments (#7), LoRA shapes, identity init, merge and adapter round-trip (#5), and a public API surface test (#10)
+- Development dependency minitest bumped to ~> 6.0
+- CONTRIBUTING.md with LibTorch setup, SECURITY.md, README badges and a known-limitations section (#11)
+
+## 0.1.1 (2026-06-09)
+
+- Fix gemspec homepage, source and changelog URLs to point to github.com/Selectus2/trainers-rb
+- No code changes. The published gem reports `Trainers::VERSION` as "0.1.0"; the constant is corrected here
+
 ## 0.1.0 (2026-06-07)
 
 - Initial release

data/README.md

@@ -1,5 +1,10 @@
 # trainers-rb
 
+[![Gem Version](https://img.shields.io/gem/v/trainers-rb)](https://rubygems.org/gems/trainers-rb)
+[![CI](https://github.com/Selectus2/trainers-rb/actions/workflows/ci.yml/badge.svg)](https://github.com/Selectus2/trainers-rb/actions/workflows/ci.yml)
+[![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE.txt)
+[![Ruby](https://img.shields.io/badge/ruby-3.3%20%7C%203.4%20%7C%204.0-red.svg)](.github/workflows/ci.yml)
+
 Fine-tune transformer models in Ruby.
 
 trainers-rb provides a training loop, LoRA (Low-Rank Adaptation), learning rate scheduling, and model serialization for HuggingFace transformer models loaded via [transformers-rb](https://github.com/ankane/transformers-ruby). It builds on [torch-rb](https://github.com/ankane/torch.rb) for autograd, optimizers, and tensor operations.
@@ -22,15 +27,16 @@ gem install trainers-rb
 
 ### Prerequisites
 
-trainers-rb depends on torch-rb, which requires LibTorch:
+trainers-rb depends on torch-rb, which compiles against LibTorch. Use the LibTorch version that matches your torch-rb: 2.15.0 for torch-rb 0.27, or 2.4.0 for torch-rb 0.17.
 
 ```bash
 # macOS arm64
-curl -L -o /tmp/libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-2.4.0.zip
-unzip /tmp/libtorch.zip -d ~/libtorch
-bundle config set build.torch-rb --with-torch-dir=$HOME/libtorch/libtorch
+curl -L -o /tmp/libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-2.15.0.zip
+unzip /tmp/libtorch.zip -d ~/libtorch-2.15.0
+bundle config set build.torch-rb --with-torch-dir=$HOME/libtorch-2.15.0/libtorch
 ```
 
+For Linux, and for running the tests, see [CONTRIBUTING.md](CONTRIBUTING.md).
 ## Quick Start
 
 ```ruby
@@ -128,7 +134,7 @@ Trainers::LoraModel.load_adapters(model, "./lora-adapters")
 | `learning_rate` | `5e-5` | Peak learning rate for AdamW |
 | `weight_decay` | `0.0` | Weight decay (applied to non-bias, non-norm params) |
 | `max_grad_norm` | `1.0` | Max gradient norm for clipping |
-| `gradient_accumulation_steps` | `1` | Accumulate gradients over N steps |
+| `gradient_accumulation_steps` | `1` | Accumulate gradients over N batches per optimizer step |
 | `warmup_steps` | `0` | Linear warmup steps |
 | `warmup_ratio` | `0.0` | Warmup as fraction of total steps (alternative to warmup_steps) |
 | `lr_scheduler_type` | `:linear` | `:linear`, `:cosine`, or `:constant` |
@@ -136,10 +142,60 @@ Trainers::LoraModel.load_adapters(model, "./lora-adapters")
 | `eval_steps` | `nil` | Evaluate every N steps (when `eval_strategy: :steps`) |
 | `save_strategy` | `:epoch` | When to save: `:no`, `:epoch`, or `:steps` |
 | `save_total_limit` | `nil` | Keep only the last N checkpoints |
+| `save_steps` | `500` | Save a checkpoint every N steps (when `save_strategy: :steps`) |
 | `logging_steps` | `500` | Log every N steps |
-| `seed` | `42` | Random seed |
+| `seed` | `42` | Seeds Torch, Ruby's RNG and data shuffling at the start of `train`; `nil` disables seeding |
+| `dataloader_drop_last` | `false` | Skip the final batch when it is smaller than the batch size |
+| `label_names` | `["labels"]` | Batch keys that hold labels; the first is used for the built-in cross-entropy loss |
 | `no_mps` | `false` | Force CPU even if MPS is available |
 
+**A step is one optimizer update.** With `gradient_accumulation_steps: 4`, one step covers four batches, and `max_steps`, `global_step`, `logging_steps`, `eval_steps`, `save_steps`, `warmup_steps` and the learning-rate schedule all count optimizer updates. When an epoch's batch count is not a multiple of the accumulation steps, the last, shorter window still makes an update, with its loss averaged over the batches it contains, so no gradient carries into the next epoch.
+
+Invalid values, such as a negative learning rate, a zero batch size or an unknown `lr_scheduler_type`, raise `ArgumentError` when `TrainingArguments` is created.
+
+## Reproducibility
+
+`Trainer#train` seeds Torch, Ruby's global RNG and its own shuffling from `args.seed`. The model's initial weights are created before training starts, so seed those yourself:
+
+```ruby
+Trainers.set_seed(42)
+model, tokenizer = Trainers.from_pretrained("distilbert-base-uncased", num_labels: 2)
+```
+
+Two runs with the same seed, data and arguments produce identical loss curves on CPU. On Apple Silicon (MPS), some kernels are not deterministic, so expect small differences between runs.
+
+## Logging
+
+trainers-rb reports progress through `Trainers.logger`, which writes to `$stderr` at `INFO` by default. Nothing is written to stdout.
+
+```ruby
+Trainers.logger = Rails.logger                  # use your app's logger
+Trainers.logger.level = Logger::WARN            # only warnings and errors
+Trainers.logger = nil                           # completely silent
+```
+
+To print progress to stdout instead, add `Trainers::PrinterCallback.new` to `callbacks:`.
+
+## Errors
+
+Every error trainers-rb raises descends from `Trainers::Error`:
+
+| Error | Raised when |
+|---|---|
+| `Trainers::ConfigurationError` | The model has no trainable parameters, returns no loss and no labels are present, or the training data yields no batches |
+| `Trainers::AdapterError` | No module matches `target_modules`, an adapter file is missing, or `save_adapters` finds no adapters |
+| `Trainers::CheckpointError` | Reserved for checkpoint save and load errors |
+
+```ruby
+begin
+  trainer.train
+rescue Trainers::Error => e
+  # anything trainers-rb raised
+end
+```
+
+Invalid `TrainingArguments` values raise `ArgumentError`.
+
 ## Callbacks
 
 Built-in callbacks:
@@ -274,6 +330,14 @@ trainers-rb
 
 trainers-rb adds the training layer that transformers-rb intentionally omits. Both gems call into the same LibTorch C++ kernels for the actual computation.
 
+## Known limitations
+
+These are open correctness issues. They are tracked, and listed here so nobody discovers them the hard way:
+
+- **A saved model cannot be loaded back.** No `config.json` is written. ([#17](https://github.com/Selectus2/trainers-rb/issues/17))
+- **Checkpoints cannot resume training.** Optimizer, scheduler and trainer state are not saved. ([#18](https://github.com/Selectus2/trainers-rb/issues/18))
+- **LoRA checkpoints save the full state dict** rather than just the adapters. ([#19](https://github.com/Selectus2/trainers-rb/issues/19))
+
 ## Roadmap
 
 - [ ] More model architectures in transformers-rb (GPT-2, Llama for text generation)
@@ -286,7 +350,7 @@ trainers-rb adds the training layer that transformers-rb intentionally omits. Bo
 
 ## Contributing
 
-Bug reports and pull requests are welcome on GitHub.
+Bug reports and pull requests are welcome. [CONTRIBUTING.md](CONTRIBUTING.md) covers LibTorch setup, running the tests against each supported torch-rb, and what a good PR looks like. Report security issues privately, as described in [SECURITY.md](SECURITY.md).
 
 ## License
 

data/lib/trainers/callbacks.rb

@@ -44,32 +44,48 @@ module Trainers
     def on_save(args, state, control, **kwargs); end
   end
 
-  # Default callback: prints training progress to stdout
+  # Prints training progress to stdout. Opt in with callbacks: [PrinterCallback.new]
   class PrinterCallback < TrainerCallback
     def on_log(args, state, control, logs: nil, **kwargs)
       return unless logs
       output = logs.map { |k, v| "#{k}: #{format_value(v)}" }.join("  ")
-      puts "[step #{state.global_step}] #{output}"
+      emit "[step #{state.global_step}] #{output}"
     end
 
     def on_train_begin(args, state, control, **kwargs)
-      puts "Starting training: #{state.num_train_epochs} epochs, #{state.max_steps} total steps"
+      emit "Starting training: #{state.num_train_epochs} epochs, #{state.max_steps} total steps"
     end
 
     def on_train_end(args, state, control, **kwargs)
-      puts "Training complete. Total steps: #{state.global_step}"
+      emit "Training complete. Total steps: #{state.global_step}"
     end
 
     def on_evaluate(args, state, control, metrics: nil, **kwargs)
       return unless metrics
       output = metrics.map { |k, v| "#{k}: #{format_value(v)}" }.join("  ")
-      puts "[eval step #{state.global_step}] #{output}"
+      emit "[eval step #{state.global_step}] #{output}"
     end
 
     private
 
+    def emit(message)
+      puts message
+    end
+
+    # Small floats such as learning rates (2e-5) would print as 0.0000
     def format_value(v)
-      v.is_a?(Float) ? format("%.4f", v) : v.to_s
+      return v.to_s unless v.is_a?(Float)
+
+      v.nonzero? && v.abs < 1e-3 ? format("%.3e", v) : format("%.4f", v)
+    end
+  end
+
+  # Default callback: reports the same progress through Trainers.logger
+  class LoggingCallback < PrinterCallback
+    private
+
+    def emit(message)
+      Trainers.log(:info, message)
     end
   end
 
@@ -95,7 +111,7 @@ module Trainers
       else
         @wait_count += 1
         if @wait_count >= @patience
-          puts "Early stopping triggered after #{@wait_count} evaluations without improvement"
+          Trainers.log(:info, "Early stopping triggered after #{@wait_count} evaluations without improvement")
           control.should_training_stop = true
         end
       end

data/lib/trainers/errors.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Trainers
+  # Base class for every error trainers-rb raises. Rescue this to catch them all.
+  class Error < StandardError; end
+
+  # The model, data or arguments cannot be trained as configured
+  class ConfigurationError < Error; end
+
+  # LoRA adapters are missing, cannot be applied, or cannot be loaded
+  class AdapterError < Error; end
+
+  # A checkpoint cannot be written or read
+  class CheckpointError < Error; end
+end

data/lib/trainers/logging.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require "logger"
+
+module Trainers
+  class << self
+    # Where trainers-rb reports progress. Defaults to $stderr at INFO.
+    # Set to any Logger-compatible object, or nil for a silent library.
+    attr_accessor :logger
+
+    def default_logger
+      Logger.new($stderr, level: Logger::INFO, progname: "trainers").tap do |logger|
+        logger.formatter = proc do |severity, _time, progname, message|
+          prefix = severity == "INFO" ? "" : "#{severity}: "
+          "[#{progname}] #{prefix}#{message}\n"
+        end
+      end
+    end
+
+    # Internal: log through Trainers.logger when one is set
+    def log(level, message)
+      logger&.public_send(level, message)
+    end
+  end
+
+  self.logger = default_logger
+end

data/lib/trainers/lora/lora_linear.rb

@@ -36,12 +36,20 @@ module Trainers
       @lora_B = Torch::NN::Parameter.new(Torch.zeros(@out_features, @r))
 
       @lora_dropout = lora_dropout > 0 ? Torch::NN::Dropout.new(p: lora_dropout) : nil
+      @merged       = false
+    end
+
+    def merged?
+      @merged
     end
 
     def forward(x)
       # Original linear
       base_output = Torch::NN::F.linear(x, @weight, @bias)
 
+      # After merge! the adapter is already folded into @weight
+      return base_output if @merged
+
       # LoRA path
       lora_input = @lora_dropout ? @lora_dropout.call(x) : x
       lora_output = lora_input.matmul(@lora_A.t).matmul(@lora_B.t) * @scaling
@@ -49,14 +57,30 @@ module Trainers
       base_output + lora_output
     end
 
-    # Merge LoRA weights into the base weight matrix (for inference)
+    # Merge LoRA weights into the base weight matrix (for inference).
+    # Idempotent: merging twice would add the adapter twice.
     def merge!
+      return self if @merged
+
       Torch.no_grad do
         @weight.add!(@lora_B.matmul(@lora_A) * @scaling)
       end
+      @merged = true
       self
     end
 
+    # A plain Linear holding the merged weights, so a merged model has the
+    # same parameter names as the base model
+    def to_linear
+      merge!
+      linear = Torch::NN::Linear.new(@in_features, @out_features, bias: !@bias.nil?)
+      Torch.no_grad do
+        linear.weight.copy!(@weight)
+        linear.bias.copy!(@bias) if @bias
+      end
+      linear
+    end
+
     # Extract LoRA adapter state for saving
     def lora_state_dict
       { "lora_A" => @lora_A.data, "lora_B" => @lora_B.data }

data/lib/trainers/lora/lora_model.rb

@@ -20,7 +20,7 @@ module Trainers
       targets = LoraUtils.find_target_modules(model, config.target_modules)
 
       if targets.empty?
-        raise "No Linear modules found matching target_modules: #{config.target_modules.inspect}. " \
+        raise AdapterError, "No Linear modules found matching target_modules: #{config.target_modules.inspect}. " \
               "Available modules: #{model.named_modules.map(&:first).join(', ')}"
       end
 
@@ -50,22 +50,20 @@ module Trainers
       end
       # :none — biases stay frozen (default)
 
-      puts "LoRA applied to #{targets.size} modules: #{targets.keys.join(', ')}"
+      Trainers.log(:info, "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)
+    # Merge all LoRA weights back into base weights (for inference).
+    # Each LoraLinear is replaced by a plain Linear with the merged weights.
     def self.merge(model)
-      count = 0
-      model.named_modules.each do |_, mod|
-        if mod.is_a?(LoraLinear)
-          mod.merge!
-          count += 1
-        end
+      lora_modules = model.named_modules.select { |_, mod| mod.is_a?(LoraLinear) }
+      lora_modules.each do |name, mod|
+        LoraUtils.replace_module(model, name, mod.to_linear)
       end
-      puts "Merged #{count} LoRA adapters into base model"
+      Trainers.log(:info, "Merged #{lora_modules.size} LoRA adapters into base model")
       model
     end
 

data/lib/trainers/lora/lora_utils.rb

@@ -30,24 +30,28 @@ module Trainers
 
       # 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
+        parent = child_module(parent, part)
+        raise AdapterError, "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
+      registered = parent.instance_variable_get(:@modules)
+      if registered&.key?(child_name)
+        registered[child_name] = new_module
       else
         parent.instance_variable_set(:"@#{child_name}", new_module)
       end
     end
 
+    # ModuleList and add_module children live in @modules, a Hash keyed by
+    # name ("0", "1", ...); everything else is an instance variable.
+    def self.child_module(parent, name)
+      registered = parent.instance_variable_get(:@modules)
+      return registered[name] if registered&.key?(name)
+
+      parent.instance_variable_get(:"@#{name}") if name.match?(/\A[a-z_]\w*\z/i)
+    end
+
     # Count total and trainable parameters
     def self.print_trainable_parameters(model)
       total    = 0
@@ -59,9 +63,9 @@ module Trainers
       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)}%"
+      Trainers.log(:info, "trainable params: #{format_number(trainable)} || " \
+                          "all params: #{format_number(total)} || " \
+                          "trainable%: #{format('%.4f', pct)}%")
 
       { total: total, trainable: trainable, percentage: pct }
     end

data/lib/trainers/optimization/optimizer.rb

@@ -35,7 +35,7 @@ module Trainers
       end
 
       if param_groups.empty?
-        raise "No trainable parameters found. Did you forget to unfreeze the model or apply LoRA?"
+        raise ConfigurationError, "No trainable parameters found. Did you forget to unfreeze the model or apply LoRA?"
       end
 
       Torch::Optim::AdamW.new(

data/lib/trainers/save_utils.rb

@@ -18,7 +18,7 @@ module Trainers
       model_path = File.join(output_dir, "model.safetensors")
       require "safetensors"
       Safetensors::Torch.save_file(tensors, model_path)
-      puts "Model saved to #{model_path}"
+      Trainers.log(:info, "Model saved to #{model_path}")
 
       # Save tokenizer
       if tokenizer
@@ -49,21 +49,20 @@ module Trainers
       end
 
       if lora_state.empty?
-        puts "No LoRA adapters found in model"
-        return
+        raise AdapterError, "No LoRA adapters found in model. Apply them with LoraModel.apply before saving."
       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)"
+      Trainers.log(:info, "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)
+      raise AdapterError, "Adapter file not found: #{adapter_path}" unless File.exist?(adapter_path)
 
       require "safetensors"
       lora_state = Safetensors::Torch.load_file(adapter_path)
@@ -78,7 +77,7 @@ module Trainers
         end
       end
 
-      puts "LoRA adapters loaded from #{adapter_path}"
+      Trainers.log(:info, "LoRA adapters loaded from #{adapter_path}")
     end
   end
 end

data/lib/trainers/trainer.rb

@@ -25,25 +25,35 @@ module Trainers
       @state           = TrainerState.new
       @control         = TrainerControl.new
 
-      all_callbacks = [PrinterCallback.new] + callbacks
+      all_callbacks = [LoggingCallback.new] + callbacks
       @callback_handler = CallbackHandler.new(all_callbacks)
     end
 
     def train
       device = @args.resolved_device
+      apply_seed
       @model.to(device)
       @model.train
 
-      num_examples   = @train_dataset.size
-      batch_size     = @args.per_device_train_batch_size
-      steps_per_epoch = (num_examples.to_f / batch_size).ceil
-      total_steps    = steps_per_epoch * @args.num_train_epochs
+      batch_size        = @args.per_device_train_batch_size
+      accum             = @args.gradient_accumulation_steps
+      batches_per_epoch = num_batches(@train_dataset, batch_size)
+      if batches_per_epoch.zero?
+        raise ConfigurationError, "train_dataset yields no batches: #{@train_dataset.size} examples, " \
+                                  "batch size #{batch_size}, dataloader_drop_last #{@args.dataloader_drop_last}"
+      end
+
+      # One step is one optimizer update. The last accumulation window of an
+      # epoch may be shorter than gradient_accumulation_steps; it still steps.
+      steps_per_epoch = (batches_per_epoch.to_f / accum).ceil
+      total_steps     = steps_per_epoch * @args.num_train_epochs
 
-      @state.max_steps       = total_steps
+      @state.max_steps        = total_steps
       @state.num_train_epochs = @args.num_train_epochs
 
       @optimizer    = create_optimizer
       @lr_scheduler = create_scheduler(total_steps)
+      @optimizer.zero_grad
 
       @callback_handler.fire(:on_train_begin, @args, @state, @control)
 
@@ -52,38 +62,38 @@ module Trainers
         @callback_handler.fire(:on_epoch_begin, @args, @state, @control)
         @model.train
 
-        epoch_loss   = 0.0
-        epoch_steps  = 0
+        epoch_loss    = 0.0
+        epoch_batches = 0
+
+        each_batch(@train_dataset, batch_size, shuffle: true).with_index do |batch, index|
+          position     = index % accum
+          window_start = index - position
+          window_size  = [accum, batches_per_epoch - window_start].min
 
-        each_batch(@train_dataset, batch_size, shuffle: true) do |batch|
-          @callback_handler.fire(:on_step_begin, @args, @state, @control)
+          @callback_handler.fire(:on_step_begin, @args, @state, @control) if position.zero?
 
           batch = move_to_device(batch, device)
           loss  = compute_loss(batch)
 
-          scaled_loss = if @args.gradient_accumulation_steps > 1
-                          loss / @args.gradient_accumulation_steps
-                        else
-                          loss
-                        end
-
+          # Average the gradient over the micro-batches in this window
+          scaled_loss = window_size > 1 ? loss / window_size : loss
           scaled_loss.backward
 
-          epoch_loss  += loss.item
-          epoch_steps += 1
-          @state.global_step += 1
+          epoch_loss    += loss.item
+          epoch_batches += 1
 
-          if @state.global_step % @args.gradient_accumulation_steps == 0
-            clip_grad_norm!(@model.parameters, @args.max_grad_norm)
-            @optimizer.step
-            @lr_scheduler.step
-            @optimizer.zero_grad
-          end
+          next unless position == window_size - 1
+
+          clip_grad_norm!(@model.parameters, @args.max_grad_norm)
+          @optimizer.step
+          @lr_scheduler.step
+          @optimizer.zero_grad
+          @state.global_step += 1
 
           # Logging
           if should_log?
             logs = {
-              loss:          epoch_loss / epoch_steps,
+              loss:          epoch_loss / epoch_batches,
               learning_rate: current_lr,
               epoch:         @state.epoch
             }
@@ -109,11 +119,13 @@ module Trainers
           break if @control.should_training_stop || @control.should_epoch_stop
         end
 
-        # Epoch-level logging
-        epoch_avg_loss = epoch_steps > 0 ? epoch_loss / epoch_steps : 0.0
-        logs = { loss: epoch_avg_loss, learning_rate: current_lr, epoch: @state.epoch }
-        @state.log_history << logs.merge(step: @state.global_step)
-        @callback_handler.fire(:on_log, @args, @state, @control, logs: logs)
+        # Epoch-level logging, unless the last step of the epoch already logged
+        unless @state.log_history.last&.fetch(:step, nil) == @state.global_step
+          epoch_avg_loss = epoch_batches > 0 ? epoch_loss / epoch_batches : 0.0
+          logs = { loss: epoch_avg_loss, learning_rate: current_lr, epoch: @state.epoch }
+          @state.log_history << logs.merge(step: @state.global_step)
+          @callback_handler.fire(:on_log, @args, @state, @control, logs: logs)
+        end
 
         # Epoch-based evaluation
         if @args.eval_strategy == :epoch && @eval_dataset
@@ -151,7 +163,7 @@ module Trainers
       Torch.no_grad do
         each_batch(dataset, @args.per_device_eval_batch_size) do |batch|
           batch  = move_to_device(batch, device)
-          labels = batch.delete(:labels) || batch.delete("labels")
+          labels = pop_labels(batch).values.first
 
           output = forward(batch)
 
@@ -209,11 +221,12 @@ module Trainers
     private
 
     def compute_loss(batch)
-      labels = batch.delete(:labels) || batch.delete("labels")
+      label_values = pop_labels(batch)
+      labels       = label_values.values.first
 
       # Try passing labels to the model (some models compute loss internally)
       output = begin
-        forward(labels ? batch.merge(labels: labels) : batch)
+        forward(batch.merge(label_values))
       rescue => e
         # If the model doesn't support labels kwarg (e.g. transformers-rb Todo),
         # fall back to forward without labels + external loss
@@ -225,7 +238,7 @@ module Trainers
       end
 
       # Restore labels to batch for downstream use
-      batch[:labels] = labels if labels
+      batch.merge!(label_values)
 
       if output.respond_to?(:loss) && output.loss
         output.loss
@@ -233,7 +246,7 @@ module Trainers
         logits = output.respond_to?(:logits) ? output.logits : output
         Torch::NN::F.cross_entropy(logits, labels)
       else
-        raise "Model did not return a loss and no labels found in batch. " \
+        raise ConfigurationError, "Model did not return a loss and no labels found in batch. " \
               "Either pass labels in your dataset or use a model that computes loss."
       end
     end
@@ -268,20 +281,40 @@ module Trainers
       )
     end
 
+    # Removes every key named in args.label_names (symbol or string) from the
+    # batch and returns them as { name_symbol => tensor }, in label_names order
+    def pop_labels(batch)
+      @args.label_names.each_with_object({}) do |name, labels|
+        value = batch.key?(name.to_sym) ? batch.delete(name.to_sym) : batch.delete(name.to_s)
+        labels[name.to_sym] = value unless value.nil?
+      end
+    end
+
+    def num_batches(dataset, batch_size)
+      full, rest = dataset.size.divmod(batch_size)
+      rest.positive? && !@args.dataloader_drop_last ? full + 1 : full
+    end
+
     def each_batch(dataset, batch_size, shuffle: false)
-      indices = (0...dataset.size).to_a
-      indices.shuffle! if shuffle
+      return enum_for(:each_batch, dataset, batch_size, shuffle: shuffle) unless block_given?
 
-      (0...dataset.size).step(batch_size) do |start|
-        batch_indices = indices[start, batch_size]
-        next if batch_indices.nil? || batch_indices.empty?
+      indices = (0...dataset.size).to_a
+      indices.shuffle!(random: @rng || Random.new) if shuffle
 
-        features = batch_indices.map { |i| dataset[i] }
-        batch = @data_collator.call(features)
-        yield batch
+      num_batches(dataset, batch_size).times do |i|
+        batch_indices = indices[i * batch_size, batch_size]
+        features = batch_indices.map { |j| dataset[j] }
+        yield @data_collator.call(features)
       end
     end
 
+    # Seeds Torch, Kernel#rand and the shuffle RNG from args.seed (nil skips)
+    def apply_seed
+      seed = @args.seed
+      @rng = seed.nil? ? Random.new : Random.new(seed)
+      Trainers.set_seed(seed) unless seed.nil?
+    end
+
     def move_to_device(batch, device)
       batch.each_with_object({}) do |(key, value), result|
         result[key] = if value.is_a?(Torch::Tensor)

data/lib/trainers/training_arguments.rb

@@ -31,18 +31,28 @@ module Trainers
       label_names:                   ["labels"]
     }.freeze
 
+    SCHEDULER_TYPES = %i[linear cosine constant].freeze
+    STRATEGIES      = %i[no epoch steps].freeze
+
+    POSITIVE_NUMBERS  = %i[learning_rate].freeze
+    POSITIVE_INTEGERS = %i[num_train_epochs per_device_train_batch_size per_device_eval_batch_size
+                           gradient_accumulation_steps logging_steps save_steps].freeze
+    NON_NEGATIVE      = %i[weight_decay max_grad_norm].freeze
+
     DEFAULTS.each_key do |key|
       attr_accessor key
     end
 
     def initialize(**kwargs)
+      unknown = kwargs.keys - DEFAULTS.keys
+      raise ArgumentError, "Unknown arguments: #{unknown.join(', ')}" unless unknown.empty?
+
       DEFAULTS.each do |key, default|
         value = kwargs.fetch(key, default)
         instance_variable_set(:"@#{key}", value)
       end
 
-      unknown = kwargs.keys - DEFAULTS.keys
-      raise ArgumentError, "Unknown arguments: #{unknown.join(', ')}" unless unknown.empty?
+      validate!
     end
 
     def resolved_device
@@ -60,5 +70,37 @@ module Trainers
         hash[key] = send(key)
       end
     end
+
+    private
+
+    # Rejects values that cannot produce a sensible run. Errors name the
+    # argument, the value received, and what was expected.
+    def validate!
+      POSITIVE_NUMBERS.each do |key|
+        check(key, "a positive number") { |v| v.is_a?(Numeric) && v.positive? }
+      end
+      POSITIVE_INTEGERS.each do |key|
+        check(key, "a positive Integer") { |v| v.is_a?(Integer) && v.positive? }
+      end
+      NON_NEGATIVE.each do |key|
+        check(key, "a non-negative number") { |v| v.is_a?(Numeric) && !v.negative? }
+      end
+      check(:warmup_steps, "a non-negative Integer") { |v| v.is_a?(Integer) && !v.negative? }
+      check(:warmup_ratio, "a number between 0 and 1") { |v| v.is_a?(Numeric) && v.between?(0, 1) }
+      check(:eval_steps, "nil or a positive Integer") { |v| v.nil? || (v.is_a?(Integer) && v.positive?) }
+      check(:seed, "nil or a non-negative Integer") { |v| v.nil? || (v.is_a?(Integer) && !v.negative?) }
+      check(:label_names, "an Array of Strings or Symbols") { |v| v.is_a?(Array) && v.all? { |n| n.is_a?(String) || n.is_a?(Symbol) } }
+
+      check(:lr_scheduler_type, "one of #{SCHEDULER_TYPES.map(&:inspect).join(', ')}") { |v| SCHEDULER_TYPES.include?(v) }
+      check(:eval_strategy, "one of #{STRATEGIES.map(&:inspect).join(', ')}") { |v| STRATEGIES.include?(v) }
+      check(:save_strategy, "one of #{STRATEGIES.map(&:inspect).join(', ')}") { |v| STRATEGIES.include?(v) }
+    end
+
+    def check(key, expected)
+      value = send(key)
+      return if yield(value)
+
+      raise ArgumentError, "#{key} must be #{expected}, got #{value.inspect}"
+    end
   end
 end

data/lib/trainers/version.rb

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

data/lib/trainers.rb

@@ -5,6 +5,8 @@ require "json"
 require "fileutils"
 
 require_relative "trainers/version"
+require_relative "trainers/errors"
+require_relative "trainers/logging"
 require_relative "trainers/trainer_utils"
 require_relative "trainers/training_arguments"
 require_relative "trainers/data/dataset"
@@ -20,6 +22,14 @@ require_relative "trainers/lora/lora_model"
 require_relative "trainers/trainer"
 
 module Trainers
+  # Seeds Torch and Ruby's global RNG. Call before building a model to make
+  # its initial weights reproducible; Trainer#train also calls it with args.seed.
+  def self.set_seed(seed)
+    Torch.manual_seed(seed)
+    srand(seed)
+    seed
+  end
+
   # Convenience method: load model + tokenizer and prepare for training
   def self.from_pretrained(model_name, task: :sequence_classification, num_labels: 2)
     require "transformers-rb"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: trainers-rb
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.2.0
 platform: ruby
 authors:
 - Vishwajeetsingh Desurkar
@@ -65,6 +65,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: 0.5.3
+- !ruby/object:Gem::Dependency
+  name: logger
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.5'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.5'
 - !ruby/object:Gem::Dependency
   name: rake
   requirement: !ruby/object:Gem::Requirement
@@ -85,20 +99,20 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '5.0'
+        version: '6.0'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '5.0'
+        version: '6.0'
 description: Training loop, LoRA, and optimization utilities for fine-tuning HuggingFace
   transformer models using torch-rb and transformers-rb. Supports full fine-tuning,
   LoRA adapters, learning rate scheduling, callbacks, and model serialization via
   safetensors.
 email:
-- selectus2@users.noreply.rubygems.org
+- vishwajeetsinghd@gmail.com
 executables: []
 extensions: []
 extra_rdoc_files: []
@@ -111,6 +125,8 @@ files:
 - lib/trainers/callbacks.rb
 - lib/trainers/data/data_collator.rb
 - lib/trainers/data/dataset.rb
+- lib/trainers/errors.rb
+- lib/trainers/logging.rb
 - lib/trainers/lora/lora_config.rb
 - lib/trainers/lora/lora_linear.rb
 - lib/trainers/lora/lora_model.rb
@@ -129,6 +145,8 @@ metadata:
   homepage_uri: https://github.com/Selectus2/trainers-rb
   source_code_uri: https://github.com/Selectus2/trainers-rb
   changelog_uri: https://github.com/Selectus2/trainers-rb/blob/main/CHANGELOG.md
+  bug_tracker_uri: https://github.com/Selectus2/trainers-rb/issues
+  rubygems_mfa_required: 'true'
 rdoc_options: []
 require_paths:
 - lib
@@ -136,14 +154,14 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: 3.1.0
+      version: 3.3.0
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.6.7
+rubygems_version: 4.0.20
 specification_version: 4
 summary: Fine-tune transformer models in Ruby
 test_files: []