|
|
|
from typing import Optional |
|
from typing import Tuple |
|
import torch |
|
from torch import Tensor |
|
from torch.nn import Linear |
|
from torch.nn import Module |
|
from torch.nn.init import constant_ |
|
from torch.nn.init import xavier_normal_ |
|
from torch.nn.init import xavier_uniform_ |
|
from torch.nn.modules.linear import NonDynamicallyQuantizableLinear |
|
from torch.nn.parameter import Parameter |
|
|
|
from torch.nn import functional as F |
|
from AR.modules.patched_mha_with_cache_onnx import multi_head_attention_forward_patched |
|
|
|
|
|
class MultiheadAttention(Module): |
|
__constants__ = ["batch_first"] |
|
bias_k: Optional[torch.Tensor] |
|
bias_v: Optional[torch.Tensor] |
|
|
|
def __init__( |
|
self, |
|
embed_dim, |
|
num_heads, |
|
dropout=0.0, |
|
bias=True, |
|
add_bias_kv=False, |
|
add_zero_attn=False, |
|
kdim=None, |
|
vdim=None, |
|
batch_first=False, |
|
linear1_cls=Linear, |
|
linear2_cls=Linear, |
|
device=None, |
|
dtype=None, |
|
) -> None: |
|
factory_kwargs = {"device": device, "dtype": dtype} |
|
super(MultiheadAttention, self).__init__() |
|
self.embed_dim = embed_dim |
|
self.kdim = kdim if kdim is not None else embed_dim |
|
self.vdim = vdim if vdim is not None else embed_dim |
|
self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim |
|
|
|
self.num_heads = num_heads |
|
self.dropout = dropout |
|
self.batch_first = batch_first |
|
self.head_dim = embed_dim // num_heads |
|
assert ( |
|
self.head_dim * num_heads == self.embed_dim |
|
), "embed_dim must be divisible by num_heads" |
|
|
|
if add_bias_kv: |
|
self.bias_k = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs)) |
|
self.bias_v = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs)) |
|
else: |
|
self.bias_k = self.bias_v = None |
|
|
|
if linear1_cls == Linear: |
|
if not self._qkv_same_embed_dim: |
|
self.q_proj_weight = Parameter( |
|
torch.empty((embed_dim, embed_dim), **factory_kwargs) |
|
) |
|
self.k_proj_weight = Parameter( |
|
torch.empty((embed_dim, self.kdim), **factory_kwargs) |
|
) |
|
self.v_proj_weight = Parameter( |
|
torch.empty((embed_dim, self.vdim), **factory_kwargs) |
|
) |
|
self.register_parameter("in_proj_weight", None) |
|
else: |
|
self.in_proj_weight = Parameter( |
|
torch.empty((3 * embed_dim, embed_dim), **factory_kwargs) |
|
) |
|
self.register_parameter("q_proj_weight", None) |
|
self.register_parameter("k_proj_weight", None) |
|
self.register_parameter("v_proj_weight", None) |
|
|
|
if bias: |
|
self.in_proj_bias = Parameter( |
|
torch.empty(3 * embed_dim, **factory_kwargs) |
|
) |
|
else: |
|
self.register_parameter("in_proj_bias", None) |
|
self.out_proj = NonDynamicallyQuantizableLinear( |
|
embed_dim, embed_dim, bias=bias, **factory_kwargs |
|
) |
|
|
|
self._reset_parameters() |
|
else: |
|
if not self._qkv_same_embed_dim: |
|
raise NotImplementedError |
|
else: |
|
self.in_proj_linear = linear1_cls( |
|
embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs |
|
) |
|
self.in_proj_weight = self.in_proj_linear.weight |
|
|
|
self.register_parameter("q_proj_weight", None) |
|
self.register_parameter("k_proj_weight", None) |
|
self.register_parameter("v_proj_weight", None) |
|
|
|
if bias: |
|
self.in_proj_bias = self.in_proj_linear.bias |
|
else: |
|
self.register_parameter("in_proj_bias", None) |
|
|
|
self.out_proj = linear2_cls( |
|
embed_dim, embed_dim, bias=bias, **factory_kwargs |
|
) |
|
|
|
if self.bias_k is not None: |
|
xavier_normal_(self.bias_k) |
|
if self.bias_v is not None: |
|
xavier_normal_(self.bias_v) |
|
|
|
self.add_zero_attn = add_zero_attn |
|
|
|
def _reset_parameters(self): |
|
if self._qkv_same_embed_dim: |
|
xavier_uniform_(self.in_proj_weight) |
|
else: |
|
xavier_uniform_(self.q_proj_weight) |
|
xavier_uniform_(self.k_proj_weight) |
|
xavier_uniform_(self.v_proj_weight) |
|
|
|
if self.in_proj_bias is not None: |
|
constant_(self.in_proj_bias, 0.0) |
|
constant_(self.out_proj.bias, 0.0) |
|
|
|
if self.bias_k is not None: |
|
xavier_normal_(self.bias_k) |
|
if self.bias_v is not None: |
|
xavier_normal_(self.bias_v) |
|
|
|
def __setstate__(self, state): |
|
|
|
if "_qkv_same_embed_dim" not in state: |
|
state["_qkv_same_embed_dim"] = True |
|
|
|
super(MultiheadAttention, self).__setstate__(state) |
|
|
|
def forward( |
|
self, |
|
query: Tensor, |
|
key: Tensor, |
|
value: Tensor, |
|
key_padding_mask: Optional[Tensor] = None, |
|
need_weights: bool = True, |
|
attn_mask: Optional[Tensor] = None, |
|
average_attn_weights: bool = True, |
|
cache=None, |
|
) -> Tuple[Tensor, Optional[Tensor]]: |
|
any_nested = query.is_nested or key.is_nested or value.is_nested |
|
query = key = value = query.transpose(1, 0) |
|
attn_output = multi_head_attention_forward_patched( |
|
query, |
|
key, |
|
value, |
|
self.embed_dim, |
|
self.num_heads, |
|
self.in_proj_weight, |
|
self.in_proj_bias, |
|
self.bias_k, |
|
self.bias_v, |
|
self.add_zero_attn, |
|
self.dropout, |
|
self.out_proj.weight, |
|
self.out_proj.bias, |
|
training=self.training, |
|
key_padding_mask=key_padding_mask, |
|
need_weights=need_weights, |
|
attn_mask=attn_mask, |
|
average_attn_weights=average_attn_weights, |
|
cache=cache, |
|
) |
|
return attn_output.transpose(1, 0) |
|
|
