Raincleared commited on
Commit
c5d0768
0 Parent(s):

Duplicate from SparseLLM/ProSparse-MiniCPM-1B-sft

Browse files
.gitattributes ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,167 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ language:
3
+ - en
4
+ - zh
5
+ tags:
6
+ - MiniCPM
7
+ - ModelBest
8
+ - THUNLP
9
+ license: apache-2.0
10
+ ---
11
+
12
+
13
+ # ProSparse-MiniCPM-1B-sft
14
+
15
+ - Original model: [MiniCPM-1B-sft-bf16](https://huggingface.co/openbmb/MiniCPM-1B-sft-bf16)
16
+ - Model creator and fine-tuned by: [ModelBest](https://modelbest.cn/), [OpenBMB](https://huggingface.co/openbmb), and [THUNLP](https://nlp.csai.tsinghua.edu.cn/)
17
+ - Paper: [link](https://arxiv.org/pdf/2402.13516.pdf)
18
+
19
+ ### Introduction
20
+
21
+ The utilization of activation sparsity, namely the existence of considerable weakly-contributed elements among activation outputs, is a promising method for inference acceleration of large language models (LLMs) ([Liu et al., 2023](https://proceedings.mlr.press/v202/liu23am/liu23am.pdf); [Song et al., 2023](https://arxiv.org/pdf/2312.12456.pdf)). Concretely, acceleration methods based on activation sparsity usually achieve higher inference speed by making wiser resource allocation and computation policies to avoid resource waste on these weakly-contributed parameters.
22
+
23
+ Adopting ReLU as the activation function is a straightforward method to achieve activation sparsity. However, most recent mainstream LLMs adopt activation functions without intrinsic sparsity (e.g., GELU and Swish). Some efforts ([Zhang et al., 2022](https://aclanthology.org/2022.findings-acl.71.pdf); [Mirzadeh et al., 2023](https://arxiv.org/pdf/2310.04564.pdf); [Zhang et al., 2024](https://arxiv.org/pdf/2402.03804.pdf)) introduce ReLU or its variants as the substitutive activation function to help non-ReLU LLMs achieve activation sparsity and inference acceleration, but few can concurrently obtain high sparsity and comparable task-specific performance.
24
+
25
+ In this work, we introduce a simple and effective sparsification method named "ProSparse" to push LLMs for higher activation sparsity while maintaining comparable performance. By applying ProSparse to Swish-activated LLaMA2-7B, LLaMA2-13B, and MiniCPM-1B, we obtain ReLU-activated models with high sparsity of 89.32%, 88.80%, and 87.89%, respectively, while their performance is comparable to the original version. These present the most sparsely activated models among open-source LLaMA versions and competitive end-size models, considerably surpassing ReluLLaMA-7B (66.98%) and ReluLLaMA-13B (71.56%). Further inference acceleration experiments demonstrate the practical speedup effects of higher sparsity on both [PowerInfer](https://arxiv.org/pdf/2312.12456.pdf) and our two sparse GPU [operators](https://github.com/Raincleared-Song/sparse_gpu_operator).
26
+
27
+ ### Training Dataset
28
+
29
+ We train the 1B model on about 473.02 billion tokens within 101,000 steps. These consist of 35,000 steps for standard ProSparse pre-training, 60,000 steps for decay, and 6,000 steps for SFT. Except for ProSparse, other training settings are highly consistent with the original [MiniCPM-1B](https://huggingface.co/openbmb/MiniCPM-1B-sft-bf16). Refer to our [paper](https://arxiv.org/pdf/2402.13516.pdf) and [MiniCPM technical report](https://arxiv.org/pdf/2404.06395) for more details.
30
+
31
+ Intuitively, training the model with even more tokens or with data of a wider coverage and higher quality will obtain better task-specific performance.
32
+
33
+ ### ProSparse: Training Methodology
34
+
35
+ The training process of ProSparse consists of three steps (refer to Section 3.2 of [paper](https://arxiv.org/pdf/2402.13516.pdf) for more details):
36
+
37
+ 1. **Activation Function Substitution**: We substitute the activation function of FFNs with ReLU and apply continual training;
38
+ 2. **Progressive Sparsity Regularization**: We jointly optimize the model on the conventional next-token prediction loss and \\(L_1\\) regularization loss. The regularization is applied to the sparse intermediate outputs of FFNs with a regularization factor increasing progressively in multiple stages. Specifically, the regularization factor \\(\lambda\\) is set to a small constant for the warmup stage, and then increases along a smooth sine curve for each of the subsequent incremental stages. Each stage is accompanied by certain steps of training. In this way, the model can have more time to adapt to the increasing regularization without radical activation shifts, thus alleviating performance degradation.
39
+ 3. **Activation Threshold Shifting**: We finally replace ReLU with FATReLU ([Kurtz et al., 2020](https://proceedings.mlr.press/v119/kurtz20a/kurtz20a.pdf)), a ReLU variant with a positive threshold. This can prune those non-zero weakly-contributed elements in activation outputs and further boost sparsity.
40
+
41
+ The hyper-parameters for each stage (including the regularization factor \\(\lambda_i\\), the accumulated training steps \\(T_i\\), and the accumulated training tokens) are shown as follows:
42
+
43
+ | Step Number \\(i\\) | \\(\lambda_i\\) | \\(T_i\\) | Accumulated Tokens (B) |
44
+ | :-------------: | :---------: | :----: | :--------------------: |
45
+ | 0 | 0 | 10,000 | 49.15 |
46
+ | 1 | \\(1e-3\\) | 15,000 | 73.73 |
47
+ | 2 | \\(5e-3\\) | 20,000 | 98.30 |
48
+ | 3 | \\(5e-3\\) | 25,000 | 122.88 |
49
+ | 4 | \\(5e-2\\) | 35,000 | 172.03 |
50
+ | decay | \\(5e-2\\) (fixed) | 95,000 | 466.94 |
51
+ | SFT | \\(1e-2\\) (fixed) | 101,000 | 473.02 |
52
+
53
+ ### Evaluation Results
54
+
55
+ The evaluation results on the above benchmarks demonstrate the advantage of ProSparse, which is the only method achieving high sparsity and comparable performance to the original Swish-activated LLaMA2. Note that models under all settings are trained with the same number of tokens on the same mixed dataset. Our evaluation is based on the framework [UltraEval](https://github.com/OpenBMB/UltraEval). The evaluation details are listed as follows:
56
+
57
+ - **Code Generation**: We compute the average pass@1 scores on HumanEval (0-shot) and MBPP (3-shot).
58
+
59
+ - **Commonsense Reasoning**: We report the average 0-shot accuracies on PIQA, SIQA, HellaSwag, WinoGrande, and COPA.
60
+
61
+ - **Reading Comprehension**: We compute the average 0-shot accuracies on BoolQ, LAMBADA, and TyDi QA.
62
+
63
+ - **Other Popular Benchmarks**: We report the average accuracies on GSM8K (8-shot), MMLU (5-shot), Big Bench Hard (BBH) (3-shot), and AGI-Eval (0-shot).
64
+
65
+ **Notes**: For PIQA, SIQA, HellaSwag, WinoGrande, COPA, BoolQ, LAMBADA, TyDi QA, and AGI-Eval, we obtain the predicted answers based on maximized perplexity. For GSM8K, MMLU, and BBH, the predicted answers are directly generated.
66
+
67
+ | Setting | Average<br>Sparsity | Average<br>Performance | Code<br>Generation | Commonsense<br>Reasoning | Reading<br>Comprehension | GSM8K | MMLU | BBH | AGI Eval |
68
+ | :-------------------: | :----------------: | :----------------------: | :----------------------: | :---: | :---: | :---: | :---------: | :-----: | :-----------------: |
69
+ | LLaMA2-7B | - | 37.96 | 16.37 | 69.59 | 61.87 | 12.96 | 44.45 | 32.96 | 27.53 |
70
+ | ReluLLaMA-7B | 66.98 | 37.62 | 15.85 | 69.64 | 70.54 | 5.84 | 38.64 | 35.07 | 27.73 |
71
+ | **ProSparse-7B**\* | 88.11 | 38.31 | 19.47 | 66.29 | 63.33 | 12.74 | 45.21 | 33.59 | 27.55 |
72
+ | **ProSparse-7B** | **89.32** | **38.46** | 19.42 | 66.27 | 63.50 | 12.13 | 45.48 | 34.99 | 27.46 |
73
+ | LLaMA2-13B | - | 44.06 | 20.19 | 72.58 | 71.55 | 22.21 | 54.69 | 37.89 | 29.33 |
74
+ | ReluLLaMA-13B | 71.56 | 42.74 | 20.19 | 70.44 | 73.29 | 18.50 | 50.58 | 37.97 | 28.22 |
75
+ | **ProSparse-13B**\* | 87.97 | **45.07** | 29.03 | 69.75 | 67.54 | 25.40 | 54.78 | 40.20 | 28.76 |
76
+ | **ProSparse-13B** | **88.80** | 44.90 | 28.42 | 69.76 | 66.91 | 26.31 | 54.35 | 39.90 | 28.67 |
77
+ | MiniCPM-1B | - | 44.44 | 36.85 | 63.67 | 60.90 | 35.48 | 50.44 | 35.03 | 28.71 |
78
+ | **ProSparse-1B**\* | 86.25 | **44.72** | 41.38 | 64.55 | 60.69 | 34.72 | 49.36 | 34.04 | 28.27 |
79
+ | **ProSparse-1B** | **87.89** | **44.72** | 42.04 | 64.37 | 60.73 | 34.57 | 49.51 | 34.08 | 27.77 |
80
+
81
+ **Notes**: "Original" refers to the original Swish-activated LLaMA2 versions. ReluLLaMA-7B and ReluLLaMA-13B are available at [7B](https://huggingface.co/SparseLLM/ReluLLaMA-7B) and [13B](https://huggingface.co/SparseLLM/ReluLLaMA-13B) respectively. MiniCPM-1B is available at [1B](https://huggingface.co/openbmb/MiniCPM-1B-sft-bf16). "ProSparse-7B\*", "ProSparse-13B\*", and "ProSparse-1B\*" denote the ProSparse versions without activation threshold shifting.
82
+
83
+ ### Evaluation Issues with LM-Eval
84
+
85
+ The above results can be replicated with [UltraEval](https://github.com/OpenBMB/UltraEval). Some abnormal results obtained with other popular frameworks such as [LM-Eval](https://github.com/EleutherAI/lm-evaluation-harness) are probably attributed to the absence of the cls token `<s>`, which is not added by default in LM-Eval. A quick temporary fix is shown in the following codes. Other differences in evaluation results may be caused by other reasons, including the few-shot settings, data pre-processing, and extra prompts.
86
+
87
+ ```python
88
+ # https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/models/huggingface.py#L945
89
+ for _, context_enc, continuation_enc in chunk:
90
+ # sanity check
91
+ assert len(context_enc) > 0
92
+ # Note: a trivial fix here
93
+ if context_enc[0] != 1:
94
+ context_enc = [1] + context_enc
95
+ assert len(continuation_enc) > 0
96
+ assert len(continuation_enc) <= self.max_length
97
+ ```
98
+
99
+ Here are the steps to adapting the original [vLLM](https://github.com/vllm-project/vllm) to ProSparse LLaMA models.
100
+
101
+ 1. Replace the file [vllm/model_executor/models/llama.py](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/llama.py) in original vLLM with this [file](https://github.com/Raincleared-Song/DejaVu_predictor/blob/main/llama.py).
102
+ 2. Replace the contents of the original [config.json](https://huggingface.co/SparseLLM/prosparse-llama-2-7b/blob/main/config.json) with this [file](https://github.com/Raincleared-Song/DejaVu_predictor/blob/main/config.json).
103
+ 3. Set the environment variable `ACT_INFO`. To test the version without activation threshold shifting, `export ACT_INFO=relu`. To test the version with activation threshold shifting, `export ACT_INFO=fatrelu_0.01`.
104
+
105
+ ### Inference Acceleration Effects
106
+
107
+ First, we utilize [PowerInfer](https://arxiv.org/pdf/2312.12456.pdf), a state-of-the-art acceleration framework leveraging activation sparsity. As its inference speed and accuracy heavily rely on the performance of activation predictors, we report the activation recall and predicted sparsity (i.e., two key metrics for evaluating the activation predictor) as well as the number of tokens generated per second by PowerInfer (with one A100 GPU and sufficient CPUs). The GGUF files and activation predictors are also available for ProSparse LLaMA models.
108
+
109
+ Moreover, considering the potential inference inaccuracies caused by wrong predictions of activation predictors, we implement two sparse GPU [operators](https://github.com/Raincleared-Song/sparse_gpu_operator) for faster accurate inference utilizing activation sparsity. They are responsible for the speedup of two key steps in a gated FFN:
110
+
111
+ - Step (2) (`S2`): a fused operator of ReLU and \\(\mathbf{s} \odot (\mathbf{x} \mathbf{W}_1^T)\\);
112
+ - Step (3) (`S3`): a sparse matrix-vector multiplication operator \\(\mathbf{x}_1 \mathbf{W}_2^T\\).
113
+
114
+ where \\(\mathbf{s}\\), \\(\mathbf{x}\\), \\(\mathbf{x}_1\\), and \\(\odot\\) denote the gating scores, the FFN input hidden states, the intermediate outputs, and the element-wise multiplication respectively. \\(\mathbf{W}_1\\) and \\(\mathbf{W}_2\\) are FFN weight matrices.
115
+
116
+ The acceleration effects of LLMs with different sparsity are displayed as follows. ProSparse, which reaches a high sparsity without performance degradation, can gain the most benefits among all the settings concerned. Refer to Section 4.3 of [paper](https://arxiv.org/pdf/2402.13516.pdf) for more details.
117
+
118
+ | Setting | Average<br>Sparsity | Activation<br>Recall | Predicted<br>Sparsity | PowerInfer<br>Speed | Speedup<br>to Dense | `S2`<br>Time | Speedup<br>to Dense | `S3`<br/>Time | Speedup<br/>to Dense |
119
+ | :-------------------: | :-----------------: | :------------------: | :-------------------: | :-----------------: | :-----------------: | :--------------: | :-----------------: | :---------------: | :------------------: |
120
+ | Dense-7B | - | - | - | 3.67 | 1.00 | 90.55 | 1.00 | 82.92 | 1.00 |
121
+ | ReluLLaMA-7B | 66.98 | 90.89 | 58.95 | 11.37 | 3.10 | 67.12 | 1.35 | 63.00 | 1.32 |
122
+ | **ProSparse-7B**\* | 88.11 | **93.46** | 75.24 | **16.30** | **4.44** | 46.66 | 1.94 | 55.56 | 1.49 |
123
+ | **ProSparse-7B** | **89.32** | 92.34 | **78.75** | - | - | **45.38** | **2.00** | **55.05** | **1.51** |
124
+ | Dense-13B | - | - | - | 1.92 | 1.00 | 131.36 | 1.00 | 113.68 | 1.00 |
125
+ | ReluLLaMA-13B | 71.56 | 86.41 | 71.93 | 6.59 | 3.43 | 69.92 | 1.88 | 75.47 | 1.51 |
126
+ | **ProSparse-13B**\* | 87.97 | 91.02 | 77.93 | **8.67** | **4.52** | 55.29 | 2.38 | 67.50 | 1.68 |
127
+ | **ProSparse-13B** | **88.80** | **91.11** | **78.28** | - | - | **53.78** | **2.44** | **66.73** | **1.70** |
128
+
129
+ **Notes**: For "Dense" settings, the "Inference Speed" (token/sec) is obtained by [llama.cpp](https://github.com/ggerganov/llama.cpp), and the time (us) for steps (2) and (3) is measured without sparse GPU operators. For other sparse settings, the "Inference Speed" is obtained by [PowerInfer](https://arxiv.org/pdf/2312.12456.pdf), and sparse GPU operators are applied. ProSparse settings with activation threshold shifting and the MiniCPM architecture are not supported by PowerInfer at present.
130
+
131
+ ### Citation
132
+
133
+ Please kindly cite using the following BibTeX:
134
+
135
+ ```bibtex
136
+ @article{song2024prosparse,
137
+ title={{ProSparse}: Introducing and Enhancing Intrinsic Activation Sparsity within Large Language Models},
138
+ author={Song, Chenyang and Han, Xu and Zhang, Zhengyan and Hu, Shengding and Shi, Xiyu and Li, Kuai and Chen, Chen and Liu, Zhiyuan and Li, Guangli and Yang, Tao and Sun, Maosong},
139
+ year={2024},
140
+ journal={arXiv preprint arXiv:2402.13516},
141
+ url={https://arxiv.org/pdf/2402.13516.pdf}
142
+ }
143
+ ```
144
+
145
+ ### License
146
+
147
+ This repository is released under the [Apache-2.0](https://github.com/OpenBMB/MiniCPM/blob/main/LICENSE) License.
148
+
149
+ The usage of MiniCPM model weights must strictly follow [the General Model License (GML)](https://github.com/OpenBMB/General-Model-License/blob/main/%E9%80%9A%E7%94%A8%E6%A8%A1%E5%9E%8B%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE-%E6%9D%A5%E6%BA%90%E8%AF%B4%E6%98%8E-%E5%AE%A3%E4%BC%A0%E9%99%90%E5%88%B6-%E5%95%86%E4%B8%9A%E6%8E%88%E6%9D%83.md).
150
+
151
+ The models and weights of MiniCPM are completely free for academic research.
152
+
153
+ If you intend to utilize the model for commercial purposes, please reach out to cpm@modelbest.cn to obtain the certificate of authorization.
154
+
155
+ ### Statement
156
+
157
+ As a language model, MiniCPM generates content by learning from a vast amount of text.
158
+
159
+ However, it does not possess the ability to comprehend or express personal opinions or value judgments.
160
+
161
+ Any content generated by MiniCPM does not represent the viewpoints or positions of the model developers.
162
+
163
+ Therefore, when using content generated by MiniCPM, users should take full responsibility for evaluating and verifying it on their own.
164
+
165
+ #### Acknowledgments
166
+
167
+ The model card is modified from [ReluLLaMA-7B](https://huggingface.co/SparseLLM/ReluLLaMA-7B) and [MiniCPM-1B](https://huggingface.co/openbmb/MiniCPM-1B-sft-bf16).
config.json ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "SparseLLM/ProSparse-MiniCPM-1B-sft",
3
+ "architectures": [
4
+ "MiniCPMForCausalLM"
5
+ ],
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_minicpm.MiniCPMConfig",
8
+ "AutoModel": "modeling_minicpm.MiniCPMModel",
9
+ "AutoModelForCausalLM": "modeling_minicpm.MiniCPMForCausalLM",
10
+ "AutoModelForSeq2SeqLM": "modeling_minicpm.MiniCPMForCausalLM",
11
+ "AutoModelForSequenceClassification": "modeling_minicpm.MiniCPMForSequenceClassification"
12
+ },
13
+ "bos_token_id": 1,
14
+ "eos_token_id": 2,
15
+ "hidden_act": "fatrelu",
16
+ "hidden_act_param": 0.03,
17
+ "hidden_size": 1536,
18
+ "initializer_range": 0.1,
19
+ "intermediate_size": 3840,
20
+ "max_position_embeddings": 4096,
21
+ "num_attention_heads": 24,
22
+ "num_hidden_layers": 52,
23
+ "num_key_value_heads": 8,
24
+ "rms_norm_eps": 1e-05,
25
+ "rope_scaling": null,
26
+ "torch_dtype": "bfloat16",
27
+ "transformers_version": "4.36.0",
28
+ "use_cache": true,
29
+ "vocab_size": 73440,
30
+ "scale_emb": 12,
31
+ "dim_model_base": 256,
32
+ "scale_depth": 1.4
33
+ }
configuration_minicpm.py ADDED
@@ -0,0 +1,206 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """ MiniCPM model configuration"""
21
+
22
+ from transformers.configuration_utils import PretrainedConfig
23
+ from transformers.utils import logging
24
+
25
+
26
+ logger = logging.get_logger(__name__)
27
+
28
+ MINICPM_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
29
+
30
+
31
+ class MiniCPMConfig(PretrainedConfig):
32
+ r"""
33
+ This is the configuration class to store the configuration of a [`MiniCPMModel`]. It is used to instantiate an MiniCPM
34
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
35
+ defaults will yield a similar configuration to that of the MiniCPM-7B.
36
+
37
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
38
+ documentation from [`PretrainedConfig`] for more information.
39
+
40
+
41
+ Args:
42
+ vocab_size (`int`, *optional*, defaults to 32000):
43
+ Vocabulary size of the MiniCPM model. Defines the number of different tokens that can be represented by the
44
+ `inputs_ids` passed when calling [`MiniCPMModel`]
45
+ hidden_size (`int`, *optional*, defaults to 4096):
46
+ Dimension of the hidden representations.
47
+ intermediate_size (`int`, *optional*, defaults to 11008):
48
+ Dimension of the MLP representations.
49
+ num_hidden_layers (`int`, *optional*, defaults to 32):
50
+ Number of hidden layers in the Transformer decoder.
51
+ num_attention_heads (`int`, *optional*, defaults to 32):
52
+ Number of attention heads for each attention layer in the Transformer decoder.
53
+ num_key_value_heads (`int`, *optional*):
54
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
55
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
56
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
57
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
58
+ by meanpooling all the original heads within that group. For more details checkout [this
59
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
60
+ `num_attention_heads`.
61
+ hidden_act (`str` or `function`, *optional*, defaults to `"relu"`):
62
+ The non-linear activation function (function or string) in the decoder.
63
+ hidden_act_param (`float`, *optional*, defaults to 0.):
64
+ The bias for shiftrelu or threshold for fatrelu.
65
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
66
+ The maximum sequence length that this model might ever be used with. MiniCPM 1 supports up to 2048 tokens,
67
+ MiniCPM 2 up to 4096, CodeMiniCPM up to 16384.
68
+ initializer_range (`float`, *optional*, defaults to 0.02):
69
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
70
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
71
+ The epsilon used by the rms normalization layers.
72
+ use_cache (`bool`, *optional*, defaults to `True`):
73
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
74
+ relevant if `config.is_decoder=True`.
75
+ pad_token_id (`int`, *optional*):
76
+ Padding token id.
77
+ bos_token_id (`int`, *optional*, defaults to 1):
78
+ Beginning of stream token id.
79
+ eos_token_id (`int`, *optional*, defaults to 2):
80
+ End of stream token id.
81
+ pretraining_tp (`int`, *optional*, defaults to 1):
82
+ Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
83
+ document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
84
+ necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
85
+ issue](https://github.com/pytorch/pytorch/issues/76232).
86
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
87
+ Whether to tie weight embeddings
88
+ rope_theta (`float`, *optional*, defaults to 10000.0):
89
+ The base period of the RoPE embeddings.
90
+ rope_scaling (`Dict`, *optional*):
91
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
92
+ strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
93
+ `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
94
+ `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
95
+ these scaling strategies behave:
96
+ https://www.reddit.com/r/LocalMiniCPM/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
97
+ experimental feature, subject to breaking API changes in future versions.
98
+ attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
99
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
100
+ attention_dropout (`float`, *optional*, defaults to 0.0):
101
+ The dropout ratio for the attention probabilities.
102
+
103
+ ```python
104
+ >>> from transformers import MiniCPMModel, MiniCPMConfig
105
+
106
+ >>> # Initializing a MiniCPM minicpm-7b style configuration
107
+ >>> configuration = MiniCPMConfig()
108
+
109
+ >>> # Initializing a model from the minicpm-7b style configuration
110
+ >>> model = MiniCPMModel(configuration)
111
+
112
+ >>> # Accessing the model configuration
113
+ >>> configuration = model.config
114
+ ```"""
115
+
116
+ model_type = "minicpm"
117
+ keys_to_ignore_at_inference = ["past_key_values"]
118
+
119
+ def __init__(
120
+ self,
121
+ vocab_size=32000,
122
+ hidden_size=4096,
123
+ intermediate_size=11008,
124
+ num_hidden_layers=32,
125
+ num_attention_heads=32,
126
+ num_key_value_heads=None,
127
+ hidden_act="relu",
128
+ hidden_act_param=0.,
129
+ max_position_embeddings=2048,
130
+ initializer_range=0.02,
131
+ rms_norm_eps=1e-6,
132
+ use_cache=True,
133
+ pad_token_id=None,
134
+ bos_token_id=1,
135
+ eos_token_id=2,
136
+ pretraining_tp=1,
137
+ tie_word_embeddings=True,
138
+ rope_theta=10000.0,
139
+ rope_scaling=None,
140
+ attention_bias=False,
141
+ attention_dropout=0.0,
142
+ scale_emb=1,
143
+ dim_model_base=1,
144
+ scale_depth=1,
145
+ **kwargs,
146
+ ):
147
+ self.vocab_size = vocab_size
148
+ self.max_position_embeddings = max_position_embeddings
149
+ self.hidden_size = hidden_size
150
+ self.intermediate_size = intermediate_size
151
+ self.num_hidden_layers = num_hidden_layers
152
+ self.num_attention_heads = num_attention_heads
153
+
154
+ # for backward compatibility
155
+ if num_key_value_heads is None:
156
+ num_key_value_heads = num_attention_heads
157
+
158
+ self.num_key_value_heads = num_key_value_heads
159
+ self.hidden_act = hidden_act
160
+ self.hidden_act_param = hidden_act_param
161
+ self.initializer_range = initializer_range
162
+ self.rms_norm_eps = rms_norm_eps
163
+ self.pretraining_tp = pretraining_tp
164
+ self.use_cache = use_cache
165
+ self.rope_theta = rope_theta
166
+ self.rope_scaling = rope_scaling
167
+ self._rope_scaling_validation()
168
+ self.attention_bias = attention_bias
169
+ self.attention_dropout = attention_dropout
170
+ self.scale_emb = scale_emb
171
+ self.dim_model_base = dim_model_base
172
+ self.scale_depth = scale_depth
173
+
174
+ super().__init__(
175
+ pad_token_id=pad_token_id,
176
+ bos_token_id=bos_token_id,
177
+ eos_token_id=eos_token_id,
178
+ tie_word_embeddings=tie_word_embeddings,
179
+ **kwargs,
180
+ )
181
+ try:
182
+ import flash_attn
183
+ self._attn_implementation = "flash_attention_2"
184
+ except:
185
+ pass
186
+
187
+ def _rope_scaling_validation(self):
188
+ """
189
+ Validate the `rope_scaling` configuration.
190
+ """
191
+ if self.rope_scaling is None:
192
+ return
193
+
194
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
195
+ raise ValueError(
196
+ "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
197
+ f"got {self.rope_scaling}"
198
+ )
199
+ rope_scaling_type = self.rope_scaling.get("type", None)
200
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
201
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
202
+ raise ValueError(
203
+ f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
204
+ )
205
+ if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
206
+ raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
generation_config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "do_sample": true,
3
+ "top_p": 0.8,
4
+ "temperature": 0.8,
5
+ "bos_token_id": 1,
6
+ "eos_token_id": 2
7
+ }
modeling_minicpm.py ADDED
@@ -0,0 +1,1467 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """ PyTorch MiniCPM model."""
21
+ import math
22
+ import warnings
23
+ from typing import List, Optional, Tuple, Union, Dict
24
+
25
+ import torch
26
+ import torch.nn.functional as F
27
+ import torch.utils.checkpoint
28
+ from torch import nn
29
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
30
+
31
+ from transformers.activations import ACT2FN
32
+ from transformers.cache_utils import Cache, DynamicCache
33
+ from transformers.modeling_attn_mask_utils import (
34
+ AttentionMaskConverter,
35
+ _prepare_4d_attention_mask,
36
+ _prepare_4d_causal_attention_mask,
37
+ _prepare_4d_causal_attention_mask_for_sdpa,
38
+ )
39
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
40
+ from transformers.modeling_utils import PreTrainedModel
41
+ from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
42
+ from transformers.utils import (
43
+ add_start_docstrings,
44
+ add_start_docstrings_to_model_forward,
45
+ is_flash_attn_2_available,
46
+ is_flash_attn_greater_or_equal_2_10,
47
+ logging,
48
+ replace_return_docstrings,
49
+ )
50
+ from transformers.utils.import_utils import is_torch_fx_available
51
+ from .configuration_minicpm import MiniCPMConfig
52
+ import re
53
+
54
+ try:
55
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
56
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
57
+ except:
58
+ pass
59
+
60
+
61
+ # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
62
+ # It means that the function will not be traced through and simply appear as a node in the graph.
63
+ if is_torch_fx_available():
64
+ if not is_torch_greater_or_equal_than_1_13:
65
+ import torch.fx
66
+
67
+ _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
68
+
69
+
70
+ logger = logging.get_logger(__name__)
71
+
72
+ _CONFIG_FOR_DOC = "MiniCPMConfig"
73
+
74
+
75
+ def _get_unpad_data(attention_mask):
76
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
77
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
78
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
79
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
80
+ return (
81
+ indices,
82
+ cu_seqlens,
83
+ max_seqlen_in_batch,
84
+ )
85
+
86
+
87
+ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
88
+ warnings.warn(
89
+ "Calling `transformers.models.minicpm.modeling_minicpm._prepare_4d_attention_mask` is deprecated and will be removed in v4.37. Use `transformers.modeling_attn_mask_utils._prepare_4d_attention_mask"
90
+ )
91
+ return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
92
+
93
+
94
+ def _make_causal_mask(
95
+ input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
96
+ ):
97
+ warnings.warn(
98
+ "Calling `transformers.models.minicpm.modeling_minicpm._make_causal_mask` is deprecated and will be removed in v4.37. Use `transformers.models.minicpm.modeling_minicpm.AttentionMaskConverter._make_causal_mask"
99
+ )
100
+ return AttentionMaskConverter._make_causal_mask(
101
+ input_ids_shape=input_ids_shape, dtype=dtype, device=device, past_key_values_length=past_key_values_length
102
+ )
103
+
104
+ # @torch.jit.script # type: ignore
105
+ def rms_layernorm(hidden: torch.Tensor, weight: torch.Tensor, eps: float):
106
+ old_dtype = hidden.dtype
107
+ variance = hidden.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
108
+ hidden = (hidden * torch.rsqrt(variance + eps)).to(old_dtype)
109
+ return hidden * weight
110
+
111
+
112
+ class MiniCPMRMSNorm(nn.Module):
113
+ def __init__(self, hidden_size, eps=1e-6):
114
+ """
115
+ MiniCPMRMSNorm is equivalent to T5LayerNorm
116
+ """
117
+ super().__init__()
118
+ self.weight = nn.Parameter(torch.ones(hidden_size))
119
+ self.variance_epsilon = eps
120
+
121
+ def forward(self, hidden_states):
122
+ return rms_layernorm(hidden_states, self.weight, self.variance_epsilon)
123
+
124
+
125
+ ALL_LAYERNORM_LAYERS.append(MiniCPMRMSNorm)
126
+
127
+
128
+ class MiniCPMRotaryEmbedding(nn.Module):
129
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
130
+ super().__init__()
131
+
132
+ self.dim = dim
133
+ self.max_position_embeddings = max_position_embeddings
134
+ self.base = base
135
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
136
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
137
+
138
+ # Build here to make `torch.jit.trace` work.
139
+ self._set_cos_sin_cache(
140
+ # seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
141
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.float32
142
+ )
143
+
144
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
145
+ self.max_seq_len_cached = seq_len
146
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
147
+ freqs = torch.outer(t, self.inv_freq)
148
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
149
+ emb = torch.cat((freqs, freqs), dim=-1)
150
+
151
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
152
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
153
+
154
+ def forward(self, x, seq_len=None):
155
+ # x: [bs, num_attention_heads, seq_len, head_size]
156
+ if seq_len > self.max_seq_len_cached:
157
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
158
+
159
+ return (
160
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
161
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
162
+ )
163
+
164
+
165
+ class MiniCPMLinearScalingRotaryEmbedding(MiniCPMRotaryEmbedding):
166
+ """MiniCPMRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
167
+
168
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
169
+ self.scaling_factor = scaling_factor
170
+ super().__init__(dim, max_position_embeddings, base, device)
171
+
172
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
173
+ self.max_seq_len_cached = seq_len
174
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
175
+ t = t / self.scaling_factor
176
+
177
+ freqs = torch.outer(t, self.inv_freq)
178
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
179
+ emb = torch.cat((freqs, freqs), dim=-1)
180
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
181
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
182
+
183
+
184
+ class MiniCPMDynamicNTKScalingRotaryEmbedding(MiniCPMRotaryEmbedding):
185
+ """MiniCPMRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
186
+
187
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
188
+ self.scaling_factor = scaling_factor
189
+ super().__init__(dim, max_position_embeddings, base, device)
190
+
191
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
192
+ self.max_seq_len_cached = seq_len
193
+
194
+ if seq_len > self.max_position_embeddings:
195
+ base = self.base * (
196
+ (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
197
+ ) ** (self.dim / (self.dim - 2))
198
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
199
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
200
+
201
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
202
+
203
+ freqs = torch.outer(t, self.inv_freq)
204
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
205
+ emb = torch.cat((freqs, freqs), dim=-1)
206
+
207
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
208
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
209
+
210
+
211
+ def rotate_half(x):
212
+ """Rotates half the hidden dims of the input."""
213
+ x1 = x[..., : x.shape[-1] // 2]
214
+ x2 = x[..., x.shape[-1] // 2 :]
215
+ return torch.cat((-x2, x1), dim=-1)
216
+
217
+
218
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
219
+ """Applies Rotary Position Embedding to the query and key tensors.
220
+
221
+ Args:
222
+ q (`torch.Tensor`): The query tensor.
223
+ k (`torch.Tensor`): The key tensor.
224
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
225
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
226
+ position_ids (`torch.Tensor`):
227
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
228
+ used to pass offsetted position ids when working with a KV-cache.
229
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
230
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
231
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
232
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
233
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
234
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
235
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
236
+ Returns:
237
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
238
+ """
239
+ # cos = cos[position_ids].unsqueeze(unsqueeze_dim)
240
+ # sin = sin[position_ids].unsqueeze(unsqueeze_dim)
241
+ # q_embed = (q * cos) + (rotate_half(q) * sin)
242
+ # k_embed = (k * cos) + (rotate_half(k) * sin)
243
+ orig_dtype = k.dtype
244
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]
245
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]
246
+ q_fp32 = q.to(dtype=torch.float32, device=q.device)
247
+ k_fp32 = k.to(dtype=torch.float32, device=k.device)
248
+ q_embed = (q_fp32 * cos) + (rotate_half(q_fp32) * sin)
249
+ k_embed = (k_fp32 * cos) + (rotate_half(k_fp32) * sin)
250
+ return q_embed.to(dtype=orig_dtype), k_embed.to(dtype=orig_dtype)
251
+
252
+ class MiniCPMMLP(nn.Module):
253
+ def __init__(self, config):
254
+ super().__init__()
255
+ self.config = config
256
+ self.hidden_size = config.hidden_size
257
+ self.intermediate_size = config.intermediate_size
258
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
259
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
260
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
261
+ if config.hidden_act in ACT2FN:
262
+ self.act_fn = ACT2FN[config.hidden_act]
263
+ elif config.hidden_act == "shiftrelu":
264
+ def shifted_relu(x):
265
+ return torch.nn.functional.relu(x - config.hidden_act_param)
266
+ self.act_fn = shifted_relu
267
+ elif config.hidden_act == "fatrelu":
268
+ def fat_relu(x):
269
+ new_x = torch.zeros_like(x)
270
+ mask = torch.ge(x, config.hidden_act_param)
271
+ new_x[mask] = x[mask]
272
+ return new_x
273
+ self.act_fn = fat_relu
274
+ else:
275
+ raise NotImplementedError(f"Unsupported activation function: {config.hidden_act}")
276
+
277
+ def forward(self, x):
278
+ if self.config.pretraining_tp > 1:
279
+ slice = self.intermediate_size // self.config.pretraining_tp
280
+ gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)
281
+ up_proj_slices = self.up_proj.weight.split(slice, dim=0)
282
+ down_proj_slices = self.down_proj.weight.split(slice, dim=1)
283
+
284
+ gate_proj = torch.cat(
285
+ [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1
286
+ )
287
+ up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)
288
+
289
+ intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)
290
+ down_proj = [
291
+ F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)
292
+ ]
293
+ down_proj = sum(down_proj)
294
+ else:
295
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
296
+
297
+ return down_proj
298
+
299
+
300
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
301
+ """
302
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
303
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
304
+ """
305
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
306
+ if n_rep == 1:
307
+ return hidden_states
308
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
309
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
310
+
311
+
312
+
313
+ class MiniCPMAttention(nn.Module):
314
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
315
+
316
+ def __init__(self, config: MiniCPMConfig, layer_idx: Optional[int] = None):
317
+ super().__init__()
318
+ self.config = config
319
+ self.layer_idx = layer_idx
320
+ if layer_idx is None:
321
+ logger.warning_once(
322
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
323
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
324
+ "when creating this class."
325
+ )
326
+
327
+ self.attention_dropout = config.attention_dropout
328
+ self.hidden_size = config.hidden_size
329
+ self.num_heads = config.num_attention_heads
330
+ self.head_dim = self.hidden_size // self.num_heads
331
+ self.num_key_value_heads = config.num_key_value_heads
332
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
333
+ self.max_position_embeddings = config.max_position_embeddings
334
+ self.rope_theta = config.rope_theta
335
+ self.is_causal = True
336
+
337
+ if (self.head_dim * self.num_heads) != self.hidden_size:
338
+ raise ValueError(
339
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
340
+ f" and `num_heads`: {self.num_heads})."
341
+ )
342
+
343
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
344
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
345
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
346
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
347
+ self._init_rope()
348
+
349
+ def _init_rope(self):
350
+ if self.config.rope_scaling is None:
351
+ self.rotary_emb = MiniCPMRotaryEmbedding(
352
+ self.head_dim,
353
+ max_position_embeddings=self.max_position_embeddings,
354
+ base=self.rope_theta,
355
+ )
356
+ else:
357
+ scaling_type = self.config.rope_scaling["type"]
358
+ scaling_factor = self.config.rope_scaling["factor"]
359
+ if scaling_type == "linear":
360
+ self.rotary_emb = MiniCPMLinearScalingRotaryEmbedding(
361
+ self.head_dim,
362
+ max_position_embeddings=self.max_position_embeddings,
363
+ scaling_factor=scaling_factor,
364
+ base=self.rope_theta,
365
+ )
366
+ elif scaling_type == "dynamic":
367
+ self.rotary_emb = MiniCPMDynamicNTKScalingRotaryEmbedding(
368
+ self.head_dim,
369
+ max_position_embeddings=self.max_position_embeddings,
370
+ scaling_factor=scaling_factor,
371
+ base=self.rope_theta,
372
+ )
373
+ else:
374
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
375
+
376
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
377
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
378
+
379
+ def forward(
380
+ self,
381
+ hidden_states: torch.Tensor,
382
+ attention_mask: Optional[torch.Tensor] = None,
383
+ position_ids: Optional[torch.LongTensor] = None,
384
+ past_key_value: Optional[Cache] = None,
385
+ output_attentions: bool = False,
386
+ use_cache: bool = False,
387
+ **kwargs,
388
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
389
+ if "padding_mask" in kwargs:
390
+ warnings.warn(
391
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
392
+ )
393
+
394
+ bsz, q_len, _ = hidden_states.size()
395
+
396
+ if self.config.pretraining_tp > 1:
397
+ key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
398
+ query_slices = self.q_proj.weight.split(
399
+ (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
400
+ )
401
+ key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
402
+ value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
403
+
404
+ query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]
405
+ query_states = torch.cat(query_states, dim=-1)
406
+
407
+ key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
408
+ key_states = torch.cat(key_states, dim=-1)
409
+
410
+ value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
411
+ value_states = torch.cat(value_states, dim=-1)
412
+
413
+ else:
414
+ query_states = self.q_proj(hidden_states)
415
+ key_states = self.k_proj(hidden_states)
416
+ value_states = self.v_proj(hidden_states)
417
+
418
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
419
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
420
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
421
+
422
+ kv_seq_len = key_states.shape[-2]
423
+ if past_key_value is not None:
424
+ if self.layer_idx is None:
425
+ raise ValueError(
426
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
427
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
428
+ "with a layer index."
429
+ )
430
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
431
+ cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)
432
+
433
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
434
+
435
+ if past_key_value is not None:
436
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
437
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
438
+
439
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
440
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
441
+
442
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
443
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
444
+ raise ValueError(
445
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
446
+ f" {attn_weights.size()}"
447
+ )
448
+
449
+ if attention_mask is not None:
450
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
451
+ raise ValueError(
452
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
453
+ )
454
+ attn_weights = attn_weights + attention_mask
455
+
456
+ # upcast attention to fp32
457
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
458
+ attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
459
+ attn_output = torch.matmul(attn_weights, value_states)
460
+
461
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
462
+ raise ValueError(
463
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
464
+ f" {attn_output.size()}"
465
+ )
466
+
467
+ attn_output = attn_output.transpose(1, 2).contiguous()
468
+
469
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
470
+
471
+ if self.config.pretraining_tp > 1:
472
+ attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
473
+ o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
474
+ attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
475
+ else:
476
+ attn_output = self.o_proj(attn_output)
477
+
478
+ if not output_attentions:
479
+ attn_weights = None
480
+
481
+ return attn_output, attn_weights, past_key_value
482
+
483
+
484
+ class MiniCPMFlashAttention2(MiniCPMAttention):
485
+ """
486
+ MiniCPM flash attention module. This module inherits from `MiniCPMAttention` as the weights of the module stays
487
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
488
+ flash attention and deal with padding tokens in case the input contains any of them.
489
+ """
490
+
491
+ def __init__(self, *args, **kwargs):
492
+ super().__init__(*args, **kwargs)
493
+
494
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
495
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
496
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
497
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
498
+
499
+ def forward(
500
+ self,
501
+ hidden_states: torch.Tensor,
502
+ attention_mask: Optional[torch.LongTensor] = None,
503
+ position_ids: Optional[torch.LongTensor] = None,
504
+ past_key_value: Optional[Cache] = None,
505
+ output_attentions: bool = False,
506
+ use_cache: bool = False,
507
+ **kwargs,
508
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
509
+ # MiniCPMFlashAttention2 attention does not support output_attentions
510
+ if "padding_mask" in kwargs:
511
+ warnings.warn(
512
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
513
+ )
514
+
515
+ # overwrite attention_mask with padding_mask
516
+ attention_mask = kwargs.pop("padding_mask")
517
+
518
+ output_attentions = False
519
+
520
+ bsz, q_len, _ = hidden_states.size()
521
+
522
+ query_states = self.q_proj(hidden_states)
523
+ key_states = self.k_proj(hidden_states)
524
+ value_states = self.v_proj(hidden_states)
525
+
526
+ # Flash attention requires the input to have the shape
527
+ # batch_size x seq_length x head_dim x hidden_dim
528
+ # therefore we just need to keep the original shape
529
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
530
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
531
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
532
+
533
+ kv_seq_len = key_states.shape[-2]
534
+ if past_key_value is not None:
535
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
536
+ cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)
537
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
538
+
539
+ if past_key_value is not None:
540
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
541
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
542
+
543
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
544
+ # to be able to avoid many of these transpose/reshape/view.
545
+ query_states = query_states.transpose(1, 2)
546
+ key_states = key_states.transpose(1, 2)
547
+ value_states = value_states.transpose(1, 2)
548
+
549
+ dropout_rate = self.attention_dropout if self.training else 0.0
550
+
551
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
552
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
553
+ # cast them back in the correct dtype just to be sure everything works as expected.
554
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
555
+ # in fp32. (MiniCPMRMSNorm handles it correctly)
556
+
557
+ input_dtype = query_states.dtype
558
+ if input_dtype == torch.float32:
559
+ # Handle the case where the model is quantized
560
+ if hasattr(self.config, "_pre_quantization_dtype"):
561
+ target_dtype = self.config._pre_quantization_dtype
562
+ else:
563
+ target_dtype = self.q_proj.weight.dtype
564
+
565
+ logger.warning_once(
566
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
567
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
568
+ f" {target_dtype}."
569
+ )
570
+
571
+ query_states = query_states.to(target_dtype)
572
+ key_states = key_states.to(target_dtype)
573
+ value_states = value_states.to(target_dtype)
574
+
575
+ attn_output = self._flash_attention_forward(
576
+ query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
577
+ )
578
+
579
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
580
+ attn_output = self.o_proj(attn_output)
581
+
582
+ if not output_attentions:
583
+ attn_weights = None
584
+
585
+ return attn_output, attn_weights, past_key_value
586
+
587
+ def _flash_attention_forward(
588
+ self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
589
+ ):
590
+ """
591
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
592
+ first unpad the input, then computes the attention scores and pad the final attention scores.
593
+
594
+ Args:
595
+ query_states (`torch.Tensor`):
596
+ Input query states to be passed to Flash Attention API
597
+ key_states (`torch.Tensor`):
598
+ Input key states to be passed to Flash Attention API
599
+ value_states (`torch.Tensor`):
600
+ Input value states to be passed to Flash Attention API
601
+ attention_mask (`torch.Tensor`):
602
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
603
+ position of padding tokens and 1 for the position of non-padding tokens.
604
+ dropout (`int`, *optional*):
605
+ Attention dropout
606
+ softmax_scale (`float`, *optional*):
607
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
608
+ """
609
+ if not self._flash_attn_uses_top_left_mask:
610
+ causal = self.is_causal
611
+ else:
612
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in MiniCPMFlashAttention2 __init__.
613
+ causal = self.is_causal and query_length != 1
614
+ # Contains at least one padding token in the sequence
615
+ if attention_mask is not None:
616
+ batch_size = query_states.shape[0]
617
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
618
+ query_states, key_states, value_states, attention_mask, query_length
619
+ )
620
+
621
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
622
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
623
+ attn_output_unpad = flash_attn_varlen_func(
624
+ query_states,
625
+ key_states,
626
+ value_states,
627
+ cu_seqlens_q=cu_seqlens_q,
628
+ cu_seqlens_k=cu_seqlens_k,
629
+ max_seqlen_q=max_seqlen_in_batch_q,
630
+ max_seqlen_k=max_seqlen_in_batch_k,
631
+ dropout_p=dropout,
632
+ softmax_scale=softmax_scale,
633
+ causal=causal,
634
+ )
635
+
636
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
637
+ else:
638
+ attn_output = flash_attn_func(
639
+ query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
640
+ )
641
+
642
+ return attn_output
643
+
644
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
645
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
646
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
647
+
648
+ key_layer = index_first_axis(
649
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
650
+ )
651
+ value_layer = index_first_axis(
652
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
653
+ )
654
+ if query_length == kv_seq_len:
655
+ query_layer = index_first_axis(
656
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
657
+ )
658
+ cu_seqlens_q = cu_seqlens_k
659
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
660
+ indices_q = indices_k
661
+ elif query_length == 1:
662
+ max_seqlen_in_batch_q = 1
663
+ cu_seqlens_q = torch.arange(
664
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
665
+ ) # There is a memcpy here, that is very bad.
666
+ indices_q = cu_seqlens_q[:-1]
667
+ query_layer = query_layer.squeeze(1)
668
+ else:
669
+ # The -q_len: slice assumes left padding.
670
+ attention_mask = attention_mask[:, -query_length:]
671
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
672
+
673
+ return (
674
+ query_layer,
675
+ key_layer,
676
+ value_layer,
677
+ indices_q,
678
+ (cu_seqlens_q, cu_seqlens_k),
679
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
680
+ )
681
+
682
+
683
+ class MiniCPMSdpaAttention(MiniCPMAttention):
684
+ """
685
+ MiniCPM attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
686
+ `MiniCPMAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
687
+ SDPA API.
688
+ """
689
+
690
+ # Adapted from MiniCPMAttention.forward
691
+ def forward(
692
+ self,
693
+ hidden_states: torch.Tensor,
694
+ attention_mask: Optional[torch.Tensor] = None,
695
+ position_ids: Optional[torch.LongTensor] = None,
696
+ past_key_value: Optional[Cache] = None,
697
+ output_attentions: bool = False,
698
+ use_cache: bool = False,
699
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
700
+ if output_attentions:
701
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
702
+ logger.warning_once(
703
+ "MiniCPMModel is using MiniCPMSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
704
+ 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
705
+ )
706
+ return super().forward(
707
+ hidden_states=hidden_states,
708
+ attention_mask=attention_mask,
709
+ position_ids=position_ids,
710
+ past_key_value=past_key_value,
711
+ output_attentions=output_attentions,
712
+ use_cache=use_cache,
713
+ )
714
+
715
+ bsz, q_len, _ = hidden_states.size()
716
+
717
+ query_states = self.q_proj(hidden_states)
718
+ key_states = self.k_proj(hidden_states)
719
+ value_states = self.v_proj(hidden_states)
720
+
721
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
722
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
723
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
724
+
725
+ kv_seq_len = key_states.shape[-2]
726
+ if past_key_value is not None:
727
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
728
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
729
+
730
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
731
+
732
+ if past_key_value is not None:
733
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
734
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
735
+
736
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
737
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
738
+
739
+ if attention_mask is not None:
740
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
741
+ raise ValueError(
742
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
743
+ )
744
+
745
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
746
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
747
+ if query_states.device.type == "cuda" and attention_mask is not None:
748
+ query_states = query_states.contiguous()
749
+ key_states = key_states.contiguous()
750
+ value_states = value_states.contiguous()
751
+
752
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
753
+ query_states,
754
+ key_states,
755
+ value_states,
756
+ attn_mask=attention_mask,
757
+ dropout_p=self.attention_dropout if self.training else 0.0,
758
+ # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
759
+ is_causal=self.is_causal and attention_mask is None and q_len > 1,
760
+ )
761
+
762
+ attn_output = attn_output.transpose(1, 2).contiguous()
763
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
764
+
765
+ attn_output = self.o_proj(attn_output)
766
+
767
+ return attn_output, None, past_key_value
768
+
769
+
770
+ MINICPM_ATTENTION_CLASSES = {
771
+ "eager": MiniCPMAttention,
772
+ "flash_attention_2": MiniCPMFlashAttention2,
773
+ "sdpa": MiniCPMSdpaAttention,
774
+ }
775
+
776
+
777
+ class MiniCPMDecoderLayer(nn.Module):
778
+ def __init__(self, config: MiniCPMConfig, layer_idx: int):
779
+ super().__init__()
780
+ self.hidden_size = config.hidden_size
781
+ self.self_attn = MINICPM_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
782
+
783
+ self.mlp = MiniCPMMLP(config)
784
+ self.input_layernorm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
785
+ self.post_attention_layernorm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
786
+
787
+ self.scale_depth = config.scale_depth
788
+ self.num_hidden_layers = config.num_hidden_layers
789
+
790
+ def forward(
791
+ self,
792
+ hidden_states: torch.Tensor,
793
+ attention_mask: Optional[torch.Tensor] = None,
794
+ position_ids: Optional[torch.LongTensor] = None,
795
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
796
+ output_attentions: Optional[bool] = False,
797
+ use_cache: Optional[bool] = False,
798
+ **kwargs,
799
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
800
+ """
801
+ Args:
802
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
803
+ attention_mask (`torch.FloatTensor`, *optional*):
804
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
805
+ query_sequence_length, key_sequence_length)` if default attention is used.
806
+ output_attentions (`bool`, *optional*):
807
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
808
+ returned tensors for more detail.
809
+ use_cache (`bool`, *optional*):
810
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
811
+ (see `past_key_values`).
812
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
813
+ """
814
+ if "padding_mask" in kwargs:
815
+ warnings.warn(
816
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
817
+ )
818
+
819
+ residual = hidden_states
820
+ hidden_states = self.input_layernorm(hidden_states)
821
+ # Self Attention
822
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
823
+ hidden_states=hidden_states,
824
+ attention_mask=attention_mask,
825
+ position_ids=position_ids,
826
+ past_key_value=past_key_value,
827
+ output_attentions=output_attentions,
828
+ use_cache=use_cache,
829
+ **kwargs,
830
+ )
831
+
832
+ hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))
833
+
834
+ # Fully Connected
835
+ residual = hidden_states
836
+ hidden_states = self.post_attention_layernorm(hidden_states)
837
+
838
+ hidden_states = self.mlp(hidden_states)
839
+ hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))
840
+
841
+ outputs = (hidden_states,)
842
+
843
+ if output_attentions:
844
+ outputs += (self_attn_weights,)
845
+
846
+ if use_cache:
847
+ outputs += (present_key_value,)
848
+
849
+ return outputs
850
+
851
+
852
+ MINICPM_START_DOCSTRING = r"""
853
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
854
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
855
+ etc.)
856
+
857
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
858
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
859
+ and behavior.
860
+
861
+ Parameters:
862
+ config ([`MiniCPMConfig`]):
863
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
864
+ load the weights associated with the model, only the configuration. Check out the
865
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
866
+ """
867
+
868
+
869
+ @add_start_docstrings(
870
+ "The bare MiniCPM Model outputting raw hidden-states without any specific head on top.",
871
+ MINICPM_START_DOCSTRING,
872
+ )
873
+ class MiniCPMPreTrainedModel(PreTrainedModel):
874
+ config_class = MiniCPMConfig
875
+ base_model_prefix = "model"
876
+ supports_gradient_checkpointing = True
877
+ _no_split_modules = ["MiniCPMDecoderLayer"]
878
+ _skip_keys_device_placement = "past_key_values"
879
+ _supports_flash_attn_2 = True
880
+ _supports_sdpa = True
881
+ _supports_cache_class = True
882
+
883
+ def _init_weights(self, module):
884
+ std = self.config.initializer_range
885
+ if isinstance(module, nn.Linear):
886
+ module.weight.data.normal_(mean=0.0, std=std)
887
+ if module.bias is not None:
888
+ module.bias.data.zero_()
889
+ elif isinstance(module, nn.Embedding):
890
+ module.weight.data.normal_(mean=0.0, std=std)
891
+ if module.padding_idx is not None:
892
+ module.weight.data[module.padding_idx].zero_()
893
+
894
+
895
+ MINICPM_INPUTS_DOCSTRING = r"""
896
+ Args:
897
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
898
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
899
+ it.
900
+
901
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
902
+ [`PreTrainedTokenizer.__call__`] for details.
903
+
904
+ [What are input IDs?](../glossary#input-ids)
905
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
906
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
907
+
908
+ - 1 for tokens that are **not masked**,
909
+ - 0 for tokens that are **masked**.
910
+
911
+ [What are attention masks?](../glossary#attention-mask)
912
+
913
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
914
+ [`PreTrainedTokenizer.__call__`] for details.
915
+
916
+ If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
917
+ `past_key_values`).
918
+
919
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
920
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
921
+ information on the default strategy.
922
+
923
+ - 1 indicates the head is **not masked**,
924
+ - 0 indicates the head is **masked**.
925
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
926
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
927
+ config.n_positions - 1]`.
928
+
929
+ [What are position IDs?](../glossary#position-ids)
930
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
931
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
932
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
933
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
934
+
935
+ Two formats are allowed:
936
+ - a [`~cache_utils.Cache`] instance;
937
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
938
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
939
+ cache format.
940
+
941
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
942
+ legacy cache format will be returned.
943
+
944
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
945
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
946
+ of shape `(batch_size, sequence_length)`.
947
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
948
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
949
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
950
+ model's internal embedding lookup matrix.
951
+ use_cache (`bool`, *optional*):
952
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
953
+ `past_key_values`).
954
+ output_attentions (`bool`, *optional*):
955
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
956
+ tensors for more detail.
957
+ output_hidden_states (`bool`, *optional*):
958
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
959
+ more detail.
960
+ return_dict (`bool`, *optional*):
961
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
962
+ """
963
+
964
+
965
+ @add_start_docstrings(
966
+ "The bare MiniCPM Model outputting raw hidden-states without any specific head on top.",
967
+ MINICPM_START_DOCSTRING,
968
+ )
969
+ class MiniCPMModel(MiniCPMPreTrainedModel):
970
+ """
971
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MiniCPMDecoderLayer`]
972
+
973
+ Args:
974
+ config: MiniCPMConfig
975
+ """
976
+
977
+ def __init__(self, config: MiniCPMConfig):
978
+ super().__init__(config)
979
+ self.padding_idx = config.pad_token_id
980
+ self.vocab_size = config.vocab_size
981
+
982
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
983
+ self.layers = nn.ModuleList(
984
+ [MiniCPMDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
985
+ )
986
+ self._use_sdpa = config._attn_implementation == "sdpa"
987
+ self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
988
+
989
+ self.norm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
990
+
991
+ self.gradient_checkpointing = False
992
+ # Initialize weights and apply final processing
993
+ self.post_init()
994
+
995
+ def get_input_embeddings(self):
996
+ return self.embed_tokens
997
+
998
+ def set_input_embeddings(self, value):
999
+ self.embed_tokens = value
1000
+
1001
+ @add_start_docstrings_to_model_forward(MINICPM_INPUTS_DOCSTRING)
1002
+ def forward(
1003
+ self,
1004
+ input_ids: torch.LongTensor = None,
1005
+ attention_mask: Optional[torch.Tensor] = None,
1006
+ position_ids: Optional[torch.LongTensor] = None,
1007
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1008
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1009
+ use_cache: Optional[bool] = None,
1010
+ output_attentions: Optional[bool] = None,
1011
+ output_hidden_states: Optional[bool] = None,
1012
+ return_dict: Optional[bool] = None,
1013
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
1014
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1015
+ output_hidden_states = (
1016
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1017
+ )
1018
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1019
+
1020
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1021
+
1022
+ # retrieve input_ids and inputs_embeds
1023
+ if input_ids is not None and inputs_embeds is not None:
1024
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
1025
+ elif input_ids is not None:
1026
+ batch_size, seq_length = input_ids.shape[:2]
1027
+ elif inputs_embeds is not None:
1028
+ batch_size, seq_length = inputs_embeds.shape[:2]
1029
+ else:
1030
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
1031
+
1032
+ if self.gradient_checkpointing and self.training:
1033
+ if use_cache:
1034
+ logger.warning_once(
1035
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1036
+ )
1037
+ use_cache = False
1038
+
1039
+ past_key_values_length = 0
1040
+ if use_cache:
1041
+ use_legacy_cache = not isinstance(past_key_values, Cache)
1042
+ if use_legacy_cache:
1043
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1044
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
1045
+
1046
+ if position_ids is None:
1047
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1048
+ position_ids = torch.arange(
1049
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
1050
+ )
1051
+ position_ids = position_ids.unsqueeze(0)
1052
+
1053
+ if inputs_embeds is None:
1054
+ inputs_embeds = self.embed_tokens(input_ids) * self.config.scale_emb
1055
+
1056
+ if self._use_flash_attention_2:
1057
+ # 2d mask is passed through the layers
1058
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
1059
+ elif self._use_sdpa and not output_attentions:
1060
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
1061
+ # the manual implementation that requires a 4D causal mask in all cases.
1062
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
1063
+ attention_mask,
1064
+ (batch_size, seq_length),
1065
+ inputs_embeds,
1066
+ past_key_values_length,
1067
+ )
1068
+ else:
1069
+ # 4d mask is passed through the layers
1070
+ attention_mask = _prepare_4d_causal_attention_mask(
1071
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
1072
+ )
1073
+
1074
+ # embed positions
1075
+ hidden_states = inputs_embeds
1076
+
1077
+ # decoder layers
1078
+ all_hidden_states = () if output_hidden_states else None
1079
+ all_self_attns = () if output_attentions else None
1080
+ next_decoder_cache = None
1081
+
1082
+ for decoder_layer in self.layers:
1083
+ if output_hidden_states:
1084
+ all_hidden_states += (hidden_states,)
1085
+
1086
+ if self.gradient_checkpointing and self.training:
1087
+ layer_outputs = self._gradient_checkpointing_func(
1088
+ decoder_layer.__call__,
1089
+ hidden_states,
1090
+ attention_mask,
1091
+ position_ids,
1092
+ past_key_values,
1093
+ output_attentions,
1094
+ use_cache,
1095
+ )
1096
+ else:
1097
+ layer_outputs = decoder_layer(
1098
+ hidden_states,
1099
+ attention_mask=attention_mask,
1100
+ position_ids=position_ids,
1101
+ past_key_value=past_key_values,
1102
+ output_attentions=output_attentions,
1103
+ use_cache=use_cache,
1104
+ )
1105
+
1106
+ hidden_states = layer_outputs[0]
1107
+
1108
+ if use_cache:
1109
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1110
+
1111
+ if output_attentions:
1112
+ all_self_attns += (layer_outputs[1],)
1113
+
1114
+ hidden_states = self.norm(hidden_states)
1115
+
1116
+ # add hidden states from the last decoder layer
1117
+ if output_hidden_states:
1118
+ all_hidden_states += (hidden_states,)
1119
+
1120
+ next_cache = None
1121
+ if use_cache:
1122
+ next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
1123
+ if not return_dict:
1124
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
1125
+ return BaseModelOutputWithPast(
1126
+ last_hidden_state=hidden_states,
1127
+ past_key_values=next_cache,
1128
+ hidden_states=all_hidden_states,
1129
+ attentions=all_self_attns,
1130
+ )
1131
+
1132
+
1133
+ class MiniCPMForCausalLM(MiniCPMPreTrainedModel):
1134
+ _tied_weights_keys = ["lm_head.weight"]
1135
+
1136
+ def __init__(self, config):
1137
+ super().__init__(config)
1138
+ self.model = MiniCPMModel(config)
1139
+ self.vocab_size = config.vocab_size
1140
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1141
+
1142
+ # Initialize weights and apply final processing
1143
+ self.post_init()
1144
+
1145
+ def get_input_embeddings(self):
1146
+ return self.model.embed_tokens
1147
+
1148
+ def set_input_embeddings(self, value):
1149
+ self.model.embed_tokens = value
1150
+
1151
+ def get_output_embeddings(self):
1152
+ return self.lm_head
1153
+
1154
+ def set_output_embeddings(self, new_embeddings):
1155
+ self.lm_head = new_embeddings
1156
+
1157
+ def set_decoder(self, decoder):
1158
+ self.model = decoder
1159
+
1160
+ def get_decoder(self):
1161
+ return self.model
1162
+
1163
+ @add_start_docstrings_to_model_forward(MINICPM_INPUTS_DOCSTRING)
1164
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1165
+ def forward(
1166
+ self,
1167
+ input_ids: torch.LongTensor = None,
1168
+ attention_mask: Optional[torch.Tensor] = None,
1169
+ position_ids: Optional[torch.LongTensor] = None,
1170
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1171
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1172
+ labels: Optional[torch.LongTensor] = None,
1173
+ use_cache: Optional[bool] = None,
1174
+ output_attentions: Optional[bool] = None,
1175
+ output_hidden_states: Optional[bool] = None,
1176
+ return_dict: Optional[bool] = None,
1177
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1178
+ r"""
1179
+ Args:
1180
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1181
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1182
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1183
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1184
+
1185
+ Returns:
1186
+
1187
+ Example:
1188
+
1189
+ ```python
1190
+ >>> from transformers import AutoTokenizer, MiniCPMForCausalLM
1191
+
1192
+ >>> model = MiniCPMForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
1193
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
1194
+
1195
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1196
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1197
+
1198
+ >>> # Generate
1199
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1200
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1201
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1202
+ ```"""
1203
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1204
+ output_hidden_states = (
1205
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1206
+ )
1207
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1208
+
1209
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1210
+ outputs = self.model(
1211
+ input_ids=input_ids,
1212
+ attention_mask=attention_mask,
1213
+ position_ids=position_ids,
1214
+ past_key_values=past_key_values,
1215
+ inputs_embeds=inputs_embeds,
1216
+ use_cache=use_cache,
1217
+ output_attentions=output_attentions,
1218
+ output_hidden_states=output_hidden_states,
1219
+ return_dict=return_dict,
1220
+ )
1221
+
1222
+ hidden_states = outputs[0]
1223
+ if self.config.pretraining_tp > 1:
1224
+ lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
1225
+ logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
1226
+ logits = torch.cat(logits, dim=-1)
1227
+ else:
1228
+ logits = self.lm_head(hidden_states / (self.config.hidden_size / self.config.dim_model_base))
1229
+ logits = logits.float()
1230
+
1231
+ loss = None
1232
+ if labels is not None:
1233
+ # Shift so that tokens < n predict n
1234
+ shift_logits = logits[..., :-1, :].contiguous()
1235
+ shift_labels = labels[..., 1:].contiguous()
1236
+ # Flatten the tokens
1237
+ loss_fct = CrossEntropyLoss()
1238
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1239
+ shift_labels = shift_labels.view(-1)
1240
+ # Enable model parallelism
1241
+ shift_labels = shift_labels.to(shift_logits.device)
1242
+ loss = loss_fct(shift_logits, shift_labels)
1243
+
1244
+ if not return_dict:
1245
+ output = (logits,) + outputs[1:]
1246
+ return (loss,) + output if loss is not None else output
1247
+
1248
+ return CausalLMOutputWithPast(
1249
+ loss=loss,
1250
+ logits=logits,
1251
+ past_key_values=outputs.past_key_values,
1252
+ hidden_states=outputs.hidden_states,
1253
+ attentions=outputs.attentions,
1254
+ )
1255
+
1256
+ def prepare_inputs_for_generation(
1257
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
1258
+ ):
1259
+ if past_key_values is not None:
1260
+ if isinstance(past_key_values, Cache):
1261
+ cache_length = past_key_values.get_seq_length()
1262
+ past_length = past_key_values.seen_tokens
1263
+ max_cache_length = past_key_values.get_max_length()
1264
+ else:
1265
+ cache_length = past_length = past_key_values[0][0].shape[2]
1266
+ max_cache_length = None
1267
+
1268
+ # Keep only the unprocessed tokens:
1269
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1270
+ # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as
1271
+ # input)
1272
+ if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
1273
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1274
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1275
+ # input_ids based on the past_length.
1276
+ elif past_length < input_ids.shape[1]:
1277
+ input_ids = input_ids[:, past_length:]
1278
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1279
+
1280
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1281
+ if (
1282
+ max_cache_length is not None
1283
+ and attention_mask is not None
1284
+ and cache_length + input_ids.shape[1] > max_cache_length
1285
+ ):
1286
+ attention_mask = attention_mask[:, -max_cache_length:]
1287
+
1288
+ position_ids = kwargs.get("position_ids", None)
1289
+ if attention_mask is not None and position_ids is None:
1290
+ # create position_ids on the fly for batch generation
1291
+ position_ids = attention_mask.long().cumsum(-1) - 1
1292
+ position_ids.masked_fill_(attention_mask == 0, 1)
1293
+ if past_key_values:
1294
+ position_ids = position_ids[:, -input_ids.shape[1] :]
1295
+
1296
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1297
+ if inputs_embeds is not None and past_key_values is None:
1298
+ model_inputs = {"inputs_embeds": inputs_embeds}
1299
+ else:
1300
+ model_inputs = {"input_ids": input_ids}
1301
+
1302
+ model_inputs.update(
1303
+ {
1304
+ "position_ids": position_ids,
1305
+ "past_key_values": past_key_values,
1306
+ "use_cache": kwargs.get("use_cache"),
1307
+ "attention_mask": attention_mask,
1308
+ }
1309
+ )
1310
+ return model_inputs
1311
+
1312
+ @staticmethod
1313
+ def _reorder_cache(past_key_values, beam_idx):
1314
+ reordered_past = ()
1315
+ for layer_past in past_key_values:
1316
+ reordered_past += (
1317
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
1318
+ )
1319
+ return reordered_past
1320
+
1321
+ @torch.inference_mode()
1322
+ def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
1323
+ max_length: int = 4096, num_beams=1, do_sample=True, top_p=0.8, temperature=0.3, logits_processor=None,
1324
+ **kwargs):
1325
+ if history is None:
1326
+ history = []
1327
+ if logits_processor:
1328
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
1329
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1330
+ else:
1331
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
1332
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1333
+
1334
+ history.append({"role": role, "content": query})
1335
+ history_str = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=False)
1336
+ inputs = tokenizer(history_str, return_tensors='pt').to(self.device)
1337
+ outputs = self.generate(**inputs, **gen_kwargs)
1338
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]
1339
+ response = tokenizer.decode(outputs)
1340
+ pattern = re.compile(r".*?(?=<AI>|<用户>)", re.DOTALL)
1341
+ matches = pattern.findall(response)
1342
+ if len(matches) > 0:
1343
+ response = matches[0]
1344
+ history.append({"role": "assistant", "content": response})
1345
+ return response, history
1346
+
1347
+
1348
+ @add_start_docstrings(
1349
+ """
1350
+ The MiniCPM Model transformer with a sequence classification head on top (linear layer).
1351
+
1352
+ [`MiniCPMForSequenceClassification`] uses the last token in order to do the classification, as other causal models
1353
+ (e.g. GPT-2) do.
1354
+
1355
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1356
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1357
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1358
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1359
+ each row of the batch).
1360
+ """,
1361
+ MINICPM_START_DOCSTRING,
1362
+ )
1363
+ class MiniCPMForSequenceClassification(MiniCPMPreTrainedModel):
1364
+ def __init__(self, config):
1365
+ super().__init__(config)
1366
+ self.num_labels = config.num_labels
1367
+ self.model = MiniCPMModel(config)
1368
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1369
+
1370
+ # Initialize weights and apply final processing
1371
+ self.post_init()
1372
+
1373
+ def get_input_embeddings(self):
1374
+ return self.model.embed_tokens
1375
+
1376
+ def set_input_embeddings(self, value):
1377
+ self.model.embed_tokens = value
1378
+
1379
+ @add_start_docstrings_to_model_forward(MINICPM_INPUTS_DOCSTRING)
1380
+ def forward(
1381
+ self,
1382
+ input_ids: torch.LongTensor = None,
1383
+ attention_mask: Optional[torch.Tensor] = None,
1384
+ position_ids: Optional[torch.LongTensor] = None,
1385
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1386
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1387
+ labels: Optional[torch.LongTensor] = None,
1388
+ use_cache: Optional[bool] = None,
1389
+ output_attentions: Optional[bool] = None,
1390
+ output_hidden_states: Optional[bool] = None,
1391
+ return_dict: Optional[bool] = None,
1392
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1393
+ r"""
1394
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1395
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1396
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1397
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1398
+ """
1399
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1400
+
1401
+ transformer_outputs = self.model(
1402
+ input_ids,
1403
+ attention_mask=attention_mask,
1404
+ position_ids=position_ids,
1405
+ past_key_values=past_key_values,
1406
+ inputs_embeds=inputs_embeds,
1407
+ use_cache=use_cache,
1408
+ output_attentions=output_attentions,
1409
+ output_hidden_states=output_hidden_states,
1410
+ return_dict=return_dict,
1411
+ )
1412
+ hidden_states = transformer_outputs[0]
1413
+ logits = self.score(hidden_states)
1414
+
1415
+ if input_ids is not None:
1416
+ batch_size = input_ids.shape[0]
1417
+ else:
1418
+ batch_size = inputs_embeds.shape[0]
1419
+
1420
+ if self.config.pad_token_id is None and batch_size != 1:
1421
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1422
+ if self.config.pad_token_id is None:
1423
+ sequence_lengths = -1
1424
+ else:
1425
+ if input_ids is not None:
1426
+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
1427
+ logits.device
1428
+ )
1429
+ else:
1430
+ sequence_lengths = -1
1431
+
1432
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1433
+
1434
+ loss = None
1435
+ if labels is not None:
1436
+ labels = labels.to(logits.device)
1437
+ if self.config.problem_type is None:
1438
+ if self.num_labels == 1:
1439
+ self.config.problem_type = "regression"
1440
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1441
+ self.config.problem_type = "single_label_classification"
1442
+ else:
1443
+ self.config.problem_type = "multi_label_classification"
1444
+
1445
+ if self.config.problem_type == "regression":
1446
+ loss_fct = MSELoss()
1447
+ if self.num_labels == 1:
1448
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1449
+ else:
1450
+ loss = loss_fct(pooled_logits, labels)
1451
+ elif self.config.problem_type == "single_label_classification":
1452
+ loss_fct = CrossEntropyLoss()
1453
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1454
+ elif self.config.problem_type == "multi_label_classification":
1455
+ loss_fct = BCEWithLogitsLoss()
1456
+ loss = loss_fct(pooled_logits, labels)
1457
+ if not return_dict:
1458
+ output = (pooled_logits,) + transformer_outputs[1:]
1459
+ return ((loss,) + output) if loss is not None else output
1460
+
1461
+ return SequenceClassifierOutputWithPast(
1462
+ loss=loss,
1463
+ logits=pooled_logits,
1464
+ past_key_values=transformer_outputs.past_key_values,
1465
+ hidden_states=transformer_outputs.hidden_states,
1466
+ attentions=transformer_outputs.attentions,
1467
+ )
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6dd932a34bea7eb041f248d38ef10b21c66fba6ee7689ec5049eb490606b7fd8
3
+ size 2720644645
special_tokens_map.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "</s>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "unk_token": {
17
+ "content": "<unk>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ }
23
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bb74d51116831c3bf65db812c553f94ab0c88dcf97a5bbb37e3504f6d359c530
3
+ size 1181204
tokenizer_config.json ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": true,
3
+ "add_eos_token": false,
4
+ "added_tokens_decoder": {
5
+ "0": {
6
+ "content": "<unk>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ },
13
+ "1": {
14
+ "content": "<s>",
15
+ "lstrip": false,
16
+ "normalized": false,
17
+ "rstrip": false,
18
+ "single_word": false,
19
+ "special": true
20
+ },
21
+ "2": {
22
+ "content": "</s>",
23
+ "lstrip": false,
24
+ "normalized": false,
25
+ "rstrip": false,
26
+ "single_word": false,
27
+ "special": true
28
+ }
29
+ },
30
+ "bos_token": "<s>",
31
+ "clean_up_tokenization_spaces": false,
32
+ "eos_token": "</s>",
33
+ "legacy": true,
34
+ "model_max_length": 1000000000000000019884624838656,
35
+ "pad_token": null,
36
+ "sp_model_kwargs": {},
37
+ "spaces_between_special_tokens": false,
38
+ "tokenizer_class": "LlamaTokenizer",
39
+ "unk_token": "<unk>",
40
+ "use_default_system_prompt": false
41
+ }