qwix quantize WAN transformer

Author: susanbao

src/maxdiffusion/configs/base_wan_14b.yml

@@ -287,4 +287,7 @@ quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
+use_qwix_quantization: False # Whether to use qwix for quantization. If set to True, the transformer of WAN will be quantized using qwix.
+# Quantization calibration method used for weights and activations. Supported methods can be found in https://github.com/google/qwix/blob/dc2a0770351c740e5ab3cce7c0efe9f7beacce9e/qwix/qconfig.py#L70-L80
+quantization_calibration_method: "absmax"
 

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -33,9 +33,11 @@
 from ...schedulers.scheduling_unipc_multistep_flax import FlaxUniPCMultistepScheduler, UniPCMultistepSchedulerState
 from transformers import AutoTokenizer, UMT5EncoderModel
 from maxdiffusion.utils.import_utils import is_ftfy_available
+from ...maxdiffusion_utils import get_dummy_wan_inputs
 import html
 import re
 import torch
+import qwix
 
 
 def basic_clean(text):
@@ -225,6 +227,68 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
     vae_cache = AutoencoderKLWanCache(wan_vae)
     return wan_vae, vae_cache
 
+  @classmethod
+  def get_basic_config(cls, dtype):
+    rules = [
+        qwix.QtRule(
+          module_path='.*',  # Apply to all modules
+          weight_qtype=dtype,
+          act_qtype=dtype,
+        )
+      ]
+    return rules
+
+  @classmethod
+  def get_fp8_config(cls, quantization_calibration_method: str):
+    """ fp8 config rules with per-tensor calibration.
+    """
+    rules = [
+        qwix.QtRule(
+          module_path='.*',  # Apply to all modules
+          weight_qtype=jnp.float8_e4m3fn,
+          act_qtype=jnp.float8_e4m3fn,
+          bwd_qtype=jnp.float8_e5m2,
+          bwd_use_original_residuals=True,
+          disable_channelwise_axes=True, # per_tensor calibration
+          weight_calibration_method = quantization_calibration_method,
+          act_calibration_method = quantization_calibration_method,
+          bwd_calibration_method = quantization_calibration_method,
+        )
+      ]
+    return rules
+
+  @classmethod
+  def get_qt_provider(cls, config: HyperParameters) -> Optional[qwix.QtProvider]:
+    """Get quantization rules based on the config."""
+    if not getattr(config, "use_qwix_quantization", False):
+      return None
+
+    quantization_calibration_method = getattr(config, "quantization_calibration_method", "absmax")
+    match config.quantization:
+      case "int8":
+        return qwix.QtProvider(cls.get_basic_config(jnp.int8))
+      case "fp8":
+        return qwix.QtProvider(cls.get_basic_config(jnp.float8_e4m3fn))
+      case "fp8_full":
+        return qwix.QtProvider(cls.get_fp8_config(quantization_calibration_method))
+    return None
+
+  @classmethod
+  def quantize_transformer(cls, config: HyperParameters, model: WanModel, pipeline: "WanPipeline", mesh: Mesh):
+    """Quantizes the transformer model."""
+    q_rules = cls.get_qt_provider(config)
+    if not q_rules:
+      return model
+    max_logging.log("Quantizing transformer with Qwix.")
+
+    batch_size = int(config.per_device_batch_size * jax.local_device_count())
+    latents, prompt_embeds, timesteps = get_dummy_wan_inputs(config, pipeline, batch_size)
+    model_inputs= (latents, timesteps, prompt_embeds)
+    with mesh:
+      quantized_model = qwix.quantize_model(model, q_rules, *model_inputs)
+    max_logging.log("Qwix Quantization complete.")
+    return quantized_model
+
   @classmethod
   def load_transformer(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: HyperParameters):
     with mesh:
@@ -264,7 +328,7 @@ def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transform
     with mesh:
       wan_vae, vae_cache = cls.load_vae(devices_array=devices_array, mesh=mesh, rngs=rngs, config=config)
 
-    return WanPipeline(
+    pipeline = WanPipeline(
         tokenizer=tokenizer,
         text_encoder=text_encoder,
         transformer=transformer,
@@ -277,6 +341,10 @@ def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transform
         config=config,
     )
 
+    pipeline.transformer = cls.quantize_transformer(config, pipeline.transformer, pipeline, mesh)
+    return pipeline
+
+
   def _get_t5_prompt_embeds(
       self,
       prompt: Union[str, List[str]] = None,