saving and loading checkpoints do not work on train_dreambooth.py when using text_encoder

[amitgurintelcom]

Describe the bug

There are 2 issues when trying to start from a checkpoint, when using --train_text_encoder.
First issue is #2480 . I wrote there how to fix it in train_dreambooth.py. It fixes how the checkpoint is saved.
Second issue is similar. It is in load_model_hook (how the text_encoder checkpoint is loaded).
type(model) is not equal to type(text_encoder) in my env.

In current code, it is written:
def load_model_hook(models, input_dir):
while len(models) > 0:
# pop models so that they are not loaded again
model = models.pop()

            if type(model) == type(text_encoder):
                # load transformers style into model
                load_model = text_encoder_cls.from_pretrained(input_dir, subfolder="text_encoder")
                model.config = load_model.config
            else:
                # load diffusers style into model
                load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
                model.register_to_config(**load_model.config)

            model.load_state_dict(load_model.state_dict())
            del load_model

The fix should be:
def load_model_hook(models, input_dir):
while len(models) > 0:
# pop models so that they are not loaded again
model = models.pop()

            if "CLIPTextModel" in str(type(model)):
                # load transformers style into model
                load_model = text_encoder_cls.from_pretrained(input_dir, subfolder="text_encoder")
                model.config = load_model.config
            else:
                # load diffusers style into model
                load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
                model.register_to_config(**load_model.config)

            model.load_state_dict(load_model.state_dict())
            del load_model

Reproduction

Command line:
accelerate launch train_dreambooth.py
--mixed_precision="fp16"
--revision="fp16"
--instance_data_dir="./XYZ"
--output_dir="checkpoints"
--instance_prompt="XYZ"
--resolution=512
--train_batch_size=2
--gradient_accumulation_steps=1
--learning_rate=1.5e-6
--lr_scheduler="constant"
--lr_warmup_steps=0
--max_train_steps=1200
--num_class_images=260
--with_prior_preservation
--prior_loss_weight=1.0
--class_prompt="man"
--class_data_dir="./man"
--train_text_encoder
--validation_prompt="Portrait of XYZ"
--validation_steps=100
--checkpointing_steps=200
--pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4"
--resume_from_checkpoint="latest"

Logs

