Skip to content

Cuda OOM during training #308

@hychaochao

Description

@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!!

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions