Gemma3 slightly alters hidden state when input_ids is batched
Author: lwiklendt
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.50.3
- Platform: Windows-11-10.0.26100-SP0
- Python version: 3.13.2
- Huggingface_hub version: 0.30.1
- Safetensors version: 0.5.3
- Accelerate version: 1.8.1
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: NO
        - mixed_precision: bf16
        - use_cpu: False
        - debug: False
        - num_processes: 1
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: all
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - enable_cpu_affinity: False
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA GeForce RTX 3090

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

```python
import torch
from transformers import AutoTokenizer, Gemma3ForCausalLM
import torch.nn.functional as F

ckpt = 'google/gemma-3-1b-pt'
hf_token = None
tokenizer = AutoTokenizer.from_pretrained(ckpt, token=hf_token)
model = Gemma3ForCausalLM.from_pretrained(
    ckpt,
    torch_dtype=torch.bfloat16,
    device_map='auto',
    token=hf_token,
).model

prompt = 'There is a problem where batching alters the hidden state in Gemma3'
token_ids = tokenizer(prompt, return_tensors='pt')['input_ids'].to(model.device)
hidden_state_single = model(token_ids, use_cache=False)['last_hidden_state']
hidden_state_batch = model(torch.concat((token_ids, token_ids)), use_cache=False)['last_hidden_state']

print(hidden_state_single[0])
print(hidden_state_batch[0])
print(torch.allclose(hidden_state_single[0], hidden_state_batch[0]))
```
prints
```
tensor([[ 0.6523, -1.8047,  0.7031,  ..., -1.4375, -0.2988, -0.4707],
        [ 1.3359,  0.2520, -2.5156,  ..., -3.1719,  1.6172, -0.9648],
        [ 2.0156, -3.5781, -1.2891,  ..., -1.5547, -0.2246, -0.7969],
        ...,
        [ 0.0505,  0.1768, -1.1328,  ...,  0.1777, -0.2246, -0.4766],
        [ 0.5195, -0.4102, -0.7148,  ..., -0.6367, -4.2188, -0.6602],
        [-0.3887, -0.7188, -1.5391,  ..., -0.3574, -0.6094, -3.0000]],
       device='cuda:0', dtype=torch.bfloat16, grad_fn=<SelectBackward0>)
tensor([[ 0.6641, -1.8281,  0.6992,  ..., -1.4688, -0.2988, -0.4707],
        [ 1.3125,  0.2041, -2.4688,  ..., -3.1875,  1.6328, -0.9805],
        [ 2.0625, -3.5781, -1.2656,  ..., -1.5547, -0.2246, -0.7773],
        ...,
        [ 0.0728,  0.1816, -1.1250,  ...,  0.2188, -0.2168, -0.4648],
        [ 0.5391, -0.3672, -0.6836,  ..., -0.6289, -4.2188, -0.6719],
        [-0.3574, -0.7500, -1.5078,  ..., -0.3848, -0.6016, -2.9375]],
       device='cuda:0', dtype=torch.bfloat16, grad_fn=<SelectBackward0>)
False
```

### Expected behavior

I expect `hidden_state_single[0]` to be equal to `hidden_state_batch[0]`.

Comments:
Hi @lwiklendt ! This comment may be helpful: https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535. Batching uses different matrix multiplication shapes, which can often lead to variations in the results.

Thanks @bvantuan, that explains it quite well.