huggingface / huggingface/diffusers

bug in load lora weights when add align_device_hook to model

オープン
#7,539 コメント 19 件 リアクション 0 件 担当者 0 名 GitHub で見る
bug wip
主要言語
Python
スター
34.5k
フォーク
7.3k
平均マージ
3日 3時間
マージ済み PR(30日)
91

説明

### Describe the bug

i noticed that when i add ```align_device_hook``` to module in pipeline manually, then ```load_lora_weights``` function will enable the sequential cpu offload. so i dig deeper and find that ```load_lora_weights``` function use ```_optionally_disable_offloading``` function to decide whether to sequentially cpu offload. this use ```_optionally_disable_offloading``` function was:
```python
def _optionally_disable_offloading(cls, _pipeline):
"""
Optionally removes offloading in case the pipeline has been already sequentially offloaded to CPU.

Args:
_pipeline (`DiffusionPipeline`):
The pipeline to disable offloading for.

Returns:
tuple:
A tuple indicating if `is_model_cpu_offload` or `is_sequential_cpu_offload` is True.
"""
is_model_cpu_offload = False
is_sequential_cpu_offload = False

if _pipeline is not None:
for _, component in _pipeline.components.items():
if isinstance(component, nn.Module) and hasattr(component, "_hf_hook"):
if not is_model_cpu_offload:
is_model_cpu_offload = isinstance(component._hf_hook, CpuOffload)
if not is_sequential_cpu_offload:
is_sequential_cpu_offload = isinstance(component._hf_hook, AlignDevicesHook)

logger.info(
"Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again."
)
remove_hook_from_module(component, recurse=is_sequential_cpu_offload)

return (is_model_cpu_offload, is_sequential_cpu_offload)
```
so i was curious that why ```is_sequential_cpu_offload = True``` when component has AlignDevicesHook? Shouldn't it be True only when the component device is CPU?

### Reproduction

```
from diffusers import StableDiffusionControlNetImg2ImgPipeline,ControlNetModel
from accelerate.hooks import attach_align_device_hook_on_blocks

pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained("/media/74nvme/checkpoints/diffusers_models/stable-diffusion-v1-5/",controlnet=[controlnet1,controlnet2],torch_dtype=torch.float16).to('cuda:0')

module_names, _ = pipe._get_signature_keys(pipe)
modules = [getattr(pipe, n, None) for n in module_names]
module_names = [name for m,name in zip(modules,module_names) if isinstance(m, torch.nn.Module)]
modules = [m for m in modules if isinstance(m, torch.nn.Module)]
print(module_names)

for module,name in zip(modules,module_names):
if name == 'unet' or name == 'controlnet':
module.to('cuda:0')
attach_align_device_hook_on_blocks(
module,
execution_device=module.device,
)
else:
module.to('cuda:1')
attach_align_device_hook_on_blocks(
module,
execution_device=module.device,
)
```
and then
```pipe.load_lora_weights(lora_weights_path)``` will change all component device

### Logs

_No response_

### System Info

diffusers:0.25.1
torch:2.2.0+cu118

### Who can help?

_No response_

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

LoRA の読み込みパスにある `_optionally_disable_offloading` から始め、`AlignDevicesHook` がどのように分類され、削除されるかを追跡します。提供されているパイプラインと `attach_align_device_hook_on_blocks` の例を使って動作を再現し、その後、オフロードの再適用パスを調べます。`load_lora_weights` の実行時に、手動でアタッチしたフックによってコンポーネントのデバイスが予期せず変更されなくなり、このシナリオがテストでカバーされていれば完了です。

索引モデルが issue の本文から書いたものです。

評価

技術スタック
python, pytorch
領域
backend, machine-learning
issue の種類
バグ
難易度
4/5
見積もり時間
3〜5日
活発さ
停滞
明瞭さ
おおむね明確
初心者へのやさしさ
42/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。