CUDA out of memory when training deepseek-coder-7b-base-v1.5 on 8 NVIDIA A100 GPUs
## problem description:
Hi! While training the original code using deepseek-ai/deepseek-coder-7b-base-v1.5 as base model on 8 A100 GPUs, I encountered the OOM issues below:
`Traceback (most recent call last): File "/data/intern/liujiayaobj/GEO/detikzify/DeTikZify/examples/train.py", line 57, in <module> train( File "/data/intern/liujiayaobj/GEO/detikzify/DeTikZify/detikzify/train/train.py", line 207, in train trainer.train(resume_from_checkpoint=last_checkpoint) File "/data/intern/liujiayaobj/anaconda3/envs/detikzify/lib/python3.10/site-packages/transformers/trainer.py", line 1624, in train return inner_training_loop( File "/data/intern/liujiayaobj/anaconda3/envs/detikzify/lib/python3.10/site-packages/transformers/trainer.py", line 1961, in _inner_training_loop tr_loss_step = self.training_step(model, inputs) File "/data/intern/liujiayaobj/anaconda3/envs/detikzify/lib/python3.10/site-packages/transformers/trainer.py", line 2911, in training_step self.accelerator.backward(loss) File"/data/intern/liujiayaobj/anaconda3/envs/detikzify/lib/python3.10/site-packages/accelerate/accelerator.py", line 2159, in backward loss.backward(**kwargs) File "/data/intern/liujiayaobj/anaconda3/envs/detikzify/lib/python3.10/site-packages/torch/_tensor.py", line 522, in backward torch.autograd.backward( File "/data/intern/liujiayaobj/anaconda3/envs/detikzify/lib/python3.10/site-packages/torch/autograd/init.py", line 266, in backward Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass File "/data/intern/liujiayaobj/anaconda3/envs/detikzify/lib/python3.10/site-packages/torch/autograd/function.py", line 289, in apply torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 172.00 MiB. GPU 4 has a total capacity of 79.15 GiB of which 118.38 MiB is free. Including non-PyTorch memory, this process has 79.04 GiB memory in use. Of the allocated memory 74.43 GiB is allocated by PyTorch, and 1.88 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)CUDA out of memory. Tried to allocate 64.00 MiB. GPU 7 has a total capacity of 79.15 GiB of which 30.38 MiB is free. Including non-PyTorch memory, this process has 79.12 GiB memory in use. Of the allocated memory 74.63 GiB is allocated by PyTorch, and 1.90 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)`

## question:
Could you please provide the training configuration details for training the 7B and 1.3B models? And I am wondering if the Out of Memory issue I encountered  is related to the DeepSpeed JSON config file. I noticed that this configuration file is not available in the code repository. Would it be possible for you to share it? Thanks a lot.
P.S. I can complete the training normally based on deepseek-ai/deepseek-coder-1.3b-base on single A100 GPU.

## Steps to Reproduce:
`sh examples/train.sh`

```
set -x

GPUS=${GPUS:-8}
# BATCH_SIZE=${BATCH_SIZE:-64}
# PER_DEVICE_BATCH_SIZE=${PER_DEVICE_BATCH_SIZE:-4}
# GRADIENT_ACC=$((BATCH_SIZE / PER_DEVICE_BATCH_SIZE / GPUS))


export PYTHONPATH="${PYTHONPATH}:$(pwd)"
export MASTER_PORT=34229
export TF_CPP_MIN_LOG_LEVEL=3
export LAUNCHER=pytorch

# export WORLD_SIZE=8

torchrun \
  --nnodes=1 \
  --node_rank=0 \
  --master_addr=127.0.0.1 \
  --nproc_per_node=${GPUS} \
  --master_port=${MASTER_PORT} \
  examples/train.py \
  --base_model deepseek-ai/deepseek-coder-7b-base-v1.5 \
  --datikz nllg/datikz-v2 \
  --gradient_checkpointing

```