mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/solar_open.rb

@@ -0,0 +1,79 @@
+require_relative "deepseek"
+
+module MlxLm
+  module Models
+    module SolarOpen
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "solar_open"
+        field :vocab_size
+        field :hidden_size
+        field :intermediate_size
+        field :moe_intermediate_size
+        field :num_hidden_layers
+        field :num_attention_heads
+        field :num_key_value_heads
+        field :head_dim
+        field :n_shared_experts
+        field :n_routed_experts
+        field :routed_scaling_factor
+        field :num_experts_per_tok
+        field :first_k_dense_replace
+        field :norm_topk_prob
+        field :max_position_embeddings
+        field :rms_norm_eps
+        field :rope_theta
+        field :tie_word_embeddings
+        field :partial_rotary_factor
+        field :rope_scaling, default: nil
+        field :attention_bias, default: false
+        field :use_qk_norm, default: false
+        field :n_group, default: 1
+        field :topk_group, default: 1
+        field :scoring_func, default: "sigmoid"
+        field :topk_method, default: "noaux_tc"
+      end
+
+      class Model < DeepSeek::Model
+        def initialize(args)
+          super(DeepSeek::ModelArgs.from_dict(_to_deepseek_config(args)))
+          self.model_type = args.model_type
+        end
+
+        def sanitize(weights)
+          sanitized = super(weights)
+          mpt_prefix = "model.layers.#{@args.num_hidden_layers}"
+          sanitized.reject do |k, _|
+            k == mpt_prefix || k.start_with?("#{mpt_prefix}.")
+          end
+        end
+
+        private
+
+        def _to_deepseek_config(args)
+          {
+            "model_type" => args.model_type,
+            "vocab_size" => args.vocab_size,
+            "hidden_size" => args.hidden_size,
+            "intermediate_size" => args.intermediate_size,
+            "moe_intermediate_size" => args.moe_intermediate_size,
+            "num_hidden_layers" => args.num_hidden_layers,
+            "num_attention_heads" => args.num_attention_heads,
+            "num_key_value_heads" => args.num_key_value_heads,
+            "n_shared_experts" => args.n_shared_experts,
+            "n_routed_experts" => args.n_routed_experts,
+            "num_experts_per_tok" => args.num_experts_per_tok,
+            "first_k_dense_replace" => args.first_k_dense_replace,
+            "moe_layer_freq" => 1,
+            "max_position_embeddings" => args.max_position_embeddings,
+            "rms_norm_eps" => args.rms_norm_eps,
+            "rope_theta" => args.rope_theta,
+            "rope_scaling" => args.rope_scaling,
+            "attention_bias" => args.attention_bias,
+          }
+        end
+      end
+
+      Models.register("solar_open", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/ssm.rb

@@ -0,0 +1,162 @@
+module MlxLm
+  module Models
+    module SSM
+      module_function
+
+      def compute_dt(dt, dt_bias, time_step_limit = [0.001, 100.0])
+        dt = MLX::NN.softplus(dt + dt_bias)
+        MLX::Core.clip(dt, time_step_limit[0], time_step_limit[1])
+      end
+
+      def segsum(x, mask: nil)
+        mx = MLX::Core
+        l = x.shape[-1]
+
+        unless mask.nil?
+          mask_e = mx.expand_dims(mask, 1)
+          x = x * mask_e
+        end
+
+        x = mx.repeat(mx.expand_dims(x, -1), l, -1)
+        x = mx.tril(x, -1)
+        x_segsum = mx.cumsum(x, -2)
+
+        unless mask.nil?
+          mask_e = mx.expand_dims(mask, 1)
+          valid = mx.multiply(mx.expand_dims(mask_e, -1), mx.expand_dims(mask_e, -2))
+          x_segsum = mx.where(valid, x_segsum, -Float::INFINITY)
+        end
+
+        x_segsum
+      end
+
+      # Baseline implementation for SSD-SSM using explicit recurrence.
+      def ssm_attn(
+        x,
+        a_log,
+        b,
+        c,
+        d,
+        dt,
+        dt_bias,
+        state: nil,
+        time_step_limit: [0.001, 100.0],
+        mask: nil,
+        lengths: nil,
+        step: 256
+      )
+        _ = step
+        raise NotImplementedError, "length-aware SSM path is not implemented yet" unless lengths.nil?
+
+        mx = MLX::Core
+        batch_size, seq_len, num_heads, head_dim = x.shape
+        _, _, num_groups, state_dim = b.shape
+
+        repeats = num_heads / num_groups
+        dt = compute_dt(dt, dt_bias, time_step_limit)
+        dt = mx.expand_dims(dt, 0) if dt.ndim == 2
+        a = mx.multiply(-1.0, mx.exp(a_log).astype(dt.dtype))
+
+        state ||= mx.zeros([batch_size, num_heads, head_dim, state_dim], x.dtype)
+
+        ys = []
+        seq_len.times do |t|
+          x_t = _slice_step(x, t)
+          dt_t = _slice_step(dt, t)
+          b_t = _slice_step(b, t)
+          c_t = _slice_step(c, t)
+
+          if repeats > 1
+            b_t = mx.repeat(b_t, repeats, 1)
+            c_t = mx.repeat(c_t, repeats, 1)
+          end
+
+          decay = mx.exp(dt_t * a.reshape([1, num_heads]))
+          prev_state = state
+          state = state * decay.reshape([batch_size, num_heads, 1, 1])
+
+          dB = dt_t.reshape([batch_size, num_heads, 1, 1]) * b_t.reshape([batch_size, num_heads, 1, state_dim])
+          state = state + x_t.reshape([batch_size, num_heads, head_dim, 1]) * dB
+
+          y_t = (state * c_t.reshape([batch_size, num_heads, 1, state_dim])).sum(-1)
+          y_t = y_t + x_t * d.reshape([1, num_heads, 1])
+
+          unless mask.nil?
+            m_t = _slice_step(mask, t)
+            m_t = m_t.reshape([batch_size, 1, 1])
+            state = mx.where(m_t, state, prev_state)
+            y_t = mx.where(m_t, y_t, mx.zeros(y_t.shape, y_t.dtype))
+          end
+
+          ys << y_t
+        end
+
+        [mx.stack(ys, 1), state]
+      end
+
+      def ssm_update_kernel(*_args, **_kwargs)
+        raise NotImplementedError,
+              "SSM metal kernel path is not implemented in mlx-ruby-lm yet"
+      end
+
+      def ssm_update(
+        hidden_states,
+        a_log,
+        b,
+        c,
+        d,
+        dt,
+        dt_bias,
+        state: nil,
+        time_step_limit: [0.001, 100.0],
+        mask: nil,
+        lengths: nil
+      )
+        mx = MLX::Core
+        seq_len = hidden_states.shape[1]
+
+        use_attn_path = seq_len > 1 ||
+          state.nil? ||
+          !mx.respond_to?(:metal_is_available) ||
+          !mx.metal_is_available ||
+          !mx.respond_to?(:default_device) ||
+          (mx.default_device.respond_to?(:type) && mx.default_device.type != :gpu)
+
+        if use_attn_path
+          return ssm_attn(
+            hidden_states,
+            a_log,
+            b,
+            c,
+            d,
+            dt,
+            dt_bias,
+            state: state,
+            time_step_limit: time_step_limit,
+            mask: mask,
+            lengths: lengths
+          )
+        end
+
+        ssm_update_kernel(
+          hidden_states,
+          a_log,
+          b,
+          c,
+          d,
+          dt,
+          dt_bias,
+          state,
+          time_step_limit
+        )
+      end
+
+      def _slice_step(array, idx)
+        mx = MLX::Core
+        tail = idx.zero? ? array : mx.split(array, [idx], 1)[1]
+        mx.squeeze(mx.split(tail, [1], 1)[0], 1)
+      end
+      private_class_method :_slice_step
+    end
+  end
+end

data/lib/mlx_lm/models/stablelm.rb

@@ -0,0 +1,160 @@
+module MlxLm
+  module Models
+    module StableLM
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "stablelm"
+        field :hidden_size, default: 2048
+        field :num_hidden_layers, default: 24
+        field :num_attention_heads, default: 32
+        field :num_key_value_heads, default: 32
+        field :intermediate_size, default: 5632
+        field :vocab_size, default: 50304
+        field :rope_theta, default: 10000.0
+        field :use_qkv_bias, default: false
+        field :partial_rotary_factor, default: 0.25
+        field :layer_norm_eps, default: 1e-5
+        field :use_parallel_residual, default: false
+        field :qk_layernorm, default: false
+
+        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)
+
+          bias = args.use_qkv_bias
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: bias)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: bias)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: bias)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: false)
+
+          # Partial rotary: only apply RoPE to a fraction of head_dim
+          rope_dim = (args.partial_rotary_factor * @head_dim).to_i
+          self.rope = MLX::NN::RoPE.new(
+            rope_dim,
+            traditional: false,
+            base: args.rope_theta
+          )
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          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])
+
+          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
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(dim, hidden_dim)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: false)
+          self.down_proj = MLX::NN::Linear.new(hidden_dim, dim, bias: false)
+          self.up_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: false)
+        end
+
+        def call(x)
+          down_proj.call(MLX::NN.silu(gate_proj.call(x)) * up_proj.call(x))
+        end
+      end
+
+      class DecoderLayer < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Attention.new(args)
+          self.mlp = MLP.new(args.hidden_size, args.intermediate_size)
+          self.input_layernorm = MLX::NN::LayerNorm.new(args.hidden_size, eps: args.layer_norm_eps)
+          @use_parallel_residual = args.use_parallel_residual
+          unless @use_parallel_residual
+            self.post_attention_layernorm = MLX::NN::LayerNorm.new(args.hidden_size, eps: args.layer_norm_eps)
+          end
+        end
+
+        def call(x, mask: nil, cache: nil)
+          h = input_layernorm.call(x)
+          r = self_attn.call(h, mask: mask, cache: cache)
+
+          if @use_parallel_residual
+            x + r + mlp.call(h)
+          else
+            h = x + r
+            r = mlp.call(post_attention_layernorm.call(h))
+            h + r
+          end
+        end
+      end
+
+      class StableLMModel < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) { DecoderLayer.new(args) }
+          self.norm = MLX::NN::LayerNorm.new(args.hidden_size, eps: args.layer_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 = StableLMModel.new(args)
+          self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          lm_head.call(out)
+        end
+
+        def sanitize(weights)
+          weights.reject { |k, _| k.include?("self_attn.rotary_emb.inv_freq") }
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("stablelm", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/starcoder2.rb

@@ -0,0 +1,161 @@
+module MlxLm
+  module Models
+    module Starcoder2
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "starcoder2"
+        field :hidden_size, default: 3072
+        field :num_hidden_layers, default: 30
+        field :num_attention_heads, default: 24
+        field :num_key_value_heads, default: 2
+        field :intermediate_size, default: 12288
+        field :vocab_size, default: 49152
+        field :norm_epsilon, default: 1e-5
+        field :rope_theta, default: 100000.0
+        field :rope_traditional, default: false
+        field :tie_word_embeddings, default: true
+        field :head_dim, default: nil
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+          @head_dim ||= @hidden_size / @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 = args.head_dim
+          @scale = @head_dim**(-0.5)
+
+          # StarCoder2: all projections have bias
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: true)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: true)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: true)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: true)
+
+          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
+
+          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])
+
+          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
+      end
+
+      # StarCoder2: simple gelu MLP (no gating)
+      class MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          hidden_dim = args.intermediate_size
+
+          self.c_fc = MLX::NN::Linear.new(dim, hidden_dim, bias: true)
+          self.c_proj = MLX::NN::Linear.new(hidden_dim, dim, bias: true)
+        end
+
+        def call(x)
+          c_proj.call(MLX::NN.gelu(c_fc.call(x)))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Attention.new(args)
+          self.mlp = MLP.new(args)
+          # StarCoder2 uses LayerNorm, not RMSNorm
+          self.input_layernorm = MLX::NN::LayerNorm.new(args.hidden_size, eps: args.norm_epsilon)
+          self.post_attention_layernorm = MLX::NN::LayerNorm.new(args.hidden_size, eps: args.norm_epsilon)
+        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 Starcoder2Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          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::LayerNorm.new(args.hidden_size, eps: args.norm_epsilon)
+        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 = Starcoder2Model.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)
+          result = weights.reject { |k, _| k.include?("self_attn.rotary_emb.inv_freq") }
+          result.delete("lm_head.weight") if @args.tie_word_embeddings
+          result
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("starcoder2", Model, ModelArgs)
+    end
+  end
+end