"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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### 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="</s>", 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.

Hi, any updates on this?<|||||>Thanks for reporting this, I'm looking into it.<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.<|||||>Unsale<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.<|||||>Running into this same issue when trying to load t5-3b as a Flax model from the Pytorch version on TPU3.8<|||||>Working on a feature that should fix this issue. This is probably because the model is initialized randomly and the weights are on the device, and then the pre-trained weights are also loaded directly on the device. So working on a feature that allows initializing the model only abstractly to consume less memory. Should be available in a couple of weeks :) <|||||>Any update on this @patil-suraj 
We are experiencing this when trying to load RoBERTa using Flax with TPU v3-8