Upload 26 files

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,71 @@
+MICROSOFT RESEARCH LICENSE TERMS
+
+IF YOU LIVE IN THE UNITED STATES, PLEASE READ THE “BINDING ARBITRATION AND CLASS ACTION WAIVER” SECTION BELOW. IT AFFECTS HOW DISPUTES ARE RESOLVED.
+
+These license terms are an agreement between you and Microsoft Corporation (or one of its affiliates). They apply to the source code, object code, machine learning models, or data (collectively “Materials”) that accompany this license. IF YOU COMPLY WITH THESE LICENSE TERMS, YOU HAVE THE RIGHTS BELOW. BY USING THE MATERIALS, YOU ACCEPT THESE TERMS.
+
+1) INSTALLATION AND USE RIGHTS TO THE MATERIALS.
+
+Subject to the terms of this agreement, you have the below rights, if applicable, to use the Materials solely for non-commercial, non-revenue generating, research purposes:
+
+a) Source Code. If source code is included, you may use and modify the source code, but you may not distribute the source code.
+
+b) Object Code. If object code is included, you may use the object code, but you may not distribute the object code.
+
+c) Models. If machine learning model(s) are included, you may use the model(s), but you may not distribute the models.
+
+d) Data. If data is included, you may use and modify the data, but your use and modification must be consistent with the consent under which the data was provided and/or gathered and you may not distribute the data or your modifications to the data.
+
+2) SCOPE OF LICENSE. The Materials are licensed, not sold. Microsoft reserves all other rights. Unless applicable law gives you more rights despite this limitation, you will not (and have no right to):
+
+a) work around any technical limitations in the Materials that only allow you to use it in certain ways;
+
+b) reverse engineer, decompile or disassemble the Materials;
+
+c) remove, minimize, block, or modify any notices of Microsoft or its suppliers in the Materials;
+
+d) use the Materials in any way that is against the law or to create or propagate malware; or
+
+e) share, publish, distribute or lend the Materials, provide the Materials as a stand-alone hosted solution for others to use, or transfer the Materials or this agreement to any third party.
+
+3) PERSONAL DATA. If the data (set forth in Section 1(c) above) includes or is found to include any data that enables any ability to identify an individual (“Personal Data”), you will not use such Personal Data for any purpose other than was authorized and consented to by the data subject/research participant. You will not use Personal Data to contact any person. You will keep Personal Data in strict confidence. You will not share any Personal Data that is collected or in your possession with any third party for any reason and as required under the original consent agreement. Further, you will destroy the Personal Data and any backup or copies, immediately upon the completion of your research.
+
+4) LICENSE TO MICROSOFT. Notwithstanding the limitations in Section 1, you may distribute your modifications back to Microsoft, and if you do provide Microsoft with modifications of the Materials, you hereby grant Microsoft, without any restrictions or limitations, a non-exclusive, perpetual, irrevocable, royalty-free, assignable and sub-licensable license, to reproduce, publicly perform or display, install, use, modify, post, distribute, make and have made, sell and transfer such modifications and derivatives for any purpose.
+
+5) PUBLICATION. You may publish (or present papers or articles) on your results from using the Materials provided that no material or substantial portion of the Materials is included in any such publication or presentation.
+
+6) FEEDBACK. Any feedback about the Materials provided by you to us is voluntarily given, and Microsoft shall be free to use the feedback as it sees fit without obligation or restriction of any kind, even if the
+
+feedback is designated by you as confidential. Such feedback shall be considered a contribution and licensed to Microsoft under the terms of Section 4 above.
+
+7) EXPORT RESTRICTIONS. You must comply with all domestic and international export laws and regulations that apply to the Materials, which include restrictions on destinations, end users, and end use. For further information on export restrictions, visit (aka.ms/exporting).
+
+8) SUPPORT SERVICES. Microsoft is not obligated under this agreement to provide any support services for the Materials. Any support provided is “as is”, “with all faults”, and without warranty of any kind.
+
+9) BINDING ARBITRATION AND CLASS ACTION WAIVER. This Section applies if you live in (or, if a business, your principal place of business is in) the United States. If you and Microsoft have a dispute, you and Microsoft agree to try for 60 days to resolve it informally. If you and Microsoft can’t, you and Microsoft agree to binding individual arbitration before the American Arbitration Association under the Federal Arbitration Act (“FAA”), and not to sue in court in front of a judge or jury. Instead, a neutral arbitrator will decide. Class action lawsuits, class-wide arbitrations, private attorney-general actions, and any other proceeding where someone acts in a representative capacity are not allowed; nor is combining individual proceedings without the consent of all parties. The complete Arbitration Agreement contains more terms and is at aka.ms/arb-agreement-1. You and Microsoft agree to these terms.
+
+10) ENTIRE AGREEMENT. This agreement, and any other terms Microsoft may provide for supplements, updates, or third-party applications, is the entire agreement for the Materials.
+
+11) APPLICABLE LAW AND PLACE TO RESOLVE DISPUTES. If you acquired the Materials in the United States or Canada, the laws of the state or province where you live (or, if a business, where your principal place of business is located) govern the interpretation of this agreement, claims for its breach, and all other claims (including consumer protection, unfair competition, and tort claims), regardless of conflict of laws principles, except that the FAA governs everything related to arbitration. If you acquired the Materials in any other country, its laws apply, except that the FAA governs everything related to arbitration. If U.S. federal jurisdiction exists, you and Microsoft consent to exclusive jurisdiction and venue in the federal court in King County, Washington for all disputes heard in court (excluding arbitration). If not, you and Microsoft consent to exclusive jurisdiction and venue in the Superior Court of King County, Washington for all disputes heard in court (excluding arbitration).
+
+12) CONSUMER RIGHTS; REGIONAL VARIATIONS. This agreement describes certain legal rights. You may have other rights, including consumer rights, under the laws of your state, province, or country. Separate and apart from your relationship with Microsoft, you may also have rights with respect to the party from which you acquired the Materials. This agreement does not change those other rights if the laws of your state, province, or country do not permit it to do so. For example, if you acquired the Materials in one of the below regions, or mandatory country law applies, then the following provisions apply to you:
+
+a) Australia. You have statutory guarantees under the Australian Consumer Law and nothing in this agreement is intended to affect those rights.
+
+b) Canada. If you acquired this software in Canada, you may stop receiving updates by turning off the automatic update feature, disconnecting your device from the Internet (if and when you re-connect to the Internet, however, the Materials will resume checking for and installing updates), or uninstalling the Materials. The product documentation, if any, may also specify how to turn off updates for your specific device or software.
+
+c) Germany and Austria.
+
+i. Warranty. The properly licensed software will perform substantially as described in any Microsoft materials that accompany the Materials. However, Microsoft gives no contractual guarantee in relation to the licensed software.
+
+ii. Limitation of Liability. In case of intentional conduct, gross negligence, claims based on the Product Liability Act, as well as, in case of death or personal or physical injury, Microsoft is liable according to the statutory law.
+
+Subject to the foregoing clause (ii), Microsoft will only be liable for slight negligence if Microsoft is in breach of such material contractual obligations, the fulfillment of which facilitate the due performance of this agreement, the breach of which would endanger the purpose of this agreement and the compliance with which a party may constantly trust in (so-called "cardinal obligations"). In other cases of slight negligence, Microsoft will not be liable for slight negligence.
+
+13) DISCLAIMER OF WARRANTY. THE MATERIALS ARE LICENSED “AS IS.” YOU BEAR THE RISK OF USING THEM. MICROSOFT GIVES NO EXPRESS WARRANTIES, GUARANTEES, OR CONDITIONS. TO THE EXTENT PERMITTED UNDER APPLICABLE LAWS, MICROSOFT EXCLUDES ALL IMPLIED WARRANTIES, INCLUDING MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND NON-INFRINGEMENT.
+
+14) LIMITATION ON AND EXCLUSION OF DAMAGES. IF YOU HAVE ANY BASIS FOR RECOVERING DAMAGES DESPITE THE PRECEDING DISCLAIMER OF WARRANTY, YOU CAN RECOVER FROM MICROSOFT AND ITS SUPPLIERS ONLY DIRECT DAMAGES UP TO U.S. $5.00. YOU CANNOT RECOVER ANY OTHER DAMAGES, INCLUDING CONSEQUENTIAL, LOST PROFITS, SPECIAL, INDIRECT OR INCIDENTAL DAMAGES.
+
+This limitation applies to (a) anything related to the Materials, services, content (including code) on third party Internet sites, or third party applications; and (b) claims for breach of contract, warranty, guarantee, or condition; strict liability, negligence, or other tort; or any other claim; in each case to the extent permitted by applicable law.
+
+It also applies even if Microsoft knew or should have known about the possibility of the damages. The above limitation or exclusion may not apply to you because your state, province, or country may not allow the exclusion or limitation of incidental, consequential, or other damages.

