Update the DeepSpeed Phi-2 impl. to work with the HF latest changes
diff --git a/deepspeed/inference/v2/engine_factory.py b/deepspeed/inference/v2/engine_factory.py
index 41cb47729237..9281640f844a 100644
--- a/deepspeed/inference/v2/engine_factory.py
+++ b/deepspeed/inference/v2/engine_factory.py
@@ -116,7 +116,7 @@ def build_hf_engine(path: str,
             policy = MixtralPolicy(model_config, checkpoint_engine=checkpoint_engine)
         elif model_config.model_type == "falcon":
             policy = FalconPolicy(model_config, checkpoint_engine=checkpoint_engine)
-        elif model_config.model_type == "phi-msft":
+        elif model_config.model_type == "phi":
             policy = PhiPolicy(model_config, checkpoint_engine=checkpoint_engine)
         elif model_config.model_type == "qwen":
             policy = QwenPolicy(model_config, checkpoint_engine=checkpoint_engine)
diff --git a/deepspeed/inference/v2/model_implementations/phi/containers.py b/deepspeed/inference/v2/model_implementations/phi/containers.py
index ab6d0181611c..21f07eb8c99a 100644
--- a/deepspeed/inference/v2/model_implementations/phi/containers.py
+++ b/deepspeed/inference/v2/model_implementations/phi/containers.py
@@ -11,41 +11,30 @@
  # HF Phi-2 model looks like this:
 
 PhiForCausalLM(
-  (transformer): PhiModel(
-    (embd): Embedding(
-      (wte): Embedding(51200, 2560)
-      (drop): Dropout(p=0.0, inplace=False)
-    )
-    (h): ModuleList(
-      (0-31): 32 x ParallelBlock(
-        (ln): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
-        (resid_dropout): Dropout(p=0.1, inplace=False)
-        (mixer): MHA(
-          (rotary_emb): RotaryEmbedding()
-          (Wqkv): Linear(in_features=2560, out_features=7680, bias=True)
-          (out_proj): Linear(in_features=2560, out_features=2560, bias=True)
-          (inner_attn): SelfAttention(
-            (drop): Dropout(p=0.0, inplace=False)
-          )
-          (inner_cross_attn): CrossAttention(
-            (drop): Dropout(p=0.0, inplace=False)
-          )
+  (model): PhiModel(
+    (embed_tokens): Embedding(51200, 2560)
+    (embed_dropout): Dropout(p=0.0, inplace=False)
+    (layers): ModuleList(
+      (0-31): 32 x PhiDecoderLayer(
+        (self_attn): PhiAttention(
+          (q_proj): Linear(in_features=2560, out_features=2560, bias=True)
+          (k_proj): Linear(in_features=2560, out_features=2560, bias=True)
+          (v_proj): Linear(in_features=2560, out_features=2560, bias=True)
+          (dense): Linear(in_features=2560, out_features=2560, bias=True)
+          (rotary_emb): PhiRotaryEmbedding()
         )
-        (mlp): MLP(
+        (mlp): PhiMLP(
+          (activation_fn): NewGELUActivation()
           (fc1): Linear(in_features=2560, out_features=10240, bias=True)
           (fc2): Linear(in_features=10240, out_features=2560, bias=True)
-          (act): NewGELUActivation()
         )
+        (input_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
+        (resid_dropout): Dropout(p=0.1, inplace=False)
       )
     )
+    (final_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
   )
-  (lm_head): CausalLMHead(
-    (ln): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
-    (linear): Linear(in_features=2560, out_features=51200, bias=True)
-  )
-  (loss): CausalLMLoss(
-    (loss_fct): CrossEntropyLoss()
-  )
+  (lm_head): Linear(in_features=2560, out_features=51200, bias=True)
 )
 '''
 
@@ -54,8 +43,8 @@ class PhiTransformerContainer(LayerContainer):
     """
         Transformer layer container for the Phi model.
     """
-    qkv_w: FusedQKVParameter
-    qkv_b: FusedQKVParameter
+    qkv_w: UnfusedQKVParameter
+    qkv_b: UnfusedQKVParameter
     attn_out_w: AttentionOutputParameter
     attn_out_b: AttentionOutputParameter
     mlp_1_w: MLP1Parameter
@@ -66,16 +55,20 @@ class PhiTransformerContainer(LayerContainer):
     ln_beta: NormParameter
 
     PARAM_MAPPING = {
-        "mixer.Wqkv.weight": "qkv_w.params",
-        "mixer.Wqkv.bias": "qkv_b.params",
-        "mixer.out_proj.weight": "attn_out_w.params",
-        "mixer.out_proj.bias": "attn_out_b.params",
+        "self_attn.q_proj.weight": "qkv_w.q_params",
+        "self_attn.k_proj.weight": "qkv_w.k_params",
+        "self_attn.v_proj.weight": "qkv_w.v_params",
+        "self_attn.q_proj.bias": "qkv_b.q_params",
+        "self_attn.k_proj.bias": "qkv_b.k_params",
+        "self_attn.v_proj.bias": "qkv_b.v_params",
+        "self_attn.dense.weight": "attn_out_w.params",
+        "self_attn.dense.bias": "attn_out_b.params",
         "mlp.fc1.weight": "mlp_1_w.params",
         "mlp.fc1.bias": "mlp_1_b.params",
         "mlp.fc2.weight": "mlp_2_w.params",
         "mlp.fc2.bias": "mlp_2_b.params",
-        "ln.weight": "ln_gamma.params",
-        "ln.bias": "ln_beta.params",
+        "input_layernorm.weight": "ln_gamma.params",
+        "input_layernorm.bias": "ln_beta.params",
     }
 
 
@@ -90,9 +83,9 @@ class PhiNonTransformerContainer(LayerContainer):
     final_norm_beta: NormParameter
 
     PARAM_MAPPING = {
-        "transformer.embd.wte.weight": "word_emb.params",
-        "lm_head.ln.weight": "final_norm_gamma.params",
-        "lm_head.ln.bias": "final_norm_beta.params",
-        "lm_head.linear.weight": "word_unembed_w.params",
-        "lm_head.linear.bias": "word_unembed_b.params",
+        "model.embed_tokens.weight": "word_emb.params",
+        "model.final_layernorm.weight": "final_norm_gamma.params",
+        "model.final_layernorm.bias": "final_norm_beta.params",
+        "lm_head.weight": "word_unembed_w.params",
+        "lm_head.bias": "word_unembed_b.params",
     }
diff --git a/deepspeed/inference/v2/model_implementations/phi/model.py b/deepspeed/inference/v2/model_implementations/phi/model.py
index 0127c87c7bff..2d5826810cb5 100644
--- a/deepspeed/inference/v2/model_implementations/phi/model.py
+++ b/deepspeed/inference/v2/model_implementations/phi/model.py
@@ -47,11 +47,11 @@ def max_sequence_length(self) -> int:
 
     @property
     def num_layers(self) -> int:
-        return self._config.n_layer
+        return self._config.num_hidden_layers
 
     @property
     def model_dim(self) -> int:
-        return self._config.n_embd
+        return self._config.hidden_size
 
     @property
     def vocab_size(self) -> int:
@@ -63,16 +63,15 @@ def head_size(self) -> int:
 
     @property
     def n_heads(self) -> int:
-        return self._config.n_head
+        return self._config.num_attention_heads
 
     @property
     def intermediate_dim(self) -> int:
-        n_inner = getattr(self._config, "n_inner", None)
-        return n_inner if n_inner is not None else 4 * self.model_dim
+        return self._config.intermediate_size
 
     @property
     def n_heads_kv(self) -> int:
-        return getattr(self._config, "n_head_kv", None) or self.n_heads
+        return self._config.num_key_value_heads
 
     @property
     def activation_dtype(self) -> DtypeEnum:
@@ -97,7 +96,8 @@ def positional_embedding_type(self) -> PositionalEmbeddingType:
 
     @property
     def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
-        return RotateHalfConfig(rotate_dim=self._config.rotary_dim)
+        rotary_dim = int(self._config.partial_rotary_factor * self.head_size)
+        return RotateHalfConfig(rotate_dim=rotary_dim, theta_base=self._config.rope_theta)
 
     """
     Forward implementations
diff --git a/deepspeed/inference/v2/model_implementations/phi/policy.py b/deepspeed/inference/v2/model_implementations/phi/policy.py
index 1e9db24022f5..4b081a8e61bd 100644
--- a/deepspeed/inference/v2/model_implementations/phi/policy.py
+++ b/deepspeed/inference/v2/model_implementations/phi/policy.py
@@ -22,7 +22,7 @@ def build_container_map(self) -> ContainerMap:
         trans_container_cls = PhiTransformerContainer
         transformer_containers = [trans_container_cls(self.model) for _ in range(self.model.num_layers)]
 
-        map.set_transformer_params(['transformer.h'], transformer_containers)
+        map.set_transformer_params(['model.layers'], transformer_containers)
 
         map.set_non_transformer_params(PhiNonTransformerContainer(self.model))
 

The latest changes in Huggingface Phi-2 implementation (https://huggingface.co/microsoft/phi-2/commit/cb2f4533604d8b67de604e7df03bfe6f3ca22869#d2h-025836)  have broken the DeepSpeed implementation. This PR address the related issues. 