mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/nanochat.rb

@@ -0,0 +1,167 @@
+module MlxLm
+  module Models
+    module Nanochat
+      module_function
+
+      def rms_norm(x, eps: 1e-5)
+        mx = MLX::Core
+        variance = mx.mean(mx.square(x), -1, true)
+        mx.multiply(x, mx.rsqrt(mx.add(variance, eps)))
+      end
+
+      def softcap(logits, cap: 15.0)
+        mx = MLX::Core
+        mx.multiply(cap, mx.tanh(mx.divide(logits, cap)))
+      end
+
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "nanochat"
+        field :hidden_size, default: 1280
+        field :num_hidden_layers, default: 20
+        field :num_attention_heads, default: 10
+        field :num_key_value_heads, default: 10
+        field :vocab_size, default: 65_536
+        field :max_position_embeddings, default: 2048
+        field :intermediate_size, default: 5120
+        field :rope_theta, default: 10_000.0
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+
+          @hidden_size = args.hidden_size
+          @num_heads = args.num_attention_heads
+          @num_kv_heads = args.num_key_value_heads
+          @head_dim = @hidden_size / @num_heads
+          @scale = @head_dim**(-0.5)
+          @rope_theta = args.rope_theta
+
+          self.c_q = MLX::NN::Linear.new(@hidden_size, @num_heads * @head_dim, bias: false)
+          self.c_k = MLX::NN::Linear.new(@hidden_size, @num_kv_heads * @head_dim, bias: false)
+          self.c_v = MLX::NN::Linear.new(@hidden_size, @num_kv_heads * @head_dim, bias: false)
+          self.c_proj = MLX::NN::Linear.new(@hidden_size, @hidden_size, bias: false)
+
+          mx = MLX::Core
+          exponent = mx.multiply(
+            mx.arange(0, @head_dim, 2, mx.float32),
+            Math.log(@rope_theta) / @head_dim.to_f
+          )
+          self._rope_freqs = mx.multiply(-1.0, mx.exp(exponent))
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          queries = c_q.call(x)
+          keys = c_k.call(x)
+          values = c_v.call(x)
+
+          queries = queries.reshape([b, l, @num_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = keys.reshape([b, l, @num_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = values.reshape([b, l, @num_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          offset = cache ? cache.offset : 0
+          queries = _apply_rotary_emb(queries, offset: offset)
+          keys = _apply_rotary_emb(keys, offset: offset)
+
+          queries = Nanochat.rms_norm(queries)
+          keys = Nanochat.rms_norm(keys)
+
+          if cache
+            keys, values = cache.update_and_fetch(keys, values)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @hidden_size])
+          c_proj.call(output)
+        end
+
+        private
+
+        def _apply_rotary_emb(x, offset:)
+          MLX::Core.rope(x, @head_dim, false, nil, 1.0, offset, _rope_freqs)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.c_fc = MLX::NN::Linear.new(args.hidden_size, args.intermediate_size, bias: false)
+          self.c_proj = MLX::NN::Linear.new(args.intermediate_size, args.hidden_size, bias: false)
+        end
+
+        def call(x)
+          c_proj.call(MLX::NN.relu2(c_fc.call(x)))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.attn = Attention.new(args)
+          self.mlp = MLP.new(args)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          h = x + attn.call(Nanochat.rms_norm(x), mask: mask, cache: cache)
+          h + mlp.call(Nanochat.rms_norm(h))
+        end
+      end
+
+      class NanoChatModel < 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) }
+        end
+
+        def call(inputs, cache: nil)
+          hidden = wte.call(inputs)
+          hidden = Nanochat.rms_norm(hidden)
+
+          layer_cache = cache || [nil] * h.length
+          mask = _create_attention_mask(hidden, layer_cache[0])
+
+          h.each_with_index do |layer, i|
+            hidden = layer.call(hidden, mask: mask, cache: layer_cache[i])
+          end
+
+          Nanochat.rms_norm(hidden)
+        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()
+          self.args = args
+          self.model_type = args.model_type
+          self.transformer = NanoChatModel.new(args)
+          self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+        end
+
+        def call(inputs, cache: nil)
+          out = transformer.call(inputs, cache: cache)
+          logits = lm_head.call(out)
+          Nanochat.softcap(logits)
+        end
+
+        def layers
+          transformer.h
+        end
+      end
+
+      Models.register("nanochat", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/nemotron.rb

@@ -0,0 +1,202 @@
+module MlxLm
+  module Models
+    module Nemotron
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "nemotron"
+        field :hidden_size
+        field :hidden_act
+        field :num_hidden_layers
+        field :intermediate_size
+        field :num_attention_heads
+        field :norm_eps
+        field :vocab_size
+        field :num_key_value_heads
+        field :head_dim, default: nil
+        field :max_position_embeddings, default: nil
+        field :attention_bias, default: false
+        field :mlp_bias, default: false
+        field :partial_rotary_factor, default: 0.5
+        field :rope_theta, default: 10_000.0
+        field :rope_traditional, default: false
+        field :rope_scaling, default: nil
+        field :tie_word_embeddings, default: false
+
+        def initialize(**kwargs)
+          super
+          @head_dim ||= @hidden_size / @num_attention_heads
+          validate_rope_scaling!
+        end
+
+        private
+
+        def rope_scaling_value(key)
+          return nil unless @rope_scaling
+
+          @rope_scaling[key] || @rope_scaling[key.to_s]
+        end
+
+        def validate_rope_scaling!
+          return unless @rope_scaling
+
+          raise ArgumentError, "rope_scaling must contain 'factor'" if rope_scaling_value(:factor).nil?
+
+          rope_type = rope_scaling_value(:type) || rope_scaling_value(:rope_type)
+          if rope_type.nil?
+            raise ArgumentError, "rope_scaling must contain either 'type' or 'rope_type'"
+          end
+          return if rope_type == "linear"
+
+          raise ArgumentError, "rope_scaling 'type' currently only supports 'linear'"
+        end
+      end
+
+      class NemotronLayerNorm1P < MLX::NN::LayerNorm
+        def call(x)
+          w = state.key?("weight") ? weight + 1.0 : nil
+          b = state.key?("bias") ? bias : nil
+          MLX::Core.layer_norm(x, w, b, @eps)
+        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
+          @partial_rotary_factor = args.partial_rotary_factor
+          @scale = @head_dim**(-0.5)
+
+          bias = args.attention_bias
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: bias)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: bias)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: bias)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: bias)
+
+          rope_scale = 1.0
+          if args.rope_scaling
+            rope_type = args.rope_scaling[:type] || args.rope_scaling["type"] ||
+              args.rope_scaling[:rope_type] || args.rope_scaling["rope_type"]
+            if rope_type == "linear"
+              factor = args.rope_scaling[:factor] || args.rope_scaling["factor"]
+              rope_scale = 1.0 / factor.to_f
+            end
+          end
+
+          self.rope = MLX::NN::RoPE.new(
+            (@partial_rotary_factor * @head_dim).to_i,
+            base: args.rope_theta,
+            scale: rope_scale
+          )
+        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
+          bias = args.mlp_bias
+          self.down_proj = MLX::NN::Linear.new(hidden_dim, dim, bias: bias)
+          self.up_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: bias)
+        end
+
+        def call(x)
+          down_proj.call(MLX::NN.relu2(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 = NemotronLayerNorm1P.new(args.hidden_size, eps: args.norm_eps)
+          self.post_attention_layernorm = NemotronLayerNorm1P.new(args.hidden_size, eps: args.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 NemotronModel < 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 = NemotronLayerNorm1P.new(args.hidden_size, eps: args.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_type = args.model_type
+          self.model = NemotronModel.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 layers
+          model.layers
+        end
+      end
+
+      Models.register("nemotron", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/nemotron_h.rb

@@ -0,0 +1,212 @@
+require_relative "falcon_h1"
+
+module MlxLm
+  module Models
+    module NemotronH
+      class ModelArgs < FalconH1::ModelArgs
+        field :model_type, default: "nemotron_h"
+        field :tie_word_embeddings, default: false
+        field :mamba_num_heads, default: nil
+        field :mamba_head_dim, default: nil
+        field :mamba_proj_bias, default: nil
+        field :ssm_state_size, default: nil
+        field :conv_kernel, default: nil
+        field :n_groups, default: nil
+        field :mlp_bias, default: nil
+        field :layer_norm_epsilon, default: nil
+        field :use_bias, default: nil
+        field :use_conv_bias, default: nil
+        field :hybrid_override_pattern, default: nil
+        field :moe_intermediate_size, default: nil
+        field :moe_shared_expert_intermediate_size, default: nil
+        field :n_group, default: nil
+        field :n_routed_experts, default: nil
+        field :n_shared_experts, default: nil
+        field :topk_group, default: nil
+        field :num_experts_per_tok, default: nil
+        field :norm_topk_prob, default: nil
+        field :routed_scaling_factor, default: nil
+        field :time_step_limit, default: nil
+        field :time_step_min, default: nil
+        field :time_step_max, default: nil
+
+        def initialize(**kwargs)
+          super
+
+          @mamba_d_conv = @conv_kernel if kwargs.key?(:conv_kernel) && !kwargs.key?(:mamba_d_conv) && !@conv_kernel.nil?
+          @rms_norm_eps = @layer_norm_epsilon if kwargs.key?(:layer_norm_epsilon) && !kwargs.key?(:rms_norm_eps) && !@layer_norm_epsilon.nil?
+          @num_attention_heads ||= @mamba_num_heads
+          @head_dim ||= @mamba_head_dim
+
+          pattern = _hybrid_pattern_array
+          @hybrid_override_pattern = pattern unless pattern.nil?
+          @hybrid_override_pattern ||= _default_hybrid_pattern
+
+          if @num_hidden_layers.nil? && @hybrid_override_pattern.is_a?(Array) && !@hybrid_override_pattern.empty?
+            @num_hidden_layers = @hybrid_override_pattern.length
+          end
+
+          @num_key_value_heads ||= @num_attention_heads
+          @mamba_d_conv ||= 4
+          @block_types ||= _to_block_types(@hybrid_override_pattern)
+        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 !@mamba_head_dim.nil?
+            @mamba_head_dim
+          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 || @moe_shared_expert_intermediate_size || hidden_size.to_i * 2,
+            "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 || @layer_norm_epsilon || 1e-5,
+            "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 _hybrid_pattern_array
+          return nil if @hybrid_override_pattern.nil?
+          return @hybrid_override_pattern if @hybrid_override_pattern.is_a?(Array)
+          return @hybrid_override_pattern.chars if @hybrid_override_pattern.is_a?(String)
+
+          nil
+        end
+
+        def _default_hybrid_pattern
+          count = @num_hidden_layers.to_i
+          return nil if count <= 0
+
+          Array.new(count) { |idx| idx.even? ? "*" : "M" }
+        end
+
+        def _to_block_types(pattern)
+          return @block_types if @block_types.is_a?(Array) && !@block_types.empty?
+          return nil unless pattern.is_a?(Array) && !pattern.empty?
+
+          pattern.map do |block_type|
+            case block_type.to_s
+            when "*"
+              "attention"
+            else
+              "recurrent"
+            end
+          end
+        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
+          _stack_experts!(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 _stack_experts!(weights)
+          mx = MLX::Core
+          grouped = Hash.new { |h, k| h[k] = [] }
+          pattern = /\A(backbone\.layers\.\d+\.mixer|model\.layers(?:\.layers)?\.\d+\.mixer)\.experts\.(\d+)\.(up_proj|down_proj)\.(weight|bias|scales|biases)\z/
+
+          weights.keys.each do |key|
+            match = pattern.match(key)
+            next unless match
+
+            prefix = match[1]
+            expert_idx = match[2].to_i
+            projection = match[3]
+            param = match[4]
+            grouped[[prefix, projection, param]] << [expert_idx, key]
+          end
+
+          grouped.each do |(prefix, projection, param), entries|
+            next if entries.empty?
+
+            stacked = entries.sort_by(&:first).map { |_, key| weights.delete(key) }
+            target = projection == "up_proj" ? "fc1" : "fc2"
+            weights["#{prefix}.switch_mlp.#{target}.#{param}"] = mx.stack(stacked)
+          end
+        end
+
+        def _remap_weight_key(key)
+          mapped = key.dup
+          mapped = mapped.gsub("backbone.embeddings.", "model.embed_tokens.")
+          mapped = mapped.gsub("backbone.norm_f.", "model.final_layernorm.")
+          mapped = mapped.gsub("backbone.layers.", "model.layers.")
+          mapped = mapped.gsub("model.layers.layers.", "model.layers.")
+
+          mapped = mapped.gsub(/\.layers\.(\d+)\.norm\./) { ".layers.#{$1}.input_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.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(".mixer.gate.", ".feed_forward.router.")
+          mapped = mapped.gsub(".mixer.switch_mlp.fc1.", ".feed_forward.switch_mlp.up_proj.")
+          mapped = mapped.gsub(".mixer.switch_mlp.fc2.", ".feed_forward.switch_mlp.down_proj.")
+          mapped = mapped.gsub(".mixer.shared_experts.up_proj.", ".feed_forward.up_proj.")
+          mapped = mapped.gsub(".mixer.shared_experts.down_proj.", ".feed_forward.down_proj.")
+          mapped = mapped.gsub(".mixer.up_proj.", ".feed_forward.up_proj.")
+          mapped = mapped.gsub(".mixer.down_proj.", ".feed_forward.down_proj.")
+          mapped
+        end
+      end
+
+      Models.register("nemotron_h", Model, ModelArgs)
+    end
+  end
+end