MLmodel

@@ -0,0 +1,24 @@
+flavors:
+  hftransformersv2:
+    code: null
+    config_hf_load_kwargs:
+      trust_remote_code: true
+    hf_config_class: AutoConfig
+    hf_pretrained_class: AutoModelForCausalLM
+    hf_tokenizer_class: CodeGenTokenizerFast
+    model_data: data
+    model_hf_load_args:
+      trust_remote_code: true
+    pytorch_version: 2.1.0+cu118
+    task_type: text-generation
+    tokenizer_hf_load_kwargs:
+      trust_remote_code: true
+    transformers_version: 4.34.0
+  python_function:
+    data: data
+    env: conda.yaml
+    loader_module: azureml.evaluate.mlflow.hftransformers
+    python_version: 3.10.11
+mlflow_version: 2.6.0
+model_uuid: 6068cffa9b034ea28c997f4538233299
+utc_time_created: '2023-11-06 18:18:55.524636'

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,41 @@
+{
+  "\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,
+  "<|endoftext|>": 50256
+}

amlignore

@@ -0,0 +1,6 @@
+## This file was auto generated by the Azure Machine Learning Studio. Please do not remove. 
+## Read more about the .amlignore file here: https://docs.microsoft.com/azure/machine-learning/how-to-save-write-experiment-files#storage-limits-of-experiment-snapshots
+        
+.ipynb_aml_checkpoints/ 
+*.amltmp 
+*.amltemp

