update special tokens

.gitattributes

@@ -43,11 +43,3 @@ pytorch_model-00001-of-00008.bin filter=lfs diff=lfs merge=lfs -text
 pytorch_model-00002-of-00008.bin filter=lfs diff=lfs merge=lfs -text
 pytorch_model-00005-of-00008.bin filter=lfs diff=lfs merge=lfs -text
 pytorch_model.bin.index.json filter=lfs diff=lfs merge=lfs -text
-model-00001-of-00008.safetensors filter=lfs diff=lfs merge=lfs -text
-model-00002-of-00008.safetensors filter=lfs diff=lfs merge=lfs -text
-model-00003-of-00008.safetensors filter=lfs diff=lfs merge=lfs -text
-model-00004-of-00008.safetensors filter=lfs diff=lfs merge=lfs -text
-model-00005-of-00008.safetensors filter=lfs diff=lfs merge=lfs -text
-model-00006-of-00008.safetensors filter=lfs diff=lfs merge=lfs -text
-model-00007-of-00008.safetensors filter=lfs diff=lfs merge=lfs -text
-model-00008-of-00008.safetensors filter=lfs diff=lfs merge=lfs -text

README.md

@@ -20,15 +20,10 @@ license: other
   
 [![evaluation](https://github.com/InternLM/InternLM/assets/22529082/f80a2a58-5ddf-471a-8da4-32ab65c8fd3b)](https://github.com/internLM/OpenCompass/)
 
-[💻Github Repo](https://github.com/InternLM/InternLM) • [🤔Reporting Issues](https://github.com/InternLM/InternLM/issues/new) • [📜Technical Report](https://arxiv.org/abs/2403.17297)
+[💻Github Repo](https://github.com/InternLM/InternLM) • [🤔Reporting Issues](https://github.com/InternLM/InternLM/issues/new)
 
 </div>
 
-<p align="center">
-    👋 join us on <a href="https://discord.gg/xa29JuW87d" target="_blank">Discord</a> and <a href="https://github.com/InternLM/InternLM/assets/25839884/a6aad896-7232-4220-ac84-9e070c2633ce" target="_blank">WeChat</a>
-</p>
-
-
 
 ## Introduction
 
@@ -99,92 +94,10 @@ for response, history in model.stream_chat(tokenizer, "Hello", history=[]):
     length = len(response)
 ```
 
-## Deployment
-
-### LMDeploy
-
-LMDeploy is a toolkit for compressing, deploying, and serving LLM, developed by the MMRazor and MMDeploy teams.
-
-```bash
-pip install lmdeploy
-```
-
-You can run batch inference locally with the following python code:
-
-```python
-import lmdeploy
-pipe = lmdeploy.pipeline("internlm/internlm2-chat-7b")
-response = pipe(["Hi, pls intro yourself", "Shanghai is"])
-print(response)
-```
-
-Or you can launch an OpenAI compatible server with the following command:
-
-```bash
-lmdeploy serve api_server internlm/internlm2-chat-7b --model-name internlm2-chat-7b --server-port 23333 
-```
-
-Then you can send a chat request to the server:
-
-```bash
-curl http://localhost:23333/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d '{
-    "model": "internlm2-chat-7b",
-    "messages": [
-    {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": "Introduce deep learning to me."}
-    ]
-    }'
-```
-
-Find more details in the [LMDeploy documentation](https://lmdeploy.readthedocs.io/en/latest/)
-
-### vLLM
-
-Launch OpenAI compatible server with `vLLM>=0.3.2`:
-
-```bash
-pip install vllm
-```
-
-```bash
-python -m vllm.entrypoints.openai.api_server --model internlm/internlm2-chat-7b --served-model-name internlm2-chat-7b --trust-remote-code
-```
-
-Then you can send a chat request to the server:
-
-```bash
-curl http://localhost:8000/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d '{
-    "model": "internlm2-chat-7b",
-    "messages": [
-    {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": "Introduce deep learning to me."}
-    ]
-    }'
-```
-
-Find more details in the [vLLM documentation](https://docs.vllm.ai/en/latest/index.html)
-
 ## Open Source License
 
 The code is licensed under Apache-2.0, while model weights are fully open for academic research and also allow **free** commercial usage. To apply for a commercial license, please fill in the [application form (English)](https://wj.qq.com/s2/12727483/5dba/)/[申请表（中文）](https://wj.qq.com/s2/12725412/f7c1/). For other questions or collaborations, please contact <internlm@pjlab.org.cn>.
 
-## Citation
-
-```
-@misc{cai2024internlm2,
-      title={InternLM2 Technical Report},
-      author={Zheng Cai and Maosong Cao and Haojiong Chen and Kai Chen and Keyu Chen and Xin Chen and Xun Chen and Zehui Chen and Zhi Chen and Pei Chu and Xiaoyi Dong and Haodong Duan and Qi Fan and Zhaoye Fei and Yang Gao and Jiaye Ge and Chenya Gu and Yuzhe Gu and Tao Gui and Aijia Guo and Qipeng Guo and Conghui He and Yingfan Hu and Ting Huang and Tao Jiang and Penglong Jiao and Zhenjiang Jin and Zhikai Lei and Jiaxing Li and Jingwen Li and Linyang Li and Shuaibin Li and Wei Li and Yining Li and Hongwei Liu and Jiangning Liu and Jiawei Hong and Kaiwen Liu and Kuikun Liu and Xiaoran Liu and Chengqi Lv and Haijun Lv and Kai Lv and Li Ma and Runyuan Ma and Zerun Ma and Wenchang Ning and Linke Ouyang and Jiantao Qiu and Yuan Qu and Fukai Shang and Yunfan Shao and Demin Song and Zifan Song and Zhihao Sui and Peng Sun and Yu Sun and Huanze Tang and Bin Wang and Guoteng Wang and Jiaqi Wang and Jiayu Wang and Rui Wang and Yudong Wang and Ziyi Wang and Xingjian Wei and Qizhen Weng and Fan Wu and Yingtong Xiong and Chao Xu and Ruiliang Xu and Hang Yan and Yirong Yan and Xiaogui Yang and Haochen Ye and Huaiyuan Ying and Jia Yu and Jing Yu and Yuhang Zang and Chuyu Zhang and Li Zhang and Pan Zhang and Peng Zhang and Ruijie Zhang and Shuo Zhang and Songyang Zhang and Wenjian Zhang and Wenwei Zhang and Xingcheng Zhang and Xinyue Zhang and Hui Zhao and Qian Zhao and Xiaomeng Zhao and Fengzhe Zhou and Zaida Zhou and Jingming Zhuo and Yicheng Zou and Xipeng Qiu and Yu Qiao and Dahua Lin},
-      year={2024},
-      eprint={2403.17297},
-      archivePrefix={arXiv},
-      primaryClass={cs.CL}
-}
-```
-
 ## 简介
 
 InternLM2 ，即书生·浦语大模型第二代，开源了面向实用场景的70亿参数基础模型与对话模型 （InternLM2-Chat-7B）。模型具有以下特点：
@@ -250,88 +163,6 @@ for response, history in model.stream_chat(tokenizer, "你好", history=[]):
     length = len(response)
 ```
 
-## 部署
-
-### LMDeploy
-
-LMDeploy 由 MMDeploy 和 MMRazor 团队联合开发，是涵盖了 LLM 任务的全套轻量化、部署和服务解决方案。
-
-```bash
-pip install lmdeploy
-```
-
-你可以使用以下 python 代码进行本地批量推理:
-
-```python
-import lmdeploy
-pipe = lmdeploy.pipeline("internlm/internlm2-chat-7b")
-response = pipe(["Hi, pls intro yourself", "Shanghai is"])
-print(response)
-```
-
-或者你可以使用以下命令启动兼容 OpenAI API 的服务:
-
-```bash
-lmdeploy serve api_server internlm/internlm2-chat-7b --server-port 23333
-```
-
-然后你可以向服务端发起一个聊天请求:
-
-```bash
-curl http://localhost:23333/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d '{
-    "model": "internlm2-chat-7b",
-    "messages": [
-    {"role": "system", "content": "你是个友善的AI助手。"},
-    {"role": "user", "content": "介绍一下深度学习。"}
-    ]
-    }'
-```
-
-更多信息请查看 [LMDeploy 文档](https://lmdeploy.readthedocs.io/en/latest/)
-
-### vLLM
-
-使用`vLLM>=0.3.2`启动兼容 OpenAI API 的服务:
-
-```bash
-pip install vllm
-```
-
-```bash
-python -m vllm.entrypoints.openai.api_server --model internlm/internlm2-chat-7b --trust-remote-code
-```
-
-然后你可以向服务端发起一个聊天请求:
-
-```bash
-curl http://localhost:8000/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d '{
-    "model": "internlm2-chat-7b",
-    "messages": [
-    {"role": "system", "content": "你是个友善的AI助手。"},
-    {"role": "user", "content": "介绍一下深度学习。"}
-    ]
-    }'
-```
-
-更多信息请查看 [vLLM 文档](https://docs.vllm.ai/en/latest/index.html)
-
 ## 开源许可证
 
-本仓库的代码依照 Apache-2.0 协议开源。模型权重对学术研究完全开放，也可申请免费的商业使用授权（[申请表](https://wj.qq.com/s2/12725412/f7c1/)）。其他问题与合作请联系 <internlm@pjlab.org.cn>。
-
-## 引用
-
-```
-@misc{cai2024internlm2,
-      title={InternLM2 Technical Report},
-      author={Zheng Cai and Maosong Cao and Haojiong Chen and Kai Chen and Keyu Chen and Xin Chen and Xun Chen and Zehui Chen and Zhi Chen and Pei Chu and Xiaoyi Dong and Haodong Duan and Qi Fan and Zhaoye Fei and Yang Gao and Jiaye Ge and Chenya Gu and Yuzhe Gu and Tao Gui and Aijia Guo and Qipeng Guo and Conghui He and Yingfan Hu and Ting Huang and Tao Jiang and Penglong Jiao and Zhenjiang Jin and Zhikai Lei and Jiaxing Li and Jingwen Li and Linyang Li and Shuaibin Li and Wei Li and Yining Li and Hongwei Liu and Jiangning Liu and Jiawei Hong and Kaiwen Liu and Kuikun Liu and Xiaoran Liu and Chengqi Lv and Haijun Lv and Kai Lv and Li Ma and Runyuan Ma and Zerun Ma and Wenchang Ning and Linke Ouyang and Jiantao Qiu and Yuan Qu and Fukai Shang and Yunfan Shao and Demin Song and Zifan Song and Zhihao Sui and Peng Sun and Yu Sun and Huanze Tang and Bin Wang and Guoteng Wang and Jiaqi Wang and Jiayu Wang and Rui Wang and Yudong Wang and Ziyi Wang and Xingjian Wei and Qizhen Weng and Fan Wu and Yingtong Xiong and Chao Xu and Ruiliang Xu and Hang Yan and Yirong Yan and Xiaogui Yang and Haochen Ye and Huaiyuan Ying and Jia Yu and Jing Yu and Yuhang Zang and Chuyu Zhang and Li Zhang and Pan Zhang and Peng Zhang and Ruijie Zhang and Shuo Zhang and Songyang Zhang and Wenjian Zhang and Wenwei Zhang and Xingcheng Zhang and Xinyue Zhang and Hui Zhao and Qian Zhao and Xiaomeng Zhao and Fengzhe Zhou and Zaida Zhou and Jingming Zhuo and Yicheng Zou and Xipeng Qiu and Yu Qiao and Dahua Lin},
-      year={2024},
-      eprint={2403.17297},
-      archivePrefix={arXiv},
-      primaryClass={cs.CL}
-}
-```
+本仓库的代码依照 Apache-2.0 协议开源。模型权重对学术研究完全开放，也可申请免费的商业使用授权（[申请表](https://wj.qq.com/s2/12725412/f7c1/)）。其他问题与合作请联系 <internlm@pjlab.org.cn>。

config.json

@@ -2,9 +2,8 @@
   "architectures": [
     "InternLM2ForCausalLM"
   ],
-  "attn_implementation": "eager",
   "auto_map": {
-    "AutoConfig": "configuration_internlm2.InternLM2Config",
+    "AutoConfig": "configuration_internlm.InternLMConfig",
     "AutoModelForCausalLM": "modeling_internlm2.InternLM2ForCausalLM",
     "AutoModel": "modeling_internlm2.InternLM2ForCausalLM"
   },
@@ -16,21 +15,20 @@
   "initializer_range": 0.02,
   "intermediate_size": 14336,
   "max_position_embeddings": 32768,
-  "model_type": "internlm2",
+  "model_type": "internlm",
   "num_attention_heads": 32,
   "num_hidden_layers": 32,
   "num_key_value_heads": 8,
   "pad_token_id": 2,
   "rms_norm_eps": 1e-05,
   "rope_scaling": {
-    "type": "dynamic",
-    "factor": 2.0
+    "factor": 1.0,
+    "type": "dynamic"
   },
   "rope_theta": 1000000,
   "tie_word_embeddings": false,
-  "torch_dtype": "bfloat16",
-  "transformers_version": "4.37.1",
+  "torch_dtype": "float16",
+  "transformers_version": "4.33.2",
   "use_cache": true,
-  "vocab_size": 92544,
-  "pretraining_tp": 1
+  "vocab_size": 92544
 }

configuration_internlm2.py → configuration_internlm.py

@@ -1,7 +1,10 @@
 # coding=utf-8
-# Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.
+# Copyright (c) InternLM. All rights reserved.
 #
-# This code is based on transformers/src/transformers/models/llama/configuration_llama.py
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,22 +17,21 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-""" InternLM2 model configuration"""
+""" InternLM model configuration"""
 
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
 logger = logging.get_logger(__name__)
 
-INTERNLM2_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+INTERNLM_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
 
 
-# Modified from transformers.model.llama.configuration_llama.LlamaConfig
-class InternLM2Config(PretrainedConfig):
+class InternLMConfig(PretrainedConfig):
     r"""
-    This is the configuration class to store the configuration of a [`InternLM2Model`]. It is used to instantiate
-    an InternLM2 model according to the specified arguments, defining the model architecture. Instantiating a
-    configuration with the defaults will yield a similar configuration to that of the InternLM2-7B.
+    This is the configuration class to store the configuration of a [`InternLMModel`]. It is used to instantiate
+    an InternLM model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the InternLM-7B.
 
     Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
     documentation from [`PretrainedConfig`] for more information.
@@ -37,16 +39,16 @@ class InternLM2Config(PretrainedConfig):
 
     Args:
         vocab_size (`int`, *optional*, defaults to 32000):
-            Vocabulary size of the InternLM2 model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed when calling [`InternLM2Model`]
+            Vocabulary size of the InternLM model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`InternLMModel`]
         hidden_size (`int`, *optional*, defaults to 4096):
             Dimension of the hidden representations.
         intermediate_size (`int`, *optional*, defaults to 11008):
             Dimension of the MLP representations.
         num_hidden_layers (`int`, *optional*, defaults to 32):
-            Number of hidden layers in the Transformer decoder.
+            Number of hidden layers in the Transformer encoder.
         num_attention_heads (`int`, *optional*, defaults to 32):
-            Number of attention heads for each attention layer in the Transformer decoder.
+            Number of attention heads for each attention layer in the Transformer encoder.
         num_key_value_heads (`int`, *optional*):
             This is the number of key_value heads that should be used to implement Grouped Query Attention. If
             `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
@@ -58,42 +60,33 @@ class InternLM2Config(PretrainedConfig):
         hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
             The non-linear activation function (function or string) in the decoder.
         max_position_embeddings (`int`, *optional*, defaults to 2048):
-            The maximum sequence length that this model might ever be used with. InternLM2 supports up to 32768 tokens.
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
         initializer_range (`float`, *optional*, defaults to 0.02):
             The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+        rms_norm_eps (`float`, *optional*, defaults to 1e-12):
             The epsilon used by the rms normalization layers.
         use_cache (`bool`, *optional*, defaults to `True`):
             Whether or not the model should return the last key/values attentions (not used by all models). Only
             relevant if `config.is_decoder=True`.
-        pad_token_id (`int`, *optional*):
-            Padding token id.
-        bos_token_id (`int`, *optional*, defaults to 1):
-            Beginning of stream token id.
-        eos_token_id (`int`, *optional*, defaults to 2):
-            End of stream token id.
-        pretraining_tp (`int`, *optional*, defaults to 1):
-            Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
-            document](https://huggingface.co/docs/transformers/main/perf_train_gpu_many#tensor-parallelism)
-            to understand more about it. This value is necessary to ensure exact reproducibility
-            of the pretraining results. Please refer to [this
-            issue](https://github.com/pytorch/pytorch/issues/76232).
-        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+        tie_word_embeddings(`bool`, *optional*, defaults to `False`):
             Whether to tie weight embeddings
-        rope_theta (`float`, *optional*, defaults to 10000.0):
-            The base period of the RoPE embeddings.
-        rope_scaling (`Dict`, *optional*):
-            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
-            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
-            `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
-            `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
-            these scaling strategies behave:
-            https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
-            experimental feature, subject to breaking API changes in future versions.
-    """
+        Example:
+
+    ```python
+    >>> from transformers import InternLMModel, InternLMConfig
+
+    >>> # Initializing a InternLM internlm-7b style configuration
+    >>> configuration = InternLMConfig()
+
+    >>> # Initializing a model from the internlm-7b style configuration
+    >>> model = InternLMModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+    model_type = "internlm"
     _auto_class = "AutoConfig"
-    model_type = "internlm2"
-    keys_to_ignore_at_inference = ["past_key_values"]
 
     def __init__(  # pylint: disable=W0102
         self,
@@ -111,12 +104,11 @@ class InternLM2Config(PretrainedConfig):
         pad_token_id=0,
         bos_token_id=1,
         eos_token_id=2,
-        pretraining_tp=1,
         tie_word_embeddings=False,
         bias=True,
         rope_theta=10000,
         rope_scaling=None,
-        attn_implementation=None,
+        attn_implementation="eager",
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -134,15 +126,14 @@ class InternLM2Config(PretrainedConfig):
         self.hidden_act = hidden_act
         self.initializer_range = initializer_range
         self.rms_norm_eps = rms_norm_eps
-        self.pretraining_tp = pretraining_tp
         self.use_cache = use_cache
         self.rope_theta = rope_theta
         self.rope_scaling = rope_scaling
         self._rope_scaling_validation()
+
         self.attn_implementation = attn_implementation
         if self.attn_implementation is None:
             self.attn_implementation = "eager"
-
         super().__init__(
             pad_token_id=pad_token_id,
             bos_token_id=bos_token_id,
@@ -169,12 +160,5 @@ class InternLM2Config(PretrainedConfig):
             raise ValueError(
                 f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
             )
-        if (
-            rope_scaling_factor is None
-            or not isinstance(rope_scaling_factor, (float, int))
-            or rope_scaling_factor < 1.0
-        ):
-            raise ValueError(
-                f"`rope_scaling`'s factor field must be a number >= 1, got {rope_scaling_factor} "
-                f"of type {type(rope_scaling_factor)}"
-            )
+        if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor < 1.0:
+            raise ValueError(f"`rope_scaling`'s factor field must be a float >= 1, got {rope_scaling_factor}")

generation_config.json

@@ -1,9 +1,7 @@
 {
+  "_from_model_config": true,
   "bos_token_id": 1,
-  "eos_token_id": [
-    2,
-    92542
-  ],
+  "eos_token_id": 2,
   "pad_token_id": 2,
-  "transformers_version": "4.37.1"
+  "transformers_version": "4.33.2"
 }

model-00005-of-00008.safetensors

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

model-00006-of-00008.safetensors

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

model-00007-of-00008.safetensors

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

model-00008-of-00008.safetensors

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

model.safetensors.index.json

@@ -1,234 +0,0 @@
-{
-  "metadata": {
-    "total_size": 15475417088
-  },
-  "weight_map": {
-    "model.layers.0.attention.wo.weight": "model-00001-of-00008.safetensors",
-    "model.layers.0.attention.wqkv.weight": "model-00001-of-00008.safetensors",
-    "model.layers.0.attention_norm.weight": "model-00001-of-00008.safetensors",
-    "model.layers.0.feed_forward.w1.weight": "model-00001-of-00008.safetensors",
-    "model.layers.0.feed_forward.w2.weight": "model-00001-of-00008.safetensors",
-    "model.layers.0.feed_forward.w3.weight": "model-00001-of-00008.safetensors",
-    "model.layers.0.ffn_norm.weight": "model-00001-of-00008.safetensors",
-    "model.layers.1.attention.wo.weight": "model-00001-of-00008.safetensors",
-    "model.layers.1.attention.wqkv.weight": "model-00001-of-00008.safetensors",
-    "model.layers.1.attention_norm.weight": "model-00001-of-00008.safetensors",
-    "model.layers.1.feed_forward.w1.weight": "model-00001-of-00008.safetensors",
-    "model.layers.1.feed_forward.w2.weight": "model-00001-of-00008.safetensors",
-    "model.layers.1.feed_forward.w3.weight": "model-00001-of-00008.safetensors",
-    "model.layers.1.ffn_norm.weight": "model-00001-of-00008.safetensors",
-    "model.layers.10.attention.wo.weight": "model-00003-of-00008.safetensors",
-    "model.layers.10.attention.wqkv.weight": "model-00003-of-00008.safetensors",
-    "model.layers.10.attention_norm.weight": "model-00003-of-00008.safetensors",
-    "model.layers.10.feed_forward.w1.weight": "model-00003-of-00008.safetensors",
-    "model.layers.10.feed_forward.w2.weight": "model-00003-of-00008.safetensors",
-    "model.layers.10.feed_forward.w3.weight": "model-00003-of-00008.safetensors",
-    "model.layers.10.ffn_norm.weight": "model-00003-of-00008.safetensors",
-    "model.layers.11.attention.wo.weight": "model-00003-of-00008.safetensors",
-    "model.layers.11.attention.wqkv.weight": "model-00003-of-00008.safetensors",
-    "model.layers.11.attention_norm.weight": "model-00004-of-00008.safetensors",
-    "model.layers.11.feed_forward.w1.weight": "model-00003-of-00008.safetensors",
-    "model.layers.11.feed_forward.w2.weight": "model-00004-of-00008.safetensors",
-    "model.layers.11.feed_forward.w3.weight": "model-00003-of-00008.safetensors",
-    "model.layers.11.ffn_norm.weight": "model-00004-of-00008.safetensors",
-    "model.layers.12.attention.wo.weight": "model-00004-of-00008.safetensors",
-    "model.layers.12.attention.wqkv.weight": "model-00004-of-00008.safetensors",
-    "model.layers.12.attention_norm.weight": "model-00004-of-00008.safetensors",
-    "model.layers.12.feed_forward.w1.weight": "model-00004-of-00008.safetensors",
-    "model.layers.12.feed_forward.w2.weight": "model-00004-of-00008.safetensors",
-    "model.layers.12.feed_forward.w3.weight": "model-00004-of-00008.safetensors",
-    "model.layers.12.ffn_norm.weight": "model-00004-of-00008.safetensors",
-    "model.layers.13.attention.wo.weight": "model-00004-of-00008.safetensors",
-    "model.layers.13.attention.wqkv.weight": "model-00004-of-00008.safetensors",
-    "model.layers.13.attention_norm.weight": "model-00004-of-00008.safetensors",
-    "model.layers.13.feed_forward.w1.weight": "model-00004-of-00008.safetensors",
-    "model.layers.13.feed_forward.w2.weight": "model-00004-of-00008.safetensors",
-    "model.layers.13.feed_forward.w3.weight": "model-00004-of-00008.safetensors",
-    "model.layers.13.ffn_norm.weight": "model-00004-of-00008.safetensors",
-    "model.layers.14.attention.wo.weight": "model-00004-of-00008.safetensors",
-    "model.layers.14.attention.wqkv.weight": "model-00004-of-00008.safetensors",
-    "model.layers.14.attention_norm.weight": "model-00004-of-00008.safetensors",
-    "model.layers.14.feed_forward.w1.weight": "model-00004-of-00008.safetensors",
-    "model.layers.14.feed_forward.w2.weight": "model-00004-of-00008.safetensors",
-    "model.layers.14.feed_forward.w3.weight": "model-00004-of-00008.safetensors",
-    "model.layers.14.ffn_norm.weight": "model-00004-of-00008.safetensors",
-    "model.layers.15.attention.wo.weight": "model-00004-of-00008.safetensors",
-    "model.layers.15.attention.wqkv.weight": "model-00004-of-00008.safetensors",
-    "model.layers.15.attention_norm.weight": "model-00004-of-00008.safetensors",
-    "model.layers.15.feed_forward.w1.weight": "model-00004-of-00008.safetensors",
-    "model.layers.15.feed_forward.w2.weight": "model-00004-of-00008.safetensors",
-    "model.layers.15.feed_forward.w3.weight": "model-00004-of-00008.safetensors",
-    "model.layers.15.ffn_norm.weight": "model-00004-of-00008.safetensors",
-    "model.layers.16.attention.wo.weight": "model-00004-of-00008.safetensors",
-    "model.layers.16.attention.wqkv.weight": "model-00004-of-00008.safetensors",
-    "model.layers.16.attention_norm.weight": "model-00005-of-00008.safetensors",
-    "model.layers.16.feed_forward.w1.weight": "model-00005-of-00008.safetensors",
-    "model.layers.16.feed_forward.w2.weight": "model-00005-of-00008.safetensors",
-    "model.layers.16.feed_forward.w3.weight": "model-00005-of-00008.safetensors",
-    "model.layers.16.ffn_norm.weight": "model-00005-of-00008.safetensors",
-    "model.layers.17.attention.wo.weight": "model-00005-of-00008.safetensors",
-    "model.layers.17.attention.wqkv.weight": "model-00005-of-00008.safetensors",
-    "model.layers.17.attention_norm.weight": "model-00005-of-00008.safetensors",
-    "model.layers.17.feed_forward.w1.weight": "model-00005-of-00008.safetensors",
-    "model.layers.17.feed_forward.w2.weight": "model-00005-of-00008.safetensors",
-    "model.layers.17.feed_forward.w3.weight": "model-00005-of-00008.safetensors",
-    "model.layers.17.ffn_norm.weight": "model-00005-of-00008.safetensors",
-    "model.layers.18.attention.wo.weight": "model-00005-of-00008.safetensors",
-    "model.layers.18.attention.wqkv.weight": "model-00005-of-00008.safetensors",
-    "model.layers.18.attention_norm.weight": "model-00005-of-00008.safetensors",
-    "model.layers.18.feed_forward.w1.weight": "model-00005-of-00008.safetensors",
-    "model.layers.18.feed_forward.w2.weight": "model-00005-of-00008.safetensors",
-    "model.layers.18.feed_forward.w3.weight": "model-00005-of-00008.safetensors",
-    "model.layers.18.ffn_norm.weight": "model-00005-of-00008.safetensors",
-    "model.layers.19.attention.wo.weight": "model-00005-of-00008.safetensors",
-    "model.layers.19.attention.wqkv.weight": "model-00005-of-00008.safetensors",
-    "model.layers.19.attention_norm.weight": "model-00005-of-00008.safetensors",
-    "model.layers.19.feed_forward.w1.weight": "model-00005-of-00008.safetensors",
-    "model.layers.19.feed_forward.w2.weight": "model-00005-of-00008.safetensors",
-    "model.layers.19.feed_forward.w3.weight": "model-00005-of-00008.safetensors",
-    "model.layers.19.ffn_norm.weight": "model-00005-of-00008.safetensors",
-    "model.layers.2.attention.wo.weight": "model-00001-of-00008.safetensors",
-    "model.layers.2.attention.wqkv.weight": "model-00001-of-00008.safetensors",
-    "model.layers.2.attention_norm.weight": "model-00002-of-00008.safetensors",
-    "model.layers.2.feed_forward.w1.weight": "model-00001-of-00008.safetensors",
-    "model.layers.2.feed_forward.w2.weight": "model-00002-of-00008.safetensors",
-    "model.layers.2.feed_forward.w3.weight": "model-00001-of-00008.safetensors",
-    "model.layers.2.ffn_norm.weight": "model-00002-of-00008.safetensors",
-    "model.layers.20.attention.wo.weight": "model-00005-of-00008.safetensors",
-    "model.layers.20.attention.wqkv.weight": "model-00005-of-00008.safetensors",
-    "model.layers.20.attention_norm.weight": "model-00006-of-00008.safetensors",
-    "model.layers.20.feed_forward.w1.weight": "model-00005-of-00008.safetensors",
-    "model.layers.20.feed_forward.w2.weight": "model-00006-of-00008.safetensors",
-    "model.layers.20.feed_forward.w3.weight": "model-00005-of-00008.safetensors",
-    "model.layers.20.ffn_norm.weight": "model-00006-of-00008.safetensors",
-    "model.layers.21.attention.wo.weight": "model-00006-of-00008.safetensors",
-    "model.layers.21.attention.wqkv.weight": "model-00006-of-00008.safetensors",
-    "model.layers.21.attention_norm.weight": "model-00006-of-00008.safetensors",
-    "model.layers.21.feed_forward.w1.weight": "model-00006-of-00008.safetensors",
-    "model.layers.21.feed_forward.w2.weight": "model-00006-of-00008.safetensors",
-    "model.layers.21.feed_forward.w3.weight": "model-00006-of-00008.safetensors",
-    "model.layers.21.ffn_norm.weight": "model-00006-of-00008.safetensors",
-    "model.layers.22.attention.wo.weight": "model-00006-of-00008.safetensors",
-    "model.layers.22.attention.wqkv.weight": "model-00006-of-00008.safetensors",
-    "model.layers.22.attention_norm.weight": "model-00006-of-00008.safetensors",
-    "model.layers.22.feed_forward.w1.weight": "model-00006-of-00008.safetensors",
-    "model.layers.22.feed_forward.w2.weight": "model-00006-of-00008.safetensors",
-    "model.layers.22.feed_forward.w3.weight": "model-00006-of-00008.safetensors",
-    "model.layers.22.ffn_norm.weight": "model-00006-of-00008.safetensors",
-    "model.layers.23.attention.wo.weight": "model-00006-of-00008.safetensors",
-    "model.layers.23.attention.wqkv.weight": "model-00006-of-00008.safetensors",
-    "model.layers.23.attention_norm.weight": "model-00006-of-00008.safetensors",
-    "model.layers.23.feed_forward.w1.weight": "model-00006-of-00008.safetensors",
-    "model.layers.23.feed_forward.w2.weight": "model-00006-of-00008.safetensors",
-    "model.layers.23.feed_forward.w3.weight": "model-00006-of-00008.safetensors",
-    "model.layers.23.ffn_norm.weight": "model-00006-of-00008.safetensors",
-    "model.layers.24.attention.wo.weight": "model-00006-of-00008.safetensors",
-    "model.layers.24.attention.wqkv.weight": "model-00006-of-00008.safetensors",
-    "model.layers.24.attention_norm.weight": "model-00006-of-00008.safetensors",
-    "model.layers.24.feed_forward.w1.weight": "model-00006-of-00008.safetensors",
-    "model.layers.24.feed_forward.w2.weight": "model-00006-of-00008.safetensors",
-    "model.layers.24.feed_forward.w3.weight": "model-00006-of-00008.safetensors",
-    "model.layers.24.ffn_norm.weight": "model-00006-of-00008.safetensors",
-    "model.layers.25.attention.wo.weight": "model-00006-of-00008.safetensors",
-    "model.layers.25.attention.wqkv.weight": "model-00006-of-00008.safetensors",
-    "model.layers.25.attention_norm.weight": "model-00007-of-00008.safetensors",
-    "model.layers.25.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
-    "model.layers.25.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
-    "model.layers.25.feed_forward.w3.weight": "model-00007-of-00008.safetensors",
-    "model.layers.25.ffn_norm.weight": "model-00007-of-00008.safetensors",
-    "model.layers.26.attention.wo.weight": "model-00007-of-00008.safetensors",
-    "model.layers.26.attention.wqkv.weight": "model-00007-of-00008.safetensors",
-    "model.layers.26.attention_norm.weight": "model-00007-of-00008.safetensors",
-    "model.layers.26.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
-    "model.layers.26.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
-    "model.layers.26.feed_forward.w3.weight": "model-00007-of-00008.safetensors",
-    "model.layers.26.ffn_norm.weight": "model-00007-of-00008.safetensors",
-    "model.layers.27.attention.wo.weight": "model-00007-of-00008.safetensors",
-    "model.layers.27.attention.wqkv.weight": "model-00007-of-00008.safetensors",
-    "model.layers.27.attention_norm.weight": "model-00007-of-00008.safetensors",
-    "model.layers.27.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
-    "model.layers.27.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
-    "model.layers.27.feed_forward.w3.weight": "model-00007-of-00008.safetensors",
-    "model.layers.27.ffn_norm.weight": "model-00007-of-00008.safetensors",
-    "model.layers.28.attention.wo.weight": "model-00007-of-00008.safetensors",
-    "model.layers.28.attention.wqkv.weight": "model-00007-of-00008.safetensors",
-    "model.layers.28.attention_norm.weight": "model-00007-of-00008.safetensors",
-    "model.layers.28.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
-    "model.layers.28.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
-    "model.layers.28.feed_forward.w3.weight": "model-00007-of-00008.safetensors",
-    "model.layers.28.ffn_norm.weight": "model-00007-of-00008.safetensors",
-    "model.layers.29.attention.wo.weight": "model-00007-of-00008.safetensors",
-    "model.layers.29.attention.wqkv.weight": "model-00007-of-00008.safetensors",
-    "model.layers.29.attention_norm.weight": "model-00008-of-00008.safetensors",
-    "model.layers.29.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
-    "model.layers.29.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
-    "model.layers.29.feed_forward.w3.weight": "model-00007-of-00008.safetensors",
-    "model.layers.29.ffn_norm.weight": "model-00008-of-00008.safetensors",
-    "model.layers.3.attention.wo.weight": "model-00002-of-00008.safetensors",
-    "model.layers.3.attention.wqkv.weight": "model-00002-of-00008.safetensors",
-    "model.layers.3.attention_norm.weight": "model-00002-of-00008.safetensors",
-    "model.layers.3.feed_forward.w1.weight": "model-00002-of-00008.safetensors",
-    "model.layers.3.feed_forward.w2.weight": "model-00002-of-00008.safetensors",
-    "model.layers.3.feed_forward.w3.weight": "model-00002-of-00008.safetensors",
-    "model.layers.3.ffn_norm.weight": "model-00002-of-00008.safetensors",
-    "model.layers.30.attention.wo.weight": "model-00008-of-00008.safetensors",
-    "model.layers.30.attention.wqkv.weight": "model-00008-of-00008.safetensors",
-    "model.layers.30.attention_norm.weight": "model-00008-of-00008.safetensors",
-    "model.layers.30.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
-    "model.layers.30.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
-    "model.layers.30.feed_forward.w3.weight": "model-00008-of-00008.safetensors",
-    "model.layers.30.ffn_norm.weight": "model-00008-of-00008.safetensors",
-    "model.layers.31.attention.wo.weight": "model-00008-of-00008.safetensors",
-    "model.layers.31.attention.wqkv.weight": "model-00008-of-00008.safetensors",
-    "model.layers.31.attention_norm.weight": "model-00008-of-00008.safetensors",
-    "model.layers.31.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
-    "model.layers.31.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
-    "model.layers.31.feed_forward.w3.weight": "model-00008-of-00008.safetensors",
-    "model.layers.31.ffn_norm.weight": "model-00008-of-00008.safetensors",
-    "model.layers.4.attention.wo.weight": "model-00002-of-00008.safetensors",
-    "model.layers.4.attention.wqkv.weight": "model-00002-of-00008.safetensors",
-    "model.layers.4.attention_norm.weight": "model-00002-of-00008.safetensors",
-    "model.layers.4.feed_forward.w1.weight": "model-00002-of-00008.safetensors",
-    "model.layers.4.feed_forward.w2.weight": "model-00002-of-00008.safetensors",
-    "model.layers.4.feed_forward.w3.weight": "model-00002-of-00008.safetensors",
-    "model.layers.4.ffn_norm.weight": "model-00002-of-00008.safetensors",
-    "model.layers.5.attention.wo.weight": "model-00002-of-00008.safetensors",
-    "model.layers.5.attention.wqkv.weight": "model-00002-of-00008.safetensors",
-    "model.layers.5.attention_norm.weight": "model-00002-of-00008.safetensors",
-    "model.layers.5.feed_forward.w1.weight": "model-00002-of-00008.safetensors",
-    "model.layers.5.feed_forward.w2.weight": "model-00002-of-00008.safetensors",
-    "model.layers.5.feed_forward.w3.weight": "model-00002-of-00008.safetensors",
-    "model.layers.5.ffn_norm.weight": "model-00002-of-00008.safetensors",
-    "model.layers.6.attention.wo.weight": "model-00002-of-00008.safetensors",
-    "model.layers.6.attention.wqkv.weight": "model-00002-of-00008.safetensors",
-    "model.layers.6.attention_norm.weight": "model-00002-of-00008.safetensors",
-    "model.layers.6.feed_forward.w1.weight": "model-00002-of-00008.safetensors",
-    "model.layers.6.feed_forward.w2.weight": "model-00002-of-00008.safetensors",
-    "model.layers.6.feed_forward.w3.weight": "model-00002-of-00008.safetensors",
-    "model.layers.6.ffn_norm.weight": "model-00002-of-00008.safetensors",
-    "model.layers.7.attention.wo.weight": "model-00002-of-00008.safetensors",
-    "model.layers.7.attention.wqkv.weight": "model-00002-of-00008.safetensors",
-    "model.layers.7.attention_norm.weight": "model-00003-of-00008.safetensors",
-    "model.layers.7.feed_forward.w1.weight": "model-00003-of-00008.safetensors",
-    "model.layers.7.feed_forward.w2.weight": "model-00003-of-00008.safetensors",
-    "model.layers.7.feed_forward.w3.weight": "model-00003-of-00008.safetensors",
-    "model.layers.7.ffn_norm.weight": "model-00003-of-00008.safetensors",
-    "model.layers.8.attention.wo.weight": "model-00003-of-00008.safetensors",
-    "model.layers.8.attention.wqkv.weight": "model-00003-of-00008.safetensors",
-    "model.layers.8.attention_norm.weight": "model-00003-of-00008.safetensors",
-    "model.layers.8.feed_forward.w1.weight": "model-00003-of-00008.safetensors",
-    "model.layers.8.feed_forward.w2.weight": "model-00003-of-00008.safetensors",
-    "model.layers.8.feed_forward.w3.weight": "model-00003-of-00008.safetensors",
-    "model.layers.8.ffn_norm.weight": "model-00003-of-00008.safetensors",
-    "model.layers.9.attention.wo.weight": "model-00003-of-00008.safetensors",
-    "model.layers.9.attention.wqkv.weight": "model-00003-of-00008.safetensors",
-    "model.layers.9.attention_norm.weight": "model-00003-of-00008.safetensors",
-    "model.layers.9.feed_forward.w1.weight": "model-00003-of-00008.safetensors",
-    "model.layers.9.feed_forward.w2.weight": "model-00003-of-00008.safetensors",
-    "model.layers.9.feed_forward.w3.weight": "model-00003-of-00008.safetensors",
-    "model.layers.9.ffn_norm.weight": "model-00003-of-00008.safetensors",
-    "model.norm.weight": "model-00008-of-00008.safetensors",
-    "model.tok_embeddings.weight": "model-00001-of-00008.safetensors",
-    "output.weight": "model-00008-of-00008.safetensors"
-  }
-}

modeling_internlm2.py

@@ -13,10 +13,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""PyTorch InternLM2 model."""
+""" PyTorch InternLM2 model."""
 import math
 import queue
 import threading
+import warnings
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -26,50 +27,49 @@ from einops import rearrange
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, DynamicCache, StaticCache
-from transformers.modeling_attn_mask_utils import AttentionMaskConverter
 from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
     CausalLMOutputWithPast,
-    QuestionAnsweringModelOutput,
     SequenceClassifierOutputWithPast,
-    TokenClassifierOutput,
 )
 from transformers.modeling_utils import PreTrainedModel
-from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 from transformers.utils import (
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
-    is_flash_attn_greater_or_equal_2_10,
     logging,
     replace_return_docstrings,
 )
 
 try:
     from transformers.generation.streamers import BaseStreamer
-except Exception:
+except:  # noqa # pylint: disable=bare-except
     BaseStreamer = None
 
-from .configuration_internlm2 import InternLM2Config
-
-
-try:
-    from flash_attn import flash_attn_func, flash_attn_varlen_func
-    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
-except:
-    pass
-
+from .configuration_internlm import InternLMConfig as InternLM2Config
 
 logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "InternLM2Config"
 
-
+flash_attn_func, flash_attn_varlen_func = None, None
+pad_input, index_first_axis, unpad_input = None, None, None
+def _import_flash_attn():
+    global flash_attn_func, flash_attn_varlen_func
+    global pad_input, index_first_axis, unpad_input
+    try:
+        from flash_attn import flash_attn_func as _flash_attn_func, flash_attn_varlen_func as _flash_attn_varlen_func
+        from flash_attn.bert_padding import pad_input as _pad_input, index_first_axis as _index_first_axis, unpad_input as _unpad_input
+        flash_attn_func, flash_attn_varlen_func = _flash_attn_func, _flash_attn_varlen_func
+        pad_input, index_first_axis, unpad_input = _pad_input, _index_first_axis, _unpad_input
+    except ImportError:
+        raise ImportError("flash_attn is not installed.")
+
+# Copied from transformers.models.llama.modeling_llama._get_unpad_data
 def _get_unpad_data(attention_mask):
     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
     indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
     max_seqlen_in_batch = seqlens_in_batch.max().item()
-    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))  # pylint: disable=E1102
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
     return (
         indices,
         cu_seqlens,
@@ -77,10 +77,45 @@ def _get_unpad_data(attention_mask):
     )
 
 
-class InternLM2RMSNorm(nn.Module):
-    """InternLM2RMSNorm is equivalent to T5LayerNorm."""
+# Copied from transformers.models.bart.modeling_bart._make_causal_mask
+def _make_causal_mask(
+    input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
+):
+    """
+    Make causal mask used for bi-directional self-attention.
+    """
+    bsz, tgt_len = input_ids_shape
+    mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
+    mask_cond = torch.arange(mask.size(-1), device=device)
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
+    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
+
+
+# Copied from transformers.models.bart.modeling_bart._expand_mask
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+    """
+    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+    """
+    bsz, src_len = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else src_len
+
+    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
+
 
+# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->InternLM2
+class InternLM2RMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-6):
+        """
+        InternLM2RMSNorm is equivalent to T5LayerNorm
+        """
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
@@ -93,68 +128,93 @@ class InternLM2RMSNorm(nn.Module):
         return self.weight * hidden_states.to(input_dtype)
 
 
-ALL_LAYERNORM_LAYERS.append(InternLM2RMSNorm)
-
-
+# Copied from transformers.model.llama.modeling_llama.LlamaRotaryEmbedding with Llama->InternLM2
 class InternLM2RotaryEmbedding(nn.Module):
-    """Rotary Position Embedding for the InternLM2 model. Credits to the Reddit user /u/lucidrains."""
-
-    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
         super().__init__()
-        self.scaling_factor = scaling_factor
+
         self.dim = dim
         self.max_position_embeddings = max_position_embeddings
         self.base = base
-        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
-        # For BC we register cos and sin cached
-        self.max_seq_len_cached = max_position_embeddings
 
-    @torch.no_grad()
-    def forward(self, x, position_ids):
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+        )
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+    def forward(self, x, seq_len=None):
         # x: [bs, num_attention_heads, seq_len, head_size]
-        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
-        position_ids_expanded = position_ids[:, None, :].float()
-        # Force float32 since bfloat16 loses precision on long contexts
-        # See https://github.com/huggingface/transformers/pull/29285
-        device_type = x.device.type
-        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):
-            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos()
-            sin = emb.sin()
-        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+        if seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=torch.float32)
 
+        return (
+            self.cos_cached[:seq_len].to(dtype=x.dtype),
+            self.sin_cached[:seq_len].to(dtype=x.dtype),
+        )
 
+
+# Copied from transformers.model.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->InternLM2
 class InternLM2LinearScalingRotaryEmbedding(InternLM2RotaryEmbedding):
     """InternLM2RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
 
-    def forward(self, x, position_ids):
-        # difference to the original RoPE: a scaling factor is aplied to the position ids
-        position_ids = position_ids.float() / self.scaling_factor
-        cos, sin = super().forward(x, position_ids)
-        return cos, sin
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+        t = t / self.scaling_factor
+
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
 
 
+# Copied from transformers.model.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->InternLM2
 class InternLM2DynamicNTKScalingRotaryEmbedding(InternLM2RotaryEmbedding):
     """InternLM2RotaryEmbedding extended with Dynamic NTK scaling.
-    Credits to the Reddit users /u/bloc97 and /u/emozilla"""
+    Credits to the Reddit users /u/bloc97 and /u/emozilla.
+    """
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
 
-    def forward(self, x, position_ids):
-        # difference to the original RoPE: inv_freq is recomputed when the sequence length > original length
-        seq_len = torch.max(position_ids) + 1
         if seq_len > self.max_position_embeddings:
             base = self.base * (
                 (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
             ) ** (self.dim / (self.dim - 2))
-            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(x.device) / self.dim))
-            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: this may break with compilation
+            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
 
-        cos, sin = super().forward(x, position_ids)
-        return cos, sin
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
 
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
 
+
+# Copied from transformers.model.llama.modeling_llama.rotate_half
 def rotate_half(x):
     """Rotates half the hidden dims of the input."""
     x1 = x[..., : x.shape[-1] // 2]
@@ -162,36 +222,17 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):  # pylint: disable=unused-argument
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`, *optional*):
-            Deprecated and unused.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
+# Copied from transformers.model.llama.modeling_llama.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors."""
+    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
     q_embed = (q * cos) + (rotate_half(q) * sin)
     k_embed = (k * cos) + (rotate_half(k) * sin)
     return q_embed, k_embed
 
 
 class InternLM2MLP(nn.Module):
-    """MLP for InternLM2 model."""
-
     def __init__(self, config):
         super().__init__()
         self.config = config
@@ -208,6 +249,7 @@ class InternLM2MLP(nn.Module):
         return down_proj
 
 
+# Copied from transformers.model.llama.modeling_llama.repeat_kv
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     """
     This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
@@ -220,27 +262,19 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
 
 
+# Modified from transformers.model.llama.modeling_llama.LlamaAttention
 class InternLM2Attention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
-    def __init__(self, config: InternLM2Config, layer_idx: Optional[int] = None):
+    def __init__(self, config: InternLM2Config):
         super().__init__()
         self.config = config
-        self.layer_idx = layer_idx
-        if layer_idx is None:
-            logger.warning_once(
-                f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
-                "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
-                "when creating this class."
-            )
-
         self.hidden_size = config.hidden_size
         self.num_heads = config.num_attention_heads
         self.head_dim = self.hidden_size // self.num_heads
         self.num_key_value_heads = config.num_key_value_heads
         self.num_key_value_groups = self.num_heads // self.num_key_value_heads
         self.max_position_embeddings = config.max_position_embeddings
-        self.rope_theta = config.rope_theta
         self.is_causal = True
 
         if (self.head_dim * self.num_heads) != self.hidden_size:
@@ -254,8 +288,8 @@ class InternLM2Attention(nn.Module):
             (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim,
             bias=config.bias,
         )
-        self.wo = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.bias)
 
+        self.wo = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.bias)
         self._init_rope()
 
     def _init_rope(self):
@@ -263,49 +297,51 @@ class InternLM2Attention(nn.Module):
             self.rotary_emb = InternLM2RotaryEmbedding(
                 self.head_dim,
                 max_position_embeddings=self.max_position_embeddings,
-                base=self.rope_theta,
+                base=self.config.rope_theta,
             )
         else:
             scaling_type = self.config.rope_scaling["type"]
             scaling_factor = self.config.rope_scaling["factor"]
-            if scaling_type == "linear":
-                self.rotary_emb = InternLM2LinearScalingRotaryEmbedding(
+            if scaling_type == "dynamic":
+                self.rotary_emb = InternLM2DynamicNTKScalingRotaryEmbedding(
                     self.head_dim,
                     max_position_embeddings=self.max_position_embeddings,
+                    base=self.config.rope_theta,
                     scaling_factor=scaling_factor,
-                    base=self.rope_theta,
                 )
-            elif scaling_type == "dynamic":
-                self.rotary_emb = InternLM2DynamicNTKScalingRotaryEmbedding(
+            elif scaling_type == "linear":
+                self.rotary_emb = InternLM2LinearScalingRotaryEmbedding(
                     self.head_dim,
                     max_position_embeddings=self.max_position_embeddings,
+                    base=self.config.rope_theta,
                     scaling_factor=scaling_factor,
-                    base=self.rope_theta,
                 )
             else:
-                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
+                raise ValueError("Currently we only support rotary embedding's type being 'dynamic' or 'linear'.")
+        return self.rotary_emb
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
 
     def forward(
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Cache] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: bool = False,
-        use_cache: bool = False,  # pylint: disable=unused-argument
-        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: bool = False,
+        **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. "
+                "Please make sure use `attention_mask` instead.`"
+            )
+
         bsz, q_len, _ = hidden_states.size()
 
-        if self.config.pretraining_tp > 1:
-            # split qkv_states by tp size
-            key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
-            qkv_slices = self.wqkv.weight.split(key_value_slicing, dim=0)
-            qkv_states = torch.cat(
-                [F.linear(hidden_states, qkv_slice) for qkv_slice in qkv_slices], dim=-1  # pylint: disable=E1102
-            )
-        else:
-            qkv_states = self.wqkv(hidden_states)
+        qkv_states = self.wqkv(hidden_states)
 
         qkv_states = rearrange(
             qkv_states,
@@ -315,26 +351,44 @@ class InternLM2Attention(nn.Module):
         )
 
         query_states = qkv_states[..., : self.num_key_value_groups, :]
-        query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d").transpose(1, 2)
-        key_states = qkv_states[..., -2, :].transpose(1, 2)
-        value_states = qkv_states[..., -1, :].transpose(1, 2)
+        query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d")
+        key_states = qkv_states[..., -2, :]
+        value_states = qkv_states[..., -1, :]
 
-        cos, sin = self.rotary_emb(value_states, position_ids)
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 
         if past_key_value is not None:
-            # sin and cos are specific to RoPE models; cache_position needed for the static cache
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
 
         key_states = repeat_kv(key_states, self.num_key_value_groups)
         value_states = repeat_kv(value_states, self.num_key_value_groups)
 
         attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
 
-        if attention_mask is not None:  # no matter the length, we just slice it
-            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
-            attn_weights = attn_weights + causal_mask
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
 
         # upcast attention to fp32
         attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
@@ -347,20 +401,9 @@ class InternLM2Attention(nn.Module):
             )
 
         attn_output = attn_output.transpose(1, 2).contiguous()
-
         attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
 
-        if self.config.pretraining_tp > 1:
-            attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
-            o_proj_slices = self.wo.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
-            attn_output = sum(
-                [
-                    F.linear(attn_output[i], o_proj_slices[i])  # pylint: disable=E1102
-                    for i in range(self.config.pretraining_tp)
-                ]
-            )
-        else:
-            attn_output = self.wo(attn_output)
+        attn_output = self.wo(attn_output)
 
         if not output_attentions:
             attn_weights = None
@@ -368,6 +411,7 @@ class InternLM2Attention(nn.Module):
         return attn_output, attn_weights, past_key_value
 
 
+# Modified from transformers.model.llama.modeling_llama.InternLM2FlashAttention2
 class InternLM2FlashAttention2(InternLM2Attention):
     """
     InternLM2 flash attention module. This module inherits from `InternLM2Attention` as the weights of the module stays
@@ -375,34 +419,26 @@ class InternLM2FlashAttention2(InternLM2Attention):
     flash attention and deal with padding tokens in case the input contains any of them.
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
-        # 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.
-        # 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).
-        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
     def forward(
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.LongTensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Cache] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
-        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        if isinstance(past_key_value, StaticCache):
-            raise ValueError(
-                "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "
-                "make sure to use `sdpa` in the mean time, and open an issue at "
-                "https://github.com/huggingface/transformers"
+        # InternLM2FlashAttention2 attention does not support output_attentions
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. "
+                "Please make sure use `attention_mask` instead.`"
             )
 
+            # overwrite attention_mask with padding_mask
+            attention_mask = kwargs.pop("padding_mask")
+
         output_attentions = False
 
         bsz, q_len, _ = hidden_states.size()
@@ -425,61 +461,37 @@ class InternLM2FlashAttention2(InternLM2Attention):
         key_states = key_states.transpose(1, 2)
         value_states = value_states.transpose(1, 2)
 
-        cos, sin = self.rotary_emb(value_states, position_ids)
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 
         if past_key_value is not None:
-            # sin and cos are specific to RoPE models; cache_position needed for the static cache
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
 
-        # 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
-        # to be able to avoid many of these transpose/reshape/view.
         query_states = query_states.transpose(1, 2)
         key_states = key_states.transpose(1, 2)
         value_states = value_states.transpose(1, 2)
 
-        # dropout_rate = self.attention_dropout if self.training else 0.0
-        dropout_rate = 0.0
-
-        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
-        # therefore the input hidden states gets silently casted in float32. Hence, we need
-        # cast them back in the correct dtype just to be sure everything works as expected.
-        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
-        # in fp32. (InternLM2RMSNorm handles it correctly)
-
-        input_dtype = query_states.dtype
-        if input_dtype == torch.float32:
-            if torch.is_autocast_enabled():
-                target_dtype = torch.get_autocast_gpu_dtype()
-            # Handle the case where the model is quantized
-            elif hasattr(self.config, "_pre_quantization_dtype"):
-                target_dtype = self.config._pre_quantization_dtype
-            else:
-                target_dtype = self.wqkv.weight.dtype
-
-            logger.warning_once(
-                f"The input hidden states seems to be silently casted in float32, this might be related to"
-                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
-                f" {target_dtype}."
-            )
-
-            query_states = query_states.to(target_dtype)
-            key_states = key_states.to(target_dtype)
-            value_states = value_states.to(target_dtype)
+        dropout_rate = 0.0 if not self.training else self.attention_dropout
 
         attn_output = self._flash_attention_forward(
             query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
         )
-
         attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
         attn_output = self.wo(attn_output)
 
         if not output_attentions:
             attn_weights = None
 
-        return attn_output, attn_weights, past_key_value  # pylint: disable=E0606
+        return attn_output, attn_weights, past_key_value
 
     def _flash_attention_forward(
         self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
@@ -498,29 +510,23 @@ class InternLM2FlashAttention2(InternLM2Attention):
             attention_mask (`torch.Tensor`):
                 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
                 position of padding tokens and 1 for the position of non-padding tokens.
-            dropout (`float`):
+            dropout (`int`, *optional*):
                 Attention dropout
             softmax_scale (`float`, *optional*):
                 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
         """
-        if not self._flash_attn_uses_top_left_mask:
-            causal = self.is_causal
-        else:
-            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1.
-            # For details, please see the comment in InternLM2FlashAttention2 __init__.
-            causal = self.is_causal and query_length != 1
-
         # Contains at least one padding token in the sequence
+        causal = self.is_causal and query_length != 1
         if attention_mask is not None:
             batch_size = query_states.shape[0]
-            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._unpad_input(
                 query_states, key_states, value_states, attention_mask, query_length
             )
 
             cu_seqlens_q, cu_seqlens_k = cu_seq_lens
             max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
 
-            attn_output_unpad = flash_attn_varlen_func(  # pylint: disable=E0606
+            attn_output_unpad = flash_attn_varlen_func(
                 query_states,
                 key_states,
                 value_states,
@@ -533,26 +539,27 @@ class InternLM2FlashAttention2(InternLM2Attention):
                 causal=causal,
             )
 
-            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)  # pylint: disable=E0606
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
         else:
-            attn_output = flash_attn_func(  # pylint: disable=E0606
+            attn_output = flash_attn_func(
                 query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
             )
 
         return attn_output
 
-    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+    def _unpad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
         indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
         batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
 
-        key_layer = index_first_axis(  # pylint: disable=E0606
+        key_layer = index_first_axis(
             key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
         )
-        value_layer = index_first_axis(  # pylint: disable=E0606
+        value_layer = index_first_axis(
             value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
         )
+
         if query_length == kv_seq_len:
-            query_layer = index_first_axis(  # pylint: disable=E0606
+            query_layer = index_first_axis(
                 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
             )
             cu_seqlens_q = cu_seqlens_k
@@ -568,139 +575,29 @@ class InternLM2FlashAttention2(InternLM2Attention):
         else:
             # The -q_len: slice assumes left padding.
             attention_mask = attention_mask[:, -query_length:]
-            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(  # pylint: disable=E0606
-                query_layer, attention_mask
-            )
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
 
         return (
             query_layer,
             key_layer,
             value_layer,
-            indices_q,
+            indices_q.to(torch.int64),
             (cu_seqlens_q, cu_seqlens_k),
             (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
         )
 
-
-# Copied from transformers.models.llama.modeling_llama.LllamaSdpaAttention with Llama->InternLM2
-class InternLM2SdpaAttention(InternLM2Attention):
-    """
-    InternLM2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
-    `InternLM2Attention` as the weights of the module stays untouched. The only changes are on the forward pass
-    to adapt to SDPA API.
-    """
-
-    # Adapted from InternLM2Attention.forward
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Cache] = None,
-        output_attentions: bool = False,
-        use_cache: bool = False,
-        cache_position: Optional[torch.LongTensor] = None,
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        if output_attentions:
-            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"`
-            # once this is implemented.
-            logger.warning_once(
-                "InternLM2Model uses InternLM2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` "
-                "does not support `output_attentions=True`. Falling back to the manual attention implementation, "
-                "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.'
-            )
-            return super().forward(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=past_key_value,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                cache_position=cache_position,
-            )
-
-        bsz, q_len, _ = hidden_states.size()
-
-        qkv_states = self.wqkv(hidden_states)
-
-        qkv_states = rearrange(
-            qkv_states,
-            "b q (h gs d) -> b q h gs d",
-            gs=2 + self.num_key_value_groups,
-            d=self.head_dim,
-        )
-
-        query_states = qkv_states[..., : self.num_key_value_groups, :]
-        query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d")
-        key_states = qkv_states[..., -2, :]
-        value_states = qkv_states[..., -1, :]
-
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
-
-        cos, sin = self.rotary_emb(value_states, position_ids)
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        if past_key_value is not None:
-            # sin and cos are specific to RoPE models; cache_position needed for the static cache
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
-        key_states = repeat_kv(key_states, self.num_key_value_groups)
-        value_states = repeat_kv(value_states, self.num_key_value_groups)
-
-        causal_mask = attention_mask
-        if attention_mask is not None:
-            causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
-
-        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
-        # custom attn_mask, Reference: https://github.com/pytorch/pytorch/issues/112577.
-        if query_states.device.type == "cuda" and causal_mask is not None:
-            query_states = query_states.contiguous()
-            key_states = key_states.contiguous()
-            value_states = value_states.contiguous()
-
-        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of
-        # an inline conditional assignment in SDPA to support both torch.compile's dynamic shapes and full graph
-        # options. An inline conditional prevents dynamic shapes from compiling.
-        is_causal = bool(causal_mask is None and q_len > 1)
-
-        attn_output = torch.nn.functional.scaled_dot_product_attention(  # pylint: disable=E1102
-            query_states,
-            key_states,
-            value_states,
-            attn_mask=causal_mask,
-            dropout_p=0.0,
-            is_causal=is_causal,
-        )
-
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.view(bsz, q_len, self.hidden_size)
-
-        attn_output = self.wo(attn_output)
-
-        return attn_output, None, past_key_value
-
-
 INTERNLM2_ATTENTION_CLASSES = {
     "eager": InternLM2Attention,
     "flash_attention_2": InternLM2FlashAttention2,
-    "sdpa": InternLM2SdpaAttention,
 }
 
-
-# Modified from transformers.models.llama.modeling_llama.LlamaDecoderLayer with Llama->InternLM2
+# Modified from transformers.model.llama.modeling_llama.LlamaDecoderLayer
 class InternLM2DecoderLayer(nn.Module):
-    """InternLM2 Decoder Layer. This module is a single layer of the InternLM2 model."""
-
-    def __init__(self, config: InternLM2Config, layer_idx: int):
+    def __init__(self, config: InternLM2Config):
         super().__init__()
         self.hidden_size = config.hidden_size
-        self.layer_idx = layer_idx
 
-        self.attention = INTERNLM2_ATTENTION_CLASSES[config.attn_implementation](config=config, layer_idx=layer_idx)
+        self.attention = INTERNLM2_ATTENTION_CLASSES[config.attn_implementation](config=config)
 
         self.feed_forward = InternLM2MLP(config)
         self.attention_norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -711,10 +608,10 @@ class InternLM2DecoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Cache] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: Optional[bool] = False,
         use_cache: Optional[bool] = False,
-        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
         """
         Args:
@@ -730,6 +627,12 @@ class InternLM2DecoderLayer(nn.Module):
                 (see `past_key_values`).
             past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
         """
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. "
+                "Please make sure use `attention_mask` instead.`"
+            )
+
         residual = hidden_states
 
         hidden_states = self.attention_norm(hidden_states)
@@ -742,7 +645,7 @@ class InternLM2DecoderLayer(nn.Module):
             past_key_value=past_key_value,
             output_attentions=output_attentions,
             use_cache=use_cache,
-            cache_position=cache_position,
+            **kwargs,
         )
         hidden_states = residual + hidden_states
 
@@ -786,20 +689,11 @@ InternLM2_START_DOCSTRING = r"""
     InternLM2_START_DOCSTRING,
 )
 class InternLM2PreTrainedModel(PreTrainedModel):
-    """
-    InternLM2 pretraiend model's base class.
-    """
-
     config_class = InternLM2Config
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
     _no_split_modules = ["InternLM2DecoderLayer"]
-    _skip_keys_device_placement = ["past_key_values"]
-    _supports_flash_attn_2 = True
-    _supports_sdpa = True
-    _supports_cache_class = True
-    _supports_quantized_cache = True
-    _supports_static_cache = True
+    _skip_keys_device_placement = "past_key_values"
 
     def _init_weights(self, module):
         std = self.config.initializer_range
@@ -848,19 +742,14 @@ InternLM2_INPUTS_DOCSTRING = r"""
             config.n_positions - 1]`.
 
             [What are position IDs?](../glossary#position-ids)
-        past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
-            Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
-            blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
-            returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
-
-            Two formats are allowed:
-            - a [`~cache_utils.Cache`] instance;
-            - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
-            shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
-            cache format.
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or
+            when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
+            `(batch_size, num_heads, decoder_sequence_length, embed_size_per_head)`.
 
-            The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
-            legacy cache format will be returned.
+            Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
 
             If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
             have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
@@ -880,14 +769,10 @@ InternLM2_INPUTS_DOCSTRING = r"""
             more detail.
         return_dict (`bool`, *optional*):
             Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
-            Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
-            this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
-            the complete sequence length.
 """
 
 
-# Modified from transformers.models.llama.modeling_llama.LlamaModel with Llama->InternLM2
+# Modified from transformers.model.llama.modeling_llama.LlamaModel
 @add_start_docstrings(
     "The bare InternLM2 Model outputting raw hidden-states without any specific head on top.",
     InternLM2_START_DOCSTRING,
@@ -910,9 +795,7 @@ class InternLM2Model(InternLM2PreTrainedModel):
 
         self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
 
-        self.layers = nn.ModuleList(
-            [InternLM2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-        )
+        self.layers = nn.ModuleList([InternLM2DecoderLayer(config) for _ in range(config.num_hidden_layers)])
         self.norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
         self.gradient_checkpointing = False
@@ -925,96 +808,142 @@ class InternLM2Model(InternLM2PreTrainedModel):
     def set_input_embeddings(self, value):
         self.tok_embeddings = value
 
+    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+        # create causal mask
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                inputs_embeds.dtype,
+                device=inputs_embeds.device,
+                past_key_values_length=past_key_values_length,
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
+                inputs_embeds.device
+            )
+            combined_attention_mask = (
+                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
     @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
     def forward(
         self,
         input_ids: torch.LongTensor = None,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         use_cache: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         use_cache = use_cache if use_cache is not None else self.config.use_cache
+
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError(
-                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
-            )
+        if self.config.attn_implementation == "flash_attention_2":
+            _import_flash_attn()
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
 
-        if self.gradient_checkpointing and self.training and use_cache:
-            logger.warning_once(
-                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
             )
-            use_cache = False
+            position_ids = position_ids.unsqueeze(0)
 
         if inputs_embeds is None:
             inputs_embeds = self.tok_embeddings(input_ids)
 
-        return_legacy_cache = False
-        if use_cache and not isinstance(past_key_values, Cache):  # kept for BC (non `Cache` `past_key_values` inputs)
-            return_legacy_cache = True
-            past_key_values = DynamicCache.from_legacy_cache(past_key_values)
-
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+        if self.config.attn_implementation == "flash_attention_2":
+            # 2d mask is passed through the layers
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        else:
+            if attention_mask is None:
+                attention_mask = torch.ones(
+                    (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
+                )
+            attention_mask = self._prepare_decoder_attention_mask(
+                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
             )
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
-
-        causal_mask = self._update_causal_mask(
-            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
-        )
 
         # embed positions
         hidden_states = inputs_embeds
 
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
         # decoder layers
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None
-        next_decoder_cache = None
+        next_decoder_cache = () if use_cache else None
 
-        for decoder_layer in self.layers:
+        for idx, decoder_layer in enumerate(self.layers):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+
             if self.gradient_checkpointing and self.training:
-                layer_outputs = self._gradient_checkpointing_func(
-                    decoder_layer.__call__,
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, output_attentions, None)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(decoder_layer),
                     hidden_states,
-                    causal_mask,
+                    attention_mask,
                     position_ids,
-                    past_key_values,
-                    output_attentions,
-                    use_cache,
-                    cache_position,
+                    None,
                 )
             else:
                 layer_outputs = decoder_layer(
                     hidden_states,
-                    attention_mask=causal_mask,
+                    attention_mask=attention_mask,
                     position_ids=position_ids,
-                    past_key_value=past_key_values,
+                    past_key_value=past_key_value,
                     output_attentions=output_attentions,
                     use_cache=use_cache,
-                    cache_position=cache_position,
                 )
 
             hidden_states = layer_outputs[0]
 
             if use_cache:
-                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
 
             if output_attentions:
                 all_self_attns += (layer_outputs[1],)
@@ -1026,9 +955,6 @@ class InternLM2Model(InternLM2PreTrainedModel):
             all_hidden_states += (hidden_states,)
 
         next_cache = next_decoder_cache if use_cache else None
-        if return_legacy_cache:
-            next_cache = next_cache.to_legacy_cache()
-
         if not return_dict:
             return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
         return BaseModelOutputWithPast(
@@ -1038,91 +964,11 @@ class InternLM2Model(InternLM2PreTrainedModel):
             attentions=all_self_attns,
         )
 
-    def _update_causal_mask(
-        self,
-        attention_mask: torch.Tensor,
-        input_tensor: torch.Tensor,
-        cache_position: torch.Tensor,
-        past_key_values: Cache,
-        output_attentions: bool,
-    ):
-        # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length
-        # even when the static KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at
-        # each decode steps due to the dynamic shapes. (`recording cudagraph tree for symint key 13`, etc.), which is
-        # VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using `fullgraph=True`.
-        # See more context in https://github.com/huggingface/transformers/pull/29114
-
-        if self.config.attn_implementation == "flash_attention_2":
-            if attention_mask is not None and 0.0 in attention_mask:
-                return attention_mask
-            return None
-
-        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
-        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
-        # to infer the attention mask.
-        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-        using_static_cache = isinstance(past_key_values, StaticCache)
-
-        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
-        if self.config.attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
-            if AttentionMaskConverter._ignore_causal_mask_sdpa(
-                attention_mask,
-                inputs_embeds=input_tensor,
-                past_key_values_length=past_seen_tokens,
-                is_training=self.training,
-            ):
-                return None
-
-        dtype, device = input_tensor.dtype, input_tensor.device
-        min_dtype = torch.finfo(dtype).min
-        sequence_length = input_tensor.shape[1]
-        if using_static_cache:
-            target_length = past_key_values.get_max_length()
-        else:
-            target_length = (
-                attention_mask.shape[-1]
-                if isinstance(attention_mask, torch.Tensor)
-                else past_seen_tokens + sequence_length + 1
-            )
 
-        if attention_mask is not None and attention_mask.dim() == 4:
-            # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing
-            if attention_mask.max() != 0:
-                raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`")
-            causal_mask = attention_mask
-        else:
-            causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
-            if sequence_length != 1:
-                causal_mask = torch.triu(causal_mask, diagonal=1)
-            causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
-            causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
-            if attention_mask is not None:
-                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
-                mask_length = attention_mask.shape[-1]
-                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
-                padding_mask = padding_mask == 0
-                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
-                    padding_mask, min_dtype
-                )
-        if (
-            self.config.attn_implementation == "sdpa"
-            and attention_mask is not None
-            and attention_mask.device.type == "cuda"
-            and not output_attentions
-        ):
-            # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
-            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
-            # Details: https://github.com/pytorch/pytorch/issues/110213
-            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)  # pylint: disable=E1120
-
-        return causal_mask
-
-
-# Modified from transformers.models.llama.modeling_llama.LlamaForCausalLM
+# Modified from transformers.model.llama.modeling_llama.LlamaForCausalLM
 class InternLM2ForCausalLM(InternLM2PreTrainedModel):
-    """Causal language model (CLM) for InternLM2."""
-
     _auto_class = "AutoModelForCausalLM"
+
     _tied_weights_keys = ["output.weight"]
 
     def __init__(self, config):
@@ -1159,14 +1005,13 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
         input_ids: torch.LongTensor = None,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
         use_cache: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         r"""
         Args:
@@ -1182,8 +1027,8 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
         ```python
         >>> from transformers import AutoTokenizer, InternLM2ForCausalLM
 
-        >>> model = InternLM2ForCausalLM.from_pretrained("meta-InternLM2/InternLM2-2-7b-hf")
-        >>> tokenizer = AutoTokenizer.from_pretrained("meta-InternLM2/InternLM2-2-7b-hf")
+        >>> model = InternLM2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
 
         >>> prompt = "Hey, are you conscious? Can you talk to me?"
         >>> inputs = tokenizer(prompt, return_tensors="pt")
@@ -1211,19 +1056,10 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
-            cache_position=cache_position,
         )
 
         hidden_states = outputs[0]
-        if self.config.pretraining_tp > 1:
-            output_slices = self.output.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
-            logits = [
-                F.linear(hidden_states, output_slices[i])  # pylint: disable=not-callable
-                for i in range(self.config.pretraining_tp)
-            ]
-            logits = torch.cat(logits, dim=-1)
-        else:
-            logits = self.output(hidden_states)
+        logits = self.output(hidden_states)
         logits = logits.float()
 
         loss = None
@@ -1252,48 +1088,19 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
         )
 
     def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        past_key_values=None,
-        attention_mask=None,
-        inputs_embeds=None,
-        cache_position=None,
-        use_cache=True,
-        **kwargs,
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
     ):
-        past_length = 0
         if past_key_values is not None:
-            if isinstance(past_key_values, Cache):
-                past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
-                max_cache_length = (
-                    torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
-                    if past_key_values.get_max_length() is not None
-                    else None
-                )
-                cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
-            # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects
+            past_length = past_key_values[0][0].shape[2]
+
+            # Some generation methods already pass only the last input ID
+            if input_ids.shape[1] > past_length:
+                remove_prefix_length = past_length
             else:
-                cache_length = past_length = past_key_values[0][0].shape[2]
-                max_cache_length = None
-
-            # Keep only the unprocessed tokens:
-            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
-            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as input)
-            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
-                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
-            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
-            # input_ids based on the past_length.
-            elif past_length < input_ids.shape[1]:
-                input_ids = input_ids[:, past_length:]
-            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
-
-            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
-            if (
-                max_cache_length is not None
-                and attention_mask is not None
-                and cache_length + input_ids.shape[1] > max_cache_length
-            ):
-                attention_mask = attention_mask[:, -max_cache_length:]  # pylint: disable=E1130
+                # Default to old behavior: keep only final ID
+                remove_prefix_length = input_ids.shape[1] - 1
+
+            input_ids = input_ids[:, remove_prefix_length:]
 
         position_ids = kwargs.get("position_ids", None)
         if attention_mask is not None and position_ids is None:
@@ -1307,24 +1114,13 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
         if inputs_embeds is not None and past_key_values is None:
             model_inputs = {"inputs_embeds": inputs_embeds}
         else:
-            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
-            # recompiles graphs as the stride of the inputs is a guard.
-            # Ref: https://github.com/huggingface/transformers/pull/29114
-            # TODO: use `next_tokens` directly instead.
-            model_inputs = {"input_ids": input_ids.contiguous()}
-
-        input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
-        if cache_position is None:
-            cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
-        elif use_cache:
-            cache_position = cache_position[-input_length:]
+            model_inputs = {"input_ids": input_ids}
 
         model_inputs.update(
             {
                 "position_ids": position_ids,
-                "cache_position": cache_position,
                 "past_key_values": past_key_values,
-                "use_cache": use_cache,
+                "use_cache": kwargs.get("use_cache"),
                 "attention_mask": attention_mask,
             }
         )
@@ -1339,18 +1135,15 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
             )
         return reordered_past
 
-    def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, meta_instruction=""):
-        if history is None:
-            history = []
-        if tokenizer.add_bos_token:
-            prompt = ""
-        else:
-            prompt = tokenizer.bos_token
+    def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = [], meta_instruction=""):
+        prompt = ""
         if meta_instruction:
-            prompt += f"""<|im_start|>system\n{meta_instruction}<|im_end|>\n"""
+            prompt += f"""<s>[UNUSED_TOKEN_146]system\n{meta_instruction}[UNUSED_TOKEN_145]\n"""
+        else:
+            prompt += "<s>"
         for record in history:
-            prompt += f"""<|im_start|>user\n{record[0]}<|im_end|>\n<|im_start|>assistant\n{record[1]}<|im_end|>\n"""
-        prompt += f"""<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"""
+            prompt += f"""[UNUSED_TOKEN_146]user\n{record[0]}[UNUSED_TOKEN_145]\n[UNUSED_TOKEN_146]assistant\n{record[1]}[UNUSED_TOKEN_145]\n"""
+        prompt += f"""[UNUSED_TOKEN_146]user\n{query}[UNUSED_TOKEN_145]\n[UNUSED_TOKEN_146]assistant\n"""
         return tokenizer([prompt], return_tensors="pt")
 
     @torch.no_grad()
@@ -1358,25 +1151,21 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
         self,
         tokenizer,
         query: str,
-        history: Optional[List[Tuple[str, str]]] = None,
+        history: List[Tuple[str, str]] = [],
         streamer: Optional[BaseStreamer] = None,
         max_new_tokens: int = 1024,
         do_sample: bool = True,
         temperature: float = 0.8,
         top_p: float = 0.8,
         meta_instruction: str = "You are an AI assistant whose name is InternLM (书生·浦语).\n"
-        "- InternLM (书生·浦语) is a conversational language model that is developed by Shanghai AI Laboratory "
-        "(上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n"
-        "- InternLM (书生·浦语) can understand and communicate fluently in the language chosen by the user such "
-        "as English and 中文.",
+        "- InternLM (书生·浦语) is a conversational language model that is developed by Shanghai AI Laboratory (上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n"
+        "- InternLM (书生·浦语) can understand and communicate fluently in the language chosen by the user such as English and 中文.",
         **kwargs,
     ):
-        if history is None:
-            history = []
         inputs = self.build_inputs(tokenizer, query, history, meta_instruction)
         inputs = {k: v.to(self.device) for k, v in inputs.items() if torch.is_tensor(v)}
         # also add end-of-assistant token in eos token id to avoid unnecessary generation
-        eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids(["<|im_end|>"])[0]]
+        eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids(["[UNUSED_TOKEN_145]"])[0]]
         outputs = self.generate(
             **inputs,
             streamer=streamer,
@@ -1389,7 +1178,7 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
         )
         outputs = outputs[0].cpu().tolist()[len(inputs["input_ids"][0]) :]
         response = tokenizer.decode(outputs, skip_special_tokens=True)
-        response = response.split("<|im_end|>")[0]
+        response = response.split("[UNUSED_TOKEN_145]")[0]
         history = history + [(query, response)]
         return response, history
 
@@ -1398,15 +1187,13 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
         self,
         tokenizer,
         query: str,
-        history: List[Tuple[str, str]] = None,
+        history: List[Tuple[str, str]] = [],
         max_new_tokens: int = 1024,
         do_sample: bool = True,
         temperature: float = 0.8,
         top_p: float = 0.8,
         **kwargs,
     ):
-        if history is None:
-            history = []
         """
         Return a generator in format: (response, history)
         Eg.
@@ -1422,10 +1209,6 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
         response_queue = queue.Queue(maxsize=20)
 
         class ChatStreamer(BaseStreamer):
-            """
-            Streamer used in generate to print words one by one.
-            """
-
             def __init__(self, tokenizer) -> None:
                 super().__init__()
                 self.tokenizer = tokenizer
@@ -1433,7 +1216,6 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
                 self.query = query
                 self.history = history
                 self.response = ""
-                self.cache = []
                 self.received_inputs = False
                 self.queue.put((self.response, history + [(self.query, self.response)]))
 
@@ -1448,15 +1230,11 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
                     self.received_inputs = True
                     return
 
-                self.cache.extend(value.tolist())
-                token = self.tokenizer.decode(self.cache, skip_special_tokens=True)
-                if token.strip() != "<|im_end|>":
+                token = self.tokenizer.decode([value[-1]], skip_special_tokens=True)
+                if token.strip() != "[UNUSED_TOKEN_145]":
                     self.response = self.response + token
                     history = self.history + [(self.query, self.response)]
                     self.queue.put((self.response, history))
-                    self.cache = []
-                else:
-                    self.end()
 
             def end(self):
                 self.queue.put(None)
@@ -1486,13 +1264,13 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
         return consumer()
 
 
-# Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->InternLM2
+# Copied from transformers.model.llama.modeling_llama.LlamaForSequenceClassification with Llama->InternLM2
 @add_start_docstrings(
     """
     The InternLM2 Model transformer with a sequence classification head on top (linear layer).
 
-    [`InternLM2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
-    (e.g. GPT-2) do.
+    [`InternLM2ForSequenceClassification`] uses the last token in order to do the classification,
+    as other causal models (e.g. GPT-2) do.
 
     Since it does classification on the last token, it requires to know the position of the last token. If a
     `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
@@ -1503,8 +1281,6 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
     InternLM2_START_DOCSTRING,
 )
 class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
-    """Sequence Classification Head for InternLM2 Model."""
-
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
@@ -1526,7 +1302,7 @@ class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
         input_ids: torch.LongTensor = None,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
         use_cache: Optional[bool] = None,
@@ -1567,10 +1343,9 @@ class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
-                sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
-                sequence_lengths = sequence_lengths % input_ids.shape[-1]
-                sequence_lengths = sequence_lengths.to(logits.device)
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
+                    logits.device
+                )
             else:
                 sequence_lengths = -1
 
@@ -1582,7 +1357,7 @@ class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
             if self.config.problem_type is None:
                 if self.num_labels == 1:
                     self.config.problem_type = "regression"
-                elif self.num_labels > 1 and (labels.dtype in (torch.long, torch.int)):
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                     self.config.problem_type = "single_label_classification"
                 else:
                     self.config.problem_type = "multi_label_classification"
@@ -1610,191 +1385,3 @@ class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
             hidden_states=transformer_outputs.hidden_states,
             attentions=transformer_outputs.attentions,
         )
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaForQuestionAnswering with Llama->InternLM2
-@add_start_docstrings(
-    """
-The InternLM2 Model transformer with a span classification head on top for extractive question-answering tasks like
-SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
-    """,
-    InternLM2_START_DOCSTRING,
-)
-class InternLM2ForQuestionAnswering(InternLM2PreTrainedModel):
-    """Question Answering model for InternLM2."""
-
-    base_model_prefix = "transformer"
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.transformer = InternLM2Model(config)
-        self.qa_outputs = nn.Linear(config.hidden_size, 2)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.transformer.tok_embeddings
-
-    def set_input_embeddings(self, value):
-        self.transformer.tok_embeddings = value
-
-    @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        start_positions: Optional[torch.LongTensor] = None,
-        end_positions: Optional[torch.LongTensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, QuestionAnsweringModelOutput]:
-        r"""
-        start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for position (index) of the start of the labelled span for computing the token classification loss.
-            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
-            are not taken into account for computing the loss.
-        end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for position (index) of the end of the labelled span for computing the token classification loss.
-            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
-            are not taken into account for computing the loss.
-        """
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        outputs = self.transformer(
-            input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        sequence_output = outputs[0]
-
-        logits = self.qa_outputs(sequence_output)
-        start_logits, end_logits = logits.split(1, dim=-1)
-        start_logits = start_logits.squeeze(-1).contiguous()
-        end_logits = end_logits.squeeze(-1).contiguous()
-
-        total_loss = None
-        if start_positions is not None and end_positions is not None:
-            # If we are on multi-GPU, split add a dimension
-            if len(start_positions.size()) > 1:
-                start_positions = start_positions.squeeze(-1).to(start_logits.device)
-            if len(end_positions.size()) > 1:
-                end_positions = end_positions.squeeze(-1).to(end_logits.device)
-            # sometimes the start/end positions are outside our model inputs, we ignore these terms
-            ignored_index = start_logits.size(1)
-            start_positions = start_positions.clamp(0, ignored_index)
-            end_positions = end_positions.clamp(0, ignored_index)
-
-            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
-            start_loss = loss_fct(start_logits, start_positions)
-            end_loss = loss_fct(end_logits, end_positions)
-            total_loss = (start_loss + end_loss) / 2
-
-        if not return_dict:
-            output = (start_logits, end_logits) + outputs[2:]
-            return ((total_loss,) + output) if total_loss is not None else output
-
-        return QuestionAnsweringModelOutput(
-            loss=total_loss,
-            start_logits=start_logits,
-            end_logits=end_logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaForTokenClassification with Llama->InternLM2
-@add_start_docstrings(
-    """
-    The InternLM2 Model transformer with a token classification head on top (a linear layer on top of the hidden-states
-    output) e.g. for Named-Entity-Recognition (NER) tasks.
-    """,
-    InternLM2_START_DOCSTRING,
-)
-class InternLM2ForTokenClassification(InternLM2PreTrainedModel):
-    """Token classification model for InternLM2."""
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.num_labels = config.num_labels
-        self.model = InternLM2Model(config)
-        if getattr(config, "classifier_dropout", None) is not None:
-            classifier_dropout = config.classifier_dropout
-        elif getattr(config, "hidden_dropout", None) is not None:
-            classifier_dropout = config.hidden_dropout
-        else:
-            classifier_dropout = 0.1
-        self.dropout = nn.Dropout(classifier_dropout)
-        self.score = nn.Linear(config.hidden_size, config.num_labels)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.model.tok_embeddings
-
-    def set_input_embeddings(self, value):
-        self.model.tok_embeddings = value
-
-    @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
-            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
-            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-        """
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        outputs = self.model(
-            input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-        sequence_output = outputs[0]
-        sequence_output = self.dropout(sequence_output)
-        logits = self.score(sequence_output)
-
-        loss = None
-        if labels is not None:
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
-        if not return_dict:
-            output = (logits,) + outputs[2:]
-            return ((loss,) + output) if loss is not None else output
-
-        return TokenClassifierOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )

model-00004-of-00008.safetensors → pytorch_model-00001-of-00008.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2dd631ca87a5e516780832e747156e304a3c9bb53f85c4304550bff970d2db26
-size 1946242728
+oid sha256:fdd03c1ea040b17ce992eb303d0824f5694932902d0c3fa57b5a8d1df2dc8082
+size 1949342053

model-00001-of-00008.safetensors → pytorch_model-00002-of-00008.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9cec9c25c5f54c545deb16123fc72a0b265fd6bed4c467b4519380a8b215b62b
-size 1949337704
+oid sha256:377d66327bcc2dc4b3e5f4e90f24b428c55832ece07194990248647ea65c2db5
+size 1946249825

model-00002-of-00008.safetensors → pytorch_model-00003-of-00008.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4735b4b15a670c33a6aefe70cf39e4ba08468a496167b5233fa43ebcfbd42895
-size 1946242696
+oid sha256:98aa01b7ffd7901699b4588b5e6b8f753e66c1874f8b855c36b23de2c95ae9d8
+size 1979786923

model-00003-of-00008.safetensors → pytorch_model-00004-of-00008.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:523c7292140bc4c37d5fbd517ec7f9d0aa27b7d22822fdfc730f94f3bf02ebac
-size 1979780440
+oid sha256:9362abdb2afd077e579a7860ad570769255c2d64af2b0831413578c77a44d779
+size 1946249889

pytorch_model-00005-of-00008.bin

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

pytorch_model-00006-of-00008.bin

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

pytorch_model-00007-of-00008.bin

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

pytorch_model-00008-of-00008.bin

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

pytorch_model.bin.index.json

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

special_tokens_map.json

@@ -1,38 +1,6 @@
 {
-  "additional_special_tokens": [
-    "<|im_start|>",
-    "<|im_end|>",
-    "<|action_start|>",
-    "<|action_end|>",
-    "<|interpreter|>",
-    "<|plugin|>"
-  ],
-  "bos_token": {
-    "content": "<s>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "eos_token": {
-    "content": "</s>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "pad_token": {
-    "content": "</s>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "unk_token": {
-    "content": "<unk>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  }
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "pad_token": "</s>",
+  "unk_token": "<unk>"
 }

tokenization_internlm2.py → tokenization_internlm.py

@@ -1,7 +1,10 @@
 # coding=utf-8
-# Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.
+# Copyright (c) InternLM. All rights reserved.
 #
-# This code is based on transformers/src/transformers/models/llama/tokenization_llama.py
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,7 +18,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tokenization classes for InternLM."""
+"""Tokenization classes for IntermLM."""
 import os
 from shutil import copyfile
 from typing import Any, Dict, List, Optional, Tuple
@@ -31,10 +34,9 @@ VOCAB_FILES_NAMES = {"vocab_file": "./tokenizer.model"}
 PRETRAINED_VOCAB_FILES_MAP = {}
 
 
-# Modified from transformers.model.llama.tokenization_llama.LlamaTokenizer
-class InternLM2Tokenizer(PreTrainedTokenizer):
+class InternLMTokenizer(PreTrainedTokenizer):
     """
-    Construct a InternLM2 tokenizer. Based on byte-level Byte-Pair-Encoding.
+    Construct a InternLM tokenizer. Based on byte-level Byte-Pair-Encoding.
 
     Args:
         vocab_file (`str`):
@@ -77,6 +79,8 @@ class InternLM2Tokenizer(PreTrainedTokenizer):
             **kwargs,
         )
 
+        """ Initialization"""
+
     @property
     def no_prefix_space_tokens(self):
         if self._no_prefix_space_tokens is None:

tokenization_internlm2_fast.py

@@ -1,214 +0,0 @@
-# coding=utf-8
-# Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.
-#
-# This code is based on transformers/src/transformers/models/llama/tokenization_llama_fast.py
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tokenization Fast class for InternLM."""
-import os
-from shutil import copyfile
-from typing import Any, Dict, Optional, Tuple
-
-from tokenizers import processors, decoders, Tokenizer, normalizers
-from tokenizers.models import BPE
-
-from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
-from transformers.utils import logging
-
-from transformers.convert_slow_tokenizer import (
-    SLOW_TO_FAST_CONVERTERS,
-    SpmConverter,
-    SentencePieceExtractor,
-)
-
-from .tokenization_internlm2 import InternLM2Tokenizer
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "./tokenizer.model"}
-
-# Modified from transformers.convert_slow_tokenizer.LlamaConverter
-class InternLM2Converter(SpmConverter):
-    handle_byte_fallback = True
-
-    def vocab(self, proto):
-        vocab = [
-            ("<unk>", 0.0),
-            ("<s>", 0.0),
-            ("</s>", 0.0),
-        ]
-        vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
-        return vocab
-
-    def unk_id(self, proto):
-        unk_id = 0
-        return unk_id
-
-    def decoder(self, replacement, add_prefix_space):
-        decoders_sequence = [
-            decoders.Replace("▁", " "),
-            decoders.ByteFallback(),
-            decoders.Fuse(),
-        ]
-        if self.proto.normalizer_spec.add_dummy_prefix:
-            decoders_sequence.append(decoders.Strip(content=" ", left=1))
-        return decoders.Sequence(decoders_sequence)
-
-    def tokenizer(self, proto):
-        model_type = proto.trainer_spec.model_type
-        vocab_scores = self.vocab(proto)
-        # special tokens
-        added_tokens = self.original_tokenizer.added_tokens_decoder
-        for i in range(len(vocab_scores)):
-            piece, score = vocab_scores[i]
-            if i in added_tokens:
-                vocab_scores[i] = (added_tokens[i].content, score)
-        if model_type == 1:
-            raise RuntimeError("InternLM2 is supposed to be a BPE model!")
-
-        elif model_type == 2:
-            _, merges = SentencePieceExtractor(self.original_tokenizer.vocab_file).extract(vocab_scores)
-            bpe_vocab = {word: i for i, (word, _score) in enumerate(vocab_scores)}
-            tokenizer = Tokenizer(
-                BPE(bpe_vocab, merges, unk_token=proto.trainer_spec.unk_piece, fuse_unk=True, byte_fallback=True)
-            )
-            tokenizer.add_special_tokens(
-                [ added_token for index, added_token in added_tokens.items()]
-            )
-        else:
-            raise Exception(
-                "You're trying to run a `Unigram` model but you're file was trained with a different algorithm"
-            )
-
-        return tokenizer
-
-    def normalizer(self, proto):
-        normalizers_list = []
-        if proto.normalizer_spec.add_dummy_prefix:
-            normalizers_list.append(normalizers.Prepend(prepend="▁"))
-        normalizers_list.append(normalizers.Replace(pattern=" ", content="▁"))
-        return normalizers.Sequence(normalizers_list)
-
-    def pre_tokenizer(self, replacement, add_prefix_space):
-        return None
-
-SLOW_TO_FAST_CONVERTERS["InternLM2Tokenizer"] = InternLM2Converter
-
-
-# Modified from transformers.model.llama.tokenization_llama_fast.LlamaTokenizerFast -> InternLM2TokenizerFast
-class InternLM2TokenizerFast(PreTrainedTokenizerFast):
-    vocab_files_names = VOCAB_FILES_NAMES
-    slow_tokenizer_class = InternLM2Tokenizer
-    padding_side = "left"
-    model_input_names = ["input_ids", "attention_mask"]
-    _auto_class = "AutoTokenizer"
-
-    def __init__(
-        self,
-        vocab_file,
-        unk_token="<unk>",
-        bos_token="<s>",
-        eos_token="</s>",
-        pad_token="</s>",
-        sp_model_kwargs: Optional[Dict[str, Any]] = None,
-        add_bos_token=True,
-        add_eos_token=False,
-        decode_with_prefix_space=False,
-        clean_up_tokenization_spaces=False,
-        **kwargs,
-    ):
-        super().__init__(
-            vocab_file=vocab_file,
-            unk_token=unk_token,
-            bos_token=bos_token,
-            eos_token=eos_token,
-            pad_token=pad_token,
-            sp_model_kwargs=sp_model_kwargs,
-            add_bos_token=add_bos_token,
-            add_eos_token=add_eos_token,
-            decode_with_prefix_space=decode_with_prefix_space,
-            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
-            **kwargs,
-        )
-        self._add_bos_token = add_bos_token
-        self._add_eos_token = add_eos_token
-        self.update_post_processor()
-        self.vocab_file = vocab_file
-
-    @property
-    def can_save_slow_tokenizer(self) -> bool:
-        return os.path.isfile(self.vocab_file) if self.vocab_file else False
-
-    def update_post_processor(self):
-        """
-        Updates the underlying post processor with the current `bos_token` and `eos_token`.
-        """
-        bos = self.bos_token
-        bos_token_id = self.bos_token_id
-        if bos is None and self.add_bos_token:
-            raise ValueError("add_bos_token = True but bos_token = None")
-
-        eos = self.eos_token
-        eos_token_id = self.eos_token_id
-        if eos is None and self.add_eos_token:
-            raise ValueError("add_eos_token = True but eos_token = None")
-
-        single = f"{(bos+':0 ') if self.add_bos_token else ''}$A:0{(' '+eos+':0') if self.add_eos_token else ''}"
-        pair = f"{single}{(' '+bos+':1') if self.add_bos_token else ''} $B:1{(' '+eos+':1') if self.add_eos_token else ''}"
-
-        special_tokens = []
-        if self.add_bos_token:
-            special_tokens.append((bos, bos_token_id))
-        if self.add_eos_token:
-            special_tokens.append((eos, eos_token_id))
-        self._tokenizer.post_processor = processors.TemplateProcessing(
-            single=single, pair=pair, special_tokens=special_tokens
-        )
-
-    @property
-    def add_eos_token(self):
-        return self._add_eos_token
-
-    @property
-    def add_bos_token(self):
-        return self._add_bos_token
-
-    @add_eos_token.setter
-    def add_eos_token(self, value):
-        self._add_eos_token = value
-        self.update_post_processor()
-
-    @add_bos_token.setter
-    def add_bos_token(self, value):
-        self._add_bos_token = value
-        self.update_post_processor()
-
-    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
-        if not self.can_save_slow_tokenizer:
-            raise ValueError(
-                "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
-                "tokenizer."
-            )
-
-        if not os.path.isdir(save_directory):
-            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
-            return
-        out_vocab_file = os.path.join(
-            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
-        )
-
-        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
-            copyfile(self.vocab_file, out_vocab_file)
-
-        return (out_vocab_file,)

tokenizer_config.json

@@ -1,102 +1,15 @@
 {
-  "add_bos_token": true,
-  "add_eos_token": false,
-  "added_tokens_decoder": {
-    "0": {
-      "content": "<unk>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "1": {
-      "content": "<s>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "2": {
-      "content": "</s>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "92538": {
-      "content": "<|plugin|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "92539": {
-      "content": "<|interpreter|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "92540": {
-      "content": "<|action_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "92541": {
-      "content": "<|action_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "92542": {
-      "content": "<|im_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "92543": {
-      "content": "<|im_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    }
-  },
-  "additional_special_tokens": [
-    "<|im_start|>",
-    "<|im_end|>",
-    "<|action_start|>",
-    "<|action_end|>",
-    "<|interpreter|>",
-    "<|plugin|>"
-  ],
   "auto_map": {
     "AutoTokenizer": [
-      "tokenization_internlm2.InternLM2Tokenizer",
-      "tokenization_internlm2_fast.InternLM2TokenizerFast"
+      "tokenization_internlm.InternLMTokenizer",
+      null
     ]
   },
   "bos_token": "<s>",
-  "chat_template": "{{ bos_token }}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
   "clean_up_tokenization_spaces": false,
-  "decode_with_prefix_space": false,
   "eos_token": "</s>",
   "model_max_length": 1000000000000000019884624838656,
   "pad_token": "</s>",
-  "sp_model_kwargs": null,
-  "tokenizer_class": "InternLM2Tokenizer",
+  "tokenizer_class": "InternLMTokenizer",
   "unk_token": "<unk>"
 }