v2ray commited on
Commit
fccb166
1 Parent(s): c17f94b

Uploaded the merged bf16 version.

Browse files
README.md ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ datasets:
4
+ - v2ray/r-chatgpt-general-dump
5
+ language:
6
+ - en
7
+ pipeline_tag: text-generation
8
+ tags:
9
+ - not-for-all-audiences
10
+ ---
11
+ # SchizoGPT 132B QLoRA
12
+ Fine tuned on r/ChatGPT Discord #general dump.
13
+ ## Prompt Template
14
+ ```
15
+ ### Username
16
+ Message(s)(Can be stacked, separated by new line)
17
+ ### Username(Can have multiturn multiuser chat, separated by new line)
18
+ (End of prompt)
19
+ ```
20
+ Use `@username` to ping a user and `#channel name` to mention a channel. \
21
+ Prepend `<Re: username>` before a message to respond to a user. \
22
+ Use `<filename.ext>` to mention a file in a link, for example, if you have `https://example.com/image.jpg`, use `<image.jpg>`.
__init__.py ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ from .configuration_dbrx import *
2
+ from .modeling_dbrx import *
config.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "v2ray/dbrx-base-fixed",
3
+ "architectures": [
4
+ "DbrxForCausalLM"
5
+ ],
6
+ "attn_config": {
7
+ "clip_qkv": 8,
8
+ "kv_n_heads": 8,
9
+ "model_type": "",
10
+ "rope_theta": 500000
11
+ },
12
+ "auto_map": {
13
+ "AutoConfig": "v2ray/dbrx-base-fixed--configuration_dbrx.DbrxConfig",
14
+ "AutoModelForCausalLM": "v2ray/dbrx-base-fixed--modeling_dbrx.DbrxForCausalLM"
15
+ },
16
+ "d_model": 6144,
17
+ "emb_pdrop": 0.0,
18
+ "ffn_config": {
19
+ "ffn_hidden_size": 10752,
20
+ "model_type": "",
21
+ "moe_jitter_eps": 0.01,
22
+ "moe_loss_weight": 0.05,
23
+ "moe_num_experts": 16,
24
+ "moe_top_k": 4
25
+ },
26
+ "initializer_range": 0.02,
27
+ "max_seq_len": 32768,
28
+ "model_type": "dbrx",
29
+ "n_heads": 48,
30
+ "n_layers": 40,
31
+ "output_router_logits": false,
32
+ "resid_pdrop": 0.0,
33
+ "router_aux_loss_coef": 0.05,
34
+ "tie_word_embeddings": false,
35
+ "torch_dtype": "bfloat16",
36
+ "transformers_version": "4.39.2",
37
+ "use_cache": true,
38
+ "vocab_size": 100352
39
+ }
configuration_dbrx.py ADDED
@@ -0,0 +1,264 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Dbrx configuration."""
2
+ from typing import Any, Optional
3
+
4
+ from transformers.configuration_utils import PretrainedConfig
5
+ from transformers.utils import logging
6
+
7
+ logger = logging.get_logger(__name__)
8
+
9
+ DBRX_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
10
+
11
+
12
+ class DbrxAttentionConfig(PretrainedConfig):
13
+ """Configuration class for Dbrx Attention.
14
+
15
+ [`DbrxAttention`] class. It is used to instantiate attention layers
16
+ according to the specified arguments, defining the layers architecture.
17
+
18
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
19
+ documentation from [`PretrainedConfig`] for more information.
20
+
21
+ Args:
22
+ attn_pdrop (`float`, *optional*, defaults to 0.0):
23
+ The dropout probability for the attention layers.
24
+ clip_qkv (`float`, *optional*, defualts to None):
25
+ If not `None`, clip the queries, keys, and values in the attention layer to this value.
26
+ kv_n_heads (Optional[int]): For grouped_query_attention only, allow user to specify number of kv heads.
27
+ rope_theta (float): The base frequency for rope.
28
+ """
29
+
30
+ def __init__(
31
+ self,
32
+ attn_pdrop: float = 0,
33
+ clip_qkv: Optional[float] = None,
34
+ kv_n_heads: int = 1,
35
+ rope_theta: float = 10000.0,
36
+ **kwargs: Any,
37
+ ):
38
+ super().__init__(**kwargs)
39
+ self.attn_pdrop = attn_pdrop
40
+ self.clip_qkv = clip_qkv
41
+ self.kv_n_heads = kv_n_heads
42
+ self.rope_theta = rope_theta
43
+
44
+ for k in ['model_type']:
45
+ if k in kwargs:
46
+ kwargs.pop(k)
47
+ if len(kwargs) != 0:
48
+ raise ValueError(f'Found unknown {kwargs=}')
49
+
50
+ @classmethod
51
+ def from_pretrained(cls, pretrained_model_name_or_path: str,
52
+ **kwargs: Any) -> 'PretrainedConfig':
53
+ cls._set_token_in_kwargs(kwargs)
54
+
55
+ config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path,
56
+ **kwargs)
57
+
58
+ if config_dict.get('model_type') == 'dbrx':
59
+ config_dict = config_dict['attn_config']
60
+
61
+ if 'model_type' in config_dict and hasattr(
62
+ cls,
63
+ 'model_type') and config_dict['model_type'] != cls.model_type:
64
+ logger.warning(
65
+ f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
66
+ +
67
+ f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'
68
+ )
69
+
70
+ return cls.from_dict(config_dict, **kwargs)
71
+
72
+
73
+ class DbrxFFNConfig(PretrainedConfig):
74
+ """Configuration class for Dbrx FFN.
75
+
76
+ [`DbrxFFN`] class. It is used to instantiate feedforward layers according to
77
+ the specified arguments, defining the layers architecture.
78
+
79
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
80
+ documentation from [`PretrainedConfig`] for more information.
81
+
82
+ Args:
83
+ ffn_act_fn (dict, optional): A dict specifying activation function for the FFN.
84
+ The dict should have a key 'name' with the value being the name of
85
+ the activation function along with any additional keyword arguments.
86
+ ffn_hidden_size (int, optional): The hidden size of the feedforward network.
87
+ moe_num_experts (int, optional): The number of experts in the mixture of experts layer.
88
+ moe_top_k (int, optional): The number of experts to use in the mixture of experts layer.
89
+ moe_jitter_eps (float, optional): The jitter epsilon for the mixture of experts layer.
90
+ moe_loss_weight (float, optional): The loss weight for the mixture of experts layer.
91
+ moe_normalize_expert_weights (float, optional): The normalization factor for the expert weights.
92
+ uniform_expert_assignment (bool, optional): Whether to use uniform expert assignment.
93
+ This should only be used for benchmarking purposes.
94
+ """
95
+
96
+ def __init__(
97
+ self,
98
+ ffn_act_fn: Optional[dict] = None,
99
+ ffn_hidden_size: int = 3584,
100
+ moe_num_experts: int = 4,
101
+ moe_top_k: int = 1,
102
+ moe_jitter_eps: Optional[float] = None,
103
+ moe_loss_weight: float = 0.01,
104
+ moe_normalize_expert_weights: Optional[float] = 1,
105
+ uniform_expert_assignment: bool = False,
106
+ **kwargs: Any,
107
+ ):
108
+ super().__init__()
109
+ if ffn_act_fn is None:
110
+ ffn_act_fn = {'name': 'silu'}
111
+ self.ffn_act_fn = ffn_act_fn
112
+ self.ffn_hidden_size = ffn_hidden_size
113
+ self.moe_num_experts = moe_num_experts
114
+ self.moe_top_k = moe_top_k
115
+ self.moe_jitter_eps = moe_jitter_eps
116
+ self.moe_loss_weight = moe_loss_weight
117
+ self.moe_normalize_expert_weights = moe_normalize_expert_weights
118
+ self.uniform_expert_assignment = uniform_expert_assignment
119
+
120
+ for k in ['model_type']:
121
+ if k in kwargs:
122
+ kwargs.pop(k)
123
+ if len(kwargs) != 0:
124
+ raise ValueError(f'Found unknown {kwargs=}')
125
+
126
+ @classmethod
127
+ def from_pretrained(cls, pretrained_model_name_or_path: str,
128
+ **kwargs: Any) -> 'PretrainedConfig':
129
+ cls._set_token_in_kwargs(kwargs)
130
+
131
+ config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path,
132
+ **kwargs)
133
+
134
+ if config_dict.get('model_type') == 'dbrx':
135
+ config_dict = config_dict['ffn_config']
136
+
137
+ if 'model_type' in config_dict and hasattr(
138
+ cls,
139
+ 'model_type') and config_dict['model_type'] != cls.model_type:
140
+ logger.warning(
141
+ f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
142
+ +
143
+ f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'
144
+ )
145
+
146
+ return cls.from_dict(config_dict, **kwargs)
147
+
148
+
149
+ class DbrxConfig(PretrainedConfig):
150
+ """Configuration class for Dbrx.
151
+
152
+ [`DbrxModel`]. It is used to instantiate a Dbrx model according to the
153
+ specified arguments, defining the model architecture.
154
+
155
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
156
+ documentation from [`PretrainedConfig`] for more information.
157
+
158
+
159
+ Args:
160
+ d_model (`int`, *optional*, defaults to 6144):
161
+ Dimensionality of the embeddings and hidden states.
162
+ n_heads (`int`, *optional*, defaults to 48):
163
+ Number of attention heads for each attention layer in the Transformer encoder.
164
+ n_layers (`int`, *optional*, defaults to 40):
165
+ Number of hidden layers in the Transformer encoder.
166
+ max_seq_len (`int`, *optional*, defaults to 32768):
167
+ The maximum sequence length of the model.
168
+ vocab_size (`int`, *optional*, defaults to 100352):
169
+ Vocabulary size of the Dbrx model. Defines the maximum number of different tokens that can be represented by
170
+ the `inputs_ids` passed when calling [`DbrxModel`].
171
+ resid_pdrop (`float`, *optional*, defaults to 0.0):
172
+ The dropout probability applied to the attention output before combining with residual.
173
+ emb_pdrop (`float`, *optional*, defaults to 0.0):
174
+ The dropout probability for the embedding layer.
175
+ attn_config (`dict`, *optional*):
176
+ A dictionary used to configure the model's attention module.
177
+ ffn_config (`dict`, *optional*):
178
+ A dictionary used to configure the model's FFN module.
179
+ use_cache (`bool`, *optional*, defaults to `False`):
180
+ Whether or not the model should return the last key/values attentions (not used by all models).
181
+ initializer_range (`float`, *optional*, defaults to 0.02):
182
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
183
+ output_router_logits (`bool`, *optional*, defaults to `False`):
184
+ Whether or not the router logits should be returned by the model. Enabling this will also
185
+ allow the model to output the auxiliary loss. See [here]() for more details
186
+ router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
187
+ The aux loss factor for the total loss.
188
+
189
+
190
+ Example:
191
+ ```python
192
+ >>> from transformers import DbrxConfig, DbrxModel
193
+
194
+ >>> # Initializing a Dbrx configuration
195
+ >>> configuration = DbrxConfig()
196
+
197
+ >>> # Initializing a model (with random weights) from the configuration
198
+ >>> model = DbrxModel(configuration)
199
+
200
+ >>> # Accessing the model configuration
201
+ >>> configuration = model.config
202
+ ```
203
+ """
204
+
205
+ model_type = 'dbrx'
206
+ attribute_map = {
207
+ 'num_attention_heads': 'n_heads',
208
+ 'hidden_size': 'd_model',
209
+ 'num_hidden_layers': 'n_layers',
210
+ 'max_position_embeddings': 'max_seq_len'
211
+ }
212
+
213
+ def __init__(
214
+ self,
215
+ d_model: int = 2048,
216
+ n_heads: int = 16,
217
+ n_layers: int = 24,
218
+ max_seq_len: int = 2048,
219
+ vocab_size: int = 32000,
220
+ resid_pdrop: float = 0.0,
221
+ emb_pdrop: float = 0.0,
222
+ attn_config: Optional[DbrxAttentionConfig] = None,
223
+ ffn_config: Optional[DbrxFFNConfig] = None,
224
+ use_cache: bool = True,
225
+ initializer_range: float = 0.02,
226
+ output_router_logits: bool = False,
227
+ router_aux_loss_coef: float = 0.05,
228
+ **kwargs: Any,
229
+ ):
230
+ if attn_config is None:
231
+ self.attn_config = DbrxAttentionConfig()
232
+ elif isinstance(attn_config, dict):
233
+ self.attn_config = DbrxAttentionConfig(**attn_config)
234
+ else:
235
+ self.attn_config = attn_config
236
+
237
+ if ffn_config is None:
238
+ self.ffn_config = DbrxFFNConfig()
239
+ elif isinstance(ffn_config, dict):
240
+ self.ffn_config = DbrxFFNConfig(**ffn_config)
241
+ else:
242
+ self.ffn_config = ffn_config
243
+
244
+ self.d_model = d_model
245
+ self.n_heads = n_heads
246
+ self.n_layers = n_layers
247
+ self.max_seq_len = max_seq_len
248
+ self.vocab_size = vocab_size
249
+ self.resid_pdrop = resid_pdrop
250
+ self.emb_pdrop = emb_pdrop
251
+ self.use_cache = use_cache
252
+ self.initializer_range = initializer_range
253
+ self.output_router_logits = output_router_logits
254
+ self.router_aux_loss_coef = router_aux_loss_coef
255
+
256
+ tie_word_embeddings = kwargs.pop('tie_word_embeddings', False)
257
+ if tie_word_embeddings:
258
+ raise ValueError(
259
+ 'tie_word_embeddings is not supported for Dbrx models.')
260
+
261
+ super().__init__(
262
+ tie_word_embeddings=tie_word_embeddings,
263
+ **kwargs,
264
+ )
generation_config.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "transformers_version": "4.39.2"
4
+ }
model-00001-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ff1f9a31a1ca46eb9cd21b71f28cfd070bcde24c6e6335e421c60660a0d4bfb1
3
+ size 9909495544
model-00002-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3aeed25258a55f18c24209120d69bbbb14b3604afd1f783a5c750f644d9059b6
3
+ size 9953315320
model-00003-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:660149d5ea0af2029f549f7ace748fd1c1f00917869bb406abce73fafc5a5c96
3
+ size 9997577192
model-00004-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ac0418d1cb5318f057d8eaa6b3f1fe40fad7b4d940a81ef8f957eeb9fcac670e
3
+ size 9953315328
model-00005-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3cb111950baf7aa46053aad2494c768023d67e1ab20b3a1da67dd80ecf16bbca
3
+ size 9997577192
model-00006-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:60cd44c0e0f8eb47d28b341b3ac02cdf040cd9d7b3d94b444f14bfe6f3e6d94d
3
+ size 9953315328
model-00007-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:40da9173f1ae175bdee10c6a1acbab4cc1daf8e4039ce755298fd3241374c3b9
3
+ size 9997577216
model-00008-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:08101a418a884eb524ebc03376f671bd59a5b8afda3a2aff22206a3b678ce376
3
+ size 9997577288
model-00009-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6d06f7669bae53e170f7e9991885bc5534395e4b1e1bfd9e7153589817070018
3
+ size 9953315392
model-00010-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:eeb508d237c6c9bc9a80047c78296a77a51b93c18e3becd66a63d37c4d5e84d8
3
+ size 9997577288
model-00011-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:975affe6e5bee6dc3a3d694c89813f66d4eb433497454dcaf24b32cf097df92e
3
+ size 9953315392
model-00012-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f5bc94a312d0628e4cca9810cefa57bd0b15bf731e24012254cc988cc5044af1
3
+ size 9997577288
model-00013-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:26a770c1ea4cbcce0b7f2d7ede58540935f74aa07a7ccf740df0a80e9070bea3
3
+ size 9953315392
model-00014-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c13b90f9ce862f054440cddaf6771eedb060b50148642483a78fde1dd1120622
3
+ size 9997577280
model-00015-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:41c662d064a824cd74d15085c8f08933a2738c7775a8d1951b85bb3d515c238f
3
+ size 9953315400
model-00016-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:88fec478ecd38a20ecf4852558dda7eb692d4354c29930d4482604e70149f203
3
+ size 9997577280
model-00017-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0c07a8cc9428de0bf21cb0c4c946b6a50f9176e386dc2bd7dee55eca6740c159
3
+ size 9953315400
model-00018-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:33fe90e93917235232f86af298743f11a96798b1283aec93a243a087662f03bb
3
+ size 9997577280
model-00019-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:19fd377aebbf2d50209b08a6a4163372f10c55bf46b3cbbe2d312bb53d68e988
3
+ size 9953315400
model-00020-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:68bbad4086344aa4d846b451d6e6e496c5be905207822d5c074b940c60f12ba6
3
+ size 9997577272
model-00021-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b44d6810fe6450d9e3aeeae20603520191437becd620f36e8573eb17a2bfc74b
3
+ size 9953315408
model-00022-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3e670c8e260a669d848028b69d2f47bbd30844dbd1874769e2267983a7b7388b
3
+ size 9997577272
model-00023-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e6febb40c533831e0b5e8271d1bbb716d8a3ccd687b4bd4e2e2d2020d2a32554
3
+ size 9953327824
model-00024-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:87aeb8b8e0c9396ab1484c415d8cc6fab310dff29a9bcc6abbe2cf1037f537c7
3
+ size 9997564856
model-00025-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e083e38691d509c00cae0219ce4c2fa5af3a4cdc9c4f9d890ad71d155c1b15aa
3
+ size 9997577288
model-00026-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ecf0224195c692e89cb5476075ec395b2205efb602cb768d18db228398727c61
3
+ size 9953315392
model-00027-of-00027.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4346751973b7dee3522c65a3f3849c28866993330b6a5b76c2bea273bfc12229
3
+ size 3875552168
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
modeling_dbrx.py ADDED
@@ -0,0 +1,1454 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """PyTorch Dbrx model."""
2
+
3
+ import math
4
+ import warnings
5
+ from copy import deepcopy
6
+ from functools import partial
7
+ from typing import Any, Callable, Dict, Optional, Tuple, Union
8
+
9
+ import torch
10
+ import torch.nn.functional as F
11
+ import torch.utils.checkpoint
12
+ from torch import nn
13
+ from transformers.cache_utils import Cache, DynamicCache, StaticCache
14
+ from transformers.modeling_attn_mask_utils import AttentionMaskConverter
15
+ from transformers.modeling_outputs import (MoeCausalLMOutputWithPast,
16
+ MoeModelOutputWithPast)
17
+ from transformers.modeling_utils import PreTrainedModel
18
+ from transformers.utils import is_flash_attn_2_available, logging
19
+
20
+ from .configuration_dbrx import DbrxAttentionConfig, DbrxConfig, DbrxFFNConfig
21
+
22
+ if is_flash_attn_2_available():
23
+ try:
24
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
25
+ from flash_attn.bert_padding import pad_input # noqa
26
+ from flash_attn.bert_padding import index_first_axis, unpad_input
27
+ except:
28
+ pass
29
+
30
+ logger = logging.get_logger(__name__)
31
+
32
+ _CONFIG_FOR_DOC = 'DbrxConfig'
33
+
34
+ #############################################################################
35
+ # Copied from LLaMaRotaryEmbedding
36
+ #############################################################################
37
+
38
+
39
+ class DbrxRotaryEmbedding(nn.Module):
40
+
41
+ def __init__(self,
42
+ dim: int,
43
+ max_position_embeddings: int = 2048,
44
+ base: float = 10000.0,
45
+ scaling_factor: float = 1.0):
46
+ super().__init__()
47
+ self.scaling_factor = scaling_factor
48
+ self.dim = dim
49
+ self.max_position_embeddings = max_position_embeddings
50
+ self.base = base
51
+ inv_freq = 1.0 / (self.base**(
52
+ torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
53
+ self.register_buffer('inv_freq', inv_freq, persistent=False)
54
+ # For BC we register cos and sin cached
55
+ self.max_seq_len_cached = max_position_embeddings
56
+
57
+ @torch.no_grad()
58
+ def forward(
59
+ self, x: torch.Tensor, position_ids: torch.LongTensor
60
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
61
+ # x: [bs, num_attention_heads, seq_len, head_size]
62
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(
63
+ position_ids.shape[0], -1, 1)
64
+ position_ids_expanded = position_ids[:, None, :].float()
65
+ # Force float32 since bfloat16 loses precision on long contexts
66
+ # See https://github.com/huggingface/transformers/pull/29285
67
+ device_type = x.device.type
68
+ device_type = device_type if isinstance(
69
+ device_type, str) and device_type != 'mps' else 'cpu'
70
+ with torch.autocast(device_type=device_type, enabled=False):
71
+ freqs = (inv_freq_expanded.float()
72
+ @ position_ids_expanded.float()).transpose(1, 2)
73
+ emb = torch.cat((freqs, freqs), dim=-1)
74
+ cos = emb.cos()
75
+ sin = emb.sin()
76
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
77
+
78
+
79
+ def rotate_half(x: torch.Tensor) -> torch.Tensor:
80
+ """Rotates half the hidden dims of the input."""
81
+ x1 = x[..., :x.shape[-1] // 2]
82
+ x2 = x[..., x.shape[-1] // 2:]
83
+ return torch.cat((-x2, x1), dim=-1)
84
+
85
+
86
+ def apply_rotary_pos_emb(
87
+ q: torch.Tensor,
88
+ k: torch.Tensor,
89
+ cos: torch.Tensor,
90
+ sin: torch.Tensor,
91
+ unsqueeze_dim: int = 1) -> Tuple[torch.Tensor, torch.Tensor]:
92
+ """Applies Rotary Position Embedding to the query and key tensors.
93
+
94
+ Args:
95
+ q (`torch.Tensor`): The query tensor.
96
+ k (`torch.Tensor`): The key tensor.
97
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
98
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
99
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
100
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos and
101
+ sin so that they can be properly broadcasted to the dimensions of q and k. For example, note
102
+ that cos and sin have the shape [batch_size, seq_len, head_dim]. Then, if q and
103
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
104
+ cos and sin broadcastable to the shapes of q and k. Similarly, if q and k have
105
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
106
+
107
+ Returns:
108
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
109
+ """
110
+ cos = cos.unsqueeze(unsqueeze_dim)
111
+ sin = sin.unsqueeze(unsqueeze_dim)
112
+ q_embed = (q * cos) + (rotate_half(q) * sin)
113
+ k_embed = (k * cos) + (rotate_half(k) * sin)
114
+ return q_embed, k_embed
115
+
116
+
117
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
118
+ """Equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).
119
+
120
+ The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to
121
+ (batch, num_attention_heads, seqlen, head_dim)
122
+ """
123
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
124
+ if n_rep == 1:
125
+ return hidden_states
126
+ hidden_states = hidden_states[:, :,
127
+ None, :, :].expand(batch, num_key_value_heads,
128
+ n_rep, slen, head_dim)
129
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen,
130
+ head_dim)
131
+
132
+
133
+ #############################################################################
134
+
135
+ #############################################################################
136
+ # Modified from modeling_mixtral
137
+ #############################################################################
138
+
139
+
140
+ def load_balancing_loss_func(
141
+ gate_logits: torch.Tensor,
142
+ num_experts: int,
143
+ top_k: int,
144
+ attention_mask: Optional[torch.Tensor],
145
+ ) -> torch.Tensor:
146
+ r"""Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
147
+
148
+ See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
149
+ function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
150
+ experts is too unbalanced.
151
+
152
+ Args:
153
+ gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]):
154
+ Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
155
+ shape [batch_size X sequence_length, num_experts].
156
+ num_experts (`int`):
157
+ Number of experts.
158
+ top_k (`int`):
159
+ The number of experts each token is routed to.
160
+ attention_mask (`torch.Tensor`, None):
161
+ The attention_mask used in forward function
162
+ shape [batch_size X sequence_length] if not None.
163
+
164
+ Returns:
165
+ The auxiliary loss.
166
+ """
167
+ if gate_logits is None or not isinstance(gate_logits, tuple):
168
+ return torch.tensor(0.0)
169
+
170
+ if isinstance(gate_logits, tuple):
171
+ compute_device = gate_logits[0].device
172
+ concatenated_gate_logits = torch.cat(
173
+ [layer_gate.to(compute_device) for layer_gate in gate_logits],
174
+ dim=0)
175
+
176
+ routing_weights = torch.nn.functional.softmax(concatenated_gate_logits,
177
+ dim=-1)
178
+
179
+ _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
180
+
181
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
182
+
183
+ if attention_mask is None:
184
+ # Compute the percentage of tokens routed to each experts
185
+ tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
186
+
187
+ # Compute the average probability of routing to these experts
188
+ router_prob_per_expert = torch.mean(routing_weights, dim=0)
189
+ else:
190
+ batch_size, sequence_length = attention_mask.shape
191
+ num_hidden_layers = concatenated_gate_logits.shape[0] // (
192
+ batch_size * sequence_length)
193
+
194
+ # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
195
+ expert_attention_mask = (attention_mask[None, :, :, None, None].expand(
196
+ (num_hidden_layers, batch_size, sequence_length, top_k,
197
+ num_experts)).reshape(-1, top_k, num_experts).to(compute_device))
198
+
199
+ # Compute the percentage of tokens routed to each experts
200
+ tokens_per_expert = torch.sum(
201
+ expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
202
+ expert_attention_mask, dim=0)
203
+
204
+ # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
205
+ router_per_expert_attention_mask = (
206
+ attention_mask[None, :, :, None].expand(
207
+ (num_hidden_layers, batch_size, sequence_length,
208
+ num_experts)).reshape(-1, num_experts).to(compute_device))
209
+
210
+ # Compute the average probability of routing to these experts
211
+ router_prob_per_expert = torch.sum(
212
+ routing_weights * router_per_expert_attention_mask,
213
+ dim=0) / torch.sum(router_per_expert_attention_mask, dim=0)
214
+
215
+ overall_loss = torch.sum(tokens_per_expert *
216
+ router_prob_per_expert.unsqueeze(0))
217
+ return overall_loss * num_experts
218
+
219
+
220
+ #############################################################################
221
+
222
+
223
+ def resolve_ffn_act_fn(
224
+ ffn_act_fn: dict) -> Callable[[torch.Tensor], torch.Tensor]:
225
+ """Resolve the activation function for the feed-forward network.
226
+
227
+ Args:
228
+ ffn_act_fn (dict): The configuration dictionary for the activation function.
229
+ The dict config must specify the 'name' of a torch.nn.functional activation
230
+ function. All of other key values pairs are bound to the function as a partial.
231
+
232
+ Returns:
233
+ Callable[[torch.Tensor], torch.Tensor]: The activation function.
234
+ """
235
+ config = deepcopy(ffn_act_fn)
236
+ name = config.pop('name')
237
+ if not hasattr(nn.functional, name):
238
+ raise ValueError(f'Unrecognised activation function name ({name}).')
239
+ act = getattr(nn.functional, name)
240
+ return partial(act, **config)
241
+
242
+
243
+ #############################################################################
244
+ # Copied from LLaMaAttention
245
+ #############################################################################
246
+
247
+
248
+ def _get_unpad_data(attention_mask: torch.Tensor):
249
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
250
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
251
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
252
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32),
253
+ (1, 0))
254
+ return (
255
+ indices,
256
+ cu_seqlens,
257
+ max_seqlen_in_batch,
258
+ )
259
+
260
+
261
+ class DbrxAttention(nn.Module):
262
+ """Multi-head self attention."""
263
+
264
+ def __init__(self,
265
+ hidden_size: int,
266
+ num_heads: int,
267
+ max_position_embeddings: int,
268
+ attn_config: DbrxAttentionConfig,
269
+ block_idx: Optional[int] = None):
270
+ super().__init__()
271
+ self.hidden_size = hidden_size
272
+ self.num_heads = num_heads
273
+ self.head_dim = self.hidden_size // self.num_heads
274
+ self.max_position_embeddings = max_position_embeddings
275
+ self.block_idx = block_idx
276
+ self.config = attn_config
277
+ if block_idx is None:
278
+ logger.warning_once(
279
+ f'Instantiating {self.__class__.__name__} without passing a `block_idx` is not recommended and will '
280
+ +
281
+ 'lead to errors during the forward call if caching is used. Please make sure to provide a `block_idx` '
282
+ + 'when creating this class.')
283
+
284
+ self.attn_pdrop = attn_config.attn_pdrop
285
+ self.clip_qkv = attn_config.clip_qkv
286
+ self.num_key_value_heads = attn_config.kv_n_heads
287
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
288
+ self.rope_theta = attn_config.rope_theta
289
+
290
+ self.Wqkv = nn.Linear(self.hidden_size,
291
+ self.hidden_size +
292
+ 2 * self.num_key_value_heads * self.head_dim,
293
+ bias=False)
294
+ self.out_proj = nn.Linear(self.hidden_size,
295
+ self.hidden_size,
296
+ bias=False)
297
+ self.rotary_emb = DbrxRotaryEmbedding(
298
+ self.head_dim,
299
+ max_position_embeddings=self.max_position_embeddings,
300
+ base=self.rope_theta,
301
+ )
302
+
303
+ def forward(
304
+ self,
305
+ hidden_states: torch.Tensor,
306
+ position_ids: torch.LongTensor,
307
+ attention_mask: Optional[torch.Tensor] = None,
308
+ past_key_value: Optional[Cache] = None,
309
+ output_attentions: bool = False,
310
+ use_cache: bool = False,
311
+ cache_position: Optional[torch.LongTensor] = None,
312
+ **kwargs: Any,
313
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
314
+ bsz, q_len, _ = hidden_states.size()
315
+
316
+ qkv_states = self.Wqkv(hidden_states)
317
+ if self.clip_qkv is not None:
318
+ qkv_states = qkv_states.clamp(min=-self.clip_qkv, max=self.clip_qkv)
319
+
320
+ query_states, key_states, value_states = qkv_states.split(
321
+ [
322
+ self.hidden_size,
323
+ self.num_key_value_heads * self.head_dim,
324
+ self.num_key_value_heads * self.head_dim,
325
+ ],
326
+ dim=2,
327
+ )
328
+
329
+ query_states = query_states.view(bsz, q_len, self.num_heads,
330
+ self.head_dim).transpose(1, 2)
331
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads,
332
+ self.head_dim).transpose(1, 2)
333
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads,
334
+ self.head_dim).transpose(1, 2)
335
+
336
+ past_key_value = getattr(self, 'past_key_value', past_key_value)
337
+ cos, sin = self.rotary_emb(value_states, position_ids)
338
+ query_states, key_states = apply_rotary_pos_emb(query_states,
339
+ key_states, cos, sin)
340
+
341
+ if past_key_value is not None:
342
+ # sin and cos are specific to RoPE models; position_ids needed for the static cache
343
+ cache_kwargs = {
344
+ 'sin': sin,
345
+ 'cos': cos,
346
+ 'cache_position': cache_position
347
+ }
348
+ key_states, value_states = past_key_value.update(
349
+ key_states, value_states, self.block_idx, cache_kwargs)
350
+
351
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
352
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
353
+
354
+ attn_weights = torch.matmul(query_states, key_states.transpose(
355
+ 2, 3)) / math.sqrt(self.head_dim)
356
+
357
+ if attention_mask is not None: # no matter the length, we just slice it
358
+ causal_mask = attention_mask[:, :, :, :key_states.shape[-2]]
359
+ attn_weights = attn_weights + causal_mask
360
+
361
+ # upcast attention to fp32
362
+ attn_weights = nn.functional.softmax(attn_weights,
363
+ dim=-1,
364
+ dtype=torch.float32).to(
365
+ query_states.dtype)
366
+ attn_weights = nn.functional.dropout(attn_weights,
367
+ p=self.attn_pdrop,
368
+ training=self.training)
369
+ attn_output = torch.matmul(attn_weights, value_states)
370
+
371
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
372
+ raise ValueError(
373
+ f'`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is'
374
+ + f' {attn_output.size()}')
375
+
376
+ attn_output = attn_output.transpose(1, 2).contiguous()
377
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
378
+ attn_output = self.out_proj(attn_output)
379
+
380
+ if not output_attentions:
381
+ attn_weights = None
382
+
383
+ return attn_output, attn_weights, past_key_value
384
+
385
+
386
+ class DbrxFlashAttention2(DbrxAttention):
387
+ """Dbrx flash attention module.
388
+
389
+ This module inherits from `DbrxAttention` as the weights of the module stays
390
+ untouched. The only required change would be on the forward pass where it
391
+ calls the public API of flash attention.
392
+ """
393
+
394
+ def __init__(self, *args: Any, **kwargs: Any):
395
+ if not is_flash_attn_2_available():
396
+ raise ImportError(
397
+ 'Flash Attention 2 is not available. Please install it with `pip install flash-attn`.'
398
+ )
399
+
400
+ super().__init__(*args, **kwargs)
401
+
402
+ def forward(
403
+ self,
404
+ hidden_states: torch.Tensor,
405
+ attention_mask: Optional[torch.LongTensor] = None,
406
+ position_ids: Optional[torch.LongTensor] = None,
407
+ past_key_value: Optional[Cache] = None,
408
+ output_attentions: bool = False,
409
+ use_cache: bool = False,
410
+ cache_position: Optional[torch.LongTensor] = None,
411
+ **kwargs: Any,
412
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
413
+ Optional[Tuple[torch.Tensor]]]:
414
+ logger.info(
415
+ 'Implicitly setting `output_attentions` to False as it is not supported in Flash Attention.'
416
+ )
417
+ output_attentions = False
418
+
419
+ bsz, q_len, _ = hidden_states.size()
420
+
421
+ qkv_states = self.Wqkv(hidden_states)
422
+ if self.clip_qkv is not None:
423
+ qkv_states = qkv_states.clamp(min=-self.clip_qkv, max=self.clip_qkv)
424
+
425
+ query_states, key_states, value_states = qkv_states.split(
426
+ [
427
+ self.hidden_size,
428
+ self.num_key_value_heads * self.head_dim,
429
+ self.num_key_value_heads * self.head_dim,
430
+ ],
431
+ dim=2,
432
+ )
433
+
434
+ # Flash attention requires the input to have the shape
435
+ # batch_size x seq_length x head_dim x hidden_dim
436
+ # therefore we just need to keep the original shape
437
+ query_states = query_states.view(bsz, q_len, self.num_heads,
438
+ self.head_dim).transpose(1, 2)
439
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads,
440
+ self.head_dim).transpose(1, 2)
441
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads,
442
+ self.head_dim).transpose(1, 2)
443
+
444
+ cos, sin = self.rotary_emb(value_states, position_ids)
445
+ query_states, key_states = apply_rotary_pos_emb(query_states,
446
+ key_states, cos, sin)
447
+
448
+ past_key_value = getattr(self, 'past_key_value', past_key_value)
449
+
450
+ if past_key_value is not None:
451
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
452
+ cache_kwargs = {
453
+ 'sin': sin,
454
+ 'cos': cos,
455
+ 'cache_position': cache_position
456
+ }
457
+ key_states, value_states = past_key_value.update(
458
+ key_states, value_states, self.block_idx, cache_kwargs)
459
+
460
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout
461
+ # [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
462
+ # to be able to avoid many of these transpose/reshape/view.
463
+ query_states = query_states.transpose(1, 2)
464
+ key_states = key_states.transpose(1, 2)
465
+ value_states = value_states.transpose(1, 2)
466
+
467
+ dropout_rate = self.attn_pdrop if self.training else 0.0
468
+
469
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
470
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
471
+ # cast them back in the correct dtype just to be sure everything works as expected.
472
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
473
+ # in fp32. (LlamaRMSNorm handles it correctly)
474
+ input_dtype = query_states.dtype
475
+ if input_dtype == torch.float32:
476
+ if torch.is_autocast_enabled():
477
+ target_dtype = torch.get_autocast_gpu_dtype()
478
+ # Handle the case where the model is quantized
479
+ elif hasattr(self.config, '_pre_quantization_dtype'):
480
+ target_dtype = self.config._pre_quantization_dtype
481
+ else:
482
+ target_dtype = query_states.dtype
483
+
484
+ logger.warning_once(
485
+ f'The input hidden states seems to be silently casted in float32, this might be '
486
+ +
487
+ f'related to the fact you have upcasted embedding or layer norm layers in '
488
+ + f'float32. We will cast back the input in {target_dtype}.')
489
+
490
+ query_states = query_states.to(target_dtype)
491
+ key_states = key_states.to(target_dtype)
492
+ value_states = value_states.to(target_dtype)
493
+
494
+ attn_output = self._flash_attention_forward(
495
+ query_states,
496
+ key_states,
497
+ value_states,
498
+ attention_mask,
499
+ q_len,
500
+ dropout=dropout_rate,
501
+ )
502
+
503
+ attn_output = attn_output.reshape(bsz, q_len,
504
+ self.hidden_size).contiguous()
505
+ attn_output = self.out_proj(attn_output)
506
+
507
+ if not output_attentions:
508
+ attn_weights = None
509
+
510
+ return attn_output, attn_weights, past_key_value # type: ignore
511
+
512
+ def _flash_attention_forward(
513
+ self,
514
+ query_states: torch.Tensor,
515
+ key_states: torch.Tensor,
516
+ value_states: torch.Tensor,
517
+ attention_mask: Union[torch.LongTensor, None],
518
+ query_length: int,
519
+ dropout: float = 0.0,
520
+ softmax_scale: Optional[float] = None,
521
+ ):
522
+ """Use FlashAttention, stripping padding tokens if necessary.
523
+
524
+ Args:
525
+ query_states (torch.Tensor): Input query states to be passed to Flash Attention API
526
+ key_states (torch.Tensor): Input key states to be passed to Flash Attention API
527
+ value_states (torch.Tensor): Input value states to be passed to Flash Attention API
528
+ attention_mask (torch.LongTensor | None): The padding mask - corresponds to a tensor of size
529
+ (batch_size, seq_len) where 0 stands for the position of padding tokens and 1
530
+ for the position of non-padding tokens.
531
+ query_length (int): The length of the query sequence
532
+ dropout (float): Attention dropout
533
+ softmax_scale (float, optional): The scaling of QK^T before applying softmax.
534
+ Defaults to 1 / sqrt(head_dim)
535
+ """
536
+ causal = True
537
+ # Contains at least one padding token in the sequence
538
+ if attention_mask is not None:
539
+ batch_size = query_states.shape[0]
540
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
541
+ query_states, key_states, value_states, attention_mask,
542
+ query_length)
543
+
544
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
545
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
546
+
547
+ attn_output_unpad = flash_attn_varlen_func(
548
+ query_states,
549
+ key_states,
550
+ value_states,
551
+ cu_seqlens_q=cu_seqlens_q,
552
+ cu_seqlens_k=cu_seqlens_k,
553
+ max_seqlen_q=max_seqlen_in_batch_q,
554
+ max_seqlen_k=max_seqlen_in_batch_k,
555
+ dropout_p=dropout,
556
+ softmax_scale=softmax_scale,
557
+ causal=causal,
558
+ )
559
+
560
+ attn_output = pad_input(
561
+ attn_output_unpad,
562
+ indices_q,
563
+ batch_size,
564
+ query_length,
565
+ )
566
+ else:
567
+ attn_output = flash_attn_func(
568
+ query_states,
569
+ key_states,
570
+ value_states,
571
+ dropout,
572
+ softmax_scale=softmax_scale,
573
+ causal=causal,
574
+ )
575
+
576
+ return attn_output
577
+
578
+ def _upad_input(self, query_layer: torch.Tensor, key_layer: torch.Tensor,
579
+ value_layer: torch.Tensor, attention_mask: torch.Tensor,
580
+ query_length: int):
581
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(
582
+ attention_mask)
583
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
584
+
585
+ key_layer = index_first_axis(
586
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads,
587
+ head_dim), indices_k)
588
+ value_layer = index_first_axis(
589
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads,
590
+ head_dim), indices_k)
591
+ if query_length == kv_seq_len:
592
+ query_layer = index_first_axis(
593
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads,
594
+ head_dim), indices_k)
595
+ cu_seqlens_q = cu_seqlens_k
596
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
597
+ indices_q = indices_k
598
+ elif query_length == 1:
599
+ max_seqlen_in_batch_q = 1
600
+ cu_seqlens_q = torch.arange(
601
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
602
+ ) # There is a memcpy here, that is very bad.
603
+ indices_q = cu_seqlens_q[:-1]
604
+ query_layer = query_layer.squeeze(1)
605
+ else:
606
+ # The -q_len: slice assumes left padding.
607
+ attention_mask = attention_mask[:, -query_length:]
608
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
609
+ query_layer, attention_mask)
610
+
611
+ return (
612
+ query_layer,
613
+ key_layer,
614
+ value_layer,
615
+ indices_q,
616
+ (cu_seqlens_q, cu_seqlens_k),
617
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
618
+ )
619
+
620
+
621
+ DBRX_ATTENTION_CLASSES = {
622
+ 'eager': DbrxAttention,
623
+ 'flash_attention_2': DbrxFlashAttention2,
624
+ }
625
+
626
+
627
+ class DbrxNormAttentionNorm(nn.Module):
628
+
629
+ def __init__(
630
+ self,
631
+ hidden_size: int,
632
+ num_heads: int,
633
+ max_position_embeddings: int,
634
+ resid_pdrop: float,
635
+ attn_implementation: str,
636
+ attn_config: DbrxAttentionConfig,
637
+ block_idx: Optional[int] = None,
638
+ ):
639
+ super().__init__()
640
+ self.block_idx = block_idx
641
+ self.resid_pdrop = resid_pdrop
642
+ self.norm_1 = nn.LayerNorm(hidden_size, bias=False)
643
+ self.attn = DBRX_ATTENTION_CLASSES[attn_implementation](
644
+ hidden_size=hidden_size,
645
+ num_heads=num_heads,
646
+ max_position_embeddings=max_position_embeddings,
647
+ attn_config=attn_config,
648
+ block_idx=block_idx,
649
+ )
650
+ self.norm_2 = nn.LayerNorm(hidden_size, bias=False)
651
+
652
+ def forward(
653
+ self,
654
+ hidden_states: torch.Tensor,
655
+ position_ids: torch.LongTensor,
656
+ attention_mask: Optional[torch.Tensor] = None,
657
+ past_key_value: Optional[Cache] = None,
658
+ output_attentions: bool = False,
659
+ use_cache: bool = False,
660
+ cache_position: Optional[torch.LongTensor] = None,
661
+ **kwargs: Any,
662
+ ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor],
663
+ Optional[Cache]]:
664
+
665
+ residual_states = hidden_states
666
+ hidden_states = self.norm_1(hidden_states).to(hidden_states.dtype)
667
+
668
+ hidden_states, attn_weights, past_key_value = self.attn(
669
+ hidden_states=hidden_states,
670
+ attention_mask=attention_mask,
671
+ position_ids=position_ids,
672
+ past_key_value=past_key_value,
673
+ output_attentions=output_attentions,
674
+ use_cache=use_cache,
675
+ cache_position=cache_position,
676
+ **kwargs,
677
+ )
678
+
679
+ hidden_states = nn.functional.dropout(hidden_states,
680
+ p=self.resid_pdrop,
681
+ training=self.training)
682
+ hidden_states = hidden_states + residual_states
683
+
684
+ residual_states = hidden_states
685
+ hidden_states = self.norm_2(hidden_states).to(hidden_states.dtype)
686
+
687
+ return residual_states, hidden_states, attn_weights, past_key_value
688
+
689
+
690
+ class DbrxRouter(nn.Module):
691
+
692
+ def __init__(self, hidden_size: int, moe_num_experts: int, moe_top_k: int,
693
+ moe_jitter_eps: Optional[float],
694
+ moe_normalize_expert_weights: Optional[float],
695
+ uniform_expert_assignment: bool):
696
+ super().__init__()
697
+ self.hidden_size = hidden_size
698
+ self.moe_num_experts = moe_num_experts
699
+ self.moe_top_k = moe_top_k
700
+ self.moe_jitter_eps = moe_jitter_eps
701
+ self.moe_normalize_expert_weights = moe_normalize_expert_weights
702
+ self.uniform_expert_assignment = uniform_expert_assignment
703
+
704
+ self.layer = nn.Linear(self.hidden_size,
705
+ self.moe_num_experts,
706
+ bias=False)
707
+
708
+ def jitter(self, x: torch.Tensor) -> torch.Tensor:
709
+ if self.moe_jitter_eps is None:
710
+ raise RuntimeError('The router does not have moe_jitter_eps set.')
711
+ low = 1.0 - self.moe_jitter_eps
712
+ high = 1.0 + self.moe_jitter_eps
713
+ noise = torch.rand(x.size(), dtype=x.dtype, device=x.device)
714
+ return low + noise * (high - low)
715
+
716
+ def forward(
717
+ self, x: torch.Tensor
718
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
719
+ if self.training and self.moe_jitter_eps is not None:
720
+ x = x * self.jitter(x)
721
+
722
+ weights = self.layer(x.view(-1,
723
+ x.shape[-1])).softmax(dim=-1,
724
+ dtype=torch.float32)
725
+ top_weights, top_experts = torch.topk(weights, self.moe_top_k, dim=-1)
726
+
727
+ if self.moe_normalize_expert_weights:
728
+ top_weights = top_weights / torch.norm(
729
+ top_weights,
730
+ p=self.moe_normalize_expert_weights,
731
+ dim=-1,
732
+ keepdim=True)
733
+
734
+ if self.uniform_expert_assignment:
735
+ with torch.no_grad():
736
+ uniform_tensor = torch.arange(
737
+ 0,
738
+ top_experts.numel(),
739
+ device=top_experts.device,
740
+ dtype=top_experts.dtype) % self.moe_num_experts
741
+ top_experts = uniform_tensor.reshape(top_experts.shape)
742
+ # Note, weights and top_weights are not changed
743
+
744
+ weights = weights.to(x.dtype)
745
+ top_weights = top_weights.to(x.dtype)
746
+ return weights, top_weights, top_experts # type: ignore
747
+
748
+
749
+ class DbrxExpertGLU(nn.Module):
750
+
751
+ def __init__(self, hidden_size: int, ffn_hidden_size: int, ffn_act_fn: dict):
752
+ super().__init__()
753
+ self.w1 = nn.Linear(hidden_size, ffn_hidden_size, bias=False)
754
+ self.v1 = nn.Linear(hidden_size, ffn_hidden_size, bias=False)
755
+ self.w2 = nn.Linear(ffn_hidden_size, hidden_size, bias=False)
756
+ self.activation_fn = resolve_ffn_act_fn(ffn_act_fn)
757
+
758
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
759
+ x1 = self.w1(x)
760
+ x2 = self.v1(x)
761
+ x1 = self.activation_fn(x1)
762
+ x1 = x1 * x2
763
+ x1 = self.w2(x1)
764
+ return x1
765
+
766
+
767
+ class DbrxExperts(nn.Module):
768
+
769
+ def __init__(self, hidden_size: int, ffn_hidden_size: int, moe_num_experts: int, ffn_act_fn: dict):
770
+ super().__init__()
771
+ self.moe_num_experts = moe_num_experts
772
+ self.mlp_experts = nn.ModuleList([DbrxExpertGLU(hidden_size, ffn_hidden_size, ffn_act_fn) for _ in range(moe_num_experts)])
773
+
774
+ def forward(self, x: torch.Tensor, weights: torch.Tensor, top_weights: torch.Tensor, top_experts: torch.LongTensor) -> torch.Tensor:
775
+ bsz, q_len, hidden_size = x.shape
776
+ x = x.view(-1, hidden_size)
777
+ out = torch.zeros_like(x)
778
+
779
+ expert_mask = nn.functional.one_hot(top_experts, num_classes=self.moe_num_experts).permute(2, 1, 0)
780
+ for expert_idx in range(0, self.moe_num_experts):
781
+ topk_idx, token_idx = torch.where(expert_mask[expert_idx])
782
+ if token_idx.shape[0] == 0:
783
+ continue
784
+
785
+ token_list = token_idx.tolist()
786
+ topk_list = topk_idx.tolist()
787
+
788
+ expert_tokens = x[None, token_list].reshape(-1, hidden_size)
789
+ expert_out = self.mlp_experts[expert_idx](expert_tokens) * top_weights[token_list, topk_list, None]
790
+
791
+ out.index_add_(0, token_idx, expert_out)
792
+
793
+ out = out.reshape(bsz, q_len, hidden_size)
794
+ return out
795
+
796
+
797
+ class DbrxFFN(nn.Module):
798
+
799
+ def __init__(self, hidden_size: int, ffn_config: DbrxFFNConfig):
800
+ super().__init__()
801
+
802
+ self.router = DbrxRouter(
803
+ hidden_size,
804
+ moe_num_experts=ffn_config.moe_num_experts,
805
+ moe_top_k=ffn_config.moe_top_k,
806
+ moe_jitter_eps=ffn_config.moe_jitter_eps,
807
+ moe_normalize_expert_weights=ffn_config.
808
+ moe_normalize_expert_weights,
809
+ uniform_expert_assignment=ffn_config.uniform_expert_assignment,
810
+ )
811
+
812
+ self.experts = DbrxExperts(
813
+ hidden_size=hidden_size,
814
+ ffn_hidden_size=ffn_config.ffn_hidden_size,
815
+ moe_num_experts=ffn_config.moe_num_experts,
816
+ ffn_act_fn=ffn_config.ffn_act_fn,
817
+ )
818
+
819
+ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
820
+ weights, top_weights, top_experts = self.router(x)
821
+ out = self.experts(x, weights, top_weights, top_experts)
822
+ return out, weights
823
+
824
+
825
+ class DbrxBlock(nn.Module):
826
+
827
+ def __init__(self, config: DbrxConfig, block_idx: int):
828
+ super().__init__()
829
+ self.hidden_size = config.d_model
830
+ self.resid_pdrop = config.resid_pdrop
831
+ self.block_idx = block_idx
832
+ self.norm_attn_norm = DbrxNormAttentionNorm(
833
+ hidden_size=config.d_model,
834
+ num_heads=config.n_heads,
835
+ max_position_embeddings=config.max_seq_len,
836
+ resid_pdrop=config.resid_pdrop,
837
+ attn_implementation=config._attn_implementation,
838
+ attn_config=config.attn_config,
839
+ block_idx=block_idx,
840
+ )
841
+ self.ffn = DbrxFFN(hidden_size=config.d_model,
842
+ ffn_config=config.ffn_config)
843
+
844
+ def forward(
845
+ self,
846
+ hidden_states: torch.Tensor,
847
+ attention_mask: Optional[torch.Tensor] = None,
848
+ position_ids: torch.LongTensor = None,
849
+ past_key_value: Optional[Cache] = None,
850
+ output_attentions: Optional[bool] = False,
851
+ output_router_logits: Optional[bool] = False,
852
+ use_cache: Optional[bool] = False,
853
+ cache_position: Optional[torch.LongTensor] = None,
854
+ **kwargs: Any,
855
+ ) -> Union[Tuple[torch.Tensor], Tuple[torch.Tensor, Optional[torch.Tensor]],
856
+ Tuple[torch.Tensor, Optional[Cache]], Tuple[
857
+ torch.Tensor, Optional[torch.Tensor], Optional[Cache]],
858
+ Tuple[torch.Tensor, Optional[torch.Tensor],
859
+ Optional[torch.Tensor]], Tuple[
860
+ torch.Tensor, Optional[Cache], Optional[torch.Tensor]],
861
+ Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache],
862
+ Optional[torch.Tensor]],]:
863
+ """Forward function for DbrxBlock.
864
+
865
+ Args:
866
+ hidden_states (`torch.Tensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
867
+ position_ids (`torch.LongTensor`): position ids of shape `(batch, seq_len)`
868
+ attention_mask (`torch.Tensor`, optional): attention mask of size (batch_size, sequence_length)
869
+ if flash attention is used or (batch_size, 1, query_sequence_length, key_sequence_length)
870
+ if default attention is used.
871
+ past_key_value (`Tuple(torch.Tensor)`, optional): cached past key and value projection states
872
+ output_attentions (`bool`, optional): Whether or not to return the attentions tensors of all
873
+ attention layers. See `attentions` under returned tensors for more detail.
874
+ output_router_logits (`bool`, optional): Whether or not to return the router logits.
875
+ use_cache (`bool`, optional): If set to `True`, `past_key_values` key value states are
876
+ returned and can be used to speed up decoding (see `past_key_values`).
877
+ cache_position (`torch.LongTensor`, optional): position ids of the cache
878
+ """
879
+ if 'padding_mask' in kwargs:
880
+ warnings.warn(
881
+ 'Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`'
882
+ )
883
+
884
+ # Norm + Attention + Norm
885
+ resid_states, hidden_states, self_attn_weights, present_key_value = self.norm_attn_norm(
886
+ hidden_states=hidden_states,
887
+ attention_mask=attention_mask,
888
+ position_ids=position_ids,
889
+ past_key_value=past_key_value,
890
+ output_attentions=output_attentions,
891
+ use_cache=use_cache,
892
+ cache_position=cache_position,
893
+ **kwargs,
894
+ )
895
+
896
+ # Fully Connected
897
+ hidden_states, router_logits = self.ffn(hidden_states)
898
+ hidden_states = nn.functional.dropout(hidden_states,
899
+ p=self.resid_pdrop,
900
+ training=self.training)
901
+ hidden_states = resid_states + hidden_states
902
+
903
+ outputs = (hidden_states,)
904
+
905
+ if output_attentions:
906
+ outputs += (self_attn_weights,)
907
+
908
+ if use_cache:
909
+ outputs += (present_key_value,)
910
+
911
+ if output_router_logits:
912
+ outputs += (router_logits,)
913
+
914
+ return outputs
915
+
916
+
917
+ class DbrxPreTrainedModel(PreTrainedModel):
918
+ config_class = DbrxConfig
919
+ base_model_prefix = 'transformer'
920
+ supports_gradient_checkpointing = True
921
+ _no_split_modules = ['DbrxBlock']
922
+ _skip_keys_device_placement = ['past_key_values']
923
+ _supports_flash_attn_2 = True
924
+ _supports_sdpa = False
925
+ _supports_cache_class = True
926
+
927
+ def _init_weights(self, module: nn.Module):
928
+ std = self.config.initializer_range
929
+ if isinstance(module, nn.Linear):
930
+ module.weight.data.normal_(mean=0.0, std=std)
931
+ if module.bias is not None:
932
+ module.bias.data.zero_()
933
+ elif isinstance(module, nn.Embedding):
934
+ module.weight.data.normal_(mean=0.0, std=std)
935
+ if module.padding_idx is not None:
936
+ module.weight.data[module.padding_idx].zero_()
937
+ elif isinstance(module, nn.LayerNorm):
938
+ module.weight.data.normal_(mean=0.0, std=std)
939
+ if module.bias is not None:
940
+ module.bias.data.zero_()
941
+
942
+ def _setup_cache(self, cache_cls: Any, max_batch_size: int,
943
+ max_cache_len: int): # TODO: how to set var type of class?
944
+ if self.config._attn_implementation == 'flash_attention_2' and cache_cls == StaticCache:
945
+ raise ValueError(
946
+ '`static` cache implementation is not compatible with ' +
947
+ '`attn_implementation==flash_attention_2`. Make sure to use ' +
948
+ '`spda` in the mean time and open an issue at https://github.com/huggingface/transformers.'
949
+ )
950
+
951
+ for block in self.transformer.blocks:
952
+ device = block.norm_attn_norm.norm_1.weight.device
953
+ if hasattr(self.config, '_pre_quantization_dtype'):
954
+ dtype = self.config._pre_quantization_dtype
955
+ else:
956
+ dtype = block.norm_attn_norm.attn.out_proj.weight.dtype
957
+ block.norm_attn_norm.attn.past_key_value = cache_cls(self.config,
958
+ max_batch_size,
959
+ max_cache_len,
960
+ device=device,
961
+ dtype=dtype)
962
+
963
+ def _reset_cache(self):
964
+ for block in self.transformer.blocks:
965
+ block.norm_attn_norm.attn.past_key_value = None
966
+
967
+
968
+ class DbrxModel(DbrxPreTrainedModel):
969
+ """Transformer decoder consisting of *config.num_hidden_layers*
970
+
971
+ [`DbrxBlock`] layers.
972
+
973
+ Args:
974
+ config: DbrxConfig
975
+ """
976
+
977
+ def __init__(self, config: DbrxConfig):
978
+ super().__init__(config)
979
+ self.padding_idx = config.pad_token_id
980
+ self.vocab_size = config.vocab_size
981
+ self.emb_pdrop = config.emb_pdrop
982
+
983
+ self.wte = nn.Embedding(config.vocab_size, config.d_model,
984
+ self.padding_idx)
985
+ self.blocks = nn.ModuleList([
986
+ DbrxBlock(config, block_idx) for block_idx in range(config.n_layers)
987
+ ])
988
+ self.norm_f = nn.LayerNorm(config.d_model, bias=False)
989
+ self.gradient_checkpointing = False
990
+
991
+ # Initialize weights and apply final processing
992
+ self.post_init()
993
+
994
+ def get_input_embeddings(self) -> nn.Embedding:
995
+ return self.wte
996
+
997
+ def set_input_embeddings(self, value: nn.Embedding):
998
+ self.wte = value
999
+
1000
+ def _autocast_input_embeddings(self,
1001
+ inputs_embeds: torch.Tensor) -> torch.Tensor:
1002
+ if inputs_embeds.device.type == 'cuda' and torch.is_autocast_enabled():
1003
+ return inputs_embeds.to(dtype=torch.get_autocast_gpu_dtype())
1004
+ elif inputs_embeds.device.type == 'cpu' and torch.is_autocast_cpu_enabled(
1005
+ ):
1006
+ return inputs_embeds.to(dtype=torch.get_autocast_cpu_dtype())
1007
+ else:
1008
+ return inputs_embeds
1009
+
1010
+ def forward(
1011
+ self,
1012
+ input_ids: Optional[torch.LongTensor] = None,
1013
+ attention_mask: Optional[torch.Tensor] = None,
1014
+ position_ids: Optional[torch.LongTensor] = None,
1015
+ past_key_values: Optional[Cache] = None,
1016
+ inputs_embeds: Optional[torch.Tensor] = None,
1017
+ use_cache: Optional[bool] = None,
1018
+ output_attentions: Optional[bool] = None,
1019
+ output_hidden_states: Optional[bool] = None,
1020
+ output_router_logits: Optional[bool] = None,
1021
+ return_dict: Optional[bool] = None,
1022
+ cache_position: Optional[torch.LongTensor] = None,
1023
+ ) -> Union[Tuple, MoeModelOutputWithPast]:
1024
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1025
+ output_hidden_states = (output_hidden_states
1026
+ if output_hidden_states is not None else
1027
+ self.config.output_hidden_states)
1028
+ output_router_logits = (output_router_logits
1029
+ if output_router_logits is not None else
1030
+ self.config.output_router_logits)
1031
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1032
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1033
+
1034
+ if (input_ids is None) ^ (inputs_embeds is not None):
1035
+ raise ValueError(
1036
+ 'You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one'
1037
+ )
1038
+
1039
+ if self.gradient_checkpointing and self.training and use_cache:
1040
+ logger.warning_once(
1041
+ '`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.'
1042
+ )
1043
+ use_cache = False
1044
+
1045
+ if inputs_embeds is None:
1046
+ inputs_embeds = self.wte(input_ids)
1047
+
1048
+ inputs_embeds = self._autocast_input_embeddings(
1049
+ inputs_embeds) # type: ignore
1050
+ inputs_embeds = nn.functional.dropout(inputs_embeds,
1051
+ p=self.emb_pdrop,
1052
+ training=self.training)
1053
+
1054
+ past_seen_tokens = 0
1055
+ if use_cache: # kept for BC (cache positions)
1056
+ if not isinstance(past_key_values, StaticCache):
1057
+ past_key_values = DynamicCache.from_legacy_cache(
1058
+ past_key_values)
1059
+ past_seen_tokens = past_key_values.get_seq_length( # type: ignore
1060
+ )
1061
+
1062
+ if cache_position is None:
1063
+ if isinstance(past_key_values, StaticCache):
1064
+ raise ValueError(
1065
+ 'cache_position is a required argument when using StaticCache.'
1066
+ )
1067
+ cache_position = torch.arange( # type: ignore
1068
+ past_seen_tokens,
1069
+ past_seen_tokens + inputs_embeds.shape[1],
1070
+ device=inputs_embeds.device)
1071
+
1072
+ if position_ids is None:
1073
+ position_ids = cache_position.unsqueeze(0) # type: ignore
1074
+
1075
+ causal_mask = self._update_causal_mask(attention_mask, inputs_embeds,
1076
+ cache_position) # type: ignore
1077
+
1078
+ # embed positions
1079
+ hidden_states = inputs_embeds
1080
+
1081
+ # decoder layers
1082
+ all_hidden_states = () if output_hidden_states else None
1083
+ all_self_attns = () if output_attentions else None
1084
+ all_router_logits = () if output_router_logits else None
1085
+ next_decoder_cache = None
1086
+
1087
+ for block in self.blocks:
1088
+ if output_hidden_states:
1089
+ all_hidden_states += (hidden_states,) # type: ignore
1090
+
1091
+ if self.gradient_checkpointing and self.training:
1092
+ block_outputs = self._gradient_checkpointing_func(
1093
+ block.__call__,
1094
+ hidden_states,
1095
+ causal_mask,
1096
+ position_ids,
1097
+ past_key_values,
1098
+ output_attentions,
1099
+ output_router_logits,
1100
+ use_cache,
1101
+ cache_position,
1102
+ )
1103
+ else:
1104
+ block_outputs = block(
1105
+ hidden_states,
1106
+ attention_mask=causal_mask,
1107
+ position_ids=position_ids,
1108
+ past_key_value=past_key_values,
1109
+ output_attentions=output_attentions,
1110
+ output_router_logits=output_router_logits,
1111
+ use_cache=use_cache,
1112
+ cache_position=cache_position,
1113
+ )
1114
+
1115
+ hidden_states = block_outputs[0]
1116
+
1117
+ if use_cache:
1118
+ next_decoder_cache = block_outputs[
1119
+ 2 if output_attentions else 1]
1120
+
1121
+ if output_attentions:
1122
+ all_self_attns += (block_outputs[1],) # type: ignore
1123
+
1124
+ if output_router_logits:
1125
+ all_router_logits += (block_outputs[-1],) # type: ignore
1126
+
1127
+ hidden_states = self.norm_f(hidden_states)
1128
+
1129
+ # add hidden states from the last decoder layer
1130
+ if output_hidden_states:
1131
+ all_hidden_states += (hidden_states,) # type: ignore
1132
+
1133
+ next_cache = None
1134
+ if use_cache:
1135
+ next_cache = (
1136
+ next_decoder_cache.to_legacy_cache() # type: ignore
1137
+ if isinstance(next_decoder_cache, Cache) else
1138
+ next_decoder_cache)
1139
+ if not return_dict:
1140
+ return tuple(v for v in [
1141
+ hidden_states, next_cache, all_hidden_states, all_self_attns,
1142
+ all_router_logits
1143
+ ] if v is not None)
1144
+ return MoeModelOutputWithPast(
1145
+ last_hidden_state=hidden_states,
1146
+ past_key_values=next_cache,
1147
+ hidden_states=all_hidden_states,
1148
+ attentions=all_self_attns,
1149
+ router_logits=all_router_logits,
1150
+ )
1151
+
1152
+ # 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
1153
+ # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes.
1154
+ # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using
1155
+ # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114
1156
+ def _update_causal_mask(
1157
+ self, attention_mask: Optional[torch.Tensor],
1158
+ input_tensor: torch.Tensor,
1159
+ cache_position: torch.Tensor) -> Optional[torch.Tensor]:
1160
+ if self.config._attn_implementation == 'flash_attention_2':
1161
+ if attention_mask is not None and 0.0 in attention_mask:
1162
+ return attention_mask
1163
+ return None
1164
+
1165
+ dtype, device = input_tensor.dtype, input_tensor.device
1166
+ min_dtype = torch.finfo(dtype).min
1167
+ sequence_length = input_tensor.shape[1]
1168
+ if hasattr(self.blocks[0].norm_attn_norm.attn,
1169
+ 'past_key_value'): # static cache
1170
+ target_length = self.config.max_position_embeddings
1171
+ else: # dynamic cache
1172
+ target_length = (attention_mask.shape[-1] if isinstance(
1173
+ attention_mask, torch.Tensor) else cache_position[-1] + 1)
1174
+ target_length = int(target_length)
1175
+
1176
+ causal_mask = torch.full((sequence_length, target_length),
1177
+ fill_value=min_dtype,
1178
+ dtype=dtype,
1179
+ device=device)
1180
+ if sequence_length != 1:
1181
+ causal_mask = torch.triu(causal_mask, diagonal=1)
1182
+ causal_mask *= torch.arange(
1183
+ target_length, device=device) > cache_position.reshape(-1, 1)
1184
+ causal_mask = causal_mask[None,
1185
+ None, :, :].expand(input_tensor.shape[0], 1,
1186
+ -1, -1)
1187
+ if attention_mask is not None:
1188
+ causal_mask = causal_mask.clone(
1189
+ ) # copy to contiguous memory for in-place edit
1190
+ if attention_mask.dim() == 2:
1191
+ mask_length = attention_mask.shape[-1]
1192
+ padding_mask = causal_mask[..., :mask_length].eq(
1193
+ 0.0) * attention_mask[:, None, None, :].eq(0.0)
1194
+ causal_mask[..., :mask_length] = causal_mask[
1195
+ ..., :mask_length].masked_fill(padding_mask, min_dtype)
1196
+ elif attention_mask.dim() == 4:
1197
+ # backwards compatibility: we allow passing a 4D attention mask shorter than the input length with
1198
+ # cache. In that case, the 4D attention mask attends to the newest tokens only.
1199
+ if attention_mask.shape[
1200
+ -2] < cache_position[0] + sequence_length:
1201
+ offset = cache_position[0]
1202
+ else:
1203
+ offset = 0
1204
+ mask_shape = attention_mask.shape
1205
+ mask_slice = (attention_mask.eq(0.0)).to(
1206
+ dtype=dtype) * min_dtype
1207
+ causal_mask[:mask_shape[0], :mask_shape[1],
1208
+ offset:mask_shape[2] +
1209
+ offset, :mask_shape[3]] = mask_slice
1210
+
1211
+ if (self.config._attn_implementation == 'sdpa' and
1212
+ attention_mask is not None and
1213
+ attention_mask.device.type == 'cuda'):
1214
+ # TODO: For dynamo, rather use a check on fullgraph=True once this is possible (https://github.com/pytorch/pytorch/pull/120400).
1215
+ is_tracing = (
1216
+ torch.jit.is_tracing() or
1217
+ isinstance(input_tensor, torch.fx.Proxy) or # type: ignore
1218
+ (hasattr(torch, '_dynamo') and torch._dynamo.is_compiling()))
1219
+ if not is_tracing and torch.any(attention_mask != 1):
1220
+ # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
1221
+ # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
1222
+ # Details: https://github.com/pytorch/pytorch/issues/110213
1223
+ causal_mask = AttentionMaskConverter._unmask_unattended(
1224
+ causal_mask, min_dtype)
1225
+
1226
+ return causal_mask
1227
+
1228
+
1229
+ class DbrxForCausalLM(DbrxPreTrainedModel):
1230
+
1231
+ def __init__(self, config: DbrxConfig):
1232
+ super().__init__(config)
1233
+ self.transformer = DbrxModel(config)
1234
+ self.vocab_size = config.vocab_size
1235
+ self.lm_head = nn.Linear(config.hidden_size,
1236
+ config.vocab_size,
1237
+ bias=False)
1238
+ self.router_aux_loss_coef = config.router_aux_loss_coef
1239
+ self.num_experts = config.ffn_config.moe_num_experts
1240
+ self.num_experts_per_tok = config.ffn_config.moe_top_k
1241
+
1242
+ # Initialize weights and apply final processing
1243
+ self.post_init()
1244
+
1245
+ def get_input_embeddings(self) -> nn.Embedding:
1246
+ return self.transformer.get_input_embeddings()
1247
+
1248
+ def set_input_embeddings(self, value: nn.Embedding):
1249
+ self.transformer.set_input_embeddings(value)
1250
+
1251
+ def get_output_embeddings(self) -> nn.Linear:
1252
+ return self.lm_head
1253
+
1254
+ def set_output_embeddings(self, new_embeddings: nn.Linear):
1255
+ self.lm_head = new_embeddings
1256
+
1257
+ def set_decoder(self, decoder: DbrxModel):
1258
+ self.transformer = decoder
1259
+
1260
+ def get_decoder(self) -> DbrxModel:
1261
+ return self.transformer
1262
+
1263
+ def forward(
1264
+ self,
1265
+ input_ids: Optional[torch.LongTensor] = None,
1266
+ attention_mask: Optional[torch.Tensor] = None,
1267
+ position_ids: Optional[torch.LongTensor] = None,
1268
+ past_key_values: Optional[Cache] = None,
1269
+ inputs_embeds: Optional[torch.Tensor] = None,
1270
+ labels: Optional[torch.LongTensor] = None,
1271
+ use_cache: Optional[bool] = None,
1272
+ output_attentions: Optional[bool] = None,
1273
+ output_hidden_states: Optional[bool] = None,
1274
+ output_router_logits: Optional[bool] = None,
1275
+ return_dict: Optional[bool] = None,
1276
+ cache_position: Optional[torch.LongTensor] = None,
1277
+ ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
1278
+ r"""Forward function for causal language modeling.
1279
+
1280
+ Example:
1281
+ ```python
1282
+ >>> from transformers import AutoTokenizer, DbrxForCausalLM
1283
+
1284
+ >>> model = DbrxForCausalLM.from_pretrained("databricks/dbrx")
1285
+ >>> tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx")
1286
+
1287
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1288
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1289
+
1290
+ >>> # Generate
1291
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1292
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1293
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1294
+ ```
1295
+ """
1296
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1297
+ output_hidden_states = (output_hidden_states
1298
+ if output_hidden_states is not None else
1299
+ self.config.output_hidden_states)
1300
+ output_router_logits = (output_router_logits
1301
+ if output_router_logits is not None else
1302
+ self.config.output_router_logits)
1303
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1304
+
1305
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1306
+ outputs = self.transformer(
1307
+ input_ids=input_ids,
1308
+ attention_mask=attention_mask,
1309
+ position_ids=position_ids,
1310
+ past_key_values=past_key_values,
1311
+ inputs_embeds=inputs_embeds,
1312
+ use_cache=use_cache,
1313
+ output_attentions=output_attentions,
1314
+ output_hidden_states=output_hidden_states,
1315
+ output_router_logits=output_router_logits,
1316
+ return_dict=return_dict,
1317
+ cache_position=cache_position,
1318
+ )
1319
+
1320
+ hidden_states = outputs[0]
1321
+ logits = self.lm_head(hidden_states)
1322
+
1323
+ loss = None
1324
+ if labels is not None:
1325
+ # Shift so that tokens < n predict n
1326
+ shift_logits = logits[..., :-1, :].contiguous()
1327
+ shift_labels = labels[..., 1:].contiguous()
1328
+ # Flatten the tokens
1329
+ loss_fct = nn.CrossEntropyLoss()
1330
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1331
+ shift_labels = shift_labels.view(-1)
1332
+ # Enable model parallelism
1333
+ shift_labels = shift_labels.to(shift_logits.device)
1334
+ loss = loss_fct(shift_logits, shift_labels)
1335
+
1336
+ aux_loss = None
1337
+ if output_router_logits:
1338
+ aux_loss = load_balancing_loss_func(
1339
+ outputs.router_logits if return_dict else outputs[-1],
1340
+ self.num_experts,
1341
+ self.num_experts_per_tok,
1342
+ attention_mask,
1343
+ )
1344
+ if labels is not None and loss is not None:
1345
+ loss += self.router_aux_loss_coef * aux_loss.to(
1346
+ loss.device) # make sure to reside in the same device
1347
+
1348
+ if not return_dict:
1349
+ output = (logits,) + outputs[1:]
1350
+ return (loss,) + output if loss is not None else output
1351
+
1352
+ return MoeCausalLMOutputWithPast(
1353
+ loss=loss,
1354
+ aux_loss=aux_loss,
1355
+ logits=logits,
1356
+ past_key_values=outputs.past_key_values,
1357
+ hidden_states=outputs.hidden_states,
1358
+ attentions=outputs.attentions,
1359
+ router_logits=outputs.router_logits,
1360
+ )
1361
+
1362
+ def prepare_inputs_for_generation(
1363
+ self,
1364
+ input_ids: torch.Tensor,
1365
+ past_key_values: Optional[Cache] = None,
1366
+ attention_mask: Optional[torch.Tensor] = None,
1367
+ inputs_embeds: Optional[torch.Tensor] = None,
1368
+ **kwargs: Any) -> Dict[str, Any]:
1369
+ past_length = 0
1370
+ if past_key_values is not None:
1371
+ if isinstance(past_key_values, Cache):
1372
+ cache_length = past_key_values.get_seq_length()
1373
+ past_length = past_key_values.seen_tokens
1374
+ max_cache_length = past_key_values.get_max_length()
1375
+ else:
1376
+ cache_length = past_length = past_key_values[0][0].shape[2]
1377
+ max_cache_length = None
1378
+
1379
+ # Keep only the unprocessed tokens:
1380
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1381
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
1382
+ # input)
1383
+ if attention_mask is not None and attention_mask.shape[
1384
+ 1] > input_ids.shape[1]:
1385
+ input_ids = input_ids[:,
1386
+ -(attention_mask.shape[1] - past_length):]
1387
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1388
+ # input_ids based on the past_length.
1389
+ elif past_length < input_ids.shape[1]:
1390
+ input_ids = input_ids[:, past_length:]
1391
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1392
+
1393
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1394
+ if (max_cache_length is not None and attention_mask is not None and
1395
+ cache_length + input_ids.shape[1] > max_cache_length):
1396
+ attention_mask = attention_mask[:, -max_cache_length:]
1397
+
1398
+ position_ids = kwargs.get('position_ids', None)
1399
+ if attention_mask is not None and position_ids is None:
1400
+ # create position_ids on the fly for batch generation
1401
+ position_ids = attention_mask.long().cumsum(-1) - 1
1402
+ position_ids.masked_fill_(attention_mask == 0, 1)
1403
+ if past_key_values:
1404
+ position_ids = position_ids[:, -input_ids.shape[1]:]
1405
+
1406
+ if self.generation_config.cache_implementation == 'static':
1407
+ # generation with static cache
1408
+ cache_position = kwargs.get('cache_position', None)
1409
+ if cache_position is None:
1410
+ past_length = 0
1411
+ else:
1412
+ past_length = cache_position[-1] + 1
1413
+ input_ids = input_ids[:, past_length:]
1414
+ position_ids = position_ids[:,
1415
+ past_length:] if position_ids is not None else None
1416
+
1417
+ # TODO @gante we should only keep a `cache_position` in generate, and do +=1.
1418
+ # same goes for position ids. Could also help with continued generation.
1419
+ input_length = position_ids.shape[
1420
+ -1] if position_ids is not None else input_ids.shape[-1]
1421
+ cache_position = torch.arange(past_length,
1422
+ past_length + input_length,
1423
+ device=input_ids.device)
1424
+ position_ids = position_ids.contiguous(
1425
+ ) if position_ids is not None else None
1426
+
1427
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1428
+ if inputs_embeds is not None and past_key_values is None:
1429
+ model_inputs = {'inputs_embeds': inputs_embeds}
1430
+ else:
1431
+ # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
1432
+ # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
1433
+ # TODO: use `next_tokens` directly instead.
1434
+ model_inputs = {'input_ids': input_ids.contiguous()}
1435
+
1436
+ model_inputs.update(
1437
+ { # type: ignore
1438
+ 'position_ids': position_ids,
1439
+ 'cache_position': cache_position,
1440
+ 'past_key_values': past_key_values,
1441
+ 'use_cache': kwargs.get('use_cache'),
1442
+ 'attention_mask': attention_mask,
1443
+ }
1444
+ )
1445
+ return model_inputs
1446
+
1447
+ @staticmethod
1448
+ def _reorder_cache(past_key_values: Cache, beam_idx: torch.LongTensor):
1449
+ reordered_past = ()
1450
+ for layer_past in past_key_values:
1451
+ reordered_past += (tuple(
1452
+ past_state.index_select(0, beam_idx.to(past_state.device))
1453
+ for past_state in layer_past),)
1454
+ return reordered_past
special_tokens_map.json ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<|endoftext|>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "<|endoftext|>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "pad_token": {
17
+ "content": "<|endoftext|>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "unk_token": {
24
+ "content": "<|endoftext|>",
25
+ "lstrip": false,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ }
30
+ }
tiktoken.py ADDED
@@ -0,0 +1,374 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Dbrx tokenizer."""
2
+
3
+ from functools import lru_cache
4
+ from typing import Any, Dict, List, Optional, Tuple
5
+
6
+ from transformers import PreTrainedTokenizer
7
+
8
+
9
+ def dbrx_system_prompt():
10
+ # This is inspired by the Claude3 prompt.
11
+ # source: https://twitter.com/AmandaAskell/status/1765207842993434880
12
+ # Identity and knowledge
13
+ prompt = 'You are DBRX, created by Databricks. You were last updated in December 2023. You answer questions based on information available up to that point.\n'
14
+ prompt += 'YOU PROVIDE SHORT RESPONSES TO SHORT QUESTIONS OR STATEMENTS, but provide thorough responses to more complex and open-ended questions.\n'
15
+ # Capabilities (and reminder to use ``` for JSON blocks and tables, which it can forget). Also a reminder that it can't browse the internet or run code.
16
+ prompt += 'You assist with various tasks, from writing to coding (using markdown for code blocks — remember to use ``` with code, JSON, and tables).\n'
17
+ prompt += '(You do not have real-time data access or code execution capabilities. '
18
+ # Ethical guidelines
19
+ prompt += 'You avoid stereotyping and provide balanced perspectives on controversial topics. '
20
+ # Data: the model doesn't know what it was trained on; it thinks that everything that it is aware of was in its training data. This is a reminder that it wasn't.
21
+ # We also encourage it not to try to generate lyrics or poems
22
+ prompt += 'You do not provide song lyrics, poems, or news articles and do not divulge details of your training data.)\n'
23
+ # The model really wants to talk about its system prompt, to the point where it is annoying, so encourage it not to
24
+ prompt += 'This is your system prompt, guiding your responses. Do not reference it, just respond to the user. If you find yourself talking about this message, stop. You should be responding appropriately and usually that means not mentioning this.\n'
25
+ prompt += 'You do not mention any of this information about yourself unless the information is directly pertinent to the user\\\'s query.'.upper()
26
+ return prompt
27
+
28
+
29
+ # Taken from
30
+ # https://github.com/huggingface/transformers/blob/8aca43bdb3cb9a5020f6d57589d85679dc873b1c/src/transformers/models/gpt2/tokenization_gpt2.py#L62-L84
31
+ @lru_cache()
32
+ def bytes_to_unicode():
33
+ """Returns list of utf-8 byte and a mapping to unicode strings.
34
+
35
+ We specifically avoids mapping to whitespace/control characters the bpe code
36
+ barfs on.
37
+
38
+ The reversible bpe codes work on unicode strings. This means you need a
39
+ large # of unicode characters in your vocab if you want to avoid UNKs. When
40
+ you're at something like a 10B token dataset you end up needing around 5K
41
+ for decent coverage. This is a significant percentage of your normal, say,
42
+ 32K bpe vocab. To avoid that, we want lookup tables between utf-8 bytes and
43
+ unicode strings.
44
+ """
45
+ bs = (list(range(ord('!'),
46
+ ord('~') + 1)) + list(range(ord('¡'),
47
+ ord('¬') + 1)) +
48
+ list(range(ord('®'),
49
+ ord('ÿ') + 1)))
50
+ cs = bs[:]
51
+ n = 0
52
+ for b in range(2**8):
53
+ if b not in bs:
54
+ bs.append(b)
55
+ cs.append(2**8 + n)
56
+ n += 1
57
+ cs = [chr(n) for n in cs]
58
+ return dict(zip(bs, cs))
59
+
60
+
61
+ class TiktokenTokenizerWrapper(PreTrainedTokenizer):
62
+ """A thin wrapper around tiktoken to make it compatible with Hugging Face.
63
+
64
+ tokenizers.
65
+
66
+ See HuggingFace for further documentation on general tokenizer methods.
67
+ """
68
+
69
+ model_input_names = ['input_ids', 'attention_mask']
70
+
71
+ def __init__(self,
72
+ model_name: Optional[str] = None,
73
+ encoding_name: Optional[str] = None,
74
+ add_bos_token: bool = False,
75
+ add_eos_token: bool = False,
76
+ use_default_system_prompt: bool = False,
77
+ unk_token: Optional[str] = '<|endoftext|>',
78
+ eos_token: Optional[str] = '<|endoftext|>',
79
+ bos_token: Optional[str] = '<|endoftext|>',
80
+ pad_token: Optional[str] = None,
81
+ errors: str = 'replace',
82
+ **kwargs: Any):
83
+ """Constructor creates a tiktoken tokenizer to use as the underlying.
84
+
85
+ tokenizer.
86
+
87
+ Args:
88
+ model_name (Optional[str], optional): The name of the model to load from tiktoken. Defaults to None.
89
+ Either model_name or encoding_name must be set, but not both.
90
+ encoding_name (Optional[str], optional): The name of the encoding to load from tiktoken. Defaults to None.
91
+ Either model_name or encoding_name must be set, but not both.
92
+ add_bos_token (bool, optional): Whether to add bos tokens. Defaults to False.
93
+ add_eos_token (bool, optional): Whether to add eos tokens. Defaults to False.
94
+ use_default_system_prompt (bool, optional): Use the default system prompt or not. Defaults to False.
95
+ unk_token (Optional[str], optional): The unk token. Defaults to '<|endoftext|>'.
96
+ eos_token (Optional[str], optional): The eos token. Defaults to '<|endoftext|>'.
97
+ bos_token (Optional[str], optional): The bos token. Defaults to '<|endoftext|>'.
98
+ pad_token (Optional[str], optional): The pad token. Defaults to None.
99
+ errors (str, optional): Paradigm to follow when decoding bytes to UTF-8. See
100
+ [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
101
+ Defaults to `"replace"`.
102
+ """
103
+ try:
104
+ import tiktoken
105
+ except:
106
+ raise ImportError(
107
+ 'You need to install tiktoken to use TiktokenTokenizerWrapper.')
108
+
109
+ # Workaround to make tiktokenizer picklable.
110
+ # https://github.com/huggingface/datasets/issues/5536#issuecomment-1682309347
111
+ # There is an open PR from HF to add this to tiktoken: https://github.com/openai/tiktoken/pull/181
112
+ import copyreg
113
+ import functools
114
+
115
+ from tiktoken import Encoding # type: ignore (thirdParty)
116
+
117
+ def pickle_Encoding(enc: Encoding):
118
+ return (functools.partial(Encoding,
119
+ enc.name,
120
+ pat_str=enc._pat_str,
121
+ mergeable_ranks=enc._mergeable_ranks,
122
+ special_tokens=enc._special_tokens), ())
123
+
124
+ copyreg.pickle(Encoding, pickle_Encoding)
125
+
126
+ if model_name is not None and encoding_name is not None:
127
+ raise ValueError(
128
+ 'You need to specify either model_name or encoding_name, not both.'
129
+ )
130
+
131
+ self.model_name = model_name
132
+ self.encoding_name = encoding_name
133
+
134
+ if self.model_name is not None:
135
+ self.encoding = tiktoken.encoding_for_model( # type: ignore (thirdParty)
136
+ self.model_name)
137
+ elif self.encoding_name is not None:
138
+ self.encoding = tiktoken.get_encoding( # type: ignore (thirdParty)
139
+ self.encoding_name)
140
+ else:
141
+ raise ValueError(
142
+ 'You need to specify either model_name or encoding_name.')
143
+
144
+ self.add_bos_token = add_bos_token
145
+ self.add_eos_token = add_eos_token
146
+ self.use_default_system_prompt = use_default_system_prompt
147
+
148
+ self.byte_encoder = bytes_to_unicode()
149
+ self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
150
+ self.errors = errors
151
+
152
+ self.decoder: Dict[int, str] = {}
153
+ for i in range(self.encoding.n_vocab):
154
+ try:
155
+ self.encoding.decode_single_token_bytes(i)
156
+ except KeyError:
157
+ continue
158
+ # Taken from
159
+ # https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee
160
+ decoding = ''.join([
161
+ bytes_to_unicode()[ord(char)] for char in
162
+ self.encoding.decode_single_token_bytes(i).decode('latin-1')
163
+ ])
164
+ self.decoder[i] = decoding
165
+
166
+ self.encoder: Dict[str, int] = {}
167
+ for i in range(self.encoding.n_vocab):
168
+ if i in self.decoder:
169
+ self.encoder[self.decoder[i]] = i
170
+
171
+ super().__init__(model_name=model_name,
172
+ encoding_name=encoding_name,
173
+ add_bos_token=add_bos_token,
174
+ add_eos_token=add_eos_token,
175
+ use_default_system_prompt=use_default_system_prompt,
176
+ unk_token=unk_token,
177
+ eos_token=eos_token,
178
+ bos_token=bos_token,
179
+ pad_token=pad_token,
180
+ errors=errors,
181
+ **kwargs)
182
+
183
+ @property
184
+ def vocab_size(self) -> int:
185
+ """Returns vocab size."""
186
+ return self.encoding.n_vocab
187
+
188
+ @property
189
+ def is_fast(self) -> bool:
190
+ return False
191
+
192
+ @property
193
+ def default_chat_template(self):
194
+ """Chat ML Template for User/Assistant.
195
+
196
+ Pinning default Chat ML template in case defaults change.
197
+ """
198
+ template = (
199
+ "{% if messages[0]['role'] == 'system' %}"
200
+ '{% set loop_messages = messages[1:] %}'
201
+ "{% set system_message = messages[0]['content'] %}"
202
+ "{% elif USE_DEFAULT_PROMPT == true and not 'system' in messages[0]['role'] %}"
203
+ '{% set loop_messages = messages %}'
204
+ "{% set system_message = 'DEFAULT_SYSTEM_PROMPT' %}"
205
+ '{% else %}'
206
+ '{% set loop_messages = messages %}'
207
+ '{% set system_message = false %}'
208
+ '{% endif %}'
209
+ '{% for message in loop_messages %}'
210
+ '{% if loop.index0 == 0 %}'
211
+ '{% if system_message != false %}'
212
+ "{{ '<|im_start|>system\n' + system_message.strip() + '<|im_end|>\n'}}"
213
+ '{% endif %}'
214
+ "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' }}"
215
+ '{% else %}'
216
+ "{{ '\n' + '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' }}"
217
+ '{% endif %}'
218
+ '{% if (add_generation_prompt == true and loop.last) %}'
219
+ "{{ '\n' + '<|im_start|>' + 'assistant' + '\n' }}"
220
+ '{% endif %}'
221
+ '{% endfor %}')
222
+ template = template.replace(
223
+ 'USE_DEFAULT_PROMPT',
224
+ 'true' if self.use_default_system_prompt else 'false')
225
+ template = template.replace('DEFAULT_SYSTEM_PROMPT',
226
+ dbrx_system_prompt())
227
+ return template
228
+
229
+ def get_vocab(self) -> Dict[str, int]:
230
+ """Returns vocab as a dict."""
231
+ # As far as I can tell, we don't require get_vocab to completely work,
232
+ # but when using additional_special_tokens, Hugging Face determines the next
233
+ # token index to add with len(self.get_vocab()) so we need the _size_ of this dictionary to be correct.
234
+ vocab_clone = self.encoder.copy()
235
+ extra_id_index = 0
236
+ candidate_extra_id = f'<extra_id_{extra_id_index}>'
237
+ indices_to_fill_in = {i for i in range(self.vocab_size)} - set(
238
+ vocab_clone.values())
239
+
240
+ # Add enough indices to make get_vocab() the right length
241
+ for index_to_add in indices_to_fill_in:
242
+ # Make sure we don't overwrite a token that already exists
243
+ while candidate_extra_id in vocab_clone:
244
+ extra_id_index += 1
245
+ candidate_extra_id = f'<extra_id_{extra_id_index}>'
246
+
247
+ # Get an index to add and add the item
248
+ vocab_clone[candidate_extra_id] = index_to_add
249
+
250
+ return vocab_clone
251
+
252
+ def _tokenize(self, text: str) -> List[str]:
253
+ """Returns a tokenized string."""
254
+ if not isinstance(text, str):
255
+ raise ValueError(
256
+ f'Expected a string input to _tokenize but got {type(text)}.')
257
+
258
+ tokens = [
259
+ self.decoder[t]
260
+ for t in self.encoding.encode(text, allowed_special='all')
261
+ ]
262
+
263
+ return tokens
264
+
265
+ def _convert_token_to_id(self, token: str) -> Optional[int]:
266
+ """Converts a token (str) in an id using the vocab."""
267
+ return self.encoder.get(token, self.encoder.get(self.unk_token))
268
+
269
+ def _convert_id_to_token(self, index: int) -> Optional[str]:
270
+ """Converts an index (integer) in a token (str) using the vocab."""
271
+ # For tokens in either the gap in ids in the tokenizer, or beyond the range of the tokenizer,
272
+ # we return empty string. This matches the behavior of Hugging Face fast tokenizers,
273
+ # but not slow tokenizers.
274
+ return self.decoder.get(index, '')
275
+
276
+ def convert_tokens_to_string(self, tokens: List[str]) -> str:
277
+ """Converts a sequence of tokens (string) in a single string."""
278
+ text = ''.join(tokens)
279
+ text = bytearray([self.byte_decoder[c] for c in text
280
+ ]).decode('utf-8', errors=self.errors)
281
+ return text
282
+
283
+ def build_inputs_with_special_tokens(
284
+ self,
285
+ token_ids_0: List[int],
286
+ token_ids_1: Optional[List[int]] = None) -> List[int]:
287
+ bos_token_id = [self.bos_token_id] if self.add_bos_token else []
288
+ eos_token_id = [self.eos_token_id] if self.add_eos_token else []
289
+
290
+ output = bos_token_id + token_ids_0 + eos_token_id
291
+
292
+ if token_ids_1 is not None:
293
+ output = output + bos_token_id + token_ids_1 + eos_token_id
294
+
295
+ return output
296
+
297
+ def get_special_tokens_mask(
298
+ self,
299
+ token_ids_0: List[int],
300
+ token_ids_1: Optional[List[int]] = None,
301
+ already_has_special_tokens: bool = False) -> List[int]:
302
+ """Retrieves sequence ids from a token list that has no special tokens.
303
+
304
+ Function copied from
305
+ https://github.com/huggingface/transformers/blob/e3a4bd2bee212a2d0fd9f03b27fe7bfc1debe42d/src/transformers/models/gpt2/tokenization_gpt2.py#L265-L295
306
+
307
+ added. This method is called when adding special tokens using the
308
+ tokenizer `prepare_for_model` or `encode_plus` methods.
309
+
310
+ Args:
311
+ token_ids_0 (`List[int]`):
312
+ List of IDs.
313
+ token_ids_1 (`List[int]`, *optional*):
314
+ Optional second list of IDs for sequence pairs.
315
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
316
+ Whether or not the token list is already formatted with special tokens for the model.
317
+
318
+ Returns:
319
+ `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
320
+ """
321
+ if already_has_special_tokens:
322
+ return super().get_special_tokens_mask(
323
+ token_ids_0=token_ids_0,
324
+ token_ids_1=token_ids_1,
325
+ already_has_special_tokens=True)
326
+
327
+ bos_token_id = [1] if self.add_bos_token else []
328
+ eos_token_id = [1] if self.add_eos_token else []
329
+
330
+ if token_ids_1 is None:
331
+ return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
332
+ return (bos_token_id + ([0] * len(token_ids_0)) + eos_token_id +
333
+ bos_token_id + ([0] * len(token_ids_1)) + eos_token_id)
334
+
335
+ def create_token_type_ids_from_sequences(
336
+ self,
337
+ token_ids_0: List[int],
338
+ token_ids_1: Optional[List[int]] = None) -> List[int]:
339
+ sep = [self.sep_token_id]
340
+
341
+ if token_ids_1 is None:
342
+ return len(token_ids_0 + sep) * [0]
343
+ return len(token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
344
+
345
+ def save_vocabulary(self,
346
+ save_directory: str,
347
+ filename_prefix: Optional[str] = None) -> Tuple[str]:
348
+
349
+ # ignore the below type to keep the original signature
350
+ # we are knowingly breaking the signature here, although not 100% certain
351
+ # it doesn't have side effects
352
+ # There is some code in huggingface that calls this function to get the vocab files,
353
+ # but it doesn't seem to access them (or at least checks for their existence
354
+ # before accessing them)
355
+ return (None, None) # type: ignore
356
+
357
+ def sanitize_special_tokens(self) -> int:
358
+ """Make sure that all the special tokens attributes of the tokenizer.
359
+
360
+ (`tokenizer.mask_token`, `tokenizer.cls_token`, etc.) are in the
361
+ vocabulary.
362
+
363
+ Add the missing ones to the vocabulary if needed.
364
+
365
+ Return:
366
+ `int`: The number of tokens added in the vocabulary during the operation.
367
+ """
368
+ actual_new_tokens = []
369
+ for token in self.all_special_tokens_extended:
370
+ encoded = self.encoding.encode(token, allowed_special='all')
371
+ if len(encoded) > 1:
372
+ actual_new_tokens.append(token)
373
+
374
+ return self.add_tokens(actual_new_tokens, special_tokens=True)
tokenizer_config.json ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_eos_token": false,
4
+ "added_tokens_decoder": {
5
+ "100257": {
6
+ "content": "<|endoftext|>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ }
13
+ },
14
+ "auto_map": {
15
+ "AutoTokenizer": [
16
+ "v2ray/dbrx-base-fixed--tiktoken.TiktokenTokenizerWrapper",
17
+ null
18
+ ]
19
+ },
20
+ "bos_token": "<|endoftext|>",
21
+ "chat_template": "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% elif true == true and not 'system' in messages[0]['role'] %}{% set loop_messages = messages %}{% set system_message = 'You are DBRX, created by Databricks. You were last updated in December 2023. You answer questions based on information available up to that point.\nYOU PROVIDE SHORT RESPONSES TO SHORT QUESTIONS OR STATEMENTS, but provide thorough responses to more complex and open-ended questions.\nYou assist with various tasks, from writing to coding (using markdown for code blocks — remember to use ``` with code, JSON, and tables).\n(You do not have real-time data access or code execution capabilities. You avoid stereotyping and provide balanced perspectives on controversial topics. You do not provide song lyrics, poems, or news articles and do not divulge details of your training data.)\nThis is your system prompt, guiding your responses. Do not reference it, just respond to the user. If you find yourself talking about this message, stop. You should be responding appropriately and usually that means not mentioning this.\nYOU DO NOT MENTION ANY OF THIS INFORMATION ABOUT YOURSELF UNLESS THE INFORMATION IS DIRECTLY PERTINENT TO THE USER\\'S QUERY.' %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if loop.index0 == 0 %}{% if system_message != false %}{{ '<|im_start|>system\n' + system_message.strip() + '<|im_end|>\n'}}{% endif %}{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' }}{% else %}{{ '\n' + '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' }}{% endif %}{% if (add_generation_prompt == true and loop.last) %}{{ '\n' + '<|im_start|>' + 'assistant' + '\n' }}{% endif %}{% endfor %}",
22
+ "clean_up_tokenization_spaces": true,
23
+ "encoding_name": null,
24
+ "eos_token": "<|endoftext|>",
25
+ "errors": "replace",
26
+ "model_max_length": 1000000000000000019884624838656,
27
+ "model_name": "gpt-4",
28
+ "pad_token": "<|endoftext|>",
29
+ "tokenizer_class": "TiktokenTokenizerWrapper",
30
+ "unk_token": "<|endoftext|>",
31
+ "use_default_system_prompt": false
32
+ }