From 3f313a5b547f1f700a57ae4c4747196ba8fff88b Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Tue, 14 Jul 2026 09:59:41 -0700
Subject: [PATCH 1/2] Fix ModelOpt pre-quantized loading

---
 src/diffusers/quantizers/modelopt/modelopt_quantizer.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/diffusers/quantizers/modelopt/modelopt_quantizer.py b/src/diffusers/quantizers/modelopt/modelopt_quantizer.py
index 5c9e198c82b2..b7cadb50e54b 100644
--- a/src/diffusers/quantizers/modelopt/modelopt_quantizer.py
+++ b/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)

From 20f76263d131afc7eab3d6226b7d5d3ea3d7b34d Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Thu, 16 Jul 2026 07:13:00 -0700
Subject: [PATCH 2/2] Add regression catching failure before the fix

---
 src/diffusers/utils/testing_utils.py         |  2 +-
 tests/quantization/modelopt/test_modelopt.py | 36 ++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py
index eefe52c477a6..afe462c8d557 100644
--- a/src/diffusers/utils/testing_utils.py
+++ b/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}."
diff --git a/tests/quantization/modelopt/test_modelopt.py b/tests/quantization/modelopt/test_modelopt.py
index 6b0624a28083..113a4d43d180 100644
--- a/tests/quantization/modelopt/test_modelopt.py
+++ b/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