Confusion about model memory usage (DeepSpeed+PEFT)
Author: IYIAscension
Pull request: no
Labels: name: DeepSpeed

name: PEFT

### System Info

- `transformers` version: 4.40.1
- Platform: Linux-4.18.0-372.80.1.el8_6.x86_64-x86_64-with-glibc2.28
- Python version: 3.11.3
- Huggingface_hub version: 0.22.2
- Safetensors version: 0.4.3
- Accelerate version: 0.29.3
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2 (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

DS report:
```
[2024-06-28 15:30:42,418] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
^[[93m [WARNING] ^[[0m async_io requires the dev libaio .so object and headers but these were not found.
^[[93m [WARNING] ^[[0m async_io: please install the libaio-devel package with yum
^[[93m [WARNING] ^[[0m If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
^[[93m [WARNING] ^[[0m Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
^[[93m [WARNING] ^[[0m sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
^[[93m [WARNING] ^[[0m using untested triton version (2.2.0), only 1.0.0 is known to be compatible
Could not import _export_to_torch_ir. Please make sure your PyTorch version is newer than 2.2.0.
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. ^[[92m[OKAY]^[[0m
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
^[[93m [WARNING] ^[[0m async_io requires the dev libaio .so object and headers but these were not found.
^[[93m [WARNING] ^[[0m async_io: please install the libaio-devel package with yum
^[[93m [WARNING] ^[[0m If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... ^[[93m[NO]^[[0m ....... ^[[93m[NO]^[[0m
fused_adam ............. ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
cpu_adam ............... ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
cpu_adagrad ............ ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
cpu_lion ............... ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
^[[93m [WARNING] ^[[0m Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... ^[[93m[NO]^[[0m ....... ^[[93m[NO]^[[0m
fp_quantizer ........... ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
fused_lamb ............. ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
fused_lion ............. ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
inference_core_ops ..... ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
cutlass_ops ............ ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
transformer_inference .. ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
quantizer .............. ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
ragged_device_ops ...... ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
ragged_ops ............. ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
random_ltd ............. ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
^[[93m [WARNING] ^[[0m sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
^[[93m [WARNING] ^[[0m using untested triton version (2.2.0), only 1.0.0 is known to be compatible
sparse_attn ............ ^[[93m[NO]^[[0m ....... ^[[93m[NO]^[[0m
spatial_inference ...... ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
transformer ............ ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
stochastic_transformer . ^[[93m[NO]^[[0m ....... ^[[92m[OKAY]^[[0m
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/sw/arch/RHEL8/EB_production/2023/software/PyTorch/2.1.2-foss-2023a-CUDA-12.1.1/lib/python3.11/site-packages/torch']
torch version .................... 2.1.2
deepspeed install path ........... ['/home/aabdalla/.local/lib/python3.11/site-packages/deepspeed']
deepspeed info ................... 0.14.2, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 12.1
deepspeed wheel compiled w. ...... torch 2.1, cuda 12.1
shared memory (/dev/shm) size .... 251.76 GB
```

### Who can help?

@muellerz @SunMarc @ArthurZucker 

### 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

Scripts can be found here: https://github.com/IYIAscension/MasterThesis

This issue is about model_finetuner.py.
The script I use to run it in a certain SLURM environment is original_finetuner.sh.
In original_finetuner.sh I use:
srun --jobid $SLURM_JOBID bash -c 'python -m torch.distributed.run --nnodes=8 --nproc_per_node=4 --master_addr=$HOSTNAME --master_port=$MASTER_PORT --node_rank=$SLURM_PROCID $HOME/original/model_finetuner.py --train_d $TMP_DIR/train_examples --dev_d $TMP_DIR/dev_examples --cp_d $TMP_DIR/checkpoints_0'

Sadly, I can't provide the data I use for finetuning as it is not mine, nor open source, but any CSV with "input, output" headers should work for reproduction purposes.

### Expected behavior

I expected that the google/mt5-xxl model with an input size of 2048, using ZeRO-3(bf16, no offload) and PEFT(LoRA r=4), would fit in 8 nodes of 4 GPUs (40 GB each, so 1280 GB total VRAM). The model is a 13B parameter model, but with an an input size of 2048, the memory requirements are much larger. (16x for some parts of the model, reddit.com/r/learnmachinelearning/comments/1327oz8). However, with PEFT the model should fit even if the 16x applied to the entire model.

Deepspeed and PEFT both seem to work, but DeepSpeed goes OOM faster than expected and PEFT seems to reduce memory usage less than expected. I understand that DeepSpeed can't perfectly divide the memory requirements across GPUs, but even so. 

Perhaps there is something wrong with my DeepSpeed/LoRA configs (found in model_finetuner.py) or installation? Or perhaps the model is sharded inefficiently because it does not expect the increase to a 2048 input size?

### Experiment details
Inference can be done with just 80gb of VRAM without PEFT. I am assuming that means the requirement of fitting the largest layer on 1 GPU is met during training as well.

A 1024 input size, which should only be a 4x multiplier, also didn't fit in the 1280 GB VRAM (currently trying with PEFT).

The google/flan-t5-xxl did fit with 2048 input size and so did an mt5-xxl with 512 input size (both without PEFT), so it seems Deepspeed does work for multi-node parallelization. PEFT also seems to work as a mt5-xl did not fit on 2 GPUs without PEFT, but it did with PEFT. However, I needed to test many different models/input sizes to find one where PEFT made the difference.

Comments:
Hey! I am not seiing the bf16 in the srun job, but it is important to make sure that this is properly used. Appart from that, It's a bit hard for use to debug in your stead, I would make sure that you are not using cache, that the input size is what you expect. Indeed PEFT should induce a very small overhead when finetuning. 

Would also recommend you to use a more recent model than mt5 ! 

Hey, thanks for your response. The BF16 is in the deepspeed config and TrainingArguments inside of the python file, which should do the job, I believe. use_cache was set to True, so that might be the issue. Sadly, I can not test it currently, so I will close this issue.