

TITLE
OOM during saving step

COMMENTS
8

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
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'

