run train_dreambooth_lora.py failed with accelerate

[webliupeng]

Describe the bug

Thanks for this awesome project!
When I run the script "train_dreambooth_lora.py" without acceleration, it works fine. But when I use acceleration launch, it fails when the number of steps reaches "checkpointing_steps".
I am running the script in a Docker with 4 * 3090 vGPUs. And I ran accelerate test, it's successed.
I am new to this and would appreciate any guidance or suggestions you can offer.

Reproduction

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export INSTANCE_DIR="/diffusers/examples/dreambooth/dunhuang512"
export OUTPUT_DIR="path-to-save-model"
cd /diffusers/examples/dreambooth/
accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --logging_dir='./logs' \
  --instance_prompt="dhstyle_test" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --checkpointing_steps=100 \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="dhstyle_test" \
  --validation_epochs=50 \
  --seed="0"\
  --enable_xformers_memory_efficient_attention \
  --use_8bit_adam

Logs

  File "/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 1093, in <module>
    main(args)
  File "/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 972, in main
    LoraLoaderMixin.save_lora_weights(
  File "/diffusers/src/diffusers/loaders.py", line 1111, in save_lora_weights
    for module_name, param in unet_lora_layers.state_dict().items()
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1818, in state_dict
    module.state_dict(destination=destination, prefix=prefix + name + '.', keep_vars=keep_vars)
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1820, in state_dict
    hook_result = hook(self, destination, prefix, local_metadata)
  File "/diffusers/src/diffusers/loaders.py", line 74, in map_to
    num = int(key.split(".")[1])  # 0 is always "layers"
ValueError: invalid literal for int() with base 10: 'layers'
Steps:  20%|████████████████████▊                                                                                   | 100/500 [03:35<14:20,  2.15s/it, loss=0.217, lr=0.0001]
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 63642 closing signal SIGTERM
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 63643 closing signal SIGTERM
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 63644 closing signal SIGTERM
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 63641) of binary: /usr/local/bin/python
Traceback (most recent call last):
  File "/usr/local/bin/accelerate", line 8, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py", line 45, in main
    args.func(args)
  File "/usr/local/lib/python3.10/site-packages/accelerate/commands/launch.py", line 914, in launch_command
    multi_gpu_launcher(args)
  File "/usr/local/lib/python3.10/site-packages/accelerate/commands/launch.py", line 603, in multi_gpu_launcher
    distrib_run.run(args)
  File "/usr/local/lib/python3.10/site-packages/torch/distributed/run.py", line 785, in run
    elastic_launch(
  File "/usr/local/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 134, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/usr/local/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 250, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
============================================================
train_dreambooth_lora.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2023-04-29_00:59:00
  host      : sd-5b564dfd58-7v76h
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 63641)
  error_file: <N/A>

System Info

• diffusers version: 0.17.0.dev0

• Platform: Linux-5.4.0-146-generic-x86_64-with-glibc2.31

• Python version: 3.10.9

• PyTorch version (GPU?): 2.0.0+cu117 (True)

• Huggingface_hub version: 0.14.0

• Transformers version: 4.25.1

• Accelerate version: 0.18.0

• xFormers version: 0.0.19

• Using GPU in script?:

• Using distributed or parallel set-up in script?:

• Accelerate default config:
  - compute_environment: LOCAL_MACHINE
  - distributed_type: MULTI_GPU
  - mixed_precision: no
  - use_cpu: False
  - num_processes: 4
  - machine_rank: 0
  - num_machines: 1
  - gpu_ids: all
  - rdzv_backend: static
  - same_network: True
  - main_training_function: main
  - downcast_bf16: no
  - tpu_use_cluster: False
  - tpu_use_sudo: False
  - tpu_env: []

[SeunghyunSEO]

i got same error

and i also got problem when i run train_custom_diffuson.py

it also happens with accelerate and when i run with python train_custom_diffuson.py, error does not occur

