Add Phi-3-mini-4k-instruct checkpoint

litgpt/config.py

@@ -54,7 +54,7 @@ class Config:
     shared_attention_norm: bool = False
     norm_class_name: Literal["LayerNorm", "RMSNorm"] = "LayerNorm"
     norm_eps: float = 1e-5
-    mlp_class_name: Literal["GptNeoxMLP", "LLaMAMLP", "GemmaMLP", "LLaMAMoE"] = "GptNeoxMLP"
+    mlp_class_name: Literal["GptNeoxMLP", "LLaMAMLP", "GemmaMLP", "LLaMAMoE", "Phi3MLP"] = "GptNeoxMLP"
     gelu_approximate: str = "none"
     intermediate_size: Optional[int] = None
     rope_condense_ratio: int = 1
@@ -836,7 +836,7 @@ def norm_class(self) -> Type:
         copy["name"] = c["name"].format(kind)
         copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
         configs.append(copy)
-        
+
 
 ###############
 # Meta LLaMA 3
@@ -1413,6 +1413,21 @@ def norm_class(self) -> Type:
         lm_head_bias=True,
         gelu_approximate="tanh",
     ),
+    # https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/config.json
+    dict(
+        name="Phi-3-mini-4k-instruct",
+        hf_config=dict(org="microsoft", name="microsoft/Phi-3-mini-4k-instruct"),
+        vocab_size=32064,
+        padded_vocab_size=32064,
+        block_size=4096,
+        n_embd=3072,
+        n_layer=32,
+        rotary_percentage=1.0,
+        bias=False,
+        norm_class_name="RMSNorm",
+        intermediate_size=16384,
+        mlp_class_name="Phi3MLP",
+    ),
 ]
 configs.extend(phi)
 

litgpt/model.py

@@ -298,6 +298,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.proj(x)
 
 
+class Phi3MLP(nn.Module):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        self.gate_up_proj = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
+        self.down_proj = nn.Linear(config.intermediate_size//2, config.n_embd, bias=config.bias)
+        self.config = config
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.gate_up_proj(x)
+        gate, y = y.chunk(2, dim=-1)
+        y = y * torch.nn.functional.silu(gate)
+        return self.down_proj(y)
+
+
 class LLaMAMLP(nn.Module):
     def __init__(self, config: Config) -> None:
         super().__init__()

litgpt/scripts/convert_hf_checkpoint.py

@@ -121,6 +121,7 @@ def copy_weights_hf_llama(
         "model.layers.{}.self_attn.q_proj.weight": None,
         "model.layers.{}.self_attn.k_proj.weight": None,
         "model.layers.{}.self_attn.v_proj.weight": None,
+        "model.layers.{}.self_attn.qkv_proj.weight": None,
         "model.layers.{}.self_attn.o_proj.weight": "transformer.h.{l}.attn.proj.weight",
         "model.layers.{}.self_attn.rotary_emb.inv_freq": None,
         "model.layers.{}.post_attention_layernorm.weight": "transformer.h.{l}.norm_2.weight",
@@ -146,6 +147,13 @@ def copy_weights_hf_llama(
                 "model.layers.{}.mlp.down_proj.weight": "transformer.h.{l}.mlp.proj.weight",
             }
         )
+    elif config.mlp_class_name in ("Phi3MLP",):
+        weight_map.update(
+            {
+                "model.layers.{}.mlp.gate_up_proj.weight": "transformer.h.{l}.mlp.gate_up_proj.weight",
+                "model.layers.{}.mlp.down_proj.weight": "transformer.h.{l}.mlp.down_proj.weight",
+            }
+        )
     else:
         raise NotImplementedError
 
@@ -156,7 +164,9 @@ def copy_weights_hf_llama(
             if "block_sparse_moe.experts" in name:
                 from_name, e = layer_template(from_name, 5)
             qkv = qkv_weights.setdefault(l, [None, None, None])
-            if "q_proj" in name:
+            if "qkv_proj" in name:
+                state_dict[f"transformer.h.{l}.attn.attn.weight"] = load_param(param, f"layer {l} qkv", dtype)
+            elif "q_proj" in name:
                 qkv[0] = param
             elif "k_proj" in name:
                 qkv[1] = param
@@ -177,21 +187,22 @@ def copy_weights_hf_llama(
         state_dict["lm_head.weight"] = state_dict["transformer.wte.weight"]
 
     # convert separate q, k, v matrices into an interleaved qkv
-    for i, (q, k, v) in list(qkv_weights.items()):
-        if q is None or k is None or v is None:
-            # split across different .bin files
-            continue
-        q = load_param(q, f"layer {i} q", dtype)
-        k = load_param(k, f"layer {i} k", dtype)
-        v = load_param(v, f"layer {i} v", dtype)
-        q_per_kv = config.n_head // config.n_query_groups
-        qs = torch.split(q, config.head_size * q_per_kv)
-        ks = torch.split(k, config.head_size)
-        vs = torch.split(v, config.head_size)
-        cycled = [t for group in zip(qs, ks, vs) for t in group]
-        qkv = torch.cat(cycled)
-        state_dict[f"transformer.h.{i}.attn.attn.weight"] = qkv
-        del qkv_weights[i]
+    if "qkv_proj" not in name:
+        for i, (q, k, v) in list(qkv_weights.items()):
+            if q is None or k is None or v is None:
+                # split across different .bin files
+                continue
+            q = load_param(q, f"layer {i} q", dtype)
+            k = load_param(k, f"layer {i} k", dtype)
+            v = load_param(v, f"layer {i} v", dtype)
+            q_per_kv = config.n_head // config.n_query_groups
+            qs = torch.split(q, config.head_size * q_per_kv)
+            ks = torch.split(k, config.head_size)
+            vs = torch.split(v, config.head_size)
+            cycled = [t for group in zip(qs, ks, vs) for t in group]
+            qkv = torch.cat(cycled)
+            state_dict[f"transformer.h.{i}.attn.attn.weight"] = qkv
+            del qkv_weights[i]
 
 
 def copy_weights_phi(
@@ -312,7 +323,7 @@ def convert_hf_checkpoint(
 
     if "falcon" in model_name:
         copy_fn = partial(copy_weights_falcon, model_name)
-    elif config.mlp_class_name in ("LLaMAMLP", "GemmaMLP", "LLaMAMoE"):
+    elif config.mlp_class_name in ("LLaMAMLP", "GemmaMLP", "LLaMAMoE", "Phi3MLP"):
         # holder to reconstitute the split q, k, v
         qkv_weights = {}
         copy_fn = partial(copy_weights_hf_llama, config, qkv_weights)
@@ -354,4 +365,4 @@ def convert_hf_checkpoint(
 if __name__ == "__main__":
     from jsonargparse import CLI
 
-    CLI(convert_hf_checkpoint)
+    CLI(convert_hf_checkpoint)

tutorials/download_model_weights.md

@@ -112,6 +112,7 @@ meta-llama/Meta-Llama-3-8B
 meta-llama/Meta-Llama-3-8B-Instruct
 microsoft/phi-1_5
 microsoft/phi-2
+microsoft/Phi-3-mini-4k-instruct
 mistralai/Mistral-7B-Instruct-v0.1
 mistralai/Mistral-7B-Instruct-v0.2
 mistralai/Mistral-7B-v0.1