[tests] refactor pipeline-level quantization tests

.github/workflows/nightly_tests.yml

@@ -343,19 +343,19 @@ jobs:
       matrix:
         config:
           - backend: "bitsandbytes"
-            test_location: "bnb"
+            marker: "bitsandbytes"
             additional_deps: ["peft"]
           - backend: "gguf"
-            test_location: "gguf"
+            marker: "gguf"
             additional_deps: ["peft", "kernels"]
           - backend: "torchao"
-            test_location: "torchao"
-            additional_deps: []
+            marker: "torchao"
+            additional_deps: ["mslk"]
           - backend: "optimum_quanto"
-            test_location: "quanto"
+            marker: "quanto"
             additional_deps: []
           - backend: "nvidia_modelopt"
-            test_location: "modelopt"
+            marker: "modelopt"
             additional_deps: []
     runs-on:
       group: aws-g6e-xlarge-plus
@@ -390,9 +390,12 @@ jobs:
           BIG_GPU_MEMORY: 40
         run: |
           pytest -n 1 --max-worker-restart=0 --dist=loadfile \
+            -m "${{ matrix.config.marker }}" \
             --make-reports=tests_${{ matrix.config.backend }}_torch_cuda \
             --report-log=tests_${{ matrix.config.backend }}_torch_cuda.log \
-            tests/quantization/${{ matrix.config.test_location }}
+            tests/models \
+            tests/quantization \
+            tests/pipelines/testing_utils/quantization.py
       - name: Failure short reports
         if: ${{ failure() }}
         run: |
@@ -440,9 +443,10 @@ jobs:
           BIG_GPU_MEMORY: 40
         run: |
           pytest -n 1 --max-worker-restart=0 --dist=loadfile \
+            -k "TestPipelineQuantization" \
             --make-reports=tests_pipeline_level_quant_torch_cuda \
             --report-log=tests_pipeline_level_quant_torch_cuda.log \
-            tests/quantization/test_pipeline_level_quantization.py
+            tests/pipelines/testing_utils/quantization.py
       - name: Failure short reports
         if: ${{ failure() }}
         run: |

src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py

@@ -345,6 +345,39 @@ def __init__(self, quantization_config, **kwargs):
         if self.quantization_config.llm_int8_skip_modules is not None:
             self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules
 
+        self._checkpoint_keys = set()
+        self._pending_quantized_state = {}
+
+    def maybe_update_loaded_keys(self, loaded_keys: list[str], checkpoint_files: list[str]) -> list[str]:
+        self._checkpoint_keys = set(loaded_keys)
+        return loaded_keys
+
+    def maybe_update_state_dict(self, state_dict: dict[str, Any]) -> dict[str, Any]:
+        if not self.pre_quantized:
+            return state_dict
+
+        # A sharded checkpoint can split an 8-bit weight from its `SCB` statistics, which must be
+        # materialized together. Hold the incomplete half back until its counterpart arrives with a
+        # later shard.
+        merged = {**self._pending_quantized_state, **state_dict}
+        pending = {}
+        for name in list(merged.keys()):
+            if name.endswith(".weight"):
+                partner = name[: -len("weight")] + "SCB"
+            elif name.endswith(".SCB"):
+                partner = name[: -len("SCB")] + "weight"
+            else:
+                continue
+            if partner in self._checkpoint_keys and partner not in merged:
+                pending[name] = merged.pop(name)
+        self._pending_quantized_state = pending
+        return merged
+
+    @property
+    def supports_parallel_loading(self) -> bool:
+        # Deferred SCB reconstruction carries incomplete weight/SCB pairs from one shard to the next.
+        return not self.pre_quantized
+
     def validate_environment(self, *args, **kwargs):
         if not (torch.cuda.is_available() or torch.xpu.is_available()):
             raise RuntimeError("No GPU found. A GPU is needed for quantization.")

src/diffusers/quantizers/torchao/torchao_quantizer.py

@@ -302,9 +302,19 @@ def maybe_update_state_dict(self, state_dict: dict[str, Any]) -> dict[str, Any]:
             return state_dict
 
         merged_state_dict = {**self._pending_flattened_state_dict, **state_dict}
+        # Tensors at the model root (e.g. Wan's `scale_shift_table`) have no module prefix and are never
+        # flattened tensor-subclass parts; torchao's unflatten helper cannot parse their names, so route
+        # them (and their metadata entries) around the reconstruction.
+        root_tensors = {k: v for k, v in merged_state_dict.items() if "." not in k}
+        merged_state_dict = {k: v for k, v in merged_state_dict.items() if "." in k}
+        metadata = self._metadata
+        tensor_names = json.loads(metadata["tensor_names"])
+        if any("." not in name for name in tensor_names):
+            metadata = {**metadata, "tensor_names": json.dumps([name for name in tensor_names if "." in name])}
         reconstructed_state_dict, self._pending_flattened_state_dict = unflatten_tensor_state_dict(
-            merged_state_dict, self._metadata
+            merged_state_dict, metadata
         )
+        reconstructed_state_dict.update(root_tensors)
 
         return reconstructed_state_dict