| 589 | |
| 590 | |
| 591 | class GLMBlock(torch.nn.Module): |
| 592 | def __init__( |
| 593 | self, |
| 594 | hidden_size, |
| 595 | num_attention_heads, |
| 596 | layernorm_epsilon, |
| 597 | layer_id, |
| 598 | inner_hidden_size=None, |
| 599 | hidden_size_per_attention_head=None, |
| 600 | layernorm=LayerNorm, |
| 601 | use_bias=True, |
| 602 | params_dtype=torch.float, |
| 603 | num_layers=28, |
| 604 | position_encoding_2d=True |
| 605 | ): |
| 606 | super(GLMBlock, self).__init__() |
| 607 | # Set output layer initialization if not provided. |
| 608 | |
| 609 | self.layer_id = layer_id |
| 610 | |
| 611 | # Layernorm on the input data. |
| 612 | self.input_layernorm = layernorm(hidden_size, eps=layernorm_epsilon) |
| 613 | |
| 614 | self.position_encoding_2d = position_encoding_2d |
| 615 | |
| 616 | # Self attention. |
| 617 | self.attention = SelfAttention( |
| 618 | hidden_size, |
| 619 | num_attention_heads, |
| 620 | layer_id, |
| 621 | hidden_size_per_attention_head=hidden_size_per_attention_head, |
| 622 | bias=use_bias, |
| 623 | params_dtype=params_dtype, |
| 624 | position_encoding_2d=self.position_encoding_2d |
| 625 | ) |
| 626 | |
| 627 | # Layernorm on the input data. |
| 628 | self.post_attention_layernorm = layernorm(hidden_size, eps=layernorm_epsilon) |
| 629 | |
| 630 | self.num_layers = num_layers |
| 631 | |
| 632 | # GLU |
| 633 | self.mlp = GLU( |
| 634 | hidden_size, |
| 635 | inner_hidden_size=inner_hidden_size, |
| 636 | bias=use_bias, |
| 637 | layer_id=layer_id, |
| 638 | params_dtype=params_dtype, |
| 639 | ) |
| 640 | |
| 641 | def forward( |
| 642 | self, |
| 643 | hidden_states: torch.Tensor, |
| 644 | position_ids, |
| 645 | attention_mask: torch.Tensor, |
| 646 | layer_id, |
| 647 | layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, |
| 648 | use_cache: bool = False, |