Skip to content

Legacy Attention Block conversion doesn't work with device_map "sequential" from accelerate #3740

Description

@stefan-canva

Describe the bug

As far as I understand, checkpoints from before this refactor of the Attention module have to be converted with this function, because the child modules within the Attention module got renamed (e.g. "query" -> "to_q").

Apparently this function doesn't get called when loading with the "sequential" device_map argument, which requires accelerate to be installed.

Reproduction

from diffusers import StableDiffusionInpaintPipeline
import torch

StableDiffusionInpaintPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-inpainting", revision="fp16", torch_dtype=torch.float16, device_map="sequential")

Logs

╭───────────────────── Traceback (most recent call last) ──────────────────────╮
│ /home/stefan/.config/JetBrains/PyCharmCE2023.1/scratches/scratch_84.py:4 in  │
│ <module>                                                                     │
│                                                                              │
│   1 from diffusers import StableDiffusionInpaintPipeline                     │
│   2 import torch                                                             │
│   3                                                                          │
│ ❱ 4 StableDiffusionInpaintPipeline.from_pretrained(                          │
│   5 │   │   │   │   "stabilityai/stable-diffusion-2-inpainting", revision="f │
│   6 │   │   │   )                                                            │
│   7                                                                          │
│                                                                              │
│ /home/stefan/anaconda3/envs/kaleidolib3/lib/python3.9/site-packages/diffuser │
│ s/pipelines/pipeline_utils.py:1063 in from_pretrained                        │
│                                                                              │
│   1060 │   │   │   │   loaded_sub_model = passed_class_obj[name]             │
│   1061 │   │   │   else:                                                     │
│   1062 │   │   │   │   # load sub model                                      │
│ ❱ 1063 │   │   │   │   loaded_sub_model = load_sub_model(                    │
│   1064 │   │   │   │   │   library_name=library_name,                        │
│   1065 │   │   │   │   │   class_name=class_name,                            │
│   1066 │   │   │   │   │   importable_classes=importable_classes,            │
│                                                                              │
│ /home/stefan/anaconda3/envs/kaleidolib3/lib/python3.9/site-packages/diffuser │
│ s/pipelines/pipeline_utils.py:451 in load_sub_model                          │
│                                                                              │
│    448 │                                                                     │
│    449 │   # check if the module is in a subdirectory                        │
│    450 │   if os.path.isdir(os.path.join(cached_folder, name)):              │
│ ❱  451 │   │   loaded_sub_model = load_method(os.path.join(cached_folder, na │
│    452 │   else:                                                             │
│    453 │   │   # else load from the root directory                           │
│    454 │   │   loaded_sub_model = load_method(cached_folder, **loading_kwarg │
│                                                                              │
│ /home/stefan/anaconda3/envs/kaleidolib3/lib/python3.9/site-packages/diffuser │
│ s/models/modeling_utils.py:649 in from_pretrained                            │
│                                                                              │
│   646 │   │   │   │   else:  # else let accelerate handle loading and dispat │
│   647 │   │   │   │   │   # Load weights and dispatch according to the devic │
│   648 │   │   │   │   │   # by default the device_map is None and the weight │
│ ❱ 649 │   │   │   │   │   accelerate.load_checkpoint_and_dispatch(           │
│   650 │   │   │   │   │   │   model,                                         │
│   651 │   │   │   │   │   │   model_file,                                    │
│   652 │   │   │   │   │   │   device_map,                                    │
│                                                                              │
│ /home/stefan/anaconda3/envs/kaleidolib3/lib/python3.9/site-packages/accelera │
│ te/big_modeling.py:486 in load_checkpoint_and_dispatch                       │
│                                                                              │
│   483 │   │   )                                                              │
│   484 │   if offload_state_dict is None and device_map is not None and "disk │
│   485 │   │   offload_state_dict = True                                      │
│ ❱ 486 │   load_checkpoint_in_model(                                          │
│   487 │   │   model,                                                         │
│   488 │   │   checkpoint,                                                    │
│   489 │   │   device_map=device_map,                                         │
│                                                                              │
│ /home/stefan/anaconda3/envs/kaleidolib3/lib/python3.9/site-packages/accelera │
│ te/utils/modeling.py:1116 in load_checkpoint_in_model                        │
│                                                                              │
│   1113 │   │   │   │   │   set_module_tensor_to_device(model, param_name, "m │
│   1114 │   │   │   │   │   offload_weight(param, param_name, state_dict_fold │
│   1115 │   │   │   │   else:                                                 │
│ ❱ 1116 │   │   │   │   │   set_module_tensor_to_device(model, param_name, pa │
│   1117 │   │                                                                 │
│   1118 │   │   # Force Python to clean up.                                   │
│   1119 │   │   del checkpoint                                                │
│                                                                              │
│ /home/stefan/anaconda3/envs/kaleidolib3/lib/python3.9/site-packages/accelera │
│ te/utils/modeling.py:142 in set_module_tensor_to_device                      │
│                                                                              │
│    139 │   if "." in tensor_name:                                            │
│    140 │   │   splits = tensor_name.split(".")                               │
│    141 │   │   for split in splits[:-1]:                                     │
│ ❱  142 │   │   │   new_module = getattr(module, split)                       │
│    143 │   │   │   if new_module is None:                                    │
│    144 │   │   │   │   raise ValueError(f"{module} has no attribute {split}. │
│    145 │   │   │   module = new_module                                       │
│                                                                              │
│ /home/stefan/anaconda3/envs/kaleidolib3/lib/python3.9/site-packages/torch/nn │
│ /modules/module.py:1269 in __getattr__                                       │
│                                                                              │
│   1266 │   │   │   modules = self.__dict__['_modules']                       │
│   1267 │   │   │   if name in modules:                                       │
│   1268 │   │   │   │   return modules[name]                                  │
│ ❱ 1269 │   │   raise AttributeError("'{}' object has no attribute '{}'".form │
│   1270 │   │   │   type(self).__name__, name))                               │
│   1271 │                                                                     │
│   1272 │   def __setattr__(self, name: str, value: Union[Tensor, 'Module'])  │
╰──────────────────────────────────────────────────────────────────────────────╯
AttributeError: 'Attention' object has no attribute 'query'

System Info

  • diffusers version: 0.17.0
  • Platform: Linux-5.15.0-73-generic-x86_64-with-glibc2.31
  • Python version: 3.9.15
  • PyTorch version (GPU?): 1.13.1+cu117 (True)
  • Huggingface_hub version: 0.15.1
  • Transformers version: 4.30.1
  • Accelerate version: 0.20.3
  • xFormers version: not installed
  • Using GPU in script?: yes
  • Using distributed or parallel set-up in script?: no

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions