--sharded_ddp "zero_dp_3 offload" fails with AssertionError 
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.5.1
- Platform: Linux-5.4.0-1043-aws-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.7.1+cu110 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: 8 x A100 (AWS p4d.24xlarge)
- Using distributed or parallel set-up in script?: python -m torch.distributed.launch 


### Who can help
<!-- Your issue 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.
-->
Library:

- deepspeed: @stas00
- trainer: @sgugger

 

## Information

Model I am using (Bert, XLNet ...): roberta-base

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

I want to perform distributed training using the example `run_mlm.py` script on the wikitext dataset.  Specifically, I'm trying to use sharded_ddp zero_dp_3 (i.e., fairscale) and **with offloading enabled**.    When I run _without_ offloading, it works.  But if I use the "offload" option, an AssertionError is thrown, as shown in the stack trace below. 

Steps to reproduce the behavior:

1. Install fairscale

        pip install fairscale==0.3.4
 
2. Run the example run_mlm.py as follows:

        export OMP_NUM_THREADS=11; 
        export TOKENIZERS_PARALLELISM=true; 
        python -m torch.distributed.launch --nproc_per_node=8 run_mlm.py --model_name_or_path roberta-base \
         --use_fast_tokenizer \
         --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1  \
         --do_train --do_eval --num_train_epochs 5 \
         --output_dir ./experiments/wikitext --sharded_ddp "zero_dp_3 offload" --fp16


<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

    Traceback (most recent call last): 
    File "run_mlm.py", line 492, in <module> main() File "run_mlm.py", line 458, in main train_result = trainer.train(resume_from_checkpoint=checkpoint) 
    File "/home/me/ve/lib/python3.6/site-packages/transformers/trainer.py", line 1120, in train tr_loss += self.training_step(model, inputs) 
    File "/home/me/ve/lib/python3.6/site-packages/transformers/trainer.py", line 1522, in training_step loss = self.compute_loss(model, inputs) 
    File "/home/me/ve/lib/python3.6/site-packages/transformers/trainer.py", line 1556, in compute_loss outputs = model(**inputs) 
    File "/home/me/ve/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl result = self.forward(*input, **kwargs) 
    File "/home/me/ve/lib/python3.6/site-packages/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 902, in forward self._lazy_init() 
    File "/home/me/ve/lib/python3.6/site-packages/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 739, in _lazy_init self._init_param_attributes(p)
    File "/home/me/ve/lib/python3.6/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context return func(*args, **kwargs) 
    File "/home/me/ve/lib/python3.6/site-packages/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 796, in _init_param_attributes assert p._fp32_shard.device == torch.device("cpu") 
    AssertionError

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
It should proceed to train.


As replied on the forums, you should rather use `--deepspeed` for Zero-offload. We will investigate this bug, but there is another one for the gradient scaler that will block you either way.<|||||>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/master/CONTRIBUTING.md) are likely to be ignored.