diff --git a/tests/conftest.py b/tests/conftest.py index fd76d1c84ee7..0f7b9ef984ff 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -32,6 +32,21 @@ def pytest_configure(config): config.addinivalue_line("markers", "big_accelerator: marks tests as requiring big accelerator resources") + config.addinivalue_line("markers", "lora: marks tests for LoRA/PEFT functionality") + config.addinivalue_line("markers", "ip_adapter: marks tests for IP Adapter functionality") + config.addinivalue_line("markers", "training: marks tests for training functionality") + config.addinivalue_line("markers", "attention: marks tests for attention processor functionality") + config.addinivalue_line("markers", "memory: marks tests for memory optimization functionality") + config.addinivalue_line("markers", "cpu_offload: marks tests for CPU offloading functionality") + config.addinivalue_line("markers", "group_offload: marks tests for group offloading functionality") + config.addinivalue_line("markers", "compile: marks tests for torch.compile functionality") + config.addinivalue_line("markers", "single_file: marks tests for single file checkpoint loading") + config.addinivalue_line("markers", "bitsandbytes: marks tests for BitsAndBytes quantization functionality") + config.addinivalue_line("markers", "quanto: marks tests for Quanto quantization functionality") + config.addinivalue_line("markers", "torchao: marks tests for TorchAO quantization functionality") + config.addinivalue_line("markers", "gguf: marks tests for GGUF quantization functionality") + config.addinivalue_line("markers", "modelopt: marks tests for NVIDIA ModelOpt quantization functionality") + config.addinivalue_line("markers", "context_parallel: marks tests for context parallel inference functionality") def pytest_addoption(parser): diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py new file mode 100644 index 000000000000..e72a3c928b64 --- /dev/null +++ b/tests/models/testing_utils/__init__.py @@ -0,0 +1,38 @@ +from .attention import AttentionTesterMixin, ContextParallelTesterMixin +from .common import ModelTesterMixin +from .compile import TorchCompileTesterMixin +from .ip_adapter import IPAdapterTesterMixin +from .lora import LoraTesterMixin +from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin +from .quantization import ( + BitsAndBytesTesterMixin, + GGUFTesterMixin, + ModelOptTesterMixin, + QuantizationTesterMixin, + QuantoTesterMixin, + TorchAoTesterMixin, +) +from .single_file import SingleFileTesterMixin +from .training import TrainingTesterMixin + + +__all__ = [ + "ContextParallelTesterMixin", + "AttentionTesterMixin", + "BitsAndBytesTesterMixin", + "CPUOffloadTesterMixin", + "GGUFTesterMixin", + "GroupOffloadTesterMixin", + "IPAdapterTesterMixin", + "LayerwiseCastingTesterMixin", + "LoraTesterMixin", + "MemoryTesterMixin", + "ModelOptTesterMixin", + "ModelTesterMixin", + "QuantizationTesterMixin", + "QuantoTesterMixin", + "SingleFileTesterMixin", + "TorchAoTesterMixin", + "TorchCompileTesterMixin", + "TrainingTesterMixin", +] diff --git a/tests/models/testing_utils/attention.py b/tests/models/testing_utils/attention.py new file mode 100644 index 000000000000..f794a7a0aa4a --- /dev/null +++ b/tests/models/testing_utils/attention.py @@ -0,0 +1,263 @@ +# 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 os + +import pytest +import torch +import torch.multiprocessing as mp + +from diffusers.models._modeling_parallel import ContextParallelConfig +from diffusers.models.attention import AttentionModuleMixin +from diffusers.models.attention_processor import ( + AttnProcessor, +) + +from ...testing_utils import is_attention, is_context_parallel, require_torch_multi_accelerator, torch_device + + +@is_attention +class AttentionTesterMixin: + """ + Mixin class for testing attention processor and module functionality on models. + + Tests functionality from AttentionModuleMixin including: + - Attention processor management (set/get) + - QKV projection fusion/unfusion + - Attention backends (XFormers, NPU, etc.) + + Expected class attributes to be set by subclasses: + - model_class: The model class to test + - base_precision: Tolerance for floating point comparisons (default: 1e-3) + - uses_custom_attn_processor: Whether model uses custom attention processors (default: False) + + Expected methods to be implemented by subclasses: + - 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: attention + Use `pytest -m "not attention"` to skip these tests + """ + + base_precision = 1e-3 + + def test_fuse_unfuse_qkv_projections(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict) + model.to(torch_device) + model.eval() + + if not hasattr(model, "fuse_qkv_projections"): + pytest.skip("Model does not support QKV projection fusion.") + + # Get output before fusion + with torch.no_grad(): + output_before_fusion = model(**inputs_dict) + if isinstance(output_before_fusion, dict): + output_before_fusion = output_before_fusion.to_tuple()[0] + + # Fuse projections + model.fuse_qkv_projections() + + # Verify fusion occurred by checking for fused attributes + has_fused_projections = False + for module in model.modules(): + if isinstance(module, AttentionModuleMixin): + if hasattr(module, "to_qkv") or hasattr(module, "to_kv"): + has_fused_projections = True + assert module.fused_projections, "fused_projections flag should be True" + break + + if has_fused_projections: + # Get output after fusion + with torch.no_grad(): + output_after_fusion = model(**inputs_dict) + if isinstance(output_after_fusion, dict): + output_after_fusion = output_after_fusion.to_tuple()[0] + + # Verify outputs match + assert torch.allclose(output_before_fusion, output_after_fusion, atol=self.base_precision), ( + "Output should not change after fusing projections" + ) + + # Unfuse projections + model.unfuse_qkv_projections() + + # Verify unfusion occurred + for module in model.modules(): + if isinstance(module, AttentionModuleMixin): + assert not hasattr(module, "to_qkv"), "to_qkv should be removed after unfusing" + assert not hasattr(module, "to_kv"), "to_kv should be removed after unfusing" + assert not module.fused_projections, "fused_projections flag should be False" + + # Get output after unfusion + with torch.no_grad(): + output_after_unfusion = model(**inputs_dict) + if isinstance(output_after_unfusion, dict): + output_after_unfusion = output_after_unfusion.to_tuple()[0] + + # Verify outputs still match + assert torch.allclose(output_before_fusion, output_after_unfusion, atol=self.base_precision), ( + "Output should match original after unfusing projections" + ) + + def test_get_set_processor(self): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + model.to(torch_device) + + # Check if model has attention processors + if not hasattr(model, "attn_processors"): + pytest.skip("Model does not have attention processors.") + + # Test getting processors + processors = model.attn_processors + assert isinstance(processors, dict), "attn_processors should return a dict" + assert len(processors) > 0, "Model should have at least one attention processor" + + # Test that all processors can be retrieved via get_processor + for module in model.modules(): + if isinstance(module, AttentionModuleMixin): + processor = module.get_processor() + assert processor is not None, "get_processor should return a processor" + + # Test setting a new processor + new_processor = AttnProcessor() + module.set_processor(new_processor) + retrieved_processor = module.get_processor() + assert retrieved_processor is new_processor, "Retrieved processor should be the same as the one set" + + def test_attention_processor_dict(self): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + model.to(torch_device) + + if not hasattr(model, "set_attn_processor"): + pytest.skip("Model does not support setting attention processors.") + + # Get current processors + current_processors = model.attn_processors + + # Create a dict of new processors + new_processors = {key: AttnProcessor() for key in current_processors.keys()} + + # Set processors using dict + model.set_attn_processor(new_processors) + + # Verify all processors were set + updated_processors = model.attn_processors + for key in current_processors.keys(): + assert type(updated_processors[key]) == AttnProcessor, f"Processor {key} should be AttnProcessor" + + def test_attention_processor_count_mismatch_raises_error(self): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + model.to(torch_device) + + if not hasattr(model, "set_attn_processor"): + pytest.skip("Model does not support setting attention processors.") + + # Get current processors + current_processors = model.attn_processors + + # Create a dict with wrong number of processors + wrong_processors = {list(current_processors.keys())[0]: AttnProcessor()} + + # Verify error is raised + with pytest.raises(ValueError) as exc_info: + model.set_attn_processor(wrong_processors) + + assert "number of processors" in str(exc_info.value).lower(), "Error should mention processor count mismatch" + + +def _context_parallel_worker(rank, world_size, model_class, init_dict, cp_dict, inputs_dict, result_queue): + try: + # Setup distributed environment + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = "12355" + + torch.distributed.init_process_group( + backend="nccl", + init_method="env://", + world_size=world_size, + rank=rank, + ) + torch.cuda.set_device(rank) + device = torch.device(f"cuda:{rank}") + + model = model_class(**init_dict) + model.to(device) + model.eval() + + inputs_on_device = {} + for key, value in inputs_dict.items(): + if isinstance(value, torch.Tensor): + inputs_on_device[key] = value.to(device) + else: + inputs_on_device[key] = value + + cp_config = ContextParallelConfig(**cp_dict) + model.enable_parallelism(config=cp_config) + + with torch.no_grad(): + output = model(**inputs_on_device) + if isinstance(output, dict): + output = output.to_tuple()[0] + + if rank == 0: + result_queue.put(("success", output.shape)) + + except Exception as e: + if rank == 0: + result_queue.put(("error", str(e))) + finally: + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +@is_context_parallel +@require_torch_multi_accelerator +class ContextParallelTesterMixin: + base_precision = 1e-3 + + @pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"]) + def test_context_parallel_inference(self, cp_type): + if not torch.distributed.is_available(): + pytest.skip("torch.distributed is not available.") + + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + pytest.skip("Context parallel requires at least 2 CUDA devices.") + + if not hasattr(self.model_class, "_cp_plan") or self.model_class._cp_plan is None: + pytest.skip("Model does not have a _cp_plan defined for context parallel inference.") + + world_size = 2 + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + cp_dict = {cp_type: world_size} + + ctx = mp.get_context("spawn") + result_queue = ctx.Queue() + + mp.spawn( + _context_parallel_worker, + args=(world_size, self.model_class, init_dict, cp_dict, inputs_dict, result_queue), + nprocs=world_size, + join=True, + ) + + status, result = result_queue.get(timeout=60) + assert status == "success", f"Context parallel inference failed: {result}" diff --git a/tests/models/testing_utils/common.py b/tests/models/testing_utils/common.py new file mode 100644 index 000000000000..9f4ae271f97f --- /dev/null +++ b/tests/models/testing_utils/common.py @@ -0,0 +1,600 @@ +# 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 json +import os +import tempfile +from collections import defaultdict + +import pytest +import torch +import torch.nn as nn +from accelerate.utils.modeling import _get_proper_dtype, compute_module_sizes, dtype_byte_size + +from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME, _add_variant, logging +from diffusers.utils.testing_utils import require_accelerator, require_torch_multi_accelerator + +from ...testing_utils import CaptureLogger, torch_device + + +def named_persistent_module_tensors( + module: nn.Module, + recurse: bool = False, +): + """ + A helper function that gathers all the tensors (parameters + persistent buffers) of a given module. + + Args: + module (`torch.nn.Module`): + The module we want the tensors on. + recurse (`bool`, *optional`, defaults to `False`): + Whether or not to go look in every submodule or just return the direct parameters and buffers. + """ + yield from module.named_parameters(recurse=recurse) + + for named_buffer in module.named_buffers(recurse=recurse): + name, _ = named_buffer + # Get parent by splitting on dots and traversing the model + parent = module + if "." in name: + parent_name = name.rsplit(".", 1)[0] + for part in parent_name.split("."): + parent = getattr(parent, part) + name = name.split(".")[-1] + if name not in parent._non_persistent_buffers_set: + yield named_buffer + + +def compute_module_persistent_sizes( + model: nn.Module, + dtype: str | torch.device | None = None, + special_dtypes: dict[str, str | torch.device] | None = None, +): + """ + Compute the size of each submodule of a given model (parameters + persistent buffers). + """ + if dtype is not None: + dtype = _get_proper_dtype(dtype) + dtype_size = dtype_byte_size(dtype) + if special_dtypes is not None: + special_dtypes = {key: _get_proper_dtype(dtyp) for key, dtyp in special_dtypes.items()} + special_dtypes_size = {key: dtype_byte_size(dtyp) for key, dtyp in special_dtypes.items()} + module_sizes = defaultdict(int) + + module_list = [] + + module_list = named_persistent_module_tensors(model, recurse=True) + + for name, tensor in module_list: + if special_dtypes is not None and name in special_dtypes: + size = tensor.numel() * special_dtypes_size[name] + elif dtype is None: + size = tensor.numel() * dtype_byte_size(tensor.dtype) + elif str(tensor.dtype).startswith(("torch.uint", "torch.int", "torch.bool")): + # According to the code in set_module_tensor_to_device, these types won't be converted + # so use their original size here + size = tensor.numel() * dtype_byte_size(tensor.dtype) + else: + size = tensor.numel() * min(dtype_size, dtype_byte_size(tensor.dtype)) + name_parts = name.split(".") + for idx in range(len(name_parts) + 1): + module_sizes[".".join(name_parts[:idx])] += size + + return module_sizes + + +def calculate_expected_num_shards(index_map_path): + """ + Calculate expected number of shards from index file. + + Args: + index_map_path: Path to the sharded checkpoint index file + + Returns: + int: Expected number of shards + """ + with open(index_map_path) as f: + weight_map_dict = json.load(f)["weight_map"] + first_key = list(weight_map_dict.keys())[0] + weight_loc = weight_map_dict[first_key] # e.g., diffusion_pytorch_model-00001-of-00002.safetensors + expected_num_shards = int(weight_loc.split("-")[-1].split(".")[0]) + return expected_num_shards + + +def check_device_map_is_respected(model, device_map): + for param_name, param in model.named_parameters(): + # Find device in device_map + while len(param_name) > 0 and param_name not in device_map: + param_name = ".".join(param_name.split(".")[:-1]) + if param_name not in device_map: + raise ValueError("device map is incomplete, it does not contain any device for `param_name`.") + + param_device = device_map[param_name] + if param_device in ["cpu", "disk"]: + assert param.device == torch.device("meta"), f"Expected device 'meta' for {param_name}, got {param.device}" + else: + assert param.device == torch.device(param_device), ( + f"Expected device {param_device} for {param_name}, got {param.device}" + ) + + +class ModelTesterMixin: + """ + Base mixin class for model testing with common test methods. + + Expected class attributes to be set by subclasses: + - model_class: The model class to test + - main_input_name: Name of the main input tensor (e.g., "sample", "hidden_states") + - base_precision: Default tolerance for floating point comparisons (default: 1e-3) + + Expected methods to be implemented by subclasses: + - 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 + """ + + model_class = None + base_precision = 1e-3 + model_split_percents = [0.5, 0.7] + + def get_init_dict(self): + raise NotImplementedError("get_init_dict must be implemented by subclasses. ") + + def get_dummy_inputs(self): + raise NotImplementedError("get_dummy_inputs must be implemented by subclasses. It should return inputs_dict.") + + def test_from_save_pretrained(self, expected_max_diff=5e-5): + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + with tempfile.TemporaryDirectory() as tmpdirname: + model.save_pretrained(tmpdirname) + new_model = self.model_class.from_pretrained(tmpdirname) + new_model.to(torch_device) + + # check if all parameters shape are the same + for param_name in model.state_dict().keys(): + param_1 = model.state_dict()[param_name] + param_2 = new_model.state_dict()[param_name] + assert param_1.shape == param_2.shape, ( + f"Parameter shape mismatch for {param_name}. Original: {param_1.shape}, loaded: {param_2.shape}" + ) + + with torch.no_grad(): + image = model(**self.get_dummy_inputs()) + + if isinstance(image, dict): + image = image.to_tuple()[0] + + new_image = new_model(**self.get_dummy_inputs()) + + if isinstance(new_image, dict): + new_image = new_image.to_tuple()[0] + + max_diff = (image - new_image).abs().max().item() + assert max_diff <= expected_max_diff, ( + f"Models give different forward passes. Max diff: {max_diff}, expected: {expected_max_diff}" + ) + + def test_from_save_pretrained_variant(self, expected_max_diff=5e-5): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + with tempfile.TemporaryDirectory() as tmpdirname: + model.save_pretrained(tmpdirname, variant="fp16") + new_model = self.model_class.from_pretrained(tmpdirname, variant="fp16") + + # non-variant cannot be loaded + with pytest.raises(OSError) as exc_info: + self.model_class.from_pretrained(tmpdirname) + + # make sure that error message states what keys are missing + assert "Error no file named diffusion_pytorch_model.bin found in directory" in str(exc_info.value) + + new_model.to(torch_device) + + with torch.no_grad(): + image = model(**self.get_dummy_inputs()) + if isinstance(image, dict): + image = image.to_tuple()[0] + + new_image = new_model(**self.get_dummy_inputs()) + + if isinstance(new_image, dict): + new_image = new_image.to_tuple()[0] + + max_diff = (image - new_image).abs().max().item() + assert max_diff <= expected_max_diff, ( + f"Models give different forward passes. Max diff: {max_diff}, expected: {expected_max_diff}" + ) + + def test_from_save_pretrained_dtype(self): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + if torch_device == "mps" and dtype == torch.bfloat16: + continue + with tempfile.TemporaryDirectory() as tmpdirname: + model.to(dtype) + model.save_pretrained(tmpdirname) + new_model = self.model_class.from_pretrained(tmpdirname, low_cpu_mem_usage=True, torch_dtype=dtype) + assert new_model.dtype == dtype + if ( + hasattr(self.model_class, "_keep_in_fp32_modules") + and self.model_class._keep_in_fp32_modules is None + ): + # When loading without accelerate dtype == torch.float32 if _keep_in_fp32_modules is not None + new_model = self.model_class.from_pretrained( + tmpdirname, low_cpu_mem_usage=False, torch_dtype=dtype + ) + assert new_model.dtype == dtype + + def test_determinism(self, expected_max_diff=1e-5): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + with torch.no_grad(): + first = model(**self.get_dummy_inputs()) + if isinstance(first, dict): + first = first.to_tuple()[0] + + second = model(**self.get_dummy_inputs()) + if isinstance(second, dict): + second = second.to_tuple()[0] + + # Remove NaN values and compute max difference + first_flat = first.flatten() + second_flat = second.flatten() + + # Filter out NaN values + mask = ~(torch.isnan(first_flat) | torch.isnan(second_flat)) + first_filtered = first_flat[mask] + second_filtered = second_flat[mask] + + max_diff = torch.abs(first_filtered - second_filtered).max().item() + assert max_diff <= expected_max_diff, ( + f"Model outputs are not deterministic. Max diff: {max_diff}, expected: {expected_max_diff}" + ) + + def test_output(self, expected_output_shape=None): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + inputs_dict = self.get_dummy_inputs() + with torch.no_grad(): + output = model(**inputs_dict) + + if isinstance(output, dict): + output = output.to_tuple()[0] + + assert output is not None, "Model output is None" + assert output[0].shape == expected_output_shape or self.output_shape, ( + f"Output shape does not match expected. Expected {expected_output_shape}, got {output.shape}" + ) + + def test_outputs_equivalence(self): + def set_nan_tensor_to_zero(t): + # Temporary fallback until `aten::_index_put_impl_` is implemented in mps + # Track progress in https://github.com/pytorch/pytorch/issues/77764 + device = t.device + if device.type == "mps": + t = t.to("cpu") + t[t != t] = 0 + return t.to(device) + + def recursive_check(tuple_object, dict_object): + if isinstance(tuple_object, (list, tuple)): + for tuple_iterable_value, dict_iterable_value in zip(tuple_object, dict_object.values()): + recursive_check(tuple_iterable_value, dict_iterable_value) + elif isinstance(tuple_object, dict): + for tuple_iterable_value, dict_iterable_value in zip(tuple_object.values(), dict_object.values()): + recursive_check(tuple_iterable_value, dict_iterable_value) + elif tuple_object is None: + return + else: + assert torch.allclose( + set_nan_tensor_to_zero(tuple_object), set_nan_tensor_to_zero(dict_object), atol=1e-5 + ), ( + "Tuple and dict output are not equal. Difference:" + f" {torch.max(torch.abs(tuple_object - dict_object))}. Tuple has `nan`:" + f" {torch.isnan(tuple_object).any()} and `inf`: {torch.isinf(tuple_object)}. Dict has" + f" `nan`: {torch.isnan(dict_object).any()} and `inf`: {torch.isinf(dict_object)}." + ) + + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + with torch.no_grad(): + outputs_dict = model(**self.get_dummy_inputs()) + outputs_tuple = model(**self.get_dummy_inputs(), return_dict=False) + + recursive_check(outputs_tuple, outputs_dict) + + def test_getattr_is_correct(self): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + + # save some things to test + model.dummy_attribute = 5 + model.register_to_config(test_attribute=5) + + logger = logging.get_logger("diffusers.models.modeling_utils") + # 30 for warning + logger.setLevel(30) + with CaptureLogger(logger) as cap_logger: + assert hasattr(model, "dummy_attribute") + assert getattr(model, "dummy_attribute") == 5 + assert model.dummy_attribute == 5 + + # no warning should be thrown + assert cap_logger.out == "" + + logger = logging.get_logger("diffusers.models.modeling_utils") + # 30 for warning + logger.setLevel(30) + with CaptureLogger(logger) as cap_logger: + assert hasattr(model, "save_pretrained") + fn = model.save_pretrained + fn_1 = getattr(model, "save_pretrained") + + assert fn == fn_1 + # no warning should be thrown + assert cap_logger.out == "" + + # warning should be thrown for config attributes accessed directly + with pytest.warns(FutureWarning): + assert model.test_attribute == 5 + + with pytest.warns(FutureWarning): + assert getattr(model, "test_attribute") == 5 + + with pytest.raises(AttributeError) as error: + model.does_not_exist + + assert str(error.value) == f"'{type(model).__name__}' object has no attribute 'does_not_exist'" + + @require_accelerator + @pytest.mark.skipif( + torch_device not in ["cuda", "xpu"], + reason="float16 and bfloat16 can only be used with an accelerator", + ) + def test_keep_in_fp32_modules(self): + model = self.model_class(**self.get_init_dict()) + fp32_modules = model._keep_in_fp32_modules + + if fp32_modules is None or len(fp32_modules) == 0: + pytest.skip("Model does not have _keep_in_fp32_modules defined.") + + # Test with float16 + model.to(torch_device) + model.to(torch.float16) + + for name, param in model.named_parameters(): + if any(module_to_keep_in_fp32 in name.split(".") for module_to_keep_in_fp32 in fp32_modules): + assert param.dtype == torch.float32, f"Parameter {name} should be float32 but got {param.dtype}" + else: + assert param.dtype == torch.float16, f"Parameter {name} should be float16 but got {param.dtype}" + + @require_accelerator + @pytest.mark.skipif( + torch_device not in ["cuda", "xpu"], + reason="float16 and bfloat16 can only be use for inference with an accelerator", + ) + def test_from_save_pretrained_float16_bfloat16(self): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + fp32_modules = model._keep_in_fp32_modules + + with tempfile.TemporaryDirectory() as tmp_dir: + for torch_dtype in [torch.bfloat16, torch.float16]: + model.to(torch_dtype).save_pretrained(tmp_dir) + model_loaded = self.model_class.from_pretrained(tmp_dir, torch_dtype=torch_dtype).to(torch_device) + + for name, param in model_loaded.named_parameters(): + if any(module_to_keep_in_fp32 in name.split(".") for module_to_keep_in_fp32 in fp32_modules): + assert param.data.dtype == torch.float32 + else: + assert param.data.dtype == torch_dtype + + with torch.no_grad(): + output = model(**self.get_dummy_inputs()) + output_loaded = model_loaded(**self.get_dummy_inputs()) + + assert torch.allclose(output, output_loaded, atol=1e-4), ( + f"Loaded model output differs for {torch_dtype}" + ) + + @require_accelerator + def test_sharded_checkpoints(self): + torch.manual_seed(0) + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + model = model.to(torch_device) + + base_output = model(**inputs_dict) + + model_size = compute_module_persistent_sizes(model)[""] + max_shard_size = int((model_size * 0.75) / (2**10)) # Convert to KB as these test models are small + with tempfile.TemporaryDirectory() as tmp_dir: + model.cpu().save_pretrained(tmp_dir, max_shard_size=f"{max_shard_size}KB") + assert os.path.exists(os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)), "Index file should exist" + + # Check if the right number of shards exists + expected_num_shards = calculate_expected_num_shards(os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)) + actual_num_shards = len([file for file in os.listdir(tmp_dir) if file.endswith(".safetensors")]) + assert actual_num_shards == expected_num_shards, ( + f"Expected {expected_num_shards} shards, got {actual_num_shards}" + ) + + new_model = self.model_class.from_pretrained(tmp_dir).eval() + new_model = new_model.to(torch_device) + + torch.manual_seed(0) + inputs_dict_new = self.get_dummy_inputs() + new_output = new_model(**inputs_dict_new) + + assert torch.allclose(base_output[0], new_output[0], atol=1e-5), ( + "Output should match after sharded save/load" + ) + + @require_accelerator + def test_sharded_checkpoints_with_variant(self): + torch.manual_seed(0) + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + model = model.to(torch_device) + + base_output = model(**inputs_dict) + + model_size = compute_module_persistent_sizes(model)[""] + max_shard_size = int((model_size * 0.75) / (2**10)) # Convert to KB as these test models are small + variant = "fp16" + with tempfile.TemporaryDirectory() as tmp_dir: + model.cpu().save_pretrained(tmp_dir, max_shard_size=f"{max_shard_size}KB", variant=variant) + + index_filename = _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant) + assert os.path.exists(os.path.join(tmp_dir, index_filename)), ( + f"Variant index file {index_filename} should exist" + ) + + # Check if the right number of shards exists + expected_num_shards = calculate_expected_num_shards(os.path.join(tmp_dir, index_filename)) + actual_num_shards = len([file for file in os.listdir(tmp_dir) if file.endswith(".safetensors")]) + assert actual_num_shards == expected_num_shards, ( + f"Expected {expected_num_shards} shards, got {actual_num_shards}" + ) + + new_model = self.model_class.from_pretrained(tmp_dir, variant=variant).eval() + new_model = new_model.to(torch_device) + + torch.manual_seed(0) + inputs_dict_new = self.get_dummy_inputs() + new_output = new_model(**inputs_dict_new) + + assert torch.allclose(base_output[0], new_output[0], atol=1e-5), ( + "Output should match after variant sharded save/load" + ) + + def test_sharded_checkpoints_with_parallel_loading(self): + import time + + from diffusers.utils import constants + + torch.manual_seed(0) + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + model = model.to(torch_device) + + base_output = model(**inputs_dict) + + model_size = compute_module_persistent_sizes(model)[""] + max_shard_size = int((model_size * 0.75) / (2**10)) # Convert to KB as these test models are small + + # Save original values to restore after test + original_parallel_loading = constants.HF_ENABLE_PARALLEL_LOADING + original_parallel_workers = getattr(constants, "HF_PARALLEL_WORKERS", None) + + try: + with tempfile.TemporaryDirectory() as tmp_dir: + model.cpu().save_pretrained(tmp_dir, max_shard_size=f"{max_shard_size}KB") + assert os.path.exists(os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)), "Index file should exist" + + # Check if the right number of shards exists + expected_num_shards = calculate_expected_num_shards(os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)) + actual_num_shards = len([file for file in os.listdir(tmp_dir) if file.endswith(".safetensors")]) + assert actual_num_shards == expected_num_shards, ( + f"Expected {expected_num_shards} shards, got {actual_num_shards}" + ) + + # Load without parallel loading + constants.HF_ENABLE_PARALLEL_LOADING = False + start_time = time.time() + model_sequential = self.model_class.from_pretrained(tmp_dir).eval() + sequential_load_time = time.time() - start_time + model_sequential = model_sequential.to(torch_device) + + torch.manual_seed(0) + + # Load with parallel loading + constants.HF_ENABLE_PARALLEL_LOADING = True + constants.DEFAULT_HF_PARALLEL_LOADING_WORKERS = 2 + + start_time = time.time() + model_parallel = self.model_class.from_pretrained(tmp_dir).eval() + parallel_load_time = time.time() - start_time + model_parallel = model_parallel.to(torch_device) + + torch.manual_seed(0) + inputs_dict_parallel = self.get_dummy_inputs() + output_parallel = model_parallel(**inputs_dict_parallel) + + assert torch.allclose(base_output[0], output_parallel[0], atol=1e-5), ( + "Output should match with parallel loading" + ) + + # Verify parallel loading is faster or at least not significantly slower + assert parallel_load_time < sequential_load_time, ( + f"Parallel loading took {parallel_load_time:.4f}s, sequential took {sequential_load_time:.4f}s" + ) + finally: + # Restore original values + constants.HF_ENABLE_PARALLEL_LOADING = original_parallel_loading + if original_parallel_workers is not None: + constants.HF_PARALLEL_WORKERS = original_parallel_workers + + @require_torch_multi_accelerator + def test_model_parallelism(self): + if self.model_class._no_split_modules is None: + pytest.skip("Test not supported for this model as `_no_split_modules` is not set.") + + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + + model = model.to(torch_device) + + torch.manual_seed(0) + base_output = model(**inputs_dict) + + model_size = compute_module_sizes(model)[""] + max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents] + + with tempfile.TemporaryDirectory() as tmp_dir: + model.cpu().save_pretrained(tmp_dir) + + for max_size in max_gpu_sizes: + max_memory = {0: max_size, 1: model_size * 2, "cpu": model_size * 2} + new_model = self.model_class.from_pretrained(tmp_dir, device_map="auto", max_memory=max_memory) + # Making sure part of the model will be on GPU 0 and GPU 1 + assert set(new_model.hf_device_map.values()) == {0, 1}, "Model should be split across GPUs" + + check_device_map_is_respected(new_model, new_model.hf_device_map) + + torch.manual_seed(0) + new_output = new_model(**inputs_dict) + + assert torch.allclose(base_output[0], new_output[0], atol=1e-5), ( + "Output should match with model parallelism" + ) diff --git a/tests/models/testing_utils/compile.py b/tests/models/testing_utils/compile.py new file mode 100644 index 000000000000..8ff4c097b404 --- /dev/null +++ b/tests/models/testing_utils/compile.py @@ -0,0 +1,162 @@ +# 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 os +import tempfile + +import pytest +import torch + +from ...testing_utils import ( + backend_empty_cache, + is_torch_compile, + require_accelerator, + require_torch_version_greater, + torch_device, +) + + +@is_torch_compile +@require_accelerator +@require_torch_version_greater("2.7.1") +class TorchCompileTesterMixin: + """ + Mixin class for testing torch.compile functionality on models. + + Expected class attributes to be set by subclasses: + - model_class: The model class to test + - different_shapes_for_compilation: Optional list of (height, width) tuples for dynamic shape testing + + Expected methods to be implemented by subclasses: + - 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: compile + Use `pytest -m "not compile"` to skip these tests + """ + + different_shapes_for_compilation = None + + def setup_method(self): + torch.compiler.reset() + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + torch.compiler.reset() + gc.collect() + backend_empty_cache(torch_device) + + def test_torch_compile_recompilation_and_graph_break(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict).to(torch_device) + model.eval() + model = torch.compile(model, fullgraph=True) + + with ( + torch._inductor.utils.fresh_inductor_cache(), + torch._dynamo.config.patch(error_on_recompile=True), + torch.no_grad(), + ): + _ = model(**inputs_dict) + _ = model(**inputs_dict) + + def test_torch_compile_repeated_blocks(self): + if self.model_class._repeated_blocks is None: + pytest.skip("Skipping test as the model class doesn't have `_repeated_blocks` set.") + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict).to(torch_device) + model.eval() + model.compile_repeated_blocks(fullgraph=True) + + recompile_limit = 1 + if self.model_class.__name__ == "UNet2DConditionModel": + recompile_limit = 2 + + with ( + torch._inductor.utils.fresh_inductor_cache(), + torch._dynamo.config.patch(recompile_limit=recompile_limit), + torch.no_grad(), + ): + _ = model(**inputs_dict) + _ = model(**inputs_dict) + + def test_compile_with_group_offloading(self): + if not self.model_class._supports_group_offloading: + pytest.skip("Model does not support group offloading.") + + torch._dynamo.config.cache_size_limit = 10000 + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict) + model.eval() + + group_offload_kwargs = { + "onload_device": torch_device, + "offload_device": "cpu", + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": True, + "non_blocking": True, + } + model.enable_group_offload(**group_offload_kwargs) + model.compile() + + with torch.no_grad(): + _ = model(**inputs_dict) + _ = model(**inputs_dict) + + def test_compile_on_different_shapes(self): + if self.different_shapes_for_compilation is None: + pytest.skip(f"Skipping as `different_shapes_for_compilation` is not set for {self.__class__.__name__}.") + torch.fx.experimental._config.use_duck_shape = False + + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + model = torch.compile(model, fullgraph=True, dynamic=True) + + for height, width in self.different_shapes_for_compilation: + with torch._dynamo.config.patch(error_on_recompile=True), torch.no_grad(): + inputs_dict = self.get_dummy_inputs(height=height, width=width) + _ = model(**inputs_dict) + + def test_compile_works_with_aot(self): + from torch._inductor.package import load_package + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict).to(torch_device) + exported_model = torch.export.export(model, args=(), kwargs=inputs_dict) + + with tempfile.TemporaryDirectory() as tmpdir: + package_path = os.path.join(tmpdir, f"{self.model_class.__name__}.pt2") + _ = torch._inductor.aoti_compile_and_package(exported_model, package_path=package_path) + assert os.path.exists(package_path), f"Package file not created at {package_path}" + loaded_binary = load_package(package_path, run_single_threaded=True) + + model.forward = loaded_binary + + with torch.no_grad(): + _ = model(**inputs_dict) + _ = model(**inputs_dict) diff --git a/tests/models/testing_utils/hub.py b/tests/models/testing_utils/hub.py new file mode 100644 index 000000000000..40d8777c33b1 --- /dev/null +++ b/tests/models/testing_utils/hub.py @@ -0,0 +1,109 @@ +# 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 tempfile +import uuid + +import pytest +import torch +from huggingface_hub.utils import ModelCard, delete_repo, is_jinja_available + +from ...others.test_utils import TOKEN, USER, is_staging_test + + +@is_staging_test +class ModelPushToHubTesterMixin: + """ + Mixin class for testing push_to_hub functionality on models. + + Expected class attributes to be set by subclasses: + - model_class: The model class to test + + Expected methods to be implemented by subclasses: + - get_init_dict(): Returns dict of arguments to initialize the model + """ + + identifier = uuid.uuid4() + repo_id = f"test-model-{identifier}" + org_repo_id = f"valid_org/{repo_id}-org" + + def test_push_to_hub(self): + """Test pushing model to hub and loading it back.""" + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + model.push_to_hub(self.repo_id, token=TOKEN) + + new_model = self.model_class.from_pretrained(f"{USER}/{self.repo_id}") + for p1, p2 in zip(model.parameters(), new_model.parameters()): + assert torch.equal(p1, p2), "Parameters don't match after push_to_hub and from_pretrained" + + # Reset repo + delete_repo(token=TOKEN, repo_id=self.repo_id) + + # Push to hub via save_pretrained + with tempfile.TemporaryDirectory() as tmp_dir: + model.save_pretrained(tmp_dir, repo_id=self.repo_id, push_to_hub=True, token=TOKEN) + + new_model = self.model_class.from_pretrained(f"{USER}/{self.repo_id}") + for p1, p2 in zip(model.parameters(), new_model.parameters()): + assert torch.equal(p1, p2), ( + "Parameters don't match after save_pretrained with push_to_hub and from_pretrained" + ) + + # Reset repo + delete_repo(self.repo_id, token=TOKEN) + + def test_push_to_hub_in_organization(self): + """Test pushing model to hub in organization namespace.""" + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + model.push_to_hub(self.org_repo_id, token=TOKEN) + + new_model = self.model_class.from_pretrained(self.org_repo_id) + for p1, p2 in zip(model.parameters(), new_model.parameters()): + assert torch.equal(p1, p2), "Parameters don't match after push_to_hub to org and from_pretrained" + + # Reset repo + delete_repo(token=TOKEN, repo_id=self.org_repo_id) + + # Push to hub via save_pretrained + with tempfile.TemporaryDirectory() as tmp_dir: + model.save_pretrained(tmp_dir, push_to_hub=True, token=TOKEN, repo_id=self.org_repo_id) + + new_model = self.model_class.from_pretrained(self.org_repo_id) + for p1, p2 in zip(model.parameters(), new_model.parameters()): + assert torch.equal(p1, p2), ( + "Parameters don't match after save_pretrained with push_to_hub to org and from_pretrained" + ) + + # Reset repo + delete_repo(self.org_repo_id, token=TOKEN) + + def test_push_to_hub_library_name(self): + """Test that library_name in model card is set to 'diffusers'.""" + if not is_jinja_available(): + pytest.skip("Model card tests cannot be performed without Jinja installed.") + + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + model.push_to_hub(self.repo_id, token=TOKEN) + + model_card = ModelCard.load(f"{USER}/{self.repo_id}", token=TOKEN).data + assert model_card.library_name == "diffusers", ( + f"Expected library_name 'diffusers', got {model_card.library_name}" + ) + + # Reset repo + delete_repo(self.repo_id, token=TOKEN) diff --git a/tests/models/testing_utils/ip_adapter.py b/tests/models/testing_utils/ip_adapter.py new file mode 100644 index 000000000000..13e141869c3a --- /dev/null +++ b/tests/models/testing_utils/ip_adapter.py @@ -0,0 +1,178 @@ +# 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 pytest +import torch + +from diffusers.models.attention_processor import IPAdapterAttnProcessor + +from ...testing_utils import is_ip_adapter, torch_device + + +def create_ip_adapter_state_dict(model): + """ + Create a dummy IP Adapter state dict for testing. + + Args: + model: The model to create IP adapter weights for + + Returns: + dict: IP adapter state dict with to_k_ip and to_v_ip weights + """ + ip_state_dict = {} + key_id = 1 + + for name in model.attn_processors.keys(): + # Skip self-attention processors + cross_attention_dim = getattr(model.config, "cross_attention_dim", None) + if cross_attention_dim is None: + continue + + # Get hidden size based on model architecture + hidden_size = getattr(model.config, "hidden_size", cross_attention_dim) + + # Create IP adapter processor to get state dict structure + sd = IPAdapterAttnProcessor( + hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, scale=1.0 + ).state_dict() + + ip_state_dict.update( + { + f"{key_id}.to_k_ip.weight": sd["to_k_ip.0.weight"], + f"{key_id}.to_v_ip.weight": sd["to_v_ip.0.weight"], + } + ) + key_id += 2 + + return {"ip_adapter": ip_state_dict} + + +def check_if_ip_adapter_correctly_set(model, processor_cls) -> bool: + """ + Check if IP Adapter processors are correctly set in the model. + + Args: + model: The model to check + + Returns: + bool: True if IP Adapter is correctly set, False otherwise + """ + for module in model.attn_processors.values(): + if isinstance(module, processor_cls): + return True + return False + + +@is_ip_adapter +class IPAdapterTesterMixin: + """ + Mixin class for testing IP Adapter functionality on models. + + Expected class attributes to be set by subclasses: + - model_class: The model class to test + + Expected methods to be implemented by subclasses: + - 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: ip_adapter + Use `pytest -m "not ip_adapter"` to skip these tests + """ + + ip_adapter_processor_cls = None + + def create_ip_adapter_state_dict(self, model): + raise NotImplementedError("child class must implement method to create IPAdapter State Dict") + + def modify_inputs_for_ip_adapter(self, model, inputs_dict): + raise NotImplementedError("child class must implement method to create IPAdapter model inputs") + + def test_load_ip_adapter(self): + 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_adapter = model(**inputs_dict, return_dict=False)[0] + + ip_adapter_state_dict = self.create_ip_adapter_state_dict(model) + + model._load_ip_adapter_weights([ip_adapter_state_dict]) + assert check_if_ip_adapter_correctly_set(model, self.ip_adapter_processor_cls), ( + "IP Adapter processors not set correctly" + ) + + inputs_dict_with_adapter = self.modify_inputs_for_ip_adapter(model, inputs_dict.copy()) + outputs_with_adapter = model(**inputs_dict_with_adapter, return_dict=False)[0] + + assert not torch.allclose(output_no_adapter, outputs_with_adapter, atol=1e-4, rtol=1e-4), ( + "Output should differ with IP Adapter enabled" + ) + + @pytest.mark.skip( + reason="Setting IP Adapter scale is not defined at the model level. Enable this test after refactoring" + ) + def test_ip_adapter_scale(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + + ip_adapter_state_dict = self.create_ip_adapter_state_dict(model) + model._load_ip_adapter_weights([ip_adapter_state_dict]) + + inputs_dict_with_adapter = self.modify_inputs_for_ip_adapter(model, inputs_dict.copy()) + + # Test scale = 0.0 (no effect) + model.set_ip_adapter_scale(0.0) + torch.manual_seed(0) + output_scale_zero = model(**inputs_dict_with_adapter, return_dict=False)[0] + + # Test scale = 1.0 (full effect) + model.set_ip_adapter_scale(1.0) + torch.manual_seed(0) + output_scale_one = model(**inputs_dict_with_adapter, return_dict=False)[0] + + # Outputs should differ with different scales + assert not torch.allclose(output_scale_zero, output_scale_one, atol=1e-4, rtol=1e-4), ( + "Output should differ with different IP Adapter scales" + ) + + @pytest.mark.skip( + reason="Unloading IP Adapter is not defined at the model level. Enable this test after refactoring" + ) + def test_unload_ip_adapter(self): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + + # Save original processors + original_processors = {k: type(v).__name__ for k, v in model.attn_processors.items()} + + # Create and load IP adapter + ip_adapter_state_dict = self.create_ip_adapter_state_dict(model) + model._load_ip_adapter_weights([ip_adapter_state_dict]) + + assert check_if_ip_adapter_correctly_set(model, self.ip_adapter_processor_cls), "IP Adapter should be set" + + # Unload IP adapter + model.unload_ip_adapter() + + assert not check_if_ip_adapter_correctly_set(model, self.ip_adapter_processor_cls), ( + "IP Adapter should be unloaded" + ) + + # Verify processors are restored + current_processors = {k: type(v).__name__ for k, v in model.attn_processors.items()} + assert original_processors == current_processors, "Processors should be restored after unload" diff --git a/tests/models/testing_utils/lora.py b/tests/models/testing_utils/lora.py new file mode 100644 index 000000000000..6777c164f280 --- /dev/null +++ b/tests/models/testing_utils/lora.py @@ -0,0 +1,220 @@ +# 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 json +import os +import tempfile + +import pytest +import safetensors.torch +import torch + +from diffusers.utils.testing_utils import check_if_dicts_are_equal + +from ...testing_utils import is_lora, require_peft_backend, torch_device + + +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 class attributes to be set by subclasses: + - model_class: The model class to test + + Expected methods to be implemented by subclasses: + - 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): + from diffusers.loaders.peft import PeftAdapterMixin + + if not issubclass(self.model_class, PeftAdapterMixin): + pytest.skip(f"PEFT is not supported for this model ({self.model_class.__name__}).") + + def test_save_load_lora_adapter(self, 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" + ) + + with tempfile.TemporaryDirectory() as tmpdir: + model.save_lora_adapter(tmpdir) + assert os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")), ( + "LoRA weights file not created" + ) + + state_dict_loaded = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")) + + model.unload_lora() + assert not check_if_lora_correctly_set(model), "LoRA should be unloaded" + + model.load_lora_adapter(tmpdir, 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 torch.allclose(loaded_v, retrieved_v), 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 torch.allclose(outputs_with_lora, outputs_with_lora_2, atol=1e-4, rtol=1e-4), ( + "Outputs should match before and after save/load" + ) + + def test_lora_wrong_adapter_name_raises_error(self): + 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" + + with tempfile.TemporaryDirectory() as tmpdir: + wrong_name = "foo" + with pytest.raises(ValueError) as exc_info: + model.save_lora_adapter(tmpdir, 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, 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" + + with tempfile.TemporaryDirectory() as tmpdir: + model.save_lora_adapter(tmpdir) + model_file = os.path.join(tmpdir, "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(tmpdir, 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): + 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" + + with tempfile.TemporaryDirectory() as tmpdir: + model.save_lora_adapter(tmpdir) + model_file = os.path.join(tmpdir, "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(tmpdir, prefix=None, use_safetensors=True) + assert "`LoraConfig` class could not be instantiated" in str(exc_info.value) diff --git a/tests/models/testing_utils/memory.py b/tests/models/testing_utils/memory.py new file mode 100644 index 000000000000..6cdc72b004c9 --- /dev/null +++ b/tests/models/testing_utils/memory.py @@ -0,0 +1,473 @@ +# 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 glob +import inspect +import tempfile +from functools import wraps + +import pytest +import torch +from accelerate.utils.modeling import compute_module_sizes + +from diffusers.utils.testing_utils import _check_safetensors_serialization +from diffusers.utils.torch_utils import get_torch_cuda_device_capability + +from ...testing_utils import ( + backend_empty_cache, + backend_max_memory_allocated, + backend_reset_peak_memory_stats, + backend_synchronize, + is_cpu_offload, + is_group_offload, + is_memory, + require_accelerator, + torch_device, +) +from .common import check_device_map_is_respected + + +def cast_maybe_tensor_dtype(inputs_dict, from_dtype, to_dtype): + """Helper to cast tensor inputs from one dtype to another.""" + for key, value in inputs_dict.items(): + if isinstance(value, torch.Tensor) and value.dtype == from_dtype: + inputs_dict[key] = value.to(to_dtype) + return inputs_dict + + +def require_offload_support(func): + """ + Decorator to skip tests if model doesn't support offloading (requires _no_split_modules). + """ + + @wraps(func) + def wrapper(self, *args, **kwargs): + if self.model_class._no_split_modules is None: + pytest.skip("Test not supported for this model as `_no_split_modules` is not set.") + return func(self, *args, **kwargs) + + return wrapper + + +def require_group_offload_support(func): + """ + Decorator to skip tests if model doesn't support group offloading. + """ + + @wraps(func) + def wrapper(self, *args, **kwargs): + if not self.model_class._supports_group_offloading: + pytest.skip("Model does not support group offloading.") + return func(self, *args, **kwargs) + + return wrapper + + +@is_cpu_offload +class CPUOffloadTesterMixin: + """ + Mixin class for testing CPU offloading functionality. + + Expected class attributes to be set by subclasses: + - model_class: The model class to test + - model_split_percents: List of percentages for splitting model across devices + + Expected methods to be implemented by subclasses: + - 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: cpu_offload + Use `pytest -m "not cpu_offload"` to skip these tests + """ + + model_split_percents = [0.5, 0.7] + + @require_offload_support + def test_cpu_offload(self): + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + + model = model.to(torch_device) + + torch.manual_seed(0) + base_output = model(**inputs_dict) + + model_size = compute_module_sizes(model)[""] + # We test several splits of sizes to make sure it works + max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents] + with tempfile.TemporaryDirectory() as tmp_dir: + model.cpu().save_pretrained(tmp_dir) + + for max_size in max_gpu_sizes: + max_memory = {0: max_size, "cpu": model_size * 2} + new_model = self.model_class.from_pretrained(tmp_dir, device_map="auto", max_memory=max_memory) + # Making sure part of the model will actually end up offloaded + assert set(new_model.hf_device_map.values()) == {0, "cpu"}, "Model should be split between GPU and CPU" + + check_device_map_is_respected(new_model, new_model.hf_device_map) + torch.manual_seed(0) + new_output = new_model(**inputs_dict) + + assert torch.allclose(base_output[0], new_output[0], atol=1e-5), ( + "Output should match with CPU offloading" + ) + + @require_offload_support + def test_disk_offload_without_safetensors(self): + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + + model = model.to(torch_device) + + torch.manual_seed(0) + base_output = model(**inputs_dict) + + model_size = compute_module_sizes(model)[""] + max_size = int(self.model_split_percents[0] * model_size) + # Force disk offload by setting very small CPU memory + max_memory = {0: max_size, "cpu": int(0.1 * max_size)} + + with tempfile.TemporaryDirectory() as tmp_dir: + model.cpu().save_pretrained(tmp_dir, safe_serialization=False) + # This errors out because it's missing an offload folder + with pytest.raises(ValueError): + new_model = self.model_class.from_pretrained(tmp_dir, device_map="auto", max_memory=max_memory) + + new_model = self.model_class.from_pretrained( + tmp_dir, device_map="auto", max_memory=max_memory, offload_folder=tmp_dir + ) + + check_device_map_is_respected(new_model, new_model.hf_device_map) + torch.manual_seed(0) + new_output = new_model(**inputs_dict) + + assert torch.allclose(base_output[0], new_output[0], atol=1e-5), "Output should match with disk offloading" + + @require_offload_support + def test_disk_offload_with_safetensors(self): + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + + model = model.to(torch_device) + + torch.manual_seed(0) + base_output = model(**inputs_dict) + + model_size = compute_module_sizes(model)[""] + with tempfile.TemporaryDirectory() as tmp_dir: + model.cpu().save_pretrained(tmp_dir) + + max_size = int(self.model_split_percents[0] * model_size) + max_memory = {0: max_size, "cpu": max_size} + new_model = self.model_class.from_pretrained( + tmp_dir, device_map="auto", offload_folder=tmp_dir, max_memory=max_memory + ) + + check_device_map_is_respected(new_model, new_model.hf_device_map) + torch.manual_seed(0) + new_output = new_model(**inputs_dict) + + assert torch.allclose(base_output[0], new_output[0], atol=1e-5), ( + "Output should match with disk offloading (safetensors)" + ) + + +@is_group_offload +class GroupOffloadTesterMixin: + """ + Mixin class for testing group offloading functionality. + + Expected class attributes to be set by subclasses: + - model_class: The model class to test + + Expected methods to be implemented by subclasses: + - 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: group_offload + Use `pytest -m "not group_offload"` to skip these tests + """ + + @require_group_offload_support + def test_group_offloading(self, record_stream=False): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + torch.manual_seed(0) + + @torch.no_grad() + def run_forward(model): + assert all( + module._diffusers_hook.get_hook("group_offloading") is not None + for module in model.modules() + if hasattr(module, "_diffusers_hook") + ), "Group offloading hook should be set" + model.eval() + return model(**inputs_dict)[0] + + model = self.model_class(**init_dict) + + model.to(torch_device) + output_without_group_offloading = run_forward(model) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1) + output_with_group_offloading1 = run_forward(model) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1, non_blocking=True) + output_with_group_offloading2 = run_forward(model) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.enable_group_offload(torch_device, offload_type="leaf_level") + output_with_group_offloading3 = run_forward(model) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.enable_group_offload( + torch_device, offload_type="leaf_level", use_stream=True, record_stream=record_stream + ) + output_with_group_offloading4 = run_forward(model) + + assert torch.allclose(output_without_group_offloading, output_with_group_offloading1, atol=1e-5), ( + "Output should match with block-level offloading" + ) + assert torch.allclose(output_without_group_offloading, output_with_group_offloading2, atol=1e-5), ( + "Output should match with non-blocking block-level offloading" + ) + assert torch.allclose(output_without_group_offloading, output_with_group_offloading3, atol=1e-5), ( + "Output should match with leaf-level offloading" + ) + assert torch.allclose(output_without_group_offloading, output_with_group_offloading4, atol=1e-5), ( + "Output should match with leaf-level offloading with stream" + ) + + @require_group_offload_support + @torch.no_grad() + def test_group_offloading_with_layerwise_casting(self, record_stream=False, offload_type="block_level"): + torch.manual_seed(0) + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict) + + model.to(torch_device) + model.eval() + _ = model(**inputs_dict)[0] + + torch.manual_seed(0) + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + storage_dtype, compute_dtype = torch.float16, torch.float32 + inputs_dict = cast_maybe_tensor_dtype(inputs_dict, torch.float32, compute_dtype) + model = self.model_class(**init_dict) + model.eval() + additional_kwargs = {} if offload_type == "leaf_level" else {"num_blocks_per_group": 1} + model.enable_group_offload( + torch_device, offload_type=offload_type, use_stream=True, record_stream=record_stream, **additional_kwargs + ) + model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype) + _ = model(**inputs_dict)[0] + + @require_group_offload_support + @torch.no_grad() + @torch.inference_mode() + def test_group_offloading_with_disk(self, offload_type="block_level", record_stream=False, atol=1e-5): + def _has_generator_arg(model): + sig = inspect.signature(model.forward) + params = sig.parameters + return "generator" in params + + def _run_forward(model, inputs_dict): + accepts_generator = _has_generator_arg(model) + if accepts_generator: + inputs_dict["generator"] = torch.manual_seed(0) + torch.manual_seed(0) + return model(**inputs_dict)[0] + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + torch.manual_seed(0) + model = self.model_class(**init_dict) + + model.eval() + model.to(torch_device) + output_without_group_offloading = _run_forward(model, inputs_dict) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.eval() + + num_blocks_per_group = None if offload_type == "leaf_level" else 1 + additional_kwargs = {} if offload_type == "leaf_level" else {"num_blocks_per_group": num_blocks_per_group} + with tempfile.TemporaryDirectory() as tmpdir: + model.enable_group_offload( + torch_device, + offload_type=offload_type, + offload_to_disk_path=tmpdir, + use_stream=True, + record_stream=record_stream, + **additional_kwargs, + ) + has_safetensors = glob.glob(f"{tmpdir}/*.safetensors") + assert has_safetensors, "No safetensors found in the directory." + + # For "leaf-level", there is a prefetching hook which makes this check a bit non-deterministic + # in nature. So, skip it. + if offload_type != "leaf_level": + is_correct, extra_files, missing_files = _check_safetensors_serialization( + module=model, + offload_to_disk_path=tmpdir, + offload_type=offload_type, + num_blocks_per_group=num_blocks_per_group, + ) + if not is_correct: + if extra_files: + raise ValueError(f"Found extra files: {', '.join(extra_files)}") + elif missing_files: + raise ValueError(f"Following files are missing: {', '.join(missing_files)}") + + output_with_group_offloading = _run_forward(model, inputs_dict) + assert torch.allclose(output_without_group_offloading, output_with_group_offloading, atol=atol), ( + "Output should match with disk-based group offloading" + ) + + +class LayerwiseCastingTesterMixin: + """ + Mixin class for testing layerwise dtype casting for memory optimization. + + Expected class attributes to be set by subclasses: + - model_class: The model class to test + + Expected methods to be implemented by subclasses: + - 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 + """ + + @torch.no_grad() + def test_layerwise_casting_memory(self): + MB_TOLERANCE = 0.2 + LEAST_COMPUTE_CAPABILITY = 8.0 + + def reset_memory_stats(): + gc.collect() + backend_synchronize(torch_device) + backend_empty_cache(torch_device) + backend_reset_peak_memory_stats(torch_device) + + def get_memory_usage(storage_dtype, compute_dtype): + torch.manual_seed(0) + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + inputs_dict = cast_maybe_tensor_dtype(inputs_dict, torch.float32, compute_dtype) + model = self.model_class(**config).eval() + model = model.to(torch_device, dtype=compute_dtype) + model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype) + + reset_memory_stats() + model(**inputs_dict) + model_memory_footprint = model.get_memory_footprint() + peak_inference_memory_allocated_mb = backend_max_memory_allocated(torch_device) / 1024**2 + + return model_memory_footprint, peak_inference_memory_allocated_mb + + fp32_memory_footprint, fp32_max_memory = get_memory_usage(torch.float32, torch.float32) + fp8_e4m3_fp32_memory_footprint, fp8_e4m3_fp32_max_memory = get_memory_usage(torch.float8_e4m3fn, torch.float32) + fp8_e4m3_bf16_memory_footprint, fp8_e4m3_bf16_max_memory = get_memory_usage( + torch.float8_e4m3fn, torch.bfloat16 + ) + + compute_capability = get_torch_cuda_device_capability() if torch_device == "cuda" else None + assert fp8_e4m3_bf16_memory_footprint < fp8_e4m3_fp32_memory_footprint < fp32_memory_footprint, ( + "Memory footprint should decrease with lower precision storage" + ) + + # NOTE: the following assertion would fail on our CI (running Tesla T4) due to bf16 using more memory than fp32. + # On other devices, such as DGX (Ampere) and Audace (Ada), the test passes. So, we conditionally check it. + if compute_capability and compute_capability >= LEAST_COMPUTE_CAPABILITY: + assert fp8_e4m3_bf16_max_memory < fp8_e4m3_fp32_max_memory, ( + "Peak memory should be lower with bf16 compute on newer GPUs" + ) + + # On this dummy test case with a small model, sometimes fp8_e4m3_fp32 max memory usage is higher than fp32 by a few + # bytes. This only happens for some models, so we allow a small tolerance. + # For any real model being tested, the order would be fp8_e4m3_bf16 < fp8_e4m3_fp32 < fp32. + assert ( + fp8_e4m3_fp32_max_memory < fp32_max_memory + or abs(fp8_e4m3_fp32_max_memory - fp32_max_memory) < MB_TOLERANCE + ), "Peak memory should be lower or within tolerance with fp8 storage" + + def test_layerwise_casting_training(self): + def test_fn(storage_dtype, compute_dtype): + if torch.device(torch_device).type == "cpu" and compute_dtype == torch.bfloat16: + pytest.skip("Skipping test because CPU doesn't go well with bfloat16.") + + model = self.model_class(**self.get_init_dict()) + model = model.to(torch_device, dtype=compute_dtype) + model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype) + model.train() + + inputs_dict = self.get_inputs_dict() + inputs_dict = cast_maybe_tensor_dtype(inputs_dict, torch.float32, compute_dtype) + with torch.amp.autocast(device_type=torch.device(torch_device).type): + output = model(**inputs_dict) + + if isinstance(output, dict): + output = output.to_tuple()[0] + + input_tensor = inputs_dict[self.main_input_name] + noise = torch.randn((input_tensor.shape[0],) + self.output_shape).to(torch_device) + noise = cast_maybe_tensor_dtype(noise, torch.float32, compute_dtype) + loss = torch.nn.functional.mse_loss(output, noise) + + loss.backward() + + test_fn(torch.float16, torch.float32) + test_fn(torch.float8_e4m3fn, torch.float32) + test_fn(torch.float8_e5m2, torch.float32) + test_fn(torch.float8_e4m3fn, torch.bfloat16) + + +@is_memory +@require_accelerator +class MemoryTesterMixin(CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin): + """ + Combined mixin class for all memory optimization tests including CPU/disk offloading, + group offloading, and layerwise dtype casting. + + This mixin inherits from: + - CPUOffloadTesterMixin: CPU and disk offloading tests + - GroupOffloadTesterMixin: Group offloading tests (block-level and leaf-level) + - LayerwiseCastingTesterMixin: Layerwise dtype casting tests + + Expected class attributes to be set by subclasses: + - model_class: The model class to test + - model_split_percents: List of percentages for splitting model across devices (default: [0.5, 0.7]) + + Expected methods to be implemented by subclasses: + - 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: memory + Use `pytest -m "not memory"` to skip these tests + """ + + pass diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py new file mode 100644 index 000000000000..140f6db6d2ff --- /dev/null +++ b/tests/models/testing_utils/quantization.py @@ -0,0 +1,828 @@ +# 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 tempfile + +import pytest +import torch + +from diffusers import BitsAndBytesConfig, GGUFQuantizationConfig, NVIDIAModelOptConfig, QuantoConfig, TorchAoConfig +from diffusers.utils.import_utils import ( + is_bitsandbytes_available, + is_gguf_available, + is_nvidia_modelopt_available, + is_optimum_quanto_available, + is_torchao_available, + is_torchao_version, +) + +from ...testing_utils import ( + backend_empty_cache, + is_bitsandbytes, + is_gguf, + is_modelopt, + is_quanto, + is_torchao, + require_accelerate, + require_accelerator, + require_bitsandbytes_version_greater, + require_gguf_version_greater_or_equal, + require_modelopt_version_greater_or_equal, + require_quanto, + require_torchao_version_greater_or_equal, + torch_device, +) + + +if is_nvidia_modelopt_available(): + import modelopt.torch.quantization as mtq + +if is_bitsandbytes_available(): + import bitsandbytes as bnb + +if is_optimum_quanto_available(): + from optimum.quanto import QLinear + +if is_gguf_available(): + pass + +if is_torchao_available(): + if is_torchao_version(">=", "0.9.0"): + pass + + +@require_accelerator +class QuantizationTesterMixin: + """ + Base mixin class providing common test implementations for quantization testing. + + Backend-specific mixins should: + 1. Implement _create_quantized_model(config_kwargs) + 2. Implement _verify_if_layer_quantized(name, module, config_kwargs) + 3. Define their config dict (e.g., BNB_CONFIGS, QUANTO_WEIGHT_TYPES, etc.) + 4. Use @pytest.mark.parametrize to create tests that call the common test methods below + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods in test classes: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + """ + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + """ + Create a quantized model with the given config kwargs. + + Args: + config_kwargs: Quantization config parameters + **extra_kwargs: Additional kwargs to pass to from_pretrained (e.g., device_map, offload_folder) + """ + raise NotImplementedError("Subclass must implement _create_quantized_model") + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + raise NotImplementedError("Subclass must implement _verify_if_layer_quantized") + + def _is_module_quantized(self, module): + """ + Check if a module is quantized. Returns True if quantized, False otherwise. + Default implementation tries _verify_if_layer_quantized and catches exceptions. + Subclasses can override for more efficient checking. + """ + try: + self._verify_if_layer_quantized("", module, {}) + return True + except (AssertionError, AttributeError): + return False + + def _load_unquantized_model(self): + kwargs = getattr(self, "pretrained_model_kwargs", {}) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _test_quantization_num_parameters(self, config_kwargs): + model = self._load_unquantized_model() + num_params = model.num_parameters() + + model_quantized = self._create_quantized_model(config_kwargs) + num_params_quantized = model_quantized.num_parameters() + + assert num_params == num_params_quantized, ( + f"Parameter count mismatch: unquantized={num_params}, quantized={num_params_quantized}" + ) + + def _test_quantization_memory_footprint(self, config_kwargs, expected_memory_reduction=1.2): + model = self._load_unquantized_model() + mem = model.get_memory_footprint() + + model_quantized = self._create_quantized_model(config_kwargs) + mem_quantized = model_quantized.get_memory_footprint() + + ratio = mem / mem_quantized + assert ratio >= expected_memory_reduction, ( + f"Memory ratio {ratio:.2f} is less than expected ({expected_memory_reduction}x). unquantized={mem}, quantized={mem_quantized}" + ) + + def _test_quantization_inference(self, config_kwargs): + model_quantized = self._create_quantized_model(config_kwargs) + + with torch.no_grad(): + inputs = self.get_dummy_inputs() + output = model_quantized(**inputs) + + if isinstance(output, tuple): + output = output[0] + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + def _test_quantization_dtype_assignment(self, config_kwargs): + model = self._create_quantized_model(config_kwargs) + + with pytest.raises(ValueError): + model.to(torch.float16) + + with pytest.raises(ValueError): + device_0 = f"{torch_device}:0" + model.to(device=device_0, dtype=torch.float16) + + with pytest.raises(ValueError): + model.float() + + with pytest.raises(ValueError): + model.half() + + model.to(torch_device) + + def _test_quantization_lora_inference(self, config_kwargs): + try: + from peft import LoraConfig + except ImportError: + pytest.skip("peft is not available") + + from diffusers.loaders.peft import PeftAdapterMixin + + if not issubclass(self.model_class, PeftAdapterMixin): + pytest.skip(f"PEFT is not supported for this model ({self.model_class.__name__})") + + model = self._create_quantized_model(config_kwargs) + + lora_config = LoraConfig( + r=4, + lora_alpha=4, + target_modules=["to_q", "to_k", "to_v", "to_out.0"], + init_lora_weights=False, + ) + model.add_adapter(lora_config) + + with torch.no_grad(): + inputs = self.get_dummy_inputs() + output = model(**inputs) + + if isinstance(output, tuple): + output = output[0] + assert output is not None, "Model output is None with LoRA" + assert not torch.isnan(output).any(), "Model output contains NaN with LoRA" + + def _test_quantization_serialization(self, config_kwargs): + model = self._create_quantized_model(config_kwargs) + + with tempfile.TemporaryDirectory() as tmpdir: + model.save_pretrained(tmpdir, safe_serialization=True) + + model_loaded = self.model_class.from_pretrained(tmpdir) + + with torch.no_grad(): + inputs = self.get_dummy_inputs() + output = model_loaded(**inputs) + if isinstance(output, tuple): + output = output[0] + assert not torch.isnan(output).any(), "Loaded model output contains NaN" + + def _test_quantized_layers(self, config_kwargs): + model_fp = self._load_unquantized_model() + num_linear_layers = sum(1 for module in model_fp.modules() if isinstance(module, torch.nn.Linear)) + + model_quantized = self._create_quantized_model(config_kwargs) + + num_fp32_modules = 0 + if hasattr(model_quantized, "_keep_in_fp32_modules") and model_quantized._keep_in_fp32_modules: + for name, module in model_quantized.named_modules(): + if isinstance(module, torch.nn.Linear): + if any(fp32_name in name for fp32_name in model_quantized._keep_in_fp32_modules): + num_fp32_modules += 1 + + expected_quantized_layers = num_linear_layers - num_fp32_modules + + num_quantized_layers = 0 + for name, module in model_quantized.named_modules(): + if isinstance(module, torch.nn.Linear): + if hasattr(model_quantized, "_keep_in_fp32_modules") and model_quantized._keep_in_fp32_modules: + if any(fp32_name in name for fp32_name in model_quantized._keep_in_fp32_modules): + continue + self._verify_if_layer_quantized(name, module, config_kwargs) + num_quantized_layers += 1 + + assert num_quantized_layers > 0, ( + f"No quantized layers found in model (expected {expected_quantized_layers} linear layers, {num_fp32_modules} kept in FP32)" + ) + assert num_quantized_layers == expected_quantized_layers, ( + f"Quantized layer count mismatch: expected {expected_quantized_layers}, got {num_quantized_layers} (total linear layers: {num_linear_layers}, FP32 modules: {num_fp32_modules})" + ) + + def _test_quantization_modules_to_not_convert(self, config_kwargs, modules_to_not_convert): + """ + Test that modules specified in modules_to_not_convert are not quantized. + + Args: + config_kwargs: Base quantization config kwargs + modules_to_not_convert: List of module names to exclude from quantization + """ + # Create config with modules_to_not_convert + config_kwargs_with_exclusion = config_kwargs.copy() + config_kwargs_with_exclusion["modules_to_not_convert"] = modules_to_not_convert + + model_with_exclusion = self._create_quantized_model(config_kwargs_with_exclusion) + + # Find a module that should NOT be quantized + found_excluded = False + for name, module in model_with_exclusion.named_modules(): + if isinstance(module, torch.nn.Linear): + # Check if this module is in the exclusion list + if any(excluded in name for excluded in modules_to_not_convert): + found_excluded = True + # This module should NOT be quantized + assert not self._is_module_quantized(module), ( + f"Module {name} should not be quantized but was found to be quantized" + ) + + assert found_excluded, f"No linear layers found in excluded modules: {modules_to_not_convert}" + + # Find a module that SHOULD be quantized (not in exclusion list) + found_quantized = False + for name, module in model_with_exclusion.named_modules(): + if isinstance(module, torch.nn.Linear): + # Check if this module is NOT in the exclusion list + if not any(excluded in name for excluded in modules_to_not_convert): + if self._is_module_quantized(module): + found_quantized = True + break + + assert found_quantized, "No quantized layers found outside of excluded modules" + + # Compare memory footprint with fully quantized model + model_fully_quantized = self._create_quantized_model(config_kwargs) + + mem_with_exclusion = model_with_exclusion.get_memory_footprint() + mem_fully_quantized = model_fully_quantized.get_memory_footprint() + + assert mem_with_exclusion > mem_fully_quantized, ( + f"Model with exclusions should be larger. With exclusion: {mem_with_exclusion}, fully quantized: {mem_fully_quantized}" + ) + + def _test_quantization_device_map(self, config_kwargs): + """ + Test that quantized models work correctly with device_map="auto". + + Args: + config_kwargs: Base quantization config kwargs + """ + model = self._create_quantized_model(config_kwargs, device_map="auto") + + # Verify device map is set + assert hasattr(model, "hf_device_map"), "Model should have hf_device_map attribute" + assert model.hf_device_map is not None, "hf_device_map should not be None" + + # Verify inference works + with torch.no_grad(): + inputs = self.get_dummy_inputs() + output = model(**inputs) + if isinstance(output, tuple): + output = output[0] + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + +@is_bitsandbytes +@require_accelerator +@require_bitsandbytes_version_greater("0.43.2") +@require_accelerate +class BitsAndBytesTesterMixin(QuantizationTesterMixin): + """ + Mixin class for testing BitsAndBytes quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - BNB_CONFIGS: Dict of config name -> BitsAndBytesConfig kwargs to test + + Pytest mark: bitsandbytes + Use `pytest -m "not bitsandbytes"` to skip these tests + """ + + # Standard BnB configs tested for all models + # Subclasses can override to add or modify configs + BNB_CONFIGS = { + "4bit_nf4": { + "load_in_4bit": True, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": torch.float16, + }, + "4bit_fp4": { + "load_in_4bit": True, + "bnb_4bit_quant_type": "fp4", + "bnb_4bit_compute_dtype": torch.float16, + }, + "8bit": { + "load_in_8bit": True, + }, + } + + BNB_EXPECTED_MEMORY_REDUCTIONS = { + "4bit_nf4": 3.0, + "4bit_fp4": 3.0, + "8bit": 1.5, + } + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + config = BitsAndBytesConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + expected_weight_class = bnb.nn.Params4bit if config_kwargs.get("load_in_4bit") else bnb.nn.Int8Params + assert module.weight.__class__ == expected_weight_class, ( + f"Layer {name} has weight type {module.weight.__class__}, expected {expected_weight_class}" + ) + + @pytest.mark.parametrize("config_name", list(BNB_CONFIGS.keys())) + def test_bnb_quantization_num_parameters(self, config_name): + self._test_quantization_num_parameters(self.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", list(BNB_CONFIGS.keys())) + def test_bnb_quantization_memory_footprint(self, config_name): + expected = self.BNB_EXPECTED_MEMORY_REDUCTIONS.get(config_name, 1.2) + self._test_quantization_memory_footprint(self.BNB_CONFIGS[config_name], expected_memory_reduction=expected) + + @pytest.mark.parametrize("config_name", list(BNB_CONFIGS.keys())) + def test_bnb_quantization_inference(self, config_name): + self._test_quantization_inference(self.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["4bit_nf4"]) + def test_bnb_quantization_dtype_assignment(self, config_name): + self._test_quantization_dtype_assignment(self.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["4bit_nf4"]) + def test_bnb_quantization_lora_inference(self, config_name): + self._test_quantization_lora_inference(self.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["4bit_nf4"]) + def test_bnb_quantization_serialization(self, config_name): + self._test_quantization_serialization(self.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", list(BNB_CONFIGS.keys())) + def test_bnb_quantized_layers(self, config_name): + self._test_quantized_layers(self.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", list(BNB_CONFIGS.keys())) + def test_bnb_quantization_config_serialization(self, config_name): + model = self._create_quantized_model(self.BNB_CONFIGS[config_name]) + + assert "quantization_config" in model.config, "Missing quantization_config" + _ = model.config["quantization_config"].to_dict() + _ = model.config["quantization_config"].to_diff_dict() + _ = model.config["quantization_config"].to_json_string() + + def test_bnb_original_dtype(self): + config_name = list(self.BNB_CONFIGS.keys())[0] + config_kwargs = self.BNB_CONFIGS[config_name] + + model = self._create_quantized_model(config_kwargs) + + assert "_pre_quantization_dtype" in model.config, "Missing _pre_quantization_dtype" + assert model.config["_pre_quantization_dtype"] in [ + torch.float16, + torch.float32, + torch.bfloat16, + ], f"Unexpected dtype: {model.config['_pre_quantization_dtype']}" + + def test_bnb_keep_modules_in_fp32(self): + if not hasattr(self.model_class, "_keep_in_fp32_modules"): + pytest.skip(f"{self.model_class.__name__} does not have _keep_in_fp32_modules") + + config_kwargs = self.BNB_CONFIGS["4bit_nf4"] + + original_fp32_modules = getattr(self.model_class, "_keep_in_fp32_modules", None) + self.model_class._keep_in_fp32_modules = ["proj_out"] + + try: + model = self._create_quantized_model(config_kwargs) + + for name, module in model.named_modules(): + if isinstance(module, torch.nn.Linear): + if any(fp32_name in name for fp32_name in model._keep_in_fp32_modules): + assert module.weight.dtype == torch.float32, ( + f"Module {name} should be FP32 but is {module.weight.dtype}" + ) + else: + assert module.weight.dtype == torch.uint8, ( + f"Module {name} should be uint8 but is {module.weight.dtype}" + ) + + with torch.no_grad(): + inputs = self.get_dummy_inputs() + _ = model(**inputs) + finally: + if original_fp32_modules is not None: + self.model_class._keep_in_fp32_modules = original_fp32_modules + + def test_bnb_modules_to_not_convert(self): + """Test that modules_to_not_convert parameter works correctly.""" + modules_to_exclude = getattr(self, "modules_to_not_convert_for_test", None) + if modules_to_exclude is None: + pytest.skip("modules_to_not_convert_for_test not defined for this model") + + self._test_quantization_modules_to_not_convert(self.BNB_CONFIGS["4bit_nf4"], modules_to_exclude) + + def test_bnb_device_map(self): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(self.BNB_CONFIGS["4bit_nf4"]) + + +@is_quanto +@require_quanto +@require_accelerate +@require_accelerator +class QuantoTesterMixin(QuantizationTesterMixin): + """ + Mixin class for testing Quanto quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - QUANTO_WEIGHT_TYPES: Dict of weight_type_name -> qtype + + Pytest mark: quanto + Use `pytest -m "not quanto"` to skip these tests + """ + + QUANTO_WEIGHT_TYPES = { + "float8": {"weights_dtype": "float8"}, + "int8": {"weights_dtype": "int8"}, + "int4": {"weights_dtype": "int4"}, + "int2": {"weights_dtype": "int2"}, + } + + QUANTO_EXPECTED_MEMORY_REDUCTIONS = { + "float8": 1.5, + "int8": 1.5, + "int4": 3.0, + "int2": 7.0, + } + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + config = QuantoConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + assert isinstance(module, QLinear), f"Layer {name} is not QLinear, got {type(module)}" + + @pytest.mark.parametrize("weight_type_name", list(QUANTO_WEIGHT_TYPES.keys())) + def test_quanto_quantization_num_parameters(self, weight_type_name): + self._test_quantization_num_parameters(self.QUANTO_WEIGHT_TYPES[weight_type_name]) + + @pytest.mark.parametrize("weight_type_name", list(QUANTO_WEIGHT_TYPES.keys())) + def test_quanto_quantization_memory_footprint(self, weight_type_name): + expected = self.QUANTO_EXPECTED_MEMORY_REDUCTIONS.get(weight_type_name, 1.2) + self._test_quantization_memory_footprint( + self.QUANTO_WEIGHT_TYPES[weight_type_name], expected_memory_reduction=expected + ) + + @pytest.mark.parametrize("weight_type_name", list(QUANTO_WEIGHT_TYPES.keys())) + def test_quanto_quantization_inference(self, weight_type_name): + self._test_quantization_inference(self.QUANTO_WEIGHT_TYPES[weight_type_name]) + + @pytest.mark.parametrize("weight_type_name", ["int8"]) + def test_quanto_quantized_layers(self, weight_type_name): + self._test_quantized_layers(self.QUANTO_WEIGHT_TYPES[weight_type_name]) + + @pytest.mark.parametrize("weight_type_name", ["int8"]) + def test_quanto_quantization_lora_inference(self, weight_type_name): + self._test_quantization_lora_inference(self.QUANTO_WEIGHT_TYPES[weight_type_name]) + + @pytest.mark.parametrize("weight_type_name", ["int8"]) + def test_quanto_quantization_serialization(self, weight_type_name): + self._test_quantization_serialization(self.QUANTO_WEIGHT_TYPES[weight_type_name]) + + def test_quanto_modules_to_not_convert(self): + """Test that modules_to_not_convert parameter works correctly.""" + modules_to_exclude = getattr(self, "modules_to_not_convert_for_test", None) + if modules_to_exclude is None: + pytest.skip("modules_to_not_convert_for_test not defined for this model") + + self._test_quantization_modules_to_not_convert(self.QUANTO_WEIGHT_TYPES["int8"], modules_to_exclude) + + def test_quanto_device_map(self): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(self.QUANTO_WEIGHT_TYPES["int8"]) + + +@is_torchao +@require_accelerator +@require_torchao_version_greater_or_equal("0.7.0") +class TorchAoTesterMixin(QuantizationTesterMixin): + """ + Mixin class for testing TorchAO quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - TORCHAO_QUANT_TYPES: Dict of quantization type strings to test + + Pytest mark: torchao + Use `pytest -m "not torchao"` to skip these tests + """ + + TORCHAO_QUANT_TYPES = { + "int4wo": {"quant_type": "int4_weight_only"}, + "int8wo": {"quant_type": "int8_weight_only"}, + "int8dq": {"quant_type": "int8_dynamic_activation_int8_weight"}, + } + + TORCHAO_EXPECTED_MEMORY_REDUCTIONS = { + "int4wo": 3.0, + "int8wo": 1.5, + "int8dq": 1.5, + } + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + config = TorchAoConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + assert isinstance(module, torch.nn.Linear), f"Layer {name} is not Linear, got {type(module)}" + + @pytest.mark.parametrize("quant_type", list(TORCHAO_QUANT_TYPES.keys())) + def test_torchao_quantization_num_parameters(self, quant_type): + self._test_quantization_num_parameters(self.TORCHAO_QUANT_TYPES[quant_type]) + + @pytest.mark.parametrize("quant_type", list(TORCHAO_QUANT_TYPES.keys())) + def test_torchao_quantization_memory_footprint(self, quant_type): + expected = self.TORCHAO_EXPECTED_MEMORY_REDUCTIONS.get(quant_type, 1.2) + self._test_quantization_memory_footprint( + self.TORCHAO_QUANT_TYPES[quant_type], expected_memory_reduction=expected + ) + + @pytest.mark.parametrize("quant_type", list(TORCHAO_QUANT_TYPES.keys())) + def test_torchao_quantization_inference(self, quant_type): + self._test_quantization_inference(self.TORCHAO_QUANT_TYPES[quant_type]) + + @pytest.mark.parametrize("quant_type", ["int8wo"]) + def test_torchao_quantized_layers(self, quant_type): + self._test_quantized_layers(self.TORCHAO_QUANT_TYPES[quant_type]) + + @pytest.mark.parametrize("quant_type", ["int8wo"]) + def test_torchao_quantization_lora_inference(self, quant_type): + self._test_quantization_lora_inference(self.TORCHAO_QUANT_TYPES[quant_type]) + + @pytest.mark.parametrize("quant_type", ["int8wo"]) + def test_torchao_quantization_serialization(self, quant_type): + self._test_quantization_serialization(self.TORCHAO_QUANT_TYPES[quant_type]) + + def test_torchao_modules_to_not_convert(self): + """Test that modules_to_not_convert parameter works correctly.""" + # Get a module name that exists in the model - this needs to be set by test classes + # For now, use a generic pattern that should work with transformer models + modules_to_exclude = getattr(self, "modules_to_not_convert_for_test", None) + if modules_to_exclude is None: + pytest.skip("modules_to_not_convert_for_test not defined for this model") + + self._test_quantization_modules_to_not_convert(self.TORCHAO_QUANT_TYPES["int8wo"], modules_to_exclude) + + def test_torchao_device_map(self): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(self.TORCHAO_QUANT_TYPES["int8wo"]) + + +@is_gguf +@require_accelerate +@require_accelerator +@require_gguf_version_greater_or_equal("0.10.0") +class GGUFTesterMixin(QuantizationTesterMixin): + """ + Mixin class for testing GGUF quantization on models. + + Expected class attributes: + - model_class: The model class to test + - gguf_filename: URL or path to the GGUF file + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: gguf + Use `pytest -m "not gguf"` to skip these tests + """ + + gguf_filename = None + + def _create_quantized_model(self, config_kwargs=None, **extra_kwargs): + if config_kwargs is None: + config_kwargs = {"compute_dtype": torch.bfloat16} + + config = GGUFQuantizationConfig(**config_kwargs) + kwargs = { + "quantization_config": config, + "torch_dtype": config_kwargs.get("compute_dtype", torch.bfloat16), + } + kwargs.update(extra_kwargs) + return self.model_class.from_single_file(self.gguf_filename, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs=None): + from diffusers.quantizers.gguf.utils import GGUFParameter + + assert isinstance(module.weight, GGUFParameter), f"{name} weight is not GGUFParameter" + assert hasattr(module.weight, "quant_type"), f"{name} weight missing quant_type" + assert module.weight.dtype == torch.uint8, f"{name} weight dtype should be uint8" + + def test_gguf_quantization_inference(self): + self._test_quantization_inference({"compute_dtype": torch.bfloat16}) + + def test_gguf_keep_modules_in_fp32(self): + if not hasattr(self.model_class, "_keep_in_fp32_modules"): + pytest.skip(f"{self.model_class.__name__} does not have _keep_in_fp32_modules") + + _keep_in_fp32_modules = self.model_class._keep_in_fp32_modules + self.model_class._keep_in_fp32_modules = ["proj_out"] + + try: + model = self._create_quantized_model() + + for name, module in model.named_modules(): + if isinstance(module, torch.nn.Linear): + if any(fp32_name in name for fp32_name in model._keep_in_fp32_modules): + assert module.weight.dtype == torch.float32, f"Module {name} should be FP32" + finally: + self.model_class._keep_in_fp32_modules = _keep_in_fp32_modules + + def test_gguf_quantization_dtype_assignment(self): + self._test_quantization_dtype_assignment({"compute_dtype": torch.bfloat16}) + + def test_gguf_quantization_lora_inference(self): + self._test_quantization_lora_inference({"compute_dtype": torch.bfloat16}) + + def test_gguf_dequantize_model(self): + from diffusers.quantizers.gguf.utils import GGUFLinear, GGUFParameter + + model = self._create_quantized_model() + model.dequantize() + + def _check_for_gguf_linear(model): + has_children = list(model.children()) + if not has_children: + return + + for name, module in model.named_children(): + if isinstance(module, torch.nn.Linear): + assert not isinstance(module, GGUFLinear), f"{name} is still GGUFLinear" + assert not isinstance(module.weight, GGUFParameter), f"{name} weight is still GGUFParameter" + + for name, module in model.named_children(): + _check_for_gguf_linear(module) + + def test_gguf_quantized_layers(self): + self._test_quantized_layers({"compute_dtype": torch.bfloat16}) + + +@is_modelopt +@require_accelerator +@require_accelerate +@require_modelopt_version_greater_or_equal("0.33.1") +class ModelOptTesterMixin(QuantizationTesterMixin): + """ + Mixin class for testing NVIDIA ModelOpt quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - MODELOPT_CONFIGS: Dict of config name -> NVIDIAModelOptConfig kwargs to test + + Pytest mark: modelopt + Use `pytest -m "not modelopt"` to skip these tests + """ + + MODELOPT_CONFIGS = { + "fp8": {"quant_type": "FP8"}, + "int8": {"quant_type": "INT8"}, + "int4": {"quant_type": "INT4"}, + } + + MODELOPT_EXPECTED_MEMORY_REDUCTIONS = { + "fp8": 1.5, + "int8": 1.5, + "int4": 3.0, + } + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + config = NVIDIAModelOptConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + assert mtq.utils.is_quantized(module), f"Layer {name} does not have weight_quantizer attribute (not quantized)" + + @pytest.mark.parametrize("config_name", ["fp8"]) + def test_modelopt_quantization_num_parameters(self, config_name): + self._test_quantization_num_parameters(self.MODELOPT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", list(MODELOPT_CONFIGS.keys())) + def test_modelopt_quantization_memory_footprint(self, config_name): + expected = self.MODELOPT_EXPECTED_MEMORY_REDUCTIONS.get(config_name, 1.2) + self._test_quantization_memory_footprint( + self.MODELOPT_CONFIGS[config_name], expected_memory_reduction=expected + ) + + @pytest.mark.parametrize("config_name", list(MODELOPT_CONFIGS.keys())) + def test_modelopt_quantization_inference(self, config_name): + self._test_quantization_inference(self.MODELOPT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"]) + def test_modelopt_quantization_dtype_assignment(self, config_name): + self._test_quantization_dtype_assignment(self.MODELOPT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"]) + def test_modelopt_quantization_lora_inference(self, config_name): + self._test_quantization_lora_inference(self.MODELOPT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"]) + def test_modelopt_quantization_serialization(self, config_name): + self._test_quantization_serialization(self.MODELOPT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"]) + def test_modelopt_quantized_layers(self, config_name): + self._test_quantized_layers(self.MODELOPT_CONFIGS[config_name]) + + def test_modelopt_modules_to_not_convert(self): + """Test that modules_to_not_convert parameter works correctly.""" + modules_to_exclude = getattr(self, "modules_to_not_convert_for_test", None) + if modules_to_exclude is None: + pytest.skip("modules_to_not_convert_for_test not defined for this model") + + self._test_quantization_modules_to_not_convert(self.MODELOPT_CONFIGS["fp8"], modules_to_exclude) + + def test_modelopt_device_map(self): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(self.MODELOPT_CONFIGS["fp8"]) diff --git a/tests/models/testing_utils/single_file.py b/tests/models/testing_utils/single_file.py new file mode 100644 index 000000000000..67d770849f00 --- /dev/null +++ b/tests/models/testing_utils/single_file.py @@ -0,0 +1,247 @@ +# 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 tempfile + +import torch +from huggingface_hub import hf_hub_download, snapshot_download + +from diffusers.loaders.single_file_utils import _extract_repo_id_and_weights_name + +from ...testing_utils import ( + backend_empty_cache, + is_single_file, + nightly, + require_torch_accelerator, + torch_device, +) + + +def download_single_file_checkpoint(pretrained_model_name_or_path, filename, tmpdir): + """Download a single file checkpoint from the Hub to a temporary directory.""" + path = hf_hub_download(pretrained_model_name_or_path, filename=filename, local_dir=tmpdir) + return path + + +def download_diffusers_config(pretrained_model_name_or_path, tmpdir): + """Download diffusers config files (excluding weights) from a repository.""" + path = snapshot_download( + pretrained_model_name_or_path, + ignore_patterns=[ + "**/*.ckpt", + "*.ckpt", + "**/*.bin", + "*.bin", + "**/*.pt", + "*.pt", + "**/*.safetensors", + "*.safetensors", + ], + allow_patterns=["**/*.json", "*.json", "*.txt", "**/*.txt"], + local_dir=tmpdir, + ) + return path + + +@nightly +@require_torch_accelerator +@is_single_file +class SingleFileTesterMixin: + """ + Mixin class for testing single file loading for models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - ckpt_path: Path or Hub path to the single file checkpoint + - subfolder: (Optional) Subfolder within the repo + - torch_dtype: (Optional) torch dtype to use for testing + + Pytest mark: single_file + Use `pytest -m "not single_file"` to skip these tests + """ + + pretrained_model_name_or_path = None + ckpt_path = None + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def test_single_file_model_config(self): + pretrained_kwargs = {} + single_file_kwargs = {} + + pretrained_kwargs["device"] = torch_device + single_file_kwargs["device"] = torch_device + + if hasattr(self, "subfolder") and self.subfolder: + pretrained_kwargs["subfolder"] = self.subfolder + + if hasattr(self, "torch_dtype") and self.torch_dtype: + pretrained_kwargs["torch_dtype"] = self.torch_dtype + single_file_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_pretrained(self.pretrained_model_name_or_path, **pretrained_kwargs) + model_single_file = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs) + + PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] + for param_name, param_value in model_single_file.config.items(): + if param_name in PARAMS_TO_IGNORE: + continue + assert model.config[param_name] == param_value, ( + f"{param_name} differs between pretrained loading and single file loading: " + f"pretrained={model.config[param_name]}, single_file={param_value}" + ) + + def test_single_file_model_parameters(self): + pretrained_kwargs = {} + single_file_kwargs = {} + + pretrained_kwargs["device"] = torch_device + single_file_kwargs["device"] = torch_device + + if hasattr(self, "subfolder") and self.subfolder: + pretrained_kwargs["subfolder"] = self.subfolder + + if hasattr(self, "torch_dtype") and self.torch_dtype: + pretrained_kwargs["torch_dtype"] = self.torch_dtype + single_file_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_pretrained(self.pretrained_model_name_or_path, **pretrained_kwargs) + model_single_file = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs) + + state_dict = model.state_dict() + state_dict_single_file = model_single_file.state_dict() + + assert set(state_dict.keys()) == set(state_dict_single_file.keys()), ( + "Model parameters keys differ between pretrained and single file loading. " + f"Missing in single file: {set(state_dict.keys()) - set(state_dict_single_file.keys())}. " + f"Extra in single file: {set(state_dict_single_file.keys()) - set(state_dict.keys())}" + ) + + for key in state_dict.keys(): + param = state_dict[key] + param_single_file = state_dict_single_file[key] + + assert param.shape == param_single_file.shape, ( + f"Parameter shape mismatch for {key}: " + f"pretrained {param.shape} vs single file {param_single_file.shape}" + ) + + assert torch.allclose(param, param_single_file, rtol=1e-5, atol=1e-5), ( + f"Parameter values differ for {key}: " + f"max difference {torch.max(torch.abs(param - param_single_file)).item()}" + ) + + def test_single_file_loading_local_files_only(self): + single_file_kwargs = {} + + if hasattr(self, "torch_dtype") and self.torch_dtype: + single_file_kwargs["torch_dtype"] = self.torch_dtype + + with tempfile.TemporaryDirectory() as tmpdir: + pretrained_model_name_or_path, weight_name = _extract_repo_id_and_weights_name(self.ckpt_path) + local_ckpt_path = download_single_file_checkpoint(pretrained_model_name_or_path, weight_name, tmpdir) + + model_single_file = self.model_class.from_single_file( + local_ckpt_path, local_files_only=True, **single_file_kwargs + ) + + assert model_single_file is not None, "Failed to load model with local_files_only=True" + + def test_single_file_loading_with_diffusers_config(self): + single_file_kwargs = {} + + if hasattr(self, "torch_dtype") and self.torch_dtype: + single_file_kwargs["torch_dtype"] = self.torch_dtype + + # Load with config parameter + model_single_file = self.model_class.from_single_file( + self.ckpt_path, config=self.pretrained_model_name_or_path, **single_file_kwargs + ) + + # Load pretrained for comparison + pretrained_kwargs = {} + if hasattr(self, "subfolder") and self.subfolder: + pretrained_kwargs["subfolder"] = self.subfolder + if hasattr(self, "torch_dtype") and self.torch_dtype: + pretrained_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_pretrained(self.pretrained_model_name_or_path, **pretrained_kwargs) + + # Compare configs + PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] + for param_name, param_value in model_single_file.config.items(): + if param_name in PARAMS_TO_IGNORE: + continue + assert model.config[param_name] == param_value, ( + f"{param_name} differs: pretrained={model.config[param_name]}, single_file={param_value}" + ) + + def test_single_file_loading_with_diffusers_config_local_files_only(self): + single_file_kwargs = {} + + if hasattr(self, "torch_dtype") and self.torch_dtype: + single_file_kwargs["torch_dtype"] = self.torch_dtype + + with tempfile.TemporaryDirectory() as tmpdir: + pretrained_model_name_or_path, weight_name = _extract_repo_id_and_weights_name(self.ckpt_path) + local_ckpt_path = download_single_file_checkpoint(pretrained_model_name_or_path, weight_name, tmpdir) + local_diffusers_config = download_diffusers_config(self.pretrained_model_name_or_path, tmpdir) + + model_single_file = self.model_class.from_single_file( + local_ckpt_path, config=local_diffusers_config, local_files_only=True, **single_file_kwargs + ) + + assert model_single_file is not None, "Failed to load model with config and local_files_only=True" + + def test_single_file_loading_dtype(self): + for dtype in [torch.float32, torch.float16]: + if torch_device == "mps" and dtype == torch.bfloat16: + continue + + model_single_file = self.model_class.from_single_file(self.ckpt_path, torch_dtype=dtype) + + assert model_single_file.dtype == dtype, f"Expected dtype {dtype}, got {model_single_file.dtype}" + + # Cleanup + del model_single_file + gc.collect() + backend_empty_cache(torch_device) + + def test_checkpoint_variant_loading(self): + if not hasattr(self, "alternate_ckpt_paths") or not self.alternate_ckpt_paths: + return + + for ckpt_path in self.alternate_ckpt_paths: + backend_empty_cache(torch_device) + + single_file_kwargs = {} + if hasattr(self, "torch_dtype") and self.torch_dtype: + single_file_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_single_file(ckpt_path, **single_file_kwargs) + + assert model is not None, f"Failed to load checkpoint from {ckpt_path}" + + del model + gc.collect() + backend_empty_cache(torch_device) diff --git a/tests/models/testing_utils/training.py b/tests/models/testing_utils/training.py new file mode 100644 index 000000000000..7e4193d59e84 --- /dev/null +++ b/tests/models/testing_utils/training.py @@ -0,0 +1,223 @@ +# 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 copy + +import pytest +import torch + +from diffusers.training_utils import EMAModel + +from ...testing_utils import is_training, require_torch_accelerator_with_training, torch_all_close, torch_device + + +@is_training +@require_torch_accelerator_with_training +class TrainingTesterMixin: + """ + Mixin class for testing training functionality on models. + + Expected class attributes to be set by subclasses: + - model_class: The model class to test + + Expected methods to be implemented by subclasses: + - 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 + + Expected properties to be implemented by subclasses: + - output_shape: Tuple defining the expected output shape + + Pytest mark: training + Use `pytest -m "not training"` to skip these tests + """ + + def test_training(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict) + model.to(torch_device) + model.train() + output = model(**inputs_dict) + + if isinstance(output, dict): + output = output.to_tuple()[0] + + noise = torch.randn((output.shape[0],) + self.output_shape).to(torch_device) + loss = torch.nn.functional.mse_loss(output, noise) + loss.backward() + + def test_training_with_ema(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict) + model.to(torch_device) + model.train() + ema_model = EMAModel(model.parameters()) + + output = model(**inputs_dict) + + if isinstance(output, dict): + output = output.to_tuple()[0] + + noise = torch.randn((output.shape[0],) + self.output_shape).to(torch_device) + loss = torch.nn.functional.mse_loss(output, noise) + loss.backward() + ema_model.step(model.parameters()) + + def test_gradient_checkpointing(self): + if not self.model_class._supports_gradient_checkpointing: + pytest.skip("Gradient checkpointing is not supported.") + + init_dict = self.get_init_dict() + + # at init model should have gradient checkpointing disabled + model = self.model_class(**init_dict) + assert not model.is_gradient_checkpointing, "Gradient checkpointing should be disabled at init" + + # check enable works + model.enable_gradient_checkpointing() + assert model.is_gradient_checkpointing, "Gradient checkpointing should be enabled" + + # check disable works + model.disable_gradient_checkpointing() + assert not model.is_gradient_checkpointing, "Gradient checkpointing should be disabled" + + def test_gradient_checkpointing_is_applied(self, expected_set=None): + if not self.model_class._supports_gradient_checkpointing: + pytest.skip("Gradient checkpointing is not supported.") + + if expected_set is None: + pytest.skip("expected_set must be provided to verify gradient checkpointing is applied.") + + init_dict = self.get_init_dict() + + model_class_copy = copy.copy(self.model_class) + model = model_class_copy(**init_dict) + model.enable_gradient_checkpointing() + + modules_with_gc_enabled = {} + for submodule in model.modules(): + if hasattr(submodule, "gradient_checkpointing"): + assert submodule.gradient_checkpointing, f"{submodule.__class__.__name__} should have GC enabled" + modules_with_gc_enabled[submodule.__class__.__name__] = True + + assert set(modules_with_gc_enabled.keys()) == expected_set, ( + f"Modules with GC enabled {set(modules_with_gc_enabled.keys())} do not match expected set {expected_set}" + ) + assert all(modules_with_gc_enabled.values()), "All modules should have GC enabled" + + def test_gradient_checkpointing_equivalence(self, loss_tolerance=1e-5, param_grad_tol=5e-5, skip=None): + if not self.model_class._supports_gradient_checkpointing: + pytest.skip("Gradient checkpointing is not supported.") + + if skip is None: + skip = set() + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + inputs_dict_copy = copy.deepcopy(inputs_dict) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.to(torch_device) + + assert not model.is_gradient_checkpointing and model.training + + out = model(**inputs_dict) + if isinstance(out, dict): + out = out.sample if hasattr(out, "sample") else out.to_tuple()[0] + + # run the backwards pass on the model + model.zero_grad() + + labels = torch.randn_like(out) + loss = (out - labels).mean() + loss.backward() + + # re-instantiate the model now enabling gradient checkpointing + torch.manual_seed(0) + model_2 = self.model_class(**init_dict) + # clone model + model_2.load_state_dict(model.state_dict()) + model_2.to(torch_device) + model_2.enable_gradient_checkpointing() + + assert model_2.is_gradient_checkpointing and model_2.training + + out_2 = model_2(**inputs_dict_copy) + if isinstance(out_2, dict): + out_2 = out_2.sample if hasattr(out_2, "sample") else out_2.to_tuple()[0] + + # run the backwards pass on the model + model_2.zero_grad() + loss_2 = (out_2 - labels).mean() + loss_2.backward() + + # compare the output and parameters gradients + assert (loss - loss_2).abs() < loss_tolerance, ( + f"Loss difference {(loss - loss_2).abs()} exceeds tolerance {loss_tolerance}" + ) + + named_params = dict(model.named_parameters()) + named_params_2 = dict(model_2.named_parameters()) + + for name, param in named_params.items(): + if "post_quant_conv" in name: + continue + if name in skip: + continue + if param.grad is None: + continue + + assert torch_all_close(param.grad.data, named_params_2[name].grad.data, atol=param_grad_tol), ( + f"Gradient mismatch for {name}" + ) + + def test_mixed_precision_training(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict) + model.to(torch_device) + model.train() + + # Test with float16 + if torch.device(torch_device).type != "cpu": + with torch.amp.autocast(device_type=torch.device(torch_device).type, dtype=torch.float16): + output = model(**inputs_dict) + + if isinstance(output, dict): + output = output.to_tuple()[0] + + noise = torch.randn((output.shape[0],) + self.output_shape).to(torch_device) + loss = torch.nn.functional.mse_loss(output, noise) + + loss.backward() + + # Test with bfloat16 + if torch.device(torch_device).type != "cpu": + model.zero_grad() + with torch.amp.autocast(device_type=torch.device(torch_device).type, dtype=torch.bfloat16): + output = model(**inputs_dict) + + if isinstance(output, dict): + output = output.to_tuple()[0] + + noise = torch.randn((output.shape[0],) + self.output_shape).to(torch_device) + loss = torch.nn.functional.mse_loss(output, noise) + + loss.backward() diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 3ab02f797b5b..30193088316c 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -13,23 +13,42 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest +from typing import Any import torch from diffusers import FluxTransformer2DModel -from diffusers.models.attention_processor import FluxIPAdapterJointAttnProcessor2_0 from diffusers.models.embeddings import ImageProjection - -from ...testing_utils import enable_full_determinism, is_peft_available, torch_device -from ..test_modeling_common import LoraHotSwappingForModelTesterMixin, ModelTesterMixin, TorchCompileTesterMixin +from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..test_modeling_common import LoraHotSwappingForModelTesterMixin +from ..testing_utils import ( + AttentionTesterMixin, + BitsAndBytesTesterMixin, + ContextParallelTesterMixin, + GGUFTesterMixin, + IPAdapterTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + ModelOptTesterMixin, + ModelTesterMixin, + QuantoTesterMixin, + SingleFileTesterMixin, + TorchAoTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) enable_full_determinism() -def create_flux_ip_adapter_state_dict(model): - # "ip_adapter" (cross-attention weights) +# TODO: This standalone function maintains backward compatibility with pipeline tests +# (tests/pipelines/test_pipelines_common.py) and will be refactored. +def create_flux_ip_adapter_state_dict(model) -> dict[str, dict[str, Any]]: + """Create a dummy IP Adapter state dict for Flux transformer testing.""" ip_cross_attn_state_dict = {} key_id = 0 @@ -39,7 +58,7 @@ def create_flux_ip_adapter_state_dict(model): joint_attention_dim = model.config["joint_attention_dim"] hidden_size = model.config["num_attention_heads"] * model.config["attention_head_dim"] - sd = FluxIPAdapterJointAttnProcessor2_0( + sd = FluxIPAdapterAttnProcessor( hidden_size=hidden_size, cross_attention_dim=joint_attention_dim, scale=1.0 ).state_dict() ip_cross_attn_state_dict.update( @@ -50,11 +69,8 @@ def create_flux_ip_adapter_state_dict(model): f"{key_id}.to_v_ip.bias": sd["to_v_ip.0.bias"], } ) - key_id += 1 - # "image_proj" (ImageProjection layer weights) - image_projection = ImageProjection( cross_attention_dim=model.config["joint_attention_dim"], image_embed_dim=( @@ -75,57 +91,21 @@ def create_flux_ip_adapter_state_dict(model): ) del sd - ip_state_dict = {} - ip_state_dict.update({"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict}) - return ip_state_dict + return {"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict} -class FluxTransformerTests(ModelTesterMixin, unittest.TestCase): +class FluxTransformerTesterConfig: model_class = FluxTransformer2DModel - main_input_name = "hidden_states" - # We override the items here because the transformer under consideration is small. - model_split_percents = [0.7, 0.6, 0.6] - - # Skip setting testing with default: AttnProcessor - uses_custom_attn_processor = True - - @property - def dummy_input(self): - return self.prepare_dummy_input() + pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-pipe" + pretrained_model_kwargs = {"subfolder": "transformer"} @property - def input_shape(self): - return (16, 4) - - @property - def output_shape(self): - return (16, 4) - - def prepare_dummy_input(self, height=4, width=4): - batch_size = 1 - num_latent_channels = 4 - num_image_channels = 3 - sequence_length = 48 - embedding_dim = 32 - - hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(torch_device) - encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device) - pooled_prompt_embeds = torch.randn((batch_size, embedding_dim)).to(torch_device) - text_ids = torch.randn((sequence_length, num_image_channels)).to(torch_device) - image_ids = torch.randn((height * width, num_image_channels)).to(torch_device) - timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size) + def generator(self): + return torch.Generator("cpu").manual_seed(0) + def get_init_dict(self) -> dict[str, int | list[int]]: + """Return Flux model initialization arguments.""" return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "img_ids": image_ids, - "txt_ids": text_ids, - "pooled_projections": pooled_prompt_embeds, - "timestep": timestep, - } - - def prepare_init_args_and_inputs_for_common(self): - init_dict = { "patch_size": 1, "in_channels": 4, "num_layers": 1, @@ -137,11 +117,40 @@ def prepare_init_args_and_inputs_for_common(self): "axes_dims_rope": [4, 4, 8], } - inputs_dict = self.dummy_input - return init_dict, inputs_dict + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + batch_size = 1 + height = width = 4 + num_latent_channels = 4 + num_image_channels = 3 + sequence_length = 48 + embedding_dim = 32 + return { + "hidden_states": randn_tensor((batch_size, height * width, num_latent_channels), generator=self.generator), + "encoder_hidden_states": randn_tensor( + (batch_size, sequence_length, embedding_dim), generator=self.generator + ), + "pooled_projections": randn_tensor((batch_size, embedding_dim), generator=self.generator), + "img_ids": randn_tensor((height * width, num_image_channels), generator=self.generator), + "txt_ids": randn_tensor((sequence_length, num_image_channels), generator=self.generator), + "timestep": torch.tensor([1.0]).to(torch_device).expand(batch_size), + } + + @property + def input_shape(self) -> tuple[int, int]: + return (16, 4) + + @property + def output_shape(self) -> tuple[int, int]: + return (16, 4) + + +class TestFluxTransformer(FluxTransformerTesterConfig, ModelTesterMixin): def test_deprecated_inputs_img_txt_ids_3d(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + """Test that deprecated 3D img_ids and txt_ids still work.""" + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict) model.to(torch_device) model.eval() @@ -162,63 +171,176 @@ def test_deprecated_inputs_img_txt_ids_3d(self): with torch.no_grad(): output_2 = model(**inputs_dict).to_tuple()[0] - self.assertEqual(output_1.shape, output_2.shape) - self.assertTrue( - torch.allclose(output_1, output_2, atol=1e-5), - msg="output with deprecated inputs (img_ids and txt_ids as 3d torch tensors) are not equal as them as 2d inputs", + assert output_1.shape == output_2.shape + assert torch.allclose(output_1, output_2, atol=1e-5), ( + "output with deprecated inputs (img_ids and txt_ids as 3d torch tensors) " + "are not equal as them as 2d inputs" ) - def test_gradient_checkpointing_is_applied(self): - expected_set = {"FluxTransformer2DModel"} - super().test_gradient_checkpointing_is_applied(expected_set=expected_set) - - # The test exists for cases like - # https://github.com/huggingface/diffusers/issues/11874 - @unittest.skipIf(not is_peft_available(), "Only with PEFT") - def test_lora_exclude_modules(self): - from peft import LoraConfig, get_peft_model_state_dict, inject_adapter_in_model, set_peft_model_state_dict - - lora_rank = 4 - target_module = "single_transformer_blocks.0.proj_out" - adapter_name = "foo" - init_dict, _ = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict).to(torch_device) - - state_dict = model.state_dict() - target_mod_shape = state_dict[f"{target_module}.weight"].shape - lora_state_dict = { - f"{target_module}.lora_A.weight": torch.ones(lora_rank, target_mod_shape[1]) * 22, - f"{target_module}.lora_B.weight": torch.ones(target_mod_shape[0], lora_rank) * 33, - } - # Passing exclude_modules should no longer be necessary (or even passing target_modules, for that matter). - config = LoraConfig( - r=lora_rank, target_modules=["single_transformer_blocks.0.proj_out"], exclude_modules=["proj_out"] - ) - inject_adapter_in_model(config, model, adapter_name=adapter_name, state_dict=lora_state_dict) - set_peft_model_state_dict(model, lora_state_dict, adapter_name) - retrieved_lora_state_dict = get_peft_model_state_dict(model, adapter_name=adapter_name) - assert len(retrieved_lora_state_dict) == len(lora_state_dict) - assert (retrieved_lora_state_dict["single_transformer_blocks.0.proj_out.lora_A.weight"] == 22).all() - assert (retrieved_lora_state_dict["single_transformer_blocks.0.proj_out.lora_B.weight"] == 33).all() +class TestFluxTransformerMemory(FluxTransformerTesterConfig, MemoryTesterMixin): + """Memory optimization tests for Flux Transformer.""" + + pass + + +class TestFluxTransformerTraining(FluxTransformerTesterConfig, TrainingTesterMixin): + """Training tests for Flux Transformer.""" + + pass + + +class TestFluxTransformerAttention(FluxTransformerTesterConfig, AttentionTesterMixin): + """Attention processor tests for Flux Transformer.""" + + pass + + +class TestFluxTransformerContextParallel(FluxTransformerTesterConfig, ContextParallelTesterMixin): + """Context Parallel inference tests for Flux Transformer""" + + pass + + +class TestFluxTransformerIPAdapter(FluxTransformerTesterConfig, IPAdapterTesterMixin): + """IP Adapter tests for Flux Transformer.""" + + ip_adapter_processor_cls = FluxIPAdapterAttnProcessor + + def modify_inputs_for_ip_adapter(self, model, inputs_dict): + torch.manual_seed(0) + # Create dummy image embeds for IP adapter + cross_attention_dim = getattr(model.config, "joint_attention_dim", 32) + image_embeds = torch.randn(1, 1, cross_attention_dim).to(torch_device) + + inputs_dict.update({"joint_attention_kwargs": {"ip_adapter_image_embeds": image_embeds}}) + + return inputs_dict + + def create_ip_adapter_state_dict(self, model: Any) -> dict[str, dict[str, Any]]: + return create_flux_ip_adapter_state_dict(model) + + +class TestFluxTransformerLoRA(FluxTransformerTesterConfig, LoraTesterMixin): + """LoRA adapter tests for Flux Transformer.""" + + pass + + +class TestFluxTransformerLoRAHotSwap(FluxTransformerTesterConfig, LoraHotSwappingForModelTesterMixin): + """LoRA hot-swapping tests for Flux Transformer.""" -class FluxTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase): - model_class = FluxTransformer2DModel different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)] - def prepare_init_args_and_inputs_for_common(self): - return FluxTransformerTests().prepare_init_args_and_inputs_for_common() + def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]: + """Override to support dynamic height/width for LoRA hotswap tests.""" + batch_size = 1 + num_latent_channels = 4 + num_image_channels = 3 + sequence_length = 24 + embedding_dim = 8 - def prepare_dummy_input(self, height, width): - return FluxTransformerTests().prepare_dummy_input(height=height, width=width) + return { + "hidden_states": randn_tensor((batch_size, height * width, num_latent_channels)), + "encoder_hidden_states": randn_tensor((batch_size, sequence_length, embedding_dim)), + "pooled_projections": randn_tensor((batch_size, embedding_dim)), + "img_ids": randn_tensor((height * width, num_image_channels)), + "txt_ids": randn_tensor((sequence_length, num_image_channels)), + "timestep": torch.tensor([1.0]).to(torch_device).expand(batch_size), + } -class FluxTransformerLoRAHotSwapTests(LoraHotSwappingForModelTesterMixin, unittest.TestCase): - model_class = FluxTransformer2DModel +class TestFluxTransformerCompile(FluxTransformerTesterConfig, TorchCompileTesterMixin): different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)] - def prepare_init_args_and_inputs_for_common(self): - return FluxTransformerTests().prepare_init_args_and_inputs_for_common() + def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]: + """Override to support dynamic height/width for compilation tests.""" + batch_size = 1 + num_latent_channels = 4 + num_image_channels = 3 + sequence_length = 24 + embedding_dim = 8 - def prepare_dummy_input(self, height, width): - return FluxTransformerTests().prepare_dummy_input(height=height, width=width) + return { + "hidden_states": randn_tensor((batch_size, height * width, num_latent_channels)), + "encoder_hidden_states": randn_tensor((batch_size, sequence_length, embedding_dim)), + "pooled_projections": randn_tensor((batch_size, embedding_dim)), + "img_ids": randn_tensor((height * width, num_image_channels)), + "txt_ids": randn_tensor((sequence_length, num_image_channels)), + "timestep": torch.tensor([1.0]).to(torch_device).expand(batch_size), + } + + +class TestFluxSingleFile(FluxTransformerTesterConfig, SingleFileTesterMixin): + ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors" + alternate_keys_ckpt_paths = ["https://huggingface.co/Comfy-Org/flux1-dev/blob/main/flux1-dev-fp8.safetensors"] + pretrained_model_name_or_path = "black-forest-labs/FLUX.1-dev" + subfolder = "transformer" + pass + + +class TestFluxTransformerBitsAndBytes(FluxTransformerTesterConfig, BitsAndBytesTesterMixin): + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + return { + "hidden_states": randn_tensor((1, 4096, 64)), + "encoder_hidden_states": randn_tensor((1, 512, 4096)), + "pooled_projections": randn_tensor((1, 768)), + "timestep": torch.tensor([1.0]).to(torch_device), + "img_ids": randn_tensor((4096, 3)), + "txt_ids": randn_tensor((512, 3)), + "guidance": torch.tensor([3.5]).to(torch_device), + } + + +class TestFluxTransformerQuanto(FluxTransformerTesterConfig, QuantoTesterMixin): + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + return { + "hidden_states": randn_tensor((1, 4096, 64)), + "encoder_hidden_states": randn_tensor((1, 512, 4096)), + "pooled_projections": randn_tensor((1, 768)), + "timestep": torch.tensor([1.0]).to(torch_device), + "img_ids": randn_tensor((4096, 3)), + "txt_ids": randn_tensor((512, 3)), + "guidance": torch.tensor([3.5]).to(torch_device), + } + + +class TestFluxTransformerTorchAo(FluxTransformerTesterConfig, TorchAoTesterMixin): + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + return { + "hidden_states": randn_tensor((1, 4096, 64)), + "encoder_hidden_states": randn_tensor((1, 512, 4096)), + "pooled_projections": randn_tensor((1, 768)), + "timestep": torch.tensor([1.0]).to(torch_device), + "img_ids": randn_tensor((4096, 3)), + "txt_ids": randn_tensor((512, 3)), + "guidance": torch.tensor([3.5]).to(torch_device), + } + + +class TestFluxTransformerGGUF(FluxTransformerTesterConfig, GGUFTesterMixin): + gguf_filename = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q8_0.gguf" + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + return { + "hidden_states": randn_tensor((1, 4096, 64)), + "encoder_hidden_states": randn_tensor((1, 512, 4096)), + "pooled_projections": randn_tensor((1, 768)), + "timestep": torch.tensor([1.0]).to(torch_device), + "img_ids": randn_tensor((4096, 3)), + "txt_ids": randn_tensor((512, 3)), + "guidance": torch.tensor([3.5]).to(torch_device), + } + + +class TestFluxTransformerModelOpt(FluxTransformerTesterConfig, ModelOptTesterMixin): + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + return { + "hidden_states": randn_tensor((1, 4096, 64)), + "encoder_hidden_states": randn_tensor((1, 512, 4096)), + "pooled_projections": randn_tensor((1, 768)), + "timestep": torch.tensor([1.0]).to(torch_device), + "img_ids": randn_tensor((4096, 3)), + "txt_ids": randn_tensor((512, 3)), + "guidance": torch.tensor([3.5]).to(torch_device), + } diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 6ed7e3467d7f..9860d64dc119 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -38,6 +38,7 @@ is_gguf_available, is_kernels_available, is_note_seq_available, + is_nvidia_modelopt_available, is_onnx_available, is_opencv_available, is_optimum_quanto_available, @@ -241,7 +242,6 @@ def parse_flag_from_env(key, default=False): _run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False) _run_nightly_tests = parse_flag_from_env("RUN_NIGHTLY", default=False) -_run_compile_tests = parse_flag_from_env("RUN_COMPILE", default=False) def floats_tensor(shape, scale=1.0, rng=None, name=None): @@ -282,12 +282,137 @@ def nightly(test_case): def is_torch_compile(test_case): """ - Decorator marking a test that runs compile tests in the diffusers CI. + Decorator marking a test as a torch.compile test. These tests can be filtered using: + pytest -m "not compile" to skip + pytest -m compile to run only these tests + """ + return pytest.mark.compile(test_case) + + +def is_single_file(test_case): + """ + Decorator marking a test as a single file loading test. These tests can be filtered using: + pytest -m "not single_file" to skip + pytest -m single_file to run only these tests + """ + return pytest.mark.single_file(test_case) + + +def is_lora(test_case): + """ + Decorator marking a test as a LoRA test. These tests can be filtered using: + pytest -m "not lora" to skip + pytest -m lora to run only these tests + """ + return pytest.mark.lora(test_case) + + +def is_ip_adapter(test_case): + """ + Decorator marking a test as an IP Adapter test. These tests can be filtered using: + pytest -m "not ip_adapter" to skip + pytest -m ip_adapter to run only these tests + """ + return pytest.mark.ip_adapter(test_case) - Compile tests are skipped by default. Set the RUN_COMPILE environment variable to a truthy value to run them. +def is_training(test_case): + """ + Decorator marking a test as a training test. These tests can be filtered using: + pytest -m "not training" to skip + pytest -m training to run only these tests """ - return pytest.mark.skipif(not _run_compile_tests, reason="test is torch compile")(test_case) + return pytest.mark.training(test_case) + + +def is_attention(test_case): + """ + Decorator marking a test as an attention test. These tests can be filtered using: + pytest -m "not attention" to skip + pytest -m attention to run only these tests + """ + return pytest.mark.attention(test_case) + + +def is_memory(test_case): + """ + Decorator marking a test as a memory optimization test. These tests can be filtered using: + pytest -m "not memory" to skip + pytest -m memory to run only these tests + """ + return pytest.mark.memory(test_case) + + +def is_cpu_offload(test_case): + """ + Decorator marking a test as a CPU offload test. These tests can be filtered using: + pytest -m "not cpu_offload" to skip + pytest -m cpu_offload to run only these tests + """ + return pytest.mark.cpu_offload(test_case) + + +def is_group_offload(test_case): + """ + Decorator marking a test as a group offload test. These tests can be filtered using: + pytest -m "not group_offload" to skip + pytest -m group_offload to run only these tests + """ + return pytest.mark.group_offload(test_case) + + +def is_bitsandbytes(test_case): + """ + Decorator marking a test as a BitsAndBytes quantization test. These tests can be filtered using: + pytest -m "not bitsandbytes" to skip + pytest -m bitsandbytes to run only these tests + """ + return pytest.mark.bitsandbytes(test_case) + + +def is_quanto(test_case): + """ + Decorator marking a test as a Quanto quantization test. These tests can be filtered using: + pytest -m "not quanto" to skip + pytest -m quanto to run only these tests + """ + return pytest.mark.quanto(test_case) + + +def is_torchao(test_case): + """ + Decorator marking a test as a TorchAO quantization test. These tests can be filtered using: + pytest -m "not torchao" to skip + pytest -m torchao to run only these tests + """ + return pytest.mark.torchao(test_case) + + +def is_gguf(test_case): + """ + Decorator marking a test as a GGUF quantization test. These tests can be filtered using: + pytest -m "not gguf" to skip + pytest -m gguf to run only these tests + """ + return pytest.mark.gguf(test_case) + + +def is_modelopt(test_case): + """ + Decorator marking a test as a NVIDIA ModelOpt quantization test. These tests can be filtered using: + pytest -m "not modelopt" to skip + pytest -m modelopt to run only these tests + """ + return pytest.mark.modelopt(test_case) + + +def is_context_parallel(test_case): + """ + Decorator marking a test as a context parallel inference test. These tests can be filtered using: + pytest -m "not context_parallel" to skip + pytest -m context_parallel to run only these tests + """ + return pytest.mark.context_parallel(test_case) def require_torch(test_case): @@ -650,6 +775,19 @@ def decorator(test_case): return decorator +def require_modelopt_version_greater_or_equal(modelopt_version): + def decorator(test_case): + correct_nvidia_modelopt_version = is_nvidia_modelopt_available() and version.parse( + version.parse(importlib.metadata.version("modelopt")).base_version + ) >= version.parse(modelopt_version) + return pytest.mark.skipif( + not correct_nvidia_modelopt_version, + f"Test requires modelopt with version greater than {modelopt_version}.", + )(test_case) + + return decorator + + def deprecate_after_peft_backend(test_case): """ Decorator marking a test that will be skipped after PEFT backend diff --git a/utils/generate_model_tests.py b/utils/generate_model_tests.py new file mode 100644 index 000000000000..ffd600dfdf29 --- /dev/null +++ b/utils/generate_model_tests.py @@ -0,0 +1,489 @@ +#!/usr/bin/env 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. + +""" +Utility script to generate test suites for diffusers model classes. + +Usage: + python utils/generate_model_tests.py src/diffusers/models/transformers/transformer_flux.py + +This will analyze the model file and generate a test file with appropriate +test classes based on the model's mixins and attributes. +""" + +import argparse +import ast +import sys +from pathlib import Path + + +MIXIN_TO_TESTER = { + "ModelMixin": "ModelTesterMixin", + "PeftAdapterMixin": "LoraTesterMixin", +} + +ATTRIBUTE_TO_TESTER = { + "_cp_plan": "ContextParallelTesterMixin", + "_supports_gradient_checkpointing": "TrainingTesterMixin", +} + +ALWAYS_INCLUDE_TESTERS = [ + "ModelTesterMixin", + "MemoryTesterMixin", + "AttentionTesterMixin", + "TorchCompileTesterMixin", +] + +OPTIONAL_TESTERS = [ + ("BitsAndBytesTesterMixin", "bnb"), + ("QuantoTesterMixin", "quanto"), + ("TorchAoTesterMixin", "torchao"), + ("GGUFTesterMixin", "gguf"), + ("ModelOptTesterMixin", "modelopt"), + ("SingleFileTesterMixin", "single_file"), + ("IPAdapterTesterMixin", "ip_adapter"), +] + + +class ModelAnalyzer(ast.NodeVisitor): + def __init__(self): + self.model_classes = [] + self.current_class = None + + def visit_ClassDef(self, node: ast.ClassDef): + base_names = [] + for base in node.bases: + if isinstance(base, ast.Name): + base_names.append(base.id) + elif isinstance(base, ast.Attribute): + base_names.append(base.attr) + + if "ModelMixin" in base_names: + class_info = { + "name": node.name, + "bases": base_names, + "attributes": {}, + "has_forward": False, + "init_params": [], + } + + for item in node.body: + if isinstance(item, ast.Assign): + for target in item.targets: + if isinstance(target, ast.Name): + attr_name = target.id + if attr_name.startswith("_"): + class_info["attributes"][attr_name] = self._get_value(item.value) + + elif isinstance(item, ast.FunctionDef): + if item.name == "forward": + class_info["has_forward"] = True + class_info["forward_params"] = self._extract_func_params(item) + elif item.name == "__init__": + class_info["init_params"] = self._extract_func_params(item) + + self.model_classes.append(class_info) + + self.generic_visit(node) + + def _extract_func_params(self, func_node: ast.FunctionDef) -> list[dict]: + params = [] + args = func_node.args + + num_defaults = len(args.defaults) + num_args = len(args.args) + first_default_idx = num_args - num_defaults + + for i, arg in enumerate(args.args): + if arg.arg == "self": + continue + + param_info = {"name": arg.arg, "type": None, "default": None} + + if arg.annotation: + param_info["type"] = self._get_annotation_str(arg.annotation) + + default_idx = i - first_default_idx + if default_idx >= 0 and default_idx < len(args.defaults): + param_info["default"] = self._get_value(args.defaults[default_idx]) + + params.append(param_info) + + return params + + def _get_annotation_str(self, node) -> str: + if isinstance(node, ast.Name): + return node.id + elif isinstance(node, ast.Constant): + return repr(node.value) + elif isinstance(node, ast.Subscript): + base = self._get_annotation_str(node.value) + if isinstance(node.slice, ast.Tuple): + args = ", ".join(self._get_annotation_str(el) for el in node.slice.elts) + else: + args = self._get_annotation_str(node.slice) + return f"{base}[{args}]" + elif isinstance(node, ast.Attribute): + return f"{self._get_annotation_str(node.value)}.{node.attr}" + elif isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr): + left = self._get_annotation_str(node.left) + right = self._get_annotation_str(node.right) + return f"{left} | {right}" + elif isinstance(node, ast.Tuple): + return ", ".join(self._get_annotation_str(el) for el in node.elts) + return "Any" + + def _get_value(self, node): + if isinstance(node, ast.Constant): + return node.value + elif isinstance(node, ast.Name): + if node.id == "None": + return None + elif node.id == "True": + return True + elif node.id == "False": + return False + return node.id + elif isinstance(node, ast.List): + return [self._get_value(el) for el in node.elts] + elif isinstance(node, ast.Dict): + return {self._get_value(k): self._get_value(v) for k, v in zip(node.keys, node.values)} + return "" + + +def analyze_model_file(filepath: str) -> list[dict]: + with open(filepath) as f: + source = f.read() + + tree = ast.parse(source) + analyzer = ModelAnalyzer() + analyzer.visit(tree) + + return analyzer.model_classes + + +def determine_testers(model_info: dict, include_optional: list[str]) -> list[str]: + testers = list(ALWAYS_INCLUDE_TESTERS) + + for base in model_info["bases"]: + if base in MIXIN_TO_TESTER: + tester = MIXIN_TO_TESTER[base] + if tester not in testers: + testers.append(tester) + + for attr, tester in ATTRIBUTE_TO_TESTER.items(): + if attr in model_info["attributes"]: + value = model_info["attributes"][attr] + if value is not None and value is not False: + if tester not in testers: + testers.append(tester) + + if "_cp_plan" in model_info["attributes"] and model_info["attributes"]["_cp_plan"] is not None: + if "ContextParallelTesterMixin" not in testers: + testers.append("ContextParallelTesterMixin") + + for tester, flag in OPTIONAL_TESTERS: + if flag in include_optional: + if tester not in testers: + testers.append(tester) + + return testers + + +def generate_config_class(model_info: dict, model_name: str) -> str: + class_name = f"{model_name}TesterConfig" + model_class = model_info["name"] + forward_params = model_info.get("forward_params", []) + init_params = model_info.get("init_params", []) + + lines = [ + f"class {class_name}:", + f" model_class = {model_class}", + ' pretrained_model_name_or_path = ""', + ' pretrained_model_kwargs = {"subfolder": "transformer"}', + "", + " @property", + " def generator(self):", + ' return torch.Generator("cpu").manual_seed(0)', + "", + " def get_init_dict(self) -> dict[str, int | list[int]]:", + ] + + if init_params: + lines.append(" # __init__ parameters:") + for param in init_params: + type_str = f": {param['type']}" if param["type"] else "" + default_str = f" = {param['default']}" if param["default"] is not None else "" + lines.append(f" # {param['name']}{type_str}{default_str}") + + lines.extend( + [ + " return {}", + "", + " def get_dummy_inputs(self) -> dict[str, torch.Tensor]:", + ] + ) + + if forward_params: + lines.append(" # forward() parameters:") + for param in forward_params: + type_str = f": {param['type']}" if param["type"] else "" + default_str = f" = {param['default']}" if param["default"] is not None else "" + lines.append(f" # {param['name']}{type_str}{default_str}") + + lines.extend( + [ + " # TODO: Fill in dummy inputs", + " return {}", + "", + " @property", + " def input_shape(self) -> tuple[int, ...]:", + " return (1, 1)", + "", + " @property", + " def output_shape(self) -> tuple[int, ...]:", + " return (1, 1)", + ] + ) + + return "\n".join(lines) + + +def generate_test_class(model_name: str, config_class: str, tester: str) -> str: + tester_short = tester.replace("TesterMixin", "") + class_name = f"Test{model_name}{tester_short}" + + lines = [f"class {class_name}({config_class}, {tester}):"] + + if tester == "TorchCompileTesterMixin": + lines.extend( + [ + " different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]", + "", + " def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:", + " # TODO: Implement dynamic input generation", + " return {}", + ] + ) + elif tester == "IPAdapterTesterMixin": + lines.extend( + [ + " ip_adapter_processor_cls = None # TODO: Set processor class", + "", + " def modify_inputs_for_ip_adapter(self, model, inputs_dict):", + " # TODO: Add IP adapter image embeds to inputs", + " return inputs_dict", + "", + " def create_ip_adapter_state_dict(self, model):", + " # TODO: Create IP adapter state dict", + " return {}", + ] + ) + elif tester == "SingleFileTesterMixin": + lines.extend( + [ + ' ckpt_path = "" # TODO: Set checkpoint path', + " alternate_keys_ckpt_paths = []", + ' pretrained_model_name_or_path = ""', + ' subfolder = "transformer"', + ] + ) + elif tester == "GGUFTesterMixin": + lines.extend( + [ + ' gguf_filename = "" # TODO: Set GGUF filename', + "", + " def get_dummy_inputs(self) -> dict[str, torch.Tensor]:", + " # TODO: Override with larger inputs for quantization tests", + " return {}", + ] + ) + elif tester in ["BitsAndBytesTesterMixin", "QuantoTesterMixin", "TorchAoTesterMixin", "ModelOptTesterMixin"]: + lines.extend( + [ + " def get_dummy_inputs(self) -> dict[str, torch.Tensor]:", + " # TODO: Override with larger inputs for quantization tests", + " return {}", + ] + ) + elif tester == "LoraHotSwappingForModelTesterMixin": + lines.extend( + [ + " different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]", + "", + " def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:", + " # TODO: Implement dynamic input generation", + " return {}", + ] + ) + else: + lines.append(" pass") + + return "\n".join(lines) + + +def generate_test_file(model_info: dict, model_filepath: str, include_optional: list[str]) -> str: + model_name = model_info["name"].replace("2DModel", "").replace("3DModel", "").replace("Model", "") + testers = determine_testers(model_info, include_optional) + tester_imports = sorted(set(testers) - {"LoraHotSwappingForModelTesterMixin"}) + + lines = [ + "# 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 torch", + "", + f"from diffusers import {model_info['name']}", + "from diffusers.utils.torch_utils import randn_tensor", + "", + "from ...testing_utils import enable_full_determinism, torch_device", + ] + + if "LoraTesterMixin" in testers: + lines.append("from ..test_modeling_common import LoraHotSwappingForModelTesterMixin") + + lines.extend( + [ + "from ..testing_utils import (", + *[f" {tester}," for tester in sorted(tester_imports)], + ")", + "", + "", + "enable_full_determinism()", + "", + "", + ] + ) + + config_class = f"{model_name}TesterConfig" + lines.append(generate_config_class(model_info, model_name)) + lines.append("") + lines.append("") + + for tester in testers: + lines.append(generate_test_class(model_name, config_class, tester)) + lines.append("") + lines.append("") + + if "LoraTesterMixin" in testers: + lines.append(generate_test_class(model_name, config_class, "LoraHotSwappingForModelTesterMixin")) + lines.append("") + lines.append("") + + return "\n".join(lines).rstrip() + "\n" + + +def get_test_output_path(model_filepath: str) -> str: + path = Path(model_filepath) + model_filename = path.stem + + if "transformers" in path.parts: + return f"tests/models/transformers/test_models_{model_filename}.py" + elif "unets" in path.parts: + return f"tests/models/unets/test_models_{model_filename}.py" + elif "autoencoders" in path.parts: + return f"tests/models/autoencoders/test_models_{model_filename}.py" + else: + return f"tests/models/test_models_{model_filename}.py" + + +def main(): + parser = argparse.ArgumentParser(description="Generate test suite for a diffusers model class") + parser.add_argument( + "model_filepath", + type=str, + help="Path to the model file (e.g., src/diffusers/models/transformers/transformer_flux.py)", + ) + parser.add_argument( + "--output", "-o", type=str, default=None, help="Output file path (default: auto-generated based on model path)" + ) + parser.add_argument( + "--include", + "-i", + type=str, + nargs="*", + default=[], + choices=["compile", "bnb", "quanto", "torchao", "gguf", "modelopt", "single_file", "ip_adapter", "all"], + help="Optional testers to include", + ) + parser.add_argument( + "--class-name", + "-c", + type=str, + default=None, + help="Specific model class to generate tests for (default: first model class found)", + ) + parser.add_argument("--dry-run", action="store_true", help="Print generated code without writing to file") + + args = parser.parse_args() + + if not Path(args.model_filepath).exists(): + print(f"Error: File not found: {args.model_filepath}", file=sys.stderr) + sys.exit(1) + + model_classes = analyze_model_file(args.model_filepath) + + if not model_classes: + print(f"Error: No model classes found in {args.model_filepath}", file=sys.stderr) + sys.exit(1) + + if args.class_name: + model_info = next((m for m in model_classes if m["name"] == args.class_name), None) + if not model_info: + available = [m["name"] for m in model_classes] + print(f"Error: Class '{args.class_name}' not found. Available: {available}", file=sys.stderr) + sys.exit(1) + else: + model_info = model_classes[0] + if len(model_classes) > 1: + print(f"Multiple model classes found, using: {model_info['name']}", file=sys.stderr) + print("Use --class-name to specify a different class", file=sys.stderr) + + include_optional = args.include + if "all" in include_optional: + include_optional = [flag for _, flag in OPTIONAL_TESTERS] + + generated_code = generate_test_file(model_info, args.model_filepath, include_optional) + + if args.dry_run: + print(generated_code) + else: + output_path = args.output or get_test_output_path(args.model_filepath) + output_dir = Path(output_path).parent + output_dir.mkdir(parents=True, exist_ok=True) + + with open(output_path, "w") as f: + f.write(generated_code) + + print(f"Generated test file: {output_path}") + print(f"Model class: {model_info['name']}") + print(f"Detected attributes: {list(model_info['attributes'].keys())}") + + +if __name__ == "__main__": + main()