diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 918126fe3f13..b237a0da76ee 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -13,7 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import gc +import os import pytest import safetensors.torch @@ -72,6 +74,7 @@ from diffusers.loaders.peft import PeftAdapterMixin if is_nvidia_modelopt_available(): + import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq if is_bitsandbytes_available(): @@ -1273,9 +1276,9 @@ class ModelOptConfigMixin: """ MODELOPT_CONFIGS = { - "fp8": {"quant_type": "FP8"}, - "int8": {"quant_type": "INT8"}, - "int4": {"quant_type": "INT4"}, + "fp8": {"quant_type": "FP8", "disable_conv_quantization": True}, + "int8": {"quant_type": "INT8", "disable_conv_quantization": True}, + "int4": {"quant_type": "INT4", "disable_conv_quantization": True}, } MODELOPT_EXPECTED_MEMORY_REDUCTIONS = { @@ -1374,6 +1377,51 @@ def test_modelopt_dequantize(self): """Test that dequantize() works correctly.""" self._test_dequantize(ModelOptConfigMixin.MODELOPT_CONFIGS["fp8"]) + @torch.no_grad() + def test_modelopt_keep_modules_in_fp32(self): + fp32_modules = getattr(self.model_class, "_keep_in_fp32_modules", None) + if not fp32_modules: + pytest.skip(f"{self.model_class.__name__} does not declare _keep_in_fp32_modules") + + model = self._create_quantized_model(ModelOptConfigMixin.MODELOPT_CONFIGS["fp8"]) + model.to(torch_device) + + for name, module in model.named_modules(): + if isinstance(module, torch.nn.Linear): + if any(fp32_name in name for fp32_name in fp32_modules): + assert module.weight.dtype == torch.float32, ( + f"Module {name} should be FP32 but is {module.weight.dtype}" + ) + + def test_modelopt_training(self): + """Test that quantized models can be used for training with adapters.""" + self._test_quantization_training(ModelOptConfigMixin.MODELOPT_CONFIGS["fp8"]) + + @require_modelopt_version_greater_or_equal("0.44.0") + def test_modelopt_prequantized_serialization(self, tmp_path): + """Test that a pre-quantized ModelOpt checkpoint round-trips through save/load with a device_map.""" + mto.enable_huggingface_checkpointing() + + config_kwargs = {"quant_type": "FP8", "modelopt_config": copy.deepcopy(mtq.FP8_DEFAULT_CFG)} + model = self._create_quantized_model(config_kwargs) + + model.save_pretrained(str(tmp_path)) + assert os.path.isfile(os.path.join(str(tmp_path), "modelopt_state.pth")), ( + "Expected modelopt_state.pth in the saved checkpoint." + ) + + saved_model = self.model_class.from_pretrained(str(tmp_path), device_map=str(torch_device)) + + named_parameters = list(saved_model.named_parameters()) + named_buffers = list(saved_model.named_buffers()) + assert any(name.endswith(("_amax", "_scale")) for name, _ in named_buffers), ( + "The restored model did not contain ModelOpt quantizer buffers." + ) + + for tensor_kind, named_tensors in (("parameter", named_parameters), ("buffer", named_buffers)): + for name, tensor in named_tensors: + assert not tensor.is_meta, f"{tensor_kind} {name} was not materialized from meta." + @is_quantization @is_sdnq diff --git a/tests/models/transformers/test_models_transformer_sd3.py b/tests/models/transformers/test_models_transformer_sd3.py index 6294bf80635a..970ed465e71d 100644 --- a/tests/models/transformers/test_models_transformer_sd3.py +++ b/tests/models/transformers/test_models_transformer_sd3.py @@ -23,6 +23,8 @@ BaseModelTesterConfig, BitsAndBytesTesterMixin, LoraTesterMixin, + ModelOptCompileTesterMixin, + ModelOptTesterMixin, ModelTesterMixin, SingleFileTesterMixin, TorchAoTesterMixin, @@ -121,6 +123,14 @@ class TestSD3TransformerCompile(SD3TransformerTesterConfig, TorchCompileTesterMi pass +class TestSD3TransformerModelOpt(SD3TransformerTesterConfig, ModelOptTesterMixin): + """NVIDIA ModelOpt quantization tests for SD3 Transformer.""" + + +class TestSD3TransformerModelOptCompile(SD3TransformerTesterConfig, ModelOptCompileTesterMixin): + """torch.compile tests for NVIDIA ModelOpt-quantized SD3 Transformer.""" + + # ======================== SD3.5 Transformer ======================== diff --git a/tests/quantization/modelopt/test_modelopt.py b/tests/quantization/modelopt/test_modelopt.py new file mode 100644 index 000000000000..9002781d5517 --- /dev/null +++ b/tests/quantization/modelopt/test_modelopt.py @@ -0,0 +1,110 @@ +import gc + +import pytest + +from diffusers import NVIDIAModelOptConfig, SD3Transformer2DModel, StableDiffusion3Pipeline +from diffusers.utils import is_torch_available + +from ...testing_utils import ( + backend_empty_cache, + backend_reset_peak_memory_stats, + enable_full_determinism, + nightly, + require_accelerate, + require_big_accelerator, + require_modelopt_version_greater_or_equal, + require_torch_cuda_compatibility, + torch_device, +) + + +if is_torch_available(): + import torch + +enable_full_determinism() + + +# Model-level ModelOpt tests live in `tests/models/testing_utils/quantization.py` +# (`ModelOptTesterMixin` / `ModelOptCompileTesterMixin`), wired into model test files via concrete +# classes (e.g. `TestSD3TransformerModelOpt`). Only pipeline-level coverage remains here. +@nightly +@require_big_accelerator +@require_accelerate +@require_modelopt_version_greater_or_equal("0.33.1") +class ModelOptBaseTesterMixin: + model_id = "hf-internal-testing/tiny-sd3-pipe" + model_cls = SD3Transformer2DModel + pipeline_cls = StableDiffusion3Pipeline + torch_dtype = torch.bfloat16 + + @pytest.fixture(autouse=True) + def _setup(self): + backend_reset_peak_memory_stats(torch_device) + backend_empty_cache(torch_device) + gc.collect() + yield + backend_reset_peak_memory_stats(torch_device) + backend_empty_cache(torch_device) + gc.collect() + + def get_dummy_init_kwargs(self): + return {"quant_type": "FP8"} + + def test_model_cpu_offload(self): + init_kwargs = self.get_dummy_init_kwargs() + transformer = self.model_cls.from_pretrained( + self.model_id, + quantization_config=NVIDIAModelOptConfig(**init_kwargs), + subfolder="transformer", + torch_dtype=torch.bfloat16, + ) + pipe = self.pipeline_cls.from_pretrained(self.model_id, transformer=transformer, torch_dtype=torch.bfloat16) + pipe.enable_model_cpu_offload(device=torch_device) + _ = pipe("a cat holding a sign that says hello", num_inference_steps=2) + + +class TestSanaTransformerFP8Weights(ModelOptBaseTesterMixin): + def get_dummy_init_kwargs(self): + return {"quant_type": "FP8"} + + +class TestSanaTransformerINT8Weights(ModelOptBaseTesterMixin): + def get_dummy_init_kwargs(self): + return {"quant_type": "INT8"} + + +@require_torch_cuda_compatibility(8.0) +class TestSanaTransformerINT4Weights(ModelOptBaseTesterMixin): + def get_dummy_init_kwargs(self): + return { + "quant_type": "INT4", + "block_quantize": 128, + "channel_quantize": -1, + "disable_conv_quantization": True, + } + + +@require_torch_cuda_compatibility(8.0) +class TestSanaTransformerNF4Weights(ModelOptBaseTesterMixin): + def get_dummy_init_kwargs(self): + return { + "quant_type": "NF4", + "block_quantize": 128, + "channel_quantize": -1, + "scale_block_quantize": 8, + "scale_channel_quantize": -1, + "modules_to_not_convert": ["conv"], + } + + +@require_torch_cuda_compatibility(8.0) +class TestSanaTransformerNVFP4Weights(ModelOptBaseTesterMixin): + def get_dummy_init_kwargs(self): + return { + "quant_type": "NVFP4", + "block_quantize": 128, + "channel_quantize": -1, + "scale_block_quantize": 8, + "scale_channel_quantize": -1, + "modules_to_not_convert": ["conv"], + }