mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/cohere.rb

@@ -0,0 +1,150 @@
+module MlxLm
+  module Models
+    module Cohere
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "cohere"
+        field :hidden_size, default: 8192
+        field :num_hidden_layers, default: 40
+        field :num_attention_heads, default: 64
+        field :num_key_value_heads, default: 64
+        field :intermediate_size, default: 22528
+        field :vocab_size, default: 256000
+        field :rope_theta, default: 8000000.0
+        field :layer_norm_eps, default: 1e-5
+        field :logit_scale, default: 0.0625
+        field :attention_bias, default: false
+        field :layer_norm_bias, default: false
+        field :use_qk_norm, 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.attention_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: bias)
+
+          self.rope = MLX::NN::RoPE.new(@head_dim, traditional: true, 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 TransformerBlock < 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
+          )
+        end
+
+        def call(x, mask: nil, cache: nil)
+          # Cohere uses parallel residuals: attn + mlp + x
+          h = input_layernorm.call(x)
+          attn_h = self_attn.call(h, mask: mask, cache: cache)
+          ff_h = mlp.call(h)
+          attn_h + ff_h + x
+        end
+      end
+
+      class CohereModel < 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.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 = CohereModel.new(args)
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          # Tied embeddings + logit scaling
+          out = model.embed_tokens.as_linear(out)
+          out * @args.logit_scale
+        end
+
+        def sanitize(weights)
+          weights.reject { |k, _| k.include?("self_attn.rotary_emb.inv_freq") }
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("cohere", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/cohere2.rb

@@ -0,0 +1,224 @@
+module MlxLm
+  module Models
+    module Cohere2
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "cohere2"
+        field :hidden_size, default: 4096
+        field :head_dim, default: 128
+        field :num_hidden_layers, default: 32
+        field :intermediate_size, default: 14336
+        field :num_attention_heads, default: 32
+        field :num_key_value_heads, default: 8
+        field :rope_theta, default: 50_000.0
+        field :vocab_size, default: 256000
+        field :layer_norm_eps, default: 1e-5
+        field :logit_scale, default: 0.0625
+        field :attention_bias, default: false
+        field :layer_norm_bias, default: false
+        field :sliding_window, default: 4096
+        field :sliding_window_pattern, default: 4
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args, layer_idx)
+          super()
+          dim = args.hidden_size
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @head_dim = args.head_dim
+          if (@head_dim * @n_heads) != dim
+            raise ArgumentError,
+              "hidden_size must equal num_attention_heads * head_dim (got #{dim} and #{@n_heads} * #{@head_dim})"
+          end
+          @scale = @head_dim**(-0.5)
+
+          bias = args.attention_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: bias)
+
+          self.rope = MLX::NN::RoPE.new(@head_dim, traditional: true, base: args.rope_theta)
+          @use_sliding_window = ((layer_idx + 1) % args.sliding_window_pattern) != 0
+        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 @use_sliding_window
+            if cache
+              queries = rope.call(queries, offset: cache.offset)
+              keys = rope.call(keys, offset: cache.offset)
+            else
+              queries = rope.call(queries)
+              keys = rope.call(keys)
+            end
+          end
+
+          keys, values = cache.update_and_fetch(keys, values) if cache
+
+          sdpa_type = queries.dtype == mx.float16 ? mx.float32 : queries.dtype
+          output = mx.scaled_dot_product_attention(
+            queries.astype(sdpa_type),
+            keys,
+            values,
+            @scale,
+            mask
+          ).astype(queries.dtype)
+
+          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.up_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: false)
+          self.down_proj = MLX::NN::Linear.new(hidden_dim, dim, bias: false)
+        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, layer_idx)
+          super()
+          self.self_attn = Attention.new(args, layer_idx)
+          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,
+            bias: args.layer_norm_bias
+          )
+        end
+
+        def call(x, mask: nil, cache: nil)
+          h = input_layernorm.call(x)
+          attn_h = self_attn.call(h, mask: mask, cache: cache)
+          ff_h = mlp.call(h)
+          attn_h + ff_h + x
+        end
+      end
+
+      class Cohere2Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          @window_size = args.sliding_window
+
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) { |i| TransformerBlock.new(args, i) }
+          self.norm = MLX::NN::LayerNorm.new(
+            args.hidden_size,
+            eps: args.layer_norm_eps,
+            bias: args.layer_norm_bias
+          )
+        end
+
+        def call(inputs, cache: nil)
+          h = embed_tokens.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+
+          pattern = @args.sliding_window_pattern
+          full_mask = _create_attention_mask(h, layer_cache[pattern - 1])
+          swa_mask = _create_attention_mask(h, layer_cache[0], window_size: @window_size)
+
+          layers.each_with_index do |layer, i|
+            is_global = (i % pattern) == (pattern - 1)
+            mask = is_global ? full_mask : swa_mask
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+
+        private
+
+        def _create_attention_mask(h, cache, window_size: nil)
+          n = h.shape[1]
+          offset = cache ? cache.offset : 0
+
+          if window_size
+            if cache || n > window_size
+              return _create_causal_mask(n, offset, window_size)
+            end
+            return nil if n == 1
+
+            return "causal"
+          end
+
+          return nil if n == 1
+
+          "causal"
+        end
+
+        def _create_causal_mask(n, offset, window_size = nil)
+          mx = MLX::Core
+          rinds = mx.arange(offset + n)
+          linds = offset.zero? ? rinds : mx.arange(offset, offset + n)
+
+          linds = mx.expand_dims(linds, 1)
+          rinds = mx.expand_dims(rinds, 0)
+          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
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model = Cohere2Model.new(args)
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          out = model.embed_tokens.as_linear(out)
+          out * @args.logit_scale
+        end
+
+        def make_cache
+          caches = []
+          @args.num_hidden_layers.times do |i|
+            is_global = (i % @args.sliding_window_pattern) == (@args.sliding_window_pattern - 1)
+            if is_global
+              caches << MlxLm::KVCache.new
+            else
+              caches << MlxLm::RotatingKVCache.new(max_size: @args.sliding_window, keep: 0)
+            end
+          end
+          caches
+        end
+
+        def sanitize(weights)
+          weights.reject { |k, _| k.include?("self_attn.rotary_emb.inv_freq") }
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("cohere2", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/dbrx.rb

@@ -0,0 +1,286 @@
+require_relative "activations"
+
+module MlxLm
+  module Models
+    module Dbrx
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "dbrx"
+        field :vocab_size, default: 32_000
+        field :d_model, default: 6144
+        field :ffn_config, default: {}
+        field :attn_config, default: {}
+        field :n_layers, default: 40
+        field :n_heads, default: 48
+
+        def initialize(**kwargs)
+          super
+          @ffn_config ||= {}
+          @attn_config ||= {}
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+
+          @num_heads = args.n_heads
+          @d_model = args.d_model
+          @head_dim = @d_model / args.n_heads
+          @num_key_value_heads = _attn_value(args.attn_config, "kv_n_heads", args.n_heads).to_i
+          @clip_qkv = _attn_value(args.attn_config, "clip_qkv", 8.0).to_f
+          @rope_theta = _attn_value(args.attn_config, "rope_theta", 10_000.0).to_f
+          @scale = @head_dim**(-0.5)
+
+          self.wqkv = MLX::NN::Linear.new(
+            args.d_model,
+            (@num_key_value_heads * 2 + @num_heads) * @head_dim,
+            bias: false
+          )
+          self.out_proj = MLX::NN::Linear.new(args.d_model, args.d_model, bias: false)
+          self.rope = MLX::NN::RoPE.new(@head_dim, traditional: false, base: @rope_theta)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          qkv = wqkv.call(x)
+          qkv = mx.clip(qkv, -@clip_qkv, @clip_qkv)
+
+          splits = [@d_model, @d_model + @head_dim * @num_key_value_heads]
+          queries, keys, values = mx.split(qkv, splits, -1)
+
+          queries = queries.reshape([b, l, @num_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = keys.reshape([b, l, @num_key_value_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = values.reshape([b, l, @num_key_value_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, @d_model])
+          out_proj.call(output)
+        end
+
+        private
+
+        def _attn_value(config, key, default = nil)
+          return default if config.nil?
+          return config[key] if config.key?(key)
+
+          config.fetch(key.to_sym, default)
+        end
+      end
+
+      class NormAttnNorm < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.norm_1 = MLX::NN::LayerNorm.new(args.d_model, bias: false)
+          self.norm_2 = MLX::NN::LayerNorm.new(args.d_model, bias: false)
+          self.attn = Attention.new(args)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          h = attn.call(norm_1.call(x), mask: mask, cache: cache)
+          residual = x + h
+          [residual, norm_2.call(residual)]
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(d_model, ffn_dim)
+          super()
+          self.v1 = MLX::NN::Linear.new(d_model, ffn_dim, bias: false)
+          self.w1 = MLX::NN::Linear.new(d_model, ffn_dim, bias: false)
+          self.w2 = MLX::NN::Linear.new(ffn_dim, d_model, bias: false)
+        end
+
+        def call(x)
+          w2.call(Activations.swiglu(w1.call(x), v1.call(x)))
+        end
+      end
+
+      class Router < MLX::NN::Module
+        def initialize(d_model, num_experts)
+          super()
+          self.layer = MLX::NN::Linear.new(d_model, num_experts, bias: false)
+        end
+
+        def call(x)
+          layer.call(x)
+        end
+      end
+
+      class SparseMoeBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          @d_model = args.d_model
+          @ffn_dim = _ffn_value(args.ffn_config, "ffn_hidden_size", args.d_model * 4).to_i
+          @num_experts = _ffn_value(args.ffn_config, "moe_num_experts", 1).to_i
+          @num_experts_per_tok = _ffn_value(args.ffn_config, "moe_top_k", 1).to_i
+
+          self.router = Router.new(@d_model, @num_experts)
+          self.experts = Array.new(@num_experts) { MLP.new(@d_model, @ffn_dim) }
+        end
+
+        def call(x)
+          mx = MLX::Core
+
+          top_k = [[@num_experts_per_tok, 1].max, @num_experts].min
+          orig_shape = x.shape
+          token_count = orig_shape[0...-1].reduce(1, :*)
+          flat_x = x.reshape([token_count, orig_shape[-1]])
+
+          gates = router.call(flat_x)
+          gates = mx.softmax(gates.astype(mx.float32), -1)
+
+          inds = mx.stop_gradient(mx.argpartition(gates * -1.0, top_k - 1, -1))
+          take_ids = mx.array((0...top_k).to_a, dtype: mx.int32)
+          inds = mx.take(inds, take_ids, -1)
+          scores = mx.take_along_axis(gates, inds, -1)
+          scores = scores / mx.expand_dims(mx.sum(scores, -1), -1)
+          scores = scores.astype(flat_x.dtype)
+
+          expert_ids = inds.to_a
+          expert_scores = scores.to_a
+
+          outputs = Array.new(flat_x.shape[0]) do |token_idx|
+            token_ids = mx.array([token_idx], dtype: mx.int32)
+            token_state = mx.squeeze(mx.take(flat_x, token_ids, 0), 0)
+
+            token_out = nil
+            expert_ids[token_idx].each_with_index do |expert_idx, score_idx|
+              expert_out = experts[expert_idx.to_i].call(token_state)
+              weighted = expert_out * expert_scores[token_idx][score_idx].to_f
+              token_out = token_out.nil? ? weighted : (token_out + weighted)
+            end
+
+            token_out
+          end
+
+          mx.stack(outputs, 0).reshape(orig_shape)
+        end
+
+        private
+
+        def _ffn_value(config, key, default = nil)
+          return default if config.nil?
+          return config[key] if config.key?(key)
+
+          config.fetch(key.to_sym, default)
+        end
+      end
+
+      class DecoderLayer < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.ffn = SparseMoeBlock.new(args)
+          self.norm_attn_norm = NormAttnNorm.new(args)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          residual, hidden = norm_attn_norm.call(x, mask: mask, cache: cache)
+          ffn.call(hidden) + residual
+        end
+      end
+
+      class DbrxModel < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.wte = MLX::NN::Embedding.new(args.vocab_size, args.d_model)
+          self.blocks = Array.new(args.n_layers) { DecoderLayer.new(args) }
+          self.norm_f = MLX::NN::LayerNorm.new(args.d_model, bias: false)
+        end
+
+        def call(inputs, cache: nil)
+          h = wte.call(inputs)
+          layer_cache = cache || [nil] * blocks.length
+          mask = _create_attention_mask(h, layer_cache[0])
+
+          blocks.each_with_index do |layer, layer_idx|
+            h = layer.call(h, mask: mask, cache: layer_cache[layer_idx])
+          end
+
+          norm_f.call(h)
+        end
+
+        private
+
+        def _create_attention_mask(hidden, cache)
+          return cache.make_mask(hidden.shape[1]) if cache && cache.respond_to?(:make_mask)
+          return nil if hidden.shape[1] == 1
+
+          "causal"
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.transformer = DbrxModel.new(args)
+          self.lm_head = MLX::NN::Linear.new(args.d_model, args.vocab_size, bias: false)
+        end
+
+        def call(inputs, cache: nil)
+          out = transformer.call(inputs, cache: cache)
+          lm_head.call(out)
+        end
+
+        def layers
+          transformer.blocks
+        end
+
+        def sanitize(weights)
+          mx = MLX::Core
+          num_experts = _ffn_value(@args.ffn_config, "moe_num_experts", 0).to_i
+          return weights if num_experts <= 0
+
+          pattern = "experts.mlp"
+          sanitized = {}
+
+          weights.each do |key, value|
+            unless key.include?(pattern)
+              sanitized[key] = value
+              next
+            end
+
+            split_weights = mx.split(value, num_experts, 0)
+            split_weights.each_with_index do |slice, expert_idx|
+              expert_key = _expert_weight_key(key, expert_idx)
+              if key.end_with?("w2") || key.end_with?("w2.weight")
+                slice = slice.transpose([1, 0])
+              end
+              sanitized[expert_key] = slice
+            end
+          end
+
+          sanitized
+        end
+
+        private
+
+        def _expert_weight_key(key, expert_idx)
+          base = key.end_with?(".weight") ? key.sub(/\.weight\z/, "") : key
+          "#{base.sub('.mlp', ".#{expert_idx}")}.weight"
+        end
+
+        def _ffn_value(config, key, default = nil)
+          return default if config.nil?
+          return config[key] if config.key?(key)
+
+          config.fetch(key.to_sym, default)
+        end
+      end
+
+      Models.register("dbrx", Model, ModelArgs)
+    end
+  end
+end