Validation Errors When Training ControlNet with FSDP

[liming-ai]

Describe the bug

Adding FSDP into training the given ControlNet example training code leads to an unexpected bug, with the following config:

compute_environment: LOCAL_MACHINE
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: SIZE_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_min_num_params: 100000000
  fsdp_offload_params: false
  fsdp_sharding_strategy: 1
  fsdp_state_dict_type: FULL_STATE_DICT
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

This error is caused by this line cannot unwrap FSDP into the original class

diffusers/examples/controlnet/train_controlnet.py

Line 78 in f74d5e1

controlnet = accelerator.unwrap_model(controlnet)

# type(controlnet): <class 'torch.distributed.fsdp.fully_sharded_data_parallel.FullyShardedDataParallel'
controlnet = accelerator.unwrap_model(controlnet)
# type(controlnet): <class 'torch.distributed.fsdp.fully_sharded_data_parallel.FullyShardedDataParallel'

To this end, we cannot use the default diffusers.StableDiffusionControlNetPipeline to run the inference/validation:

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /home/tiger/diffusers/examples/controlnet/train_controlnet.py:1127 in <module>                   │
│                                                                                                  │
│   1124                                                                                           │
│   1125 if __name__ == "__main__":                                                                │
│   1126 │   args = parse_args()                                                                   │
│ ❱ 1127 │   main(args)                                                                            │
│   1128                                                                                           │
│                                                                                                  │
│ /home/tiger/diffusers/examples/controlnet/train_controlnet.py:1083 in main                       │
│                                                                                                  │
│   1080 │   │   │   │   │   │   │   │   │   shutil.rmtree(removing_checkpoint)                    │
│   1081 │   │   │   │   │                                                                         │
│   1082 │   │   │   │   │   if args.validation_prompt is not None and global_step % args.validat  │
│ ❱ 1083 │   │   │   │   │   │   image_logs = log_validation(                                      │
│   1084 │   │   │   │   │   │   │   vae,                                                          │
│   1085 │   │   │   │   │   │   │   text_encoder,                                                 │
│   1086 │   │   │   │   │   │   │   tokenizer,                                                    │
│                                                                                                  │
│ /home/tiger/diffusers/examples/controlnet/train_controlnet.py:126 in log_validation              │
│                                                                                                  │
│    123 │   │                                                                                     │
│    124 │   │   for _ in range(args.num_validation_images):                                       │
│    125 │   │   │   with torch.autocast("cuda"):                                                  │
│ ❱  126 │   │   │   │   image = pipeline(                                                         │
│    127 │   │   │   │   │   validation_prompt, validation_image, num_inference_steps=20, generat  │
│    128 │   │   │   │   ).images[0]                                                               │
│    129                                                                                           │
│                                                                                                  │
│ /usr/local/lib/python3.9/dist-packages/torch/utils/_contextlib.py:115 in decorate_context        │
│                                                                                                  │
│   112 │   @functools.wraps(func)                                                                 │
│   113 │   def decorate_context(*args, **kwargs):                                                 │
│   114 │   │   with ctx_factory():                                                                │
│ ❱ 115 │   │   │   return func(*args, **kwargs)                                                   │
│   116 │                                                                                          │
│   117 │   return decorate_context                                                                │
│   118                                                                                            │
│                                                                                                  │
│ /home/tiger/.local/lib/python3.9/site-packages/diffusers/pipelines/controlnet/pipeline_controlne │
│ t.py:840 in __call__                                                                             │
│                                                                                                  │
│    837 │   │   │   ]                                                                             │
│    838 │   │                                                                                     │
│    839 │   │   # 1. Check inputs. Raise error if not correct                                     │
│ ❱  840 │   │   self.check_inputs(                                                                │
│    841 │   │   │   prompt,                                                                       │
│    842 │   │   │   image,                                                                        │
│    843 │   │   │   callback_steps,                                                               │
│                                                                                                  │
│ /home/tiger/.local/lib/python3.9/site-packages/diffusers/pipelines/controlnet/pipeline_controlne │
│ t.py:570 in check_inputs                                                                         │
│                                                                                                  │
│    567 │   │   │   for image_ in image:                                                          │
│    568 │   │   │   │   self.check_image(image_, prompt, prompt_embeds)                           │
│    569 │   │   else:                                                                             │
│ ❱  570 │   │   │   assert False                                                                  │
│    571 │   │                                                                                     │
│    572 │   │   # Check `controlnet_conditioning_scale`                                           │
│    573 │   │   if (                                                                              │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
AssertionError

Reproduction

Please use FSDP, and run the given ControlNet example training code.

Logs

No response

System Info

• diffusers version: 0.18.0.dev0
• Platform: Linux-5.4.56.bsk.10-amd64-x86_64-with-glibc2.31
• Python version: 3.9.2
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Huggingface_hub version: 0.15.1
• Transformers version: 4.27.4
• Accelerate version: 0.20.3
• xFormers version: 0.0.18
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: Yes

Who can help?

No response

[sayakpaul]

Sorry, we don't have the cycles to look into FSDP. But we're happy to welcome any patches to review.

[liming-ai]

Sorry, we don't have the cycles to look into FSDP. But we're happy to welcome any patches to review.

Does this mean we cannot do the default inference procedure (diffusers.StableDiffusionControlNetPipeline) if we use FSDP during training?

[sayakpaul]

That inference part should be regardless of that. If you have the compatible checkpoints already then performing inference should be error-free.

[liming-ai]

That inference part should be regardless of that. If you have the compatible checkpoints already then performing inference should be error-free.

This is the biggest problem I have encountered. How can I correctly save the model weight after FSDP wrapping as the weight normally used by diffusers? Could you please give me some examples?

[sayakpaul]

I believe FSDP wraps the model with another internal class.

When you do controlnet = accelerator.unwrap_model(controlnet), I'd imagine the FSDP wrapper would go away, but your logs suggest differently.

@muellerzr do you have any advice?

[liming-ai]

I believe FSDP wraps the model with another internal class.

When you do controlnet = accelerator.unwrap_model(controlnet), I'd imagine the FSDP wrapper would go away, but your logs suggest differently.

@muellerzr do you have any advice?

Unfortunately, the accelerator.unwrap_model() function cannot remove FSDP wrapper, only support DP and DDP
https://github.com/huggingface/accelerate/blob/180ef7c41576bc55cb1265fcdeb7da287e95c9b7/src/accelerate/utils/other.py#L58

[muellerzr]

CC @pacman100

[pacman100]

Hello @liming-ai, even if the unwrap gets you the controlnet, you will still run into issues if there are modules as part of the global FSDP unit as the parameters of those modules won't be gathered if you aren't calling global FSDP wrapper's forward call. So, those checks need to change for FSDP during eval which are in diffusers. You can try PR huggingface/accelerate#1753 but I think you will run into FSDP issues that i stated above.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[gunshi]

hey @liming-ai @pacman100, was wondering if you ended up being able to use FSDP with diffusers? Would be very grateful for some pointers on how to go about it.