pulled

Author: Serenagu525

src/maxdiffusion/configs/ltx_video.yml

@@ -27,33 +27,29 @@ output_dir: 'ltx-video-output'
 save_config_to_gcs: False
 
 #parallelism
-mesh_axes: ['data', 'fsdp', 'tensor', 'fsdp_transpose', 'expert', 'tensor_transpose', 'tensor_sequence', 'sequence']
+mesh_axes: ['data', 'fsdp', 'tensor']
 logical_axis_rules: [
                       ['batch', 'data'],
+                      ['activation_heads', 'fsdp'],
                       ['activation_batch', ['data','fsdp']],
-                      ['activation_heads', 'tensor'],
                       ['activation_kv', 'tensor'],
                       ['mlp','tensor'],
                       ['embed','fsdp'],
                       ['heads', 'tensor'],
+                      ['norm', 'fsdp'],
                       ['conv_batch', ['data','fsdp']],
                       ['out_channels', 'tensor'],
                       ['conv_out', 'fsdp'],
+                      ['conv_in', 'fsdp']
                     ]
-data_sharding: [['data', 'fsdp', 'tensor', 'fsdp_transpose', 'expert', 'tensor_transpose', 'tensor_sequence', 'sequence']]
+data_sharding: [['data', 'fsdp', 'tensor']]
 dcn_data_parallelism: 1  # recommended DCN axis to be auto-sharded
 dcn_fsdp_parallelism: -1
 dcn_tensor_parallelism: 1
-ici_data_parallelism: -1
-ici_fsdp_parallelism: 1  # recommended ICI axis to be auto-sharded
+ici_data_parallelism: 1
+ici_fsdp_parallelism: -1  # recommended ICI axis to be auto-sharded
 ici_tensor_parallelism: 1
 
-ici_fsdp_transpose_parallelism: 1
-ici_sequence_parallelism: 1
-ici_tensor_transpose_parallelism: 1
-ici_expert_parallelism: 1
-ici_sequence_parallelism: 1
-
 
 
 

src/maxdiffusion/max_utils.py

@@ -257,21 +257,6 @@ def create_device_mesh(config, devices=None, logging=True):
   if devices is None:
     devices = jax.devices()
   num_devices = len(devices)
-  ##special case for ltx-video
-  if "fsdp_transpose" in config.mesh_axes:
-    num_slices = 1
-    # if config.inference_benchmark_test else config.num_slices
-    num_devices_per_slice = num_devices // num_slices
-    # Find possible unspecified parallelisms
-    ici_parallelism = fill_unspecified_mesh_axes(config.ici_parallelism.copy(), num_devices_per_slice, "ICI")
-    mesh = mesh_utils.create_device_mesh(
-        ici_parallelism,
-        devices,
-    )
-    max_logging.log(f"Num_devices: {num_devices}, shape {mesh.shape}")
-
-    return mesh
-
   try:
     num_slices = 1 + max([d.slice_index for d in devices])
   except:
@@ -417,11 +402,7 @@ def setup_initial_state(
           config.enable_single_replica_ckpt_restoring,
       )
       if state:
