diff --git a/src/maxdiffusion/configs/ltx2_3_video.yml b/src/maxdiffusion/configs/ltx2_3_video.yml index a6275cfae..eb8c56862 100644 --- a/src/maxdiffusion/configs/ltx2_3_video.yml +++ b/src/maxdiffusion/configs/ltx2_3_video.yml @@ -1,7 +1,13 @@ #hardware hardware: 'tpu' skip_jax_distributed_system: False -attention: 'flash' +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +use_base2_exp: False +use_experimental_scheduler: False +# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. +ulysses_shards: -1 +# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. +ulysses_attention_chunks: 1 a2v_attention_kernel: 'flash' v2a_attention_kernel: 'dot_product' attention_sharding_uniform: True @@ -106,6 +112,10 @@ enable_ondemand_xprof: True skip_first_n_steps_for_profiler: 0 profiler_steps: 5 +# Enable JAX named scopes for detailed profiling and debugging +# When enabled, adds named scopes around key operations in transformer and attention layers +enable_jax_named_scopes: False + replicate_vae: False run_text_encoder_on_tpu: False @@ -173,4 +183,16 @@ upsampler_temporal_patch_size: 1 upsampler_adain_factor: 0.0 upsampler_tone_map_compression_ratio: 0.0 upsampler_rational_spatial_scale: 2.0 -upsampler_output_type: "pil" \ No newline at end of file +upsampler_output_type: "pil" + +aot_cache_dir: '' +converted_weights_dir: '' + + +# Tile-size auto-tuning. When enable_tile_search: True, generate_ltx2 runs a fast one-DiT-block +# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites +# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op). +enable_tile_search: False +tile_search_mode: 'smart' # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep) +tile_search_iters: 10 +tile_search_out: '' # dir for the results CSV; '' -> print only diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml index 8ddfaba1b..fc92f6c6e 100644 --- a/src/maxdiffusion/configs/ltx2_video.yml +++ b/src/maxdiffusion/configs/ltx2_video.yml @@ -1,7 +1,13 @@ #hardware hardware: 'tpu' skip_jax_distributed_system: False -attention: 'flash' +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +use_base2_exp: False +use_experimental_scheduler: False +# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. +ulysses_shards: -1 +# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. +ulysses_attention_chunks: 1 a2v_attention_kernel: 'dot_product' v2a_attention_kernel: 'dot_product' attention_sharding_uniform: True @@ -111,6 +117,10 @@ enable_ondemand_xprof: True skip_first_n_steps_for_profiler: 0 profiler_steps: 5 +# Enable JAX named scopes for detailed profiling and debugging +# When enabled, adds named scopes around key operations in transformer and attention layers +enable_jax_named_scopes: False + replicate_vae: False use_bwe: False @@ -171,3 +181,15 @@ upsampler_adain_factor: 0.0 upsampler_tone_map_compression_ratio: 0.0 upsampler_rational_spatial_scale: 2.0 upsampler_output_type: "pil" + +aot_cache_dir: '' +converted_weights_dir: '' + + +# Tile-size auto-tuning. When enable_tile_search: True, generate_ltx2 runs a fast one-DiT-block +# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites +# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op). +enable_tile_search: False +tile_search_mode: 'smart' # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep) +tile_search_iters: 10 +tile_search_out: '' # dir for the results CSV; '' -> print only diff --git a/src/maxdiffusion/generate_ltx2.py b/src/maxdiffusion/generate_ltx2.py index 0445913ef..eb9e57cf2 100644 --- a/src/maxdiffusion/generate_ltx2.py +++ b/src/maxdiffusion/generate_ltx2.py @@ -19,7 +19,7 @@ import os import subprocess from maxdiffusion.checkpointing.ltx2_checkpointer import LTX2Checkpointer -from maxdiffusion import pyconfig, max_logging, max_utils +from maxdiffusion import aot_cache, pyconfig, max_logging, max_utils from absl import app from google.cloud import storage from google.api_core.exceptions import GoogleAPIError @@ -113,7 +113,57 @@ def call_pipeline(config, pipeline, prompt, negative_prompt): return out +def maybe_tune_block_sizes(config): + """If enable_tile_search, run a fast one-DiT-block tile-size grid search and overwrite + flash_block_sizes' block_q/block_kv/block_kv_compute with the winner IN PLACE. + """ + keys = config.get_keys() + val = keys.get("enable_tile_search", False) + if str(val).lower() not in ("true", "1", "yes"): + return + from maxdiffusion.utils.tile_size_grid_search import grid_search + from maxdiffusion.utils.ltx2_block_benchmark import LTX2BlockBenchmark + + vmem = config.flash_block_sizes.get("vmem_limit_bytes", None) if config.flash_block_sizes else None + if vmem is None: + import os + import re + + m = re.search(r"--xla_tpu_scoped_vmem_limit_kib=(\d+)", os.environ.get("LIBTPU_INIT_ARGS", "")) + vmem = int(m.group(1)) * 1024 if m else 32 * 1024 * 1024 + + mesh = jax.sharding.Mesh(max_utils.create_device_mesh(config), config.mesh_axes) + bench = LTX2BlockBenchmark.from_config(config, mesh, vmem_limit_bytes=vmem) + max_logging.log(f"[tile-search] tuning block sizes for {bench.label} (vmem={vmem/1024/1024:.1f}MB) before inference...") + result = grid_search( + bench, + mode=keys.get("tile_search_mode", "smart"), + iters=keys.get("tile_search_iters", 10), + out_dir=(keys.get("tile_search_out", "") or None), + log=max_logging.log, + ) + if result.best is None: + max_logging.log("[tile-search] no config succeeded; keeping configured flash_block_sizes") + return + fbs = dict(config.flash_block_sizes) if config.flash_block_sizes else {} + fbs.update({ + "block_q": result.best.bq, + "block_kv": result.best.bkv, + "block_kv_compute": result.best.bkv_compute, + "block_kv_compute_in": result.best.bkv_compute, + "vmem_limit_bytes": vmem, + }) + config.get_keys()["flash_block_sizes"] = fbs + max_logging.log( + f"[tile-search] using block_q={result.best.bq} block_kv={result.best.bkv} " + f"(block-bench {result.best.mean_ms:.2f} ms)" + ) + + def run(config, pipeline=None, filename_prefix="", commit_hash=None): + if pipeline is None: + maybe_tune_block_sizes(config) + writer = max_utils.initialize_summary_writer(config) if jax.process_index() == 0 and writer: max_logging.log(f"TensorBoard logs will be written to: {config.tensorboard_dir}") @@ -189,14 +239,37 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): original_enable_mld = config.get_keys().get("enable_ml_diagnostics", False) original_num_steps = config.get_keys().get("num_inference_steps", 40) + # Per-shape AOT executable cache + aot_cache.install( + getattr(config, "aot_cache_dir", ""), + meta={ + "model": config.pretrained_model_name_or_path, + "attention": getattr(config, "attention", ""), + "flash_block_sizes": str(getattr(config, "flash_block_sizes", "")), + "mesh_shape": str(pipeline.mesh.shape) if pipeline and hasattr(pipeline, "mesh") and pipeline.mesh else "", + "weights_dtype": str(getattr(config, "weights_dtype", "bfloat16")), + "activations_dtype": str(getattr(config, "activations_dtype", "bfloat16")), + "scan_layers": str(getattr(config, "scan_layers", True)), + "jax": jax.__version__, + }, + mesh=pipeline.mesh if pipeline else None, + ) + aot_cache.wait_for_loads() + # --------------------------------------------------------- # Run 1: Warmup Compilation (Original steps, NO profiling) # --------------------------------------------------------- config.get_keys()["enable_profiler"] = False config.get_keys()["enable_ml_diagnostics"] = False + warmup_steps = min(2, original_num_steps) + config.get_keys()["num_inference_steps"] = warmup_steps + + max_logging.log(f"🚀 Starting warmup compilation pass ({warmup_steps} steps)...") + with aot_cache.warmup_mode(): + _ = call_pipeline(config, pipeline, prompt, negative_prompt) - max_logging.log(f"🚀 Starting warmup compilation pass ({original_num_steps} steps)...") - _ = call_pipeline(config, pipeline, prompt, negative_prompt) + aot_cache.save_pending() + config.get_keys()["num_inference_steps"] = original_num_steps compile_time = time.perf_counter() - s0 max_logging.log(f"compile_time: {compile_time}") @@ -237,7 +310,9 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): # Export videos for i in range(len(videos)): - video_path = f"{filename_prefix}ltx2_output_{getattr(config, 'seed', 0)}_{i}.mp4" + model_name = getattr(config, "model_name", "ltx2") or "ltx2" + model_name_prefix = model_name.replace(".", "_") + video_path = f"{filename_prefix}{model_name_prefix}_output_{getattr(config, 'seed', 0)}_{i}.mp4" audio_i = audios[i] if audios is not None else None audio_format = getattr(config, "audio_format", "s16") diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index c56d8c64c..a76da7dea 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -1914,7 +1914,7 @@ def __init__( self, mesh: Mesh, attention_kernel: str, - scale: int, + scale: float, heads: int, dim_head: int, use_memory_efficient_attention: bool = False, @@ -2009,7 +2009,7 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask class AttentionOp(nn.Module): mesh: Mesh attention_kernel: str - scale: int + scale: float heads: int dim_head: int use_memory_efficient_attention: bool = False diff --git a/src/maxdiffusion/models/ltx2/attention_ltx2.py b/src/maxdiffusion/models/ltx2/attention_ltx2.py index 214690c98..d9b01ed3e 100644 --- a/src/maxdiffusion/models/ltx2/attention_ltx2.py +++ b/src/maxdiffusion/models/ltx2/attention_ltx2.py @@ -88,13 +88,8 @@ def apply_split_rotary_emb(x: Array, freqs: Tuple[Array, Array]) -> Array: first_x = split_x[..., 0, :] second_x = split_x[..., 1, :] - cos_u = jnp.expand_dims(cos, axis=-2) - sin_u = jnp.expand_dims(sin, axis=-2) - - out = split_x * cos_u - - out_first = out[..., 0, :] - second_x * sin_u.squeeze(-2) - out_second = out[..., 1, :] + first_x * sin_u.squeeze(-2) + out_first = first_x * cos - second_x * sin + out_second = second_x * cos + first_x * sin out = jnp.stack([out_first, out_second], axis=-2) out = out.reshape(*out.shape[:-2], last_dim) @@ -176,12 +171,6 @@ def prepare_video_coords( patch_ends = grid + patch_size_delta # Combine start and end coordinates - latent_coords = jnp.stack([grid, patch_ends], axis=-1) # [3, N_F, N_H, N_W, 2] - latent_coords = latent_coords.transpose(1, 2, 3, 0, 4) # [N_F, N_H, N_W, 3, 2] - latent_coords = latent_coords.reshape(-1, 3, 2) # [num_patches, 3, 2] - latent_coords = jnp.expand_dims(latent_coords, 0) # [1, num_patches, 3, 2] - latent_coords = jnp.tile(latent_coords, (batch_size, 1, 1, 1)) # [B, num_patches, 3, 2] - latent_coords = jnp.stack([grid, patch_ends], axis=-1) # [3, N_F, N_H, N_W, 2] latent_coords = latent_coords.reshape(3, -1, 2) # [3, num_patches, 2] latent_coords = jnp.expand_dims(latent_coords, 0) # [1, 3, num_patches, 2] @@ -352,6 +341,8 @@ def __init__( flash_min_seq_length: int = 4096, sharding_specs: Optional[LTX2DiTShardingSpecs] = None, gated_attn: bool = False, + ulysses_shards: int = -1, + ulysses_attention_chunks: int = 1, ): self.heads = heads self.rope_type = rope_type @@ -445,6 +436,24 @@ def __init__( dtype=dtype, ) + is_self_attention = context_dim is None + if not is_self_attention: + if attention_kernel in ("tokamax_ring", "tokamax_ring_custom", "ulysses_ring"): + attention_kernel = "tokamax_flash" # do not use ring attention for cross attention + if attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir"): + attention_kernel = "ulysses_custom" # plain ulysses (no ring) for cross attention + + axis_names_q = (common_types.BATCH, common_types.CROSS_ATTN_HEAD, common_types.CROSS_ATTN_Q_LENGTH, common_types.D_KV) + axis_names_kv = ( + common_types.BATCH, + common_types.CROSS_ATTN_HEAD, + common_types.CROSS_ATTN_KV_LENGTH, + common_types.D_KV, + ) + else: + axis_names_q = (common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_Q_LENGTH, common_types.D_KV) + axis_names_kv = (common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_KV_LENGTH, common_types.D_KV) + self.attention_op = NNXAttentionOp( mesh=mesh, attention_kernel=attention_kernel, @@ -452,10 +461,12 @@ def __init__( heads=heads, dim_head=dim_head, dtype=dtype, - axis_names_q=(common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_Q_LENGTH, common_types.D_KV), - axis_names_kv=(common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_KV_LENGTH, common_types.D_KV), + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, flash_block_sizes=flash_block_sizes, flash_min_seq_length=flash_min_seq_length, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) def __call__( @@ -485,7 +496,7 @@ def __call__( # 3. Apply RoPE with jax.named_scope("Apply RoPE"): if rotary_emb is not None: - if hasattr(self, "rope_type") and self.rope_type == "split": + if self.rope_type == "split": # Split RoPE: passing full freqs [B, H, S, D//2] # apply_split_rotary_emb handles reshaping query/key diff --git a/src/maxdiffusion/models/ltx2/ltx2_utils.py b/src/maxdiffusion/models/ltx2/ltx2_utils.py index 7cf7fa761..1d4153700 100644 --- a/src/maxdiffusion/models/ltx2/ltx2_utils.py +++ b/src/maxdiffusion/models/ltx2/ltx2_utils.py @@ -17,6 +17,7 @@ import json from typing import Optional import torch +import numpy as np import jax import jax.numpy as jnp from maxdiffusion import max_logging @@ -184,6 +185,14 @@ def load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device, fi return tensors +def _torch_tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray: + import ml_dtypes + + if tensor.dtype == torch.bfloat16: + return tensor.view(torch.uint16).numpy().view(ml_dtypes.bfloat16) + return tensor.numpy() + + def load_transformer_weights( pretrained_model_name_or_path: str, eval_shapes: dict, @@ -193,38 +202,110 @@ def load_transformer_weights( scan_layers: bool = True, subfolder: str = "transformer", ): + import threading + import concurrent.futures + import time + device = jax.local_devices(backend=device)[0] max_logging.log(f"Load and port {pretrained_model_name_or_path} {subfolder} on {device}") - with jax.default_device(device): - # Support sharded loading - tensors = load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device) + index_file = "diffusion_pytorch_model.safetensors.index.json" + try: + index_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=index_file) + with open(index_path, "r") as f: + index_data = json.load(f) + weight_map = index_data["weight_map"] + shards = sorted(set(weight_map.values())) - flax_state_dict = {} - cpu = jax.local_devices(backend="cpu")[0] - flattened_dict = flatten_dict(eval_shapes) + def resolve_shard_path(model_file): + return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=model_file) - random_flax_state_dict = {} - for key in flattened_dict: - random_flax_state_dict[tuple(str(item) for item in key)] = flattened_dict[key] + except EntryNotFoundError: + shards = ["diffusion_pytorch_model.safetensors"] - for pt_key, tensor in tensors.items(): - renamed_pt_key = rename_key(pt_key) - renamed_pt_key = rename_for_ltx2_transformer(renamed_pt_key) + def resolve_shard_path(model_file): + try: + return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=model_file) + except EntryNotFoundError: + return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename="diffusion_pytorch_model.bin") - pt_tuple_key = tuple(renamed_pt_key.split(".")) + t_start = time.perf_counter() - flax_key, flax_tensor = get_key_and_value( - pt_tuple_key, tensor, flax_state_dict, random_flax_state_dict, scan_layers, num_layers - ) + flattened_dict = flatten_dict(eval_shapes) + random_flax_state_dict = {} + for key in flattened_dict: + random_flax_state_dict[tuple(str(item) for item in key)] = flattened_dict[key] - flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu) + flax_state_dict = {} + dict_lock = threading.Lock() + + def convert_chunk(ckpt_shard_path, chunk_keys): + if ckpt_shard_path.endswith(".safetensors"): + with safe_open(ckpt_shard_path, framework="pt") as f: + for pt_key in chunk_keys: + tensor = _torch_tensor_to_numpy(f.get_tensor(pt_key)) + process_tensor(pt_key, tensor) + else: + loaded_state_dict = torch.load(ckpt_shard_path, map_location="cpu") + for pt_key in chunk_keys: + tensor = _torch_tensor_to_numpy(loaded_state_dict[pt_key]) + process_tensor(pt_key, tensor) + + def process_tensor(pt_key, tensor): + renamed_pt_key = rename_key(pt_key) + renamed_pt_key = rename_for_ltx2_transformer(renamed_pt_key) + pt_tuple_key = tuple(renamed_pt_key.split(".")) + + block_index = None + if scan_layers and len(pt_tuple_key) > 0 and "transformer_blocks_" in pt_tuple_key[0]: + import re + + m = re.match(r"transformer_blocks_(\d+)", pt_tuple_key[0]) + if m: + block_index = int(m.group(1)) + pt_tuple_key = ("transformer_blocks",) + pt_tuple_key[1:] - validate_flax_state_dict(eval_shapes, flax_state_dict) - flax_state_dict = unflatten_dict(flax_state_dict) - del tensors - jax.clear_caches() - return flax_state_dict + flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict, scan_layers) + flax_key_str = [str(k) for k in flax_key] + if "scale_shift_table" in flax_key_str and flax_key_str[-1] in ["kernel", "weight"]: + flax_key_str.pop() + flax_key = tuple(flax_key_str) + flax_key = _tuple_str_to_int(flax_key) + + if block_index is not None: + with dict_lock: + stacked = flax_state_dict.get(flax_key) + if stacked is None: + stacked = np.empty((num_layers,) + flax_tensor.shape, dtype=flax_tensor.dtype) + flax_state_dict[flax_key] = stacked + stacked[block_index] = flax_tensor + else: + value = np.array(flax_tensor, dtype=flax_tensor.dtype, copy=True, order="C") + with dict_lock: + flax_state_dict[flax_key] = value + + chunk_size = 32 + tasks = [] + for model_file in shards: + ckpt_shard_path = resolve_shard_path(model_file) + if ckpt_shard_path.endswith(".safetensors"): + with safe_open(ckpt_shard_path, framework="pt") as f: + shard_keys = list(f.keys()) + else: + loaded_state_dict = torch.load(ckpt_shard_path, map_location="cpu") + shard_keys = list(loaded_state_dict.keys()) + for i in range(0, len(shard_keys), chunk_size): + tasks.append((ckpt_shard_path, shard_keys[i : i + chunk_size])) + + with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor: + futures = [executor.submit(convert_chunk, path, keys) for path, keys in tasks] + for future in concurrent.futures.as_completed(futures): + future.result() + + validate_flax_state_dict(eval_shapes, flax_state_dict) + flax_state_dict = unflatten_dict(flax_state_dict) + max_logging.log(f"Converted weights in {time.perf_counter() - t_start:.1f}s") + return flax_state_dict def load_vae_weights( diff --git a/src/maxdiffusion/models/ltx2/text_encoders/torchax_text_encoder.py b/src/maxdiffusion/models/ltx2/text_encoders/torchax_text_encoder.py index 882c65a25..e18cf90c4 100644 --- a/src/maxdiffusion/models/ltx2/text_encoders/torchax_text_encoder.py +++ b/src/maxdiffusion/models/ltx2/text_encoders/torchax_text_encoder.py @@ -20,11 +20,9 @@ import jax from torchax import interop, default_env -# --- Monkeypatch transformers masking_utils to avoid torchax integer tracing bug --- +import contextlib import transformers.masking_utils -_orig_sliding_window_overlay = transformers.masking_utils.sliding_window_overlay - def _patched_sliding_window_overlay(sliding_window: int): # pylint: disable=unused-argument @@ -44,6 +42,16 @@ def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: return inner_mask +@contextlib.contextmanager +def patch_sliding_window_overlay(): + orig = transformers.masking_utils.sliding_window_overlay + transformers.masking_utils.sliding_window_overlay = _patched_sliding_window_overlay + try: + yield + finally: + transformers.masking_utils.sliding_window_overlay = orig + + class TorchaxGemma3TextEncoder(interop.JittableModule): """ A jittable Torchax module for wrapping the HuggingFace PyTorch @@ -57,8 +65,7 @@ def __call__( self, input_ids: jax.Array, attention_mask: jax.Array, output_hidden_states: bool = True ) -> Tuple[jax.Array, ...]: # Dynamically patch transformers.masking_utils only during the duration of this call - transformers.masking_utils.sliding_window_overlay = _patched_sliding_window_overlay - try: + with patch_sliding_window_overlay(): with default_env(): input_ids = interop.torch_view(input_ids) attention_mask = interop.torch_view(attention_mask) @@ -72,9 +79,6 @@ def __call__( output_hidden_states=output_hidden_states, ) return interop.jax_view(output) - finally: - # Restore original behavior to prevent side effects on other potential models in same env - transformers.masking_utils.sliding_window_overlay = _orig_sliding_window_overlay @staticmethod def _forward_inner(model, input_ids, attention_mask, output_hidden_states=True): diff --git a/src/maxdiffusion/models/ltx2/transformer_ltx2.py b/src/maxdiffusion/models/ltx2/transformer_ltx2.py index f2511ed07..4df3c8163 100644 --- a/src/maxdiffusion/models/ltx2/transformer_ltx2.py +++ b/src/maxdiffusion/models/ltx2/transformer_ltx2.py @@ -27,6 +27,33 @@ from maxdiffusion.configuration_utils import ConfigMixin, register_to_config from maxdiffusion.common_types import BlockSizes from .logical_sharding_ltx2 import get_sharding_specs, LTX2DiTShardingSpecs +from flax import struct + + +@struct.dataclass +class LTX2BlockContext: + hidden_states: jax.Array + audio_hidden_states: jax.Array + encoder_hidden_states: jax.Array + audio_encoder_hidden_states: jax.Array + temb: jax.Array + temb_audio: jax.Array + temb_ca_scale_shift: jax.Array + temb_ca_audio_scale_shift: jax.Array + temb_ca_gate: jax.Array + temb_ca_audio_gate: jax.Array + temb_prompt: Optional[jax.Array] = None + temb_prompt_audio: Optional[jax.Array] = None + modality_mask: Optional[jax.Array] = None + video_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None + audio_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None + ca_video_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None + ca_audio_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None + encoder_attention_mask: Optional[jax.Array] = None + audio_encoder_attention_mask: Optional[jax.Array] = None + a2v_cross_attention_mask: Optional[jax.Array] = None + v2a_cross_attention_mask: Optional[jax.Array] = None + perturbation_mask: Optional[jax.Array] = None class LTX2AdaLayerNormSingle(nnx.Module): @@ -128,6 +155,8 @@ def __init__( flash_min_seq_length: int = 4096, sharding_specs: Optional[LTX2DiTShardingSpecs] = None, perturbed_attn: bool = False, + ulysses_shards: int = -1, + ulysses_attention_chunks: int = 1, ): self.dim = dim self.norm_eps = norm_eps @@ -166,6 +195,8 @@ def __init__( flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) self.audio_norm1 = nnx.RMSNorm( @@ -194,6 +225,8 @@ def __init__( flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) # 2. Prompt Cross-Attention @@ -223,6 +256,8 @@ def __init__( flash_block_sizes=flash_block_sizes, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) self.audio_norm2 = nnx.RMSNorm( @@ -252,6 +287,8 @@ def __init__( flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention @@ -279,9 +316,11 @@ def __init__( attention_kernel=a2v_attention_kernel, rope_type=rope_type, flash_block_sizes=flash_block_sizes, - flash_min_seq_length=0, + flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) self.video_to_audio_norm = nnx.RMSNorm( @@ -311,6 +350,8 @@ def __init__( flash_min_seq_length=flash_min_seq_length, sharding_specs=self.sharding_specs, gated_attn=gated_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) # 4. Feed Forward @@ -399,31 +440,49 @@ def __init__( def __call__( self, - hidden_states: jax.Array, # Video - audio_hidden_states: jax.Array, # Audio - encoder_hidden_states: jax.Array, # Context (Text) - audio_encoder_hidden_states: jax.Array, # Audio Context - # Timestep embeddings for AdaLN - temb: jax.Array, - temb_audio: jax.Array, - temb_ca_scale_shift: jax.Array, - temb_ca_audio_scale_shift: jax.Array, - temb_ca_gate: jax.Array, - temb_ca_audio_gate: jax.Array, - temb_prompt: Optional[jax.Array] = None, - temb_prompt_audio: Optional[jax.Array] = None, - modality_mask: Optional[jax.Array] = None, - # RoPE - video_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None, - audio_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None, - ca_video_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None, - ca_audio_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None, - encoder_attention_mask: Optional[jax.Array] = None, - audio_encoder_attention_mask: Optional[jax.Array] = None, - a2v_cross_attention_mask: Optional[jax.Array] = None, - v2a_cross_attention_mask: Optional[jax.Array] = None, - perturbation_mask: Optional[jax.Array] = None, + ctx: "LTX2BlockContext", ) -> Tuple[jax.Array, jax.Array]: + """ + Forward pass of the LTX2 video/audio transformer block. + + This block handles complex multi-modal attention including: + - Video Self-Attention (video -> video) + - Audio Self-Attention (audio -> audio) + - Video Cross-Attention (video -> text caption) + - Audio Cross-Attention (audio -> text caption) + - Video-to-Audio Cross-Attention + - Audio-to-Video Cross-Attention + + Args: + ctx: An `LTX2BlockContext` object containing all hidden states, timestep + embeddings, attention masks, rotary embeddings, and modulation + parameters needed for this layer's forward pass. + + Returns: + A tuple of `(output_hidden_states, output_audio_hidden_states)`. + """ + hidden_states = ctx.hidden_states + audio_hidden_states = ctx.audio_hidden_states + encoder_hidden_states = ctx.encoder_hidden_states + audio_encoder_hidden_states = ctx.audio_encoder_hidden_states + temb = ctx.temb + temb_audio = ctx.temb_audio + temb_ca_scale_shift = ctx.temb_ca_scale_shift + temb_ca_audio_scale_shift = ctx.temb_ca_audio_scale_shift + temb_ca_gate = ctx.temb_ca_gate + temb_ca_audio_gate = ctx.temb_ca_audio_gate + temb_prompt = ctx.temb_prompt + temb_prompt_audio = ctx.temb_prompt_audio + modality_mask = ctx.modality_mask + video_rotary_emb = ctx.video_rotary_emb + audio_rotary_emb = ctx.audio_rotary_emb + ca_video_rotary_emb = ctx.ca_video_rotary_emb + ca_audio_rotary_emb = ctx.ca_audio_rotary_emb + encoder_attention_mask = ctx.encoder_attention_mask + audio_encoder_attention_mask = ctx.audio_encoder_attention_mask + a2v_cross_attention_mask = ctx.a2v_cross_attention_mask + v2a_cross_attention_mask = ctx.v2a_cross_attention_mask + perturbation_mask = ctx.perturbation_mask batch_size = hidden_states.shape[0] axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) @@ -695,6 +754,8 @@ def __init__( use_prompt_embeddings: bool = True, perturbed_attn: bool = False, spatio_temporal_guidance_blocks: Tuple[int, ...] = (), + ulysses_shards: int = -1, + ulysses_attention_chunks: int = 1, **kwargs, ): self.spatio_temporal_guidance_blocks = spatio_temporal_guidance_blocks @@ -879,6 +940,7 @@ def __init__( # 3. Output Layer Scale/Shift Modulation parameters param_rng = rngs.params() + audio_param_rng = rngs.params() table_sharding = self.sharding_specs.scale_shift_table self.scale_shift_table = nnx.Param( nnx.with_partitioning( @@ -889,7 +951,7 @@ def __init__( nnx.with_partitioning( lambda key, shape: jax.random.normal(key, shape, dtype=self.weights_dtype) / jnp.sqrt(audio_inner_dim), table_sharding, - )(param_rng, (2, audio_inner_dim)) + )(audio_param_rng, (2, audio_inner_dim)) ) # 4. Rotary Positional Embeddings (RoPE) @@ -990,6 +1052,8 @@ def init_block(rngs): flash_block_sizes=flash_block_sizes, flash_min_seq_length=self.flash_min_seq_length, perturbed_attn=self.perturbed_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) if self.scan_layers: @@ -999,6 +1063,7 @@ def init_block(rngs): for _ in range(self.num_layers): block = LTX2VideoTransformerBlock( rngs=rngs, + sharding_specs=self.sharding_specs, dim=inner_dim, num_attention_heads=self.num_attention_heads, attention_head_dim=self.attention_head_dim, @@ -1028,6 +1093,8 @@ def init_block(rngs): flash_block_sizes=flash_block_sizes, flash_min_seq_length=self.flash_min_seq_length, perturbed_attn=self.perturbed_attn, + ulysses_shards=ulysses_shards, + ulysses_attention_chunks=ulysses_attention_chunks, ) blocks.append(block) self.transformer_blocks = nnx.List(blocks) @@ -1085,9 +1152,45 @@ def __call__( return_dict: bool = True, perturbation_mask: Optional[jax.Array] = None, ) -> Any: + """ + Forward pass for the full LTX2 Video/Audio Diffusion Transformer. + + Args: + hidden_states: Video latent patches of shape `(batch, seq_len, in_channels)`. + audio_hidden_states: Audio latent patches of shape `(batch, audio_seq_len, audio_in_channels)`. + encoder_hidden_states: Text embeddings for video generation. + audio_encoder_hidden_states: Text embeddings for audio generation. + timestep: Timestep array for video diffusion. + audio_timestep: Optional timestep array for audio diffusion. If None, uses `timestep`. + sigma: Optional noise scale for video (for flow matching). + audio_sigma: Optional noise scale for audio. + encoder_attention_mask: Mask for video text embeddings. + audio_encoder_attention_mask: Mask for audio text embeddings. + num_frames: Number of video frames. + height: Height of the video frames. + width: Width of the video frames. + fps: Frames per second. + audio_num_frames: Number of audio frames. + video_coords: Optional pre-computed 3D coordinates for video RoPE. + audio_coords: Optional pre-computed 1D coordinates for audio RoPE. + attention_kwargs: Additional kwargs for the attention mechanisms. + use_cross_timestep: Whether to use a cross-modal timestep interaction. + modality_mask: Mask indicating which modality to drop/keep. + return_dict: If True, returns a dictionary. Otherwise, returns a tuple. + perturbation_mask: Optional mask for perturbing attention. + + Returns: + Output dict containing `sample` (video) and `audio_sample` (audio). + """ # Determine timestep for audio. audio_timestep = audio_timestep if audio_timestep is not None else timestep + a2v_cross_attention_mask = None + v2a_cross_attention_mask = None + if attention_kwargs is not None: + a2v_cross_attention_mask = attention_kwargs.get("a2v_cross_attention_mask", None) + v2a_cross_attention_mask = attention_kwargs.get("v2a_cross_attention_mask", None) + if self.attention_kernel == "dot_product": if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: encoder_attention_mask = (1 - encoder_attention_mask.astype(self.dtype)) * -10000.0 @@ -1097,6 +1200,14 @@ def __call__( audio_encoder_attention_mask = (1 - audio_encoder_attention_mask.astype(self.dtype)) * -10000.0 audio_encoder_attention_mask = jnp.expand_dims(audio_encoder_attention_mask, axis=1) + if a2v_cross_attention_mask is not None and a2v_cross_attention_mask.ndim == 2: + a2v_cross_attention_mask = (1 - a2v_cross_attention_mask.astype(self.dtype)) * -10000.0 + a2v_cross_attention_mask = jnp.expand_dims(a2v_cross_attention_mask, axis=1) + + if v2a_cross_attention_mask is not None and v2a_cross_attention_mask.ndim == 2: + v2a_cross_attention_mask = (1 - v2a_cross_attention_mask.astype(self.dtype)) * -10000.0 + v2a_cross_attention_mask = jnp.expand_dims(v2a_cross_attention_mask, axis=1) + batch_size = hidden_states.shape[0] # 1. Prepare RoPE positional embeddings @@ -1152,9 +1263,8 @@ def __call__( temb_prompt_audio = None if use_cross_timestep: - assert ( - sigma is not None and audio_sigma is not None - ), "sigma and audio_sigma must be provided when use_cross_timestep is True" + if sigma is None or audio_sigma is None: + raise ValueError("sigma and audio_sigma must be provided when use_cross_timestep is True") video_ca_timestep = audio_sigma.flatten() audio_ca_timestep = sigma.flatten() else: @@ -1195,38 +1305,48 @@ def __call__( audio_encoder_hidden_states = audio_encoder_hidden_states.reshape(batch_size, -1, audio_hidden_states.shape[-1]) # 5. Run transformer blocks with jax.named_scope("Transformer Blocks"): + base_context = LTX2BlockContext( + hidden_states=hidden_states, + audio_hidden_states=audio_hidden_states, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + temb=temb, + temb_audio=temb_audio, + temb_ca_scale_shift=video_cross_attn_scale_shift, + temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, + temb_ca_gate=video_cross_attn_a2v_gate, + temb_ca_audio_gate=audio_cross_attn_v2a_gate, + temb_prompt=temb_prompt, + temb_prompt_audio=temb_prompt_audio, + video_rotary_emb=video_rotary_emb, + audio_rotary_emb=audio_rotary_emb, + ca_video_rotary_emb=video_cross_attn_rotary_emb, + ca_audio_rotary_emb=audio_cross_attn_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + audio_encoder_attention_mask=audio_encoder_attention_mask, + a2v_cross_attention_mask=a2v_cross_attention_mask, + v2a_cross_attention_mask=v2a_cross_attention_mask, + modality_mask=modality_mask, + ) + + def apply_block(block, context: LTX2BlockContext, mask): + orig_perturbation_mask = context.perturbation_mask + context = context.replace(perturbation_mask=mask) + with jax.named_scope("Transformer Layer"): + hidden_states_out, audio_hidden_states_out = block(context) + context = context.replace( + hidden_states=hidden_states_out.astype(context.hidden_states.dtype), + audio_hidden_states=audio_hidden_states_out.astype(context.audio_hidden_states.dtype), + perturbation_mask=orig_perturbation_mask, + ) + return context + if perturbation_mask is None: - # Fast-path: No perturbation masking (standard LTX-2 or disabled STG) + def scan_fn_ltx2(carry, block): - hidden_states, audio_hidden_states, rngs_carry = carry - with jax.named_scope("Transformer Layer"): - hidden_states_out, audio_hidden_states_out = block( - hidden_states=hidden_states, - audio_hidden_states=audio_hidden_states, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - temb=temb, - temb_audio=temb_audio, - temb_ca_scale_shift=video_cross_attn_scale_shift, - temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, - temb_ca_gate=video_cross_attn_a2v_gate, - temb_ca_audio_gate=audio_cross_attn_v2a_gate, - temb_prompt=temb_prompt, - temb_prompt_audio=temb_prompt_audio, - video_rotary_emb=video_rotary_emb, - audio_rotary_emb=audio_rotary_emb, - ca_video_rotary_emb=video_cross_attn_rotary_emb, - ca_audio_rotary_emb=audio_cross_attn_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=audio_encoder_attention_mask, - perturbation_mask=None, - modality_mask=modality_mask, - ) - return ( - hidden_states_out.astype(hidden_states.dtype), - audio_hidden_states_out.astype(audio_hidden_states.dtype), - rngs_carry, - ), None + context, rngs_carry = carry + context = apply_block(block, context, None) + return (context, rngs_carry), None if self.scan_layers: rematted_scan_fn = self.gradient_checkpoint.apply( @@ -1235,40 +1355,23 @@ def scan_fn_ltx2(carry, block): self.names_which_can_be_offloaded, prevent_cse=not self.scan_layers, ) - carry = (hidden_states, audio_hidden_states, nnx.Rngs(0)) - (hidden_states, audio_hidden_states, _), _ = nnx.scan( + carry = (base_context, nnx.Rngs(0)) + (final_context, _), _ = nnx.scan( rematted_scan_fn, length=self.num_layers, in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, 0), transform_metadata={nnx.PARTITION_NAME: "layers"}, )(carry, self.transformer_blocks) + hidden_states = final_context.hidden_states + audio_hidden_states = final_context.audio_hidden_states else: + current_context = base_context for block in self.transformer_blocks: - hidden_states, audio_hidden_states = block( - hidden_states=hidden_states, - audio_hidden_states=audio_hidden_states, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - temb=temb, - temb_audio=temb_audio, - temb_ca_scale_shift=video_cross_attn_scale_shift, - temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, - temb_ca_gate=video_cross_attn_a2v_gate, - temb_ca_audio_gate=audio_cross_attn_v2a_gate, - temb_prompt=temb_prompt, - temb_prompt_audio=temb_prompt_audio, - video_rotary_emb=video_rotary_emb, - audio_rotary_emb=audio_rotary_emb, - ca_video_rotary_emb=video_cross_attn_rotary_emb, - ca_audio_rotary_emb=audio_cross_attn_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=audio_encoder_attention_mask, - perturbation_mask=None, - modality_mask=modality_mask, - ) + current_context = apply_block(block, current_context, None) + hidden_states = current_context.hidden_states + audio_hidden_states = current_context.audio_hidden_states else: - # Slow-path: Dynamic perturbation masking (LTX-2.3 STG enabled) masks = jnp.ones((self.num_layers, batch_size, 1, 1), dtype=self.dtype) for i in self.spatio_temporal_guidance_blocks: if i < self.num_layers: @@ -1277,35 +1380,9 @@ def scan_fn_ltx2(carry, block): def scan_fn_ltx23(carry, block_and_mask): block, mask = block_and_mask - hidden_states, audio_hidden_states, rngs_carry = carry - with jax.named_scope("Transformer Layer"): - hidden_states_out, audio_hidden_states_out = block( - hidden_states=hidden_states, - audio_hidden_states=audio_hidden_states, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - temb=temb, - temb_audio=temb_audio, - temb_ca_scale_shift=video_cross_attn_scale_shift, - temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, - temb_ca_gate=video_cross_attn_a2v_gate, - temb_ca_audio_gate=audio_cross_attn_v2a_gate, - temb_prompt=temb_prompt, - temb_prompt_audio=temb_prompt_audio, - video_rotary_emb=video_rotary_emb, - audio_rotary_emb=audio_rotary_emb, - ca_video_rotary_emb=video_cross_attn_rotary_emb, - ca_audio_rotary_emb=audio_cross_attn_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=audio_encoder_attention_mask, - perturbation_mask=mask, - modality_mask=modality_mask, - ) - return ( - hidden_states_out.astype(hidden_states.dtype), - audio_hidden_states_out.astype(audio_hidden_states.dtype), - rngs_carry, - ), None + context, rngs_carry = carry + context = apply_block(block, context, mask) + return (context, rngs_carry), None if self.scan_layers: rematted_scan_fn = self.gradient_checkpoint.apply( @@ -1314,39 +1391,23 @@ def scan_fn_ltx23(carry, block_and_mask): self.names_which_can_be_offloaded, prevent_cse=not self.scan_layers, ) - carry = (hidden_states, audio_hidden_states, nnx.Rngs(0)) - (hidden_states, audio_hidden_states, _), _ = nnx.scan( + carry = (base_context, nnx.Rngs(0)) + (final_context, _), _ = nnx.scan( rematted_scan_fn, length=self.num_layers, in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, 0), transform_metadata={nnx.PARTITION_NAME: "layers"}, )(carry, (self.transformer_blocks, perturbation_mask_per_layer)) + hidden_states = final_context.hidden_states + audio_hidden_states = final_context.audio_hidden_states else: + current_context = base_context for i, block in enumerate(self.transformer_blocks): mask = perturbation_mask_per_layer[i] if perturbation_mask_per_layer is not None else None - hidden_states, audio_hidden_states = block( - hidden_states=hidden_states, - audio_hidden_states=audio_hidden_states, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - temb=temb, - temb_audio=temb_audio, - temb_ca_scale_shift=video_cross_attn_scale_shift, - temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, - temb_ca_gate=video_cross_attn_a2v_gate, - temb_ca_audio_gate=audio_cross_attn_v2a_gate, - temb_prompt=temb_prompt, - temb_prompt_audio=temb_prompt_audio, - video_rotary_emb=video_rotary_emb, - audio_rotary_emb=audio_rotary_emb, - ca_video_rotary_emb=video_cross_attn_rotary_emb, - ca_audio_rotary_emb=audio_cross_attn_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=audio_encoder_attention_mask, - perturbation_mask=mask, - modality_mask=modality_mask, - ) + current_context = apply_block(block, current_context, mask) + hidden_states = current_context.hidden_states + audio_hidden_states = current_context.audio_hidden_states # 6. Output layers with jax.named_scope("Output Projection & Norm"): diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py index 19069b6f9..542c8cfad 100644 --- a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py @@ -56,9 +56,15 @@ from ...pyconfig import HyperParameters from ... import max_logging from ... import max_utils +from ... import aot_cache from ...max_utils import get_precision, device_put_replicated, get_flash_block_sizes +@partial(jax.jit, static_argnums=(1,)) +def _enforce_layout(x, axes): + return jax.lax.with_sharding_constraint(x, axes) + + TORCH_DTYPE_MAP = { "bfloat16": torch.bfloat16, "float16": torch.float16, @@ -81,6 +87,8 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): """ std_text = jnp.std(noise_pred_text, axis=list(range(1, noise_pred_text.ndim)), keepdims=True) std_cfg = jnp.std(noise_cfg, axis=list(range(1, noise_cfg.ndim)), keepdims=True) + # Prevent division by zero + std_cfg = jnp.maximum(std_cfg, 1e-5) # rescale the results from guidance (fixes overexposure) noise_pred_rescaled = noise_cfg * (std_text / std_cfg) # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images @@ -148,6 +156,8 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict): ltx2_config["precision"] = get_precision(config) ltx2_config["flash_block_sizes"] = get_flash_block_sizes(config) ltx2_config["flash_min_seq_length"] = getattr(config, "flash_min_seq_length", 4096) + ltx2_config["ulysses_shards"] = getattr(config, "ulysses_shards", -1) + ltx2_config["ulysses_attention_chunks"] = getattr(config, "ulysses_attention_chunks", 1) ltx2_config["remat_policy"] = config.remat_policy ltx2_config["names_which_can_be_saved"] = config.names_which_can_be_saved ltx2_config["names_which_can_be_offloaded"] = config.names_which_can_be_offloaded @@ -188,11 +198,58 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict): params = jax.tree_util.tree_map_with_path( lambda path, x: cast_with_exclusion(path, x, dtype_to_cast=config.weights_dtype), params ) + + t_put_start = time.perf_counter() + put_specs = [] for path, val in flax.traverse_util.flatten_dict(params).items(): if restored_checkpoint: path = path[:-1] sharding = logical_state_sharding[path].get_value() - state[path].set_value(device_put_replicated(val, sharding)) + put_specs.append((path, val, sharding)) + + if jax.process_count() == 1: + n_devices = mesh.devices.size + dim0_sharding = NamedSharding(mesh, P(mesh.axis_names)) + + def stages_via_ici(val, sharding) -> bool: + return ( + sharding.is_fully_replicated + and val.ndim > 0 + and val.shape[0] % n_devices == 0 + and val.nbytes >= 1 << 26 # 64MB: below this, staging overhead wins + ) + + staged_ids = [i for i, (_, val, sharding) in enumerate(put_specs) if stages_via_ici(val, sharding)] + direct_ids = [i for i in range(len(put_specs)) if i not in set(staged_ids)] + + put_arrays = [None] * len(put_specs) + if staged_ids: + staged = jax.device_put([put_specs[i][1] for i in staged_ids], [dim0_sharding] * len(staged_ids)) + replicate_fn = jax.jit(lambda xs: xs, out_shardings=[put_specs[i][2] for i in staged_ids]) + for i, replicated in zip(staged_ids, replicate_fn(staged)): + put_arrays[i] = replicated + if direct_ids: + for i, put_array in zip( + direct_ids, + jax.device_put([put_specs[i][1] for i in direct_ids], [put_specs[i][2] for i in direct_ids]), + ): + put_arrays[i] = put_array + for (path, _, _), put_array in zip(put_specs, put_arrays): + state[path].set_value(put_array) + else: + for path, val, sharding in put_specs: + try: + state[path].set_value(device_put_replicated(val, sharding)) + except Exception as e: + max_logging.log(f"Failed to device_put_replicated {path}: {e}") + max_logging.log(f"Trying to use process_allgather for {path}") + val_on_host = jax.experimental.multihost_utils.process_allgather(val, tiled=True) + state[path].set_value(device_put_replicated(val_on_host, sharding)) + del val_on_host + + jax.block_until_ready([state[path].value for path, _, _ in put_specs]) + max_logging.log(f"Transformer {subfolder or 'transformer'} weights on device in {time.perf_counter() - t_put_start:.1f}s") + state = nnx.from_flat_state(state) transformer = nnx.merge(graphdef, state, rest_of_state) @@ -716,25 +773,38 @@ def _create_common_components(cls, config: HyperParameters, vae_only=False): def _load_and_init( cls, config: HyperParameters, restored_checkpoint, vae_only=False, load_transformer=True, load_upsampler=False ): - components = cls._create_common_components(config, vae_only) + import concurrent.futures - transformer = None - if load_transformer: - max_logging.log("Loading Transformer...") - transformer = cls.load_transformer( - devices_array=components["devices_array"], - mesh=components["mesh"], - rngs=components["rngs"], - config=config, - restored_checkpoint=restored_checkpoint, - ) + common_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + common_future = common_executor.submit(cls._create_common_components, config, vae_only) + transformer = None latent_upsampler = None latent_upsampler_params = None - if load_upsampler: - latent_upsampler, latent_upsampler_params = cls.load_upsampler( - devices_array=components["devices_array"], mesh=components["mesh"], rngs=components["rngs"], config=config - ) + + if load_transformer or load_upsampler: + # Need mesh and devices to load these + devices_array = max_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + rngs = nnx.Rngs(jax.random.key(config.seed)) + + if load_transformer: + max_logging.log("Loading Transformer...") + transformer = cls.load_transformer( + devices_array=devices_array, + mesh=mesh, + rngs=rngs, + config=config, + restored_checkpoint=restored_checkpoint, + ) + + if load_upsampler: + latent_upsampler, latent_upsampler_params = cls.load_upsampler( + devices_array=devices_array, mesh=mesh, rngs=rngs, config=config + ) + + components = common_future.result() + common_executor.shutdown() pipeline = cls( scheduler=components["scheduler"], @@ -855,6 +925,8 @@ def _get_gemma_prompt_embeds( prompt = [p.strip() for p in prompt] + target_dtype = dtype if dtype is not None else jnp.bfloat16 + if self.text_encoder is not None: run_text_encoder_on_tpu = getattr(self.config, "run_text_encoder_on_tpu", False) if hasattr(self, "config") else False if run_text_encoder_on_tpu: @@ -872,30 +944,31 @@ def _get_gemma_prompt_embeds( # Distribute the batch dimension across available TPUs to prevent Softmax OOM # (reduces 512MB allocation down to 64MB per TPU for batch size 16) - devices = np.array(jax.devices()) - num_shards = 1 - for i in range(len(devices), 0, -1): - if text_input_ids.shape[0] % i == 0: - num_shards = i - break - - if num_shards > 1: - mesh = Mesh(devices[:num_shards], axis_names=("batch",)) - sharding = NamedSharding(mesh, P("batch")) + if hasattr(self, "mesh") and self.mesh is not None: + data_axis = self.mesh.axis_names[0] + sharding = NamedSharding(self.mesh, P(data_axis)) text_input_ids = jax.device_put(text_input_ids, sharding) prompt_attention_mask = jax.device_put(prompt_attention_mask, sharding) + else: + devices = np.array(jax.devices()) + num_shards = 1 + for i in range(len(devices), 0, -1): + if text_input_ids.shape[0] % i == 0: + num_shards = i + break + + if num_shards > 1: + mesh = Mesh(devices[:num_shards], axis_names=("batch",)) + sharding = NamedSharding(mesh, P("batch")) + text_input_ids = jax.device_put(text_input_ids, sharding) + prompt_attention_mask = jax.device_put(prompt_attention_mask, sharding) # Torchax wrapper returns tuple of hidden states natively text_encoder_hidden_states = self.text_encoder( input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True ) - prompt_embeds_list = [] - # Iterate instead of stacking eagerly to avoid 5.7+ GB HBM allocations outside JIT - for state in text_encoder_hidden_states: - prompt_embeds_list.append(state.astype(jnp.bfloat16)) - - prompt_embeds = prompt_embeds_list + prompt_embeds = jax.tree.map(lambda x: x.astype(target_dtype), list(text_encoder_hidden_states)) del text_encoder_hidden_states # Free memory prompt_attention_mask = prompt_attention_mask.astype(jnp.bool_) @@ -923,25 +996,16 @@ def _get_gemma_prompt_embeds( text_encoder_hidden_states = text_encoder_outputs.hidden_states del text_encoder_outputs # Free memory - prompt_embeds_list = [] - # Iterate instead of stacking eagerly to avoid 5.7+ GB HBM allocations outside JIT - for state in text_encoder_hidden_states: - state_np = state.cpu().to(torch.float32).numpy() - prompt_embeds_list.append(jnp.array(state_np, dtype=jnp.bfloat16)) - - prompt_embeds = prompt_embeds_list + prompt_embeds = jax.tree.map( + lambda state: jnp.array(state.cpu().to(torch.float32).numpy(), dtype=target_dtype), + list(text_encoder_hidden_states), + ) del text_encoder_hidden_states # Free PyTorch tensor memory prompt_attention_mask = jnp.array(prompt_attention_mask.cpu().to(torch.float32).numpy(), dtype=jnp.bool_) else: raise ValueError("`text_encoder` is required to encode prompts.") - if dtype is not None: - if isinstance(prompt_embeds, list): - prompt_embeds = [state.astype(dtype) for state in prompt_embeds] - else: - prompt_embeds = prompt_embeds.astype(dtype) - if isinstance(prompt_embeds, list): _, seq_len, _ = prompt_embeds[0].shape prompt_embeds = [ @@ -1235,6 +1299,9 @@ def _create_noised_state(latents: jax.Array, noise_scale: float, generator: Opti else: # Fallback or expect noise to be handled otherwise? # pipeline prepare_latents typically generates noise. + max_logging.log( + "WARNING: No PRNG generator provided. Falling back to deterministic zero-seed noise (jax.random.key(0))." + ) noise = jax.random.normal(jax.random.key(0), latents.shape, dtype=latents.dtype) # Default fallback noised_latents = noise_scale * noise + (1 - noise_scale) * latents @@ -1617,12 +1684,14 @@ def __call__( audio_embeds_sharded = audio_embeds if not self.transformer.scan_layers: - activation_axes = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) - activation_axes_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed")) - spec = NamedSharding(self.mesh, P(*activation_axes)) - spec_audio = NamedSharding(self.mesh, P(*activation_axes_audio)) - video_embeds_sharded = jax.device_put(video_embeds, spec) - audio_embeds_sharded = jax.device_put(audio_embeds, spec_audio) + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + activation_axes = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) + activation_axes_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed")) + + latents_jax = _enforce_layout(latents_jax, P(*activation_axes)) + audio_latents_jax = _enforce_layout(audio_latents_jax, P(*activation_axes_audio)) + video_embeds_sharded = _enforce_layout(video_embeds_sharded, P(*activation_axes)) + audio_embeds_sharded = _enforce_layout(audio_embeds_sharded, P(*activation_axes_audio)) timesteps_jax = jnp.array(timesteps, dtype=jnp.float32) t0_denoise = time.perf_counter() @@ -1657,6 +1726,8 @@ def __call__( self.scheduler.step, tuple(tuple(rule) if isinstance(rule, list) else rule for rule in self.config.logical_axis_rules), use_cross_timestep=use_cross_timestep, + do_cfg=do_cfg, + do_stg=do_stg, ) else: # Old Python loop path @@ -1668,8 +1739,9 @@ def __call__( audio_latents_jax_sharded = audio_latents_jax if not self.transformer.scan_layers: - activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) - activation_axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed")) + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) + activation_axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed")) latents_jax_sharded = jax.lax.with_sharding_constraint(latents_jax, activation_axis_names) audio_latents_jax_sharded = jax.lax.with_sharding_constraint(audio_latents_jax, activation_axis_names_audio) @@ -1956,7 +2028,7 @@ def convert_to_vel(lat, x0, sig): @partial( - jax.jit, + aot_cache.cached_jit, static_argnames=( "latent_num_frames", "latent_height", @@ -2045,16 +2117,8 @@ def transformer_forward_pass( @partial( - jax.jit, + aot_cache.cached_jit, static_argnames=( - "guidance_scale", - "stg_scale", - "modality_scale", - "guidance_rescale", - "audio_guidance_scale", - "audio_stg_scale", - "audio_modality_scale", - "audio_guidance_rescale", "latent_num_frames", "latent_height", "latent_width", @@ -2065,6 +2129,8 @@ def transformer_forward_pass( "scheduler_step", "logical_axis_rules", "use_cross_timestep", + "do_cfg", + "do_stg", ), ) def run_diffusion_loop( @@ -2097,15 +2163,14 @@ def run_diffusion_loop( logical_axis_rules, perturbation_mask=None, use_cross_timestep=False, + do_cfg=False, + do_stg=False, ): """Runs the diffusion loop.""" # pylint: disable=too-many-positional-arguments latents_jax = latents_jax.astype(jnp.float32) audio_latents_jax = audio_latents_jax.astype(jnp.float32) - do_cfg = guidance_scale > 1.0 - do_stg = stg_scale > 0.0 - # Helper functions matching Diffusers Delta formulation def convert_to_x0(lat, vel, sigma_t): return lat - vel * sigma_t @@ -2117,133 +2182,139 @@ def scan_body(carry, inputs): t, sigma_t = inputs latents, audio_latents, s_state = carry - with nn_partitioning.axis_rules(logical_axis_rules): - latents_sharded = latents - audio_latents_sharded = audio_latents + latents_sharded = latents + audio_latents_sharded = audio_latents - if not scan_layers: + if not scan_layers: + with nn_partitioning.axis_rules(logical_axis_rules): activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed")) - latents_sharded = jax.lax.with_sharding_constraint(latents, activation_axis_names) - audio_latents_sharded = jax.lax.with_sharding_constraint(audio_latents, activation_axis_names) - - # Forward Pass - noise_pred, noise_pred_audio = transformer_forward_pass( - graphdef, - state, - latents_sharded, - audio_latents_sharded, - t, - video_embeds_sharded, - audio_embeds_sharded, - new_attention_mask, - new_attention_mask, - latent_num_frames=latent_num_frames, - latent_height=latent_height, - latent_width=latent_width, - audio_num_frames=audio_num_frames, - fps=fps, - global_batch_size=batch_size, - sigma=t, - audio_sigma=t, - use_cross_timestep=use_cross_timestep, - is_cfg_stg_mode=do_cfg and do_stg, - ) + activation_axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed")) + latents_sharded = jax.lax.with_sharding_constraint(latents, activation_axis_names) + audio_latents_sharded = jax.lax.with_sharding_constraint(audio_latents, activation_axis_names_audio) + video_embeds_sharded_constrained = jax.lax.with_sharding_constraint(video_embeds_sharded, activation_axis_names) + audio_embeds_sharded_constrained = jax.lax.with_sharding_constraint(audio_embeds_sharded, activation_axis_names_audio) + else: + video_embeds_sharded_constrained = video_embeds_sharded + audio_embeds_sharded_constrained = audio_embeds_sharded + + # Forward Pass + noise_pred, noise_pred_audio = transformer_forward_pass( + graphdef, + state, + latents_sharded, + audio_latents_sharded, + t, + video_embeds_sharded_constrained, + audio_embeds_sharded_constrained, + new_attention_mask, + new_attention_mask, + latent_num_frames=latent_num_frames, + latent_height=latent_height, + latent_width=latent_width, + audio_num_frames=audio_num_frames, + fps=fps, + global_batch_size=batch_size, + sigma=t, + audio_sigma=t, + use_cross_timestep=use_cross_timestep, + is_cfg_stg_mode=do_cfg and do_stg, + ) - # Extract latents_step based on stacking strategy - if do_cfg and do_stg: - latents_step = latents[batch_size : 2 * batch_size] - audio_latents_step = audio_latents[batch_size : 2 * batch_size] - elif do_cfg: - latents_step = latents[batch_size:] - audio_latents_step = audio_latents[batch_size:] - else: - latents_step = latents - audio_latents_step = audio_latents + # Extract latents_step based on stacking strategy + if do_cfg and do_stg: + latents_step = latents[batch_size : 2 * batch_size] + audio_latents_step = audio_latents[batch_size : 2 * batch_size] + elif do_cfg: + latents_step = latents[batch_size:] + audio_latents_step = audio_latents[batch_size:] + else: + latents_step = latents + audio_latents_step = audio_latents - # Apply Diffusers STG + CFG + Modality Delta Logic - if do_cfg and do_stg: - noise_pred_uncond, noise_pred_text, noise_pred_perturb, noise_pred_isolated = jnp.split(noise_pred, 4, axis=0) + # Apply Diffusers STG + CFG + Modality Delta Logic + if do_cfg and do_stg: + noise_pred_uncond, noise_pred_text, noise_pred_perturb, noise_pred_isolated = jnp.split(noise_pred, 4, axis=0) - # Convert to x0 - x0_uncond = convert_to_x0(latents_step, noise_pred_uncond, sigma_t) - x0_text = convert_to_x0(latents_step, noise_pred_text, sigma_t) - x0_perturb = convert_to_x0(latents_step, noise_pred_perturb, sigma_t) - x0_isolated = convert_to_x0(latents_step, noise_pred_isolated, sigma_t) + # Convert to x0 + x0_uncond = convert_to_x0(latents_step, noise_pred_uncond, sigma_t) + x0_text = convert_to_x0(latents_step, noise_pred_text, sigma_t) + x0_perturb = convert_to_x0(latents_step, noise_pred_perturb, sigma_t) + x0_isolated = convert_to_x0(latents_step, noise_pred_isolated, sigma_t) - # Delta formulation - cfg_delta = (guidance_scale - 1) * (x0_text - x0_uncond) - stg_delta = stg_scale * (x0_text - x0_perturb) - video_modality_delta = (modality_scale - 1) * (x0_text - x0_isolated) + # Delta formulation + cfg_delta = (guidance_scale - 1) * (x0_text - x0_uncond) + stg_delta = stg_scale * (x0_text - x0_perturb) + video_modality_delta = (modality_scale - 1) * (x0_text - x0_isolated) - x0_combined = x0_text + cfg_delta + stg_delta + video_modality_delta + x0_combined = x0_text + cfg_delta + stg_delta + video_modality_delta - if guidance_rescale > 0: - x0_combined = rescale_noise_cfg(x0_combined, x0_text, guidance_rescale=guidance_rescale) + x0_combined = jax.lax.cond( + guidance_rescale > 0, + lambda: rescale_noise_cfg(x0_combined, x0_text, guidance_rescale=guidance_rescale), + lambda: x0_combined, + ) - noise_pred = convert_to_vel(latents_step, x0_combined, sigma_t) + noise_pred = convert_to_vel(latents_step, x0_combined, sigma_t) - # Audio guidance - noise_pred_audio_uncond, noise_pred_audio_text, noise_pred_audio_perturb, noise_pred_audio_isolated = jnp.split( - noise_pred_audio, 4, axis=0 - ) + # Audio guidance + noise_pred_audio_uncond, noise_pred_audio_text, noise_pred_audio_perturb, noise_pred_audio_isolated = jnp.split( + noise_pred_audio, 4, axis=0 + ) - x0_audio_uncond = convert_to_x0(audio_latents_step, noise_pred_audio_uncond, sigma_t) - x0_audio_text = convert_to_x0(audio_latents_step, noise_pred_audio_text, sigma_t) - x0_audio_perturb = convert_to_x0(audio_latents_step, noise_pred_audio_perturb, sigma_t) - x0_audio_isolated = convert_to_x0(audio_latents_step, noise_pred_audio_isolated, sigma_t) + x0_audio_uncond = convert_to_x0(audio_latents_step, noise_pred_audio_uncond, sigma_t) + x0_audio_text = convert_to_x0(audio_latents_step, noise_pred_audio_text, sigma_t) + x0_audio_perturb = convert_to_x0(audio_latents_step, noise_pred_audio_perturb, sigma_t) + x0_audio_isolated = convert_to_x0(audio_latents_step, noise_pred_audio_isolated, sigma_t) - cfg_audio_delta = (audio_guidance_scale - 1 if audio_guidance_scale is not None else guidance_scale - 1) * ( - x0_audio_text - x0_audio_uncond - ) - stg_audio_delta = (audio_stg_scale if audio_stg_scale is not None else stg_scale) * ( - x0_audio_text - x0_audio_perturb - ) - audio_modality_delta = (audio_modality_scale - 1 if audio_modality_scale is not None else modality_scale - 1) * ( - x0_audio_text - x0_audio_isolated - ) + cfg_audio_delta = (audio_guidance_scale - 1 if audio_guidance_scale is not None else guidance_scale - 1) * ( + x0_audio_text - x0_audio_uncond + ) + stg_audio_delta = (audio_stg_scale if audio_stg_scale is not None else stg_scale) * (x0_audio_text - x0_audio_perturb) + audio_modality_delta = (audio_modality_scale - 1 if audio_modality_scale is not None else modality_scale - 1) * ( + x0_audio_text - x0_audio_isolated + ) - x0_audio_combined = x0_audio_text + cfg_audio_delta + stg_audio_delta + audio_modality_delta + x0_audio_combined = x0_audio_text + cfg_audio_delta + stg_audio_delta + audio_modality_delta - if audio_guidance_rescale > 0: - x0_audio_combined = rescale_noise_cfg(x0_audio_combined, x0_audio_text, guidance_rescale=audio_guidance_rescale) + x0_audio_combined = jax.lax.cond( + audio_guidance_rescale > 0, + lambda: rescale_noise_cfg(x0_audio_combined, x0_audio_text, guidance_rescale=audio_guidance_rescale), + lambda: x0_audio_combined, + ) - noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined, sigma_t) + noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined, sigma_t) - # ... (Standard CFG paths can be added here, but for brevity and since LTX2.3 runs with STG this handles the core logic) - elif do_cfg: - noise_pred_uncond, noise_pred_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + # ... (Standard CFG paths can be added here, but for brevity and since LTX2.3 runs with STG this handles the core logic) + elif do_cfg: + noise_pred_uncond, noise_pred_text = jnp.split(noise_pred, 2, axis=0) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred_audio_uncond, noise_pred_audio_text = jnp.split(noise_pred_audio, 2, axis=0) - noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (noise_pred_audio_text - noise_pred_audio_uncond) + noise_pred_audio_uncond, noise_pred_audio_text = jnp.split(noise_pred_audio, 2, axis=0) + noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (noise_pred_audio_text - noise_pred_audio_uncond) - # Step scheduler - latents_step, _ = scheduler_step(s_state, noise_pred, t, latents_step, return_dict=False) - latents_step = latents_step.astype(latents.dtype) + # Step scheduler + latents_step, _ = scheduler_step(s_state, noise_pred, t, latents_step, return_dict=False) + latents_step = latents_step.astype(latents.dtype) - audio_latents_step, _ = scheduler_step(s_state, noise_pred_audio, t, audio_latents_step, return_dict=False) - audio_latents_step = audio_latents_step.astype(audio_latents.dtype) + audio_latents_step, _ = scheduler_step(s_state, noise_pred_audio, t, audio_latents_step, return_dict=False) + audio_latents_step = audio_latents_step.astype(audio_latents.dtype) - if do_cfg and do_stg: - latents_next = jnp.concatenate([latents_step] * 4, axis=0) - audio_latents_next = jnp.concatenate([audio_latents_step] * 4, axis=0) - elif do_cfg: - latents_next = jnp.concatenate([latents_step] * 2, axis=0) - audio_latents_next = jnp.concatenate([audio_latents_step] * 2, axis=0) - else: - latents_next = latents_step - audio_latents_next = audio_latents_step + if do_cfg and do_stg: + latents_next = jnp.concatenate([latents_step] * 4, axis=0) + audio_latents_next = jnp.concatenate([audio_latents_step] * 4, axis=0) + elif do_cfg: + latents_next = jnp.concatenate([latents_step] * 2, axis=0) + audio_latents_next = jnp.concatenate([audio_latents_step] * 2, axis=0) + else: + latents_next = latents_step + audio_latents_next = audio_latents_step - new_carry = (latents_next, audio_latents_next, s_state) - return new_carry, None + new_carry = (latents_next, audio_latents_next, s_state) + return new_carry, None initial_carry = (latents_jax, audio_latents_jax, scheduler_state) scan_inputs = (timesteps_jax, sigmas) - final_carry, _ = nnx.scan( - scan_body, - in_axes=(nnx.Carry, 0), - out_axes=(nnx.Carry, 0), - )(initial_carry, scan_inputs) + final_carry, _ = jax.lax.scan(scan_body, initial_carry, scan_inputs) return final_carry[0], final_carry[1] diff --git a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py index 831699474..4e4d86634 100644 --- a/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py +++ b/src/maxdiffusion/tests/ltx2/test_pipeline_ltx2.py @@ -188,13 +188,16 @@ def test_encode_prompt_no_cfg(self, list_embed_mock): self.assertIsNone(n_e) self.assertIsNone(n_a) + @patch("maxdiffusion.pipelines.ltx2.ltx2_pipeline.Mesh") + @patch("maxdiffusion.pipelines.ltx2.ltx2_pipeline.max_utils.create_device_mesh") @patch("maxdiffusion.pipelines.ltx2.ltx2_pipeline.LTX2Pipeline.load_transformer") @patch("maxdiffusion.pipelines.ltx2.ltx2_pipeline.LTX2Pipeline._create_common_components") @patch("maxdiffusion.pipelines.ltx2.ltx2_pipeline.LTX2Pipeline.quantize_transformer") - def test_load_and_init(self, mock_quantize, mock_create_common, mock_load_transformer): + def test_load_and_init(self, mock_quantize, mock_create_common, mock_load_transformer, mock_create_device_mesh, mock_Mesh): """Test that pipeline loading correctly wires all the dependencies down to __init__.""" mock_config = MagicMock() mock_mesh = MagicMock() + mock_Mesh.return_value = mock_mesh mock_common = { "vae": MagicMock(), diff --git a/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py b/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py index f8c5b8d62..281b82784 100644 --- a/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py +++ b/src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py @@ -25,6 +25,7 @@ from maxdiffusion import pyconfig from maxdiffusion.max_utils import create_device_mesh from maxdiffusion.models.ltx2.transformer_ltx2 import ( + LTX2BlockContext, LTX2VideoTransformerBlock, LTX2VideoTransformer3DModel, LTX2AdaLayerNormSingle, @@ -203,7 +204,7 @@ def test_ltx2_transformer_block(self): temb_ca_gate = jnp.zeros((batch_size, 1 * dim)) temb_ca_audio_gate = jnp.zeros((batch_size, 1 * audio_dim)) - output_hidden, output_audio = block( + ctx = LTX2BlockContext( hidden_states=hidden_states, audio_hidden_states=audio_hidden_states, encoder_hidden_states=encoder_hidden_states, @@ -216,6 +217,8 @@ def test_ltx2_transformer_block(self): temb_ca_audio_gate=temb_ca_audio_gate, ) + output_hidden, output_audio = block(ctx) + self.assertEqual(output_hidden.shape, hidden_states.shape) self.assertEqual(output_audio.shape, audio_hidden_states.shape) diff --git a/src/maxdiffusion/utils/ltx2_block_benchmark.py b/src/maxdiffusion/utils/ltx2_block_benchmark.py new file mode 100644 index 000000000..3aefd7ee4 --- /dev/null +++ b/src/maxdiffusion/utils/ltx2_block_benchmark.py @@ -0,0 +1,323 @@ +""" +LTX2 implementation of the model-agnostic `BlockBenchmark` (see tile_size_grid_search.py). +""" +import os +import functools + +os.environ.setdefault( + "LIBTPU_INIT_ARGS", + " ".join([ + "--xla_tpu_dvfs_p_state=7", + "--xla_tpu_spmd_rng_bit_generator_unsafe=true", + "--xla_tpu_enable_dot_strength_reduction=true", + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true", + "--xla_enable_async_collective_permute=true", + "--xla_tpu_enable_async_collective_fusion=true", + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true", + "--xla_tpu_overlap_compute_collective_tc=true", + "--xla_enable_async_all_gather=true", + "--xla_tpu_scoped_vmem_limit_kib=65536", + "--xla_tpu_enable_async_all_to_all=true", + "--xla_tpu_enable_all_experimental_scheduler_features=true", + "--xla_tpu_enable_latency_hiding_scheduler=true", + "--xla_tpu_enable_megacore_fusion=true", + ]), +) +os.environ.setdefault("JAX_DEFAULT_MATMUL_PRECISION", "bfloat16") +os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.95") + +import jax +import jax.numpy as jnp +from flax import linen as nn +from flax import nnx +from flax.linen import partitioning as nn_partitioning + +from maxdiffusion import max_logging, max_utils, pyconfig +from maxdiffusion.models.ltx2.transformer_ltx2 import LTX2VideoTransformer3DModel +from maxdiffusion.utils.tile_size_grid_search import ( + BenchResult, + BlockBenchmark, + grid_search, + time_callable, +) + +_IN_CHANNELS = 128 # LTX2 uses 128 channels in latent space +_VAE_T, _VAE_S = ( + 8, + 32, +) # LTX2 spatial/temporal compression: patch_size_t=1, patch_size=1 (it actually operates on 32x32 compressed) Wait, LTX2 spatial downsampling is 32x32 in pixel space, so in latent space it's 1x1 patch? No, let's just use the latent shape. + + +def latent_seq_len(num_frames: int, height: int, width: int) -> int: + lf = (num_frames - 1) // 8 + 1 + lh, lw = height // 32, width // 32 + return lf * lh * lw + + +def tiled_seq_len(full_seq: int, attention: str, context_shards: int, ulysses_shards: int) -> int: + _RING_VARIANTS = {"ulysses_ring_custom", "ulysses_ring_custom_bidir", "tokamax_ring", "ring"} + if attention in _RING_VARIANTS and ulysses_shards >= 1 and context_shards >= 1: + ring_shards = max(1, context_shards // max(1, ulysses_shards)) + return full_seq // ring_shards + return full_seq + + +@functools.partial(nnx.jit, static_argnames=("num_frames", "height", "width")) +def _forward( + model, + latents, + timestep, + prompt_embeds, + prompt_attention_mask, + audio_latents, + audio_prompt_embeds, + audio_prompt_attention_mask, + *, + num_frames, + height, + width, +): + return model( + hidden_states=latents, + audio_hidden_states=audio_latents, + encoder_hidden_states=prompt_embeds, + audio_encoder_hidden_states=audio_prompt_embeds, + timestep=timestep, + audio_timestep=None, + sigma=None, + audio_sigma=None, + encoder_attention_mask=prompt_attention_mask, + audio_encoder_attention_mask=audio_prompt_attention_mask, + num_frames=num_frames, + height=height, + width=width, + fps=24.0, + audio_num_frames=128, + video_coords=None, + audio_coords=None, + attention_kwargs=None, + use_cross_timestep=False, + ) + + +class LTX2BlockBenchmark(BlockBenchmark): + + def __init__( + self, + config, + mesh, + *, + num_frames, + height, + width, + batch=None, + vmem_limit_bytes=64 * 1024 * 1024, + ): + self._config = config + self._mesh = mesh + self._rules = config.logical_axis_rules + self._attention = getattr(config, "attention", "flash") + self._context_shards = int(mesh.shape.get("context", 1)) + self._ulysses_shards = int(getattr(config, "ulysses_shards", 1) or 1) + self._vmem = int(vmem_limit_bytes) + self.label = f"ltx2/{self._attention}/u{self._ulysses_shards}" + + self._lf = (num_frames - 1) // 8 + 1 + self._lh, self._lw = height // 32, width // 32 + self._full_seq = latent_seq_len(num_frames, height, width) + self._num_frames_orig = num_frames + self._height_orig = height + self._width_orig = width + + data_shards = int(mesh.shape.get("data", 1)) * int(mesh.shape.get("fsdp", 1)) + self._batch = batch if batch is not None else max(1, data_shards) + self._hf_cfg = LTX2VideoTransformer3DModel.load_config(config.pretrained_model_name_or_path, subfolder="transformer") + self._inputs = self._make_inputs() + + @classmethod + def from_config(cls, config, mesh, **overrides): + return cls( + config, + mesh, + num_frames=config.num_frames, + height=config.height, + width=config.width, + **overrides, + ) + + def tiled_seq_lens(self): + s = tiled_seq_len(self._full_seq, self._attention, self._context_shards, self._ulysses_shards) + return (s, s) + + def vmem_bytes(self): + return self._vmem + + def run(self, bq, bkv, *, bkv_compute=None, iters=10, warmup=2): + cmp = bkv_compute or bkv + try: + with self._mesh: + model = self._build_model(bq, bkv, cmp) + ( + latents, + timestep, + prompt_embeds, + prompt_attention_mask, + audio_latents, + audio_prompt_embeds, + audio_prompt_attention_mask, + ) = self._inputs + with self._mesh, nn_partitioning.axis_rules(self._rules): + mean, std, times, compile_ms = time_callable( + lambda: _forward( + model, + latents, + timestep, + prompt_embeds, + prompt_attention_mask, + audio_latents, + audio_prompt_embeds, + audio_prompt_attention_mask, + num_frames=self._lf, + height=self._lh, + width=self._lw, + ), + iters=iters, + warmup=warmup, + sync=jax.block_until_ready, + ) + return BenchResult( + bq, + bkv, + cmp, + "ok", + mean_ms=mean, + std_ms=std, + times_ms=times, + compile_ms=compile_ms, + ) + except Exception as e: + import traceback + + traceback.print_exc() + msg = str(e) + oom = any(t in msg for t in ("RESOURCE_EXHAUSTED", "out of memory", "Mosaic", "VMEM")) + return BenchResult(bq, bkv, cmp, "oom" if oom else "error", detail=msg[:200]) + + def _flash_block_sizes(self, bq, bkv, cmp): + from maxdiffusion.max_utils import CustomFlashBlockSizes + + return CustomFlashBlockSizes( + block_q=bq, + block_kv=bkv, + block_kv_compute=cmp, + block_kv_compute_in=cmp, + heads_per_tile=1, + vmem_limit_bytes=self._vmem, + ) + + def _build_model(self, bq, bkv, cmp): + c = self._config + ltx2_config = dict(self._hf_cfg) + if ltx2_config.get("activation_fn") == "gelu-approximate": + ltx2_config["activation_fn"] = "gelu" + ltx2_config.update( + mesh=self._mesh, + dtype=c.activations_dtype, + weights_dtype=c.weights_dtype, + attention_kernel=c.attention, + a2v_attention_kernel=getattr(c, "a2v_attention_kernel", "flash"), + v2a_attention_kernel=getattr(c, "v2a_attention_kernel", "dot_product"), + precision=max_utils.get_precision(c), + flash_block_sizes=self._flash_block_sizes(bq, bkv, cmp), + flash_min_seq_length=getattr(c, "flash_min_seq_length", 4096), + ulysses_shards=getattr(c, "ulysses_shards", -1), + ulysses_attention_chunks=getattr(c, "ulysses_attention_chunks", 1), + remat_policy=getattr(c, "remat_policy", "NONE"), + scan_layers=False, + num_layers=1, + ) + model = LTX2VideoTransformer3DModel(**ltx2_config, rngs=nnx.Rngs(params=0)) + gd, state, rest = nnx.split(model, nnx.Param, ...) + shardings = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), self._mesh, self._rules) + return nnx.merge(gd, jax.device_put(state, shardings), rest) + + def _make_inputs(self): + dtype = self._config.activations_dtype + k1, k2, k3, k4 = jax.random.split(jax.random.key(0), 4) + seq_len = self._lf * self._lh * self._lw + latents = jax.random.normal(k1, (self._batch, seq_len, _IN_CHANNELS), dtype) + + # Audio latents + audio_seq_len = 128 + audio_latents = jax.random.normal(k3, (self._batch, audio_seq_len, _IN_CHANNELS), dtype) + + # Prompts + prompt_embeds = jax.random.normal(k2, (self._batch, 1024, 3840), dtype) + prompt_attention_mask = jnp.ones((self._batch, 1024), dtype=jnp.int32) + + audio_prompt_embeds = jax.random.normal(k4, (self._batch, 1024, 3840), dtype) + audio_prompt_attention_mask = jnp.ones((self._batch, 1024), dtype=jnp.int32) + + timestep = jnp.zeros((self._batch,), jnp.float32) + repl = jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec()) + return ( + jax.device_put(latents, repl), + jax.device_put(timestep, repl), + jax.device_put(prompt_embeds, repl), + jax.device_put(prompt_attention_mask, repl), + jax.device_put(audio_latents, repl), + jax.device_put(audio_prompt_embeds, repl), + jax.device_put(audio_prompt_attention_mask, repl), + ) + + +def _build_cli_config(argv_yaml, attention, ulysses_shards, num_frames, height, width): + pyconfig.initialize([ + "ltx2_block_benchmark", + argv_yaml, + f"attention={attention}", + f"ulysses_shards={ulysses_shards}", + "skip_jax_distributed_system=True", + "weights_dtype=bfloat16", + "activations_dtype=bfloat16", + "per_device_batch_size=1", + f"num_frames={num_frames}", + f"height={height}", + f"width={width}", + "ici_data_parallelism=1", + "ici_fsdp_parallelism=1", + f"ici_context_parallelism={jax.device_count()}", + "ici_tensor_parallelism=1", + "allow_split_physical_axes=True", + "use_base2_exp=true", + "use_experimental_scheduler=true", + "enable_jax_named_scopes=false", + ]) + return pyconfig.config + + +def main(): + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("config", help="Path to yaml config (e.g. configs/ltx2_video.yml)") + parser.add_argument("--attention", default="ulysses_ring_custom") + parser.add_argument("--ulysses-shards", type=int, default=1) + parser.add_argument("--num-frames", type=int, default=161) + parser.add_argument("--height", type=int, default=512) + parser.add_argument("--width", type=int, default=768) + parser.add_argument("--smart-search", action="store_true") + parser.add_argument("--full-search", action="store_true") + parser.add_argument("--out-dir", default="") + args = parser.parse_args() + + config = _build_cli_config(args.config, args.attention, args.ulysses_shards, args.num_frames, args.height, args.width) + mesh = jax.sharding.Mesh(max_utils.create_device_mesh(config), config.mesh_axes) + bench = LTX2BlockBenchmark.from_config(config, mesh) + + mode = "full" if args.full_search else "smart" + grid_search(bench, mode=mode, out_dir=args.out_dir or None, log=max_logging.log) + + +if __name__ == "__main__": + main()