Merge branch 'main' into update_readme_wan

Author: entrpn

src/maxdiffusion/configs/base_flux_schnell.yml

@@ -236,6 +236,7 @@ enable_profiler: False
 # the iteration time a chance to stabilize.
 skip_first_n_steps_for_profiler: 5
 profiler_steps: 10
+profiler: ""
 
 # Generation parameters
 prompt: "A magical castle in the middle of a forest, artistic drawing"
@@ -284,3 +285,5 @@ quantization: ''
 quantization_local_shard_count: -1
 use_qwix_quantization: False 
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
+
+save_final_checkpoint: False

src/maxdiffusion/configs/base_wan_14b.yml

@@ -313,10 +313,9 @@ compile_topology_num_slices: -1 # Number of target slices, set to a positive int
 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"
+qwix_module_path: ".*"
 
 # Eval model on per eval_every steps. -1 means don't eval.
 eval_every: -1
 eval_data_dir: ""
 enable_generate_video_for_eval: False # This will increase the used TPU memory.
-eval_max_number_of_samples_in_bucket: 60
-eval_max_processed_batch_size: 8 # This is the max batch size per device for eval step. If the global eval batch size is larger than this, the eval step will be run multiple times.

src/maxdiffusion/generate_wan.py

@@ -21,6 +21,7 @@
 from absl import app
 from maxdiffusion.utils import export_to_video
 from google.cloud import storage
+import flax
 
 
 def upload_video_to_gcs(output_dir: str, video_path: str):
@@ -161,6 +162,7 @@ def run(config, pipeline=None, filename_prefix=""):
 
 def main(argv: Sequence[str]) -> None:
   pyconfig.initialize(argv)
+  flax.config.update('flax_always_shard_variable', False)
   run(pyconfig.config)
 
 

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -243,34 +243,48 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
     return wan_vae, vae_cache
 
   @classmethod
-  def get_basic_config(cls, dtype):
+  def get_basic_config(cls, dtype, config: HyperParameters):
     rules = [
         qwix.QtRule(
-            module_path=".*",  # Apply to all modules
+            module_path=config.qwix_module_path,
             weight_qtype=dtype,
             act_qtype=dtype,
+            op_names=("dot_general","einsum", "conv_general_dilated"),
         )
     ]
     return rules
 
   @classmethod
-  def get_fp8_config(cls, quantization_calibration_method: str):
+  def get_fp8_config(cls, config: HyperParameters):
     """
     fp8 config rules with per-tensor calibration.
     FLAX API (https://flax-linen.readthedocs.io/en/v0.10.6/guides/quantization/fp8_basics.html#flax-low-level-api):
     The autodiff does not automatically use E5M2 for gradients and E4M3 for activations/weights during training, which is the recommended practice.
     """
     rules = [
         qwix.QtRule(
-            module_path=".*",  # Apply to all modules
+            module_path=config.qwix_module_path,
             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=config.quantization_calibration_method,
+            act_calibration_method=config.quantization_calibration_method,
+            bwd_calibration_method=config.quantization_calibration_method,
+            op_names=("dot_general","einsum"),
+        ),
+        qwix.QtRule(
+            module_path=config.qwix_module_path,
+            weight_qtype=jnp.float8_e4m3fn, # conv_general_dilated requires the same dtypes
+            act_qtype=jnp.float8_e4m3fn,
             bwd_qtype=jnp.float8_e4m3fn,
             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,
+            weight_calibration_method=config.quantization_calibration_method,
+            act_calibration_method=config.quantization_calibration_method,
+            bwd_calibration_method=config.quantization_calibration_method,
+            op_names=("conv_general_dilated"),
         )
     ]
     return rules
@@ -281,14 +295,13 @@ def get_qt_provider(cls, config: HyperParameters) -> Optional[qwix.QtProvider]:
     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))
+        return qwix.QtProvider(cls.get_basic_config(jnp.int8, config))
       case "fp8":
