Upload folder using huggingface_hub (#1)
Browse files- d8ec4feed7ff02230d96e38591bb267ac25f5d9b55f98e08ebf7d6449890f7ce (c8eaaa6eeea63314527020178f732f03666b5634)
- README.md +85 -0
- added_tokens.json +42 -0
- config.json +50 -0
- configuration_phi.py +62 -0
- generation_config.json +4 -0
- merges.txt +0 -0
- model.safetensors +3 -0
- modeling_phi.py +967 -0
- smash_config.json +31 -0
- special_tokens_map.json +30 -0
- tokenizer.json +0 -0
- tokenizer_config.json +343 -0
- vocab.json +0 -0
README.md
ADDED
@@ -0,0 +1,85 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
thumbnail: "https://assets-global.website-files.com/646b351987a8d8ce158d1940/64ec9e96b4334c0e1ac41504_Logo%20with%20white%20text.svg"
|
3 |
+
base_model: cognitivecomputations/dolphin-2_6-phi-2
|
4 |
+
metrics:
|
5 |
+
- memory_disk
|
6 |
+
- memory_inference
|
7 |
+
- inference_latency
|
8 |
+
- inference_throughput
|
9 |
+
- inference_CO2_emissions
|
10 |
+
- inference_energy_consumption
|
11 |
+
tags:
|
12 |
+
- pruna-ai
|
13 |
+
---
|
14 |
+
<!-- header start -->
|
15 |
+
<!-- 200823 -->
|
16 |
+
<div style="width: auto; margin-left: auto; margin-right: auto">
|
17 |
+
<a href="https://www.pruna.ai/" target="_blank" rel="noopener noreferrer">
|
18 |
+
<img src="https://i.imgur.com/eDAlcgk.png" alt="PrunaAI" style="width: 100%; min-width: 400px; display: block; margin: auto;">
|
19 |
+
</a>
|
20 |
+
</div>
|
21 |
+
<!-- header end -->
|
22 |
+
|
23 |
+
[![Twitter](https://img.shields.io/twitter/follow/PrunaAI?style=social)](https://twitter.com/PrunaAI)
|
24 |
+
[![GitHub](https://img.shields.io/github/followers/PrunaAI?label=Follow%20%40PrunaAI&style=social)](https://github.com/PrunaAI)
|
25 |
+
[![LinkedIn](https://img.shields.io/badge/LinkedIn-Connect-blue)](https://www.linkedin.com/company/93832878/admin/feed/posts/?feedType=following)
|
26 |
+
[![Discord](https://img.shields.io/badge/Discord-Join%20Us-blue?style=social&logo=discord)](https://discord.gg/CP4VSgck)
|
27 |
+
|
28 |
+
# Simply make AI models cheaper, smaller, faster, and greener!
|
29 |
+
|
30 |
+
- Give a thumbs up if you like this model!
|
31 |
+
- Contact us and tell us which model to compress next [here](https://www.pruna.ai/contact).
|
32 |
+
- Request access to easily compress your *own* AI models [here](https://z0halsaff74.typeform.com/pruna-access?typeform-source=www.pruna.ai).
|
33 |
+
- Read the documentations to know more [here](https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/)
|
34 |
+
- Join Pruna AI community on Discord [here](https://discord.gg/CP4VSgck) to share feedback/suggestions or get help.
|
35 |
+
|
36 |
+
## Results
|
37 |
+
|
38 |
+
![image info](./plots.png)
|
39 |
+
|
40 |
+
**Frequently Asked Questions**
|
41 |
+
- ***How does the compression work?*** The model is compressed with llm-int8.
|
42 |
+
- ***How does the model quality change?*** The quality of the model output might vary compared to the base model.
|
43 |
+
- ***How is the model efficiency evaluated?*** These results were obtained on HARDWARE_NAME with configuration described in `model/smash_config.json` and are obtained after a hardware warmup. The smashed model is directly compared to the original base model. Efficiency results may vary in other settings (e.g. other hardware, image size, batch size, ...). We recommend to directly run them in the use-case conditions to know if the smashed model can benefit you.
|
44 |
+
- ***What is the model format?*** We use safetensors.
|
45 |
+
- ***What calibration data has been used?*** If needed by the compression method, we used WikiText as the calibration data.
|
46 |
+
- ***What is the naming convention for Pruna Huggingface models?*** We take the original model name and append "turbo", "tiny", or "green" if the smashed model has a measured inference speed, inference memory, or inference energy consumption which is less than 90% of the original base model.
|
47 |
+
- ***How to compress my own models?*** You can request premium access to more compression methods and tech support for your specific use-cases [here](https://z0halsaff74.typeform.com/pruna-access?typeform-source=www.pruna.ai).
|
48 |
+
- ***What are "first" metrics?*** Results mentioning "first" are obtained after the first run of the model. The first run might take more memory or be slower than the subsequent runs due cuda overheads.
|
49 |
+
- ***What are "Sync" and "Async" metrics?*** "Sync" metrics are obtained by syncing all GPU processes and stop measurement when all of them are executed. "Async" metrics are obtained without syncing all GPU processes and stop when the model output can be used by the CPU. We provide both metrics since both could be relevant depending on the use-case. We recommend to test the efficiency gains directly in your use-cases.
|
50 |
+
|
51 |
+
## Setup
|
52 |
+
|
53 |
+
You can run the smashed model with these steps:
|
54 |
+
|
55 |
+
0. Check requirements from the original repo cognitivecomputations/dolphin-2_6-phi-2 installed. In particular, check python, cuda, and transformers versions.
|
56 |
+
1. Make sure that you have installed quantization related packages.
|
57 |
+
```bash
|
58 |
+
pip install transformers accelerate bitsandbytes>0.37.0
|
59 |
+
```
|
60 |
+
2. Load & run the model.
|
61 |
+
```python
|
62 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
63 |
+
|
64 |
+
|
65 |
+
model = AutoModelForCausalLM.from_pretrained("PrunaAI/cognitivecomputations-dolphin-2_6-phi-2-bnb-4bit-smashed", trust_remote_code=True, device_map='auto')
|
66 |
+
tokenizer = AutoTokenizer.from_pretrained("cognitivecomputations/dolphin-2_6-phi-2")
|
67 |
+
|
68 |
+
input_ids = tokenizer("What is the color of prunes?,", return_tensors='pt').to(model.device)["input_ids"]
|
69 |
+
|
70 |
+
outputs = model.generate(input_ids, max_new_tokens=216)
|
71 |
+
tokenizer.decode(outputs[0])
|
72 |
+
```
|
73 |
+
|
74 |
+
## Configurations
|
75 |
+
|
76 |
+
The configuration info are in `smash_config.json`.
|
77 |
+
|
78 |
+
## Credits & License
|
79 |
+
|
80 |
+
The license of the smashed model follows the license of the original model. Please check the license of the original model cognitivecomputations/dolphin-2_6-phi-2 before using this model which provided the base model. The license of the `pruna-engine` is [here](https://pypi.org/project/pruna-engine/) on Pypi.
|
81 |
+
|
82 |
+
## Want to compress other models?
|
83 |
+
|
84 |
+
- Contact us and tell us which model to compress next [here](https://www.pruna.ai/contact).
|
85 |
+
- Request access to easily compress your own AI models [here](https://z0halsaff74.typeform.com/pruna-access?typeform-source=www.pruna.ai).
|
added_tokens.json
ADDED
@@ -0,0 +1,42 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"\t\t": 50294,
|
3 |
+
"\t\t\t": 50293,
|
4 |
+
"\t\t\t\t": 50292,
|
5 |
+
"\t\t\t\t\t": 50291,
|
6 |
+
"\t\t\t\t\t\t": 50290,
|
7 |
+
"\t\t\t\t\t\t\t": 50289,
|
8 |
+
"\t\t\t\t\t\t\t\t": 50288,
|
9 |
+
"\t\t\t\t\t\t\t\t\t": 50287,
|
10 |
+
" ": 50286,
|
11 |
+
" ": 50285,
|
12 |
+
" ": 50284,
|
13 |
+
" ": 50283,
|
14 |
+
" ": 50282,
|
15 |
+
" ": 50281,
|
16 |
+
" ": 50280,
|
17 |
+
" ": 50279,
|
18 |
+
" ": 50278,
|
19 |
+
" ": 50277,
|
20 |
+
" ": 50276,
|
21 |
+
" ": 50275,
|
22 |
+
" ": 50274,
|
23 |
+
" ": 50273,
|
24 |
+
" ": 50272,
|
25 |
+
" ": 50271,
|
26 |
+
" ": 50270,
|
27 |
+
" ": 50269,
|
28 |
+
" ": 50268,
|
29 |
+
" ": 50267,
|
30 |
+
" ": 50266,
|
31 |
+
" ": 50265,
|
32 |
+
" ": 50264,
|
33 |
+
" ": 50263,
|
34 |
+
" ": 50262,
|
35 |
+
" ": 50261,
|
36 |
+
" ": 50260,
|
37 |
+
" ": 50259,
|
38 |
+
" ": 50258,
|
39 |
+
" ": 50257,
|
40 |
+
"<|im_end|>": 50295,
|
41 |
+
"<|im_start|>": 50296
|
42 |
+
}
|
config.json
ADDED
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"_name_or_path": "/ceph/hdd/staff/charpent/.cache/modelsl07ahlnoi8xqlpdb",
|
3 |
+
"activation_function": "gelu_new",
|
4 |
+
"architectures": [
|
5 |
+
"PhiForCausalLM"
|
6 |
+
],
|
7 |
+
"attn_pdrop": 0.0,
|
8 |
+
"auto_map": {
|
9 |
+
"AutoConfig": "configuration_phi.PhiConfig",
|
10 |
+
"AutoModelForCausalLM": "modeling_phi.PhiForCausalLM"
|
11 |
+
},
|
12 |
+
"embd_pdrop": 0.0,
|
13 |
+
"flash_attn": false,
|
14 |
+
"flash_rotary": false,
|
15 |
+
"fused_dense": false,
|
16 |
+
"img_processor": null,
|
17 |
+
"initializer_range": 0.02,
|
18 |
+
"layer_norm_epsilon": 1e-05,
|
19 |
+
"model_type": "phi-msft",
|
20 |
+
"n_embd": 2560,
|
21 |
+
"n_head": 32,
|
22 |
+
"n_head_kv": null,
|
23 |
+
"n_inner": null,
|
24 |
+
"n_layer": 32,
|
25 |
+
"n_positions": 2048,
|
26 |
+
"quantization_config": {
|
27 |
+
"_load_in_4bit": true,
|
28 |
+
"_load_in_8bit": false,
|
29 |
+
"bnb_4bit_compute_dtype": "bfloat16",
|
30 |
+
"bnb_4bit_quant_storage": "uint8",
|
31 |
+
"bnb_4bit_quant_type": "fp4",
|
32 |
+
"bnb_4bit_use_double_quant": false,
|
33 |
+
"llm_int8_enable_fp32_cpu_offload": false,
|
34 |
+
"llm_int8_has_fp16_weight": false,
|
35 |
+
"llm_int8_skip_modules": [
|
36 |
+
"lm_head"
|
37 |
+
],
|
38 |
+
"llm_int8_threshold": 6.0,
|
39 |
+
"load_in_4bit": true,
|
40 |
+
"load_in_8bit": false,
|
41 |
+
"quant_method": "bitsandbytes"
|
42 |
+
},
|
43 |
+
"resid_pdrop": 0.1,
|
44 |
+
"rotary_dim": 32,
|
45 |
+
"tie_word_embeddings": false,
|
46 |
+
"torch_dtype": "float16",
|
47 |
+
"transformers_version": "4.40.0",
|
48 |
+
"use_cache": false,
|
49 |
+
"vocab_size": 51200
|
50 |
+
}
|
configuration_phi.py
ADDED
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Microsoft Corporation.
|
2 |
+
# Licensed under the MIT license.
|
3 |
+
|
4 |
+
import math
|
5 |
+
from typing import Optional
|
6 |
+
|
7 |
+
from transformers import PretrainedConfig
|
8 |
+
|
9 |
+
|
10 |
+
class PhiConfig(PretrainedConfig):
|
11 |
+
"""Phi configuration."""
|
12 |
+
|
13 |
+
model_type = "phi-msft"
|
14 |
+
attribute_map = {
|
15 |
+
"max_position_embeddings": "n_positions",
|
16 |
+
"hidden_size": "n_embd",
|
17 |
+
"num_attention_heads": "n_head",
|
18 |
+
"num_hidden_layers": "n_layer",
|
19 |
+
}
|
20 |
+
|
21 |
+
def __init__(
|
22 |
+
self,
|
23 |
+
vocab_size: int = 50304,
|
24 |
+
n_positions: int = 2048,
|
25 |
+
n_embd: int = 1024,
|
26 |
+
n_layer: int = 20,
|
27 |
+
n_inner: Optional[int] = None,
|
28 |
+
n_head: int = 16,
|
29 |
+
n_head_kv: Optional[int] = None,
|
30 |
+
rotary_dim: Optional[int] = 32,
|
31 |
+
activation_function: Optional[str] = "gelu_new",
|
32 |
+
flash_attn: bool = False,
|
33 |
+
flash_rotary: bool = False,
|
34 |
+
fused_dense: bool = False,
|
35 |
+
attn_pdrop: float = 0.0,
|
36 |
+
embd_pdrop: float = 0.0,
|
37 |
+
resid_pdrop: float = 0.0,
|
38 |
+
layer_norm_epsilon: float = 1e-5,
|
39 |
+
initializer_range: float = 0.02,
|
40 |
+
tie_word_embeddings: bool = False,
|
41 |
+
pad_vocab_size_multiple: int = 64,
|
42 |
+
**kwargs
|
43 |
+
) -> None:
|
44 |
+
self.vocab_size = int(math.ceil(vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple)
|
45 |
+
self.n_positions = n_positions
|
46 |
+
self.n_embd = n_embd
|
47 |
+
self.n_layer = n_layer
|
48 |
+
self.n_inner = n_inner
|
49 |
+
self.n_head = n_head
|
50 |
+
self.n_head_kv = n_head_kv
|
51 |
+
self.rotary_dim = min(rotary_dim, n_embd // n_head)
|
52 |
+
self.activation_function = activation_function
|
53 |
+
self.flash_attn = flash_attn
|
54 |
+
self.flash_rotary = flash_rotary
|
55 |
+
self.fused_dense = fused_dense
|
56 |
+
self.attn_pdrop = attn_pdrop
|
57 |
+
self.embd_pdrop = embd_pdrop
|
58 |
+
self.resid_pdrop = resid_pdrop
|
59 |
+
self.layer_norm_epsilon = layer_norm_epsilon
|
60 |
+
self.initializer_range = initializer_range
|
61 |
+
|
62 |
+
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
|
generation_config.json
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"_from_model_config": true,
|
3 |
+
"transformers_version": "4.40.0"
|
4 |
+
}
|
merges.txt
ADDED
The diff for this file is too large to render.
See raw diff
|
|
model.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:40532af2099f9fae7d04cda8a9cedb884b9c53c254910d0a121af37762c73530
|
3 |
+
size 1941878416
|
modeling_phi.py
ADDED
@@ -0,0 +1,967 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Microsoft Corporation.
|
2 |
+
# Licensed under the MIT license.
|
3 |
+
#
|
4 |
+
# Copyright (c) 2022, Tri Dao, [email protected].
|
5 |
+
# Licensed under the BSD 3-Clause License.
|
6 |
+
|
7 |
+
from __future__ import annotations
|
8 |
+
|
9 |
+
import math
|
10 |
+
from dataclasses import dataclass, field
|
11 |
+
from typing import Any, Dict, Optional, Tuple, Union
|
12 |
+
|
13 |
+
import torch
|
14 |
+
import torch.nn as nn
|
15 |
+
from einops import rearrange, repeat
|
16 |
+
from transformers import PretrainedConfig, PreTrainedModel
|
17 |
+
from transformers.activations import ACT2FN
|
18 |
+
from transformers.modeling_outputs import CausalLMOutputWithPast
|
19 |
+
|
20 |
+
from .configuration_phi import PhiConfig
|
21 |
+
|
22 |
+
try:
|
23 |
+
from flash_attn.bert_padding import pad_input, unpad_input
|
24 |
+
from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
|
25 |
+
from flash_attn.modules.mha import FlashCrossAttention, FlashSelfAttention
|
26 |
+
from flash_attn.ops.fused_dense import FusedDense
|
27 |
+
except:
|
28 |
+
pad_input, unpad_input = None, None
|
29 |
+
FlashRotaryEmbedding = None
|
30 |
+
FlashSelfAttention, FlashCrossAttention = None, None
|
31 |
+
FusedDense = None
|
32 |
+
|
33 |
+
|
34 |
+
@dataclass
|
35 |
+
class InferenceParams:
|
36 |
+
"""Inference parameters passed to model to efficiently calculate
|
37 |
+
and store context during inference.
|
38 |
+
|
39 |
+
Reference:
|
40 |
+
https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/utils/generation.py.
|
41 |
+
|
42 |
+
Args:
|
43 |
+
max_seqlen: Maximum sequence length.
|
44 |
+
max_batch_size: Maximum batch size.
|
45 |
+
seqlen_offset: Sequence length offset.
|
46 |
+
batch_size_offset: Batch size offset.
|
47 |
+
key_value_memory_dict: Key value memory dictionary.
|
48 |
+
lengths_per_sample: Lengths per sample.
|
49 |
+
|
50 |
+
"""
|
51 |
+
|
52 |
+
max_seqlen: int = field(metadata={"help": "Maximum sequence length."})
|
53 |
+
|
54 |
+
max_batch_size: int = field(metadata={"help": "Maximum batch size."})
|
55 |
+
|
56 |
+
seqlen_offset: int = field(default=0, metadata={"help": "Sequence length offset."})
|
57 |
+
|
58 |
+
batch_size_offset: int = field(default=0, metadata={"help": "Batch size offset."})
|
59 |
+
|
60 |
+
key_value_memory_dict: Dict[str, Any] = field(
|
61 |
+
default_factory=dict, metadata={"help": "Key value memory dictionary."}
|
62 |
+
)
|
63 |
+
|
64 |
+
lengths_per_sample: torch.Tensor = field(default=None, metadata={"help": "Lengths per sample."})
|
65 |
+
|
66 |
+
|
67 |
+
class Embedding(nn.Module):
|
68 |
+
"""Token embedding with dropout."""
|
69 |
+
|
70 |
+
def __init__(self, config: PretrainedConfig) -> None:
|
71 |
+
super().__init__()
|
72 |
+
|
73 |
+
self.wte = nn.Embedding(config.vocab_size, config.n_embd)
|
74 |
+
self.drop = nn.Dropout(config.embd_pdrop)
|
75 |
+
|
76 |
+
def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
|
77 |
+
input_shape = input_ids.size()
|
78 |
+
input_ids = input_ids.view(-1, input_shape[-1])
|
79 |
+
|
80 |
+
hidden_states = self.wte(input_ids)
|
81 |
+
hidden_states = self.drop(hidden_states)
|
82 |
+
|
83 |
+
return hidden_states
|
84 |
+
|
85 |
+
|
86 |
+
def _apply_rotary_emb(
|
87 |
+
x: torch.FloatTensor,
|
88 |
+
cos: torch.FloatTensor,
|
89 |
+
sin: torch.FloatTensor,
|
90 |
+
) -> torch.FloatTensor:
|
91 |
+
_, seqlen, _, _ = x.shape
|
92 |
+
_, rotary_dim = cos.shape
|
93 |
+
rotary_dim *= 2
|
94 |
+
|
95 |
+
x_rot = x[:, :, :, :rotary_dim]
|
96 |
+
x_pass = x[:, :, :, rotary_dim:]
|
97 |
+
|
98 |
+
x1, x2 = x_rot.chunk(2, dim=-1)
|
99 |
+
c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
|
100 |
+
x1, x2, c, s = [t.to(dtype=torch.float32) for t in [x1, x2, c, s]]
|
101 |
+
|
102 |
+
x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], axis=-1).to(x.dtype)
|
103 |
+
|
104 |
+
return torch.cat([x_rot, x_pass], axis=-1)
|
105 |
+
|
106 |
+
|
107 |
+
def _apply_rotary_emb_kv(
|
108 |
+
kv: torch.FloatTensor,
|
109 |
+
cos: torch.FloatTensor,
|
110 |
+
sin: torch.FloatTensor,
|
111 |
+
cos_k: Optional[torch.FloatTensor] = None,
|
112 |
+
sin_k: Optional[torch.FloatTensor] = None,
|
113 |
+
) -> torch.FloatTensor:
|
114 |
+
_, seqlen, _, _, _ = kv.shape
|
115 |
+
_, rotary_dim = cos.shape
|
116 |
+
rotary_dim *= 2
|
117 |
+
|
118 |
+
k_rot = kv[:, :, 0, :, :rotary_dim]
|
119 |
+
k_pass = kv[:, :, 0, :, rotary_dim:]
|
120 |
+
|
121 |
+
k1, k2 = k_rot.chunk(2, dim=-1)
|
122 |
+
c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
|
123 |
+
k1, k2, c, s = [t.to(dtype=torch.float32) for t in [k1, k2, c, s]]
|
124 |
+
|
125 |
+
k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(kv.dtype)
|
126 |
+
|
127 |
+
return torch.cat(
|
128 |
+
[
|
129 |
+
torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
|
130 |
+
kv[:, :, 1:2, :, :],
|
131 |
+
],
|
132 |
+
axis=2,
|
133 |
+
)
|
134 |
+
|
135 |
+
|
136 |
+
def _apply_rotary_emb_qkv(
|
137 |
+
qkv: torch.FloatTensor,
|
138 |
+
cos: torch.FloatTensor,
|
139 |
+
sin: torch.FloatTensor,
|
140 |
+
cos_k: Optional[torch.FloatTensor] = None,
|
141 |
+
sin_k: Optional[torch.FloatTensor] = None,
|
142 |
+
) -> torch.FloatTensor:
|
143 |
+
_, seqlen, _, _, _ = qkv.shape
|
144 |
+
_, rotary_dim = cos.shape
|
145 |
+
rotary_dim *= 2
|
146 |
+
|
147 |
+
q_rot = qkv[:, :, 0, :, :rotary_dim]
|
148 |
+
q_pass = qkv[:, :, 0, :, rotary_dim:]
|
149 |
+
|
150 |
+
k_rot = qkv[:, :, 1, :, :rotary_dim]
|
151 |
+
k_pass = qkv[:, :, 1, :, rotary_dim:]
|
152 |
+
|
153 |
+
q1, q2 = q_rot.chunk(2, dim=-1)
|
154 |
+
k1, k2 = k_rot.chunk(2, dim=-1)
|
155 |
+
c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
|
156 |
+
q1, q2, k1, k2, c, s = [t.to(dtype=torch.float32) for t in [q1, q2, k1, k2, c, s]]
|
157 |
+
|
158 |
+
q_rot = torch.cat([q1 * c - q2 * s, q1 * s + q2 * c], axis=-1).to(qkv.dtype)
|
159 |
+
k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(qkv.dtype)
|
160 |
+
|
161 |
+
return torch.cat(
|
162 |
+
[
|
163 |
+
torch.cat([q_rot, q_pass], axis=-1).unsqueeze(2),
|
164 |
+
torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
|
165 |
+
qkv[:, :, 2:3, :, :],
|
166 |
+
],
|
167 |
+
axis=2,
|
168 |
+
)
|
169 |
+
|
170 |
+
|
171 |
+
class RotaryEmbedding(nn.Module):
|
172 |
+
"""Rotary positional embedding (RoPE).
|
173 |
+
|
174 |
+
Reference:
|
175 |
+
RoFormer: Enhanced Transformer with Rotary Position Embedding.
|
176 |
+
https://arxiv.org/pdf/2104.09864.pdf.
|
177 |
+
|
178 |
+
"""
|
179 |
+
|
180 |
+
def __init__(
|
181 |
+
self,
|
182 |
+
dim: int,
|
183 |
+
base: int = 10000,
|
184 |
+
scale_base: Optional[float] = None,
|
185 |
+
pos_idx_in_fp32: bool = True,
|
186 |
+
max_position_embeddings: int = 2048,
|
187 |
+
device: Optional[str] = None,
|
188 |
+
**kwargs,
|
189 |
+
) -> None:
|
190 |
+
super().__init__()
|
191 |
+
|
192 |
+
if scale_base is not None:
|
193 |
+
raise NotImplementedError
|
194 |
+
|
195 |
+
self.dim = dim
|
196 |
+
self.base = float(base)
|
197 |
+
self.scale_base = scale_base
|
198 |
+
self.pos_idx_in_fp32 = pos_idx_in_fp32
|
199 |
+
self.max_position_embeddings = max_position_embeddings
|
200 |
+
self.device = device
|
201 |
+
|
202 |
+
# Generate and save the inverse frequency buffer (non-trainable)
|
203 |
+
inv_freq = self._compute_inv_freq(device)
|
204 |
+
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
205 |
+
|
206 |
+
# Generate and save the scale buffer (non-trainable)
|
207 |
+
scale = (
|
208 |
+
(torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
|
209 |
+
if scale_base is not None
|
210 |
+
else None
|
211 |
+
)
|
212 |
+
self.register_buffer("scale", scale, persistent=False)
|
213 |
+
|
214 |
+
# Initialize cached attributes since ONNX can't rely on dynamic initialization
|
215 |
+
self._update_cos_sin_cache(max_position_embeddings, device=device, dtype=torch.float32)
|
216 |
+
|
217 |
+
def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor:
|
218 |
+
return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
|
219 |
+
|
220 |
+
def _update_cos_sin_cache(
|
221 |
+
self,
|
222 |
+
seqlen: int,
|
223 |
+
device: Optional[str] = None,
|
224 |
+
dtype: Optional[torch.dtype] = None,
|
225 |
+
) -> None:
|
226 |
+
self._seq_len_cached = seqlen
|
227 |
+
|
228 |
+
# fp32 is preferred since the output of `torch.arange` can be quite large
|
229 |
+
# and bf16 would lose a lot of precision
|
230 |
+
if self.pos_idx_in_fp32:
|
231 |
+
t = torch.arange(seqlen, device=device, dtype=torch.float32)
|
232 |
+
if self.inv_freq.dtype != torch.float32:
|
233 |
+
inv_freq = self._compute_inv_freq(device=device)
|
234 |
+
else:
|
235 |
+
inv_freq = self.inv_freq
|
236 |
+
else:
|
237 |
+
t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
|
238 |
+
inv_freq = self.inv_freq
|
239 |
+
|
240 |
+
# `torch.outer` is preferred since `torch.einsum` converts from fp32 to fp16 if used with AMP
|
241 |
+
freqs = torch.outer(t, inv_freq)
|
242 |
+
if self.scale is None:
|
243 |
+
self._cos_cached = torch.cos(freqs).to(dtype)
|
244 |
+
self._sin_cached = torch.sin(freqs).to(dtype)
|
245 |
+
else:
|
246 |
+
power = (
|
247 |
+
torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
|
248 |
+
) / self.scale_base
|
249 |
+
scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
|
250 |
+
|
251 |
+
# Force the scale multiplication to happen in fp32
|
252 |
+
self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
|
253 |
+
self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
|
254 |
+
self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
|
255 |
+
self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
|
256 |
+
|
257 |
+
def forward(
|
258 |
+
self,
|
259 |
+
qkv: torch.Tensor,
|
260 |
+
kv: Optional[torch.Tensor] = None,
|
261 |
+
seqlen_offset: int = 0,
|
262 |
+
**kwargs,
|
263 |
+
) -> Tuple[torch.Tensor, torch.Tensor]:
|
264 |
+
if (
|
265 |
+
self._seq_len_cached < qkv.shape[1] + seqlen_offset
|
266 |
+
or self._cos_cached.device != qkv.device
|
267 |
+
or self._cos_cached.dtype != qkv.dtype
|
268 |
+
or (self.training and self._cos_cached.is_inference())
|
269 |
+
):
|
270 |
+
self._update_cos_sin_cache(qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype)
|
271 |
+
|
272 |
+
if kv is None:
|
273 |
+
return _apply_rotary_emb_qkv(
|
274 |
+
qkv,
|
275 |
+
self._cos_cached[seqlen_offset:],
|
276 |
+
self._sin_cached[seqlen_offset:],
|
277 |
+
)
|
278 |
+
else:
|
279 |
+
q = _apply_rotary_emb(
|
280 |
+
qkv,
|
281 |
+
self._cos_cached[seqlen_offset:],
|
282 |
+
self._sin_cached[seqlen_offset:],
|
283 |
+
)
|
284 |
+
kv = _apply_rotary_emb_kv(
|
285 |
+
kv,
|
286 |
+
self._cos_cached[seqlen_offset:],
|
287 |
+
self._sin_cached[seqlen_offset:],
|
288 |
+
)
|
289 |
+
|
290 |
+
return q, kv
|
291 |
+
|
292 |
+
|
293 |
+
class MLP(nn.Module):
|
294 |
+
"""Multi-Layer Perceptron.
|
295 |
+
|
296 |
+
Reference:
|
297 |
+
Attention Is All You Need.
|
298 |
+
https://arxiv.org/pdf/1706.03762.pdf.
|
299 |
+
|
300 |
+
"""
|
301 |
+
|
302 |
+
def __init__(
|
303 |
+
self,
|
304 |
+
config: PretrainedConfig,
|
305 |
+
n_inner: Optional[int] = None,
|
306 |
+
act_fn: Optional[str] = None,
|
307 |
+
) -> None:
|
308 |
+
super().__init__()
|
309 |
+
|
310 |
+
act_fn = config.activation_function if act_fn is None else act_fn
|
311 |
+
|
312 |
+
n_inner = getattr(config, "n_inner", None) if n_inner is None else n_inner
|
313 |
+
n_inner = n_inner if n_inner is not None else 4 * config.n_embd
|
314 |
+
|
315 |
+
self.fc1 = nn.Linear(config.n_embd, n_inner)
|
316 |
+
self.fc2 = nn.Linear(n_inner, config.n_embd)
|
317 |
+
self.act = ACT2FN[act_fn]
|
318 |
+
|
319 |
+
def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
|
320 |
+
hidden_states = self.fc1(hidden_states)
|
321 |
+
hidden_states = self.act(hidden_states)
|
322 |
+
hidden_states = self.fc2(hidden_states)
|
323 |
+
|
324 |
+
return hidden_states
|
325 |
+
|
326 |
+
|
327 |
+
class SelfAttention(nn.Module):
|
328 |
+
"""Self-attention layer (compatible with PyTorch).
|
329 |
+
|
330 |
+
Reference:
|
331 |
+
https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
|
332 |
+
|
333 |
+
"""
|
334 |
+
|
335 |
+
def __init__(
|
336 |
+
self,
|
337 |
+
causal: bool = True,
|
338 |
+
softmax_scale: Optional[float] = None,
|
339 |
+
attention_dropout: float = 0.0,
|
340 |
+
) -> None:
|
341 |
+
super().__init__()
|
342 |
+
|
343 |
+
self.causal = causal
|
344 |
+
self.softmax_scale = softmax_scale
|
345 |
+
self.drop = nn.Dropout(attention_dropout)
|
346 |
+
|
347 |
+
@torch.autocast("cpu", enabled=False)
|
348 |
+
@torch.autocast("cuda", enabled=False)
|
349 |
+
def forward(
|
350 |
+
self,
|
351 |
+
qkv: torch.FloatTensor,
|
352 |
+
causal: bool = None,
|
353 |
+
key_padding_mask: Optional[torch.BoolTensor] = None,
|
354 |
+
**kwargs,
|
355 |
+
) -> torch.FloatTensor:
|
356 |
+
batch_size, seqlen = qkv.shape[0], qkv.shape[1]
|
357 |
+
q, k, v = qkv.unbind(dim=2)
|
358 |
+
|
359 |
+
q = q.to(torch.float32)
|
360 |
+
k = k.to(torch.float32)
|
361 |
+
|
362 |
+
causal = self.causal if causal is None else causal
|
363 |
+
softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
|
364 |
+
|
365 |
+
# Autocast is manually disabled to avoid `torch.einsum` performing the operation
|
366 |
+
# using float16, which might lead to overflow
|
367 |
+
scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
|
368 |
+
|
369 |
+
if key_padding_mask is not None:
|
370 |
+
padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device)
|
371 |
+
padding_mask.masked_fill_(key_padding_mask, 0.0)
|
372 |
+
|
373 |
+
scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
|
374 |
+
|
375 |
+
if causal:
|
376 |
+
causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1)
|
377 |
+
scores = scores + causal_mask.to(dtype=scores.dtype)
|
378 |
+
|
379 |
+
attention = torch.softmax(scores, dim=-1).to(v.dtype)
|
380 |
+
attention = self.drop(attention)
|
381 |
+
|
382 |
+
output = torch.einsum("bhts,bshd->bthd", attention, v)
|
383 |
+
|
384 |
+
return output
|
385 |
+
|
386 |
+
|
387 |
+
class CrossAttention(nn.Module):
|
388 |
+
"""Cross-attention layer (compatible with PyTorch).
|
389 |
+
|
390 |
+
Reference:
|
391 |
+
https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
|
392 |
+
|
393 |
+
"""
|
394 |
+
|
395 |
+
def __init__(
|
396 |
+
self,
|
397 |
+
causal: bool = True,
|
398 |
+
softmax_scale: Optional[float] = None,
|
399 |
+
attention_dropout: float = 0.0,
|
400 |
+
) -> None:
|
401 |
+
super().__init__()
|
402 |
+
|
403 |
+
self.causal = causal
|
404 |
+
self.softmax_scale = softmax_scale
|
405 |
+
self.drop = nn.Dropout(attention_dropout)
|
406 |
+
|
407 |
+
@torch.autocast("cpu", enabled=False)
|
408 |
+
@torch.autocast("cuda", enabled=False)
|
409 |
+
def forward(
|
410 |
+
self,
|
411 |
+
q: torch.FloatTensor,
|
412 |
+
kv: torch.FloatTensor,
|
413 |
+
causal: bool = None,
|
414 |
+
key_padding_mask: Optional[torch.BoolTensor] = None,
|
415 |
+
**kwargs,
|
416 |
+
) -> torch.FloatTensor:
|
417 |
+
batch_size, seqlen_q = q.shape[0], q.shape[1]
|
418 |
+
seqlen_k = kv.shape[1]
|
419 |
+
|
420 |
+
if kv.shape[3] != q.shape[2]:
|
421 |
+
kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3])
|
422 |
+
k, v = kv.unbind(dim=2)
|
423 |
+
|
424 |
+
q = q.to(torch.float32)
|
425 |
+
k = k.to(torch.float32)
|
426 |
+
|
427 |
+
causal = self.causal if causal is None else causal
|
428 |
+
softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
|
429 |
+
|
430 |
+
# Autocast is manually disabled to avoid `torch.einsum` performing the operation
|
431 |
+
# using float16, which might lead to overflow
|
432 |
+
scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
|
433 |
+
|
434 |
+
if key_padding_mask is not None:
|
435 |
+
padding_mask = torch.full(
|
436 |
+
(batch_size, seqlen_k),
|
437 |
+
-10000.0,
|
438 |
+
dtype=scores.dtype,
|
439 |
+
device=scores.device,
|
440 |
+
)
|
441 |
+
padding_mask.masked_fill_(key_padding_mask, 0.0)
|
442 |
+
|
443 |
+
scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
|
444 |
+
|
445 |
+
if causal:
|
446 |
+
rows = rearrange(torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1")
|
447 |
+
cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long)
|
448 |
+
causal_mask = cols > rows + seqlen_k - seqlen_q
|
449 |
+
|
450 |
+
scores = scores.masked_fill(causal_mask, -10000.0)
|
451 |
+
|
452 |
+
attention = torch.softmax(scores, dim=-1).to(v.dtype)
|
453 |
+
attention = self.drop(attention)
|
454 |
+
|
455 |
+
output = torch.einsum("bhts,bshd->bthd", attention, v)
|
456 |
+
|
457 |
+
return output
|
458 |
+
|
459 |
+
|
460 |
+
def _find_mha_dims(
|
461 |
+
config: PretrainedConfig,
|
462 |
+
n_head: Optional[int] = None,
|
463 |
+
n_head_kv: Optional[int] = None,
|
464 |
+
head_dim: Optional[int] = None,
|
465 |
+
) -> Tuple[int, int]:
|
466 |
+
if n_head is None and head_dim is None:
|
467 |
+
head_dim = config.n_embd // config.n_head
|
468 |
+
n_head = config.n_head
|
469 |
+
elif n_head is None or head_dim is None:
|
470 |
+
raise ValueError("`n_head` and `head_dim` must be both specified or `None`.")
|
471 |
+
|
472 |
+
if n_head_kv is None:
|
473 |
+
n_head_kv = getattr(config, "n_head_kv", None) or n_head
|
474 |
+
|
475 |
+
return n_head, n_head_kv, head_dim
|
476 |
+
|
477 |
+
|
478 |
+
def _update_kv_cache(kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int) -> torch.FloatTensor:
|
479 |
+
num_heads, head_dim = kv.shape[-2:]
|
480 |
+
|
481 |
+
if layer_idx not in inference_params.key_value_memory_dict:
|
482 |
+
inference_params.key_value_memory_dict[layer_idx] = torch.empty(
|
483 |
+
inference_params.max_batch_size,
|
484 |
+
inference_params.max_seqlen,
|
485 |
+
2,
|
486 |
+
num_heads,
|
487 |
+
head_dim,
|
488 |
+
dtype=kv.dtype,
|
489 |
+
device=kv.device,
|
490 |
+
)
|
491 |
+
|
492 |
+
batch_start = inference_params.batch_size_offset
|
493 |
+
batch_end = batch_start + kv.shape[0]
|
494 |
+
|
495 |
+
sequence_start = inference_params.seqlen_offset
|
496 |
+
sequence_end = sequence_start + kv.shape[1]
|
497 |
+
|
498 |
+
# When the current sequence length is equal to or larger than the maximum sequence length,
|
499 |
+
# we need to concatenate the current `kv` with the cached `kv` to expand its length
|
500 |
+
if sequence_end >= inference_params.max_seqlen:
|
501 |
+
inference_params.key_value_memory_dict[layer_idx] = torch.concatenate((inference_params.key_value_memory_dict[layer_idx], kv), dim=1)
|
502 |
+
|
503 |
+
inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...] = kv
|
504 |
+
kv = inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, :sequence_end, ...]
|
505 |
+
|
506 |
+
return kv
|
507 |
+
|
508 |
+
|
509 |
+
class MHA(nn.Module):
|
510 |
+
"""Multi-head attention layer."""
|
511 |
+
|
512 |
+
def __init__(
|
513 |
+
self,
|
514 |
+
config: PretrainedConfig,
|
515 |
+
dtype: Optional[torch.dtype] = None,
|
516 |
+
device: Optional[str] = None,
|
517 |
+
rotary_dim: Optional[int] = None,
|
518 |
+
rotary_base: float = 10000.0,
|
519 |
+
rotary_scale_base: Optional[float] = None,
|
520 |
+
n_head: Optional[int] = None,
|
521 |
+
n_head_kv: Optional[int] = None,
|
522 |
+
head_dim: Optional[int] = None,
|
523 |
+
bias: bool = True,
|
524 |
+
causal: bool = True,
|
525 |
+
softmax_scale: Optional[float] = None,
|
526 |
+
layer_idx: Optional[int] = None,
|
527 |
+
return_residual: bool = False,
|
528 |
+
checkpointing: bool = False,
|
529 |
+
) -> None:
|
530 |
+
super().__init__()
|
531 |
+
|
532 |
+
# Rotary embedding
|
533 |
+
self.rotary_dim = rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0)
|
534 |
+
if self.rotary_dim > 0:
|
535 |
+
rotary_cls = FlashRotaryEmbedding if config.flash_rotary else RotaryEmbedding
|
536 |
+
if rotary_cls is None:
|
537 |
+
rotary_cls = RotaryEmbedding
|
538 |
+
|
539 |
+
rotary_kwargs = {}
|
540 |
+
if rotary_cls is RotaryEmbedding:
|
541 |
+
rotary_kwargs["max_position_embeddings"] = config.n_positions
|
542 |
+
|
543 |
+
self.rotary_emb = rotary_cls(
|
544 |
+
self.rotary_dim,
|
545 |
+
base=rotary_base,
|
546 |
+
scale_base=rotary_scale_base,
|
547 |
+
device=device,
|
548 |
+
**rotary_kwargs,
|
549 |
+
)
|
550 |
+
|
551 |
+
# MLP
|
552 |
+
self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims(
|
553 |
+
config, n_head=n_head, n_head_kv=n_head_kv, head_dim=head_dim
|
554 |
+
)
|
555 |
+
op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv)
|
556 |
+
hidden_size = config.n_embd
|
557 |
+
|
558 |
+
linear_cls = FusedDense if config.fused_dense else nn.Linear
|
559 |
+
if linear_cls is None:
|
560 |
+
linear_cls = nn.Linear
|
561 |
+
|
562 |
+
self.Wqkv = linear_cls(hidden_size, op_size, bias=bias, device=device, dtype=dtype)
|
563 |
+
self.out_proj = linear_cls(hidden_size, hidden_size, bias=bias, device=device, dtype=dtype)
|
564 |
+
|
565 |
+
# Attention
|
566 |
+
attn_cls = FlashSelfAttention if config.flash_attn else SelfAttention
|
567 |
+
if attn_cls is None:
|
568 |
+
attn_cls = SelfAttention
|
569 |
+
|
570 |
+
cross_attn_cls = FlashCrossAttention if config.flash_attn else CrossAttention
|
571 |
+
if cross_attn_cls is None:
|
572 |
+
cross_attn_cls = CrossAttention
|
573 |
+
|
574 |
+
self.inner_attn = attn_cls(
|
575 |
+
causal=causal,
|
576 |
+
softmax_scale=softmax_scale,
|
577 |
+
attention_dropout=config.attn_pdrop,
|
578 |
+
)
|
579 |
+
self.inner_cross_attn = cross_attn_cls(
|
580 |
+
causal=causal,
|
581 |
+
softmax_scale=softmax_scale,
|
582 |
+
attention_dropout=config.attn_pdrop,
|
583 |
+
)
|
584 |
+
|
585 |
+
self.flash_attn = config.flash_attn and attn_cls is FlashSelfAttention
|
586 |
+
self.layer_idx = layer_idx
|
587 |
+
self.return_residual = return_residual
|
588 |
+
self.checkpointing = checkpointing
|
589 |
+
|
590 |
+
def _forward_self_attn(
|
591 |
+
self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor]
|
592 |
+
) -> torch.FloatTensor:
|
593 |
+
qkv = self.Wqkv(x)
|
594 |
+
qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim)
|
595 |
+
|
596 |
+
if self.rotary_dim > 0:
|
597 |
+
qkv = self.rotary_emb(qkv)
|
598 |
+
|
599 |
+
if self.flash_attn:
|
600 |
+
batch_size, seqlen = qkv.shape[0], qkv.shape[1]
|
601 |
+
|
602 |
+
cu_seqlens, max_seqlen = None, None
|
603 |
+
if key_padding_mask is not None:
|
604 |
+
# If `key_padding_mask` is supplied, we need to unpad the input and retrieve
|
605 |
+
# the `cu_seqlens` and `max_seqlen` to be used by `flash-attn`
|
606 |
+
qkv, indices, cu_seqlens, max_seqlen = unpad_input(qkv, key_padding_mask)
|
607 |
+
|
608 |
+
if self.checkpointing and self.training:
|
609 |
+
attn_output = torch.utils.checkpoint.checkpoint(
|
610 |
+
self.inner_attn, qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
|
611 |
+
)
|
612 |
+
else:
|
613 |
+
attn_output = self.inner_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen).to(qkv.device)
|
614 |
+
|
615 |
+
# If `key_padding_mask` is supplied, we need to pad the output back to the original shape
|
616 |
+
return pad_input(attn_output, indices, batch_size, seqlen) if key_padding_mask is not None else attn_output
|
617 |
+
|
618 |
+
if self.checkpointing and self.training:
|
619 |
+
return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, key_padding_mask=key_padding_mask, use_reentrant=False)
|
620 |
+
|
621 |
+
return self.inner_attn(qkv, key_padding_mask=key_padding_mask)
|
622 |
+
|
623 |
+
def _forward_cross_attn(
|
624 |
+
self,
|
625 |
+
x: torch.FloatTensor,
|
626 |
+
past_key_values: Optional[InferenceParams],
|
627 |
+
key_padding_mask: Optional[torch.BoolTensor],
|
628 |
+
) -> torch.FloatTensor:
|
629 |
+
batch_size = x.shape[0]
|
630 |
+
|
631 |
+
qkv = self.Wqkv(x)
|
632 |
+
|
633 |
+
q = qkv[..., : self.n_head * self.head_dim]
|
634 |
+
q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
|
635 |
+
|
636 |
+
kv = qkv[..., self.n_head * self.head_dim :]
|
637 |
+
kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)
|
638 |
+
|
639 |
+
seqlen_offset = past_key_values.seqlen_offset if past_key_values is not None else 0
|
640 |
+
causal = None if seqlen_offset == 0 else False
|
641 |
+
if self.rotary_dim > 0:
|
642 |
+
q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset)
|
643 |
+
|
644 |
+
if past_key_values is not None:
|
645 |
+
kv = _update_kv_cache(kv, past_key_values, self.layer_idx)
|
646 |
+
|
647 |
+
if self.flash_attn:
|
648 |
+
batch_size, seqlen_q = q.shape[0], q.shape[1]
|
649 |
+
seqlen_k = kv.shape[1]
|
650 |
+
|
651 |
+
cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = (
|
652 |
+
None,
|
653 |
+
None,
|
654 |
+
None,
|
655 |
+
None,
|
656 |
+
)
|
657 |
+
if key_padding_mask is not None:
|
658 |
+
kv, _, cu_seqlens_k, max_seqlen_k = unpad_input(kv, key_padding_mask)
|
659 |
+
|
660 |
+
if seqlen_q == 1:
|
661 |
+
key_padding_mask = torch.ones(batch_size, 1, device=q.device)
|
662 |
+
elif seqlen_q != seqlen_k:
|
663 |
+
key_padding_mask = key_padding_mask[:, -seqlen_q:]
|
664 |
+
|
665 |
+
q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, key_padding_mask)
|
666 |
+
|
667 |
+
if self.checkpointing and self.training:
|
668 |
+
attn_output = torch.utils.checkpoint.checkpoint(
|
669 |
+
self.inner_cross_attn,
|
670 |
+
q,
|
671 |
+
kv,
|
672 |
+
causal=causal,
|
673 |
+
cu_seqlens=cu_seqlens_q,
|
674 |
+
max_seqlen=max_seqlen_q,
|
675 |
+
cu_seqlens_k=cu_seqlens_k,
|
676 |
+
max_seqlen_k=max_seqlen_k,
|
677 |
+
use_reentrant=False
|
678 |
+
)
|
679 |
+
else:
|
680 |
+
attn_output = self.inner_cross_attn(
|
681 |
+
q,
|
682 |
+
kv,
|
683 |
+
causal=causal,
|
684 |
+
cu_seqlens=cu_seqlens_q,
|
685 |
+
max_seqlen=max_seqlen_q,
|
686 |
+
cu_seqlens_k=cu_seqlens_k,
|
687 |
+
max_seqlen_k=max_seqlen_k,
|
688 |
+
)
|
689 |
+
|
690 |
+
return (
|
691 |
+
pad_input(attn_output, indices_q, batch_size, max_seqlen_q)
|
692 |
+
if key_padding_mask is not None
|
693 |
+
else attn_output
|
694 |
+
)
|
695 |
+
|
696 |
+
if self.checkpointing and self.training:
|
697 |
+
return torch.utils.checkpoint.checkpoint(
|
698 |
+
self.inner_cross_attn,
|
699 |
+
q,
|
700 |
+
kv,
|
701 |
+
key_padding_mask=key_padding_mask,
|
702 |
+
causal=causal,
|
703 |
+
use_reentrant=False
|
704 |
+
)
|
705 |
+
|
706 |
+
return self.inner_cross_attn(q, kv, key_padding_mask=key_padding_mask, causal=causal)
|
707 |
+
|
708 |
+
def forward(
|
709 |
+
self,
|
710 |
+
x: torch.FloatTensor,
|
711 |
+
past_key_values: Optional[InferenceParams] = None,
|
712 |
+
attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
|
713 |
+
**kwargs,
|
714 |
+
) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
|
715 |
+
if attention_mask is not None:
|
716 |
+
attention_mask = attention_mask.bool()
|
717 |
+
else:
|
718 |
+
attention_mask = None
|
719 |
+
|
720 |
+
# MHA
|
721 |
+
if self.n_head == self.n_head_kv:
|
722 |
+
if past_key_values is None:
|
723 |
+
# If `past_key_values` are not supplied, we run self-attention
|
724 |
+
attn_output = self._forward_self_attn(x, attention_mask)
|
725 |
+
else:
|
726 |
+
# If `past_key_values` are supplied, it means that we might have cached values and
|
727 |
+
# could take advantage of cross-attention
|
728 |
+
attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)
|
729 |
+
# MQA / GQA
|
730 |
+
else:
|
731 |
+
# Regardless of `past_key_values` being supplied or not, it always use cross-attention
|
732 |
+
# because `q` and `kv` lengths might be different
|
733 |
+
attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)
|
734 |
+
|
735 |
+
output = rearrange(attn_output, "... h d -> ... (h d)")
|
736 |
+
output = self.out_proj(output)
|
737 |
+
|
738 |
+
return output if not self.return_residual else (output, x)
|
739 |
+
|
740 |
+
|
741 |
+
class ParallelBlock(nn.Module):
|
742 |
+
"""Parallel block.
|
743 |
+
|
744 |
+
This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen).
|
745 |
+
|
746 |
+
"""
|
747 |
+
|
748 |
+
def __init__(
|
749 |
+
self,
|
750 |
+
config: PretrainedConfig,
|
751 |
+
block_idx: Optional[int] = None,
|
752 |
+
) -> None:
|
753 |
+
super().__init__()
|
754 |
+
|
755 |
+
self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
|
756 |
+
self.resid_dropout = nn.Dropout(config.resid_pdrop)
|
757 |
+
self.block_idx = block_idx
|
758 |
+
|
759 |
+
self.mixer = MHA(config, layer_idx=block_idx)
|
760 |
+
self.mlp = MLP(config)
|
761 |
+
|
762 |
+
def forward(
|
763 |
+
self,
|
764 |
+
hidden_states: torch.FloatTensor,
|
765 |
+
past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
|
766 |
+
attention_mask: Optional[torch.BoolTensor] = None,
|
767 |
+
**kwargs,
|
768 |
+
) -> torch.FloatTensor:
|
769 |
+
residual = hidden_states
|
770 |
+
hidden_states = self.ln(hidden_states)
|
771 |
+
|
772 |
+
attn_outputs = self.mixer(
|
773 |
+
hidden_states,
|
774 |
+
past_key_values=past_key_values,
|
775 |
+
attention_mask=attention_mask,
|
776 |
+
)
|
777 |
+
if isinstance(attn_outputs, tuple):
|
778 |
+
attn_outputs = attn_outputs[0]
|
779 |
+
|
780 |
+
attn_outputs = self.resid_dropout(attn_outputs)
|
781 |
+
feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
|
782 |
+
|
783 |
+
hidden_states = attn_outputs + feed_forward_hidden_states + residual
|
784 |
+
|
785 |
+
return hidden_states
|
786 |
+
|
787 |
+
|
788 |
+
class CausalLMHead(nn.Module):
|
789 |
+
"""Causal Language Modeling head.
|
790 |
+
|
791 |
+
Reference:
|
792 |
+
Improving Language Understanding by Generative Pre-Training.
|
793 |
+
https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
|
794 |
+
|
795 |
+
"""
|
796 |
+
|
797 |
+
def __init__(self, config: PretrainedConfig) -> None:
|
798 |
+
super().__init__()
|
799 |
+
|
800 |
+
self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
|
801 |
+
self.linear = nn.Linear(config.n_embd, config.vocab_size)
|
802 |
+
|
803 |
+
def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
|
804 |
+
hidden_states = self.ln(hidden_states)
|
805 |
+
logits = self.linear(hidden_states).to(torch.float32)
|
806 |
+
|
807 |
+
return logits
|
808 |
+
|
809 |
+
|
810 |
+
class CausalLMLoss(nn.Module):
|
811 |
+
"""Causal Language Modeling loss.
|
812 |
+
|
813 |
+
Reference:
|
814 |
+
Improving Language Understanding by Generative Pre-Training.
|
815 |
+
https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
|
816 |
+
|
817 |
+
"""
|
818 |
+
|
819 |
+
def __init__(self, shift_labels: bool = True) -> None:
|
820 |
+
super().__init__()
|
821 |
+
|
822 |
+
self.shift_labels = shift_labels
|
823 |
+
self.loss_fct = nn.CrossEntropyLoss()
|
824 |
+
|
825 |
+
def forward(self, logits: torch.FloatTensor, labels: torch.LongTensor) -> torch.FloatTensor:
|
826 |
+
if self.shift_labels:
|
827 |
+
logits = logits[..., :-1, :].contiguous()
|
828 |
+
labels = labels[..., 1:].contiguous()
|
829 |
+
|
830 |
+
loss = self.loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
|
831 |
+
|
832 |
+
return loss
|
833 |
+
|
834 |
+
|
835 |
+
class PhiPreTrainedModel(PreTrainedModel):
|
836 |
+
"""Phi pre-trained model."""
|
837 |
+
|
838 |
+
config_class = PhiConfig
|
839 |
+
base_model_prefix = "transformer"
|
840 |
+
supports_gradient_checkpointing = True
|
841 |
+
_no_split_modules = ["ParallelBlock"]
|
842 |
+
|
843 |
+
def __init__(self, *inputs, **kwargs) -> None:
|
844 |
+
super().__init__(*inputs, **kwargs)
|
845 |
+
|
846 |
+
def _init_weights(self, module: nn.Module) -> None:
|
847 |
+
if isinstance(module, (nn.Linear,)):
|
848 |
+
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
|
849 |
+
if module.bias is not None:
|
850 |
+
module.bias.data.zero_()
|
851 |
+
elif isinstance(module, nn.Embedding):
|
852 |
+
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
|
853 |
+
if module.padding_idx is not None:
|
854 |
+
module.weight.data[module.padding_idx].zero_()
|
855 |
+
elif isinstance(module, nn.LayerNorm):
|
856 |
+
if module.bias is not None:
|
857 |
+
module.bias.data.zero_()
|
858 |
+
module.weight.data.fill_(1.0)
|
859 |
+
|
860 |
+
|
861 |
+
def _set_gradient_checkpointing(self, module, value=False):
|
862 |
+
if isinstance(module, MHA):
|
863 |
+
module.checkpointing = value
|
864 |
+
|
865 |
+
def prepare_inputs_for_generation(
|
866 |
+
self,
|
867 |
+
input_ids: torch.LongTensor,
|
868 |
+
past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
|
869 |
+
attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
|
870 |
+
**kwargs,
|
871 |
+
) -> Dict[str, Any]:
|
872 |
+
if past_key_values is None or not (isinstance(past_key_values, InferenceParams)):
|
873 |
+
past_key_values = InferenceParams(
|
874 |
+
max_seqlen=self.config.n_positions,
|
875 |
+
max_batch_size=input_ids.shape[0],
|
876 |
+
seqlen_offset=0,
|
877 |
+
batch_size_offset=0,
|
878 |
+
key_value_memory_dict={},
|
879 |
+
lengths_per_sample=None,
|
880 |
+
)
|
881 |
+
else:
|
882 |
+
# Assume that `past_key_values` has cached all tokens up to the last token in `input_ids`
|
883 |
+
past_key_values.seqlen_offset = input_ids.shape[1] - 1
|
884 |
+
input_ids = input_ids[:, -1].unsqueeze(-1)
|
885 |
+
|
886 |
+
return {
|
887 |
+
"input_ids": input_ids,
|
888 |
+
"past_key_values": past_key_values,
|
889 |
+
"attention_mask": attention_mask,
|
890 |
+
}
|
891 |
+
|
892 |
+
|
893 |
+
class PhiModel(PhiPreTrainedModel):
|
894 |
+
"""Phi model."""
|
895 |
+
|
896 |
+
_keys_to_ignore_on_load_missing = [""]
|
897 |
+
_keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
|
898 |
+
|
899 |
+
def __init__(self, config: PhiConfig) -> None:
|
900 |
+
super().__init__(config)
|
901 |
+
|
902 |
+
self.embd = Embedding(config)
|
903 |
+
self.h = nn.ModuleList([ParallelBlock(config, block_idx=i) for i in range(config.n_layer)])
|
904 |
+
self.gradient_checkpointing = False
|
905 |
+
self.post_init()
|
906 |
+
|
907 |
+
def get_input_embeddings(self) -> nn.Embedding:
|
908 |
+
return self.embd.wte
|
909 |
+
|
910 |
+
def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
|
911 |
+
self.embd.wte = new_embeddings
|
912 |
+
|
913 |
+
def forward(
|
914 |
+
self,
|
915 |
+
input_ids: torch.LongTensor,
|
916 |
+
past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
|
917 |
+
attention_mask: Optional[torch.BoolTensor] = None,
|
918 |
+
) -> torch.FloatTensor:
|
919 |
+
hidden_states = self.embd(input_ids)
|
920 |
+
|
921 |
+
for layer in self.h:
|
922 |
+
hidden_states = layer(
|
923 |
+
hidden_states,
|
924 |
+
past_key_values=past_key_values,
|
925 |
+
attention_mask=attention_mask,
|
926 |
+
)
|
927 |
+
|
928 |
+
return hidden_states
|
929 |
+
|
930 |
+
|
931 |
+
class PhiForCausalLM(PhiPreTrainedModel):
|
932 |
+
"""Phi for Causal Language Modeling."""
|
933 |
+
|
934 |
+
_keys_to_ignore_on_load_missing = [""]
|
935 |
+
_keys_to_ignore_on_load_unexpected = [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
|
936 |
+
|
937 |
+
def __init__(self, config: PhiConfig) -> None:
|
938 |
+
super().__init__(config)
|
939 |
+
|
940 |
+
self.transformer = PhiModel(config)
|
941 |
+
self.lm_head = CausalLMHead(config)
|
942 |
+
self.loss = CausalLMLoss()
|
943 |
+
|
944 |
+
self.post_init()
|
945 |
+
|
946 |
+
def get_output_embeddings(self) -> nn.Linear:
|
947 |
+
return self.lm_head.linear
|
948 |
+
|
949 |
+
def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
|
950 |
+
self.lm_head.linear = new_embeddings
|
951 |
+
|
952 |
+
def forward(
|
953 |
+
self,
|
954 |
+
input_ids: torch.LongTensor,
|
955 |
+
past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
|
956 |
+
attention_mask: Optional[torch.BoolTensor] = None,
|
957 |
+
labels: Optional[torch.LongTensor] = None,
|
958 |
+
**kwargs,
|
959 |
+
) -> CausalLMOutputWithPast:
|
960 |
+
hidden_states = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask)
|
961 |
+
lm_logits = self.lm_head(hidden_states)
|
962 |
+
|
963 |
+
loss = None
|
964 |
+
if labels is not None:
|
965 |
+
loss = self.loss(lm_logits, labels)
|
966 |
+
|
967 |
+
return CausalLMOutputWithPast(loss=loss, logits=lm_logits, past_key_values=past_key_values)
|
smash_config.json
ADDED
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"api_key": null,
|
3 |
+
"verify_url": "http://johnrachwan.pythonanywhere.com",
|
4 |
+
"smash_config": {
|
5 |
+
"pruners": "None",
|
6 |
+
"pruning_ratio": 0.0,
|
7 |
+
"factorizers": "None",
|
8 |
+
"quantizers": "['llm-int8']",
|
9 |
+
"weight_quantization_bits": 4,
|
10 |
+
"output_deviation": 0.005,
|
11 |
+
"compilers": "None",
|
12 |
+
"static_batch": true,
|
13 |
+
"static_shape": true,
|
14 |
+
"controlnet": "None",
|
15 |
+
"unet_dim": 4,
|
16 |
+
"device": "cuda",
|
17 |
+
"cache_dir": "/ceph/hdd/staff/charpent/.cache/modelsl07ahlno",
|
18 |
+
"batch_size": 1,
|
19 |
+
"model_name": "cognitivecomputations/dolphin-2_6-phi-2",
|
20 |
+
"task": "text_text_generation",
|
21 |
+
"max_batch_size": 1,
|
22 |
+
"qtype_weight": "torch.qint8",
|
23 |
+
"qtype_activation": "torch.quint8",
|
24 |
+
"qobserver": "<class 'torch.ao.quantization.observer.MinMaxObserver'>",
|
25 |
+
"qscheme": "torch.per_tensor_symmetric",
|
26 |
+
"qconfig": "x86",
|
27 |
+
"group_size": 128,
|
28 |
+
"damp_percent": 0.1,
|
29 |
+
"save_load_fn": "bitsandbytes"
|
30 |
+
}
|
31 |
+
}
|
special_tokens_map.json
ADDED
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"bos_token": {
|
3 |
+
"content": "<|endoftext|>",
|
4 |
+
"lstrip": false,
|
5 |
+
"normalized": false,
|
6 |
+
"rstrip": false,
|
7 |
+
"single_word": false
|
8 |
+
},
|
9 |
+
"eos_token": {
|
10 |
+
"content": "<|im_end|>",
|
11 |
+
"lstrip": false,
|
12 |
+
"normalized": false,
|
13 |
+
"rstrip": false,
|
14 |
+
"single_word": false
|
15 |
+
},
|
16 |
+
"pad_token": {
|
17 |
+
"content": "<|endoftext|>",
|
18 |
+
"lstrip": false,
|
19 |
+
"normalized": false,
|
20 |
+
"rstrip": false,
|
21 |
+
"single_word": false
|
22 |
+
},
|
23 |
+
"unk_token": {
|
24 |
+
"content": "<|endoftext|>",
|
25 |
+
"lstrip": false,
|
26 |
+
"normalized": false,
|
27 |
+
"rstrip": false,
|
28 |
+
"single_word": false
|
29 |
+
}
|
30 |
+
}
|
tokenizer.json
ADDED
The diff for this file is too large to render.
See raw diff
|
|
tokenizer_config.json
ADDED
@@ -0,0 +1,343 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"add_prefix_space": false,
|
3 |
+
"added_tokens_decoder": {
|
4 |
+
"50256": {
|
5 |
+
"content": "<|endoftext|>",
|
6 |
+
"lstrip": false,
|
7 |
+
"normalized": false,
|
8 |
+
"rstrip": false,
|
9 |
+
"single_word": false,
|
10 |
+
"special": true
|
11 |
+
},
|
12 |
+
"50257": {
|
13 |
+
"content": " ",
|
14 |
+
"lstrip": false,
|
15 |
+
"normalized": true,
|
16 |
+
"rstrip": false,
|
17 |
+
"single_word": false,
|
18 |
+
"special": false
|
19 |
+
},
|
20 |
+
"50258": {
|
21 |
+
"content": " ",
|
22 |
+
"lstrip": false,
|
23 |
+
"normalized": true,
|
24 |
+
"rstrip": false,
|
25 |
+
"single_word": false,
|
26 |
+
"special": false
|
27 |
+
},
|
28 |
+
"50259": {
|
29 |
+
"content": " ",
|
30 |
+
"lstrip": false,
|
31 |
+
"normalized": true,
|
32 |
+
"rstrip": false,
|
33 |
+
"single_word": false,
|
34 |
+
"special": false
|
35 |
+
},
|
36 |
+
"50260": {
|
37 |
+
"content": " ",
|
38 |
+
"lstrip": false,
|
39 |
+
"normalized": true,
|
40 |
+
"rstrip": false,
|
41 |
+
"single_word": false,
|
42 |
+
"special": false
|
43 |
+
},
|
44 |
+
"50261": {
|
45 |
+
"content": " ",
|
46 |
+
"lstrip": false,
|
47 |
+
"normalized": true,
|
48 |
+
"rstrip": false,
|
49 |
+
"single_word": false,
|
50 |
+
"special": false
|
51 |
+
},
|
52 |
+
"50262": {
|
53 |
+
"content": " ",
|
54 |
+
"lstrip": false,
|
55 |
+
"normalized": true,
|
56 |
+
"rstrip": false,
|
57 |
+
"single_word": false,
|
58 |
+
"special": false
|
59 |
+
},
|
60 |
+
"50263": {
|
61 |
+
"content": " ",
|
62 |
+
"lstrip": false,
|
63 |
+
"normalized": true,
|
64 |
+
"rstrip": false,
|
65 |
+
"single_word": false,
|
66 |
+
"special": false
|
67 |
+
},
|
68 |
+
"50264": {
|
69 |
+
"content": " ",
|
70 |
+
"lstrip": false,
|
71 |
+
"normalized": true,
|
72 |
+
"rstrip": false,
|
73 |
+
"single_word": false,
|
74 |
+
"special": false
|
75 |
+
},
|
76 |
+
"50265": {
|
77 |
+
"content": " ",
|
78 |
+
"lstrip": false,
|
79 |
+
"normalized": true,
|
80 |
+
"rstrip": false,
|
81 |
+
"single_word": false,
|
82 |
+
"special": false
|
83 |
+
},
|
84 |
+
"50266": {
|
85 |
+
"content": " ",
|
86 |
+
"lstrip": false,
|
87 |
+
"normalized": true,
|
88 |
+
"rstrip": false,
|
89 |
+
"single_word": false,
|
90 |
+
"special": false
|
91 |
+
},
|
92 |
+
"50267": {
|
93 |
+
"content": " ",
|
94 |
+
"lstrip": false,
|
95 |
+
"normalized": true,
|
96 |
+
"rstrip": false,
|
97 |
+
"single_word": false,
|
98 |
+
"special": false
|
99 |
+
},
|
100 |
+
"50268": {
|
101 |
+
"content": " ",
|
102 |
+
"lstrip": false,
|
103 |
+
"normalized": true,
|
104 |
+
"rstrip": false,
|
105 |
+
"single_word": false,
|
106 |
+
"special": false
|
107 |
+
},
|
108 |
+
"50269": {
|
109 |
+
"content": " ",
|
110 |
+
"lstrip": false,
|
111 |
+
"normalized": true,
|
112 |
+
"rstrip": false,
|
113 |
+
"single_word": false,
|
114 |
+
"special": false
|
115 |
+
},
|
116 |
+
"50270": {
|
117 |
+
"content": " ",
|
118 |
+
"lstrip": false,
|
119 |
+
"normalized": true,
|
120 |
+
"rstrip": false,
|
121 |
+
"single_word": false,
|
122 |
+
"special": false
|
123 |
+
},
|
124 |
+
"50271": {
|
125 |
+
"content": " ",
|
126 |
+
"lstrip": false,
|
127 |
+
"normalized": true,
|
128 |
+
"rstrip": false,
|
129 |
+
"single_word": false,
|
130 |
+
"special": false
|
131 |
+
},
|
132 |
+
"50272": {
|
133 |
+
"content": " ",
|
134 |
+
"lstrip": false,
|
135 |
+
"normalized": true,
|
136 |
+
"rstrip": false,
|
137 |
+
"single_word": false,
|
138 |
+
"special": false
|
139 |
+
},
|
140 |
+
"50273": {
|
141 |
+
"content": " ",
|
142 |
+
"lstrip": false,
|
143 |
+
"normalized": true,
|
144 |
+
"rstrip": false,
|
145 |
+
"single_word": false,
|
146 |
+
"special": false
|
147 |
+
},
|
148 |
+
"50274": {
|
149 |
+
"content": " ",
|
150 |
+
"lstrip": false,
|
151 |
+
"normalized": true,
|
152 |
+
"rstrip": false,
|
153 |
+
"single_word": false,
|
154 |
+
"special": false
|
155 |
+
},
|
156 |
+
"50275": {
|
157 |
+
"content": " ",
|
158 |
+
"lstrip": false,
|
159 |
+
"normalized": true,
|
160 |
+
"rstrip": false,
|
161 |
+
"single_word": false,
|
162 |
+
"special": false
|
163 |
+
},
|
164 |
+
"50276": {
|
165 |
+
"content": " ",
|
166 |
+
"lstrip": false,
|
167 |
+
"normalized": true,
|
168 |
+
"rstrip": false,
|
169 |
+
"single_word": false,
|
170 |
+
"special": false
|
171 |
+
},
|
172 |
+
"50277": {
|
173 |
+
"content": " ",
|
174 |
+
"lstrip": false,
|
175 |
+
"normalized": true,
|
176 |
+
"rstrip": false,
|
177 |
+
"single_word": false,
|
178 |
+
"special": false
|
179 |
+
},
|
180 |
+
"50278": {
|
181 |
+
"content": " ",
|
182 |
+
"lstrip": false,
|
183 |
+
"normalized": true,
|
184 |
+
"rstrip": false,
|
185 |
+
"single_word": false,
|
186 |
+
"special": false
|
187 |
+
},
|
188 |
+
"50279": {
|
189 |
+
"content": " ",
|
190 |
+
"lstrip": false,
|
191 |
+
"normalized": true,
|
192 |
+
"rstrip": false,
|
193 |
+
"single_word": false,
|
194 |
+
"special": false
|
195 |
+
},
|
196 |
+
"50280": {
|
197 |
+
"content": " ",
|
198 |
+
"lstrip": false,
|
199 |
+
"normalized": true,
|
200 |
+
"rstrip": false,
|
201 |
+
"single_word": false,
|
202 |
+
"special": false
|
203 |
+
},
|
204 |
+
"50281": {
|
205 |
+
"content": " ",
|
206 |
+
"lstrip": false,
|
207 |
+
"normalized": true,
|
208 |
+
"rstrip": false,
|
209 |
+
"single_word": false,
|
210 |
+
"special": false
|
211 |
+
},
|
212 |
+
"50282": {
|
213 |
+
"content": " ",
|
214 |
+
"lstrip": false,
|
215 |
+
"normalized": true,
|
216 |
+
"rstrip": false,
|
217 |
+
"single_word": false,
|
218 |
+
"special": false
|
219 |
+
},
|
220 |
+
"50283": {
|
221 |
+
"content": " ",
|
222 |
+
"lstrip": false,
|
223 |
+
"normalized": true,
|
224 |
+
"rstrip": false,
|
225 |
+
"single_word": false,
|
226 |
+
"special": false
|
227 |
+
},
|
228 |
+
"50284": {
|
229 |
+
"content": " ",
|
230 |
+
"lstrip": false,
|
231 |
+
"normalized": true,
|
232 |
+
"rstrip": false,
|
233 |
+
"single_word": false,
|
234 |
+
"special": false
|
235 |
+
},
|
236 |
+
"50285": {
|
237 |
+
"content": " ",
|
238 |
+
"lstrip": false,
|
239 |
+
"normalized": true,
|
240 |
+
"rstrip": false,
|
241 |
+
"single_word": false,
|
242 |
+
"special": false
|
243 |
+
},
|
244 |
+
"50286": {
|
245 |
+
"content": " ",
|
246 |
+
"lstrip": false,
|
247 |
+
"normalized": true,
|
248 |
+
"rstrip": false,
|
249 |
+
"single_word": false,
|
250 |
+
"special": false
|
251 |
+
},
|
252 |
+
"50287": {
|
253 |
+
"content": "\t\t\t\t\t\t\t\t\t",
|
254 |
+
"lstrip": false,
|
255 |
+
"normalized": true,
|
256 |
+
"rstrip": false,
|
257 |
+
"single_word": false,
|
258 |
+
"special": false
|
259 |
+
},
|
260 |
+
"50288": {
|
261 |
+
"content": "\t\t\t\t\t\t\t\t",
|
262 |
+
"lstrip": false,
|
263 |
+
"normalized": true,
|
264 |
+
"rstrip": false,
|
265 |
+
"single_word": false,
|
266 |
+
"special": false
|
267 |
+
},
|
268 |
+
"50289": {
|
269 |
+
"content": "\t\t\t\t\t\t\t",
|
270 |
+
"lstrip": false,
|
271 |
+
"normalized": true,
|
272 |
+
"rstrip": false,
|
273 |
+
"single_word": false,
|
274 |
+
"special": false
|
275 |
+
},
|
276 |
+
"50290": {
|
277 |
+
"content": "\t\t\t\t\t\t",
|
278 |
+
"lstrip": false,
|
279 |
+
"normalized": true,
|
280 |
+
"rstrip": false,
|
281 |
+
"single_word": false,
|
282 |
+
"special": false
|
283 |
+
},
|
284 |
+
"50291": {
|
285 |
+
"content": "\t\t\t\t\t",
|
286 |
+
"lstrip": false,
|
287 |
+
"normalized": true,
|
288 |
+
"rstrip": false,
|
289 |
+
"single_word": false,
|
290 |
+
"special": false
|
291 |
+
},
|
292 |
+
"50292": {
|
293 |
+
"content": "\t\t\t\t",
|
294 |
+
"lstrip": false,
|
295 |
+
"normalized": true,
|
296 |
+
"rstrip": false,
|
297 |
+
"single_word": false,
|
298 |
+
"special": false
|
299 |
+
},
|
300 |
+
"50293": {
|
301 |
+
"content": "\t\t\t",
|
302 |
+
"lstrip": false,
|
303 |
+
"normalized": true,
|
304 |
+
"rstrip": false,
|
305 |
+
"single_word": false,
|
306 |
+
"special": false
|
307 |
+
},
|
308 |
+
"50294": {
|
309 |
+
"content": "\t\t",
|
310 |
+
"lstrip": false,
|
311 |
+
"normalized": true,
|
312 |
+
"rstrip": false,
|
313 |
+
"single_word": false,
|
314 |
+
"special": false
|
315 |
+
},
|
316 |
+
"50295": {
|
317 |
+
"content": "<|im_end|>",
|
318 |
+
"lstrip": false,
|
319 |
+
"normalized": false,
|
320 |
+
"rstrip": false,
|
321 |
+
"single_word": false,
|
322 |
+
"special": true
|
323 |
+
},
|
324 |
+
"50296": {
|
325 |
+
"content": "<|im_start|>",
|
326 |
+
"lstrip": false,
|
327 |
+
"normalized": false,
|
328 |
+
"rstrip": false,
|
329 |
+
"single_word": false,
|
330 |
+
"special": false
|
331 |
+
}
|
332 |
+
},
|
333 |
+
"bos_token": "<|endoftext|>",
|
334 |
+
"chat_template": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
|
335 |
+
"clean_up_tokenization_spaces": true,
|
336 |
+
"eos_token": "<|im_end|>",
|
337 |
+
"legacy": false,
|
338 |
+
"model_max_length": 2048,
|
339 |
+
"pad_token": "<|endoftext|>",
|
340 |
+
"return_token_type_ids": false,
|
341 |
+
"tokenizer_class": "CodeGenTokenizer",
|
342 |
+
"unk_token": "<|endoftext|>"
|
343 |
+
}
|
vocab.json
ADDED
The diff for this file is too large to render.
See raw diff
|
|