From d52f2773ef03295eec356e4c1376730717845960 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 21 Jul 2026 05:03:57 +0000 Subject: [PATCH] [tests] migrate model-level ModelOpt quantization tests to model test mixins Move the model-level tests out of tests/quantization/modelopt/test_modelopt.py (the ModelOptBaseTesterMixin subclasses) into the standardized model-test mixin structure: - Extend ModelOptTesterMixin with keep_modules_in_fp32, training and prequantized-serialization coverage, and set disable_conv_quantization in the shared MODELOPT_CONFIGS so conv-containing models (SD3 patch embed) don't break. - Wire concrete SD3 classes: TestSD3TransformerModelOpt and TestSD3TransformerModelOptCompile. - Trim test_modelopt.py to the remaining pipeline-level test_model_cpu_offload (kept per quant type, incl. NF4/NVFP4). Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/models/testing_utils/quantization.py | 54 +++- .../test_models_transformer_sd3.py | 10 + tests/quantization/modelopt/test_modelopt.py | 241 +----------------- 3 files changed, 65 insertions(+), 240 deletions(-) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index a458fc7b2bef..b5b3814b1165 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 torch @@ -62,6 +64,7 @@ if is_nvidia_modelopt_available(): + import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq if is_bitsandbytes_available(): @@ -1103,9 +1106,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 = { @@ -1204,6 +1207,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_torch_compile diff --git a/tests/models/transformers/test_models_transformer_sd3.py b/tests/models/transformers/test_models_transformer_sd3.py index e38c7853a613..1094ec3ab740 100644 --- a/tests/models/transformers/test_models_transformer_sd3.py +++ b/tests/models/transformers/test_models_transformer_sd3.py @@ -22,6 +22,8 @@ from ..testing_utils import ( BaseModelTesterConfig, BitsAndBytesTesterMixin, + ModelOptCompileTesterMixin, + ModelOptTesterMixin, ModelTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, @@ -119,6 +121,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 index 8f104c0bb898..9002781d5517 100644 --- a/tests/quantization/modelopt/test_modelopt.py +++ b/tests/quantization/modelopt/test_modelopt.py @@ -1,19 +1,15 @@ -import copy import gc -import os -import tempfile import pytest from diffusers import NVIDIAModelOptConfig, SD3Transformer2DModel, StableDiffusion3Pipeline -from diffusers.utils import is_nvidia_modelopt_available, is_torch_available +from diffusers.utils import is_torch_available from ...testing_utils import ( backend_empty_cache, backend_reset_peak_memory_stats, enable_full_determinism, nightly, - numpy_cosine_similarity_distance, require_accelerate, require_big_accelerator, require_modelopt_version_greater_or_equal, @@ -22,18 +18,15 @@ ) -if is_nvidia_modelopt_available(): - import modelopt.torch.opt as mto - import modelopt.torch.quantization as mtq - if is_torch_available(): import torch - from ..utils import LoRALayer, get_memory_consumption_stat - 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 @@ -43,10 +36,6 @@ class ModelOptBaseTesterMixin: model_cls = SD3Transformer2DModel pipeline_cls = StableDiffusion3Pipeline torch_dtype = torch.bfloat16 - expected_memory_reduction = 0.0 - keep_in_fp32_module = "" - modules_to_not_convert = "" - _test_torch_compile = False @pytest.fixture(autouse=True) def _setup(self): @@ -61,154 +50,6 @@ def _setup(self): def get_dummy_init_kwargs(self): return {"quant_type": "FP8"} - def get_dummy_model_init_kwargs(self): - return { - "pretrained_model_name_or_path": self.model_id, - "torch_dtype": self.torch_dtype, - "quantization_config": NVIDIAModelOptConfig(**self.get_dummy_init_kwargs()), - "subfolder": "transformer", - } - - def test_modelopt_layers(self): - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear): - assert mtq.utils.is_quantized(module) - - def test_modelopt_memory_usage(self): - inputs = self.get_dummy_inputs() - inputs = { - k: v.to(device=torch_device, dtype=torch.bfloat16) for k, v in inputs.items() if not isinstance(v, bool) - } - - unquantized_model = self.model_cls.from_pretrained( - self.model_id, torch_dtype=self.torch_dtype, subfolder="transformer" - ) - unquantized_model.to(torch_device) - unquantized_model_memory = get_memory_consumption_stat(unquantized_model, inputs) - - quantized_model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - quantized_model.to(torch_device) - quantized_model_memory = get_memory_consumption_stat(quantized_model, inputs) - - assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_reduction - - def test_keep_modules_in_fp32(self): - _keep_in_fp32_modules = self.model_cls._keep_in_fp32_modules - self.model_cls._keep_in_fp32_modules = self.keep_in_fp32_module - - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - model.to(torch_device) - - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear): - if name in model._keep_in_fp32_modules: - assert module.weight.dtype == torch.float32 - self.model_cls._keep_in_fp32_modules = _keep_in_fp32_modules - - def test_modules_to_not_convert(self): - init_kwargs = self.get_dummy_model_init_kwargs() - quantization_config_kwargs = self.get_dummy_init_kwargs() - quantization_config_kwargs.update({"modules_to_not_convert": self.modules_to_not_convert}) - quantization_config = NVIDIAModelOptConfig(**quantization_config_kwargs) - init_kwargs.update({"quantization_config": quantization_config}) - - model = self.model_cls.from_pretrained(**init_kwargs) - model.to(torch_device) - - for name, module in model.named_modules(): - if name in self.modules_to_not_convert: - assert not mtq.utils.is_quantized(module) - - def test_dtype_assignment(self): - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_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_serialization(self): - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - inputs = self.get_dummy_inputs() - - model.to(torch_device) - with torch.no_grad(): - model_output = model(**inputs) - - with tempfile.TemporaryDirectory() as tmp_dir: - model.save_pretrained(tmp_dir) - saved_model = self.model_cls.from_pretrained( - tmp_dir, - torch_dtype=torch.bfloat16, - ) - - saved_model.to(torch_device) - with torch.no_grad(): - saved_model_output = saved_model(**inputs) - - assert torch.allclose(model_output.sample, saved_model_output.sample, rtol=1e-5, atol=1e-5) - - def test_torch_compile(self): - if not self._test_torch_compile: - return - - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - compiled_model = torch.compile(model, mode="max-autotune", fullgraph=True, dynamic=False) - - model.to(torch_device) - with torch.no_grad(): - model_output = model(**self.get_dummy_inputs()).sample - - compiled_model.to(torch_device) - with torch.no_grad(): - compiled_model_output = compiled_model(**self.get_dummy_inputs()).sample - - model_output = model_output.detach().float().cpu().numpy() - compiled_model_output = compiled_model_output.detach().float().cpu().numpy() - - max_diff = numpy_cosine_similarity_distance(model_output.flatten(), compiled_model_output.flatten()) - assert max_diff < 1e-3 - - def test_device_map_error(self): - with pytest.raises(ValueError): - _ = self.model_cls.from_pretrained( - **self.get_dummy_model_init_kwargs(), - device_map={0: "8GB", "cpu": "16GB"}, - ) - - def get_dummy_inputs(self): - batch_size = 1 - seq_len = 16 - height = width = 32 - num_latent_channels = 4 - caption_channels = 8 - - torch.manual_seed(0) - hidden_states = torch.randn((batch_size, num_latent_channels, height, width)).to( - torch_device, dtype=torch.bfloat16 - ) - encoder_hidden_states = torch.randn((batch_size, seq_len, caption_channels)).to( - torch_device, dtype=torch.bfloat16 - ) - timestep = torch.tensor([1.0]).to(torch_device, dtype=torch.bfloat16).expand(batch_size) - - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "timestep": timestep, - } - def test_model_cpu_offload(self): init_kwargs = self.get_dummy_init_kwargs() transformer = self.model_cls.from_pretrained( @@ -221,89 +62,19 @@ def test_model_cpu_offload(self): pipe.enable_model_cpu_offload(device=torch_device) _ = pipe("a cat holding a sign that says hello", num_inference_steps=2) - def test_training(self): - quantization_config = NVIDIAModelOptConfig(**self.get_dummy_init_kwargs()) - quantized_model = self.model_cls.from_pretrained( - self.model_id, - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=torch.bfloat16, - ).to(torch_device) - - for param in quantized_model.parameters(): - param.requires_grad = False - if param.ndim == 1: - param.data = param.data.to(torch.float32) - - for _, module in quantized_model.named_modules(): - if hasattr(module, "to_q"): - module.to_q = LoRALayer(module.to_q, rank=4) - if hasattr(module, "to_k"): - module.to_k = LoRALayer(module.to_k, rank=4) - if hasattr(module, "to_v"): - module.to_v = LoRALayer(module.to_v, rank=4) - - with torch.amp.autocast(str(torch_device), dtype=torch.bfloat16): - inputs = self.get_dummy_inputs() - output = quantized_model(**inputs)[0] - output.norm().backward() - - for module in quantized_model.modules(): - if isinstance(module, LoRALayer): - assert module.adapter[1].weight.grad is not None - class TestSanaTransformerFP8Weights(ModelOptBaseTesterMixin): - expected_memory_reduction = 0.6 - def get_dummy_init_kwargs(self): return {"quant_type": "FP8"} - @require_modelopt_version_greater_or_equal("0.44.0") - def test_prequantized_serialization_with_device_map(self): - mto.enable_huggingface_checkpointing() - model = self.model_cls.from_pretrained( - self.model_id, - subfolder="transformer", - torch_dtype=self.torch_dtype, - quantization_config=NVIDIAModelOptConfig( - quant_type="FP8", modelopt_config=copy.deepcopy(mtq.FP8_DEFAULT_CFG) - ), - ) - model.to(torch_device) - - with tempfile.TemporaryDirectory() as tmp_dir: - model.save_pretrained(tmp_dir) - assert os.path.isfile(os.path.join(tmp_dir, "modelopt_state.pth")) - saved_model = self.model_cls.from_pretrained( - tmp_dir, - torch_dtype=self.torch_dtype, - device_map=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." - class TestSanaTransformerINT8Weights(ModelOptBaseTesterMixin): - expected_memory_reduction = 0.6 - _test_torch_compile = True - def get_dummy_init_kwargs(self): return {"quant_type": "INT8"} @require_torch_cuda_compatibility(8.0) class TestSanaTransformerINT4Weights(ModelOptBaseTesterMixin): - expected_memory_reduction = 0.55 - def get_dummy_init_kwargs(self): return { "quant_type": "INT4", @@ -315,8 +86,6 @@ def get_dummy_init_kwargs(self): @require_torch_cuda_compatibility(8.0) class TestSanaTransformerNF4Weights(ModelOptBaseTesterMixin): - expected_memory_reduction = 0.65 - def get_dummy_init_kwargs(self): return { "quant_type": "NF4", @@ -330,8 +99,6 @@ def get_dummy_init_kwargs(self): @require_torch_cuda_compatibility(8.0) class TestSanaTransformerNVFP4Weights(ModelOptBaseTesterMixin): - expected_memory_reduction = 0.65 - def get_dummy_init_kwargs(self): return { "quant_type": "NVFP4",