ignis-dl 0.0.1

data/lib/nnw/ai/transformer/modern.rb

@@ -0,0 +1,191 @@
+# frozen_string_literal: true
+
+module Ignis
+  module AI
+    module Transformer
+      # Modern decoder components: RoPE + GQA attention, and a Llama/Qwen/SmolLM-style
+      # block (pre-RMSNorm, RoPE-GQA attention, SwiGLU MLP) and model (no learned
+      # position embedding — RoPE supplies position). These target Qwen3, SmolLM3,
+      # Phi, Llama-family architectures.
+
+      # Compute the RoPE inverse-frequency table [head_dim/2], optionally applying
+      # a scaling scheme. Supports the "llama3" rope_type (Llama-3.1/3.2): low
+      # frequencies are divided by `factor`, high frequencies are kept, and a smooth
+      # interpolation bridges the two — matching HF's _compute_llama3_parameters.
+      # @param head_dim [Integer]
+      # @param base [Float] rope_theta
+      # @param scaling [Hash, nil] e.g. {rope_type:"llama3", factor:, low_freq_factor:,
+      #   high_freq_factor:, original_max_position_embeddings:} (symbol or string keys)
+      # @return [Array<Float>] length head_dim/2
+      def self.compute_inv_freq(head_dim, base, scaling = nil)
+        half = head_dim / 2
+        freqs = (0...half).map { |i| base.to_f**(-2.0 * i / head_dim) }
+        return freqs unless scaling
+        g = ->(k) { scaling[k] || scaling[k.to_s] }
+        return freqs unless (g.call(:rope_type) || g.call(:type)) == "llama3"
+
+        factor   = g.call(:factor).to_f
+        low_ff   = g.call(:low_freq_factor).to_f
+        high_ff  = g.call(:high_freq_factor).to_f
+        old_ctx  = g.call(:original_max_position_embeddings).to_f
+        low_wl   = old_ctx / low_ff
+        high_wl  = old_ctx / high_ff
+        freqs.map do |f|
+          wl = 2.0 * Math::PI / f
+          if wl > low_wl
+            f / factor
+          elsif wl < high_wl
+            f
+          else
+            smooth = (old_ctx / wl - low_ff) / (high_ff - low_ff)
+            (1.0 - smooth) * (f / factor) + smooth * f
+          end
+        end
+      end
+
+      # Attention with rotary embeddings (RoPE) and grouped-query attention (GQA).
+      # No bias on projections (Llama/Qwen convention).
+      class RopeGqaAttention < NN::Module
+        # @param embed_dim [Integer]
+        # @param num_heads [Integer] query heads
+        # @param num_kv_heads [Integer] key/value heads (== num_heads ⇒ plain MHA)
+        # @param head_dim [Integer, nil] per-head dim (default embed_dim/num_heads)
+        # @param rope_base [Float] RoPE theta
+        # @param bias [Boolean]
+        # @param device_id [Integer]
+        def initialize(embed_dim, num_heads, num_kv_heads:, head_dim: nil,
+                       rope_base: 10000.0, rope_scaling: nil, bias: false, device_id: 0)
+          super()
+          raise ArgumentError, "num_heads must be a multiple of num_kv_heads" unless (num_heads % num_kv_heads).zero?
+          @embed_dim = embed_dim
+          @num_heads = num_heads
+          @num_kv_heads = num_kv_heads
+          @head_dim = head_dim || (embed_dim / num_heads)
+          # Fail early (at construction) rather than silently miscompute later:
+          # RoPE needs an even head_dim; the flash kernels cap head_dim at 128.
+          raise ArgumentError, "head_dim must be even for RoPE (got #{@head_dim})" unless @head_dim.even?
+          raise ArgumentError, "head_dim #{@head_dim} exceeds flash-attention HEAD_DIM_MAX (128)" if @head_dim > 128
+          @rope_base = rope_base
+          # Precompute the (optionally scaled) inv_freq table once; reused every layer/step.
+          @inv_freq = Transformer.compute_inv_freq(@head_dim, rope_base, rope_scaling)
+          q_out  = num_heads * @head_dim
+          kv_out = num_kv_heads * @head_dim
+          @q_proj = register_module("q_proj", NN::Linear.new(embed_dim, q_out, bias: bias, device_id: device_id))
+          @k_proj = register_module("k_proj", NN::Linear.new(embed_dim, kv_out, bias: bias, device_id: device_id))
+          @v_proj = register_module("v_proj", NN::Linear.new(embed_dim, kv_out, bias: bias, device_id: device_id))
+          @o_proj = register_module("o_proj", NN::Linear.new(q_out, embed_dim, bias: bias, device_id: device_id))
+        end
+
+        # @param x [Tensor] [seq, embed_dim]
+        # @param pos_offset [Integer] absolute position of row 0 (for KV-cache decode)
+        # @return [Tensor] [seq, embed_dim]
+        def forward(x, pos_offset: 0)
+          # RoPE is applied to Q and K (not V); q has num_heads, k has num_kv_heads.
+          q = @q_proj.call(x).rope(num_heads: @num_heads, pos_offset: pos_offset, inv_freq: @inv_freq)
+          k = @k_proj.call(x).rope(num_heads: @num_kv_heads, pos_offset: pos_offset, inv_freq: @inv_freq)
+          v = @v_proj.call(x)
+          ctx = q.sdpa(k, v, num_heads: @num_heads, num_kv_heads: @num_kv_heads, causal: true)
+          @o_proj.call(ctx)
+        end
+
+        # @return [String]
+        def to_s
+          "RopeGqaAttention(embed=#{@embed_dim}, q_heads=#{@num_heads}, kv_heads=#{@num_kv_heads}, head_dim=#{@head_dim})"
+        end
+      end
+
+      # Llama/Qwen-style block: x += attn(rmsnorm(x)); x += swiglu(rmsnorm(x)).
+      class ModernBlock < NN::Module
+        # @param embed_dim [Integer]
+        # @param num_heads [Integer]
+        # @param num_kv_heads [Integer]
+        # @param ff_dim [Integer]
+        # @param rope_base [Float]
+        # @param eps [Float] RMSNorm epsilon
+        # @param device_id [Integer]
+        def initialize(embed_dim, num_heads, num_kv_heads:, ff_dim:,
+                       rope_base: 10000.0, rope_scaling: nil, head_dim: nil, eps: 1e-6, device_id: 0)
+          super()
+          @attn  = register_module("attn",
+                    RopeGqaAttention.new(embed_dim, num_heads, num_kv_heads: num_kv_heads,
+                                         head_dim: head_dim, rope_base: rope_base,
+                                         rope_scaling: rope_scaling, device_id: device_id))
+          @mlp   = register_module("mlp", SwiGLU.new(embed_dim, ff_dim, device_id: device_id))
+          @norm1 = register_module("norm1", NN::RMSNorm.new(embed_dim, eps: eps, device_id: device_id))
+          @norm2 = register_module("norm2", NN::RMSNorm.new(embed_dim, eps: eps, device_id: device_id))
+        end
+
+        # @param x [Tensor] [seq, embed]
+        # @return [Tensor]
+        def forward(x)
+          x = x + @attn.call(@norm1.call(x))
+          x + @mlp.call(@norm2.call(x))
+        end
+
+        # @return [String]
+        def to_s
+          "ModernBlock(#{@attn}, #{@mlp})"
+        end
+      end
+
+      # Llama/Qwen/SmolLM-style decoder LM: token embedding → N ModernBlocks →
+      # final RMSNorm → LM head. No learned position embedding — RoPE supplies it.
+      class ModernModel < NN::Module
+        # @return [Integer]
+        attr_reader :vocab_size, :embed_dim, :num_heads, :num_kv_heads, :num_layers, :max_seq_len
+
+        # @param vocab_size [Integer]
+        # @param embed_dim [Integer]
+        # @param num_heads [Integer]
+        # @param num_kv_heads [Integer]
+        # @param num_layers [Integer]
+        # @param ff_dim [Integer]
+        # @param max_seq_len [Integer]
+        # @param rope_base [Float]
+        # @param eps [Float]
+        # @param device_id [Integer]
+        def initialize(vocab_size:, embed_dim:, num_heads:, num_kv_heads:, num_layers:,
+                       ff_dim:, max_seq_len:, rope_base: 10000.0, rope_scaling: nil,
+                       head_dim: nil, eps: 1e-6, device_id: 0)
+          super()
+          @vocab_size = vocab_size
+          @embed_dim = embed_dim
+          @num_heads = num_heads
+          @num_kv_heads = num_kv_heads
+          @num_layers = num_layers
+          @max_seq_len = max_seq_len
+          @device_id = device_id
+          @head_dim = head_dim || (embed_dim / num_heads)
+
+          @token_embedding = register_module("token_embedding",
+                              NN::Embedding.new(vocab_size, embed_dim, device_id: device_id))
+          @blocks = []
+          num_layers.times do |i|
+            blk = ModernBlock.new(embed_dim, num_heads, num_kv_heads: num_kv_heads,
+                                  ff_dim: ff_dim, rope_base: rope_base, rope_scaling: rope_scaling,
+                                  head_dim: head_dim, eps: eps, device_id: device_id)
+            @blocks << register_module("blocks.#{i}", blk)
+          end
+          @norm = register_module("norm", NN::RMSNorm.new(embed_dim, eps: eps, device_id: device_id))
+          @head = register_module("head", NN::Linear.new(embed_dim, vocab_size, bias: false, device_id: device_id))
+        end
+
+        # @param input_ids [Tensor] token indices [seq] (int32)
+        # @param mask [Tensor, nil] unused (attention is causal)
+        # @return [Tensor] logits [seq, vocab]
+        def forward(input_ids, mask: nil)
+          x = @token_embedding.call(input_ids) # [seq, embed]; RoPE (in attn) supplies position
+          @blocks.each { |block| x = block.call(x) }
+          x = @norm.call(x)
+          @head.call(x)
+        end
+
+        # @return [String]
+        def to_s
+          "ModernModel(vocab=#{@vocab_size}, embed=#{@embed_dim}, q_heads=#{@num_heads}, " \
+          "kv_heads=#{@num_kv_heads}, layers=#{@num_layers})"
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/transformer/swiglu.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Ignis
+  module AI
+    module Transformer
+      # SwiGLU feed-forward (Llama / Qwen / SmolLM / Mistral):
+      #   down( silu(gate(x)) ⊙ up(x) )
+      # Two input projections (gate, up) into the hidden dim and one output (down)
+      # back to embed. silu(z) = z·sigmoid(z). Llama-style uses no bias. This is a
+      # pure composition of verified ops (Linear, silu, elementwise mul), so autograd
+      # produces the backward automatically.
+      class SwiGLU < NN::Module
+        # @param embed_dim [Integer]
+        # @param ff_dim [Integer] hidden dim (Llama uses ~8/3·embed, rounded)
+        # @param bias [Boolean] include projection biases (Llama/Qwen: false)
+        # @param device_id [Integer]
+        def initialize(embed_dim, ff_dim, bias: false, device_id: 0)
+          super()
+          @embed_dim = embed_dim
+          @ff_dim = ff_dim
+          @gate = register_module("gate", NN::Linear.new(embed_dim, ff_dim, bias: bias, device_id: device_id))
+          @up   = register_module("up",   NN::Linear.new(embed_dim, ff_dim, bias: bias, device_id: device_id))
+          @down = register_module("down", NN::Linear.new(ff_dim, embed_dim, bias: bias, device_id: device_id))
+        end
+
+        # @param x [Tensor] [*, embed_dim]
+        # @return [Tensor] [*, embed_dim]
+        def forward(x)
+          @down.call(@gate.call(x).silu * @up.call(x))
+        end
+
+        # @return [String]
+        def to_s
+          "SwiGLU(embed=#{@embed_dim}, ff=#{@ff_dim})"
+        end
+      end
+    end
+  end
+end

