Fixed weight loading from original Phi2 model

config.json

@@ -13,7 +13,7 @@
     "torch_dtype": "float16",
     "transformers_version": "4.29.0",
 
-    "vocab_size": 50304,
+    "vocab_size": 51200,
     "vocab_chunk_for_gpu_efficiency": 64,
     "initial_cos_sin_cache_len": 2048,
     "d_embedding": 2560,

phi2_model.py

@@ -13,11 +13,6 @@ class Phi2PreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = False
     # _no_split_modules = ["ParallelAttentionBlock"]
 
-    # weight loading
-    # base_model_prefix = "transformer"
-    # _keys_to_ignore_on_load_missing = [""]
-    # _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
-
     def __init__(self, config: Phi2Config):
         super().__init__(config)
         self.config = config
@@ -42,6 +37,7 @@ class Phi2PreTrainedModel(PreTrainedModel):
         input_ids: torch.LongTensor,  # dim: (batch_size, seq_len)
         kv_cache: KVCache | None = None,
         key_padding_mask: torch.LongTensor | torch.BoolTensor | None = None,
+        **kwargs,
     ) -> dict[str, Any]:
         if not kv_cache:
             kv_cache = KVCache(
@@ -142,7 +138,7 @@ class Phi2Model(Phi2PreTrainedModel):
 class Phi2ModelForCausalLM(Phi2PreTrainedModel):
     def __init__(self, config: Phi2Config) -> None:
         super().__init__(config)
-        self.pretrained_model = Phi2Model(config)
+        self.model = Phi2Model(config)
         self.lm_head_layer_norm = nn.LayerNorm(config.d_embedding, eps=config.layer_norm_epsilon)
         self.lm_head_linear = nn.Linear(config.d_embedding, config.vocab_size)
         self.loss_fn = nn.CrossEntropyLoss()
@@ -154,8 +150,9 @@ class Phi2ModelForCausalLM(Phi2PreTrainedModel):
         kv_cache: KVCache | None = None,
         key_padding_mask: torch.BoolTensor | None = None,
         labels: torch.LongTensor | None = None,
+        **kwargs,
     ) -> CausalLMOutputWithPast:
-        x = self.pretrained_model(input_ids, kv_cache=kv_cache, key_padding_mask=key_padding_mask)
+        x = self.model(input_ids, kv_cache=kv_cache, key_padding_mask=key_padding_mask)
         x = self.lm_head_layer_norm(x)
         logits = self.lm_head_linear(x).to(torch.float32)
         loss = (

streaming_inference.py

@@ -20,22 +20,23 @@ if __name__ == "__main__":
     phi_model_state_dict = phi_model.state_dict()
     model_state_dict = {}
     for key, value in phi_model_state_dict.items():
-        # transformer.embd.wte.weight -> model.rotary_embedding.embeddings.weight
-        # transformer.h.0.mlp.fc1.weight -> pretrained_model.parallel_blocks.0.mlp.fc1.weight
-        # transformer.h.0.ln.weight -> pretrained_model.parallel_blocks.0.layer_norm.weight
-        # transformer.h.0.mixer.Wqkv.weight -> pretrained_model.parallel_blocks.0.multi_head_attention.Wqkv.weight
-        # transformer.h.0.mixer.out_proj.weight -> pretrained_model.parallel_blocks.0.multi_head_attention.fc_out.weight
         # lm_head.ln.weight -> lm_head_layer_norm.weight
         # lm_head.linear.weight -> lm_head_linear.weight
+        # transformer.embd.wte.weight -> model.rotary_embedding.embeddings.weight
+        # transformer.h.0.mlp.fc1.weight -> model.parallel_blocks.0.mlp.fc1.weight
+        # transformer.h.0.ln.weight -> model.parallel_blocks.0.layer_norm.weight
+        # transformer.h.0.mixer.Wqkv.weight -> model.parallel_blocks.0.multi_head_attention.Wqkv.weight
+        # transformer.h.0.mixer.out_proj.weight -> model.parallel_blocks.0.multi_head_attention.fc_out.weight
         if key.startswith("transformer"):
-            key.replace("transformer.", "model.")
-            key.replace(".embd.wte.", ".rotary_embedding.embeddings.")
-            key.replace(".h.", ".parallel_blocks")
-            key.replace(".ln.", ".layer_norm.")
-            key.replace(".mixer.Wqkv.", ".multi_head_attention.Wqkv.")
-            key.replace(".mixer.out_proj.", ".multi_head_attention.fc_out.")
-            key.replace(".lm_head.ln.", ".lm_head_layer_norm.")
-            key.replace(".lm_head.linear.", ".lm_head_linear.")
+            key = key.replace("transformer.", "model.")
+            key = key.replace(".embd.wte.", ".rotary_embedding.embeddings.")
+            key = key.replace(".h.", ".parallel_blocks.")
+            key = key.replace(".ln.", ".layer_norm.")
+            key = key.replace(".mixer.Wqkv.", ".multi_head_attention.Wqkv.")
+            key = key.replace(".mixer.out_proj.", ".multi_head_attention.fc_out.")
+        else:
+            key = key.replace("lm_head.ln.", "lm_head_layer_norm.")
+            key = key.replace("lm_head.linear.", "lm_head_linear.")
         model_state_dict[key] = value
     model.load_state_dict(model_state_dict)