Currently this is only possible with 2D-mask when sdap is enabled.
# modeling bert
# Expand the attention mask
if use_sdpa_attention_masks:
# Expand the attention mask for SDPA.
# [bsz, seq_len] -> [bsz, 1, seq_len, seq_len]
if self.config.is_decoder:
extended_attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
attention_mask,
input_shape,
embedding_output,
past_key_values_length,
)
else:
extended_attention_mask = _prepare_4d_attention_mask_for_sdpa(
attention_mask, embedding_output.dtype, tgt_len=seq_length
)
# print(f"sdpa attn mask:{embedding_output.dtype=} {extended_attention_mask.sum(axis=(2,3))=}")
else:
# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
# ourselves in which case we just need to make it broadcastable to all heads.
extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape)
Motivation
4D attention mask is widely used such as padding bert.
We construct 4D attention mask before the model forward and pass it to the forward call.
I think it is easy to exend the current implement of _expand_mask to support 3D/4D attention mask.
Feature request
Currently this is only possible with 2D-mask when sdap is enabled.
Motivation
4D attention mask is widely used such as padding bert.
We construct 4D attention mask before the model forward and pass it to the forward call.
I think it is easy to exend the current implement of _expand_mask to support 3D/4D attention mask.
Your contribution
I can do PR and contribution.