data/lib/nnw/ai/weight_map.rb

@@ -0,0 +1,139 @@
+# frozen_string_literal: true
+
+module Ignis
+  module AI
+    # Weight mapping from HuggingFace model naming conventions to Ignis parameter names.
+    # Used by Safetensors.load_model to map HF checkpoint keys to Ignis module keys.
+    module WeightMap
+      class << self
+        # Get the weight map for a given model architecture.
+        # @param architecture [Symbol] :gpt2, :llama, :bert
+        # @return [Hash{String => String}] HF key → Ignis key
+        def for(architecture)
+          case architecture
+          when :gpt2    then gpt2_map
+          when :llama   then llama_map
+          when :bert    then bert_map
+          else
+            raise ArgumentError, "Unknown architecture: #{architecture}. Supported: :gpt2, :llama, :bert"
+          end
+        end
+
+        private
+
+        # GPT-2 weight mapping: HuggingFace → Ignis
+        # HF naming: transformer.h.{i}.attn.c_attn.weight
+        # Ignis naming: blocks.{i}.attention.q_proj.weight (split)
+        # @return [Hash{String => String}]
+        def gpt2_map
+          map = {}
+
+          # Embeddings
+          map["wte.weight"] = "token_embedding.weight"
+          map["wpe.weight"] = "position_embedding.weight"
+
+          # Final layer norm
+          map["ln_f.weight"] = "norm.weight"
+          map["ln_f.bias"] = "norm.bias"
+
+          # LM head (GPT-2 ties wte and lm_head)
+          # map["lm_head.weight"] = "head.weight"  # Usually tied to wte
+
+          # Per-layer mappings
+          12.times do |i|
+            hf_prefix = "h.#{i}"
+            nnw_prefix = "blocks.#{i}"
+
+            # Attention
+            map["#{hf_prefix}.attn.c_attn.weight"] = "#{nnw_prefix}.attention._c_attn_weight"
+            map["#{hf_prefix}.attn.c_attn.bias"] = "#{nnw_prefix}.attention._c_attn_bias"
+            map["#{hf_prefix}.attn.c_proj.weight"] = "#{nnw_prefix}.attention.out_proj.weight"
+            map["#{hf_prefix}.attn.c_proj.bias"] = "#{nnw_prefix}.attention.out_proj.bias"
+
+            # Layer norms
+            map["#{hf_prefix}.ln_1.weight"] = "#{nnw_prefix}.norm1.weight"
+            map["#{hf_prefix}.ln_1.bias"] = "#{nnw_prefix}.norm1.bias"
+            map["#{hf_prefix}.ln_2.weight"] = "#{nnw_prefix}.norm2.weight"
+            map["#{hf_prefix}.ln_2.bias"] = "#{nnw_prefix}.norm2.bias"
+
+            # MLP / Feed-forward
+            map["#{hf_prefix}.mlp.c_fc.weight"] = "#{nnw_prefix}.feed_forward.fc1.weight"
+            map["#{hf_prefix}.mlp.c_fc.bias"] = "#{nnw_prefix}.feed_forward.fc1.bias"
+            map["#{hf_prefix}.mlp.c_proj.weight"] = "#{nnw_prefix}.feed_forward.fc2.weight"
+            map["#{hf_prefix}.mlp.c_proj.bias"] = "#{nnw_prefix}.feed_forward.fc2.bias"
+          end
+
+          map
+        end
+
+        # LLaMA weight mapping: HuggingFace → Ignis
+        # @return [Hash{String => String}]
+        def llama_map
+          map = {}
+
+          map["model.embed_tokens.weight"] = "token_embedding.weight"
+          map["model.norm.weight"] = "norm.weight"
+          map["lm_head.weight"] = "head.weight"
+
+          32.times do |i|
+            hf_prefix = "model.layers.#{i}"
+            nnw_prefix = "blocks.#{i}"
+
+            # Self-attention
+            map["#{hf_prefix}.self_attn.q_proj.weight"] = "#{nnw_prefix}.attention.q_proj.weight"
+            map["#{hf_prefix}.self_attn.k_proj.weight"] = "#{nnw_prefix}.attention.k_proj.weight"
+            map["#{hf_prefix}.self_attn.v_proj.weight"] = "#{nnw_prefix}.attention.v_proj.weight"
+            map["#{hf_prefix}.self_attn.o_proj.weight"] = "#{nnw_prefix}.attention.out_proj.weight"
+
+            # RMS norms
+            map["#{hf_prefix}.input_layernorm.weight"] = "#{nnw_prefix}.norm1.weight"
+            map["#{hf_prefix}.post_attention_layernorm.weight"] = "#{nnw_prefix}.norm2.weight"
+
+            # MLP (SwiGLU)
+            map["#{hf_prefix}.mlp.gate_proj.weight"] = "#{nnw_prefix}.feed_forward.fc1.weight"
+            map["#{hf_prefix}.mlp.up_proj.weight"] = "#{nnw_prefix}.feed_forward._up_proj_weight"
+            map["#{hf_prefix}.mlp.down_proj.weight"] = "#{nnw_prefix}.feed_forward.fc2.weight"
+          end
+
+          map
+        end
+
+        # BERT weight mapping: HuggingFace → Ignis
+        # @return [Hash{String => String}]
+        def bert_map
+          map = {}
+
+          map["embeddings.word_embeddings.weight"] = "token_embedding.weight"
+          map["embeddings.position_embeddings.weight"] = "position_embedding.weight"
+          map["embeddings.LayerNorm.weight"] = "embedding_norm.weight"
+          map["embeddings.LayerNorm.bias"] = "embedding_norm.bias"
+
+          12.times do |i|
+            hf_prefix = "encoder.layer.#{i}"
+            nnw_prefix = "blocks.#{i}"
+
+            map["#{hf_prefix}.attention.self.query.weight"] = "#{nnw_prefix}.attention.q_proj.weight"
+            map["#{hf_prefix}.attention.self.query.bias"] = "#{nnw_prefix}.attention.q_proj.bias"
+            map["#{hf_prefix}.attention.self.key.weight"] = "#{nnw_prefix}.attention.k_proj.weight"
+            map["#{hf_prefix}.attention.self.key.bias"] = "#{nnw_prefix}.attention.k_proj.bias"
+            map["#{hf_prefix}.attention.self.value.weight"] = "#{nnw_prefix}.attention.v_proj.weight"
+            map["#{hf_prefix}.attention.self.value.bias"] = "#{nnw_prefix}.attention.v_proj.bias"
+            map["#{hf_prefix}.attention.output.dense.weight"] = "#{nnw_prefix}.attention.out_proj.weight"
+            map["#{hf_prefix}.attention.output.dense.bias"] = "#{nnw_prefix}.attention.out_proj.bias"
+            map["#{hf_prefix}.attention.output.LayerNorm.weight"] = "#{nnw_prefix}.norm1.weight"
+            map["#{hf_prefix}.attention.output.LayerNorm.bias"] = "#{nnw_prefix}.norm1.bias"
+
+            map["#{hf_prefix}.intermediate.dense.weight"] = "#{nnw_prefix}.feed_forward.fc1.weight"
+            map["#{hf_prefix}.intermediate.dense.bias"] = "#{nnw_prefix}.feed_forward.fc1.bias"
+            map["#{hf_prefix}.output.dense.weight"] = "#{nnw_prefix}.feed_forward.fc2.weight"
+            map["#{hf_prefix}.output.dense.bias"] = "#{nnw_prefix}.feed_forward.fc2.bias"
+            map["#{hf_prefix}.output.LayerNorm.weight"] = "#{nnw_prefix}.norm2.weight"
+            map["#{hf_prefix}.output.LayerNorm.bias"] = "#{nnw_prefix}.norm2.bias"
+          end
+
+          map
+        end
+      end
+    end
+  end
+end

