[Discussion] Loading and initialising large models with Flax
Opening up this issue to discuss how to initialise and load large models with Flax.

# Goals:
1. Be able to initialise and pjit really large models that don’t fit on a single device.
2. Avoid memory fragmentation when loading pre-trained weights in `from_pretrained`.

# Issues:
- With the current API it’s not possible to initialise a model which does not fit on a single device, since the model
weights are always initialised on the available device when we init the model class.

```python
model = FlaxPretrainedModel(conifg) # weights are always initialised in this call. OOMs when the model can't fit on the device. 
```

It’s also not easy to use with pjit or init on CPU as described by @borisdayma  in this [colab](https://colab.research.google.com/drive/1S06sYeuJ8n7P2EiHunYI6SPPq-Ycfum8?usp=sharing)

- The `from_pretrained` method creates memory fragmentation and also fails to load large models. Because in `from_pretrained` the model is first initialised randomly to get names and shapes of parameters and then the pre-trained weights are loaded. Here both the random weights and pre-trained weights are on the device so we have two copies of weights in device memory.

# Possible solutions:

Posting the two options proposed by @borisdayma

## Option 1: params returned separately

```python
model, params = FlaxPretrainedModel(…)
predictions = model(params, input)
```

Cons:
- Would be a big breaking change and doesn’t really fit well with the rest of the Transformers API.
- Still doesn’t solve the problem of initialising really large models.

## Option 2: params created separately

```python
model = FlaxPretrainedModel(…)
params = model.init_weights(…)
```

- here we can pjit the `init_weights` method or jit it with cpu backend to initialise the weights on CPU.

How would we implement this?
- introduce a new flag `do_init` (need a better name) which will default to `True`, and the API will stay the same as it is now.
- if it’s `False` then don’t init the params. In this case `model.params` will always be `None` and the users will always have to pass
    params to the forward call. We use `eval_shape` to get the params tree with shapes and save it on `model.params_shape_tree` (better name ?)

It would look something like this:

```python
# existing API
model = FlaxPretrainedModel(…)
model.params # this will store the model params

# with do_init=False
model = FlaxPretrainedModel(do_init=False)
model.params # this will be None

model(inputs) # this will raise error since params are not initialised

params = model.init_weights()
# always pass params to forward
model(inputs, params=params)
```

### How to handle `from_pretrained` in this case (`do_init=False`)?

This is also related to the second goal. To avoid fragmentation we could use `jax.eval_shape` here to get the params shape info that is required when loading pre-trained weights.

But there are a few issues with this:

**1. How to load a model with a head with pre-trained weights from the base model?**

In this case, we need to randomly initialise only the head weights and load rest from pre-trained. But there seems to be no way of only initialising some parameters and not others.
One possible solution is to add a method called `init_head` to every module with a head and when head weights are missing we call that method. See #15584 for what that would look like.
But this adds a lot of boilerplate code since we’ll need to do this for every head module in every flax model.

It’s also not possible to do this when some other intermediate weights are missing and need to be initialised randomly.

@marcvanzee @jheek @Avital is there any way in Flax where we could just initialise certain parameters without initialising the whole model?

**2.  What if the pre-trained weights don’t fit on a single device?**

In this case, we’ll need to load the weights on the CPU and then let the user shard those across devices.
Should we introduce an argument called `device` or `backend` which will specify where to load the weights? (cc @patrickvonplaten )


**3. Should we return the loaded weights or assign them to `model.params` ?**

When `do_init` is `False` we always set `model.params` to `None` and keep the params external as described above. But then the `from_pretrained` API will look something like below.

```python
model, params = FlaxPretrainedModel.from_pretrained(do_init=False)
```

which again is not well aligned with the rest of the API. What do we think about this? Looking forward to your suggestions @patrickvonplaten @borisdayma @LysandreJik 

@jheek @marcvanzee @avital It would be awesome if you could take a look at this and share how large model loading is handled in Flax.
Had an offline discussion with @patrickvonplaten and here's what we discussed.

We can go with option 2 with the following behaviour:

- introduce a new argument `FlaxPretrainedModel.__init__` called `do_init`, default to `True`. The API stays as it is right now.
`do_init` will be completely optional behaviour only to be used initialising and/or loading large models.
- If `do_init=False`, don't initialise the params, set `model.params=None`. And let the user call `model.init_weights`. Here the user has control so he can shard the init or init on CPU.


**For `from_pretrained`:**
- if `do_init=True`, which is the default behaviour, nothing changes.
- if `do_init=False`, don't do random initialisation and only load the pre-trained weights. And _**return weights along with the model**_.
- The user here will still have to call the `model.init_weights` method. If `do_init=False` the user will always have to call `model.init_weights` even when loading pre-trained weights. This will ensure to initialise any missing head params (if there are any) otherwise it'll be a no-op.
- if `do_init=False` we always load pre-trained params on CPU. We can then transfer these to the device by pjitting or jitting `init_weights` method. ~~This is a breaking change, but since this is an edge case and only required for large models here we think the user should be aware of these things.~~

This is how it would look like

```python
model, params = FlaxPretrainedModel.from_pretrained(do_init=False) # params on CPU
params = model.init_weights(params) # randomly initialise missing weights. 
# We could also `pjit` this method to shard the params which will also take care of moving these to the device from CPU.
```

This should help resolve both issues while still offering compatibility with the API. Curious to hear what you think @borisdayma and others.<|||||>One small addition:

> This is a breaking change, but since this is an edge case and only required for large models here we think the user should be aware of these things.

It's not really a breaking change since we introduce a new behavior with `do_init=False` so this can't by definition break anything existing. Very interested in hearing others opinion on this!<|||||>This is great, I like it!!!

I'm just wondering if `from_pretrained` could even be pjitted directly though not sure if it works with loading weights from an external file.<|||||>Hi, sorry for the late reply on this!

For the single-host case, initializing parameters simply comes down to pjitting the init function with a PartitioningSpec, there is nothing special about it. Since Flax Modules have a functional API, calling `init` just returns the parameters (we don't have a `model.params`), so I think the problems you are running into in this issue do not apply in our examples.

For the multi-host case, we recommend using GlobalDeviceArray (GDA)'s serialization library. (there unfortunately is not much documentation on this yet, but [these tests give some examples](https://github.com/google/jax/blob/main/jax/experimental/gda_serialization/serialization_test.py)). However, I think you are mainly talking about the single-host case here.<|||||>Thanks for the answer @marcvanzee !

> I think you are mainly talking about the single-host case here.

Yes, the example posted above is single-host. But we also want to enable this for multi-host. If I'm not wrong pjiting
`init` should also work on multi-host, no ?<|||||>> `init` should also work on multi-host, no ?

I guess the main issue would be when loading a checkpoint if it doesn't fit on CPU (which we have not yet reached so maybe can keep it for later?)<|||||>`pjit` is meant to work for both single- and multi-host settings, with the latter using the newer GDA abstraction should help automate some tedious multihost details when later dealing w. checkpointing, etc.<|||||>Thanks a lot for the feedback guys! Think we're going in the right direction then by fully seperating the `params` from the model architecture by doing `_do_init=False`. By setting this flag to False, we'll not allow to set `model.params` at all, so one can be sure that `params` are fully separated from the model architecture. 

I think this is very natural when initializing a model from scratch, e.g.:

```py
model = FlaxBertModel(config, _do_init=False)  # calling or setting model.params will error out now
params = model.init_weighs()
```

but not very natural when using `from_pretrained(...)`:

```py
model, params = FlaxBertForQuestionAnswering.from_pretrained("bert-base-cased", _do_init=False)  # calling or setting model.params will also error out, returned params are *only* the pretrained weights
params = model.init_weighs(params=params)  # in case of missing uninitialized layers - *e.g.* the QA head, we'll pass params to the model init function
```

Flax will then probably be the only framework the returns both model parameters and the model architecture when `from_pretrained(...)` is called with `_do_init=False`. This is not very "transformers-like", but it's the only way in our opinion we can ensure a more JAX-like approach to dealing with pretrained models and weights.<|||||>I think the proposed API fits the JAX philosophy nicely, we've experienced that maintaining a state will eventually bite you, so it seems like a good decision to choose this functional API.

It is somewhat unfortunate though you had to make an exception in your API, hopefully it will not be too confusing for your users.

@cgarciae