mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/plamo2.rb

@@ -0,0 +1,173 @@
+require_relative "falcon_h1"
+
+module MlxLm
+  module Models
+    module Plamo2
+      class ModelArgs < FalconH1::ModelArgs
+        field :model_type, default: "plamo2"
+        field :rope_theta, default: 10_000.0
+        field :tie_word_embeddings, default: true
+        field :hidden_size_per_head, default: nil
+        field :full_attention_idx, default: nil
+        field :mamba_d_state, default: nil
+        field :mamba_num_heads, default: nil
+        field :mamba_step, default: 2
+        field :mamba_chunk_size, default: nil
+        field :mamba_enabled, default: true
+
+        def initialize(**kwargs)
+          super
+          @head_dim = @hidden_size_per_head if kwargs.key?(:hidden_size_per_head) && !kwargs.key?(:head_dim) && !@hidden_size_per_head.nil?
+          @num_attention_heads ||= @mamba_num_heads
+          @num_key_value_heads ||= @num_attention_heads
+          @mamba_d_conv ||= 4
+          @attention_window_size ||= @max_position_embeddings
+          @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 !@hidden_size_per_head.nil?
+            @hidden_size_per_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,
+            "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 _to_block_types
+          return @block_types if @block_types.is_a?(Array) && !@block_types.empty?
+
+          count = @num_hidden_layers.to_i
+          return nil if count <= 0
+
+          if @full_attention_idx.is_a?(Array) && !@full_attention_idx.empty?
+            full_attention = @full_attention_idx.map(&:to_i)
+            return Array.new(count) { |i| full_attention.include?(i) ? "attention" : "recurrent" }
+          end
+
+          return Array.new(count, "attention") unless @mamba_enabled
+
+          step = @mamba_step.to_i
+          step = 2 if step <= 1
+          midpoint = step / 2
+
+          if count <= midpoint
+            return Array.new(count) { |i| i == count - 1 ? "attention" : "recurrent" }
+          end
+
+          Array.new(count) { |i| (i % step) == midpoint ? "attention" : "recurrent" }
+        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.is_a?(Hash) ? weights.dup : weights.to_h
+          _split_gate_up_proj!(normalized)
+
+          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 _split_gate_up_proj!(weights)
+          mx = MLX::Core
+          pattern = /\A(model\.layers(?:\.layers)?\.\d+\.mlp)\.gate_up_proj\.(weight|bias|scales|biases)\z/
+
+          weights.keys.each do |key|
+            match = pattern.match(key)
+            next unless match
+
+            prefix = match[1]
+            param = match[2]
+            gate_up = weights.delete(key)
+            mid = gate_up.shape[0] / 2
+            next if mid <= 0
+
+            gate_proj, up_proj = mx.split(gate_up, [mid], 0)
+            weights["#{prefix}.gate_proj.#{param}"] = gate_proj
+            weights["#{prefix}.up_proj.#{param}"] = up_proj
+          end
+        end
+
+        def _remap_weight_key(key)
+          mapped = key.dup
+          mapped = mapped.gsub("model.layers.layers.", "model.layers.")
+          mapped = mapped.gsub("model.norm.", "model.final_layernorm.")
+
+          mapped = mapped.gsub(/\.layers\.(\d+)\.pre_mixer_norm\./) { ".layers.#{$1}.input_layernorm." }
+          mapped = mapped.gsub(/\.layers\.(\d+)\.pre_mlp_norm\./) { ".layers.#{$1}.pre_ff_layernorm." }
+
+          mapped = mapped.gsub(".mixer.conv1d.", ".mamba.conv1d.")
+          mapped = mapped.gsub(".mixer.in_proj.", ".mamba.in_proj.")
+          mapped = mapped.gsub(".mixer.out_proj.", ".mamba.out_proj.")
+          mapped = mapped.gsub(".mixer.qkv_proj.", ".self_attn.q_proj.")
+          mapped = mapped.gsub(".mixer.q_proj.", ".self_attn.q_proj.")
+          mapped = mapped.gsub(".mixer.k_proj.", ".self_attn.k_proj.")
+          mapped = mapped.gsub(".mixer.v_proj.", ".self_attn.v_proj.")
+          mapped = mapped.gsub(".mixer.o_proj.", ".self_attn.o_proj.")
+          mapped = mapped.gsub(".mlp.gate_up_proj.", ".feed_forward.gate_proj.")
+          mapped = mapped.gsub(".mlp.gate_proj.", ".feed_forward.gate_proj.")
+          mapped = mapped.gsub(".mlp.up_proj.", ".feed_forward.up_proj.")
+          mapped = mapped.gsub(".mlp.down_proj.", ".feed_forward.down_proj.")
+          mapped
+        end
+      end
+
+      Models.register("plamo2", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/qwen.rb

@@ -0,0 +1,175 @@
+module MlxLm
+  module Models
+    module Qwen
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "qwen"
+        field :hidden_size, default: 2048
+        field :num_attention_heads, default: 16
+        field :num_hidden_layers, default: 24
+        field :kv_channels, default: 128
+        field :max_position_embeddings, default: 8192
+        field :layer_norm_epsilon, default: 1e-6
+        field :intermediate_size, default: 11008
+        field :no_bias, default: true
+        field :vocab_size, default: 151936
+        field :num_key_value_heads, default: nil
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+
+          hidden_size = args.hidden_size
+          @num_attention_heads = args.num_attention_heads
+          hidden_size_per_attention_head = hidden_size / @num_attention_heads
+
+          self.rotary_emb = MLX::NN::RoPE.new(
+            hidden_size_per_attention_head,
+            traditional: false
+          )
+
+          @proj_size = args.kv_channels * @num_attention_heads
+
+          self.c_attn = MLX::NN::Linear.new(hidden_size, @proj_size * 3, bias: true)
+          self.c_proj = MLX::NN::Linear.new(hidden_size, @proj_size, bias: !args.no_bias)
+
+          @head_dim = args.kv_channels
+          @scale = hidden_size_per_attention_head**(-0.5)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+
+          qkv = c_attn.call(x)
+          q, k, v = mx.split(qkv, [@proj_size, 2 * @proj_size], -1)
+
+          b, l, _ = q.shape
+
+          queries = q.reshape([b, l, @num_attention_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = k.reshape([b, l, @num_attention_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = v.reshape([b, l, @num_attention_heads, @head_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
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @proj_size])
+
+          c_proj.call(output)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+
+          self.w1 = MLX::NN::Linear.new(
+            args.hidden_size,
+            args.intermediate_size / 2,
+            bias: !args.no_bias
+          )
+          self.w2 = MLX::NN::Linear.new(
+            args.hidden_size,
+            args.intermediate_size / 2,
+            bias: !args.no_bias
+          )
+          self.c_proj = MLX::NN::Linear.new(
+            args.intermediate_size / 2,
+            args.hidden_size,
+            bias: !args.no_bias
+          )
+        end
+
+        def call(x)
+          a1 = w1.call(x)
+          a2 = w2.call(x)
+          c_proj.call(Activations.swiglu(a2, a1))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+
+          self.ln_1 = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.layer_norm_epsilon)
+          self.attn = Attention.new(args)
+          self.ln_2 = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.layer_norm_epsilon)
+          self.mlp = MLP.new(args)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          residual = x
+          x = ln_1.call(x)
+          x = attn.call(x, mask: mask, cache: cache)
+          residual = x + residual
+          x = ln_2.call(residual)
+          x = mlp.call(x)
+          x + residual
+        end
+      end
+
+      class QwenModel < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.wte = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.h = Array.new(args.num_hidden_layers) { TransformerBlock.new(args) }
+          self.ln_f = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.layer_norm_epsilon)
+        end
+
+        def call(inputs, cache: nil)
+          x = wte.call(inputs)
+          layer_cache = cache || [nil] * h.length
+
+          mask = nil
+          mask = "causal" if x.shape[1] > 1
+
+          h.each_with_index do |layer, i|
+            x = layer.call(x, mask: mask, cache: layer_cache[i])
+          end
+
+          ln_f.call(x)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(config)
+          super()
+          @args = config
+          self.model_type = config.model_type
+          self.transformer = QwenModel.new(config)
+          self.lm_head = MLX::NN::Linear.new(
+            config.hidden_size,
+            config.vocab_size,
+            bias: !config.no_bias
+          )
+        end
+
+        def call(x, cache: nil)
+          y = transformer.call(x, cache: cache)
+          lm_head.call(y)
+        end
+
+        def sanitize(weights)
+          weights.reject { |k, _| k.include?("rotary_emb.inv_freq") }
+        end
+
+        def layers
+          transformer.h
+        end
+      end
+
+      Models.register("qwen", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/qwen2.rb

@@ -0,0 +1,162 @@
+module MlxLm
+  module Models
+    module Qwen2
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "qwen2"
+        field :hidden_size, default: 4096
+        field :num_hidden_layers, default: 32
+        field :num_attention_heads, default: 32
+        field :num_key_value_heads, default: 32
+        field :intermediate_size, default: 22016
+        field :vocab_size, default: 151936
+        field :rms_norm_eps, default: 1e-6
+        field :rope_theta, default: 1000000.0
+        field :rope_traditional, default: false
+        field :rope_scaling, default: nil
+        field :tie_word_embeddings, default: true
+        field :head_dim, default: nil
+        field :max_position_embeddings, default: 32768
+
+        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)
+
+          # Qwen2: Q, K, V have bias; O does not
+          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: false)
+
+          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
+
+      class MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          hidden_dim = args.intermediate_size
+
+          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)
+          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 Qwen2Model < 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)
+          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 = Qwen2Model.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("qwen2", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/qwen2_moe.rb

@@ -0,0 +1,189 @@
+require_relative "activations"
+require_relative "qwen2"
+require_relative "switch_layers"
+
+module MlxLm
+  module Models
+    module Qwen2Moe
+      class ModelArgs < Qwen2::ModelArgs
+        field :model_type, default: "qwen2_moe"
+        field :num_key_value_heads, default: nil
+        field :num_experts_per_tok
+        field :num_experts
+        field :moe_intermediate_size
+        field :shared_expert_intermediate_size
+        field :tie_word_embeddings, default: false
+
+        def initialize(**kwargs)
+          super
+          validate_rope_scaling!
+        end
+
+        private
+
+        def validate_rope_scaling!
+          return unless @rope_scaling
+
+          required_keys = %w[factor type]
+          unless required_keys.all? { |key| _rope_scaling_has_key?(key) }
+            raise ArgumentError, "rope_scaling must contain keys #{required_keys}"
+          end
+
+          return if _rope_scaling_value("type") == "linear"
+
+          raise ArgumentError, "rope_scaling 'type' currently only supports 'linear'"
+        end
+
+        def _rope_scaling_has_key?(key)
+          @rope_scaling.key?(key) || @rope_scaling.key?(key.to_sym)
+        end
+
+        def _rope_scaling_value(key)
+          return @rope_scaling[key] if @rope_scaling.key?(key)
+
+          @rope_scaling[key.to_sym]
+        end
+      end
+
+      class SharedExpertMLP < 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 SparseMoeBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          intermediate_size = args.moe_intermediate_size
+          shared_expert_intermediate_size = args.shared_expert_intermediate_size
+
+          @num_experts = args.num_experts
+          @top_k = args.num_experts_per_tok
+
+          self.gate = MLX::NN::Linear.new(dim, @num_experts, bias: false)
+          self.switch_mlp = SwitchLayers::SwitchGLU.new(dim, intermediate_size, @num_experts)
+
+          self.shared_expert = SharedExpertMLP.new(dim, shared_expert_intermediate_size)
+          self.shared_expert_gate = MLX::NN::Linear.new(dim, 1, bias: false)
+        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)
+
+          y = switch_mlp.call(x, inds)
+          y = mx.sum(y * mx.expand_dims(scores, -1), -2)
+
+          shared_expert_output = shared_expert.call(x)
+          shared_expert_output = mx.sigmoid(shared_expert_gate.call(x)) * shared_expert_output
+
+          y + shared_expert_output
+        end
+      end
+
+      class DecoderLayer < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Qwen2::Attention.new(args)
+          self.mlp = SparseMoeBlock.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 Qwen2MoeModel < 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::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, 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 = Qwen2MoeModel.new(args)
+          self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+        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.mlp.experts.0.up_proj.weight")
+
+          mx = MLX::Core
+          result = weights.dup
+
+          @args.num_hidden_layers.times do |layer_idx|
+            prefix = "model.layers.#{layer_idx}"
+            %w[up_proj down_proj gate_proj].each do |projection|
+              %w[weight scales biases].each do |param|
+                first_key = "#{prefix}.mlp.experts.0.#{projection}.#{param}"
+                next unless result.key?(first_key)
+
+                expert_keys = (0...@args.num_experts).map do |expert_idx|
+                  "#{prefix}.mlp.experts.#{expert_idx}.#{projection}.#{param}"
+                end
+                next unless expert_keys.all? { |key| result.key?(key) }
+
+                stacked = expert_keys.map { |key| result.delete(key) }
+                result["#{prefix}.mlp.switch_mlp.#{projection}.#{param}"] = mx.stack(stacked)
+              end
+            end
+          end
+
+          result
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("qwen2_moe", Model, ModelArgs)
+    end
+  end
+end