Cuda OOM during training

[hychaochao]

When I use 4 A100 40G to train the alpaca, I encountered an oom error during training. this is my training arguments:

#!/bin/bash
module load  compilers/cuda/11.8 compilers/gcc/9.3.0 cudnn/8.4.0.27_cuda11.x anaconda
source activate torch_new
export PYTHONUNBUFFERED=1 
torchrun --nproc_per_node=4 --master_port=11223 train.py \
    --model_name_or_path llama2 \
    --data_path alpaca_data.json \
    --output_dir alpaca_llama2 \
    --num_train_epochs 3 \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 1 \
    --gradient_accumulation_steps 1 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 1000 \
    --save_total_limit 1 \
    --learning_rate 2e-5 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --bf16 true\
    --fsdp "full_shard auto_wrap" \
    --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer'

and this is the error:

22%|██▏       | 8402/39003 [10:12:57<36:10:57,  4.26s/it]Traceback (most recent call last):
  File "/home/bingxing2/home/scx6203/luckychao/stanford_alpaca/train.py", line 220, in <module>
    train()
'''
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 208.00 MiB. GPU 1 has a total capacty of 39.41 GiB of which 183.50 MiB is free. Including non-PyTorch memory, this process has 39.23 GiB memory in use. Of the allocated memory 37.37 GiB is allocated by PyTorch, and 539.20 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

anyone know the reason? :( thanks again!!