mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/phimoe.rb

@@ -0,0 +1,206 @@
+require_relative "rope_utils"
+require_relative "switch_layers"
+
+module MlxLm
+  module Models
+    module PhiMoe
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "phimoe"
+        field :vocab_size, default: 32064
+        field :hidden_size, default: 4096
+        field :intermediate_size, default: 6400
+        field :num_hidden_layers, default: 32
+        field :num_attention_heads, default: 32
+        field :num_key_value_heads, default: 8
+        field :max_position_embeddings, default: 131072
+        field :original_max_position_embeddings, default: 4096
+        field :rms_norm_eps, default: 1e-6
+        field :rope_scaling, default: nil
+        field :num_local_experts, default: 16
+        field :num_experts_per_tok, default: 2
+        field :rope_theta, default: 10_000.0
+      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)
+
+          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)
+
+          scaling = args.rope_scaling || {}
+          self.rope = SuScaledRoPE.new(
+            @head_dim,
+            base: args.rope_theta,
+            max_position_embeddings: args.max_position_embeddings,
+            original_max_position_embeddings: args.original_max_position_embeddings,
+            short_factor: _config_value(scaling, "short_factor", 1.0),
+            long_factor: _config_value(scaling, "long_factor", 1.0),
+            short_mscale: _config_value(scaling, "short_mscale"),
+            long_mscale: _config_value(scaling, "long_mscale")
+          )
+        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
+
+        private
+
+        def _config_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 PhiMoESparseMoeBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+
+          @hidden_dim = args.hidden_size
+          @ffn_dim = args.intermediate_size
+          @num_experts = args.num_local_experts
+          @top_k = args.num_experts_per_tok
+
+          self.gate = MLX::NN::Linear.new(@hidden_dim, @num_experts, bias: false)
+          self.switch_mlp = SwitchLayers::SwitchGLU.new(@hidden_dim, @ffn_dim, @num_experts)
+        end
+
+        def call(x)
+          mx = MLX::Core
+
+          k = [@top_k, @num_experts].min
+          gates = gate.call(x)
+          inds = mx.stop_gradient(mx.argpartition(gates * -1.0, k - 1, -1))
+          take_ids = mx.array((0...k).to_a, dtype: mx.int32)
+          inds = mx.take(inds, take_ids, -1)
+          scores = mx.take_along_axis(gates, inds, -1)
+          scores = mx.softmax(scores.astype(mx.float32), -1).astype(gates.dtype)
+
+          y = switch_mlp.call(x, inds)
+          mx.sum(y * mx.expand_dims(scores, -1), -2)
+        end
+      end
+
+      class DecoderLayer < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Attention.new(args)
+          self.block_sparse_moe = PhiMoESparseMoeBlock.new(args)
+          self.input_layernorm = MLX::NN::LayerNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_attention_layernorm = MLX::NN::LayerNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          residual = x
+          hidden_states = input_layernorm.call(x)
+          hidden_states = self_attn.call(hidden_states, mask: mask, cache: cache)
+          hidden_states = residual + hidden_states
+
+          residual = hidden_states
+          hidden_states = post_attention_layernorm.call(hidden_states)
+          hidden_states = block_sparse_moe.call(hidden_states)
+          residual + hidden_states
+        end
+      end
+
+      class PhiMoEModel < 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::LayerNorm.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, layer_idx|
+            h = layer.call(h, mask: mask, cache: layer_cache[layer_idx])
+          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 = PhiMoEModel.new(args)
+          self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: true)
+        end
+
+        def call(inputs, cache: nil)
+          lm_head.call(model.call(inputs, cache: cache))
+        end
+
+        def sanitize(weights)
+          return weights unless weights.key?("model.layers.0.block_sparse_moe.experts.0.w1.weight")
+
+          mx = MLX::Core
+          result = weights.dup
+
+          @args.num_hidden_layers.times do |layer_idx|
+            prefix = "model.layers.#{layer_idx}"
+            [["w1", "gate_proj"], ["w2", "down_proj"], ["w3", "up_proj"]].each do |source, target|
+              %w[weight scales biases].each do |param|
+                first_key = "#{prefix}.block_sparse_moe.experts.0.#{source}.#{param}"
+                next unless result.key?(first_key)
+
+                expert_keys = (0...@args.num_local_experts).map do |expert_idx|
+                  "#{prefix}.block_sparse_moe.experts.#{expert_idx}.#{source}.#{param}"
+                end
+                next unless expert_keys.all? { |key| result.key?(key) }
+
+                stacked = expert_keys.map { |key| result.delete(key) }
+                result["#{prefix}.block_sparse_moe.switch_mlp.#{target}.#{param}"] = mx.stack(stacked)
+              end
+            end
+          end
+
+          result
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("phimoe", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/phixtral.rb

@@ -0,0 +1,208 @@
+require_relative "switch_layers"
+
+module MlxLm
+  module Models
+    module Phixtral
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "phixtral"
+        field :num_vocab, default: 51_200
+        field :model_dim, default: 2_560
+        field :num_heads, default: 32
+        field :num_layers, default: 32
+        field :rotary_dim, default: 32
+        field :num_experts_per_tok, default: 2
+        field :num_local_experts, default: 4
+      end
+
+      class RoPEAttention < MLX::NN::Module
+        def initialize(dims, num_heads, rotary_dim)
+          super()
+          @num_heads = num_heads
+          @head_dim = dims / num_heads
+          @scale = @head_dim**(-0.5)
+
+          self.rope = MLX::NN::RoPE.new(rotary_dim, traditional: false)
+          self.wqkv = MLX::NN::Linear.new(dims, 3 * dims)
+          self.out_proj = MLX::NN::Linear.new(dims, dims)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, d = x.shape
+
+          qkv = wqkv.call(x)
+          queries, keys, values = mx.split(qkv, [d, 2 * d], -1)
+
+          queries = queries.reshape([b, l, @num_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = keys.reshape([b, l, @num_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = values.reshape([b, l, @num_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
+
+          queries = queries.astype(mx.float32)
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask).astype(values.dtype)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, d])
+          out_proj.call(output)
+        end
+      end
+
+      class MOE < MLX::NN::Module
+        def initialize(args, dim, hidden_dim)
+          super()
+          @num_experts = args.num_local_experts
+          @num_experts_per_tok = args.num_experts_per_tok
+
+          self.switch_mlp = SwitchLayers::SwitchMLP.new(
+            dim,
+            hidden_dim,
+            @num_experts,
+            bias: true
+          )
+          self.gate = MLX::NN::Linear.new(args.model_dim, @num_experts, bias: false)
+        end
+
+        def call(x)
+          mx = MLX::Core
+          k = @num_experts_per_tok
+
+          gates = gate.call(x)
+          inds = mx.stop_gradient(mx.argpartition(gates * -1.0, k - 1, -1))
+          take_ids = mx.array((0...k).to_a, dtype: mx.int32)
+          inds = mx.take(inds, take_ids, -1)
+
+          scores = mx.take_along_axis(gates, inds, -1)
+          scores = mx.softmax(scores.astype(mx.float32), -1).astype(gates.dtype)
+
+          y = switch_mlp.call(x, inds)
+          mx.sum(y * mx.expand_dims(scores, -1), -2)
+        end
+      end
+
+      class ParallelBlock < MLX::NN::Module
+        def initialize(config)
+          super()
+          dims = config.model_dim
+          mlp_dims = dims * 4
+
+          self.mixer = RoPEAttention.new(dims, config.num_heads, config.rotary_dim)
+          self.ln = MLX::NN::LayerNorm.new(dims)
+          self.moe = MOE.new(config, dims, mlp_dims)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          h = ln.call(x)
+          attn_h = mixer.call(h, mask: mask, cache: cache)
+          ff_h = moe.call(h)
+          attn_h + ff_h + x
+        end
+      end
+
+      class Embd < MLX::NN::Module
+        def initialize(config)
+          super()
+          self.wte = MLX::NN::Embedding.new(config.num_vocab, config.model_dim)
+        end
+
+        def call(x)
+          wte.call(x)
+        end
+      end
+
+      class TransformerDecoder < MLX::NN::Module
+        def initialize(config)
+          super()
+          self.embd = Embd.new(config)
+          self.h = Array.new(config.num_layers) { ParallelBlock.new(config) }
+        end
+
+        def call(x, mask: nil, cache: nil)
+          hidden = embd.call(x)
+          layer_cache = cache || [nil] * h.length
+
+          h.each_with_index do |layer, i|
+            hidden = layer.call(hidden, mask: mask, cache: layer_cache[i])
+          end
+
+          hidden
+        end
+      end
+
+      class OutputHead < MLX::NN::Module
+        def initialize(config)
+          super()
+          self.ln = MLX::NN::LayerNorm.new(config.model_dim)
+          self.linear = MLX::NN::Linear.new(config.model_dim, config.num_vocab)
+        end
+
+        def call(inputs)
+          linear.call(ln.call(inputs))
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(config)
+          super()
+          @args = config
+
+          self.model_type = config.model_type
+          self.transformer = TransformerDecoder.new(config)
+          self.lm_head = OutputHead.new(config)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          local_mask = mask || _create_attention_mask(x, cache)
+          y = transformer.call(x, mask: local_mask, cache: cache)
+          lm_head.call(y)
+        end
+
+        def sanitize(weights)
+          first_key = "transformer.h.0.moe.mlp.0.fc1.weight"
+          return weights unless weights.key?(first_key)
+
+          mx = MLX::Core
+          result = weights.dup
+
+          @args.num_layers.times do |layer_idx|
+            prefix = "transformer.h.#{layer_idx}"
+            %w[fc1 fc2].each do |proj|
+              %w[weight scales biases bias].each do |suffix|
+                expert_keys = (0...@args.num_local_experts).map do |expert_idx|
+                  "#{prefix}.moe.mlp.#{expert_idx}.#{proj}.#{suffix}"
+                end
+                next unless expert_keys.all? { |k| result.key?(k) }
+
+                stacked = expert_keys.map { |k| result.delete(k) }
+                result["#{prefix}.moe.switch_mlp.#{proj}.#{suffix}"] = mx.stack(stacked)
+              end
+            end
+          end
+
+          result
+        end
+
+        def layers
+          transformer.h
+        end
+
+        private
+
+        def _create_attention_mask(tokens, cache)
+          first_cache = cache.is_a?(Array) ? cache[0] : cache
+          return first_cache.make_mask(tokens.shape[1]) if first_cache && first_cache.respond_to?(:make_mask)
+          return nil if tokens.shape[1] == 1
+
+          "causal"
+        end
+      end
+
+      Models.register("phixtral", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/pipeline.rb

@@ -0,0 +1,37 @@
+module MlxLm
+  module Models
+    module PipelineMixin
+      attr_accessor :pipeline_rank, :pipeline_size, :start_idx, :end_idx
+
+      def initialize(*args, **kwargs)
+        super(*args, **kwargs)
+        @pipeline_rank = 0
+        @pipeline_size = 1
+        @start_idx = 0
+        @end_idx = nil
+      end
+
+      def pipeline_layers
+        layers[@start_idx...@end_idx]
+      end
+
+      def pipeline(group)
+        # Split layers in reverse so rank=0 gets the last layers and
+        # rank=pipeline_size-1 gets the first.
+        @pipeline_rank = group.rank
+        @pipeline_size = group.size
+        layers_per_rank = layers.length / @pipeline_size
+        extra = layers.length - (layers_per_rank * @pipeline_size)
+        layers_per_rank += 1 if @pipeline_rank < extra
+
+        @start_idx = (@pipeline_size - @pipeline_rank - 1) * layers_per_rank
+        @end_idx = @start_idx + layers_per_rank
+
+        self.layers = layers[0...@end_idx]
+        # Keep layer numbering stable for checkpoint loading.
+        self.layers[0...@start_idx] = Array.new(@start_idx, nil)
+        self
+      end
+    end
+  end
+end

data/lib/mlx_lm/models/pixtral.rb

@@ -0,0 +1,47 @@
+module MlxLm
+  module Models
+    module Pixtral
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "pixtral"
+        field :text_config
+
+        def initialize(**kwargs)
+          super
+          @text_config ||= {}
+          @text_config["tie_word_embeddings"] = false
+          unless @text_config.key?("num_attention_heads") || @text_config.key?(:num_attention_heads)
+            @text_config["num_attention_heads"] = 32
+          end
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.language_model = Llama::Model.new(Llama::ModelArgs.from_dict(args.text_config))
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          language_model.call(inputs, cache: cache)
+        end
+
+        def sanitize(weights)
+          weights.reject do |key, _|
+            key == "vision_tower" ||
+              key.start_with?("vision_tower.") ||
+              key == "multi_modal_projector" ||
+              key.start_with?("multi_modal_projector.")
+          end
+        end
+
+        def layers
+          language_model.model.layers
+        end
+      end
+
+      Models.register("pixtral", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/plamo.rb

@@ -0,0 +1,169 @@
+module MlxLm
+  module Models
+    module Plamo
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "plamo"
+        field :hidden_size
+        field :num_hidden_layers
+        field :intermediate_size
+        field :num_attention_heads
+        field :rms_norm_eps
+        field :vocab_size
+        field :n_shared_head, default: 8
+        field :rope_theta, default: 10_000.0
+        field :rope_traditional, default: false
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(config)
+          super()
+          @config = config
+          @hidden_size = config.hidden_size
+          @q_num_heads = config.num_attention_heads
+          @head_dim = @hidden_size / @q_num_heads
+          @qk_dim = @head_dim
+          @v_dim = @head_dim
+          @k_num_heads = (@q_num_heads.to_f / config.n_shared_head).ceil
+          @v_num_heads = @k_num_heads
+          @scale = @head_dim**(-0.5)
+
+          self.q_proj = MLX::NN::Linear.new(@hidden_size, @q_num_heads * @qk_dim, bias: false)
+          self.k_proj = MLX::NN::Linear.new(@hidden_size, @k_num_heads * @qk_dim, bias: false)
+          self.v_proj = MLX::NN::Linear.new(@hidden_size, @v_num_heads * @v_dim, bias: false)
+          self.o_proj = MLX::NN::Linear.new(@q_num_heads * @v_dim, @hidden_size, bias: false)
+          self.rotary_emb = MLX::NN::RoPE.new(
+            @head_dim,
+            traditional: config.rope_traditional,
+            base: config.rope_theta,
+            scale: 1.0
+          )
+        end
+
+        def call(hidden_states, attention_mask: nil, cache: nil)
+          mx = MLX::Core
+          bsz, q_len, _d = hidden_states.shape
+
+          queries = q_proj.call(hidden_states)
+          keys = k_proj.call(hidden_states)
+          values = v_proj.call(hidden_states)
+
+          queries = queries.reshape([bsz, q_len, @q_num_heads, @qk_dim]).transpose([0, 2, 1, 3])
+          keys = keys.reshape([bsz, q_len, @k_num_heads, @qk_dim]).transpose([0, 2, 1, 3])
+          values = values.reshape([bsz, q_len, @v_num_heads, @v_dim]).transpose([0, 2, 1, 3])
+
+          if cache
+            queries = rotary_emb.call(queries, offset: cache.offset)
+            keys = rotary_emb.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rotary_emb.call(queries)
+            keys = rotary_emb.call(keys)
+          end
+
+          keys = mx.tile(keys, [1, @config.n_shared_head, 1, 1])
+          values = mx.tile(values, [1, @config.n_shared_head, 1, 1])
+
+          output = mx.scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            @scale,
+            attention_mask
+          )
+          output = output.transpose([0, 2, 1, 3]).reshape([bsz, q_len, @q_num_heads * @v_dim])
+          o_proj.call(output)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(config)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(config.hidden_size, config.intermediate_size, bias: false)
+          self.up_proj = MLX::NN::Linear.new(config.hidden_size, config.intermediate_size, bias: false)
+          self.down_proj = MLX::NN::Linear.new(config.intermediate_size, config.hidden_size, bias: false)
+        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(config)
+          super()
+          self.self_attn = Attention.new(config)
+          self.mlp = MLP.new(config)
+          self.norm = MLX::NN::RMSNorm.new(config.hidden_size, eps: config.rms_norm_eps)
+        end
+
+        def call(hidden_states, attention_mask: nil, cache: nil)
+          residual = hidden_states
+          hidden_states = norm.call(hidden_states)
+
+          hidden_states_sa = self_attn.call(
+            hidden_states,
+            attention_mask: attention_mask,
+            cache: cache
+          )
+          hidden_states_mlp = mlp.call(hidden_states)
+
+          residual + hidden_states_sa + hidden_states_mlp
+        end
+      end
+
+      class PlamoModel < MLX::NN::Module
+        def initialize(config)
+          super()
+          self.embed_tokens = MLX::NN::Embedding.new(config.vocab_size, config.hidden_size)
+          self.layers = Array.new(config.num_hidden_layers) { DecoderLayer.new(config) }
+          self.norm = MLX::NN::RMSNorm.new(config.hidden_size, eps: config.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil)
+          h = embed_tokens.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+          mask = _create_attention_mask(h, layer_cache[0])
+
+          layers.each_with_index do |layer, i|
+            h = layer.call(h, attention_mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+
+        private
+
+        def _create_attention_mask(h, cache)
+          return cache.make_mask(h.shape[1]) if cache && cache.respond_to?(:make_mask)
+          return nil if h.shape[1] == 1
+
+          "causal"
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.model_type = args.model_type
+          self.model = PlamoModel.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("plamo", Model, ModelArgs)
+    end
+  end
+end