OOM during saving step
I'm trying to train the Blenderbot-9B model using the Deepspeed integration on 8 GPUs, each of them has 16GB VRAM (one node).

Script:
`deepspeed --hostfile myhostfile \
    ${_PATH}/examples/pytorch/summarization/run_summarization.py \
    --model_name_or_path hyunwoongko/blenderbot-9B \
    --do_train \
    --do_eval \
    --dataset_name cnn_dailymail \
    --dataset_config "3.0.0" \
    --source_prefix "summarize: " \
    --output_dir /tmp/tst-summarization \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 8 \
    --deepspeed ${_PATH}/tests/deepspeed/ds_config_zero3.json \
    --logging_steps 1 \
    --fp16 \
    --overwrite_output_dir \
    --save_steps 10 \
    --gradient_accumulation_steps 1 \
    --evaluation_strategy="steps" \
    --max_train_samples 10024 \
    --max_eval_samples 32 \
    --max_source_length 128 --max_target_length 128 \
    --eval_steps 5
`


Training and evaluation seem to run fine, I see about 10GB of VRAM occupied on each GPU, so there is even free space left on the GPUs. However afterwards during the saving step I get OOM, which I don't understand.

Log:
[log.txt](https://github.com/huggingface/transformers/files/6785035/log.txt)

Deespeed: 0.4.3+c9fee82
torch 1.8, cuda 11.1

Transformers:
'4.9.0.dev0'


cc @sgugger <|||||>I think this is more on the DeepSpeed side so cc-ing @stas00 to confirm.<|||||>Thank you for the full log.

Yes, it's on the deepspeed side. 

As you can see in https://huggingface.co/transformers/master/main_classes/deepspeed.html#getting-the-model-weights-out if you use:

```
{
    "zero_optimization": {
        "stage3_gather_fp16_weights_on_model_save": true
    }
}
```

then it reconsolidates the whole fp16 model on cpu, while gathering one layer at a time on GPU (and then moving to cpu).

You can see the code here: https://github.com/microsoft/DeepSpeed/blob/5652072e5451077da4179e5398b1c0c71c752c34/deepspeed/runtime/engine.py#L1991

So to first unblock you disable the above setting in `ds_config.json` by setting it to `false` and then use `zero_to_fp32.py` as explained here https://huggingface.co/transformers/master/main_classes/deepspeed.html#getting-the-model-weights-out if you need to extract the weights - as a bonus you get fp32 weights then.

Meanwhile let me have a look and see if I can make that code more memory tight - in theory if the training had enough gpu  memory it should too - e.g. can iterate over each param, rather full layers. I will experiment and get back to you.<|||||>OK the proper fix is here https://github.com/microsoft/DeepSpeed/pull/1220 if you want to try that branch, but should be merged into deepspeed master shortly and hopefully a new release will be made soon.
<|||||>sorry, looks like more work is needed there. will keep you posted.<|||||>This version should do the right thing as all the tests now pass: https://github.com/microsoft/DeepSpeed/pull/1223

Unfortunately missed the new deepspeed release, so will enter the next one.

Do let me know if you encounter any issues with this PR branch.

Thank you.<|||||>Thank you @stas00 ! Here is what I did:
- I tried with your PR, no OOM anymore during save step. So the original problem is solved. 
- However when trying to resume from that checkpoint via `--resume_from_checkpoint /tmp/tst-summarization/checkpoint-10` I ran out of cpu ram (512GB in my case). 

Just some further comments:
- Setting the option `"stage3_gather_fp16_weights_on_model_save": false` works as well (HF model is simply not saved).
- Exporting the Deepspeed checkpoint offline using the script as you said works and I can also resume training using this exported model via `--model_name_or_path`. <|||||>Closing as original problem was solved.