CUDA OOM error when loading sharded checkpoint
Author: abarbet
Pull request: no
Labels: name: solved

### System Info

* `transformers` version: 4.27.1
* Platform: Linux-5.19.0-41-generic-x86_64-with-glibc2.35
* Python version: 3.9.12
* Huggingface_hub version: 0.13.2
* PyTorch version (GPU?): 2.0.0+cu117 (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?: Yes
* Using distributed or parallel set-up in script?: Yes, parallel (accelerate auto-mapping)

### Who can help?

@sgugger @pacman100 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

This is a port-over from an issue I wrote on the PyTorch forums [here](https://discuss.pytorch.org/t/cuda-oom-error-when-loading-sharded-checkpoint/180710). I received some help from the folks on the PyTorch side, but unfortunately, they seem to be suggesting that there may be an error in the way `Trainer` saves FSDP models. I will rehash the issue here with the additional context:

> We fine-tuned Stability’s StableLM-7b using Huggingface’s Trainer API (with FSDP) and then saved the resulting checkpoints in the sharded format that is typical for large language models. Quite surprisingly, however, attempting to load the model for inference leads to a strange error when loading one of the checkpoints (`Unable to load weights from pytorch checkpoint file`)
> 
> We took some further investigative steps by making a simple `torch.load` call on the problem shard, and got a CUDA OOM error. The exceedingly strange thing about this OOM error is that we are working with a node with 8xA100s (80GB), and the given state dict is only 171kB (comprising only 7 layers of the model). So, you can imagine seeing the following error was quite a shock:
> 
> ```
> torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 29.31 GiB (GPU 0; 79.19 GiB total capacity; 55.76 GiB already allocated; 22.48 GiB free; 55.76 GiB 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
> ```
> 
> After looking into this further, I discovered a few threads discussing this issue, like [this one 3](https://discuss.pytorch.org/t/cuda-error-out-of-memory-when-load-models/38011), and attempted some of the fixes, namely loading the state dict on CPU first. After doing so, I received the following error:
> `RuntimeError: Trying to resize storage that is not resizable`
> 
> So it seems that approach is out of the question. As I previously said, the strange thing here is that the first two shards load without issue, while the third and fourth cannot be loaded. Additionally, nothing seems particularly out of place in the shard-layer mapping JSON. I am stumped here.

The folks at PyTorch let us know that with FSDP models should _not_ be saved using `torch.save` and provided an example script of how they should be saved [here](https://github.com/pytorch/pytorch/blob/e71ab214226af1f9dbded944e939c6447e0e8f09/torch/distributed/checkpoint/examples/fsdp_checkpoint_example.py#L59). Does `Trainer` properly handle these larger models, or is there an extra step we should be taking here?

### Expected behavior

Typically, I would expect `save_model` to process the model shards in a way that allows them to be reloaded without issue using `from_pretrained` along with `accelerate`'s auto device mapping.

Comments:
cc @pacman100 

Hello, looking into this. In the meantime, could you try the main branch of transformers and accelerate and let us know if that works as expected? 


Hi @pacman100, thanks for taking a look! When you say `main` branch, do you mean bumping the versions of `transformers` and `accelerate`?

Hello, can you do `pip install git+https://github.com/huggingface/transformers` and `pip install git+https://github.com/huggingface/accelerate`? The above PR adds functionality for `SHARDED_STATE_DICT`.

Use Accelerate launcher with Trainer. More info here: [Using Accelerate Launcher with Trainer](https://huggingface.co/docs/transformers/main/en/main_classes/trainer#using-accelerate-launcher-with-trainer). Choose state_dict_type as `SHARDED_STATE_DICT` when answering questionnaire post running the command  `accelerate config`.

Please let us know if this solves the issue.

Thank you for this update! We utilized the new `SHARDED_STATE_DICT` functionality, but it looks like there may have been a small typo in the `trainer.py` code where the `full_osd` variable isn't saved in fsdp mode. I proposed a fix on the PR below, which allowed me to successfully save a model locally in fsdp mode:
https://github.com/huggingface/transformers/pull/24328