| 283 | """ |
| 284 | |
| 285 | def __init__( |
| 286 | self, |
| 287 | hidden_size, |
| 288 | num_attention_heads, |
| 289 | layer_number, |
| 290 | fp16=True, |
| 291 | attention_softmax_in_fp32=True, |
| 292 | ): |
| 293 | super(TopQuerySelfAttention, self).__init__() |
| 294 | self.hidden_size = hidden_size |
| 295 | self.num_attention_heads = num_attention_heads |
| 296 | self.fp16 = fp16 |
| 297 | self.attention_softmax_in_fp32 = attention_softmax_in_fp32 |
| 298 | self.layer_number = max(1, layer_number) |
| 299 | |
| 300 | assert self.hidden_size % self.num_attention_heads == 0 |
| 301 | self.hidden_size_per_attention_head = int(self.hidden_size // self.num_attention_heads) |
| 302 | |
| 303 | self.query = torch.nn.Linear(self.hidden_size, self.hidden_size) |
| 304 | self.key = torch.nn.Linear(self.hidden_size, self.hidden_size) |
| 305 | self.value = torch.nn.Linear(self.hidden_size, self.hidden_size) |
| 306 | |
| 307 | self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) |
| 308 | self.softmax = torch.nn.Softmax(dim=-1) |
| 309 | |
| 310 | self.dense = torch.nn.Linear(self.hidden_size, self.hidden_size) |
| 311 | |
| 312 | def forward( |
| 313 | self, |