transformer step and test

Author: Serenagu525

src/maxdiffusion/configs/ltx_video.yml

@@ -22,12 +22,25 @@ weights_dtype: 'bfloat16'
 activations_dtype: 'bfloat16'
 
 
+run_name: ''
+output_dir: 'ltx-video-output'
+save_config_to_gcs: False
+
+#hardware
+hardware: 'tpu'
+skip_jax_distributed_system: False
+
+jax_cache_dir: ''
+weights_dtype: 'bfloat16'
+activations_dtype: 'bfloat16'
+
+
 run_name: ''
 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,13 +53,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
 
 
 
@@ -63,3 +82,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

src/maxdiffusion/generate_ltx_video.py

@@ -1,63 +1,180 @@
-"""
- Copyright 2025 Google LLC
-
- Licensed under the Apache License, Version 2.0 (the "License");
- you may not use this file except in compliance with the License.
- You may obtain a copy of the License at
-
-      https://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing, software
- distributed under the License is distributed on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- See the License for the specific language governing permissions and
- limitations under the License.
-"""
-
 from absl import app
 from typing import Sequence
 import jax
 import json
+from flax.linen import partitioning as nn_partitioning
 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
+from jax.sharding import Mesh, PartitionSpec as P
+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 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,
+  )
+  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    noise_pred, transformer_state, _ = jax.lax.fori_loop(0, 1, loop_body_p, (latents, transformer_state, timestep))
+  return noise_pred
 
 
 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)  # noqa F841
+  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)
+  relative_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(in_channels, key, model_config["caption_channels"], eval_only=True)  # noqa F841
+  weights_init_fn = functools.partial(
+      transformer.init_weights, in_channels, key, model_config["caption_channels"], eval_only=True
+  )
 
-  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)  # noqa F841
+  absolute_ckpt_path = os.path.abspath(relative_ckpt_path)
+
+  checkpoint_manager = ocp.CheckpointManager(absolute_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,
+  )
 
-  key, split_key = jax.random.split(key)
-  weights_init_fn = functools.partial(  # noqa F841
-      transformer.init_weights, split_key, batch_size, text_tokens, num_tokens, features, eval_only=True
+  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:
   pyconfig.initialize(argv)

src/maxdiffusion/max_utils.py

@@ -257,6 +257,21 @@ 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 config.ici_fsdp_transpose_parallelism:
+    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:
@@ -402,7 +417,11 @@ def setup_initial_state(
           config.enable_single_replica_ckpt_restoring,
       )
       if state:
-        state = state[checkpoint_item]
+        ###!Edited
+        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/xora_v1.2-13B-balanced-128.json

@@ -1,4 +1,5 @@
 {
+    "ckpt_path": "",
     "activation_fn": "gelu-approximate",
     "attention_bias": true,
     "attention_head_dim": 128,

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,8 @@ 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):