| 552 | |
| 553 | |
| 554 | class BertEncoder(nn.Module): |
| 555 | def __init__(self, config): |
| 556 | super().__init__() |
| 557 | self.config = config |
| 558 | self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)]) |
| 559 | self.gradient_checkpointing = False |
| 560 | |
| 561 | def forward( |
| 562 | self, |
| 563 | hidden_states: torch.Tensor, |
| 564 | attention_mask: Optional[torch.FloatTensor] = None, |
| 565 | head_mask: Optional[torch.FloatTensor] = None, |
| 566 | encoder_hidden_states: Optional[torch.FloatTensor] = None, |
| 567 | encoder_attention_mask: Optional[torch.FloatTensor] = None, |
| 568 | past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, |
| 569 | use_cache: Optional[bool] = None, |
| 570 | output_attentions: Optional[bool] = False, |
| 571 | output_hidden_states: Optional[bool] = False, |
| 572 | return_dict: Optional[bool] = True, |
| 573 | ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]: |
| 574 | all_hidden_states = () if output_hidden_states else None |
| 575 | all_self_attentions = () if output_attentions else None |
| 576 | all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None |
| 577 | |
| 578 | next_decoder_cache = () if use_cache else None |
| 579 | for i, layer_module in enumerate(self.layer): |
| 580 | if output_hidden_states: |
| 581 | all_hidden_states = all_hidden_states + (hidden_states,) |
| 582 | |
| 583 | layer_head_mask = head_mask[i] if head_mask is not None else None |
| 584 | past_key_value = past_key_values[i] if past_key_values is not None else None |
| 585 | |
| 586 | if self.gradient_checkpointing and self.training: |
| 587 | |
| 588 | if use_cache: |
| 589 | logger.warning( |
| 590 | "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." |
| 591 | ) |
| 592 | use_cache = False |
| 593 | |
| 594 | def create_custom_forward(module): |
| 595 | def custom_forward(*inputs): |
| 596 | return module(*inputs, past_key_value, output_attentions) |
| 597 | |
| 598 | return custom_forward |
| 599 | |
| 600 | layer_outputs = torch.utils.checkpoint.checkpoint( |
| 601 | create_custom_forward(layer_module), |
| 602 | hidden_states, |
| 603 | attention_mask, |
| 604 | layer_head_mask, |
| 605 | encoder_hidden_states, |
| 606 | encoder_attention_mask, |
| 607 | ) |
| 608 | else: |
| 609 | layer_outputs = layer_module( |
| 610 | hidden_states, |
| 611 | attention_mask, |