flax model parallel training
Author: patil-suraj
Pull request: yes

# What does this PR do?

Adds model parallel training example for GPTNeo using jax's [`pjit`](https://jax.readthedocs.io/en/latest/jax.experimental.pjit.html) transformation.
(This example probably just works on a single TPU v3-8).

This should enable training bigger models like 1.3B GPTNeo on a single TPU V3-8.


The `partition.py` file defines the `PyTree` of the `PartitionSpec` file which describes how the model parameters will be sharded. The actual sharding is automatically handled by `pjit`. 

The key idea is to `pjit` the entire training step function. To do that we
- Define the mesh structure.
- Define `PartitionSpec` for every input argument and return value of the pjitted function. The axis names that are used here should match the axis names used in `PartitionSpec`. This means we need the spec for our parameter and optimizer state PyTrees
- The structure of the `PyTree` of `PartitionSpec` needs to match the structure of the `PyTree` of the actual values.
- Call the pijitted fun in a mesh context.



Below is a not-so minimal code-snippet that describes the approach

```python

# init our model
model = FlaxGPTNeoForCausalLM.from_pretrained("gpt-neo-125M")

# get the partition spec for model params
param_spec = set_partitions(unfreeze(model.params))

# get optimizer
optim = optax.adamw(learning_rate=decay_fn)

# mesh defination
mesh_devices = np.array(jax.devices()).reshape(1, jax.local_device_count())

def get_initial_state(params):
    state = optim.init(params)
    return tuple(state), params

# init optim in abstract way, this just returns the PyTree of opt_state with shapes
# so we can get the PartitionSpec for opt_state using the tree
shapes = jax.tree_map(lambda x: x.shape, model.params)
state = jax.eval_shape(get_initial_state, shapes)

# Get the opt spec
def get_opt_spec(x):
    if isinstance(x, dict):
        return param_spec
    return None

opt_state_spec, param_spec = jax.tree_map(
    get_opt_spec, state, is_leaf=lambda x: isinstance(x, (dict, optax.EmptyState))
)

# Now actually initialize the opt state
# this also takes care of sharding the opt and param state according to the spec.
p_get_initial_state = pjit(
    get_initial_state,
    in_axis_resources=None,
    out_axis_resources=(opt_state_spec, param_spec),
)
with mesh(mesh_devices, ("dp", "mp")):
    opt_state, params = p_get_initial_state(freeze(model.params))

# define out train step
def train_step(params, opt_state, dropout_rng, batch):
	....
	return new_params, tuple(new_opt_state), new_dropout_rng, metrics

# pjit the train step
# in_axis_resources and out_axis_resources expect the PartitionSpec 
# for every input argument and return values
p_train_step = pjit(
    train_step,
    in_axis_resources=(param_spec, opt_state_spec, None, None),
    out_axis_resources=(param_spec, opt_state_spec, None, None),
)

# do the training
with mesh(mesh_devices, ("dp", "mp")):
    params, state, loss, rng = p_train_step(params, opt_state, ...)
```

As we can see above, all the sharding logic is outside of the model definition, so ideally we don't need to modify the modeling code. This also means it should be possible to apply this to any other model by defining the right `PyTree` of `PartitionSpec`. 

A few things to consider for future work.

- A convenient way to get the PyTree of model parameters, so we can define the partition spec.
- Currently,  model weights are initialized when the model class is instantiated. This can cause problems for models that cannot fit on one device. There should be an option to abstractly initialize the model without having to initialize the weights. This will also allow a convenient way to get the PyTree.
- The `from_pretrained` method also directly puts the weights on the device, we need to consider either sharded loading or initially loading the weights on the CPU then sharding them on the devices, to avoid OOM with huge models.



Comments:
This is typically done using using various load balancing methods, e.g. deepspeed pipe has:
https://www.deepspeed.ai/tutorials/pipeline/#load-balancing-pipeline-modules

pytorch has these too but I can't find any mentions of these in their docs. 
Have to go to the source:
https://github.com/pytorch/pytorch/blob/58adaaba60441c1ed59f35389598aabf91a772dd/torch/distributed/pipeline/sync/_balance/__init__.py
```
def balance_cost
def balance_by_time
def balance_by_size
```
Is that what you're referring to, @patrickvonplaten