copy of phi-2 repo

CODE_OF_CONDUCT.md

@@ -0,0 +1,9 @@
+# Microsoft Open Source Code of Conduct
+
+This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
+
+Resources:
+
+- [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/)
+- [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/)
+- Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns

LICENSE

@@ -0,0 +1,22 @@
+PhyAGI.
+Copyright (c) Microsoft Corporation.
+
+MIT License
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

NOTICE.md

@@ -0,0 +1,38 @@
+NOTICES AND INFORMATION
+Do Not Translate or Localize
+
+This software incorporates material from third parties.
+
+**Component.** https://github.com/Dao-AILab/flash-attention
+
+**Open Source License/Copyright Notice.**
+
+BSD 3-Clause License
+
+Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -0,0 +1,166 @@
+---
+inference: false
+license: mit
+license_link: https://huggingface.co/microsoft/phi-2/resolve/main/LICENSE
+language:
+- en
+pipeline_tag: text-generation
+tags:
+- nlp
+- code
+---
+
+## Model Summary
+
+Phi-2 is a Transformer with **2.7 billion** parameters. It was trained using the same data sources as [Phi-1.5](https://huggingface.co/microsoft/phi-1.5), augmented with a new data source that consists of various NLP synthetic texts and filtered websites (for safety and educational value). When assessed against benchmarks testing common sense, language understanding, and logical reasoning, Phi-2 showcased a nearly state-of-the-art performance among models with less than 13 billion parameters.
+
+Our model hasn't been fine-tuned through reinforcement learning from human feedback. The intention behind crafting this open-source model is to provide the research community with a non-restricted small model to explore vital safety challenges, such as reducing toxicity, understanding societal biases, enhancing controllability, and more.
+
+## Intended Uses
+
+Given the nature of the training data, the Phi-2 model is best suited for prompts using the QA format, the chat format, and the code format.
+
+### QA Format:
+
+You can provide the prompt as a standalone question as follows:
+
+```markdown
+Write a detailed analogy between mathematics and a lighthouse.
+```
+where the model generates the text after "." . 
+To encourage the model to write more concise answers, you can also try the following QA format using "Instruct: \<prompt\>\nOutput:"
+```markdown
+Instruct: Write a detailed analogy between mathematics and a lighthouse.
+Output: Mathematics is like a lighthouse. Just as a lighthouse guides ships safely to shore, mathematics provides a guiding light in the world of numbers and logic. It helps us navigate through complex problems and find solutions. Just as a lighthouse emits a steady beam of light, mathematics provides a consistent framework for reasoning and problem-solving. It illuminates the path to understanding and helps us make sense of the world around us.
+```
+where the model generates the text after "Output:".
+
+### Chat Format:
+
+```markdown
+Alice: I don't know why, I'm struggling to maintain focus while studying. Any suggestions?
+Bob: Well, have you tried creating a study schedule and sticking to it?
+Alice: Yes, I have, but it doesn't seem to help much.
+Bob: Hmm, maybe you should try studying in a quiet environment, like the library.
+Alice: ...
+```
+
+where the model generates the text after the first "Bob:".
+
+### Code Format:
+
+```python
+def print_prime(n):
+   """
+   Print all primes between 1 and n
+   """
+   primes = []
+   for num in range(2, n+1):
+       is_prime = True
+       for i in range(2, int(math.sqrt(num))+1):
+           if num % i == 0:
+               is_prime = False
+               break
+       if is_prime:
+           primes.append(num)
+   print(primes)
+```
+where the model generates the text after the comments.
+
+**Notes:**
+* Phi-2 is intended for QA, chat, and code purposes. The model-generated text/code should be treated as a starting point rather than a definitive solution for potential use cases. Users should be cautious when employing these models in their applications.
+* Direct adoption for production tasks without evaluation is out of scope of this project. As a result, the Phi-2 model has not been tested to ensure that it performs adequately for any production-level application. Please refer to the limitation sections of this document for more details.
+* If you are using `transformers>=4.36.0`, always load the model with `trust_remote_code=True` to prevent side-effects.
+
+## Sample Code
+
+There are four types of execution mode:
+
+1. FP16 / Flash-Attention / CUDA:
+   ```python
+   model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", flash_attn=True, flash_rotary=True, fused_dense=True, device_map="cuda", trust_remote_code=True)
+   ```
+2. FP16 / CUDA:
+   ```python
+   model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", device_map="cuda", trust_remote_code=True)
+   ```
+3. FP32 / CUDA:
+   ```python
+   model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype=torch.float32, device_map="cuda", trust_remote_code=True)
+   ```
+4. FP32 / CPU:
+   ```python
+   model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype=torch.float32, device_map="cpu", trust_remote_code=True)
+   ```
+
+To ensure the maximum compatibility, we recommend using the second execution mode (FP16 / CUDA), as follows:
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+torch.set_default_device("cuda")
+
+model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
+
+inputs = tokenizer('''def print_prime(n):
+   """
+   Print all primes between 1 and n
+   """''', return_tensors="pt", return_attention_mask=False)
+
+outputs = model.generate(**inputs, max_length=200)
+text = tokenizer.batch_decode(outputs)[0]
+print(text)
+```
+
+**Remark:** In the generation function, our model currently does not support beam search (`num_beams > 1`).
+Furthermore, in the forward pass of the model, we currently do not support outputting hidden states or attention values, or using custom input embeddings.
+
+## Limitations of Phi-2
+
+* Generate Inaccurate Code and Facts: The model may produce incorrect code snippets and statements. Users should treat these outputs as suggestions or starting points, not as definitive or accurate solutions.
+
+* Limited Scope for code: Majority of Phi-2 training data is based in Python and use common packages such as "typing, math, random, collections, datetime, itertools". If the model generates Python scripts that utilize other packages or scripts in other languages, we strongly recommend users manually verify all API uses.
+
+* Unreliable Responses to Instruction: The model has not undergone instruction fine-tuning. As a result, it may struggle or fail to adhere to intricate or nuanced instructions provided by users.
+
+* Language Limitations: The model is primarily designed to understand standard English. Informal English, slang, or any other languages might pose challenges to its comprehension, leading to potential misinterpretations or errors in response.
+
+* Potential Societal Biases: Phi-2 is not entirely free from societal biases despite efforts in assuring training data safety. There's a possibility it may generate content that mirrors these societal biases, particularly if prompted or instructed to do so. We urge users to be aware of this and to exercise caution and critical thinking when interpreting model outputs.
+
+* Toxicity: Despite being trained with carefully selected data, the model can still produce harmful content if explicitly prompted or instructed to do so. We chose to release the model to help the open-source community develop the most effective ways to reduce the toxicity of a model directly after pretraining.
+
+* Verbosity: Phi-2 being a base model often produces irrelevant or extra text and responses following its first answer to user prompts within a single turn. This is due to its training dataset being primarily textbooks, which results in textbook-like responses.
+
+## Training
+
+### Model
+
+* Architecture: a Transformer-based model with next-word prediction objective
+
+* Context length: 2048 tokens
+
+* Dataset size: 250B tokens, combination of NLP synthetic data created by AOAI GPT-3.5 and filtered web data from Falcon RefinedWeb and SlimPajama, which was assessed by AOAI GPT-4.
+
+* Training tokens: 1.4T tokens
+
+* GPUs: 96xA100-80G
+
+* Training time: 14 days
+
+### Software
+
+* [PyTorch](https://github.com/pytorch/pytorch)
+
+* [DeepSpeed](https://github.com/microsoft/DeepSpeed)
+
+* [Flash-Attention](https://github.com/HazyResearch/flash-attention)
+
+### License
+
+The model is licensed under the [MIT license](https://huggingface.co/microsoft/phi-2/resolve/main/LICENSE).
+
+## Trademarks
+
+This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft trademarks or logos is subject to and must follow [Microsoft’s Trademark & Brand Guidelines](https://www.microsoft.com/en-us/legal/intellectualproperty/trademarks). Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship. Any use of third-party trademarks or logos are subject to those third-party’s policies.

SECURITY.md

@@ -0,0 +1,41 @@
+<!-- BEGIN MICROSOFT SECURITY.MD V0.0.9 BLOCK -->
+
+## Security
+
+Microsoft takes the security of our software products and services seriously, which includes all source code repositories managed through our GitHub organizations, which include [Microsoft](https://github.com/Microsoft), [Azure](https://github.com/Azure), [DotNet](https://github.com/dotnet), [AspNet](https://github.com/aspnet) and [Xamarin](https://github.com/xamarin).
+
+If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](https://aka.ms/security.md/definition), please report it to us as described below.
+
+## Reporting Security Issues
+
+**Please do not report security vulnerabilities through public GitHub issues.**
+
+Instead, please report them to the Microsoft Security Response Center (MSRC) at [https://msrc.microsoft.com/create-report](https://aka.ms/security.md/msrc/create-report).
+
+If you prefer to submit without logging in, send email to [secure@microsoft.com](mailto:secure@microsoft.com).  If possible, encrypt your message with our PGP key; please download it from the [Microsoft Security Response Center PGP Key page](https://aka.ms/security.md/msrc/pgp).
+
+You should receive a response within 24 hours. If for some reason you do not, please follow up via email to ensure we received your original message. Additional information can be found at [microsoft.com/msrc](https://www.microsoft.com/msrc). 
+
+Please include the requested information listed below (as much as you can provide) to help us better understand the nature and scope of the possible issue:
+
+  * Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.)
+  * Full paths of source file(s) related to the manifestation of the issue
+  * The location of the affected source code (tag/branch/commit or direct URL)
+  * Any special configuration required to reproduce the issue
+  * Step-by-step instructions to reproduce the issue
+  * Proof-of-concept or exploit code (if possible)
+  * Impact of the issue, including how an attacker might exploit the issue
+
+This information will help us triage your report more quickly.
+
+If you are reporting for a bug bounty, more complete reports can contribute to a higher bounty award. Please visit our [Microsoft Bug Bounty Program](https://aka.ms/security.md/msrc/bounty) page for more details about our active programs.
+
+## Preferred Languages
+
+We prefer all communications to be in English.
+
+## Policy
+
+Microsoft follows the principle of [Coordinated Vulnerability Disclosure](https://aka.ms/security.md/cvd).
+
+<!-- END MICROSOFT SECURITY.MD BLOCK -->

added_tokens.json

@@ -0,0 +1,40 @@
+{
+  "\t\t": 50294,
+  "\t\t\t": 50293,
+  "\t\t\t\t": 50292,
+  "\t\t\t\t\t": 50291,
+  "\t\t\t\t\t\t": 50290,
+  "\t\t\t\t\t\t\t": 50289,
+  "\t\t\t\t\t\t\t\t": 50288,
+  "\t\t\t\t\t\t\t\t\t": 50287,
+  "  ": 50286,
+  "   ": 50285,
+  "    ": 50284,
+  "     ": 50283,
+  "      ": 50282,
+  "       ": 50281,
+  "        ": 50280,
+  "         ": 50279,
+  "          ": 50278,
+  "           ": 50277,
+  "            ": 50276,
+  "             ": 50275,
+  "              ": 50274,
+  "               ": 50273,
+  "                ": 50272,
+  "                 ": 50271,
+  "                  ": 50270,
+  "                   ": 50269,
+  "                    ": 50268,
+  "                     ": 50267,
+  "                      ": 50266,
+  "                       ": 50265,
+  "                        ": 50264,
+  "                         ": 50263,
+  "                          ": 50262,
+  "                           ": 50261,
+  "                            ": 50260,
+  "                             ": 50259,
+  "                              ": 50258,
+  "                               ": 50257
+}

config.json

@@ -0,0 +1,32 @@
+{
+  "_name_or_path": "microsoft/phi-2",
+  "activation_function": "gelu_new",
+  "architectures": [
+    "PhiForCausalLM"
+  ],
+  "attn_pdrop": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_phi.PhiConfig",
+    "AutoModelForCausalLM": "modeling_phi.PhiForCausalLM"
+  },
+  "embd_pdrop": 0.0,
+  "flash_attn": false,
+  "flash_rotary": false,
+  "fused_dense": false,
+  "img_processor": null,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "phi-msft",
+  "n_embd": 2560,
+  "n_head": 32,
+  "n_head_kv": null,
+  "n_inner": null,
+  "n_layer": 32,
+  "n_positions": 2048,
+  "resid_pdrop": 0.1,
+  "rotary_dim": 32,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float16",
+  "transformers_version": "4.35.2",
+  "vocab_size": 51200
+}

configuration_phi.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import math
+from typing import Optional
+
+from transformers import PretrainedConfig
+
+
+class PhiConfig(PretrainedConfig):
+    """Phi configuration."""
+
+    model_type = "phi-msft"
+    attribute_map = {
+        "max_position_embeddings": "n_positions",
+        "hidden_size": "n_embd",
+        "num_attention_heads": "n_head",
+        "num_hidden_layers": "n_layer",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 50304,
+        n_positions: int = 2048,
+        n_embd: int = 1024,
+        n_layer: int = 20,
+        n_inner: Optional[int] = None,
+        n_head: int = 16,
+        n_head_kv: Optional[int] = None,
+        rotary_dim: Optional[int] = 32,
+        activation_function: Optional[str] = "gelu_new",
+        flash_attn: bool = False,
+        flash_rotary: bool = False,
+        fused_dense: bool = False,
+        attn_pdrop: float = 0.0,
+        embd_pdrop: float = 0.0,
+        resid_pdrop: float = 0.0,
+        layer_norm_epsilon: float = 1e-5,
+        initializer_range: float = 0.02,
+        tie_word_embeddings: bool = False,
+        pad_vocab_size_multiple: int = 64,
+        **kwargs
+    ) -> None:
+        self.vocab_size = int(math.ceil(vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple)
+        self.n_positions = n_positions
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_inner = n_inner
+        self.n_head = n_head
+        self.n_head_kv = n_head_kv
+        self.rotary_dim = min(rotary_dim, n_embd // n_head)
+        self.activation_function = activation_function
+        self.flash_attn = flash_attn
+        self.flash_rotary = flash_rotary
+        self.fused_dense = fused_dense
+        self.attn_pdrop = attn_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.resid_pdrop = resid_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.35.2"
+}

model.safetensors.index.json

@@ -0,0 +1,332 @@
+{
+  "metadata": {
+    "total_size": 5559367680
+  },
+  "weight_map": {
+    "lm_head.linear.bias": "model-00002-of-00002.safetensors",
+    "lm_head.linear.weight": "model-00002-of-00002.safetensors",
+    "lm_head.ln.bias": "model-00002-of-00002.safetensors",
+    "lm_head.ln.weight": "model-00002-of-00002.safetensors",
+    "transformer.embd.wte.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.0.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.0.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.0.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.0.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.0.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.0.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.0.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.0.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.0.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.0.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.1.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.1.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.1.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.1.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.1.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.1.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.1.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.1.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.1.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.1.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.10.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.10.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.10.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.10.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.10.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.10.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.10.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.10.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.10.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.10.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.11.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.11.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.11.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.11.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.11.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.11.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.11.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.11.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.11.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.11.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.12.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.12.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.12.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.12.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.12.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.12.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.12.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.12.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.12.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.12.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.13.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.13.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.13.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.13.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.13.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.13.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.13.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.13.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.13.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.13.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.14.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.14.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.14.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.14.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.14.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.14.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.14.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.14.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.14.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.14.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.15.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.15.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.15.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.15.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.15.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.15.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.15.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.15.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.15.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.15.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.16.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.16.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.16.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.16.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.16.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.16.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.16.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.16.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.16.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.16.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.17.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.17.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.17.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.17.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.17.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.17.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.17.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.17.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.17.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.17.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.18.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.18.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.18.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.18.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.18.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.18.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.18.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.18.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.18.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.18.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.19.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.19.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.19.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.19.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.19.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.19.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.19.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.19.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.19.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.19.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.2.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.2.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.2.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.2.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.2.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.2.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.2.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.2.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.2.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.2.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.20.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.20.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.20.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.20.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.20.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.20.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.20.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.20.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.20.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.20.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.21.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.21.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.21.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.21.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.21.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.21.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.21.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.21.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.21.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.21.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.22.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.22.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.22.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.22.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.22.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.22.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.22.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.22.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.22.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.22.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.23.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.23.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.23.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.23.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.23.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.23.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.23.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.23.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.23.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.23.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.24.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.24.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.24.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.24.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.24.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.24.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.24.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.24.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.24.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.24.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.25.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.25.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.25.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.25.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.25.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.25.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.25.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.25.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.25.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.25.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.26.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.26.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.26.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.26.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.26.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.26.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.26.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.26.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.26.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.26.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.27.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.27.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.27.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.27.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.27.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.27.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.27.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.27.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.27.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.27.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.28.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.28.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.28.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.28.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.28.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.28.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.28.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.28.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.28.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.28.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.29.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.29.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.29.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.29.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.29.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.29.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.29.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.29.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.29.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.29.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.3.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.3.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.3.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.3.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.3.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.3.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.3.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.3.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.3.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.3.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.30.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.30.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.30.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "transformer.h.30.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "transformer.h.30.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "transformer.h.30.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "transformer.h.30.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "transformer.h.30.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "transformer.h.30.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "transformer.h.30.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "transformer.h.31.ln.bias": "model-00002-of-00002.safetensors",
+    "transformer.h.31.ln.weight": "model-00002-of-00002.safetensors",
+    "transformer.h.31.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "transformer.h.31.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "transformer.h.31.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "transformer.h.31.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "transformer.h.31.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "transformer.h.31.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "transformer.h.31.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "transformer.h.31.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "transformer.h.4.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.4.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.4.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.4.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.4.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.4.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.4.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.4.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.4.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.4.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.5.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.5.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.5.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.5.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.5.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.5.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.5.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.5.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.5.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.5.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.6.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.6.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.6.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.6.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.6.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.6.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.6.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.6.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.6.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.6.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.7.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.7.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.7.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.7.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.7.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.7.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.7.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.7.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.7.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.7.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.8.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.8.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.8.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.8.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.8.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.8.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.8.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.8.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.8.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.8.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.9.ln.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.9.ln.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.9.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.9.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.9.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.9.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.9.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.9.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "transformer.h.9.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "transformer.h.9.mlp.fc2.weight": "model-00001-of-00002.safetensors"
+  }
+}

modeling_phi.py

@@ -0,0 +1,980 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+#
+# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu.
+# Licensed under the BSD 3-Clause License.
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithNoAttention
+
+from .configuration_phi import PhiConfig
+
+try:
+    from flash_attn.bert_padding import pad_input, unpad_input
+    from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
+    from flash_attn.modules.mha import FlashCrossAttention, FlashSelfAttention
+    from flash_attn.ops.fused_dense import FusedDense
+except:
+    pad_input, unpad_input = None, None
+    FlashRotaryEmbedding = None
+    FlashSelfAttention, FlashCrossAttention = None, None
+    FusedDense = None
+
+
+@dataclass
+class InferenceParams:
+    """Inference parameters passed to model to efficiently calculate
+    and store context during inference.
+
+    Reference:
+        https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/utils/generation.py.
+
+    Args:
+        max_seqlen: Maximum sequence length.
+        max_batch_size: Maximum batch size.
+        seqlen_offset: Sequence length offset.
+        batch_size_offset: Batch size offset.
+        key_value_memory_dict: Key value memory dictionary.
+        lengths_per_sample: Lengths per sample.
+
+    """
+
+    max_seqlen: int = field(metadata={"help": "Maximum sequence length."})
+
+    max_batch_size: int = field(metadata={"help": "Maximum batch size."})
+
+    seqlen_offset: int = field(default=0, metadata={"help": "Sequence length offset."})
+
+    batch_size_offset: int = field(default=0, metadata={"help": "Batch size offset."})
+
+    key_value_memory_dict: Dict[str, Any] = field(
+        default_factory=dict, metadata={"help": "Key value memory dictionary."}
+    )
+
+    lengths_per_sample: torch.Tensor = field(default=None, metadata={"help": "Lengths per sample."})
+
+
+class Embedding(nn.Module):
+    """Token embedding with dropout."""
+
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+
+        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
+        self.drop = nn.Dropout(config.embd_pdrop)
+
+    def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
+        input_shape = input_ids.size()
+        input_ids = input_ids.view(-1, input_shape[-1])
+
+        hidden_states = self.wte(input_ids)
+        hidden_states = self.drop(hidden_states)
+
+        return hidden_states
+
+
+def _apply_rotary_emb(
+    x: torch.FloatTensor,
+    cos: torch.FloatTensor,
+    sin: torch.FloatTensor,
+) -> torch.FloatTensor:
+    _, seqlen, _, _ = x.shape
+    _, rotary_dim = cos.shape
+    rotary_dim *= 2
+
+    x_rot = x[:, :, :, :rotary_dim]
+    x_pass = x[:, :, :, rotary_dim:]
+
+    x1, x2 = x_rot.chunk(2, dim=-1)
+    c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
+    x1, x2, c, s = [t.to(dtype=torch.float32) for t in [x1, x2, c, s]]
+
+    x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], axis=-1).to(x.dtype)
+
+    return torch.cat([x_rot, x_pass], axis=-1)
+
+
+def _apply_rotary_emb_kv(
+    kv: torch.FloatTensor,
+    cos: torch.FloatTensor,
+    sin: torch.FloatTensor,
+    cos_k: Optional[torch.FloatTensor] = None,
+    sin_k: Optional[torch.FloatTensor] = None,
+) -> torch.FloatTensor:
+    _, seqlen, _, _, _ = kv.shape
+    _, rotary_dim = cos.shape
+    rotary_dim *= 2
+
+    k_rot = kv[:, :, 0, :, :rotary_dim]
+    k_pass = kv[:, :, 0, :, rotary_dim:]
+
+    k1, k2 = k_rot.chunk(2, dim=-1)
+    c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
+    k1, k2, c, s = [t.to(dtype=torch.float32) for t in [k1, k2, c, s]]
+
+    k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(kv.dtype)
+
+    return torch.cat(
+        [
+            torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
+            kv[:, :, 1:2, :, :],
+        ],
+        axis=2,
+    )
+
+
+def _apply_rotary_emb_qkv(
+    qkv: torch.FloatTensor,
+    cos: torch.FloatTensor,
+    sin: torch.FloatTensor,
+    cos_k: Optional[torch.FloatTensor] = None,
+    sin_k: Optional[torch.FloatTensor] = None,
+) -> torch.FloatTensor:
+    _, seqlen, _, _, _ = qkv.shape
+    _, rotary_dim = cos.shape
+    rotary_dim *= 2
+
+    q_rot = qkv[:, :, 0, :, :rotary_dim]
+    q_pass = qkv[:, :, 0, :, rotary_dim:]
+
+    k_rot = qkv[:, :, 1, :, :rotary_dim]
+    k_pass = qkv[:, :, 1, :, rotary_dim:]
+
+    q1, q2 = q_rot.chunk(2, dim=-1)
+    k1, k2 = k_rot.chunk(2, dim=-1)
+    c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
+    q1, q2, k1, k2, c, s = [t.to(dtype=torch.float32) for t in [q1, q2, k1, k2, c, s]]
+
+    q_rot = torch.cat([q1 * c - q2 * s, q1 * s + q2 * c], axis=-1).to(qkv.dtype)
+    k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(qkv.dtype)
+
+    return torch.cat(
+        [
+            torch.cat([q_rot, q_pass], axis=-1).unsqueeze(2),
+            torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
+            qkv[:, :, 2:3, :, :],
+        ],
+        axis=2,
+    )
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary positional embedding (RoPE).
+
+    Reference:
+        RoFormer: Enhanced Transformer with Rotary Position Embedding.
+        https://arxiv.org/pdf/2104.09864.pdf.
+
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        base: int = 10000,
+        scale_base: Optional[float] = None,
+        pos_idx_in_fp32: bool = True,
+        max_position_embeddings: int = 2048,
+        device: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+
+        if scale_base is not None:
+            raise NotImplementedError
+
+        self.dim = dim
+        self.base = float(base)
+        self.scale_base = scale_base
+        self.pos_idx_in_fp32 = pos_idx_in_fp32
+        self.max_position_embeddings = max_position_embeddings
+        self.device = device
+
+        # Generate and save the inverse frequency buffer (non-trainable)
+        inv_freq = self._compute_inv_freq(device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Generate and save the scale buffer (non-trainable)
+        scale = (
+            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
+            if scale_base is not None
+            else None
+        )
+        self.register_buffer("scale", scale, persistent=False)
+
+        # Initialize cached attributes since ONNX can't rely on dynamic initialization
+        self._update_cos_sin_cache(max_position_embeddings, device=device, dtype=torch.float32)
+
+    def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor:
+        return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
+
+    def _update_cos_sin_cache(
+        self,
+        seqlen: int,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> None:
+        self._seq_len_cached = seqlen
+
+        # fp32 is preferred since the output of `torch.arange` can be quite large
+        # and bf16 would lose a lot of precision
+        if self.pos_idx_in_fp32:
+            t = torch.arange(seqlen, device=device, dtype=torch.float32)
+            if self.inv_freq.dtype != torch.float32:
+                inv_freq = self._compute_inv_freq(device=device)
+            else:
+                inv_freq = self.inv_freq
+        else:
+            t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
+            inv_freq = self.inv_freq
+
+        # `torch.outer` is preferred since `torch.einsum` converts from fp32 to fp16 if used with AMP
+        freqs = torch.outer(t, inv_freq)
+        if self.scale is None:
+            self._cos_cached = torch.cos(freqs).to(dtype)
+            self._sin_cached = torch.sin(freqs).to(dtype)
+        else:
+            power = (
+                torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
+            ) / self.scale_base
+            scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
+
+            # Force the scale multiplication to happen in fp32
+            self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
+            self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
+            self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
+            self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
+
+    def forward(
+        self,
+        qkv: torch.Tensor,
+        kv: Optional[torch.Tensor] = None,
+        seqlen_offset: int = 0,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if (
+            self._seq_len_cached < qkv.shape[1] + seqlen_offset
+            or self._cos_cached.device != qkv.device
+            or self._cos_cached.dtype != qkv.dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            self._update_cos_sin_cache(qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype)
+
+        if kv is None:
+            return _apply_rotary_emb_qkv(
+                qkv,
+                self._cos_cached[seqlen_offset:],
+                self._sin_cached[seqlen_offset:],
+            )
+        else:
+            q = _apply_rotary_emb(
+                qkv,
+                self._cos_cached[seqlen_offset:],
+                self._sin_cached[seqlen_offset:],
+            )
+            kv = _apply_rotary_emb_kv(
+                kv,
+                self._cos_cached[seqlen_offset:],
+                self._sin_cached[seqlen_offset:],
+            )
+
+            return q, kv
+
+
+class MLP(nn.Module):
+    """Multi-Layer Perceptron.
+
+    Reference:
+        Attention Is All You Need.
+        https://arxiv.org/pdf/1706.03762.pdf.
+
+    """
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        n_inner: Optional[int] = None,
+        act_fn: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+
+        act_fn = config.activation_function if act_fn is None else act_fn
+
+        n_inner = getattr(config, "n_inner", None) if n_inner is None else n_inner
+        n_inner = n_inner if n_inner is not None else 4 * config.n_embd
+
+        self.fc1 = nn.Linear(config.n_embd, n_inner)
+        self.fc2 = nn.Linear(n_inner, config.n_embd)
+        self.act = ACT2FN[act_fn]
+
+    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+
+        return hidden_states
+
+
+class SelfAttention(nn.Module):
+    """Self-attention layer (compatible with PyTorch).
+
+    Reference:
+        https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
+
+    """
+
+    def __init__(
+        self,
+        causal: bool = True,
+        softmax_scale: Optional[float] = None,
+        attention_dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.drop = nn.Dropout(attention_dropout)
+
+    @torch.autocast("cpu", enabled=False)
+    @torch.autocast("cuda", enabled=False)
+    def forward(
+        self,
+        qkv: torch.FloatTensor,
+        causal: bool = None,
+        key_padding_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        batch_size, seqlen = qkv.shape[0], qkv.shape[1]
+        q, k, v = qkv.unbind(dim=2)
+
+        q = q.to(torch.float32)
+        k = k.to(torch.float32)
+
+        causal = self.causal if causal is None else causal
+        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+
+        # Autocast is manually disabled to avoid `torch.einsum` performing the operation
+        # using float16, which might lead to overflow
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+        if key_padding_mask is not None:
+            padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device)
+            padding_mask.masked_fill_(key_padding_mask, 0.0)
+
+            scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
+
+        if causal:
+            causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1)
+            scores = scores + causal_mask.to(dtype=scores.dtype)
+
+        attention = torch.softmax(scores, dim=-1).to(v.dtype)
+        attention = self.drop(attention)
+
+        output = torch.einsum("bhts,bshd->bthd", attention, v)
+
+        return output
+
+
+class CrossAttention(nn.Module):
+    """Cross-attention layer (compatible with PyTorch).
+
+    Reference:
+        https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
+
+    """
+
+    def __init__(
+        self,
+        causal: bool = True,
+        softmax_scale: Optional[float] = None,
+        attention_dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.drop = nn.Dropout(attention_dropout)
+
+    @torch.autocast("cpu", enabled=False)
+    @torch.autocast("cuda", enabled=False)
+    def forward(
+        self,
+        q: torch.FloatTensor,
+        kv: torch.FloatTensor,
+        causal: bool = None,
+        key_padding_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        batch_size, seqlen_q = q.shape[0], q.shape[1]
+        seqlen_k = kv.shape[1]
+
+        if kv.shape[3] != q.shape[2]:
+            kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3])
+        k, v = kv.unbind(dim=2)
+
+        q = q.to(torch.float32)
+        k = k.to(torch.float32)
+
+        causal = self.causal if causal is None else causal
+        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+
+        # Autocast is manually disabled to avoid `torch.einsum` performing the operation
+        # using float16, which might lead to overflow
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+        if key_padding_mask is not None:
+            padding_mask = torch.full(
+                (batch_size, seqlen_k),
+                -10000.0,
+                dtype=scores.dtype,
+                device=scores.device,
+            )
+            padding_mask.masked_fill_(key_padding_mask, 0.0)
+
+            scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
+
+        if causal:
+            rows = rearrange(torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1")
+            cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long)
+            causal_mask = cols > rows + seqlen_k - seqlen_q
+
+            scores = scores.masked_fill(causal_mask, -10000.0)
+
+        attention = torch.softmax(scores, dim=-1).to(v.dtype)
+        attention = self.drop(attention)
+
+        output = torch.einsum("bhts,bshd->bthd", attention, v)
+
+        return output
+
+
+def _find_mha_dims(
+    config: PretrainedConfig,
+    n_head: Optional[int] = None,
+    n_head_kv: Optional[int] = None,
+    head_dim: Optional[int] = None,
+) -> Tuple[int, int]:
+    if n_head is None and head_dim is None:
+        head_dim = config.n_embd // config.n_head
+        n_head = config.n_head
+    elif n_head is None or head_dim is None:
+        raise ValueError("`n_head` and `head_dim` must be both specified or `None`.")
+
+    if n_head_kv is None:
+        n_head_kv = getattr(config, "n_head_kv", None) or n_head
+
+    return n_head, n_head_kv, head_dim
+
+
+def _update_kv_cache(kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int) -> torch.FloatTensor:
+    num_heads, head_dim = kv.shape[-2:]
+
+    if layer_idx not in inference_params.key_value_memory_dict:
+        inference_params.key_value_memory_dict[layer_idx] = torch.empty(
+            inference_params.max_batch_size,
+            inference_params.max_seqlen,
+            2,
+            num_heads,
+            head_dim,
+            dtype=kv.dtype,
+            device=kv.device,
+        )
+
+    batch_start = inference_params.batch_size_offset
+    batch_end = batch_start + kv.shape[0]
+
+    sequence_start = inference_params.seqlen_offset
+    sequence_end = sequence_start + kv.shape[1]
+
+    # When the current sequence length is equal to or larger than the maximum sequence length,
+    # we need to concatenate the current `kv` with the cached `kv` to expand its length
+    if sequence_end >= inference_params.max_seqlen:
+        inference_params.key_value_memory_dict[layer_idx] = torch.concatenate((inference_params.key_value_memory_dict[layer_idx], kv), dim=1)
+
+    inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...] = kv
+    kv = inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, :sequence_end, ...]
+        
+    return kv
+
+
+class MHA(nn.Module):
+    """Multi-head attention layer."""
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[str] = None,
+        rotary_dim: Optional[int] = None,
+        rotary_base: float = 10000.0,
+        rotary_scale_base: Optional[float] = None,
+        n_head: Optional[int] = None,
+        n_head_kv: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        bias: bool = True,
+        causal: bool = True,
+        softmax_scale: Optional[float] = None,
+        layer_idx: Optional[int] = None,
+        return_residual: bool = False,
+        checkpointing: bool = False,
+    ) -> None:
+        super().__init__()
+
+        # Rotary embedding
+        self.rotary_dim = rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0)
+        if self.rotary_dim > 0:
+            rotary_cls = FlashRotaryEmbedding if config.flash_rotary else RotaryEmbedding
+            if rotary_cls is None:
+                rotary_cls = RotaryEmbedding
+
+            rotary_kwargs = {}
+            if rotary_cls is RotaryEmbedding:
+                rotary_kwargs["max_position_embeddings"] = config.n_positions
+
+            self.rotary_emb = rotary_cls(
+                self.rotary_dim,
+                base=rotary_base,
+                scale_base=rotary_scale_base,
+                device=device,
+                **rotary_kwargs,
+            )
+
+        # MLP
+        self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims(
+            config, n_head=n_head, n_head_kv=n_head_kv, head_dim=head_dim
+        )
+        op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv)
+        hidden_size = config.n_embd
+
+        linear_cls = FusedDense if config.fused_dense else nn.Linear
+        if linear_cls is None:
+            linear_cls = nn.Linear
+
+        self.Wqkv = linear_cls(hidden_size, op_size, bias=bias, device=device, dtype=dtype)
+        self.out_proj = linear_cls(hidden_size, hidden_size, bias=bias, device=device, dtype=dtype)
+
+        # Attention
+        attn_cls = FlashSelfAttention if config.flash_attn else SelfAttention
+        if attn_cls is None:
+            attn_cls = SelfAttention
+
+        cross_attn_cls = FlashCrossAttention if config.flash_attn else CrossAttention
+        if cross_attn_cls is None:
+            cross_attn_cls = CrossAttention
+
+        self.inner_attn = attn_cls(
+            causal=causal,
+            softmax_scale=softmax_scale,
+            attention_dropout=config.attn_pdrop,
+        )
+        self.inner_cross_attn = cross_attn_cls(
+            causal=causal,
+            softmax_scale=softmax_scale,
+            attention_dropout=config.attn_pdrop,
+        )
+
+        self.flash_attn = config.flash_attn and attn_cls is FlashSelfAttention
+        self.layer_idx = layer_idx
+        self.return_residual = return_residual
+        self.checkpointing = checkpointing
+
+    def _forward_self_attn(
+        self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor]
+    ) -> torch.FloatTensor:
+        qkv = self.Wqkv(x)
+        qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim)
+
+        if self.rotary_dim > 0:
+            qkv = self.rotary_emb(qkv)
+
+        if self.flash_attn:
+            batch_size, seqlen = qkv.shape[0], qkv.shape[1]
+
+            cu_seqlens, max_seqlen = None, None
+            if key_padding_mask is not None:
+                # If `key_padding_mask` is supplied, we need to unpad the input and retrieve
+                # the `cu_seqlens` and `max_seqlen` to be used by `flash-attn`
+                qkv, indices, cu_seqlens, max_seqlen = unpad_input(qkv, key_padding_mask)
+
+            if self.checkpointing:
+                attn_output = torch.utils.checkpoint.checkpoint(
+                    self.inner_attn, qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
+                )
+            else:
+                attn_output = self.inner_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen).to(qkv.device)
+
+            # If `key_padding_mask` is supplied, we need to pad the output back to the original shape
+            return pad_input(attn_output, indices, batch_size, seqlen) if key_padding_mask is not None else attn_output
+
+        if self.checkpointing:
+            return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, key_padding_mask=key_padding_mask)
+
+        return self.inner_attn(qkv, key_padding_mask=key_padding_mask)
+
+    def _forward_cross_attn(
+        self,
+        x: torch.FloatTensor,
+        past_key_values: Optional[InferenceParams],
+        key_padding_mask: Optional[torch.BoolTensor],
+    ) -> torch.FloatTensor:
+        batch_size = x.shape[0]
+
+        qkv = self.Wqkv(x)
+
+        q = qkv[..., : self.n_head * self.head_dim]
+        q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
+
+        kv = qkv[..., self.n_head * self.head_dim :]
+        kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)
+
+        seqlen_offset = past_key_values.seqlen_offset if past_key_values is not None else 0
+        causal = None if seqlen_offset == 0 else False
+        if self.rotary_dim > 0:
+            q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset)
+
+        if past_key_values is not None:
+            kv = _update_kv_cache(kv, past_key_values, self.layer_idx)
+
+        if self.flash_attn:
+            batch_size, seqlen_q = q.shape[0], q.shape[1]
+            seqlen_k = kv.shape[1]
+
+            cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = (
+                None,
+                None,
+                None,
+                None,
+            )
+            if key_padding_mask is not None:
+                kv, _, cu_seqlens_k, max_seqlen_k = unpad_input(kv, key_padding_mask)
+
+                if seqlen_q == 1:
+                    key_padding_mask = torch.ones(batch_size, 1, device=q.device)
+                elif seqlen_q != seqlen_k:
+                    key_padding_mask = key_padding_mask[:, -seqlen_q:]
+
+                q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, key_padding_mask)
+
+            if self.checkpointing:
+                attn_output = torch.utils.checkpoint.checkpoint(
+                    self.inner_cross_attn,
+                    q,
+                    kv,
+                    causal=causal,
+                    cu_seqlens=cu_seqlens_q,
+                    max_seqlen=max_seqlen_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_k=max_seqlen_k,
+                )
+            else:
+                attn_output = self.inner_cross_attn(
+                    q,
+                    kv,
+                    causal=causal,
+                    cu_seqlens=cu_seqlens_q,
+                    max_seqlen=max_seqlen_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_k=max_seqlen_k,
+                )
+
+            return (
+                pad_input(attn_output, indices_q, batch_size, max_seqlen_q)
+                if key_padding_mask is not None
+                else attn_output
+            )
+
+        if self.checkpointing:
+            return torch.utils.checkpoint.checkpoint(
+                self.inner_cross_attn,
+                q,
+                kv,
+                key_padding_mask=key_padding_mask,
+                causal=causal,
+            )
+
+        return self.inner_cross_attn(q, kv, key_padding_mask=key_padding_mask, causal=causal)
+
+    def forward(
+        self,
+        x: torch.FloatTensor,
+        past_key_values: Optional[InferenceParams] = None,
+        attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+        if attention_mask is not None:
+            attention_mask = attention_mask.bool()
+        else:
+            attention_mask = None
+
+        # MHA
+        if self.n_head == self.n_head_kv:
+            if past_key_values is None:
+                # If `past_key_values` are not supplied, we run self-attention
+                attn_output = self._forward_self_attn(x, attention_mask)
+            else:
+                # If `past_key_values` are supplied, it means that we might have cached values and
+                # could take advantage of cross-attention
+                attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)
+        # MQA / GQA
+        else:
+            # Regardless of `past_key_values` being supplied or not, it always use cross-attention
+            # because `q` and `kv` lengths might be different
+            attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)
+
+        output = rearrange(attn_output, "... h d -> ... (h d)")
+        output = self.out_proj(output)
+
+        return output if not self.return_residual else (output, x)
+
+
+class ParallelBlock(nn.Module):
+    """Parallel block.
+
+    This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen).
+
+    """
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        block_idx: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+
+        self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+        self.block_idx = block_idx
+
+        self.mixer = MHA(config, layer_idx=block_idx)
+        self.mlp = MLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        residual = hidden_states
+        hidden_states = self.ln(hidden_states)
+
+        attn_outputs = self.mixer(
+            hidden_states,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+        )
+        if isinstance(attn_outputs, tuple):
+            attn_outputs = attn_outputs[0]
+
+        attn_outputs = self.resid_dropout(attn_outputs)
+        feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
+
+        hidden_states = attn_outputs + feed_forward_hidden_states + residual
+
+        return hidden_states
+
+
+class CausalLMHead(nn.Module):
+    """Causal Language Modeling head.
+
+    Reference:
+        Improving Language Understanding by Generative Pre-Training.
+        https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
+
+    """
+
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+
+        self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.linear = nn.Linear(config.n_embd, config.vocab_size)
+
+    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+        hidden_states = self.ln(hidden_states)
+        logits = self.linear(hidden_states).to(torch.float32)
+
+        return logits
+
+
+class CausalLMLoss(nn.Module):
+    """Causal Language Modeling loss.
+
+    Reference:
+        Improving Language Understanding by Generative Pre-Training.
+        https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
+
+    """
+
+    def __init__(self, shift_labels: bool = True) -> None:
+        super().__init__()
+
+        self.shift_labels = shift_labels
+        self.loss_fct = nn.CrossEntropyLoss()
+
+    def forward(self, logits: torch.FloatTensor, labels: torch.LongTensor) -> torch.FloatTensor:
+        if self.shift_labels:
+            logits = logits[..., :-1, :].contiguous()
+            labels = labels[..., 1:].contiguous()
+
+        loss = self.loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        return loss
+
+
+class PhiPreTrainedModel(PreTrainedModel):
+    """Phi pre-trained model."""
+
+    config_class = PhiConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["ParallelBlock"]
+
+    def __init__(self, *inputs, **kwargs) -> None:
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, (nn.Linear,)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            if module.bias is not None:
+                module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        if past_key_values is None or not (isinstance(past_key_values, InferenceParams)):
+            past_key_values = InferenceParams(
+                max_seqlen=self.config.n_positions,
+                max_batch_size=input_ids.shape[0],
+                seqlen_offset=0,
+                batch_size_offset=0,
+                key_value_memory_dict={},
+                lengths_per_sample=None,
+            )
+        else:
+            # Assume that `past_key_values` has cached all tokens up to the last token in `input_ids`
+            past_key_values.seqlen_offset = input_ids.shape[1] - 1
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+        }
+
+
+class PhiModel(PhiPreTrainedModel):
+    """Phi model."""
+
+    _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: PhiConfig) -> None:
+        super().__init__(config)
+
+        self.embd = Embedding(config)
+        self.h = nn.ModuleList([ParallelBlock(config, block_idx=i) for i in range(config.n_layer)])
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embd.wte
+
+    def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
+        self.embd.wte = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+    ) -> torch.FloatTensor:
+        
+        all_hidden_states = () if output_hidden_states else None
+
+        hidden_states = self.embd(input_ids)
+
+        for layer in self.h:
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            hidden_states = layer(
+                hidden_states,
+                past_key_values=past_key_values,
+                attention_mask=attention_mask,
+            )
+        if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+        return BaseModelOutputWithNoAttention(
+            last_hidden_state=hidden_states,
+            hidden_states = all_hidden_states
+        )
+
+
+class PhiForCausalLM(PhiPreTrainedModel):
+    """Phi for Causal Language Modeling."""
+
+    _keys_to_ignore_on_load_missing = [""]
+    _keys_to_ignore_on_load_unexpected = [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
+
+    def __init__(self, config: PhiConfig) -> None:
+        super().__init__(config)
+
+        self.transformer = PhiModel(config)
+        self.lm_head = CausalLMHead(config)
+        self.loss = CausalLMLoss()
+
+        self.post_init()
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head.linear
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.lm_head.linear = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        output = self.transformer(
+            input_ids, past_key_values=past_key_values,
+              attention_mask=attention_mask, output_hidden_states=output_hidden_states)
+        hidden_states = output.last_hidden_state
+        lm_logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss(lm_logits, labels)
+
+        return CausalLMOutputWithPast(
+            loss=loss, 
+            logits=lm_logits,
+            past_key_values=past_key_values,
+            hidden_states=output.hidden_states
+        )

special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

tokenizer_config.json

@@ -0,0 +1,323 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50257": {
+      "content": "                               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50258": {
+      "content": "                              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50259": {
+      "content": "                             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50260": {
+      "content": "                            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50261": {
+      "content": "                           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50262": {
+      "content": "                          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50263": {
+      "content": "                         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50264": {
+      "content": "                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50265": {
+      "content": "                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50266": {
+      "content": "                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50267": {
+      "content": "                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50268": {
+      "content": "                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50269": {
+      "content": "                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50270": {
+      "content": "                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50271": {
+      "content": "                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50272": {
+      "content": "                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50273": {
+      "content": "               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50274": {
+      "content": "              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50275": {
+      "content": "             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50276": {
+      "content": "            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50277": {
+      "content": "           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50278": {
+      "content": "          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50279": {
+      "content": "         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50280": {
+      "content": "        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50281": {
+      "content": "       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50282": {
+      "content": "      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50283": {
+      "content": "     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50284": {
+      "content": "    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50285": {
+      "content": "   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50286": {
+      "content": "  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50287": {
+      "content": "\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50288": {
+      "content": "\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50289": {
+      "content": "\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50290": {
+      "content": "\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50291": {
+      "content": "\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50292": {
+      "content": "\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50293": {
+      "content": "\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50294": {
+      "content": "\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 2048,
+  "tokenizer_class": "CodeGenTokenizer",
+  "unk_token": "<|endoftext|>"
+}