| 451 | |
| 452 | |
| 453 | class CoreAttention(torch.nn.Module): |
| 454 | def __init__(self, config: ChatGLMConfig, layer_number): |
| 455 | super(CoreAttention, self).__init__() |
| 456 | |
| 457 | self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling |
| 458 | self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32 |
| 459 | if self.apply_query_key_layer_scaling: |
| 460 | self.attention_softmax_in_fp32 = True |
| 461 | self.layer_number = max(1, layer_number) |
| 462 | |
| 463 | projection_size = config.kv_channels * config.num_attention_heads |
| 464 | |
| 465 | # Per attention head and per partition values. |
| 466 | self.hidden_size_per_partition = projection_size |
| 467 | self.hidden_size_per_attention_head = projection_size // config.num_attention_heads |
| 468 | self.num_attention_heads_per_partition = config.num_attention_heads |
| 469 | |
| 470 | coeff = None |
| 471 | self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) |
| 472 | if self.apply_query_key_layer_scaling: |
| 473 | coeff = self.layer_number |
| 474 | self.norm_factor *= coeff |
| 475 | self.coeff = coeff |
| 476 | |
| 477 | self.attention_dropout = torch.nn.Dropout(config.attention_dropout) |
| 478 | |
| 479 | def forward(self, query_layer, key_layer, value_layer, attention_mask): |
| 480 | pytorch_major_version = int(torch.__version__.split('.')[0]) |
| 481 | if pytorch_major_version >= 2: |
| 482 | query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]] |
| 483 | if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]: |
| 484 | context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer, |
| 485 | is_causal=True) |
| 486 | else: |
| 487 | if attention_mask is not None: |
| 488 | attention_mask = ~attention_mask |
| 489 | context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer, |
| 490 | attention_mask) |
| 491 | context_layer = context_layer.permute(2, 0, 1, 3) |
| 492 | new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) |
| 493 | context_layer = context_layer.reshape(*new_context_layer_shape) |
| 494 | else: |
| 495 | # Raw attention scores |
| 496 | |
| 497 | # [b, np, sq, sk] |
| 498 | output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0)) |
| 499 | |
| 500 | # [sq, b, np, hn] -> [sq, b * np, hn] |
| 501 | query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1) |
| 502 | # [sk, b, np, hn] -> [sk, b * np, hn] |
| 503 | key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1) |
| 504 | |
| 505 | # preallocting input tensor: [b * np, sq, sk] |
| 506 | matmul_input_buffer = torch.empty( |
| 507 | output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype, |
| 508 | device=query_layer.device |
| 509 | ) |
| 510 | |