fixe unit test

Author: susanbao

src/maxdiffusion/tests/wan_transformer_test.py

@@ -332,7 +332,9 @@ def create_real_rule_instance(*args, **kwargs):
     config_fp8_full = Mock(spec=HyperParameters)
     config_fp8_full.use_qwix_quantization = True
     config_fp8_full.quantization = "fp8_full"
-    config_fp8_full.quantization_calibration_method = "absmax"
+    config_fp8_full.weight_quantization_calibration_method = "absmax"
+    config_fp8_full.act_quantization_calibration_method = "absmax"
+    config_fp8_full.bwd_quantization_calibration_method = "absmax"
     config_fp8_full.qwix_module_path = ".*"
     provider_fp8_full = WanPipeline.get_qt_provider(config_fp8_full)
     self.assertIsNotNone(provider_fp8_full)
@@ -343,9 +345,9 @@ def create_real_rule_instance(*args, **kwargs):
             act_qtype=jnp.float8_e4m3fn,
             bwd_qtype=jnp.float8_e5m2,
             disable_channelwise_axes=True,  # per_tensor calibration
-            weight_calibration_method=config_fp8_full.quantization_calibration_method,
-            act_calibration_method=config_fp8_full.quantization_calibration_method,
-            bwd_calibration_method=config_fp8_full.quantization_calibration_method,
+            weight_calibration_method=config_fp8_full.weight_quantization_calibration_method,
+            act_calibration_method=config_fp8_full.act_quantization_calibration_method,
+            bwd_calibration_method=config_fp8_full.bwd_quantization_calibration_method,
             op_names=("dot_general", "einsum"),
         ),
         call(
@@ -354,9 +356,9 @@ def create_real_rule_instance(*args, **kwargs):
             act_qtype=jnp.float8_e4m3fn,
             bwd_qtype=jnp.float8_e4m3fn,
             disable_channelwise_axes=True,  # per_tensor calibration
-            weight_calibration_method=config_fp8_full.quantization_calibration_method,
-            act_calibration_method=config_fp8_full.quantization_calibration_method,
-            bwd_calibration_method=config_fp8_full.quantization_calibration_method,
+            weight_calibration_method=config_fp8_full.weight_quantization_calibration_method,
+            act_calibration_method=config_fp8_full.act_quantization_calibration_method,
+            bwd_calibration_method=config_fp8_full.bwd_quantization_calibration_method,
             op_names=("conv_general_dilated"),
         ),
     ]