Repository: huggingface/transformers
Pull request: fix: loss computation after embeddings resize - mllama

# What does this PR do?

background: https://github.com/huggingface/transformers/pull/36591

Fixes loss computation when vocab is resized by resize_embeddings. Only vocab size of parent conditionalgenerationclass is modified, hence loss has to be calculated there. Chosen approach until bigger refactor. Solution discussed with @zucchini-nlp in above PR 

Fixes # (issue)
https://github.com/huggingface/transformers/issues/36590

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@zucchini-nlp as discussed solution

Changed files:
filename: src/transformers/models/mllama/modeling_mllama.py
patch: @@ -2056,6 +2056,7 @@ def forward(
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
         logits_to_keep: Union[int, torch.Tensor] = 0,
+        **loss_kwargs,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         r"""
             labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
@@ -2157,15 +2158,31 @@ def forward(
             past_key_values=past_key_values,
             use_cache=use_cache,
             inputs_embeds=inputs_embeds,
-            labels=labels,
             output_hidden_states=output_hidden_states,
             output_attentions=output_attentions,
             return_dict=return_dict,
             cache_position=cache_position,
             logits_to_keep=logits_to_keep,
+            **loss_kwargs,
         )
 
-        return outputs
+        # Temporary fix to calculate the loss in main class, as the model's vocab size may be resized
+        loss = None
+        logits = outputs[0]
+
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.config.get_text_config().vocab_size, **loss_kwargs)
+
+        if not return_dict:
+            return (loss,) + outputs if loss is not None else outputs
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=outputs.logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
 
     def prepare_inputs_for_generation(
         self,

filename: tests/models/mllama/test_modeling_mllama.py
patch: @@ -321,6 +321,24 @@ def test_inputs_embeds_matches_input_ids(self):
                 out_embeds = model(inputs_embeds=inputs_embeds, **inputs)[0]
             torch.testing.assert_close(out_embeds, out_ids)
 
+    def test_resize_embeddings_results_in_successful_loss(self):
+        # resizing embeddings should result in successful loss computation
+        config, inputs = self.model_tester.prepare_config_and_inputs_for_common()
+
+        for model_class in self.all_model_classes:
+            model = model_class(config)
+            model_vocab_size = config.get_text_config().vocab_size
+            inputs = self._prepare_for_class(inputs, model_class, return_labels=True)
+            # Resize embeddings and call forward
+            model.resize_token_embeddings(model_vocab_size + 10)
+            output = model(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                labels=inputs["labels"],
+                return_dict=True,
+            )
+            self.assertTrue("loss" in output)
+
     def _check_attentions_for_generate(
         self, batch_size, attentions, prompt_length, output_length, config, decoder_past_key_values
     ):