weight conversion wip

Author: kanpuriyanawab

keras_hub/api/models/__init__.py

@@ -292,9 +292,7 @@
     QwenTokenizer as Qwen2Tokenizer,
 )
 from keras_hub.src.models.qwen_moe.qwen_moe_backbone import QwenMoeBackbone
-from keras_hub.src.models.qwen_moe.qwen_moe_causal_lm import (
-    QwenmOECausalLM as QwenMoeCausalLM,
-)
+from keras_hub.src.models.qwen_moe.qwen_moe_causal_lm import QwenMoeCausalLM
 from keras_hub.src.models.qwen_moe.qwen_moe_causal_lm_preprocessor import (
     QwenMoeCausalLMPreprocessor,
 )

keras_hub/api/tokenizers/__init__.py

@@ -34,6 +34,7 @@
 from keras_hub.src.models.qwen.qwen_tokenizer import (
     QwenTokenizer as Qwen2Tokenizer,
 )
+from keras_hub.src.models.qwen_moe.qwen_moe_tokenizer import QwenMoeTokenizer
 from keras_hub.src.models.roberta.roberta_tokenizer import RobertaTokenizer
 from keras_hub.src.models.siglip.siglip_tokenizer import SigLIPTokenizer
 from keras_hub.src.models.t5.t5_tokenizer import T5Tokenizer

keras_hub/src/models/qwen/qwen_decoder.py

@@ -7,8 +7,8 @@
 from keras_hub.src.layers.modeling.transformer_layer_utils import (
     merge_padding_and_attention_mask,
 )
-from keras_hub.src.models.qwen.qwen_attention import QwenAttention
 from keras_hub.src.models.qwen.qwen_layernorm import QwenLayerNorm
+from keras_hub.src.models.qwen_moe.qwen_moe_attention import QwenMoeAttention
 from keras_hub.src.utils.keras_utils import clone_initializer
 
 
@@ -79,7 +79,7 @@ def build(self, decoder_sequence_shape):
         self.hidden_dim = decoder_sequence_shape[-1]
 
         # Self attention layer.
