xdecoder/Instruct-X-Decoder
163
1# Code copy from PyTorch, modified by Xueyan Zou2 3import warnings4from typing import Optional, Tuple5 6import torch7import torch.nn as nn8from torch import Tensor9from torch.nn.init import constant_, xavier_normal_, xavier_uniform_10from torch.nn.parameter import Parameter11from torch.overrides import has_torch_function, handle_torch_function12from torch.nn.functional import pad, linear, softmax, dropout13 14 15def multi_head_attention_forward(16 query: Tensor,17 key: Tensor,18 value: Tensor,19 embed_dim_to_check: int,20 num_heads: int,21 in_proj_weight: Tensor,22 in_proj_bias: Tensor,23 bias_k: Optional[Tensor],24 bias_v: Optional[Tensor],25 add_zero_attn: bool,26 dropout_p: float,27 out_proj_weight: Tensor,28 out_proj_bias: Tensor,29 training: bool = True,30 key_padding_mask: Optional[Tensor] = None,31 need_weights: bool = True,32 attn_mask: Optional[Tensor] = None,33 use_separate_proj_weight: bool = False,34 q_proj_weight: Optional[Tensor] = None,35 k_proj_weight: Optional[Tensor] = None,36 v_proj_weight: Optional[Tensor] = None,37 static_k: Optional[Tensor] = None,38 static_v: Optional[Tensor] = None,39) -> Tuple[Tensor, Optional[Tensor]]:40 r"""41 Args:42 query, key, value: map a query and a set of key-value pairs to an output.43 See "Attention Is All You Need" for more details.44 embed_dim_to_check: total dimension of the model.45 num_heads: parallel attention heads.46 in_proj_weight, in_proj_bias: input projection weight and bias.47 bias_k, bias_v: bias of the key and value sequences to be added at dim=0.48 add_zero_attn: add a new batch of zeros to the key and49 value sequences at dim=1.50 dropout_p: probability of an element to be zeroed.51 out_proj_weight, out_proj_bias: the output projection weight and bias.52 training: apply dropout if is ``True``.53 key_padding_mask: if provided, specified padding elements in the key will54 be ignored by the attention. This is an binary mask. When the value is True,55 the corresponding value on the attention layer will be filled with -inf.56 need_weights: output attn_output_weights.57 attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all58 the batches while a 3D mask allows to specify a different mask for the entries of each batch.59 use_separate_proj_weight: the function accept the proj. weights for query, key,60 and value in different forms. If false, in_proj_weight will be used, which is61 a combination of q_proj_weight, k_proj_weight, v_proj_weight.62 q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.63 static_k, static_v: static key and value used for attention operators.64 65 66 Shape:67 Inputs:68 - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is69 the embedding dimension.70 - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is71 the embedding dimension.72 - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is73 the embedding dimension.74 - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.75 If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions76 will be unchanged. If a BoolTensor is provided, the positions with the77 value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.78 - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.79 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,80 S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked81 positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend82 while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``83 are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor84 is provided, it will be added to the attention weight.85 - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,86 N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.87 - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,88 N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.89 90 Outputs:91 - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,92 E is the embedding dimension.93 - attn_output_weights: :math:`(N, L, S)` where N is the batch size,94 L is the target sequence length, S is the source sequence length.95 """96 tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v, out_proj_weight, out_proj_bias)97 if has_torch_function(tens_ops):98 return handle_torch_function(99 multi_head_attention_forward,100 tens_ops,101 query,102 key,103 value,104 embed_dim_to_check,105 num_heads,106 in_proj_weight,107 in_proj_bias,108 bias_k,109 bias_v,110 add_zero_attn,111 dropout_p,112 out_proj_weight,113 out_proj_bias,114 training=training,115 key_padding_mask=key_padding_mask,116 need_weights=need_weights,117 attn_mask=attn_mask,118 use_separate_proj_weight=use_separate_proj_weight,119 q_proj_weight=q_proj_weight,120 k_proj_weight=k_proj_weight,121 v_proj_weight=v_proj_weight,122 static_k=static_k,123 static_v=static_v,124 )125 tgt_len, bsz, embed_dim = query.size()126 assert embed_dim == embed_dim_to_check127 # allow MHA to have different sizes for the feature dimension128 assert key.size(0) == value.size(0) and key.size(1) == value.size(1)129 130 head_dim = embed_dim // num_heads131 assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"132 scaling = float(head_dim) ** -0.5133 134 if not use_separate_proj_weight:135 if (query is key or torch.equal(query, key)) and (key is value or torch.equal(key, value)):136 # self-attention137 q, k, v = linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1)138 139 elif key is value or torch.equal(key, value):140 # encoder-decoder attention141 # This is inline in_proj function with in_proj_weight and in_proj_bias142 _b = in_proj_bias143 _start = 0144 _end = embed_dim145 _w = in_proj_weight[_start:_end, :]146 if _b is not None:147 _b = _b[_start:_end]148 q = linear(query, _w, _b)149 150 if key is None:151 assert value is None152 k = None153 v = None154 else:155 156 # This is inline in_proj function with in_proj_weight and in_proj_bias157 _b = in_proj_bias158 _start = embed_dim159 _end = None160 _w = in_proj_weight[_start:, :]161 if _b is not None:162 _b = _b[_start:]163 k, v = linear(key, _w, _b).chunk(2, dim=-1)164 165 else:166 # This is inline in_proj function with in_proj_weight and in_proj_bias167 _b = in_proj_bias168 _start = 0169 _end = embed_dim170 _w = in_proj_weight[_start:_end, :]171 if _b is not None:172 _b = _b[_start:_end]173 q = linear(query, _w, _b)174 175 # This is inline in_proj function with in_proj_weight and in_proj_bias176 _b = in_proj_bias177 _start = embed_dim178 _end = embed_dim * 2179 _w = in_proj_weight[_start:_end, :]180 if _b is not None:181 _b = _b[_start:_end]182 k = linear(key, _w, _b)183 184 # This is inline in_proj function with in_proj_weight and in_proj_bias185 _b = in_proj_bias186 _start = embed_dim * 2187 _end = None188 _w = in_proj_weight[_start:, :]189 if _b is not None:190 _b = _b[_start:]191 v = linear(value, _w, _b)192 else:193 q_proj_weight_non_opt = torch.jit._unwrap_optional(q_proj_weight)194 len1, len2 = q_proj_weight_non_opt.size()195 assert len1 == embed_dim and len2 == query.size(-1)196 197 k_proj_weight_non_opt = torch.jit._unwrap_optional(k_proj_weight)198 len1, len2 = k_proj_weight_non_opt.size()199 assert len1 == embed_dim and len2 == key.size(-1)200 201 v_proj_weight_non_opt = torch.jit._unwrap_optional(v_proj_weight)202 len1, len2 = v_proj_weight_non_opt.size()203 assert len1 == embed_dim and len2 == value.size(-1)204 205 if in_proj_bias is not None:206 q = linear(query, q_proj_weight_non_opt, in_proj_bias[0:embed_dim])207 k = linear(key, k_proj_weight_non_opt, in_proj_bias[embed_dim : (embed_dim * 2)])208 v = linear(value, v_proj_weight_non_opt, in_proj_bias[(embed_dim * 2) :])209 else:210 q = linear(query, q_proj_weight_non_opt, in_proj_bias)211 k = linear(key, k_proj_weight_non_opt, in_proj_bias)212 v = linear(value, v_proj_weight_non_opt, in_proj_bias)213 q = q * scaling214 215 if attn_mask is not None:216 assert (217 attn_mask.dtype == torch.float32218 or attn_mask.dtype == torch.float64219 or attn_mask.dtype == torch.float16220 or attn_mask.dtype == torch.uint8221 or attn_mask.dtype == torch.bool222 ), "Only float, byte, and bool types are supported for attn_mask, not {}".format(attn_mask.dtype)223 if attn_mask.dtype == torch.uint8:224 warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")225 attn_mask = attn_mask.to(torch.bool)226 227 if attn_mask.dim() == 2:228 attn_mask = attn_mask.unsqueeze(0)229 if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:230 raise RuntimeError("The size of the 2D attn_mask is not correct.")231 elif attn_mask.dim() == 3:232 if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:233 raise RuntimeError("The size of the 3D attn_mask is not correct.")234 else:235 raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))236 # attn_mask's dim is 3 now.237 238 # convert ByteTensor key_padding_mask to bool239 if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:240 warnings.warn(241 "Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead."242 )243 key_padding_mask = key_padding_mask.to(torch.bool)244 245 if bias_k is not None and bias_v is not None:246 if static_k is None and static_v is None:247 k = torch.cat([k, bias_k.repeat(1, bsz, 1)])248 v = torch.cat([v, bias_v.repeat(1, bsz, 1)])249 if attn_mask is not None:250 attn_mask = pad(attn_mask, (0, 1))251 if key_padding_mask is not None:252 key_padding_mask = pad(key_padding_mask, (0, 1))253 else:254 assert static_k is None, "bias cannot be added to static key."255 assert static_v is None, "bias cannot be added to static value."256 else:257 assert bias_k is None258 assert bias_v is None259 260 q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)261 if k is not None:262 k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)263 if v is not None:264 v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)265 266 if static_k is not None:267 assert static_k.size(0) == bsz * num_heads268 assert static_k.size(2) == head_dim269 k = static_k270 271 if static_v is not None:272 assert static_v.size(0) == bsz * num_heads273 assert static_v.size(2) == head_dim274 v = static_v275 276 src_len = k.size(1)277 278 if key_padding_mask is not None:279 # assert key_padding_mask.size(0) == bsz280 assert key_padding_mask.size(1) == src_len281 282 if add_zero_attn:283 src_len += 1284 k = torch.cat([k, torch.zeros((k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)], dim=1)285 v = torch.cat([v, torch.zeros((v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)], dim=1)286 if attn_mask is not None:287 attn_mask = pad(attn_mask, (0, 1))288 if key_padding_mask is not None:289 key_padding_mask = pad(key_padding_mask, (0, 1))290 291 attn_output_weights = torch.bmm(q, k.transpose(1, 2))292 assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]293 294 if attn_mask is not None:295 if attn_mask.dtype == torch.bool:296 attn_output_weights.masked_fill_(attn_mask, float("-inf"))297 else:298 attn_output_weights += attn_mask299 300 if key_padding_mask is not None:301 attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)302 attn_output_weights = attn_output_weights.masked_fill(303 key_padding_mask.unsqueeze(1),304 float("-inf"),305 )306 attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len)307 308 attn_output_weights = softmax(attn_output_weights, dim=-1)309 attn_output_weights = dropout(attn_output_weights, p=dropout_p, training=training)310 311 attn_output = torch.bmm(attn_output_weights, v)312 assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]313 attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)314 attn_output = linear(attn_output, out_proj_weight, out_proj_bias)315 316 if need_weights:317 # average attention weights over heads318 attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)319 return attn_output, attn_output_weights.sum(dim=1) / num_heads320 else:321 return attn_output, None322 323 324# This class exists solely for Transformer; it has an annotation stating325# that bias is never None, which appeases TorchScript326class _LinearWithBias(nn.Linear):327 bias: Tensor # type: ignore328 329 def __init__(self, in_features: int, out_features: int) -> None:330 super().__init__(in_features, out_features, bias=True) # type: ignore331 332 333class MultiheadAttention(nn.Module):334 r"""Allows the model to jointly attend to information335 from different representation subspaces.336 See `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_337 338 .. math::339 \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O340 341 where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.342 343 Args:344 embed_dim: total dimension of the model.345 num_heads: parallel attention heads.346 dropout: a Dropout layer on attn_output_weights. Default: 0.0.347 bias: add bias as module parameter. Default: True.348 add_bias_kv: add bias to the key and value sequences at dim=0.349 add_zero_attn: add a new batch of zeros to the key and350 value sequences at dim=1.351 kdim: total number of features in key. Default: None.352 vdim: total number of features in value. Default: None.353 354 Note that if :attr:`kdim` and :attr:`vdim` are None, they will be set355 to :attr:`embed_dim` such that query, key, and value have the same356 number of features.357 358 Examples::359 360 >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)361 >>> attn_output, attn_output_weights = multihead_attn(query, key, value)362 """363 bias_k: Optional[torch.Tensor]364 bias_v: Optional[torch.Tensor]365 366 def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):367 super(MultiheadAttention, self).__init__()368 self.embed_dim = embed_dim369 self.kdim = kdim if kdim is not None else embed_dim370 self.vdim = vdim if vdim is not None else embed_dim371 self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim372 373 self.num_heads = num_heads374 self.dropout = dropout375 self.head_dim = embed_dim // num_heads376 assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"377 378 if self._qkv_same_embed_dim is False:379 self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))380 self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))381 self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))382 self.register_parameter('in_proj_weight', None)383 else:384 self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))385 self.register_parameter('q_proj_weight', None)386 self.register_parameter('k_proj_weight', None)387 self.register_parameter('v_proj_weight', None)388 389 if bias:390 self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))391 else:392 self.register_parameter('in_proj_bias', None)393 self.out_proj = _LinearWithBias(embed_dim, embed_dim)394 395 if add_bias_kv:396 self.bias_k = Parameter(torch.empty(1, 1, embed_dim))397 self.bias_v = Parameter(torch.empty(1, 1, embed_dim))398 else:399 self.bias_k = self.bias_v = None400 401 self.add_zero_attn = add_zero_attn402 403 self._reset_parameters()404 405 def _reset_parameters(self):406 if self._qkv_same_embed_dim:407 xavier_uniform_(self.in_proj_weight)408 else:409 xavier_uniform_(self.q_proj_weight)410 xavier_uniform_(self.k_proj_weight)411 xavier_uniform_(self.v_proj_weight)412 413 if self.in_proj_bias is not None:414 constant_(self.in_proj_bias, 0.)415 constant_(self.out_proj.bias, 0.)416 if self.bias_k is not None:417 xavier_normal_(self.bias_k)418 if self.bias_v is not None:419 xavier_normal_(self.bias_v)420 421 def __setstate__(self, state):422 # Support loading old MultiheadAttention checkpoints generated by v1.1.0423 if '_qkv_same_embed_dim' not in state:424 state['_qkv_same_embed_dim'] = True425 426 super(MultiheadAttention, self).__setstate__(state)427 428 def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: Optional[Tensor] = None,429 need_weights: bool = True, attn_mask: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]:430 r"""431 Args:432 query, key, value: map a query and a set of key-value pairs to an output.433 See "Attention Is All You Need" for more details.434 key_padding_mask: if provided, specified padding elements in the key will435 be ignored by the attention. When given a binary mask and a value is True,436 the corresponding value on the attention layer will be ignored. When given437 a byte mask and a value is non-zero, the corresponding value on the attention438 layer will be ignored439 need_weights: output attn_output_weights.440 attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all441 the batches while a 3D mask allows to specify a different mask for the entries of each batch.442 443 Shapes for inputs:444 - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is445 the embedding dimension.446 - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is447 the embedding dimension.448 - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is449 the embedding dimension.450 - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.451 If a ByteTensor is provided, the non-zero positions will be ignored while the position452 with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the453 value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.454 - attn_mask: if a 2D mask: :math:`(L, S)` where L is the target sequence length, S is the455 source sequence length.456 457 If a 3D mask: :math:`(N\cdot\text{num\_heads}, L, S)` where N is the batch size, L is the target sequence458 length, S is the source sequence length. ``attn_mask`` ensure that position i is allowed to attend459 the unmasked positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend460 while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``461 is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor462 is provided, it will be added to the attention weight.463 464 Shapes for outputs:465 - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,466 E is the embedding dimension.467 - attn_output_weights: :math:`(N, L, S)` where N is the batch size,468 L is the target sequence length, S is the source sequence length.469 """470 if not self._qkv_same_embed_dim:471 return multi_head_attention_forward(472 query, key, value, self.embed_dim, self.num_heads,473 self.in_proj_weight, self.in_proj_bias,474 self.bias_k, self.bias_v, self.add_zero_attn,475 self.dropout, self.out_proj.weight, self.out_proj.bias,476 training=self.training,477 key_padding_mask=key_padding_mask, need_weights=need_weights,478 attn_mask=attn_mask, use_separate_proj_weight=True,479 q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,480 v_proj_weight=self.v_proj_weight)481 else:482 return multi_head_attention_forward(483 query, key, value, self.embed_dim, self.num_heads,484 self.in_proj_weight, self.in_proj_bias,485 self.bias_k, self.bias_v, self.add_zero_attn,486 self.dropout, self.out_proj.weight, self.out_proj.bias,487 training=self.training,488 key_padding_mask=key_padding_mask, need_weights=need_weights,489 attn_mask=attn_mask)