[Trainer][Eval] Why the model output for the first element in eval batch is skipped in logits?

[konradkalita]

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:

transformers/src/transformers/trainer.py

Line 4370 in ca541bd

logits = outputs[1:]

[Rajatavaa]

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:]

[LysandreJik]

cc @muellerzr @SunMarc

[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 .

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.

[konradkalita]

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:

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))

[SunMarc]

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

    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

[konradkalita]

The caveat is that I have a custom compute_loss (subclassed from Trainer as suggested in the doc) which looks like this:

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?

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 #34198 ?

[github-actions]

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 are likely to be ignored.

[SunMarc]

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

[github-actions]

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 are likely to be ignored.