-        return qwix.QtProvider(cls.get_basic_config(jnp.float8_e4m3fn))
+        return qwix.QtProvider(cls.get_basic_config(jnp.float8_e4m3fn, config))
       case "fp8_full":
-        return qwix.QtProvider(cls.get_fp8_config(quantization_calibration_method))
+        return qwix.QtProvider(cls.get_fp8_config(config))
     return None
 
   @classmethod

src/maxdiffusion/tests/wan_transformer_test.py

@@ -19,7 +19,7 @@
 import jax.numpy as jnp
 import pytest
 import unittest
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, patch, call
 from absl.testing import absltest
 from flax import nnx
 from jax.sharding import Mesh
@@ -37,6 +37,11 @@
 from ..models.attention_flax import FlaxWanAttention
 from maxdiffusion.pyconfig import HyperParameters
 from maxdiffusion.pipelines.wan.wan_pipeline import WanPipeline
+import qwix
+import flax
+
+flax.config.update('flax_always_shard_variable', False)
+RealQtRule = qwix.QtRule
 
 
 IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
@@ -282,6 +287,10 @@ def test_get_qt_provider(self, mock_qt_rule):
     """
     Tests the provider logic for all config branches.
     """
+    def create_real_rule_instance(*args, **kwargs):
+        return RealQtRule(*args, **kwargs)
+    mock_qt_rule.side_effect = create_real_rule_instance
+
     # Case 1: Quantization disabled
     config_disabled = Mock(spec=HyperParameters)
     config_disabled.use_qwix_quantization = False
@@ -291,28 +300,43 @@ def test_get_qt_provider(self, mock_qt_rule):
     config_int8 = Mock(spec=HyperParameters)
     config_int8.use_qwix_quantization = True
     config_int8.quantization = "int8"
+    config_int8.qwix_module_path = ".*"
     provider_int8 = WanPipeline.get_qt_provider(config_int8)
     self.assertIsNotNone(provider_int8)
-    mock_qt_rule.assert_called_once_with(module_path=".*", weight_qtype=jnp.int8, act_qtype=jnp.int8)
+    mock_qt_rule.assert_called_once_with(module_path=".*", weight_qtype=jnp.int8, act_qtype=jnp.int8, op_names=("dot_general","einsum", "conv_general_dilated"))
 
     # Case 3: Quantization enabled, type 'fp8'
     mock_qt_rule.reset_mock()
     config_fp8 = Mock(spec=HyperParameters)
     config_fp8.use_qwix_quantization = True
     config_fp8.quantization = "fp8"
+    config_fp8.qwix_module_path = ".*"
     provider_fp8 = WanPipeline.get_qt_provider(config_fp8)
     self.assertIsNotNone(provider_fp8)
-    mock_qt_rule.assert_called_once_with(module_path=".*", weight_qtype=jnp.float8_e4m3fn, act_qtype=jnp.float8_e4m3fn)
+    mock_qt_rule.assert_called_once_with(module_path=".*", weight_qtype=jnp.float8_e4m3fn, act_qtype=jnp.float8_e4m3fn, op_names=("dot_general","einsum", "conv_general_dilated"))
 
     # Case 4: Quantization enabled, type 'fp8_full'
     mock_qt_rule.reset_mock()
     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.qwix_module_path = ".*"
     provider_fp8_full = WanPipeline.get_qt_provider(config_fp8_full)
     self.assertIsNotNone(provider_fp8_full)
-    mock_qt_rule.assert_called_once_with(
+    expected_calls = [
+      call(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=config_fp8_full.quantization_calibration_method,
+        act_calibration_method=config_fp8_full.quantization_calibration_method,
+        bwd_calibration_method=config_fp8_full.quantization_calibration_method,
+        op_names=("dot_general","einsum"),
+        ),
+      call(
         module_path=".*",  # Apply to all modules
         weight_qtype=jnp.float8_e4m3fn,
         act_qtype=jnp.float8_e4m3fn,
@@ -322,7 +346,10 @@ def test_get_qt_provider(self, mock_qt_rule):
         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,
-    )
+        op_names=("conv_general_dilated"),
+      )
+    ]
+    mock_qt_rule.assert_has_calls(expected_calls, any_order=True)
 
     # Case 5: Invalid quantization type
     config_invalid = Mock(spec=HyperParameters)
