mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/baichuan_m1.rb

@@ -0,0 +1,306 @@
+module MlxLm
+  module Models
+    module BaichuanM1
+      class ModelArgs < BaseModelArgs
+        field :vocab_size
+        field :hidden_size
+        field :intermediate_size
+        field :num_hidden_layers
+        field :num_attention_heads
+        field :num_key_value_heads
+        field :rope_theta
+        field :sliding_window
+        field :sliding_window_layers
+        field :conv_window
+        field :rms_norm_eps
+        field :model_type, default: "baichuan_m1"
+        field :num_swa_attention_heads, default: nil
+        field :num_swa_key_value_heads, default: nil
+        field :tie_word_embeddings, default: false
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(config, layer_idx: nil)
+          super()
+
+          raise ArgumentError, "Layer index must be provided to Attention module." if layer_idx.nil?
+
+          swa_layers = config.sliding_window_layers || []
+          @is_swa = swa_layers.include?(layer_idx)
+
+          @num_heads = if @is_swa && config.num_swa_attention_heads
+            config.num_swa_attention_heads
+          else
+            config.num_attention_heads
+          end
+
+          @num_kv_heads = if @is_swa && config.num_swa_key_value_heads
+            config.num_swa_key_value_heads
+          else
+            config.num_key_value_heads
+          end
+
+          @hidden_size = config.hidden_size
+          @head_dim = @hidden_size / @num_heads
+
+          unless (@head_dim * @num_heads) == @hidden_size
+            raise ArgumentError, "hidden_size must be divisible by num_heads"
+          end
+
+          @scale = @head_dim**(-0.5)
+
+          self.w_pack = MLX::NN::Linear.new(
+            config.hidden_size,
+            @hidden_size + 2 * @num_kv_heads * @head_dim,
+            bias: false
+          )
+          self.o_proj = MLX::NN::Linear.new(
+            @num_heads * @head_dim,
+            config.hidden_size,
+            bias: false
+          )
+
+          self.rope = MLX::NN::RoPE.new(@head_dim, traditional: false, base: config.rope_theta)
+
+          @conv_window = config.conv_window
+          raise ArgumentError, "conv_window must be 2" unless @conv_window == 2
+
+          mx = MLX::Core
+          self.conv_k = mx.zeros([1, 1, @num_kv_heads, 1, @conv_window])
+          self.conv_v = mx.zeros([1, 1, @num_kv_heads, 1, @conv_window])
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, d = x.shape
+
+          proj = w_pack.call(x)
+          q, k, v = mx.split(proj, [d, d + @num_kv_heads * @head_dim], -1)
+
+          q = q.reshape([b, l, @num_heads, @head_dim]).transpose([0, 2, 1, 3])
+          k = k.reshape([b, l, @num_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          v = v.reshape([b, l, @num_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          layer_cache = cache || [nil, nil]
+          conv_cache = layer_cache[0]
+          kv_cache = layer_cache[1]
+
+          if conv_cache
+            offset = kv_cache.offset
+            last_k = conv_cache[0]
+            last_v = conv_cache[1]
+          else
+            offset = 0
+            last_k = nil
+            last_v = nil
+          end
+
+          k_init = k
+          v_init = v
+
+          k = _custom_convolution(k, conv_k, state: last_k)
+          v = _custom_convolution(v, conv_v, state: last_v)
+          q = rope.call(q, offset: offset)
+          k = rope.call(k, offset: offset)
+
+          if conv_cache
+            k, v = kv_cache.update_and_fetch(k, v)
+            if l > 0
+              conv_cache[0] = mx.split(k_init, [l - 1], 2)[1]
+              conv_cache[1] = mx.split(v_init, [l - 1], 2)[1]
+            end
+          end
+
+          out = mx.scaled_dot_product_attention(q, k, v, @scale, mask)
+          out = out.transpose([0, 2, 1, 3]).reshape([b, l, @num_heads * @head_dim])
+          o_proj.call(out)
+        end
+
+        private
+
+        def _custom_convolution(u, weights, state: nil)
+          mx = MLX::Core
+          b, h, l, d = u.shape
+
+          weights = weights.reshape([1, h, @conv_window, 1, 1])
+          w0 = mx.take(weights, 0, 2)
+          w1 = mx.take(weights, 1, 2)
+
+          state ||= mx.zeros([b, h, 1, d], u.dtype)
+          if l > 1
+            u_prev = mx.concatenate([state, mx.split(u, [l - 1], 2)[0]], 2)
+          else
+            u_prev = state
+          end
+
+          mx.add(mx.multiply(u_prev, w0), mx.multiply(u, w1))
+        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, layer_idx)
+          super()
+          self.self_attn = Attention.new(config, layer_idx: layer_idx)
+          self.mlp = MLP.new(config)
+          self.input_layernorm = MLX::NN::RMSNorm.new(config.hidden_size, eps: config.rms_norm_eps)
+          self.post_attention_layernorm = MLX::NN::RMSNorm.new(config.hidden_size, eps: config.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 BaichuanModel < MLX::NN::Module
+        def initialize(config)
+          super()
+          @config = config
+          @sliding_window = config.sliding_window
+          @swa_layers = config.sliding_window_layers || []
+
+          self.embed_tokens = MLX::NN::Embedding.new(config.vocab_size, config.hidden_size)
+          self.layers = Array.new(config.num_hidden_layers) { |i| DecoderLayer.new(config, i) }
+          self.norm = MLX::NN::RMSNorm.new(config.hidden_size, eps: config.rms_norm_eps)
+
+          self.first_swa_idx = @swa_layers.empty? ? nil : @swa_layers[0]
+          self.first_global_idx = nil
+          config.num_hidden_layers.times do |i|
+            next if @swa_layers.include?(i)
+
+            self.first_global_idx = i
+            break
+          end
+        end
+
+        def call(inputs, cache: nil)
+          x = embed_tokens.call(inputs)
+          layer_cache = cache || Array.new(layers.length) { [nil, nil] }
+
+          c_global = first_global_idx.nil? ? nil : layer_cache[first_global_idx][1]
+          c_swa = first_swa_idx.nil? ? nil : layer_cache[first_swa_idx][1]
+
+          global_mask = _create_attention_mask(x, c_global)
+          swa_mask = _create_attention_mask(x, c_swa, window_size: @sliding_window)
+
+          layers.each_with_index do |layer, i|
+            mask = @swa_layers.include?(i) ? swa_mask : global_mask
+            x = layer.call(x, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(x)
+        end
+
+        private
+
+        def _create_attention_mask(x, cache = nil, window_size: nil)
+          n = x.shape[1]
+          return cache.make_mask(n, window_size: window_size) if cache && cache.respond_to?(:make_mask)
+          return nil if n == 1
+          return _create_causal_mask(n, window_size: window_size) if window_size && n > window_size
+
+          "causal"
+        end
+
+        def _create_causal_mask(n, offset: 0, window_size: nil)
+          mx = MLX::Core
+          rinds = mx.arange(0, offset + n, 1, mx.int32).reshape([1, offset + n])
+          linds = mx.arange(offset, offset + n, 1, mx.int32).reshape([n, 1])
+
+          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(config)
+          super()
+          @config = config
+          self.model_type = config.model_type
+          self.model = BaichuanModel.new(config)
+          @tie_word_embeddings = config.tie_word_embeddings
+          unless @tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(config.hidden_size, config.vocab_size, bias: false)
+          end
+        end
+
+        def make_cache
+          caches = []
+          swa_layers = @config.sliding_window_layers || []
+          @config.num_hidden_layers.times do |i|
+            is_swa = swa_layers.include?(i)
+            conv_cache = MlxLm::ArraysCache.new(2)
+            kv_cache = if is_swa
+              MlxLm::RotatingKVCache.new(max_size: @config.sliding_window)
+            else
+              MlxLm::KVCache.new
+            end
+            caches << MlxLm::CacheList.new(conv_cache, kv_cache)
+          end
+          caches
+        end
+
+        def sanitize(weights)
+          mx = MLX::Core
+          is_quantized = weights.key?("lm_head.scales")
+
+          if !is_quantized && weights.key?("lm_head.weight")
+            w = weights["lm_head.weight"]
+            dtype = w.dtype
+            w = w.astype(mx.float32)
+            norm = mx.norm(w, nil, -1, true)
+            w = (w / (norm + 1e-7)).astype(dtype)
+            weights["lm_head.weight"] = w
+          end
+
+          weights
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          if @tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("baichuan_m1", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/bailing_moe.rb

@@ -0,0 +1,399 @@
+require_relative "activations"
+require_relative "rope_utils"
+require_relative "switch_layers"
+
+module MlxLm
+  module Models
+    module BailingMoe
+      class ModelArgs < BaseModelArgs
+        field :model_type
+        field :hidden_size
+        field :intermediate_size
+        field :max_position_embeddings
+        field :moe_intermediate_size
+        field :num_experts
+        field :num_shared_experts
+        field :norm_topk_prob
+        field :num_attention_heads
+        field :num_experts_per_tok
+        field :num_hidden_layers
+        field :num_key_value_heads
+        field :rms_norm_eps
+        field :rope_theta
+        field :vocab_size
+        field :first_k_dense_replace
+        field :rope_scaling, default: nil
+        field :use_bias, default: false
+        field :use_qkv_bias, default: false
+        field :norm_head, default: false
+        field :norm_softmax, default: false
+        field :use_qk_norm, default: false
+        field :tie_word_embeddings, default: false
+        field :partial_rotary_factor, default: 1.0
+        field :rotary_dim, default: nil
+        field :moe_router_enable_expert_bias, default: false
+        field :moe_router_enable_routed_scaling, default: true
+        field :routed_scaling_factor, default: 1.0
+        field :score_function, default: "softmax"
+        field :n_group, default: 1
+        field :topk_group, default: 4
+        field :moe_shared_expert_intermediate_size, default: nil
+        field :moe_router_enable_shared_expert, default: true
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+        end
+      end
+
+      module_function
+
+      def aggregate_expert_outputs(expert_outputs, scores)
+        mx = MLX::Core
+        mx.sum(expert_outputs * mx.expand_dims(scores, -1), -2).astype(expert_outputs.dtype)
+      end
+
+      def group_expert_select(
+        gates,
+        e_score_correction_bias,
+        top_k,
+        n_group,
+        topk_group,
+        routed_scaling_factor,
+        norm_topk_prob,
+        score_function
+      )
+        mx = MLX::Core
+        in_type = gates.dtype
+
+        scores = if score_function == "sigmoid"
+          mx.sigmoid(gates.astype(mx.float32))
+        else
+          mx.softmax(gates.astype(mx.float32), -1)
+        end
+        orig_scores = scores
+        scores = scores + e_score_correction_bias if e_score_correction_bias
+
+        if n_group.to_i > 1
+          experts_per_group = scores.shape[-1] / n_group
+          scores = mx.unflatten(scores, -1, [n_group, experts_per_group])
+          group_scores = mx.topk(scores, 2, -1)
+          group_scores = mx.expand_dims(mx.sum(group_scores, -1), -1)
+
+          drop_count = n_group - topk_group.to_i
+          if drop_count > 0
+            group_idx = mx.argpartition(group_scores, drop_count - 1, -2)
+            take_ids = mx.array((0...drop_count).to_a, dtype: mx.int32)
+            group_idx = mx.take(group_idx, take_ids, -2)
+            scores = mx.put_along_axis(
+              scores,
+              mx.stop_gradient(group_idx),
+              mx.array(0.0),
+              -2
+            )
+          end
+
+          scores = mx.flatten(scores, -2, -1)
+        end
+
+        k = [top_k.to_i, scores.shape[-1]].min
+        inds = mx.argpartition(scores * -1.0, k - 1, -1)
+        take_ids = mx.array((0...k).to_a, dtype: mx.int32)
+        inds = mx.take(inds, take_ids, -1)
+
+        selected_scores = mx.take_along_axis(orig_scores, inds, -1)
+        if k > 1 && norm_topk_prob
+          denominator = mx.expand_dims(mx.sum(selected_scores, -1), -1) + 1e-20
+          selected_scores = selected_scores / denominator
+        end
+
+        selected_scores = selected_scores * routed_scaling_factor.to_f
+        [inds, selected_scores.astype(in_type)]
+      end
+
+      class BailingMoeMLP < MLX::NN::Module
+        def initialize(args, intermediate_size: nil)
+          super()
+          hidden_dim = intermediate_size || args.intermediate_size
+
+          self.gate_proj = MLX::NN::Linear.new(args.hidden_size, hidden_dim, bias: args.use_bias)
+          self.down_proj = MLX::NN::Linear.new(hidden_dim, args.hidden_size, bias: args.use_bias)
+          self.up_proj = MLX::NN::Linear.new(args.hidden_size, hidden_dim, bias: args.use_bias)
+        end
+
+        def call(x)
+          down_proj.call(Activations.swiglu(gate_proj.call(x), up_proj.call(x)))
+        end
+      end
+
+      class BailingMoeAttention < MLX::NN::Module
+        def initialize(args)
+          super()
+          @use_qk_norm = args.use_qk_norm
+          @num_attention_heads = args.num_attention_heads
+          @num_key_value_heads = args.num_key_value_heads
+          @head_dim = args.hidden_size / @num_attention_heads
+          @scale = @head_dim**(-0.5)
+
+          self.query_key_value = MLX::NN::Linear.new(
+            args.hidden_size,
+            (@num_attention_heads + 2 * @num_key_value_heads) * @head_dim,
+            bias: args.use_qkv_bias
+          )
+          self.dense = MLX::NN::Linear.new(
+            @num_attention_heads * @head_dim,
+            args.hidden_size,
+            bias: args.use_bias
+          )
+
+          if @use_qk_norm
+            self.key_layernorm = MLX::NN::RMSNorm.new(@head_dim, eps: args.rms_norm_eps)
+            self.query_layernorm = MLX::NN::RMSNorm.new(@head_dim, eps: args.rms_norm_eps)
+          end
+
+          rope_dim = args.rotary_dim || (@head_dim * args.partial_rotary_factor.to_f).to_i
+          rope_dim = [rope_dim, 1].max
+          self.rope = MlxLm::Models.initialize_rope(
+            rope_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
+
+          qkv = query_key_value.call(x)
+
+          q_size = @num_attention_heads * @head_dim
+          kv_size = @num_key_value_heads * @head_dim
+          q, k, v = mx.split(qkv, [q_size, q_size + kv_size], -1)
+
+          queries = q.reshape([b, l, @num_attention_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = k.reshape([b, l, @num_key_value_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = v.reshape([b, l, @num_key_value_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          if @use_qk_norm
+            queries = query_layernorm.call(queries)
+            keys = key_layernorm.call(keys)
+          end
+
+          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, @num_attention_heads * @head_dim])
+          dense.call(output)
+        end
+      end
+
+      class BailingMoeGate < MLX::NN::Module
+        def initialize(args)
+          super()
+          @norm_topk_prob = args.norm_topk_prob
+          @top_k = args.num_experts_per_tok
+          @n_group = args.n_group
+          @topk_group = args.topk_group
+          @routed_scaling_factor = args.routed_scaling_factor
+          @score_function = args.score_function
+
+          self.gate_proj = MLX::NN::Linear.new(args.hidden_size, args.num_experts, bias: false)
+          self.expert_bias = if args.moe_router_enable_expert_bias
+            MLX::Core.zeros([args.num_experts])
+          else
+            nil
+          end
+        end
+
+        def call(x)
+          BailingMoe.group_expert_select(
+            gate_proj.call(x),
+            expert_bias,
+            @top_k,
+            @n_group,
+            @topk_group,
+            @routed_scaling_factor,
+            @norm_topk_prob,
+            @score_function
+          )
+        end
+      end
+
+      class BailingMoeSparseMoeBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.switch_mlp = SwitchLayers::SwitchGLU.new(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.num_experts,
+            bias: args.use_bias
+          )
+          self.gate = BailingMoeGate.new(args)
+
+          shared_dim = args.moe_shared_expert_intermediate_size || args.moe_intermediate_size
+          self.shared_experts = if args.num_shared_experts.to_i > 0 && args.moe_router_enable_shared_expert
+            BailingMoeMLP.new(
+              args,
+              intermediate_size: shared_dim * args.num_shared_experts
+            )
+          end
+        end
+
+        def call(x)
+          topk_idx, topk_weight = gate.call(x)
+          out = switch_mlp.call(x, topk_idx)
+          out = BailingMoe.aggregate_expert_outputs(out, topk_weight)
+          out = out + shared_experts.call(x) if respond_to?(:shared_experts)
+          out
+        end
+      end
+
+      class BailingMoeDecoderLayer < MLX::NN::Module
+        def initialize(args, layer_idx:)
+          super()
+          self.attention = BailingMoeAttention.new(args)
+          self.mlp = if !args.num_experts.nil? && layer_idx >= args.first_k_dense_replace
+            BailingMoeSparseMoeBlock.new(args)
+          else
+            BailingMoeMLP.new(args)
+          end
+          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 = attention.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 BailingMoeModel < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.word_embeddings = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) do |layer_idx|
+            BailingMoeDecoderLayer.new(args, layer_idx: layer_idx)
+          end
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil)
+          h = word_embeddings.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+          mask = _create_attention_mask(h, layer_cache[0])
+
+          layers.each_with_index do |layer, layer_idx|
+            h = layer.call(h, mask: mask, cache: layer_cache[layer_idx])
+          end
+          norm.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
+          @norm_head = args.norm_head
+          self.model_type = args.model_type
+          self.model = BailingMoeModel.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.word_embeddings.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
+
+          if @norm_head && result.key?("lm_head.weight")
+            w = result["lm_head.weight"]
+            dtype = w.dtype
+            w_fp32 = w.astype(mx.float32)
+            weight_norm = mx.sqrt(mx.sum(mx.square(w_fp32), 0, true)) + 1e-7
+            result["lm_head.weight"] = (w_fp32 / weight_norm).astype(dtype)
+          end
+
+          @args.num_hidden_layers.times do |layer_idx|
+            next if layer_idx < @args.first_k_dense_replace.to_i
+
+            prefix = "model.layers.#{layer_idx}"
+            %w[gate_proj down_proj up_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
+
+            if result.key?("#{prefix}.mlp.gate.weight")
+              result["#{prefix}.mlp.gate.gate_proj.weight"] = result.delete("#{prefix}.mlp.gate.weight")
+            end
+            if result.key?("#{prefix}.mlp.gate.bias")
+              result["#{prefix}.mlp.gate.gate_proj.bias"] = result.delete("#{prefix}.mlp.gate.bias")
+            end
+          end
+
+          result
+        end
+
+        def quant_predicate
+          lambda do |path, _|
+            if path.to_s.end_with?("mlp.gate.gate_proj")
+              { group_size: 64, bits: 8 }
+            else
+              true
+            end
+          end
+        end
+
+        def cast_predicate
+          lambda { |key| !key.to_s.include?("expert_bias") }
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("bailing_moe", Model, ModelArgs)
+    end
+  end
+end