Fix TF Causal LM models' returned logits
Author: ydshieh
Pull request: yes

# What does this PR do?

Fix TF causal LM models' returned logits

## Details

In TF causal LM models, the returned logits is the one being cut

```
        if inputs["labels"] is not None:
            # shift labels to the left and cut last logit token
            **logits = logits[:, :-1]**
            labels = inputs["labels"][:, 1:]
            loss = self.hf_compute_loss(labels=labels, logits=logits)

        return TFCausalLMOutputWithPast(
            loss=loss,
            **logits=logits,**
```

While for PyTorch causal LM models, the original logits is returned

```
        if labels is not None:
            # Shift so that tokens < n predict n
            **shift_logits = lm_logits[..., :-1, :].contiguous()**
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

        return CausalLMOutputWithPast(
            loss=loss,
            **logits=lm_logits,**
```

This PR fixes this inconsistency + test the cases where `labels` is passed in PT/TF equivalence test.

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Could @Rocketknight1 or @gante give this a look?

To summarize: In causal LM models, logits and labels are offset by 1 before loss computations. In PT, the pre-shift logits are returned in the output, but in TF the post-shift logits are returned instead. This seems like a huge bug - I don't understand how several tf-pt equivalence tests didn't fail before, and how things like `generate()` didn't completely fail as a result too. I need to investigate the code more deeply to understand how this ever worked.

Ah, actually, `generate()` would not pass labels and so the shift would never happen.

> tf-pt equivalence tests

 tf-pt equivalence tests  don't test with `labels` as far as I know. And this PR tries to add it too.

@ydshieh that makes sense! And yes, in most use-cases when people are passing `labels`, they're usually interested in the `loss` output and not the `logits` output. So that probably explains how this was invisible for so long. Great work!

I make this PR as draft again, because there are other **potential PT/TF inconsistency** uncaught.
(The test inputs are randomly generated, so we don't always have the same outputs each run)

Here is the failed output in the test

```
                # Some models require extra condition to return loss. For example, `BertForPreTraining` requires both
                # `labels` and `next_sentence_label`.
                # Moreover, some PT models return loss while the corresponding TF/Flax models don't.
                if tf_loss is not None and pt_loss is not None:
    
                    tf_loss = tf.math.reduce_mean(tf_loss).numpy()
                    pt_loss = pt_loss.numpy()
    
                    tf_nans = np.copy(np.isnan(tf_loss))
                    pt_nans = np.copy(np.isnan(pt_loss))
                    # the 2 losses need to be both nan or both not nan
                    # (`TapasForQuestionAnswering` gives nan loss here)
>                   self.assertEqual(tf_nans, pt_nans)
E                   AssertionError: array(False) != array(True)
```

> I make this PR as draft again, because there are other **potential PT/TF inconsistency** uncaught. (The test inputs are randomly generated, so we don't always have the same outputs each run)

Having random tests might actually uncover interesting corner cases, but failures risk being overlooked if we can't get the random seed back 🤔  @sgugger , what's our go-to policy with respect to tests -- as deterministic as possible, or is randomness embraced?

It would be best to keep those random indeed.

> It would be best to keep those random indeed.

We can discuss about this for another PR. Here the issues I can identify are more deterministic -- one issue is coming from a sequence with all attention mask being 0 (although randomly generated in this case).

So for that particular case, we can  just add an extra test case (with such an attention mask).

See #15326 if you are interested - but it is not TF related.

Woow, great catch @ydshieh !

Hi, I reverted the change in `test_pt_tf_model_equivalence` - it requires some time to fix the issues I could find. I don't want to this PR being blocked due to this new version of test. If you feel OK, we can merge this PR for now. Thank you!

As always, thank you @ydshieh, this was a great contribution. Will you be working on the improved equivalence test? (If not, I might pick it up :) )

> As always, thank you @ydshieh, this was a great contribution. Will you be working on the improved equivalence test? (If not, I might pick it up :) )

I already have one version (that is used for hunting the bugs :-) ). Would definitely open a PR to include it once other issues are fixed (so the test won't fail)

basically, the new version checks in a recursively way:

tuple --> check each element inside it
tensor --> check shape, and the abs max diff of tensor elements

A common issue is that the tuples returned by PT/TF have different length. Another one is the tensors have different shape.

If you want to go ahead, I am fine with it though :-) 

No, by all means, we are very grateful that you are working on it ❤️