Initialize Flax model params on CPU 
Author: gianlucadetommaso
Pull request: no
Labels: name: Good Second Issue

name: Flax

### Feature request

Currently, the `from_pretrained` method of Flax models automatically puts model parameters on a single GPU device, if available. For very large models, this is not great, as the model parameters may just not fit on GPU memory. 

In contrast, when passing `_do_init=False` to `from_pretrained`, the parameters are returned on CPU, outside the model.

I would love to have a feature that allows me to initialize model parameters on the device I want - in this case, on CPU - but at the same time initialize the model parameters within the model. Right now I have to call `_do_init=False` to avoid out-of-memory, but this causes inconsistencies with my API.

The feature could be either implemented as just another type (if we detect a numpy type, we initialize on CPU; otherwise on GPU) or as an additional argument, e.g. `initialize_on_cpu: bool = False`.

### Motivation

Described above. Another reason is to be more consistent with the PyTorch behaviour, where parameters are initialized (as a generator) on CPU.

### Your contribution

If we agree on on the design, I am happy to add this myself.

Comments:
Related to this, the `init_weights` method should not initialize all random parameters on GPU when `params` are actually passed (see, for example, [Flax GPT-J](https://github.com/huggingface/transformers/blob/495729427045c7a58e040fa9bf6df81c16f54208/src/transformers/models/gptj/modeling_flax_gptj.py#L401)). This makes me go out-of-memory even if I pass all the parameters to the method, initialized on CPU using `_do_init=False`.

cc @sanchit-gandhi 

Hey @gianlucadetommaso!

For reference, the PR to add the `_do_init` flag was added in this PR: #16148. Feel free to have a read through on what the motivations behind this PR and design were. Think you'd find it interesting!

I think it would be a nice design to have the params loaded on CPU by default. There's an open PR for this here: #15295. You're more than welcome to pick-up where Boris left off here and finish the PR! The comments detailing the proposed design are quite thorough, but feel free to ping me if you have any other questions or want to clarify something

@sanchit-gandhi thanks for the links! As soon as I have time, I can try and do this.

By the way, to make sure I am not just doing something wrong, it would help me a lot if you could have a look and comment on [this](https://github.com/google/jax/discussions/16659) discussion. It regards memory consumption of initializing a sharded state using pjit. I think you had a discussion related to it in [this](https://github.com/huggingface/transformers/issues/22224) thread before, thus it'd be great hearing your thoughts.

Awesome, sounds great! Had a look at the linked discussion - not entirely sure why we see this behaviour (think it's one for the JAX team to answer), but what you can do is use a few helper functions from the [T5x codebase](https://github.com/google-research/t5x/tree/main) to assist you here.

You can load your model into a T5x `Checkpointer` with `use_gda=True` (use global device arrays) on the CPU: https://github.com/huggingface/bloom-jax-inference/blob/2a04aa519d262729d54adef3d19d63879f81ea89/bloom_inference/generator.py#L86

And then save this `Checkpointer` state to a Google Cloud bucket (use the built in save function).

When you then come to loading your state, you can load each shard of your weights onto the mapped devices (so if shard 1 goes on device 1, it'll be loaded straight there, so you won't blow up your memory trying to load a sharded model onto your accelerator device):
https://github.com/huggingface/bloom-jax-inference/blob/2a04aa519d262729d54adef3d19d63879f81ea89/bloom_inference/generator.py#L95

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/main/CONTRIBUTING.md) are likely to be ignored.

Hey @gianlucadetommaso - going to re-open this one since we have a draft PR in progress at https://github.com/huggingface/transformers/pull/15295

Would you like to have a go at seeing the PR through to completion?

Hi @sanchit-gandhi, right now I would not have the time to drive this myself. Sorry.

No worries! Leaving this one open to the community :)