update

examples/training/consistency_finetune/ode_init/finetune_ode_init.sh

@@ -7,7 +7,8 @@ export TOKENIZERS_PARALLELISM=false
 
 # MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
+# DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-extended-t2v-1-3b/"
 VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
 NUM_GPUS=1
 # export CUDA_VISIBLE_DEVICES=4,5
@@ -16,11 +17,11 @@ NUM_GPUS=1
 # Training arguments
 training_args=(
   --tracker_project_name "wan_ode_init"
-  --output_dir "wan_ode_init"
+  --output_dir "wan_ode_init_70k"
   --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "fixed_wan_ode_init_6e-6"
+  --wandb_run_name "fixed_wan_ode_init_70k_6e-6"
   # --resume_from_checkpoint "ode_init_diffusers/"
-  --max_train_steps 2000
+  --max_train_steps 6000
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 1
@@ -28,6 +29,7 @@ training_args=(
   --num_height 480
   --num_width 832
   --num_frames 77
+  --warp_denoising_step
   # --enable_gradient_checkpointing_type "full"
 )
 
@@ -56,7 +58,7 @@ dataset_args=(
 validation_args=(
   --log_validation
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 20
+  --validation_steps 50
   --validation_sampling_steps "50"
   --validation_guidance_scale "6.0"
 )
@@ -65,7 +67,7 @@ validation_args=(
 optimizer_args=(
   --learning_rate 6e-6
   --mixed_precision "bf16"
-  --checkpointing_steps 2000
+  --checkpointing_steps 1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/consistency_finetune/ode_init/finetune_ode_init_2k_vidprom.slurm

@@ -39,7 +39,7 @@ echo "NODE_RANK: $NODE_RANK"
 # MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing/"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
 NUM_GPUS=8
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

examples/training/consistency_finetune/ode_init/finetune_ode_init_70k_vidprom.slurm

@@ -0,0 +1,132 @@
+#!/bin/bash
+#SBATCH --job-name=4e6_ode_vidprom2k
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=ode_vidprom2k_outputtest/ode_vidprom2k_4e-6.out
+#SBATCH --error=ode_vidprom2k_outputtest/ode_vidprom2k_4e-6.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate will-fv2
+
+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 WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+
+# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-extended-t2v-1-3b/"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
+NUM_GPUS=2
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_vidprom_70k_test"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "vidprom_70k_wan_ode_init_4e-6_test"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --max_train_steps 6001
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --warp_denoising_step
+  # --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 2
+  --hsdp_shard_dim 1
+)
+
+# 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 "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 4e-6
+  --mixed_precision "bf16"
+  --checkpointing_steps 2000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --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
+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/training/ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm

@@ -39,7 +39,7 @@ echo "NODE_RANK: $NODE_RANK"
 # MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
 NUM_GPUS=8
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

fastvideo/configs/pipelines/wan.py

@@ -146,5 +146,6 @@ class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
 @dataclass
 class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
     is_causal: bool = True
+    flow_shift: int = 5
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 750, 500, 250])

fastvideo/fastvideo_args.py

@@ -681,6 +681,7 @@ class TrainingArgs(FastVideoArgs):
     log_visualization: bool = False
     # simulate generator forward to match inference
     simulate_generator_forward: bool = False
+    warp_denoising_step: bool = False
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -1025,6 +1026,10 @@ class TrainingArgs(FastVideoArgs):
             "--simulate-generator-forward",
             action=StoreBoolean,
             help="Whether to simulate generator forward to match inference")
+        parser.add_argument(
+            "--warp-denoising-step",
+            action=StoreBoolean,
+            help="Whether to warp denoising step according to the scheduler time shift")
 
         return parser
 

fastvideo/training/ode_causal_pipeline.py

@@ -53,6 +53,17 @@ class ODEInitTrainingPipeline(TrainingPipeline):
             shift=self.timestep_shift)
 
         self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250], dtype=torch.long, device=get_local_torch_device())
+        if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
+            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))).cuda()
+            logger.info(f"timesteps: {timesteps}")
+            self.dmd_denoising_steps = timesteps[1000 - self.dmd_denoising_steps]
+            logger.info(f"warped self.dmd_denoising_steps: {self.dmd_denoising_steps}")
+        else:
+            logger.info(f"not warped")
+        self.dmd_denoising_steps = self.dmd_denoising_steps.to(get_local_torch_device())
+
+        logger.info(f"denoising_step_list: {self.dmd_denoising_steps}")
+
         logger.info("Initialized ODE-init training pipeline with %s denoising steps", len(self.dmd_denoising_steps))
         # Cache for nearest trajectory index per DMD step (computed lazily on first batch)
         self._cached_closest_idx_per_dmd = None
@@ -117,6 +128,7 @@ class ODEInitTrainingPipeline(TrainingPipeline):
         else:
             raise ValueError(
                 f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}")
+        # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing
         trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5)
 
         # Move to device