oom when using adafactor optimizer in deepspeed

[zhangvia]

System Info

- `transformers` version: 4.44.2
- Platform: Linux-5.15.0-105-generic-x86_64-with-glibc2.31
- Python version: 3.10.0
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.2
- Accelerate version: 0.33.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.0+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A800 80GB PCIe

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

i'm running train_xl.sh in this repo. and i change the 8bit adam optimizer to adafactor optimizer using transformers.optimization.Adafactor. i'm using two 40GB a100, deepspeed stage 2, batchsize=1,VTON-HD dataset.

the adafactor optimizer should use less gpu memory, because of less optimizer states than 8bit adam, but it get oom in this line

and oom happens after 10 steps, i don't know what happen in 10th step, i call the accelerate.backward() and optimizer.step() every step.

and in 10th step, the memory usage increased from 29GB to 39GB when using 8bit adam optimizer, and get oom when using adafactor optimizer

Expected behavior

could anybody explain this phenomenon

[LysandreJik]

cc @SunMarc and @muellerzr