ignis-dl 0.0.1

data/lib/nnw/ai/nn/sequential.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require_relative "module"
+
+module Ignis
+  module AI
+    module NN
+      # Sequential container: chains modules in order.
+      # forward(x) passes x through each module sequentially.
+      class Sequential < Module
+        # @param modules [Array<Module>] modules to chain
+        def initialize(*modules)
+          super()
+          modules.each_with_index do |mod, i|
+            register_module(i.to_s, mod)
+          end
+          @layers = modules
+        end
+
+        # Forward pass: reduce input through all modules.
+        # @param x [Tensor]
+        # @return [Tensor]
+        def forward(x)
+          @layers.reduce(x) { |h, m| m.call(h) }
+        end
+
+        # Number of layers.
+        # @return [Integer]
+        def length
+          @layers.length
+        end
+
+        # Access layer by index.
+        # @param idx [Integer]
+        # @return [Module]
+        def [](idx)
+          @layers[idx]
+        end
+
+        # @return [String]
+        def to_s
+          lines = ["Sequential("]
+          @layers.each_with_index do |layer, i|
+            lines << "  (#{i}): #{layer}"
+          end
+          lines << ")"
+          lines.join("\n")
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/optim/adam.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require_relative "base"
+
+module Ignis
+  module AI
+    module Optim
+      # Adam optimizer (Kingma & Ba, 2014).
+      # Fused single-kernel per parameter: updates m, v, and param in one launch.
+      class Adam < Base
+        # @param params [Array<Tensor>]
+        # @param lr [Float] learning rate
+        # @param beta1 [Float] exponential decay rate for first moment
+        # @param beta2 [Float] exponential decay rate for second moment
+        # @param eps [Float] numerical stability
+        # @param weight_decay [Float] L2 regularization
+        def initialize(params, lr: 1e-3, beta1: 0.9, beta2: 0.999,
+                       eps: 1e-8, weight_decay: 0.0)
+          super(params, lr: lr)
+          @beta1 = beta1.to_f
+          @beta2 = beta2.to_f
+          @eps = eps.to_f
+          @weight_decay = weight_decay.to_f
+
+          # Initialize first and second moment estimates to zero
+          @m_states = {}
+          @v_states = {}
+          params.each do |p|
+            m = Ignis::Shared::NvArray.new(shape: p.shape, dtype: p.dtype, device_id: p.device_id)
+            m.from_host(Array.new(p.numel, 0.0))
+            @m_states[p.object_id] = m
+
+            v = Ignis::Shared::NvArray.new(shape: p.shape, dtype: p.dtype, device_id: p.device_id)
+            v.from_host(Array.new(p.numel, 0.0))
+            @v_states[p.object_id] = v
+          end
+        end
+
+        # @return [void]
+        def step
+          @step_count += 1
+          kernel = Ignis::JIT::Kernels::Optimizer.adam_step
+
+          bias_correction1 = 1.0 - @beta1**@step_count
+          bias_correction2 = 1.0 - @beta2**@step_count
+
+          @params.each do |p|
+            next unless p.grad
+            n = p.numel
+            m = @m_states[p.object_id]
+            v = @v_states[p.object_id]
+
+            kernel.launch(grid: [(n + 255) / 256], block: [256],
+                          args: [p.data, p.grad, m, v,
+                                 @lr, @beta1, @beta2, @eps, @weight_decay,
+                                 bias_correction1.to_f, bias_correction2.to_f, n])
+          end
+          Ignis.synchronize
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/optim/adamw.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require_relative "base"
+
+module Ignis
+  module AI
+    module Optim
+      # AdamW optimizer (Loshchilov & Hutter, 2019).
+      # Key difference from Adam: weight decay is *decoupled* from gradient update.
+      # This means WD is applied directly to parameters, not through the gradient.
+      class AdamW < Base
+        # @param params [Array<Tensor>]
+        # @param lr [Float]
+        # @param beta1 [Float]
+        # @param beta2 [Float]
+        # @param eps [Float]
+        # @param weight_decay [Float] decoupled weight decay
+        def initialize(params, lr: 1e-3, beta1: 0.9, beta2: 0.999,
+                       eps: 1e-8, weight_decay: 0.01)
+          super(params, lr: lr)
+          @beta1 = beta1.to_f
+          @beta2 = beta2.to_f
+          @eps = eps.to_f
+          @weight_decay = weight_decay.to_f
+
+          @m_states = {}
+          @v_states = {}
+          params.each do |p|
+            m = Ignis::Shared::NvArray.new(shape: p.shape, dtype: p.dtype, device_id: p.device_id)
+            m.from_host(Array.new(p.numel, 0.0))
+            @m_states[p.object_id] = m
+
+            v = Ignis::Shared::NvArray.new(shape: p.shape, dtype: p.dtype, device_id: p.device_id)
+            v.from_host(Array.new(p.numel, 0.0))
+            @v_states[p.object_id] = v
+          end
+        end
+
+        # @return [void]
+        def step
+          @step_count += 1
+          kernel = Ignis::JIT::Kernels::Optimizer.adamw_step
+
+          bias_correction1 = 1.0 - @beta1**@step_count
+          bias_correction2 = 1.0 - @beta2**@step_count
+
+          @params.each do |p|
+            next unless p.grad
+            n = p.numel
+            m = @m_states[p.object_id]
+            v = @v_states[p.object_id]
+
+            kernel.launch(grid: [(n + 255) / 256], block: [256],
+                          args: [p.data, p.grad, m, v,
+                                 @lr, @beta1, @beta2, @eps, @weight_decay,
+                                 bias_correction1.to_f, bias_correction2.to_f, n])
+          end
+          Ignis.synchronize
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/optim/base.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+module Ignis
+  module AI
+    module Optim
+      # Base optimizer class.
+      # Provides shared infrastructure for all optimizers.
+      class Base
+        # @return [Array<Tensor>] parameters being optimized
+        attr_reader :params
+
+        # @return [Integer] current step count for LR scheduling
+        attr_reader :step_count
+
+        # @param params [Array<Tensor>] parameters to optimize
+        # @param lr [Float] learning rate
+        def initialize(params, lr:)
+          @params = params
+          @lr = lr.to_f
+          @step_count = 0
+        end
+
+        # Perform one optimization step using accumulated gradients.
+        # @return [void]
+        def step
+          raise NotImplementedError, "Subclass must implement #step"
+        end
+
+        # Zero all parameter gradients.
+        # @return [void]
+        def zero_grad!
+          @params.each(&:zero_grad!)
+        end
+
+        # Get current learning rate.
+        # @return [Float]
+        def lr
+          @lr
+        end
+
+        # Set learning rate (for schedulers).
+        # @param new_lr [Float]
+        # @return [void]
+        def lr=(new_lr)
+          @lr = new_lr.to_f
+        end
+
+        # Clip gradient global norm across all parameters.
+        # @param max_norm [Float]
+        # @return [Float] the actual total norm (before clipping)
+        def clip_grad_norm!(max_norm)
+          # Phase 1: compute total squared sum across all params
+          total_norm_nv = Ignis::Shared::NvArray.new(shape: [1], dtype: :float32,
+                                                    device_id: @params.first.device_id)
+          total_norm_nv.from_host([0.0])
+
+          @params.each do |p|
+            next unless p.grad
+
+            sq_kernel = Ignis::JIT::Kernels::Optimizer.grad_squared_sum
+            n = p.numel
+            block_size = 256
+            grid_size = (n + block_size - 1) / block_size
+            sq_kernel.launch(grid: [grid_size], block: [block_size],
+                             shared_mem: block_size * 4,
+                             args: [p.grad, total_norm_nv, n])
+          end
+          Ignis.synchronize
+
+          total_norm = Math.sqrt(total_norm_nv.to_host[0])
+
+          if total_norm > max_norm
+            clip_factor = max_norm / (total_norm + 1e-6)
+            clip_kernel = Ignis::JIT::Kernels::Optimizer.grad_clip_scale
+
+            @params.each do |p|
+              next unless p.grad
+              n = p.numel
+              clip_kernel.launch(grid: [(n + 255) / 256], block: [256],
+                                 args: [p.grad, clip_factor.to_f, n])
+            end
+            Ignis.synchronize
+          end
+
+          total_norm
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/optim/lr_scheduler.rb

