run_clm.py exits with error -9 on checkpoint restart
Author: randywreed
Pull request: no
Labels: name: DeepSpeed

name: bug

### System Info

```shell
- `transformers` version: 4.20.0.dev0
- Platform: Linux-5.13.0-30-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.6.0
- PyTorch version (GPU?): 1.10.1 (True)
- Tensorflow version (GPU?): 2.7.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: 2
- Using distributed or parallel set-up in script?: deepspeed
```


### Who can help?

@patil-suraj @sgugger   



### Information

- [X] The official example scripts
- [ ] 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

I was running run_clm.py on my own custom data on GPT-J6B. I ran out of disk space, and restarted from latest checkpoint. All seemed to restart appropriately and then the script crashed with an error -9. No error message. Full log attached.
```
Using /home/ubuntu/.cache/torch_extensions/py38_cu111 as PyTorch extensions root...
No modifications detected for re-loaded extension module utils, skipping build step...
Loading extension module utils...
Time to load utils op: 0.00033783912658691406 seconds
[INFO|deepspeed.py:449] 2022-05-15 00:12:20,447 >> Attempting to resume from finetuned/checkpoint-22
[2022-05-15 00:12:51,292] [INFO] [engine.py:2754:_get_all_zero_checkpoint_state_dicts] successfully read 2 ZeRO state_dicts for rank 0
[2022-05-15 00:12:57,785] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 97277
[2022-05-15 00:12:57,786] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 97278
[2022-05-15 00:12:57,786] [ERROR] [launch.py:184:sigkill_handler] ['/usr/bin/python3', '-u', 'run_clm.py', '--local_rank=1', '--deepspeed', './Finetune_GPTNEO_GPTJ6B/finetuning_repo/ds_config_gptj6b.json', '--model_name_or_path', 'EleutherAI/gpt-j-6B', '--train_file', 'Jesus_sayings.txt', '--do_train', '--fp16', '--overwrite_cache', '--evaluation_strategy=steps', '--output_dir', 'finetuned', '--num_train_epochs', '5', '--eval_steps', '1', '--gradient_accumulation_steps', '32', '--per_device_train_batch_size', '1', '--use_fast_tokenizer', 'False', '--learning_rate', '5e-06', '--warmup_steps', '10', '--save_total_limit', '2', '--save_steps', '2', '--save_strategy', 'steps', '--tokenizer_name', 'gpt2'] exits with return code = -9
```
This is running on two 48gb GPUS using deepspeed. Was training without problem until crash and then on restart got the error.   
Original Command:  
```
deepspeed --num_gpus=2 run_clm.py --deepspeed \
./Finetune_GPTNEO_GPTJ6B/finetuning_repo/ds_config_gptj6b.json \
--model_name_or_path EleutherAI/gpt-j-6B --train_file Jesus_sayings.txt \
--do_train --fp16 --overwrite_cache --evaluation_strategy=steps --output_dir \
finetuned --num_train_epochs 5 --eval_steps 1 --gradient_accumulation_steps 32 \
--per_device_train_batch_size 1 --use_fast_tokenizer False --learning_rate \
5e-06 --warmup_steps 10 --save_total_limit 2 --save_steps 2 --save_strategy \
steps --tokenizer_name gpt2
```

