[MODEL] Make constant vocab size for models instead dynamic

python/mlc_llm/model/gemma/gemma_model.py

@@ -195,7 +195,7 @@ class GemmaModel(nn.Module):
     def __init__(self, config: GemmaConfig):
         self.hidden_size = config.hidden_size
         assert config.hidden_size % config.num_attention_heads == 0
-        self.embed_tokens = GemmaEmbedding("vocab_size", config.hidden_size)
+        self.embed_tokens = GemmaEmbedding(config.vocab_size, config.hidden_size)
         self.layers = nn.ModuleList(
             [GemmaDecoderLayer(config) for _ in range(config.num_hidden_layers)]
         )

python/mlc_llm/model/llama/llama_model.py

@@ -178,7 +178,7 @@ def _apply_residual(self, out, residual):
 class LlamaModel(nn.Module):
     def __init__(self, config: LlamaConfig):
         assert config.hidden_size % config.num_attention_heads == 0
-        self.embed_tokens = nn.Embedding("vocab_size", config.hidden_size)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
         self.layers = nn.ModuleList(
             [LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)]
         )
@@ -195,7 +195,7 @@ def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
 class LlamaForCasualLM(nn.Module):  # pylint: disable=too-many-instance-attributes
     def __init__(self, config: LlamaConfig):
         self.model = LlamaModel(config)
-        self.lm_head = nn.Linear(config.hidden_size, "vocab_size", bias=False)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.num_hidden_layers = config.num_hidden_layers
         self.num_attention_heads = config.num_attention_heads
         self.num_key_value_heads = config.num_key_value_heads

python/mlc_llm/model/mistral/mistral_model.py

@@ -156,7 +156,7 @@ class MistralModel(nn.Module):
 
     def __init__(self, config: MistralConfig):
         assert config.hidden_size % config.num_attention_heads == 0
-        self.embed_tokens = nn.Embedding("vocab_size", config.hidden_size)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
         self.layers = nn.ModuleList(
             [MistralDecoderLayer(config) for _ in range(config.num_hidden_layers)]
         )
@@ -176,7 +176,7 @@ class MistralForCasualLM(nn.Module):  # pylint: disable=too-many-instance-attrib
 
     def __init__(self, config: MistralConfig):
         self.model = MistralModel(config)
-        self.lm_head = nn.Linear(config.hidden_size, "vocab_size", bias=False)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.num_hidden_layers = config.num_hidden_layers
         self.num_attention_heads = config.num_attention_heads
         self.num_key_value_heads = config.num_key_value_heads

python/mlc_llm/model/phi/phi_model.py

@@ -280,7 +280,7 @@ def __init__(self, config: PhiConfig) -> None:
         super().__init__()
 
         self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
-        self.linear = nn.Linear(config.n_embd, "vocab_size")
+        self.linear = nn.Linear(config.n_embd, config.vocab_size)
 
     def forward(self, hidden_states: Tensor):
         hidden_states = self.ln(hidden_states)