Transformer class.
| 487 | |
| 488 | |
| 489 | class GLMTransformer(torch.nn.Module): |
| 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] |
| 517 | |
| 518 | def forward( |
| 519 | self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None, |
| 520 | use_cache: Optional[bool] = True, |
| 521 | output_hidden_states: Optional[bool] = False, |
| 522 | ): |
| 523 | if not kv_caches: |
| 524 | kv_caches = [None for _ in range(self.num_layers)] |
| 525 | presents = () if use_cache else None |
| 526 | if self.gradient_checkpointing and self.training: |
| 527 | if use_cache: |
| 528 | # logger.warning_once( |
| 529 | # "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." |
| 530 | # ) |
| 531 | use_cache = False |
| 532 | |
| 533 | all_self_attentions = None |
| 534 | all_hidden_states = () if output_hidden_states else None |
| 535 | for index in range(self.num_layers): |
| 536 | if output_hidden_states: |
| 537 | all_hidden_states = all_hidden_states + (hidden_states,) |
| 538 | |
| 539 | layer = self._get_layer(index) |
| 540 | if self.gradient_checkpointing and self.training: |
| 541 | layer_ret = torch.utils.checkpoint.checkpoint( |
| 542 | layer, |
| 543 | hidden_states, |
| 544 | attention_mask, |
| 545 | rotary_pos_emb, |
| 546 | kv_caches[index], |
nothing calls this directly
no outgoing calls
no test coverage detected