A single top query layer. Top query layer takes input with size [b, s, h] and returns an output of the same size.
| 730 | |
| 731 | |
| 732 | class ParallelTopQueryLayer(MegatronModule): |
| 733 | """A single top query layer. |
| 734 | |
| 735 | Top query layer takes input with size [b, s, h] and returns an |
| 736 | output of the same size. |
| 737 | """ |
| 738 | |
| 739 | def __init__(self, init_method, |
| 740 | output_layer_init_method, layer_number): |
| 741 | args = get_args() |
| 742 | |
| 743 | super(ParallelTopQueryLayer, self).__init__() |
| 744 | self.layer_number = layer_number |
| 745 | |
| 746 | self.apply_residual_connection_post_layernorm \ |
| 747 | = args.apply_residual_connection_post_layernorm |
| 748 | |
| 749 | # Layernorm on the input data. |
| 750 | self.input_layernorm = LayerNorm( |
| 751 | args.hidden_size, |
| 752 | eps=args.layernorm_epsilon) |
| 753 | |
| 754 | # Self attention. |
| 755 | self.attention = ParallelTopQuerySelfAttention(init_method, |
| 756 | output_layer_init_method, |
| 757 | layer_number) |
| 758 | |
| 759 | self.hidden_dropout = args.hidden_dropout |
| 760 | self.bias_dropout_fusion = args.bias_dropout_fusion |
| 761 | |
| 762 | # Layernorm on the input data. |
| 763 | self.post_attention_layernorm = LayerNorm( |
| 764 | args.hidden_size, |
| 765 | eps=args.layernorm_epsilon) |
| 766 | |
| 767 | if hasattr(args, 'ln_fp16'): |
| 768 | self.ln_fp16 = args.ln_fp16 |
| 769 | else: |
| 770 | self.ln_fp16 = False |
| 771 | |
| 772 | # MLP |
| 773 | self.mlp = ParallelMLP(init_method, |
| 774 | output_layer_init_method) |
| 775 | |
| 776 | def forward( |
| 777 | self, |
| 778 | hidden_states, |
| 779 | query_hidden_state, |
| 780 | attention_mask, |
| 781 | layer_past=None, |
| 782 | get_key_value=False, |
| 783 | prompt_length=None, |
| 784 | context_length=None, |
| 785 | ): |
| 786 | # hidden_states: [b, s, h] |
| 787 | assert query_hidden_state != None |
| 788 | |
| 789 | # Layer norm at the beginning of the transformer layer. |