Add tie_word_embedding option for Qwen2 model (#2535)

Author: rickzx

python/mlc_llm/model/qwen2/qwen2_model.py

@@ -33,6 +33,7 @@ class QWen2Config(ConfigBase):  # pylint: disable=too-many-instance-attributes
     rms_norm_eps: float
     rope_theta: int
     vocab_size: int
+    tie_word_embeddings: bool = False
     context_window_size: int = 0
     prefill_chunk_size: int = 0
     tensor_parallel_shards: int = 1
@@ -120,6 +121,19 @@ def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id:
 }
 
 
+class Qwen2Embedding(nn.Embedding):
+    """The embedding module specialized for Qwen2 so that
+    it can be shared with the final lm_head.
+    """
+
+    def lm_head_forward(self, x: nn.Tensor):
+        """The lm_head forwarding, which transposes the weight and multiplies
+        with the input tensor.
+        """
+        weight = nn.op.permute_dims(self.weight)
+        return nn.op.matmul(x, weight, out_dtype="float32")
+
+
 class QWen2MLP(nn.Module):
     def __init__(self, config: QWen2Config):
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
@@ -185,7 +199,7 @@ def _apply_residual(self, out, residual):
 
 class QWen2Model(nn.Module):
     def __init__(self, config: QWen2Config):
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.embed_tokens = Qwen2Embedding(config.vocab_size, config.hidden_size)
         self.layers = nn.ModuleList(
             [QWen2DecoderLayer(config) for _ in range(config.num_hidden_layers)]
         )
@@ -202,7 +216,9 @@ def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache):
 class QWen2LMHeadModel(nn.Module):  # pylint: disable=too-many-instance-attributes
     def __init__(self, config: QWen2Config):
         self.model = QWen2Model(config)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.tie_word_embeddings = config.tie_word_embeddings
+        if not config.tie_word_embeddings:
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.dtype = config.dtype
         self.hidden_size = config.hidden_size
         self.num_hidden_layers = config.num_hidden_layers
@@ -231,7 +247,11 @@ def batch_forward(
         hidden_states = self.model(input_embeds, paged_kv_cache)
         if logit_positions is not None:
             hidden_states = op.take(hidden_states, logit_positions, axis=1)
-        logits = self.lm_head(hidden_states)
+
+        if self.tie_word_embeddings:
+            logits = self.model.embed_tokens.lm_head_forward(hidden_states)
+        else:
+            logits = self.lm_head(hidden_states)
         if logits.dtype != "float32":
             logits = logits.astype("float32")
         return logits
@@ -250,7 +270,10 @@ def _index(x: te.Tensor):  # x[:-1,:]
 
         hidden_states = self.model(input_embed, paged_kv_cache)
         hidden_states = op.tensor_expr_op(_index, name_hint="index", args=[hidden_states])
-        logits = self.lm_head(hidden_states)
+        if self.tie_word_embeddings:
+            logits = self.model.embed_tokens.lm_head_forward(hidden_states)
+        else:
+            logits = self.lm_head(hidden_states)
         if logits.dtype != "float32":
             logits = logits.astype("float32")
         return logits, paged_kv_cache
@@ -259,7 +282,10 @@ def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
         op_ext.configure()
 
         hidden_states = self.model(input_embed, paged_kv_cache)
-        logits = self.lm_head(hidden_states)
+        if self.tie_word_embeddings:
+            logits = self.model.embed_tokens.lm_head_forward(hidden_states)
+        else:
+            logits = self.lm_head(hidden_states)
         if logits.dtype != "float32":
             logits = logits.astype("float32")
         return logits, paged_kv_cache