Renaming state dict keys from Phi2

phi2_model.py

@@ -91,7 +91,7 @@ class Embedding(nn.Module):
 class Phi2Model(Phi2PreTrainedModel):
     def __init__(self, config: Phi2Config) -> None:
         super().__init__(config)
-        self.embedding = Embedding(
+        self.rotary_embedding = Embedding(
             vocab_size=config.vocab_size,
             d_embedding=config.d_embedding,
             embd_pdrop=config.embd_pdrop,
@@ -117,10 +117,10 @@ class Phi2Model(Phi2PreTrainedModel):
 
     """
     def get_input_embeddings(self) -> nn.Embedding:
-        return self.embedding.embeddings
+        return self.rotary_embedding.embeddings
 
     def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
-        self.embedding.embeddings = new_embeddings
+        self.rotary_embedding.embeddings = new_embeddings
     """
 
     def forward(
@@ -129,7 +129,7 @@ class Phi2Model(Phi2PreTrainedModel):
         kv_cache: KVCache | None = None,
         key_padding_mask: torch.BoolTensor | None = None,
     ) -> torch.FloatTensor:
-        x = self.embedding(input_ids)
+        x = self.rotary_embedding(input_ids)
         for block in self.parallel_blocks:
             x = block(
                 x,
@@ -143,8 +143,8 @@ class Phi2ModelForCausalLM(Phi2PreTrainedModel):
     def __init__(self, config: Phi2Config) -> None:
         super().__init__(config)
         self.pretrained_model = Phi2Model(config)
-        self.layer_norm = nn.LayerNorm(config.d_embedding, eps=config.layer_norm_epsilon)
-        self.linear = nn.Linear(config.d_embedding, config.vocab_size)
+        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()
         self.post_init()  # calls self._init_weights() for all modules
 
@@ -156,8 +156,8 @@ class Phi2ModelForCausalLM(Phi2PreTrainedModel):
         labels: torch.LongTensor | None = None,
     ) -> CausalLMOutputWithPast:
         x = self.pretrained_model(input_ids, kv_cache=kv_cache, key_padding_mask=key_padding_mask)
-        x = self.layer_norm(x)
-        logits = self.linear(x).to(torch.float32)
+        x = self.lm_head_layer_norm(x)
+        logits = self.lm_head_linear(x).to(torch.float32)
         loss = (
             self.loss_fn(logits.view(-1, logits.size(-1)), labels.view(-1))
             if labels is not None

streaming_inference.py

@@ -1,43 +1,11 @@
 import json
 from threading import Thread
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-import torch
 
 from .phi2_configuration import Phi2Config
 from .phi2_model import Phi2ModelForCausalLM
 
 
-# This works, but is not streaming
-"""
-if __name__ == "__main__":
-    device = "cuda"
-
-    model_config = Phi2Config(**json.load(open("simplified_phi2/config.json")))
-    model = Phi2ModelForCausalLM(model_config).to(device)
-    phi_model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", trust_remote_code=True)
-    model.load_state_dict(phi_model.state_dict())
-
-    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
-
-    text = "Write an essay on sea monkeys: "
-    tokens = tokenizer(text, return_tensors="pt", return_attention_mask=False).to(device)
-    outputs = model.generate(**tokens, max_length=200)
-    text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-    print(text)
-"""
-
-
-# This is streaming, but does not work because you can't set trust_remote_code=True
-"""
-if __name__ == "__main__":
-    client = InferenceClient(model="microsoft/phi-2")
-    text = "How do you make cheese?"
-    for token in client.text_generation(text, max_new_tokens=500, stream=True):
-        print(token, end="")
-"""
-
-
-# This is trying the TextIteratorStreamer class
 if __name__ == "__main__":
     # make and load tokenizer, use tokenizer to initialize token_streamer
     tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
@@ -48,7 +16,29 @@ if __name__ == "__main__":
     model_config = Phi2Config(**json.load(open("simplified_phi2/config.json")))
     model = Phi2ModelForCausalLM(model_config).to(device)
     phi_model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", trust_remote_code=True)
-    model.load_state_dict(phi_model.state_dict())
+
+    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
+        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.")
+        model_state_dict[key] = value
+    model.load_state_dict(model_state_dict)
+
     thread = Thread(
         target=model.generate,
         kwargs=dict(