| 490 | """Transformer class.""" |
| 491 | |
| 492 | def __init__(self, config: ChatGLMConfig, device=None): |
| 493 | super(GLMTransformer, self).__init__() |
| 494 | |
| 495 | self.fp32_residual_connection = config.fp32_residual_connection |
| 496 | self.post_layer_norm = config.post_layer_norm |
| 497 | |
| 498 | # Number of layers. |
| 499 | self.num_layers = config.num_layers |
| 500 | |
| 501 | # Transformer layers. |
| 502 | def build_layer(layer_number): |
| 503 | return GLMBlock(config, layer_number, device=device) |
| 504 | |
| 505 | self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)]) |
| 506 | |
| 507 | if self.post_layer_norm: |
| 508 | LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm |
| 509 | # Final layer norm before output. |
| 510 | self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, |
| 511 | dtype=config.torch_dtype) |
| 512 | |
| 513 | self.gradient_checkpointing = False |
| 514 | |
| 515 | def _get_layer(self, layer_number): |
| 516 | return self.layers[layer_number] |