one transformer inference pass done

Author: Serenagu525

src/maxdiffusion/configs/ltx_video.yml

@@ -12,7 +12,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']],
@@ -25,13 +25,19 @@ 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
 
 
 
@@ -48,4 +54,4 @@ per_device_batch_size: 1
 compile_topology_num_slices: -1 
 quantization_local_shard_count: -1
 jit_initializers: True 
-enable_single_replica_ckpt_restoring: False
+enable_single_replica_ckpt_restoring: False

src/maxdiffusion/generate_ltx_video.py

@@ -28,6 +28,48 @@ def validate_transformer_inputs(prompt_embeds, fractional_coords, latents, noise
   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 loop_body(
+    step,
+    args,
+    transformer,
+    fractional_cords,
+    prompt_embeds,
+    segment_ids,
+    encoder_attention_segment_ids
+):
+  latents, state, noise_cond = args
+  noise_pred = transformer.apply(
+      {"params": state.params},
+      hidden_states=latents,
+      indices_grid=fractional_cords,
+      encoder_hidden_states=prompt_embeds,
+      timestep=noise_cond,
+      segment_ids=segment_ids,
+      encoder_attention_segment_ids=encoder_attention_segment_ids
+  )
+  return noise_pred, state, noise_cond 
+
+
+
+def run_inference(
+    states, transformer, config, mesh, latents, fractional_cords, prompt_embeds, timestep, segment_ids, encoder_attention_segment_ids
+):
+  transformer_state = states["transformer"]
+  loop_body_p = functools.partial(
+      loop_body,
+      transformer=transformer,
+      fractional_cords=fractional_cords,
+      prompt_embeds=prompt_embeds,
+      segment_ids=segment_ids,
+      encoder_attention_segment_ids=encoder_attention_segment_ids
+  )
+  ## TODO: add vae decode step
+  ## TODO: add loop
+  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    latents, transformer_state, _ = jax.lax.fori_loop(0, 1, loop_body_p, (latents, transformer_state, timestep))   
+  return latents
+  
 def run(config):
   key = jax.random.PRNGKey(0)
 
@@ -50,7 +92,7 @@ def run(config):
 
 
   transformer = Transformer3DModel(**model_config, dtype=jnp.float32, gradient_checkpointing="matmul_without_batch", sharding_mesh=mesh)
-  transformer_param_shapes = transformer.init_weights(in_channels, model_config['caption_channels'], eval_only = True)
+  transformer_param_shapes = transformer.init_weights(in_channels, model_config['caption_channels'], eval_only = True) 
 
   weights_init_fn = functools.partial(
       transformer.init_weights, 
@@ -75,7 +117,61 @@ def run(config):
   )
 
 
+  
+  
+  transformer_state = jax.device_put(transformer_state, transformer_state_shardings)
+  get_memory_allocations()
+
+  states = {}
+  state_shardings = {}
+
+  state_shardings["transformer"] = transformer_state_shardings
+  states["transformer"] = transformer_state
+
+  #create dummy inputs:
+  example_inputs = {}
+  batch_size, num_tokens = 4, 256
+  input_shapes = {
+    "latents": (batch_size, num_tokens, in_channels),
+    "fractional_coords": (batch_size, 3, num_tokens),
+    "prompt_embeds": (batch_size, 128, model_config["caption_channels"]),
+    "timestep": (batch_size, 256), 
+    "segment_ids": (batch_size, 256),
+    "encoder_attention_segment_ids": (batch_size, 128),
+  }
+  for name, shape in input_shapes.items():
+    example_inputs[name] = jnp.ones(
+      shape, dtype=jnp.float32 if name not in ["attention_mask", "encoder_attention_mask"] else jnp.bool
+    )
+
+  data_sharding = jax.sharding.NamedSharding(mesh, P(*config.data_sharding))
+  latents = jax.device_put(example_inputs["latents"], data_sharding)
+  prompt_embeds = jax.device_put(example_inputs["prompt_embeds"], data_sharding)
+  fractional_coords = jax.device_put(example_inputs["fractional_coords"], data_sharding)
+  noise_cond = jax.device_put(example_inputs["timestep"], data_sharding)
+  segment_ids = jax.device_put(example_inputs["segment_ids"], data_sharding)
+  encoder_attention_segment_ids = jax.device_put(example_inputs["encoder_attention_segment_ids"], data_sharding)
+
+  validate_transformer_inputs(prompt_embeds, fractional_coords, latents, noise_cond, segment_ids, encoder_attention_segment_ids)
+  p_run_inference = jax.jit(
+      functools.partial(
+          run_inference,
+          transformer=transformer,
+          config=config,
+          mesh=mesh,
+          latents=latents,
+          fractional_cords=fractional_coords,
+          prompt_embeds=prompt_embeds,
+          timestep = noise_cond,
+          segment_ids=segment_ids,
+          encoder_attention_segment_ids=encoder_attention_segment_ids
+      ),
+      in_shardings=(state_shardings,),
+      out_shardings=None,
+  )
 
+  noise_pred = p_run_inference(states).block_until_ready()
+  print(noise_pred)  #(4, 256, 128)
 
 
 def main(argv: Sequence[str]) -> None:
@@ -89,4 +185,14 @@ def main(argv: Sequence[str]) -> None:
 
 
 
+  
+
+
+
+
+
+
+
+
+  
 

src/maxdiffusion/max_utils.py

@@ -251,46 +251,24 @@ 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
   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}")
+  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
 
-
 def unbox_logicallypartioned_trainstate(boxed_train_state: train_state.TrainState):
   """Unboxes the flax.LogicallyPartitioned pieces in a train state.
 
@@ -612,4 +590,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

@@ -631,27 +631,13 @@ 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,
+                ("data", "fsdp", "fsdp_transpose", "expert"),
+                ("tensor", "tensor_transpose", "sequence", "tensor_sequence"),
                 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)
+            qkv_segment_ids_spec = jax.sharding.PartitionSpec(("data", "fsdp", "fsdp_transpose", "expert"), "sequence")
             wrapped_flash_attention = shard_map(
                 partial_flash_attention,
                 mesh=sharding_mesh,
@@ -910,4 +896,4 @@ def apply_rotary_emb(input_tensor: jax.Array, freqs_cis: Tuple[jax.Array, jax.Ar
     # Apply rotary embeddings
     out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs
 
-    return out
+    return out

src/maxdiffusion/pyconfig.py

@@ -41,6 +41,21 @@ 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__}")
@@ -154,6 +169,7 @@ 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)
+    raw_keys = create_parallelisms_list(raw_keys)
 
 
 def get_num_slices(raw_keys):
@@ -204,4 +220,4 @@ def initialize(argv, **kwargs):
 if __name__ == "__main__":
   initialize(sys.argv)
   print(config.steps)
-  r = range(config.steps)
+  r = range(config.steps)