diff --git a/.gitignore b/.gitignore index 4f3a2088e..657c99fb8 100644 --- a/.gitignore +++ b/.gitignore @@ -33,11 +33,11 @@ env **.txt *.log weights/ -official_weights/ -converted_weights/ # SSIM test outputs fastvideo/tests/ssim/generated_videos/ +**/.cache/** + # Distribution / packaging build/ diff --git a/examples/inference/basic/basic_ltx2.py b/examples/inference/basic/basic_ltx2.py index 8ee0ee09f..c7196e291 100644 --- a/examples/inference/basic/basic_ltx2.py +++ b/examples/inference/basic/basic_ltx2.py @@ -20,10 +20,10 @@ def main() -> None: # Uses FastVideo default sampling settings for LTX2 base. generator = VideoGenerator.from_pretrained( "Davids048/LTX2-Base-Diffusers", - num_gpus=8, + num_gpus=1, ) - output_path = "outputs_video/ltx2_basic/output_ltx2_base_t2v_1088_1920_1.4.mp4" + output_path = "outputs_video/ltx2_basic/output_ltx2_base_t2v_1088_1920_1.1.mp4" generator.generate_video( prompt=PROMPT, output_path=output_path, @@ -31,15 +31,9 @@ def main() -> None: num_frames=121, height=1088, width=1920, - # LTX2 uses these parameters for multi-modal CFG instead of guidance_scale - # ltx2_cfg_scale_video=3.0, - # ltx2_cfg_scale_audio=7.0, - # ltx2_modality_scale_video=3.0, - # ltx2_modality_scale_audio=3.0, - # ltx2_rescale_scale=0.7, ) generator.shutdown() if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/examples/inference/basic/basic_ltx2_distilled.py b/examples/inference/basic/basic_ltx2_distilled.py index d8acd56bc..1ac732eaf 100644 --- a/examples/inference/basic/basic_ltx2_distilled.py +++ b/examples/inference/basic/basic_ltx2_distilled.py @@ -13,12 +13,13 @@ PROMPT = ( "The woman covers her face, on the verge of tears. The tone is deadpan, " "absurd, and quietly tragic." ) - +import os +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" def main() -> None: generator = VideoGenerator.from_pretrained( "FastVideo/LTX2-Distilled-Diffusers", - num_gpus=1, + num_gpus=4, ) output_path = "outputs_video/ltx2_basic/output_ltx2_distilled_t2v.mp4" diff --git a/examples/training/finetune/ltx2/.gitignore b/examples/training/finetune/ltx2/.gitignore new file mode 100644 index 000000000..1082d908d --- /dev/null +++ b/examples/training/finetune/ltx2/.gitignore @@ -0,0 +1,3 @@ +*.err +*.out +*slurm_logs/* \ No newline at end of file diff --git a/examples/training/finetune/ltx2/overfit/README.md b/examples/training/finetune/ltx2/README.md similarity index 68% rename from examples/training/finetune/ltx2/overfit/README.md rename to examples/training/finetune/ltx2/README.md index 2e19da33e..9e279887a 100644 --- a/examples/training/finetune/ltx2/overfit/README.md +++ b/examples/training/finetune/ltx2/README.md @@ -7,15 +7,17 @@ These are e2e example scripts for finetuning LTX-2 on the crush-smol dataset. ### Download crush-smol dataset: -`bash examples/training/finetune/ltx2/overfit/download_dataset.sh` +`bash examples/training/finetune/ltx2/download_dataset.sh` + +#### Or use the scripts at scripts/dataset_preparation to download and prepare the dataset ### Preprocess the videos and captions into latents: -`bash examples/training/finetune/ltx2/overfit/preprocess_ltx2_data_t2v_new.sh` +`bash examples/training/finetune/ltx2/preprocess_ltx2_data_t2v_new.sh` ### Edit the following file and run finetuning: -`bash examples/training/finetune/ltx2/overfit/finetune_t2v.sh` +`bash examples/training/finetune/ltx2/finetune_t2v.sh` Notes: - Update `DATASET_PATH` in the preprocess script to point to your merged dataset root (`videos/` + `videos2caption.json`). diff --git a/examples/training/finetune/ltx2/overfit/download_dataset.sh b/examples/training/finetune/ltx2/download_dataset.sh similarity index 100% rename from examples/training/finetune/ltx2/overfit/download_dataset.sh rename to examples/training/finetune/ltx2/download_dataset.sh diff --git a/examples/training/finetune/ltx2/overfit/finetune_t2v.sh b/examples/training/finetune/ltx2/finetune_t2v.sh similarity index 75% rename from examples/training/finetune/ltx2/overfit/finetune_t2v.sh rename to examples/training/finetune/ltx2/finetune_t2v.sh index f05bb640c..c69bb598d 100644 --- a/examples/training/finetune/ltx2/overfit/finetune_t2v.sh +++ b/examples/training/finetune/ltx2/finetune_t2v.sh @@ -4,16 +4,14 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -MODEL_PATH="Davids048/LTX2-Base-Diffusers" -# Also can use simple 1 video for overfitting experiments. -# DATA_DIR="/home/hal-jundas/codes/FastVideo/data/crush-smol" -DATA_DIR="" +MODEL_PATH="FastVideo/LTX2-Distilled-Diffusers" +DATA_DIR="data/crush-smol" VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" echo VALIDATION_DATASET_FILE: $VALIDATION_DATASET_FILE NUM_GPUS=4 -OVERFIT_HEIGHT=480 -OVERFIT_WIDTH=832 -OVERFIT_FRAMES=73 +HEIGHT=1088 +WIDTH=1920 +FRAMES=121 training_args=( --tracker_project_name "ltx2_t2v_finetune" @@ -22,18 +20,17 @@ training_args=( --train_batch_size 1 --train_sp_batch_size 1 --gradient_accumulation_steps 1 - --num_latent_t 10 - --num_height $OVERFIT_HEIGHT - --num_width $OVERFIT_WIDTH - --num_frames $OVERFIT_FRAMES - --ltx2-first-frame-conditioning-p 0.1 + --num_latent_t 16 + --num_height $HEIGHT + --num_width $WIDTH + --num_frames $FRAMES --enable_gradient_checkpointing_type "full" --mode "finetuning" ) parallel_args=( --num_gpus $NUM_GPUS - --sp_size $NUM_GPUS + --sp_size 2 --tp_size 1 --hsdp_replicate_dim 1 --hsdp_shard_dim $NUM_GPUS @@ -52,9 +49,9 @@ dataset_args=( validation_args=( --log_validation --validation_dataset_file $VALIDATION_DATASET_FILE - --validation_steps 50 - --validation_sampling_steps "50" - --validation_guidance_scale "3.0" + --validation_steps 20 + --validation_sampling_steps "8" + --validation_guidance_scale "1.0" ) optimizer_args=( @@ -76,14 +73,15 @@ miscellaneous_args=( --text_encoder_cpu_offload False --image_encoder_cpu_offload False --vae_cpu_offload False + --ltx2-first-frame-conditioning-p 0.0 ) # NOTE: Setting this environment variable to TORCH_SDPA to avoid the issue of stacking that failed in flash attn. -export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA torchrun \ --nnodes 1 \ + --master_port 29501 \ --nproc_per_node $NUM_GPUS \ fastvideo/training/ltx2_training_pipeline.py \ "${parallel_args[@]}" \ diff --git a/examples/training/finetune/ltx2/overfit/finetune_t2v_lora.sh b/examples/training/finetune/ltx2/finetune_t2v_lora.sh similarity index 97% rename from examples/training/finetune/ltx2/overfit/finetune_t2v_lora.sh rename to examples/training/finetune/ltx2/finetune_t2v_lora.sh index d59a43466..65695a204 100644 --- a/examples/training/finetune/ltx2/overfit/finetune_t2v_lora.sh +++ b/examples/training/finetune/ltx2/finetune_t2v_lora.sh @@ -4,7 +4,7 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -MODEL_PATH="/path/to/LTX-2" +MODEL_PATH="FastVideo/LTX2-Distilled-Diffusers" DATA_DIR="data/crush-smol" VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" NUM_GPUS=1 diff --git a/examples/training/finetune/ltx2/overfit/validation.frog.json b/examples/training/finetune/ltx2/overfit/validation.frog.json deleted file mode 100644 index 047c0802d..000000000 --- a/examples/training/finetune/ltx2/overfit/validation.frog.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "data": [ - { - "caption": "The camera opens in a calm, sunlit frog yoga studio. Warm morning light washes over the wooden floor as incense smoke drifts lazily in the air. The senior frog instructor sits cross-legged at the center, eyes closed, voice deep and calm. “We are one with the pond.” All the frogs answer softly: “Ommm...” “We are one with the mud.” “Ommm...” He smiles faintly. “We are one with the flies.” A quiet pause. The camera slowly pans to the side — one frog twitches, eyes darting. Suddenly — *thwip!* — its tongue snaps out, catching a fly mid-air and pulling it into its mouth. The master exhales slowly, still serene.", - "image_path": null, - "video_path": null, - "num_inference_steps": 50, - "height": 1088, - "width": 1920, - "num_frames": 121 - } - ] -} \ No newline at end of file diff --git a/examples/training/finetune/ltx2/overfit/validation.json b/examples/training/finetune/ltx2/overfit/validation.json deleted file mode 100644 index 674c3cb62..000000000 --- a/examples/training/finetune/ltx2/overfit/validation.json +++ /dev/null @@ -1,31 +0,0 @@ -{ - "data": [ - { - "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.", - "image_path": null, - "video_path": null, - "num_inference_steps": 50, - "height": 480, - "width": 832, - "num_frames": 77 - }, - { - "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.", - "image_path": null, - "video_path": null, - "num_inference_steps": 50, - "height": 480, - "width": 832, - "num_frames": 77 - }, - { - "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.", - "image_path": null, - "video_path": null, - "num_inference_steps": 50, - "height": 480, - "width": 832, - "num_frames": 77 - } - ] -} diff --git a/examples/training/finetune/ltx2/overfit/preprocess_ltx2_data_t2v_new.sh b/examples/training/finetune/ltx2/preprocess_ltx2_data_t2v_new.sh similarity index 52% rename from examples/training/finetune/ltx2/overfit/preprocess_ltx2_data_t2v_new.sh rename to examples/training/finetune/ltx2/preprocess_ltx2_data_t2v_new.sh index 911b8d6b5..e878bc059 100644 --- a/examples/training/finetune/ltx2/overfit/preprocess_ltx2_data_t2v_new.sh +++ b/examples/training/finetune/ltx2/preprocess_ltx2_data_t2v_new.sh @@ -1,21 +1,13 @@ #!/bin/bash -GPU_NUM=1 -MODEL_PATH="Davids048/LTX2-Base-Diffusers" -# DATASET_PATH="data/overfit" +GPU_NUM=4 +MODEL_PATH="FastVideo/LTX2-Distilled-Diffusers" DATASET_PATH="data/crush-smol" OUTPUT_DIR="$DATASET_PATH" WITH_AUDIO=true -# Convert one-file overfit metadata into merged format if needed. -if [ ! -f "$DATASET_PATH/videos2caption.json" ] && [ -f "$DATASET_PATH/overfit.json" ]; then - python scripts/dataset_preparation/convert_to_merged_dataset.py \ - --items-json "$DATASET_PATH/overfit.json" \ - --output-dir "$DATASET_PATH" -fi - -torchrun --nproc_per_node=$GPU_NUM \ +torchrun --nproc_per_node=$GPU_NUM \ --master_port=29513 \ -m fastvideo.pipelines.preprocess.v1_preprocessing_new \ --model_path $MODEL_PATH \ @@ -28,8 +20,8 @@ torchrun --nproc_per_node=$GPU_NUM \ --preprocess.with_audio $WITH_AUDIO \ --preprocess.preprocess_video_batch_size 1 \ --preprocess.dataloader_num_workers 0 \ - --preprocess.max_height 480 \ - --preprocess.max_width 832 \ - --preprocess.num_frames 73 \ - --preprocess.train_fps 16 \ + --preprocess.max_height 1088 \ + --preprocess.max_width 1920 \ + --preprocess.num_frames 121 \ + --preprocess.train_fps 24 \ --preprocess.video_length_tolerance_range 5 diff --git a/examples/training/finetune/ltx2/validation.json b/examples/training/finetune/ltx2/validation.json new file mode 100644 index 000000000..14e1a7134 --- /dev/null +++ b/examples/training/finetune/ltx2/validation.json @@ -0,0 +1,16 @@ +{ + "data": [ + { + "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press." + }, + { + "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press." + }, + { + "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table." + }, + { + "caption": "A massive steel piston descends onto a stack of chocolate chip cookies, crushing them into crumbs as though they are being compressed by a hydraulic press." + } + ] +} \ No newline at end of file diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index c1281b8ae..0f978d64e 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/fastvideo/attention/backends/flash_attn.py @@ -92,11 +92,59 @@ class FlashAttentionImpl(AttentionImpl): value: torch.Tensor, attn_metadata: FlashAttnMetadata, ): + + def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor, + key_len: int) -> torch.Tensor: + # Normalize attn_mask to [B, key_len] where True means valid token. + if attn_mask.dim() == 4: + attn_mask = attn_mask[:, 0, 0, :] + elif attn_mask.dim() == 3: + attn_mask = attn_mask[:, 0, :] + elif attn_mask.dim() != 2: + raise ValueError( + f"Unsupported attn_mask shape for FLASH_ATTN: {attn_mask.shape}" + ) + + if attn_mask.dtype == torch.bool: + key_padding_mask = attn_mask + else: + # SDPA additive mask convention: valid=0, masked=-inf/large negative. + key_padding_mask = attn_mask >= 0 + + if key_padding_mask.shape[-1] != key_len: + raise ValueError( + "Invalid key padding mask length for FLASH_ATTN: " + f"expected {key_len}, got {key_padding_mask.shape[-1]}") + return key_padding_mask + if attn_metadata is not None and hasattr( attn_metadata, "attn_mask") and attn_metadata.attn_mask is not None: - from fastvideo.attention.utils.flash_attn_no_pad import flash_attn_no_pad + from fastvideo.attention.utils.flash_attn_no_pad import ( + flash_attn_no_pad, flash_attn_varlen_qk_no_pad) attn_mask = attn_metadata.attn_mask + + # flash_attn_no_pad packs q/k/v as one tensor and assumes equal q/k + # sequence lengths. Cross-attention can violate this. + if query.shape[1] != key.shape[1]: + query_padding_mask = torch.ones( + (query.shape[0], query.shape[1]), + dtype=torch.bool, + device=query.device, + ) + key_padding_mask = _key_padding_mask_from_attn_mask( + attn_mask, key.shape[1]).to(device=key.device) + return flash_attn_varlen_qk_no_pad( + query, + key, + value, + query_padding_mask=query_padding_mask, + key_padding_mask=key_padding_mask, + causal=self.causal, + dropout_p=0.0, + softmax_scale=self.softmax_scale, + ) + qkv = torch.stack([query, key, value], dim=2) attn_mask = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), diff --git a/fastvideo/attention/utils/flash_attn_no_pad.py b/fastvideo/attention/utils/flash_attn_no_pad.py index 1983f8d38..60bd96651 100644 --- a/fastvideo/attention/utils/flash_attn_no_pad.py +++ b/fastvideo/attention/utils/flash_attn_no_pad.py @@ -97,3 +97,60 @@ def flash_attn_no_pad_v3(qkv, "b s (h d) -> b s h d", h=nheads) return output + + +def flash_attn_varlen_qk_no_pad( + query, + key, + value, + query_padding_mask, + key_padding_mask, + causal=False, + dropout_p=0.0, + softmax_scale=None, + deterministic=False, +): + from flash_attn.bert_padding import pad_input, unpad_input + + try: + from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_impl + except ImportError: + from flash_attn import flash_attn_varlen_func as flash_attn_varlen_func_impl + + if flash_attn_varlen_func_impl is None: + raise ImportError("FlashAttention varlen backend not available") + + batch_size, q_seqlen, nheads, _ = query.shape + + query_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input( + rearrange(query, "b s h d -> b s (h d)"), query_padding_mask) + key_unpad, _, cu_seqlens_k, max_seqlen_k, _ = unpad_input( + rearrange(key, "b s h d -> b s (h d)"), key_padding_mask) + value_unpad, _, _, _, _ = unpad_input( + rearrange(value, "b s h d -> b s (h d)"), key_padding_mask) + + query_unpad = rearrange(query_unpad, "nnz (h d) -> nnz h d", h=nheads) + key_unpad = rearrange(key_unpad, "nnz (h d) -> nnz h d", h=nheads) + value_unpad = rearrange(value_unpad, "nnz (h d) -> nnz h d", h=nheads) + + output_unpad = flash_attn_varlen_func_impl( + query_unpad, + key_unpad, + value_unpad, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + ) + + output = rearrange( + pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), q_indices, + batch_size, q_seqlen), + "b s (h d) -> b s h d", + h=nheads, + ) + return output diff --git a/fastvideo/configs/models/encoders/gemma.py b/fastvideo/configs/models/encoders/gemma.py index 73c51dd03..5205120ec 100644 --- a/fastvideo/configs/models/encoders/gemma.py +++ b/fastvideo/configs/models/encoders/gemma.py @@ -7,6 +7,19 @@ from fastvideo.configs.models.encoders.base import ( ) +def _is_feature_extractor_linear(n: str, m) -> bool: + return n.endswith("feature_extractor_linear") + + +def _is_embeddings(n: str, m) -> bool: + return n.endswith("embeddings_connector") or n.endswith( + "audio_embeddings_connector") + + +def _is_gemma_model(n: str, m) -> bool: + return "_gemma_model" in n + + @dataclass class LTX2GemmaArchConfig(TextEncoderArchConfig): architectures: list[str] = field( @@ -35,6 +48,10 @@ class LTX2GemmaArchConfig(TextEncoderArchConfig): connector_double_precision_rope: bool = False connector_num_learnable_registers: int | None = 128 + _fsdp_shard_conditions: list = field( + default_factory=lambda: + [_is_feature_extractor_linear, _is_embeddings, _is_gemma_model]) + def __post_init__(self) -> None: super().__post_init__() self.tokenizer_kwargs["padding"] = "max_length" diff --git a/fastvideo/configs/models/vaes/ltx2vae.py b/fastvideo/configs/models/vaes/ltx2vae.py index d659b9d2a..c4ac77f7c 100644 --- a/fastvideo/configs/models/vaes/ltx2vae.py +++ b/fastvideo/configs/models/vaes/ltx2vae.py @@ -16,6 +16,7 @@ class LTX2VAEArchConfig(VAEArchConfig): in_channels: int = 3 out_channels: int = 3 latent_channels: int = 128 + z_dim: int = 128 # follow num_channels_latents encoder_blocks: list = field(default_factory=list) decoder_blocks: list = field(default_factory=list) patch_size: int = 4 diff --git a/fastvideo/dataset/parquet_dataset_map_style.py b/fastvideo/dataset/parquet_dataset_map_style.py index faffd1df8..dac622497 100644 --- a/fastvideo/dataset/parquet_dataset_map_style.py +++ b/fastvideo/dataset/parquet_dataset_map_style.py @@ -95,28 +95,9 @@ class DP_SP_BatchSampler(Sampler[list[int]]): def get_parquet_files_and_length(path: str): - path = os.path.abspath(path) - if not os.path.exists(path): - raise FileNotFoundError( - f"Dataset path does not exist: {path}. " - "Expected a directory containing *.parquet files.") - if not os.path.isdir(path): - raise NotADirectoryError( - f"Dataset path is not a directory: {path}. " - "Expected a directory containing *.parquet files.") - has_parquet_file = False - for _, _, files in os.walk(path): - if any(file.endswith(".parquet") for file in files): - has_parquet_file = True - break - if not has_parquet_file: - raise FileNotFoundError( - "No parquet files found under dataset path: " - f"{path}. " - "Please verify this path points to preprocessed parquet data.") - + dataset_root = os.path.realpath(os.path.expanduser(path)) # Check if cached info exists - cache_dir = os.path.join(path, "map_style_cache") + cache_dir = os.path.join(dataset_root, "map_style_cache") cache_file = os.path.join(cache_dir, "file_info.pkl") # Only rank 0 checks for cache and scans files if needed @@ -131,12 +112,39 @@ def get_parquet_files_and_length(path: str): try: with open(cache_file, "rb") as f: file_names_sorted, lengths_sorted = pickle.load(f) - if len(file_names_sorted) == 0: - raise ValueError( - "Cached parquet metadata is empty in " - f"{cache_file}.") - cache_loaded = True - logger.info("Successfully loaded cached file info") + file_names_sorted = tuple( + os.path.realpath( + os.path.join(os.getcwd(), p) + if not os.path.isabs(p) else p) + for p in file_names_sorted) + files_outside_dataset_root = [ + file_path for file_path in file_names_sorted + if os.path.commonpath([dataset_root, file_path + ]) != dataset_root + ] + missing_files = [ + file_path for file_path in file_names_sorted + if not os.path.exists(file_path) + ] + if files_outside_dataset_root: + logger.warning( + "Cached parquet file list points outside dataset root " + "(%s). Cache will be rebuilt. First out-of-root file: %s", + dataset_root, + files_outside_dataset_root[0], + ) + cache_loaded = False + elif missing_files: + logger.warning( + "Cached parquet file list contains %d missing files. " + "Cache will be rebuilt. First missing file: %s", + len(missing_files), + missing_files[0], + ) + cache_loaded = False + else: + cache_loaded = True + logger.info("Successfully loaded cached file info") except Exception as e: logger.error("Error loading cached file info: %s", str(e)) logger.info("Falling back to scanning files") @@ -147,10 +155,10 @@ def get_parquet_files_and_length(path: str): logger.info("Scanning parquet files to get lengths") lengths = [] file_names = [] - for root, _, files in os.walk(path): + for root, _, files in os.walk(dataset_root): for file in sorted(files): if file.endswith('.parquet'): - file_path = os.path.join(root, file) + file_path = os.path.realpath(os.path.join(root, file)) file_names.append(file_path) if len(file_names) == 0: raise FileNotFoundError( diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index db9b9a21c..cb041c11e 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -855,7 +855,13 @@ class TransformerArgsPreprocessor: ) def prepare(self, modality: Modality) -> TransformerArgs: - x = self.patchify_proj(modality.latent) + # `rearrange`-produced latent tokens can be non-contiguous with large + # leading strides. Materialize a contiguous view before Linear to avoid + # invalid BF16 GEMM descriptors on some CUDA/cuBLAS stacks. + latent = modality.latent + if not latent.is_contiguous(): + latent = latent.contiguous() + x = self.patchify_proj(latent) timestep, embedded_timestep = self._prepare_timestep(modality.timesteps, x.shape[0], modality.latent.dtype) context, attention_mask = self._prepare_context(modality.context, x, modality.context_mask) attention_mask = self._prepare_attention_mask(attention_mask, modality.latent.dtype) @@ -1017,6 +1023,7 @@ class LTXDistributedAttention(DistributedAttention): replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, + gate_compress: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: @@ -1043,6 +1050,56 @@ class LTXDistributedAttention(DistributedAttention): forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata + use_vsa = self.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN + + if use_vsa: + assert ( + replicated_q is None and replicated_k is None + and replicated_v is None + ), "Replicated QKV is not supported for VSA now" + if gate_compress is None: + raise ValueError( + "gate_compress must be provided when using VIDEO_SPARSE_ATTN" + ) + + qkvg = torch.cat([q, k, v, gate_compress], + dim=0) # [4*batch, seq_len, heads, head_dim] + qkvg = sequence_model_parallel_all_to_all_4D(qkvg, + scatter_dim=2, + gather_dim=1) + + valid_seq_len = None + if attention_mask is not None: + valid_seq_len = (attention_mask[0] == 1).sum().item() + qkvg = qkvg[:, :valid_seq_len, :, :] + + if ltx_freqs_cis is not None: + cos, sin = ltx_freqs_cis + heads_per_rank = num_heads // world_size + head_start = local_rank * heads_per_rank + head_end = head_start + heads_per_rank + cos_local = cos[:, head_start:head_end, :valid_seq_len, :] + sin_local = sin[:, head_start:head_end, :valid_seq_len, :] + qk_part = qkvg[:batch_size * 2].transpose(1, 2) + qk_part = apply_ltx_rotary_emb_4d( + qk_part, (cos_local, sin_local), self.rope_type) + qkvg[:batch_size * 2] = qk_part.transpose(1, 2) + + qkvg = self.attn_impl.preprocess_qkv(qkvg, ctx_attn_metadata) + q, k, v, gate_compress = qkvg.chunk(4, dim=0) + output = self.attn_impl.forward( # type: ignore[call-arg] + q, k, v, gate_compress, ctx_attn_metadata) + output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) + + if attention_mask is not None: + pad_len = (attention_mask[0] == 0).sum().item() + output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len)) + + output = sequence_model_parallel_all_to_all_4D(output, + scatter_dim=1, + gather_dim=2) + return output, None + # Stack QKV qkv = torch.cat([q, k, v], dim=0) # [3*batch, seq_len, num_heads, head_dim] @@ -1143,6 +1200,7 @@ class LTXLocalAttention(LocalAttention): q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + gate_compress: torch.Tensor | None = None, ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, ltx_k_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: @@ -1176,7 +1234,15 @@ class LTXLocalAttention(LocalAttention): q = q.transpose(1, 2) k = k.transpose(1, 2) - output = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + if self.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: + if gate_compress is None: + raise ValueError( + "gate_compress must be provided when using VIDEO_SPARSE_ATTN" + ) + output = self.attn_impl.forward( # type: ignore[call-arg] + q, k, v, gate_compress, ctx_attn_metadata) + else: + output = self.attn_impl.forward(q, k, v, ctx_attn_metadata) return output @@ -1216,6 +1282,9 @@ class LTXSelfAttention(nn.Module): causal=False, supported_attention_backends=supported_attention_backends, ) + self.to_gate_compress: nn.Linear | None = None + if self.attn.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: + self.to_gate_compress = nn.Linear(context_dim, inner_dim, bias=True) self.attn_masked = LTXLocalAttention( num_heads=heads, head_size=dim_head, @@ -1238,6 +1307,8 @@ class LTXSelfAttention(nn.Module): context = x if context is None else context k = self.to_k(context) v = self.to_v(context) + gate_compress = (self.to_gate_compress(context) + if self.to_gate_compress is not None else None) q = self.q_norm(q) k = self.k_norm(k) @@ -1248,6 +1319,9 @@ class LTXSelfAttention(nn.Module): q = q.view(b, q_len, self.heads, self.dim_head) k = k.view(b, k_len, self.heads, self.dim_head) v = v.view(b, k_len, self.heads, self.dim_head) + if gate_compress is not None: + gate_compress = gate_compress.view(b, k_len, self.heads, + self.dim_head) if mask is not None: if mask.ndim == 2: @@ -1270,9 +1344,18 @@ class LTXSelfAttention(nn.Module): attn_metadata=attn_metadata, forward_batch=forward_batch, ): - out = self.attn_masked(q, k, v, ltx_freqs_cis=pe, ltx_k_freqs_cis=k_pe) + out = self.attn_masked(q, + k, + v, + ltx_freqs_cis=pe, + ltx_k_freqs_cis=k_pe) else: - out = self.attn(q, k, v, ltx_freqs_cis=pe, ltx_k_freqs_cis=k_pe) + out = self.attn(q, + k, + v, + gate_compress=gate_compress, + ltx_freqs_cis=pe, + ltx_k_freqs_cis=k_pe) out = out.reshape(b, q_len, -1) return self.to_out(out) @@ -1314,6 +1397,9 @@ class LTXDistributedSelfAttention(nn.Module): supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn", ) + self.to_gate_compress: nn.Linear | None = None + if self.attn.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: + self.to_gate_compress = nn.Linear(context_dim, inner_dim, bias=True) def forward( self, @@ -1337,6 +1423,8 @@ class LTXDistributedSelfAttention(nn.Module): context = x if context is None else context k = self.to_k(context) v = self.to_v(context) + gate_compress = (self.to_gate_compress(context) + if self.to_gate_compress is not None else None) q = self.q_norm(q) k = self.k_norm(k) @@ -1348,11 +1436,15 @@ class LTXDistributedSelfAttention(nn.Module): q = q.view(b, q_len, self.heads, self.dim_head) k = k.view(b, k_len, self.heads, self.dim_head) v = v.view(b, k_len, self.heads, self.dim_head) + if gate_compress is not None: + gate_compress = gate_compress.view(b, k_len, self.heads, + self.dim_head) # Pass full RoPE to distributed attention - it will apply after all-to-all # and slice to this rank's heads out, _ = self.attn( q, k, v, + gate_compress=gate_compress, attention_mask=attention_mask, ltx_freqs_cis=pe, ) @@ -1383,6 +1475,15 @@ class BasicAVTransformerBlock(torch.nn.Module): # Text cross-attention uses LocalAttention (text embeddings are replicated) SelfAttnCls = LTXDistributedSelfAttention if use_distributed_attention else LTXSelfAttention CrossAttnCls = LTXSelfAttention # Text cross-attention is always local + video_self_attn_backends = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.VIDEO_SPARSE_ATTN, + ) + dense_attn_backends = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) if video is not None: # Video self-attention - use distributed when SP > 1 @@ -1393,7 +1494,7 @@ class BasicAVTransformerBlock(torch.nn.Module): dim_head=video.d_head, norm_eps=norm_eps, rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + supported_attention_backends=video_self_attn_backends, prefix=f"{prefix}.blocks.{idx}.attn1" if use_distributed_attention else "", ) if use_distributed_attention else LTXSelfAttention( query_dim=video.dim, @@ -1402,7 +1503,7 @@ class BasicAVTransformerBlock(torch.nn.Module): dim_head=video.d_head, norm_eps=norm_eps, rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + supported_attention_backends=video_self_attn_backends, ) # Text cross-attention - always local (text is replicated) self.attn2 = CrossAttnCls( @@ -1412,7 +1513,7 @@ class BasicAVTransformerBlock(torch.nn.Module): dim_head=video.d_head, norm_eps=norm_eps, rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + supported_attention_backends=dense_attn_backends, ) self.ff = FeedForward(video.dim, dim_out=video.dim) self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim)) @@ -1426,7 +1527,7 @@ class BasicAVTransformerBlock(torch.nn.Module): dim_head=audio.d_head, norm_eps=norm_eps, rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + supported_attention_backends=dense_attn_backends, prefix=f"{prefix}.blocks.{idx}.audio_attn1" if use_distributed_attention else "", ) if use_distributed_attention else LTXSelfAttention( query_dim=audio.dim, @@ -1435,7 +1536,7 @@ class BasicAVTransformerBlock(torch.nn.Module): dim_head=audio.d_head, norm_eps=norm_eps, rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + supported_attention_backends=dense_attn_backends, ) # Text cross-attention - always local (text is replicated) self.audio_attn2 = CrossAttnCls( @@ -1445,7 +1546,7 @@ class BasicAVTransformerBlock(torch.nn.Module): dim_head=audio.d_head, norm_eps=norm_eps, rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + supported_attention_backends=dense_attn_backends, ) self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim) self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim)) @@ -1460,7 +1561,7 @@ class BasicAVTransformerBlock(torch.nn.Module): dim_head=audio.d_head, norm_eps=norm_eps, rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + supported_attention_backends=dense_attn_backends, ) # Video-to-audio cross-attention # Uses local attention - context is gathered from all SP ranks in forward() @@ -1471,7 +1572,7 @@ class BasicAVTransformerBlock(torch.nn.Module): dim_head=audio.d_head, norm_eps=norm_eps, rope_type=rope_type, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + supported_attention_backends=dense_attn_backends, ) self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim)) self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim)) @@ -2120,7 +2221,9 @@ class LTX2Transformer3DModel(CachableDiT): # Get SP world size for distributed attention sp_world_size = get_sp_world_size() - use_distributed_attention = sp_world_size > 1 + use_vsa_backend = os.getenv("FASTVIDEO_ATTENTION_BACKEND", + "") == "VIDEO_SPARSE_ATTN" + use_distributed_attention = sp_world_size > 1 or use_vsa_backend # Validate that attention heads are divisible by SP world size if sp_world_size > 1: @@ -2135,6 +2238,10 @@ class LTX2Transformer3DModel(CachableDiT): logger.info( f"LTX2 sequence parallelism enabled with SP world size {sp_world_size}" ) + elif use_vsa_backend: + logger.info( + "LTX2 VSA enabled with SP world size 1; using distributed " + "attention path for VSA compatibility") model_type = LTXModelType.AudioVideo self.model = LTXModel( diff --git a/fastvideo/models/encoders/gemma.py b/fastvideo/models/encoders/gemma.py index efc78847b..6741a3473 100644 --- a/fastvideo/models/encoders/gemma.py +++ b/fastvideo/models/encoders/gemma.py @@ -21,6 +21,7 @@ from fastvideo.models.dits.ltx2 import ( ) from fastvideo.models.loader.weight_utils import default_weight_loader from fastvideo.platforms import AttentionBackendEnum +from fastvideo.distributed import get_local_torch_device def _debug_log_line(message: str) -> None: @@ -516,15 +517,18 @@ class LTX2GemmaTextEncoderModel(TextEncoder): attention_mask = torch.ones_like(input_ids) model = self.gemma_model - input_ids = input_ids.to(device=model.device) - attention_mask = attention_mask.to(device=model.device) + orig_device = model.device + model.to(device=get_local_torch_device()) + # input_ids = input_ids.to(device=model.device) + # attention_mask = attention_mask.to(device=model.device) outputs = model( input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, return_dict=True, ) - + model.to(device=orig_device) + encoded_inputs = self._run_feature_extractor( outputs.hidden_states, attention_mask, diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 9afe4a1c0..ead93ee89 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -518,10 +518,25 @@ class TokenizerLoader(ComponentLoader): def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the tokenizer based on the model path, and inference args.""" logger.info("Loading tokenizer from %s", model_path) + resolved_model_path = model_path + + # LTX2 checkpoints may not ship a top-level tokenizer/ directory. + # In that case, tokenizer assets live under text_encoder/gemma/. + if not os.path.isdir(resolved_model_path): + ltx2_gemma_path = os.path.normpath( + os.path.join(resolved_model_path, "..", "text_encoder", + "gemma")) + if os.path.isdir(ltx2_gemma_path): + logger.info( + "Tokenizer directory %s missing; falling back to %s", + resolved_model_path, + ltx2_gemma_path, + ) + resolved_model_path = ltx2_gemma_path # Cosmos2.5 stores an AutoProcessor config in `tokenizer/config.json` (not a tokenizer # config). Use its `_name_or_path` (e.g. Qwen/Qwen2.5-VL-7B-Instruct) as the source. - tokenizer_cfg_path = os.path.join(model_path, "config.json") + tokenizer_cfg_path = os.path.join(resolved_model_path, "config.json") if os.path.exists(tokenizer_cfg_path): try: with open(tokenizer_cfg_path, "r") as f: @@ -547,10 +562,11 @@ class TokenizerLoader(ComponentLoader): pass tokenizer = AutoTokenizer.from_pretrained( - model_path, # "/tokenizer" + resolved_model_path, # "/tokenizer" # in v0, this was same string as encoder_name "ClipTextModel" # TODO(will): pass these tokenizer kwargs from inference args? Maybe # other method of config? + local_files_only=os.path.isdir(resolved_model_path), ) padding_side = None if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"): diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index ca92f1343..a274034ea 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -374,8 +374,21 @@ class ComposedPipelineBase(ABC): logger.info("Loading required modules: %s", required_modules) modules = {} - for module_name, (transformers_or_diffusers, - architecture) in model_index.items(): + for module_name, module_spec in model_index.items(): + if not isinstance(module_spec, list | tuple): + logger.info( + "Skipping non-module config entry %s=%s", + module_name, + module_spec, + ) + continue + if len(module_spec) < 1: + logger.warning( + "Skipping module %s due to invalid empty spec in model_index.json", + module_name, + ) + continue + transformers_or_diffusers = module_spec[0] if transformers_or_diffusers is None: logger.warning( "Module %s in model_index.json has null value, removing from required_config_modules", diff --git a/fastvideo/pipelines/preprocess/wan/__init__.py b/fastvideo/pipelines/preprocess/ltx2/__init__.py similarity index 100% rename from fastvideo/pipelines/preprocess/wan/__init__.py rename to fastvideo/pipelines/preprocess/ltx2/__init__.py diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/pipelines/stages/base.py index 97d818044..22891e7bb 100644 --- a/fastvideo/pipelines/stages/base.py +++ b/fastvideo/pipelines/stages/base.py @@ -148,16 +148,19 @@ class PipelineStage(ABC): # Execute the actual stage logic if envs.FASTVIDEO_STAGE_LOGGING: logger.info("[%s] Starting execution", stage_name) + torch.cuda.synchronize() start_time = time.perf_counter() try: result = self.forward(batch, fastvideo_args) + torch.cuda.synchronize() execution_time = time.perf_counter() - start_time logger.info("[%s] Execution completed in %s ms", stage_name, execution_time * 1000) batch.logging_info.add_stage_execution_time( stage_name, execution_time) except Exception as e: + torch.cuda.synchronize() execution_time = time.perf_counter() - start_time logger.error("[%s] Error during execution after %s ms: %s", stage_name, execution_time * 1000, e) diff --git a/fastvideo/pipelines/stages/ltx2_denoising.py b/fastvideo/pipelines/stages/ltx2_denoising.py index 707d8ae82..6ac573f49 100644 --- a/fastvideo/pipelines/stages/ltx2_denoising.py +++ b/fastvideo/pipelines/stages/ltx2_denoising.py @@ -11,6 +11,9 @@ import os import torch from tqdm.auto import tqdm +import fastvideo.envs as envs +from fastvideo.attention.backends.video_sparse_attn import ( + VideoSparseAttentionMetadataBuilder) from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context from fastvideo.pipelines.pipeline_batch_info import ForwardBatch @@ -23,6 +26,7 @@ from fastvideo.models.dits.ltx2 import ( DEFAULT_LTX2_AUDIO_DOWNSAMPLE, DEFAULT_LTX2_AUDIO_HOP_LENGTH, DEFAULT_LTX2_AUDIO_MEL_BINS, DEFAULT_LTX2_AUDIO_SAMPLE_RATE, VideoLatentShape) +from fastvideo.utils import is_vsa_available BASE_SHIFT_ANCHOR = 1024 MAX_SHIFT_ANCHOR = 4096 @@ -35,6 +39,11 @@ DISTILLED_SIGMA_VALUES = [ logger = init_logger(__name__) +try: + vsa_available = is_vsa_available() +except ImportError: + vsa_available = False + def _ltx2_sigmas( steps: int, @@ -234,6 +243,10 @@ class LTX2DenoisingStage(PipelineStage): tuple(sigmas.shape), tuple(latents.shape), ) + use_vsa = (vsa_available + and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN") + vsa_metadata_builder = (VideoSparseAttentionMetadataBuilder() + if use_vsa else None) for step_index in tqdm(range(len(sigmas) - 1)): sigma = sigmas[step_index] @@ -241,6 +254,16 @@ class LTX2DenoisingStage(PipelineStage): timestep = timestep_template * sigma audio_timestep = (audio_timestep_template * sigma if audio_timestep_template is not None else None) + attn_metadata = None + if vsa_metadata_builder is not None: + attn_metadata = vsa_metadata_builder.build( + current_timestep=step_index, + raw_latent_shape=latents.shape[2:5], + patch_size=fastvideo_args.pipeline_config.dit_config. + patch_size, + VSA_sparsity=fastvideo_args.VSA_sparsity, + device=latents.device, + ) with torch.autocast( device_type="cuda", @@ -248,7 +271,7 @@ class LTX2DenoisingStage(PipelineStage): enabled=autocast_enabled, ), set_forward_context( current_timestep=sigma.item(), - attn_metadata=None, + attn_metadata=attn_metadata, forward_batch=batch, ): # Pass 1: Full conditioning (text + cross-modal) diff --git a/fastvideo/training/activation_checkpoint.py b/fastvideo/training/activation_checkpoint.py index 874d013f7..00ceb9603 100644 --- a/fastvideo/training/activation_checkpoint.py +++ b/fastvideo/training/activation_checkpoint.py @@ -51,6 +51,7 @@ def apply_activation_checkpointing( def _apply_activation_checkpointing_blocks(module: torch.nn.Module, n_layer: int | None = None ) -> torch.nn.Module: + applied = False for transformer_block_name in TRANSFORMER_BLOCK_NAMES: blocks: torch.nn.Module = getattr(module, transformer_block_name, None) if blocks is None: @@ -59,6 +60,9 @@ def _apply_activation_checkpointing_blocks(module: torch.nn.Module, if n_layer is None or index % n_layer == 0: block = checkpoint_wrapper(block, preserve_rng_state=False) blocks.register_module(layer_id, block) + applied = True + if not applied: + raise ValueError("Activation checkpointing is not applied successfully") return module diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index e65b7cef2..8abcfa955 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -75,6 +75,30 @@ class DistillationPipeline(TrainingPipeline): video_latent_shape_sp: tuple[int, ...] train_fake_score_transformer_2: bool = False + @staticmethod + def _clone_batch_value(value: Any) -> Any: + """Clone values in a TrainingBatch without tensor deepcopy.""" + if isinstance(value, torch.Tensor): + # Avoid torch tensor deepcopy limitations on non-leaf tensors. + return value.detach().clone() + if isinstance(value, dict): + return { + k: DistillationPipeline._clone_batch_value(v) + for k, v in value.items() + } + if isinstance(value, list): + return [DistillationPipeline._clone_batch_value(v) for v in value] + if isinstance(value, tuple): + return tuple( + DistillationPipeline._clone_batch_value(v) for v in value) + return copy.deepcopy(value) + + def _clone_training_batch(self, batch: TrainingBatch) -> TrainingBatch: + cloned_batch = TrainingBatch() + for key, value in batch.__dict__.items(): + setattr(cloned_batch, key, self._clone_batch_value(value)) + return cloned_batch + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): raise RuntimeError( "create_pipeline_stages should not be called for training pipeline") @@ -971,13 +995,13 @@ class DistillationPipeline(TrainingPipeline): batch.attn_metadata.VSA_sparsity = 0.0 # type: ignore batches.append(batch) - self.optimizer.zero_grad() + self.optimizer.zero_grad(set_to_none=True) total_dmd_loss = 0.0 dmd_latent_vis_dict = {} fake_score_latent_vis_dict = {} if (self.current_trainstep % self.generator_update_interval == 0): for batch in batches: - batch_gen = copy.deepcopy(batch) + batch_gen = self._clone_training_batch(batch) with set_forward_context( current_timestep=batch_gen.timesteps, @@ -1022,12 +1046,12 @@ class DistillationPipeline(TrainingPipeline): else: training_batch.generator_loss = 0.0 - self.fake_score_optimizer.zero_grad() + self.fake_score_optimizer.zero_grad(set_to_none=True) if self.fake_score_transformer_2 is not None: - self.fake_score_optimizer_2.zero_grad() + self.fake_score_optimizer_2.zero_grad(set_to_none=True) total_fake_score_loss = 0.0 for batch in batches: - batch_fake = copy.deepcopy(batch) + batch_fake = self._clone_training_batch(batch) batch_fake, fake_score_loss = self.faker_score_forward(batch_fake) with set_forward_context(current_timestep=batch_fake.timesteps, attn_metadata=batch_fake.attn_metadata): @@ -1236,9 +1260,12 @@ class DistillationPipeline(TrainingPipeline): # Helper function to run validation with optional EMA contexts def run_validation_with_ema( - steps: int) -> tuple[list[np.ndarray], list[str]]: + steps: int + ) -> tuple[list[np.ndarray], list[str], list[Any], list[Any]]: videos: list[np.ndarray] = [] captions: list[str] = [] + audios: list[Any] = [] + audio_sample_rates: list[Any] = [] for validation_batch in validation_dataloader: batch = self._prepare_validation_batch( sampling_param, training_args, validation_batch, steps) @@ -1282,25 +1309,32 @@ class DistillationPipeline(TrainingPipeline): x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) frames.append((x * 255).numpy().astype(np.uint8)) videos.append(frames) + audios.append(output_batch.extra.get("audio")) + audio_sample_rates.append( + output_batch.extra.get("audio_sample_rate")) - return videos, captions + return videos, captions, audios, audio_sample_rates # Apply EMA contexts if available (nested context managers) if ema_context is not None and ema_2_context is not None: with ema_context, ema_2_context: - step_videos, step_captions = run_validation_with_ema( - num_inference_steps) + (step_videos, step_captions, step_audios, + step_audio_sample_rates + ) = run_validation_with_ema(num_inference_steps) elif ema_context is not None: with ema_context: - step_videos, step_captions = run_validation_with_ema( - num_inference_steps) + (step_videos, step_captions, step_audios, + step_audio_sample_rates + ) = run_validation_with_ema(num_inference_steps) elif ema_2_context is not None: with ema_2_context: - step_videos, step_captions = run_validation_with_ema( - num_inference_steps) + (step_videos, step_captions, step_audios, + step_audio_sample_rates + ) = run_validation_with_ema(num_inference_steps) else: - step_videos, step_captions = run_validation_with_ema( - num_inference_steps) + (step_videos, step_captions, step_audios, + step_audio_sample_rates + ) = run_validation_with_ema(num_inference_steps) # Log validation results for this step world_group = get_world_group() @@ -1313,24 +1347,42 @@ class DistillationPipeline(TrainingPipeline): # Global rank 0 collects results from all sp_group leaders all_videos = step_videos # Start with own results all_captions = step_captions + all_audios = step_audios + all_audio_sample_rates = step_audio_sample_rates # Receive from other sp_group leaders for sp_group_idx in range(1, num_sp_groups): src_rank = sp_group_idx * self.sp_world_size # Global rank of other sp_group leaders recv_videos = world_group.recv_object(src=src_rank) recv_captions = world_group.recv_object(src=src_rank) + recv_audios = world_group.recv_object(src=src_rank) + recv_audio_sample_rates = world_group.recv_object( + src=src_rank) all_videos.extend(recv_videos) all_captions.extend(recv_captions) + all_audios.extend(recv_audios) + all_audio_sample_rates.extend(recv_audio_sample_rates) video_filenames = [] - for i, (video, caption) in enumerate( - zip(all_videos, all_captions, strict=True)): + for i, (video, caption, audio, + audio_sample_rate) in enumerate( + zip(all_videos, + all_captions, + all_audios, + all_audio_sample_rates, + strict=True)): os.makedirs(training_args.output_dir, exist_ok=True) filename = os.path.join( training_args.output_dir, f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4" ) imageio.mimsave(filename, video, fps=sampling_param.fps) + if (audio is not None and audio_sample_rate is not None + and not self._mux_audio(filename, audio, + audio_sample_rate)): + logger.warning( + "Audio mux failed for validation video %s; saved video without audio.", + filename) video_filenames.append(filename) artifacts = [] @@ -1351,16 +1403,38 @@ class DistillationPipeline(TrainingPipeline): # Other sp_group leaders send their results to global rank 0 world_group.send_object(step_videos, dst=0) world_group.send_object(step_captions, dst=0) + world_group.send_object(step_audios, dst=0) + world_group.send_object(step_audio_sample_rates, dst=0) # Re-enable gradients for training - set both transformers back to train mode transformer.train() if hasattr(self, 'transformer_2') and self.transformer_2 is not None: self.transformer_2.train() gc.collect() + torch.cuda.empty_cache() def visualize_intermediate_latents(self, training_batch: TrainingBatch, training_args: TrainingArgs, step: int): """Add visualization data to tracker logging and save frames to disk.""" + + def _prepare_vae_latents(latents: torch.Tensor) -> torch.Tensor: + # Most training paths store latents as [B, C, T, H, W]. Older + # visualization code permuted to [B, T, C, H, W] for WAN-like VAEs. + # LTX2 VAE expects [B, C, T, H, W], so keep native layout there. + if hasattr(self.vae, "TIME_SCALE") and hasattr( + self.vae, "SPATIAL_SCALE"): + return latents + return latents.permute(0, 2, 1, 3, 4) + + def _apply_vae_scale(latents: torch.Tensor) -> torch.Tensor: + scaling_factor = getattr(self.vae, "scaling_factor", None) + if scaling_factor is None: + return latents + if isinstance(scaling_factor, torch.Tensor): + return latents / scaling_factor.to(latents.device, + latents.dtype) + return latents / scaling_factor + tracker_loss_dict: dict[str, Any] = {} dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict @@ -1369,13 +1443,8 @@ class DistillationPipeline(TrainingPipeline): for latent_key in fake_score_log_keys: latents = fake_score_latents_vis_dict[latent_key] - latents = latents.permute(0, 2, 1, 3, 4) - - if isinstance(self.vae.scaling_factor, torch.Tensor): - latents = latents / self.vae.scaling_factor.to( - latents.device, latents.dtype) - else: - latents = latents / self.vae.scaling_factor + latents = _prepare_vae_latents(latents) + latents = _apply_vae_scale(latents) # Apply shifting if needed if (hasattr(self.vae, "shift_factor") @@ -1402,13 +1471,9 @@ class DistillationPipeline(TrainingPipeline): if 'generator_pred_video' in dmd_latents_vis_dict: for latent_key in dmd_log_keys: latents = dmd_latents_vis_dict[latent_key] - latents = latents.permute(0, 2, 1, 3, 4) + latents = _prepare_vae_latents(latents) # decoded_latent = decode_stage(ForwardBatch(data_type="video", latents=latents), training_args) - if isinstance(self.vae.scaling_factor, torch.Tensor): - latents = latents / self.vae.scaling_factor.to( - latents.device, latents.dtype) - else: - latents = latents / self.vae.scaling_factor + latents = _apply_vae_scale(latents) # Apply shifting if needed if (hasattr(self.vae, "shift_factor") @@ -1451,13 +1516,14 @@ class DistillationPipeline(TrainingPipeline): set_random_seed(seed + self.global_rank) # Set random seeds for deterministic training - self.noise_random_generator = torch.Generator(device="cpu").manual_seed( - self.seed) - self.noise_gen_cuda = torch.Generator(device="cuda").manual_seed( - self.seed) + self.noise_random_generator = torch.Generator( + device="cpu").manual_seed(self.seed + self.global_rank) + self.noise_gen_cuda = torch.Generator( + device="cuda").manual_seed(self.seed + self.global_rank) self.validation_random_generator = torch.Generator( - device="cpu").manual_seed(self.seed) - logger.info("Initialized random seeds with seed: %s", seed) + device="cpu").manual_seed(self.seed + self.global_rank) + logger.info("Initialized random seeds with seed: %s", + seed + self.global_rank) # Initialize current_trainstep for EMA ready checks #TODO: check if needed @@ -1488,6 +1554,9 @@ class DistillationPipeline(TrainingPipeline): use_vsa = vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" for step in range(self.init_steps + 1, self.training_args.max_train_steps + 1): + if step % 5 == 0: + gc.collect() + torch.cuda.empty_cache() start_time = time.perf_counter() if use_vsa: vsa_sparsity = self.training_args.VSA_sparsity diff --git a/fastvideo/training/ltx2_training_pipeline.py b/fastvideo/training/ltx2_training_pipeline.py index 3d77aee8a..e79f9a0c5 100644 --- a/fastvideo/training/ltx2_training_pipeline.py +++ b/fastvideo/training/ltx2_training_pipeline.py @@ -6,7 +6,9 @@ from pathlib import Path import torch import torch.distributed as dist -from fastvideo.dataset import build_ltx2_precomputed_dataloader +from fastvideo.dataset import (build_ltx2_precomputed_dataloader, + build_parquet_map_style_dataloader) +from fastvideo.dataset.dataloader.schema import pyarrow_schema_text_only from fastvideo.distributed import get_local_torch_device, get_sp_group, get_world_group from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.forward_context import set_forward_context @@ -64,14 +66,14 @@ class LTX2TrainingPipeline(TrainingPipeline): self.seed = training_args.seed assert self.seed is not None, "seed must be set" - set_random_seed(self.seed) + set_random_seed(self.seed + self.global_rank) self.transformer.train() if training_args.enable_gradient_checkpointing_type is not None: from fastvideo.training.activation_checkpoint import ( apply_activation_checkpointing) - self.transformer = apply_activation_checkpointing( - self.transformer, + self.transformer.model = apply_activation_checkpointing( + self.transformer.model, checkpointing_type=training_args. enable_gradient_checkpointing_type) @@ -103,17 +105,37 @@ class LTX2TrainingPipeline(TrainingPipeline): last_epoch=self.init_steps - 1, ) - data_sources = self._get_ltx2_data_sources(training_args.data_path) - self.with_audio = "audio_latents" in data_sources - self.train_dataset, self.train_dataloader = ( - build_ltx2_precomputed_dataloader( - training_args.data_path, - training_args.train_batch_size, - num_data_workers=training_args.dataloader_num_workers, - data_sources=data_sources, - drop_last=True, - seed=self.seed, - )) + if self._has_precomputed_pt_data(training_args.data_path): + data_sources = self._get_ltx2_data_sources(training_args.data_path) + self.with_audio = "audio_latents" in data_sources + self.train_dataset, self.train_dataloader = ( + build_ltx2_precomputed_dataloader( + training_args.data_path, + training_args.train_batch_size, + num_data_workers=training_args.dataloader_num_workers, + data_sources=data_sources, + drop_last=True, + seed=self.seed, + )) + self._use_parquet_data = False + logger.info("Using precomputed .pt data") + else: + text_padding_length = (training_args.pipeline_config. + text_encoder_configs[0].arch_config.text_len) + self.with_audio = False + self.train_dataset, self.train_dataloader = ( + build_parquet_map_style_dataloader( + training_args.data_path, + training_args.train_batch_size, + num_data_workers=training_args.dataloader_num_workers, + parquet_schema=pyarrow_schema_text_only, + cfg_rate=training_args.training_cfg_rate, + drop_last=True, + text_padding_length=text_padding_length, + seed=self.seed, + )) + self._use_parquet_data = True + logger.info("No precomputed .pt data found, using parquet data") self.num_update_steps_per_epoch = max( 1, @@ -168,6 +190,17 @@ class LTX2TrainingPipeline(TrainingPipeline): ) self.validation_pipeline = validation_pipeline + def _has_precomputed_pt_data(self, data_path: str) -> bool: + """Check whether precomputed .pt data exists for LTX2.""" + data_root = Path(data_path).expanduser().resolve() + if (data_root / ".precomputed").exists(): + data_root = data_root / ".precomputed" + latents_dir = data_root / "latents" + conditions_dir = data_root / "conditions" + return (latents_dir.exists() and conditions_dir.exists() + and any(latents_dir.rglob("*.pt")) + and any(conditions_dir.rglob("*.pt"))) + def _get_ltx2_data_sources(self, data_path: str) -> dict[str, str]: data_root = Path(data_path).expanduser().resolve() if (data_root / ".precomputed").exists(): @@ -187,52 +220,104 @@ class LTX2TrainingPipeline(TrainingPipeline): self.train_loader_iter = iter(self.train_dataloader) batch = next(self.train_loader_iter) - latents = batch["latents"]["latents"].to(get_local_torch_device(), - dtype=torch.bfloat16) - latents = latents[:, :, :self.training_args.num_latent_t] - conditions = batch["conditions"] - if ("video_prompt_embeds" in conditions - and "audio_prompt_embeds" in conditions - and "prompt_attention_mask" in conditions): - video_embeds = conditions["video_prompt_embeds"].to( - get_local_torch_device()) - audio_embeds = conditions["audio_prompt_embeds"].to( - get_local_torch_device()) - attention_mask = conditions["prompt_attention_mask"].to( - get_local_torch_device(), dtype=torch.int64) + if self._use_parquet_data: + training_batch = self._get_next_batch_parquet( + batch, training_batch) else: - prompt_embeds = conditions["prompt_embeds"].to( - get_local_torch_device()) - prompt_attention_mask = conditions["prompt_attention_mask"].to( - get_local_torch_device(), dtype=torch.int64) + training_batch = self._get_next_batch_pt(batch, training_batch) - video_embeds, audio_embeds, attention_mask = ( - self.text_encoder.run_connectors(prompt_embeds, - prompt_attention_mask)) + return training_batch - training_batch.latents = latents.to(get_local_torch_device(), - dtype=torch.bfloat16) - training_batch.encoder_hidden_states = video_embeds.to( - get_local_torch_device(), dtype=torch.bfloat16) - training_batch.encoder_attention_mask = attention_mask.to( + def _get_next_batch_parquet( + self, + batch: dict, + training_batch: TrainingBatch, + ) -> TrainingBatch: + """Extract training data from a parquet-format batch.""" + device = get_local_torch_device() + prompt_embeds = batch["text_embedding"].to(device) + prompt_attention_mask = batch["text_attention_mask"].to( + device, dtype=torch.int64) + + video_embeds, audio_embeds, attention_mask = ( + self.text_encoder.run_connectors(prompt_embeds, + prompt_attention_mask)) + + # Parquet text-only data has no latents; create zeros. + batch_size = video_embeds.shape[0] + vae_cfg = (self.training_args.pipeline_config.vae_config.arch_config) + num_channels = vae_cfg.z_dim + scr = vae_cfg.spatial_compression_ratio + latent_h = self.training_args.num_height // scr + latent_w = self.training_args.num_width // scr + latents = torch.zeros(batch_size, + num_channels, + self.training_args.num_latent_t, + latent_h, + latent_w, + device=device, + dtype=torch.bfloat16) + + training_batch.latents = latents + training_batch.encoder_hidden_states = video_embeds.to( + device, dtype=torch.bfloat16) + training_batch.encoder_attention_mask = attention_mask.to( + device, dtype=torch.bfloat16) + training_batch.infos = [] + training_batch.raw_latent_shape = latents.shape + return training_batch + + def _get_next_batch_pt( + self, + batch: dict, + training_batch: TrainingBatch, + ) -> TrainingBatch: + """Extract training data from a precomputed .pt batch.""" + latents = batch["latents"]["latents"].to(get_local_torch_device(), + dtype=torch.bfloat16) + latents = latents[:, :, :self.training_args.num_latent_t] + conditions = batch["conditions"] + if ("video_prompt_embeds" in conditions + and "audio_prompt_embeds" in conditions + and "prompt_attention_mask" in conditions): + video_embeds = conditions["video_prompt_embeds"].to( + get_local_torch_device()) + audio_embeds = conditions["audio_prompt_embeds"].to( + get_local_torch_device()) + attention_mask = conditions["prompt_attention_mask"].to( + get_local_torch_device(), dtype=torch.int64) + else: + prompt_embeds = conditions["prompt_embeds"].to( + get_local_torch_device()) + prompt_attention_mask = conditions["prompt_attention_mask"].to( + get_local_torch_device(), dtype=torch.int64) + + video_embeds, audio_embeds, attention_mask = ( + self.text_encoder.run_connectors(prompt_embeds, + prompt_attention_mask)) + + training_batch.latents = latents.to(get_local_torch_device(), + dtype=torch.bfloat16) + training_batch.encoder_hidden_states = video_embeds.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.encoder_attention_mask = attention_mask.to( + get_local_torch_device(), dtype=torch.bfloat16) + + if self.with_audio and "audio_latents" in batch: + audio_latents = batch["audio_latents"]["latents"].to( get_local_torch_device(), dtype=torch.bfloat16) + training_batch.audio_latents = audio_latents + training_batch.audio_encoder_hidden_states = (audio_embeds.to( + get_local_torch_device(), dtype=torch.bfloat16)) + training_batch.audio_encoder_attention_mask = (attention_mask.to( + get_local_torch_device())) - if self.with_audio and "audio_latents" in batch: - audio_latents = batch["audio_latents"]["latents"].to( - get_local_torch_device(), dtype=torch.bfloat16) - training_batch.audio_latents = audio_latents - training_batch.audio_encoder_hidden_states = audio_embeds.to( - get_local_torch_device(), dtype=torch.bfloat16) - training_batch.audio_encoder_attention_mask = attention_mask.to( - get_local_torch_device()) - - idxs = batch.get("idx") - if idxs is not None and torch.is_tensor(idxs): - training_batch.infos = [{"idx": int(i)} for i in idxs.tolist()] - else: - training_batch.infos = [] - training_batch.raw_latent_shape = latents.shape - + idxs = batch.get("idx") + if idxs is not None and torch.is_tensor(idxs): + training_batch.infos = [{"idx": int(i)} for i in idxs.tolist()] + else: + training_batch.infos = [] + training_batch.raw_latent_shape = latents.shape return training_batch def _normalize_dit_input(self, @@ -339,8 +424,7 @@ class LTX2TrainingPipeline(TrainingPipeline): def _build_attention_metadata( self, training_batch: TrainingBatch) -> TrainingBatch: - training_batch.attn_metadata = None - return training_batch + return super()._build_attention_metadata(training_batch) def _build_input_kwargs(self, training_batch: TrainingBatch) -> TrainingBatch: diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index ea3856f7a..d9aa6e227 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -856,6 +856,9 @@ class TrainingPipeline(LoRAPipeline, ABC): step_videos: list[np.ndarray] = [] step_captions: list[str] = [] + step_audio: list[np.ndarray | None] = [] + step_sample_rates: list[int | None] = [] + for validation_batch in validation_dataloader: batch = self._prepare_validation_batch(sampling_param, training_args, @@ -876,6 +879,16 @@ class TrainingPipeline(LoRAPipeline, ABC): batch, training_args) samples = output_batch.output.cpu() + # Capture audio if available + audio = output_batch.extra.get("audio") + sample_rate = output_batch.extra.get("audio_sample_rate") + + if audio is not None and torch.is_tensor(audio): + audio = audio.detach().cpu().float().numpy() + + step_audio.append(audio) + step_sample_rates.append(sample_rate) + if self.rank_in_sp_group != 0: continue @@ -894,18 +907,29 @@ class TrainingPipeline(LoRAPipeline, ABC): # Global rank 0 collects results from all sp_group leaders all_videos = step_videos # Start with own results all_captions = step_captions + all_audios = step_audio + all_sample_rates = step_sample_rates # Receive from other sp_group leaders for sp_group_idx in range(1, num_sp_groups): src_rank = sp_group_idx * self.sp_world_size # Global rank of other sp_group leaders recv_videos = world_group.recv_object(src=src_rank) recv_captions = world_group.recv_object(src=src_rank) + recv_audios = world_group.recv_object(src=src_rank) + recv_sample_rates = world_group.recv_object(src=src_rank) + all_videos.extend(recv_videos) all_captions.extend(recv_captions) + all_audios.extend(recv_audios) + all_sample_rates.extend(recv_sample_rates) video_filenames = [] - for i, (video, caption) in enumerate( - zip(all_videos, all_captions, strict=True)): + for i, (video, caption, audio, sample_rate) in enumerate( + zip(all_videos, + all_captions, + all_audios, + all_sample_rates, + strict=True)): os.makedirs(training_args.output_dir, exist_ok=True) filename = os.path.join( training_args.output_dir, @@ -913,14 +937,11 @@ class TrainingPipeline(LoRAPipeline, ABC): ) imageio.mimsave(filename, video, fps=sampling_param.fps) # Mux audio if available - audio = output_batch.extra.get("audio") - audio_sample_rate = output_batch.extra.get( - "audio_sample_rate") - if (audio is not None and audio_sample_rate is not None + if (audio is not None and sample_rate is not None and not self._mux_audio( filename, audio, - audio_sample_rate, + sample_rate, )): logger.warning( "Audio mux failed for validation video %s; saved video without audio.", @@ -945,6 +966,8 @@ class TrainingPipeline(LoRAPipeline, ABC): # Other sp_group leaders send their results to global rank 0 world_group.send_object(step_videos, dst=0) world_group.send_object(step_captions, dst=0) + world_group.send_object(step_audio, dst=0) + world_group.send_object(step_sample_rates, dst=0) # Re-enable gradients for training training_args.inference_mode = False diff --git a/pyproject.toml b/pyproject.toml index 3e82dd061..932a8ebc8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,7 +76,7 @@ dependencies = [ "av", # Preprocessing Dependencies - "torchcodec==0.5.0", + "torchcodec", "ray>=2.49.1", "ftfy==6.3.1", ] diff --git a/pyproject_other.toml b/pyproject_other.toml index b48c48778..77390040b 100644 --- a/pyproject_other.toml +++ b/pyproject_other.toml @@ -75,7 +75,7 @@ dependencies = [ "av", # Preprocessing Dependencies - "torchcodec==0.5.0", + "torchcodec", "ray>=2.49.1", ]