(self, config: ChatGLMConfig, layer_number, device=None)
| 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: |
nothing calls this directly
no test coverage detected