CUDA OOM at `self.optimizer.consolidate_state_dict()` in Trainer when using sharded_ddp 
Author: yana-xuyan
Pull request: no

## 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.12.3
- Platform: Linux-5.4.0-1057-aws-x86_64-with-debian-buster-sid
- Python version: 3.7.10
- PyTorch version (GPU?): 1.7.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: 8 GPUs
- Using distributed or parallel set-up in script?: sharded_ddp (fairscale 0.4.2)

### Who can help

@sgugger


## Information

Model I am using (Bert, XLNet ...): BART-base

The problem arises when using:
* my own modified scripts: (give details below)
* I'm using my own code which is mainly modified from `run_mlm.py`(https://github.com/huggingface/transformers/blob/v4.12.3/examples/pytorch/language-modeling/run_mlm.py) for pretraining with huggingface trainer

The tasks I am working on is:
* my own task or dataset: (give details below)
* I'm using wikipedia corpus.

## To reproduce

Steps to reproduce the behavior:

1. run the script `run_mlm.py`(https://github.com/huggingface/transformers/blob/v4.12.3/examples/pytorch/language-modeling/run_mlm.py)
2. run the script with the following command line
```
python -m torch.distributed.launch --nproc_per_node=8 --master_port=10000 run_mlm.py \
    --model_name_or_path roberta-base \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
    --do_eval \
    --cache_dir /tmp/test-mlm \
    --output_dir /tmp/test-mlm \
    --sharded_ddp simple \
    --overwrite_output_dir \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 4
```

<!-- 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 538, in <module>
    main()
  File "run_mlm.py", line 487, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/transformers/trainer.py", line 1383, in train
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/transformers/trainer.py", line 1495, in _maybe_log_save_evaluate
    self._save_checkpoint(model, trial, metrics=metrics)
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/transformers/trainer.py", line 1565, in _save_checkpoint
    self.optimizer.consolidate_state_dict()
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/fairscale/optim/oss.py", line 358, in consolidate_state_dict
    obj_list, src=self._local_to_global_rank[rank], group=self.group,
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 1403, in broadcast_object_list
    object_list[i] = _tensor_to_object(obj_view, obj_size)
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 1187, in _tensor_to_object
    out = pickle.loads(buf)
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/torch/storage.py", line 141, in _load_from_bytes
    return torch.load(io.BytesIO(b))
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/torch/serialization.py", line 595, in load
    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/torch/serialization.py", line 774, in _legacy_load
    result = unpickler.load()
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/torch/serialization.py", line 730, in persistent_load
    deserialized_objects[root_key] = restore_location(obj, location)
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/torch/serialization.py", line 175, in default_restore_location
    result = fn(storage, location)
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/torch/serialization.py", line 155, in _cuda_deserialize
    return storage_type(obj.size())
  File "/home/ubuntu/anaconda3/envs/pytorch_p37/lib/python3.7/site-packages/torch/cuda/__init__.py", line 462, in _lazy_new
    return super(_CudaBase, cls).__new__(cls, *args, **kwargs)
RuntimeError: CUDA error: out of memory



## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
Could you please tell me how to fix this issue?

Comments:
You should remove the checkpointing by using `save_strategy="no"` to avoid getting out of memory: since the optimizer states are sharded to avoid spending too much GPU memory, you get the OOM error when the `Trainer` tries to save the optimizer state.

Thank you for your reply! But if I remove the checkpointing, then I cannot get any checkpoint until the end of the training or resume training from the checkpoint. Do you have any idea if I still want to the checkpoints?

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.

> Thank you for your reply! But if I remove the checkpointing, then I cannot get any checkpoint until the end of the training or resume training from the checkpoint. Do you have any idea if I still want to the checkpoints?

So are there any solutions?

At the end, I rewrite the Trainer saving function to skip saving the optimizer weights. Please let me know if there's any better solution! Thx!

Same here. I just commented that line which saves the optimizer

Any update?  @sgugger 

Probably need to add an option to `training_argument` that we want to disable the saving of optimizer state_dict, and default is `True` 

> 

Yeah, I think that could be the "unsatisfactory" solution for us, but at least it is a simple solution to implement. And I assume it wont affect much if we dont save optimizer right?

@allanj  @sgugger  I found the solution.

You need to update fairscale==0.4.6, then in the huggingface transformers trainer code (function `create_optimizer`), add `force_broadcast_object=True` , then problem solved. Now it can perfectly store the optimizer state_dict compared to just skipping the save of optimizer weights which @yana-xuyan  suggested.

You can find the `force_broadcast_object=True` in the latest code in [fairscale](https://github.com/facebookresearch/fairscale/blob/main/fairscale/optim/oss.py)


Do you want to make a PR with the fix?

> Do you want to make a PR with the fix?

Yeah sure, with pleasure. Will do it before this weekend. Cheers!