"Resource exhausted" when loading Flax GPT-Neo 2.7B

## Environment info

- `transformers` version: 4.10.0.dev0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.5
- PyTorch version (GPU?): 1.8.1+cu102 (False)
- Tensorflow version (GPU?): 2.6.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
- Jax version: 0.2.19
- JaxLib version: 0.1.70
- Using GPU in script?:
- Using distributed or parallel set-up in script?:

### Who can help
@patrickvonplaten @patil-suraj @LysandreJik

## Information

I am not able to load the Flax GPT-Neo 2.7B model in my TPU VM v3-8 instance.
```python
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B", pad_token="", padding_side="left")
model = FlaxAutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-2.7B", pad_token_id=tokenizer.eos_token_id)
```
The model will download but will fail to load with
```
RuntimeError: Resource exhausted: Failed to allocate request for 100.00MiB (104857600B) on device ordinal 0
```
However, the pytorch version will load and run just fine.