(self, config)
| 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)) |
nothing calls this directly
no test coverage detected