Llama 4 Scout w/ Flex Attention Fails
Author: ardapekis
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.51.0
- Platform: Linux-6.8.0-53-generic-x86_64-with-glibc2.35
- Python version: 3.11.10
- Huggingface_hub version: 0.30.1
- Safetensors version: 0.5.3
- Accelerate version: 1.6.0
- Accelerate config:    not found
- DeepSpeed version: 0.15.4
- 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?: Y
- Using GPU in script?: Y
- GPU type: NVIDIA H200

### Who can help?

@ArthurZucker @zucchini-nlp @SunMarc 

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

Run the official code example for Llama 4 Scout: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct#how-to-use-with-transformers

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

The following error results:

```python
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[1], line 35
     16 messages = [
     17     {
     18         "role": "user",
   (...)     24     },
     25 ]
     27 inputs = processor.apply_chat_template(
     28     messages,
     29     add_generation_prompt=True,
   (...)     32     return_tensors="pt",
     33 ).to(model.device)
---> 35 outputs = model.generate(
     36     **inputs,
     37     max_new_tokens=256,
     38 )
     40 response = processor.batch_decode(outputs[:, inputs["input_ids"].shape[-1]:])[0]
     41 print(response)

File /workspace/.venv/lib/python3.11/site-packages/torch/utils/_contextlib.py:116, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    113 @functools.wraps(func)
    114 def decorate_context(*args, **kwargs):
    115     with ctx_factory():
--> 116         return func(*args, **kwargs)

File /workspace/.venv/lib/python3.11/site-packages/transformers/generation/utils.py:2460, in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, use_model_defaults, **kwargs)
   2452     input_ids, model_kwargs = self._expand_inputs_for_generation(
   2453         input_ids=input_ids,
   2454         expand_size=generation_config.num_return_sequences,
   2455         is_encoder_decoder=self.config.is_encoder_decoder,
   2456         **model_kwargs,
   2457     )
   2459     # 12. run sample (it degenerates to greedy search when `generation_config.do_sample=False`)
-> 2460     result = self._sample(
   2461         input_ids,
   2462         logits_processor=prepared_logits_processor,
   2463         stopping_criteria=prepared_stopping_criteria,
   2464         generation_config=generation_config,
   2465         synced_gpus=synced_gpus,
   2466         streamer=streamer,
   2467         **model_kwargs,
   2468     )
   2470 elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH):
   2471     # 11. interleave input_ids with `num_beams` additional sequences per batch
   2472     input_ids, model_kwargs = self._expand_inputs_for_generation(
   2473         input_ids=input_ids,
   2474         expand_size=generation_config.num_beams,
   2475         is_encoder_decoder=self.config.is_encoder_decoder,
   2476         **model_kwargs,
   2477     )

File /workspace/.venv/lib/python3.11/site-packages/transformers/generation/utils.py:3426, in GenerationMixin._sample(self, input_ids, logits_processor, stopping_criteria, generation_config, synced_gpus, streamer, **model_kwargs)
   3423 model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
   3425 if is_prefill:
-> 3426     outputs = self(**model_inputs, return_dict=True)
   3427     is_prefill = False
   3428 else:

File /workspace/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1739, in Module._wrapped_call_impl(self, *args, **kwargs)
   1737     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1738 else:
-> 1739     return self._call_impl(*args, **kwargs)

File /workspace/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1750, in Module._call_impl(self, *args, **kwargs)
   1745 # If we don't have any hooks, we want to skip the rest of the logic in
   1746 # this function, and just call forward.
   1747 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1748         or _global_backward_pre_hooks or _global_backward_hooks
   1749         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1750     return forward_call(*args, **kwargs)
   1752 result = None
   1753 called_always_called_hooks = set()

File /workspace/.venv/lib/python3.11/site-packages/accelerate/hooks.py:176, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    174         output = module._old_forward(*args, **kwargs)
    175 else:
--> 176     output = module._old_forward(*args, **kwargs)
    177 return module._hf_hook.post_forward(module, output)

File /workspace/.venv/lib/python3.11/site-packages/transformers/models/llama4/modeling_llama4.py:1761, in Llama4ForConditionalGeneration.forward(self, input_ids, pixel_values, attention_mask, position_ids, past_key_values, inputs_embeds, vision_feature_layer, vision_feature_select_strategy, labels, use_cache, output_attentions, output_hidden_states, return_dict, cache_position, logits_to_keep, image_sizes, **lm_kwargs)
   1757     inputs_embeds.masked_scatter_(expanded_mask, projected_vision_flat)
   1759     inputs_embeds = inputs_embeds.view(original_inputs_embeds_shape)
-> 1761 outputs = self.language_model(
   1762     attention_mask=attention_mask,
   1763     position_ids=position_ids,
   1764     past_key_values=past_key_values,
   1765     inputs_embeds=inputs_embeds,
   1766     use_cache=use_cache,
   1767     output_attentions=output_attentions,
   1768     output_hidden_states=output_hidden_states,
   1769     return_dict=return_dict,
   1770     cache_position=cache_position,
   1771     logits_to_keep=logits_to_keep,
   1772     **lm_kwargs,
   1773 )
   1775 logits = outputs[0]
   1777 loss = None

File /workspace/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1739, in Module._wrapped_call_impl(self, *args, **kwargs)
   1737     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1738 else:
-> 1739     return self._call_impl(*args, **kwargs)

File /workspace/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1750, in Module._call_impl(self, *args, **kwargs)
   1745 # If we don't have any hooks, we want to skip the rest of the logic in
   1746 # this function, and just call forward.
   1747 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1748         or _global_backward_pre_hooks or _global_backward_hooks
   1749         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1750     return forward_call(*args, **kwargs)
   1752 result = None
   1753 called_always_called_hooks = set()

File /workspace/.venv/lib/python3.11/site-packages/transformers/models/llama4/modeling_llama4.py:1015, in Llama4ForCausalLM.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, cache_position, logits_to_keep, **kwargs)
   1012 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
   1014 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-> 1015 outputs = self.model(
   1016     input_ids=input_ids,
   1017     attention_mask=attention_mask,
   1018     position_ids=position_ids,
   1019     past_key_values=past_key_values,
   1020     inputs_embeds=inputs_embeds,
   1021     use_cache=use_cache,
   1022     output_attentions=output_attentions,
   1023     output_hidden_states=output_hidden_states,
   1024     return_dict=return_dict,
   1025     cache_position=cache_position,
   1026     **kwargs,
   1027 )
   1029 hidden_states = outputs[0]
   1030 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss

File /workspace/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1739, in Module._wrapped_call_impl(self, *args, **kwargs)
   1737     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1738 else:
-> 1739     return self._call_impl(*args, **kwargs)

File /workspace/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1750, in Module._call_impl(self, *args, **kwargs)
   1745 # If we don't have any hooks, we want to skip the rest of the logic in
   1746 # this function, and just call forward.
   1747 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1748         or _global_backward_pre_hooks or _global_backward_hooks
   1749         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1750     return forward_call(*args, **kwargs)
   1752 result = None
   1753 called_always_called_hooks = set()

File /workspace/.venv/lib/python3.11/site-packages/transformers/models/llama4/modeling_llama4.py:669, in Llama4TextModel.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, cache_position, **flash_attn_kwargs)
    666 if position_ids is None:
    667     position_ids = cache_position.unsqueeze(0)
--> 669 causal_mask, chunk_causal_mask = self._update_causal_mask(
    670     attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
    671 )
    673 hidden_states = inputs_embeds
    675 # create position embeddings to be shared across the decoder layers

File /workspace/.venv/lib/python3.11/site-packages/transformers/models/llama4/modeling_llama4.py:779, in Llama4TextModel._update_causal_mask(self, attention_mask, input_tensor, cache_position, past_key_values, output_attentions, chunked_attention_mask)
    775     offsets = (first_cache_position, max(last_cache_position - key_length, 0))
    776     chunked_attention_mask = make_flex_block_causal_mask(
    777         attention_mask, self.config.attention_chunk_size, sequence_length, key_length, offsets=offsets
    778     )
--> 779     attention_mask = make_flex_block_causal_mask(
    780         attention_mask,
    781         query_length=sequence_length,
    782         key_length=past_key_values.get_max_cache_shape(),
    783         offsets=None if sequence_length != 1 else (first_cache_position, 0),
    784     )
    785     return attention_mask, chunked_attention_mask
    786 if isinstance(attention_mask, BlockMask):

File /workspace/.venv/lib/python3.11/site-packages/transformers/integrations/flex_attention.py:103, in make_flex_block_causal_mask(attention_mask_2d, attention_chunk_size, query_length, key_length, offsets)
     74 def make_flex_block_causal_mask(
     75     attention_mask_2d: torch.Tensor,
     76     attention_chunk_size: Optional[int] = None,
   (...)     79     offsets: Optional[Tuple[Offset, Offset]] = None,
     80 ) -> "BlockMask":
     81     """
     82     Create a block causal document mask for a batch of sequences, both packed and unpacked.
     83     Create Block causal logic and passing it into :func:`torch.nn.attention.flex_attention.create_block_mask`.
   (...)    101         BlockMask
    102     """
--> 103     attention_mask_2d = torch.nn.functional.pad(attention_mask_2d, value=0, pad=(0, key_length))
    104     device = attention_mask_2d.device
    105     document_ids = attention_mask_2d.clone()

File /workspace/.venv/lib/python3.11/site-packages/torch/nn/functional.py:5209, in pad(input, pad, mode, value)
   5202         if mode == "replicate":
   5203             # Use slow decomp whose backward will be in terms of index_put.
   5204             # importlib is required because the import cannot be top level
   5205             # (cycle) and cannot be nested (TS doesn't support)
   5206             return importlib.import_module(
   5207                 "torch._decomp.decompositions"
   5208             )._replication_pad(input, pad)
-> 5209 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

I compared running with `attn_implementation="sdpa"` which seems to work. The output was:

```
The two images feature anthropomorphic animals, a rabbit and a cat, dressed in human-like attire. The rabbit is depicted standing on a dirt path with flowers on either side, while the cat sits on a stone pathway surrounded by greenery.

**Similarities:**

* Both characters are dressed in formal clothing, including jackets, vests, shirts, and bow ties.
* They both have large eyes and appear to be looking directly at the viewer.
* The backgrounds of both images feature natural settings, such as gardens or countryside scenes.

**Differences:**

* **Species:** The most obvious difference is that one character is a rabbit and the other is a cat.
* **Pose:** The rabbit stands upright on two legs, while the cat sits on all fours.
* **Attire:** Although both characters wear similar outfits, there are some differences in their attire. The rabbit's jacket appears longer and more worn, while the cat's jacket has a more tailored fit.
* **Background:** While both backgrounds feature natural settings, the rabbit's background includes a house and a mountain range in the distance, whereas the cat's background focuses more on the immediate surroundings, with a house visible in the distance.
* **Color scheme:** The overall color palette of the two images differs slightly,
tensor([200000, 200005,   1556,  ...,  55738,  17432,     24], device='cuda:0')
```

Comments:
Fixes as this is on main! 🤗