Repository: huggingface/transformers
Pull request: fix gemma3 grad acc

# What does this PR do?

This PR fixes the grad acc issue with gemma3 model. The issue was that we passed **kwargs in the model forward, so we were making the assumption that he was passing `**loss_kwargs` -> `num_items_in_batch` to calculate the loss. Not sure what is the best way to fix this @ArthurZucker in general as this might probably happen again. Maybe set `accepts_loss_kwargs` to `False` in general and set it to `True` for models that we fixed ? I'm fine also just setting it `False` for models that don't use the `kwargs` for the loss. 

As for why I didn't have the loss function: In the code, they are filetring the logits/labels so I decided to simply not use `num_items_in_batch` to calculate the loss. Otherwise, the loss won't be correctly calculated for one of the cases.

Also I fixed an issue related to peft as we couldn't have access to that attribute as the model was a peft model.


#### To reproduce
winglian script
https://gist.github.com/winglian/569924fe154824c8ce148f6e185cd4cd

### After fix 
grad acc 2 bs 1 and  grad acc 1 bs 2 
<img width="548" alt="Screenshot 2025-04-02 at 4 38 54 PM" src="https://github.com/user-attachments/assets/8d8dc273-4851-4c37-a33c-933ec54c47d1" />

Fixes #37197

cc @winglian 

Changed files:
filename: src/transformers/models/gemma3/modeling_gemma3.py
patch: @@ -777,6 +777,8 @@ def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
 )
 class Gemma3Model(Gemma3PreTrainedModel):
     _checkpoint_conversion_mapping = {"language_model.model": "language_model"}
+    # we are filtering the logits/labels so we shouldn't divide the loss based on num_items_in_batch
+    accepts_loss_kwargs = False
 
     def __init__(self, config: Gemma3Config):
         super().__init__(config)

filename: src/transformers/models/gemma3/modular_gemma3.py
patch: @@ -727,6 +727,9 @@ def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
 
 
 class Gemma3Model(PaliGemmaModel):
+    # we are filtering the logits/labels so we shouldn't divide the loss based on num_items_in_batch
+    accepts_loss_kwargs = False
+
     def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
         """
         Projects the last hidden state from the vision model into language model space.

filename: src/transformers/models/paligemma/modeling_paligemma.py
patch: @@ -132,6 +132,8 @@ def _init_weights(self, module):
 )
 class PaliGemmaModel(PaliGemmaPreTrainedModel):
     _checkpoint_conversion_mapping = {"language_model.model": "language_model"}
+    # we are filtering the logits/labels so we shouldn't divide the loss based on num_items_in_batch
+    accepts_loss_kwargs = False
 
     def __init__(self, config: PaliGemmaConfig):
         super().__init__(config)

filename: src/transformers/trainer.py
patch: @@ -629,18 +629,16 @@ def __init__(
 
         # Just in case the model was wrapped outside of the `Trainer`
         unwrapped_model = self.accelerator.unwrap_model(model)
-        model_forward = (
-            unwrapped_model.forward
-            if not _is_peft_model(unwrapped_model)
-            else unwrapped_model.get_base_model().forward
-        )
-        forward_params = inspect.signature(model_forward).parameters
+        # We also unwrap peft model
+        if _is_peft_model(unwrapped_model):
+            unwrapped_model = unwrapped_model.get_base_model()
 
         # Check if the model has explicit setup for loss kwargs,
         # if not, check if `**kwargs` are in model.forward
-        if hasattr(model, "accepts_loss_kwargs"):
-            self.model_accepts_loss_kwargs = model.accepts_loss_kwargs
+        if hasattr(unwrapped_model, "accepts_loss_kwargs"):
+            self.model_accepts_loss_kwargs = unwrapped_model.accepts_loss_kwargs
         else:
+            forward_params = inspect.signature(unwrapped_model.forward).parameters
             self.model_accepts_loss_kwargs = any(
                 k.kind == inspect.Parameter.VAR_KEYWORD for k in forward_params.values()
             )