chore: Add tokamax dependency and unsafe RNG utilities

Author: eltsai

.github/workflows/UnitTests.yml

@@ -58,8 +58,7 @@ jobs:
         pip show jax jaxlib flax transformers datasets tensorflow tensorflow_datasets
     - name: PyTest
       run: | #--deselect=src/maxdiffusion/tests/input_pipeline_interface_test.py
-        HF_HUB_CACHE=/mnt/disks/github-runner-disk/ HF_HOME=/mnt/disks/github-runner-disk/ TOKENIZERS_PARALLELISM=false python3 -m pytest --deselect=src/maxdiffusion/tests/ltx_transformer_step_test.py -x
-#  add_pull_ready:
+        HF_HUB_CACHE=/mnt/disks/github-runner-disk/ HF_HOME=/mnt/disks/github-runner-disk/ TOKENIZERS_PARALLELISM=false LIBTPU_INIT_ARGS="--xla_tpu_scoped_vmem_limit_kib=65472" python3 -m pytest --deselect=src/maxdiffusion/tests/ltx_transformer_step_test.py -x
 #    if: github.ref != 'refs/heads/main'
 #    permissions:
 #      checks: read

.gitignore

@@ -97,6 +97,7 @@ celerybeat-schedule
 *.sage.py
 
 # Environments
+.history
 .env
 .venv
 env/

requirements.txt

@@ -13,6 +13,7 @@ ftfy
 tensorboard>=2.17.0
 tensorboardx>=2.6.2.2
 tensorboard-plugin-profile>=2.15.2
+tokamax
 Jinja2
 scikit-image
 parameterized

src/maxdiffusion/generate_wan.py

@@ -62,6 +62,15 @@ def delete_file(file_path: str):
 
 
 jax.config.update("jax_use_shardy_partitioner", True)
+jax.config.update("jax_default_prng_impl", "unsafe_rbg")
+  # TF allocates extraneous GPU memory when using TFDS data
+  # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
+  # tf.config.set_visible_devices([], "GPU")
+if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
+  max_logging.log("Enabling unsafe RNG bit generator for TPU SPMD.")
+  os.environ["LIBTPU_INIT_ARGS"] = (
+      os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
+  )
 
 def get_pipeline(model_name: str):
   if model_name == "wan2.1":

src/maxdiffusion/max_utils.py

@@ -501,17 +501,26 @@ def get_flash_block_sizes(config):
   """Create custom flash attention BlockSizes."""
   flash_block_sizes = None
   if len(config.flash_block_sizes.keys()) > 0:
-    use_fused_bwd_kernel = config.flash_block_sizes.get("use_fused_bwd_kernel", False)
+    attention_is_tokamax = "tokamax" in config.attention
+    user_block_sizes:Dict[str, int] = config.flash_block_sizes
+    if attention_is_tokamax:
+      max_logging.log("Tokamax kernel specified, Note: Tokamax only supports fused backward kernel."
+                      "Hence following flash block properties specified will be ignored:"
+                      f"block_q: {user_block_sizes['block_q']},"
+                      f"block_q_dq: {user_block_sizes.get('block_q_dq')},"
+                      f"block_kv_dq: {user_block_sizes.get('block_kv_dq')},"
+                      f"use_fused_bwd_kernel: {user_block_sizes.get('use_fused_bwd_kernel')}"
+                      )
     flash_block_sizes = splash_attention_kernel.BlockSizes(
-        block_q=config.flash_block_sizes["block_q"],
-        block_kv_compute=config.flash_block_sizes["block_kv_compute"],
-        block_kv=config.flash_block_sizes["block_kv"],
-        block_q_dkv=config.flash_block_sizes["block_q_dkv"],
-        block_kv_dkv=config.flash_block_sizes["block_kv_dkv"],
-        block_kv_dkv_compute=config.flash_block_sizes["block_kv_dkv_compute"],
-        block_q_dq=value_or_none(config.flash_block_sizes, "block_q_dq"),
-        block_kv_dq=value_or_none(config.flash_block_sizes, "block_kv_dq"),
-        use_fused_bwd_kernel=value_or_none(config.flash_block_sizes, "use_fused_bwd_kernel"),
+        block_q=user_block_sizes.get("block_q_dkv", user_block_sizes["block_kv"]) if attention_is_tokamax else user_block_sizes["block_q"],
+        block_kv_compute=user_block_sizes["block_kv_compute"],
+        block_kv=user_block_sizes["block_kv"],
+        block_q_dkv=user_block_sizes["block_q_dkv"],
+        block_kv_dkv=user_block_sizes["block_kv_dkv"],
+        block_kv_dkv_compute=user_block_sizes["block_kv_dkv_compute"],
+        block_q_dq=None if attention_is_tokamax else value_or_none(user_block_sizes, "block_q_dq"),
+        block_kv_dq=None if attention_is_tokamax else value_or_none(user_block_sizes, "block_kv_dq"),
+        use_fused_bwd_kernel=True if attention_is_tokamax else value_or_none(user_block_sizes, "use_fused_bwd_kernel"),
     )
   return flash_block_sizes