mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/granite.rb

@@ -0,0 +1,170 @@
+module MlxLm
+  module Models
+    module Granite
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "granite"
+        field :hidden_size
+        field :num_hidden_layers
+        field :intermediate_size
+        field :num_attention_heads
+        field :rms_norm_eps
+        field :vocab_size
+        field :logits_scaling
+        field :attention_multiplier
+        field :embedding_multiplier
+        field :residual_multiplier
+        field :max_position_embeddings
+        field :num_key_value_heads
+        field :attention_bias
+        field :mlp_bias
+        field :rope_theta
+        field :rope_scaling, default: nil
+        field :tie_word_embeddings, default: true
+
+        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 = args.attention_multiplier
+
+          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 = MlxLm::Models.initialize_rope(
+            @head_dim,
+            args.rope_theta,
+            false,
+            args.rope_scaling,
+            max_position_embeddings: args.max_position_embeddings
+          )
+        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
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+
+          dim = args.hidden_size
+          hidden_dim = args.intermediate_size
+          bias = args.mlp_bias
+
+          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)
+          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)
+          @residual_multiplier = args.residual_multiplier
+        end
+
+        def call(x, mask: nil, cache: nil)
+          r = self_attn.call(input_layernorm.call(x), mask: mask, cache: cache)
+          h = x + r * @residual_multiplier
+          r = mlp.call(post_attention_layernorm.call(h))
+          h + r * @residual_multiplier
+        end
+      end
+
+      class GraniteModel < 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)
+          @embedding_multiplier = args.embedding_multiplier
+        end
+
+        def call(inputs, cache: nil)
+          h = embed_tokens.call(inputs) * @embedding_multiplier
+          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_type = args.model_type
+          self.model = GraniteModel.new(args)
+          self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false) unless args.tie_word_embeddings
+          @logits_scaling = args.logits_scaling
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          out = if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+          out / @logits_scaling
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("granite", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/granitemoe.rb

@@ -0,0 +1,58 @@
+require_relative "granite"
+
+module MlxLm
+  module Models
+    module GraniteMoe
+      class ModelArgs < Granite::ModelArgs
+        field :model_type, default: "granitemoe"
+        field :num_local_experts
+        field :num_experts_per_tok
+      end
+
+      class Model < Granite::Model
+        def sanitize(weights)
+          result = weights.dup
+          rewrite_legacy_moe_weights(result)
+          result.delete("lm_head.weight") if @args.tie_word_embeddings
+          result
+        end
+
+        private
+
+        def rewrite_legacy_moe_weights(weights)
+          mx = MLX::Core
+
+          layers.length.times do |layer_idx|
+            prefix = "model.layers.#{layer_idx}.block_sparse_moe"
+            input_key = _first_existing_key(
+              weights,
+              ["#{prefix}.input_linear.weight", "#{prefix}.input_linear"]
+            )
+            output_key = _first_existing_key(
+              weights,
+              ["#{prefix}.output_linear.weight", "#{prefix}.output_linear"]
+            )
+            next unless input_key && output_key
+
+            input_linear = weights.delete(input_key)
+            output_linear = weights.delete(output_key)
+            mid = input_linear.shape[1] / 2
+            gate_proj, up_proj = mx.split(input_linear, [mid], 1)
+
+            weights["#{prefix}.switch_mlp.gate_proj.weight"] = gate_proj
+            weights["#{prefix}.switch_mlp.up_proj.weight"] = up_proj
+            weights["#{prefix}.switch_mlp.down_proj.weight"] = output_linear
+          end
+
+          weights
+        end
+
+        def _first_existing_key(weights, candidates)
+          candidates.find { |key| weights.key?(key) }
+        end
+      end
+
+      Models.register("granitemoe", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/granitemoehybrid.rb

@@ -0,0 +1,178 @@
+require_relative "falcon_h1"
+
+module MlxLm
+  module Models
+    module GraniteMoeHybrid
+      class ModelArgs < FalconH1::ModelArgs
+        field :model_type, default: "granitemoehybrid"
+        field :embedding_multiplier, default: 1.0
+        field :attention_multiplier, default: 1.0
+        field :logits_scaling, default: 1.0
+        field :residual_multiplier, default: 1.0
+        field :num_local_experts, default: nil
+        field :num_experts_per_tok, default: nil
+        field :shared_intermediate_size, default: nil
+        field :mamba_n_heads, default: nil
+        field :mamba_d_head, default: nil
+        field :mamba_proj_bias, default: false
+        field :mamba_d_state, default: nil
+        field :mamba_n_groups, default: nil
+        field :mamba_conv_bias, default: false
+        field :layer_types, default: nil
+        field :position_embedding_type, default: "rope"
+        field :time_step_limit, default: [0.001, 100.0]
+        field :mlp_bias, default: false
+
+        def initialize(**kwargs)
+          super
+          @num_hidden_layers ||= Array(@layer_types).length
+          @num_attention_heads ||= @mamba_n_heads
+          @num_key_value_heads ||= @num_attention_heads
+          @head_dim ||= @mamba_d_head
+          @mamba_d_conv ||= 4
+          @layer_types ||= _default_layer_types
+          @block_types ||= _to_block_types
+        end
+
+        def to_falcon_h1_dict
+          hidden_size = @hidden_size
+          attention_heads = @num_attention_heads
+          inferred_head_dim = if !@head_dim.nil?
+            @head_dim
+          elsif !@mamba_d_head.nil?
+            @mamba_d_head
+          elsif !hidden_size.nil? && attention_heads.to_i > 0
+            hidden_size / attention_heads
+          else
+            64
+          end
+
+          {
+            "model_type" => @model_type,
+            "attention_bias" => @attention_bias,
+            "head_dim" => inferred_head_dim,
+            "hidden_size" => hidden_size,
+            "intermediate_size" => @intermediate_size || @shared_intermediate_size || hidden_size.to_i * 2,
+            "max_position_embeddings" => @max_position_embeddings,
+            "mamba_d_conv" => @mamba_d_conv,
+            "num_attention_heads" => attention_heads,
+            "num_hidden_layers" => @num_hidden_layers,
+            "num_key_value_heads" => @num_key_value_heads,
+            "rms_norm_eps" => @rms_norm_eps,
+            "rope_theta" => @rope_theta,
+            "vocab_size" => @vocab_size,
+            "tie_word_embeddings" => @tie_word_embeddings,
+            "attention_window_size" => @attention_window_size,
+            "block_types" => @block_types,
+          }
+        end
+
+        private
+
+        def _default_layer_types
+          count = @num_hidden_layers.to_i
+          return nil if count <= 0
+
+          Array.new(count) { |idx| idx.even? ? "mamba" : "attention" }
+        end
+
+        def _to_block_types
+          return @block_types if @block_types.is_a?(Array) && !@block_types.empty?
+          return nil unless @layer_types.is_a?(Array) && !@layer_types.empty?
+
+          @layer_types.map { |layer_type| layer_type.to_s == "mamba" ? "recurrent" : "attention" }
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.wrapped_model = FalconH1::Model.new(
+            FalconH1::ModelArgs.from_dict(args.to_falcon_h1_dict)
+          )
+        end
+
+        def call(inputs, cache: nil)
+          wrapped_model.call(inputs, cache: cache)
+        end
+
+        def sanitize(weights)
+          normalized = weights.dup
+          _rewrite_block_sparse_moe!(normalized)
+          _rewrite_shared_mlp!(normalized)
+          normalized.delete("lm_head.weight") if @args.tie_word_embeddings
+
+          remapped = {}
+          normalized.each do |key, value|
+            remapped[_remap_weight_key(key)] = value
+          end
+          wrapped_model.sanitize(remapped)
+        end
+
+        def layers
+          wrapped_model.layers
+        end
+
+        def make_cache
+          return nil unless wrapped_model.respond_to?(:make_cache)
+
+          wrapped_model.make_cache
+        end
+
+        private
+
+        def _rewrite_block_sparse_moe!(weights)
+          mx = MLX::Core
+
+          @args.num_hidden_layers.to_i.times do |layer_idx|
+            prefix = "model.layers.#{layer_idx}.block_sparse_moe"
+            input_key = "#{prefix}.input_linear.weight"
+            output_key = "#{prefix}.output_linear.weight"
+            next unless weights.key?(input_key) && weights.key?(output_key)
+
+            input_linear = weights.delete(input_key)
+            output_linear = weights.delete(output_key)
+            mid = input_linear.shape[1] / 2
+            gate_proj, up_proj = mx.split(input_linear, [mid], 1)
+
+            weights["#{prefix}.switch_mlp.gate_proj.weight"] = gate_proj
+            weights["#{prefix}.switch_mlp.up_proj.weight"] = up_proj
+            weights["#{prefix}.switch_mlp.down_proj.weight"] = output_linear
+          end
+        end
+
+        def _rewrite_shared_mlp!(weights)
+          mx = MLX::Core
+
+          @args.num_hidden_layers.to_i.times do |layer_idx|
+            prefix = "model.layers.#{layer_idx}.shared_mlp"
+            input_key = "#{prefix}.input_linear.weight"
+            output_key = "#{prefix}.output_linear.weight"
+            next unless weights.key?(input_key) && weights.key?(output_key)
+
+            input_linear = weights.delete(input_key)
+            mid = input_linear.shape[0] / 2
+            gate_proj, up_proj = mx.split(input_linear, [mid], 0)
+
+            weights["model.layers.#{layer_idx}.mlp.gate_proj.weight"] = gate_proj
+            weights["model.layers.#{layer_idx}.mlp.up_proj.weight"] = up_proj
+            weights["model.layers.#{layer_idx}.mlp.down_proj.weight"] = weights.delete(output_key)
+          end
+        end
+
+        def _remap_weight_key(key)
+          mapped = key.dup
+          mapped = mapped.gsub(".block_sparse_moe.", ".feed_forward.")
+          mapped = mapped.gsub(".shared_mlp.", ".feed_forward.")
+          mapped = mapped.gsub(".post_attention_layernorm.", ".pre_ff_layernorm.")
+          mapped = mapped.gsub("model.norm.", "model.final_layernorm.")
+          mapped
+        end
+      end
+
+      Models.register("granitemoehybrid", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/helium.rb

@@ -0,0 +1,158 @@
+module MlxLm
+  module Models
+    module Helium
+      class ModelArgs < BaseModelArgs
+        field :hidden_size, default: 256
+        field :num_hidden_layers, default: 24
+        field :intermediate_size, default: 1024
+        field :num_attention_heads, default: 4
+        field :num_key_value_heads, default: nil
+        field :rms_norm_eps, default: 1e-5
+        field :vocab_size, default: 32_000
+        field :attention_bias, default: false
+        field :head_dim, default: nil
+        field :max_position_embeddings, default: 2048
+        field :mlp_bias, default: false
+        field :model_type, default: "helium"
+        field :rope_theta, default: 10_000.0
+        field :tie_word_embeddings, default: false
+
+        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.hidden_size / @n_heads
+          @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: false)
+          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(args)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(
+            args.hidden_size,
+            args.intermediate_size,
+            bias: args.mlp_bias
+          )
+          self.up_proj = MLX::NN::Linear.new(
+            args.hidden_size,
+            args.intermediate_size,
+            bias: args.mlp_bias
+          )
+          self.down_proj = MLX::NN::Linear.new(
+            args.intermediate_size,
+            args.hidden_size,
+            bias: args.mlp_bias
+          )
+        end
+
+        def call(x)
+          down_proj.call(Activations.swiglu(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)
+          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 HeliumModel < 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) { DecoderLayer.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
+          @model_type = args.model_type
+          self.model = HeliumModel.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)
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("helium", Model, ModelArgs)
+    end
+  end
+end