Attention fwd/bwd issues with all zero masks

[RuiWang1998]

Hi,

We have been observing some issues with masking, namely that if all tokens are masked out. In this specific case, the output of attention will be twice of what native pytorch's implementation is, and the gradient of v value vector is simply garbage value.

To demonstrate the need for masking all tokens, consider the case where we land on no templates in template search but the computation graph needs to stay the same during training to avoid a bunch of issues including memory fragmentation.

Is there a way to fix this?

Best,
Rui

[DejunL]

Hi, is this for the triangle attention or attention with pair bias usage case?

For the triangle attention, we can't guarantee specific values produced from our code for those masked-out garbage elements. The user is required to mask those elements for downstream computation so they don't contribute to the loss or inference computation. If the purpose is to just avoid memory issues, you can always zero-out the output (and the d_loss/d_output for bwd) by multiplying with the mask without relying on what values being produced for the garbage elements

[hsadasivan]

@RuiWang1998 is this resolved?