BertForPreTraining should compute only <MASKED> prediction_scores
Author: songsuoyuan
Pull request: no
Labels: name: wontfix

# 🚀 Feature request
In class transformers.BertForPreTraining, the forward compute all prediction_scores. In fact, we may only calculate the prediction_scores on the <MASKED> tokens to save some computational cost.

## Motivation
In source code of class BertForPreTraining:
```python
outputs = self.bert(input_ids, attention_mask, ...)
sequence_output, pooled_output = outputs[:2]
prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
```
It compute `prediction_scores` using `sequence_output` as the input of function `self.cls()`. I was wondering if we can gather the MASKED index (not equal to -100) and only pass the MASKED sequence_output into function `self.cls` and return the MASKED prediction_scores. Then pass them into `CrossEntropy()` to compute the `masked_lm_loss`. In this way, we can save some computational cost and partially relief the OOM problem in GPU.

## My suggestion
We may change the code to
```python
outputs = self.bert(input_ids, attention_mask, ...)
sequence_output, pooled_output = outputs[:2] 
# before gather_indexes, size of sequence_output: (batch_size, sequence_length, hidden_size)
# after gather_indexes, size of sequence_output: (batch_size, masked_lm_nums, hidden_size)
sequence_output = gather_indexes(sequence_output, masked_ln_labels)
prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
```
Then pass prediction_scores and corresponding `masked_lm_labels` (also need to gather the indexes) into function CrossEntropy to compute the masked_lm_loss.

Note: if we adopt this method, the return of BertForPreTraining is changed since we will not compute prediction_scores on all tokens.

Comments:
I had little modification on the source code and solve the problem. (But I may changed the behavior and the return format of class BertForPretraining)

```python
def gather_indexes_auto(sequence_tensor, masked_lm_labels):
    """Gathers the vectors according to masked_lm_labels over a minibatch.
    Input 
        sequence_tensor:  (batch_size, sequence_length, hidden_size)
        masked_lm_labels: (batch_size, sequence_length)
    Output
        output_tensor:    (-1, hidden_size)
        output_lm_labels: (-1, )
    """
    batch_size = sequence_tensor.size(0)
    sequence_length = sequence_tensor.size(1)
    hidden_size = sequence_tensor.size(2)
    # Flatten sequence_tensor into  (-1, hidden_size)
    # Flatten masked_lm_labels into (-1, )
    sequence_tensor_flat = sequence_tensor.view(batch_size*sequence_length, hidden_size)
    masked_lm_labels_flat = masked_lm_labels.view(-1)
    # Get non -100 index 
    # Note: the input index of torch.index_select is 1-D tensor
    masked_lm_location = masked_lm_labels_flat.ge(0).nonzero().view(-1)
    # Select corresponding values 
    output_tensor = torch.index_select(sequence_tensor_flat, dim=0, index=masked_lm_location)
    output_lm_labels = torch.index_select(masked_lm_labels_flat, dim=0, index=masked_lm_location)
    return output_tensor, output_lm_labels
```
And in class `BertForPreTraining`
```
class BertForPreTraining(BertPreTrainedModel):
    ...
    def forward(...):
        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )

        sequence_output, pooled_output = outputs[:2]
        sequence_output, output_lm_labels = gather_indexes_auto(sequence_output, masked_lm_labels)

        prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

        outputs = (prediction_scores, seq_relationship_score,) + outputs[
            2:
        ]  # add hidden states and attention if they are here

        if masked_lm_labels is not None and next_sentence_label is not None:
            loss_fct = CrossEntropyLoss()
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), output_lm_labels.view(-1))
            next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
            total_loss = masked_lm_loss + next_sentence_loss
            outputs = (total_loss,) 

        return outputs  # (loss), 
```
Note that this code only compute and return the loss on <MASKED> labels and thus save a lot computations and GPU memories.
Before this change, I can only run BERT-MEDIUM (L8 H512 A8) on batch_size = 64 using P100 (16G RAM), after this change, I can run the pretraining using batch_size = 128.

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.