AttributeError: 'CLIPTextModel' object has no attribute 'register_to_config'
Traceback (most recent call last):
  File "/home/agur/.vscode-server/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/pydevd.py", line 3489, in <module>
    main()
  File "/home/agur/.vscode-server/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/pydevd.py", line 3482, in main
    globals = debugger.run(setup['file'], None, None, is_module)
  File "/home/agur/.vscode-server/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/pydevd.py", line 2510, in run
    return self._exec(is_module, entry_point_fn, module_name, file, globals, locals)
  File "/home/agur/.vscode-server/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/pydevd.py", line 2517, in _exec
    globals = pydevd_runpy.run_path(file, globals, '__main__')
  File "/home/agur/.vscode-server/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 321, in run_path
    return _run_module_code(code, init_globals, run_name,
  File "/home/agur/.vscode-server/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 135, in _run_module_code
    _run_code(code, mod_globals, init_globals,
  File "/home/agur/.vscode-server/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 124, in _run_code
    exec(code, run_globals)
  File "train_dreambooth.py", line 1085, in <module>
    main(args)
  File "train_dreambooth.py", line 944, in main
    accelerator.load_state(os.path.join(args.output_dir, path))
  File "/home/agur/.local/lib/python3.8/site-packages/accelerate/accelerator.py", line 2394, in load_state
    hook(models, input_dir)
  File "train_dreambooth.py", line 759, in load_model_hook
    model.register_to_config(**load_model.config)
  File "/home/agur/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1614, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'CLIPTextModel' object has no attribute 'register_to_config'

System Info

diffusers-cli env:

• diffusers version: 0.17.0.dev0
• Platform: Linux-5.15.0-71-generic-x86_64-with-glibc2.29
• Python version: 3.8.10
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Huggingface_hub version: 0.13.3
• Transformers version: 4.28.1
• Accelerate version: 0.18.0
• xFormers version: not installed
• Using GPU in script?: yes. nvidia-driver-530
• Using distributed or parallel set-up in script?: yes. Data parallel

[sayakpaul]

Ccing @williamberman for the accelerate hook issue.

[williamberman]

Ok, interesting thanks for the explanation @amitgurintelcom :)

In summary

1. We want to use accelerate saving model hooks to save the separate model components in the diffusers format that writes both the model config and the weights in their own separate directories.
2. We know our model components are going to support save_pretrained so we don't have an issue there
3. We have to check the class type to know which directory to save it in, however we check against the same type as the previously defined unet in the script and that unet might be wrapped by accelerate while the unet passed to the callback is not wrapped.

I'm not really a fan of "CLIPTextModel" in str(type(model)), it's a bit brittle.

Ideally there should be a recommended canonical way to dispatch on the type of the models passed to the callback. From the accelerate code, the models are passed from the list Accelerator._models. From a glance, it looks like we can usually rely on the type of the class to be the one from the class initially passed to prepare. However, in some cases it looks like the initial class is wrapped (preparing for megatron) so I would want a stronger guarantee about what the type of the class is if we can dispatch on it.

Ideally, we have a strong guarantee from accelerate on the type of the class, so we should ask the accelerate team if we can get that, will ping them in slack.

[amitgurintelcom]

Thank you. Meanwhile, the proposed workarounds (both for saving and loading the checkpoint) works for me. From: Will Berman ***@***.***> Sent: Tuesday, 2 May, 2023 22:55 To: huggingface/diffusers ***@***.***> Cc: Gur, Amit ***@***.***>; Mention ***@***.***> Subject: Re: [huggingface/diffusers] saving and loading checkpoints do not work on train_dreambooth.py when using text_encoder (Issue #3296) Ok, interesting thanks for the explanation @amitgurintelcom<https://github.com/amitgurintelcom> :) In summary 1. We want to use accelerate saving model hooks to save the separate model components in the diffusers format that writes both the model config and the weights in their own separate directories. 2. We know our model components are going to support save_pretrained so we don't have an issue there 3. We have to check the class type to know which directory to save it in, however we check against the same type as the previously defined unet in the script and that unet might be wrapped by accelerate while the unet passed to the callback is not wrapped. I'm not really a fan of "CLIPTextModel" in str(type(model)), it's a bit brittle. Ideally there should be a recommended canonical way to dispatch on the type of the models passed to the callback. From the accelerate code, the models are passed from the list Accelerator._models. From a glance, it looks like we can usually rely on the type of the class to be the one from the class initially passed to prepare. However, in some cases it looks like the initial class is wrapped (preparing for megatron) so I would want a stronger guarantee about what the type of the class is if we can dispatch on it. Ideally, we have a strong guarantee from accelerate on the type of the class, so we should ask the accelerate team if we can get that, will ping them in slack. — Reply to this email directly, view it on GitHub<#3296 (comment)>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AMEXRJXCVKOW4XWHRE4H7IDXEFRA3ANCNFSM6AAAAAAXQ67G54>. You are receiving this because you were mentioned.Message ID: ***@***.******@***.***>> --------------------------------------------------------------------- Intel Israel (74) Limited This e-mail and any attachments may contain confidential material for the sole use of the intended recipient(s). Any review or distribution by others is strictly prohibited. If you are not the intended recipient, please contact the sender and delete all copies.

[patrickvonplaten]

I don't think accelerate can make too many changes here just for diffusers -> can't we fix this directly in the training scripts? cc @williamberman

[patrickvonplaten]

To make some progress here. Can we try to first establish an easy, minimum reproducible codesnippet?

The following code:

accelerate launch train_dreambooth.py
--mixed_precision="fp16"
--revision="fp16"
--instance_data_dir="./XYZ"
--output_dir="checkpoints"
--instance_prompt="XYZ"
--resolution=512
--train_batch_size=2
--gradient_accumulation_steps=1
--learning_rate=1.5e-6
--lr_scheduler="constant"
--lr_warmup_steps=0
--max_train_steps=1200
--num_class_images=260
--with_prior_preservation
--prior_loss_weight=1.0
--class_prompt="man"
--class_data_dir="./man"
--train_text_encoder
--validation_prompt="Portrait of XYZ"
--validation_steps=100
--checkpointing_steps=200
--pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4"
--resume_from_checkpoint="latest"

cannot be run because I don't know what "./XYZ" refers to.

It would be amazing if we could get a easy to use recipe of how to reproduce the problem before we dive deep into changing logic in accelerate

[markrmiller]

Should just take a small addition to test_dreambooth_lora_with_text_encoder I believe - you just have to run the script again with resume from checkpoint and it won't be able to.

[patrickvonplaten]

cc @williamberman here, could you check this?

[patrickvonplaten]

Hmm, actually I just tried reproducing it, but couldn't.

I've done the following:

1. Download data

from huggingface_hub import snapshot_download

local_dir = "./dog"
snapshot_download(
    "diffusers/dog-example",
    local_dir=local_dir, repo_type="dataset",
    ignore_patterns=".gitattributes",
)

2. Run script for the first time

export MODEL_NAME="hf-internal-testing/tiny-stable-diffusion-pipe"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dog-example-demo"

accelerate launch train_dreambooth.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=10 \
  --resume_from_checkpoint="latest" \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --train_text_encoder \
  --lr_warmup_steps=0 \
  --max_train_steps=400 \
  --push_to_hub

3. Interrupt it in the middle so that only the -10 checkpoint folder is saved. Then start again, the same script. It correctly loads from the 10th checkpoint

[patrickvonplaten]

@markrmiller @amitgurintelcom can you check whether you still have the error ? Do you also get an error when following this: #3296 (comment)

[amitgurintelcom]

@patrickvonplaten @markrmiller Tried again. Same issue.
This is the error I get:
Traceback (most recent call last):
File "train_dreambooth.py", line 1085, in
main(args)
File "train_dreambooth.py", line 944, in main
accelerator.load_state(os.path.join(args.output_dir, path))
File "/home/agur/.local/lib/python3.8/site-packages/accelerate/accelerator.py", line 2394, in load_state
hook(models, input_dir)
File "train_dreambooth.py", line 759, in load_model_hook
model.register_to_config(**load_model.config)
File "/home/agur/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1614, in getattr
Traceback (most recent call last):
File "train_dreambooth.py", line 1085, in
raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'CLIPTextModel' object has no attribute 'register_to_config'

The issues are in 2 places:

1. https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py#L741
   sub_dir = "unet" if type(model) == type(unet) else "text_encoder"

In my env, this is what I get:
type(model): <class 'diffusers.models.unet_2d_condition.UNet2DConditionModel'>
type(unet): <class 'torch.nn.parallel.distributed.DistributedDataParallel'>

As you can see, in my env, type(unet) is a class of DistributedDataParallel
Proposed fix:
sub_dir = "unet" if "unet" in str(type(model)) else "text_encoder"

2. https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py#L752
   if type(model) == type(text_encoder):

type(model): <class 'transformers.models.clip.modeling_clip.CLIPTextModel'>
type(text_encoder): <class 'torch.nn.parallel.distributed.DistributedDataParallel'>
Proposed fix:
if "CLIPTextModel" in str(type(model)):

Issue #2480 is the same as this issue and is not resolved.

Thanks,
Amit.

[amitgurintelcom]

This is the command line I used:
export DISABLE_TELEMETRY=YES
accelerate launch train_dreambooth.py
--mixed_precision="fp16"
--revision="fp16"
--instance_data_dir="./Dog"
--output_dir="checkpoints"
--instance_prompt="sks Dog"
--resolution=512
--train_batch_size=2
--gradient_accumulation_steps=1
--learning_rate=1.5e-6
--lr_scheduler="constant"
--lr_warmup_steps=0
--max_train_steps=1400
--num_class_images=75
--with_prior_preservation
--prior_loss_weight=1.0
--class_prompt="dog"
--class_data_dir="./dog"
--train_text_encoder
--validation_prompt="Portrait of a Dog, hyper realistic face"
--validation_steps=100
--checkpointing_steps=200
--pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4"
--resume_from_checkpoint="latest"

This is the accelerate env:

• Accelerate version: 0.18.0
• Platform: Linux-5.15.0-71-generic-x86_64-with-glibc2.29
• Python version: 3.8.10
• Numpy version: 1.24.2
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Accelerate default config:
  - compute_environment: LOCAL_MACHINE
  - distributed_type: MULTI_GPU
  - mixed_precision: fp16
  - 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: []