The scipt exit when calling "pretrainmodel.save_pretrained" because of OOM, but the previous training phase is well
Author: Zcchill
Pull request: no

### System Info

NVIDIA A100 Tensor Core GPUs (80G);
python 3.8.13;
torch 1.10.0+cu113;
transformers 4.20.1;


### Who can help?

 @sgugger 

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

When I fine-tune the OPT model based on Transformer.Trainer() (deepspeed, zero2), there is no problem in training process, however, the script fails to load the final best model in the end. It's really confused since the problem. I checked the function "pretrainmodel.save_pretrained", it seems that it has doing some operations like "del state_dict"  to save memory. Thus, it's really confused that our gpu memory is enough(80G) since it can finish the training process, but it fails to load the finally best model into memory. Here is the corresponding traceback. 

### Traceback

```
[INFO|trainer.py:1834] 2022-09-07 21:46:26,324 >> Loading best model from /ssdwork/results/results_mc6k_6.7b/20220907-1655/checkpoint-68 (score: 0.77099609375).

Traceback (most recent call last):
  File "finetune.py", line 497, in <module>
    main()

  File "finetune.py", line 460, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1409, in train
    return inner_training_loop(

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1771, in _inner_training_loop
    self._load_best_model()

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1867, in _load_best_model
    load_result = load_sharded_checkpoint(model, self.state.best_model_checkpoint, strict=strict_load)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/modeling_utils.py", line 445, in load_sharded_checkpoint
    state_dict = torch.load(os.path.join(folder, shard_file))
  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 607, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 882, in _load
    result = unpickler.load()

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 857, in persistent_load
    load_tensor(data_type, size, key, _maybe_decode_ascii(location))

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 846, in load_tensor
    loaded_storages[key] = restore_location(storage, location)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 175, in default_restore_location
    result = fn(storage, location)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 157, in _cuda_deserialize
    return obj.cuda(device)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/_utils.py", line 79, in _cuda
    return new_type(self.size()).copy_(self, non_blocking)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/cuda/__init__.py", line 606, in _lazy_new
    return super(_CudaBase, cls).__new__(cls, *args, **kwargs)

RuntimeError: CUDA out of memory. Tried to allocate 12.40 GiB (GPU 0; 79.17 GiB total capacity; 0 bytes already allocated; 3.48 GiB free; 0 bytes reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

Traceback (most recent call last):
  File "finetune.py", line 497, in <module>
    main()

  File "finetune.py", line 460, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1409, in train
    return inner_training_loop(

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1771, in _inner_training_loop
    self._load_best_model()

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1867, in _load_best_model
    load_result = load_sharded_checkpoint(model, self.state.best_model_checkpoint, strict=strict_load)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/modeling_utils.py", line 445, in load_sharded_checkpoint
    state_dict = torch.load(os.path.join(folder, shard_file))

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 607, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 882, in _load
    result = unpickler.load()

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 857, in persistent_load
    load_tensor(data_type, size, key, _maybe_decode_ascii(location))

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 846, in load_tensor
    loaded_storages[key] = restore_location(storage, location)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 175, in default_restore_location
    result = fn(storage, location)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 157, in _cuda_deserialize
    return obj.cuda(device)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/_utils.py", line 79, in _cuda
    return new_type(self.size()).copy_(self, non_blocking)

  File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/cuda/__init__.py", line 606, in _lazy_new
    return super(_CudaBase, cls).__new__(cls, *args, **kwargs)

RuntimeError: CUDA out of memory. Tried to allocate 12.40 GiB (GPU 0; 79.17 GiB total capacity; 0 bytes already allocated; 3.38 GiB free; 0 bytes reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

[2022-09-07 21:46:38,873] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 859055
[2022-09-07 21:46:38,873] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 859056
[2022-09-07 21:46:38,874] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 859057
[2022-09-07 21:46:38,874] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 859058
[2022-09-07 21:46:38,874] [ERROR] [launch.py:184:sigkill_handler] ['/home/anaconda3/envs/tk-instruct/bin/python', '-u', 'finetune.py', '--local_rank=3', '--deepspeed', '/home/opt/ds_config.json', '--model_name_or_path', 'facebook/opt-6.7b', '--train_file', '/home/opt/data/train_mc6k.csv', '--validation_file', '/home/opt/data/valid_mc_300.csv', '--do_train', '--do_eval', '--fp16', '--output_dir', '/ssdwork/opt/results/results_mc6k_6.7b/20220907-1655/', '--num_train_epochs', '1000', '--per_device_train_batch_size', '4', '--evaluation_strategy', 'epoch', '--save_strategy', 'epoch', '--load_best_model_at_end', '--metric_for_best_model', 'eval_loss', '--greater_is_better', 'False', '--gradient_accumulation_steps', '32', '--use_fast_tokenizer', 'False', '--learning_rate', '1e-05', '--warmup_steps', '10', '--save_total_limit', '1', '--overwrite_cache', '--block_size', '2048'] exits with return code = 1

/home/anaconda3/envs/tk-instruct/lib/python3.8/multiprocessing/resource_tracker.py:216: UserWarning: resource_tracker: There appear to be 6 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d 
```
### Expected behavior

