A single top query layer. Top query layer takes input with size [b, s, h] and returns an output of the same size.
| 571 | |
| 572 | |
| 573 | class TopQueryLayer(torch.nn.Module): |
| 574 | """A single top query layer. |
| 575 | Top query layer takes input with size [b, s, h] and returns an |
| 576 | output of the same size. |
| 577 | """ |
| 578 | |
| 579 | def __init__( |
| 580 | self, |
| 581 | hidden_size, |
| 582 | num_attention_heads, |
| 583 | layer_number, |
| 584 | layernorm_epsilon=1e-5, |
| 585 | ): |
| 586 | super(TopQueryLayer, self).__init__() |
| 587 | self.hidden_size = hidden_size |
| 588 | self.num_attention_heads = num_attention_heads |
| 589 | self.layernorm_epsilon = layernorm_epsilon |
| 590 | self.layer_number = layer_number |
| 591 | |
| 592 | # Use FP32 for Layernorm |
| 593 | self.input_layernorm = torch.nn.LayerNorm(self.hidden_size, |
| 594 | eps=self.layernorm_epsilon) |
| 595 | |
| 596 | # Self attention. |
| 597 | self.attention = TopQuerySelfAttention(self.hidden_size, |
| 598 | self.num_attention_heads, |
| 599 | self.layer_number) |
| 600 | # Layernorm on the input data. |
| 601 | self.post_attention_layernorm = torch.nn.LayerNorm(self.hidden_size, |
| 602 | eps=self.layernorm_epsilon) |
| 603 | |
| 604 | # MLP |
| 605 | self.mlp = MLP(self.hidden_size) |
| 606 | |
| 607 | def forward( |
| 608 | self, |
| 609 | hidden_states, |
| 610 | query_hidden_state, |
| 611 | attention_mask, |
| 612 | layer_past=None, |
| 613 | get_key_value=False, |
| 614 | prompt_length=None, |
| 615 | context_length=None, |
| 616 | ): |
| 617 | # hidden_states: [b, s, h] |
| 618 | assert query_hidden_state != None |
| 619 | |
| 620 | # Use FP32 for Layernorm |
| 621 | # layernorm_output = self.input_layernorm(hidden_states.float()).half() |
| 622 | layernorm_output = self.input_layernorm(hidden_states) |
| 623 | |
| 624 | # Self attention. |
| 625 | attention_output = self.attention(layernorm_output, |
| 626 | query_hidden_state, |
| 627 | attention_mask, |
| 628 | layer_past=layer_past, |
| 629 | get_key_value=get_key_value, |
| 630 | prompt_length=prompt_length, |