@@ -0,0 +1,118 @@
+# frozen_string_literal: true
+
+module Ignis
+  module AI
+    module Optim
+      # Learning rate schedulers.
+      # Adjust the optimizer's learning rate based on training progress.
+      module LRScheduler
+        # Linear warmup: linearly increase LR from 0 to base_lr over warmup_steps.
+        # After warmup, keep constant.
+        class LinearWarmup
+          # @param optimizer [Base]
+          # @param warmup_steps [Integer]
+          def initialize(optimizer, warmup_steps:)
+            @optimizer = optimizer
+            @warmup_steps = warmup_steps
+            @base_lr = optimizer.lr
+          end
+
+          # Call after each optimizer step to update LR.
+          # @return [Float] new learning rate
+          def step
+            current = @optimizer.step_count
+            if current < @warmup_steps
+              factor = current.to_f / @warmup_steps
+              new_lr = @base_lr * factor
+              @optimizer.lr = new_lr
+              new_lr
+            else
+              @optimizer.lr = @base_lr
+              @base_lr
+            end
+          end
+
+          # @return [Float]
+          def current_lr
+            @optimizer.lr
+          end
+        end
+
+        # Cosine annealing: decay LR following a cosine curve.
+        # LR drops from base_lr to min_lr over total_steps.
+        class CosineAnnealing
+          # @param optimizer [Base]
+          # @param total_steps [Integer]
+          # @param min_lr [Float]
+          def initialize(optimizer, total_steps:, min_lr: 0.0)
+            @optimizer = optimizer
+            @total_steps = total_steps
+            @min_lr = min_lr.to_f
+            @base_lr = optimizer.lr
+          end
+
+          # @return [Float]
+          def step
+            current = @optimizer.step_count
+            if current >= @total_steps
+              @optimizer.lr = @min_lr
+              return @min_lr
+            end
+
+            progress = current.to_f / @total_steps
+            new_lr = @min_lr + 0.5 * (@base_lr - @min_lr) * (1.0 + Math.cos(Math::PI * progress))
+            @optimizer.lr = new_lr
+            new_lr
+          end
+
+          # @return [Float]
+          def current_lr
+            @optimizer.lr
+          end
+        end
+
+        # Cosine with linear warmup: common combo for Transformer training.
+        class CosineWithWarmup
+          # @param optimizer [Base]
+          # @param warmup_steps [Integer]
+          # @param total_steps [Integer]
+          # @param min_lr [Float]
+          def initialize(optimizer, warmup_steps:, total_steps:, min_lr: 0.0)
+            @optimizer = optimizer
+            @warmup_steps = warmup_steps
+            @total_steps = total_steps
+            @min_lr = min_lr.to_f
+            @base_lr = optimizer.lr
+          end
+
+          # @return [Float]
+          def step
+            current = @optimizer.step_count
+
+            if current < @warmup_steps
+              # Linear warmup phase
+              factor = current.to_f / @warmup_steps
+              new_lr = @base_lr * factor
+            elsif current >= @total_steps
+              new_lr = @min_lr
+            else
+              # Cosine decay phase
+              decay_steps = @total_steps - @warmup_steps
+              decay_current = current - @warmup_steps
+              progress = decay_current.to_f / decay_steps
+              new_lr = @min_lr + 0.5 * (@base_lr - @min_lr) * (1.0 + Math.cos(Math::PI * progress))
+            end
+
+            @optimizer.lr = new_lr
+            new_lr
+          end
+
+          # @return [Float]
+          def current_lr
+            @optimizer.lr
+          end
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/optim/sgd.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require_relative "base"
+
+module Ignis
+  module AI
+    module Optim
+      # SGD optimizer with optional momentum and weight decay.
+      class SGD < Base
+        # @param params [Array<Tensor>]
+        # @param lr [Float] learning rate
+        # @param momentum [Float] momentum factor (0.0 = no momentum)
+        # @param weight_decay [Float] L2 penalty
+        def initialize(params, lr:, momentum: 0.0, weight_decay: 0.0)
+          super(params, lr: lr)
+          @momentum = momentum.to_f
+          @weight_decay = weight_decay.to_f
+
+          @velocities = {}
+          if @momentum > 0.0
+            params.each do |p|
+              v = Ignis::Shared::NvArray.new(shape: p.shape, dtype: p.dtype, device_id: p.device_id)
+              v.from_host(Array.new(p.numel, 0.0))
+              @velocities[p.object_id] = v
+            end
+          end
+        end
+
+        # @return [void]
+        def step
+          kernel = Ignis::JIT::Kernels::Optimizer.sgd_step
+
+          @params.each do |p|
+            next unless p.grad
+            n = p.numel
+            vel = @velocities[p.object_id]
+            vel_ptr = vel || p.grad  # If no momentum, velocity unused but still passed
+
+            kernel.launch(grid: [(n + 255) / 256], block: [256],
+                          args: [p.data, p.grad, vel_ptr, @lr, @momentum, @weight_decay, n])
+          end
+          Ignis.synchronize
+
+          @step_count += 1
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/safetensors.rb