amlignore (1)

@@ -0,0 +1,6 @@
+## This file was auto generated by the Azure Machine Learning Studio. Please do not remove. 
+## Read more about the .amlignore file here: https://docs.microsoft.com/azure/machine-learning/how-to-save-write-experiment-files#storage-limits-of-experiment-snapshots
+        
+.ipynb_aml_checkpoints/ 
+*.amltmp 
+*.amltemp

conda.yaml

@@ -0,0 +1,28 @@
+channels:
+- conda-forge
+dependencies:
+- python=3.10.11
+- pip<=23.1.2
+- pip:
+  - mlflow==2.6.0
+  - cloudpickle==2.2.1
+  - jsonpickle==3.0.1
+  - mlflow-skinny==2.6.0
+  - azureml-core==1.51.0.post1
+  - azureml-mlflow==1.51.0
+  - azureml-metrics[all]==0.0.32
+  - scikit-learn==1.2.2
+  - cryptography==41.0.1
+  - python-dateutil==2.8.2
+  - datasets==2.14.6
+  - soundfile==0.12.1
+  - librosa==0.10.1
+  - diffusers==0.21.4
+  - sentencepiece==0.1.99
+  - transformers==4.34.0
+  - torch==2.1.0
+  - accelerate==0.23.0
+  - Pillow==9.4.0
+  - einops
+  - azureml-evaluate-mlflow==0.0.32
+name: mlflow-env

config (1).json

@@ -0,0 +1,36 @@
+{
+  "_name_or_path": "mojanjp/phi-2",
+  "activation_function": "gelu_new",
+  "architecture": {
+    "block_cls": "parallel",
+    "mlp": {
+      "mlp_cls": "fused_mlp"
+    }
+  },
+  "architectures": [
+    "MixFormerSequentialForCausalLM"
+  ],
+  "attn_pdrop": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_mixformer_sequential.MixFormerSequentialConfig",
+    "AutoModelForCausalLM": "modeling_mixformer_sequential.MixFormerSequentialForCausalLM"
+  },
+  "embd_pdrop": 0.0,
+  "flash_rotary": false,
+  "fused_dense": false,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "mixformer-sequential",
+  "n_embd": 2560,
+  "n_head": 32,
+  "n_head_kv": null,
+  "n_inner": null,
+  "n_layer": 32,
+  "n_positions": 2048,
+  "resid_pdrop": 0.0,
+  "rotary_dim": 32,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.34.0",
+  "vocab_size": 51200
+}

