A single transformer layer. Transformer layer takes input with size [s, b, h] and returns an output of the same size.
| 413 | |
| 414 | |
| 415 | class GLMBlock(torch.nn.Module): |
| 416 | """A single transformer layer. |
| 417 | |
| 418 | Transformer layer takes input with size [s, b, h] and returns an |
| 419 | output of the same size. |
| 420 | """ |
| 421 | |
| 422 | def __init__(self, config: ChatGLMConfig, layer_number, device=None): |
| 423 | super(GLMBlock, self).__init__() |
| 424 | self.layer_number = layer_number |
| 425 | |
| 426 | self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm |
| 427 | |
| 428 | self.fp32_residual_connection = config.fp32_residual_connection |
| 429 | |
| 430 | LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm |
| 431 | # Layernorm on the input data. |
| 432 | self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, |
| 433 | dtype=config.torch_dtype) |
| 434 | |
| 435 | # Self attention. |
| 436 | self.self_attention = SelfAttention(config, layer_number, device=device) |
| 437 | self.hidden_dropout = config.hidden_dropout |
| 438 | |
| 439 | # Layernorm on the attention output |
| 440 | self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, |
| 441 | dtype=config.torch_dtype) |
| 442 | |
| 443 | # MLP |
| 444 | self.mlp = MLP(config, device=device) |
| 445 | |
| 446 | def forward( |
| 447 | self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True, |
| 448 | ): |
| 449 | # hidden_states: [s, b, h] |
| 450 | |
| 451 | # Layer norm at the beginning of the transformer layer. |
| 452 | layernorm_output = self.input_layernorm(hidden_states) |
| 453 | # Self attention. |
| 454 | attention_output, kv_cache = self.self_attention( |
| 455 | layernorm_output, |
| 456 | attention_mask, |
| 457 | rotary_pos_emb, |
| 458 | kv_cache=kv_cache, |
| 459 | use_cache=use_cache |
| 460 | ) |
| 461 | |
| 462 | # Residual connection. |
| 463 | if self.apply_residual_connection_post_layernorm: |
| 464 | residual = layernorm_output |
| 465 | else: |
| 466 | residual = hidden_states |
| 467 | |
| 468 | layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training) |
| 469 | layernorm_input = residual + layernorm_input |
| 470 | |
| 471 | # Layer norm post the self attention. |
| 472 | layernorm_output = self.post_attention_layernorm(layernorm_input) |