Update attention.py
Browse files- attention.py +131 -93
attention.py
CHANGED
@@ -1,15 +1,30 @@
|
|
1 |
"""Attention layers."""
|
2 |
import math
|
3 |
import warnings
|
4 |
-
from typing import Optional
|
5 |
import torch
|
6 |
import torch.nn as nn
|
7 |
from einops import rearrange
|
8 |
from packaging import version
|
9 |
from torch import nn
|
10 |
-
from .
|
|
|
11 |
|
12 |
-
def
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
13 |
if original_is_causal and num_query_tokens != num_key_tokens:
|
14 |
if num_query_tokens != 1:
|
15 |
raise NotImplementedError('MPT does not support query and key with different number of tokens, unless number of query tokens is 1.')
|
@@ -17,9 +32,27 @@ def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_cau
|
|
17 |
return False
|
18 |
return original_is_causal
|
19 |
|
20 |
-
def
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
21 |
q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)
|
22 |
-
kv_n_heads = 1 if multiquery else n_heads
|
23 |
k = rearrange(key, 'b s (h d) -> b h d s', h=kv_n_heads)
|
24 |
v = rearrange(value, 'b s (h d) -> b h s d', h=kv_n_heads)
|
25 |
if past_key_value is not None:
|
@@ -29,6 +62,9 @@ def scaled_multihead_dot_product_attention(query, key, value, n_heads, past_key_
|
|
29 |
past_key_value = (k, v)
|
30 |
(b, _, s_q, d) = q.shape
|
31 |
s_k = k.size(-1)
|
|
|
|
|
|
|
32 |
if softmax_scale is None:
|
33 |
softmax_scale = 1 / math.sqrt(d)
|
34 |
attn_weight = q.matmul(k) * softmax_scale
|
@@ -42,11 +78,11 @@ def scaled_multihead_dot_product_attention(query, key, value, n_heads, past_key_
|
|
42 |
min_val = torch.finfo(q.dtype).min
|
43 |
if key_padding_mask is not None:
|
44 |
if attn_bias is not None:
|
45 |
-
warnings.warn('
|
46 |
attn_weight = attn_weight.masked_fill(~key_padding_mask.view((b, 1, 1, s_k)), min_val)
|
47 |
if is_causal and (not q.size(2) == 1):
|
48 |
s = max(s_q, s_k)
|
49 |
-
causal_mask = attn_weight.new_ones(s, s, dtype=torch.
|
50 |
causal_mask = causal_mask.tril()
|
51 |
causal_mask = causal_mask.to(torch.bool)
|
52 |
causal_mask = ~causal_mask
|
@@ -61,19 +97,27 @@ def scaled_multihead_dot_product_attention(query, key, value, n_heads, past_key_
|
|
61 |
return (out, attn_weight, past_key_value)
|
62 |
return (out, None, past_key_value)
|
63 |
|
64 |
-
def check_valid_inputs(*tensors
|
|
|
|
|
65 |
for tensor in tensors:
|
66 |
if tensor.dtype not in valid_dtypes:
|
67 |
raise TypeError(f'tensor.dtype={tensor.dtype!r} must be in valid_dtypes={valid_dtypes!r}.')
|
68 |
if not tensor.is_cuda:
|
69 |
raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
|
70 |
|
71 |
-
def flash_attn_fn(query, key, value, n_heads, past_key_value=None, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
|
72 |
try:
|
73 |
from flash_attn import bert_padding, flash_attn_interface
|
74 |
except:
|
75 |
-
raise RuntimeError('Please install flash-attn==1.0.3.
|
76 |
check_valid_inputs(query, key, value)
|
|
|
|
|
|
|
|
|
|
|
|
|
77 |
if past_key_value is not None:
|
78 |
if len(past_key_value) != 0:
|
79 |
key = torch.cat([past_key_value[0], key], dim=1)
|
@@ -92,19 +136,27 @@ def flash_attn_fn(query, key, value, n_heads, past_key_value=None, softmax_scale
|
|
92 |
(query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(query, query_padding_mask)
|
93 |
query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
|
94 |
(key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(key, key_padding_mask)
|
95 |
-
key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=
|
96 |
(value_unpad, _, _, _) = bert_padding.unpad_input(value, key_padding_mask)
|
97 |
-
value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=
|
98 |
-
if
|
99 |
key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size(-1))
|
100 |
value_unpad = value_unpad.expand(value_unpad.size(0), n_heads, value_unpad.size(-1))
|
|
|
|
|
|
|
101 |
dropout_p = dropout_p if training else 0.0
|
102 |
reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
|
103 |
-
|
|
|
|
|
|
|
|
|
|
|
104 |
output = bert_padding.pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size, seqlen)
|
105 |
return (output, None, past_key_value)
|
106 |
|
107 |
-
def triton_flash_attn_fn(query, key, value, n_heads, past_key_value=None, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
|
108 |
try:
|
109 |
from .flash_attn_triton import flash_attn_func
|
110 |
except:
|
@@ -116,8 +168,14 @@ def triton_flash_attn_fn(query, key, value, n_heads, past_key_value=None, softma
|
|
116 |
except:
|
117 |
_installed = False
|
118 |
if not _installed:
|
119 |
-
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.')
|
120 |
check_valid_inputs(query, key, value)
|
|
|
|
|
|
|
|
|
|
|
|
|
121 |
if past_key_value is not None:
|
122 |
if len(past_key_value) != 0:
|
123 |
key = torch.cat([past_key_value[0], key], dim=1)
|
@@ -129,6 +187,7 @@ def triton_flash_attn_fn(query, key, value, n_heads, past_key_value=None, softma
|
|
129 |
attn_bias = attn_bias[:, :, _s_q:, _s_k:]
|
130 |
if dropout_p:
|
131 |
raise NotImplementedError(f'Dropout not implemented for attn_impl: triton.')
|
|
|
132 |
if needs_weights:
|
133 |
raise NotImplementedError(f'attn_impl: triton cannot return attn weights.')
|
134 |
if key_padding_mask is not None:
|
@@ -138,124 +197,103 @@ def triton_flash_attn_fn(query, key, value, n_heads, past_key_value=None, softma
|
|
138 |
attn_bias = query.new_zeros(b_size, 1, 1, s_k)
|
139 |
attn_bias = attn_bias.masked_fill(~key_padding_mask.view((b_size, 1, 1, s_k)), torch.finfo(query.dtype).min)
|
140 |
query = rearrange(query, 'b s (h d) -> b s h d', h=n_heads)
|
141 |
-
key = rearrange(key, 'b s (h d) -> b s h d', h=
|
142 |
-
value = rearrange(value, 'b s (h d) -> b s h d', h=
|
143 |
-
if
|
144 |
-
key = key.
|
145 |
-
value = value.
|
|
|
|
|
|
|
146 |
reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
|
147 |
attn_output = flash_attn_func(query, key, value, attn_bias, reset_is_causal, softmax_scale)
|
148 |
output = attn_output.view(*attn_output.shape[:2], -1)
|
149 |
return (output, None, past_key_value)
|
150 |
|
151 |
-
class
|
152 |
-
"""Multi-head
|
153 |
|
154 |
-
|
155 |
-
|
|
|
|
|
|
|
156 |
"""
|
157 |
|
158 |
-
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,
|
159 |
super().__init__()
|
160 |
self.attn_impl = attn_impl
|
161 |
self.clip_qkv = clip_qkv
|
162 |
self.qk_ln = qk_ln
|
163 |
self.d_model = d_model
|
164 |
self.n_heads = n_heads
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
165 |
self.softmax_scale = softmax_scale
|
166 |
if self.softmax_scale is None:
|
167 |
self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
|
168 |
self.attn_dropout_p = attn_pdrop
|
169 |
-
|
170 |
-
|
|
|
|
|
|
|
171 |
self.Wqkv._fused = (0, fuse_splits)
|
172 |
if self.qk_ln:
|
173 |
-
|
174 |
-
self.q_ln =
|
175 |
-
self.k_ln =
|
176 |
if self.attn_impl == 'flash':
|
177 |
self.attn_fn = flash_attn_fn
|
178 |
elif self.attn_impl == 'triton':
|
179 |
self.attn_fn = triton_flash_attn_fn
|
180 |
-
if verbose:
|
181 |
-
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`.')
|
182 |
elif self.attn_impl == 'torch':
|
183 |
self.attn_fn = scaled_multihead_dot_product_attention
|
184 |
-
if torch.cuda.is_available() and verbose:
|
185 |
-
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`.')
|
186 |
else:
|
187 |
raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
|
188 |
-
self.out_proj =
|
189 |
self.out_proj._is_residual = True
|
190 |
|
191 |
-
def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
|
192 |
qkv = self.Wqkv(x)
|
193 |
if self.clip_qkv:
|
194 |
-
qkv.
|
195 |
-
(query, key, value) = qkv.
|
196 |
key_padding_mask = attention_mask
|
197 |
if self.qk_ln:
|
198 |
dtype = query.dtype
|
199 |
query = self.q_ln(query).to(dtype)
|
200 |
key = self.k_ln(key).to(dtype)
|
201 |
-
(context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, past_key_value=past_key_value, 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)
|
202 |
return (self.out_proj(context), attn_weights, past_key_value)
|
203 |
|
204 |
-
class
|
205 |
-
"""Multi-
|
206 |
|
207 |
-
Using torch or triton attention
|
208 |
additive bias.
|
209 |
"""
|
210 |
|
211 |
-
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,
|
212 |
-
super().__init__()
|
213 |
-
self.attn_impl = attn_impl
|
214 |
-
self.clip_qkv = clip_qkv
|
215 |
-
self.qk_ln = qk_ln
|
216 |
-
self.d_model = d_model
|
217 |
-
self.n_heads = n_heads
|
218 |
-
self.head_dim = d_model // n_heads
|
219 |
-
self.softmax_scale = softmax_scale
|
220 |
-
if self.softmax_scale is None:
|
221 |
-
self.softmax_scale = 1 / math.sqrt(self.head_dim)
|
222 |
-
self.attn_dropout_p = attn_pdrop
|
223 |
-
self.Wqkv = nn.Linear(d_model, d_model + 2 * self.head_dim, device=device)
|
224 |
-
fuse_splits = (d_model, d_model + self.head_dim)
|
225 |
-
self.Wqkv._fused = (0, fuse_splits)
|
226 |
-
if self.qk_ln:
|
227 |
-
layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
|
228 |
-
self.q_ln = layernorm_class(d_model, device=device)
|
229 |
-
self.k_ln = layernorm_class(self.head_dim, device=device)
|
230 |
-
if self.attn_impl == 'flash':
|
231 |
-
self.attn_fn = flash_attn_fn
|
232 |
-
elif self.attn_impl == 'triton':
|
233 |
-
self.attn_fn = triton_flash_attn_fn
|
234 |
-
if verbose:
|
235 |
-
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`.')
|
236 |
-
elif self.attn_impl == 'torch':
|
237 |
-
self.attn_fn = scaled_multihead_dot_product_attention
|
238 |
-
if torch.cuda.is_available() and verbose:
|
239 |
-
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`.')
|
240 |
-
else:
|
241 |
-
raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
|
242 |
-
self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
|
243 |
-
self.out_proj._is_residual = True
|
244 |
|
245 |
-
|
246 |
-
|
247 |
-
|
248 |
-
|
249 |
-
|
250 |
-
|
251 |
-
|
252 |
-
|
253 |
-
|
254 |
-
key = self.k_ln(key).to(dtype)
|
255 |
-
(context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, past_key_value=past_key_value, 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)
|
256 |
-
return (self.out_proj(context), attn_weights, past_key_value)
|
257 |
|
258 |
-
def attn_bias_shape(attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_sequence_id):
|
259 |
if attn_impl == 'flash':
|
260 |
return None
|
261 |
elif attn_impl in ['torch', 'triton']:
|
@@ -269,7 +307,7 @@ def attn_bias_shape(attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_s
|
|
269 |
else:
|
270 |
raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
|
271 |
|
272 |
-
def build_attn_bias(attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=False, alibi_bias_max=8):
|
273 |
if attn_impl == 'flash':
|
274 |
return None
|
275 |
elif attn_impl in ['torch', 'triton']:
|
@@ -280,7 +318,7 @@ def build_attn_bias(attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=
|
|
280 |
else:
|
281 |
raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
|
282 |
|
283 |
-
def gen_slopes(n_heads, alibi_bias_max=8, device=None):
|
284 |
_n_heads = 2 ** math.ceil(math.log2(n_heads))
|
285 |
m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)
|
286 |
m = m.mul(alibi_bias_max / _n_heads)
|
@@ -289,7 +327,7 @@ def gen_slopes(n_heads, alibi_bias_max=8, device=None):
|
|
289 |
slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
|
290 |
return slopes.view(1, n_heads, 1, 1)
|
291 |
|
292 |
-
def build_alibi_bias(n_heads, seq_len, full=False, alibi_bias_max=8, device=None, dtype=None):
|
293 |
alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, 1, seq_len)
|
294 |
if full:
|
295 |
alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, seq_len, 1)
|
@@ -297,4 +335,4 @@ def build_alibi_bias(n_heads, seq_len, full=False, alibi_bias_max=8, device=None
|
|
297 |
slopes = gen_slopes(n_heads, alibi_bias_max, device=device)
|
298 |
alibi_bias = alibi_bias * slopes
|
299 |
return alibi_bias.to(dtype=dtype)
|
300 |
-
ATTN_CLASS_REGISTRY = {'multihead_attention': MultiheadAttention, 'multiquery_attention': MultiQueryAttention}
|
|
|
1 |
"""Attention layers."""
|
2 |
import math
|
3 |
import warnings
|
4 |
+
from typing import Any, List, Optional, Tuple
|
5 |
import torch
|
6 |
import torch.nn as nn
|
7 |
from einops import rearrange
|
8 |
from packaging import version
|
9 |
from torch import nn
|
10 |
+
from .fc import FC_CLASS_REGISTRY
|
11 |
+
from .norm import NORM_CLASS_REGISTRY
|
12 |
|
13 |
+
def is_flash_v2_installed():
|
14 |
+
try:
|
15 |
+
import flash_attn as flash_attn
|
16 |
+
except:
|
17 |
+
return False
|
18 |
+
return version.parse(flash_attn.__version__) >= version.parse('2.0.0')
|
19 |
+
|
20 |
+
def is_flash_v1_installed():
|
21 |
+
try:
|
22 |
+
import flash_attn as flash_attn
|
23 |
+
except:
|
24 |
+
return False
|
25 |
+
return version.parse(flash_attn.__version__) < version.parse('2.0.0')
|
26 |
+
|
27 |
+
def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool) -> bool:
|
28 |
if original_is_causal and num_query_tokens != num_key_tokens:
|
29 |
if num_query_tokens != 1:
|
30 |
raise NotImplementedError('MPT does not support query and key with different number of tokens, unless number of query tokens is 1.')
|
|
|
32 |
return False
|
33 |
return original_is_causal
|
34 |
|
35 |
+
def repeat_kv_for_gqa(hidden: torch.Tensor, n_rep: int) -> torch.Tensor:
|
36 |
+
"""Perform repeat of kv heads along a particular dimension.
|
37 |
+
|
38 |
+
hidden.shape expected to be: (batch size, seq len, kv_n_heads, head_dim)
|
39 |
+
n_rep: amount of repetitions of kv_n_heads
|
40 |
+
Unlike torch.repeat_interleave, this function avoids allocating new memory.
|
41 |
+
"""
|
42 |
+
if n_rep == 1:
|
43 |
+
return hidden
|
44 |
+
(b, s, kv_n_heads, d) = hidden.shape
|
45 |
+
hidden = hidden[:, :, :, None, :].expand(b, s, kv_n_heads, n_rep, d)
|
46 |
+
return hidden.reshape(b, s, kv_n_heads * n_rep, d)
|
47 |
+
|
48 |
+
def scaled_multihead_dot_product_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: Optional[int]=None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False, multiquery: bool=False) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
|
49 |
+
if multiquery:
|
50 |
+
warnings.warn(DeprecationWarning('The direct use of the multiquery arg is deprecated. Setting kv_n_heads=1 automatically. Please set kv_n_heads=1 explicitly to remove this warning.'))
|
51 |
+
kv_n_heads = 1
|
52 |
+
elif kv_n_heads is None:
|
53 |
+
warnings.warn(DeprecationWarning('Not specifying a value for the kv_n_heads arg is deprecated. Setting kv_n_heads=n_heads automatically. Please set kv_n_heads=n_heads explicitly to remove this warning.'))
|
54 |
+
kv_n_heads = n_heads
|
55 |
q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)
|
|
|
56 |
k = rearrange(key, 'b s (h d) -> b h d s', h=kv_n_heads)
|
57 |
v = rearrange(value, 'b s (h d) -> b h s d', h=kv_n_heads)
|
58 |
if past_key_value is not None:
|
|
|
62 |
past_key_value = (k, v)
|
63 |
(b, _, s_q, d) = q.shape
|
64 |
s_k = k.size(-1)
|
65 |
+
if kv_n_heads > 1 and kv_n_heads < n_heads:
|
66 |
+
k = repeat_kv_for_gqa(k.transpose(1, 2), n_heads // kv_n_heads).transpose(1, 2)
|
67 |
+
v = repeat_kv_for_gqa(v.transpose(1, 2), n_heads // kv_n_heads).transpose(1, 2)
|
68 |
if softmax_scale is None:
|
69 |
softmax_scale = 1 / math.sqrt(d)
|
70 |
attn_weight = q.matmul(k) * softmax_scale
|
|
|
78 |
min_val = torch.finfo(q.dtype).min
|
79 |
if key_padding_mask is not None:
|
80 |
if attn_bias is not None:
|
81 |
+
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.')
|
82 |
attn_weight = attn_weight.masked_fill(~key_padding_mask.view((b, 1, 1, s_k)), min_val)
|
83 |
if is_causal and (not q.size(2) == 1):
|
84 |
s = max(s_q, s_k)
|
85 |
+
causal_mask = attn_weight.new_ones(s, s, dtype=torch.float32)
|
86 |
causal_mask = causal_mask.tril()
|
87 |
causal_mask = causal_mask.to(torch.bool)
|
88 |
causal_mask = ~causal_mask
|
|
|
97 |
return (out, attn_weight, past_key_value)
|
98 |
return (out, None, past_key_value)
|
99 |
|
100 |
+
def check_valid_inputs(*tensors: torch.Tensor, valid_dtypes: Optional[List[torch.dtype]]=None):
|
101 |
+
if valid_dtypes is None:
|
102 |
+
valid_dtypes = [torch.float16, torch.bfloat16]
|
103 |
for tensor in tensors:
|
104 |
if tensor.dtype not in valid_dtypes:
|
105 |
raise TypeError(f'tensor.dtype={tensor.dtype!r} must be in valid_dtypes={valid_dtypes!r}.')
|
106 |
if not tensor.is_cuda:
|
107 |
raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
|
108 |
|
109 |
+
def flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: Optional[int]=None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False, multiquery: bool=False) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
|
110 |
try:
|
111 |
from flash_attn import bert_padding, flash_attn_interface
|
112 |
except:
|
113 |
+
raise RuntimeError('Please install flash-attn==1.0.9 or flash-attn==2.3.2')
|
114 |
check_valid_inputs(query, key, value)
|
115 |
+
if multiquery:
|
116 |
+
warnings.warn(DeprecationWarning('The direct use of the multiquery arg is deprecated. Setting kv_n_heads=1 automatically. Please set kv_n_heads=1 explicitly to remove this warning.'))
|
117 |
+
kv_n_heads = 1
|
118 |
+
elif kv_n_heads is None:
|
119 |
+
warnings.warn(DeprecationWarning('Not specifying a value for the kv_n_heads arg is deprecated. Setting kv_n_heads=n_heads automatically. Please set kv_n_heads=n_heads explicitly to remove this warning.'))
|
120 |
+
kv_n_heads = n_heads
|
121 |
if past_key_value is not None:
|
122 |
if len(past_key_value) != 0:
|
123 |
key = torch.cat([past_key_value[0], key], dim=1)
|
|
|
136 |
(query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(query, query_padding_mask)
|
137 |
query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
|
138 |
(key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(key, key_padding_mask)
|
139 |
+
key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=kv_n_heads)
|
140 |
(value_unpad, _, _, _) = bert_padding.unpad_input(value, key_padding_mask)
|
141 |
+
value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=kv_n_heads)
|
142 |
+
if kv_n_heads == 1:
|
143 |
key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size(-1))
|
144 |
value_unpad = value_unpad.expand(value_unpad.size(0), n_heads, value_unpad.size(-1))
|
145 |
+
elif kv_n_heads < n_heads:
|
146 |
+
key_unpad = repeat_kv_for_gqa(key_unpad.view(batch_size, seqlen, kv_n_heads, -1), n_heads // kv_n_heads).view(batch_size * seqlen, n_heads, -1)
|
147 |
+
value_unpad = repeat_kv_for_gqa(value_unpad.view(batch_size, seqlen, kv_n_heads, -1), n_heads // kv_n_heads).view(batch_size * seqlen, n_heads, -1)
|
148 |
dropout_p = dropout_p if training else 0.0
|
149 |
reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
|
150 |
+
if is_flash_v1_installed():
|
151 |
+
output_unpad = flash_attn_interface.flash_attn_unpadded_func(q=query_unpad, k=key_unpad, v=value_unpad, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
|
152 |
+
elif is_flash_v2_installed():
|
153 |
+
output_unpad = flash_attn_interface.flash_attn_varlen_func(q=query_unpad, k=key_unpad, v=value_unpad, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
|
154 |
+
else:
|
155 |
+
raise RuntimeError('flash-attn==1.0.9 or flash-attn==2.3.2 is required.')
|
156 |
output = bert_padding.pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size, seqlen)
|
157 |
return (output, None, past_key_value)
|
158 |
|
159 |
+
def triton_flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: Optional[int]=None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False, multiquery: bool=False) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
|
160 |
try:
|
161 |
from .flash_attn_triton import flash_attn_func
|
162 |
except:
|
|
|
168 |
except:
|
169 |
_installed = False
|
170 |
if not _installed:
|
171 |
+
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.')
|
172 |
check_valid_inputs(query, key, value)
|
173 |
+
if multiquery:
|
174 |
+
warnings.warn(DeprecationWarning('The direct use of the multiquery arg is deprecated. Setting kv_n_heads=1 automatically. Please set kv_n_heads=1 explicitly to remove this warning.'))
|
175 |
+
kv_n_heads = 1
|
176 |
+
elif kv_n_heads is None:
|
177 |
+
warnings.warn(DeprecationWarning('Not specifying a value for the kv_n_heads arg is deprecated. Setting kv_n_heads=n_heads automatically. Please set kv_n_heads=n_heads explicitly to remove this warning.'))
|
178 |
+
kv_n_heads = n_heads
|
179 |
if past_key_value is not None:
|
180 |
if len(past_key_value) != 0:
|
181 |
key = torch.cat([past_key_value[0], key], dim=1)
|
|
|
187 |
attn_bias = attn_bias[:, :, _s_q:, _s_k:]
|
188 |
if dropout_p:
|
189 |
raise NotImplementedError(f'Dropout not implemented for attn_impl: triton.')
|
190 |
+
dropout_p = dropout_p if training else 0.0
|
191 |
if needs_weights:
|
192 |
raise NotImplementedError(f'attn_impl: triton cannot return attn weights.')
|
193 |
if key_padding_mask is not None:
|
|
|
197 |
attn_bias = query.new_zeros(b_size, 1, 1, s_k)
|
198 |
attn_bias = attn_bias.masked_fill(~key_padding_mask.view((b_size, 1, 1, s_k)), torch.finfo(query.dtype).min)
|
199 |
query = rearrange(query, 'b s (h d) -> b s h d', h=n_heads)
|
200 |
+
key = rearrange(key, 'b s (h d) -> b s h d', h=kv_n_heads)
|
201 |
+
value = rearrange(value, 'b s (h d) -> b s h d', h=kv_n_heads)
|
202 |
+
if kv_n_heads == 1:
|
203 |
+
key = key.repeat(1, 1, n_heads, 1)
|
204 |
+
value = value.repeat(1, 1, n_heads, 1)
|
205 |
+
elif kv_n_heads < n_heads:
|
206 |
+
key = repeat_kv_for_gqa(key, n_heads // kv_n_heads)
|
207 |
+
value = repeat_kv_for_gqa(value, n_heads // kv_n_heads)
|
208 |
reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
|
209 |
attn_output = flash_attn_func(query, key, value, attn_bias, reset_is_causal, softmax_scale)
|
210 |
output = attn_output.view(*attn_output.shape[:2], -1)
|
211 |
return (output, None, past_key_value)
|
212 |
|
213 |
+
class GroupedQueryAttention(nn.Module):
|
214 |
+
"""Grouped Query Attention (GQA) is a generalization of Multi-head (MHA).
|
215 |
|
216 |
+
and Multi-query attention (MQA).
|
217 |
+
|
218 |
+
This allows the user to set a variable of number of kv_n_heads, rather than
|
219 |
+
just n_heads or 1, as in MHA and MQA. Using torch or triton attention
|
220 |
+
implementation enables user to also use additive bias.
|
221 |
"""
|
222 |
|
223 |
+
def __init__(self, d_model: int, n_heads: int, kv_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, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True):
|
224 |
super().__init__()
|
225 |
self.attn_impl = attn_impl
|
226 |
self.clip_qkv = clip_qkv
|
227 |
self.qk_ln = qk_ln
|
228 |
self.d_model = d_model
|
229 |
self.n_heads = n_heads
|
230 |
+
self.kv_n_heads = kv_n_heads
|
231 |
+
self.head_dim = d_model // n_heads
|
232 |
+
if self.kv_n_heads <= 0:
|
233 |
+
raise ValueError('kv_n_heads should be greater than zero.')
|
234 |
+
if self.kv_n_heads > self.n_heads:
|
235 |
+
raise ValueError('The number of KV heads should be less than or equal to Q heads.')
|
236 |
+
if self.n_heads % self.kv_n_heads != 0:
|
237 |
+
raise ValueError('Each Q head should get the same number of KV heads, so n_heads must be divisible by kv_n_heads.')
|
238 |
self.softmax_scale = softmax_scale
|
239 |
if self.softmax_scale is None:
|
240 |
self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
|
241 |
self.attn_dropout_p = attn_pdrop
|
242 |
+
fc_kwargs: dict[str, Any] = {'bias': bias}
|
243 |
+
if fc_type != 'te':
|
244 |
+
fc_kwargs['device'] = device
|
245 |
+
self.Wqkv = FC_CLASS_REGISTRY[fc_type](self.d_model, self.d_model + 2 * self.kv_n_heads * self.head_dim, **fc_kwargs)
|
246 |
+
fuse_splits = [i * self.head_dim for i in range(1, self.n_heads + 2 * self.kv_n_heads)]
|
247 |
self.Wqkv._fused = (0, fuse_splits)
|
248 |
if self.qk_ln:
|
249 |
+
norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
|
250 |
+
self.q_ln = norm_class(self.d_model, device=device)
|
251 |
+
self.k_ln = norm_class(self.kv_n_heads * self.head_dim, device=device)
|
252 |
if self.attn_impl == 'flash':
|
253 |
self.attn_fn = flash_attn_fn
|
254 |
elif self.attn_impl == 'triton':
|
255 |
self.attn_fn = triton_flash_attn_fn
|
|
|
|
|
256 |
elif self.attn_impl == 'torch':
|
257 |
self.attn_fn = scaled_multihead_dot_product_attention
|
|
|
|
|
258 |
else:
|
259 |
raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
|
260 |
+
self.out_proj = FC_CLASS_REGISTRY[fc_type](self.d_model, self.d_model, **fc_kwargs)
|
261 |
self.out_proj._is_residual = True
|
262 |
|
263 |
+
def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, is_causal: bool=True, needs_weights: bool=False) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
|
264 |
qkv = self.Wqkv(x)
|
265 |
if self.clip_qkv:
|
266 |
+
qkv = qkv.clamp(min=-self.clip_qkv, max=self.clip_qkv)
|
267 |
+
(query, key, value) = qkv.split([self.d_model, self.kv_n_heads * self.head_dim, self.kv_n_heads * self.head_dim], dim=2)
|
268 |
key_padding_mask = attention_mask
|
269 |
if self.qk_ln:
|
270 |
dtype = query.dtype
|
271 |
query = self.q_ln(query).to(dtype)
|
272 |
key = self.k_ln(key).to(dtype)
|
273 |
+
(context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, self.kv_n_heads, past_key_value=past_key_value, 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)
|
274 |
return (self.out_proj(context), attn_weights, past_key_value)
|
275 |
|
276 |
+
class MultiheadAttention(GroupedQueryAttention):
|
277 |
+
"""Multi-head self attention.
|
278 |
|
279 |
+
Using torch or triton attention implementation enables user to also use
|
280 |
additive bias.
|
281 |
"""
|
282 |
|
283 |
+
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, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True):
|
284 |
+
super().__init__(d_model=d_model, n_heads=n_heads, kv_n_heads=n_heads, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device, bias=bias)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
285 |
|
286 |
+
class MultiQueryAttention(GroupedQueryAttention):
|
287 |
+
"""Multi-Query self attention.
|
288 |
+
|
289 |
+
Using torch or triton attention implementation enables user to also use
|
290 |
+
additive bias.
|
291 |
+
"""
|
292 |
+
|
293 |
+
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, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True):
|
294 |
+
super().__init__(d_model=d_model, n_heads=n_heads, kv_n_heads=1, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device, bias=bias)
|
|
|
|
|
|
|
295 |
|
296 |
+
def attn_bias_shape(attn_impl: str, n_heads: int, seq_len: int, alibi: bool, prefix_lm: bool, causal: bool, use_sequence_id: bool) -> Optional[Tuple[int, int, int, int]]:
|
297 |
if attn_impl == 'flash':
|
298 |
return None
|
299 |
elif attn_impl in ['torch', 'triton']:
|
|
|
307 |
else:
|
308 |
raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
|
309 |
|
310 |
+
def build_attn_bias(attn_impl: str, attn_bias: torch.Tensor, n_heads: int, seq_len: int, causal: bool=False, alibi: bool=False, alibi_bias_max: int=8) -> Optional[torch.Tensor]:
|
311 |
if attn_impl == 'flash':
|
312 |
return None
|
313 |
elif attn_impl in ['torch', 'triton']:
|
|
|
318 |
else:
|
319 |
raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
|
320 |
|
321 |
+
def gen_slopes(n_heads: int, alibi_bias_max: int=8, device: Optional[torch.device]=None) -> torch.Tensor:
|
322 |
_n_heads = 2 ** math.ceil(math.log2(n_heads))
|
323 |
m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)
|
324 |
m = m.mul(alibi_bias_max / _n_heads)
|
|
|
327 |
slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
|
328 |
return slopes.view(1, n_heads, 1, 1)
|
329 |
|
330 |
+
def build_alibi_bias(n_heads: int, seq_len: int, full: bool=False, alibi_bias_max: int=8, device: Optional[torch.device]=None, dtype: Optional[torch.dtype]=None) -> torch.Tensor:
|
331 |
alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, 1, seq_len)
|
332 |
if full:
|
333 |
alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, seq_len, 1)
|
|
|
335 |
slopes = gen_slopes(n_heads, alibi_bias_max, device=device)
|
336 |
alibi_bias = alibi_bias * slopes
|
337 |
return alibi_bias.to(dtype=dtype)
|
338 |
+
ATTN_CLASS_REGISTRY = {'multihead_attention': MultiheadAttention, 'multiquery_attention': MultiQueryAttention, 'grouped_query_attention': GroupedQueryAttention}
|