-        ###!Edited
-        if checkpoint_item == " ":
-          state = state
-        else:
-          state = state[checkpoint_item]
+        state = state[checkpoint_item]
     if not state:
       max_logging.log(f"Could not find the item in orbax, creating state...")
       init_train_state_partial = functools.partial(
@@ -628,4 +609,4 @@ def maybe_initialize_jax_distributed_system(raw_keys):
     initialize_jax_for_gpu()
     max_logging.log("Jax distributed system initialized on GPU!")
   else:
-    jax.distributed.initialize()
+    jax.distributed.initialize()

src/maxdiffusion/models/ltx_video/transformers/attention.py

@@ -18,6 +18,7 @@
 import math
 from typing import Any, Dict, Optional, Tuple
 from enum import Enum, auto
+
 import jax
 import jax.nn as jnn
 import jax.numpy as jnp
@@ -213,7 +214,8 @@ def __call__(
 
     # Adaptive Norm
     if self.adaptive_norm in ["single_scale_shift", "single_scale"]:
-      assert timestep.ndim == 3  # [batch, 1 or num_tokens, embedding_dim]
+      # [batch, 1 or num_tokens, embedding_dim]
+      assert timestep.ndim == 3
       num_ada_params = self.scale_shift_table.shape[0]
       ada_values = self.scale_shift_table[None, None].astype(self.weight_dtype) + timestep.reshape(
           batch_size, timestep.shape[1], num_ada_params, -1
@@ -452,7 +454,7 @@ def __call__(
       deterministic: bool = True,
       **cross_attention_kwargs,
   ) -> jnp.ndarray:
-    cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters}  # noqa: F821
+    cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters}  # noqa F821
     assert cross_attention_kwargs.get("scale", None) is None, "Not supported"
 
     input_axis_names = ("activation_batch", "activation_length", "activation_embed")
@@ -636,27 +638,20 @@ def partial_flash_attention(q, k, v, q_segment_ids, kv_segment_ids):
         raise ValueError(f"Expected mask with 2 dims, got {q_segment_ids.ndim}.")
       # Based on: ("activation_kv_batch", "activation_kv_heads", "activation_length", "activation_kv_head_dim")
       # Computation of the spec based on the logical constraints can be found in logical_axes_to_spec.py.
-      qkvo_sharding_spec = jax.sharding.PartitionSpec(
-          ("data", "fsdp", "fsdp_transpose", "expert"),
-          ("tensor", "tensor_transpose", "sequence", "tensor_sequence"),
-          None,
-          None,
-      )
       # qkvo_sharding_spec = jax.sharding.PartitionSpec(
       #     ("data", "fsdp", "fsdp_transpose", "expert"),
       #     ("tensor", "tensor_transpose", "sequence", "tensor_sequence"),
       #     None,
       #     None,
       # )
-      # qkvo_sharding_spec = jax.sharding.PartitionSpec(
-      #     None,
-      #     None,
-      #     None,
-      #     None,
-      # )
+      qkvo_sharding_spec = jax.sharding.PartitionSpec(
+          "data",
+          "fsdp",
+          None,
+          "tensor",
+      )
       # Based on: ("activation_kv_batch", "activation_length")
-      qkv_segment_ids_spec = jax.sharding.PartitionSpec(("data", "fsdp", "fsdp_transpose", "expert"), "sequence")
-      # qkv_segment_ids_spec = jax.sharding.PartitionSpec(None, None)
+      qkv_segment_ids_spec = jax.sharding.PartitionSpec("data", None)
       wrapped_flash_attention = shard_map(
           partial_flash_attention,
           mesh=sharding_mesh,
@@ -841,7 +836,8 @@ def __call__(self, hidden_states: jax.Array, scale: float = 1.0, deterministic:
       inner_dim = dim * self.mult
       if inner_dim < 256:
         raise ValueError("inner_dim must be at least 256")
-      inner_dim = round(inner_dim / 256) * 256  # round to nearest multiple of 256
+      # round to nearest multiple of 256
+      inner_dim = round(inner_dim / 256) * 256
     else:
       inner_dim = self.inner_dim
 

src/maxdiffusion/pyconfig.py

@@ -25,6 +25,7 @@
 import yaml
 from . import max_logging
 from . import max_utils
+from .models.wan.wan_utils import CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH, WAN_21_FUSION_X_MODEL_NAME_OR_PATH
 
 
 def string_to_bool(s: str) -> bool:
@@ -41,21 +42,6 @@ def string_to_bool(s: str) -> bool:
 config = None
 
 
-def create_parallelisms_list(raw_keys):
-  ici_parallelism = [
-      raw_keys["ici_data_parallelism"],
-      raw_keys["ici_fsdp_parallelism"],
-      raw_keys["ici_fsdp_transpose_parallelism"],
-      raw_keys["ici_sequence_parallelism"],
-      raw_keys["ici_tensor_parallelism"],
-      raw_keys["ici_tensor_transpose_parallelism"],
-      raw_keys["ici_expert_parallelism"],
-      raw_keys["ici_sequence_parallelism"],
-  ]
-  raw_keys["ici_parallelism"] = ici_parallelism
-  return raw_keys
-
-
 def print_system_information():
   max_logging.log(f"System Information: Jax Version: {jax.__version__}")
   max_logging.log(f"System Information: Jaxlib Version: {jax.lib.__version__}")
@@ -117,6 +103,7 @@ def __init__(self, argv: list[str], **kwargs):
       jax.config.update("jax_compilation_cache_dir", raw_keys["jax_cache_dir"])
 
     _HyperParameters.user_init(raw_keys)
+    _HyperParameters.wan_init(raw_keys)
     self.keys = raw_keys
     for k in sorted(raw_keys.keys()):
       max_logging.log(f"Config param {k}: {raw_keys[k]}")
@@ -125,6 +112,26 @@ def _load_kwargs(self, argv: list[str]):
     args_dict = dict(a.split("=", 1) for a in argv[2:])
     return args_dict
 
+  @staticmethod
+  def wan_init(raw_keys):
+    if "wan_transformer_pretrained_model_name_or_path" in raw_keys:
+      transformer_pretrained_model_name_or_path = raw_keys["wan_transformer_pretrained_model_name_or_path"]
+      if transformer_pretrained_model_name_or_path == "":
+        raw_keys["wan_transformer_pretrained_model_name_or_path"] = raw_keys["pretrained_model_name_or_path"]
+      elif (
+          transformer_pretrained_model_name_or_path == CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH
+          or transformer_pretrained_model_name_or_path == WAN_21_FUSION_X_MODEL_NAME_OR_PATH
+      ):
+        # Set correct parameters for CausVid in case of user error.
+        raw_keys["guidance_scale"] = 1.0
+        num_inference_steps = raw_keys["num_inference_steps"]
+        if num_inference_steps > 10:
+          max_logging.log(
+              f"Warning: Try setting num_inference_steps to less than 8 steps when using CausVid, currently you are setting {num_inference_steps} steps."
+          )
+      else:
+        raise ValueError(f"{transformer_pretrained_model_name_or_path} transformer model is not supported for Wan 2.1")
+
   @staticmethod
   def user_init(raw_keys):
     """Transformations between the config data and configs used at runtime"""
@@ -169,8 +176,6 @@ def user_init(raw_keys):
     raw_keys["total_train_batch_size"] = max_utils.get_global_batch_size(raw_keys["per_device_batch_size"])
     raw_keys["num_slices"] = get_num_slices(raw_keys)
     raw_keys["quantization_local_shard_count"] = get_quantization_local_shard_count(raw_keys)
-    if "ici_fsdp_transpose_parallelism" in raw_keys:
-      raw_keys = create_parallelisms_list(raw_keys)
 
 
 def get_num_slices(raw_keys):