mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/qwen2_vl.rb

@@ -0,0 +1,48 @@
+module MlxLm
+  module Models
+    module Qwen2VL
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "qwen2_vl"
+        field :text_config
+
+        def self.from_dict(params)
+          has_text_config = params.key?("text_config") || params.key?(:text_config)
+          return super if has_text_config
+
+          new(model_type: params["model_type"] || params[:model_type], text_config: params)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.language_model = Qwen2::Model.new(Qwen2::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)
+          sanitized = {}
+          weights.each do |key, value|
+            next if key == "visual" || key.start_with?("visual.")
+            next if key == "vision_tower" || key.start_with?("vision_tower.")
+
+            mapped_key = key.start_with?("language_model.") ? key : "language_model.#{key}"
+            sanitized[mapped_key] = value
+          end
+          sanitized
+        end
+
+        def layers
+          language_model.model.layers
+        end
+      end
+
+      Models.register("qwen2_vl", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/qwen3.rb

@@ -0,0 +1,167 @@
+module MlxLm
+  module Models
+    module Qwen3
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "qwen3"
+        field :hidden_size, default: 2048
+        field :num_hidden_layers, default: 24
+        field :intermediate_size, default: 11008
+        field :num_attention_heads, default: 16
+        field :rms_norm_eps, default: 1e-6
+        field :vocab_size, default: 151936
+        field :num_key_value_heads, default: nil
+        field :max_position_embeddings, default: 32768
+        field :rope_theta, default: 1_000_000.0
+        field :head_dim, default: nil
+        field :tie_word_embeddings, default: true
+        field :rope_scaling, 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)
+
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: false)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: false)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: false)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: false)
+
+          self.q_norm = MLX::NN::RMSNorm.new(@head_dim, eps: args.rms_norm_eps)
+          self.k_norm = MLX::NN::RMSNorm.new(@head_dim, eps: args.rms_norm_eps)
+          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 = q_norm.call(queries.reshape([b, l, @n_heads, @head_dim])).transpose([0, 2, 1, 3])
+          keys = k_norm.call(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(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(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)
+          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 Qwen3Model < 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::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          h = input_embeddings || 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_type = args.model_type
+          self.model = Qwen3Model.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, input_embeddings: nil)
+          out = model.call(inputs, cache: cache, input_embeddings: input_embeddings)
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          result = weights.dup
+          result.delete("lm_head.weight") if @args.tie_word_embeddings
+          result
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("qwen3", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/qwen3_5.rb

@@ -0,0 +1,69 @@
+module MlxLm
+  module Models
+    module Qwen35
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "qwen3_5"
+        field :text_config, default: nil
+
+        def self.from_dict(params)
+          has_text_config = params.key?("text_config") || params.key?(:text_config)
+          return super if has_text_config
+
+          new(model_type: params["model_type"] || params[:model_type], text_config: params)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.language_model = Qwen3::Model.new(Qwen3::ModelArgs.from_dict(_text_config_for_qwen3(args)))
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          language_model.call(inputs, cache: cache, input_embeddings: input_embeddings)
+        end
+
+        def sanitize(weights)
+          language_model.sanitize(remap_language_model_weights(weights))
+        end
+
+        def layers
+          language_model.layers
+        end
+
+        protected
+
+        def remap_language_model_weights(weights)
+          remapped = {}
+          weights.each do |key, value|
+            next if key.start_with?("model.visual")
+
+            mapped_key = if key.start_with?("model.language_model")
+              key.sub("model.language_model", "language_model.model")
+            elsif key.start_with?("language_model.")
+              key
+            else
+              "language_model.#{key}"
+            end
+            remapped[mapped_key] = value
+          end
+          remapped
+        end
+
+        private
+
+        def _text_config_for_qwen3(args)
+          config = {}
+          (args.text_config || {}).each { |key, value| config[key.to_s] = value }
+          config["model_type"] ||= args.model_type
+          config["tie_word_embeddings"] = false unless config.key?("tie_word_embeddings")
+          config
+        end
+      end
+
+      Models.register("qwen3_5", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/qwen3_5_moe.rb

@@ -0,0 +1,54 @@
+module MlxLm
+  module Models
+    module Qwen35Moe
+      class ModelArgs < Qwen35::ModelArgs
+        field :model_type, default: "qwen3_5_moe"
+      end
+
+      class Model < Qwen35::Model
+        def sanitize(weights)
+          remapped = remap_language_model_weights(weights)
+          rewrite_moe_expert_weights(remapped)
+          language_model.sanitize(remapped)
+        end
+
+        private
+
+        def rewrite_moe_expert_weights(weights)
+          mx = MLX::Core
+
+          layers.length.times do |layer_idx|
+            prefix = "language_model.model.layers.#{layer_idx}.mlp"
+            gate_up_key = _first_existing_key(
+              weights,
+              ["#{prefix}.experts.gate_up_proj", "#{prefix}.experts.gate_up_proj.weight"]
+            )
+            down_proj_key = _first_existing_key(
+              weights,
+              ["#{prefix}.experts.down_proj", "#{prefix}.experts.down_proj.weight"]
+            )
+
+            next unless gate_up_key && down_proj_key
+
+            gate_up = weights.delete(gate_up_key)
+            down_proj = weights.delete(down_proj_key)
+            mid = gate_up.shape[-2] / 2
+            gate_proj, up_proj = mx.split(gate_up, [mid], -2)
+
+            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"] = down_proj
+          end
+
+          weights
+        end
+
+        def _first_existing_key(weights, candidates)
+          candidates.find { |key| weights.key?(key) }
+        end
+      end
+
+      Models.register("qwen3_5_moe", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/qwen3_moe.rb

@@ -0,0 +1,166 @@
+require_relative "qwen3"
+require_relative "switch_layers"
+
+module MlxLm
+  module Models
+    module Qwen3Moe
+      class ModelArgs < Qwen3::ModelArgs
+        field :model_type, default: "qwen3_moe"
+        field :num_experts, default: 128
+        field :num_experts_per_tok, default: 8
+        field :decoder_sparse_step, default: 1
+        field :mlp_only_layers, default: []
+        field :moe_intermediate_size, default: 1408
+        field :norm_topk_prob, default: false
+
+        def initialize(**kwargs)
+          super
+          @mlp_only_layers ||= []
+        end
+      end
+
+      class SparseMoeBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          @top_k = [args.num_experts_per_tok.to_i, 1].max
+          @num_experts = args.num_experts
+          @norm_topk_prob = args.norm_topk_prob
+
+          dim = args.hidden_size
+          hidden_dim = args.moe_intermediate_size
+
+          self.gate = MLX::NN::Linear.new(dim, @num_experts, bias: false)
+          self.switch_mlp = SwitchLayers::SwitchGLU.new(dim, hidden_dim, @num_experts)
+        end
+
+        def call(x)
+          mx = MLX::Core
+
+          gates = gate.call(x)
+          gates = mx.softmax(gates.astype(mx.float32), -1).astype(gates.dtype)
+
+          k = [@top_k, @num_experts].min
+          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)
+
+          if @norm_topk_prob
+            denom = mx.expand_dims(mx.sum(scores, -1), -1)
+            scores = scores / denom
+          end
+
+          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, layer_idx)
+          super()
+          self.self_attn = Qwen3::Attention.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)
+
+          if _use_sparse_moe_layer?(args, layer_idx)
+            self.mlp = SparseMoeBlock.new(args)
+          else
+            self.mlp = Qwen3::MLP.new(args.hidden_size, args.intermediate_size)
+          end
+        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
+
+        private
+
+        def _use_sparse_moe_layer?(args, layer_idx)
+          sparse_step = [args.decoder_sparse_step.to_i, 1].max
+          mlp_only_layers = args.mlp_only_layers || []
+
+          !mlp_only_layers.include?(layer_idx) &&
+            args.num_experts.to_i > 0 &&
+            ((layer_idx + 1) % sparse_step).zero?
+        end
+      end
+
+      class Qwen3MoeModel < 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) { |layer_idx| DecoderLayer.new(args, layer_idx) }
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          h = input_embeddings || 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 = Qwen3MoeModel.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, input_embeddings: nil)
+          out = model.call(inputs, cache: cache, input_embeddings: input_embeddings)
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          mx = MLX::Core
+
+          result = weights.dup
+          result.delete("lm_head.weight") if @args.tie_word_embeddings
+          return result unless result.key?("model.layers.0.mlp.experts.0.up_proj.weight")
+
+          @args.num_hidden_layers.times do |layer_idx|
+            prefix = "model.layers.#{layer_idx}.mlp"
+            %w[up_proj down_proj gate_proj].each do |projection|
+              expert_keys = (0...@args.num_experts).map do |expert_idx|
+                "#{prefix}.experts.#{expert_idx}.#{projection}.weight"
+              end
+              next unless expert_keys.all? { |key| result.key?(key) }
+
+              stacked = expert_keys.map { |key| result.delete(key) }
+              result["#{prefix}.switch_mlp.#{projection}.weight"] = mx.stack(stacked)
+            end
+          end
+
+          result
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("qwen3_moe", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/qwen3_next.rb

@@ -0,0 +1,147 @@
+require_relative "kimi_linear"
+
+module MlxLm
+  module Models
+    module Qwen3Next
+      class ModelArgs < KimiLinear::ModelArgs
+        field :model_type, default: "qwen3_next"
+        field :linear_num_value_heads, default: nil
+        field :linear_num_key_heads, default: nil
+        field :linear_key_head_dim, default: nil
+        field :linear_value_head_dim, default: nil
+        field :linear_conv_kernel_dim, default: nil
+        field :decoder_sparse_step, default: nil
+        field :shared_expert_intermediate_size, default: nil
+        field :mlp_only_layers, default: []
+        field :full_attention_interval, default: 4
+        field :head_dim, default: nil
+        field :attention_bias, default: false
+        field :num_shared_experts, default: 1
+        field :norm_topk_prob, default: false
+        field :first_k_dense_replace, default: 0
+
+        def self.from_dict(params)
+          normalized = params.each_with_object({}) do |(key, value), out|
+            out[key.to_s] = value
+          end
+
+          {
+            "shared_expert_intermediate_size" => "moe_shared_expert_intermediate_size",
+          }.each do |source_key, target_key|
+            next unless normalized.key?(source_key)
+
+            normalized[target_key] = normalized[source_key] unless normalized.key?(target_key)
+          end
+
+          if normalized.key?("attention_bias")
+            normalized["use_bias"] = normalized["attention_bias"] unless normalized.key?("use_bias")
+            normalized["use_qkv_bias"] = normalized["attention_bias"] unless normalized.key?("use_qkv_bias")
+          end
+
+          if normalized.key?("linear_num_key_heads") && !normalized.key?("num_key_value_heads")
+            normalized["num_key_value_heads"] = normalized["linear_num_key_heads"]
+          end
+
+          if normalized.key?("mlp_only_layers") && !normalized.key?("first_k_dense_replace")
+            normalized["first_k_dense_replace"] = _dense_prefix_length(normalized["mlp_only_layers"])
+          end
+
+          normalized["num_shared_experts"] = 1 unless normalized.key?("num_shared_experts")
+          normalized["norm_topk_prob"] = false unless normalized.key?("norm_topk_prob")
+          normalized["first_k_dense_replace"] = 0 unless normalized.key?("first_k_dense_replace")
+          normalized["model_type"] ||= "qwen3_next"
+          super(normalized)
+        end
+
+        def initialize(**kwargs)
+          super
+          @moe_shared_expert_intermediate_size = @shared_expert_intermediate_size if kwargs.key?(:shared_expert_intermediate_size) && !kwargs.key?(:moe_shared_expert_intermediate_size) && !@shared_expert_intermediate_size.nil?
+
+          if kwargs.key?(:attention_bias) && !@attention_bias.nil?
+            @use_bias = @attention_bias unless kwargs.key?(:use_bias)
+            @use_qkv_bias = @attention_bias unless kwargs.key?(:use_qkv_bias)
+          end
+
+          if kwargs.key?(:mlp_only_layers) && !kwargs.key?(:first_k_dense_replace)
+            @first_k_dense_replace = self.class._dense_prefix_length(@mlp_only_layers)
+          end
+
+          @num_shared_experts = 1 if @num_shared_experts.nil?
+          @norm_topk_prob = false if @norm_topk_prob.nil?
+          @first_k_dense_replace = 0 if @first_k_dense_replace.nil?
+          @num_key_value_heads ||= @num_attention_heads
+        end
+
+        def to_kimi_linear_dict
+          dict = to_bailing_moe_linear_dict
+          dict["model_type"] = @model_type
+          dict["num_shared_experts"] = @num_shared_experts || 1
+          dict["norm_topk_prob"] = @norm_topk_prob.nil? ? false : @norm_topk_prob
+          dict["first_k_dense_replace"] = @first_k_dense_replace || 0
+          dict["use_bias"] = @use_bias
+          dict["use_qkv_bias"] = @use_qkv_bias
+          dict["moe_shared_expert_intermediate_size"] = @moe_shared_expert_intermediate_size unless @moe_shared_expert_intermediate_size.nil?
+          dict
+        end
+
+        def self._dense_prefix_length(mlp_only_layers)
+          layers = Array(mlp_only_layers).map(&:to_i)
+          count = 0
+          count += 1 while layers.include?(count)
+          count
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.wrapped_model = KimiLinear::Model.new(
+            KimiLinear::ModelArgs.from_dict(args.to_kimi_linear_dict)
+          )
+        end
+
+        def call(inputs, cache: nil)
+          wrapped_model.call(inputs, cache: cache)
+        end
+
+        def sanitize(weights)
+          remapped = {}
+          flat_weights = weights.is_a?(Hash) ? weights : weights.to_h
+          flat_weights.each do |key, value|
+            mapped = key.to_s.gsub(".mlp.shared_expert.", ".mlp.shared_experts.")
+            next if mapped.include?(".mtp.")
+
+            remapped[mapped] = value
+          end
+          wrapped_model.sanitize(remapped)
+        end
+
+        def layers
+          wrapped_model.layers
+        end
+
+        def make_cache
+          return wrapped_model.make_cache if wrapped_model.respond_to?(:make_cache)
+
+          nil
+        end
+
+        def cast_predicate
+          return wrapped_model.cast_predicate if wrapped_model.respond_to?(:cast_predicate)
+
+          lambda { |_key| true }
+        end
+
+        def quant_predicate
+          return wrapped_model.quant_predicate if wrapped_model.respond_to?(:quant_predicate)
+
+          lambda { |_key, _value| true }
+        end
+      end
+
+      Models.register("qwen3_next", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/qwen3_vl.rb

@@ -0,0 +1,48 @@
+module MlxLm
+  module Models
+    module Qwen3VL
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "qwen3_vl"
+        field :text_config, default: nil
+
+        def self.from_dict(params)
+          return super if params.key?("text_config")
+
+          new(model_type: params["model_type"], text_config: params)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.language_model = Qwen3::Model.new(Qwen3::ModelArgs.from_dict(args.text_config))
+        end
+
+        def call(inputs, cache: nil, input_embeddings: nil)
+          language_model.call(inputs, cache: cache, input_embeddings: input_embeddings)
+        end
+
+        def sanitize(weights)
+          nested = MLX::Utils.tree_unflatten(weights.to_a)
+          nested.delete("vision_tower") if nested.is_a?(Hash)
+
+          flattened = MLX::Utils.tree_flatten(nested, destination: {})
+          sanitized = {}
+          flattened.each do |key, value|
+            sanitized_key = key.start_with?("language_model.") ? key : "language_model.#{key}"
+            sanitized[sanitized_key] = value
+          end
+          sanitized
+        end
+
+        def layers
+          language_model.layers
+        end
+      end
+
+      Models.register("qwen3_vl", Model, ModelArgs)
+    end
+  end
+end