[instruct_pix2pix]: When I try to use accelerate launch train_instruct_pix2pix.py with one gpus, it report the error as below:

[elricwan]

Describe the bug

When I try to use accelerate launch train_instruct_pix2pix.py with one gpus, it report the error as below:

File "/home/xiangpeng.wan/miniconda3/envs/transformers/lib/python3.8/site-packages/accelerate/utils/dataclasses.py", line 836, in set_auto_wrap_policy
raise Exception("Could not find the transformer layer class to wrap in the model.")

File "train_instruct_pix2pix.py", line 706, in main
unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare
Exception: Could not find the transformer layer class to wrap in the model.

I did the accelerate config default

Reproduction

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export DATASET_ID="fusing/instructpix2pix-1000-samples"

accelerate launch --mixed_precision="fp16" train_instruct_pix2pix.py
--pretrained_model_name_or_path=$MODEL_NAME
--dataset_name=$DATASET_ID
--enable_xformers_memory_efficient_attention
--resolution=256 --random_flip
--train_batch_size=4 --gradient_accumulation_steps=4 --gradient_checkpointing
--max_train_steps=15000
--checkpointing_steps=5000 --checkpoints_total_limit=1
--learning_rate=5e-05 --max_grad_norm=1 --lr_warmup_steps=0
--conditioning_dropout_prob=0.05
--mixed_precision=fp16
--seed=42

Logs

No response

System Info

diffusers-0.15.0.dev0
python=3.8
torch=2.0.0
accelerate=0.18.0
ubuntu 20.04

[patrickvonplaten]

cc @sayakpaul here

[sayakpaul]

Hi!

Here's what I did:

1. Cloned diffusers with git clone https://github.com/huggingface/diffusers.
2. Then I ran the following command:

accelerate launch --mixed_precision="fp16" train_instruct_pix2pix.py \
--pretrained_model_name_or_path=runwayml/stable-diffusion-v1-5 \
--dataset_name=sayakpaul/instructpix2pix-1000-samples \
--use_ema \
--enable_xformers_memory_efficient_attention \
--resolution=512 --random_flip \
--train_batch_size=2 --gradient_accumulation_steps=4 --gradient_checkpointing \
--max_train_steps=20 \
--checkpointing_steps=10 --checkpoints_total_limit=1 \
--learning_rate=5e-05 --lr_warmup_steps=0 \
--conditioning_dropout_prob=0.05 \
--mixed_precision=fp16 \
--val_image_url="https://hf.co/datasets/diffusers/diffusers-images-docs/resolve/main/mountain.png" \
--validation_prompt="make the mountains snowy" \
--seed=42 \
--report_to=wandb

It worked on both single-GPU and multi-GPU machines. Notice that I didn't specify the --multi-gpu flag while launching training.

I installed diffusers from source by running: pip install git+https://github.com/huggingface/diffusers.

With that, the training actually went fine and I didn't face any issues.

Could you help me reproduce the error?

[elricwan]

Hi!

Thank you for the response. I find out that the problem is with the accelerate config, if you set the accelerate default_config.yaml as, you can reproduce the error. May caused by fsdp.

compute_environment: LOCAL_MACHINE
distributed_type: FSDP
downcast_bf16: 'no'
dynamo_config:
  dynamo_backend: INDUCTOR
  dynamo_mode: default
  dynamo_use_dynamic: true
  dynamo_use_fullgraph: true
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_offload_params: false
  fsdp_sharding_strategy: 2
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: ''
machine_rank: 0
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 1
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

When I use the default config, which is:

{
  "compute_environment": "LOCAL_MACHINE",
  "distributed_type": "MULTI_GPU",
  "downcast_bf16": false,
  "machine_rank": 0,
  "main_training_function": "main",
  "mixed_precision": "no",
  "num_machines": 1,
  "num_processes": 4,
  "rdzv_backend": "static",
  "same_network": false,
  "tpu_use_cluster": false,
  "tpu_use_sudo": false,
  "use_cpu": false
}

I got error:

RuntimeError: [3]: params[0] in this process with sizes [320, 4, 3, 3] appears not to match sizes of the same param in process 0.

Then I change the "num_processes" to 1, it runs to the optimization step, but raise the CUDA out of memory error. Even if I change the resolution to 16. Is there anything I can change? I am using a 3090 GPU with 24G memory.
BTW, this the the scripts I used:

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export DATASET_ID="fusing/instructpix2pix-1000-samples"

accelerate launch --mixed_precision="fp16" train_instruct_pix2pix.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--dataset_name=$DATASET_ID \
--use_ema \
--enable_xformers_memory_efficient_attention \
--resolution=16 --random_flip \
--train_batch_size=1 --gradient_accumulation_steps=4 --gradient_checkpointing \
--max_train_steps=20 \
--checkpointing_steps=10 --checkpoints_total_limit=1 \
--learning_rate=5e-05 --lr_warmup_steps=0 \
--conditioning_dropout_prob=0.05 \
--mixed_precision=fp16 \
--val_image_url="https://hf.co/datasets/diffusers/diffusers-images-docs/resolve/main/mountain.png" \
--validation_prompt="make the mountains snowy" \
--seed=42 \
--report_to=wandb

[sayakpaul]

Maybe clear all the CUDA cache and restart the process?

[sayakpaul]

You could also disable validation for dealing with lower GPU memory.

Also, maybe try with Torch 1.31.1 if possible?

[elricwan]

I clear all the cache, what accelerate config and gpu do you use?

[sayakpaul]

I am using default accelerate config with a single A100.

[elricwan]

I see, thanks

[Pierre-Wong]

I encountered the same problem. This is due to line 515-521 in the code

with torch.no_grad():
        new_conv_in = nn.Conv2d(
            in_channels, out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding
        )
        new_conv_in.weight.zero_()
        new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
        unet.conv_in = new_conv_in

where the structure of unet is changed.

In the multi-gpu case, the easiest method to deal with this issue is firstly run the script using a single GPU just a few steps to resize the Unet, then delete the following lines:

with torch.no_grad():
        new_conv_in = nn.Conv2d(
            in_channels, out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding
        )
        new_conv_in.weight.zero_()
        new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
        unet.conv_in = new_conv_in

Then set MODEL_NAME to be the output model from the above step and run the script with multi-gpus.