Llama 4: TypeError: pad(): argument 'pad' failed to unpack the object at pos 2 with error "type must be tuple of ints,but got NoneType"
Author: daniel-dona
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.51.0
- Platform: Linux-6.8.0-57-generic-x86_64-with-glibc2.39
- Python version: 3.12.3
- Huggingface_hub version: 0.30.1
- Safetensors version: 0.5.3
- Accelerate version: 1.6.0
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0+cu124 (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?: no
- Using GPU in script?: yes
- GPU type: NVIDIA GeForce RTX 3090 (Note: x4 GPU cluster)


### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

Steps:
1. Create a .venv with:
```
python3 -m venv .venv
```
2. Activate the venv with:
```
source .venv/bin/activate
```
3. Install libraries:
```
pip install transformers torch accelerate pillow torchvision flash_attn huggingface_hub huggingface_hub[hf_xet]
```
4. Run HF repo sample code
```python
from transformers import AutoProcessor, Llama4ForConditionalGeneration
import torch

model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"

processor = AutoProcessor.from_pretrained(model_id)
model = Llama4ForConditionalGeneration.from_pretrained(
    model_id,
    attn_implementation="flex_attention",
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

url1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg"
url2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png"
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "url": url1},
            {"type": "image", "url": url2},
            {"type": "text", "text": "Can you describe how these two images are similar, and how they differ?"},
        ]
    },
]

inputs = processor.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=256,
)

response = processor.batch_decode(outputs[:, inputs["input_ids"].shape[-1]:])[0]
print(response)
print(outputs[0])

```
5. Got this error:
```
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [05:30<00:00,  6.60s/it]
Some parameters are on the meta device because they were offloaded to the cpu.
Traceback (most recent call last):
  File "/home/dani/llama4.py", line 35, in <module>
    outputs = model.generate(
              ^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 2460, in generate
    result = self._sample(
             ^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 3426, in _sample
    outputs = self(**model_inputs, return_dict=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/accelerate/hooks.py", line 176, in new_forward
    output = module._old_forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/transformers/models/llama4/modeling_llama4.py", line 1761, in forward
    outputs = self.language_model(
              ^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/transformers/models/llama4/modeling_llama4.py", line 1015, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/transformers/models/llama4/modeling_llama4.py", line 669, in forward
    causal_mask, chunk_causal_mask = self._update_causal_mask(
                                     ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/transformers/models/llama4/modeling_llama4.py", line 779, in _update_causal_mask
    attention_mask = make_flex_block_causal_mask(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/transformers/integrations/flex_attention.py", line 103, in make_flex_block_causal_mask
    attention_mask_2d = torch.nn.functional.pad(attention_mask_2d, value=0, pad=(0, key_length))
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dani/.venv/lib/python3.12/site-packages/torch/nn/functional.py", line 5209, in pad
    return torch._C._nn.pad(input, pad, mode, value)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: pad(): argument 'pad' failed to unpack the object at pos 2 with error "type must be tuple of ints,but got NoneType"
```

### Expected behavior

1. Sample code to run successfully 

Comments:
@ArthurZucker  We are getting this error too with `flex attention`. Once you switch off flex attention and enable eager_attention or flash attention, we get the previously reported casual mask errors. 

> [@ArthurZucker](https://github.com/ArthurZucker) We are getting this error too with `flex attention`. Once you switch off flex attention and enable eager_attention or flash attention, we get the previously reported casual mask errors.

Just tried changing `attn_implementation="eager"` and it works with text only input, trying now with image input.

It's because of this line https://github.com/huggingface/transformers/blob/d1b92369ca193da49f9f7ecd01b08ece45c2c9aa/src/transformers/models/llama4/modeling_llama4.py#L782

It seems that by default a dynamic cache get's initialized while flex attn needs a strict size of how many more tokens can be generated (which doesn't work with dynamic cache as it's unlimited). Would generally suggest avoiding flex attn for now since it's very experimental (and hence many things are not fleshed out yet).

And, after a [deeper dive](https://github.com/vasqu/transformers/tree/fix-llama4-flex-mask), I encountered other issues (because of my lack of understanding of chunked attention and how it is supposed to interact here). Will leave it here, maybe it helps regardless.

> Just tried changing `attn_implementation="eager"` and it works with text only input, trying now with image input.

same error when I've tried the image input, but it works with `attn_implementation="eager"`

```python
from transformers import AutoProcessor, Llama4ForConditionalGeneration
import torch

model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"

processor = AutoProcessor.from_pretrained(model_id)
model = Llama4ForConditionalGeneration.from_pretrained(
    model_id,
    attn_implementation="eager",
    device_map="auto",
    torch_dtype=torch.bfloat16,
    cache_dir = cache_dir,
    use_auth_token = use_auth_token
)
```

```python
url1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg"
url2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png"
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "url": url1},
            {"type": "image", "url": url2},
            {"type": "text", "text": "Can you describe how these two images are similar, and how they differ?"},
        ]
    },
]

inputs = processor.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=256,
)

response = processor.batch_decode(outputs[:, inputs["input_ids"].shape[-1]:])[0]
print(response)
print(outputs[0])
```


```
The two images depict anthropomorphic animals, a rabbit and a cat, dressed in similar attire. The rabbit is standing on a dirt path with flowers and houses in the background, while the cat is sitting on a stone path surrounded by greenery.

**Similarities:**

* Both characters are wearing blue jackets, brown vests, and bow ties.
* They have similar facial expressions, with large eyes and endearing features.
* The backgrounds of both images feature natural settings, including paths, flowers, and houses.

**Differences:**

* **Species:** The most obvious difference is that one character is a rabbit and the other is a cat.
* **Pose:** The rabbit is standing upright on two legs, while the cat is sitting on all fours.
* **Attire:** Although their outfits share similarities, there are some differences in design and color. The rabbit's jacket appears longer and more formal, whereas the cat's jacket has a more casual look.
* **Background:** While both backgrounds feature natural settings, they differ in terms of architecture and layout. The rabbit's background includes a larger house with a steeper roof, whereas the cat's background features a smaller cottage with a thatched roof.
* **Color palette:** The overall color schemes of the two images differ slightly
tensor([200000, 200005,   1556,  ...,   8392,   1667,  17432], device='cuda:0')
```


There is a PR for the flex attention padding fix. https://github.com/huggingface/transformers/pull/37327 

For `generate` the `cache_implementation` should be `hybrid` by default! Using dynamic was indeed passing `None` key length

Fixed by #37327 closing!