mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/bailing_moe_linear.rb

@@ -0,0 +1,91 @@
+require_relative "bailing_moe"
+
+module MlxLm
+  module Models
+    module BailingMoeLinear
+      class ModelArgs < BailingMoe::ModelArgs
+        field :model_type, default: "bailing_moe_linear"
+        field :layer_group_size, default: nil
+        field :group_norm_size, default: nil
+        field :use_rmsnorm, default: nil
+        field :head_dim, default: nil
+        field :rope_traditional, default: false
+
+        def to_bailing_moe_dict
+          {
+            "model_type" => @model_type,
+            "hidden_size" => @hidden_size,
+            "intermediate_size" => @intermediate_size,
+            "max_position_embeddings" => @max_position_embeddings,
+            "moe_intermediate_size" => @moe_intermediate_size,
+            "num_experts" => @num_experts,
+            "num_shared_experts" => @num_shared_experts,
+            "norm_topk_prob" => @norm_topk_prob,
+            "num_attention_heads" => @num_attention_heads,
+            "num_experts_per_tok" => @num_experts_per_tok,
+            "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,
+            "first_k_dense_replace" => @first_k_dense_replace,
+            "rope_scaling" => @rope_scaling,
+            "use_bias" => @use_bias,
+            "use_qkv_bias" => @use_qkv_bias,
+            "norm_head" => @norm_head,
+            "norm_softmax" => @norm_softmax,
+            "use_qk_norm" => @use_qk_norm,
+            "tie_word_embeddings" => @tie_word_embeddings,
+            "partial_rotary_factor" => @partial_rotary_factor,
+            "rotary_dim" => @rotary_dim,
+            "moe_router_enable_expert_bias" => @moe_router_enable_expert_bias,
+            "moe_router_enable_routed_scaling" => @moe_router_enable_routed_scaling,
+            "routed_scaling_factor" => @routed_scaling_factor,
+            "score_function" => @score_function,
+            "n_group" => @n_group,
+            "topk_group" => @topk_group,
+            "moe_shared_expert_intermediate_size" => @moe_shared_expert_intermediate_size,
+            "moe_router_enable_shared_expert" => @moe_router_enable_shared_expert,
+          }
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.wrapped_model = BailingMoe::Model.new(BailingMoe::ModelArgs.from_dict(args.to_bailing_moe_dict))
+        end
+
+        def call(inputs, cache: nil)
+          wrapped_model.call(inputs, cache: cache)
+        end
+
+        def sanitize(weights)
+          wrapped_model.sanitize(weights)
+        end
+
+        def layers
+          wrapped_model.layers
+        end
+
+        def make_cache
+          return nil unless wrapped_model.respond_to?(:make_cache)
+
+          wrapped_model.make_cache
+        end
+
+        def cast_predicate
+          wrapped_model.cast_predicate
+        end
+
+        def quant_predicate
+          wrapped_model.quant_predicate
+        end
+      end
+
+      Models.register("bailing_moe_linear", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/bitlinear_layers.rb

@@ -0,0 +1,108 @@
+module MlxLm
+  module Models
+    class BitLinear < MLX::NN::Module
+      attr_reader :in_features, :out_features, :invert_weight_scales
+
+      def initialize(
+        in_features,
+        out_features,
+        bias: true,
+        invert_weight_scales: false
+      )
+        super()
+        mx = MLX::Core
+
+        @in_features = in_features
+        @out_features = out_features
+        @invert_weight_scales = invert_weight_scales
+
+        packed_out_features = (out_features + 3) / 4
+        self.weight = mx.zeros([packed_out_features, in_features], mx.uint8)
+        self.weight_scale = mx.array([1.0], dtype: mx.float32)
+        self.bias = mx.zeros([out_features], mx.float32) if bias
+      end
+
+      def call(x)
+        y = execute_matmul_kernel(x, weight)
+        state.key?("bias") ? MLX::Core.add(y, bias) : y
+      end
+
+      def execute_matmul_kernel(x, packed_weights)
+        # TODO(phase1e): switch to a custom Metal kernel once MLX Ruby exposes
+        # a stable fast-kernel API equivalent to Python's mx.fast.metal_kernel.
+        execute_matmul_fallback(x, packed_weights)
+      end
+
+      private
+
+      def execute_matmul_fallback(x, packed_weights)
+        input_dims = x.shape[-1]
+        unless input_dims == @in_features
+          raise ArgumentError, "Expected input features #{@in_features}, got #{input_dims}"
+        end
+
+        ternary_weight = unpack_packed_weights(packed_weights, x.dtype)
+        out = MLX::Core.matmul(x, ternary_weight.T)
+
+        scale = weight_scale.astype(x.dtype)
+        scale = MLX::Core.divide(1.0, scale) if invert_weight_scales
+        MLX::Core.multiply(out, scale)
+      end
+
+      def unpack_packed_weights(packed_weights, dtype)
+        mx = MLX::Core
+
+        w0 = (mx.bitwise_and(packed_weights, 0x03).astype(dtype) - 1.0)
+        w1 = (mx.bitwise_and(mx.right_shift(packed_weights, 2), 0x03).astype(dtype) - 1.0)
+        w2 = (mx.bitwise_and(mx.right_shift(packed_weights, 4), 0x03).astype(dtype) - 1.0)
+        w3 = (mx.bitwise_and(mx.right_shift(packed_weights, 6), 0x03).astype(dtype) - 1.0)
+
+        expanded = mx.concatenate([w0, w1, w2, w3], 0)
+        return expanded if expanded.shape[0] == @out_features
+
+        keep = mx.arange(0, @out_features, 1, mx.int32)
+        mx.take(expanded, keep, 0)
+      end
+    end
+
+    module_function
+
+    def bitnet_quantize(model, quantization_config = {})
+      modules_to_not_convert = Array(config_value(quantization_config, "modules_to_not_convert", []))
+                               .map(&:to_s)
+      invert_weight_scales = config_value(quantization_config, "linear_class", "").to_s != "autobitlinear"
+
+      replacements = []
+      leaves = model.leaf_modules
+      flat = MLX::Utils.tree_flatten(leaves, is_leaf: lambda { |node| node.is_a?(MLX::NN::Module) })
+
+      flat.each do |path, layer|
+        path_s = path.to_s
+        next if modules_to_not_convert.include?(path_s)
+        next unless layer.is_a?(MLX::NN::Linear)
+
+        out_features, in_features = layer.weight.shape
+        replacements << [
+          path_s,
+          BitLinear.new(
+            in_features,
+            out_features,
+            bias: layer.state.key?("bias"),
+            invert_weight_scales: invert_weight_scales
+          ),
+        ]
+      end
+
+      model.update_modules(MLX::Utils.tree_unflatten(replacements)) unless replacements.empty?
+      model
+    end
+
+    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
+    private_class_method :config_value
+  end
+end

data/lib/mlx_lm/models/bitnet.rb

@@ -0,0 +1,176 @@
+module MlxLm
+  module Models
+    module Bitnet
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "bitnet"
+        field :hidden_size, default: 4096
+        field :num_hidden_layers, default: 32
+        field :intermediate_size, default: 11_008
+        field :num_attention_heads, default: 32
+        field :num_key_value_heads, default: nil
+        field :rms_norm_eps, default: 1e-6
+        field :vocab_size, default: 32_000
+        field :head_dim, default: nil
+        field :max_position_embeddings, default: nil
+        field :attention_bias, default: false
+        field :mlp_bias, default: false
+        field :rope_theta, default: 10_000.0
+        field :rope_traditional, default: false
+        field :rope_scaling, default: nil
+        field :tie_word_embeddings, default: true
+
+        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)
+
+          bias = args.attention_bias
+          self.q_proj = BitLinear.new(dim, @n_heads * @head_dim, bias: bias)
+          self.k_proj = BitLinear.new(dim, @n_kv_heads * @head_dim, bias: bias)
+          self.v_proj = BitLinear.new(dim, @n_kv_heads * @head_dim, bias: bias)
+          self.o_proj = BitLinear.new(@n_heads * @head_dim, dim, bias: bias)
+
+          self.rope = MlxLm::Models.initialize_rope(
+            @head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            max_position_embeddings: args.max_position_embeddings
+          )
+          self.attn_sub_norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        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(attn_sub_norm.call(output))
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          hidden_dim = args.intermediate_size
+          bias = args.mlp_bias
+
+          self.gate_proj = BitLinear.new(dim, hidden_dim, bias: bias)
+          self.down_proj = BitLinear.new(hidden_dim, dim, bias: bias)
+          self.up_proj = BitLinear.new(dim, hidden_dim, bias: bias)
+          self.ffn_sub_norm = MLX::NN::RMSNorm.new(hidden_dim, eps: args.rms_norm_eps)
+        end
+
+        def call(x)
+          h = MLX::NN.relu2(gate_proj.call(x)) * up_proj.call(x)
+          down_proj.call(ffn_sub_norm.call(h))
+        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 BitnetModel < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) { TransformerBlock.new(args) }
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        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, 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()
+          @args = args
+          self.model_type = args.model_type
+          self.model = BitnetModel.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("bitnet", Model, ModelArgs)
+    end
+  end
+end