It should be noted that the total process is good when fine-tuning model with 2.7B parameters and below but will exit for model with 6.7B parameters and above. I think it may not simply due to OOM, since our single GPU memory is already 80G , and we mainly call the Trainer() to finish the process (based on deepspeed zero2 optimizer). I think if I ignored some significant parameters like "--save_on_each_node".

Comments:
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/main/CONTRIBUTING.md) are likely to be ignored.

> ### System Info
> NVIDIA A100 Tensor Core GPUs (80G); python 3.8.13; torch 1.10.0+cu113; transformers 4.20.1;
> 
> ### Who can help?
> @sgugger
> 
> ### Information
> * [x]  The official example scripts
> * [ ]  My own modified scripts
> 
> ### Tasks
> * [x]  An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
> * [ ]  My own task or dataset (give details below)
> 
> ### Reproduction
> When I fine-tune the OPT model based on Transformer.Trainer() (deepspeed, zero2), there is no problem in training process, however, the script fails to load the final best model in the end. It's really confused since the problem. I checked the function "pretrainmodel.save_pretrained", it seems that it has doing some operations like "del state_dict" to save memory. Thus, it's really confused that our gpu memory is enough(80G) since it can finish the training process, but it fails to load the finally best model into memory. Here is the corresponding traceback.
> 
> ### Traceback
> ```
> [INFO|trainer.py:1834] 2022-09-07 21:46:26,324 >> Loading best model from /ssdwork/results/results_mc6k_6.7b/20220907-1655/checkpoint-68 (score: 0.77099609375).
> 
> Traceback (most recent call last):
>   File "finetune.py", line 497, in <module>
>     main()
> 
>   File "finetune.py", line 460, in main
>     train_result = trainer.train(resume_from_checkpoint=checkpoint)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1409, in train
>     return inner_training_loop(
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1771, in _inner_training_loop
>     self._load_best_model()
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1867, in _load_best_model
>     load_result = load_sharded_checkpoint(model, self.state.best_model_checkpoint, strict=strict_load)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/modeling_utils.py", line 445, in load_sharded_checkpoint
>     state_dict = torch.load(os.path.join(folder, shard_file))
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 607, in load
>     return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 882, in _load
>     result = unpickler.load()
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 857, in persistent_load
>     load_tensor(data_type, size, key, _maybe_decode_ascii(location))
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 846, in load_tensor
>     loaded_storages[key] = restore_location(storage, location)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 175, in default_restore_location
>     result = fn(storage, location)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 157, in _cuda_deserialize
>     return obj.cuda(device)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/_utils.py", line 79, in _cuda
>     return new_type(self.size()).copy_(self, non_blocking)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/cuda/__init__.py", line 606, in _lazy_new
>     return super(_CudaBase, cls).__new__(cls, *args, **kwargs)
> 
> RuntimeError: CUDA out of memory. Tried to allocate 12.40 GiB (GPU 0; 79.17 GiB total capacity; 0 bytes already allocated; 3.48 GiB free; 0 bytes reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
> 
> Traceback (most recent call last):
>   File "finetune.py", line 497, in <module>
>     main()
> 
>   File "finetune.py", line 460, in main
>     train_result = trainer.train(resume_from_checkpoint=checkpoint)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1409, in train
>     return inner_training_loop(
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1771, in _inner_training_loop
>     self._load_best_model()
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/trainer.py", line 1867, in _load_best_model
>     load_result = load_sharded_checkpoint(model, self.state.best_model_checkpoint, strict=strict_load)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/transformers/modeling_utils.py", line 445, in load_sharded_checkpoint
>     state_dict = torch.load(os.path.join(folder, shard_file))
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 607, in load
>     return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 882, in _load
>     result = unpickler.load()
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 857, in persistent_load
>     load_tensor(data_type, size, key, _maybe_decode_ascii(location))
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 846, in load_tensor
>     loaded_storages[key] = restore_location(storage, location)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 175, in default_restore_location
>     result = fn(storage, location)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/serialization.py", line 157, in _cuda_deserialize
>     return obj.cuda(device)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/_utils.py", line 79, in _cuda
>     return new_type(self.size()).copy_(self, non_blocking)
> 
>   File "/home/anaconda3/envs/tk-instruct/lib/python3.8/site-packages/torch/cuda/__init__.py", line 606, in _lazy_new
>     return super(_CudaBase, cls).__new__(cls, *args, **kwargs)
> 
> RuntimeError: CUDA out of memory. Tried to allocate 12.40 GiB (GPU 0; 79.17 GiB total capacity; 0 bytes already allocated; 3.38 GiB free; 0 bytes reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
> 
> [2022-09-07 21:46:38,873] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 859055
> [2022-09-07 21:46:38,873] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 859056
> [2022-09-07 21:46:38,874] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 859057
> [2022-09-07 21:46:38,874] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 859058
> [2022-09-07 21:46:38,874] [ERROR] [launch.py:184:sigkill_handler] ['/home/anaconda3/envs/tk-instruct/bin/python', '-u', 'finetune.py', '--local_rank=3', '--deepspeed', '/home/opt/ds_config.json', '--model_name_or_path', 'facebook/opt-6.7b', '--train_file', '/home/opt/data/train_mc6k.csv', '--validation_file', '/home/opt/data/valid_mc_300.csv', '--do_train', '--do_eval', '--fp16', '--output_dir', '/ssdwork/opt/results/results_mc6k_6.7b/20220907-1655/', '--num_train_epochs', '1000', '--per_device_train_batch_size', '4', '--evaluation_strategy', 'epoch', '--save_strategy', 'epoch', '--load_best_model_at_end', '--metric_for_best_model', 'eval_loss', '--greater_is_better', 'False', '--gradient_accumulation_steps', '32', '--use_fast_tokenizer', 'False', '--learning_rate', '1e-05', '--warmup_steps', '10', '--save_total_limit', '1', '--overwrite_cache', '--block_size', '2048'] exits with return code = 1
> 
> /home/anaconda3/envs/tk-instruct/lib/python3.8/multiprocessing/resource_tracker.py:216: UserWarning: resource_tracker: There appear to be 6 leaked semaphore objects to clean up at shutdown
>   warnings.warn('resource_tracker: There appear to be %d 
> ```
> 
> ### Expected behavior
> It should be noted that the total process is good when fine-tuning model with 2.7B parameters and below but will exit for model with 6.7B parameters and above. I think it may not simply due to OOM, since our single GPU memory is already 80G , and we mainly call the Trainer() to finish the process (based on deepspeed zero2 optimizer). I think if I ignored some significant parameters like "--save_on_each_node".

I have the same problem as you. How did you solve it? 

After using deepspeed for large models（xl、xxl）, the parameters will be stored in pieces, but the loading method of the best model parameters will change, and deepspeed is not used.