Parallel self-attention layer abstract class. Self-attention layer takes input with size [s, b, h] and returns output of the same size.
| 569 | |
| 570 | |
| 571 | class SelfAttention(torch.nn.Module): |
| 572 | """Parallel self-attention layer abstract class. |
| 573 | |
| 574 | Self-attention layer takes input with size [s, b, h] |
| 575 | and returns output of the same size. |
| 576 | """ |
| 577 | |
| 578 | def __init__(self, config: ChatGLMConfig, layer_number, device=None): |
| 579 | super(SelfAttention, self).__init__() |
| 580 | self.layer_number = max(1, layer_number) |
| 581 | |
| 582 | self.projection_size = config.kv_channels * config.num_attention_heads |
| 583 | |
| 584 | # Per attention head and per partition values. |
| 585 | self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads |
| 586 | self.num_attention_heads_per_partition = config.num_attention_heads |
| 587 | |
| 588 | self.multi_query_attention = config.multi_query_attention |
| 589 | self.qkv_hidden_size = 3 * self.projection_size |
| 590 | if self.multi_query_attention: |
| 591 | self.num_multi_query_groups_per_partition = config.multi_query_group_num |
| 592 | self.qkv_hidden_size = ( |
| 593 | self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num |
| 594 | ) |
| 595 | self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size, |
| 596 | bias=config.add_bias_linear or config.add_qkv_bias, |
| 597 | device=device, **_config_to_kwargs(config) |
| 598 | ) |
| 599 | |
| 600 | self.core_attention = CoreAttention(config, self.layer_number) |
| 601 | |
| 602 | # Output. |
| 603 | self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear, |
| 604 | device=device, **_config_to_kwargs(config) |
| 605 | ) |
| 606 | |
| 607 | def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None): |
| 608 | if self.multi_query_attention: |
| 609 | num_attention_heads = self.num_multi_query_groups_per_partition |
| 610 | else: |
| 611 | num_attention_heads = self.num_attention_heads_per_partition |
| 612 | return torch.empty( |
| 613 | inference_max_sequence_len, |
| 614 | batch_size, |
| 615 | num_attention_heads, |
| 616 | self.hidden_size_per_attention_head, |
| 617 | dtype=dtype, |
| 618 | device=device, |
| 619 | ) |
| 620 | |
| 621 | def forward( |
| 622 | self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True |
| 623 | ): |
| 624 | # hidden_states: [sq, b, h] |
| 625 | |
| 626 | # ================================================= |
| 627 | # Pre-allocate memory for key-values for inference. |
| 628 | # ================================================= |