faisalashraf/abaffinity
0
1# Copyright (c) Meta Platforms, Inc. and affiliates.2#3# This source code is licensed under the MIT license found in the4# LICENSE file in the root directory of this source tree.5 6import math7from typing import Dict, Optional, Tuple8 9import torch10import torch.nn.functional as F11from torch import Tensor, nn12from torch.nn import Parameter13from .rotary_embedding import RotaryEmbedding14 15import uuid16 17 18def utils_softmax(x, dim: int, onnx_trace: bool = False):19 if onnx_trace:20 return F.softmax(x.float(), dim=dim)21 else:22 return F.softmax(x, dim=dim, dtype=torch.float32)23 24 25class FairseqIncrementalState(object):26 def __init__(self, *args, **kwargs):27 super().__init__(*args, **kwargs)28 self.init_incremental_state()29 30 def init_incremental_state(self):31 self._incremental_state_id = str(uuid.uuid4())32 33 def _get_full_incremental_state_key(self, key: str) -> str:34 return "{}.{}".format(self._incremental_state_id, key)35 36 def get_incremental_state(37 self,38 incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],39 key: str,40 ) -> Optional[Dict[str, Optional[Tensor]]]:41 """Helper for getting incremental state for an nn.Module."""42 full_key = self._get_full_incremental_state_key(key)43 if incremental_state is None or full_key not in incremental_state:44 return None45 return incremental_state[full_key]46 47 def set_incremental_state(48 self,49 incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],50 key: str,51 value: Dict[str, Optional[Tensor]],52 ) -> Optional[Dict[str, Dict[str, Optional[Tensor]]]]:53 """Helper for setting incremental state for an nn.Module."""54 if incremental_state is not None:55 full_key = self._get_full_incremental_state_key(key)56 incremental_state[full_key] = value57 return incremental_state58 59 60def with_incremental_state(cls):61 cls.__bases__ = (FairseqIncrementalState,) + tuple(62 b for b in cls.__bases__ if b != FairseqIncrementalState63 )64 return cls65 66 67@with_incremental_state68class MultiheadAttention(nn.Module):69 """Multi-headed attention.70 71 See "Attention Is All You Need" for more details.72 """73 74 def __init__(75 self,76 embed_dim,77 num_heads,78 kdim=None,79 vdim=None,80 dropout=0.0,81 bias=True,82 add_bias_kv: bool = False,83 add_zero_attn: bool = False,84 self_attention: bool = False,85 encoder_decoder_attention: bool = False,86 use_rotary_embeddings: bool = False,87 ):88 super().__init__()89 self.embed_dim = embed_dim90 self.kdim = kdim if kdim is not None else embed_dim91 self.vdim = vdim if vdim is not None else embed_dim92 self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim93 94 self.num_heads = num_heads95 self.dropout = dropout96 self.head_dim = embed_dim // num_heads97 assert (98 self.head_dim * num_heads == self.embed_dim99 ), "embed_dim must be divisible by num_heads"100 self.scaling = self.head_dim**-0.5101 102 self.self_attention = self_attention103 self.encoder_decoder_attention = encoder_decoder_attention104 105 assert not self.self_attention or self.qkv_same_dim, (106 "Self-attention requires query, key and " "value to be of the same size"107 )108 109 self.k_proj = nn.Linear(self.kdim, embed_dim, bias=bias)110 self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias)111 self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)112 113 self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)114 115 if add_bias_kv:116 self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))117 self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))118 else:119 self.bias_k = self.bias_v = None120 121 self.add_zero_attn = add_zero_attn122 123 self.reset_parameters()124 125 self.onnx_trace = False126 self.rot_emb = None127 if use_rotary_embeddings:128 self.rot_emb = RotaryEmbedding(dim=self.head_dim)129 130 self.enable_torch_version = False131 if hasattr(F, "multi_head_attention_forward"):132 self.enable_torch_version = True133 else:134 self.enable_torch_version = False135 136 def prepare_for_onnx_export_(self):137 self.onnx_trace = True138 139 def reset_parameters(self):140 if self.qkv_same_dim:141 # Empirically observed the convergence to be much better with142 # the scaled initialization143 nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))144 nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))145 nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))146 else:147 nn.init.xavier_uniform_(self.k_proj.weight)148 nn.init.xavier_uniform_(self.v_proj.weight)149 nn.init.xavier_uniform_(self.q_proj.weight)150 151 nn.init.xavier_uniform_(self.out_proj.weight)152 if self.out_proj.bias is not None:153 nn.init.constant_(self.out_proj.bias, 0.0)154 if self.bias_k is not None:155 nn.init.xavier_normal_(self.bias_k)156 if self.bias_v is not None:157 nn.init.xavier_normal_(self.bias_v)158 159 def forward(160 self,161 query,162 key: Optional[Tensor],163 value: Optional[Tensor],164 key_padding_mask: Optional[Tensor] = None,165 incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,166 need_weights: bool = True,167 static_kv: bool = False,168 attn_mask: Optional[Tensor] = None,169 before_softmax: bool = False,170 need_head_weights: bool = False,171 ) -> Tuple[Tensor, Optional[Tensor]]:172 """Input shape: Time x Batch x Channel173 174 Args:175 key_padding_mask (ByteTensor, optional): mask to exclude176 keys that are pads, of shape `(batch, src_len)`, where177 padding elements are indicated by 1s.178 need_weights (bool, optional): return the attention weights,179 averaged over heads (default: False).180 attn_mask (ByteTensor, optional): typically used to181 implement causal attention, where the mask prevents the182 attention from looking forward in time (default: None).183 before_softmax (bool, optional): return the raw attention184 weights and values before the attention softmax.185 need_head_weights (bool, optional): return the attention186 weights for each head. Implies *need_weights*. Default:187 return the average attention weights over all heads.188 """189 if need_head_weights:190 need_weights = True191 192 tgt_len, bsz, embed_dim = query.size()193 assert embed_dim == self.embed_dim194 assert list(query.size()) == [tgt_len, bsz, embed_dim]195 196 if (197 not self.rot_emb198 and self.enable_torch_version199 and not self.onnx_trace200 and incremental_state is None201 and not static_kv202 # A workaround for quantization to work. Otherwise JIT compilation203 # treats bias in linear module as method.204 and not torch.jit.is_scripting()205 and not need_head_weights206 ):207 assert key is not None and value is not None208 return F.multi_head_attention_forward(209 query,210 key,211 value,212 self.embed_dim,213 self.num_heads,214 torch.empty([0]),215 torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)),216 self.bias_k,217 self.bias_v,218 self.add_zero_attn,219 self.dropout,220 self.out_proj.weight,221 self.out_proj.bias,222 self.training,223 key_padding_mask,224 need_weights,225 attn_mask,226 use_separate_proj_weight=True,227 q_proj_weight=self.q_proj.weight,228 k_proj_weight=self.k_proj.weight,229 v_proj_weight=self.v_proj.weight,230 )231 if incremental_state is not None:232 saved_state = self._get_input_buffer(incremental_state)233 if saved_state is not None and "prev_key" in saved_state:234 # previous time steps are cached - no need to recompute235 # key and value if they are static236 if static_kv:237 assert self.encoder_decoder_attention and not self.self_attention238 key = value = None239 else:240 saved_state = None241 242 if self.self_attention:243 q = self.q_proj(query)244 k = self.k_proj(query)245 v = self.v_proj(query)246 elif self.encoder_decoder_attention:247 # encoder-decoder attention248 q = self.q_proj(query)249 if key is None:250 assert value is None251 k = v = None252 else:253 k = self.k_proj(key)254 v = self.v_proj(key)255 256 else:257 assert key is not None and value is not None258 q = self.q_proj(query)259 k = self.k_proj(key)260 v = self.v_proj(value)261 q *= self.scaling262 263 if self.bias_k is not None:264 assert self.bias_v is not None265 k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])266 v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])267 if attn_mask is not None:268 attn_mask = torch.cat(269 [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1270 )271 if key_padding_mask is not None:272 key_padding_mask = torch.cat(273 [274 key_padding_mask,275 key_padding_mask.new_zeros(key_padding_mask.size(0), 1),276 ],277 dim=1,278 )279 280 q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)281 if k is not None:282 k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)283 if v is not None:284 v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)285 286 if saved_state is not None:287 # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)288 if "prev_key" in saved_state:289 _prev_key = saved_state["prev_key"]290 assert _prev_key is not None291 prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)292 if static_kv:293 k = prev_key294 else:295 assert k is not None296 k = torch.cat([prev_key, k], dim=1)297 if "prev_value" in saved_state:298 _prev_value = saved_state["prev_value"]299 assert _prev_value is not None300 prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)301 if static_kv:302 v = prev_value303 else:304 assert v is not None305 v = torch.cat([prev_value, v], dim=1)306 prev_key_padding_mask: Optional[Tensor] = None307 if "prev_key_padding_mask" in saved_state:308 prev_key_padding_mask = saved_state["prev_key_padding_mask"]309 assert k is not None and v is not None310 key_padding_mask = MultiheadAttention._append_prev_key_padding_mask(311 key_padding_mask=key_padding_mask,312 prev_key_padding_mask=prev_key_padding_mask,313 batch_size=bsz,314 src_len=k.size(1),315 static_kv=static_kv,316 )317 318 saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim)319 saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim)320 saved_state["prev_key_padding_mask"] = key_padding_mask321 # In this branch incremental_state is never None322 assert incremental_state is not None323 incremental_state = self._set_input_buffer(incremental_state, saved_state)324 assert k is not None325 src_len = k.size(1)326 327 # This is part of a workaround to get around fork/join parallelism328 # not supporting Optional types.329 if key_padding_mask is not None and key_padding_mask.dim() == 0:330 key_padding_mask = None331 332 if key_padding_mask is not None:333 assert key_padding_mask.size(0) == bsz334 assert key_padding_mask.size(1) == src_len335 336 if self.add_zero_attn:337 assert v is not None338 src_len += 1339 k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)340 v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)341 if attn_mask is not None:342 attn_mask = torch.cat(343 [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1344 )345 if key_padding_mask is not None:346 key_padding_mask = torch.cat(347 [348 key_padding_mask,349 torch.zeros(key_padding_mask.size(0), 1).type_as(key_padding_mask),350 ],351 dim=1,352 )353 354 if self.rot_emb:355 q, k = self.rot_emb(q, k)356 357 attn_weights = torch.bmm(q, k.transpose(1, 2))358 attn_weights = MultiheadAttention.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)359 360 assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]361 362 if attn_mask is not None:363 attn_mask = attn_mask.unsqueeze(0)364 if self.onnx_trace:365 attn_mask = attn_mask.repeat(attn_weights.size(0), 1, 1)366 attn_weights += attn_mask367 368 if key_padding_mask is not None:369 # don't attend to padding symbols370 attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)371 attn_weights = attn_weights.masked_fill(372 key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf")373 )374 attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)375 376 if before_softmax:377 return attn_weights, v378 379 attn_weights_float = utils_softmax(attn_weights, dim=-1, onnx_trace=self.onnx_trace)380 attn_weights = attn_weights_float.type_as(attn_weights)381 attn_probs = F.dropout(382 attn_weights_float.type_as(attn_weights),383 p=self.dropout,384 training=self.training,385 )386 assert v is not None387 attn = torch.bmm(attn_probs, v)388 assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]389 if self.onnx_trace and attn.size(1) == 1:390 # when ONNX tracing a single decoder step (sequence length == 1)391 # the transpose is a no-op copy before view, thus unnecessary392 attn = attn.contiguous().view(tgt_len, bsz, embed_dim)393 else:394 attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)395 attn = self.out_proj(attn)396 attn_weights: Optional[Tensor] = None397 if need_weights:398 attn_weights = attn_weights_float.view(399 bsz, self.num_heads, tgt_len, src_len400 ).type_as(attn).transpose(1, 0)401 if not need_head_weights:402 # average attention weights over heads403 attn_weights = attn_weights.mean(dim=0)404 405 return attn, attn_weights406 407 @staticmethod408 def _append_prev_key_padding_mask(409 key_padding_mask: Optional[Tensor],410 prev_key_padding_mask: Optional[Tensor],411 batch_size: int,412 src_len: int,413 static_kv: bool,414 ) -> Optional[Tensor]:415 # saved key padding masks have shape (bsz, seq_len)416 if prev_key_padding_mask is not None and static_kv:417 new_key_padding_mask = prev_key_padding_mask418 elif prev_key_padding_mask is not None and key_padding_mask is not None:419 new_key_padding_mask = torch.cat(420 [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1421 )422 # During incremental decoding, as the padding token enters and423 # leaves the frame, there will be a time when prev or current424 # is None425 elif prev_key_padding_mask is not None:426 filler = torch.zeros(427 (batch_size, src_len - prev_key_padding_mask.size(1)),428 device=prev_key_padding_mask.device,429 )430 new_key_padding_mask = torch.cat(431 [prev_key_padding_mask.float(), filler.float()], dim=1432 )433 elif key_padding_mask is not None:434 filler = torch.zeros(435 (batch_size, src_len - key_padding_mask.size(1)),436 device=key_padding_mask.device,437 )438 new_key_padding_mask = torch.cat([filler.float(), key_padding_mask.float()], dim=1)439 else:440 new_key_padding_mask = prev_key_padding_mask441 return new_key_padding_mask442 443 @torch.jit.export444 def reorder_incremental_state(445 self, incremental_state: Dict[str, Dict[str, Optional[Tensor]]], new_order: Tensor446 ):447 """Reorder buffered internal state (for incremental generation)."""448 input_buffer = self._get_input_buffer(incremental_state)449 if input_buffer is not None:450 for k in input_buffer.keys():451 input_buffer_k = input_buffer[k]452 if input_buffer_k is not None:453 if self.encoder_decoder_attention and input_buffer_k.size(0) == new_order.size(454 0455 ):456 break457 input_buffer[k] = input_buffer_k.index_select(0, new_order)458 incremental_state = self._set_input_buffer(incremental_state, input_buffer)459 return incremental_state460 461 def _get_input_buffer(462 self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]]463 ) -> Dict[str, Optional[Tensor]]:464 result = self.get_incremental_state(incremental_state, "attn_state")465 if result is not None:466 return result467 else:468 empty_result: Dict[str, Optional[Tensor]] = {}469 return empty_result470 471 def _set_input_buffer(472 self,473 incremental_state: Dict[str, Dict[str, Optional[Tensor]]],474 buffer: Dict[str, Optional[Tensor]],475 ):476 return self.set_incremental_state(incremental_state, "attn_state", buffer)477 478 def apply_sparse_mask(attn_weights, tgt_len: int, src_len: int, bsz: int):479 return attn_weights480 481 def upgrade_state_dict_named(self, state_dict, name):482 prefix = name + "." if name != "" else ""483 items_to_add = {}484 keys_to_remove = []485 for k in state_dict.keys():486 if k.endswith(prefix + "in_proj_weight"):487 # in_proj_weight used to be q + k + v with same dimensions488 dim = int(state_dict[k].shape[0] / 3)489 items_to_add[prefix + "q_proj.weight"] = state_dict[k][:dim]490 items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim : 2 * dim]491 items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim :]492 493 keys_to_remove.append(k)494 495 k_bias = prefix + "in_proj_bias"496 if k_bias in state_dict.keys():497 dim = int(state_dict[k].shape[0] / 3)498 items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim]499 items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][dim : 2 * dim]500 items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :]501 502 keys_to_remove.append(prefix + "in_proj_bias")503 504 for k in keys_to_remove:505 del state_dict[k]506 507 for key, value in items_to_add.items():508 state_dict[key] = value509