Handling of full_text_row_masked_out_mask in mllama is incorrect.
Author: nshepperd
Pull request: no
Labels: name: WIP

name: bug

### System Info

N/A

### Who can help?

@ArthurZucker 

### Information

- [ ] 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

In MllamaCrossAttentionDecoderLayer:

```py
        hidden_states, attn_weights = self.cross_attn(
            hidden_states=hidden_states,
            attention_mask=cross_attention_mask,
            cross_attention_states=cross_attention_states,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            cache_position=cache_position,
            **kwargs,
        )
        hidden_states = residual + self.cross_attn_attn_gate.tanh() * hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        if full_text_row_masked_out_mask is not None:
            hidden_states = full_text_row_masked_out_mask[:, 0] * hidden_states  # type: ignore
        hidden_states = residual + self.cross_attn_mlp_gate.tanh() * hidden_states
```

the mask is not applied until after the hidden_states (which contains the result of the cross attention) has been copied to the residual. This is at odds with how it's done in the reference implementation [here](https://github.com/meta-llama/llama-models/blob/01dc8ce46fecf06b639598f715efbb4ab981fb4c/models/llama3/multimodal/model.py#L912):

```py
    def forward(
        self,
        x: torch.Tensor,
        xattn_mask: torch.Tensor,
        full_text_row_masked_out_mask: torch.Tensor,
        xattn_cache: torch.Tensor,
    ) -> torch.Tensor:
        xq = F.linear(x, self.wq.weight)
        bsz, seqlen, _ = x.shape

        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xq = self.q_norm(xq)
        xq = xq.transpose(1, 2)

        xk, xv = xattn_cache

        output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=xattn_mask, dropout_p=0.0)
        output = output * full_text_row_masked_out_mask
        output = output.transpose(1, 2).contiguous().reshape(bsz, seqlen, -1)

        out = F.linear(output, self.wo.weight)
        out = reduce_from_tensor_model_parallel_region(out)
        return out
```

The implementation is MllamaCrossAttentionDecoderLayer is clearly incorrect, as it allows text tokens that should be unable to see any image token (due to preceding them) to observe all image tokens via the residual.

This may affect any usage of llama vision models with images where an image isn't the very first token (finetuning or inference) in transformers.

### Expected behavior

N/A

Comments:
Hi @nshepperd  @ArthurZucker 
I'd like to take this issue on and submit a fix for the incorrect handling of full_text_row_masked_out_mask in MllamaCrossAttentionDecoderLayer.

I've reviewed the reference implementation and understand the masking should be applied before the residual connection to prevent unintended attention leakage.

Could you please assign this issue to me so I can proceed with a PR? 

Thanks!

Hey! Not sure its incorrect, we had 1-1 matching logits, cc @zucchini-nlp if you remember! 

Hmm, I don't really remember, since the `full_mask` was added later just before shipping. We definitely had 1-1 matching text outputs, and allclose logits. Though I have never seen the original implementation, we were working based off the draft transformers PR 😆 

Yeah, we can try to check the balmes

It took an awful lot of debugging, but I was able to make a self-contained test case demonstrating that this results in a discrepancy, and how it is fixed: https://github.com/nshepperd/llama-3.2-transformers-bug

Which caused me to realize `full_text_row_masked_out_mask` needs to be applied in two places, after both the cross_attn and the mlp. (In the reference impl it is applied after the mlp here: https://github.com/meta-llama/llama-models/blob/a9c89c471f793423afd4cc3ca8671d6e56fe64cb/models/llama3/multimodal/model.py#L998)

It's possible your tests didn't catch this because there wouldn't be a discrepancy in any case where the 
 is the first token, which is the case for most self-contained question answering prompts with hf's tokenizer.

> It's possible your tests didn't catch this because there wouldn't be a discrepancy in any case where the 
 is the first token, which is the case for most self-contained question answering prompts with hf's tokenizer.


Oh yeah, totally could be. Can you verify it and submit a PR with the edge case of images being the first token?

> > It's possible your tests didn't catch this because there wouldn't be a discrepancy in any case where the 
 is the first token, which is the case for most self-contained question answering prompts with hf's tokenizer.
> 
> Oh yeah, totally could be. Can you verify it and submit a PR with the edge case of images being the first token?

Is the issue still available to take ?


Sure, feel free to submit a PR :)

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

> This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
> 
> Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.


Yepp i am working on it


I've opened #47103 with a fix and a regression test for the image-not-first-token edge case @zucchini-nlp mentioned. The mask is now applied to the cross-attention output as well (not just the MLP), so text rows that precede their image no longer pick up image information through the residual — matching the reference implementation.

@ArkVex — apologies if you were mid-way on this; I didn't see an open PR so I went ahead. If you have one in progress I'm happy to close mine and review yours instead, or to collaborate — whatever's easiest for you.