multi-gpu training error about bk-sdm HOT 9 CLOSED

@youngwanLEE, we appreciate your kind response. After a sanity check with longer iterations, we’ve added this feature to the main branch and accordingly closed the issue.

Thank you once again for your help, and please feel free to reopen this issue if you have any questions.

from bk-sdm.

@bokyeong1015
I really appreciate your quick response and ddp code.

I successfully started to train with 8 A100 GPUs by using your ddp branch code.

For comparing with the result using single GPU, I set the effective batch size of 256 (8 per GPU x 8 GPUs) same as yours.

I'll share the re-produced result with multi-gpus as soon as the result will come out.

My training script is below:

MODEL_NAME="CompVis/stable-diffusion-v1-4"
TRAIN_DATA_DIR="./data/laion_aes/preprocessed_212k" # please adjust it if needed
UNET_CONFIG_PATH="./src/unet_config"

UNET_NAME="bk_small" # option: ["bk_base", "bk_small", "bk_tiny"]

BATCH_SIZE=32 #64
GRAD_ACCUMULATION=1 #4

NUM_GPUS=8

OUTPUT_DIR="./results/kd_"${UNET_NAME}_{$BATCH_SIZE}x${GRAD_ACCUMULATION}x${NUM_GPUS} # please adjust it if needed

StartTime=$(date +%s)

accelerate launch --multi_gpu --num_processes ${NUM_GPUS} src/kd_train_text_to_image.py
--pretrained_model_name_or_path $MODEL_NAME
--train_data_dir $TRAIN_DATA_DIR
--use_ema
--resolution 512 --center_crop --random_flip
--train_batch_size $BATCH_SIZE
--gradient_checkpointing
--mixed_precision="fp16"
--learning_rate 5e-05
--max_grad_norm 1
--lr_scheduler="constant" --lr_warmup_steps=0
--report_to="all"
--max_train_steps=400000
--seed 1234
--gradient_accumulation_steps $GRAD_ACCUMULATION
--checkpointing_steps 5000
--valid_steps 500
--lambda_sd 1.0 --lambda_kd_output 1.0 --lambda_kd_feat 1.0
--use_copy_weight_from_teacher
--unet_config_path $UNET_CONFIG_PATH --unet_config_name $UNET_NAME
--output_dir $OUTPUT_DIR

It consumed GPU memory like this:

from bk-sdm.

Thank you for sharing, we haven't tried multi-gpu experiments.
We will get back to you after testing the script you kindly provided!

from bk-sdm.

@bokyeong1015 Thanks for your quick reply.

my package versions are

pytorch==1.13.1
diffusers==0.15.0
accelerate==0.18.0
trasnformers==4.27.4
datasets==2.11.0

I tried to train with multi-gpu settings by following your command except for multi-gpu settings (8 A6000 GPUs).

Have you trained your code with multi-gpu setting?

MODEL_NAME="CompVis/stable-diffusion-v1-4"
TRAIN_DATA_DIR="./data/laion_aes/preprocessed_212k" # please adjust it if needed
UNET_CONFIG_PATH="./src/unet_config"

UNET_NAME="bk_small" # option: ["bk_base", "bk_small", "bk_tiny"]
OUTPUT_DIR="./results/kd_"$UNET_NAME # please adjust it if needed

BATCH_SIZE=8
GRAD_ACCUMULATION=1

NUM_GPUS=8

accelerate launch --multi_gpu --num_processes ${NUM_GPUS} src/kd_train_text_to_image.py
--pretrained_model_name_or_path $MODEL_NAME
--train_data_dir $TRAIN_DATA_DIR
--use_ema
--resolution 512 --center_crop --random_flip
--train_batch_size $BATCH_SIZE
--gradient_checkpointing
--mixed_precision="fp16"
--learning_rate 5e-05
--max_grad_norm 1
--lr_scheduler="constant" --lr_warmup_steps=0
--report_to="all"
--max_train_steps=400000
--seed 1234
--gradient_accumulation_steps $GRAD_ACCUMULATION
--checkpointing_steps 5000
--valid_steps 500
--lambda_sd 1.0 --lambda_kd_output 1.0 --lambda_kd_feat 1.0
--use_copy_weight_from_teacher
--unet_config_path $UNET_CONFIG_PATH --unet_config_name $UNET_NAME
--output_dir $OUTPUT_DIR

from bk-sdm.

@bokyeong1015

I want to note that when I tried to train the training script with single-GPU, it operated normally.

from bk-sdm.

@youngwanLEE thank you once again for reporting this issue and sharing your script.

We've created 10-ddp branch [deleted, Aug/19/2023], and the modified parts can be found at this link.
Could you kindly check whether src/kd_train_text_to_image.py at that branch works for your multi-gpu experiments?

===
FYI, we ran a simple experiment under the settings below and confirmed that: (i) the loss scales were similar and (ii) no problems of DDP-trained U-Net for the inference (src/generate.py)

• (1) BATCH_SIZE=4 (single gpu)
• (2) BATCH_SIZE=2 & NUM_GPUS=2 (multigpu)
• (3) BATCH_SIZE=1 & NUM_GPUS=4 (multigpu)
• Common setting:

  UNET_NAME="bk_tiny"
  GRAD_ACCUMULATION=4

from bk-sdm.

Hi, thanks for your interest :)

We tested our toy script on 3090 and our main script on A100, both using a single GPU setting and the exact package versions in requirements.txt, and no errors were reported on our end.

Would you provide more context to better understand the issue?

• the package versions related to requirements.txt
• the bash script you tried to run

from bk-sdm.

@bokyeong1015

I want to conduct zero-shot MSCOCO evaluation for my intermediate checkpoint trained with multi-GPU setting, I'm not sure how to denote my checkpoint.

Could you give me some hints for this?

In your instruction(2), you enter model_id.

Could I change the model_id to my checkpoint path?

However, I don't know which one should be denoted.

I guess the unet_ema/diffusion_pytorch_model.bin. Am I right?

Thanks in advance.

from bk-sdm.

@youngwanLEE hi, thanks for your question.
We would like to handle this as a separate issue due to a different topic and for making it easy for other people to find in the future.

Could you kindly refer to our response at this link?

from bk-sdm.