OOM during saving step
Author: thies1006
Pull request: no
Labels: name: DeepSpeed

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'



Comments:
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.