A single transformer layer. Transformore layer takes input with size [b, s, h] and returns an output of the same size.
| 571 | |
| 572 | |
| 573 | class ParallelTransformerLayer(MegatronModule): |
| 574 | """A single transformer layer. |
| 575 | |
| 576 | Transformore layer takes input with size [b, s, h] and returns an |
| 577 | output of the same size. |
| 578 | """ |
| 579 | |
| 580 | def __init__(self, init_method, |
| 581 | output_layer_init_method, layer_number): |
| 582 | args = get_args() |
| 583 | |
| 584 | super(ParallelTransformerLayer, self).__init__() |
| 585 | self.layer_number = layer_number |
| 586 | |
| 587 | self.apply_residual_connection_post_layernorm \ |
| 588 | = args.apply_residual_connection_post_layernorm |
| 589 | |
| 590 | # Layernorm on the input data. |
| 591 | self.input_layernorm = LayerNorm( |
| 592 | args.hidden_size, |
| 593 | eps=args.layernorm_epsilon) |
| 594 | |
| 595 | # Self attention. |
| 596 | self.attention = ParallelSelfAttention(init_method, |
| 597 | output_layer_init_method, |
| 598 | layer_number) |
| 599 | self.hidden_dropout = args.hidden_dropout |
| 600 | self.bias_dropout_fusion = args.bias_dropout_fusion |
| 601 | |
| 602 | # Layernorm on the input data. |
| 603 | self.post_attention_layernorm = LayerNorm( |
| 604 | args.hidden_size, |
| 605 | eps=args.layernorm_epsilon) |
| 606 | if hasattr(args, 'attention_upweight'): |
| 607 | self.attention_upweight = args.attention_upweight |
| 608 | else: |
| 609 | self.attention_upweight = None |
| 610 | if hasattr(args, 'ln_fp16'): |
| 611 | self.ln_fp16 = args.ln_fp16 |
| 612 | else: |
| 613 | self.ln_fp16 = False |
| 614 | # MLP |
| 615 | self.mlp = ParallelMLP(init_method, |
| 616 | output_layer_init_method, |
| 617 | scale=2 if args.compress else 4) |
| 618 | |
| 619 | def forward( |
| 620 | self, |
| 621 | hidden_states, |
| 622 | attention_mask, |
| 623 | layer_past=None, |
| 624 | get_key_value=False, |
| 625 | prompt_length=None, |
| 626 | context_length=None, |
| 627 | ): |
| 628 | # hidden_states: [b, s, h] |
| 629 | if self.ln_fp16: |
| 630 | layernorm_output = self.input_layernorm(hidden_states) |