mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/internlm2.rb

@@ -0,0 +1,160 @@
+module MlxLm
+  module Models
+    module InternLM2
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "internlm2"
+        field :hidden_size, default: 4096
+        field :num_hidden_layers, default: 32
+        field :num_attention_heads, default: 32
+        field :num_key_value_heads, default: nil
+        field :intermediate_size, default: 11008
+        field :vocab_size, default: 103168
+        field :rms_norm_eps, default: 1e-6
+        field :rope_theta, default: 10000.0
+        field :rope_traditional, default: false
+        field :rope_scaling, default: nil
+        field :bias, default: true
+        field :tie_word_embeddings, default: false
+        field :max_position_embeddings, default: 32768
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @head_dim = dim / @n_heads
+          @scale = @head_dim**(-0.5)
+
+          # Combined QKV projection
+          total_qkv = (@n_heads + 2 * @n_kv_heads) * @head_dim
+          self.wqkv = MLX::NN::Linear.new(dim, total_qkv, bias: args.bias)
+          self.wo = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: args.bias)
+
+          self.rope = MLX::NN::RoPE.new(
+            @head_dim,
+            traditional: args.rope_traditional,
+            base: args.rope_theta
+          )
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          qkv = wqkv.call(x)
+          q_size = @n_heads * @head_dim
+          kv_size = @n_kv_heads * @head_dim
+          queries, keys, values = mx.split(qkv, [q_size, q_size + kv_size], 2)
+
+          queries = queries.reshape([b, l, @n_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = keys.reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = values.reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @n_heads * @head_dim])
+          wo.call(output)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(dim, hidden_dim)
+          super()
+          self.w1 = MLX::NN::Linear.new(dim, hidden_dim, bias: false)
+          self.w2 = MLX::NN::Linear.new(hidden_dim, dim, bias: false)
+          self.w3 = MLX::NN::Linear.new(dim, hidden_dim, bias: false)
+        end
+
+        def call(x)
+          w2.call(MLX::NN.silu(w1.call(x)) * w3.call(x))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.attention = Attention.new(args)
+          self.feed_forward = MLP.new(args.hidden_size, args.intermediate_size)
+          self.attention_norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.ffn_norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          r = attention.call(attention_norm.call(x), mask: mask, cache: cache)
+          h = x + r
+          r = feed_forward.call(ffn_norm.call(h))
+          h + r
+        end
+      end
+
+      class InternLM2Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.tok_embeddings = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) { TransformerBlock.new(args) }
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil)
+          h = tok_embeddings.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+
+          mask = nil
+          mask = "causal" if h.shape[1] > 1
+
+          layers.each_with_index do |layer, i|
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model = InternLM2Model.new(args)
+          unless args.tie_word_embeddings
+            self.output = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          if @args.tie_word_embeddings
+            model.tok_embeddings.as_linear(out)
+          else
+            output.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          weights.reject { |k, _| k.include?("attention.rope.inv_freq") }
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("internlm2", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/internlm3.rb

@@ -0,0 +1,237 @@
+module MlxLm
+  module Models
+    module InternLM3
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "internlm3"
+        field :hidden_size, default: 4096
+        field :num_hidden_layers, default: 32
+        field :intermediate_size, default: 11008
+        field :num_attention_heads, default: 32
+        field :rms_norm_eps, default: 1e-6
+        field :vocab_size, default: 103168
+        field :bias, default: false
+        field :qkv_bias, default: false
+        field :max_position_embeddings, default: 32768
+        field :num_key_value_heads, default: nil
+        field :rope_theta, default: 10_000.0
+        field :rope_traditional, default: false
+        field :rope_scaling, default: nil
+        field :tie_word_embeddings, default: false
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+
+          return if @rope_scaling.nil?
+
+          required_keys = %w[factor rope_type]
+          missing = required_keys.reject { |k| _config_has_key?(k) }
+          unless missing.empty?
+            raise ArgumentError, "rope_scaling must contain keys #{required_keys}"
+          end
+
+          rope_type = _config_value("rope_type")
+          unless %w[linear dynamic].include?(rope_type)
+            raise ArgumentError, "rope_scaling 'rope_type' only supports 'linear' or 'dynamic'"
+          end
+        end
+
+        private
+
+        def _config_has_key?(key)
+          return false unless @rope_scaling.respond_to?(:key?)
+
+          @rope_scaling.key?(key) || @rope_scaling.key?(key.to_sym)
+        end
+
+        def _config_value(key, default = nil)
+          return default unless _config_has_key?(key)
+
+          if @rope_scaling.key?(key)
+            @rope_scaling[key]
+          else
+            @rope_scaling[key.to_sym]
+          end
+        end
+      end
+
+      class DynamicNTKScalingRoPE < MLX::NN::Module
+        def initialize(
+          dims,
+          max_position_embeddings: 2048,
+          traditional: false,
+          base: 10_000.0,
+          scale: 1.0
+        )
+          super()
+          @max_position_embeddings = max_position_embeddings
+          @original_base = base
+          @dims = dims
+          @traditional = traditional
+          @scale = scale
+        end
+
+        def call(x, offset: 0)
+          seq_len = x.shape[-2] + offset
+          if seq_len > @max_position_embeddings
+            scaled_ctx = (@scale * seq_len.to_f / @max_position_embeddings) - (@scale - 1.0)
+            base = @original_base * (scaled_ctx**(@dims.to_f / (@dims - 2)))
+          else
+            base = @original_base
+          end
+
+          MLX::Core.rope(x, @dims, @traditional, base, @scale, offset)
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          qkv_bias = args.qkv_bias
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @head_dim = args.hidden_size / @n_heads
+          @scale = @head_dim**(-0.5)
+
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: qkv_bias)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: qkv_bias)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: qkv_bias)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: qkv_bias)
+
+          rope_scale = if args.rope_scaling && _config_value(args.rope_scaling, "rope_type") == "linear"
+            1.0 / _config_value(args.rope_scaling, "factor").to_f
+          else
+            2.0
+          end
+
+          self.rope = DynamicNTKScalingRoPE.new(
+            @head_dim,
+            max_position_embeddings: args.max_position_embeddings,
+            traditional: args.rope_traditional,
+            base: args.rope_theta,
+            scale: rope_scale
+          )
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          queries = q_proj.call(x)
+          keys = k_proj.call(x)
+          values = v_proj.call(x)
+
+          queries = queries.reshape([b, l, @n_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = keys.reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = values.reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @n_heads * @head_dim])
+          o_proj.call(output)
+        end
+
+        private
+
+        def _config_value(config, key, default = nil)
+          return default if config.nil? || !config.respond_to?(:key?)
+          return config[key] if config.key?(key)
+
+          config.fetch(key.to_sym, default)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(dim, hidden_dim, bias)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: bias)
+          self.down_proj = MLX::NN::Linear.new(hidden_dim, dim, bias: bias)
+          self.up_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: bias)
+        end
+
+        def call(x)
+          down_proj.call(Activations.swiglu(gate_proj.call(x), up_proj.call(x)))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Attention.new(args)
+          self.mlp = MLP.new(args.hidden_size, args.intermediate_size, args.bias)
+          self.input_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_attention_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          r = self_attn.call(input_layernorm.call(x), mask: mask, cache: cache)
+          h = x + r
+          r = mlp.call(post_attention_layernorm.call(h))
+          h + r
+        end
+      end
+
+      class InternLM3Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) { TransformerBlock.new(args) }
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil)
+          h = embed_tokens.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+
+          mask = nil
+          mask = "causal" if h.shape[1] > 1
+
+          layers.each_with_index do |layer, i|
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model = InternLM3Model.new(args)
+          unless args.tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          weights.reject { |k, _| k.include?("attention.rope.inv_freq") }
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("internlm3", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/iquestloopcoder.rb

@@ -0,0 +1,261 @@
+require_relative "cache"
+require_relative "rope_utils"
+
+module MlxLm
+  module Models
+    module Iquestloopcoder
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "iquestloopcoder"
+        field :hidden_size
+        field :num_hidden_layers
+        field :intermediate_size
+        field :num_attention_heads
+        field :rms_norm_eps
+        field :vocab_size
+        field :head_dim
+        field :num_key_value_heads
+        field :max_position_embeddings, default: 131_072
+        field :attention_bias, default: false
+        field :mlp_bias, default: false
+        field :rope_theta, default: 500_000.0
+        field :rope_scaling, default: nil
+        field :tie_word_embeddings, default: false
+        field :loop_num, default: 2
+        field :loop_window_size, default: 64
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+          @head_dim ||= @hidden_size / @num_attention_heads
+        end
+      end
+
+      class LoopGateProjection < MLX::NN::Module
+        def initialize(num_heads, head_dim)
+          super()
+          @num_heads = num_heads
+          @head_dim = head_dim
+
+          mx = MLX::Core
+          self.weight = mx.zeros([num_heads, head_dim])
+          self.bias = mx.zeros([num_heads])
+        end
+
+        def call(query)
+          mx = MLX::Core
+          projection = weight.reshape([@num_heads, @head_dim, 1])
+          gate_logits = mx.matmul(query, projection)
+          gate_logits = gate_logits + bias.reshape([1, @num_heads, 1, 1])
+          mx.sigmoid(gate_logits)
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @head_dim = args.head_dim
+          @scale = @head_dim**(-0.5)
+
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: args.attention_bias)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: args.attention_bias)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: args.attention_bias)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: args.attention_bias)
+
+          self.rope = MlxLm::Models.initialize_rope(
+            @head_dim,
+            args.rope_theta,
+            false,
+            args.rope_scaling,
+            max_position_embeddings: args.max_position_embeddings
+          )
+        end
+
+        def get_qkv(x, offset: 0)
+          b, l, _d = x.shape
+
+          queries = q_proj.call(x).reshape([b, l, @n_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = k_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = v_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          queries = rope.call(queries, offset: offset)
+          keys = rope.call(keys, offset: offset)
+
+          [queries, keys, values]
+        end
+
+        def attention(queries, keys, values, mask: nil, cache: nil)
+          _cache = cache
+          MLX::Core.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          hidden_dim = args.intermediate_size
+          self.gate_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: args.mlp_bias)
+          self.down_proj = MLX::NN::Linear.new(hidden_dim, dim, bias: args.mlp_bias)
+          self.up_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: args.mlp_bias)
+        end
+
+        def call(x)
+          down_proj.call(MLX::NN.silu(gate_proj.call(x)) * up_proj.call(x))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Attention.new(args)
+          self.mlp = MLP.new(args)
+          self.input_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_attention_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+      end
+
+      class IQuestLoopCoderModel < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          unless args.loop_num == 2
+            raise ArgumentError, "Only loop_num=2 is supported, got #{args.loop_num}"
+          end
+
+          self.vocab_size = args.vocab_size
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) { TransformerBlock.new(args) }
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.gate_projections = Array.new(args.num_hidden_layers) do
+            LoopGateProjection.new(args.num_attention_heads, args.head_dim)
+          end
+          self.loop_num = args.loop_num
+          self.loop_window_size = args.loop_window_size
+        end
+
+        def call(inputs, cache: nil)
+          mx = MLX::Core
+          b, l = inputs.shape[0], inputs.shape[1]
+
+          h = embed_tokens.call(inputs)
+          layer_count = layers.length
+          layer_cache = cache || [nil] * (2 * layer_count)
+
+          mask = _create_attention_mask(h, layer_cache[0])
+          window_mask = _create_attention_mask(h, layer_cache[layer_count], window_size: loop_window_size)
+
+          loop1_kv = []
+          layers.each_with_index do |layer, idx|
+            c = layer_cache[idx]
+            h_norm = layer.input_layernorm.call(h)
+            offset = c ? c.offset : 0
+            q1, k1, v1 = layer.self_attn.get_qkv(h_norm, offset: offset)
+
+            if c
+              k1, v1 = c.update_and_fetch(k1, v1)
+            end
+            loop1_kv << [k1, v1]
+
+            out = layer.self_attn.attention(q1, k1, v1, mask: mask, cache: c)
+            r = layer.self_attn.o_proj.call(out.transpose([0, 2, 1, 3]).reshape([b, l, @args.hidden_size]))
+            h = h + r
+            r = layer.mlp.call(layer.post_attention_layernorm.call(h))
+            h = h + r
+          end
+
+          layers.each_with_index do |layer, idx|
+            gate_proj = gate_projections[idx]
+            c = layer_cache[layer_count + idx]
+            k1, v1 = loop1_kv[idx]
+
+            h_norm = layer.input_layernorm.call(h)
+            offset = c ? c.offset : 0
+            q2, k2, v2 = layer.self_attn.get_qkv(h_norm, offset: offset)
+
+            gate = gate_proj.call(q2)
+            attn_global = layer.self_attn.attention(q2, k1, v1, mask: mask, cache: c)
+
+            if c
+              k2, v2 = c.update_and_fetch(k2, v2)
+            end
+
+            attn_local = layer.self_attn.attention(q2, k2, v2, mask: window_mask, cache: c)
+            mixed = _mix_attention(gate, attn_global, attn_local)
+
+            r = layer.self_attn.o_proj.call(mixed.transpose([0, 2, 1, 3]).reshape([b, l, @args.hidden_size]))
+            h = h + r
+            r = layer.mlp.call(layer.post_attention_layernorm.call(h))
+            h = h + r
+          end
+
+          norm.call(h)
+        end
+
+        private
+
+        def _create_attention_mask(h, cache = nil, window_size: nil)
+          n = h.shape[1]
+          return cache.make_mask(n) if cache && cache.respond_to?(:make_mask)
+          return nil if n == 1
+          return _create_causal_mask(n, window_size: window_size) if window_size && n > window_size
+
+          "causal"
+        end
+
+        def _create_causal_mask(n, offset: 0, window_size: nil)
+          mx = MLX::Core
+          rinds = mx.arange(0, offset + n, 1, mx.int32).reshape([1, offset + n])
+          linds = mx.arange(offset, offset + n, 1, mx.int32).reshape([n, 1])
+
+          mask = mx.greater_equal(linds, rinds)
+          if window_size
+            mask = mx.logical_and(mask, mx.less(linds, mx.add(rinds, window_size)))
+          end
+          mask
+        end
+
+        def _mix_attention(gate, attn_global, attn_local)
+          gate = gate.astype(attn_global.dtype)
+          (gate * attn_global) + ((1.0 - gate) * attn_local)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.model = IQuestLoopCoderModel.new(args)
+          unless args.tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def layers
+          model.layers
+        end
+
+        def make_cache
+          Array.new(layers.length) { MlxLm::KVCache.new } +
+            Array.new(layers.length) { MlxLm::RotatingKVCache.new(max_size: @args.loop_window_size) }
+        end
+      end
+
+      Models.register("iquestloopcoder", Model, ModelArgs)
+    end
+
+    IQuestLoopCoder = Iquestloopcoder unless const_defined?(:IQuestLoopCoder)
+  end
+end