| 506 | |
| 507 | |
| 508 | class TelechatBlock(nn.Module): |
| 509 | def __init__(self, config: TelechatConfig, layer_idx): |
| 510 | super().__init__() |
| 511 | hidden_size = config.hidden_size |
| 512 | |
| 513 | self.input_layernorm = MixedFusedRMSNorm(hidden_size, eps=config.layer_norm_epsilon) |
| 514 | self.num_heads = config.n_head |
| 515 | self.layer_idx = layer_idx |
| 516 | self.self_attention = TelechatAttention(config, layer_idx) |
| 517 | self.post_attention_layernorm = MixedFusedRMSNorm(hidden_size, eps=config.layer_norm_epsilon) |
| 518 | |
| 519 | self.mlp = TelechatMLP(config) |
| 520 | |
| 521 | self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm |
| 522 | self.hidden_dropout = config.hidden_dropout |
| 523 | |
| 524 | def forward( |
| 525 | self, |
| 526 | hidden_states: torch.Tensor, |
| 527 | attention_mask: torch.Tensor, |
| 528 | layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, |
| 529 | use_cache: bool = False, |
| 530 | output_attentions: bool = False, |
| 531 | ): |
| 532 | layernorm_output = self.input_layernorm(hidden_states) |
| 533 | if self.apply_residual_connection_post_layernorm: |
| 534 | residual = layernorm_output |
| 535 | else: |
| 536 | residual = hidden_states |
| 537 | |
| 538 | attn_outputs = self.self_attention( |
| 539 | layernorm_output, |
| 540 | residual, |
| 541 | layer_past=layer_past, |
| 542 | attention_mask=attention_mask, |
| 543 | use_cache=use_cache, |
| 544 | output_attentions=output_attentions, |
| 545 | ) |
| 546 | |
| 547 | attention_output = attn_outputs[0] |
| 548 | outputs = attn_outputs[1:] |
| 549 | layernorm_output = self.post_attention_layernorm(attention_output) |
| 550 | |
| 551 | if self.apply_residual_connection_post_layernorm: |
| 552 | residual = layernorm_output |
| 553 | else: |
| 554 | residual = attention_output |
| 555 | output = self.mlp(layernorm_output, residual) |
| 556 | |
| 557 | if use_cache: |
| 558 | outputs = (output,) + outputs |
| 559 | else: |
| 560 | outputs = (output,) + outputs[1:] |
| 561 | |
| 562 | return outputs |
| 563 | |
| 564 | |
| 565 | class TelechatPreTrainedModel(PreTrainedModel): |