config.json

@@ -0,0 +1,36 @@
+{
+  "_name_or_path": "mojanjp/phi-2",
+  "activation_function": "gelu_new",
+  "architecture": {
+    "block_cls": "parallel",
+    "mlp": {
+      "mlp_cls": "fused_mlp"
+    }
+  },
+  "architectures": [
+    "MixFormerSequentialForCausalLM"
+  ],
+  "attn_pdrop": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_mixformer_sequential.MixFormerSequentialConfig",
+    "AutoModelForCausalLM": "modeling_mixformer_sequential.MixFormerSequentialForCausalLM"
+  },
+  "embd_pdrop": 0.0,
+  "flash_rotary": false,
+  "fused_dense": false,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "mixformer-sequential",
+  "n_embd": 2560,
+  "n_head": 32,
+  "n_head_kv": null,
+  "n_inner": null,
+  "n_layer": 32,
+  "n_positions": 2048,
+  "resid_pdrop": 0.0,
+  "rotary_dim": 32,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.34.0",
+  "vocab_size": 51200
+}

configuration_mixformer_sequential (1).py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import math
+from typing import Optional
+
+from transformers import PretrainedConfig
+
+
+class MixFormerSequentialConfig(PretrainedConfig):
+    """MixFormer (sequential for DeepSpeed) configuration."""
+
+    model_type = "mixformer-sequential"
+
+    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_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_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)

configuration_mixformer_sequential.py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import math
+from typing import Optional
+
+from transformers import PretrainedConfig
+
+
+class MixFormerSequentialConfig(PretrainedConfig):
+    """MixFormer (sequential for DeepSpeed) configuration."""
+
+    model_type = "mixformer-sequential"
+
+    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_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_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)

finetune_config.json

@@ -0,0 +1,30 @@
+{
+  "load_config_kwargs": {
+    "trust_remote_code": true
+  },
+  "load_tokenizer_kwargs": {
+    "pad_token": "<|endoftext|>",
+    "trust_remote_code": true
+  },
+  "finetune_args": {},
+  "mlflow_ft_conf": {
+    "mlflow_hftransformers_misc_conf": {
+      "config_hf_load_kwargs": {
+        "trust_remote_code": true
+      },
+      "tokenizer_hf_load_kwargs": {
+        "return_token_type_ids": false
+      },
+      "model_hf_load_kwargs": {
+        "trust_remote_code": true,
+        "ignore_mismatched_sizes": true
+      },
+      "hf_predict_module": "phi_predict"
+    },
+    "mlflow_save_model_kwargs": {
+      "extra_pip_requirements": [
+        "einops"
+      ]
+    }
+  }
+}

generation_config.json

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

modeling_mixformer_sequential (1).py

