multi-node training out of memory during ViT training

[frankxyy]

I use 2node, 4gpus per node. The same training batch size work for single node while when applying to multi-node, this error is printed out:

[zhisbug]

@frankxyy are you saying the exactly same training setting can work on 1x4 GPUs but not on 2x4 GPUs?

[frankxyy]

@frankxyy are you saying the exactly same training setting can work on 1x4 GPUs but not on 2x4 GPUs?

@zhisbug yes

The same command of

NCCL_SOCKET_IFNAME=bond0 CUDA_VISIBLE_DEVICES=4,5,6,7  python -m pdb run_image_classification.py \
    --output_dir ./vit-base-patch16-imagenette \
    --train_dir="/home/xuyangyang/imagenette2/train" \
    --validation_dir="/home/xuyangyang/imagenette2/val" \
    --num_train_epochs 50 \
    --num_micro_batches 2 \
    --learning_rate 1e-3 \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 2 \
    --overwrite_output_dir \
    --preprocessing_num_workers 2 \
    --num_of_nodes 1