conversion script checked

Author: Serenagu525

src/maxdiffusion/checkpointing/checkpointing_utils.py

@@ -213,7 +213,10 @@ def load_state_if_possible(
     max_logging.log(f"restoring from this run's directory latest step {latest_step}")
     try:
       if not enable_single_replica_ckpt_restoring:
-        item = {checkpoint_item: orbax.checkpoint.args.PyTreeRestore(item=abstract_unboxed_pre_state)}
+        if checkpoint_item == " ":
+          return checkpoint_manager.restore(latest_step, args=ocp.args.StandardRestore(abstract_unboxed_pre_state))
+        else:
+          item = {checkpoint_item: orbax.checkpoint.args.PyTreeRestore(item=abstract_unboxed_pre_state)}
         return checkpoint_manager.restore(latest_step, args=orbax.checkpoint.args.Composite(**item))
 
       def map_to_pspec(data):

src/maxdiffusion/configs/ltx_video.yml

@@ -27,7 +27,7 @@ output_dir: 'ltx-video-output'
 save_config_to_gcs: False
 
 #parallelism
-mesh_axes: ['data', 'fsdp', 'tensor']
+mesh_axes: ['data', 'fsdp', 'tensor', 'fsdp_transpose', 'expert', 'tensor_transpose', 'tensor_sequence', 'sequence']
 logical_axis_rules: [
                       ['batch', 'data'],
                       ['activation_batch', ['data','fsdp']],
@@ -40,14 +40,20 @@ logical_axis_rules: [
                       ['out_channels', 'tensor'],
                       ['conv_out', 'fsdp'],
                     ]
-data_sharding: [['data', 'fsdp', 'tensor']]
+data_sharding: [['data', 'fsdp', 'tensor', 'fsdp_transpose', 'expert', 'tensor_transpose', 'tensor_sequence', 'sequence']]
 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_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
+
 
 
 
@@ -62,4 +68,5 @@ cache_latents_text_encoder_outputs: True
 per_device_batch_size: 1
 compile_topology_num_slices: -1 
 quantization_local_shard_count: -1
-jit_initializers: True 
+jit_initializers: True 
+enable_single_replica_ckpt_restoring: False

src/maxdiffusion/generate_ltx_video.py

@@ -20,43 +20,90 @@
 import json
 from maxdiffusion.models.ltx_video.transformers.transformer3d import Transformer3DModel
 import os
+import functools
 import jax.numpy as jnp
 from maxdiffusion import pyconfig
 from maxdiffusion.max_utils import (
     create_device_mesh,
+    setup_initial_state,
+    get_memory_allocations,
 )
+from jax.sharding import Mesh
+import orbax.checkpoint as ocp
 
 
-def validate_transformer_inputs(prompt_embeds, fractional_coords, latents, noise_cond):
+def validate_transformer_inputs(
+    prompt_embeds, fractional_coords, latents, noise_cond, segment_ids, encoder_attention_segment_ids
+):
   print("prompts_embeds.shape: ", prompt_embeds.shape, prompt_embeds.dtype)
   print("fractional_coords.shape: ", fractional_coords.shape, fractional_coords.dtype)
   print("latents.shape: ", latents.shape, latents.dtype)
   print("noise_cond.shape: ", noise_cond.shape, noise_cond.dtype)
+  print("noise_cond.shape: ", noise_cond.shape, noise_cond.dtype)
+  print("segment_ids.shape: ", segment_ids.shape, segment_ids.dtype)
+  print("encoder_attention_segment_ids.shape: ", encoder_attention_segment_ids.shape, encoder_attention_segment_ids.dtype)
 
 
 def run(config):
-  key = jax.random.PRNGKey(0)
+
+  key = jax.random.PRNGKey(42)
 
   devices_array = create_device_mesh(config)
   mesh = Mesh(devices_array, config.mesh_axes)
 
-  batch_size, text_tokens, num_tokens, features = 4, 256, 2048, 128
   base_dir = os.path.dirname(__file__)
 
-  # load in model config
+  ##load in model config
   config_path = os.path.join(base_dir, "models/ltx_video/xora_v1.2-13B-balanced-128.json")
   with open(config_path, "r") as f:
     model_config = json.load(f)
+  ckpt_path = model_config["ckpt_path"]
+
+  ignored_keys = [
+      "_class_name",
+      "_diffusers_version",
+      "_name_or_path",
+      "causal_temporal_positioning",
+      "in_channels",
+      "ckpt_path",
+  ]
+  in_channels = model_config["in_channels"]
+  for name in ignored_keys:
+    if name in model_config:
+      del model_config[name]
+
+  transformer = Transformer3DModel(
+      **model_config, dtype=jnp.float32, gradient_checkpointing="matmul_without_batch", sharding_mesh=mesh
+  )
+  transformer_param_shapes = transformer.init_weights(  # noqa: F841
+      in_channels, key, model_config["caption_channels"], eval_only=True
+  )  # use this to test!
 
-  transformer = Transformer3DModel(**model_config, dtype=jnp.bfloat16, gradient_checkpointing="matmul_without_batch")
-  transformer_param_shapes = transformer.init_weights(key, batch_size, text_tokens, num_tokens, features, eval_only=False)
+  weights_init_fn = functools.partial(
+      transformer.init_weights, in_channels, key, model_config["caption_channels"], eval_only=True
+  )
 
-  key, split_key = jax.random.split(key)
+  checkpoint_manager = ocp.CheckpointManager(ckpt_path)
+  transformer_state, transformer_state_shardings = setup_initial_state(
+      model=transformer,
+      tx=None,
+      config=config,
+      mesh=mesh,
+      weights_init_fn=weights_init_fn,
+      checkpoint_manager=checkpoint_manager,
+      checkpoint_item=" ",
+      model_params=None,
+      training=False,
+  )
 
+  transformer_state = jax.device_put(transformer_state, transformer_state_shardings)
+  get_memory_allocations()
 
-  weights_init_fn = functools.partial(
-      transformer.init_weights, split_key, batch_size, text_tokens, num_tokens, features, eval_only=True
-  )
+  states = {}
+  state_shardings = {}
+
+  state_shardings["transformer"] = transformer_state_shardings
+  states["transformer"] = transformer_state
 
 
 def main(argv: Sequence[str]) -> None:
@@ -66,3 +113,9 @@ def main(argv: Sequence[str]) -> None:
 
 if __name__ == "__main__":
   app.run(main)
+
+
+###setup_initial_state, can optionally load from checkpoint
+
+
+# end to end steps from ltx repo: pipeline_ltx_video.py

src/maxdiffusion/max_utils.py

@@ -252,45 +252,88 @@ def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_typ
   return parallelism_vals
 
 
-def create_device_mesh(config, devices=None, logging=True):
+def create_device_mesh(config, devices=None):
   """Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas"""
   if devices is None:
     devices = jax.devices()
   num_devices = len(devices)
-  try:
-    num_slices = 1 + max([d.slice_index for d in devices])
-  except:
-    num_slices = 1
+  num_slices = 1
+  # if config.inference_benchmark_test else config.num_slices
   num_devices_per_slice = num_devices // num_slices
-  max_logging.log(f"Devices: {devices} (num_devices: {num_devices})")
 
-  multi_slice_env = num_slices > 1
-
-  dcn_parallelism = [
-      config.dcn_data_parallelism,
-      config.dcn_fsdp_parallelism,
-      config.dcn_tensor_parallelism,
-  ]
-  ici_parallelism = [
-      config.ici_data_parallelism,
-      config.ici_fsdp_parallelism,
-      config.ici_tensor_parallelism,
-  ]
+  # multi_slice_env = num_slices > 1
 
   # Find possible unspecified parallelisms
-  ici_parallelism = fill_unspecified_mesh_axes(ici_parallelism, num_devices_per_slice, "ICI")
-  if multi_slice_env:
-    dcn_parallelism = fill_unspecified_mesh_axes(dcn_parallelism, num_slices, "DCN")
-    mesh = mesh_utils.create_hybrid_device_mesh(ici_parallelism, dcn_parallelism, devices)
-  else:
-    mesh = mesh_utils.create_device_mesh(ici_parallelism, devices)
-
-  if logging:
-    max_logging.log(f"Decided on mesh: {mesh}")
+  ici_parallelism = fill_unspecified_mesh_axes(config.ici_parallelism.copy(), num_devices_per_slice, "ICI")
+
+  # allow_split_physical_axes = config.allow_split_physical_axes if config.allow_split_physical_axes else False
+
+  # if allow_split_physical_axes:
+  #   if max_utils.is_valid_custom_mesh(ici_parallelism, config.custom_mesh):
+  #     mesh = mesh_utils.create_device_mesh(
+  #         [16, 16],
+  #         devices,
+  #         contiguous_submeshes=False,
+  #         allow_split_physical_axes=False,
+  #     )
+  #     mesh = max_utils.reshape_mesh_to_rings(mesh, config.custom_mesh)
+  #     mesh = np.reshape(mesh, ici_parallelism)
+  #   else:
+  #     mesh = mesh_utils.create_device_mesh(
+  #         ici_parallelism,
+  #         devices,
+  #         contiguous_submeshes=False,
+  #         allow_split_physical_axes=allow_split_physical_axes,
+  #     )
+  # else:
+  mesh = mesh_utils.create_device_mesh(
+      ici_parallelism,
+      devices,
+  )
+  max_logging.log(f"Num_devices: {num_devices}, shape {mesh.shape}")
 
   return mesh
 
 
+# def create_device_mesh(config, devices=None, logging=True):
+#   """Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas"""
+#   if devices is None:
+#     devices = jax.devices()
+#   num_devices = len(devices)
+#   try:
+#     num_slices = 1 + max([d.slice_index for d in devices])
+#   except:
+#     num_slices = 1
+#   num_devices_per_slice = num_devices // num_slices
+#   max_logging.log(f"Devices: {devices} (num_devices: {num_devices})")
+
+#   multi_slice_env = num_slices > 1
+
+#   dcn_parallelism = [
+#       config.dcn_data_parallelism,
+#       config.dcn_fsdp_parallelism,
+#       config.dcn_tensor_parallelism,
+#   ]
+#   ici_parallelism = [
+#       config.ici_data_parallelism,
+#       config.ici_fsdp_parallelism,
+#       config.ici_tensor_parallelism,
+#   ]
+
+#   # Find possible unspecified parallelisms
+#   ici_parallelism = fill_unspecified_mesh_axes(ici_parallelism, num_devices_per_slice, "ICI")
+#   if multi_slice_env:
+#     dcn_parallelism = fill_unspecified_mesh_axes(dcn_parallelism, num_slices, "DCN")
+#     mesh = mesh_utils.create_hybrid_device_mesh(ici_parallelism, dcn_parallelism, devices)
+#   else:
+#     mesh = mesh_utils.create_device_mesh(ici_parallelism, devices)
+
+#   if logging:
+#     max_logging.log(f"Decided on mesh: {mesh}")
+
+#   return mesh
+
+
 def unbox_logicallypartioned_trainstate(boxed_train_state: train_state.TrainState):
   """Unboxes the flax.LogicallyPartitioned pieces in a train state.
 
@@ -402,7 +445,10 @@ def setup_initial_state(
           config.enable_single_replica_ckpt_restoring,
       )
       if state:
-        state = state[checkpoint_item]
+        if checkpoint_item == " ":
+          state = state
+        else:
+          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(

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

@@ -2,7 +2,6 @@
 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
@@ -198,8 +197,7 @@ def __call__(
 
     # Adaptive Norm
     if self.adaptive_norm in ["single_scale_shift", "single_scale"]:
-      # [batch, 1 or num_tokens, embedding_dim]
-      assert timestep.ndim == 3
+      assert timestep.ndim == 3  # [batch, 1 or num_tokens, embedding_dim]
       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
@@ -438,7 +436,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}
+    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")
@@ -628,8 +626,21 @@ def partial_flash_attention(q, k, v, q_segment_ids, kv_segment_ids):
           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,
+      # )
       # 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)
       wrapped_flash_attention = shard_map(
           partial_flash_attention,
           mesh=sharding_mesh,
@@ -814,8 +825,7 @@ 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")
-      # round to nearest multiple of 256
-      inner_dim = round(inner_dim / 256) * 256
+      inner_dim = round(inner_dim / 256) * 256  # round to nearest multiple of 256
     else:
       inner_dim = self.inner_dim