[Trainer][Eval] Why the model output for the first element in eval batch is skipped in logits?
Author: konradkalita
Pull request: no
Labels: name: trainer

I wonder why this is written this way. It causes a bug in my eval pipeline as I loose a one element from every batch in eval loop:

https://github.com/huggingface/transformers/blob/ca541bd4f4d932f486a4116deba833b4ffaebd15/src/transformers/trainer.py#L4370

Comments:
logits = outputs[1:] will return only the logits of all the dimension from the second if you don't want to loose any element then start from outputs[0:]

cc @muellerzr @SunMarc 

Hey @konradkalita, thanks for the report. Are you using transformers from source ? Could you share a reproducer of the issue ? The error might be due to this [PR](https://github.com/huggingface/transformers/pull/34198) .

We probably needs to update the condition to get either output or output[1:] (because the loss can sometimes be in output[0] so we need to do the sclicing) with the following: 
`if (self.label_smoother is not None or self.compute_loss_func is not None) and "labels" in inputs:`

cc @muellerzr as you might refactor a bit the trainer to fix the grad acc issue. 

Thanks for the response @SunMarc. 

>  Are you using transformers from source ?

I use version `4.45.2`

> Could you share a reproducer of the issue ?

Unfortunately I don't have a minimal reproduction code which I could share with you. In my case I train a text encoder with contrastive loss function so I don't have labels passed as separate variable and I go into this `has_labels or loss_without_labels` condition. However, the PR (#34198) you mentioned does not change the most important part of the code  i.e. how the `prediction_step` method is implemented, so not sure if it is relevant.
In this method we have:
```python
if has_labels or loss_without_labels:
    with self.compute_loss_context_manager():
        loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
    loss = loss.mean().detach()

    if isinstance(outputs, dict):
        logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"])
    else:
        logits = outputs[1:]
```
Now as the loss is explicitly returned by `self.compute_loss(...)` in the first element of the returned tuple I don't now if there are conventions where it also return in first element of the tensor (in my case it does not make sense as output is of the shape `(batch_size, seq_length, embedding_size)`)


Can you check what's happening in the compute_loss function and the forward of your model ? 
If we pass an input that contrains `labels` to the model, it will compute the loss and we can fetch it in `outputs[0]`. Hence, we need to do the slicing after. 

Taken from transformers from source
```python 
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.

        Subclass and override for custom behavior.
        """
        if (self.label_smoother is not None or self.compute_loss_func is not None) and "labels" in inputs:
            labels = inputs.pop("labels")
        else:
            labels = None
        # if num_items_in_batch is not None:
        #     inputs["num_items_in_batch"] = num_items_in_batch
        outputs = model(**inputs)
        # Save past state if it exists
        # TODO: this needs to be fixed and made cleaner later.
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

        if labels is not None:
            unwrapped_model = self.accelerator.unwrap_model(model)
            if _is_peft_model(unwrapped_model):
                model_name = unwrapped_model.base_model.model._get_name()
            else:
                model_name = unwrapped_model._get_name()
            # User-defined compute_loss function
            if self.compute_loss_func is not None:
                loss = self.compute_loss_func(outputs, labels, num_items_in_batch=num_items_in_batch)
            elif model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
                loss = self.label_smoother(outputs, labels, shift_labels=True)
            else:
                loss = self.label_smoother(outputs, labels)
        else:
            if isinstance(outputs, dict) and "loss" not in outputs:
                raise ValueError(
                    "The model did not return a loss from the inputs, only the following keys: "
                    f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
                )
            # We don't use .loss here since the model may return tuples instead of ModelOutput.
            ######## Here as you can see, we get it with outputs[0] #######
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]

        return (loss, outputs) if return_outputs else loss
```

The caveat is that I have a custom `compute_loss` (subclassed from Trainer as suggested in the doc) which looks like this:
```python
def compute_loss(self, model, inputs, return_outputs=False):
        outputs = model(**inputs)
        q, p, n = self._splitter.split_embeddings_into_queries_positives_negatives(outputs)
        loss = self._loss(q, p, n)
        return (loss, outputs) if return_outputs else loss
```
Now if I understand you to be consistent with actual `compute_loss` contract I should do it rather like this?
```python
def compute_loss(self, model, inputs, return_outputs=False):
        outputs = model(**inputs)
        q, p, n = self._splitter.split_embeddings_into_queries_positives_negatives(outputs)
        loss = self._loss(q, p, n)
        if return_outputs:
            return loss, {"outputs": outputs, "loss": loss} 
        return loss
```

Then I will not go into this unfortunate branch in `predict_step`. Maybe it could be described  a little more precisely in a docstring of `compute_loss` or maybe it will be done differently after release of https://github.com/huggingface/transformers/pull/34198 ?

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hey @konradkalita , we have added a `compute_loss_func` arg in Trainer, so you don't need to overwrite `computer_loss` anymore ! Would love to get your input on that. Maybe this will also fix your issue 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.