add TE context to sdxl

Author: hx89

src/maxdiffusion/train_flux.py

@@ -19,10 +19,10 @@
 import jax
 from absl import app
 from maxdiffusion import (max_logging, pyconfig)
-from contextlib import contextmanager
 
 from maxdiffusion.train_utils import (
     validate_train_config,
+    transformer_engine_context,
 )
 
 
@@ -40,24 +40,6 @@ def main(argv: Sequence[str]) -> None:
   max_logging.log(f"Found {jax.device_count()} devices.")
   train(config)
 
-@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
-
 if __name__ == "__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