Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,211 @@
+---
+language:
+- multilingual
+- ar
+- zh
+- cs
+- da
+- nl
+- en
+- fi
+- fr
+- de
+- he
+- hu
+- it
+- ja
+- ko
+- 'no'
+- pl
+- pt
+- ru
+- es
+- sv
+- th
+- tr
+- uk
+library_name: mlx
+license: mit
+license_link: https://huggingface.co/microsoft/Phi-4-mini-instruct/resolve/main/LICENSE
+pipeline_tag: text-generation
+tags:
+- nlp
+- code
+- mlx
+- phi-4
+- fine-tuned
+- dolly-15k
+- instruction-following
+- lora
+widget:
+- messages:
+  - role: user
+    content: Can you provide ways to eat combinations of bananas and dragonfruits?
+- messages:
+  - role: user
+    content: Explain machine learning in simple terms.
+- messages:
+  - role: user
+    content: What are the benefits of renewable energy?
+base_model: microsoft/Phi-4-mini-instruct
+datasets:
+- databricks/databricks-dolly-15k
+model-index:
+- name: Phi-4-mini-instruct-dolly-15k-mlx
+  results: []
+---
+
+# Phi-4 Mini Instruct - Fine-tuned on Dolly 15K (MLX)
+
+This model is a fine-tuned version of [microsoft/Phi-4-mini-instruct](https://huggingface.co/microsoft/Phi-4-mini-instruct) using Apple's MLX framework, trained on the [Databricks Dolly 15K](https://huggingface.co/datasets/databricks/databricks-dolly-15k) instruction dataset.
+
+## Model Description
+
+This model enhances Phi-4 Mini's instruction-following capabilities by fine-tuning it on 15,000 high-quality human-generated instruction-following examples from the Dolly dataset. The fine-tuning was performed using LoRA (Low-Rank Adaptation) on Apple Silicon hardware, making it optimized for deployment on Mac devices.
+
+### Key Features
+- **Base Model**: Phi-4-mini-instruct (3.8B parameters)
+- **Fine-tuning Method**: LoRA with MLX
+- **Training Dataset**: Dolly 15K (12,000 training, 3,000 validation examples)
+- **Optimized for**: Apple Silicon (M1, M2, M3, M4)
+- **License**: MIT
+
+## Intended Uses & Limitations
+
+### Intended Uses
+- General instruction following and question answering
+- Educational applications
+- Content generation
+- Code assistance
+- Creative writing tasks
+
+### Limitations
+- The model inherits limitations from the base Phi-4 model
+- Performance is optimized for Apple Silicon; may have different characteristics on other hardware
+- Should not be used for critical decision-making without human oversight
+- May exhibit biases present in the training data
+
+## Training Details
+
+### Training Configuration
+- **LoRA Rank**: 64
+- **LoRA Alpha**: 16
+- **LoRA Dropout**: 0.1
+- **Target Layers**: 16
+- **Learning Rate**: 1e-4
+- **Batch Size**: 2
+- **Training Iterations**: 1,000
+- **Max Sequence Length**: 2,048
+
+### Training Hardware
+- Trained on Apple Silicon (M-series chip)
+- Peak memory usage: ~8GB
+- Training time: 25-60 minutes depending on specific chip
+
+### Training Data
+The model was fine-tuned on the Databricks Dolly 15K dataset, which consists of:
+- **Total examples**: 15,000 instruction-following records
+- **Training split**: 12,000 examples (80%)
+- **Validation split**: 3,000 examples (20%)
+- **Categories**: Brainstorming, classification, closed QA, generation, information extraction, open QA, and summarization
+
+## Performance
+
+### Inference Speed (on Apple Silicon)
+- **Prompt processing**: ~240-315 tokens/second
+- **Generation**: ~66-73 tokens/second
+- **Peak memory usage**: 7.7-8.0 GB
+
+### Example Outputs
+
+**Machine Learning Explanation:**
+> "Machine learning is when our computers learn from experience by looking at examples of things that have already been done, and then apply that knowledge to new situations. For example, if we want to teach a computer to recognize cats, we can show it many pictures of cats and tell it 'this is a cat', and it will eventually be able to recognize when it sees a cat in a picture."
+
+**Renewable Energy Benefits:**
+> "Renewable energy is sustainable and minimizes carbon emissions. It is readily available in most locations and can be exploited in the smallest of communities. It can be used for various purposes, such as generating heat and electricity and powering electric vehicles. Renewable energy is a clean fuel source. It helps to reduce air pollution and the risk of climate change."
+
+## How to Use
+
+### With MLX
+
+```python
+from mlx_lm import load, generate
+
+# Load the model
+model, tokenizer = load("your-username/phi-4-mini-instruct-dolly-15k-mlx")
+
+# Generate text
+prompt = "<|user|>\nExplain quantum computing in simple terms<|end|>\n<|assistant|>\n"
+response = generate(model, tokenizer, prompt=prompt, max_tokens=300)
+print(response)
+```
+
+### Chat Format
+
+The model uses Phi-4's chat template with special tokens:
+- `<|user|>` - User message start
+- `<|assistant|>` - Assistant message start  
+- `<|end|>` - Message end
+
+Example:
+```
+<|user|>
+What is machine learning?
+<|end|>
+<|assistant|>
+Machine learning is...
+<|end|>
+```
+
+## Training Procedure
+
+### Data Preprocessing
+1. Dolly 15K dataset was downloaded and formatted for Phi-4's chat template
+2. Instructions and responses were wrapped with appropriate special tokens
+3. Data was split 80/20 for training/validation
+4. Saved as JSONL files for MLX compatibility
+
+### LoRA Fine-tuning
+1. Applied LoRA adapters to 16 transformer layers
+2. Trained for 1,000 iterations with validation every 200 steps
+3. Saved checkpoints every 500 iterations
+4. Final adapter weights were fused with the base model
+
+### Post-processing
+The LoRA adapters were merged with the base model weights to create this standalone model, eliminating the need for adapter loading during inference.
+
+## Evaluation
+
+Validation loss decreased from ~3.0 to ~1.5-2.0 during training, indicating successful learning of the instruction-following patterns in the Dolly dataset.
+
+## Environmental Impact
+
+This model was trained on energy-efficient Apple Silicon hardware, resulting in lower power consumption compared to traditional GPU training. Estimated carbon footprint is minimal due to:
+- Short training time (< 1 hour)
+- Efficient LoRA method (only 0.082% of parameters trained)
+- Apple Silicon's power efficiency
+
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@misc{phi4-mini-dolly-mlx,
+  author = {Your Name},
+  title = {Phi-4 Mini Instruct Fine-tuned on Dolly 15K for MLX},
+  year = {2025},
+  publisher = {HuggingFace},
+  url = {https://huggingface.co/your-username/phi-4-mini-instruct-dolly-15k-mlx}
+}
+```
+
+## Acknowledgments
+
+- Microsoft for the Phi-4 base model
+- Databricks for the Dolly 15K dataset
+- Apple MLX team for the framework
+- The open-source community
+
+## Model Card Contact
+
+For questions or issues with this model, please open an issue on the [GitHub repository](https://github.com/yourusername/phi4-mlx-training) or contact via Hugging Face.

added_tokens.json

@@ -0,0 +1,12 @@
+{
+  "<|/tool_call|>": 200026,
+  "<|/tool|>": 200024,
+  "<|assistant|>": 200019,
+  "<|end|>": 200020,
+  "<|system|>": 200022,
+  "<|tag|>": 200028,
+  "<|tool_call|>": 200025,
+  "<|tool_response|>": 200027,
+  "<|tool|>": 200023,
+  "<|user|>": 200021
+}

chat_template.jinja

@@ -0,0 +1 @@
+{% for message in messages %}{% if message['role'] == 'system' and 'tools' in message and message['tools'] is not none %}{{ '<|' + message['role'] + '|>' + message['content'] + '<|tool|>' + message['tools'] + '<|/tool|>' + '<|end|>' }}{% else %}{{ '<|' + message['role'] + '|>' + message['content'] + '<|end|>' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|assistant|>' }}{% else %}{{ eos_token }}{% endif %}

config.json

@@ -0,0 +1,143 @@
+{
+    "architectures": [
+        "Phi3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "auto_map": {
+        "AutoConfig": "configuration_phi3.Phi3Config",
+        "AutoModelForCausalLM": "modeling_phi3.Phi3ForCausalLM",
+        "AutoTokenizer": "Xenova/gpt-4o"
+    },
+    "bos_token_id": 199999,
+    "embd_pdrop": 0.0,
+    "eos_token_id": 199999,
+    "full_attn_mod": 1,
+    "hidden_act": "silu",
+    "hidden_size": 3072,
+    "initializer_range": 0.02,
+    "intermediate_size": 8192,
+    "interpolate_factor": 1,
+    "lm_head_bias": false,
+    "max_position_embeddings": 131072,
+    "mlp_bias": false,
+    "model_type": "phi3",
+    "num_attention_heads": 24,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 8,
+    "original_max_position_embeddings": 4096,
+    "pad_token_id": 199999,
+    "partial_rotary_factor": 0.75,
+    "resid_pdrop": 0.0,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": {
+        "long_factor": [
+            1,
+            1.118320672,
+            1.250641126,
+            1.398617824,
+            1.564103225,
+            1.74916897,
+            1.956131817,
+            2.187582649,
+            2.446418898,
+            2.735880826,
+            3.059592084,
+            3.421605075,
+            3.826451687,
+            4.279200023,
+            4.785517845,
+            5.351743533,
+            5.984965424,
+            6.693110555,
+            7.485043894,
+            8.370679318,
+            9.36110372,
+            10.4687158,
+            11.70738129,
+            13.09260651,
+            14.64173252,
+            16.37415215,
+            18.31155283,
+            20.47818807,
+            22.90118105,
+            25.61086418,
+            28.64115884,
+            32.03,
+            32.1,
+            32.13,
+            32.23,
+            32.6,
+            32.61,
+            32.64,
+            32.66,
+            32.7,
+            32.71,
+            32.93,
+            32.97,
+            33.28,
+            33.49,
+            33.5,
+            44.16,
+            47.77
+        ],
+        "short_factor": [
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0,
+            1.0
+        ],
+        "type": "longrope"
+    },
+    "rope_theta": 10000.0,
+    "sliding_window": 262144,
+    "tie_word_embeddings": true,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.45.0",
+    "use_cache": true,
+    "vocab_size": 200064
+}

configuration_phi3.py

@@ -0,0 +1,226 @@
+# coding=utf-8
+# Copyright 2024 Microsoft and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Phi-3 model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class Phi3Config(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Phi3Model`]. It is used to instantiate a Phi-3
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the
+    [microsoft/Phi-3-mini-4k-instruct](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 32064):
+            Vocabulary size of the Phi-3 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`Phi3Model`].
+        hidden_size (`int`, *optional*, defaults to 3072):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 8192):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        resid_pdrop (`float`, *optional*, defaults to 0.0):
+            Dropout probability for mlp outputs.
+        embd_pdrop (`int`, *optional*, defaults to 0.0):
+            The dropout ratio for the embeddings.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio after computing the attention scores.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 4096):
+            The maximum sequence length that this model might ever be used with.
+        original_max_position_embeddings (`int`, *optional*, defaults to 4096):
+            The maximum sequence length that this model was trained with. This is used to determine the size of the
+            original RoPE embeddings when using long scaling.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon value used for the RMSNorm.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`. Whether to tie weight embeddings or not.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`dict`, *optional*):
+            The scaling strategy for the RoPE embeddings. If `None`, no scaling is applied. If a dictionary, it must
+            contain the following keys: `type`, `short_factor` and `long_factor`. The `type` must be `longrope` and
+            the `short_factor` and `long_factor` must be lists of numbers with the same length as the hidden size
+            divided by the number of attention heads divided by 2.
+        partial_rotary_factor (`float`, *optional*, defaults to 1.0):
+            Percentage of the query and keys which will have rotary embedding. Must be between 0.0 and 1.0.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            The id of the "beginning-of-sequence" token.
+        eos_token_id (`int`, *optional*, defaults to 32000):
+            The id of the "end-of-sequence" token.
+        pad_token_id (`int`, *optional*, defaults to 32000):
+            The id of the padding token.
+        sliding_window (`int`, *optional*):
+            Sliding window attention window size. If `None`, no sliding window is applied.
+
+    Example:
+
+    ```python
+    >>> from transformers import Phi3Model, Phi3Config
+
+    >>> # Initializing a Phi-3 style configuration
+    >>> configuration = Phi3Config.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
+
+    >>> # Initializing a model from the configuration
+    >>> model = Phi3Model(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "phi3"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=32064,
+        hidden_size=3072,
+        intermediate_size=8192,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        resid_pdrop=0.0,
+        embd_pdrop=0.0,
+        attention_dropout=0.0,
+        hidden_act="silu",
+        max_position_embeddings=4096,
+        original_max_position_embeddings=4096,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        partial_rotary_factor=1.0,
+        bos_token_id=1,
+        eos_token_id=32000,
+        pad_token_id=32000,
+        sliding_window=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.resid_pdrop = resid_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.attention_dropout = attention_dropout
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.original_max_position_embeddings = original_max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.partial_rotary_factor = partial_rotary_factor
+        self._rope_scaling_adjustment()
+        self._rope_scaling_validation()
+        self.sliding_window = sliding_window
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    def _rope_scaling_adjustment(self):
+        """
+        Adjust the `type` of the `rope_scaling` configuration for backward compatibility.
+        """
+        if self.rope_scaling is None:
+            return
+
+        rope_scaling_type = self.rope_scaling.get("type", None)
+
+        # For backward compatibility if previous version used "su" or "yarn"
+        if rope_scaling_type is not None and rope_scaling_type in ["su", "yarn"]:
+            self.rope_scaling["type"] = "longrope"
+
+    def _rope_scaling_validation(self):
+        """
+        Validate the `rope_scaling` configuration.
+        """
+        if self.rope_scaling is None:
+            return
+
+        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 3:
+            raise ValueError(
+                "`rope_scaling` must be a dictionary with three fields, `type`, `short_factor` and `long_factor`, "
+                f"got {self.rope_scaling}"
+            )
+        rope_scaling_type = self.rope_scaling.get("type", None)
+        rope_scaling_short_factor = self.rope_scaling.get("short_factor", None)
+        rope_scaling_long_factor = self.rope_scaling.get("long_factor", None)
+        if rope_scaling_type is None or rope_scaling_type not in ["longrope"]:
+            raise ValueError(f"`rope_scaling`'s type field must be one of ['longrope'], got {rope_scaling_type}")
+        if not (
+            isinstance(rope_scaling_short_factor, list)
+            and all(isinstance(x, (int, float)) for x in rope_scaling_short_factor)
+        ):
+            raise ValueError(
+                f"`rope_scaling`'s short_factor field must be a list of numbers, got {rope_scaling_short_factor}"
+            )
+        rotary_ndims = int(self.hidden_size // self.num_attention_heads * self.partial_rotary_factor)
+        if not len(rope_scaling_short_factor) == rotary_ndims // 2:
+            raise ValueError(
+                f"`rope_scaling`'s short_factor field must have length {rotary_ndims // 2}, got {len(rope_scaling_short_factor)}"
+            )
+        if not (
+            isinstance(rope_scaling_long_factor, list)
+            and all(isinstance(x, (int, float)) for x in rope_scaling_long_factor)
+        ):
+            raise ValueError(
+                f"`rope_scaling`'s long_factor field must be a list of numbers, got {rope_scaling_long_factor}"
+            )
+        if not len(rope_scaling_long_factor) == rotary_ndims // 2:
+            raise ValueError(
+                f"`rope_scaling`'s long_factor field must have length {rotary_ndims // 2}, got {len(rope_scaling_long_factor)}"
+            )

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 199999,
+  "eos_token_id": [
+    200020,
+    199999
+  ],
+  "pad_token_id": 199999,
+  "transformers_version": "4.45.0"
+}

model-00001-of-00002.safetensors

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

model-00002-of-00002.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,201 @@
+{
+    "metadata": {
+        "total_size": 7672043520
+    },
+    "weight_map": {
+        "model.embed_tokens.weight": "model-00001-of-00002.safetensors",
+        "model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.0.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.0.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.1.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.1.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.1.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.10.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.10.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.10.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.10.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.10.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.10.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.11.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.11.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.11.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.11.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.11.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.11.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.12.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.12.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.12.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.12.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.12.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.12.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.13.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.13.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.13.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.13.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.13.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.13.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.14.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.14.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.14.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.14.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.14.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.14.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.15.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.15.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.15.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.15.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.15.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.15.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.16.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.16.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.16.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.16.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.16.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.16.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.17.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.17.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.17.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.17.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.17.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.17.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.18.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.18.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.18.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.18.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.18.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.18.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.19.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.19.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.19.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.19.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.19.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.19.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.2.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.2.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.2.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.20.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.20.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.20.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.20.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.20.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.20.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.21.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.21.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.21.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.21.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.21.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.21.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.22.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.22.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.22.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.22.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.22.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.22.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.23.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.23.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.23.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.23.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.23.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.23.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.24.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.24.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.24.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.24.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.24.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.24.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.25.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.25.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.25.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.25.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.25.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.25.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.26.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.26.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.26.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.26.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.26.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.26.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.27.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.27.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.27.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.27.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.27.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.27.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.28.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.28.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.28.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.28.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.28.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.28.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.29.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.29.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.29.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.29.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.29.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.29.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.3.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.3.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.3.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.3.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.30.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.30.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.30.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.30.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.30.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.30.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.31.input_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.31.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.31.mlp.gate_up_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.31.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+        "model.layers.31.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.31.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
+        "model.layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.4.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.4.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.4.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.5.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.5.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.5.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.5.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.6.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.6.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.6.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.6.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.7.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.7.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.7.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.7.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.8.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.8.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.8.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.8.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.9.input_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.9.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.9.mlp.gate_up_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.9.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+        "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+        "model.layers.9.self_attn.qkv_proj.weight": "model-00001-of-00002.safetensors",
+        "model.norm.weight": "model-00002-of-00002.safetensors"
+    }
+}

modeling_phi3.py

@@ -0,0 +1,1180 @@
+# coding=utf-8
+# Copyright 2024 Microsoft and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""PyTorch Phi-3 model."""
+
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache
+from transformers.generation import GenerationMixin
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
+)
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    LossKwargs,
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.utils.deprecation import deprecate_kwarg
+from .configuration_phi3 import Phi3Config
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "microsoft/Phi-3-mini-4k-instruct"
+_CONFIG_FOR_DOC = "Phi3Config"
+
+
+class Phi3MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.config = config
+        self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.activation_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+        up_states = self.gate_up_proj(hidden_states)
+
+        gate, up_states = up_states.chunk(2, dim=-1)
+        up_states = up_states * self.activation_fn(gate)
+
+        return self.down_proj(up_states)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+
+    rotary_dim = cos.shape[-1]
+    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+    q_embed = torch.cat([(q_rot * cos) + (rotate_half(q_rot) * sin), q_pass], dim=-1)
+    k_embed = torch.cat([(k_rot * cos) + (rotate_half(k_rot) * sin), k_pass], dim=-1)
+    return q_embed, k_embed
+
+
+class Phi3Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: Phi3Config, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        op_size = config.num_attention_heads * self.head_dim + 2 * (config.num_key_value_heads * self.head_dim)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+        self.qkv_proj = nn.Linear(config.hidden_size, op_size, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        qkv = self.qkv_proj(hidden_states)
+        query_pos = self.config.num_attention_heads * self.head_dim
+        query_states = qkv[..., :query_pos]
+        key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]
+        value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]
+
+        query_states = query_states.view(hidden_shape).transpose(1, 2)
+        key_states = key_states.view(hidden_shape).transpose(1, 2)
+        value_states = value_states.view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
+                logger.warning_once(
+                    "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
+                    'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+                )
+            else:
+                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=getattr(self.config, "sliding_window", None),
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class Phi3RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        Phi3RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Phi3DecoderLayer(nn.Module):
+    def __init__(self, config: Phi3Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = Phi3Attention(config=config, layer_idx=layer_idx)
+        self.mlp = Phi3MLP(config)
+        self.input_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.config = config
+        self.resid_attn_dropout = nn.Dropout(config.resid_pdrop)
+        self.resid_mlp_dropout = nn.Dropout(config.resid_pdrop)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`):
+                input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
+                Indices of positions of each input sequence tokens in the position embeddings. Selected in the range
+                `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)
+            past_key_value (`Cache`, *optional*): cached past key and value projection states
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+                Indices depicting the position of the input sequence tokens in the sequence
+            kwargs (`dict`, *optional*):
+                Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+                into the model
+        """
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + self.resid_attn_dropout(hidden_states)  # main diff with Llama
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + self.resid_mlp_dropout(hidden_states)  # main diff with Llama
+
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        return outputs
+
+
+class Phi3RotaryEmbedding(nn.Module):
+    def __init__(self, config: Phi3Config, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    def _dynamic_frequency_update(self, position_ids, device):
+        """
+        dynamic RoPE layers should recompute `inv_freq` in the following situations:
+        1 - growing beyond the cached sequence length (allow scaling)
+        2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
+        """
+        seq_len = torch.max(position_ids) + 1
+        if seq_len > self.max_seq_len_cached:  # growth
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len)
+            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: may break with compilation
+            self.max_seq_len_cached = seq_len
+
+        if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset
+            # This .to() is needed if the model has been moved to a device after being initialized (because
+            # the buffer is automatically moved, but not the original copy)
+            self.original_inv_freq = self.original_inv_freq.to(device)
+            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
+            self.max_seq_len_cached = self.original_max_seq_len
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        if "dynamic" in self.rope_type:
+            self._dynamic_frequency_update(position_ids, device=x.device)
+        elif self.rope_type == "longrope":
+            self._longrope_frequency_update(position_ids, device=x.device)
+
+        # Core RoPE block
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+        # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
+        device_type = x.device.type
+        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+
+        # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
+        cos = cos * self.attention_scaling
+        sin = sin * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+    def _longrope_frequency_update(self, position_ids, device):
+        """Longrope uses long factor if sequence is larger than original pretraining length, short otherwise."""
+        seq_len = torch.max(position_ids) + 1
+        if hasattr(self.config, "original_max_position_embeddings"):
+            original_max_position_embeddings = self.config.original_max_position_embeddings
+        else:
+            original_max_position_embeddings = self.config.max_position_embeddings
+        if seq_len > original_max_position_embeddings:
+            if not hasattr(self, "long_inv_freq"):
+                self.long_inv_freq, _ = self.rope_init_fn(
+                    self.config, device, seq_len=original_max_position_embeddings + 1
+                )
+            self.register_buffer("inv_freq", self.long_inv_freq, persistent=False)
+        else:
+            # This .to() is needed if the model has been moved to a device after being initialized (because
+            # the buffer is automatically moved, but not the original copy)
+            self.original_inv_freq = self.original_inv_freq.to(device)
+            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
+
+
+PHI3_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`Phi3Config`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare Phi3 Model outputting raw hidden-states without any specific head on top.",
+    PHI3_START_DOCSTRING,
+)
+class Phi3PreTrainedModel(PreTrainedModel):
+    config_class = Phi3Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Phi3DecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+    _version = "0.0.5"
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+PHI3_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+            returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+            Two formats are allowed:
+            - a [`~cache_utils.Cache`] instance, see our
+            [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache);
+            - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+            shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+            cache format.
+
+            The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+            legacy cache format will be returned.
+
+            If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+            of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+            Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
+            this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
+            the complete sequence length.
+"""
+
+
+@add_start_docstrings(
+    "The bare Phi3 Model outputting raw hidden-states without any specific head on top.",
+    PHI3_START_DOCSTRING,
+)
+class Phi3Model(Phi3PreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`]
+
+    Args:
+        config: Phi3Config
+    """
+
+    def __init__(self, config: Phi3Config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Phi3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Phi3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(PHI3_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+        )
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                    position_embeddings,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    position_embeddings=position_embeddings,
+                    **flash_attn_kwargs,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        output = BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+        return output if return_dict else output.to_tuple()
+
+    def _update_causal_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_tensor: torch.Tensor,
+        cache_position: torch.Tensor,
+        past_key_values: Cache,
+        output_attentions: bool,
+    ):
+        if self.config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and past_key_values is not None:
+                is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0]
+                if is_padding_right:
+                    raise ValueError(
+                        "You are attempting to perform batched generation with padding_side='right'"
+                        " this may lead to unexpected behaviour for Flash Attention version of Phi3. Make sure to "
+                        " call `tokenizer.padding_side  = 'left'` before tokenizing the input. "
+                    )
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+
+        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
+        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
+        # to infer the attention mask.
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        using_static_cache = isinstance(past_key_values, StaticCache)
+        using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)
+
+        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+        if (
+            self.config._attn_implementation == "sdpa"
+            and not (using_static_cache or using_sliding_window_cache)
+            and not output_attentions
+        ):
+            if AttentionMaskConverter._ignore_causal_mask_sdpa(
+                attention_mask,
+                inputs_embeds=input_tensor,
+                past_key_values_length=past_seen_tokens,
+                sliding_window=self.config.sliding_window,
+                is_training=self.training,
+            ):
+                return None
+
+        dtype, device = input_tensor.dtype, input_tensor.device
+        min_dtype = torch.finfo(dtype).min
+        sequence_length = input_tensor.shape[1]
+        # SlidingWindowCache or StaticCache
+        if using_sliding_window_cache or using_static_cache:
+            target_length = past_key_values.get_max_cache_shape()
+        # DynamicCache or no cache
+        else:
+            target_length = (
+                attention_mask.shape[-1]
+                if isinstance(attention_mask, torch.Tensor)
+                else past_seen_tokens + sequence_length + 1
+            )
+
+        # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+        causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
+            attention_mask,
+            sequence_length=sequence_length,
+            target_length=target_length,
+            dtype=dtype,
+            device=device,
+            cache_position=cache_position,
+            batch_size=input_tensor.shape[0],
+            config=self.config,
+            past_key_values=past_key_values,
+        )
+
+        if (
+            self.config._attn_implementation == "sdpa"
+            and attention_mask is not None
+            and attention_mask.device.type in ["cuda", "xpu"]
+            and not output_attentions
+        ):
+            # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
+            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
+            # Details: https://github.com/pytorch/pytorch/issues/110213
+            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
+
+        return causal_mask
+
+    @staticmethod
+    def _prepare_4d_causal_attention_mask_with_cache_position(
+        attention_mask: torch.Tensor,
+        sequence_length: int,
+        target_length: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_position: torch.Tensor,
+        batch_size: int,
+        config: Phi3Config,
+        past_key_values: Cache,
+    ):
+        """
+        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+        Args:
+            attention_mask (`torch.Tensor`):
+                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
+            sequence_length (`int`):
+                The sequence length being processed.
+            target_length (`int`):
+                The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
+            dtype (`torch.dtype`):
+                The dtype to use for the 4D attention mask.
+            device (`torch.device`):
+                The device to plcae the 4D attention mask on.
+            cache_position (`torch.Tensor`):
+                Indices depicting the position of the input sequence tokens in the sequence.
+            batch_size (`torch.Tensor`):
+                Batch size.
+            config (`Phi3Config`):
+                The model's configuration class
+            past_key_values (`Cache`):
+                The cache class that is being used currently to generate
+        """
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+            causal_mask = attention_mask
+        else:
+            min_dtype = torch.finfo(dtype).min
+            causal_mask = torch.full(
+                (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
+            )
+            diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
+            if config.sliding_window is not None:
+                # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also
+                # the check is needed to verify is current checkpoint was trained with sliding window or not
+                if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length:
+                    sliding_attend_mask = torch.arange(target_length, device=device) <= (
+                        cache_position.reshape(-1, 1) - config.sliding_window
+                    )
+                    diagonal_attend_mask.bitwise_or_(sliding_attend_mask)
+            causal_mask *= diagonal_attend_mask
+            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+            if attention_mask is not None:
+                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
+                if attention_mask.shape[-1] > target_length:
+                    attention_mask = attention_mask[:, :target_length]
+                mask_length = attention_mask.shape[-1]
+                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+                    causal_mask.device
+                )
+                padding_mask = padding_mask == 0
+                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                    padding_mask, min_dtype
+                )
+        return causal_mask
+
+
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
+
+
+class Phi3ForCausalLM(Phi3PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Phi3Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+    @add_start_docstrings_to_model_forward(PHI3_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[KwargsForCausalLM],
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+            logits_to_keep (`int` or `torch.Tensor`, *optional*):
+                If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+                `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+                token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+                If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+                This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, Phi3ForCausalLM
+
+        >>> model = Phi3ForCausalLM.from_pretrained("meta-phi3/Phi3-2-7b-hf")
+        >>> tokenizer = AutoTokenizer.from_pretrained("meta-phi3/Phi3-2-7b-hf")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        logits_to_keep=None,
+        **kwargs,
+    ):
+        # Overwritten -- this model may need to switch between short and long rope, invalidating the cache in the
+        # process
+
+        # When the first time input length reached long and short factor switching point, enforce re-compute cache
+        # It will cause downside of slower at this single token position, however, better than current failure.
+        if (
+            past_key_values
+            and self.config.rope_scaling
+            and input_ids.shape[1] >= self.config.original_max_position_embeddings + 1
+        ):
+            past_length = cache_position[0]
+            if past_length <= self.config.original_max_position_embeddings:
+                past_key_values = None
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            position_ids=position_ids,
+            use_cache=use_cache,
+            logits_to_keep=logits_to_keep,
+            **kwargs,
+        )
+        return model_inputs
+
+
+@add_start_docstrings(
+    """
+    The Phi3 Model transformer with a sequence classification head on top (linear layer).
+
+    [`Phi3ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+    (e.g. GPT-2) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+    each row of the batch).
+    """,
+    PHI3_START_DOCSTRING,
+)
+class Phi3ForSequenceClassification(Phi3PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = Phi3Model(config)
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(PHI3_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        transformer_outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+        if self.config.pad_token_id is None:
+            last_non_pad_token = -1
+        elif input_ids is not None:
+            # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+            non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+            token_indices = torch.arange(input_ids.shape[-1], device=logits.device)
+            last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+        else:
+            last_non_pad_token = -1
+            logger.warning_once(
+                f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+                "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+            )
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)
+
+        if not return_dict:
+            output = (pooled_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    The Phi3 Model transformer with a token classification head on top (a linear layer on top of the hidden-states
+    output) e.g. for Named-Entity-Recognition (NER) tasks.
+    """,
+    PHI3_START_DOCSTRING,
+)
+class Phi3ForTokenClassification(Phi3PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = Phi3Model(config)
+        if getattr(config, "classifier_dropout", None) is not None:
+            classifier_dropout = config.classifier_dropout
+        elif getattr(config, "hidden_dropout", None) is not None:
+            classifier_dropout = config.hidden_dropout
+        else:
+            classifier_dropout = 0.1
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.score = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(PHI3_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=TokenClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, TokenClassifierOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        logits = self.score(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.config)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

sample_finetune.py

@@ -0,0 +1,214 @@
+import sys
+import logging
+
+import datasets
+from datasets import load_dataset
+from peft import LoraConfig
+import torch
+import transformers
+from trl import SFTTrainer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig
+
+"""
+A simple example on using SFTTrainer and Accelerate to finetune Phi-4-Mini-Instruct model. For
+a more advanced example, please follow HF alignment-handbook/scripts/run_sft.py.
+This example has utilized DeepSpeed ZeRO3 offload to reduce the memory usage. The
+script can be run on V100 or later generation GPUs. Here are some suggestions on 
+futher reducing memory consumption:
+    - reduce batch size
+    - decrease lora dimension
+    - restrict lora target modules
+Please follow these steps to run the script:
+1. Install dependencies: 
+    conda install -c conda-forge accelerate=1.3.0
+    pip3 install -i https://pypi.org/simple/ bitsandbytes
+    pip3 install peft==0.14.0
+    pip3 install transformers==4.48.1
+    pip3 install trl datasets
+    pip3 install deepspeed
+2. Setup accelerate and deepspeed config based on the machine used:
+    accelerate config
+Here is a sample config for deepspeed zero3:
+    compute_environment: LOCAL_MACHINE
+    debug: false
+    deepspeed_config:
+      gradient_accumulation_steps: 1
+      offload_optimizer_device: none
+      offload_param_device: none
+      zero3_init_flag: true
+      zero3_save_16bit_model: true
+      zero_stage: 3
+    distributed_type: DEEPSPEED
+    downcast_bf16: 'no'
+    enable_cpu_affinity: false
+    machine_rank: 0
+    main_training_function: main
+    mixed_precision: bf16
+    num_machines: 1
+    num_processes: 4
+    rdzv_backend: static
+    same_network: true
+    tpu_env: []
+    tpu_use_cluster: false
+    tpu_use_sudo: false
+    use_cpu: false
+3. check accelerate config:
+    accelerate env
+4. Run the code:
+    accelerate launch sample_finetune.py
+"""
+
+logger = logging.getLogger(__name__)
+
+
+###################
+# Hyper-parameters
+###################
+training_config = {
+    "bf16": True,
+    "do_eval": False,
+    "learning_rate": 5.0e-06,
+    "log_level": "info",
+    "logging_steps": 20,
+    "logging_strategy": "steps",
+    "lr_scheduler_type": "cosine",
+    "num_train_epochs": 1,
+    "max_steps": -1,
+    "output_dir": "./checkpoint_dir",
+    "overwrite_output_dir": True,
+    "per_device_eval_batch_size": 4,
+    "per_device_train_batch_size": 4,
+    "remove_unused_columns": True,
+    "save_steps": 100,
+    "save_total_limit": 1,
+    "seed": 0,
+    "gradient_checkpointing": True,
+    "gradient_checkpointing_kwargs":{"use_reentrant": False},
+    "gradient_accumulation_steps": 1,
+    "warmup_ratio": 0.2,
+    }
+
+peft_config = {
+    "r": 16,
+    "lora_alpha": 32,
+    "lora_dropout": 0.05,
+    "bias": "none",
+    "task_type": "CAUSAL_LM",
+    "target_modules": "all-linear",
+    "modules_to_save": None,
+}
+train_conf = TrainingArguments(**training_config)
+peft_conf = LoraConfig(**peft_config)
+
+
+###############
+# Setup logging
+###############
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    handlers=[logging.StreamHandler(sys.stdout)],
+)
+log_level = train_conf.get_process_log_level()
+logger.setLevel(log_level)
+datasets.utils.logging.set_verbosity(log_level)
+transformers.utils.logging.set_verbosity(log_level)
+transformers.utils.logging.enable_default_handler()
+transformers.utils.logging.enable_explicit_format()
+
+# Log on each process a small summary
+logger.warning(
+    f"Process rank: {train_conf.local_rank}, device: {train_conf.device}, n_gpu: {train_conf.n_gpu}"
+    + f" distributed training: {bool(train_conf.local_rank != -1)}, 16-bits training: {train_conf.fp16}"
+)
+logger.info(f"Training/evaluation parameters {train_conf}")
+logger.info(f"PEFT parameters {peft_conf}")
+
+
+################
+# Model Loading
+################
+checkpoint_path = "microsoft/Phi-4-mini-instruct"
+model_kwargs = dict(
+    use_cache=False,
+    trust_remote_code=True,
+    attn_implementation="flash_attention_2",  # loading the model with flash-attention support
+    torch_dtype=torch.bfloat16,
+    device_map=None
+)
+model = AutoModelForCausalLM.from_pretrained(checkpoint_path, **model_kwargs)
+tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
+tokenizer.model_max_length = 2048
+tokenizer.pad_token = tokenizer.unk_token  # use unk rather than eos token to prevent endless generation
+tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
+tokenizer.padding_side = 'right'
+
+
+##################
+# Data Processing
+##################
+def apply_chat_template(
+    example,
+    tokenizer,
+):
+    messages = example["messages"]
+    example["text"] = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=False)
+    return example
+
+
+train_dataset, test_dataset = load_dataset("HuggingFaceH4/ultrachat_200k", split=["train_sft", "test_sft"])
+column_names = list(train_dataset.features)
+
+processed_train_dataset = train_dataset.map(
+    apply_chat_template,
+    fn_kwargs={"tokenizer": tokenizer},
+    num_proc=10,
+    remove_columns=column_names,
+    desc="Applying chat template to train_sft",
+)
+
+processed_test_dataset = test_dataset.map(
+    apply_chat_template,
+    fn_kwargs={"tokenizer": tokenizer},
+    num_proc=10,
+    remove_columns=column_names,
+    desc="Applying chat template to test_sft",
+)
+
+
+###########
+# Training
+###########
+trainer = SFTTrainer(
+    model=model,
+    args=train_conf,
+    peft_config=peft_conf,
+    train_dataset=processed_train_dataset,
+    eval_dataset=processed_test_dataset,
+    max_seq_length=2048,
+    dataset_text_field="text",
+    tokenizer=tokenizer,
+    packing=True
+)
+train_result = trainer.train()
+metrics = train_result.metrics
+trainer.log_metrics("train", metrics)
+trainer.save_metrics("train", metrics)
+trainer.save_state()
+
+
+#############
+# Evaluation
+#############
+tokenizer.padding_side = 'left'
+metrics = trainer.evaluate()
+metrics["eval_samples"] = len(processed_test_dataset)
+trainer.log_metrics("eval", metrics)
+trainer.save_metrics("eval", metrics)
+
+
+# ############
+# # Save model
+# ############
+trainer.save_model(train_conf.output_dir)

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,111 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "199999": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200018": {
+      "content": "<|endofprompt|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200019": {
+      "content": "<|assistant|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "200020": {
+      "content": "<|end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "200021": {
+      "content": "<|user|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "200022": {
+      "content": "<|system|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "200023": {
+      "content": "<|tool|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": false
+    },
+    "200024": {
+      "content": "<|/tool|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": false
+    },
+    "200025": {
+      "content": "<|tool_call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": false
+    },
+    "200026": {
+      "content": "<|/tool_call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": false
+    },
+    "200027": {
+      "content": "<|tool_response|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": false
+    },
+    "200028": {
+      "content": "<|tag|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}