diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/README.md b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/README.md new file mode 100644 index 000000000..4ec75167d --- /dev/null +++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/README.md @@ -0,0 +1,10 @@ +This directory contain e2e examples scripts for finetuning Wan2.1 I2V. + +Execute the following commands from `FastVideo/` to run training: + +- Download crush-smol dataset: +`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh` +- Preprocess the videos and captions into latents: +`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh` +- Edit the following file and run finetuning: +`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh` \ No newline at end of file diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/download_dataset.sh b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/download_dataset.sh new file mode 100644 index 000000000..f2942b7da --- /dev/null +++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/download_dataset.sh @@ -0,0 +1,3 @@ +#!/bin/bash + +python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset" \ No newline at end of file diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh new file mode 100644 index 000000000..18fe7915b --- /dev/null +++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh @@ -0,0 +1,93 @@ +#!/bin/bash + +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers" +DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/" +VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" +NUM_GPUS=8 +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_i2v_finetune" + --output_dir "$DATA_DIR/outputs/wan_i2v_finetune" + --max_train_steps 2000 + --train_batch_size 4 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 8 + --num_height 480 + --num_width 832 + --num_frames 77 +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 4 + --tp_size 4 + --hsdp_replicate_dim 2 + --hsdp_shard_dim 4 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 1 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 100 + --validation_sampling_steps "40" + --validation_guidance_scale "1.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 2e-5 + --mixed_precision "bf16" + --checkpointing_steps 2000 + --weight_decay 1e-4 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --allow_tf32 + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + --enable_gradient_checkpointing_type "full" +) + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +torchrun \ + --nnodes 1 \ + --nproc_per_node $NUM_GPUS \ + fastvideo/v1/training/wan_i2v_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm new file mode 100644 index 000000000..61eb25c8a --- /dev/null +++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm @@ -0,0 +1,131 @@ +#!/bin/bash +#SBATCH --job-name=i2v +#SBATCH --partition=main +#SBATCH --qos=hao +#SBATCH --nodes=4 +#SBATCH --ntasks=4 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --nodelist=fs-mbz-gpu-[100-850] +#SBATCH --mem=1440G +#SBATCH --output=i2v_output/i2v_%j.out +#SBATCH --error=i2v_output/i2v_%j.err +#SBATCH --exclusive +set -e -x + +# Environment Setup +source ~/conda/miniconda/bin/activate +conda activate will-fv + +# Basic Info +export WANDB_MODE="online" +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +# different cache dir for different processes +export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +echo "MASTER_ADDR: $MASTER_ADDR" +echo "NODE_RANK: $NODE_RANK" + +# Configs +NUM_GPUS=8 + +MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers" +DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/" +VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t + +# Training arguments +training_args=( + --tracker_project_name wan_i2v_finetune + --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n" + --max_train_steps=2000 + --train_batch_size=2 + --train_sp_batch_size 1 + --gradient_accumulation_steps=1 + --num_latent_t 8 + --num_height 480 + --num_width 832 + --num_frames 77 +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size $NUM_GPUS + --tp_size $NUM_GPUS + --hsdp_replicate_dim $SLURM_JOB_NUM_NODES + --hsdp_shard_dim $NUM_GPUS +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 10 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 100 + --validation_sampling_steps "40" + --validation_guidance_scale "1.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate=1e-5 + --mixed_precision="bf16" + --checkpointing_steps=1000 + --weight_decay 1e-4 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --allow_tf32 + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + --enable_gradient_checkpointing_type "full" +) + +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node $NUM_GPUS \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/v1/training/wan_i2v_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ No newline at end of file diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/preprocess_wan_data_i2v.sh b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/preprocess_wan_data_i2v.sh new file mode 100644 index 000000000..9ac91a2c3 --- /dev/null +++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/preprocess_wan_data_i2v.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +GPU_NUM=1 # 2,4,8 +MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers" +MODEL_TYPE="wan" +DATA_MERGE_PATH="data/crush-smol/merge.txt" +OUTPUT_DIR="data/crush-smol_processed_i2v_1_3b_inp/" +VALIDATION_PATH="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" + +torchrun --nproc_per_node=$GPU_NUM \ + fastvideo/v1/pipelines/preprocess/v1_preprocess.py \ + --model_path $MODEL_PATH \ + --data_merge_path $DATA_MERGE_PATH \ + --preprocess_video_batch_size 8 \ + --max_height 480 \ + --max_width 832 \ + --num_frames 77 \ + --dataloader_num_workers 0 \ + --output_dir=$OUTPUT_DIR \ + --model_type $MODEL_TYPE \ + --train_fps 16 \ + --validation_dataset_file $VALIDATION_PATH \ + --samples_per_file 8 \ + --flush_frequency 8 \ + --video_length_tolerance_range 5 \ + --preprocess_task "i2v" \ No newline at end of file diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json new file mode 100644 index 000000000..592166e98 --- /dev/null +++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json @@ -0,0 +1,31 @@ +{ + "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": "validation_dataset/yYcK4nANZz4-Scene-034.mp4", + "num_inference_steps": 40, + "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": "validation_dataset/yYcK4nANZz4-Scene-027.mp4", + "num_inference_steps": 40, + "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": "validation_dataset/yYcK4nANZz4-Scene-030.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 new file mode 100644 index 000000000..8e416c6b0 Binary files /dev/null and b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 differ diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 new file mode 100644 index 000000000..c2c781b51 Binary files /dev/null and b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 differ diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 new file mode 100644 index 000000000..abb42a2a8 Binary files /dev/null and b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 differ diff --git a/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh b/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh index 1281089e5..299c132d4 100644 --- a/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh +++ b/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh @@ -2,11 +2,12 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/" -VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/" +VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json" NUM_GPUS=8 # export CUDA_VISIBLE_DEVICES=4,5 # IP=[MASTER NODE IP] @@ -49,7 +50,7 @@ dataset_args=( # Validation arguments validation_args=( --log_validation - --validation_preprocessed_path "$VALIDATION_DIR" + --validation_dataset_file "$VALIDATION_DATASET_FILE" --validation_steps 100 --validation_sampling_steps "40" --validation_guidance_scale "1.0" @@ -75,6 +76,7 @@ miscellaneous_args=( --dit_precision "fp32" --num_euler_timesteps 50 --ema_start_step 0 + --enable_gradient_checkpointing_type "full" ) # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t diff --git a/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm b/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm index 63da7ec91..7cce5925f 100644 --- a/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm +++ b/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm @@ -42,7 +42,7 @@ NUM_GPUS=8 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/" -VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/" +VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json" # export CUDA_VISIBLE_DEVICES=4,5 # IP=[MASTER NODE IP] @@ -86,7 +86,7 @@ dataset_args=( # Validation arguments validation_args=( --log_validation - --validation_preprocessed_path "$VALIDATION_DIR" + --validation_dataset_file "$VALIDATION_DATASET_FILE" --validation_steps 100 --validation_sampling_steps "40" --validation_guidance_scale "1.0" @@ -112,6 +112,7 @@ miscellaneous_args=( --dit_precision "fp32" --num_euler_timesteps 50 --ema_start_step 0 + --enable_gradient_checkpointing_type "full" ) srun torchrun \ diff --git a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh index cd443fbc4..f1a79d688 100644 --- a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh +++ b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh @@ -2,11 +2,12 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" -VALIDATION_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/" +VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json" NUM_GPUS=4 # export CUDA_VISIBLE_DEVICES=4,5 @@ -49,7 +50,7 @@ dataset_args=( # Validation arguments validation_args=( --log_validation - --validation_preprocessed_path $VALIDATION_DIR + --validation_dataset_file $VALIDATION_DATASET_FILE --validation_steps 50 --validation_sampling_steps "50" --validation_guidance_scale "1.0" @@ -75,6 +76,7 @@ miscellaneous_args=( --dit_precision "fp32" --num_euler_timesteps 50 --ema_start_step 0 + --enable_gradient_checkpointing_type "full" ) torchrun \ diff --git a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm index ad31610b0..f439a634b 100644 --- a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm +++ b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm @@ -41,7 +41,7 @@ echo "NODE_RANK: $NODE_RANK" NUM_GPUS=8 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" -VALIDATION_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/" +VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json" # export CUDA_VISIBLE_DEVICES=4,5 # IP=[MASTER NODE IP] @@ -83,7 +83,7 @@ dataset_args=( # Validation arguments validation_args=( --log_validation - --validation_preprocessed_path "$VALIDATION_DIR" + --validation_dataset_file "$VALIDATION_DATASET_FILE" --validation_steps 100 --validation_sampling_steps "50" --validation_guidance_scale "1.0" @@ -109,6 +109,7 @@ miscellaneous_args=( --dit_precision "fp32" --num_euler_timesteps 50 --ema_start_step 0 + --enable_gradient_checkpointing_type "full" ) srun torchrun \ diff --git a/fastvideo/v1/configs/pipelines/registry.py b/fastvideo/v1/configs/pipelines/registry.py index ab873b3ba..2239cea30 100644 --- a/fastvideo/v1/configs/pipelines/registry.py +++ b/fastvideo/v1/configs/pipelines/registry.py @@ -23,6 +23,7 @@ PIPE_NAME_TO_CONFIG: Dict[str, Type[PipelineConfig]] = { "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig, "hunyuanvideo-community/HunyuanVideo": HunyuanConfig, "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig, + "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig, "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig, "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig, "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig, diff --git a/fastvideo/v1/dataset/__init__.py b/fastvideo/v1/dataset/__init__.py index e20ae425c..e1e9bdb53 100644 --- a/fastvideo/v1/dataset/__init__.py +++ b/fastvideo/v1/dataset/__init__.py @@ -12,7 +12,10 @@ from fastvideo.v1.dataset.validation_dataset import ValidationDataset def getdataset(args) -> VideoCaptionMergedDataset: - temporal_sample = TemporalRandomCrop(args.num_frames) # 16 x + if args.do_temporal_sample: + temporal_sample = TemporalRandomCrop(args.num_frames) # 16 x + else: + temporal_sample = None norm_fun = Lambda(lambda x: 2.0 * x - 1.0) resize_topcrop = [ CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True), diff --git a/fastvideo/v1/dataset/preprocessing_datasets.py b/fastvideo/v1/dataset/preprocessing_datasets.py index c586e52f7..a3078e9fc 100644 --- a/fastvideo/v1/dataset/preprocessing_datasets.py +++ b/fastvideo/v1/dataset/preprocessing_datasets.py @@ -281,10 +281,12 @@ class FrameSamplingStage(DatasetFilterStage): frame_interval).astype(int) # Temporal crop if too long - if len(frame_indices - ) > self.num_frames and temporal_sample_fn is not None: - begin_index, end_index = temporal_sample_fn(len(frame_indices)) - frame_indices = frame_indices[begin_index:end_index] + if len(frame_indices) > self.num_frames: + if temporal_sample_fn is not None: + begin_index, end_index = temporal_sample_fn(len(frame_indices)) + frame_indices = frame_indices[begin_index:end_index] + else: + frame_indices = frame_indices[:self.num_frames] batch.sample_frame_index = frame_indices.tolist() batch.sample_num_frames = len(frame_indices) diff --git a/fastvideo/v1/dataset/validation_dataset.py b/fastvideo/v1/dataset/validation_dataset.py index 5fceb1ea3..3795dde29 100644 --- a/fastvideo/v1/dataset/validation_dataset.py +++ b/fastvideo/v1/dataset/validation_dataset.py @@ -4,15 +4,17 @@ import os import pathlib import datasets -import torch +from torch.utils.data import IterableDataset +from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank, + get_world_size) from fastvideo.v1.logger import init_logger from fastvideo.v1.models.vision_utils import load_image, load_video logger = init_logger(__name__) -class ValidationDataset(torch.utils.data.IterableDataset): +class ValidationDataset(IterableDataset): def __init__(self, filename: str): super().__init__() @@ -48,10 +50,65 @@ class ValidationDataset(torch.utils.data.IterableDataset): f"Unsupported file format {self.filename.suffix} for validation dataset. Supported formats are: {_SUPPORTED_FILE_FORMATS}" ) - self._data = data.to_iterable_dataset() + # Get distributed training info + self.global_rank = get_world_rank() + self.world_size = get_world_size() + self.sp_world_size = get_sp_world_size() + self.num_sp_groups = self.world_size // self.sp_world_size + + # Convert to list to get total samples + self.all_samples = list(data) + self.original_total_samples = len(self.all_samples) + + # Extend samples to be a multiple of DP degree (num_sp_groups) + remainder = self.original_total_samples % self.num_sp_groups + if remainder != 0: + samples_to_add = self.num_sp_groups - remainder + + # Duplicate samples cyclically to reach the target + additional_samples = [] + for i in range(samples_to_add): + additional_samples.append( + self.all_samples[i % self.original_total_samples]) + + self.all_samples.extend(additional_samples) + + self.total_samples = len(self.all_samples) + + # Calculate which SP group this rank belongs to + self.sp_group_id = self.global_rank // self.sp_world_size + + # Now all SP groups will have equal number of samples + self.samples_per_sp_group = self.total_samples // self.num_sp_groups + + # Calculate start and end indices for this SP group + self.start_idx = self.sp_group_id * self.samples_per_sp_group + self.end_idx = self.start_idx + self.samples_per_sp_group + + # Get samples for this SP group + self.sp_group_samples = self.all_samples[self.start_idx:self.end_idx] + + logger.info( + "Rank %s (SP group %s): " + "Original samples: %s, " + "Extended samples: %s, " + "SP group samples: %s, " + "Range: [%s:%s]", + self.global_rank, + self.sp_group_id, + self.original_total_samples, + self.total_samples, + len(self.sp_group_samples), + self.start_idx, + self.end_idx, + local_main_process_only=False) + + def __len__(self): + """Return the number of samples for this SP group.""" + return len(self.sp_group_samples) def __iter__(self): - for sample in self._data: + for sample in self.sp_group_samples: # For consistency reasons, we mandate that "caption" is always present in the validation dataset. # However, since the model specifications use "prompt", we create an alias here. sample["prompt"] = sample["caption"] @@ -64,6 +121,7 @@ class ValidationDataset(torch.utils.data.IterableDataset): if sample.get("image_path", None) is not None: image_path = sample["image_path"] image_path = os.path.join(self.dir, image_path) + sample["image_path"] = image_path if not pathlib.Path(image_path).is_file( ) and not image_path.startswith("http"): logger.warning("Image file %s does not exist.", image_path) @@ -73,6 +131,7 @@ class ValidationDataset(torch.utils.data.IterableDataset): if sample.get("video_path", None) is not None: video_path = sample["video_path"] video_path = os.path.join(self.dir, video_path) + sample["video_path"] = video_path if not pathlib.Path(video_path).is_file( ) and not video_path.startswith("http"): logger.warning("Video file %s does not exist.", video_path) @@ -82,6 +141,7 @@ class ValidationDataset(torch.utils.data.IterableDataset): if sample.get("control_image_path", None) is not None: control_image_path = sample["control_image_path"] control_image_path = os.path.join(self.dir, control_image_path) + sample["control_image_path"] = control_image_path if not pathlib.Path(control_image_path).is_file( ) and not control_image_path.startswith("http"): logger.warning("Control Image file %s does not exist.", @@ -92,6 +152,7 @@ class ValidationDataset(torch.utils.data.IterableDataset): if sample.get("control_video_path", None) is not None: control_video_path = sample["control_video_path"] control_video_path = os.path.join(self.dir, control_video_path) + sample["control_video_path"] = control_video_path if not pathlib.Path(control_video_path).is_file( ) and not control_video_path.startswith("http"): logger.warning("Control Video file %s does not exist.", diff --git a/fastvideo/v1/distributed/parallel_state.py b/fastvideo/v1/distributed/parallel_state.py index 8c992f2a5..ff511d58f 100644 --- a/fastvideo/v1/distributed/parallel_state.py +++ b/fastvideo/v1/distributed/parallel_state.py @@ -750,6 +750,7 @@ def init_distributed_environment( distributed_init_method: str = "env://", local_rank: int = 0, backend: str = "nccl", + device_id: Optional[torch.device] = None, ): logger.debug( "world_size=%d rank=%d local_rank=%d " @@ -764,7 +765,8 @@ def init_distributed_environment( backend=backend, init_method=distributed_init_method, world_size=world_size, - rank=rank) + rank=rank, + device_id=device_id) # set the local rank # local_rank is not available in torch ProcessGroup, # see https://github.com/pytorch/pytorch/issues/122816 @@ -921,13 +923,15 @@ def maybe_init_distributed_environment_and_model_parallel( local_rank = int(os.environ.get("LOCAL_RANK", 0)) world_size = int(os.environ.get("WORLD_SIZE", 1)) rank = int(os.environ.get("RANK", 0)) + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) - torch.cuda.set_device(local_rank) init_distributed_environment( world_size=world_size, rank=rank, local_rank=local_rank, - distributed_init_method=distributed_init_method) + distributed_init_method=distributed_init_method, + device_id=device) initialize_model_parallel(tensor_model_parallel_size=tp_size, sequence_model_parallel_size=sp_size) diff --git a/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py b/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py index 286aa9dfd..d3b3f9505 100644 --- a/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py +++ b/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py @@ -89,11 +89,10 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline): """Get CLIP features from the first frame of each video.""" first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute( 0, 2, 3, 1) # (B, C, T, H, W) -> (B, H, W, C) - batch_size, _, num_frames, height, width = valid_data[ - "pixel_values"].shape - latent_height = height // self.get_module( - "vae").spatial_compression_ratio - latent_width = width // self.get_module("vae").spatial_compression_ratio + _, _, num_frames, height, width = valid_data["pixel_values"].shape + # latent_height = height // self.get_module( + # "vae").spatial_compression_ratio + # latent_width = width // self.get_module("vae").spatial_compression_ratio processed_images = [] # Frame has values between -1 and 1 @@ -157,26 +156,26 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline): latent_condition = latent_condition * self.get_module( "vae").scaling_factor - mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, - latent_width) - mask_lat_size[:, :, list(range(1, num_frames))] = 0 - first_frame_mask = mask_lat_size[:, :, 0:1] - first_frame_mask = torch.repeat_interleave( - first_frame_mask, - dim=2, - repeats=self.get_module("vae").temporal_compression_ratio) - mask_lat_size = torch.concat( - [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) - mask_lat_size = mask_lat_size.view( - batch_size, -1, - self.get_module("vae").temporal_compression_ratio, latent_height, - latent_width) - mask_lat_size = mask_lat_size.transpose(1, 2) - mask_lat_size = mask_lat_size.to(latent_condition.device) + # mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, + # latent_width) + # mask_lat_size[:, :, list(range(1, num_frames))] = 0 + # first_frame_mask = mask_lat_size[:, :, 0:1] + # first_frame_mask = torch.repeat_interleave( + # first_frame_mask, + # dim=2, + # repeats=self.get_module("vae").temporal_compression_ratio) + # mask_lat_size = torch.concat( + # [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) + # mask_lat_size = mask_lat_size.view( + # batch_size, -1, + # self.get_module("vae").temporal_compression_ratio, latent_height, + # latent_width) + # mask_lat_size = mask_lat_size.transpose(1, 2) + # mask_lat_size = mask_lat_size.to(latent_condition.device) - image_latent = torch.concat([mask_lat_size, latent_condition], dim=1) + # image_latent = torch.concat([mask_lat_size, latent_condition], dim=1) - features["first_frame_latent"] = image_latent + features["first_frame_latent"] = latent_condition return features diff --git a/fastvideo/v1/pipelines/preprocess/v1_preprocess.py b/fastvideo/v1/pipelines/preprocess/v1_preprocess.py index a8b562226..0d297cb81 100644 --- a/fastvideo/v1/pipelines/preprocess/v1_preprocess.py +++ b/fastvideo/v1/pipelines/preprocess/v1_preprocess.py @@ -79,6 +79,9 @@ if __name__ == "__main__": parser.add_argument("--text_max_length", type=int, default=256) parser.add_argument("--speed_factor", type=float, default=1.0) parser.add_argument("--drop_short_ratio", type=float, default=1.0) + parser.add_argument("--do_temporal_sample", + default=False, + action="store_true") # text encoder & vae & diffusion model parser.add_argument("--text_encoder_name", type=str, diff --git a/fastvideo/v1/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py index 6fc7b79c2..711a2ecf3 100644 --- a/fastvideo/v1/pipelines/stages/denoising.py +++ b/fastvideo/v1/pipelines/stages/denoising.py @@ -300,10 +300,6 @@ class DenoisingStage(PipelineStage): if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.STA_mode == STA_Mode.STA_SEARCHING: self.save_sta_search_results(batch) - if fastvideo_args.use_cpu_offload: - self.transformer.to('cpu') - torch.cuda.empty_cache() - return batch def prepare_extra_func_kwargs(self, func, kwargs) -> Dict[str, Any]: diff --git a/fastvideo/v1/pipelines/stages/encoding.py b/fastvideo/v1/pipelines/stages/encoding.py index 410dc2aee..5dc701212 100644 --- a/fastvideo/v1/pipelines/stages/encoding.py +++ b/fastvideo/v1/pipelines/stages/encoding.py @@ -56,23 +56,16 @@ class EncodingStage(PipelineStage): latent_height = batch.height // self.vae.spatial_compression_ratio latent_width = batch.width // self.vae.spatial_compression_ratio - image = batch.preprocessed_image - # TODO(will) - if image is None: - assert batch.pil_image is not None - image = batch.pil_image - image = self.preprocess( - image, - vae_scale_factor=self.vae.spatial_compression_ratio, - height=batch.height, - width=batch.width).to(get_local_torch_device(), - dtype=torch.float32) + assert batch.pil_image is not None + image = batch.pil_image + image = self.preprocess( + image, + vae_scale_factor=self.vae.spatial_compression_ratio, + height=batch.height, + width=batch.width).to(get_local_torch_device(), dtype=torch.float32) + + image = image.unsqueeze(2) - image = image.unsqueeze(2) - else: - # assumes image is loaded from parquet file and used for validation - image = image.transpose(1, 2) - logger.info("image: %s", image.shape) video_condition = torch.cat([ image, image.new_zeros(image.shape[0], image.shape[1], @@ -143,6 +136,9 @@ class EncodingStage(PipelineStage): if hasattr(self, 'maybe_free_model_hooks'): self.maybe_free_model_hooks() + self.vae.to("cpu") + torch.cuda.empty_cache() + return batch def retrieve_latents(self, diff --git a/fastvideo/v1/pipelines/stages/image_encoding.py b/fastvideo/v1/pipelines/stages/image_encoding.py index 1dd3f87ab..11e9bbf77 100644 --- a/fastvideo/v1/pipelines/stages/image_encoding.py +++ b/fastvideo/v1/pipelines/stages/image_encoding.py @@ -54,8 +54,7 @@ class ImageEncodingStage(PipelineStage): Returns: The batch with encoded prompt embeddings. """ - if fastvideo_args.use_cpu_offload: - self.image_encoder = self.image_encoder.to(get_local_torch_device()) + self.image_encoder = self.image_encoder.to(get_local_torch_device()) image = batch.pil_image diff --git a/fastvideo/v1/pipelines/stages/input_validation.py b/fastvideo/v1/pipelines/stages/input_validation.py index 07f705ed1..31127d06e 100644 --- a/fastvideo/v1/pipelines/stages/input_validation.py +++ b/fastvideo/v1/pipelines/stages/input_validation.py @@ -7,7 +7,7 @@ import torch from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.logger import init_logger -from fastvideo.v1.models.vision_utils import load_image +from fastvideo.v1.models.vision_utils import load_image, load_video from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.pipelines.stages.base import PipelineStage from fastvideo.v1.pipelines.stages.validators import (StageValidators, @@ -94,7 +94,10 @@ class InputValidationStage(PipelineStage): # for i2v, get image from image_path if batch.image_path is not None: - image = load_image(batch.image_path) + if batch.image_path.endswith(".mp4"): + image = load_video(batch.image_path)[0] + else: + image = load_image(batch.image_path) batch.pil_image = image return batch diff --git a/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py b/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py index 2b4278bbe..7344b11f7 100644 --- a/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py +++ b/fastvideo/v1/pipelines/wan/wan_i2v_pipeline.py @@ -76,36 +76,4 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase): stage=DecodingStage(vae=self.get_module("vae"))) -class WanImageToVideoValidationPipeline(ComposedPipelineBase): - """ - I2V Validation pipeline for Wan2.1, assumes that the input are preprocess latents. - """ - _required_config_modules = ["vae", "scheduler", "transformer"] - - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - self.modules["scheduler"] = FlowUniPCMultistepScheduler( - shift=fastvideo_args.pipeline_config.flow_shift) - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - self.add_stage(stage_name="timestep_preparation_stage", - stage=TimestepPreparationStage( - scheduler=self.get_module("scheduler"))) - - self.add_stage(stage_name="latent_preparation_stage", - stage=LatentPreparationStage( - scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer"))) - - self.add_stage(stage_name="image_latent_preparation_stage", - stage=EncodingStage(vae=self.get_module("vae"))) - - self.add_stage(stage_name="denoising_stage", - stage=DenoisingStage( - transformer=self.get_module("transformer"), - scheduler=self.get_module("scheduler"))) - - self.add_stage(stage_name="decoding_stage", - stage=DecodingStage(vae=self.get_module("vae"))) - - EntryClass = WanImageToVideoPipeline diff --git a/fastvideo/v1/pipelines/wan/wan_pipeline.py b/fastvideo/v1/pipelines/wan/wan_pipeline.py index 67ece217f..1b3a51786 100644 --- a/fastvideo/v1/pipelines/wan/wan_pipeline.py +++ b/fastvideo/v1/pipelines/wan/wan_pipeline.py @@ -67,34 +67,4 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase): stage=DecodingStage(vae=self.get_module("vae"))) -class WanValidationPipeline(ComposedPipelineBase): - """ - Validation pipeline for Wan2.1, assumes that the input are preprocess latents. - """ - _required_config_modules = ["vae", "scheduler", "transformer"] - - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - self.modules["scheduler"] = FlowUniPCMultistepScheduler( - shift=fastvideo_args.pipeline_config.flow_shift) - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - """Set up pipeline stages with proper dependency injection.""" - self.add_stage(stage_name="timestep_preparation_stage", - stage=TimestepPreparationStage( - scheduler=self.get_module("scheduler"))) - - self.add_stage(stage_name="latent_preparation_stage", - stage=LatentPreparationStage( - scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer", None))) - - self.add_stage(stage_name="denoising_stage", - stage=DenoisingStage( - transformer=self.get_module("transformer"), - scheduler=self.get_module("scheduler"))) - - self.add_stage(stage_name="decoding_stage", - stage=DecodingStage(vae=self.get_module("vae"))) - - EntryClass = WanPipeline diff --git a/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py b/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py index 5e98f7247..670cca47c 100644 --- a/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py +++ b/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py @@ -32,7 +32,7 @@ def run_worker(): "--inference_mode", "False", "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--data_path", "data/mini_dataset_i2v_VSA/combined_parquet_dataset", - "--validation_preprocessed_path", "data/mini_dataset_i2v_VSA/validation_parquet_dataset", + "--validation_dataset_file", "examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json", "--train_batch_size", "1", "--num_latent_t", "4", "--num_gpus", "2", diff --git a/fastvideo/v1/tests/training/Vanilla/test_training_loss.py b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py index b49ba603b..7f1459cc0 100644 --- a/fastvideo/v1/tests/training/Vanilla/test_training_loss.py +++ b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py @@ -38,7 +38,7 @@ def run_worker(): "--inference_mode", "False", "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset", - "--validation_preprocessed_path", "data/crush-smol_processed_t2v/validation_parquet_dataset", + "--validation_dataset_file", "examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json", "--train_batch_size", "2", "--num_latent_t", "4", "--num_gpus", "4", diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index 0596f33f4..4bfa18e73 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -14,6 +14,7 @@ import torchvision from diffusers import FlowMatchEulerDiscreteScheduler from diffusers.optimization import get_scheduler from einops import rearrange +from torch.utils.data import DataLoader from torchdata.stateful_dataloader import StatefulDataLoader from tqdm.auto import tqdm @@ -24,6 +25,7 @@ from fastvideo.v1.configs.sample import SamplingParam from fastvideo.v1.dataset import build_parquet_map_style_dataloader from fastvideo.v1.dataset.dataloader.schema import ( pyarrow_schema_t2v, pyarrow_schema_t2v_validation) +from fastvideo.v1.dataset.validation_dataset import ValidationDataset from fastvideo.v1.distributed import (cleanup_dist_env_and_memory, get_local_torch_device, get_sp_group, get_world_group) @@ -38,7 +40,7 @@ from fastvideo.v1.training.training_utils import ( clip_grad_norm_while_handling_failing_dtensor_cases, compute_density_for_timestep_sampling, get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint, shard_latents_across_sp) -from fastvideo.v1.utils import is_vsa_available, set_random_seed +from fastvideo.v1.utils import is_vsa_available, set_random_seed, shallow_asdict import wandb # isort: skip @@ -416,6 +418,9 @@ class TrainingPipeline(ComposedPipelineBase, ABC): self.init_steps = 0 def train(self) -> None: + logger.info('rank: %s: start training', + self.global_rank, + local_main_process_only=False) assert self.training_args is not None # Set random seeds for deterministic training @@ -548,60 +553,40 @@ class TrainingPipeline(ComposedPipelineBase, ABC): logger.info("VSA validation sparsity: %s", self.training_args.VSA_sparsity) - def _prepare_validation_inputs( - self, sampling_param: SamplingParam, training_args: TrainingArgs, - validation_batch: Dict[str, Any], num_inference_steps: int, - negative_prompt_embeds: torch.Tensor | None, - negative_prompt_attention_mask: torch.Tensor | None - ) -> ForwardBatch: + def _prepare_validation_batch(self, sampling_param: SamplingParam, + training_args: TrainingArgs, + validation_batch: Dict[str, Any], + num_inference_steps: int) -> ForwardBatch: + sampling_param.prompt = validation_batch['prompt'] + sampling_param.height = training_args.num_height + sampling_param.width = training_args.num_width + sampling_param.num_inference_steps = num_inference_steps + sampling_param.data_type = "video" + sampling_param.seed = self.seed - assert len(validation_batch['info_list'] - ) == 1, "Only batch size 1 is supported for validation" - prompt = validation_batch['info_list'][0]['prompt'] - prompt_embeds = validation_batch['text_embedding'] - prompt_attention_mask = validation_batch['text_attention_mask'] - - prompt_embeds = prompt_embeds.to(get_local_torch_device()) - prompt_attention_mask = prompt_attention_mask.to( - get_local_torch_device()) - - # Calculate sizes latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8] n_tokens = latents_size[0] * latents_size[1] * latents_size[2] - temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio num_frames = (training_args.num_latent_t - 1) * temporal_compression_factor + 1 - - # Prepare batch for validation + sampling_param.num_frames = num_frames batch = ForwardBatch( - prompt=prompt, - data_type="video", + **shallow_asdict(sampling_param), latents=None, - seed=self.seed, # Use deterministic seed generator=torch.Generator(device="cpu").manual_seed(self.seed), - prompt_embeds=[prompt_embeds], - prompt_attention_mask=[prompt_attention_mask], - negative_prompt_embeds=[negative_prompt_embeds], - negative_attention_mask=[negative_prompt_attention_mask], - height=training_args.num_height, - width=training_args.num_width, - num_frames=num_frames, - num_inference_steps= - num_inference_steps, # Use the current validation step - guidance_scale=sampling_param.guidance_scale, n_tokens=n_tokens, eta=0.0, VSA_sparsity=training_args.VSA_sparsity, ) + return batch @torch.no_grad() def _log_validation(self, transformer, training_args, global_step) -> None: assert training_args is not None training_args.inference_mode = True - training_args.use_cpu_offload = False + training_args.use_cpu_offload = True if not training_args.log_validation: return if self.validation_pipeline is None: @@ -617,48 +602,53 @@ class TrainingPipeline(ComposedPipelineBase, ABC): logger.info("Using validation seed: %s", self.seed) # Prepare validation prompts - logger.info('fastvideo_args.validation_preprocessed_path: %s', - training_args.validation_preprocessed_path) - validation_dataset, validation_dataloader = build_parquet_map_style_dataloader( - training_args.validation_preprocessed_path, - batch_size=1, - parquet_schema=self.validation_dataset_schema, - num_data_workers=0, - cfg_rate=0.0, - drop_last=False, - drop_first_row=sampling_param.negative_prompt is not None) - if sampling_param.negative_prompt: - negative_prompt_embeds, negative_prompt_attention_mask, negative_prompt = validation_dataset.get_validation_negative_prompt( - ) - logger.info("Using negative_prompt: %s", negative_prompt) + logger.info('rank: %s: fastvideo_args.validation_dataset_file: %s', + self.global_rank, + training_args.validation_dataset_file, + local_main_process_only=False) + validation_dataset = ValidationDataset( + training_args.validation_dataset_file) + validation_dataloader = DataLoader(validation_dataset, + batch_size=None, + num_workers=0) transformer.eval() validation_steps = training_args.validation_sampling_steps.split(",") validation_steps = [int(step) for step in validation_steps] validation_steps = [step for step in validation_steps if step > 0] + # Log validation results for this step + world_group = get_world_group() + num_sp_groups = world_group.world_size // self.sp_group.world_size # Process each validation prompt for each validation step for num_inference_steps in validation_steps: + logger.info("rank: %s: num_inference_steps: %s", + self.global_rank, + num_inference_steps, + local_main_process_only=False) step_videos: List[np.ndarray] = [] step_captions: List[str] = [] for validation_batch in validation_dataloader: - batch = self._prepare_validation_inputs( - sampling_param, training_args, validation_batch, - num_inference_steps, negative_prompt_embeds, - negative_prompt_attention_mask) + batch = self._prepare_validation_batch(sampling_param, + training_args, + validation_batch, + num_inference_steps) + logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s", + self.global_rank, + self.rank_in_sp_group, + batch.prompt, + local_main_process_only=False) assert batch.prompt is not None and isinstance( batch.prompt, str) step_captions.append(batch.prompt) # Run validation inference - with torch.no_grad(), torch.autocast("cuda", - dtype=torch.bfloat16): - output_batch = self.validation_pipeline.forward( - batch, training_args) - samples = output_batch.output + output_batch = self.validation_pipeline.forward( + batch, training_args) + samples = output_batch.output if self.rank_in_sp_group != 0: continue @@ -672,10 +662,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC): frames.append((x * 255).numpy().astype(np.uint8)) step_videos.append(frames) - # Log validation results for this step - world_group = get_world_group() - num_sp_groups = world_group.world_size // self.sp_group.world_size - + torch.distributed.barrier() # Only sp_group leaders (rank_in_sp_group == 0) need to send their # results to global rank 0 if self.rank_in_sp_group == 0: @@ -716,8 +703,10 @@ class TrainingPipeline(ComposedPipelineBase, 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) + torch.distributed.barrier() # Re-enable gradients for training + training_args.inference_mode = False transformer.train() gc.collect() torch.cuda.empty_cache() diff --git a/fastvideo/v1/training/wan_i2v_training_pipeline.py b/fastvideo/v1/training/wan_i2v_training_pipeline.py index 1c5475e75..e6f482668 100644 --- a/fastvideo/v1/training/wan_i2v_training_pipeline.py +++ b/fastvideo/v1/training/wan_i2v_training_pipeline.py @@ -15,10 +15,9 @@ from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler) from fastvideo.v1.pipelines.pipeline_batch_info import (ForwardBatch, TrainingBatch) -from fastvideo.v1.pipelines.wan.wan_i2v_pipeline import ( - WanImageToVideoValidationPipeline) +from fastvideo.v1.pipelines.wan.wan_i2v_pipeline import WanImageToVideoPipeline from fastvideo.v1.training.training_pipeline import TrainingPipeline -from fastvideo.v1.utils import is_vsa_available +from fastvideo.v1.utils import is_vsa_available, shallow_asdict vsa_available = is_vsa_available() @@ -50,15 +49,18 @@ class WanI2VTrainingPipeline(TrainingPipeline): args_copy = deepcopy(training_args) args_copy.inference_mode = True - args_copy.pipeline_config.vae_config.load_encoder = False - validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained( + args_copy.use_cpu_offload = True + # args_copy.pipeline_config.vae_config.load_encoder = False + # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained( + validation_pipeline = WanImageToVideoPipeline.from_pretrained( training_args.model_path, args=None, inference_mode=True, loaded_modules={"transformer": self.get_module("transformer")}, tp_size=training_args.tp_size, sp_size=training_args.sp_size, - num_gpus=training_args.num_gpus) + num_gpus=training_args.num_gpus, + use_cpu_offload=True) self.validation_pipeline = validation_pipeline @@ -117,8 +119,29 @@ class WanI2VTrainingPipeline(TrainingPipeline): image_latents = training_batch.image_latents.to( get_local_torch_device(), dtype=torch.bfloat16) + temporal_compression_ratio = 4 + num_frames = (self.training_args.num_latent_t - + 1) * temporal_compression_ratio + 1 + batch_size, num_channels, _, latent_height, latent_width = image_latents.shape + mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, + latent_width) + mask_lat_size[:, :, 1:] = 0 + + first_frame_mask = mask_lat_size[:, :, :1] + first_frame_mask = torch.repeat_interleave( + first_frame_mask, dim=2, repeats=temporal_compression_ratio) + mask_lat_size = torch.cat([first_frame_mask, mask_lat_size[:, :, 1:]], + dim=2) + mask_lat_size = mask_lat_size.view(batch_size, -1, + temporal_compression_ratio, + latent_height, latent_width) + mask_lat_size = mask_lat_size.transpose(1, 2) + mask_lat_size = mask_lat_size.to( + image_latents.device).to(dtype=torch.bfloat16) + training_batch.noisy_model_input = torch.cat( - [training_batch.noisy_model_input, image_latents], dim=1) + [training_batch.noisy_model_input, mask_lat_size, image_latents], + dim=1) return training_batch @@ -156,55 +179,34 @@ class WanI2VTrainingPipeline(TrainingPipeline): } return training_batch - def _prepare_validation_inputs( - self, sampling_param: SamplingParam, training_args: TrainingArgs, - validation_batch: Dict[str, Any], num_inference_steps: int, - negative_prompt_embeds: torch.Tensor | None, - negative_prompt_attention_mask: torch.Tensor | None - ) -> ForwardBatch: - embeddings = validation_batch['text_embedding'] - masks = validation_batch['text_attention_mask'] - clip_features = validation_batch['clip_feature'] - pil_image = validation_batch['pil_image'] - infos = validation_batch['info_list'] - prompt = infos[0]['prompt'] + def _prepare_validation_batch(self, sampling_param: SamplingParam, + training_args: TrainingArgs, + validation_batch: Dict[str, Any], + num_inference_steps: int) -> ForwardBatch: + sampling_param.prompt = validation_batch['prompt'] + sampling_param.height = training_args.num_height + sampling_param.width = training_args.num_width + sampling_param.image_path = validation_batch['video_path'] + sampling_param.num_inference_steps = num_inference_steps + sampling_param.data_type = "video" + sampling_param.seed = self.seed - prompt_embeds = embeddings.to(get_local_torch_device()) - prompt_attention_mask = masks.to(get_local_torch_device()) - clip_features = clip_features.to(get_local_torch_device()) - - # Calculate sizes latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8] n_tokens = latents_size[0] * latents_size[1] * latents_size[2] - temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio num_frames = (training_args.num_latent_t - 1) * temporal_compression_factor + 1 - - # Prepare batch for validation + sampling_param.num_frames = num_frames batch = ForwardBatch( - prompt=prompt, - data_type="video", + **shallow_asdict(sampling_param), latents=None, - seed=self.seed, # Use deterministic seed generator=torch.Generator(device="cpu").manual_seed(self.seed), - prompt_embeds=[prompt_embeds], - prompt_attention_mask=[prompt_attention_mask], - negative_prompt_embeds=[negative_prompt_embeds], - negative_attention_mask=[negative_prompt_attention_mask], - image_embeds=[clip_features], - preprocessed_image=pil_image, - height=training_args.num_height, - width=training_args.num_width, - num_frames=num_frames, - num_inference_steps= - num_inference_steps, # Use the current validation step - guidance_scale=sampling_param.guidance_scale, n_tokens=n_tokens, eta=0.0, VSA_sparsity=training_args.VSA_sparsity, ) + return batch diff --git a/fastvideo/v1/training/wan_training_pipeline.py b/fastvideo/v1/training/wan_training_pipeline.py index 6657a1d80..971aa7845 100644 --- a/fastvideo/v1/training/wan_training_pipeline.py +++ b/fastvideo/v1/training/wan_training_pipeline.py @@ -6,7 +6,7 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.v1.logger import init_logger from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler) -from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline +from fastvideo.v1.pipelines.wan.wan_pipeline import WanPipeline from fastvideo.v1.training.training_pipeline import TrainingPipeline from fastvideo.v1.utils import is_vsa_available @@ -36,8 +36,9 @@ class WanTrainingPipeline(TrainingPipeline): args_copy = deepcopy(training_args) args_copy.inference_mode = True + args_copy.use_cpu_offload = True args_copy.pipeline_config.vae_config.load_encoder = False - validation_pipeline = WanValidationPipeline.from_pretrained( + validation_pipeline = WanPipeline.from_pretrained( training_args.model_path, args=None, inference_mode=True, diff --git a/pyproject.toml b/pyproject.toml index ae3736a9f..e22db6a91 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,7 +79,6 @@ exclude = ["assets*", "docker*", "docs", "scripts*"] exclude = ["assets*", "docker*", "docs", "scripts*"] [tool.mypy] -warn_return_any = true warn_unused_configs = true ignore_missing_imports = true disallow_untyped_calls = true diff --git a/scripts/finetune/finetune_v1.sh b/scripts/finetune/finetune_v1.sh index f4b280e76..3f97a029b 100644 --- a/scripts/finetune/finetune_v1.sh +++ b/scripts/finetune/finetune_v1.sh @@ -1,9 +1,10 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA DATA_DIR=[your data dir] -VALIDATION_DIR=[your validation dir] +VALIDATION_DATASET_FILE=[your validation dataset file] NUM_GPUS=4 # export CUDA_VISIBLE_DEVICES=4,5 # IP=[MASTER NODE IP] @@ -15,7 +16,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\ --inference_mode False\ --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ --data_path "$DATA_DIR"\ - --validation_preprocessed_path "$VALIDATION_DIR"\ + --validation_dataset_file "$VALIDATION_DATASET_FILE"\ --train_batch_size=4 \ --num_latent_t 20 \ --sp_size 4 \ diff --git a/scripts/finetune/finetune_v1_VSA.sh b/scripts/finetune/finetune_v1_VSA.sh index cebd57df4..1229dc773 100644 --- a/scripts/finetune/finetune_v1_VSA.sh +++ b/scripts/finetune/finetune_v1_VSA.sh @@ -1,10 +1,11 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=offline export WANDB_API_KEY='your_wandb_api_key' +export TOKENIZERS_PARALLELISM=false # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA export TRITON_CACHE_DIR=/tmp/triton_cache DATA_DIR=~/train/ -VALIDATION_DIR=~latents/test/ +VALIDATION_DATASET_FILE=[your validation dataset file] NUM_GPUS=8 export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN # export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN @@ -21,7 +22,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \ --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ --cache_dir "/home/ray/.cache" \ --data_path "$DATA_DIR" \ - --validation_preprocessed_path "$VALIDATION_DIR" \ + --validation_dataset_file "$VALIDATION_DATASET_FILE" \ --train_batch_size 1 \ --num_latent_t 16 \ --sp_size 1 \