A single transformer layer. Transformer layer takes input with size [s, b, h] and returns an output of the same size.
| 503 | |
| 504 | |
| 505 | class GLMBlock(torch.nn.Module): |
| 506 | """A single transformer layer. |
| 507 | |
| 508 | Transformer layer takes input with size [s, b, h] and returns an |
| 509 | output of the same size. |
| 510 | """ |
| 511 | |
| 512 | def __init__(self, config: ChatGLMConfig, layer_number, device=None): |
| 513 | super(GLMBlock, self).__init__() |
| 514 | self.layer_number = layer_number |
| 515 | |
| 516 | self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm |
| 517 | |
| 518 | self.fp32_residual_connection = config.fp32_residual_connection |
| 519 | |
| 520 | LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm |
| 521 | # Layernorm on the input data. |
| 522 | self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, |
| 523 | dtype=config.torch_dtype) |
| 524 | |
| 525 | # Self attention. |
| 526 | self.self_attention = SelfAttention(config, layer_number, device=device) |
| 527 | self.hidden_dropout = config.hidden_dropout |
| 528 | |
| 529 | # Layernorm on the attention output |
| 530 | self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, |
| 531 | dtype=config.torch_dtype) |
| 532 | |
| 533 | # MLP |
| 534 | self.mlp = MLP(config, device=device) |
| 535 | |
| 536 | def forward( |
| 537 | self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True, |
| 538 | ): |
| 539 | # hidden_states: [s, b, h] |
| 540 | |
| 541 | # Layer norm at the beginning of the transformer layer. |
| 542 | layernorm_output = self.input_layernorm(hidden_states) |
| 543 | # Self attention. |
| 544 | attention_output, kv_cache = self.self_attention( |
| 545 | layernorm_output, |
| 546 | attention_mask, |
| 547 | rotary_pos_emb, |
| 548 | kv_cache=kv_cache, |
| 549 | use_cache=use_cache |
| 550 | ) |
| 551 | |
| 552 | # Residual connection. |
| 553 | if self.apply_residual_connection_post_layernorm: |
| 554 | residual = layernorm_output |
| 555 | else: |
| 556 | residual = hidden_states |
| 557 | |
| 558 | layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training) |
| 559 | layernorm_input = residual + layernorm_input |
| 560 | |
| 561 | # Layer norm post the self attention. |
| 562 | layernorm_output = self.post_attention_layernorm(layernorm_input) |