feat(flax): leave restored weights on CPU
Author: borisdayma
Pull request: yes
Labels: name: WIP

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

When restoring a flax model with `.from_pretrained()`, leave the weights on CPU.

The removed section was linked to issue https://github.com/google/flax/issues/1261 which is now closed.
When calling `jnp.array()`, the tensors are converted from numpy arrays and placed on default device (typically GPU or TPU), which can cause issues when loading very large models that don't fit on one single instance.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@patil-suraj @patrickvonplaten 
<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_15295). All of your documentation changes will be reflected on that endpoint.

Thanks @borisdayma,

This looks good to me! @patil-suraj - could you take a look? :-)

Yeah maybe we can have a `load_on_cpu` argument somewhere.

Actually for training large models my approach is currently the following:
* force loading the model on CPU (whether pre-trained or not)
* overwrite `model._params` with their frozen version so they are compatible with `pjit` (and 2 copies would be excessive)
* create `params` from `model._params` by moving them to the correct devices (potentially with model parallelism)
* remove `model._params` to save space

I do something similar when restoring my optimizer state.

Maybe I could create a separate issue to discuss the proper strategy.
I'm thinking that loading a model should return both model and weights separately and follow JAX stateless approach…

What do you think of this approach: https://github.com/huggingface/transformers/compare/master...borisdayma:flax-cpu2

I could add this in this PR if you like it.

Basically following @patil-suraj experimental branch:
* we have an argument `abstract_init` so we don't initialize the parameters (mainly for the case where we will load a pretrained model)
* we have an argument `load_on_cpu` so we can pjit the model later if desired

This is the approach I'm currently using.

A few notes on `load_on_cpu`:

* it is based on [this comment](https://github.com/google/flax/discussions/1690#discussioncomment-1715766)
* it works but is a bit slow due to compilation

I took some time to think about it and I think we kind of have two problems here:

1) User should be able to force weights to be initialized on CPU. I fully agree that we should allow this especially since large model training is one of the main goals of Flax in Transformers. This applies to both **pretraining from scratch** and **fine-tuning**
2) We should not create allocate memory for weights that are loaded anyways when doing fine-tuning. This only applies to **fine-tuning**. 

To tackle 1) I think the best solution is actually to add not initialize the weights at all with a flag `do_init_weights` which we'll default to `True`, but that users can disable. I think in general the current design is not really in the spirit of JAX anyways as usually the user has full control over how to initialize the weights and allocating memory is normally not done when just "configuring the model" in JAX. So it's not really intuitive for JAX users that `bert = FlaxBertModel(config)` allocates memory for random weights anyways I think. However we had to also align Flax with PyTorch in the library which is why the current design makes sense from a Transformers point of view. I think the best solution here would be to allow the following:

```python
bert = FlaxBertModel(config, init_weights=False)  # it should be possible to infer the expected shapes here without allocating any memory no?
random_weights = jax.jit(bert.init_weights, static_argnums=(1,), backend="cpu")(input_shape)
bert.params = random_weights  # this setter method should verify that the params have the correct shape
```
=> this design should work no? Also for `from_pretrained(...)` no? Also there is a function in JAX that allows one to infer the expected shapes with 0 Flops and without allocating memory no so that we can get the expected shapes.

2) For 2 I think we should all Flax models would benefit from a mechanism that only inits weights that we need to actually init. Think we already talked about this @patil-suraj no. Not sure how feasible / easy to do this is, but I'm sure there is a way (worst case we add `init_head(...)` functions to all heads that are usually not included in the pretrained weighs.


Keen to hear your opinions here @borisdayma @patil-suraj 

> there is a function in JAX that allows one to infer the expected shapes with 0 Flops and without allocating memory no so that we can get the expected shapes.

Yes, `jax.eval_shape(fn, *fn_args…)` which will return the output with only shape and dtype.
What is cool is the inputs only need shape and dtype as well.

When using it we may not need to use `jax.jit(…, backend='cpu')` which I believe may add extra time through compilation (To be confirmed) so it's really just when training from scratch.

> We should not create allocate memory for weights that are loaded anyways when doing fine-tuning. This only applies to **fine-tuning**.

Solved with `jax.eval_shape()` which returns you the required structure to use `from_bytes()` later.

> I think in general the current design is not really in the spirit of JAX anyways as usually the user has full control over how to initialize the weights and allocating memory is normally not done when just "configuring the model" in JAX.

I was thinking that we could have the params separated from the model (in JAX spirit). I would for example do this:

```python
# load model & params
model, params = bert.from_pretrained(model_name, load_on_cpu=True/False, init_weigths=True/False)

# create state (loaded on TPU)
state = TrainState.create(model, params, optimizer)

# create state (sharded)
pjit_create = pjit(TrainState.create, in_axis_resources=…, out_axis_resources=…, donate_argnums=(…))
state = pjit_create(model, params, optimizer)
```

We could even directly load a sharded model and not go through the CPU loading with something like:

```
pjit_from_pretrained = pjit(bert.from_pretrained…, static_argnums=...)
```

I would need to test but the most efficient way is probably to pjit as a single function model loading + state creation and have params only within the state.

Separating params and model gives more control because sometimes we need `freeze(params)` with sharded spec so then you could have 2 copies of the params and I often end up doing hacky things like:
```
# use frozen dict instead of regular dict
model._params = freeze(model.params)

# create state
state = TrainState.create(model._params)

# save RAM
del model._params
```

I think the current approach (using non-frozen dict + adding params to model) is due making it more accessible to Pytorch users but as you scale up, it may add complexity to JAX logic.