[Quantization] Fix ModelOpt pre-quantized loading

src/diffusers/quantizers/modelopt/modelopt_quantizer.py

@@ -69,7 +69,11 @@ def check_if_quantized_param(
 
         module, tensor_name = get_module_from_name(model, param_name)
         if self.pre_quantized:
-            return True
+            # ModelOpt restoration recreates quantizer state such as `_amax` and
+            # `_scale` as buffers. Let the regular low-memory loader materialize
+            # those buffers on their target device. Only parameters need the
+            # wrapper-aware assignment below.
+            return tensor_name in module._parameters
         elif is_quantized(module) and "weight" in tensor_name:
             return True
         return False
@@ -92,6 +96,8 @@ def create_quantized_param(
         dtype = kwargs.get("dtype", torch.float32)
         module, tensor_name = get_module_from_name(model, param_name)
         if self.pre_quantized:
+            if tensor_name not in module._parameters:
+                raise ValueError(f"Expected {param_name} to be a parameter in the restored ModelOpt graph.")
             module._parameters[tensor_name] = torch.nn.Parameter(param_value.to(device=target_device))
         else:
             set_module_tensor_to_device(model, param_name, target_device, param_value, dtype)

src/diffusers/utils/testing_utils.py

@@ -648,7 +648,7 @@ def decorator(test_case):
 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(importlib.metadata.version("nvidia-modelopt")).base_version
         ) >= version.parse(modelopt_version)
         return unittest.skipUnless(
             correct_nvidia_modelopt_version, f"Test requires modelopt with version greater than {modelopt_version}."

tests/quantization/modelopt/test_modelopt.py

@@ -1,4 +1,6 @@
+import copy
 import gc
+import os
 import tempfile
 import unittest
 
@@ -19,6 +21,7 @@
 
 
 if is_nvidia_modelopt_available():
+    import modelopt.torch.opt as mto
     import modelopt.torch.quantization as mtq
 
 if is_torch_available():
@@ -254,6 +257,39 @@ class SanaTransformerFP8WeightsTest(ModelOptBaseTesterMixin, unittest.TestCase):
     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)
+            self.assertTrue(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())
+        self.assertTrue(
+            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:
+                self.assertFalse(tensor.is_meta, f"{tensor_kind} {name} was not materialized from meta.")
+
 
 class SanaTransformerINT8WeightsTest(ModelOptBaseTesterMixin, unittest.TestCase):
     expected_memory_reduction = 0.6