│ /workspace/diffusers/src/diffusers/models/attention_processor.py:546 in __call__                 │                                
│                                                                                                  │                                
│    543 │   │   key = attn.head_to_batch_dim(key)                                                 │                                
│    544 │   │   value = attn.head_to_batch_dim(value)                                             │                                
│    545 │   │                                                                                     │                                
│ ❱  546 │   │   attention_probs = attn.get_attention_scores(query, key, attention_mask)           │                                
│    547 │   │   hidden_states = torch.bmm(attention_probs, value)                                 │                                
│    548 │   │   hidden_states = attn.batch_to_head_dim(hidden_states)                             │                                
│    549                                                                                           │                                
│                                                                                                  │                               
│ /workspace/diffusers/src/diffusers/models/attention_processor.py:308 in get_attention_scores     │                               
│                                                                                                  │                               
│    305 │   │   │   baddbmm_input = attention_mask                                                │                               
│    306 │   │   │   beta = 1                                                                      │                               
│    307 │   │                                                                                     │                               
│ ❱  308 │   │   attention_scores = torch.baddbmm(                                                 │                               
│    309 │   │   │   baddbmm_input,                                                                │                               
│    310 │   │   │   query,                                                                        │                               
│    311 │   │   │   key.transpose(-1, -2),                                                        │                               
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯                               
RuntimeError: expected scalar type Half but found Float
Steps:   0%|    

[SeunghyunSEO]

i'm not sure why error occurs with accelerate, but it can be fixed by modfiy following lines of here

class AttnProcsLayers(torch.nn.Module):
    def __init__(self, state_dict: Dict[str, torch.Tensor]):
        super().__init__()
        self.layers = torch.nn.ModuleList(state_dict.values())
        self.mapping = dict(enumerate(state_dict.keys()))
        self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}

        # we add a hook to state_dict() and load_state_dict() so that the
        # naming fits with `unet.attn_processors`
        def map_to(module, state_dict, *args, **kwargs):
            new_state_dict = {}
            for key, value in state_dict.items():
                layer_index = 2 if 'module' in key else 1 ## you should add this line
                num = int(key.split(".")[layer_index])  # 0 is always "layers"
                new_key = key.replace(f"layers.{num}", module.mapping[num])
                new_state_dict[new_key] = value

            return new_state_dict

this is because key of state_dict is like module.layers.0.to_q_lora.down.weight not layers.0.to_q_lora.down.weight. so layers, the second elements of key could not be int.
i guess latest version did not check this with accelerate

this function is called by here when we run the code with accelerate

                    if accelerator.is_main_process:
                        save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                        # We combine the text encoder and UNet LoRA parameters with a simple
                        # custom logic. `accelerator.save_state()` won't know that. So,
                        # use `LoraLoaderMixin.save_lora_weights()`.
                        LoraLoaderMixin.save_lora_weights(
                            save_directory=save_path,
                            unet_lora_layers=unet_lora_layers,
                            text_encoder_lora_layers=text_encoder_lora_layers,
                        )

---

should be

        def map_to(module, state_dict, *args, **kwargs):
            new_state_dict = {}
            for key, value in state_dict.items():
                # num = int(key.split(".")[layer_index])  # 0 is always "layers"
                # new_key = key.replace(f"layers.{num}", module.mapping[num])
                if 'module' in key:
                    num = int(key.split(".")[2]) 
                    replace_key = f"module.layers.{num}"
                else: 
                    num = int(key.split(".")[1]) 
                    replace_key = f"layers.{num}"
                new_key = key.replace(replace_key, module.mapping[num])
                new_state_dict[new_key] = value

so you can load pytorch_lora_weights.bin correctly

[webliupeng]

i'm not sure why error occurs with accelerate, but it can be fixed by modfiy following lines of here

class AttnProcsLayers(torch.nn.Module):
    def __init__(self, state_dict: Dict[str, torch.Tensor]):
        super().__init__()
        self.layers = torch.nn.ModuleList(state_dict.values())
        self.mapping = dict(enumerate(state_dict.keys()))
        self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}

        # we add a hook to state_dict() and load_state_dict() so that the
        # naming fits with `unet.attn_processors`
        def map_to(module, state_dict, *args, **kwargs):
            new_state_dict = {}
            for key, value in state_dict.items():
                layer_index = 2 if 'module' in key else 1 ## you should add this line
                num = int(key.split(".")[layer_index])  # 0 is always "layers"
                new_key = key.replace(f"layers.{num}", module.mapping[num])
                new_state_dict[new_key] = value

            return new_state_dict

