| 59 | |
| 60 | |
| 61 | class MossAttention(nn.Module): |
| 62 | def __init__(self, config): |
| 63 | super().__init__() |
| 64 | |
| 65 | max_positions = config.max_position_embeddings |
| 66 | self.register_buffer( |
| 67 | "causal_mask", |
| 68 | torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view( |
| 69 | 1, 1, max_positions, max_positions |
| 70 | ), |
| 71 | ) |
| 72 | |
| 73 | self.attn_dropout = nn.Dropout(config.attn_pdrop) |
| 74 | self.resid_dropout = nn.Dropout(config.resid_pdrop) |
| 75 | |
| 76 | self.embed_dim = config.hidden_size |
| 77 | self.num_attention_heads = config.num_attention_heads |
| 78 | self.head_dim = self.embed_dim // self.num_attention_heads |
| 79 | if self.head_dim * self.num_attention_heads != self.embed_dim: |
| 80 | raise ValueError( |
| 81 | f"embed_dim must be divisible by num_attention_heads (got `embed_dim`: {self.embed_dim} and" |
| 82 | f" `num_attention_heads`: {self.num_attention_heads})." |
| 83 | ) |
| 84 | self.scale_attn = torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype()) |
| 85 | self.qkv_proj = nn.Linear(self.embed_dim, self.embed_dim * 3, bias=False) |
| 86 | |
| 87 | self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False) |
| 88 | self.rotary_dim = config.rotary_dim |
| 89 | pos_embd_dim = self.rotary_dim or self.embed_dim |
| 90 | self.embed_positions = create_sinusoidal_positions(max_positions, pos_embd_dim) |
| 91 | |
| 92 | def _split_heads(self, x, n_head, dim_head, mp_num): |
| 93 | reshaped = x.reshape(x.shape[:-1] + (n_head // mp_num, dim_head)) |
| 94 | reshaped = reshaped.reshape(x.shape[:-2] + (-1,) + reshaped.shape[-1:]) |
| 95 | return reshaped |
| 96 | |
| 97 | def _merge_heads(self, tensor, num_attention_heads, attn_head_size): |
| 98 | """ |
| 99 | Merges attn_head_size dim and num_attn_heads dim into n_ctx |
| 100 | """ |
| 101 | if len(tensor.shape) == 5: |
| 102 | tensor = tensor.permute(0, 1, 3, 2, 4).contiguous() |
| 103 | elif len(tensor.shape) == 4: |
| 104 | tensor = tensor.permute(0, 2, 1, 3).contiguous() |
| 105 | else: |
| 106 | raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}") |
| 107 | new_shape = tensor.size()[:-2] + (num_attention_heads * attn_head_size,) |
| 108 | return tensor.view(new_shape) |
| 109 | |
| 110 | def _attn( |
| 111 | self, |
| 112 | query, |
| 113 | key, |
| 114 | value, |
| 115 | attention_mask=None, |
| 116 | head_mask=None, |
| 117 | ): |
| 118 | # compute causal mask from causal mask buffer |