| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torch.utils.checkpoint | |
| import einops | |
| from einops import rearrange, repeat | |
| from inspect import isfunction | |
| from .rotary import RotaryEmbedding | |
| from .modules import RMSNorm | |
| if hasattr(nn.functional, 'scaled_dot_product_attention'): | |
| ATTENTION_MODE = 'flash' | |
| else: | |
| ATTENTION_MODE = 'math' | |
| print(f'attention mode is {ATTENTION_MODE}') | |
| def add_mask(sim, mask): | |
| b, ndim = sim.shape[0], mask.ndim | |
| if ndim == 3: | |
| mask = rearrange(mask, "b n m -> b 1 n m") | |
| if ndim == 2: | |
| mask = repeat(mask, "n m -> b 1 n m", b=b) | |
| max_neg_value = -torch.finfo(sim.dtype).max | |
| sim = sim.masked_fill(~mask, max_neg_value) | |
| return sim | |
| def create_mask(q_shape, k_shape, device, q_mask=None, k_mask=None): | |
| def default(val, d): | |
| return val if val is not None else (d() if isfunction(d) else d) | |
| b, i, j, device = q_shape[0], q_shape[-2], k_shape[-2], device | |
| q_mask = default(q_mask, torch.ones((b, i), device=device, dtype=torch.bool)) | |
| k_mask = default(k_mask, torch.ones((b, j), device=device, dtype=torch.bool)) | |
| attn_mask = rearrange(q_mask, 'b i -> b 1 i 1') * rearrange(k_mask, 'b j -> b 1 1 j') | |
| return attn_mask | |
| class Attention(nn.Module): | |
| def __init__(self, dim, context_dim=None, num_heads=8, | |
| qkv_bias=False, qk_scale=None, qk_norm=None, | |
| attn_drop=0., proj_drop=0., rope_mode='none'): | |
| super().__init__() | |
| self.num_heads = num_heads | |
| head_dim = dim // num_heads | |
| self.scale = qk_scale or head_dim ** -0.5 | |
| if context_dim is None: | |
| self.cross_attn = False | |
| else: | |
| self.cross_attn = True | |
| context_dim = dim if context_dim is None else context_dim | |
| self.to_q = nn.Linear(dim, dim, bias=qkv_bias) | |
| self.to_k = nn.Linear(context_dim, dim, bias=qkv_bias) | |
| self.to_v = nn.Linear(context_dim, dim, bias=qkv_bias) | |
| if qk_norm is None: | |
| self.norm_q = nn.Identity() | |
| self.norm_k = nn.Identity() | |
| elif qk_norm == 'layernorm': | |
| self.norm_q = nn.LayerNorm(head_dim) | |
| self.norm_k = nn.LayerNorm(head_dim) | |
| elif qk_norm == 'rmsnorm': | |
| self.norm_q = RMSNorm(head_dim) | |
| self.norm_k = RMSNorm(head_dim) | |
| else: | |
| raise NotImplementedError | |
| self.attn_drop_p = attn_drop | |
| self.attn_drop = nn.Dropout(attn_drop) | |
| self.proj = nn.Linear(dim, dim) | |
| self.proj_drop = nn.Dropout(proj_drop) | |
| if self.cross_attn: | |
| assert rope_mode == 'none' | |
| self.rope_mode = rope_mode | |
| if self.rope_mode == 'shared' or self.rope_mode == 'x_only': | |
| self.rotary = RotaryEmbedding(dim=head_dim) | |
| elif self.rope_mode == 'dual': | |
| self.rotary_x = RotaryEmbedding(dim=head_dim) | |
| self.rotary_c = RotaryEmbedding(dim=head_dim) | |
| def _rotary(self, q, k, extras): | |
| if self.rope_mode == 'shared': | |
| q, k = self.rotary(q=q, k=k) | |
| elif self.rope_mode == 'x_only': | |
| q_x, k_x = self.rotary(q=q[:, :, extras:, :], k=k[:, :, extras:, :]) | |
| q_c, k_c = q[:, :, :extras, :], k[:, :, :extras, :] | |
| q = torch.cat((q_c, q_x), dim=2) | |
| k = torch.cat((k_c, k_x), dim=2) | |
| elif self.rope_mode == 'dual': | |
| q_x, k_x = self.rotary_x(q=q[:, :, extras:, :], k=k[:, :, extras:, :]) | |
| q_c, k_c = self.rotary_c(q=q[:, :, :extras, :], k=k[:, :, :extras, :]) | |
| q = torch.cat((q_c, q_x), dim=2) | |
| k = torch.cat((k_c, k_x), dim=2) | |
| elif self.rope_mode == 'none': | |
| pass | |
| else: | |
| raise NotImplementedError | |
| return q, k | |
| def _attn(self, q, k, v, mask_binary): | |
| if ATTENTION_MODE == 'flash': | |
| x = F.scaled_dot_product_attention(q, k, v, | |
| dropout_p=self.attn_drop_p, | |
| attn_mask=mask_binary) | |
| x = einops.rearrange(x, 'B H L D -> B L (H D)') | |
| elif ATTENTION_MODE == 'math': | |
| attn = (q @ k.transpose(-2, -1)) * self.scale | |
| attn = add_mask(attn, mask_binary) if mask_binary is not None else attn | |
| attn = attn.softmax(dim=-1) | |
| attn = self.attn_drop(attn) | |
| x = (attn @ v).transpose(1, 2) | |
| x = einops.rearrange(x, 'B H L D -> B L (H D)') | |
| else: | |
| raise NotImplementedError | |
| return x | |
| def forward(self, x, context=None, context_mask=None, extras=0): | |
| B, L, C = x.shape | |
| if context is None: | |
| context = x | |
| q = self.to_q(x) | |
| k = self.to_k(context) | |
| v = self.to_v(context) | |
| if context_mask is not None: | |
| mask_binary = create_mask(x.shape, context.shape, | |
| x.device, None, context_mask) | |
| else: | |
| mask_binary = None | |
| q = einops.rearrange(q, 'B L (H D) -> B H L D', H=self.num_heads) | |
| k = einops.rearrange(k, 'B L (H D) -> B H L D', H=self.num_heads) | |
| v = einops.rearrange(v, 'B L (H D) -> B H L D', H=self.num_heads) | |
| q = self.norm_q(q) | |
| k = self.norm_k(k) | |
| q, k = self._rotary(q, k, extras) | |
| x = self._attn(q, k, v, mask_binary) | |
| x = self.proj(x) | |
| x = self.proj_drop(x) | |
| return x | |
| class JointAttention(nn.Module): | |
| def __init__(self, dim, num_heads=8, | |
| qkv_bias=False, qk_scale=None, qk_norm=None, | |
| attn_drop=0., proj_drop=0., | |
| rope_mode='none'): | |
| super().__init__() | |
| self.num_heads = num_heads | |
| head_dim = dim // num_heads | |
| self.scale = qk_scale or head_dim ** -0.5 | |
| self.to_qx, self.to_kx, self.to_vx = self._make_qkv_layers(dim, qkv_bias) | |
| self.to_qc, self.to_kc, self.to_vc = self._make_qkv_layers(dim, qkv_bias) | |
| self.norm_qx, self.norm_kx = self._make_norm_layers(qk_norm, head_dim) | |
| self.norm_qc, self.norm_kc = self._make_norm_layers(qk_norm, head_dim) | |
| self.attn_drop_p = attn_drop | |
| self.attn_drop = nn.Dropout(attn_drop) | |
| self.proj_x = nn.Linear(dim, dim) | |
| self.proj_drop_x = nn.Dropout(proj_drop) | |
| self.proj_c = nn.Linear(dim, dim) | |
| self.proj_drop_c = nn.Dropout(proj_drop) | |
| self.rope_mode = rope_mode | |
| if self.rope_mode == 'shared' or self.rope_mode == 'x_only': | |
| self.rotary = RotaryEmbedding(dim=head_dim) | |
| elif self.rope_mode == 'dual': | |
| self.rotary_x = RotaryEmbedding(dim=head_dim) | |
| self.rotary_c = RotaryEmbedding(dim=head_dim) | |
| def _make_qkv_layers(self, dim, qkv_bias): | |
| return (nn.Linear(dim, dim, bias=qkv_bias), | |
| nn.Linear(dim, dim, bias=qkv_bias), | |
| nn.Linear(dim, dim, bias=qkv_bias)) | |
| def _make_norm_layers(self, qk_norm, head_dim): | |
| if qk_norm is None: | |
| norm_q = nn.Identity() | |
| norm_k = nn.Identity() | |
| elif qk_norm == 'layernorm': | |
| norm_q = nn.LayerNorm(head_dim) | |
| norm_k = nn.LayerNorm(head_dim) | |
| elif qk_norm == 'rmsnorm': | |
| norm_q = RMSNorm(head_dim) | |
| norm_k = RMSNorm(head_dim) | |
| else: | |
| raise NotImplementedError | |
| return norm_q, norm_k | |
| def _rotary(self, q, k, extras): | |
| if self.rope_mode == 'shared': | |
| q, k = self.rotary(q=q, k=k) | |
| elif self.rope_mode == 'x_only': | |
| q_x, k_x = self.rotary(q=q[:, :, extras:, :], k=k[:, :, extras:, :]) | |
| q_c, k_c = q[:, :, :extras, :], k[:, :, :extras, :] | |
| q = torch.cat((q_c, q_x), dim=2) | |
| k = torch.cat((k_c, k_x), dim=2) | |
| elif self.rope_mode == 'dual': | |
| q_x, k_x = self.rotary_x(q=q[:, :, extras:, :], k=k[:, :, extras:, :]) | |
| q_c, k_c = self.rotary_c(q=q[:, :, :extras, :], k=k[:, :, :extras, :]) | |
| q = torch.cat((q_c, q_x), dim=2) | |
| k = torch.cat((k_c, k_x), dim=2) | |
| elif self.rope_mode == 'none': | |
| pass | |
| else: | |
| raise NotImplementedError | |
| return q, k | |
| def _attn(self, q, k, v, mask_binary): | |
| if ATTENTION_MODE == 'flash': | |
| x = F.scaled_dot_product_attention(q, k, v, | |
| dropout_p=self.attn_drop_p, | |
| attn_mask=mask_binary) | |
| x = einops.rearrange(x, 'B H L D -> B L (H D)') | |
| elif ATTENTION_MODE == 'math': | |
| attn = (q @ k.transpose(-2, -1)) * self.scale | |
| attn = add_mask(attn, mask_binary) if mask_binary is not None else attn | |
| attn = attn.softmax(dim=-1) | |
| attn = self.attn_drop(attn) | |
| x = (attn @ v).transpose(1, 2) | |
| x = einops.rearrange(x, 'B H L D -> B L (H D)') | |
| else: | |
| raise NotImplementedError | |
| return x | |
| def _cat_mask(self, x, context, x_mask=None, context_mask=None): | |
| B = x.shape[0] | |
| if x_mask is None: | |
| x_mask = torch.ones(B, x.shape[-2], device=x.device).bool() | |
| if context_mask is None: | |
| context_mask = torch.ones(B, context.shape[-2], device=context.device).bool() | |
| mask = torch.cat([context_mask, x_mask], dim=1) | |
| return mask | |
| def forward(self, x, context, x_mask=None, context_mask=None, extras=0): | |
| B, Lx, C = x.shape | |
| _, Lc, _ = context.shape | |
| if x_mask is not None or context_mask is not None: | |
| mask = self._cat_mask(x, context, | |
| x_mask=x_mask, | |
| context_mask=context_mask) | |
| shape = [B, Lx+Lc, C] | |
| mask_binary = create_mask(q_shape=shape, k_shape=shape, | |
| device=x.device, | |
| q_mask=None, k_mask=mask) | |
| else: | |
| mask_binary = None | |
| qx, kx, vx = self.to_qx(x), self.to_kx(x), self.to_vx(x) | |
| qc, kc, vc = self.to_qc(context), self.to_kc(context), self.to_vc(context) | |
| qx, kx, vx = map(lambda t: einops.rearrange(t, 'B L (H D) -> B H L D', | |
| H=self.num_heads), [qx, kx, vx]) | |
| qc, kc, vc = map(lambda t: einops.rearrange(t, 'B L (H D) -> B H L D', | |
| H=self.num_heads), [qc, kc, vc]) | |
| qx, kx = self.norm_qx(qx), self.norm_kx(kx) | |
| qc, kc = self.norm_qc(qc), self.norm_kc(kc) | |
| q, k, v = (torch.cat([qc, qx], dim=2), | |
| torch.cat([kc, kx], dim=2), | |
| torch.cat([vc, vx], dim=2)) | |
| q, k = self._rotary(q, k, extras) | |
| x = self._attn(q, k, v, mask_binary) | |
| context, x = x[:, :Lc, :], x[:, Lc:, :] | |
| x = self.proj_x(x) | |
| x = self.proj_drop_x(x) | |
| context = self.proj_c(context) | |
| context = self.proj_drop_c(context) | |
| return x, context |