Szymon Tworkowski commited on
Commit
4552b82
1 Parent(s): 27390d6

init release

Browse files
config.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LongLlamaForCausalLM"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "configuration_longllama.LongLlamaConfig",
7
+ "AutoModel": "modeling_longllama.LongLlamaModel",
8
+ "AutoModelForCausalLM": "modeling_longllama.LongLlamaForCausalLM",
9
+ "AutoModelForSequenceClassification": "modeling_longllama.LongLlamaForSequenceClassification"
10
+ },
11
+ "bos_token_id": 1,
12
+ "eos_token_id": 2,
13
+ "gradient_checkpoint_every_ith": 1,
14
+ "hidden_act": "silu",
15
+ "hidden_size": 3200,
16
+ "initializer_range": 0.02,
17
+ "intermediate_size": 8640,
18
+ "last_context_length": 1024,
19
+ "max_position_embeddings": 2048,
20
+ "mem_attention_grouping": null,
21
+ "mem_dtype": "bfloat16",
22
+ "mem_layers": [
23
+ 6,
24
+ 12,
25
+ 18
26
+ ],
27
+ "mem_positionals": true,
28
+ "model_type": "longllama",
29
+ "num_attention_heads": 32,
30
+ "num_hidden_layers": 26,
31
+ "pad_token_id": 0,
32
+ "rms_norm_eps": 1e-06,
33
+ "tie_word_embeddings": false,
34
+ "torch_attention": false,
35
+ "torch_dtype": "bfloat16",
36
+ "transformers_version": "4.30.0",
37
+ "use_cache": true,
38
+ "vocab_size": 32000
39
+ }
configuration_longllama.py ADDED
@@ -0,0 +1,150 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 EleutherAI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """ LongLLaMA model configuration"""
21
+
22
+ from transformers.configuration_utils import PretrainedConfig
23
+ from transformers.utils import logging
24
+
25
+
26
+ logger = logging.get_logger(__name__)
27
+
28
+ LONGLLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP = {
29
+ "syzymon/long_llama_3b": "https://huggingface.co/syzymon/long_llama_3b/resolve/main/config.json",
30
+ }
31
+
32
+
33
+ class LongLlamaConfig(PretrainedConfig):
34
+ r"""
35
+ This is the configuration class to store the configuration of a [`LongLlamaModel`]. It is used to instantiate an LongLLaMA
36
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
37
+ defaults will yield a similar configuration to that of the LongLLaMA-7B.
38
+
39
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
40
+ documentation from [`PretrainedConfig`] for more information.
41
+
42
+
43
+ Args:
44
+ vocab_size (`int`, *optional*, defaults to 32000):
45
+ Vocabulary size of the LongLLaMA model. Defines the number of different tokens that can be represented by the
46
+ `inputs_ids` passed when calling [`LongLlamaModel`]
47
+ hidden_size (`int`, *optional*, defaults to 4096):
48
+ Dimension of the hidden representations.
49
+ intermediate_size (`int`, *optional*, defaults to 11008):
50
+ Dimension of the MLP representations.
51
+ num_hidden_layers (`int`, *optional*, defaults to 32):
52
+ Number of hidden layers in the Transformer encoder.
53
+ num_attention_heads (`int`, *optional*, defaults to 32):
54
+ Number of attention heads for each attention layer in the Transformer encoder.
55
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
56
+ The non-linear activation function (function or string) in the decoder.
57
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
58
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
59
+ just in case (e.g., 512 or 1024 or 2048).
60
+ initializer_range (`float`, *optional*, defaults to 0.02):
61
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
62
+ rms_norm_eps (`float`, *optional*, defaults to 1e-12):
63
+ The epsilon used by the rms normalization layers.
64
+ use_cache (`bool`, *optional*, defaults to `True`):
65
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
66
+ relevant if `config.is_decoder=True`.
67
+ tie_word_embeddings(`bool`, *optional*, defaults to `False`):
68
+ Whether to tie weight embeddings
69
+ mem_layers (`List[int]`, defaults to `[]`):
70
+ Layers with memory
71
+ mem_positionals (`bool`, *optional*, defaults to `True`):
72
+ Whether to use positional embeddings in memory layers
73
+ mem_dtype (`str`, *optional*, defaults to `"bfloat16"`):
74
+ Type for keys and values stored in memory
75
+ mem_attention_grouping (`Tuple[int, int]`, *optional*, defaults to `None`):
76
+ One can trade speed for memory by performing attention
77
+ in memory layers sequentially.
78
+ When equal to `(4, 2048)` the memory layers will process at most 4 heads and 2048 queries from each head at once.
79
+ That is at most 4*2048 queries at once.
80
+ torch_attention (`bool`, *optional*, defaults to `False`):
81
+ Whether to use torch scaled_dot_product_attention
82
+ gradient_checkpoint_every_ith (`int`, *optional*, defaults to `1`):
83
+ When gradient checkpointing is enabled checkpoint every ith layer
84
+
85
+ Example:
86
+
87
+ ```python
88
+ >>> from transformers import LongLlamaModel, LongLlamaConfig
89
+
90
+ >>> # Initializing a LongLLaMA longllama-7b style configuration
91
+ >>> configuration = LongLlamaConfig()
92
+
93
+ >>> # Initializing a model from the longllama-7b style configuration
94
+ >>> model = LongLlamaModel(configuration)
95
+
96
+ >>> # Accessing the model configuration
97
+ >>> configuration = model.config
98
+ ```"""
99
+ model_type = "longllama"
100
+ keys_to_ignore_at_inference = ["past_key_values"]
101
+
102
+ def __init__(
103
+ self,
104
+ vocab_size=32000,
105
+ hidden_size=4096,
106
+ intermediate_size=11008,
107
+ num_hidden_layers=32,
108
+ num_attention_heads=32,
109
+ hidden_act="silu",
110
+ max_position_embeddings=2048,
111
+ initializer_range=0.02,
112
+ rms_norm_eps=1e-6,
113
+ use_cache=True,
114
+ pad_token_id=0,
115
+ bos_token_id=1,
116
+ eos_token_id=2,
117
+ tie_word_embeddings=False,
118
+ last_context_length=1024,
119
+ mem_layers=[],
120
+ mem_positionals=True,
121
+ mem_dtype="bfloat16",
122
+ mem_attention_grouping=None,
123
+ torch_attention=False,
124
+ gradient_checkpoint_every_ith=1,
125
+ **kwargs,
126
+ ):
127
+ self.vocab_size = vocab_size
128
+ self.max_position_embeddings = max_position_embeddings
129
+ self.hidden_size = hidden_size
130
+ self.intermediate_size = intermediate_size
131
+ self.num_hidden_layers = num_hidden_layers
132
+ self.num_attention_heads = num_attention_heads
133
+ self.hidden_act = hidden_act
134
+ self.initializer_range = initializer_range
135
+ self.rms_norm_eps = rms_norm_eps
136
+ self.use_cache = use_cache
137
+ self.last_context_length = last_context_length
138
+ self.mem_layers = mem_layers
139
+ self.mem_positionals = mem_positionals
140
+ self.mem_dtype = mem_dtype
141
+ self.mem_attention_grouping = mem_attention_grouping
142
+ self.torch_attention = torch_attention
143
+ self.gradient_checkpoint_every_ith = gradient_checkpoint_every_ith
144
+ super().__init__(
145
+ pad_token_id=pad_token_id,
146
+ bos_token_id=bos_token_id,
147
+ eos_token_id=eos_token_id,
148
+ tie_word_embeddings=tie_word_embeddings,
149
+ **kwargs,
150
+ )
generation_config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 1,
4
+ "eos_token_id": 2,
5
+ "pad_token_id": 0,
6
+ "transformers_version": "4.30.0"
7
+ }
longllama_utils.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from collections import namedtuple
2
+ from dataclasses import dataclass
3
+ import torch
4
+ from typing import Tuple, Optional
5
+
6
+
7
+ @dataclass
8
+ class LongLlamaMemConfig:
9
+ """
10
+ Class for configuring memory caches for LongLlama model.
11
+
12
+ Args:
13
+ positionals (`boolean`)
14
+ Whether to use positional embeddings in memory layer
15
+ cache_dtype (`torch.dtype`)
16
+ Specifies storing type for keys and values
17
+ attention_grouping (`Tuple[int, int]`, *optional*)
18
+ One can trade speed for memory by performing attention
19
+ in memory layers sequentially.
20
+ When equal to `(4, 128)` the memory layers will process at most 4 heads and 128 queries
21
+ from each head at once. That is at most 512 queries at once.
22
+ """
23
+
24
+ positionals: bool = True
25
+ cache_dtype: torch.dtype = torch.bfloat16
26
+ attention_grouping: Optional[Tuple[int, int]] = None
27
+
28
+
29
+ @dataclass
30
+ class LongLlamaMemCache:
31
+ """
32
+ Class with LongLlama's memory cache
33
+
34
+ Args:
35
+ keys (`torch.FloatTensor` of shape `(batch_size, num_heads, mem_length, embed_size_per_head)`)
36
+ values (`torch.FloatTensor` of shape `(batch_size, num_heads, mem_length, embed_size_per_head)`)
37
+ masks (`torch.FloatTensor` of shape `(batch_size, 1, mem_length, 1)`)
38
+ For masking out parts of memory
39
+ """
40
+
41
+ keys: torch.FloatTensor
42
+ values: torch.FloatTensor
43
+ masks: torch.FloatTensor
44
+
45
+
46
+ def mem_apply_update(
47
+ prev_mem_cache: LongLlamaMemCache, new_mem_content: LongLlamaMemCache, mem_config: LongLlamaMemConfig
48
+ ):
49
+ def update_one(prev, new):
50
+ if len(prev.shape) != 4 or len(new.shape) != 4:
51
+ raise ValueError(f"Memory cache content should be consistent in shape got {prev.shape} {new.shape}")
52
+
53
+ return torch.concat([prev, new], dim=-2)
54
+
55
+ insert_size = new_mem_content.keys.shape[-2]
56
+
57
+ if new_mem_content.values.shape[-2] != insert_size or new_mem_content.masks.shape[-2] != insert_size:
58
+ raise ValueError(f"Inconsistent mem_length in new_mem_content")
59
+
60
+ return LongLlamaMemCache(
61
+ keys=update_one(prev_mem_cache.keys, new_mem_content.keys),
62
+ values=update_one(prev_mem_cache.values, new_mem_content.values),
63
+ masks=update_one(prev_mem_cache.masks, new_mem_content.masks),
64
+ )
modeling_longllama.py ADDED
@@ -0,0 +1,1455 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 EleutherAI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """ PyTorch LongLLaMA model."""
21
+ from dataclasses import dataclass
22
+ import math
23
+ from typing import List, Optional, Tuple, Union
24
+
25
+ import torch
26
+ import torch.utils.checkpoint
27
+ from torch import nn
28
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
29
+
30
+ from transformers.activations import ACT2FN
31
+ from transformers.modeling_outputs import (
32
+ BaseModelOutputWithPast,
33
+ CausalLMOutputWithPast,
34
+ SequenceClassifierOutputWithPast,
35
+ )
36
+ from transformers.modeling_utils import PreTrainedModel
37
+ from transformers.utils import (
38
+ add_start_docstrings,
39
+ add_start_docstrings_to_model_forward,
40
+ logging,
41
+ replace_return_docstrings,
42
+ )
43
+ from .configuration_longllama import LongLlamaConfig
44
+ from .longllama_utils import mem_apply_update, LongLlamaMemCache, LongLlamaMemConfig
45
+
46
+
47
+ logger = logging.get_logger(__name__)
48
+
49
+ _CONFIG_FOR_DOC = "LongLlamaConfig"
50
+
51
+
52
+ @dataclass
53
+ class LongLlamaModelOutputWithPast(BaseModelOutputWithPast):
54
+ """
55
+ Based on BaseModelOutputWithPast
56
+
57
+ Args:
58
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
59
+ past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
60
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
61
+ attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
62
+ mem_caches (`tuple(LongLlamaMemCache))`, *optional*, returned for layers with memory cache enabled):
63
+ For the layers without memory None is returned
64
+ """
65
+
66
+ mem_caches: Optional[LongLlamaMemCache] = None
67
+
68
+
69
+ # Copied from transformers.models.bart.modeling_bart._make_causal_mask
70
+ def _make_causal_mask(
71
+ input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
72
+ ):
73
+ """
74
+ Make causal mask used for bi-directional self-attention.
75
+ """
76
+ bsz, tgt_len = input_ids_shape
77
+ mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
78
+ mask_cond = torch.arange(mask.size(-1), device=device)
79
+ mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
80
+ mask = mask.to(dtype)
81
+
82
+ if past_key_values_length > 0:
83
+ mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
84
+ return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
85
+
86
+
87
+ # Copied from transformers.models.bart.modeling_bart._expand_mask
88
+ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
89
+ """
90
+ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
91
+ """
92
+ bsz, src_len = mask.size()
93
+ tgt_len = tgt_len if tgt_len is not None else src_len
94
+
95
+ expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
96
+
97
+ inverted_mask = 1.0 - expanded_mask
98
+
99
+ return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
100
+
101
+
102
+ # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->LongLlama
103
+ class LongLlamaRMSNorm(nn.Module):
104
+ def __init__(self, hidden_size, eps=1e-6):
105
+ """
106
+ LongLlamaRMSNorm is equivalent to T5LayerNorm
107
+ """
108
+ super().__init__()
109
+ self.weight = nn.Parameter(torch.ones(hidden_size))
110
+ self.variance_epsilon = eps
111
+
112
+ def forward(self, hidden_states):
113
+ input_dtype = hidden_states.dtype
114
+ variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
115
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
116
+
117
+ return (self.weight * hidden_states).to(input_dtype)
118
+
119
+
120
+ # Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->LongLlama
121
+ class LongLlamaRotaryEmbedding(torch.nn.Module):
122
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
123
+ super().__init__()
124
+ inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
125
+ self.register_buffer("inv_freq", inv_freq)
126
+
127
+ # Build here to make `torch.jit.trace` work.
128
+ self.max_seq_len_cached = max_position_embeddings
129
+ t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
130
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
131
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
132
+ emb = torch.cat((freqs, freqs), dim=-1)
133
+ dtype = torch.get_default_dtype()
134
+ self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
135
+ self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
136
+
137
+ def forward(self, x, seq_len=None):
138
+ # x: [bs, num_attention_heads, seq_len, head_size]
139
+ # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
140
+ if seq_len > self.max_seq_len_cached:
141
+ self.max_seq_len_cached = seq_len
142
+ t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
143
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
144
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
145
+ emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
146
+ self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(x.dtype), persistent=False)
147
+ self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(x.dtype), persistent=False)
148
+ return (
149
+ self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
150
+ self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
151
+ )
152
+
153
+
154
+ def rotate_half(x):
155
+ """Rotates half the hidden dims of the input."""
156
+ x1 = x[..., : x.shape[-1] // 2]
157
+ x2 = x[..., x.shape[-1] // 2 :]
158
+ return torch.cat((-x2, x1), dim=-1)
159
+
160
+
161
+ # Based on transformers.models.llama.modeling_llama.apply_rotary_pos_emb
162
+ def rotate_one(x, cos, sin, position_ids):
163
+ if len(position_ids.shape) != 2 or x.shape[0] != position_ids.shape[0] or x.shape[-2] != position_ids.shape[1]:
164
+ raise ValueError(f"Position ids shoud have shape [bsz, seq_len] got {position_ids.shape}")
165
+ # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
166
+ cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
167
+ sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
168
+ cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
169
+ sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
170
+ x_embed = (x * cos) + (rotate_half(x) * sin)
171
+ return x_embed
172
+
173
+
174
+ def rotate_as_if_first(x, rotary_emb):
175
+ # x: [bs, num_attention_heads, seq_len, head_size]
176
+ # apply rotary as if all elements were first in the sequence
177
+ cos, sin = rotary_emb(x, x.shape[-2])
178
+ return rotate_one(x, cos, sin, torch.zeros(x.shape[0], x.shape[-2], dtype=torch.long, device=cos.device))
179
+
180
+
181
+ # Copied from transformers.models.llama.modeling_llama.LlamaMLP with Llama->LongLlama
182
+ class LongLlamaMLP(nn.Module):
183
+ def __init__(
184
+ self,
185
+ hidden_size: int,
186
+ intermediate_size: int,
187
+ hidden_act: str,
188
+ ):
189
+ super().__init__()
190
+ self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
191
+ self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
192
+ self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
193
+ self.act_fn = ACT2FN[hidden_act]
194
+
195
+ def forward(self, x):
196
+ return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
197
+
198
+
199
+ # Modified transformers.models.llama.modeling_llama.LlamaAttention
200
+ class LongLlamaAttention(nn.Module):
201
+ """Multi-headed attention from 'Attention Is All You Need' paper with FoT modifications"""
202
+
203
+ def __init__(self, config: LongLlamaConfig, mem_config: Optional[LongLlamaMemConfig] = None):
204
+ super().__init__()
205
+ self.config = config
206
+ self.hidden_size = config.hidden_size
207
+ self.num_heads = config.num_attention_heads
208
+ self.head_dim = self.hidden_size // self.num_heads
209
+ self.max_position_embeddings = config.max_position_embeddings
210
+ self.max_cache = self.max_position_embeddings
211
+
212
+ if (self.head_dim * self.num_heads) != self.hidden_size:
213
+ raise ValueError(
214
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
215
+ f" and `num_heads`: {self.num_heads})."
216
+ )
217
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
218
+ self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
219
+ self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
220
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
221
+ self.rotary_emb = LongLlamaRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
222
+ self.mem_config = mem_config
223
+
224
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
225
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
226
+
227
+ def forward(
228
+ self,
229
+ hidden_states: torch.Tensor,
230
+ attention_mask: Optional[torch.Tensor] = None,
231
+ position_ids: Optional[torch.LongTensor] = None,
232
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
233
+ output_attentions: bool = False,
234
+ use_cache: bool = False,
235
+ mem_cache: Optional[LongLlamaMemCache] = None,
236
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
237
+ if attention_mask is None:
238
+ tgt_seq_len = hidden_states.shape[-2]
239
+ if past_key_value is not None:
240
+ src_seq_len = past_key_value[0].shape[-2] + tgt_seq_len
241
+ else:
242
+ src_seq_len = tgt_seq_len
243
+
244
+ attention_mask = torch.zeros(
245
+ hidden_states.shape[0],
246
+ 1,
247
+ tgt_seq_len,
248
+ src_seq_len,
249
+ device=hidden_states.device,
250
+ dtype=hidden_states.dtype,
251
+ )
252
+ bsz, q_len, _ = hidden_states.size()
253
+
254
+ query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
255
+ key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
256
+ value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
257
+ position_ids = position_ids[:, None, :, None]
258
+
259
+ if position_ids.shape != (key_states.shape[0], 1, key_states.shape[-2], 1):
260
+ raise ValueError("position_ids should match batch and seq_len of the input")
261
+
262
+ mem_no_local_cache = self.mem_config is not None and past_key_value is None and (not use_cache)
263
+ mem_and_local_cache = self.mem_config is not None and use_cache
264
+ # positonal embeddings can be disabled for memory layers
265
+ use_positionals = self.mem_config is None or self.mem_config.positionals
266
+
267
+ if mem_no_local_cache:
268
+ # the whole context window will be moved to memory cache after the attention
269
+ if use_positionals:
270
+ # positionally embedd memory content as first token in the sequence
271
+ rfst_key_states = rotate_as_if_first(key_states, self.rotary_emb)
272
+ else:
273
+ rfst_key_states = key_states
274
+ # attention_mask [bsz, 1, tgt_seq_len, src_seq_len]
275
+ # we base the mask on the last token in the context window
276
+ mem_update = LongLlamaMemCache(
277
+ keys=rfst_key_states.to(self.mem_config.cache_dtype),
278
+ values=value_states.to(self.mem_config.cache_dtype),
279
+ masks=attention_mask[..., -1, :, None],
280
+ )
281
+
282
+ if past_key_value is not None:
283
+ past_local_cache_size = past_key_value[0].shape[-2]
284
+ key_states = torch.cat([past_key_value[0], key_states], dim=-2)
285
+ value_states = torch.cat([past_key_value[1], value_states], dim=-2)
286
+ # FoT additionally stores position_ids to support long inputs
287
+ position_ids = torch.cat([past_key_value[2], position_ids], dim=-2)
288
+
289
+ if attention_mask.shape[-1] != key_states.shape[-2] and attention_mask.shape[-2] != query_states.shape[-2]:
290
+ raise ValueError("attention_mask should be provided for all key_states in local context")
291
+
292
+ # local cache is maintained so that it is <= self.max_cache
293
+ # remaining elements are either dropped or go to memory cache
294
+ if key_states.shape[-2] > self.max_cache:
295
+ num_elems_to_drop = past_local_cache_size
296
+
297
+ if mem_and_local_cache:
298
+ drop_keys = key_states[:, :, :num_elems_to_drop, :]
299
+ drop_values = value_states[:, :, :num_elems_to_drop, :]
300
+ # as memory mask use the masking of the last key in context
301
+ # attention_mask [bsz, 1, tgt_seq_len, src_seq_len]
302
+ drop_masks = attention_mask[..., -1, :, None]
303
+ drop_masks = drop_masks[:, :, :num_elems_to_drop, :]
304
+
305
+ if use_positionals:
306
+ rfst_drop_keys = rotate_as_if_first(drop_keys, self.rotary_emb)
307
+ else:
308
+ rfst_drop_keys = drop_keys
309
+ mem_update = LongLlamaMemCache(
310
+ keys=rfst_drop_keys.to(self.mem_config.cache_dtype),
311
+ values=drop_values.to(self.mem_config.cache_dtype),
312
+ masks=drop_masks,
313
+ )
314
+ if mem_cache is None:
315
+ mem_cache = mem_update
316
+ else:
317
+ mem_cache = mem_apply_update(
318
+ prev_mem_cache=mem_cache, new_mem_content=mem_update, mem_config=self.mem_config
319
+ )
320
+
321
+ key_states = key_states[:, :, num_elems_to_drop:, :]
322
+ value_states = value_states[:, :, num_elems_to_drop:, :]
323
+ position_ids = position_ids[:, :, num_elems_to_drop:, :]
324
+ attention_mask = attention_mask[..., num_elems_to_drop:]
325
+
326
+ # FoT additionally stores position_ids to support long inputs
327
+ past_key_value = (key_states, value_states, position_ids) if use_cache else None
328
+
329
+ kv_seq_len = key_states.shape[-2]
330
+
331
+ if use_positionals:
332
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
333
+ rel_pos_ids = position_ids - torch.min(position_ids, dim=-2, keepdim=True)[0]
334
+ rel_pos_ids = rel_pos_ids.squeeze(3).squeeze(1)
335
+
336
+ query_states = rotate_one(query_states, cos, sin, rel_pos_ids[:, -query_states.shape[-2] :])
337
+ key_states = rotate_one(key_states, cos, sin, rel_pos_ids)
338
+
339
+ if self.mem_config is not None and self.mem_config.attention_grouping is not None:
340
+ attn_grouping_h, attn_grouping_q = self.mem_config.attention_grouping
341
+ if attn_grouping_h <= 0 or attn_grouping_q <= 0:
342
+ raise ValueError("Attention grouping should be positive")
343
+ else:
344
+ attn_grouping_h, attn_grouping_q = self.num_heads, q_len
345
+
346
+ attn_output_h = []
347
+ for beg_h in range(0, self.num_heads, attn_grouping_h):
348
+ end_h = min(beg_h + attn_grouping_h, self.num_heads)
349
+
350
+ attn_output_q = []
351
+ for beg_q in range(0, q_len, attn_grouping_q):
352
+ end_q = min(beg_q + attn_grouping_q, q_len)
353
+
354
+ if self.config.torch_attention:
355
+ if mem_cache is not None:
356
+ attn_keys = torch.concat(
357
+ [key_states[:, beg_h:end_h], mem_cache.keys[:, beg_h:end_h].to(key_states.dtype)], dim=-2
358
+ )
359
+ attn_values = torch.concat(
360
+ [value_states[:, beg_h:end_h], mem_cache.values[:, beg_h:end_h].to(value_states.dtype)],
361
+ dim=-2,
362
+ )
363
+ mem_mask = mem_cache.masks.squeeze(-1).unsqueeze(-2)
364
+ assert len(mem_mask.shape) == 4
365
+ assert mem_mask.shape[2] == 1
366
+ assert mem_mask.shape[3] == mem_cache.keys.shape[-2]
367
+ mem_mask = torch.broadcast_to(
368
+ mem_mask, (mem_mask.shape[0], mem_mask.shape[1], end_q - beg_q, mem_mask.shape[3])
369
+ )
370
+ attn_mask = torch.concat([attention_mask[:, :, beg_q:end_q], mem_mask], dim=-1)
371
+ assert attn_mask.shape[-1] == attn_keys.shape[-2]
372
+ else:
373
+ attn_keys = key_states[:, beg_h:end_h]
374
+ attn_values = value_states[:, beg_h:end_h]
375
+ attn_mask = attention_mask[:, :, beg_q:end_q]
376
+
377
+ attn_queries = query_states[:, beg_h:end_h, beg_q:end_q]
378
+
379
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
380
+ query=attn_queries, key=attn_keys, value=attn_values, attn_mask=attn_mask
381
+ )
382
+ attn_output_q.append(attn_output)
383
+ else:
384
+ attn_weights = torch.matmul(
385
+ query_states[:, beg_h:end_h, beg_q:end_q], key_states[:, beg_h:end_h].transpose(2, 3)
386
+ ) / math.sqrt(self.head_dim)
387
+
388
+ if attn_weights.size() != (bsz, end_h - beg_h, end_q - beg_q, kv_seq_len):
389
+ raise ValueError(
390
+ f"Attention weights should be of size {(bsz, end_h - beg_h, end_q - beg_q, kv_seq_len)}, but is"
391
+ f" {attn_weights.size()}"
392
+ )
393
+
394
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
395
+ raise ValueError(
396
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
397
+ )
398
+ attn_weights = attn_weights + attention_mask[:, :, beg_q:end_q]
399
+ min_value = (
400
+ torch.finfo(attn_weights.dtype).min
401
+ if -1000000.0 < torch.finfo(attn_weights.dtype).min
402
+ else -1000000.0
403
+ )
404
+ attn_weights = torch.max(
405
+ attn_weights, torch.tensor(min_value, device=attn_weights.device, dtype=attn_weights.dtype)
406
+ )
407
+
408
+ if mem_cache is not None:
409
+ mem_mask = mem_cache.masks.squeeze(-1).unsqueeze(-2)
410
+ mem_attn_weights = torch.matmul(
411
+ query_states[:, beg_h:end_h, beg_q:end_q],
412
+ mem_cache.keys[:, beg_h:end_h].transpose(2, 3).to(key_states.dtype),
413
+ ) / math.sqrt(self.head_dim)
414
+
415
+ assert mem_mask.shape[2] == 1
416
+ mem_attn_weights = mem_attn_weights + mem_mask
417
+ min_value = (
418
+ torch.finfo(mem_attn_weights.dtype).min
419
+ if -1000000.0 < torch.finfo(mem_attn_weights.dtype).min
420
+ else -1000000.0
421
+ )
422
+ mem_attn_weights = torch.max(
423
+ mem_attn_weights,
424
+ torch.tensor(min_value, device=mem_attn_weights.device, dtype=mem_attn_weights.dtype),
425
+ )
426
+
427
+ attn_weights = torch.concat([attn_weights, mem_attn_weights], dim=-1)
428
+ combined_value_states = torch.concat(
429
+ [value_states[:, beg_h:end_h], mem_cache.values[:, beg_h:end_h].to(value_states.dtype)],
430
+ dim=-2,
431
+ )
432
+ else:
433
+ combined_value_states = value_states[:, beg_h:end_h]
434
+ # upcast attention to fp32
435
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
436
+ query_states.dtype
437
+ )
438
+ attn_output = torch.matmul(attn_weights, combined_value_states)
439
+ assert attn_output.shape[-2] == end_q - beg_q
440
+ attn_output_q.append(attn_output)
441
+ attn_output_h.append(torch.concat(attn_output_q, dim=-2))
442
+
443
+ attn_output = torch.concat(attn_output_h, dim=-3)
444
+
445
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
446
+ raise ValueError(
447
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
448
+ f" {attn_output.size()}"
449
+ )
450
+
451
+ attn_output = attn_output.transpose(1, 2)
452
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
453
+
454
+ attn_output = self.o_proj(attn_output)
455
+
456
+ if not output_attentions:
457
+ attn_weights = None
458
+
459
+ if mem_no_local_cache:
460
+ if mem_cache is not None:
461
+ mem_cache = mem_apply_update(
462
+ prev_mem_cache=mem_cache, new_mem_content=mem_update, mem_config=self.mem_config
463
+ )
464
+ else:
465
+ mem_cache = mem_update
466
+
467
+ return attn_output, attn_weights, past_key_value, mem_cache
468
+
469
+
470
+ # Modified transformers.models.llama.modeling_llama.LlamaDecoderLayer
471
+ class LongLlamaDecoderLayer(nn.Module):
472
+ def __init__(self, config: LongLlamaConfig, mem_config: Optional[LongLlamaMemConfig] = None):
473
+ super().__init__()
474
+ self.hidden_size = config.hidden_size
475
+ self.self_attn = LongLlamaAttention(config=config, mem_config=mem_config)
476
+ self.mlp = LongLlamaMLP(
477
+ hidden_size=self.hidden_size,
478
+ intermediate_size=config.intermediate_size,
479
+ hidden_act=config.hidden_act,
480
+ )
481
+ self.input_layernorm = LongLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
482
+ self.post_attention_layernorm = LongLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
483
+
484
+ def forward(
485
+ self,
486
+ hidden_states: torch.Tensor,
487
+ attention_mask: Optional[torch.Tensor] = None,
488
+ position_ids: Optional[torch.LongTensor] = None,
489
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
490
+ output_attentions: Optional[bool] = False,
491
+ use_cache: Optional[bool] = False,
492
+ mem_cache: Optional[LongLlamaMemCache] = None,
493
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
494
+ """
495
+ Args:
496
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
497
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
498
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
499
+ output_attentions (`bool`, *optional*):
500
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
501
+ returned tensors for more detail.
502
+ use_cache (`bool`, *optional*):
503
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
504
+ (see `past_key_values`).
505
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
506
+ along with information about positions
507
+ mem_cache (`LongLlamaMemCache`, *optional*): memory cache for specific layers
508
+ """
509
+
510
+ residual = hidden_states
511
+
512
+ hidden_states = self.input_layernorm(hidden_states)
513
+
514
+ # Self Attention
515
+ hidden_states, self_attn_weights, present_key_value, mem_cache = self.self_attn(
516
+ hidden_states=hidden_states,
517
+ attention_mask=attention_mask,
518
+ position_ids=position_ids,
519
+ past_key_value=past_key_value,
520
+ output_attentions=output_attentions,
521
+ use_cache=use_cache,
522
+ mem_cache=mem_cache,
523
+ )
524
+ hidden_states = residual + hidden_states
525
+
526
+ # Fully Connected
527
+ residual = hidden_states
528
+ hidden_states = self.post_attention_layernorm(hidden_states)
529
+ hidden_states = self.mlp(hidden_states)
530
+ hidden_states = residual + hidden_states
531
+
532
+ outputs = (hidden_states,)
533
+
534
+ if output_attentions:
535
+ outputs += (self_attn_weights,)
536
+
537
+ if use_cache:
538
+ outputs += (present_key_value,)
539
+
540
+ return outputs + (mem_cache,)
541
+
542
+
543
+ LONGLLAMA_START_DOCSTRING = r"""
544
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
545
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
546
+ etc.)
547
+
548
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
549
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
550
+ and behavior.
551
+
552
+ Parameters:
553
+ config ([`LongLlamaConfig`]):
554
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
555
+ load the weights associated with the model, only the configuration. Check out the
556
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
557
+ """
558
+ LONGLLAMA_MEML_DOCSTRING = r"""
559
+ mem_layers ([`int`], *optional*):
560
+ Indices of layers to be augmented with memory, if None then parameters from config will be used
561
+ mem_dtype (`str`, *optional*):
562
+ Keys and values will be casted to this type for storage.
563
+
564
+ """
565
+
566
+
567
+ @add_start_docstrings(
568
+ "The bare LongLLaMA Model outputting raw hidden-states without any specific head on top.",
569
+ LONGLLAMA_START_DOCSTRING,
570
+ )
571
+ # Copied from transformers.models.llama.modeling_llama.LlamaPreTrainedModel with Llama->LongLlama
572
+ class LongLlamaPreTrainedModel(PreTrainedModel):
573
+ config_class = LongLlamaConfig
574
+ base_model_prefix = "model"
575
+ supports_gradient_checkpointing = True
576
+ _no_split_modules = ["LongLlamaDecoderLayer"]
577
+ _skip_keys_device_placement = "past_key_values"
578
+ _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
579
+
580
+ def _init_weights(self, module):
581
+ std = self.config.initializer_range
582
+ if isinstance(module, nn.Linear):
583
+ module.weight.data.normal_(mean=0.0, std=std)
584
+ if module.bias is not None:
585
+ module.bias.data.zero_()
586
+ elif isinstance(module, nn.Embedding):
587
+ module.weight.data.normal_(mean=0.0, std=std)
588
+ if module.padding_idx is not None:
589
+ module.weight.data[module.padding_idx].zero_()
590
+
591
+ def _set_gradient_checkpointing(self, module, value=False):
592
+ if isinstance(module, LongLlamaModel):
593
+ module.gradient_checkpointing = value
594
+
595
+
596
+ LONGLLAMA_COMMON_INPUTS_DOCSTRING = r"""
597
+ Args:
598
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
599
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
600
+ it.
601
+
602
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
603
+ [`PreTrainedTokenizer.__call__`] for details.
604
+
605
+ [What are input IDs?](../glossary#input-ids)
606
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
607
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
608
+
609
+ - 1 for tokens that are **not masked**,
610
+ - 0 for tokens that are **masked**.
611
+
612
+ [What are attention masks?](../glossary#attention-mask)
613
+
614
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
615
+ [`PreTrainedTokenizer.__call__`] for details.
616
+
617
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
618
+ `past_key_values`).
619
+
620
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
621
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
622
+ information on the default strategy.
623
+
624
+ - 1 indicates the head is **not masked**,
625
+ - 0 indicates the head is **masked**.
626
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
627
+ Indices of positions of each input sequence tokens in the position embeddings.
628
+
629
+ [What are position IDs?](../glossary#position-ids)
630
+ past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`
631
+ or memory cache is enabled):
632
+ Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
633
+ `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 1 additional tensor of shape
634
+ `(batch_size, 1, sequence_length, 1)`. For memory enriched layers it also contains content of memory cache.
635
+ It is padded with empty tensors so when returned it alwyas has 6 elements.
636
+
637
+ Contains pre-computed hidden-states (key and values in the self-attention blocks)
638
+ that can be used (see `past_key_values` input) to speed up sequential decoding.
639
+
640
+ If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
641
+ don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
642
+ `decoder_input_ids` of shape `(batch_size, sequence_length)`.
643
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
644
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
645
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
646
+ model's internal embedding lookup matrix.
647
+ use_cache (`bool`, *optional*):
648
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
649
+ `past_key_values`).
650
+ output_attentions (`bool`, *optional*):
651
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
652
+ tensors for more detail. This is NOT supported in LongLlamaForCausalLM and LongLlamaForSequenceClassification
653
+ due to the specific input processing.
654
+ output_hidden_states (`bool`, *optional*):
655
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
656
+ more detail.
657
+ return_dict (`bool`, *optional*):
658
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
659
+ """
660
+ LONGLLAMA_MODEL_INPUTS_DOCSTRING = r"""
661
+ mem_caches (`tuple(LongLlamaMemCache)`, *optional*)
662
+ Memory caches for specified layers, None for others
663
+ """
664
+
665
+ LONGLLAMA_ADD_INPUTS_DOCSTRING = r"""
666
+ last_context_length (`int`, *optional*)
667
+ Useful for generation, specifies number of tokens that won't be loaded to memory and
668
+ will be left for generation cache
669
+ """
670
+
671
+
672
+ def _prepare_pos_ids(past_key_values, batch_size, input_length, device):
673
+ if past_key_values is not None:
674
+ # take previous max pos_id + 1
675
+ if past_key_values[0][2].shape[0] != batch_size:
676
+ raise ValueError(
677
+ f"first dimension of past_key_values should match batch size: {batch_size}"
678
+ f"but got {past_key_values[0][2].shape[0]}"
679
+ )
680
+ next_pos = torch.max(past_key_values[0][2].view(batch_size, -1), dim=-1)[0] + 1
681
+ next_pos = next_pos.view(batch_size, 1)
682
+ else:
683
+ next_pos = torch.zeros(batch_size, 1, device=device, dtype=torch.long)
684
+
685
+ position_ids = torch.arange(0, input_length, dtype=torch.long, device=device).view(1, input_length)
686
+ position_ids = position_ids + next_pos
687
+ return position_ids
688
+
689
+
690
+ @add_start_docstrings(
691
+ "The bare LongLLaMA Model outputting raw hidden-states without any specific head on top.",
692
+ LONGLLAMA_START_DOCSTRING,
693
+ LONGLLAMA_MEML_DOCSTRING,
694
+ )
695
+ # Modified transformers.models.llama.modeling_llama.LlamaModel
696
+ class LongLlamaModel(LongLlamaPreTrainedModel):
697
+ """
698
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LongLlamaDecoderLayer`]
699
+
700
+ Args:
701
+ config: LlamaConfig
702
+ """
703
+
704
+ def __init__(self, config: LongLlamaConfig):
705
+ super().__init__(config)
706
+ self.mem_layers = config.mem_layers
707
+ self.mem_config = LongLlamaMemConfig(
708
+ positionals=config.mem_positionals,
709
+ cache_dtype=getattr(torch, config.mem_dtype),
710
+ attention_grouping=config.mem_attention_grouping,
711
+ )
712
+ self.padding_idx = config.pad_token_id
713
+ self.vocab_size = config.vocab_size
714
+
715
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
716
+
717
+ for mem_layer_id in self.mem_layers:
718
+ if mem_layer_id < 0 or mem_layer_id >= config.num_hidden_layers:
719
+ raise ValueError(
720
+ f"Memory layer ids should be between 0 and {config.num_hidden_layers}, got {mem_layer_id}"
721
+ )
722
+
723
+ layers = []
724
+ for layer_id in range(config.num_hidden_layers):
725
+ if layer_id in self.mem_layers:
726
+ layer = LongLlamaDecoderLayer(config, mem_config=self.mem_config)
727
+ else:
728
+ layer = LongLlamaDecoderLayer(config, mem_config=None)
729
+ layers.append(layer)
730
+
731
+ self.layers = nn.ModuleList(layers)
732
+ self.norm = LongLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
733
+
734
+ self.gradient_checkpointing = False
735
+
736
+ # Initialize weights and apply final processing
737
+ self.post_init()
738
+
739
+ def get_input_embeddings(self):
740
+ return self.embed_tokens
741
+
742
+ def set_input_embeddings(self, value):
743
+ self.embed_tokens = value
744
+
745
+ # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
746
+ def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
747
+ # create causal mask
748
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
749
+ combined_attention_mask = None
750
+ if input_shape[-1] > 1:
751
+ combined_attention_mask = _make_causal_mask(
752
+ input_shape,
753
+ inputs_embeds.dtype,
754
+ device=inputs_embeds.device,
755
+ past_key_values_length=past_key_values_length,
756
+ )
757
+
758
+ if attention_mask is not None:
759
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
760
+ expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
761
+ inputs_embeds.device
762
+ )
763
+ combined_attention_mask = (
764
+ expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
765
+ )
766
+
767
+ return combined_attention_mask
768
+
769
+ @add_start_docstrings_to_model_forward(LONGLLAMA_COMMON_INPUTS_DOCSTRING, LONGLLAMA_MODEL_INPUTS_DOCSTRING)
770
+ def forward(
771
+ self,
772
+ input_ids: torch.LongTensor = None,
773
+ attention_mask: Optional[torch.Tensor] = None,
774
+ position_ids: Optional[torch.LongTensor] = None,
775
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
776
+ inputs_embeds: Optional[torch.FloatTensor] = None,
777
+ use_cache: Optional[bool] = None,
778
+ output_attentions: Optional[bool] = None,
779
+ output_hidden_states: Optional[bool] = None,
780
+ return_dict: Optional[bool] = None,
781
+ mem_caches: Optional[Tuple[Optional[LongLlamaMemCache]]] = None,
782
+ ) -> Union[Tuple, LongLlamaModelOutputWithPast]:
783
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
784
+ output_hidden_states = (
785
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
786
+ )
787
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
788
+
789
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
790
+
791
+ # retrieve input_ids and inputs_embeds
792
+ if input_ids is not None and inputs_embeds is not None:
793
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
794
+ elif input_ids is not None:
795
+ batch_size, seq_length = input_ids.shape
796
+ elif inputs_embeds is not None:
797
+ batch_size, seq_length, _ = inputs_embeds.shape
798
+ else:
799
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
800
+
801
+ seq_length_with_past = seq_length
802
+ past_key_values_length = 0
803
+
804
+ if past_key_values is not None:
805
+ past_key_values_length = past_key_values[0][0].shape[-2]
806
+ seq_length_with_past = seq_length_with_past + past_key_values_length
807
+
808
+ if position_ids is None:
809
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
810
+ position_ids = _prepare_pos_ids(past_key_values, batch_size, seq_length, device)
811
+ else:
812
+ position_ids = position_ids.view(-1, seq_length).long()
813
+
814
+ if inputs_embeds is None:
815
+ inputs_embeds = self.embed_tokens(input_ids)
816
+ # embed positions
817
+ if attention_mask is None:
818
+ attention_mask = torch.ones(
819
+ (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
820
+ )
821
+ attention_mask = self._prepare_decoder_attention_mask(
822
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
823
+ )
824
+
825
+ hidden_states = inputs_embeds
826
+
827
+ if self.gradient_checkpointing and self.training:
828
+ if use_cache:
829
+ logger.warning_once(
830
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
831
+ )
832
+ use_cache = False
833
+
834
+ # decoder layers
835
+ all_hidden_states = () if output_hidden_states else None
836
+ all_self_attns = () if output_attentions else None
837
+ next_decoder_cache = ()
838
+ next_mem_caches = ()
839
+ for idx, decoder_layer in enumerate(self.layers):
840
+ if output_hidden_states:
841
+ all_hidden_states += (hidden_states,)
842
+
843
+ past_key_value = past_key_values[idx] if past_key_values is not None else None
844
+ mem_cache = mem_caches[idx] if mem_caches else None
845
+
846
+ if mem_cache is not None and idx not in self.mem_layers:
847
+ raise ValueError("Memory cache provided for a non-memory leayer")
848
+
849
+ if (
850
+ self.gradient_checkpointing
851
+ and self.training
852
+ and mem_cache is None
853
+ and idx % self.config.gradient_checkpoint_every_ith == 0
854
+ ):
855
+
856
+ def create_custom_forward(module):
857
+ def custom_forward(*inputs):
858
+ # None for past_key_value
859
+ return module(*inputs, output_attentions, None, mem_cache=None)
860
+
861
+ return custom_forward
862
+
863
+ layer_outputs = torch.utils.checkpoint.checkpoint(
864
+ create_custom_forward(decoder_layer),
865
+ hidden_states,
866
+ attention_mask,
867
+ position_ids,
868
+ None,
869
+ )
870
+ else:
871
+ layer_outputs = decoder_layer(
872
+ hidden_states,
873
+ attention_mask=attention_mask,
874
+ position_ids=position_ids,
875
+ past_key_value=past_key_value,
876
+ output_attentions=output_attentions,
877
+ use_cache=use_cache,
878
+ mem_cache=mem_cache,
879
+ )
880
+
881
+ new_mem_cache = layer_outputs[-1]
882
+ layer_outputs = layer_outputs[:-1]
883
+ next_mem_caches += (new_mem_cache,)
884
+
885
+ hidden_states = layer_outputs[0]
886
+
887
+ if use_cache:
888
+ next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
889
+ else:
890
+ next_decoder_cache += (None,)
891
+
892
+ if output_attentions:
893
+ all_self_attns += (layer_outputs[1],)
894
+
895
+ hidden_states = self.norm(hidden_states)
896
+
897
+ # add hidden states from the last decoder layer
898
+ if output_hidden_states:
899
+ all_hidden_states += (hidden_states,)
900
+
901
+ next_cache = next_decoder_cache if use_cache else None
902
+
903
+ mem_cache_returned = False
904
+ for mem_cache in next_mem_caches:
905
+ if mem_cache is not None:
906
+ mem_cache_returned = True
907
+ next_mem_caches = next_mem_caches if mem_cache_returned else None
908
+
909
+ if not return_dict:
910
+ return tuple(
911
+ v
912
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, next_mem_caches]
913
+ if v is not None
914
+ )
915
+ return LongLlamaModelOutputWithPast(
916
+ last_hidden_state=hidden_states,
917
+ past_key_values=next_cache,
918
+ hidden_states=all_hidden_states,
919
+ attentions=all_self_attns,
920
+ mem_caches=next_mem_caches,
921
+ )
922
+
923
+
924
+ def _handle_output_of_past_key_values(outputs):
925
+ # merges local caches and memory caches into one single tuple of past_key_values
926
+ # in order to support generation
927
+ batch_size = outputs.last_hidden_state.shape[0]
928
+ if outputs.past_key_values is None and outputs.mem_caches is None:
929
+ return None
930
+
931
+ if outputs.past_key_values is None:
932
+ out_past_key_values = (None,) * len(outputs.mem_caches)
933
+ else:
934
+ out_past_key_values = outputs.past_key_values
935
+
936
+ if outputs.mem_caches is None:
937
+ out_mem_caches = (None,) * len(outputs.past_key_values)
938
+ else:
939
+ out_mem_caches = outputs.mem_caches
940
+
941
+ device = outputs.last_hidden_state.device
942
+ past_key_values = ()
943
+ for local_cache, mem_cache in zip(out_past_key_values, out_mem_caches):
944
+ layer = ()
945
+ if local_cache is not None:
946
+ assert len(local_cache) == 3
947
+ layer += local_cache
948
+ else:
949
+ layer += (torch.empty(batch_size, 0, 0, 0, device=device),) * 3
950
+
951
+ if mem_cache is not None:
952
+ layer += (mem_cache.keys, mem_cache.values, mem_cache.masks)
953
+ else:
954
+ layer += (torch.empty(batch_size, 0, 0, 0, device=device),) * 3
955
+
956
+ assert len(layer) == 6
957
+
958
+ past_key_values += (layer,)
959
+
960
+ return past_key_values
961
+
962
+
963
+ def _split_past_key_values(past_key_values):
964
+ # splits past_key_values to local cache and memory cache
965
+ local_cache_preset = False
966
+ mem_caches_present = False
967
+ if past_key_values is not None:
968
+ local_caches = ()
969
+ mem_caches = ()
970
+ for layer in past_key_values:
971
+ if len(layer) != 6:
972
+ raise ValueError(
973
+ "Expected elements of past_key_values to contain 6 elements."
974
+ "First 3 describing local cache and last 3 describing memory cache."
975
+ f"Instead got {len(layer)} elements"
976
+ )
977
+ else:
978
+ lk, lv, li, memk, memv, memm = layer
979
+ if lk.shape[-2] != 0:
980
+ local_cache_preset = True
981
+ local_caches += ((lk, lv, li),)
982
+ else:
983
+ local_caches += (None,)
984
+
985
+ if memk.shape[-2] != 0:
986
+ mem_caches_present = True
987
+ mem_caches += (LongLlamaMemCache(keys=memk, values=memv, masks=memm),)
988
+ else:
989
+ mem_caches += (None,)
990
+
991
+ local_caches = local_caches if local_cache_preset else None
992
+ mem_caches = mem_caches if mem_caches_present else None
993
+
994
+ return local_caches, mem_caches
995
+
996
+
997
+ def _handle_long_input(
998
+ model,
999
+ input_ids,
1000
+ attention_mask,
1001
+ position_ids,
1002
+ past_key_values,
1003
+ inputs_embeds,
1004
+ use_cache,
1005
+ output_attentions,
1006
+ output_hidden_states,
1007
+ return_dict,
1008
+ context_window_length,
1009
+ last_context_length,
1010
+ ):
1011
+ if output_attentions:
1012
+ logger.warning(
1013
+ f"Outputing attentions is not supported in LongLlamaForCausalLM and LongLlamaForSequenceClassification. "
1014
+ f"Attention of the last window will be returned"
1015
+ )
1016
+
1017
+ past_key_values, mem_caches = _split_past_key_values(past_key_values)
1018
+
1019
+ if past_key_values is not None and use_cache is False:
1020
+ raise ValueError("past_key_values it not None should imply use_cache == True")
1021
+
1022
+ if past_key_values is not None:
1023
+ initial_past_key_values_length = past_key_values[0][0].shape[-2]
1024
+ else:
1025
+ initial_past_key_values_length = 0
1026
+
1027
+ if input_ids is not None:
1028
+ batch_size, input_length = input_ids.shape
1029
+ else:
1030
+ batch_size, input_length, _ = inputs_embeds.shape
1031
+
1032
+ if position_ids is None:
1033
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1034
+ position_ids = _prepare_pos_ids(past_key_values, batch_size, input_length, device)
1035
+
1036
+ if position_ids.shape != (batch_size, input_length):
1037
+ raise ValueError(f"Shape of position_ids [{position_ids}] should match [{batch_size, input_length}]")
1038
+
1039
+ if attention_mask is not None:
1040
+ attention_mask = attention_mask[..., -(initial_past_key_values_length + input_length) :]
1041
+ if attention_mask is not None and (
1042
+ attention_mask.shape != (batch_size, initial_past_key_values_length + input_length)
1043
+ ):
1044
+ raise ValueError(
1045
+ "Attention mask should be provided for both the local cache and the input",
1046
+ f"Expected shape {(batch_size, initial_past_key_values_length + input_length)},"
1047
+ f"got {attention_mask.shape}.",
1048
+ )
1049
+
1050
+ # First we load prefix to memory cache
1051
+ mem_input_length = max(input_length - last_context_length, 0)
1052
+ outputs_list = []
1053
+ attn_offset = initial_past_key_values_length
1054
+ if mem_input_length > 0:
1055
+ for i in range(0, mem_input_length, context_window_length):
1056
+ beg, end = i, min(mem_input_length, i + context_window_length)
1057
+
1058
+ if attention_mask is not None:
1059
+ if past_key_values is not None:
1060
+ local_cache_size = past_key_values[0][0].shape[-2]
1061
+ else:
1062
+ local_cache_size = 0
1063
+ attn_length = attention_mask.shape[-1]
1064
+ attn_beg = beg + attn_offset - local_cache_size
1065
+ attn_end = end + attn_offset
1066
+ assert attn_end <= attn_length
1067
+ assert attn_beg >= 0 and attn_end > attn_beg
1068
+
1069
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn, mem_caches)
1070
+ outputs = model(
1071
+ input_ids=input_ids[..., beg:end] if input_ids is not None else None,
1072
+ attention_mask=attention_mask[..., attn_beg:attn_end] if attention_mask is not None else None,
1073
+ position_ids=position_ids[..., beg:end],
1074
+ past_key_values=past_key_values,
1075
+ inputs_embeds=inputs_embeds[..., beg:end, :] if inputs_embeds is not None else None,
1076
+ use_cache=False if past_key_values is None else use_cache,
1077
+ output_attentions=output_attentions,
1078
+ output_hidden_states=output_hidden_states,
1079
+ return_dict=True,
1080
+ mem_caches=mem_caches,
1081
+ )
1082
+ if i > 0:
1083
+ if mem_caches is not None and past_key_values is None:
1084
+ for mc_layer in mem_caches:
1085
+ if mc_layer is not None:
1086
+ del mc_layer.keys
1087
+ del mc_layer.values
1088
+ del mc_layer.masks
1089
+
1090
+ mem_caches = outputs.mem_caches
1091
+ outputs.mem_caches = None
1092
+ past_key_values = outputs.past_key_values
1093
+ outputs.past_key_values = None
1094
+ outputs_list.append(outputs)
1095
+
1096
+ remaining_input_length = input_length - mem_input_length
1097
+ beg = mem_input_length
1098
+ attn_length = remaining_input_length
1099
+ if past_key_values is not None:
1100
+ attn_length += past_key_values[0][0].shape[-2]
1101
+ attention_mask = attention_mask[..., -attn_length:] if attention_mask is not None else None
1102
+
1103
+ outputs = model(
1104
+ input_ids=input_ids[..., beg:] if input_ids is not None else None,
1105
+ attention_mask=attention_mask,
1106
+ position_ids=position_ids[..., beg:],
1107
+ past_key_values=past_key_values,
1108
+ inputs_embeds=inputs_embeds[..., beg:, :] if inputs_embeds is not None else None,
1109
+ use_cache=use_cache,
1110
+ output_attentions=output_attentions,
1111
+ output_hidden_states=output_hidden_states,
1112
+ return_dict=True,
1113
+ mem_caches=mem_caches,
1114
+ )
1115
+
1116
+ outputs_list.append(outputs)
1117
+
1118
+ past_key_values = _handle_output_of_past_key_values(outputs_list[-1])
1119
+
1120
+ if output_hidden_states:
1121
+ hidden_states = ()
1122
+ for hd in zip(*[x.hidden_states for x in outputs_list]):
1123
+ hidden_states += (torch.cat(hd, dim=-2),)
1124
+ else:
1125
+ hidden_states = None
1126
+
1127
+ outputs = BaseModelOutputWithPast(
1128
+ last_hidden_state=torch.concat([x.last_hidden_state for x in outputs_list], dim=-2),
1129
+ past_key_values=past_key_values,
1130
+ hidden_states=hidden_states,
1131
+ attentions=outputs_list[-1].attentions,
1132
+ )
1133
+
1134
+ if not return_dict:
1135
+ outputs = tuple(
1136
+ v
1137
+ for v in [outputs.last_hidden_state, outputs.past_key_values, outputs.hidden_states, outputs.attentions]
1138
+ if v is not None
1139
+ )
1140
+ return outputs
1141
+
1142
+
1143
+ # Modified transformers.models.llama.modeling_llama.LlamaForCausalLM
1144
+ class LongLlamaForCausalLM(LongLlamaPreTrainedModel):
1145
+ _tied_weights_keys = ["lm_head.weight"]
1146
+
1147
+ def __init__(self, config):
1148
+ super().__init__(config)
1149
+ self.context_window_length = config.max_position_embeddings
1150
+
1151
+ self.model = LongLlamaModel(config)
1152
+
1153
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1154
+
1155
+ # Initialize weights and apply final processing
1156
+ self.post_init()
1157
+
1158
+ def get_input_embeddings(self):
1159
+ return self.model.embed_tokens
1160
+
1161
+ def set_input_embeddings(self, value):
1162
+ self.model.embed_tokens = value
1163
+
1164
+ def get_output_embeddings(self):
1165
+ return self.lm_head
1166
+
1167
+ def set_output_embeddings(self, new_embeddings):
1168
+ self.lm_head = new_embeddings
1169
+
1170
+ def set_decoder(self, decoder):
1171
+ self.model = decoder
1172
+
1173
+ def get_decoder(self):
1174
+ return self.model
1175
+
1176
+ def _has_generation_cache(self, past_key_values):
1177
+ if past_key_values is not None:
1178
+ assert len(past_key_values[0]) == 6
1179
+ return past_key_values[0][0].shape[-2] != 0
1180
+
1181
+ return False
1182
+
1183
+ @add_start_docstrings_to_model_forward(LONGLLAMA_COMMON_INPUTS_DOCSTRING, LONGLLAMA_ADD_INPUTS_DOCSTRING)
1184
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1185
+ def forward(
1186
+ self,
1187
+ input_ids: torch.LongTensor = None,
1188
+ attention_mask: Optional[torch.Tensor] = None,
1189
+ position_ids: Optional[torch.LongTensor] = None,
1190
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1191
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1192
+ labels: Optional[torch.LongTensor] = None,
1193
+ use_cache: Optional[bool] = None,
1194
+ output_attentions: Optional[bool] = None,
1195
+ output_hidden_states: Optional[bool] = None,
1196
+ return_dict: Optional[bool] = None,
1197
+ last_context_length: Optional[int] = None,
1198
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1199
+ r"""
1200
+ Args:
1201
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1202
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1203
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1204
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1205
+
1206
+ Returns:
1207
+
1208
+ Example:
1209
+
1210
+ ```python
1211
+ >>> from transformers import AutoTokenizer, LlamaForCausalLM
1212
+
1213
+ >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
1214
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
1215
+
1216
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1217
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1218
+
1219
+ >>> # Generate
1220
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1221
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1222
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1223
+ ```"""
1224
+ last_context_length = (
1225
+ last_context_length if last_context_length is not None else self.config.last_context_length
1226
+ )
1227
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1228
+ output_hidden_states = (
1229
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1230
+ )
1231
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1232
+
1233
+ outputs = _handle_long_input(
1234
+ model=self.model,
1235
+ input_ids=input_ids,
1236
+ attention_mask=attention_mask,
1237
+ position_ids=position_ids,
1238
+ past_key_values=past_key_values,
1239
+ inputs_embeds=inputs_embeds,
1240
+ use_cache=use_cache,
1241
+ output_attentions=output_attentions,
1242
+ output_hidden_states=output_hidden_states,
1243
+ return_dict=return_dict,
1244
+ context_window_length=self.context_window_length,
1245
+ last_context_length=last_context_length,
1246
+ )
1247
+
1248
+ hidden_states = outputs[0]
1249
+ logits = self.lm_head(hidden_states)
1250
+
1251
+ loss = None
1252
+ if labels is not None:
1253
+ # Shift so that tokens < n predict n
1254
+ shift_logits = logits[..., :-1, :].contiguous()
1255
+ shift_labels = labels[..., 1:].contiguous()
1256
+ # Flatten the tokens
1257
+ loss_fct = CrossEntropyLoss()
1258
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1259
+ shift_labels = shift_labels.view(-1)
1260
+ # Enable model parallelism
1261
+ shift_labels = shift_labels.to(shift_logits.device)
1262
+ loss = loss_fct(shift_logits, shift_labels)
1263
+
1264
+ if not return_dict:
1265
+ output = (logits,) + outputs[1:]
1266
+ return (loss,) + output if loss is not None else output
1267
+
1268
+ return CausalLMOutputWithPast(
1269
+ loss=loss,
1270
+ logits=logits,
1271
+ past_key_values=outputs.past_key_values,
1272
+ hidden_states=outputs.hidden_states,
1273
+ attentions=outputs.attentions,
1274
+ )
1275
+
1276
+ def prepare_inputs_for_generation(
1277
+ self,
1278
+ input_ids,
1279
+ past_key_values=None,
1280
+ attention_mask=None,
1281
+ inputs_embeds=None,
1282
+ last_context_length=None,
1283
+ **kwargs,
1284
+ ):
1285
+ if self._has_generation_cache(past_key_values):
1286
+ input_ids = input_ids[:, -1:]
1287
+
1288
+ position_ids = kwargs.get("position_ids", None)
1289
+ if attention_mask is not None and position_ids is None:
1290
+ # create position_ids on the fly for batch generation
1291
+ position_ids = attention_mask.long().cumsum(-1) - 1
1292
+ position_ids.masked_fill(position_ids < 0, 0)
1293
+ if self._has_generation_cache(past_key_values):
1294
+ position_ids = position_ids[:, -1].unsqueeze(-1)
1295
+
1296
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1297
+ if inputs_embeds is not None and past_key_values is None:
1298
+ model_inputs = {"inputs_embeds": inputs_embeds}
1299
+ else:
1300
+ model_inputs = {"input_ids": input_ids}
1301
+
1302
+ model_inputs.update(
1303
+ {
1304
+ "position_ids": position_ids,
1305
+ "past_key_values": past_key_values,
1306
+ "use_cache": kwargs.get("use_cache"),
1307
+ "attention_mask": attention_mask,
1308
+ "last_context_length": last_context_length,
1309
+ }
1310
+ )
1311
+ return model_inputs
1312
+
1313
+ @staticmethod
1314
+ def _reorder_cache(past_key_values, beam_idx):
1315
+ reordered_past = ()
1316
+ for layer_past in past_key_values:
1317
+ reordered_past += (
1318
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
1319
+ )
1320
+ return reordered_past
1321
+
1322
+
1323
+ @add_start_docstrings(
1324
+ """
1325
+ The LongLLaMA Model transformer with a sequence classification head on top (linear layer).
1326
+
1327
+ [`LongLlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal models
1328
+ (e.g. GPT-2) do.
1329
+
1330
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1331
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1332
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1333
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1334
+ each row of the batch).
1335
+ """,
1336
+ LONGLLAMA_START_DOCSTRING,
1337
+ LONGLLAMA_MEML_DOCSTRING,
1338
+ )
1339
+ # Modified from transformers.models.llama.modeling_llama.LlamaForSequenceClassification
1340
+ class LongLlamaForSequenceClassification(LongLlamaPreTrainedModel):
1341
+ _keys_to_ignore_on_load_missing = [r"lm_head.weight"]
1342
+
1343
+ def __init__(self, config):
1344
+ super().__init__(config)
1345
+ self.num_labels = config.num_labels
1346
+ self.context_window_length = config.max_position_embeddings
1347
+ self.model = LongLlamaModel(config)
1348
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1349
+
1350
+ # Initialize weights and apply final processing
1351
+ self.post_init()
1352
+
1353
+ def get_input_embeddings(self):
1354
+ return self.model.embed_tokens
1355
+
1356
+ def set_input_embeddings(self, value):
1357
+ self.model.embed_tokens = value
1358
+
1359
+ @add_start_docstrings_to_model_forward(LONGLLAMA_COMMON_INPUTS_DOCSTRING, LONGLLAMA_ADD_INPUTS_DOCSTRING)
1360
+ def forward(
1361
+ self,
1362
+ input_ids: torch.LongTensor = None,
1363
+ attention_mask: Optional[torch.Tensor] = None,
1364
+ position_ids: Optional[torch.LongTensor] = None,
1365
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1366
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1367
+ labels: Optional[torch.LongTensor] = None,
1368
+ use_cache: Optional[bool] = None,
1369
+ output_attentions: Optional[bool] = None,
1370
+ output_hidden_states: Optional[bool] = None,
1371
+ return_dict: Optional[bool] = None,
1372
+ last_context_length: Optional[int] = None,
1373
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1374
+ r"""
1375
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1376
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1377
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1378
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1379
+ """
1380
+ last_context_length = (
1381
+ last_context_length if last_context_length is not None else self.config.last_context_length
1382
+ )
1383
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1384
+ output_hidden_states = (
1385
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1386
+ )
1387
+ transformer_outputs = _handle_long_input(
1388
+ model=self.model,
1389
+ input_ids=input_ids,
1390
+ attention_mask=attention_mask,
1391
+ position_ids=position_ids,
1392
+ past_key_values=past_key_values,
1393
+ inputs_embeds=inputs_embeds,
1394
+ use_cache=use_cache,
1395
+ output_attentions=output_attentions,
1396
+ output_hidden_states=output_hidden_states,
1397
+ return_dict=return_dict,
1398
+ context_window_length=self.context_window_length,
1399
+ last_context_length=last_context_length,
1400
+ )
1401
+
1402
+ hidden_states = transformer_outputs[0]
1403
+ logits = self.score(hidden_states)
1404
+
1405
+ if input_ids is not None:
1406
+ batch_size = input_ids.shape[0]
1407
+ else:
1408
+ batch_size = inputs_embeds.shape[0]
1409
+
1410
+ if self.config.pad_token_id is None and batch_size != 1:
1411
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1412
+ if self.config.pad_token_id is None:
1413
+ sequence_lengths = -1
1414
+ else:
1415
+ if input_ids is not None:
1416
+ sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device)
1417
+ else:
1418
+ sequence_lengths = -1
1419
+
1420
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1421
+
1422
+ loss = None
1423
+ if labels is not None:
1424
+ labels = labels.to(logits.device)
1425
+ if self.config.problem_type is None:
1426
+ if self.num_labels == 1:
1427
+ self.config.problem_type = "regression"
1428
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1429
+ self.config.problem_type = "single_label_classification"
1430
+ else:
1431
+ self.config.problem_type = "multi_label_classification"
1432
+
1433
+ if self.config.problem_type == "regression":
1434
+ loss_fct = MSELoss()
1435
+ if self.num_labels == 1:
1436
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1437
+ else:
1438
+ loss = loss_fct(pooled_logits, labels)
1439
+ elif self.config.problem_type == "single_label_classification":
1440
+ loss_fct = CrossEntropyLoss()
1441
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1442
+ elif self.config.problem_type == "multi_label_classification":
1443
+ loss_fct = BCEWithLogitsLoss()
1444
+ loss = loss_fct(pooled_logits, labels)
1445
+ if not return_dict:
1446
+ output = (pooled_logits,) + transformer_outputs[1:]
1447
+ return ((loss,) + output) if loss is not None else output
1448
+
1449
+ return SequenceClassifierOutputWithPast(
1450
+ loss=loss,
1451
+ logits=pooled_logits,
1452
+ past_key_values=transformer_outputs.past_key_values,
1453
+ hidden_states=transformer_outputs.hidden_states,
1454
+ attentions=transformer_outputs.attentions,
1455
+ )
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f7717a6a554bbe997a12ae743d784057ba93d18b34add9e28d5f8c2f5064246d
3
+ size 6853041805
special_tokens_map.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": true,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "</s>",
11
+ "lstrip": false,
12
+ "normalized": true,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "unk_token": {
17
+ "content": "<unk>",
18
+ "lstrip": false,
19
+ "normalized": true,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ }
23
+ }
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:91b289e85fa20fd375d8b33dc12f77616f18abc6359804471d1fafcb425fecb8
3
+ size 511574
tokenizer_config.json ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": true,
3
+ "add_eos_token": false,
4
+ "bos_token": {
5
+ "__type": "AddedToken",
6
+ "content": "",
7
+ "lstrip": false,
8
+ "normalized": true,
9
+ "rstrip": false,
10
+ "single_word": false
11
+ },
12
+ "clean_up_tokenization_spaces": {
13
+ "__type": "AddedToken",
14
+ "content": "<s>",
15
+ "lstrip": false,
16
+ "normalized": true,
17
+ "rstrip": false,
18
+ "single_word": false
19
+ },
20
+ "eos_token": {
21
+ "__type": "AddedToken",
22
+ "content": "</s>",
23
+ "lstrip": false,
24
+ "normalized": true,
25
+ "rstrip": false,
26
+ "single_word": false
27
+ },
28
+ "model_max_length": 1000000000000000019884624838656,
29
+ "pad_token": null,
30
+ "padding_side": "right",
31
+ "sp_model_kwargs": {},
32
+ "tokenizer_class": "LlamaTokenizer",
33
+ "unk_token": {
34
+ "__type": "AddedToken",
35
+ "content": "<unk>",
36
+ "lstrip": false,
37
+ "normalized": true,
38
+ "rstrip": false,
39
+ "single_word": false
40
+ },
41
+ "use_fast": false
42
+ }