@@ -0,0 +1,220 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Ignis
+  module AI
+    # Safetensors — pure Ruby loader/saver for the HuggingFace safetensors format.
+    #
+    # Format:
+    #   8 bytes: LE uint64 header length
+    #   N bytes: JSON header {"tensor_name": {"dtype": "F16", "shape": [d1,d2], "data_offsets": [begin, end]}}
+    #   Remaining: raw tensor data (contiguous, little-endian, row-major)
+    #
+    # Zero-copy: uses IO.read + direct cudaMemcpy to GPU.
+    module Safetensors
+      # Dtype mapping: safetensors string → Ignis::Shared::NvArray symbol + Ruby unpack format
+      DTYPE_MAP = {
+        "F16"   => { symbol: :float16,  bytes: 2, pack: "v*" },
+        "BF16"  => { symbol: :bfloat16, bytes: 2, pack: "v*" },
+        "F32"   => { symbol: :float32,  bytes: 4, pack: "e*" },
+        "F64"   => { symbol: :float64,  bytes: 8, pack: "E*" },
+        "I8"    => { symbol: :uint8,    bytes: 1, pack: "c*" },
+        "I16"   => { symbol: :int32,    bytes: 2, pack: "s<*" },
+        "I32"   => { symbol: :int32,    bytes: 4, pack: "l<*" },
+        "I64"   => { symbol: :int64,    bytes: 8, pack: "q<*" },
+        "U8"    => { symbol: :uint8,    bytes: 1, pack: "C*" },
+        "BOOL"  => { symbol: :uint8,    bytes: 1, pack: "C*" }
+      }.freeze
+
+      # Maximum header size (100MB — DOS protection per spec)
+      MAX_HEADER_SIZE = 100 * 1024 * 1024
+
+      class << self
+        # Load all tensors from a safetensors file.
+        # @param path [String] path to .safetensors file
+        # @param device_id [Integer] target GPU device
+        # @return [Hash{String => Tensor}] name → Tensor mapping
+        def load(path, device_id: 0)
+          raise ArgumentError, "File not found: #{path}" unless File.exist?(path)
+
+          File.open(path, "rb") do |f|
+            header_size_bytes = f.read(8)
+            header_size = header_size_bytes.unpack1("Q<")
+
+            raise "Header size #{header_size} exceeds maximum #{MAX_HEADER_SIZE}" if header_size > MAX_HEADER_SIZE
+
+            header_json = f.read(header_size)
+            header = JSON.parse(header_json)
+
+            # Data section starts after 8 + header_size bytes
+            data_offset = 8 + header_size
+
+            tensors = {}
+
+            header.each do |name, meta|
+              next if name == "__metadata__"
+
+              dtype_str = meta["dtype"]
+              shape = meta["shape"]
+              offsets = meta["data_offsets"]
+
+              dtype_info = DTYPE_MAP[dtype_str]
+              raise "Unsupported dtype: #{dtype_str}" unless dtype_info
+
+              begin_offset = offsets[0]
+              end_offset = offsets[1]
+              byte_count = end_offset - begin_offset
+
+              # Read raw bytes from file
+              f.seek(data_offset + begin_offset)
+              raw_bytes = f.read(byte_count)
+              raise "Short read for #{name}: expected #{byte_count}, got #{raw_bytes&.length}" unless raw_bytes&.length == byte_count
+
+              # Convert to NvArray on GPU
+              nv_array = bytes_to_nv_array(raw_bytes, shape, dtype_info, device_id)
+              tensors[name] = Tensor.new(data: nv_array, requires_grad: false)
+            end
+
+            tensors
+          end
+        end
+
+        # Load tensors into a model using a weight map.
+        # @param model [NN::Module] the model to load into
+        # @param path [String] path to .safetensors file
+        # @param weight_map [Hash{String => String}, nil] HF name → Ignis name mapping
+        # @param strict [Boolean] fail on missing/extra keys
+        # @param device_id [Integer]
+        # @return [NN::Module]
+        def load_model(model, path, weight_map: nil, strict: false, device_id: 0)
+          tensors = load(path, device_id: device_id)
+          model_params = model.named_parameters
+
+          loaded_count = 0
+          skipped = []
+
+          tensors.each do |name, tensor|
+            mapped_name = weight_map ? (weight_map[name] || name) : name
+
+            if model_params.key?(mapped_name)
+              param = model_params[mapped_name]
+              if param.shape == tensor.shape
+                param.data.from_host(tensor.to_host)
+                loaded_count += 1
+              else
+                Ignis.logger.warn("Shape mismatch for #{mapped_name}: " \
+                                   "model=#{param.shape} file=#{tensor.shape}")
+                skipped << name
+              end
+            else
+              skipped << name
+            end
+          end
+
+          if strict
+            missing = model_params.keys - tensors.keys.map { |n| weight_map ? (weight_map[n] || n) : n }
+            unless missing.empty?
+              raise KeyError, "Missing weights: #{missing.join(', ')}"
+            end
+          end
+
+          Ignis.logger.info("Loaded #{loaded_count}/#{tensors.size} tensors, skipped #{skipped.size}")
+          model
+        end
+
+        # Save tensors to safetensors format.
+        # @param tensors [Hash{String => Tensor}] name → Tensor
+        # @param path [String] output file path
+        # @param metadata [Hash, nil] optional metadata
+        # @return [void]
+        def save(tensors, path, metadata: nil)
+          header = {}
+          header["__metadata__"] = metadata if metadata
+
+          # Calculate offsets
+          current_offset = 0
+          tensor_data = []
+
+          tensors.each do |name, tensor|
+            host_data = tensor.to_host
+            dtype_str = nv_dtype_to_safetensors(tensor.dtype)
+            bytes = host_values_to_bytes(host_data, tensor.dtype)
+            byte_count = bytes.length
+
+            header[name] = {
+              "dtype"        => dtype_str,
+              "shape"        => tensor.shape,
+              "data_offsets" => [current_offset, current_offset + byte_count]
+            }
+
+            tensor_data << bytes
+            current_offset += byte_count
+          end
+
+          header_json = JSON.generate(header)
+          # Pad header to 8-byte alignment
+          padding = (8 - (header_json.length % 8)) % 8
+          header_json += " " * padding
+
+          File.open(path, "wb") do |f|
+            f.write([header_json.length].pack("Q<"))
+            f.write(header_json)
+            tensor_data.each { |d| f.write(d) }
+          end
+        end
+
+        private
+
+        # Convert raw bytes to NvArray on GPU.
+        # @param raw_bytes [String]
+        # @param shape [Array<Integer>]
+        # @param dtype_info [Hash]
+        # @param device_id [Integer]
+        # @return [Ignis::Shared::NvArray]
+        def bytes_to_nv_array(raw_bytes, shape, dtype_info, device_id)
+          nv = Ignis::Shared::NvArray.new(shape: shape, dtype: dtype_info[:symbol], device_id: device_id)
+
+          # The safetensors payload is already stored in the tensor's native
+          # little-endian row-major layout — including IEEE-754 F16/BF16 — so we
+          # copy the raw bytes straight to the device. This is both correct
+          # (no lossy float<->half round trip) and zero-conversion.
+          nv.from_host_raw(raw_bytes)
+        end
+
+        # Convert Ignis dtype symbol to safetensors string.
+        # @param dtype [Symbol]
+        # @return [String]
+        def nv_dtype_to_safetensors(dtype)
+          case dtype
+          when :float16  then "F16"
+          when :bfloat16 then "BF16"
+          when :float32  then "F32"
+          when :float64  then "F64"
+          when :int32    then "I32"
+          when :int64    then "I64"
+          when :uint8    then "U8"
+          else raise "Unknown dtype: #{dtype}"
+          end
+        end
+
+        # Convert host values to raw bytes for saving.
+        # @param values [Array<Numeric>]
+        # @param dtype [Symbol]
+        # @return [String]
+        def host_values_to_bytes(values, dtype)
+          case dtype
+          when :float32  then values.pack("e*")
+          when :float64  then values.pack("E*")
+          when :int32    then values.pack("l<*")
+          when :int64    then values.pack("q<*")
+          when :uint8    then values.pack("C*")
+          when :float16  then values.map { |v| ::Ignis::Half.f32_to_f16(v) }.pack("v*")
+          when :bfloat16 then values.map { |v| ::Ignis::Half.f32_to_bf16(v) }.pack("v*")
+          else raise "Cannot save dtype: #{dtype}"
+          end
+        end
+      end
+    end
+  end
+end