ds_config_gpt6b.json
```
{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
[full_traceback_run_CLM_error.txt](https://github.com/huggingface/transformers/files/8694097/full_traceback_run_CLM_error.txt)

        "initial_scale_power": 12,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": false
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": false
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_fp16_weights_on_model_save": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
```
[full_traceback_run_CLM_error.txt](https://github.com/huggingface/transformers/files/8694101/full_traceback_run_CLM_error.txt)



### Expected behavior

```shell
Should restart and continue training.
```


Comments:
Looks like the error was thrown by DeepSpeed reloading the checkpoint, so maybe your issue would be better suited in their repo? Also cc @stas00 for information.

Based on your report I don't think it has anything to do with Deepspeed.

```
[2022-05-15 00:12:57,785] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 97277
[2022-05-15 00:12:57,786] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 97278
```

It looks like `cgroups` or oom killer or whatever resource control system you use killed the launcher process, which killed the other processes. Check `dmesg` to see if you get a report of a process killed by kernel or a helper util.

Do you have enough CPU memory to load the checkpoint? It's possible that there was more CPU RAM at the start and then it got reduced when you restarted?

You could activate `--skip_memory_metrics 0` with just a few steps and get stats on how much CPU memory your script is using in each stage, then comparing to your free CPU memory.

Adding swap memory often can help with the situation. Let me know if you want instructions for that.

How much CPU memory do you have on this host? Do you use it for training only or is it part of a desktop that you use for other things.

@stas00 So this is a cloud system used just for finetuning. It has 200gb of Ram, 2 48gb GPUs. But I set up a ram tracker and it hits 100% when it fails. So that's clearly the problem. Is the swap you are talking about different than a normal swap file in ubuntu?


As I suspected. Those resource controlling tools aren't very user-friendly. I just run a lot into this `Killed` w/o any explanation use-case on HPC so I know to suspect these.

Yes, normal swap. The question is whether the `cgroups` is set to have swap help out. Doesn't hurt to try.

`cgroups` typically monitors the residential memory, so the unused memory will go to swap if there is one.

Here is how I normally add it (of course edit the paths):

```

### Add a new swap file or extend one ###

# turn off all swap processes
sudo swapoff -a

# add 128GB file (or resize it if it already exists)
sudo dd if=/dev/zero of=/mnt/nvme0/swapfile bs=1G count=128

# prep as swap
sudo chmod 600 /mnt/nvme0/swapfile
sudo chown root.root /mnt/nvme0/swapfile
sudo mkswap /mnt/nvme0/swapfile

# activate the swap file
sudo swapon /mnt/nvme0/swapfile

# check the amount of swap available
grep SwapTotal /proc/meminfo

# to make permanent add to /etc/fstab if it isn’t already there
/mnt/nvme0/swapfile none swap sw 0 0
```

Now the more interesting question is why 200GB of RAM is not enough. Can you tell how much is available when you start the program?

To load the model on each process would be 2x - so 2*24GB = just 48 GB, which is temp memory and is freed once you moved models to gpu. 

Then you have 6*18=108GB just for the weights, optim states and grads and then some for activations. So you can't fully fit those into 2x 48GB gpus

And so you're using the CPU offload, and that's where you run out of memory as you're offloading both optim states and the params.

But given that you have 2x 48GB GPU - you probably don't need to offload both, params and optim stages - how about just offloading the optim states, i.e. set:

```
        "offload_param": {
            "device": "none",
            "pin_memory": false
        },
```

Also monitor `nvidia-smi` and watch your gpu memory usage - I bet at the moment it's barely being used. (I use an alias `wn=watch -n 1 nvidia-smi`)

And additionally you can play with the buffer sizes, please have a look at the discussion here:
https://huggingface.co/docs/transformers/main/main_classes/deepspeed#zero3-config

I'm talking about tweaking these param:
```
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
```

Additionally this is an expensive config for when you save the checkpoint as it has to reconstruct the model on one gpu, so you can turn it to `False`:

```
        "stage3_gather_16bit_weights_on_model_save": true
```
and then use zero_to_fp32 to extrace the fp32 weights:
https://huggingface.co/docs/transformers/main/main_classes/deepspeed#getting-the-model-weights-out

The config in the docs is sort of a generic one that fits many cases and requires a finetuning in some cases to fit a specific use case.

I know this can appear complex, so please don't hesitate to ask questions and I'm sure we will figure out how to fit your model finetuning into 200GB CPU RAM and 96GB GPU RAM.



@stas00 So I ran it again, just as it was and this time mapped the gpu and the cpu memory. And you are right (again), the gpu sits at 93% free while the checkpoint is loaded completely in to cpu ram which when the checkpoint load starts is 76% used. Is it weird that this happens only in restarting from a checkpoint and not during training? Uninterrupted the model trains without ever exceeding the cpu memory max.

Here's what I tried in response to your helpful comments:

The real answer was adding a 200gb swap file. That got it over the hump.

Changing offload_param to "none" meant that 60% of the gpu was used. But it still maxed out cpu memory. 
 
I Changed the weights_on_model to false and I looked at the documentation on the stage3_max parameters, but it was unclear what was a significant change. I increased it to 2e9 and then 3e9, but it still maxed out the cpu memory. Seemed like with 3e9 on of the gpu's increased use slightly 60%->62%. Still not of that prevented getting to maxed out cpu memory.


I'm glad you found a workaround, @randywreed 

It's odd that you see a different pattern during initial training vs. same but loaded from a checkpoint - perhaps a model is leaked somewhere - or may b e a bug in deepspeed where it allocates everything on CPU even when it shouldn't?

Let me see if I can try to analyze the memory usage during different stages. At the very least we will have a map of what we should expect.

OK, I was able to investigate this some more and found the explanation for the problem you're experiencing.

tldr; the problem is the overhead of `torch.load` for the optim states at resume which don't exist when you finetune the first time. the file is huge and thus requires a ton of additional CPU peak memory.

The full analysis:

I'm going to offload only optimizer states, that is:

```
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "none",
            "pin_memory": true
        },
        [...]
```

Let's take a smaller gpt2-large model so it's faster to run:

```
# 1. create checkpoint

deepspeed --num_gpus=1 examples/pytorch/language-modeling/run_clm.py \
--model_name_or_path gpt2-large --train_file tests/fixtures/sample_text.txt \
--do_train --fp16 --evaluation_strategy=steps --output_dir xxx \
--num_train_epochs 1 --eval_steps 1 --gradient_accumulation_steps 1 \
--per_device_train_batch_size 2 --use_fast_tokenizer False --learning_rate \
5e-06 --warmup_steps 10 --save_steps 1 --save_strategy steps --tokenizer_name \
gpt2 --max_train_samples 2 --max_eval_samples 2 --deepspeed \
tests/deepspeed/ds_config_zero3.json --skip_memory_metrics 0 \
--overwrite_output_dir



***** train metrics *****
  before_init_mem_cpu        =     3745MB
  before_init_mem_gpu        =     1786MB
  epoch                      =        1.0
  init_mem_cpu_alloc_delta   =        0MB
  init_mem_cpu_peaked_delta  =        0MB
  init_mem_gpu_alloc_delta   =        0MB
  init_mem_gpu_peaked_delta  =        0MB
  train_loss                 =      3.793
  train_mem_cpu_alloc_delta  =    17578MB
  train_mem_cpu_peaked_delta =     2952MB
  train_mem_gpu_alloc_delta  =     -148MB
  train_mem_gpu_peaked_delta =     6372MB
  train_runtime              = 0:00:13.66
  train_samples              =          1
  train_samples_per_second   =      0.073
  train_steps_per_second     =      0.073

***** eval metrics *****
  epoch                     =        1.0
  eval_accuracy             =     0.2344
  eval_loss                 =     4.0664
  eval_mem_cpu_alloc_delta  =        0MB
  eval_mem_cpu_peaked_delta =        0MB
  eval_mem_gpu_alloc_delta  =        0MB
  eval_mem_gpu_peaked_delta =     1377MB
  eval_runtime              = 0:00:01.58
  eval_samples              =          1
  eval_samples_per_second   =      0.632
  eval_steps_per_second     =      0.632
  perplexity                =    58.3469

# 2. resume from checkpoint

deepspeed --num_gpus=1 examples/pytorch/language-modeling/run_clm.py \
--model_name_or_path gpt2-large --train_file tests/fixtures/sample_text.txt \
--do_train --fp16 --evaluation_strategy=steps --output_dir xxx \
--num_train_epochs 1 --eval_steps 1 --gradient_accumulation_steps 1 \
--per_device_train_batch_size 2 --use_fast_tokenizer False --learning_rate \
5e-06 --warmup_steps 10 --save_steps 1 --save_strategy steps --tokenizer_name \
gpt2 --max_train_samples 2 --max_eval_samples 2 --deepspeed \
tests/deepspeed/ds_config_zero3.json --skip_memory_metrics 0 \
--resume_from_checkpoint xxx/checkpoint-1


***** train metrics *****
  before_init_mem_cpu        =     4618MB
  before_init_mem_gpu        =     1786MB
  epoch                      =        1.0
  init_mem_cpu_alloc_delta   =        0MB
  init_mem_cpu_peaked_delta  =        0MB
  init_mem_gpu_alloc_delta   =        0MB
  init_mem_gpu_peaked_delta  =        0MB
  train_loss                 =        0.0
  train_mem_cpu_alloc_delta  =    16071MB
  train_mem_cpu_peaked_delta =    14762MB
  train_mem_gpu_alloc_delta  =     -148MB
  train_mem_gpu_peaked_delta =        0MB
  train_runtime              = 0:00:00.00
  train_samples              =          1
  train_samples_per_second   =    135.318
  train_steps_per_second     =    135.318

***** eval metrics *****
  epoch                     =        1.0
  eval_accuracy             =     0.2344
  eval_loss                 =     4.0664
  eval_mem_cpu_alloc_delta  =        9MB
  eval_mem_cpu_peaked_delta =        0MB
  eval_mem_gpu_alloc_delta  =       -2MB
  eval_mem_gpu_peaked_delta =      146MB
  eval_runtime              = 0:00:00.41
  eval_samples              =          1
  eval_samples_per_second   =      2.399
  eval_steps_per_second     =      2.399
  perplexity                =    58.3469

```

As you can see in `train_mem_cpu_alloc_delta+train_mem_cpu_peaked_delta` numbers - the resuming one took more than 10GB extra of CPU peak memory.

I was able to reproduce your issue with GPT-J-6 on a somewhat similar setup, except using one 80GB gpu.

And I forced max 100GB CPU and 50GB RAM with:

```
systemd-run --user --scope -p MemoryHigh=100G -p MemoryMax=100G -p MemorySwapMax=50G bash
```

Now the same commands but with `EleutherAI/gpt-j-6B`:
```

deepspeed --num_gpus=1 examples/pytorch/language-modeling/run_clm.py \
--model_name_or_path EleutherAI/gpt-j-6B --train_file \
tests/fixtures/sample_text.txt --do_train --fp16 --evaluation_strategy=steps \
--output_dir xxx --num_train_epochs 1 --eval_steps 1 \
--gradient_accumulation_steps 1 --per_device_train_batch_size 2 \
--use_fast_tokenizer False --learning_rate 5e-06 --warmup_steps 10 \
--save_steps 1 --save_strategy steps --tokenizer_name gpt2 --max_train_samples \
2 --max_eval_samples 2 --deepspeed tests/deepspeed/ds_config_zero3.json \
--skip_memory_metrics 0 --overwrite_output_dir

***** train metrics *****
  before_init_mem_cpu        =     3714MB
  before_init_mem_gpu        =    12438MB
  epoch                      =        1.0
  init_mem_cpu_alloc_delta   =        0MB
  init_mem_cpu_peaked_delta  =        0MB
  init_mem_gpu_alloc_delta   =        0MB
  init_mem_gpu_peaked_delta  =        0MB
  train_loss                 =     2.6777
  train_mem_cpu_alloc_delta  =    66213MB
  train_mem_cpu_peaked_delta =    32635MB
  train_mem_gpu_alloc_delta  =       33MB
  train_mem_gpu_peaked_delta =    11879MB
  train_runtime              = 0:07:44.12
  train_samples              =          1
  train_samples_per_second   =      0.002
  train_steps_per_second     =      0.002

***** eval metrics *****
  epoch                     =        1.0
  eval_accuracy             =     0.4531
  eval_loss                 =     2.5234
  eval_mem_cpu_alloc_delta  =       25MB
  eval_mem_cpu_peaked_delta =        0MB
  eval_mem_gpu_alloc_delta  =        0MB
  eval_mem_gpu_peaked_delta =     2319MB
  eval_runtime              = 0:00:00.58
  eval_samples              =          1
  eval_samples_per_second   =      1.708
  eval_steps_per_second     =      1.708
  perplexity                =    12.4714

note the huge size of the checkpoint it needs to load into cpu:

$ ls -l xxx/checkpoint-1/global_step1/
total 68G
-rw-rw-r-- 1 stas stas 113M May 20 17:51 zero_pp_rank_0_mp_rank_00_model_states.pt
-rw-rw-r-- 1 stas stas  68G May 20 17:56 zero_pp_rank_0_mp_rank_00_optim_states.pt


# 2. resume from checkpoint

deepspeed --num_gpus=1 examples/pytorch/language-modeling/run_clm.py \
--model_name_or_path EleutherAI/gpt-j-6B --train_file \
tests/fixtures/sample_text.txt --do_train --fp16 --evaluation_strategy=steps \
--output_dir xxx --num_train_epochs 1 --eval_steps 1 \
--gradient_accumulation_steps 1 --per_device_train_batch_size 2 \
--use_fast_tokenizer False --learning_rate 5e-06 --warmup_steps 10 \
--save_steps 1 --save_strategy steps --tokenizer_name gpt2 --max_train_samples \
2 --max_eval_samples 2 --deepspeed tests/deepspeed/ds_config_zero3.json \
--skip_memory_metrics 0 --resume_from_checkpoint xxx/checkpoint-1


[....]
[INFO|deepspeed.py:449] 2022-05-20 17:58:56,691 >> Attempting to resume from xxx/checkpoint-1
[2022-05-20 18:00:08,066] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 121940
```

The problem is that `torch.save` saves one param at a time, but `torch.load` loads the whole thing to CPU memory at once. That's why the first finetuning works, but resuming doesn't.

(edit: this is an incorrect statement as I show in the next comment)

The second stage needs some additional 70GB of CPU RAM.

In the past Deepspeed devs and Deepspeed users all used huge expensive DGX servers which had TBs of CPU RAM so nobody was worried about "normal" users with little CPU RAM. Things are shifting and slowly slowly the Deepspeed team is moving towards embracing the low resource stage.

Tunji and I are working on an universal checkpoint format where each param and optim states are saved as separate files and thus can be loaded on a tiny amount of CPU memory. Currently we are working on the Megatron-Deepspeed checkpoints since we need it for manipulating the 176B checkpoint, which is much bigger than 6B of GPT-J-6. If all goes well this work will eventually end up in normal ZeRO stages as well. The current `torch.load()` to cpu is simply not an option we can continue with.

So for now please use the swap memory workaround, it shouldn't impact anything other than making the startup a bit slower.

--------------

The other issue this Issue has shined light to is not having much flexibility about how much is offloaded to CPU - same historical observation applies here. Currently one can only offload 12x or 14x params (8+4 for optim states and 2 for half precision params) and the GPU remains mainly empty, which is far from good utilization of resources. I have passed to Tunji a  request to support more flexible offloading in the future. Let's see what comes out of it.

--------------

Both issues are Deepspeed's core issues so there is not much we can do at the HF side to make thing better.

If you have any other questions and want me to explain anything please don't hesitate to ask. and if all is clear please feel free to close this issue.

Thank you for your patience, @randywreed 


Also cc: @tjruwase for awareness.



I was digging some more into this and noticed that `torch.load` is symmetrical to `torch.save` when it comes to a model on gpu <=> disc if `map_location="cuda"` is used - it doesn't copy it fully to CPU memory first - it will use CPU peak memory of the size of the largest entry in the state_dict. 

We can see that empirically through the following test:

Here is a 12GB checkpoint.

```
$ ls -l xxx/checkpoint-1/pytorch_model.bin
-rw-rw-r-- 1 stas stas 12G May 20 17:51 xxx/checkpoint-1/pytorch_model.bin
```

Let's load it to gpu:
```
$ /usr/bin/time -f %M python -c 'import torch; _=torch.load
("xxx/checkpoint-1/pytorch_model.bin", map_location="cuda")
3279196
```

It used only 3GB of CPU RAM in total. (Largest key)

and of course, let's check the baseline of loading to cpu:

```
$ /usr/bin/time -f %M python -c 'import torch; _=torch.load
("xxx/checkpoint-1/pytorch_model.bin", map_location="cpu")'
12167976
```

It used 12GB of CPU RAM in total as the size of the checkpoint.


Looking deeper it appears that the issue is on the deepspeed side. It loads the checkpoint into cpu first:

https://github.com/microsoft/DeepSpeed/blob/5208eb73da5269034ded69c4dd7c4bff81df81e7/deepspeed/runtime/engine.py#L2748

and hence the huge additional peak memory usage.

I filed an issue https://github.com/microsoft/DeepSpeed/issues/1971

Thanks for the help on this. I appreciate it.