Training text_to_image on multi-GPU with Megatron-LM or DeepSpeed does not work

[zipliucc]

I have an EC2 instance with 4 GPUs and each GPU has 16GB memory.

By default, hugging face uses Data Parallelism and tries to fit one copy of the model in each GPU. Of course, 16GB is not enough and I ran into CUDA Out of Memory error.

I tried to use Tensor Parallelism and Pipeline Parallelism through Megatron-LM . For example, with Pipeline Parallel Degree of 4 using Megatron-LM, it is supposed to split the model over 4 GPUs and totally 64 GB GPU memory. However, I still ran into CUDA Out of Memory error.

I also tried to use ZeRO Stage 2 or 3 through DeepSpeed. However, I still ran into CUDA Out of Memory error.

In both cases, I used hugging face accelerate to configure Megatron-LM or DeepSpeed.

I wonder if the diffuser repo works out of box with Megatron-LM and DeepSpeed. Or, do I need to make code change to make it work with Tensor/Pipeline parallelism?

Thank you very much,

[zipliucc]

Here is one example.

Accelerate config

compute_environment: LOCAL_MACHINE
distributed_type: MEGATRON_LM
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
megatron_lm_config:
megatron_lm_gradient_clipping: 1.0
megatron_lm_num_micro_batches: 1
megatron_lm_pp_degree: 4
megatron_lm_recompute_activations: true
megatron_lm_tp_degree: 1
megatron_lm_use_distributed_optimizer: true
mixed_precision: fp16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

Accelerate command & arguments.

`export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export TRAIN_DIR="/home/ubuntu/code/labweek23-itg-support/Nexus/"
mkdir "mymodel"
export OUTPUT_DIR="mymodel"

accelerate launch --mixed_precision="fp16" --multi_gpu train_text_to_image.py
--pretrained_model_name_or_path=$MODEL_NAME
--train_data_dir=$TRAIN_DIR
--use_ema
--resolution=512 --center_crop --random_flip
--train_batch_size=1
--gradient_checkpointing
--gradient_accumulation_steps=4
--mixed_precision="fp16"
--max_train_steps=5000
--learning_rate=1e-05
--max_grad_norm=1
--lr_scheduler="constant" --lr_warmup_steps=0
--output_dir=${OUTPUT_DIR} `

With this configuration, I was trying to run Pipeline Parallelization, one model over 4 GPUs (64 GB GPU memory in total). However, I got torch.cuda.OutOfMemoryErrorCUDA out of memory

[patrickvonplaten]

Hey @zipliucc,

We don't yet have explored good support for Megatron-LM or DeepSpeed. We recommend to instead use torch.compile, see PR here: #3403

[zipliucc]

Hi @patrickvonplaten ,

Thank you so much for the suggestions. Really appreciate your insights and feedback .

Regards,

[sayakpaul]

Closing this issue then. In case problems persist, feel free to re-open :)

[chaojiewang94]

@zipliucc have you find the solution to split the model over 4 GPUs?

[zipliucc]

Hi @chaojiewang94 , sorry for missing your message.

No I did not figure out how to make the Tensor Parallelism or Pipeline Parallelism work over multiple GPU.

I managed to make the fine-tune work on single GPU with "--use_8bit_adam"