System Info
N/A
Who can help?
@ArthurZucker
Information
Tasks
Reproduction
In MllamaCrossAttentionDecoderLayer:
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_statesthe 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:
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 outThe 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
System Info
N/A
Who can help?
@ArthurZucker
Information
Tasks
examplesfolder (such as GLUE/SQuAD, ...)Reproduction
In MllamaCrossAttentionDecoderLayer:
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:
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