this is because key of state_dict is like module.layers.0.to_q_lora.down.weight not layers.0.to_q_lora.down.weight. so layers, the second elements of key could not be int. i guess latest version did not check this with accelerate

this function is called by here when we run the code with accelerate

                    if accelerator.is_main_process:
                        save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                        # We combine the text encoder and UNet LoRA parameters with a simple
                        # custom logic. `accelerator.save_state()` won't know that. So,
                        # use `LoraLoaderMixin.save_lora_weights()`.
                        LoraLoaderMixin.save_lora_weights(
                            save_directory=save_path,
                            unet_lora_layers=unet_lora_layers,
                            text_encoder_lora_layers=text_encoder_lora_layers,
                        )

should be

        def map_to(module, state_dict, *args, **kwargs):
            new_state_dict = {}
            for key, value in state_dict.items():
                # num = int(key.split(".")[layer_index])  # 0 is always "layers"
                # new_key = key.replace(f"layers.{num}", module.mapping[num])
                if 'module' in key:
                    num = int(key.split(".")[2]) 
                    replace_key = f"module.layers.{num}"
                else: 
                    num = int(key.split(".")[1]) 
                    replace_key = f"layers.{num}"
                new_key = key.replace(replace_key, module.mapping[num])
                new_state_dict[new_key] = value

so you can load pytorch_lora_weights.bin correctly
Nice , I have changed this code. and then this error is gone.
But it still can't train successfully. I@SeunghyunSEO .

these are errors:

Traceback (most recent call last):
  File "/data/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 1093, in <module>
    main(args)
  File "/data/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 1048, in main
    pipeline.load_lora_weights(args.output_dir)
  File "/data/diffusers/src/diffusers/loaders.py", line 851, in load_lora_weights
    self.unet.load_attn_procs(unet_lora_state_dict)
  File "/data/diffusers/src/diffusers/loaders.py", line 305, in load_attn_procs
    self.set_attn_processor(attn_processors)
  File "/data/diffusers/src/diffusers/models/unet_2d_condition.py", line 533, in set_attn_processor
    fn_recursive_attn_processor(name, module, processor)
  File "/data/diffusers/src/diffusers/models/unet_2d_condition.py", line 530, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  File "/data/diffusers/src/diffusers/models/unet_2d_condition.py", line 530, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  File "/data/diffusers/src/diffusers/models/unet_2d_condition.py", line 530, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  [Previous line repeated 3 more times]
  File "/data/diffusers/src/diffusers/models/unet_2d_condition.py", line 527, in fn_recursive_attn_processor
    module.set_processor(processor.pop(f"{name}.processor"))
KeyError: 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.processor'
Steps: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 500/500 

[SeunghyunSEO]

i'm not sure why error occurs with accelerate, but it can be fixed by modfiy following lines of here

class AttnProcsLayers(torch.nn.Module):

def __init__(self, state_dict: Dict[str, torch.Tensor]):

    super().__init__()

    self.layers = torch.nn.ModuleList(state_dict.values())

    self.mapping = dict(enumerate(state_dict.keys()))

    self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}

    # we add a hook to state_dict() and load_state_dict() so that the

    # naming fits with `unet.attn_processors`

    def map_to(module, state_dict, *args, **kwargs):

        new_state_dict = {}

        for key, value in state_dict.items():

            layer_index = 2 if 'module' in key else 1 ## you should add this line

            num = int(key.split(".")[layer_index])  # 0 is always "layers"

            new_key = key.replace(f"layers.{num}", module.mapping[num])

            new_state_dict[new_key] = value

        return new_state_dict

this is because key of state_dict is like module.layers.0.to_q_lora.down.weight not layers.0.to_q_lora.down.weight. so layers, the second elements of key could not be int. i guess latest version did not check this with accelerate

