[GPTNeo] create local attention mask ones
Author: patil-suraj
Pull request: yes

# What does this PR do?

This PR refactors GPT Neo such that the causal attention mask for the local attention layer is only computed once per batch in the `GPTNeoModel` class and then shared between the layers instead of re-computing it in each layer.

I've verified that all slow tests are passing.

#### Benchmarks

This PR does not change memory/speed for the forward pass

On master
```
====================       INFERENCE - SPEED - RESULT       ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length     Time in s   
--------------------------------------------------------------------------------
/home/suraj/projects/gpt-neo-c       2               32            0.021     
/home/suraj/projects/gpt-neo-c       2              128            0.059     
/home/suraj/projects/gpt-neo-c       2              512            0.227
/home/suraj/projects/gpt-neo-c       2              1024           0.464        
/home/suraj/projects/gpt-neo-c       4               32            0.033     
/home/suraj/projects/gpt-neo-c       4              128            0.113     
/home/suraj/projects/gpt-neo-c       4              512            0.449
/home/suraj/projects/gpt-neo-c       2              1024           N/A        
--------------------------------------------------------------------------------

====================      INFERENCE - MEMORY - RESULT       ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
/home/suraj/projects/gpt-neo-c       2               32             6136     
/home/suraj/projects/gpt-neo-c       2              128             6268     
/home/suraj/projects/gpt-neo-c       2              512             6790
/home/suraj/projects/gpt-neo-c       2              1024            7472          
/home/suraj/projects/gpt-neo-c       4               32             6204     
/home/suraj/projects/gpt-neo-c       4              128             6428     
/home/suraj/projects/gpt-neo-c       4              512             7456
/home/suraj/projects/gpt-neo-c       4              1024            N/A         
--------------------------------------------------------------------------------
```

On this PR
```
===================       INFERENCE - SPEED - RESULT       ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length     Time in s   
--------------------------------------------------------------------------------
/home/suraj/projects/gpt-neo-c       2               32            0.021     
/home/suraj/projects/gpt-neo-c       2              128            0.058     
/home/suraj/projects/gpt-neo-c       2              512            0.222
/home/suraj/projects/gpt-neo-c       2              1024           0.453      
/home/suraj/projects/gpt-neo-c       4               32            0.032     
/home/suraj/projects/gpt-neo-c       4              128            0.11     
/home/suraj/projects/gpt-neo-c       4              512            0.439
/home/suraj/projects/gpt-neo-c       4              1024          N/A           
--------------------------------------------------------------------------------

====================      INFERENCE - MEMORY - RESULT       ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
/home/suraj/projects/gpt-neo-c       2               32             6136     
/home/suraj/projects/gpt-neo-c       2              128             6268     
/home/suraj/projects/gpt-neo-c       2              512             6790
/home/suraj/projects/gpt-neo-c       2              1024            7476     
/home/suraj/projects/gpt-neo-c       4               32             6204     
/home/suraj/projects/gpt-neo-c       4              128             6428     
/home/suraj/projects/gpt-neo-c       4              512             7460 
/home/suraj/projects/gpt-neo-c       4              1024            N/A    
--------------------------------------------------------------------------------
```

I did a micro-benchmark using the 125M model for generation and this PR does give a small speed-up when generating longer sequences.

On master 

```
%timeit model.generate(**enc, do_sample=False, max_length=512, min_length=512)
4.63 s ± 25.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit -n 10 model.generate(**enc, do_sample=False, max_length=1024, min_length=1024)
9.63 s ± 549 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

On this PR
```
%timeit model.generate(**enc, do_sample=False, max_length=512, min_length=512)
4.25 s ± 189 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit -n 10 model.generate(**enc, do_sample=False, max_length=1024, min_length=1024)

9 s ± 437 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

 

Comments: