Corrected param name

phi2_model.py

@@ -87,7 +87,7 @@ class Embedding(nn.Module):
 class Phi2Model(Phi2PreTrainedModel):
     def __init__(self, config: Phi2Config) -> None:
         super().__init__(config)
-        self.rotary_embedding = Embedding(
+        self.embedding = Embedding(
             vocab_size=config.vocab_size,
             d_embedding=config.d_embedding,
             embd_pdrop=config.embd_pdrop,
@@ -113,10 +113,10 @@ class Phi2Model(Phi2PreTrainedModel):
 
     """
     def get_input_embeddings(self) -> nn.Embedding:
-        return self.rotary_embedding.embeddings
+        return self.embedding.embeddings
 
     def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
-        self.rotary_embedding.embeddings = new_embeddings
+        self.embedding.embeddings = new_embeddings
     """
 
     def forward(
@@ -125,7 +125,7 @@ class Phi2Model(Phi2PreTrainedModel):
         kv_cache: KVCache | None = None,
         key_padding_mask: torch.BoolTensor | None = None,
     ) -> torch.FloatTensor:
-        x = self.rotary_embedding(input_ids)
+        x = self.embedding(input_ids)
         for block in self.parallel_blocks:
             x = block(
                 x,

streaming_inference.py

@@ -22,14 +22,14 @@ if __name__ == "__main__":
     for key, value in phi_model_state_dict.items():
         # 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.embd.wte.weight -> model.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 = key.replace("transformer.", "model.")
-            key = key.replace(".embd.wte.", ".rotary_embedding.embeddings.")
+            key = key.replace(".embd.wte.", ".embedding.embeddings.")
             key = key.replace(".h.", ".parallel_blocks.")
             key = key.replace(".ln.", ".layer_norm.")
             key = key.replace(".mixer.Wqkv.", ".multi_head_attention.Wqkv.")