@@ -0,0 +1,855 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+#
+# BSD 3-Clause License
+# 
+# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu.
+# 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.
+
+from __future__ import annotations
+
+import math
+from typing import Any, Dict, Optional, Tuple, Union
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+
+from einops import rearrange, repeat
+from transformers.activations import ACT2FN
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_mixformer_sequential import MixFormerSequentialConfig
+
+
+try:
+    from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
+    from flash_attn.ops.fused_dense import FusedDense
+except:
+    FlashRotaryEmbedding = 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, _, head_dim = x.shape
+    rotary_seqlen, rotary_dim = cos.shape
+    rotary_dim *= 2
+
+    assert rotary_dim <= head_dim
+    assert seqlen <= rotary_seqlen
+    assert cos.shape == sin.shape == (rotary_seqlen, 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, two, _, head_dim = kv.shape
+    assert two == 2
+
+    rotary_seqlen, rotary_dim = cos.shape
+    rotary_dim *= 2
+    assert rotary_dim <= head_dim
+    assert seqlen <= rotary_seqlen
+    assert cos.shape == sin.shape == (rotary_seqlen, 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, three, _, head_dim = qkv.shape
+    assert three == 3
+
+    rotary_seqlen, rotary_dim = cos.shape
+    rotary_dim *= 2
+    assert rotary_dim <= head_dim
+    assert seqlen <= rotary_seqlen
+    assert cos.shape == sin.shape == (rotary_seqlen, 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,
+        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.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)
+
+        self._seq_len_cached = 0
+        self._cos_cached = None
+        self._sin_cached = None
+        self._cos_k_cached = None
+        self._sin_k_cached = None
+
+    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:
+        # Reset the tables if sequence length has been chaned, if we are on a
+        # new device or if we are switching from inference mode to training
+        if (
+            seqlen > self._seq_len_cached
+            or self._cos_cached is None
+            or self._cos_cached.device != device
+            or self._cos_cached.dtype != dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            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,
+        max_seqlen: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        seqlen = qkv.shape[1]
+
+        if max_seqlen is not None:
+            self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
+        else:
+            self._update_cos_sin_cache(seqlen + 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
+        assert act_fn in ACT2FN.keys(), f"`act_fn` must be one of: {ACT2FN.keys()}."
+
+        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)
+
+    def forward(
+        self,
+        qkv: torch.FloatTensor,
+        causal: bool = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        batch_size, seqlen = qkv.shape[0], qkv.shape[1]
+        q, k, v = qkv.unbind(dim=2)
+
+        causal = self.causal if causal is None else causal
+        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+        if attention_mask is not None:
+            padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device)
+            padding_mask.masked_fill_(attention_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, dtype=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)
+
+    def forward(
+        self,
+        q: torch.FloatTensor,
+        kv: torch.FloatTensor,
+        causal: bool = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        batch_size, seqlen_q = q.shape[0], q.shape[1]
+        seqlen_k = kv.shape[1]
+        assert kv.shape[0] == batch_size and kv.shape[4] == q.shape[3]
+
+        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)
+
+        causal = self.causal if causal is None else causal
+        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+        if attention_mask is not None:
+            padding_mask = torch.full((batch_size, seqlen_k), -10000.0, dtype=scores.dtype, device=scores.device)
+            padding_mask.masked_fill_(attention_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, dtype=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]:
+    assert all(
+        hasattr(config, attr) for attr in ["n_embd", "n_head"]
+    ), "`config` must have `n_embd` and `n_head` attributes."
+
+    if head_dim is None:
+        assert (
+            config.n_embd % config.n_head == 0
+        ), f"Hidden size ({config.n_embd}) must be divisible by the number of heads ({config.n_head})."
+
+    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
+    assert n_head % n_head_kv == 0, "`n_head` must be divisible by `n_head_kv`."
+
+    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:
+        kv_cache = torch.empty(
+            inference_params.max_batch_size,
+            inference_params.max_seqlen,
+            2,
+            num_heads,
+            head_dim,
+            dtype=kv.dtype,
+            device=kv.device,
+        )
+        inference_params.key_value_memory_dict[layer_idx] = kv_cache
+    else:
+        kv_cache = inference_params.key_value_memory_dict[layer_idx]
+
+    batch_start = inference_params.batch_size_offset
+    batch_end = batch_start + kv.shape[0]
+    assert batch_end <= kv_cache.shape[0]
+
+    sequence_start = inference_params.seqlen_offset
+    sequence_end = sequence_start + kv.shape[1]
+    assert sequence_end <= kv_cache.shape[1]
+
+    assert kv_cache is not None
+    kv_cache[batch_start:batch_end, sequence_start:sequence_end, ...] = kv
+    kv = kv_cache[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_emb_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_emb_dim = rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0)
+        if self.rotary_emb_dim > 0:
+            rotary_kwargs = {"device": device}
+            if rotary_emb_scale_base is not None and rotary_emb_scale_base > 0.0:
+                rotary_kwargs["scale_base"] = rotary_emb_scale_base
+            
+            rotary_cls = FlashRotaryEmbedding if config.flash_rotary else RotaryEmbedding
+            if rotary_cls is None:
+                rotary_cls = RotaryEmbedding
+            self.rotary_emb = rotary_cls(self.rotary_emb_dim, **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
+        self.inner_attn = SelfAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=config.attn_pdrop)
+        self.inner_cross_attn = CrossAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=config.attn_pdrop)
+
+        self.layer_idx = layer_idx
+        self.return_residual = return_residual
+        self.checkpointing = checkpointing
+
+    def _forward_self_attn(
+        self, x: torch.FloatTensor, attention_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_emb_dim > 0:
+            qkv = self.rotary_emb(qkv)
+
+        if self.checkpointing:
+            return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, attention_mask=attention_mask)
+
+        return self.inner_attn(qkv, attention_mask=attention_mask)
+
+    def _forward_cross_attn(
+        self,
+        x: torch.FloatTensor,
+        past_key_values: Optional[InferenceParams],
+        attention_mask: Optional[torch.BoolTensor],
+    ) -> torch.FloatTensor:
+        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_emb_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.checkpointing:
+            return torch.utils.checkpoint.checkpoint(
+                self.inner_cross_attn, q, kv, attention_mask=attention_mask, causal=causal
+            )
+
+        return self.inner_cross_attn(q, kv, attention_mask=attention_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 and torch.any(~attention_mask.bool()):
+            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 MixFormerSequentialPreTrainedModel(PreTrainedModel):
+    """MixFormer (sequential for DeepSpeed) pre-trained model."""
+
+    config_class = MixFormerSequentialConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+
+    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 = len(input_ids[0]) - 1
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+        }
+    
+    def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False) -> None:
+        if isinstance(module, MixFormerSequentialPreTrainedModel):
+            module.gradient_checkpointing = value
+
+
+class MixFormerSequentialForCausalLM(MixFormerSequentialPreTrainedModel):
+    """MixFormer (sequential for DeepSpeed) for Causal Language Modeling."""
+
+    _keys_to_ignore_on_load_missing = [""]
+    _keys_to_ignore_on_load_unexpected = [r"layers\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
+    _no_split_modules = ["ParallelBlock"]
+
+    def __init__(self, config: MixFormerSequentialConfig) -> None:
+        super().__init__(config)
+
+        modules = [Embedding(config)]
+        modules += [ParallelBlock(config, block_idx=i) for i in range(config.n_layer)]
+        modules.append(CausalLMHead(config))
+
+        self.layers = nn.Sequential(*modules)
+        self.loss = CausalLMLoss()
+
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.layers[0].wte
+
+    def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
+        self.layers[0].wte = new_embeddings
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.layers[-1].linear
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.layers[-1].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,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden_layer = self.layers[0](input_ids)
+        for module in self.layers[1:-1]:
+            hidden_layer = module(hidden_layer, past_key_values=past_key_values, attention_mask=attention_mask)
+        lm_logits = self.layers[-1](hidden_layer)
+
+        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)

modeling_mixformer_sequential.py

@@ -0,0 +1,855 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+#
+# BSD 3-Clause License
+# 
+# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu.
+# 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.
+
+from __future__ import annotations
+
+import math
+from typing import Any, Dict, Optional, Tuple, Union
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+
+from einops import rearrange, repeat
+from transformers.activations import ACT2FN
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_mixformer_sequential import MixFormerSequentialConfig
+
+
+try:
+    from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
+    from flash_attn.ops.fused_dense import FusedDense
+except:
+    FlashRotaryEmbedding = 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, _, head_dim = x.shape
+    rotary_seqlen, rotary_dim = cos.shape
+    rotary_dim *= 2
+
+    assert rotary_dim <= head_dim
+    assert seqlen <= rotary_seqlen
+    assert cos.shape == sin.shape == (rotary_seqlen, 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, two, _, head_dim = kv.shape
+    assert two == 2
+
+    rotary_seqlen, rotary_dim = cos.shape
+    rotary_dim *= 2
+    assert rotary_dim <= head_dim
+    assert seqlen <= rotary_seqlen
+    assert cos.shape == sin.shape == (rotary_seqlen, 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, three, _, head_dim = qkv.shape
+    assert three == 3
+
+    rotary_seqlen, rotary_dim = cos.shape
+    rotary_dim *= 2
+    assert rotary_dim <= head_dim
+    assert seqlen <= rotary_seqlen
+    assert cos.shape == sin.shape == (rotary_seqlen, 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,
+        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.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)
+
+        self._seq_len_cached = 0
+        self._cos_cached = None
+        self._sin_cached = None
+        self._cos_k_cached = None
+        self._sin_k_cached = None
+
+    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:
+        # Reset the tables if sequence length has been chaned, if we are on a
+        # new device or if we are switching from inference mode to training
+        if (
+            seqlen > self._seq_len_cached
+            or self._cos_cached is None
+            or self._cos_cached.device != device
+            or self._cos_cached.dtype != dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            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,
+        max_seqlen: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        seqlen = qkv.shape[1]
+
+        if max_seqlen is not None:
+            self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
+        else:
+            self._update_cos_sin_cache(seqlen + 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
+        assert act_fn in ACT2FN.keys(), f"`act_fn` must be one of: {ACT2FN.keys()}."
+
+        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)
+
+    def forward(
+        self,
+        qkv: torch.FloatTensor,
+        causal: bool = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        batch_size, seqlen = qkv.shape[0], qkv.shape[1]
+        q, k, v = qkv.unbind(dim=2)
+
+        causal = self.causal if causal is None else causal
+        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+        if attention_mask is not None:
+            padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device)
+            padding_mask.masked_fill_(attention_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, dtype=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)
+
+    def forward(
+        self,
+        q: torch.FloatTensor,
+        kv: torch.FloatTensor,
+        causal: bool = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        batch_size, seqlen_q = q.shape[0], q.shape[1]
+        seqlen_k = kv.shape[1]
+        assert kv.shape[0] == batch_size and kv.shape[4] == q.shape[3]
+
+        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)
+
+        causal = self.causal if causal is None else causal
+        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+        if attention_mask is not None:
+            padding_mask = torch.full((batch_size, seqlen_k), -10000.0, dtype=scores.dtype, device=scores.device)
+            padding_mask.masked_fill_(attention_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, dtype=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]:
+    assert all(
+        hasattr(config, attr) for attr in ["n_embd", "n_head"]
+    ), "`config` must have `n_embd` and `n_head` attributes."
+
+    if head_dim is None:
+        assert (
+            config.n_embd % config.n_head == 0
+        ), f"Hidden size ({config.n_embd}) must be divisible by the number of heads ({config.n_head})."
+
+    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
+    assert n_head % n_head_kv == 0, "`n_head` must be divisible by `n_head_kv`."
+
+    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:
+        kv_cache = torch.empty(
+            inference_params.max_batch_size,
+            inference_params.max_seqlen,
+            2,
+            num_heads,
+            head_dim,
+            dtype=kv.dtype,
+            device=kv.device,
+        )
+        inference_params.key_value_memory_dict[layer_idx] = kv_cache
+    else:
+        kv_cache = inference_params.key_value_memory_dict[layer_idx]
+
+    batch_start = inference_params.batch_size_offset
+    batch_end = batch_start + kv.shape[0]
+    assert batch_end <= kv_cache.shape[0]
+
+    sequence_start = inference_params.seqlen_offset
+    sequence_end = sequence_start + kv.shape[1]
+    assert sequence_end <= kv_cache.shape[1]
+
+    assert kv_cache is not None
+    kv_cache[batch_start:batch_end, sequence_start:sequence_end, ...] = kv
+    kv = kv_cache[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_emb_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_emb_dim = rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0)
+        if self.rotary_emb_dim > 0:
+            rotary_kwargs = {"device": device}
+            if rotary_emb_scale_base is not None and rotary_emb_scale_base > 0.0:
+                rotary_kwargs["scale_base"] = rotary_emb_scale_base
+            
+            rotary_cls = FlashRotaryEmbedding if config.flash_rotary else RotaryEmbedding
+            if rotary_cls is None:
+                rotary_cls = RotaryEmbedding
+            self.rotary_emb = rotary_cls(self.rotary_emb_dim, **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
+        self.inner_attn = SelfAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=config.attn_pdrop)
+        self.inner_cross_attn = CrossAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=config.attn_pdrop)
+
+        self.layer_idx = layer_idx
+        self.return_residual = return_residual
+        self.checkpointing = checkpointing
+
+    def _forward_self_attn(
+        self, x: torch.FloatTensor, attention_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_emb_dim > 0:
+            qkv = self.rotary_emb(qkv)
+
+        if self.checkpointing:
+            return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, attention_mask=attention_mask)
+
+        return self.inner_attn(qkv, attention_mask=attention_mask)
+
+    def _forward_cross_attn(
+        self,
+        x: torch.FloatTensor,
+        past_key_values: Optional[InferenceParams],
+        attention_mask: Optional[torch.BoolTensor],
+    ) -> torch.FloatTensor:
+        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_emb_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.checkpointing:
+            return torch.utils.checkpoint.checkpoint(
+                self.inner_cross_attn, q, kv, attention_mask=attention_mask, causal=causal
+            )
+
+        return self.inner_cross_attn(q, kv, attention_mask=attention_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 and torch.any(~attention_mask.bool()):
+            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 MixFormerSequentialPreTrainedModel(PreTrainedModel):
+    """MixFormer (sequential for DeepSpeed) pre-trained model."""
+
+    config_class = MixFormerSequentialConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+
+    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 = len(input_ids[0]) - 1
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+        }
+    
+    def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False) -> None:
+        if isinstance(module, MixFormerSequentialPreTrainedModel):
+            module.gradient_checkpointing = value
+
+
+class MixFormerSequentialForCausalLM(MixFormerSequentialPreTrainedModel):
+    """MixFormer (sequential for DeepSpeed) for Causal Language Modeling."""
+
+    _keys_to_ignore_on_load_missing = [""]
+    _keys_to_ignore_on_load_unexpected = [r"layers\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
+    _no_split_modules = ["ParallelBlock"]
+
+    def __init__(self, config: MixFormerSequentialConfig) -> None:
+        super().__init__(config)
+
+        modules = [Embedding(config)]
+        modules += [ParallelBlock(config, block_idx=i) for i in range(config.n_layer)]
+        modules.append(CausalLMHead(config))
+
+        self.layers = nn.Sequential(*modules)
+        self.loss = CausalLMLoss()
+
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.layers[0].wte
+
+    def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
+        self.layers[0].wte = new_embeddings
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.layers[-1].linear
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.layers[-1].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,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden_layer = self.layers[0](input_ids)
+        for module in self.layers[1:-1]:
+            hidden_layer = module(hidden_layer, past_key_values=past_key_values, attention_mask=attention_mask)
+        lm_logits = self.layers[-1](hidden_layer)
+
+        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)

phi_predict.py

@@ -0,0 +1,16 @@
+import pandas as pd
+
+def predict(data, task, model, tokenizer, config, **kwargs):
+    if isinstance(data, pd.DataFrame):
+        data = data[data.columns[0]].tolist()
+        is_df = True
+    results = []
+    addn_args = kwargs.get("addn_args", {})
+    for d in data:
+        inputs = tokenizer(d, return_tensors="pt", return_attention_mask=False)
+        outputs = model.generate(**inputs, **addn_args, max_length=50)
+        text = tokenizer.batch_decode(outputs)[0]
+        results.append(text)
+    if is_df:
+        return pd.DataFrame(results,columns =['output'])
+    return {"output": results}

python_env.yaml

@@ -0,0 +1,7 @@
+python: 3.10.11
+build_dependencies:
+- pip==23.1.2
+- setuptools==67.8.0
+- wheel==0.38.4
+dependencies:
+- -r requirements.txt

pytorch_model-00001-of-00002.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7b27189d6c06238a304388792e15a72810274bc8873e50bd5427b99973d561e
+size 9964971527

pytorch_model-00002-of-00002 (1).bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c93d72cf5143e888da71bcdbd7f105ae91138583e52407ccbb133f41eb34139
+size 1153871286

pytorch_model.bin.index.json

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

requirements.txt

@@ -0,0 +1,21 @@
+mlflow==2.6.0
+cloudpickle==2.2.1
+jsonpickle==3.0.1
+mlflow-skinny==2.6.0
+azureml-core==1.51.0.post1
+azureml-mlflow==1.51.0
+azureml-metrics[all]==0.0.32
+scikit-learn==1.2.2
+cryptography==41.0.1
+python-dateutil==2.8.2
+datasets==2.14.6
+soundfile==0.12.1
+librosa==0.10.1
+diffusers==0.21.4
+sentencepiece==0.1.99
+transformers==4.34.0
+torch==2.1.0
+accelerate==0.23.0
+Pillow==9.4.0
+einops
+azureml-evaluate-mlflow==0.0.32

special_tokens_map.json

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