@@ -341,7 +368,9 @@ def test_quantize_transformer_enabled(self, mock_get_dummy_inputs, mock_quantize
     mock_config = Mock(spec=HyperParameters)
     mock_config.use_qwix_quantization = True
     mock_config.quantization = "fp8_full"
+    mock_config.qwix_module_path = ".*"
     mock_config.per_device_batch_size = 1
+    mock_config.quantization_calibration_method = "absmax"
 
     mock_model = Mock(spec=WanModel)
     mock_pipeline = Mock()

src/maxdiffusion/tests/wan_vae_test.py

@@ -46,12 +46,13 @@
 from ..models.wan.wan_utils import load_wan_vae
 from ..utils import load_video
 from ..video_processor import VideoProcessor
+import flax
 
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
 
 CACHE_T = 2
 
-
+flax.config.update('flax_always_shard_variable', False)
 class TorchWanRMS_norm(nn.Module):
   r"""
   A custom RMS normalization layer.

src/maxdiffusion/train_flux.py

@@ -22,6 +22,7 @@
 
 from maxdiffusion.train_utils import (
     validate_train_config,
+    transformer_engine_context,
 )
 
 
@@ -39,6 +40,6 @@ def main(argv: Sequence[str]) -> None:
   max_logging.log(f"Found {jax.device_count()} devices.")
   train(config)
 
-
 if __name__ == "__main__":
-  app.run(main)
+  with transformer_engine_context():
+    app.run(main)

src/maxdiffusion/train_sdxl.py

@@ -27,6 +27,7 @@
 
 from maxdiffusion.train_utils import (
     validate_train_config,
+    transformer_engine_context,
 )
 
 
@@ -51,4 +52,5 @@ def main(argv: Sequence[str]) -> None:
   tf.config.set_visible_devices([], "GPU")
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
   torch.set_default_device("cpu")
-  app.run(main)
+  with transformer_engine_context():
+    app.run(main)

src/maxdiffusion/train_utils.py

@@ -20,7 +20,7 @@
 import queue
 
 from maxdiffusion import max_utils, max_logging
-
+from contextlib import contextmanager
 
 def get_first_step(state):
   return int(state.step)
@@ -196,3 +196,22 @@ def generate_timestep_weights(config, num_timesteps):
   weights[bias_indices] *= timestep_bias_config["multiplier"]
   weights /= weights.sum()
   return jnp.array(weights)
+
+
+@contextmanager
+def transformer_engine_context():
+  """ If TransformerEngine is available, this context manager will provide the library with MaxDiffusion-specific details needed for correcct operation. """
+  try:
+    from transformer_engine.jax.sharding import global_shard_guard, MeshResource
+    # Inform TransformerEngine of MaxDiffusion's physical mesh resources.
+    mesh_resource = MeshResource(
+      dp_resource = "data",
+      tp_resource = "tensor",
+      fsdp_resource = "fsdp",
+      pp_resource = None,
+      cp_resource = None,
+    )
+    with global_shard_guard(mesh_resource):
+      yield
+  except ImportError:
+    yield

src/maxdiffusion/train_wan.py

@@ -20,6 +20,7 @@
 from absl import app
 from maxdiffusion import max_logging, pyconfig
 from maxdiffusion.train_utils import validate_train_config
+import flax
 
 
 def train(config):
@@ -34,6 +35,7 @@ def main(argv: Sequence[str]) -> None:
   config = pyconfig.config
   validate_train_config(config)
   max_logging.log(f"Found {jax.device_count()} devices.")
+  flax.config.update('flax_always_shard_variable', False)
   train(config)