emozilla commited on
Commit
b9eb078
1 Parent(s): 4aafd9c

initial commit

Browse files
adapt_tokenizer.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ from typing import Union
3
+ from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
4
+ Tokenizer = Union[(PreTrainedTokenizer, PreTrainedTokenizerFast)]
5
+ NUM_SENTINEL_TOKENS: int = 100
6
+
7
+ def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):
8
+ 'Adds sentinel tokens and padding token (if missing).\n\n Expands the tokenizer vocabulary to include sentinel tokens\n used in mixture-of-denoiser tasks as well as a padding token.\n\n All added tokens are added as special tokens. No tokens are\n added if sentinel tokens and padding token already exist.\n '
9
+ sentinels_to_add = [f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)]
10
+ tokenizer.add_tokens(sentinels_to_add, special_tokens=True)
11
+ if (tokenizer.pad_token is None):
12
+ tokenizer.add_tokens('<pad>', special_tokens=True)
13
+ tokenizer.pad_token = '<pad>'
14
+ assert (tokenizer.pad_token_id is not None)
15
+ sentinels = ''.join([f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)])
16
+ _sentinel_token_ids = tokenizer(sentinels, add_special_tokens=False).input_ids
17
+ tokenizer.sentinel_token_ids = _sentinel_token_ids
18
+
19
+ class AutoTokenizerForMOD(AutoTokenizer):
20
+ 'AutoTokenizer + Adaptation for MOD.\n\n A simple wrapper around AutoTokenizer to make instantiating\n an MOD-adapted tokenizer a bit easier.\n\n MOD-adapted tokenizers have sentinel tokens (e.g., <extra_id_0>),\n a padding token, and a property to get the token ids of the\n sentinel tokens.\n '
21
+
22
+ @classmethod
23
+ def from_pretrained(cls, *args, **kwargs):
24
+ 'See `AutoTokenizer.from_pretrained` docstring.'
25
+ tokenizer = super().from_pretrained(*args, **kwargs)
26
+ adapt_tokenizer_for_denoising(tokenizer)
27
+ return tokenizer
attention.py ADDED
@@ -0,0 +1,280 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ 'Attention layers.'
3
+ import math
4
+ import warnings
5
+ from typing import Optional
6
+ import torch
7
+ import torch.nn as nn
8
+ from einops import rearrange
9
+ from packaging import version
10
+ from torch import nn
11
+ from .norm import LPLayerNorm
12
+
13
+ def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool):
14
+ if (original_is_causal and (num_query_tokens != num_key_tokens)):
15
+ if (num_query_tokens != 1):
16
+ raise NotImplementedError('MPT does not support query and key with different number of tokens, unless number of query tokens is 1.')
17
+ else:
18
+ return False
19
+ return original_is_causal
20
+
21
+ def scaled_multihead_dot_product_attention(query, key, value, n_heads, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
22
+ q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)
23
+ k = rearrange(key, 'b s (h d) -> b h d s', h=(1 if multiquery else n_heads))
24
+ v = rearrange(value, 'b s (h d) -> b h s d', h=(1 if multiquery else n_heads))
25
+ min_val = torch.finfo(q.dtype).min
26
+ (b, _, s_q, d) = q.shape
27
+ s_k = k.size((- 1))
28
+ if (softmax_scale is None):
29
+ softmax_scale = (1 / math.sqrt(d))
30
+ attn_weight = (q.matmul(k) * softmax_scale)
31
+ if (attn_bias is not None):
32
+ if (((attn_bias.size((- 1)) != 1) and (attn_bias.size((- 1)) != s_k)) or ((attn_bias.size((- 2)) != 1) and (attn_bias.size((- 2)) != s_q))):
33
+ raise RuntimeError(f'attn_bias (shape: {attn_bias.shape}) is expected to broadcast to shape: {attn_weight.shape}.')
34
+ attn_weight = (attn_weight + attn_bias)
35
+ if (key_padding_mask is not None):
36
+ if (attn_bias is not None):
37
+ warnings.warn((((('Propogating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ') + 'unneccessary computation/memory usage. Consider integrating ') + 'into attn_bias once and passing that to each attention ') + 'module instead.'))
38
+ attn_weight = attn_weight.masked_fill((~ key_padding_mask.view((b, 1, 1, s_k))), min_val)
39
+ if is_causal:
40
+ s = max(s_q, s_k)
41
+ causal_mask = attn_weight.new_ones(s, s, dtype=torch.float16)
42
+ causal_mask = causal_mask.tril()
43
+ causal_mask = causal_mask.to(torch.bool)
44
+ causal_mask = (~ causal_mask)
45
+ causal_mask = causal_mask[(- s_q):, (- s_k):]
46
+ attn_weight = attn_weight.masked_fill(causal_mask.view(1, 1, s_q, s_k), min_val)
47
+ attn_weight = torch.softmax(attn_weight, dim=(- 1))
48
+ if dropout_p:
49
+ attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p, training=training, inplace=True)
50
+ out = attn_weight.matmul(v)
51
+ out = rearrange(out, 'b h s d -> b s (h d)')
52
+ if needs_weights:
53
+ return (out, attn_weight)
54
+ return (out, None)
55
+
56
+ def check_valid_inputs(*tensors, valid_dtypes=[torch.float16, torch.bfloat16]):
57
+ for tensor in tensors:
58
+ if (tensor.dtype not in valid_dtypes):
59
+ raise TypeError(f'tensor.dtype={tensor.dtype!r} must be in valid_dtypes={valid_dtypes!r}.')
60
+ if (not tensor.is_cuda):
61
+ raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
62
+
63
+ def flash_attn_fn(query, key, value, n_heads, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
64
+ try:
65
+ from flash_attn import bert_padding, flash_attn_interface
66
+ except:
67
+ raise RuntimeError('Please install flash-attn==1.0.3.post0')
68
+ check_valid_inputs(query, key, value)
69
+ if (attn_bias is not None):
70
+ raise NotImplementedError(f'attn_bias not implemented for flash attn.')
71
+ (batch_size, seqlen) = query.shape[:2]
72
+ if (key_padding_mask is None):
73
+ key_padding_mask = torch.ones_like(key[:, :, 0], dtype=torch.bool)
74
+ query_padding_mask = key_padding_mask[:, (- query.size(1)):]
75
+ (query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(query, query_padding_mask)
76
+ query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
77
+ (key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(key, key_padding_mask)
78
+ key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=(1 if multiquery else n_heads))
79
+ (value_unpad, _, _, _) = bert_padding.unpad_input(value, key_padding_mask)
80
+ value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=(1 if multiquery else n_heads))
81
+ if multiquery:
82
+ key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size((- 1)))
83
+ value_unpad = value_unpad.expand(value_unpad.size(0), n_heads, value_unpad.size((- 1)))
84
+ dropout_p = (dropout_p if training else 0.0)
85
+ reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
86
+ output_unpad = flash_attn_interface.flash_attn_unpadded_func(query_unpad, key_unpad, value_unpad, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
87
+ output = bert_padding.pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size, seqlen)
88
+ return (output, None)
89
+
90
+ def triton_flash_attn_fn(query, key, value, n_heads, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
91
+ try:
92
+ from .flash_attn_triton import flash_attn_func
93
+ except:
94
+ _installed = False
95
+ if (version.parse(torch.__version__) < version.parse('2.0.0')):
96
+ _installed = True
97
+ try:
98
+ from flash_attn.flash_attn_triton import flash_attn_func
99
+ except:
100
+ _installed = False
101
+ if (not _installed):
102
+ raise RuntimeError('Requirements for `attn_impl: triton` not installed. Either (1) have a CUDA-compatible GPU and `pip install .[gpu]` if installing from llm-foundry source or `pip install triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir#subdirectory=python` if installing from pypi, or (2) use torch attn model.attn_config.attn_impl=torch (torch attn_impl will be slow). Note: (1) requires you have CMake and PyTorch already installed.')
103
+ check_valid_inputs(query, key, value)
104
+ if dropout_p:
105
+ raise NotImplementedError(f'Dropout not implemented for attn_impl: triton.')
106
+ if needs_weights:
107
+ raise NotImplementedError(f'attn_impl: triton cannot return attn weights.')
108
+ if (key_padding_mask is not None):
109
+ warnings.warn((((('Propagating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ') + 'unnecessary computation/memory usage. Consider integrating ') + 'into attn_bias once and passing that to each attention ') + 'module instead.'))
110
+ (b_size, s_k) = key_padding_mask.shape[:2]
111
+ if (attn_bias is None):
112
+ attn_bias = query.new_zeros(b_size, 1, 1, s_k)
113
+ attn_bias = attn_bias.masked_fill((~ key_padding_mask.view((b_size, 1, 1, s_k))), torch.finfo(query.dtype).min)
114
+ query = rearrange(query, 'b s (h d) -> b s h d', h=n_heads)
115
+ key = rearrange(key, 'b s (h d) -> b s h d', h=(1 if multiquery else n_heads))
116
+ value = rearrange(value, 'b s (h d) -> b s h d', h=(1 if multiquery else n_heads))
117
+ if multiquery:
118
+ key = key.expand(*key.shape[:2], n_heads, key.size((- 1)))
119
+ value = value.expand(*value.shape[:2], n_heads, value.size((- 1)))
120
+ reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
121
+ attn_output = flash_attn_func(query, key, value, attn_bias, reset_is_causal, softmax_scale)
122
+ output = attn_output.view(*attn_output.shape[:2], (- 1))
123
+ return (output, None)
124
+
125
+ class MultiheadAttention(nn.Module):
126
+ 'Multi-head self attention.\n\n Using torch or triton attention implemetation enables user to also use\n additive bias.\n '
127
+
128
+ def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, low_precision_layernorm: bool=False, verbose: int=0, device: Optional[str]=None):
129
+ super().__init__()
130
+ self.attn_impl = attn_impl
131
+ self.clip_qkv = clip_qkv
132
+ self.qk_ln = qk_ln
133
+ self.d_model = d_model
134
+ self.n_heads = n_heads
135
+ self.softmax_scale = softmax_scale
136
+ if (self.softmax_scale is None):
137
+ self.softmax_scale = (1 / math.sqrt((self.d_model / self.n_heads)))
138
+ self.attn_dropout_p = attn_pdrop
139
+ self.Wqkv = nn.Linear(self.d_model, (3 * self.d_model), device=device)
140
+ fuse_splits = (d_model, (2 * d_model))
141
+ self.Wqkv._fused = (0, fuse_splits)
142
+ if self.qk_ln:
143
+ layernorm_class = (LPLayerNorm if low_precision_layernorm else nn.LayerNorm)
144
+ self.q_ln = layernorm_class(self.d_model, device=device)
145
+ self.k_ln = layernorm_class(self.d_model, device=device)
146
+ if (self.attn_impl == 'flash'):
147
+ self.attn_fn = flash_attn_fn
148
+ elif (self.attn_impl == 'triton'):
149
+ self.attn_fn = triton_flash_attn_fn
150
+ if verbose:
151
+ warnings.warn(((('While `attn_impl: triton` can be faster than `attn_impl: flash` ' + 'it uses more memory. When training larger models this can trigger ') + 'alloc retries which hurts performance. If encountered, we recommend ') + 'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.'))
152
+ elif (self.attn_impl == 'torch'):
153
+ self.attn_fn = scaled_multihead_dot_product_attention
154
+ if (torch.cuda.is_available() and verbose):
155
+ warnings.warn((('Using `attn_impl: torch`. If your model does not use `alibi` or ' + '`prefix_lm` we recommend using `attn_impl: flash` otherwise ') + 'we recommend using `attn_impl: triton`.'))
156
+ else:
157
+ raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
158
+ self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
159
+ self.out_proj._is_residual = True
160
+
161
+ def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
162
+ qkv = self.Wqkv(x)
163
+ if self.clip_qkv:
164
+ qkv.clamp_(min=(- self.clip_qkv), max=self.clip_qkv)
165
+ (query, key, value) = qkv.chunk(3, dim=2)
166
+ key_padding_mask = attention_mask
167
+ if self.qk_ln:
168
+ dtype = query.dtype
169
+ query = self.q_ln(query).to(dtype)
170
+ key = self.k_ln(key).to(dtype)
171
+ if (past_key_value is not None):
172
+ if (len(past_key_value) != 0):
173
+ key = torch.cat([past_key_value[0], key], dim=1)
174
+ value = torch.cat([past_key_value[1], value], dim=1)
175
+ past_key_value = (key, value)
176
+ if (attn_bias is not None):
177
+ attn_bias = attn_bias[:, :, (- query.size(1)):, (- key.size(1)):]
178
+ (context, attn_weights) = self.attn_fn(query, key, value, self.n_heads, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights)
179
+ return (self.out_proj(context), attn_weights, past_key_value)
180
+
181
+ class MultiQueryAttention(nn.Module):
182
+ 'Multi-Query self attention.\n\n Using torch or triton attention implemetation enables user to also use\n additive bias.\n '
183
+
184
+ def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, low_precision_layernorm: bool=False, verbose: int=0, device: Optional[str]=None):
185
+ super().__init__()
186
+ self.attn_impl = attn_impl
187
+ self.clip_qkv = clip_qkv
188
+ self.qk_ln = qk_ln
189
+ self.d_model = d_model
190
+ self.n_heads = n_heads
191
+ self.head_dim = (d_model // n_heads)
192
+ self.softmax_scale = softmax_scale
193
+ if (self.softmax_scale is None):
194
+ self.softmax_scale = (1 / math.sqrt(self.head_dim))
195
+ self.attn_dropout_p = attn_pdrop
196
+ self.Wqkv = nn.Linear(d_model, (d_model + (2 * self.head_dim)), device=device)
197
+ fuse_splits = (d_model, (d_model + self.head_dim))
198
+ self.Wqkv._fused = (0, fuse_splits)
199
+ if self.qk_ln:
200
+ layernorm_class = (LPLayerNorm if low_precision_layernorm else nn.LayerNorm)
201
+ self.q_ln = layernorm_class(d_model, device=device)
202
+ self.k_ln = layernorm_class(self.head_dim, device=device)
203
+ if (self.attn_impl == 'flash'):
204
+ self.attn_fn = flash_attn_fn
205
+ elif (self.attn_impl == 'triton'):
206
+ self.attn_fn = triton_flash_attn_fn
207
+ if verbose:
208
+ warnings.warn(((('While `attn_impl: triton` can be faster than `attn_impl: flash` ' + 'it uses more memory. When training larger models this can trigger ') + 'alloc retries which hurts performance. If encountered, we recommend ') + 'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.'))
209
+ elif (self.attn_impl == 'torch'):
210
+ self.attn_fn = scaled_multihead_dot_product_attention
211
+ if (torch.cuda.is_available() and verbose):
212
+ warnings.warn((('Using `attn_impl: torch`. If your model does not use `alibi` or ' + '`prefix_lm` we recommend using `attn_impl: flash` otherwise ') + 'we recommend using `attn_impl: triton`.'))
213
+ else:
214
+ raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
215
+ self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
216
+ self.out_proj._is_residual = True
217
+
218
+ def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
219
+ qkv = self.Wqkv(x)
220
+ if self.clip_qkv:
221
+ qkv.clamp_(min=(- self.clip_qkv), max=self.clip_qkv)
222
+ (query, key, value) = qkv.split([self.d_model, self.head_dim, self.head_dim], dim=2)
223
+ key_padding_mask = attention_mask
224
+ if self.qk_ln:
225
+ dtype = query.dtype
226
+ query = self.q_ln(query).to(dtype)
227
+ key = self.k_ln(key).to(dtype)
228
+ if (past_key_value is not None):
229
+ if (len(past_key_value) != 0):
230
+ key = torch.cat([past_key_value[0], key], dim=1)
231
+ value = torch.cat([past_key_value[1], value], dim=1)
232
+ past_key_value = (key, value)
233
+ if (attn_bias is not None):
234
+ attn_bias = attn_bias[:, :, (- query.size(1)):, (- key.size(1)):]
235
+ (context, attn_weights) = self.attn_fn(query, key, value, self.n_heads, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights, multiquery=True)
236
+ return (self.out_proj(context), attn_weights, past_key_value)
237
+
238
+ def attn_bias_shape(attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_sequence_id):
239
+ if (attn_impl == 'flash'):
240
+ return None
241
+ elif (attn_impl in ['torch', 'triton']):
242
+ if alibi:
243
+ if ((prefix_lm or (not causal)) or use_sequence_id):
244
+ return (1, n_heads, seq_len, seq_len)
245
+ return (1, n_heads, 1, seq_len)
246
+ elif (prefix_lm or use_sequence_id):
247
+ return (1, 1, seq_len, seq_len)
248
+ return None
249
+ else:
250
+ raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
251
+
252
+ def build_attn_bias(attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=False, alibi_bias_max=8):
253
+ if (attn_impl == 'flash'):
254
+ return None
255
+ elif (attn_impl in ['torch', 'triton']):
256
+ if alibi:
257
+ (device, dtype) = (attn_bias.device, attn_bias.dtype)
258
+ attn_bias = attn_bias.add(build_alibi_bias(n_heads, seq_len, full=(not causal), alibi_bias_max=alibi_bias_max, device=device, dtype=dtype))
259
+ return attn_bias
260
+ else:
261
+ raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
262
+
263
+ def gen_slopes(n_heads, alibi_bias_max=8, device=None):
264
+ _n_heads = (2 ** math.ceil(math.log2(n_heads)))
265
+ m = torch.arange(1, (_n_heads + 1), dtype=torch.float32, device=device)
266
+ m = m.mul((alibi_bias_max / _n_heads))
267
+ slopes = (1.0 / torch.pow(2, m))
268
+ if (_n_heads != n_heads):
269
+ slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
270
+ return slopes.view(1, n_heads, 1, 1)
271
+
272
+ def build_alibi_bias(n_heads, seq_len, full=False, alibi_bias_max=8, device=None, dtype=None):
273
+ alibi_bias = torch.arange((1 - seq_len), 1, dtype=torch.int32, device=device).view(1, 1, 1, seq_len)
274
+ if full:
275
+ alibi_bias = (alibi_bias - torch.arange((1 - seq_len), 1, dtype=torch.int32, device=device).view(1, 1, seq_len, 1))
276
+ alibi_bias = alibi_bias.abs().mul((- 1))
277
+ slopes = gen_slopes(n_heads, alibi_bias_max, device=device)
278
+ alibi_bias = (alibi_bias * slopes)
279
+ return alibi_bias.to(dtype=dtype)
280
+ ATTN_CLASS_REGISTRY = {'multihead_attention': MultiheadAttention, 'multiquery_attention': MultiQueryAttention}
blocks.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ 'GPT Blocks used for the GPT Model.'
3
+ from typing import Dict, Optional, Tuple
4
+ import torch
5
+ import torch.nn as nn
6
+ from .attention import ATTN_CLASS_REGISTRY
7
+ from .norm import NORM_CLASS_REGISTRY
8
+
9
+ class MPTMLP(nn.Module):
10
+
11
+ def __init__(self, d_model: int, expansion_ratio: int, device: Optional[str]=None):
12
+ super().__init__()
13
+ self.up_proj = nn.Linear(d_model, (expansion_ratio * d_model), device=device)
14
+ self.act = nn.GELU(approximate='none')
15
+ self.down_proj = nn.Linear((expansion_ratio * d_model), d_model, device=device)
16
+ self.down_proj._is_residual = True
17
+
18
+ def forward(self, x):
19
+ return self.down_proj(self.act(self.up_proj(x)))
20
+
21
+ class MPTBlock(nn.Module):
22
+
23
+ def __init__(self, d_model: int, n_heads: int, expansion_ratio: int, attn_config: Dict={'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', verbose: int=0, device: Optional[str]=None, **kwargs):
24
+ del kwargs
25
+ super().__init__()
26
+ norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
27
+ attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]
28
+ self.norm_1 = norm_class(d_model, device=device)
29
+ self.attn = attn_class(attn_impl=attn_config['attn_impl'], clip_qkv=attn_config['clip_qkv'], qk_ln=attn_config['qk_ln'], softmax_scale=attn_config['softmax_scale'], attn_pdrop=attn_config['attn_pdrop'], d_model=d_model, n_heads=n_heads, verbose=verbose, device=device)
30
+ self.norm_2 = norm_class(d_model, device=device)
31
+ self.ffn = MPTMLP(d_model=d_model, expansion_ratio=expansion_ratio, device=device)
32
+ self.resid_attn_dropout = nn.Dropout(resid_pdrop)
33
+ self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
34
+
35
+ def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.ByteTensor]=None, is_causal: bool=True) -> Tuple[(torch.Tensor, Optional[Tuple[torch.Tensor]])]:
36
+ a = self.norm_1(x)
37
+ (b, _, past_key_value) = self.attn(a, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=is_causal)
38
+ x = (x + self.resid_attn_dropout(b))
39
+ m = self.norm_2(x)
40
+ n = self.ffn(m)
41
+ x = (x + self.resid_ffn_dropout(n))
42
+ return (x, past_key_value)
config.json ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "MPTForCausalLM"
4
+ ],
5
+ "attn_config": {
6
+ "alibi": true,
7
+ "alibi_bias_max": 16,
8
+ "attn_impl": "torch",
9
+ "attn_pdrop": 0,
10
+ "attn_type": "multihead_attention",
11
+ "attn_uses_sequence_id": false,
12
+ "clip_qkv": 6,
13
+ "prefix_lm": false,
14
+ "qk_ln": false,
15
+ "softmax_scale": null
16
+ },
17
+ "auto_map": {
18
+ "AutoConfig": "configuration_mpt.MPTConfig",
19
+ "AutoModelForCausalLM": "modeling_mpt.MPTForCausalLM"
20
+ },
21
+ "d_model": 4096,
22
+ "emb_pdrop": 0,
23
+ "embedding_fraction": 1.0,
24
+ "expansion_ratio": 4,
25
+ "init_config": {
26
+ "emb_init_std": null,
27
+ "emb_init_uniform_lim": null,
28
+ "fan_mode": "fan_in",
29
+ "init_div_is_residual": true,
30
+ "init_gain": 0,
31
+ "init_nonlinearity": "relu",
32
+ "init_std": 0.02,
33
+ "name": "kaiming_normal_",
34
+ "verbose": 0
35
+ },
36
+ "init_device": "cpu",
37
+ "learned_pos_emb": true,
38
+ "logit_scale": null,
39
+ "max_seq_len": 65536,
40
+ "model_type": "mpt",
41
+ "n_heads": 32,
42
+ "n_layers": 32,
43
+ "no_bias": true,
44
+ "norm_type": "low_precision_layernorm",
45
+ "resid_pdrop": 0,
46
+ "tokenizer_name": "EleutherAI/gpt-neox-20b",
47
+ "torch_dtype": "bfloat16",
48
+ "transformers_version": "4.29.2",
49
+ "use_cache": false,
50
+ "verbose": 0,
51
+ "vocab_size": 50432
52
+ }
configuration_mpt.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ 'A HuggingFace-style model configuration.'
3
+ from typing import Dict, Optional, Union
4
+ from transformers import PretrainedConfig
5
+ attn_config_defaults: Dict = {'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}
6
+ init_config_defaults: Dict = {'name': 'kaiming_normal_', 'fan_mode': 'fan_in', 'init_nonlinearity': 'relu', 'init_div_is_residual': True, 'emb_init_std': None, 'emb_init_uniform_lim': None, 'init_std': None, 'init_gain': 0.0}
7
+
8
+ class MPTConfig(PretrainedConfig):
9
+ model_type = 'mpt'
10
+
11
+ def __init__(self, d_model: int=2048, n_heads: int=16, n_layers: int=24, expansion_ratio: int=4, max_seq_len: int=2048, vocab_size: int=50368, resid_pdrop: float=0.0, emb_pdrop: float=0.0, learned_pos_emb: bool=True, attn_config: Dict=attn_config_defaults, init_device: str='cpu', logit_scale: Optional[Union[(float, str)]]=None, no_bias: bool=False, verbose: int=0, embedding_fraction: float=1.0, norm_type: str='low_precision_layernorm', use_cache: bool=False, init_config: Dict=init_config_defaults, **kwargs):
12
+ "The MPT configuration class.\n\n Args:\n d_model (int): The size of the embedding dimension of the model.\n n_heads (int): The number of attention heads.\n n_layers (int): The number of layers in the model.\n expansion_ratio (int): The ratio of the up/down scale in the MLP.\n max_seq_len (int): The maximum sequence length of the model.\n vocab_size (int): The size of the vocabulary.\n resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.\n emb_pdrop (float): The dropout probability for the embedding layer.\n learned_pos_emb (bool): Whether to use learned positional embeddings\n attn_config (Dict): A dictionary used to configure the model's attention module:\n attn_type (str): type of attention to use. Options: multihead_attention, multiquery_attention\n attn_pdrop (float): The dropout probability for the attention layers.\n attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.\n qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.\n clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to\n this value.\n softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,\n use the default scale of ``1/sqrt(d_keys)``.\n prefix_lm (Optional[bool]): Whether the model should operate as a Prefix LM. This requires passing an\n extra `prefix_mask` argument which indicates which tokens belong to the prefix. Tokens in the prefix\n can attend to one another bi-directionally. Tokens outside the prefix use causal attention.\n attn_uses_sequence_id (Optional[bool]): Whether to restrict attention to tokens that have the same sequence_id.\n When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates\n which sub-sequence each token belongs to.\n Defaults to ``False`` meaning any provided `sequence_id` will be ignored.\n alibi (bool): Whether to use the alibi bias instead of position embeddings.\n alibi_bias_max (int): The maximum value of the alibi bias.\n init_device (str): The device to use for parameter initialization.\n logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.\n no_bias (bool): Whether to use bias in all layers.\n verbose (int): The verbosity level. 0 is silent.\n embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.\n norm_type (str): choose type of norm to use\n multiquery_attention (bool): Whether to use multiquery attention implementation.\n use_cache (bool): Whether or not the model should return the last key/values attentions\n init_config (Dict): A dictionary used to configure the model initialization:\n init_config.name: The parameter initialization scheme to use. Options: 'default_', 'baseline_',\n 'kaiming_uniform_', 'kaiming_normal_', 'neox_init_', 'small_init_', 'xavier_uniform_', or\n 'xavier_normal_'. These mimic the parameter initialization methods in PyTorch.\n init_div_is_residual (Union[int, float, str, bool]): Value to divide initial weights by if ``module._is_residual`` is True.\n emb_init_std (Optional[float]): The standard deviation of the normal distribution used to initialize the embedding layer.\n emb_init_uniform_lim (Optional[Union[Tuple[float, float], float]]): The lower and upper limits of the uniform distribution\n used to initialize the embedding layer. Mutually exclusive with ``emb_init_std``.\n init_std (float): The standard deviation of the normal distribution used to initialize the model,\n if using the baseline_ parameter initialization scheme.\n init_gain (float): The gain to use for parameter initialization with kaiming or xavier initialization schemes.\n fan_mode (str): The fan mode to use for parameter initialization with kaiming initialization schemes.\n init_nonlinearity (str): The nonlinearity to use for parameter initialization with kaiming initialization schemes.\n ---\n See llmfoundry.models.utils.param_init_fns.py for info on other param init config options\n "
13
+ self.d_model = d_model
14
+ self.n_heads = n_heads
15
+ self.n_layers = n_layers
16
+ self.expansion_ratio = expansion_ratio
17
+ self.max_seq_len = max_seq_len
18
+ self.vocab_size = vocab_size
19
+ self.resid_pdrop = resid_pdrop
20
+ self.emb_pdrop = emb_pdrop
21
+ self.learned_pos_emb = learned_pos_emb
22
+ self.attn_config = attn_config
23
+ self.init_device = init_device
24
+ self.logit_scale = logit_scale
25
+ self.no_bias = no_bias
26
+ self.verbose = verbose
27
+ self.embedding_fraction = embedding_fraction
28
+ self.norm_type = norm_type
29
+ self.use_cache = use_cache
30
+ self.init_config = init_config
31
+ if ('name' in kwargs):
32
+ del kwargs['name']
33
+ if ('loss_fn' in kwargs):
34
+ del kwargs['loss_fn']
35
+ super().__init__(**kwargs)
36
+ self._validate_config()
37
+
38
+ def _set_config_defaults(self, config, config_defaults):
39
+ for (k, v) in config_defaults.items():
40
+ if (k not in config):
41
+ config[k] = v
42
+ return config
43
+
44
+ def _validate_config(self):
45
+ self.attn_config = self._set_config_defaults(self.attn_config, attn_config_defaults)
46
+ self.init_config = self._set_config_defaults(self.init_config, init_config_defaults)
47
+ if ((self.d_model % self.n_heads) != 0):
48
+ raise ValueError('d_model must be divisible by n_heads')
49
+ if any((((prob < 0) or (prob > 1)) for prob in [self.attn_config['attn_pdrop'], self.resid_pdrop, self.emb_pdrop])):
50
+ raise ValueError("self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1")
51
+ if (self.attn_config['attn_impl'] not in ['torch', 'flash', 'triton']):
52
+ raise ValueError(f"Unknown attn_impl={self.attn_config['attn_impl']}")
53
+ if (self.attn_config['prefix_lm'] and (self.attn_config['attn_impl'] not in ['torch', 'triton'])):
54
+ raise NotImplementedError('prefix_lm only implemented with torch and triton attention.')
55
+ if (self.attn_config['alibi'] and (self.attn_config['attn_impl'] not in ['torch', 'triton'])):
56
+ raise NotImplementedError('alibi only implemented with torch and triton attention.')
57
+ if (self.attn_config['attn_uses_sequence_id'] and (self.attn_config['attn_impl'] not in ['torch', 'triton'])):
58
+ raise NotImplementedError('attn_uses_sequence_id only implemented with torch and triton attention.')
59
+ if ((self.embedding_fraction > 1) or (self.embedding_fraction <= 0)):
60
+ raise ValueError('model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!')
61
+ if (isinstance(self.logit_scale, str) and (self.logit_scale != 'inv_sqrt_d_model')):
62
+ raise ValueError(f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
63
+ if (self.init_config.get('name', None) is None):
64
+ raise ValueError(f"self.init_config={self.init_config!r} 'name' needs to be set.")
65
+ if ((not self.learned_pos_emb) and (not self.attn_config['alibi'])):
66
+ raise ValueError(f'Positional information must be provided to the model using either learned_pos_emb or alibi.')
flash_attn_triton.py ADDED
@@ -0,0 +1,426 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ '\nCopied from https://github.com/HazyResearch/flash-attention/blob/eff9fe6b8076df59d64d7a3f464696738a3c7c24/flash_attn/flash_attn_triton.py\nupdate imports to use \'triton_pre_mlir\'\n\n*Experimental* implementation of FlashAttention in Triton.\nTested with triton==2.0.0.dev20221202.\nTriton 2.0 has a new backend (MLIR) but seems like it doesn\'t yet work for head dimensions\nother than 64:\nhttps://github.com/openai/triton/blob/d376020f90002757eea3ea9475d4f7cfc2ec5ead/python/triton/ops/flash_attention.py#L207\nWe\'ll update this implementation with the new Triton backend once this is fixed.\n\nWe use the FlashAttention implementation from Phil Tillet a starting point.\nhttps://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py\n\nChanges:\n- Implement both causal and non-causal attention.\n- Implement both self-attention and cross-attention.\n- Support arbitrary seqlens (not just multiples of 128), for both forward and backward.\n- Support all head dimensions up to 128 (not just 16, 32, 64, 128), for both forward and backward.\n- Support attention bias.\n- Speed up the forward pass a bit, and only store the LSE instead of m and l.\n- Make the backward for d=128 much faster by reducing register spilling.\n- Optionally parallelize the backward pass across seqlen_k, to deal with the case of\nsmall batch size * nheads.\n\nCaution:\n- This is an *experimental* implementation. The forward pass should be quite robust but\nI\'m not 100% sure that the backward pass doesn\'t have race conditions (due to the Triton compiler).\n- This implementation has only been tested on A100.\n- If you plan to use headdim other than 64 and 128, you should test for race conditions\n(due to the Triton compiler), as done in tests/test_flash_attn.py\n"test_flash_attn_triton_race_condition". I\'ve tested and fixed many race conditions\nfor different head dimensions (40, 48, 64, 128, 80, 88, 96), but I\'m still not 100% confident\nthat there are none left for other head dimensions.\n\nDifferences between this Triton version and the CUDA version:\n- Triton version doesn\'t support dropout.\n- Triton forward is generally faster than CUDA forward, while Triton backward is\ngenerally slower than CUDA backward. Overall Triton forward + backward is slightly slower\nthan CUDA forward + backward.\n- Triton version doesn\'t support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor).\n- Triton version supports attention bias, while CUDA version doesn\'t.\n'
3
+ import math
4
+ import torch
5
+ import triton_pre_mlir as triton
6
+ import triton_pre_mlir.language as tl
7
+
8
+ @triton.heuristics({'EVEN_M': (lambda args: ((args['seqlen_q'] % args['BLOCK_M']) == 0)), 'EVEN_N': (lambda args: ((args['seqlen_k'] % args['BLOCK_N']) == 0)), 'EVEN_HEADDIM': (lambda args: (args['headdim'] == args['BLOCK_HEADDIM']))})
9
+ @triton.jit
10
+ def _fwd_kernel(Q, K, V, Bias, Out, Lse, TMP, softmax_scale, stride_qb, stride_qh, stride_qm, stride_kb, stride_kh, stride_kn, stride_vb, stride_vh, stride_vn, stride_bb, stride_bh, stride_bm, stride_ob, stride_oh, stride_om, nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim, CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
11
+ start_m = tl.program_id(0)
12
+ off_hb = tl.program_id(1)
13
+ off_b = (off_hb // nheads)
14
+ off_h = (off_hb % nheads)
15
+ offs_m = ((start_m * BLOCK_M) + tl.arange(0, BLOCK_M))
16
+ offs_n = tl.arange(0, BLOCK_N)
17
+ offs_d = tl.arange(0, BLOCK_HEADDIM)
18
+ q_ptrs = (((Q + (off_b * stride_qb)) + (off_h * stride_qh)) + ((offs_m[:, None] * stride_qm) + offs_d[None, :]))
19
+ k_ptrs = (((K + (off_b * stride_kb)) + (off_h * stride_kh)) + ((offs_n[:, None] * stride_kn) + offs_d[None, :]))
20
+ v_ptrs = (((V + (off_b * stride_vb)) + (off_h * stride_vh)) + ((offs_n[:, None] * stride_vn) + offs_d[None, :]))
21
+ if (BIAS_TYPE == 'vector'):
22
+ b_ptrs = (((Bias + (off_b * stride_bb)) + (off_h * stride_bh)) + offs_n)
23
+ elif (BIAS_TYPE == 'matrix'):
24
+ b_ptrs = (((Bias + (off_b * stride_bb)) + (off_h * stride_bh)) + ((offs_m[:, None] * stride_bm) + offs_n[None, :]))
25
+ t_ptrs = ((TMP + (off_hb * seqlen_q_rounded)) + offs_m)
26
+ lse_i = (tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf'))
27
+ m_i = (tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf'))
28
+ acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
29
+ if (EVEN_M & EVEN_N):
30
+ if EVEN_HEADDIM:
31
+ q = tl.load(q_ptrs)
32
+ else:
33
+ q = tl.load(q_ptrs, mask=(offs_d[None, :] < headdim), other=0.0)
34
+ elif EVEN_HEADDIM:
35
+ q = tl.load(q_ptrs, mask=(offs_m[:, None] < seqlen_q), other=0.0)
36
+ else:
37
+ q = tl.load(q_ptrs, mask=((offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim)), other=0.0)
38
+ end_n = (seqlen_k if (not IS_CAUSAL) else tl.minimum(((start_m + 1) * BLOCK_M), seqlen_k))
39
+ for start_n in range(0, end_n, BLOCK_N):
40
+ start_n = tl.multiple_of(start_n, BLOCK_N)
41
+ if (EVEN_N & EVEN_M):
42
+ if EVEN_HEADDIM:
43
+ k = tl.load((k_ptrs + (start_n * stride_kn)))
44
+ else:
45
+ k = tl.load((k_ptrs + (start_n * stride_kn)), mask=(offs_d[None, :] < headdim), other=0.0)
46
+ elif EVEN_HEADDIM:
47
+ k = tl.load((k_ptrs + (start_n * stride_kn)), mask=((start_n + offs_n)[:, None] < seqlen_k), other=0.0)
48
+ else:
49
+ k = tl.load((k_ptrs + (start_n * stride_kn)), mask=(((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim)), other=0.0)
50
+ qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
51
+ qk += tl.dot(q, k, trans_b=True)
52
+ if (not EVEN_N):
53
+ qk += tl.where(((start_n + offs_n)[None, :] < seqlen_k), 0, float('-inf'))
54
+ if IS_CAUSAL:
55
+ qk += tl.where((offs_m[:, None] >= (start_n + offs_n)[None, :]), 0, float('-inf'))
56
+ if (BIAS_TYPE != 'none'):
57
+ if (BIAS_TYPE == 'vector'):
58
+ if EVEN_N:
59
+ bias = tl.load((b_ptrs + start_n)).to(tl.float32)
60
+ else:
61
+ bias = tl.load((b_ptrs + start_n), mask=((start_n + offs_n) < seqlen_k), other=0.0).to(tl.float32)
62
+ bias = bias[None, :]
63
+ elif (BIAS_TYPE == 'matrix'):
64
+ if (EVEN_M & EVEN_N):
65
+ bias = tl.load((b_ptrs + start_n)).to(tl.float32)
66
+ else:
67
+ bias = tl.load((b_ptrs + start_n), mask=((offs_m[:, None] < seqlen_q) & ((start_n + offs_n)[None, :] < seqlen_k)), other=0.0).to(tl.float32)
68
+ qk = ((qk * softmax_scale) + bias)
69
+ m_ij = tl.maximum(tl.max(qk, 1), lse_i)
70
+ p = tl.exp((qk - m_ij[:, None]))
71
+ else:
72
+ m_ij = tl.maximum((tl.max(qk, 1) * softmax_scale), lse_i)
73
+ p = tl.exp(((qk * softmax_scale) - m_ij[:, None]))
74
+ l_ij = tl.sum(p, 1)
75
+ acc_o_scale = tl.exp((m_i - m_ij))
76
+ tl.store(t_ptrs, acc_o_scale)
77
+ acc_o_scale = tl.load(t_ptrs)
78
+ acc_o = (acc_o * acc_o_scale[:, None])
79
+ if (EVEN_N & EVEN_M):
80
+ if EVEN_HEADDIM:
81
+ v = tl.load((v_ptrs + (start_n * stride_vn)))
82
+ else:
83
+ v = tl.load((v_ptrs + (start_n * stride_vn)), mask=(offs_d[None, :] < headdim), other=0.0)
84
+ elif EVEN_HEADDIM:
85
+ v = tl.load((v_ptrs + (start_n * stride_vn)), mask=((start_n + offs_n)[:, None] < seqlen_k), other=0.0)
86
+ else:
87
+ v = tl.load((v_ptrs + (start_n * stride_vn)), mask=(((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim)), other=0.0)
88
+ p = p.to(v.dtype)
89
+ acc_o += tl.dot(p, v)
90
+ m_i = m_ij
91
+ l_i_new = (tl.exp((lse_i - m_ij)) + l_ij)
92
+ lse_i = (m_ij + tl.log(l_i_new))
93
+ o_scale = tl.exp((m_i - lse_i))
94
+ tl.store(t_ptrs, o_scale)
95
+ o_scale = tl.load(t_ptrs)
96
+ acc_o = (acc_o * o_scale[:, None])
97
+ start_m = tl.program_id(0)
98
+ offs_m = ((start_m * BLOCK_M) + tl.arange(0, BLOCK_M))
99
+ lse_ptrs = ((Lse + (off_hb * seqlen_q_rounded)) + offs_m)
100
+ tl.store(lse_ptrs, lse_i)
101
+ offs_d = tl.arange(0, BLOCK_HEADDIM)
102
+ out_ptrs = (((Out + (off_b * stride_ob)) + (off_h * stride_oh)) + ((offs_m[:, None] * stride_om) + offs_d[None, :]))
103
+ if EVEN_M:
104
+ if EVEN_HEADDIM:
105
+ tl.store(out_ptrs, acc_o)
106
+ else:
107
+ tl.store(out_ptrs, acc_o, mask=(offs_d[None, :] < headdim))
108
+ elif EVEN_HEADDIM:
109
+ tl.store(out_ptrs, acc_o, mask=(offs_m[:, None] < seqlen_q))
110
+ else:
111
+ tl.store(out_ptrs, acc_o, mask=((offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim)))
112
+
113
+ @triton.jit
114
+ def _bwd_preprocess_do_o_dot(Out, DO, Delta, stride_ob, stride_oh, stride_om, stride_dob, stride_doh, stride_dom, nheads, seqlen_q, seqlen_q_rounded, headdim, BLOCK_M: tl.constexpr, BLOCK_HEADDIM: tl.constexpr):
115
+ start_m = tl.program_id(0)
116
+ off_hb = tl.program_id(1)
117
+ off_b = (off_hb // nheads)
118
+ off_h = (off_hb % nheads)
119
+ offs_m = ((start_m * BLOCK_M) + tl.arange(0, BLOCK_M))
120
+ offs_d = tl.arange(0, BLOCK_HEADDIM)
121
+ o = tl.load(((((Out + (off_b * stride_ob)) + (off_h * stride_oh)) + (offs_m[:, None] * stride_om)) + offs_d[None, :]), mask=((offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim)), other=0.0).to(tl.float32)
122
+ do = tl.load(((((DO + (off_b * stride_dob)) + (off_h * stride_doh)) + (offs_m[:, None] * stride_dom)) + offs_d[None, :]), mask=((offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim)), other=0.0).to(tl.float32)
123
+ delta = tl.sum((o * do), axis=1)
124
+ tl.store(((Delta + (off_hb * seqlen_q_rounded)) + offs_m), delta)
125
+
126
+ @triton.jit
127
+ def _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr):
128
+ if (EVEN_N & EVEN_M):
129
+ if EVEN_HEADDIM:
130
+ tl.store(dv_ptrs, dv)
131
+ tl.store(dk_ptrs, dk)
132
+ else:
133
+ tl.store(dv_ptrs, dv, mask=(offs_d[None, :] < headdim))
134
+ tl.store(dk_ptrs, dk, mask=(offs_d[None, :] < headdim))
135
+ elif EVEN_HEADDIM:
136
+ tl.store(dv_ptrs, dv, mask=(offs_n[:, None] < seqlen_k))
137
+ tl.store(dk_ptrs, dk, mask=(offs_n[:, None] < seqlen_k))
138
+ else:
139
+ tl.store(dv_ptrs, dv, mask=((offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim)))
140
+ tl.store(dk_ptrs, dk, mask=((offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim)))
141
+
142
+ @triton.jit
143
+ def _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD: tl.constexpr, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
144
+ begin_m = (0 if (not IS_CAUSAL) else (((start_n * BLOCK_N) // BLOCK_M) * BLOCK_M))
145
+ offs_qm = (begin_m + tl.arange(0, BLOCK_M))
146
+ offs_n = ((start_n * BLOCK_N) + tl.arange(0, BLOCK_N))
147
+ offs_m = tl.arange(0, BLOCK_M)
148
+ offs_d = tl.arange(0, BLOCK_HEADDIM)
149
+ q_ptrs = (Q + ((offs_qm[:, None] * stride_qm) + offs_d[None, :]))
150
+ k_ptrs = (K + ((offs_n[:, None] * stride_kn) + offs_d[None, :]))
151
+ v_ptrs = (V + ((offs_n[:, None] * stride_vn) + offs_d[None, :]))
152
+ do_ptrs = (DO + ((offs_qm[:, None] * stride_dom) + offs_d[None, :]))
153
+ dq_ptrs = (DQ + ((offs_qm[:, None] * stride_dqm) + offs_d[None, :]))
154
+ if (BIAS_TYPE == 'vector'):
155
+ b_ptrs = (Bias + offs_n)
156
+ elif (BIAS_TYPE == 'matrix'):
157
+ b_ptrs = (Bias + ((offs_qm[:, None] * stride_bm) + offs_n[None, :]))
158
+ dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
159
+ dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
160
+ if (begin_m >= seqlen_q):
161
+ dv_ptrs = (DV + ((offs_n[:, None] * stride_dvn) + offs_d[None, :]))
162
+ dk_ptrs = (DK + ((offs_n[:, None] * stride_dkn) + offs_d[None, :]))
163
+ _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM)
164
+ return
165
+ if (EVEN_N & EVEN_M):
166
+ if EVEN_HEADDIM:
167
+ k = tl.load(k_ptrs)
168
+ v = tl.load(v_ptrs)
169
+ else:
170
+ k = tl.load(k_ptrs, mask=(offs_d[None, :] < headdim), other=0.0)
171
+ v = tl.load(v_ptrs, mask=(offs_d[None, :] < headdim), other=0.0)
172
+ elif EVEN_HEADDIM:
173
+ k = tl.load(k_ptrs, mask=(offs_n[:, None] < seqlen_k), other=0.0)
174
+ v = tl.load(v_ptrs, mask=(offs_n[:, None] < seqlen_k), other=0.0)
175
+ else:
176
+ k = tl.load(k_ptrs, mask=((offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim)), other=0.0)
177
+ v = tl.load(v_ptrs, mask=((offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim)), other=0.0)
178
+ num_block_m = tl.cdiv(seqlen_q, BLOCK_M)
179
+ for start_m in range(begin_m, (num_block_m * BLOCK_M), BLOCK_M):
180
+ start_m = tl.multiple_of(start_m, BLOCK_M)
181
+ offs_m_curr = (start_m + offs_m)
182
+ if (EVEN_M & EVEN_HEADDIM):
183
+ q = tl.load(q_ptrs)
184
+ elif EVEN_HEADDIM:
185
+ q = tl.load(q_ptrs, mask=(offs_m_curr[:, None] < seqlen_q), other=0.0)
186
+ else:
187
+ q = tl.load(q_ptrs, mask=((offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim)), other=0.0)
188
+ qk = tl.dot(q, k, trans_b=True)
189
+ if (not EVEN_N):
190
+ qk = tl.where((offs_n[None, :] < seqlen_k), qk, float('-inf'))
191
+ if IS_CAUSAL:
192
+ qk = tl.where((offs_m_curr[:, None] >= offs_n[None, :]), qk, float('-inf'))
193
+ if (BIAS_TYPE != 'none'):
194
+ tl.debug_barrier()
195
+ if (BIAS_TYPE == 'vector'):
196
+ if EVEN_N:
197
+ bias = tl.load(b_ptrs).to(tl.float32)
198
+ else:
199
+ bias = tl.load(b_ptrs, mask=(offs_n < seqlen_k), other=0.0).to(tl.float32)
200
+ bias = bias[None, :]
201
+ elif (BIAS_TYPE == 'matrix'):
202
+ if (EVEN_M & EVEN_N):
203
+ bias = tl.load(b_ptrs).to(tl.float32)
204
+ else:
205
+ bias = tl.load(b_ptrs, mask=((offs_m_curr[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k)), other=0.0).to(tl.float32)
206
+ qk = ((qk * softmax_scale) + bias)
207
+ if (not (EVEN_M & EVEN_HEADDIM)):
208
+ tl.debug_barrier()
209
+ lse_i = tl.load((LSE + offs_m_curr))
210
+ if (BIAS_TYPE == 'none'):
211
+ p = tl.exp(((qk * softmax_scale) - lse_i[:, None]))
212
+ else:
213
+ p = tl.exp((qk - lse_i[:, None]))
214
+ if (EVEN_M & EVEN_HEADDIM):
215
+ do = tl.load(do_ptrs)
216
+ else:
217
+ do = tl.load(do_ptrs, mask=((offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim)), other=0.0)
218
+ dv += tl.dot(p.to(do.dtype), do, trans_a=True)
219
+ if (not (EVEN_M & EVEN_HEADDIM)):
220
+ tl.debug_barrier()
221
+ dp = tl.dot(do, v, trans_b=True)
222
+ if (not EVEN_HEADDIM):
223
+ tl.debug_barrier()
224
+ Di = tl.load((D + offs_m_curr))
225
+ ds = ((p * (dp - Di[:, None])) * softmax_scale).to(q.dtype)
226
+ dk += tl.dot(ds, q, trans_a=True)
227
+ if (not (EVEN_M & EVEN_HEADDIM)):
228
+ tl.debug_barrier()
229
+ if (not ATOMIC_ADD):
230
+ if (EVEN_M & EVEN_HEADDIM):
231
+ dq = tl.load(dq_ptrs, eviction_policy='evict_last')
232
+ dq += tl.dot(ds, k)
233
+ tl.store(dq_ptrs, dq, eviction_policy='evict_last')
234
+ elif EVEN_HEADDIM:
235
+ dq = tl.load(dq_ptrs, mask=(offs_m_curr[:, None] < seqlen_q), other=0.0, eviction_policy='evict_last')
236
+ dq += tl.dot(ds, k)
237
+ tl.store(dq_ptrs, dq, mask=(offs_m_curr[:, None] < seqlen_q), eviction_policy='evict_last')
238
+ else:
239
+ dq = tl.load(dq_ptrs, mask=((offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim)), other=0.0, eviction_policy='evict_last')
240
+ dq += tl.dot(ds, k)
241
+ tl.store(dq_ptrs, dq, mask=((offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim)), eviction_policy='evict_last')
242
+ else:
243
+ dq = tl.dot(ds, k)
244
+ if (EVEN_M & EVEN_HEADDIM):
245
+ tl.atomic_add(dq_ptrs, dq)
246
+ elif EVEN_HEADDIM:
247
+ tl.atomic_add(dq_ptrs, dq, mask=(offs_m_curr[:, None] < seqlen_q))
248
+ else:
249
+ tl.atomic_add(dq_ptrs, dq, mask=((offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim)))
250
+ dq_ptrs += (BLOCK_M * stride_dqm)
251
+ q_ptrs += (BLOCK_M * stride_qm)
252
+ do_ptrs += (BLOCK_M * stride_dom)
253
+ if (BIAS_TYPE == 'matrix'):
254
+ b_ptrs += (BLOCK_M * stride_bm)
255
+ dv_ptrs = (DV + ((offs_n[:, None] * stride_dvn) + offs_d[None, :]))
256
+ dk_ptrs = (DK + ((offs_n[:, None] * stride_dkn) + offs_d[None, :]))
257
+ _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM)
258
+
259
+ def init_to_zero(name):
260
+ return (lambda nargs: nargs[name].zero_())
261
+
262
+ @triton.autotune(configs=[triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'SEQUENCE_PARALLEL': False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'SEQUENCE_PARALLEL': True}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ'))], key=['CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL', 'BLOCK_HEADDIM'])
263
+ @triton.heuristics({'EVEN_M': (lambda args: ((args['seqlen_q'] % args['BLOCK_M']) == 0)), 'EVEN_N': (lambda args: ((args['seqlen_k'] % args['BLOCK_N']) == 0)), 'EVEN_HEADDIM': (lambda args: (args['headdim'] == args['BLOCK_HEADDIM']))})
264
+ @triton.jit
265
+ def _bwd_kernel(Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qb, stride_qh, stride_qm, stride_kb, stride_kh, stride_kn, stride_vb, stride_vh, stride_vn, stride_bb, stride_bh, stride_bm, stride_dob, stride_doh, stride_dom, stride_dqb, stride_dqh, stride_dqm, stride_dkb, stride_dkh, stride_dkn, stride_dvb, stride_dvh, stride_dvn, nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim, CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, SEQUENCE_PARALLEL: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
266
+ off_hb = tl.program_id(1)
267
+ off_b = (off_hb // nheads)
268
+ off_h = (off_hb % nheads)
269
+ Q += ((off_b * stride_qb) + (off_h * stride_qh))
270
+ K += ((off_b * stride_kb) + (off_h * stride_kh))
271
+ V += ((off_b * stride_vb) + (off_h * stride_vh))
272
+ DO += ((off_b * stride_dob) + (off_h * stride_doh))
273
+ DQ += ((off_b * stride_dqb) + (off_h * stride_dqh))
274
+ DK += ((off_b * stride_dkb) + (off_h * stride_dkh))
275
+ DV += ((off_b * stride_dvb) + (off_h * stride_dvh))
276
+ if (BIAS_TYPE != 'none'):
277
+ Bias += ((off_b * stride_bb) + (off_h * stride_bh))
278
+ D += (off_hb * seqlen_q_rounded)
279
+ LSE += (off_hb * seqlen_q_rounded)
280
+ if (not SEQUENCE_PARALLEL):
281
+ num_block_n = tl.cdiv(seqlen_k, BLOCK_N)
282
+ for start_n in range(0, num_block_n):
283
+ _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD=False, BIAS_TYPE=BIAS_TYPE, IS_CAUSAL=IS_CAUSAL, BLOCK_HEADDIM=BLOCK_HEADDIM, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N)
284
+ else:
285
+ start_n = tl.program_id(0)
286
+ _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD=True, BIAS_TYPE=BIAS_TYPE, IS_CAUSAL=IS_CAUSAL, BLOCK_HEADDIM=BLOCK_HEADDIM, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N)
287
+
288
+ def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):
289
+ (batch, seqlen_q, nheads, d) = q.shape
290
+ (_, seqlen_k, _, _) = k.shape
291
+ assert (k.shape == (batch, seqlen_k, nheads, d))
292
+ assert (v.shape == (batch, seqlen_k, nheads, d))
293
+ assert (d <= 128), 'FlashAttention only support head dimensions up to 128'
294
+ assert (q.dtype == k.dtype == v.dtype), 'All tensors must have the same type'
295
+ assert (q.dtype in [torch.float16, torch.bfloat16]), 'Only support fp16 and bf16'
296
+ assert (q.is_cuda and k.is_cuda and v.is_cuda)
297
+ softmax_scale = (softmax_scale or (1.0 / math.sqrt(d)))
298
+ has_bias = (bias is not None)
299
+ bias_type = 'none'
300
+ if has_bias:
301
+ assert (bias.dtype in [q.dtype, torch.float])
302
+ assert bias.is_cuda
303
+ assert (bias.dim() == 4)
304
+ if (bias.stride((- 1)) != 1):
305
+ bias = bias.contiguous()
306
+ if (bias.shape[2:] == (1, seqlen_k)):
307
+ bias_type = 'vector'
308
+ elif (bias.shape[2:] == (seqlen_q, seqlen_k)):
309
+ bias_type = 'matrix'
310
+ else:
311
+ raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k) or (seqlen_q, seqlen_k)')
312
+ bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)
313
+ bias_strides = ((bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0))
314
+ seqlen_q_rounded = (math.ceil((seqlen_q / 128)) * 128)
315
+ lse = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32)
316
+ tmp = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32)
317
+ o = torch.empty_like(q)
318
+ BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
319
+ BLOCK = 128
320
+ num_warps = (4 if (d <= 64) else 8)
321
+ grid = (lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), (batch * nheads)))
322
+ _fwd_kernel[grid](q, k, v, bias, o, lse, tmp, softmax_scale, q.stride(0), q.stride(2), q.stride(1), k.stride(0), k.stride(2), k.stride(1), v.stride(0), v.stride(2), v.stride(1), *bias_strides, o.stride(0), o.stride(2), o.stride(1), nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d, (seqlen_q // 32), (seqlen_k // 32), bias_type, causal, BLOCK_HEADDIM, BLOCK_M=BLOCK, BLOCK_N=BLOCK, num_warps=num_warps, num_stages=1)
323
+ return (o, lse, softmax_scale)
324
+
325
+ def _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv, bias=None, causal=False, softmax_scale=None):
326
+ if (do.stride((- 1)) != 1):
327
+ do = do.contiguous()
328
+ (batch, seqlen_q, nheads, d) = q.shape
329
+ (_, seqlen_k, _, _) = k.shape
330
+ assert (d <= 128)
331
+ seqlen_q_rounded = (math.ceil((seqlen_q / 128)) * 128)
332
+ assert (lse.shape == (batch, nheads, seqlen_q_rounded))
333
+ assert (q.stride((- 1)) == k.stride((- 1)) == v.stride((- 1)) == o.stride((- 1)) == 1)
334
+ assert (dq.stride((- 1)) == dk.stride((- 1)) == dv.stride((- 1)) == 1)
335
+ softmax_scale = (softmax_scale or (1.0 / math.sqrt(d)))
336
+ dq_accum = torch.empty_like(q, dtype=torch.float32)
337
+ delta = torch.empty_like(lse)
338
+ BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
339
+ grid = (lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), (batch * nheads)))
340
+ _bwd_preprocess_do_o_dot[grid](o, do, delta, o.stride(0), o.stride(2), o.stride(1), do.stride(0), do.stride(2), do.stride(1), nheads, seqlen_q, seqlen_q_rounded, d, BLOCK_M=128, BLOCK_HEADDIM=BLOCK_HEADDIM)
341
+ has_bias = (bias is not None)
342
+ bias_type = 'none'
343
+ if has_bias:
344
+ assert (bias.dtype in [q.dtype, torch.float])
345
+ assert bias.is_cuda
346
+ assert (bias.dim() == 4)
347
+ assert (bias.stride((- 1)) == 1)
348
+ if (bias.shape[2:] == (1, seqlen_k)):
349
+ bias_type = 'vector'
350
+ elif (bias.shape[2:] == (seqlen_q, seqlen_k)):
351
+ bias_type = 'matrix'
352
+ else:
353
+ raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k) or (seqlen_q, seqlen_k)')
354
+ bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)
355
+ bias_strides = ((bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0))
356
+ grid = (lambda META: ((triton.cdiv(seqlen_k, META['BLOCK_N']) if META['SEQUENCE_PARALLEL'] else 1), (batch * nheads)))
357
+ _bwd_kernel[grid](q, k, v, bias, do, dq_accum, dk, dv, lse, delta, softmax_scale, q.stride(0), q.stride(2), q.stride(1), k.stride(0), k.stride(2), k.stride(1), v.stride(0), v.stride(2), v.stride(1), *bias_strides, do.stride(0), do.stride(2), do.stride(1), dq_accum.stride(0), dq_accum.stride(2), dq_accum.stride(1), dk.stride(0), dk.stride(2), dk.stride(1), dv.stride(0), dv.stride(2), dv.stride(1), nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d, (seqlen_q // 32), (seqlen_k // 32), bias_type, causal, BLOCK_HEADDIM)
358
+ dq.copy_(dq_accum)
359
+
360
+ class FlashAttnQKVPackedFunc(torch.autograd.Function):
361
+
362
+ @staticmethod
363
+ def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):
364
+ '\n qkv: (batch, seqlen, 3, nheads, headdim)\n bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).\n For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).\n ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)\n '
365
+ if (qkv.stride((- 1)) != 1):
366
+ qkv = qkv.contiguous()
367
+ (o, lse, ctx.softmax_scale) = _flash_attn_forward(qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], bias=bias, causal=causal, softmax_scale=softmax_scale)
368
+ ctx.save_for_backward(qkv, o, lse, bias)
369
+ ctx.causal = causal
370
+ return o
371
+
372
+ @staticmethod
373
+ def backward(ctx, do):
374
+ (qkv, o, lse, bias) = ctx.saved_tensors
375
+ assert (not ctx.needs_input_grad[1]), 'FlashAttention does not support bias gradient yet'
376
+ with torch.inference_mode():
377
+ dqkv = torch.empty_like(qkv)
378
+ _flash_attn_backward(do, qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], o, lse, dqkv[:, :, 0], dqkv[:, :, 1], dqkv[:, :, 2], bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
379
+ return (dqkv, None, None, None)
380
+ flash_attn_qkvpacked_func = FlashAttnQKVPackedFunc.apply
381
+
382
+ class FlashAttnKVPackedFunc(torch.autograd.Function):
383
+
384
+ @staticmethod
385
+ def forward(ctx, q, kv, bias=None, causal=False, softmax_scale=None):
386
+ '\n q: (batch, seqlen_q, nheads, headdim)\n kv: (batch, seqlen_k, 2, nheads, headdim)\n bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).\n For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).\n ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)\n '
387
+ (q, kv) = [(x if (x.stride((- 1)) == 1) else x.contiguous()) for x in [q, kv]]
388
+ (o, lse, ctx.softmax_scale) = _flash_attn_forward(q, kv[:, :, 0], kv[:, :, 1], bias=bias, causal=causal, softmax_scale=softmax_scale)
389
+ ctx.save_for_backward(q, kv, o, lse, bias)
390
+ ctx.causal = causal
391
+ return o
392
+
393
+ @staticmethod
394
+ def backward(ctx, do):
395
+ (q, kv, o, lse, bias) = ctx.saved_tensors
396
+ if (len(ctx.needs_input_grad) >= 3):
397
+ assert (not ctx.needs_input_grad[2]), 'FlashAttention does not support bias gradient yet'
398
+ with torch.inference_mode():
399
+ dq = torch.empty_like(q)
400
+ dkv = torch.empty_like(kv)
401
+ _flash_attn_backward(do, q, kv[:, :, 0], kv[:, :, 1], o, lse, dq, dkv[:, :, 0], dkv[:, :, 1], bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
402
+ return (dq, dkv, None, None, None)
403
+ flash_attn_kvpacked_func = FlashAttnKVPackedFunc.apply
404
+
405
+ class FlashAttnFunc(torch.autograd.Function):
406
+
407
+ @staticmethod
408
+ def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None):
409
+ '\n q: (batch_size, seqlen_q, nheads, headdim)\n k, v: (batch_size, seqlen_k, nheads, headdim)\n bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).\n For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).\n ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)\n '
410
+ (q, k, v) = [(x if (x.stride((- 1)) == 1) else x.contiguous()) for x in [q, k, v]]
411
+ (o, lse, ctx.softmax_scale) = _flash_attn_forward(q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale)
412
+ ctx.save_for_backward(q, k, v, o, lse, bias)
413
+ ctx.causal = causal
414
+ return o
415
+
416
+ @staticmethod
417
+ def backward(ctx, do):
418
+ (q, k, v, o, lse, bias) = ctx.saved_tensors
419
+ assert (not ctx.needs_input_grad[3]), 'FlashAttention does not support bias gradient yet'
420
+ with torch.inference_mode():
421
+ dq = torch.empty_like(q)
422
+ dk = torch.empty_like(k)
423
+ dv = torch.empty_like(v)
424
+ _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv, bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
425
+ return (dq, dk, dv, None, None, None)
426
+ flash_attn_func = FlashAttnFunc.apply
generation_config.json ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "transformers_version": "4.29.2",
4
+ "use_cache": false
5
+ }
hf_prefixlm_converter.py ADDED
@@ -0,0 +1,325 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ 'Converts Huggingface Causal LM to Prefix LM.\n\nConversion does lightweight surgery on a HuggingFace\nCausal LM to convert it to a Prefix LM.\n\nPrefix LMs accepts a `bidirectional_mask` input in `forward`\nand treat the input prompt as the prefix in `generate`.\n'
3
+ import math
4
+ import warnings
5
+ from types import MethodType
6
+ from typing import Any, Dict, List, Optional, Tuple, Union
7
+ import torch
8
+ from transformers.models.bloom.modeling_bloom import BaseModelOutputWithPastAndCrossAttentions, BloomForCausalLM, BloomModel, CausalLMOutputWithCrossAttentions, CrossEntropyLoss
9
+ from transformers.models.bloom.modeling_bloom import _expand_mask as _expand_mask_bloom
10
+ from transformers.models.bloom.modeling_bloom import _make_causal_mask as _make_causal_mask_bloom
11
+ from transformers.models.bloom.modeling_bloom import logging
12
+ from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
13
+ from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoForCausalLM
14
+ from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
15
+ from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
16
+ from transformers.models.opt.modeling_opt import OPTForCausalLM
17
+ from transformers.models.opt.modeling_opt import _expand_mask as _expand_mask_opt
18
+ from transformers.models.opt.modeling_opt import _make_causal_mask as _make_causal_mask_opt
19
+ logger = logging.get_logger(__name__)
20
+ _SUPPORTED_GPT_MODELS = (GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM)
21
+ CAUSAL_GPT_TYPES = Union[(GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM)]
22
+
23
+ def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_TYPES:
24
+ 'Converts a GPT-style Causal LM to a Prefix LM.\n\n Supported HuggingFace model classes:\n - `GPT2LMHeadModel`\n - `GPTNeoForCausalLM`\n - `GPTNeoXForCausalLM`\n - `GPTJForCausalLM`\n\n See `convert_hf_causal_lm_to_prefix_lm` for more details.\n '
25
+ if hasattr(model, '_prefix_lm_converted'):
26
+ return model
27
+ assert isinstance(model, _SUPPORTED_GPT_MODELS)
28
+ assert (model.config.add_cross_attention == False), 'Only supports GPT-style decoder-only models'
29
+
30
+ def _get_attn_modules(model: CAUSAL_GPT_TYPES) -> List[torch.nn.Module]:
31
+ "Helper that gets a list of the model's attention modules.\n\n Each module has a `bias` buffer used for causal masking. The Prefix LM\n conversion adds logic to dynamically manipulate these biases to support\n Prefix LM attention masking.\n "
32
+ attn_modules = []
33
+ if isinstance(model, GPTNeoXForCausalLM):
34
+ blocks = model.gpt_neox.layers
35
+ else:
36
+ blocks = model.transformer.h
37
+ for block in blocks:
38
+ if isinstance(model, GPTNeoForCausalLM):
39
+ if (block.attn.attention_type != 'global'):
40
+ continue
41
+ attn_module = block.attn.attention
42
+ elif isinstance(model, GPTNeoXForCausalLM):
43
+ attn_module = block.attention
44
+ else:
45
+ attn_module = block.attn
46
+ attn_modules.append(attn_module)
47
+ return attn_modules
48
+ setattr(model, '_original_forward', getattr(model, 'forward'))
49
+ setattr(model, '_original_generate', getattr(model, 'generate'))
50
+
51
+ def forward(self: CAUSAL_GPT_TYPES, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]]=None, attention_mask: Optional[torch.FloatTensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
52
+ 'Wraps original forward to enable PrefixLM attention.'
53
+
54
+ def call_og_forward():
55
+ if isinstance(self, GPTNeoXForCausalLM):
56
+ return self._original_forward(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
57
+ else:
58
+ return self._original_forward(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
59
+ if (bidirectional_mask is None):
60
+ return call_og_forward()
61
+ assert isinstance(bidirectional_mask, torch.Tensor)
62
+ attn_modules = _get_attn_modules(model)
63
+ (b, s) = bidirectional_mask.shape
64
+ max_length = attn_modules[0].bias.shape[(- 1)]
65
+ if (s > max_length):
66
+ raise ValueError((f'bidirectional_mask sequence length (={s}) exceeds the ' + f'max length allowed by the model ({max_length}).'))
67
+ assert (s <= max_length)
68
+ if (s < max_length):
69
+ pad = torch.zeros((int(b), int((max_length - s))), dtype=bidirectional_mask.dtype, device=bidirectional_mask.device)
70
+ bidirectional_mask = torch.cat([bidirectional_mask, pad], dim=1)
71
+ bidirectional = bidirectional_mask.unsqueeze(1).unsqueeze(1)
72
+ for attn_module in attn_modules:
73
+ attn_module.bias.data = torch.logical_or(attn_module.bias.data, bidirectional)
74
+ output = call_og_forward()
75
+ for attn_module in attn_modules:
76
+ attn_module.bias.data = torch.tril(attn_module.bias.data[(0, 0)])[(None, None)]
77
+ return output
78
+
79
+ def generate(self: CAUSAL_GPT_TYPES, *args: tuple, **kwargs: Dict[(str, Any)]):
80
+ 'Wraps original generate to enable PrefixLM attention.'
81
+ attn_modules = _get_attn_modules(model)
82
+ for attn_module in attn_modules:
83
+ attn_module.bias.data[:] = 1
84
+ output = self._original_generate(*args, **kwargs)
85
+ for attn_module in attn_modules:
86
+ attn_module.bias.data = torch.tril(attn_module.bias.data[(0, 0)])[(None, None)]
87
+ return output
88
+ setattr(model, 'forward', MethodType(forward, model))
89
+ setattr(model, 'generate', MethodType(generate, model))
90
+ setattr(model, '_prefix_lm_converted', True)
91
+ return model
92
+
93
+ def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCausalLM:
94
+ 'Converts a BLOOM Causal LM to a Prefix LM.\n\n Supported HuggingFace model classes:\n - `BloomForCausalLM`\n\n See `convert_hf_causal_lm_to_prefix_lm` for more details.\n '
95
+ if hasattr(model, '_prefix_lm_converted'):
96
+ return model
97
+ assert isinstance(model, BloomForCausalLM)
98
+ assert (model.config.add_cross_attention == False), 'Only supports BLOOM decoder-only models'
99
+
100
+ def _prepare_attn_mask(self: BloomModel, attention_mask: torch.Tensor, bidirectional_mask: Optional[torch.Tensor], input_shape: Tuple[(int, int)], past_key_values_length: int) -> torch.BoolTensor:
101
+ combined_attention_mask = None
102
+ device = attention_mask.device
103
+ (_, src_length) = input_shape
104
+ if (src_length > 1):
105
+ combined_attention_mask = _make_causal_mask_bloom(input_shape, device=device, past_key_values_length=past_key_values_length)
106
+ if (bidirectional_mask is not None):
107
+ assert (attention_mask.shape == bidirectional_mask.shape)
108
+ expanded_bidirectional_mask = _expand_mask_bloom(bidirectional_mask, tgt_length=src_length)
109
+ combined_attention_mask = torch.logical_and(combined_attention_mask, expanded_bidirectional_mask)
110
+ expanded_attn_mask = _expand_mask_bloom(attention_mask, tgt_length=src_length)
111
+ combined_attention_mask = (expanded_attn_mask if (combined_attention_mask is None) else (expanded_attn_mask | combined_attention_mask))
112
+ return combined_attention_mask
113
+
114
+ def _build_alibi_tensor(self: BloomModel, batch_size: int, query_length: int, key_length: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
115
+ num_heads = self.config.n_head
116
+ closest_power_of_2 = (2 ** math.floor(math.log2(num_heads)))
117
+ base = torch.tensor((2 ** (- (2 ** (- (math.log2(closest_power_of_2) - 3))))), device=device, dtype=torch.float32)
118
+ powers = torch.arange(1, (1 + closest_power_of_2), device=device, dtype=torch.int32)
119
+ slopes = torch.pow(base, powers)
120
+ if (closest_power_of_2 != num_heads):
121
+ extra_base = torch.tensor((2 ** (- (2 ** (- (math.log2((2 * closest_power_of_2)) - 3))))), device=device, dtype=torch.float32)
122
+ num_remaining_heads = min(closest_power_of_2, (num_heads - closest_power_of_2))
123
+ extra_powers = torch.arange(1, (1 + (2 * num_remaining_heads)), 2, device=device, dtype=torch.int32)
124
+ slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
125
+ qa = torch.arange(query_length, device=device, dtype=torch.int32).view((- 1), 1)
126
+ ka = torch.arange(key_length, device=device, dtype=torch.int32).view(1, (- 1))
127
+ diffs = (((qa - ka) + key_length) - query_length)
128
+ diffs = (- diffs.abs())
129
+ alibi = (slopes.view(1, num_heads, 1, 1) * diffs.view(1, 1, query_length, key_length))
130
+ alibi = alibi.expand(batch_size, (- 1), (- 1), (- 1)).reshape((- 1), query_length, key_length)
131
+ return alibi.to(dtype)
132
+ KeyValueT = Tuple[(torch.Tensor, torch.Tensor)]
133
+
134
+ def forward(self: BloomModel, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[(KeyValueT, ...)]]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.LongTensor]=None, inputs_embeds: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **deprecated_arguments) -> Union[(Tuple[(torch.Tensor, ...)], BaseModelOutputWithPastAndCrossAttentions)]:
135
+ if (deprecated_arguments.pop('position_ids', False) is not False):
136
+ warnings.warn(('`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. ' + 'You can safely ignore passing `position_ids`.'), FutureWarning)
137
+ if (len(deprecated_arguments) > 0):
138
+ raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
139
+ output_attentions = (output_attentions if (output_attentions is not None) else self.config.output_attentions)
140
+ output_hidden_states = (output_hidden_states if (output_hidden_states is not None) else self.config.output_hidden_states)
141
+ use_cache = (use_cache if (use_cache is not None) else self.config.use_cache)
142
+ return_dict = (return_dict if (return_dict is not None) else self.config.use_return_dict)
143
+ if ((input_ids is not None) and (inputs_embeds is not None)):
144
+ raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
145
+ elif (input_ids is not None):
146
+ (batch_size, seq_length) = input_ids.shape
147
+ elif (inputs_embeds is not None):
148
+ (batch_size, seq_length, _) = inputs_embeds.shape
149
+ else:
150
+ raise ValueError('You have to specify either input_ids or inputs_embeds')
151
+ if (past_key_values is None):
152
+ past_key_values = tuple(([None] * len(self.h)))
153
+ head_mask = self.get_head_mask(head_mask, self.config.n_layer)
154
+ if (inputs_embeds is None):
155
+ inputs_embeds = self.word_embeddings(input_ids)
156
+ hidden_states = self.word_embeddings_layernorm(inputs_embeds)
157
+ presents = (() if use_cache else None)
158
+ all_self_attentions = (() if output_attentions else None)
159
+ all_hidden_states = (() if output_hidden_states else None)
160
+ seq_length_with_past = seq_length
161
+ past_key_values_length = 0
162
+ if (past_key_values[0] is not None):
163
+ tmp = past_key_values[0][0]
164
+ past_key_values_length = tmp.shape[2]
165
+ seq_length_with_past = (seq_length_with_past + past_key_values_length)
166
+ if (attention_mask is None):
167
+ attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
168
+ else:
169
+ attention_mask = attention_mask.to(hidden_states.device)
170
+ alibi = self._build_alibi_tensor(batch_size=batch_size, query_length=seq_length, key_length=seq_length_with_past, dtype=hidden_states.dtype, device=hidden_states.device)
171
+ causal_mask = self._prepare_attn_mask(attention_mask, bidirectional_mask, input_shape=(batch_size, seq_length), past_key_values_length=past_key_values_length)
172
+ for (i, (block, layer_past)) in enumerate(zip(self.h, past_key_values)):
173
+ if output_hidden_states:
174
+ hst = (hidden_states,)
175
+ all_hidden_states = (all_hidden_states + hst)
176
+ if (self.gradient_checkpointing and self.training):
177
+ if use_cache:
178
+ logger.warning('`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...')
179
+ use_cache = False
180
+
181
+ def create_custom_forward(module):
182
+
183
+ def custom_forward(*inputs):
184
+ return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)
185
+ return custom_forward
186
+ outputs = torch.utils.checkpoint.checkpoint(create_custom_forward(block), hidden_states, alibi, causal_mask, head_mask[i])
187
+ else:
188
+ outputs = block(hidden_states, layer_past=layer_past, attention_mask=causal_mask, head_mask=head_mask[i], use_cache=use_cache, output_attentions=output_attentions, alibi=alibi)
189
+ hidden_states = outputs[0]
190
+ if (use_cache is True):
191
+ presents = (presents + (outputs[1],))
192
+ if output_attentions:
193
+ oa = (outputs[(2 if use_cache else 1)],)
194
+ all_self_attentions = (all_self_attentions + oa)
195
+ hidden_states = self.ln_f(hidden_states)
196
+ if output_hidden_states:
197
+ hst = (hidden_states,)
198
+ all_hidden_states = (all_hidden_states + hst)
199
+ if (not return_dict):
200
+ return tuple((v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if (v is not None)))
201
+ return BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=hidden_states, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attentions)
202
+ setattr(model.transformer, '_prepare_attn_mask', MethodType(_prepare_attn_mask, model.transformer))
203
+ setattr(model.transformer, '_build_alibi_tensor', MethodType(_build_alibi_tensor, model.transformer))
204
+ setattr(model.transformer, 'forward', MethodType(forward, model.transformer))
205
+ KeyValueT = Tuple[(torch.Tensor, torch.Tensor)]
206
+
207
+ def forward(self: BloomForCausalLM, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[(KeyValueT, ...)]]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, labels: Optional[torch.Tensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **deprecated_arguments) -> Union[(Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions)]:
208
+ 'Replacement forward method for BloomCausalLM.'
209
+ if (deprecated_arguments.pop('position_ids', False) is not False):
210
+ warnings.warn(('`position_ids` have no functionality in BLOOM and will be removed ' + 'in v5.0.0. You can safely ignore passing `position_ids`.'), FutureWarning)
211
+ if (len(deprecated_arguments) > 0):
212
+ raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
213
+ return_dict = (return_dict if (return_dict is not None) else self.config.use_return_dict)
214
+ transformer_outputs = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask, bidirectional_mask=bidirectional_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
215
+ hidden_states = transformer_outputs[0]
216
+ lm_logits = self.lm_head(hidden_states)
217
+ loss = None
218
+ if (labels is not None):
219
+ shift_logits = lm_logits[..., :(- 1), :].contiguous()
220
+ shift_labels = labels[..., 1:].contiguous()
221
+ (batch_size, seq_length, vocab_size) = shift_logits.shape
222
+ loss_fct = CrossEntropyLoss()
223
+ loss = loss_fct(shift_logits.view((batch_size * seq_length), vocab_size), shift_labels.view((batch_size * seq_length)))
224
+ if (not return_dict):
225
+ output = ((lm_logits,) + transformer_outputs[1:])
226
+ return (((loss,) + output) if (loss is not None) else output)
227
+ return CausalLMOutputWithCrossAttentions(loss=loss, logits=lm_logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions)
228
+
229
+ def prepare_inputs_for_generation(self: BloomForCausalLM, input_ids: torch.LongTensor, past: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, **kwargs) -> dict:
230
+ if past:
231
+ input_ids = input_ids[:, (- 1)].unsqueeze((- 1))
232
+ bidirectional_mask = None
233
+ if (past[0][0].shape[0] == input_ids.shape[0]):
234
+ past = self._convert_to_bloom_cache(past)
235
+ else:
236
+ bidirectional_mask = torch.ones_like(input_ids)
237
+ return {'input_ids': input_ids, 'past_key_values': past, 'use_cache': True, 'attention_mask': attention_mask, 'bidirectional_mask': bidirectional_mask}
238
+ setattr(model, 'forward', MethodType(forward, model))
239
+ setattr(model, 'prepare_inputs_for_generation', MethodType(prepare_inputs_for_generation, model))
240
+ setattr(model, '_prefix_lm_converted', True)
241
+ return model
242
+
243
+ def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM:
244
+ 'Converts an OPT Causal LM to a Prefix LM.\n\n Supported HuggingFace model classes:\n - `OPTForCausalLM`\n\n See `convert_hf_causal_lm_to_prefix_lm` for more details.\n '
245
+ if hasattr(model, '_prefix_lm_converted'):
246
+ return model
247
+ assert isinstance(model, OPTForCausalLM)
248
+ assert (model.config.add_cross_attention == False), 'Only supports OPT decoder-only models'
249
+ setattr(model, '_original_forward', getattr(model, 'forward'))
250
+ setattr(model, '_original_generate', getattr(model, 'generate'))
251
+ model.model.decoder.bidirectional_mask = None
252
+
253
+ def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
254
+ combined_attention_mask = None
255
+ if (input_shape[(- 1)] > 1):
256
+ if (self.bidirectional_mask == 'g'):
257
+ (bsz, src_length) = input_shape
258
+ combined_attention_mask = torch.zeros((bsz, 1, src_length, (src_length + past_key_values_length)), dtype=inputs_embeds.dtype, device=inputs_embeds.device)
259
+ else:
260
+ combined_attention_mask = _make_causal_mask_opt(input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length).to(inputs_embeds.device)
261
+ if (self.bidirectional_mask is not None):
262
+ assert (attention_mask.shape == self.bidirectional_mask.shape)
263
+ expanded_bidirectional_mask = _expand_mask_opt(self.bidirectional_mask, inputs_embeds.dtype, tgt_len=input_shape[(- 1)]).to(inputs_embeds.device)
264
+ combined_attention_mask = torch.maximum(expanded_bidirectional_mask, combined_attention_mask)
265
+ if (attention_mask is not None):
266
+ expanded_attn_mask = _expand_mask_opt(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[(- 1)]).to(inputs_embeds.device)
267
+ combined_attention_mask = (expanded_attn_mask if (combined_attention_mask is None) else (expanded_attn_mask + combined_attention_mask))
268
+ return combined_attention_mask
269
+ setattr(model.model.decoder, '_prepare_decoder_attention_mask', MethodType(_prepare_decoder_attention_mask, model.model.decoder))
270
+
271
+ def forward(self: OPTForCausalLM, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.ByteTensor]=None, head_mask: Optional[torch.Tensor]=None, past_key_values: Optional[List[torch.FloatTensor]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
272
+
273
+ def call_og_forward():
274
+ return self._original_forward(input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
275
+ if (bidirectional_mask is None):
276
+ return call_og_forward()
277
+ self.model.decoder.bidirectional_mask = bidirectional_mask
278
+ try:
279
+ outputs = call_og_forward()
280
+ except:
281
+ self.model.decoder.bidirectional_mask = None
282
+ raise
283
+ self.model.decoder.bidirectional_mask = None
284
+ return outputs
285
+
286
+ def generate(self: OPTForCausalLM, *args: tuple, **kwargs: Dict[(str, Any)]):
287
+ 'Wraps original generate to enable PrefixLM-style attention.'
288
+ self.model.decoder.bidirectional_mask = 'g'
289
+ try:
290
+ output = self._original_generate(*args, **kwargs)
291
+ except:
292
+ self.model.decoder.bidirectional_mask = None
293
+ raise
294
+ self.model.decoder.bidirectional_mask = None
295
+ return output
296
+ setattr(model, 'forward', MethodType(forward, model))
297
+ setattr(model, 'generate', MethodType(generate, model))
298
+ setattr(model, '_prefix_lm_converted', True)
299
+ return model
300
+ _SUPPORTED_HF_MODELS = (_SUPPORTED_GPT_MODELS + (BloomForCausalLM, OPTForCausalLM))
301
+ CAUSAL_LM_TYPES = Union[(GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM, BloomForCausalLM, OPTForCausalLM)]
302
+
303
+ def convert_hf_causal_lm_to_prefix_lm(model: CAUSAL_LM_TYPES) -> CAUSAL_LM_TYPES:
304
+ 'Converts a HuggingFace Causal LM to a Prefix LM.\n\n Supported HuggingFace model classes:\n - `GPT2LMHeadModel`\n - `GPTNeoForCausalLM`\n - `GPTNeoXForCausalLM`\n - `GPTJForCausalLM`\n - `BloomForCausalLM`\n - `OPTForCausalLM`\n\n Conversion to a Prefix LM is done by modifying the `forward` method, and possibly also the\n `generate` method and/or select underlying methods depending on the model class.\n\n These changes preserve the model API, but add a new input to `forward`: "bidirectional_mask".\n\n Notes on training:\n To actually train the converted model as a Prefix LM, training batches will need to indicate\n the prefix/target structure by including `bidirectional_mask` as part of the batch inputs.\n\n **This is not a standard input and requires custom layers either within or after your dataloader.**\n\n In addition to adding `bidirectional_mask` to the batch, this custom code should modify `labels`\n such that `batch[\'labels\'][batch[\'bidirectional_mask\'] == 1] == -100`.\n That is, the prefix portion of the sequence should not generate any loss. Loss should only be\n generated by the target portion of the sequence.\n\n Notes on `GPTNeoForCausalLM`:\n To simplify the implementation, "global" and "local" attention layers are handled differently.\n For "global" layers, we handle conversion as described above. For "local" layers, which use a\n causal attention mask within a restricted local window, we do not alter the masking.\n\n Notes on `forward` method conversion:\n After conversion, the `forward` method will handle a new input, `bidirectional_mask`,\n which should be a [batch_size, seq_length] byte tensor, where 1 indicates token positions\n belonging to the prefix (prefix tokens can attend to one another bidirectionally), and\n 0 indicates token positions belonging to the target.\n\n The new `forward` method will incorporate `bidirectional_mask` (if supplied) into the existing\n causal mask, call the original `forward` method, and (if the causal mask is a buffer) reset\n the causal masks before returning the result.\n\n Notes on `generate` method conversion:\n After conversion, the `generate` method will have the same signature but will internally\n convert all causal masks to be purely bidirectional, call the original `generate` method, and\n (where appropriate) reset the causal masks before returning the result.\n\n This works thanks to the logic of the HuggingFace `generate` API, which first encodes the token\n "prompt" passed to `generate` (which is treated as the prefix) and then sequentially generates\n each new token. Encodings are cached as generation happens, so all prefix tokens can attend to one\n another (as expected in a Prefix LM) and generated tokens can only attend to prefix tokens and\n previously-generated tokens (also as expected in a Prefix LM).\n\n To preserve the API, the original methods are renamed to `_original_forward` and\n `_original_generate`, and replaced with new `forward` and `generate` methods that wrap\n them, respectively. Although implementation details vary by model class.\n '
305
+ if isinstance(model, _SUPPORTED_GPT_MODELS):
306
+ return _convert_gpt_causal_lm_to_prefix_lm(model)
307
+ elif isinstance(model, BloomForCausalLM):
308
+ return _convert_bloom_causal_lm_to_prefix_lm(model)
309
+ elif isinstance(model, OPTForCausalLM):
310
+ return _convert_opt_causal_lm_to_prefix_lm(model)
311
+ else:
312
+ raise TypeError(((f'Cannot convert model to Prefix LM. ' + f'Model does not belong to set of supported HF models:') + f'''
313
+ {_SUPPORTED_HF_MODELS}'''))
314
+
315
+ def add_bidirectional_mask_if_missing(batch: Dict[(str, Any)]):
316
+ "Attempts to add bidirectional_mask to batch if missing.\n\n Raises:\n KeyError if bidirectional_mask is missing and can't be inferred\n "
317
+ if ('bidirectional_mask' not in batch):
318
+ if (batch.get('mode', None) == 'icl_task'):
319
+ batch['bidirectional_mask'] = batch['attention_mask'].clone()
320
+ for (i, continuation_indices) in enumerate(batch['continuation_indices']):
321
+ batch['bidirectional_mask'][(i, continuation_indices)] = 0
322
+ elif (('labels' in batch) and ('attention_mask' in batch)):
323
+ batch['bidirectional_mask'] = torch.logical_and(torch.eq(batch['attention_mask'], 1), torch.eq(batch['labels'], (- 100))).type_as(batch['attention_mask'])
324
+ else:
325
+ raise KeyError('No bidirectional_mask in batch and not sure how to construct one.')
meta_init_context.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ from contextlib import contextmanager
3
+ import torch
4
+ import torch.nn as nn
5
+
6
+ @contextmanager
7
+ def init_empty_weights(include_buffers: bool=False):
8
+ "Meta initialization context manager.\n\n A context manager under which models are initialized with all parameters\n on the meta device, therefore creating an empty model. Useful when just\n initializing the model would blow the available RAM.\n\n Args:\n include_buffers (`bool`, *optional*, defaults to `False`): Whether or\n not to also put all buffers on the meta device while initializing.\n\n Example:\n ```python\n import torch.nn as nn\n\n # Initialize a model with 100 billions parameters in no time and without using any RAM.\n with init_empty_weights():\n tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])\n ```\n\n <Tip warning={true}>\n\n Any model created under this context manager has no weights. As such you can't do something like\n `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].\n\n </Tip>\n "
9
+ with init_on_device(torch.device('meta'), include_buffers=include_buffers) as f:
10
+ (yield f)
11
+
12
+ @contextmanager
13
+ def init_on_device(device: torch.device, include_buffers: bool=False):
14
+ 'Device initialization context manager.\n\n A context manager under which models are initialized with all parameters\n on the specified device.\n\n Args:\n device (`torch.device`): Device to initialize all parameters on.\n include_buffers (`bool`, *optional*, defaults to `False`): Whether or\n not to also put all buffers on the meta device while initializing.\n\n Example:\n ```python\n import torch.nn as nn\n\n with init_on_device(device=torch.device("cuda")):\n tst = nn.Liner(100, 100) # on `cuda` device\n ```\n '
15
+ old_register_parameter = nn.Module.register_parameter
16
+ if include_buffers:
17
+ old_register_buffer = nn.Module.register_buffer
18
+
19
+ def register_empty_parameter(module, name, param):
20
+ old_register_parameter(module, name, param)
21
+ if (param is not None):
22
+ param_cls = type(module._parameters[name])
23
+ kwargs = module._parameters[name].__dict__
24
+ module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
25
+
26
+ def register_empty_buffer(module, name, buffer):
27
+ old_register_buffer(module, name, buffer)
28
+ if (buffer is not None):
29
+ module._buffers[name] = module._buffers[name].to(device)
30
+ if include_buffers:
31
+ tensor_constructors_to_patch = {torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ['empty', 'zeros', 'ones', 'full']}
32
+ else:
33
+ tensor_constructors_to_patch = {}
34
+
35
+ def patch_tensor_constructor(fn):
36
+
37
+ def wrapper(*args, **kwargs):
38
+ kwargs['device'] = device
39
+ return fn(*args, **kwargs)
40
+ return wrapper
41
+ try:
42
+ nn.Module.register_parameter = register_empty_parameter
43
+ if include_buffers:
44
+ nn.Module.register_buffer = register_empty_buffer
45
+ for torch_function_name in tensor_constructors_to_patch.keys():
46
+ setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
47
+ (yield)
48
+ finally:
49
+ nn.Module.register_parameter = old_register_parameter
50
+ if include_buffers:
51
+ nn.Module.register_buffer = old_register_buffer
52
+ for (torch_function_name, old_torch_function) in tensor_constructors_to_patch.items():
53
+ setattr(torch, torch_function_name, old_torch_function)
modeling_mpt.py ADDED
@@ -0,0 +1,289 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ 'A simple, flexible implementation of a GPT model.\n\nInspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py\n'
3
+ import math
4
+ import warnings
5
+ from typing import List, Optional, Tuple, Union
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
10
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
11
+ from .attention import attn_bias_shape, build_attn_bias
12
+ from .blocks import MPTBlock
13
+ from .norm import NORM_CLASS_REGISTRY
14
+ from .configuration_mpt import MPTConfig
15
+ from .adapt_tokenizer import AutoTokenizerForMOD, adapt_tokenizer_for_denoising
16
+ from .hf_prefixlm_converter import add_bidirectional_mask_if_missing, convert_hf_causal_lm_to_prefix_lm
17
+ from .meta_init_context import init_empty_weights
18
+ from .param_init_fns import MODEL_INIT_REGISTRY, generic_param_init_fn_
19
+ try:
20
+ from .flash_attn_triton import flash_attn_func
21
+ except:
22
+ pass
23
+ Tokenizer = Union[(PreTrainedTokenizer, PreTrainedTokenizerFast)]
24
+
25
+ class MPTPreTrainedModel(PreTrainedModel):
26
+ config_class = MPTConfig
27
+ base_model_prefix = 'model'
28
+
29
+ class MPTModel(MPTPreTrainedModel):
30
+
31
+ def __init__(self, config: MPTConfig):
32
+ config._validate_config()
33
+ super().__init__(config)
34
+ self.attn_impl = config.attn_config['attn_impl']
35
+ self.prefix_lm = config.attn_config['prefix_lm']
36
+ self.attn_uses_sequence_id = config.attn_config['attn_uses_sequence_id']
37
+ self.alibi = config.attn_config['alibi']
38
+ self.alibi_bias_max = config.attn_config['alibi_bias_max']
39
+ if (config.norm_type.lower() not in NORM_CLASS_REGISTRY.keys()):
40
+ norm_options = ' | '.join(NORM_CLASS_REGISTRY.keys())
41
+ raise NotImplementedError(f'Requested norm type ({config.norm_type}) is not implemented within this repo (Options: {norm_options}).')
42
+ norm_class = NORM_CLASS_REGISTRY[config.norm_type.lower()]
43
+ self.embedding_fraction = config.embedding_fraction
44
+ self.wte = nn.Embedding(config.vocab_size, config.d_model, device=config.init_device)
45
+ if (not self.alibi):
46
+ self.wpe = nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
47
+ self.emb_drop = nn.Dropout(config.emb_pdrop)
48
+ self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
49
+ self.norm_f = norm_class(config.d_model, device=config.init_device)
50
+ if (config.init_device != 'meta'):
51
+ print(f'You are using config.init_device={config.init_device!r}, but you can also use config.init_device="meta" with Composer + FSDP for fast initialization.')
52
+ self.apply(self.param_init_fn)
53
+ self.is_causal = (not self.prefix_lm)
54
+ self._attn_bias_initialized = False
55
+ self.attn_bias = None
56
+ self.attn_bias_shape = attn_bias_shape(self.attn_impl, config.n_heads, config.max_seq_len, self.alibi, prefix_lm=self.prefix_lm, causal=self.is_causal, use_sequence_id=self.attn_uses_sequence_id)
57
+ if config.no_bias:
58
+ for module in self.modules():
59
+ if (hasattr(module, 'bias') and isinstance(module.bias, nn.Parameter)):
60
+ if config.verbose:
61
+ warnings.warn(f'Removing bias ({module.bias}) from {module}.')
62
+ module.register_parameter('bias', None)
63
+ if (config.verbose and (config.verbose > 2)):
64
+ print(self)
65
+ if ('verbose' not in self.config.init_config):
66
+ self.config.init_config['verbose'] = self.config.verbose
67
+ if (self.config.init_config['verbose'] > 1):
68
+ init_fn_name = self.config.init_config['name']
69
+ warnings.warn(f'Using {init_fn_name} initialization.')
70
+
71
+ def get_input_embeddings(self):
72
+ return self.wte
73
+
74
+ def set_input_embeddings(self, value):
75
+ self.wte = value
76
+
77
+ @torch.no_grad()
78
+ def _attn_bias(self, device, dtype, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None):
79
+ if (not self._attn_bias_initialized):
80
+ if self.attn_bias_shape:
81
+ self.attn_bias = torch.zeros(self.attn_bias_shape, device=device, dtype=dtype)
82
+ self.attn_bias = build_attn_bias(self.attn_impl, self.attn_bias, self.config.n_heads, self.config.max_seq_len, causal=self.is_causal, alibi=self.alibi, alibi_bias_max=self.alibi_bias_max)
83
+ self._attn_bias_initialized = True
84
+ if (self.attn_impl == 'flash'):
85
+ return (self.attn_bias, attention_mask)
86
+ if (self.attn_bias is not None):
87
+ self.attn_bias = self.attn_bias.to(dtype=dtype, device=device)
88
+ attn_bias = self.attn_bias
89
+ if self.prefix_lm:
90
+ assert isinstance(attn_bias, torch.Tensor)
91
+ assert isinstance(prefix_mask, torch.Tensor)
92
+ attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)
93
+ if (self.attn_uses_sequence_id and (sequence_id is not None)):
94
+ assert isinstance(attn_bias, torch.Tensor)
95
+ attn_bias = self._apply_sequence_id(attn_bias, sequence_id)
96
+ if (attention_mask is not None):
97
+ s_k = attention_mask.shape[(- 1)]
98
+ if (attn_bias is None):
99
+ attn_bias = torch.zeros((1, 1, 1, s_k), device=device, dtype=dtype)
100
+ else:
101
+ attn_bias = attn_bias[:, :, :, (- s_k):]
102
+ if ((prefix_mask is not None) and (attention_mask.shape != prefix_mask.shape)):
103
+ raise ValueError((f'attention_mask shape={attention_mask.shape} ' + f'and prefix_mask shape={prefix_mask.shape} are not equal.'))
104
+ min_val = torch.finfo(attn_bias.dtype).min
105
+ attn_bias = attn_bias.masked_fill((~ attention_mask.view((- 1), 1, 1, s_k)), min_val)
106
+ return (attn_bias, None)
107
+
108
+ def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor):
109
+ (s_k, s_q) = attn_bias.shape[(- 2):]
110
+ if ((s_k != self.config.max_seq_len) or (s_q != self.config.max_seq_len)):
111
+ raise ValueError((('attn_bias does not match the expected shape. ' + f'The last two dimensions should both be {self.config.max_length} ') + f'but are {s_k} and {s_q}.'))
112
+ seq_len = prefix_mask.shape[(- 1)]
113
+ if (seq_len > self.config.max_seq_len):
114
+ raise ValueError(f'prefix_mask sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
115
+ attn_bias = attn_bias[..., :seq_len, :seq_len]
116
+ causal = torch.tril(torch.ones((seq_len, seq_len), dtype=torch.bool, device=prefix_mask.device)).view(1, 1, seq_len, seq_len)
117
+ prefix = prefix_mask.view((- 1), 1, 1, seq_len)
118
+ cannot_attend = (~ torch.logical_or(causal, prefix.bool()))
119
+ min_val = torch.finfo(attn_bias.dtype).min
120
+ attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
121
+ return attn_bias
122
+
123
+ def _apply_sequence_id(self, attn_bias: torch.Tensor, sequence_id: torch.LongTensor):
124
+ seq_len = sequence_id.shape[(- 1)]
125
+ if (seq_len > self.config.max_seq_len):
126
+ raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
127
+ attn_bias = attn_bias[..., :seq_len, :seq_len]
128
+ cannot_attend = torch.logical_not(torch.eq(sequence_id.view((- 1), seq_len, 1), sequence_id.view((- 1), 1, seq_len))).unsqueeze(1)
129
+ min_val = torch.finfo(attn_bias.dtype).min
130
+ attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
131
+ return attn_bias
132
+
133
+ def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None):
134
+ return_dict = (return_dict if (return_dict is not None) else self.config.return_dict)
135
+ use_cache = (use_cache if (use_cache is not None) else self.config.use_cache)
136
+ if (attention_mask is not None):
137
+ attention_mask = attention_mask.bool()
138
+ if (prefix_mask is not None):
139
+ prefix_mask = prefix_mask.bool()
140
+ if (not return_dict):
141
+ raise NotImplementedError('return_dict False is not implemented yet for MPT')
142
+ if output_attentions:
143
+ raise NotImplementedError('output_attentions is not implemented yet for MPT')
144
+ if ((attention_mask is not None) and (attention_mask[:, 0].sum() != attention_mask.shape[0]) and self.training):
145
+ raise NotImplementedError('MPT does not support training with left padding.')
146
+ if (self.prefix_lm and (prefix_mask is None)):
147
+ raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
148
+ if self.training:
149
+ if (self.attn_uses_sequence_id and (sequence_id is None)):
150
+ raise ValueError(('sequence_id is a required argument when MPT is configured with attn_uses_sequence_id=True ' + 'and the model is in train mode.'))
151
+ elif ((self.attn_uses_sequence_id is False) and (sequence_id is not None)):
152
+ warnings.warn(('MPT received non-None input for `sequence_id` but is configured with attn_uses_sequence_id=False. ' + 'This input will be ignored. If you want the model to use `sequence_id`, set attn_uses_sequence_id to True.'))
153
+ S = input_ids.size(1)
154
+ assert (S <= self.config.max_seq_len), f'Cannot forward input with seq_len={S}, this model only supports seq_len<={self.config.max_seq_len}'
155
+ tok_emb = self.wte(input_ids)
156
+ if self.alibi:
157
+ x = tok_emb
158
+ else:
159
+ past_position = 0
160
+ if (past_key_values is not None):
161
+ if (len(past_key_values) != self.config.n_layers):
162
+ raise ValueError((f'past_key_values must provide a past_key_value for each attention ' + f'layer in the network (len(past_key_values)={len(past_key_values)!r}; self.config.n_layers={self.config.n_layers!r}).'))
163
+ past_position = past_key_values[0][0].size(1)
164
+ if ((S + past_position) > self.config.max_seq_len):
165
+ raise ValueError(f'Cannot forward input with past sequence length {past_position} and current sequence length {(S + 1)}, this model only supports total sequence length <= {self.config.max_seq_len}.')
166
+ pos = torch.arange(past_position, (S + past_position), dtype=torch.long, device=input_ids.device).unsqueeze(0)
167
+ if (attention_mask is not None):
168
+ pos = torch.clamp((pos - torch.cumsum((~ attention_mask).to(torch.int32), dim=1)[:, past_position:]), min=0)
169
+ pos_emb = self.wpe(pos)
170
+ x = (tok_emb + pos_emb)
171
+ if (self.embedding_fraction == 1):
172
+ x = self.emb_drop(x)
173
+ else:
174
+ x_shrunk = ((x * self.embedding_fraction) + (x.detach() * (1 - self.embedding_fraction)))
175
+ assert isinstance(self.emb_drop, nn.Module)
176
+ x = self.emb_drop(x_shrunk)
177
+ (attn_bias, attention_mask) = self._attn_bias(device=x.device, dtype=x.dtype, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id)
178
+ if (use_cache and (past_key_values is None)):
179
+ past_key_values = [() for _ in range(self.config.n_layers)]
180
+ all_hidden_states = (() if output_hidden_states else None)
181
+ for (b_idx, block) in enumerate(self.blocks):
182
+ if output_hidden_states:
183
+ assert (all_hidden_states is not None)
184
+ all_hidden_states = (all_hidden_states + (x,))
185
+ past_key_value = (past_key_values[b_idx] if (past_key_values is not None) else None)
186
+ (x, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=self.is_causal)
187
+ if (past_key_values is not None):
188
+ past_key_values[b_idx] = past_key_value
189
+ x = self.norm_f(x)
190
+ return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values, hidden_states=all_hidden_states)
191
+
192
+ def param_init_fn(self, module):
193
+ init_fn_name = self.config.init_config['name']
194
+ MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
195
+
196
+ def fsdp_wrap_fn(self, module):
197
+ return isinstance(module, MPTBlock)
198
+
199
+ def activation_checkpointing_fn(self, module):
200
+ return isinstance(module, MPTBlock)
201
+
202
+ class MPTForCausalLM(MPTPreTrainedModel):
203
+
204
+ def __init__(self, config: MPTConfig):
205
+ super().__init__(config)
206
+ if (not config.tie_word_embeddings):
207
+ raise ValueError('MPTForCausalLM only supports tied word embeddings')
208
+ self.transformer = MPTModel(config)
209
+ self.logit_scale = None
210
+ if (config.logit_scale is not None):
211
+ logit_scale = config.logit_scale
212
+ if isinstance(logit_scale, str):
213
+ if (logit_scale == 'inv_sqrt_d_model'):
214
+ logit_scale = (1 / math.sqrt(config.d_model))
215
+ else:
216
+ raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
217
+ self.logit_scale = logit_scale
218
+
219
+ def get_input_embeddings(self):
220
+ return self.transformer.wte
221
+
222
+ def set_input_embeddings(self, value):
223
+ self.transformer.wte = value
224
+
225
+ def get_output_embeddings(self):
226
+ return self.transformer.wte
227
+
228
+ def set_output_embeddings(self, new_embeddings):
229
+ self.transformer.wte = new_embeddings
230
+
231
+ def set_decoder(self, decoder):
232
+ self.transformer = decoder
233
+
234
+ def get_decoder(self):
235
+ return self.transformer
236
+
237
+ def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, labels: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None):
238
+ return_dict = (return_dict if (return_dict is not None) else self.config.return_dict)
239
+ use_cache = (use_cache if (use_cache is not None) else self.config.use_cache)
240
+ outputs = self.transformer(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id, return_dict=return_dict, output_attentions=output_attentions, output_hidden_states=output_hidden_states, use_cache=use_cache)
241
+ logits = F.linear(outputs.last_hidden_state, self.transformer.wte.weight)
242
+ if (self.logit_scale is not None):
243
+ if (self.logit_scale == 0):
244
+ warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
245
+ logits *= self.logit_scale
246
+ loss = None
247
+ if (labels is not None):
248
+ labels = torch.roll(labels, shifts=(- 1))
249
+ labels[:, (- 1)] = (- 100)
250
+ loss = F.cross_entropy(logits.view((- 1), logits.size((- 1))), labels.to(logits.device).view((- 1)))
251
+ return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states)
252
+
253
+ def param_init_fn(self, module):
254
+ init_fn_name = self.config.init_config['name']
255
+ MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
256
+
257
+ def fsdp_wrap_fn(self, module):
258
+ return isinstance(module, MPTBlock)
259
+
260
+ def activation_checkpointing_fn(self, module):
261
+ return isinstance(module, MPTBlock)
262
+
263
+ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
264
+ if (inputs_embeds is not None):
265
+ raise NotImplementedError('inputs_embeds is not implemented for MPT yet')
266
+ attention_mask = kwargs['attention_mask'].bool()
267
+ if (attention_mask[:, (- 1)].sum() != attention_mask.shape[0]):
268
+ raise NotImplementedError('MPT does not support generation with right padding.')
269
+ if (self.transformer.attn_uses_sequence_id and self.training):
270
+ sequence_id = torch.zeros_like(input_ids[:1])
271
+ else:
272
+ sequence_id = None
273
+ if (past_key_values is not None):
274
+ input_ids = input_ids[:, (- 1)].unsqueeze((- 1))
275
+ if self.transformer.prefix_lm:
276
+ prefix_mask = torch.ones_like(attention_mask)
277
+ if (kwargs.get('use_cache') == False):
278
+ raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
279
+ else:
280
+ prefix_mask = None
281
+ return {'input_ids': input_ids, 'attention_mask': attention_mask, 'prefix_mask': prefix_mask, 'sequence_id': sequence_id, 'past_key_values': past_key_values, 'use_cache': kwargs.get('use_cache', True)}
282
+
283
+ @staticmethod
284
+ def _reorder_cache(past_key_values, beam_idx):
285
+ 'Used by HuggingFace generate when using beam search with kv-caching.\n\n See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133\n for an example in transformers.\n '
286
+ reordered_past = []
287
+ for layer_past in past_key_values:
288
+ reordered_past += [tuple((past_state.index_select(0, beam_idx) for past_state in layer_past))]
289
+ return reordered_past
norm.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import torch
3
+
4
+ def _cast_if_autocast_enabled(tensor):
5
+ if torch.is_autocast_enabled():
6
+ if (tensor.device.type == 'cuda'):
7
+ dtype = torch.get_autocast_gpu_dtype()
8
+ elif (tensor.device.type == 'cpu'):
9
+ dtype = torch.get_autocast_cpu_dtype()
10
+ else:
11
+ raise NotImplementedError()
12
+ return tensor.to(dtype=dtype)
13
+ return tensor
14
+
15
+ class LPLayerNorm(torch.nn.LayerNorm):
16
+
17
+ def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None):
18
+ super().__init__(normalized_shape=normalized_shape, eps=eps, elementwise_affine=elementwise_affine, device=device, dtype=dtype)
19
+
20
+ def forward(self, x):
21
+ module_device = x.device
22
+ downcast_x = _cast_if_autocast_enabled(x)
23
+ downcast_weight = (_cast_if_autocast_enabled(self.weight) if (self.weight is not None) else self.weight)
24
+ downcast_bias = (_cast_if_autocast_enabled(self.bias) if (self.bias is not None) else self.bias)
25
+ with torch.autocast(enabled=False, device_type=module_device.type):
26
+ return torch.nn.functional.layer_norm(downcast_x, self.normalized_shape, downcast_weight, downcast_bias, self.eps)
27
+
28
+ def rms_norm(x, weight=None, eps=1e-05):
29
+ output = (x / torch.rsqrt((x.pow(2).mean((- 1), keepdim=True) + eps)))
30
+ if (weight is not None):
31
+ return (output * weight)
32
+ return output
33
+
34
+ class RMSNorm(torch.nn.Module):
35
+
36
+ def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
37
+ super().__init__()
38
+ self.eps = eps
39
+ if weight:
40
+ self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))
41
+ else:
42
+ self.register_parameter('weight', None)
43
+
44
+ def forward(self, x):
45
+ return rms_norm(x.float(), self.weight, self.eps).to(dtype=x.dtype)
46
+
47
+ class LPRMSNorm(RMSNorm):
48
+
49
+ def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
50
+ super().__init__(normalized_shape=normalized_shape, eps=eps, weight=weight, dtype=dtype, device=device)
51
+
52
+ def forward(self, x):
53
+ downcast_x = _cast_if_autocast_enabled(x)
54
+ downcast_weight = (_cast_if_autocast_enabled(self.weight) if (self.weight is not None) else self.weight)
55
+ with torch.autocast(enabled=False, device_type=x.device.type):
56
+ return rms_norm(downcast_x, downcast_weight, self.eps).to(dtype=x.dtype)
57
+ NORM_CLASS_REGISTRY = {'layernorm': torch.nn.LayerNorm, 'low_precision_layernorm': LPLayerNorm, 'rmsnorm': RMSNorm, 'low_precision_rmsnorm': LPRMSNorm}
param_init_fns.py ADDED
@@ -0,0 +1,177 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import math
3
+ import warnings
4
+ from collections.abc import Sequence
5
+ from functools import partial
6
+ from typing import Optional, Tuple, Union
7
+ import torch
8
+ from torch import nn
9
+ from .norm import NORM_CLASS_REGISTRY
10
+
11
+ def torch_default_param_init_fn_(module: nn.Module, verbose: int=0, **kwargs):
12
+ del kwargs
13
+ if (verbose > 1):
14
+ warnings.warn(f"Initializing network using module's reset_parameters attribute")
15
+ if hasattr(module, 'reset_parameters'):
16
+ module.reset_parameters()
17
+
18
+ def fused_init_helper_(module: nn.Module, init_fn_):
19
+ _fused = getattr(module, '_fused', None)
20
+ if (_fused is None):
21
+ raise RuntimeError(f'Internal logic error')
22
+ (dim, splits) = _fused
23
+ splits = (0, *splits, module.weight.size(dim))
24
+ for (s, e) in zip(splits[:(- 1)], splits[1:]):
25
+ slice_indices = ([slice(None)] * module.weight.ndim)
26
+ slice_indices[dim] = slice(s, e)
27
+ init_fn_(module.weight[slice_indices])
28
+
29
+ def generic_param_init_fn_(module: nn.Module, init_fn_, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[(int, float, str, bool)]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[(Tuple[(float, float)], float)]]=None, verbose: int=0, **kwargs):
30
+ del kwargs
31
+ if (verbose > 1):
32
+ warnings.warn(f'If model has bias parameters they are initialized to 0.')
33
+ init_div_is_residual = init_div_is_residual
34
+ if (init_div_is_residual is False):
35
+ div_is_residual = 1.0
36
+ elif (init_div_is_residual is True):
37
+ div_is_residual = math.sqrt((2 * n_layers))
38
+ elif (isinstance(init_div_is_residual, float) or isinstance(init_div_is_residual, int)):
39
+ div_is_residual = init_div_is_residual
40
+ elif (isinstance(init_div_is_residual, str) and init_div_is_residual.isnumeric()):
41
+ div_is_residual = float(init_div_is_residual)
42
+ else:
43
+ div_is_residual = 1.0
44
+ raise ValueError(f'Expected init_div_is_residual to be boolean or numeric, got {init_div_is_residual}')
45
+ if (init_div_is_residual is not False):
46
+ if (verbose > 1):
47
+ warnings.warn((f'Initializing _is_residual layers then dividing them by {div_is_residual:.3f}. ' + f'Set `init_div_is_residual: false` in init config to disable this.'))
48
+ if isinstance(module, nn.Linear):
49
+ if hasattr(module, '_fused'):
50
+ fused_init_helper_(module, init_fn_)
51
+ else:
52
+ init_fn_(module.weight)
53
+ if (module.bias is not None):
54
+ torch.nn.init.zeros_(module.bias)
55
+ if ((init_div_is_residual is not False) and getattr(module, '_is_residual', False)):
56
+ with torch.no_grad():
57
+ module.weight.div_(div_is_residual)
58
+ elif isinstance(module, nn.Embedding):
59
+ if (emb_init_std is not None):
60
+ std = emb_init_std
61
+ if (std == 0):
62
+ warnings.warn(f'Embedding layer initialized to 0.')
63
+ emb_init_fn_ = partial(torch.nn.init.normal_, mean=0.0, std=std)
64
+ if (verbose > 1):
65
+ warnings.warn(f'Embedding layer initialized using normal distribution with mean=0 and std={std!r}.')
66
+ elif (emb_init_uniform_lim is not None):
67
+ lim = emb_init_uniform_lim
68
+ if isinstance(lim, Sequence):
69
+ if (len(lim) > 2):
70
+ raise ValueError(f'Uniform init requires a min and a max limit. User input: {lim}.')
71
+ if (lim[0] == lim[1]):
72
+ warnings.warn(f'Embedding layer initialized to {lim[0]}.')
73
+ else:
74
+ if (lim == 0):
75
+ warnings.warn(f'Embedding layer initialized to 0.')
76
+ lim = [(- lim), lim]
77
+ (a, b) = lim
78
+ emb_init_fn_ = partial(torch.nn.init.uniform_, a=a, b=b)
79
+ if (verbose > 1):
80
+ warnings.warn(f'Embedding layer initialized using uniform distribution in range {lim}.')
81
+ else:
82
+ emb_init_fn_ = init_fn_
83
+ emb_init_fn_(module.weight)
84
+ elif isinstance(module, tuple(set(NORM_CLASS_REGISTRY.values()))):
85
+ if (verbose > 1):
86
+ warnings.warn(f'Norm weights are set to 1. If norm layer has a bias it is initialized to 0.')
87
+ if (hasattr(module, 'weight') and (module.weight is not None)):
88
+ torch.nn.init.ones_(module.weight)
89
+ if (hasattr(module, 'bias') and (module.bias is not None)):
90
+ torch.nn.init.zeros_(module.bias)
91
+ elif isinstance(module, nn.MultiheadAttention):
92
+ if module._qkv_same_embed_dim:
93
+ assert (module.in_proj_weight is not None)
94
+ assert ((module.q_proj_weight is None) and (module.k_proj_weight is None) and (module.v_proj_weight is None))
95
+ assert (d_model is not None)
96
+ _d = d_model
97
+ splits = (0, _d, (2 * _d), (3 * _d))
98
+ for (s, e) in zip(splits[:(- 1)], splits[1:]):
99
+ init_fn_(module.in_proj_weight[s:e])
100
+ else:
101
+ assert ((module.q_proj_weight is not None) and (module.k_proj_weight is not None) and (module.v_proj_weight is not None))
102
+ assert (module.in_proj_weight is None)
103
+ init_fn_(module.q_proj_weight)
104
+ init_fn_(module.k_proj_weight)
105
+ init_fn_(module.v_proj_weight)
106
+ if (module.in_proj_bias is not None):
107
+ torch.nn.init.zeros_(module.in_proj_bias)
108
+ if (module.bias_k is not None):
109
+ torch.nn.init.zeros_(module.bias_k)
110
+ if (module.bias_v is not None):
111
+ torch.nn.init.zeros_(module.bias_v)
112
+ init_fn_(module.out_proj.weight)
113
+ if ((init_div_is_residual is not False) and getattr(module.out_proj, '_is_residual', False)):
114
+ with torch.no_grad():
115
+ module.out_proj.weight.div_(div_is_residual)
116
+ if (module.out_proj.bias is not None):
117
+ torch.nn.init.zeros_(module.out_proj.bias)
118
+ else:
119
+ for _ in module.parameters(recurse=False):
120
+ raise NotImplementedError(f'{module.__class__.__name__} parameters are not initialized by param_init_fn.')
121
+
122
+ def _normal_init_(std, mean=0.0):
123
+ return partial(torch.nn.init.normal_, mean=mean, std=std)
124
+
125
+ def _normal_param_init_fn_(module: nn.Module, std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[(int, float, str, bool)]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[(Tuple[(float, float)], float)]]=None, verbose: int=0, **kwargs):
126
+ del kwargs
127
+ init_fn_ = _normal_init_(std=std)
128
+ if (verbose > 1):
129
+ warnings.warn(f'Using torch.nn.init.normal_ init fn mean=0.0, std={std}')
130
+ generic_param_init_fn_(module=module, init_fn_=init_fn_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
131
+
132
+ def baseline_param_init_fn_(module: nn.Module, init_std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[(int, float, str, bool)]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[(Tuple[(float, float)], float)]]=None, verbose: int=0, **kwargs):
133
+ del kwargs
134
+ if (init_std is None):
135
+ raise ValueError("You must set model.init_config['init_std'] to a float value to use the default initialization scheme.")
136
+ _normal_param_init_fn_(module=module, std=init_std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
137
+
138
+ def small_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, init_div_is_residual: Union[(int, float, str, bool)]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[(Tuple[(float, float)], float)]]=None, verbose: int=0, **kwargs):
139
+ del kwargs
140
+ std = math.sqrt((2 / (5 * d_model)))
141
+ _normal_param_init_fn_(module=module, std=std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
142
+
143
+ def neox_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[(Tuple[(float, float)], float)]]=None, verbose: int=0, **kwargs):
144
+ 'From section 2.3.1 of GPT-NeoX-20B:\n\n An Open-Source AutoregressiveLanguage Model — Black et. al. (2022)\n see https://github.com/EleutherAI/gpt-neox/blob/9610391ab319403cef079b438edd016a2443af54/megatron/model/init_functions.py#L151\n and https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/transformer.py\n '
145
+ del kwargs
146
+ residual_div = (n_layers / math.sqrt(10))
147
+ if (verbose > 1):
148
+ warnings.warn(f'setting init_div_is_residual to {residual_div}')
149
+ small_param_init_fn_(module=module, d_model=d_model, n_layers=n_layers, init_div_is_residual=residual_div, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
150
+
151
+ def kaiming_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[(int, float, str, bool)]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[(Tuple[(float, float)], float)]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', verbose: int=0, **kwargs):
152
+ del kwargs
153
+ if (verbose > 1):
154
+ warnings.warn((f'Using nn.init.kaiming_uniform_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}'))
155
+ kaiming_uniform_ = partial(nn.init.kaiming_uniform_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
156
+ generic_param_init_fn_(module=module, init_fn_=kaiming_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
157
+
158
+ def kaiming_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[(int, float, str, bool)]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[(Tuple[(float, float)], float)]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', verbose: int=0, **kwargs):
159
+ del kwargs
160
+ if (verbose > 1):
161
+ warnings.warn((f'Using nn.init.kaiming_normal_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}'))
162
+ kaiming_normal_ = partial(torch.nn.init.kaiming_normal_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
163
+ generic_param_init_fn_(module=module, init_fn_=kaiming_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
164
+
165
+ def xavier_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[(int, float, str, bool)]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[(Tuple[(float, float)], float)]]=None, init_gain: float=0, verbose: int=0, **kwargs):
166
+ del kwargs
167
+ xavier_uniform_ = partial(torch.nn.init.xavier_uniform_, gain=init_gain)
168
+ if (verbose > 1):
169
+ warnings.warn((f'Using torch.nn.init.xavier_uniform_ init fn with parameters: ' + f'gain={init_gain}'))
170
+ generic_param_init_fn_(module=module, init_fn_=xavier_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
171
+
172
+ def xavier_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[(int, float, str, bool)]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[(Tuple[(float, float)], float)]]=None, init_gain: float=0, verbose: int=0, **kwargs):
173
+ xavier_normal_ = partial(torch.nn.init.xavier_normal_, gain=init_gain)
174
+ if (verbose > 1):
175
+ warnings.warn((f'Using torch.nn.init.xavier_normal_ init fn with parameters: ' + f'gain={init_gain}'))
176
+ generic_param_init_fn_(module=module, init_fn_=xavier_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
177
+ MODEL_INIT_REGISTRY = {'default_': torch_default_param_init_fn_, 'baseline_': baseline_param_init_fn_, 'kaiming_uniform_': kaiming_uniform_param_init_fn_, 'kaiming_normal_': kaiming_normal_param_init_fn_, 'neox_init_': neox_param_init_fn_, 'small_init_': small_param_init_fn_, 'xavier_uniform_': xavier_uniform_param_init_fn_, 'xavier_normal_': xavier_normal_param_init_fn_}
pytorch_model-00001-of-00002.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f3df6189e65e853b8de547de86f198a1fbcaac9a1950aba87e576541f92a9e12
3
+ size 9943040275
pytorch_model-00002-of-00002.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:44eb87aa397044c43a4afc7fe19ed31248eaa999b7a50444285c19f8204daedb
3
+ size 3355599187
pytorch_model.bin.index.json ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 13298573312
4
+ },
5
+ "weight_map": {
6
+ "transformer.blocks.0.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
7
+ "transformer.blocks.0.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
8
+ "transformer.blocks.0.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
9
+ "transformer.blocks.0.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
10
+ "transformer.blocks.0.norm_1.weight": "pytorch_model-00001-of-00002.bin",
11
+ "transformer.blocks.0.norm_2.weight": "pytorch_model-00001-of-00002.bin",
12
+ "transformer.blocks.1.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
13
+ "transformer.blocks.1.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
14
+ "transformer.blocks.1.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
15
+ "transformer.blocks.1.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
16
+ "transformer.blocks.1.norm_1.weight": "pytorch_model-00001-of-00002.bin",
17
+ "transformer.blocks.1.norm_2.weight": "pytorch_model-00001-of-00002.bin",
18
+ "transformer.blocks.10.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
19
+ "transformer.blocks.10.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
20
+ "transformer.blocks.10.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
21
+ "transformer.blocks.10.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
22
+ "transformer.blocks.10.norm_1.weight": "pytorch_model-00001-of-00002.bin",
23
+ "transformer.blocks.10.norm_2.weight": "pytorch_model-00001-of-00002.bin",
24
+ "transformer.blocks.11.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
25
+ "transformer.blocks.11.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
26
+ "transformer.blocks.11.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
27
+ "transformer.blocks.11.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
28
+ "transformer.blocks.11.norm_1.weight": "pytorch_model-00001-of-00002.bin",
29
+ "transformer.blocks.11.norm_2.weight": "pytorch_model-00001-of-00002.bin",
30
+ "transformer.blocks.12.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
31
+ "transformer.blocks.12.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
32
+ "transformer.blocks.12.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
33
+ "transformer.blocks.12.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
34
+ "transformer.blocks.12.norm_1.weight": "pytorch_model-00001-of-00002.bin",
35
+ "transformer.blocks.12.norm_2.weight": "pytorch_model-00001-of-00002.bin",
36
+ "transformer.blocks.13.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
37
+ "transformer.blocks.13.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
38
+ "transformer.blocks.13.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
39
+ "transformer.blocks.13.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
40
+ "transformer.blocks.13.norm_1.weight": "pytorch_model-00001-of-00002.bin",
41
+ "transformer.blocks.13.norm_2.weight": "pytorch_model-00001-of-00002.bin",
42
+ "transformer.blocks.14.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
43
+ "transformer.blocks.14.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
44
+ "transformer.blocks.14.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
45
+ "transformer.blocks.14.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
46
+ "transformer.blocks.14.norm_1.weight": "pytorch_model-00001-of-00002.bin",
47
+ "transformer.blocks.14.norm_2.weight": "pytorch_model-00001-of-00002.bin",
48
+ "transformer.blocks.15.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
49
+ "transformer.blocks.15.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
50
+ "transformer.blocks.15.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
51
+ "transformer.blocks.15.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
52
+ "transformer.blocks.15.norm_1.weight": "pytorch_model-00001-of-00002.bin",
53
+ "transformer.blocks.15.norm_2.weight": "pytorch_model-00001-of-00002.bin",
54
+ "transformer.blocks.16.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
55
+ "transformer.blocks.16.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
56
+ "transformer.blocks.16.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
57
+ "transformer.blocks.16.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
58
+ "transformer.blocks.16.norm_1.weight": "pytorch_model-00001-of-00002.bin",
59
+ "transformer.blocks.16.norm_2.weight": "pytorch_model-00001-of-00002.bin",
60
+ "transformer.blocks.17.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
61
+ "transformer.blocks.17.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
62
+ "transformer.blocks.17.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
63
+ "transformer.blocks.17.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
64
+ "transformer.blocks.17.norm_1.weight": "pytorch_model-00001-of-00002.bin",
65
+ "transformer.blocks.17.norm_2.weight": "pytorch_model-00001-of-00002.bin",
66
+ "transformer.blocks.18.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
67
+ "transformer.blocks.18.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
68
+ "transformer.blocks.18.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
69
+ "transformer.blocks.18.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
70
+ "transformer.blocks.18.norm_1.weight": "pytorch_model-00001-of-00002.bin",
71
+ "transformer.blocks.18.norm_2.weight": "pytorch_model-00001-of-00002.bin",
72
+ "transformer.blocks.19.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
73
+ "transformer.blocks.19.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
74
+ "transformer.blocks.19.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
75
+ "transformer.blocks.19.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
76
+ "transformer.blocks.19.norm_1.weight": "pytorch_model-00001-of-00002.bin",
77
+ "transformer.blocks.19.norm_2.weight": "pytorch_model-00001-of-00002.bin",
78
+ "transformer.blocks.2.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
79
+ "transformer.blocks.2.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
80
+ "transformer.blocks.2.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
81
+ "transformer.blocks.2.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
82
+ "transformer.blocks.2.norm_1.weight": "pytorch_model-00001-of-00002.bin",
83
+ "transformer.blocks.2.norm_2.weight": "pytorch_model-00001-of-00002.bin",
84
+ "transformer.blocks.20.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
85
+ "transformer.blocks.20.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
86
+ "transformer.blocks.20.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
87
+ "transformer.blocks.20.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
88
+ "transformer.blocks.20.norm_1.weight": "pytorch_model-00001-of-00002.bin",
89
+ "transformer.blocks.20.norm_2.weight": "pytorch_model-00001-of-00002.bin",
90
+ "transformer.blocks.21.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
91
+ "transformer.blocks.21.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
92
+ "transformer.blocks.21.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
93
+ "transformer.blocks.21.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
94
+ "transformer.blocks.21.norm_1.weight": "pytorch_model-00001-of-00002.bin",
95
+ "transformer.blocks.21.norm_2.weight": "pytorch_model-00001-of-00002.bin",
96
+ "transformer.blocks.22.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
97
+ "transformer.blocks.22.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
98
+ "transformer.blocks.22.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
99
+ "transformer.blocks.22.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
100
+ "transformer.blocks.22.norm_1.weight": "pytorch_model-00001-of-00002.bin",
101
+ "transformer.blocks.22.norm_2.weight": "pytorch_model-00001-of-00002.bin",
102
+ "transformer.blocks.23.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
103
+ "transformer.blocks.23.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
104
+ "transformer.blocks.23.ffn.down_proj.weight": "pytorch_model-00002-of-00002.bin",
105
+ "transformer.blocks.23.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
106
+ "transformer.blocks.23.norm_1.weight": "pytorch_model-00001-of-00002.bin",
107
+ "transformer.blocks.23.norm_2.weight": "pytorch_model-00001-of-00002.bin",
108
+ "transformer.blocks.24.attn.Wqkv.weight": "pytorch_model-00002-of-00002.bin",
109
+ "transformer.blocks.24.attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
110
+ "transformer.blocks.24.ffn.down_proj.weight": "pytorch_model-00002-of-00002.bin",
111
+ "transformer.blocks.24.ffn.up_proj.weight": "pytorch_model-00002-of-00002.bin",
112
+ "transformer.blocks.24.norm_1.weight": "pytorch_model-00002-of-00002.bin",
113
+ "transformer.blocks.24.norm_2.weight": "pytorch_model-00002-of-00002.bin",
114
+ "transformer.blocks.25.attn.Wqkv.weight": "pytorch_model-00002-of-00002.bin",
115
+ "transformer.blocks.25.attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
116
+ "transformer.blocks.25.ffn.down_proj.weight": "pytorch_model-00002-of-00002.bin",
117
+ "transformer.blocks.25.ffn.up_proj.weight": "pytorch_model-00002-of-00002.bin",
118
+ "transformer.blocks.25.norm_1.weight": "pytorch_model-00002-of-00002.bin",
119
+ "transformer.blocks.25.norm_2.weight": "pytorch_model-00002-of-00002.bin",
120
+ "transformer.blocks.26.attn.Wqkv.weight": "pytorch_model-00002-of-00002.bin",
121
+ "transformer.blocks.26.attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
122
+ "transformer.blocks.26.ffn.down_proj.weight": "pytorch_model-00002-of-00002.bin",
123
+ "transformer.blocks.26.ffn.up_proj.weight": "pytorch_model-00002-of-00002.bin",
124
+ "transformer.blocks.26.norm_1.weight": "pytorch_model-00002-of-00002.bin",
125
+ "transformer.blocks.26.norm_2.weight": "pytorch_model-00002-of-00002.bin",
126
+ "transformer.blocks.27.attn.Wqkv.weight": "pytorch_model-00002-of-00002.bin",
127
+ "transformer.blocks.27.attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
128
+ "transformer.blocks.27.ffn.down_proj.weight": "pytorch_model-00002-of-00002.bin",
129
+ "transformer.blocks.27.ffn.up_proj.weight": "pytorch_model-00002-of-00002.bin",
130
+ "transformer.blocks.27.norm_1.weight": "pytorch_model-00002-of-00002.bin",
131
+ "transformer.blocks.27.norm_2.weight": "pytorch_model-00002-of-00002.bin",
132
+ "transformer.blocks.28.attn.Wqkv.weight": "pytorch_model-00002-of-00002.bin",
133
+ "transformer.blocks.28.attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
134
+ "transformer.blocks.28.ffn.down_proj.weight": "pytorch_model-00002-of-00002.bin",
135
+ "transformer.blocks.28.ffn.up_proj.weight": "pytorch_model-00002-of-00002.bin",
136
+ "transformer.blocks.28.norm_1.weight": "pytorch_model-00002-of-00002.bin",
137
+ "transformer.blocks.28.norm_2.weight": "pytorch_model-00002-of-00002.bin",
138
+ "transformer.blocks.29.attn.Wqkv.weight": "pytorch_model-00002-of-00002.bin",
139
+ "transformer.blocks.29.attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
140
+ "transformer.blocks.29.ffn.down_proj.weight": "pytorch_model-00002-of-00002.bin",
141
+ "transformer.blocks.29.ffn.up_proj.weight": "pytorch_model-00002-of-00002.bin",
142
+ "transformer.blocks.29.norm_1.weight": "pytorch_model-00002-of-00002.bin",
143
+ "transformer.blocks.29.norm_2.weight": "pytorch_model-00002-of-00002.bin",
144
+ "transformer.blocks.3.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
145
+ "transformer.blocks.3.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
146
+ "transformer.blocks.3.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
147
+ "transformer.blocks.3.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
148
+ "transformer.blocks.3.norm_1.weight": "pytorch_model-00001-of-00002.bin",
149
+ "transformer.blocks.3.norm_2.weight": "pytorch_model-00001-of-00002.bin",
150
+ "transformer.blocks.30.attn.Wqkv.weight": "pytorch_model-00002-of-00002.bin",
151
+ "transformer.blocks.30.attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
152
+ "transformer.blocks.30.ffn.down_proj.weight": "pytorch_model-00002-of-00002.bin",
153
+ "transformer.blocks.30.ffn.up_proj.weight": "pytorch_model-00002-of-00002.bin",
154
+ "transformer.blocks.30.norm_1.weight": "pytorch_model-00002-of-00002.bin",
155
+ "transformer.blocks.30.norm_2.weight": "pytorch_model-00002-of-00002.bin",
156
+ "transformer.blocks.31.attn.Wqkv.weight": "pytorch_model-00002-of-00002.bin",
157
+ "transformer.blocks.31.attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
158
+ "transformer.blocks.31.ffn.down_proj.weight": "pytorch_model-00002-of-00002.bin",
159
+ "transformer.blocks.31.ffn.up_proj.weight": "pytorch_model-00002-of-00002.bin",
160
+ "transformer.blocks.31.norm_1.weight": "pytorch_model-00002-of-00002.bin",
161
+ "transformer.blocks.31.norm_2.weight": "pytorch_model-00002-of-00002.bin",
162
+ "transformer.blocks.4.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
163
+ "transformer.blocks.4.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
164
+ "transformer.blocks.4.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
165
+ "transformer.blocks.4.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
166
+ "transformer.blocks.4.norm_1.weight": "pytorch_model-00001-of-00002.bin",
167
+ "transformer.blocks.4.norm_2.weight": "pytorch_model-00001-of-00002.bin",
168
+ "transformer.blocks.5.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
169
+ "transformer.blocks.5.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
170
+ "transformer.blocks.5.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
171
+ "transformer.blocks.5.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
172
+ "transformer.blocks.5.norm_1.weight": "pytorch_model-00001-of-00002.bin",
173
+ "transformer.blocks.5.norm_2.weight": "pytorch_model-00001-of-00002.bin",
174
+ "transformer.blocks.6.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
175
+ "transformer.blocks.6.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
176
+ "transformer.blocks.6.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
177
+ "transformer.blocks.6.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
178
+ "transformer.blocks.6.norm_1.weight": "pytorch_model-00001-of-00002.bin",
179
+ "transformer.blocks.6.norm_2.weight": "pytorch_model-00001-of-00002.bin",
180
+ "transformer.blocks.7.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
181
+ "transformer.blocks.7.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
182
+ "transformer.blocks.7.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
183
+ "transformer.blocks.7.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
184
+ "transformer.blocks.7.norm_1.weight": "pytorch_model-00001-of-00002.bin",
185
+ "transformer.blocks.7.norm_2.weight": "pytorch_model-00001-of-00002.bin",
186
+ "transformer.blocks.8.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
187
+ "transformer.blocks.8.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
188
+ "transformer.blocks.8.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
189
+ "transformer.blocks.8.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
190
+ "transformer.blocks.8.norm_1.weight": "pytorch_model-00001-of-00002.bin",
191
+ "transformer.blocks.8.norm_2.weight": "pytorch_model-00001-of-00002.bin",
192
+ "transformer.blocks.9.attn.Wqkv.weight": "pytorch_model-00001-of-00002.bin",
193
+ "transformer.blocks.9.attn.out_proj.weight": "pytorch_model-00001-of-00002.bin",
194
+ "transformer.blocks.9.ffn.down_proj.weight": "pytorch_model-00001-of-00002.bin",
195
+ "transformer.blocks.9.ffn.up_proj.weight": "pytorch_model-00001-of-00002.bin",
196
+ "transformer.blocks.9.norm_1.weight": "pytorch_model-00001-of-00002.bin",
197
+ "transformer.blocks.9.norm_2.weight": "pytorch_model-00001-of-00002.bin",
198
+ "transformer.norm_f.weight": "pytorch_model-00002-of-00002.bin",
199
+ "transformer.wte.weight": "pytorch_model-00001-of-00002.bin"
200
+ }
201
+ }
special_tokens_map.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": "<|endoftext|>",
3
+ "eos_token": "<|endoftext|>",
4
+ "pad_token": "<|endoftext|>",
5
+ "unk_token": "<|endoftext|>"
6
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "bos_token": "<|endoftext|>",
4
+ "clean_up_tokenization_spaces": true,
5
+ "eos_token": "<|endoftext|>",
6
+ "model_max_length": 8192,
7
+ "tokenizer_class": "GPTNeoXTokenizer",
8
+ "unk_token": "<|endoftext|>"
9
+ }