Add support for ZeRO-2/3 and ZeRO-offload in fairscale
Author: sgugger
Pull request: yes

# What does this PR do?

This PR adds support for the new `FullyShardedDataParallel` introduced in fairscale. See [this PR](https://github.com/facebookresearch/fairscale/pull/413) for more details.

The PR changes a tiny bit the behavior of the `--sharded_ddp` flag/training argument to support a list of options. You can still use the TrainingArguments class with `sharded_dpp=True` but if launching a script, `--sharded_ddp` has to be replaced with `--sharded_ddp simple`. The `--sharded_ddp` was marked as an experimental API so I think this breaking change is fine if properly documented.

Other values supported are: `zero_dp_2`, `zero_dp_2 offload`, `zero_dp_3` and `zero_dp_3 offload`. To fully take advantage of the `zero_dp_3`/`zero_dp_3 offload` the model passed to the `Trainer` will need to have its internal layers wrapped inside the `FullyShardedDataParallel`, but this out of scope for this particular PR.

For all those new modes, the model simply needs to be wrapped inside `FullyShardedDataParallel` but the optimizer needs to be created after the model wrapping (to get the parameters shards).

Note that:
- `predict_with_generate` does not work with this integration
- `cpu_offload` does not work for now due to the bug mentioned in [this issue](https://github.com/facebookresearch/fairscale/issues/421). Once the issue is fixed, the option should work with the existing code.

One thing to think further is that this integration breaks the usual convention that `self.model` is the original model (`FullyShardedDataParallel` consumes the model to use less memory).

Comments:
> Other values supported are: `zero2`, `zero2_offload`, `zero3` and `zero3_offload`. To fully take advantage of the `zero3`/`zero3_offload` the model passed to the `Trainer` will need to have its internal layers wrapped inside the `FullyShardedDataParallel`, but this out of scope for this particular PR.

Do you feel it's better to hardcode these combinations and not have a more flexible approach of:
```
--sharded_ddp "zero2;offload;future_option"
```
or
```
--sharded_ddp "zero2 offload future_option"
```
which would enable adding new features in the future, without needing to create all possible combinations of options which would double every time a new option will be added.

This is the cmd API I'm tentatively using for the pipelines `--pipeline "chunks=5 device_map=0:0-5,1:5-10 ...."`

> One thing to think further is that this integration breaks the usual convention that `self.model` is the original model (`FullyShardedDataParallel` consumes the model to use less memory).

Yes, we will need to rethink this - the trainer is getting more and more complex.

> Do you feel it's better to hardcode these combinations and not have a more flexible approach of:
>
> --sharded_ddp "zero2;offload;future_option"

Happy to explore that design as it seems more flexible and less prone to future breaking changes. Will adapt the PR accordingly once we get the wrapper to work.

Probably whitespace separation is more readable: `--sharded_ddp "zero2 offload future_option"`

Also we need to make sure that we distinguish between `FullyShardedDataParallel` and `ShardedDataParallel` since as the [commentary was made](https://github.com/facebookresearch/fairscale/pull/413#issuecomment-784168151), they aren't quite the same. Perhaps `not_full` for `ShardedDataParallel`? both should be corresponding to stage2 but they don't work in the same way.

Deepspeed has a `stage` param which goes from 0 to 3. where stage=0 doesn't enable ZeRO, and then each number matches the stage.

For the user's sake perhaps we could make things as similar as possible so it'd be more intuitive for them to switch between fairscale (and eventually pytorch) and deepspeed.

Also note that DeepSpeed exposes other params like the size of buckets, which actually are very important and need to be user-configurable. I won't be surprised that FSDP will also have those configurable down the road - i.e. more params.

Reworked the API to take your suggestion of list of options into account @stas00. I don't think we have to worry about uniformizing with deepspeed or cleaning more at this stage as:
- this API will evolve in the future (ShardedDataParallel might very well disappear if FullyShardedDataParallel is better, and this might change again on the road to be merged in PyTorch)
- we don't know yet all the options we will have between deepspeed/fairscale/PyTorch
- this is an experimental API and while we won't break it just for fun, we can make slight changes down the road.

Moving out the cl arg naming discussion from https://github.com/huggingface/transformers/pull/10354#pullrequestreview-596676591 to the open

So if's not DDP but DP, then we should probably change the cl arg to `_dp` as I suggested above so that it's consistently either DP or DDP all the way through.

Or perhaps we should just call it `--sharded`? the dp part is already inside the value anyway as in: `--sharded zero_dp_3`