this function is called by here when we run the code with accelerate

                if accelerator.is_main_process:

                    save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")

                    # We combine the text encoder and UNet LoRA parameters with a simple

                    # custom logic. `accelerator.save_state()` won't know that. So,

                    # use `LoraLoaderMixin.save_lora_weights()`.

                    LoraLoaderMixin.save_lora_weights(

                        save_directory=save_path,

                        unet_lora_layers=unet_lora_layers,

                        text_encoder_lora_layers=text_encoder_lora_layers,

                    )

should be

    def map_to(module, state_dict, *args, **kwargs):

        new_state_dict = {}

        for key, value in state_dict.items():

            # num = int(key.split(".")[layer_index])  # 0 is always "layers"

            # new_key = key.replace(f"layers.{num}", module.mapping[num])

            if 'module' in key:

                num = int(key.split(".")[2]) 

                replace_key = f"module.layers.{num}"

            else: 

                num = int(key.split(".")[1]) 

                replace_key = f"layers.{num}"

            new_key = key.replace(replace_key, module.mapping[num])

            new_state_dict[new_key] = value

so you can load pytorch_lora_weights.bin correctly

Nice , I have changed this code. and then this error is gone.

But it still can't train successfully. I@SeunghyunSEO .

these are errors:

Traceback (most recent call last):

  File "/data/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 1093, in <module>

    main(args)

  File "/data/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 1048, in main

    pipeline.load_lora_weights(args.output_dir)

  File "/data/diffusers/src/diffusers/loaders.py", line 851, in load_lora_weights

    self.unet.load_attn_procs(unet_lora_state_dict)

  File "/data/diffusers/src/diffusers/loaders.py", line 305, in load_attn_procs

    self.set_attn_processor(attn_processors)

  File "/data/diffusers/src/diffusers/models/unet_2d_condition.py", line 533, in set_attn_processor

    fn_recursive_attn_processor(name, module, processor)

  File "/data/diffusers/src/diffusers/models/unet_2d_condition.py", line 530, in fn_recursive_attn_processor

    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

  File "/data/diffusers/src/diffusers/models/unet_2d_condition.py", line 530, in fn_recursive_attn_processor

    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

  File "/data/diffusers/src/diffusers/models/unet_2d_condition.py", line 530, in fn_recursive_attn_processor

    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

  [Previous line repeated 3 more times]

  File "/data/diffusers/src/diffusers/models/unet_2d_condition.py", line 527, in fn_recursive_attn_processor

    module.set_processor(processor.pop(f"{name}.processor"))

KeyError: 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.processor'

Steps: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 500/500 

i told you that you should also edit replace_key. check my last code block

[webliupeng]

@SeunghyunSEO I mean another error occurred.

[SeunghyunSEO]

@SeunghyunSEO I mean another error occurred.

that's what i encountered error without this code

        def map_to(module, state_dict, *args, **kwargs):
            new_state_dict = {}
            for key, value in state_dict.items():
                # num = int(key.split(".")[layer_index])  # 0 is always "layers"
                # new_key = key.replace(f"layers.{num}", module.mapping[num])
                if 'module' in key:
                    num = int(key.split(".")[2]) 
                    replace_key = f"module.layers.{num}"
                else: 
                    num = int(key.split(".")[1]) 
                    replace_key = f"layers.{num}"
                new_key = key.replace(replace_key, module.mapping[num])

please check dictionary keys of lora weight with torch.load, maybe you will notice the keys are like module.layers.0. ~~~, that's what i mentioned. with my code, i got no problem to save or load weights.

[sayakpaul]

We need to be able to reproduce this error.

I conducted some training runs with the following setups. First, I installed diffusers from source: pip -U install git+https://github.com/huggingface/diffusers.

1. Without training the text encoder LoRA:

accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --checkpointing_steps=100 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=50 \
  --seed="0" \
  --push_to_hub

Final checkpoints: https://huggingface.co/sayakpaul/dog-test-lora.

2. With text encoder LoRA along with the UNet LoRA:

accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --checkpointing_steps=100 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=50 \
  --seed="0" \
  --train_text_encoder \
  --push_to_hub

