[GPT-Neo] Simplify local attention

[patil-suraj]

What does this PR do?

Co-authored-by: finetuneanon [email protected]

This PR is a continuation of #11630 which simplifies GPT Neo's local attention implementation. All credit to @finetuneanon for finding this issue, providing the fix, and a detailed explanation. Thanks a lot for working on this :)

The issue is described in #11320 and performance evaluation results are available here:
#12106 (comment)

This PR does some cleanup and updates tests on top of @finetuneanon changes.

Fixes #11320, Fixes #11787, Fixes #12964, Fixes #11096

[patrickvonplaten]

this doesn't break in fp16? It's -1e4 in modeling_gpt2.py I think

[patrickvonplaten]

Does this correctly convert to fp16?

torch.tensor(-1e9).to(torch.float16)

gives -inf -> which could later lead to problems in fp16 no?

[patil-suraj]

@patrickvonplaten @sgugger

-1e9 is the value used in the original codebase and actually, the attention weights are always computed in fp32

transformers/src/transformers/models/gpt_neo/modeling_gpt_neo.py

Lines 274 to 286 in 09549aa

	query = query.to(torch.float32)
	key = key.to(torch.float32)
	attn_weights = torch.matmul(query, key.transpose(-1, -2))
	attn_weights = torch.where(causal_mask, attn_weights, masked_bias.to(attn_weights.dtype))
	if attention_mask is not None:
	# Apply the attention mask
	attn_weights = attn_weights + attention_mask
	attn_weights = nn.Softmax(dim=-1)(attn_weights)
	attn_weights = attn_weights.to(value.dtype)
	attn_weights = attn_dropout(attn_weights)

so the masked_bias is always fp32. The attn_weights are only cast back to the original dtype after softmax.

However, as the models are trained on TPU and possibly with bf16 (see #11076 (comment)) I'm not sure we can guarantee that the models will always work with fp16. See #11076.

[patrickvonplaten]

Awesome new tests! Could we also add one/two fp16 tests? To make sure generation and forward pass works correctly :-)

[patrickvonplaten]

Great clean-up! Think this should solve the vram memory issue

The only thing that would be good to check IMO is if this implementation is 100% fp16 compatible (would be great to add some tests as well for this)

[sgugger]

Thanks for working on this! Agreed with Patrick on the FP16 tests to make sure it all works fine.

[LysandreJik]

This looks good!

[LysandreJik]

This will save to the state dict - is that voluntary? If not, you can add the persistent=False flag (introduced in pytorch 1.6)

[tianleiwu]

I had error in running:

  File "src\transformers\models\gpt_neo\configuration_gpt_neo.py", line 218, in __init__
    from .modeling_gpt_neo import GPTNeoAttentionMixin
ImportError: cannot import name 'GPTNeoAttentionMixin'

Need a patch for this PR.