-        self._self_attention_layer = QwenAttention(
+        self._self_attention_layer = QwenMoeAttention(
             num_query_heads=self.num_query_heads,
             num_key_value_heads=self.num_key_value_heads,
             rope_max_wavelength=self.rope_max_wavelength,

keras_hub/src/models/qwen_moe/README.md

@@ -10,4 +10,48 @@ Immediate TODOs
 1. What about new caching mechanism in HF?
 
 
-Reference - https://huggingface.co/docs/transformers/en/model_doc/qwen2_moe
+Reference - https://huggingface.co/docs/transformers/en/model_doc/qwen2_moe
+
+Model Architecture:
+
+```
+Qwen2MoeForCausalLM(
+  (model): Qwen2MoeModel(
+    (embed_tokens): Embedding(151936, 2048)
+    (layers): ModuleList(
+      (0-23): 24 x Qwen2MoeDecoderLayer(
+        (self_attn): Qwen2MoeSdpaAttention(
+          (q_proj): Linear(in_features=2048, out_features=2048, bias=True)
+          (k_proj): Linear(in_features=2048, out_features=2048, bias=True)
+          (v_proj): Linear(in_features=2048, out_features=2048, bias=True)
+          (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
+          (rotary_emb): Qwen2MoeRotaryEmbedding()
+        )
+        (mlp): Qwen2MoeSparseMoeBlock(
+          (gate): Linear(in_features=2048, out_features=60, bias=False)
+          (experts): ModuleList(
+            (0-59): 60 x Qwen2MoeMLP(
+              (gate_proj): Linear(in_features=2048, out_features=1408, bias=False)
+              (up_proj): Linear(in_features=2048, out_features=1408, bias=False)
+              (down_proj): Linear(in_features=1408, out_features=2048, bias=False)
+              (act_fn): SiLU()
+            )
+          )
+          (shared_expert): Qwen2MoeMLP(
+            (gate_proj): Linear(in_features=2048, out_features=5632, bias=False)
+            (up_proj): Linear(in_features=2048, out_features=5632, bias=False)
+            (down_proj): Linear(in_features=5632, out_features=2048, bias=False)
+            (act_fn): SiLU()
+          )
+          (shared_expert_gate): Linear(in_features=2048, out_features=1, bias=False)
+        )
+        (input_layernorm): Qwen2MoeRMSNorm((2048,), eps=1e-06)
+        (post_attention_layernorm): Qwen2MoeRMSNorm((2048,), eps=1e-06)
+      )
+    )
+    (norm): Qwen2MoeRMSNorm((2048,), eps=1e-06)
+    (rotary_emb): Qwen2MoeRotaryEmbedding()
+  )
+  (lm_head): Linear(in_features=2048, out_features=151936, bias=False)
+)
+```

keras_hub/src/models/qwen_moe/qwen_moe_causal_lm_preprocessor.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.causal_lm_preprocessor import CausalLMPreprocessor
-from keras_hub.src.models.qwen.qwen_tokenizer import QwenTokenizer
 from keras_hub.src.models.qwen_moe.qwen_moe_backbone import QwenMoeBackbone
+from keras_hub.src.models.qwen_moe.qwen_moe_tokenizer import QwenMoeTokenizer
 
 
 @keras_hub_export(
@@ -11,7 +11,7 @@
 )
 class QwenMoeCausalLMPreprocessor(CausalLMPreprocessor):
     backbone_cls = QwenMoeBackbone
-    tokenizer_cls = QwenTokenizer
+    tokenizer_cls = QwenMoeTokenizer
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)

keras_hub/src/models/qwen_moe/qwen_moe_decoder.py

@@ -13,10 +13,21 @@
 
 
 class QwenMoeMLP(keras.layers.Layer):
-    def __init__(self, intermediate_dim, hidden_dim, activation_fn="silu"):
+    def __init__(
+        self,
+        intermediate_dim,
+        hidden_dim,
+        activation_fn="silu",
+        layer_norm_epsilon=1e-5,
+        kernel_initializer="glorot_uniform",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
         self.intermediate_dim = intermediate_dim
         self.hidden_dim = hidden_dim
         self.activation_fn = activation_fn
+        self.kernel_initializer = kernel_initializer
+        self.layer_norm_epsilon = layer_norm_epsilon
 
     def build(self, decoder_sequence_shape):
         # Feedforward layers.
@@ -91,42 +102,59 @@ def __init__(
         num_experts,
         top_k,
         norm_topk_prob,
-        kernel_initializer,
+        kernel_initializer="glorot_uniform",
+        layer_norm_epsilon=1e-5,
+        **kwargs,
     ):
+        super().__init__(**kwargs)
         self.hidden_dim = hidden_dim
         self.moe_intermediate_dim = moe_intermediate_dim
         self.shared_expert_intermediate_dim = shared_expert_intermediate_dim
         self.num_experts = num_experts
         self.top_k = top_k
         self.norm_topk_prob = norm_topk_prob
         self.kernel_initializer = kernel_initializer
+        self.layer_norm_epsilon = layer_norm_epsilon
 
-    def build(self, input_shape):
-        self.gate_proj = keras.layers.Dense(
-            self.hidden_dim,
+    def build(self, decoder_sequence_shape):
+        self._sparse_feedforward_gate_dense = keras.layers.Dense(
+            self.num_experts,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             use_bias=False,
             dtype=self.dtype_policy,
-            name="sparse_block_gate_proj",
+            name="sparse_feedforward_gate_dense",
         )
+        self._sparse_feedforward_gate_dense.build(decoder_sequence_shape)
 
         self.experts = [
             QwenMoeMLP(
                 intermediate_dim=self.moe_intermediate_dim,
                 hidden_dim=self.hidden_dim,
+                kernel_initializer=self.kernel_initializer,
+                layer_norm_epsilon=self.layer_norm_epsilon,
             )
             for _ in range(self.num_experts)
         ]
-        self.shared_expert = QwenMoeMLP(
-            intermediate_dim=self.shared_expert_intermediate_dim
+        for expert in self.experts:
+            expert.build(decoder_sequence_shape)
+
+        self.shared_expert_dense = QwenMoeMLP(
+            intermediate_dim=self.shared_expert_intermediate_dim,
+            hidden_dim=self.hidden_dim,
+            kernel_initializer=self.kernel_initializer,
+            layer_norm_epsilon=self.layer_norm_epsilon,
         )
-        self.shared_expert_gate_proj = keras.layers.Dense(1, use_bias=False)
+        self.shared_expert_dense.build(decoder_sequence_shape)
+
+        self.shared_expert_gate_dense = keras.layers.Dense(1, use_bias=False)
+        self.shared_expert_gate_dense.build(decoder_sequence_shape)
+        self.built = True
 
     def call(self, hidden_states):
         batch_size, seq_len, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.reshape(-1, hidden_dim)
 
-        router_logits = self.gate_proj(hidden_states)
+        router_logits = self._sparse_feedforward_gate_dense(hidden_states)
 
         routing_weights = ops.softmax(router_logits, axis=1)
         routing_weights, selected_experts = ops.top_k(
@@ -175,7 +203,7 @@ def call(self, hidden_states):
 
         shared_expert_output = self.shared_expert(hidden_states)
         shared_expert_output = (
-            ops.sigmoid(self.shared_expert_gate_proj(hidden_states))
+            ops.sigmoid(self.shared_expert_gate_dense(hidden_states))
             * shared_expert_output
         )
 
@@ -210,6 +238,7 @@ def __init__(
         sliding_window_size=4096,
         layer_index=0,
         mlp_only_layers=[],
+        output_router_logits=False,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -238,6 +267,7 @@ def __init__(
         self.top_k = top_k
         self.norm_topk_prob = norm_topk_prob
         self.decoder_sparse_step = decoder_sparse_step
+        self.output_router_logits = output_router_logits
 
         self.supports_masking = True
 
@@ -287,11 +317,20 @@ def build(self, decoder_sequence_shape):
                 norm_topk_prob=self.norm_topk_prob,
                 kernel_initializer=self.kernel_initializer,
             )
+            self.mlp.build(decoder_sequence_shape)
         else:
             self.mlp = QwenMoeMLP(
                 intermediate_dim=self.intermediate_dim,
                 hidden_dim=self.hidden_dim,
             )
+            self.mlp.build(decoder_sequence_shape)
+
+        self._feedforward_layernorm = QwenLayerNorm(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="feedforward_layernorm",
+        )
+        self._feedforward_layernorm.build(decoder_sequence_shape)
 
         self.built = True
 
@@ -301,7 +340,6 @@ def call(
         decoder_padding_mask=None,
         decoder_attention_mask=None,
         self_attention_cache=None,
-        output_router_logits=False,
         self_attention_cache_update_index=None,
         training=None,
     ):
@@ -364,7 +402,7 @@ def call(
         if self_attention_cache is not None:
             output += self_attention_cache
 
-        if output_router_logits:
+        if self.output_router_logits:
             output += (router_logits,)
 
         return output

keras_hub/src/models/qwen_moe/qwen_moe_tokenizer.py

@@ -0,0 +1,46 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.qwen_moe.qwen_moe_backbone import QwenMoeBackbone
+from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+
+
+@keras_hub_export(
+    "keras_hub.tokenizers.QwenMoeTokenizer",
+)
+class QwenMoeTokenizer(BytePairTokenizer):
+    """Tokenizer for Qwen models.
+
+    This tokenizer implements byte-pair encoding (BPE) for Qwen models,
+    handling special tokens like BOS (beginning of sequence) and EOS (end of
+    sequence).
+
+    Args:
+        vocabulary: Dictionary mapping tokens to token IDs, or path to
+            vocabulary file.
+        merges: List of BPE merges, or path to merges file.
+        bos_token: Beginning of sequence token. Defaults to None.
+        eos_token: End of sequence token. Defaults to "<|endoftext|>".
+        misc_special_tokens: Set of additional special tokens. Defaults to
+            empty set.
+    """
+
+    backbone_cls = QwenMoeBackbone
+
+    def __init__(
+        self,
+        vocabulary=None,
+        merges=None,
+        **kwargs,
+    ):
+        # Add EOS token
+        eos_token = "<|endoftext|>"
+        self._add_special_token(eos_token, "end_token")
+
+        self.start_token_id = None
+        self.start_token = None
+        self.pad_token_id = 0
+
+        super().__init__(
+            vocabulary=vocabulary,
+            merges=merges,
+            **kwargs,
+        )