metadata

@@ -0,0 +1,91 @@
+--- !ruby/object:Gem::Specification
+name: ignis-dl
+version: !ruby/object:Gem::Version
+  version: 0.0.1
+platform: ruby
+authors:
+- NNW / Ignis contributors
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: ignis-autograd
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - '='
+      - !ruby/object:Gem::Version
+        version: 0.0.1
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - '='
+      - !ruby/object:Gem::Version
+        version: 0.0.1
+description: |
+  ignis-dl is the deep-learning layer of the Ignis ecosystem: NN modules
+  (Linear, Embedding, LayerNorm, RMSNorm, Dropout), optimizers (SGD/Adam/AdamW),
+  losses, and a transformer stack (multi-head + grouped-query attention, RoPE,
+  SwiGLU, KV cache) with HuggingFace weight loaders (GPT-2, Llama). Loads real
+  GPT-2 and Llama-3.2 checkpoints and matches HuggingFace logits, and trains
+  transformers from scratch — in Ruby, on native Windows. Installing this pulls
+  the whole stack (ignis + ignis-autograd), so it also serves as the meta-gem.
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- README.md
+- lib/ignis-dl.rb
+- lib/nnw/ai/gpt2_loader.rb
+- lib/nnw/ai/inference.rb
+- lib/nnw/ai/kv_cache.rb
+- lib/nnw/ai/llama_loader.rb
+- lib/nnw/ai/loss.rb
+- lib/nnw/ai/nn/dropout.rb
+- lib/nnw/ai/nn/embedding.rb
+- lib/nnw/ai/nn/layer_norm.rb
+- lib/nnw/ai/nn/linear.rb
+- lib/nnw/ai/nn/module.rb
+- lib/nnw/ai/nn/rms_norm.rb
+- lib/nnw/ai/nn/sequential.rb
+- lib/nnw/ai/optim/adam.rb
+- lib/nnw/ai/optim/adamw.rb
+- lib/nnw/ai/optim/base.rb
+- lib/nnw/ai/optim/lr_scheduler.rb
+- lib/nnw/ai/optim/sgd.rb
+- lib/nnw/ai/safetensors.rb
+- lib/nnw/ai/server.rb
+- lib/nnw/ai/tokenizer.rb
+- lib/nnw/ai/trainer.rb
+- lib/nnw/ai/transformer/attention.rb
+- lib/nnw/ai/transformer/block.rb
+- lib/nnw/ai/transformer/feed_forward.rb
+- lib/nnw/ai/transformer/model.rb
+- lib/nnw/ai/transformer/modern.rb
+- lib/nnw/ai/transformer/swiglu.rb
+- lib/nnw/ai/weight_map.rb
+homepage: https://github.com/tigel-agm/Ignis
+licenses:
+- MIT
+metadata:
+  source_code_uri: https://github.com/tigel-agm/Ignis
+  rubygems_mfa_required: 'true'
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '3.1'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.6.9
+specification_version: 4
+summary: Transformers + neural-net layers for Ruby, on the Ignis GPU/autograd stack
+test_files: []