mirror of
https://github.com/huggingface/diffusers.git
synced 2026-01-27 17:22:53 +03:00
555 lines
23 KiB
Python
555 lines
23 KiB
Python
# coding=utf-8
|
|
# Copyright 2025 HuggingFace Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import gc
|
|
import json
|
|
import os
|
|
import re
|
|
|
|
import pytest
|
|
import safetensors.torch
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from diffusers.utils.import_utils import is_peft_available
|
|
from diffusers.utils.testing_utils import check_if_dicts_are_equal
|
|
|
|
from ...testing_utils import (
|
|
assert_tensors_close,
|
|
backend_empty_cache,
|
|
is_lora,
|
|
is_torch_compile,
|
|
require_peft_backend,
|
|
require_peft_version_greater,
|
|
require_torch_accelerator,
|
|
require_torch_version_greater,
|
|
torch_device,
|
|
)
|
|
|
|
|
|
if is_peft_available():
|
|
from diffusers.loaders.peft import PeftAdapterMixin
|
|
|
|
|
|
def check_if_lora_correctly_set(model) -> bool:
|
|
"""
|
|
Check if LoRA layers are correctly set in the model.
|
|
|
|
Args:
|
|
model: The model to check
|
|
|
|
Returns:
|
|
bool: True if LoRA is correctly set, False otherwise
|
|
"""
|
|
from peft.tuners.tuners_utils import BaseTunerLayer
|
|
|
|
for module in model.modules():
|
|
if isinstance(module, BaseTunerLayer):
|
|
return True
|
|
return False
|
|
|
|
|
|
@is_lora
|
|
@require_peft_backend
|
|
class LoraTesterMixin:
|
|
"""
|
|
Mixin class for testing LoRA/PEFT functionality on models.
|
|
|
|
Expected from config mixin:
|
|
- model_class: The model class to test
|
|
|
|
Expected methods from config mixin:
|
|
- get_init_dict(): Returns dict of arguments to initialize the model
|
|
- get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass
|
|
|
|
Pytest mark: lora
|
|
Use `pytest -m "not lora"` to skip these tests
|
|
"""
|
|
|
|
def setup_method(self):
|
|
if not issubclass(self.model_class, PeftAdapterMixin):
|
|
pytest.skip(f"PEFT is not supported for this model ({self.model_class.__name__}).")
|
|
|
|
@torch.no_grad()
|
|
def test_save_load_lora_adapter(self, tmp_path, rank=4, lora_alpha=4, use_dora=False):
|
|
from peft import LoraConfig
|
|
from peft.utils import get_peft_model_state_dict
|
|
|
|
init_dict = self.get_init_dict()
|
|
inputs_dict = self.get_dummy_inputs()
|
|
model = self.model_class(**init_dict).to(torch_device)
|
|
|
|
torch.manual_seed(0)
|
|
output_no_lora = model(**inputs_dict, return_dict=False)[0]
|
|
|
|
denoiser_lora_config = LoraConfig(
|
|
r=rank,
|
|
lora_alpha=lora_alpha,
|
|
target_modules=["to_q", "to_k", "to_v", "to_out.0"],
|
|
init_lora_weights=False,
|
|
use_dora=use_dora,
|
|
)
|
|
model.add_adapter(denoiser_lora_config)
|
|
assert check_if_lora_correctly_set(model), "LoRA layers not set correctly"
|
|
|
|
torch.manual_seed(0)
|
|
outputs_with_lora = model(**inputs_dict, return_dict=False)[0]
|
|
|
|
assert not torch.allclose(output_no_lora, outputs_with_lora, atol=1e-4, rtol=1e-4), (
|
|
"Output should differ with LoRA enabled"
|
|
)
|
|
|
|
model.save_lora_adapter(tmp_path)
|
|
assert os.path.isfile(os.path.join(tmp_path, "pytorch_lora_weights.safetensors")), (
|
|
"LoRA weights file not created"
|
|
)
|
|
|
|
state_dict_loaded = safetensors.torch.load_file(os.path.join(tmp_path, "pytorch_lora_weights.safetensors"))
|
|
|
|
model.unload_lora()
|
|
assert not check_if_lora_correctly_set(model), "LoRA should be unloaded"
|
|
|
|
model.load_lora_adapter(tmp_path, prefix=None, use_safetensors=True)
|
|
state_dict_retrieved = get_peft_model_state_dict(model, adapter_name="default_0")
|
|
|
|
for k in state_dict_loaded:
|
|
loaded_v = state_dict_loaded[k]
|
|
retrieved_v = state_dict_retrieved[k].to(loaded_v.device)
|
|
assert_tensors_close(loaded_v, retrieved_v, atol=1e-5, rtol=0, msg=f"Mismatch in LoRA weight {k}")
|
|
|
|
assert check_if_lora_correctly_set(model), "LoRA layers not set correctly after reload"
|
|
|
|
torch.manual_seed(0)
|
|
outputs_with_lora_2 = model(**inputs_dict, return_dict=False)[0]
|
|
|
|
assert not torch.allclose(output_no_lora, outputs_with_lora_2, atol=1e-4, rtol=1e-4), (
|
|
"Output should differ with LoRA enabled"
|
|
)
|
|
assert_tensors_close(
|
|
outputs_with_lora,
|
|
outputs_with_lora_2,
|
|
atol=1e-4,
|
|
rtol=1e-4,
|
|
msg="Outputs should match before and after save/load",
|
|
)
|
|
|
|
def test_lora_wrong_adapter_name_raises_error(self, tmp_path):
|
|
from peft import LoraConfig
|
|
|
|
init_dict = self.get_init_dict()
|
|
model = self.model_class(**init_dict).to(torch_device)
|
|
|
|
denoiser_lora_config = LoraConfig(
|
|
r=4,
|
|
lora_alpha=4,
|
|
target_modules=["to_q", "to_k", "to_v", "to_out.0"],
|
|
init_lora_weights=False,
|
|
use_dora=False,
|
|
)
|
|
model.add_adapter(denoiser_lora_config)
|
|
assert check_if_lora_correctly_set(model), "LoRA layers not set correctly"
|
|
|
|
wrong_name = "foo"
|
|
with pytest.raises(ValueError) as exc_info:
|
|
model.save_lora_adapter(tmp_path, adapter_name=wrong_name)
|
|
|
|
assert f"Adapter name {wrong_name} not found in the model." in str(exc_info.value)
|
|
|
|
def test_lora_adapter_metadata_is_loaded_correctly(self, tmp_path, rank=4, lora_alpha=4, use_dora=False):
|
|
from peft import LoraConfig
|
|
|
|
init_dict = self.get_init_dict()
|
|
model = self.model_class(**init_dict).to(torch_device)
|
|
|
|
denoiser_lora_config = LoraConfig(
|
|
r=rank,
|
|
lora_alpha=lora_alpha,
|
|
target_modules=["to_q", "to_k", "to_v", "to_out.0"],
|
|
init_lora_weights=False,
|
|
use_dora=use_dora,
|
|
)
|
|
model.add_adapter(denoiser_lora_config)
|
|
metadata = model.peft_config["default"].to_dict()
|
|
assert check_if_lora_correctly_set(model), "LoRA layers not set correctly"
|
|
|
|
model.save_lora_adapter(tmp_path)
|
|
model_file = os.path.join(tmp_path, "pytorch_lora_weights.safetensors")
|
|
assert os.path.isfile(model_file), "LoRA weights file not created"
|
|
|
|
model.unload_lora()
|
|
assert not check_if_lora_correctly_set(model), "LoRA should be unloaded"
|
|
|
|
model.load_lora_adapter(tmp_path, prefix=None, use_safetensors=True)
|
|
parsed_metadata = model.peft_config["default_0"].to_dict()
|
|
check_if_dicts_are_equal(metadata, parsed_metadata)
|
|
|
|
def test_lora_adapter_wrong_metadata_raises_error(self, tmp_path):
|
|
from peft import LoraConfig
|
|
|
|
from diffusers.loaders.lora_base import LORA_ADAPTER_METADATA_KEY
|
|
|
|
init_dict = self.get_init_dict()
|
|
model = self.model_class(**init_dict).to(torch_device)
|
|
|
|
denoiser_lora_config = LoraConfig(
|
|
r=4,
|
|
lora_alpha=4,
|
|
target_modules=["to_q", "to_k", "to_v", "to_out.0"],
|
|
init_lora_weights=False,
|
|
use_dora=False,
|
|
)
|
|
model.add_adapter(denoiser_lora_config)
|
|
assert check_if_lora_correctly_set(model), "LoRA layers not set correctly"
|
|
|
|
model.save_lora_adapter(tmp_path)
|
|
model_file = os.path.join(tmp_path, "pytorch_lora_weights.safetensors")
|
|
assert os.path.isfile(model_file), "LoRA weights file not created"
|
|
|
|
# Perturb the metadata in the state dict
|
|
loaded_state_dict = safetensors.torch.load_file(model_file)
|
|
metadata = {"format": "pt"}
|
|
lora_adapter_metadata = denoiser_lora_config.to_dict()
|
|
lora_adapter_metadata.update({"foo": 1, "bar": 2})
|
|
for key, value in lora_adapter_metadata.items():
|
|
if isinstance(value, set):
|
|
lora_adapter_metadata[key] = list(value)
|
|
metadata[LORA_ADAPTER_METADATA_KEY] = json.dumps(lora_adapter_metadata, indent=2, sort_keys=True)
|
|
safetensors.torch.save_file(loaded_state_dict, model_file, metadata=metadata)
|
|
|
|
model.unload_lora()
|
|
assert not check_if_lora_correctly_set(model), "LoRA should be unloaded"
|
|
|
|
with pytest.raises(TypeError) as exc_info:
|
|
model.load_lora_adapter(tmp_path, prefix=None, use_safetensors=True)
|
|
assert "`LoraConfig` class could not be instantiated" in str(exc_info.value)
|
|
|
|
|
|
@is_lora
|
|
@is_torch_compile
|
|
@require_peft_backend
|
|
@require_peft_version_greater("0.14.0")
|
|
@require_torch_version_greater("2.7.1")
|
|
@require_torch_accelerator
|
|
class LoraHotSwappingForModelTesterMixin:
|
|
"""
|
|
Mixin class for testing LoRA hot swapping functionality on models.
|
|
|
|
Test that hotswapping does not result in recompilation on the model directly.
|
|
We're not extensively testing the hotswapping functionality since it is implemented in PEFT
|
|
and is extensively tested there. The goal of this test is specifically to ensure that
|
|
hotswapping with diffusers does not require recompilation.
|
|
|
|
See https://github.com/huggingface/peft/blob/eaab05e18d51fb4cce20a73c9acd82a00c013b83/tests/test_gpu_examples.py#L4252
|
|
for the analogous PEFT test.
|
|
|
|
Expected from config mixin:
|
|
- model_class: The model class to test
|
|
|
|
Optional properties:
|
|
- different_shapes_for_compilation: List of (height, width) tuples for dynamic compilation tests (default: None)
|
|
|
|
Expected methods from config mixin:
|
|
- get_init_dict(): Returns dict of arguments to initialize the model
|
|
- get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass
|
|
|
|
Pytest marks: lora, torch_compile
|
|
Use `pytest -m "not lora"` or `pytest -m "not torch_compile"` to skip these tests
|
|
"""
|
|
|
|
@property
|
|
def different_shapes_for_compilation(self) -> list[tuple[int, int]] | None:
|
|
"""Optional list of (height, width) tuples for dynamic compilation tests."""
|
|
return None
|
|
|
|
def setup_method(self):
|
|
if not issubclass(self.model_class, PeftAdapterMixin):
|
|
pytest.skip(f"PEFT is not supported for this model ({self.model_class.__name__}).")
|
|
|
|
def teardown_method(self):
|
|
# It is critical that the dynamo cache is reset for each test. Otherwise, if the test re-uses the same model,
|
|
# there will be recompilation errors, as torch caches the model when run in the same process.
|
|
torch.compiler.reset()
|
|
gc.collect()
|
|
backend_empty_cache(torch_device)
|
|
|
|
def _get_lora_config(self, lora_rank, lora_alpha, target_modules):
|
|
from peft import LoraConfig
|
|
|
|
lora_config = LoraConfig(
|
|
r=lora_rank,
|
|
lora_alpha=lora_alpha,
|
|
target_modules=target_modules,
|
|
init_lora_weights=False,
|
|
use_dora=False,
|
|
)
|
|
return lora_config
|
|
|
|
def _get_linear_module_name_other_than_attn(self, model):
|
|
linear_names = [
|
|
name for name, module in model.named_modules() if isinstance(module, nn.Linear) and "to_" not in name
|
|
]
|
|
return linear_names[0]
|
|
|
|
def _check_model_hotswap(self, tmp_path, do_compile, rank0, rank1, target_modules0, target_modules1=None):
|
|
"""
|
|
Check that hotswapping works on a model.
|
|
|
|
Steps:
|
|
- create 2 LoRA adapters and save them
|
|
- load the first adapter
|
|
- hotswap the second adapter
|
|
- check that the outputs are correct
|
|
- optionally compile the model
|
|
- optionally check if recompilations happen on different shapes
|
|
|
|
Note: We set rank == alpha here because save_lora_adapter does not save the alpha scalings, thus the test would
|
|
fail if the values are different. Since rank != alpha does not matter for the purpose of this test, this is
|
|
fine.
|
|
"""
|
|
different_shapes = self.different_shapes_for_compilation
|
|
# create 2 adapters with different ranks and alphas
|
|
torch.manual_seed(0)
|
|
init_dict = self.get_init_dict()
|
|
inputs_dict = self.get_dummy_inputs()
|
|
model = self.model_class(**init_dict).to(torch_device)
|
|
|
|
alpha0, alpha1 = rank0, rank1
|
|
max_rank = max([rank0, rank1])
|
|
if target_modules1 is None:
|
|
target_modules1 = target_modules0[:]
|
|
lora_config0 = self._get_lora_config(rank0, alpha0, target_modules0)
|
|
lora_config1 = self._get_lora_config(rank1, alpha1, target_modules1)
|
|
|
|
model.add_adapter(lora_config0, adapter_name="adapter0")
|
|
with torch.inference_mode():
|
|
torch.manual_seed(0)
|
|
output0_before = model(**inputs_dict)["sample"]
|
|
|
|
model.add_adapter(lora_config1, adapter_name="adapter1")
|
|
model.set_adapter("adapter1")
|
|
with torch.inference_mode():
|
|
torch.manual_seed(0)
|
|
output1_before = model(**inputs_dict)["sample"]
|
|
|
|
# sanity checks:
|
|
tol = 5e-3
|
|
assert not torch.allclose(output0_before, output1_before, atol=tol, rtol=tol)
|
|
assert not (output0_before == 0).all()
|
|
assert not (output1_before == 0).all()
|
|
|
|
# save the adapter checkpoints
|
|
model.save_lora_adapter(os.path.join(tmp_path, "0"), safe_serialization=True, adapter_name="adapter0")
|
|
model.save_lora_adapter(os.path.join(tmp_path, "1"), safe_serialization=True, adapter_name="adapter1")
|
|
del model
|
|
|
|
# load the first adapter
|
|
torch.manual_seed(0)
|
|
init_dict = self.get_init_dict()
|
|
model = self.model_class(**init_dict).to(torch_device)
|
|
|
|
if do_compile or (rank0 != rank1):
|
|
# no need to prepare if the model is not compiled or if the ranks are identical
|
|
model.enable_lora_hotswap(target_rank=max_rank)
|
|
|
|
file_name0 = os.path.join(os.path.join(tmp_path, "0"), "pytorch_lora_weights.safetensors")
|
|
file_name1 = os.path.join(os.path.join(tmp_path, "1"), "pytorch_lora_weights.safetensors")
|
|
model.load_lora_adapter(file_name0, safe_serialization=True, adapter_name="adapter0", prefix=None)
|
|
|
|
if do_compile:
|
|
model = torch.compile(model, mode="reduce-overhead", dynamic=different_shapes is not None)
|
|
|
|
with torch.inference_mode():
|
|
# additionally check if dynamic compilation works.
|
|
if different_shapes is not None:
|
|
for height, width in different_shapes:
|
|
new_inputs_dict = self.prepare_dummy_input(height=height, width=width)
|
|
_ = model(**new_inputs_dict)
|
|
else:
|
|
output0_after = model(**inputs_dict)["sample"]
|
|
assert_tensors_close(
|
|
output0_before, output0_after, atol=tol, rtol=tol, msg="Output mismatch after loading adapter0"
|
|
)
|
|
|
|
# hotswap the 2nd adapter
|
|
model.load_lora_adapter(file_name1, adapter_name="adapter0", hotswap=True, prefix=None)
|
|
|
|
# we need to call forward to potentially trigger recompilation
|
|
with torch.inference_mode():
|
|
if different_shapes is not None:
|
|
for height, width in different_shapes:
|
|
new_inputs_dict = self.prepare_dummy_input(height=height, width=width)
|
|
_ = model(**new_inputs_dict)
|
|
else:
|
|
output1_after = model(**inputs_dict)["sample"]
|
|
assert_tensors_close(
|
|
output1_before,
|
|
output1_after,
|
|
atol=tol,
|
|
rtol=tol,
|
|
msg="Output mismatch after hotswapping to adapter1",
|
|
)
|
|
|
|
# check error when not passing valid adapter name
|
|
name = "does-not-exist"
|
|
msg = f"Trying to hotswap LoRA adapter '{name}' but there is no existing adapter by that name"
|
|
with pytest.raises(ValueError, match=re.escape(msg)):
|
|
model.load_lora_adapter(file_name1, adapter_name=name, hotswap=True, prefix=None)
|
|
|
|
@pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)])
|
|
def test_hotswapping_model(self, tmp_path, rank0, rank1):
|
|
self._check_model_hotswap(
|
|
tmp_path, do_compile=False, rank0=rank0, rank1=rank1, target_modules0=["to_q", "to_k", "to_v", "to_out.0"]
|
|
)
|
|
|
|
@pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)])
|
|
def test_hotswapping_compiled_model_linear(self, tmp_path, rank0, rank1):
|
|
# It's important to add this context to raise an error on recompilation
|
|
target_modules = ["to_q", "to_k", "to_v", "to_out.0"]
|
|
with torch._dynamo.config.patch(error_on_recompile=True), torch._inductor.utils.fresh_inductor_cache():
|
|
self._check_model_hotswap(
|
|
tmp_path, do_compile=True, rank0=rank0, rank1=rank1, target_modules0=target_modules
|
|
)
|
|
|
|
@pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)])
|
|
def test_hotswapping_compiled_model_conv2d(self, tmp_path, rank0, rank1):
|
|
if "unet" not in self.model_class.__name__.lower():
|
|
pytest.skip("Test only applies to UNet.")
|
|
|
|
# It's important to add this context to raise an error on recompilation
|
|
target_modules = ["conv", "conv1", "conv2"]
|
|
with torch._dynamo.config.patch(error_on_recompile=True), torch._inductor.utils.fresh_inductor_cache():
|
|
self._check_model_hotswap(
|
|
tmp_path, do_compile=True, rank0=rank0, rank1=rank1, target_modules0=target_modules
|
|
)
|
|
|
|
@pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)])
|
|
def test_hotswapping_compiled_model_both_linear_and_conv2d(self, tmp_path, rank0, rank1):
|
|
if "unet" not in self.model_class.__name__.lower():
|
|
pytest.skip("Test only applies to UNet.")
|
|
|
|
# It's important to add this context to raise an error on recompilation
|
|
target_modules = ["to_q", "conv"]
|
|
with torch._dynamo.config.patch(error_on_recompile=True), torch._inductor.utils.fresh_inductor_cache():
|
|
self._check_model_hotswap(
|
|
tmp_path, do_compile=True, rank0=rank0, rank1=rank1, target_modules0=target_modules
|
|
)
|
|
|
|
@pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)])
|
|
def test_hotswapping_compiled_model_both_linear_and_other(self, tmp_path, rank0, rank1):
|
|
# In `test_hotswapping_compiled_model_both_linear_and_conv2d()`, we check if we can do hotswapping
|
|
# with `torch.compile()` for models that have both linear and conv layers. In this test, we check
|
|
# if we can target a linear layer from the transformer blocks and another linear layer from non-attention
|
|
# block.
|
|
target_modules = ["to_q"]
|
|
init_dict = self.get_init_dict()
|
|
model = self.model_class(**init_dict)
|
|
|
|
target_modules.append(self._get_linear_module_name_other_than_attn(model))
|
|
del model
|
|
|
|
# It's important to add this context to raise an error on recompilation
|
|
with torch._dynamo.config.patch(error_on_recompile=True):
|
|
self._check_model_hotswap(
|
|
tmp_path, do_compile=True, rank0=rank0, rank1=rank1, target_modules0=target_modules
|
|
)
|
|
|
|
def test_enable_lora_hotswap_called_after_adapter_added_raises(self):
|
|
# ensure that enable_lora_hotswap is called before loading the first adapter
|
|
lora_config = self._get_lora_config(8, 8, target_modules=["to_q"])
|
|
init_dict = self.get_init_dict()
|
|
model = self.model_class(**init_dict).to(torch_device)
|
|
model.add_adapter(lora_config)
|
|
|
|
msg = re.escape("Call `enable_lora_hotswap` before loading the first adapter.")
|
|
with pytest.raises(RuntimeError, match=msg):
|
|
model.enable_lora_hotswap(target_rank=32)
|
|
|
|
def test_enable_lora_hotswap_called_after_adapter_added_warning(self, caplog):
|
|
# ensure that enable_lora_hotswap is called before loading the first adapter
|
|
import logging
|
|
|
|
lora_config = self._get_lora_config(8, 8, target_modules=["to_q"])
|
|
init_dict = self.get_init_dict()
|
|
model = self.model_class(**init_dict).to(torch_device)
|
|
model.add_adapter(lora_config)
|
|
msg = (
|
|
"It is recommended to call `enable_lora_hotswap` before loading the first adapter to avoid recompilation."
|
|
)
|
|
with caplog.at_level(logging.WARNING):
|
|
model.enable_lora_hotswap(target_rank=32, check_compiled="warn")
|
|
assert any(msg in record.message for record in caplog.records)
|
|
|
|
def test_enable_lora_hotswap_called_after_adapter_added_ignore(self, caplog):
|
|
# check possibility to ignore the error/warning
|
|
import logging
|
|
|
|
lora_config = self._get_lora_config(8, 8, target_modules=["to_q"])
|
|
init_dict = self.get_init_dict()
|
|
model = self.model_class(**init_dict).to(torch_device)
|
|
model.add_adapter(lora_config)
|
|
with caplog.at_level(logging.WARNING):
|
|
model.enable_lora_hotswap(target_rank=32, check_compiled="ignore")
|
|
assert len(caplog.records) == 0
|
|
|
|
def test_enable_lora_hotswap_wrong_check_compiled_argument_raises(self):
|
|
# check that wrong argument value raises an error
|
|
lora_config = self._get_lora_config(8, 8, target_modules=["to_q"])
|
|
init_dict = self.get_init_dict()
|
|
model = self.model_class(**init_dict).to(torch_device)
|
|
model.add_adapter(lora_config)
|
|
msg = re.escape("check_compiles should be one of 'error', 'warn', or 'ignore', got 'wrong-argument' instead.")
|
|
with pytest.raises(ValueError, match=msg):
|
|
model.enable_lora_hotswap(target_rank=32, check_compiled="wrong-argument")
|
|
|
|
def test_hotswap_second_adapter_targets_more_layers_raises(self, tmp_path, caplog):
|
|
# check the error and log
|
|
import logging
|
|
|
|
# at the moment, PEFT requires the 2nd adapter to target the same or a subset of layers
|
|
target_modules0 = ["to_q"]
|
|
target_modules1 = ["to_q", "to_k"]
|
|
with pytest.raises(RuntimeError): # peft raises RuntimeError
|
|
with caplog.at_level(logging.ERROR):
|
|
self._check_model_hotswap(
|
|
tmp_path,
|
|
do_compile=True,
|
|
rank0=8,
|
|
rank1=8,
|
|
target_modules0=target_modules0,
|
|
target_modules1=target_modules1,
|
|
)
|
|
assert any("Hotswapping adapter0 was unsuccessful" in record.message for record in caplog.records)
|
|
|
|
@pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)])
|
|
@require_torch_version_greater("2.7.1")
|
|
def test_hotswapping_compile_on_different_shapes(self, tmp_path, rank0, rank1):
|
|
different_shapes_for_compilation = self.different_shapes_for_compilation
|
|
if different_shapes_for_compilation is None:
|
|
pytest.skip(f"Skipping as `different_shapes_for_compilation` is not set for {self.__class__.__name__}.")
|
|
# Specifying `use_duck_shape=False` instructs the compiler if it should use the same symbolic
|
|
# variable to represent input sizes that are the same. For more details,
|
|
# check out this [comment](https://github.com/huggingface/diffusers/pull/11327#discussion_r2047659790).
|
|
torch.fx.experimental._config.use_duck_shape = False
|
|
|
|
target_modules = ["to_q", "to_k", "to_v", "to_out.0"]
|
|
with torch._dynamo.config.patch(error_on_recompile=True):
|
|
self._check_model_hotswap(
|
|
tmp_path,
|
|
do_compile=True,
|
|
rank0=rank0,
|
|
rank1=rank1,
|
|
target_modules0=target_modules,
|
|
)
|