Final checkpoints: https://huggingface.co/sayakpaul/dreambooth-text-encoder-test.

I was able to conduct the above training runs without any failures.

What am I missing out on?

[webliupeng]

We need to be able to reproduce this error.

I conducted some training runs with the following setups. First, I installed diffusers from source: pip -U install git+https://github.com/huggingface/diffusers.

1. Without training the text encoder LoRA:

accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --checkpointing_steps=100 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=50 \
  --seed="0" \
  --push_to_hub

Final checkpoints: https://huggingface.co/sayakpaul/dog-test-lora.

2. With text encoder LoRA along with the UNet LoRA:

accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --checkpointing_steps=100 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=50 \
  --seed="0" \
  --train_text_encoder \
  --push_to_hub

Final checkpoints: https://huggingface.co/sayakpaul/dreambooth-text-encoder-test.

I was able to conduct the above training runs without any failures.

What am I missing out on?

@sayakpaul
do you use MULTI GPU to training?
I followed the examples and tried to run train_dreambooth_lora.py in two different environments, both of which had 4 GPUs. One is docker container, Another is virual host , However, I encountered the same error in both environments.

[SeunghyunSEO]

We need to be able to reproduce this error.

I conducted some training runs with the following setups. First, I installed diffusers from source: pip -U install git+https://github.com/huggingface/diffusers.

1. Without training the text encoder LoRA:

accelerate launch train_dreambooth_lora.py \

  --pretrained_model_name_or_path=$MODEL_NAME  \

  --instance_data_dir=$INSTANCE_DIR \

  --output_dir=$OUTPUT_DIR \

  --instance_prompt="a photo of sks dog" \

  --resolution=512 \

  --train_batch_size=1 \

  --gradient_accumulation_steps=1 \

  --checkpointing_steps=100 \

  --learning_rate=1e-4 \

  --report_to="wandb" \

  --lr_scheduler="constant" \

  --lr_warmup_steps=0 \

  --max_train_steps=500 \

  --validation_prompt="A photo of sks dog in a bucket" \

  --validation_epochs=50 \

  --seed="0" \

  --push_to_hub

Final checkpoints: https://huggingface.co/sayakpaul/dog-test-lora.

2. With text encoder LoRA along with the UNet LoRA:

accelerate launch train_dreambooth_lora.py \

  --pretrained_model_name_or_path=$MODEL_NAME  \

  --instance_data_dir=$INSTANCE_DIR \

  --output_dir=$OUTPUT_DIR \

  --instance_prompt="a photo of sks dog" \

  --resolution=512 \

  --train_batch_size=1 \

  --gradient_accumulation_steps=1 \

  --checkpointing_steps=100 \

  --learning_rate=1e-4 \

  --report_to="wandb" \

  --lr_scheduler="constant" \

  --lr_warmup_steps=0 \

  --max_train_steps=500 \

  --validation_prompt="A photo of sks dog in a bucket" \

  --validation_epochs=50 \

  --seed="0" \

  --train_text_encoder \

  --push_to_hub

Final checkpoints: https://huggingface.co/sayakpaul/dreambooth-text-encoder-test.

I was able to conduct the above training runs without any failures.

What am I missing out on?

i installed main branch of diffusers from source

[sayakpaul]

For train_dreambooth_lora.py single GPU be fine. Since we're only training a handful of parameters, single GPU should be fine.

[capntrips]

I'm running into the same thing with a single GPU using DeepSpeed. The above fix worked for me in that it prevented the initial error, but I get the same error as @webliupeng in load_lora_weights.

If I put a print statement in map_to after applying the fix, I see the keys are updated on every call to save_lora_weights. but the incorrect keys are still written to pytorch_lora_weights.bin.

[sayakpaul]

Could you try updating your local clone of diffusers (git checkout main && git pull) and then retry?

As mentioned in #3284 (comment), I need to be able to reproduce the error minimally.

i installed main branch of diffusers from source

@SeunghyunSEO, could you also update your local clone of diffusers (git checkout main && git pull)?