fix hardcode

examples/distill/SFWan2.1-A1.3B/distill_dmd.slurm

@@ -0,0 +1,182 @@
+#!/bin/bash
+#SBATCH --job-name=1.3B_moe_4n_sf_distill_no_2nd_update
+#SBATCH --partition=main
+#SBATCH --nodes=4
+#SBATCH --ntasks=4
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=1.3B_moe_4n_sf_distill_output/moe_sf_distill_no_2nd_update.out
+#SBATCH --error=1.3B_moe_4n_sf_distill_output/moe_sf_distill_no_2nd_update.err
+#SBATCH --exclusive
+set -e -x
+
+
+# Environment Setup
+# source ~/conda/miniconda/bin/activate
+# conda activate wei-fv
+# export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Basic Info
+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_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=32
+
+# Model paths for Self-Forcing DMD distillation:
+# GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+GENERATOR_MODEL_PATH="rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Teacher model
+# FAKE_SCORE_MODEL_PATH="rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers"  # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+# DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+# VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_2.json"
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/1.3B_MoE_SFwan_t2v_finetune"
+  --wandb_run_name "1.3B_moe_sf_distill"
+  # --use_sf_wan
+  # --sf_ode_init_path "checkpoints/ode_init.pt"
+  --max_train_steps 5000
+  --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 81  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
+  --enable_gradient_checkpointing_type "full"
+  --log_visualization
+  --simulate_generator_forward
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --generator_model_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 8
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-6
+  --fake_score_learning_rate 4e-7
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 10.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 200
+  --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
+  --init_weights_from_safetensors_2 "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
+  # --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/FastVideo2/vidprom_8b16k_1e-5_gn1/checkpoint-3500/transformer/diffusion_pytorch_model.safetensors"
+#   --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/sf2/checkpoint_v3_7k/model.safetensors"
+)
+
+# Self-forcing DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 4.0
+  --real_score_guidance_scale_2 3.0
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+# Self-forcing specific arguments
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+  --validate_cache_structure False  # Set to True for debugging KV cache issues
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+--nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"

examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.slurm

@@ -0,0 +1,180 @@
+#!/bin/bash
+#SBATCH --job-name=sf_distill_profile
+#SBATCH --partition=main
+#SBATCH --nodes=2
+#SBATCH --ntasks=2
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=sf_distill_output/sf_distill_profile.out
+#SBATCH --error=sf_distill_output/sf_distill_profile.err
+#SBATCH --exclusive
+set -e -x
+
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Basic Info
+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=29503
+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_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=offline
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+export FASTVIDEO_TORCH_PROFILER_DIR=/mnt/weka/home/hao.zhang/wei/FastVideo/traces
+export FASTVIDEO_TORCH_PROFILE_REGIONS=profiler_region_model_loading
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=16
+
+# Model paths for Self-Forcing DMD distillation:
+GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+# DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+# VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_2.json"
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/SFwan_t2v_finetune"
+  --wandb_run_name "sf_distill_2e-6_4e-7_4n"
+  # --use_sf_wan
+  # --sf_ode_init_path "checkpoints/ode_init.pt"
+  --max_train_steps 2
+  --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 81  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
+  --enable_gradient_checkpointing_type "full"
+  --log_visualization
+  --simulate_generator_forward
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --generator_model_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 8
+)
+
+# Validation arguments
+validation_args=(
+#   --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-6
+  --fake_score_learning_rate 4e-7
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 10.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 200
+  --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
+  # --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/FastVideo2/vidprom_8b16k_1e-5_gn1/checkpoint-3500/transformer/diffusion_pytorch_model.safetensors"
+#   --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/sf2/checkpoint_v3_7k/model.safetensors"
+)
+
+# Self-forcing DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 3.0
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+# Self-forcing specific arguments
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+  --validate_cache_structure False  # Set to True for debugging KV cache issues
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+--nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"

examples/distill/SFWan2.2-A14B/distill_dmd.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-#SBATCH --job-name=t2v
+#SBATCH --job-name=14B_moe_4n_sf_distill_no_2nd_update
 #SBATCH --partition=main
 #SBATCH --nodes=4
 #SBATCH --ntasks=4
@@ -7,144 +7,167 @@
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
 #SBATCH --mem=1440G
-#SBATCH --output=dmd_t2v_output/t2v_%j.out
-#SBATCH --error=dmd_t2v_output/t2v_%j.err
+#SBATCH --output=14B_moe_4n_sf_distill_output/moe_sf_distill_no_2nd_update.out
+#SBATCH --error=14B_moe_4n_sf_distill_output/moe_sf_distill_no_2nd_update.err
 #SBATCH --exclusive
+set -e -x
+
+
+# Environment Setup
+# source ~/conda/miniconda/bin/activate
+# conda activate wei-fv
+# export HOME="/mnt/weka/home/hao.zhang/wei"
 
 # Basic Info
 export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
-export NCCL_DEBUG_SUBSYS=INIT,NET
 # 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_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f"
-# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
 # Configs
 NUM_GPUS=8
+NUM_TOTAL_GPUS=32
 
-# Model paths for Self-Forcing DMD distillation with Wan2.2:
-# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
-# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
-# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
-GENERATOR_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Updated to Wan2.2
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Teacher model
-FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
+# Model paths for Self-Forcing DMD distillation:
+# GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+GENERATOR_MODEL_PATH="rand0nmr/SFWan2.2-T2V-A14B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Teacher model
+# FAKE_SCORE_MODEL_PATH="rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers"  # Critic model
 
 # DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
 DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
-# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
-# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
 # VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
 VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
-
+# VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_2.json"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
+# Training arguments
 training_args=(
-  --tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
-  --output_dir "/mnt/sharefs/users/hao.zhang/SFwan2.2_t2v_finetune"
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --max_train_steps 4000
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/14B_MoE_SFwan_t2v_finetune"
+  --wandb_run_name "14B_moe_sf_distill"
+  # --use_sf_wan
+  # --sf_ode_init_path "checkpoints/ode_init.pt"
+  --max_train_steps 5000
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 1
   --num_latent_t 21
-  --num_height 448  # Updated to match Wan2.2 config
-  --num_width 832   # Updated to match Wan2.2 config
+  --num_height 480
+  --num_width 832
+  --num_frames 81  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
   --enable_gradient_checkpointing_type "full"
+  --log_visualization
   --simulate_generator_forward
-  # --log_visualization
-  --num_frames 81
   --num_frame_per_block 3  # Frame generation block size for self-forcing
   --enable_gradient_masking
   --gradient_mask_last_n_frames 21
-  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
-  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
 )
 
+# Parallel arguments
 parallel_args=(
-  --num_gpus 32 # 64
+  --num_gpus $NUM_TOTAL_GPUS # 64
   --sp_size 1
   --tp_size 1
   --hsdp_replicate_dim 1 # 64
-  --hsdp_shard_dim 32
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
 )
 
+# Model arguments
 model_args=(
-  --model_path $GENERATOR_MODEL_PATH
+  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
   --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --generator_model_path $GENERATOR_MODEL_PATH
   --real_score_model_path $REAL_SCORE_MODEL_PATH
   --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )
 
+# Dataset arguments
 dataset_args=(
   --data_path "$DATA_DIR"
-  --dataloader_num_workers 4
+  --dataloader_num_workers 8
 )
 
+# Validation arguments
 validation_args=(
   --log_validation
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 20
+  --validation_steps 50
   --validation_sampling_steps "4"
   --validation_guidance_scale "6.0" # not used for dmd inference
 )
 
+# Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 2e-6
+  --fake_score_learning_rate 4e-7
   --mixed_precision "bf16"
   --training_state_checkpointing_steps 500
   --weight_only_checkpointing_steps 500
   --weight_decay 0.01
   --betas '0.0,0.999'
-  --max_grad_norm 1.0
+  --max_grad_norm 10.0
 )
 
+# Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
   --checkpoints_total_limit 3
   --training_cfg_rate 0.0
   --dit_precision "fp32"
-  --flow_shift 5
+  --flow_shift 12
   --seed 1000
   --use_ema True
   --ema_decay 0.99
-  --ema_start_step 100
+  --ema_start_step 200
+  --init_weights_from_safetensors "/mnt/sharefs/users/hao.zhang/wl/release/self_forcing_ode_init_wan22_high_bz128_1e-5/diffusers_2000/"
+  --init_weights_from_safetensors_2 "/mnt/sharefs/users/hao.zhang/wl/release/self_forcing_ode_init_wan22_low_bz128_1e-5/diffusers_2000/"
 )
 
+# Self-forcing DMD arguments
 dmd_args=(
   --dmd_denoising_steps '1000,750,500,250'
   --min_timestep_ratio 0.02
   --max_timestep_ratio 0.98
   --dfake_gen_update_ratio 5
-  --real_score_guidance_scale 3.0
-  --fake_score_learning_rate 8e-6
+  --real_score_guidance_scale 4.0
+  --real_score_guidance_scale_2 3.0
   --fake_score_betas '0.0,0.999'
   --warp_denoising_step
 )
 
+# Self-forcing specific arguments
 self_forcing_args=(
   --independent_first_frame False  # Whether to treat first frame independently
   --same_step_across_blocks True  # Whether to use same denoising step across all blocks
   --last_step_only False  # Whether to only use the last denoising step
   --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+  --validate_cache_structure False  # Set to True for debugging KV cache issues
 )
 
 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" \
+--nproc_per_node $NUM_GPUS \
     fastvideo/training/wan_self_forcing_distillation_pipeline.py \
     "${parallel_args[@]}" \
     "${model_args[@]}" \

examples/distill/SFWan2.2-A14B/distill_dmd.slurm

@@ -0,0 +1,165 @@
+#!/bin/bash
+#SBATCH --job-name=moe_4n_sf_distill
+#SBATCH --partition=main
+#SBATCH --nodes=4
+#SBATCH --ntasks=4
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=moe_4n_sf_distill_output/moe_sf_distill_%j.out
+#SBATCH --error=moe_4n_sf_distill_output/moe_sf_distill_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export NCCL_DEBUG_SUBSYS=INIT,NET
+# 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 TOKENIZERS_PARALLELISM=false
+# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde"
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=32
+
+# Model paths for Self-Forcing DMD distillation with Wan2.2:
+# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+GENERATOR_MODEL_PATH="rand0nmr/SFWan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
+# GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" # Updated to Wan2.2
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Teacher model
+# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
+# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+training_args=(
+  --tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_dmd"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 448  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --enable_gradient_checkpointing_type "full"
+  --simulate_generator_forward
+  --log_visualization
+  --num_frames 81
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+  --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
+  --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
+  # --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors
+)
+
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 40
+  --validation_sampling_steps "8"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+optimizer_args=(
+  --learning_rate 2e-6
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 10.0
+)
+
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 12
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 200
+)
+
+dmd_args=(
+  --dmd_denoising_steps '1000,850,700,550,350,275,200,125'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 4.0
+  --real_score_guidance_scale_2 3.0
+  --fake_score_learning_rate 4e-7
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+)
+
+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/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"

examples/distill/SFWan2.2-A14B/distill_low_dmd.sh

@@ -0,0 +1,181 @@
+#!/bin/bash
+#SBATCH --job-name=14B_moe_4n_sf_distill_no_2nd_update
+#SBATCH --partition=main
+#SBATCH --nodes=4
+#SBATCH --ntasks=4
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=14B_moe_4n_sf_distill_output/moe_sf_distill_no_2nd_update.out
+#SBATCH --error=14B_moe_4n_sf_distill_output/moe_sf_distill_no_2nd_update.err
+#SBATCH --exclusive
+set -e -x
+
+
+# Environment Setup
+# source ~/conda/miniconda/bin/activate
+# conda activate wei-fv
+# export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Basic Info
+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_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=32
+
+# Model paths for Self-Forcing DMD distillation:
+# GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+GENERATOR_MODEL_PATH="rand0nmr/SFWan2.2-T2V-A14B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Teacher model
+# FAKE_SCORE_MODEL_PATH="rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers"  # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+# DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+# VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_2.json"
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/14B_MoE_SFwan_t2v_finetune"
+  --wandb_run_name "14B_moe_sf_distill_low_dmd"
+  # --use_sf_wan
+  # --sf_ode_init_path "checkpoints/ode_init.pt"
+  --max_train_steps 5000
+  --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 81  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
+  --enable_gradient_checkpointing_type "full"
+  --log_visualization
+  --simulate_generator_forward
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --generator_model_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 8
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-6
+  --fake_score_learning_rate 4e-7
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 10.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 12
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 200
+  --init_weights_from_safetensors "/mnt/sharefs/users/hao.zhang/wl/release/self_forcing_ode_init_wan22_low_bz128_1e-5/diffusers_2000/"
+  --init_weights_from_safetensors_2 "/mnt/sharefs/users/hao.zhang/wl/release/self_forcing_ode_init_wan22_low_bz128_1e-5/diffusers_2000/"
+)
+
+# Self-forcing DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '350,275,200,125'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 4.0
+  --real_score_guidance_scale_2 3.0
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+# Self-forcing specific arguments
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+  --validate_cache_structure False  # Set to True for debugging KV cache issues
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+--nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"
+

examples/distill/Wan2.2-A14B/distill_dmd.slurm

@@ -0,0 +1,144 @@
+#!/bin/bash
+#SBATCH --job-name=moe_dmd_distill
+#SBATCH --partition=main
+#SBATCH --nodes=2
+#SBATCH --ntasks=2
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=moe_dmd_distill_output/moe_dmd_distill_%j.out
+#SBATCH --error=moe_dmd_distill_output/moe_dmd_distill_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export NCCL_DEBUG_SUBSYS=INIT,NET
+# 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 TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=16
+
+# Model paths for Self-Forcing DMD distillation with Wan2.2:
+GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
+# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
+# GENERATOR_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Updated to Wan2.2
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol_processed_t2v/combined_parquet_dataset"
+# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
+# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+training_args=(
+  --tracker_project_name SFwan2.2_t2v_distill_dmd  # Updated for Wan2.2
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_dmd"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 448  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --enable_gradient_checkpointing_type "full"
+  # --log_visualization
+  --simulate_generator_forward
+  --num_frames 81
+  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
+  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
+)
+
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 20
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --max_grad_norm 1.0
+)
+
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --ema_start_step 100
+)
+
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --generator_update_interval 5
+  --real_score_guidance_scale 3.0
+)
+
+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/wan_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}"

examples/training/finetune/Wan2.2-A14B/finetune.slurm

@@ -0,0 +1,131 @@
+#!/bin/bash
+#SBATCH --job-name=moe_finetune
+#SBATCH --partition=main
+#SBATCH --nodes=2
+#SBATCH --ntasks=2
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=moe_output/moe_%j.out
+#SBATCH --error=moe_output/moe_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export NCCL_DEBUG_SUBSYS=INIT,NET
+# 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 TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=16
+
+# Model paths for Self-Forcing DMD distillation with Wan2.2:
+GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
+# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol_processed_t2v/combined_parquet_dataset"
+# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
+# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol-single_processed_t2v/validation.json"
+
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+training_args=(
+  --tracker_project_name VSA_finetune  # Updated for Wan2.2
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/Wan2.2-MoE-finetune"
+  --max_train_steps 200
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 448  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --enable_gradient_checkpointing_type "full"
+  # --log_visualization
+  --num_frames 81
+  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
+  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
+)
+
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+#   --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "5.0" # not used for dmd inference
+)
+
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --dit_cpu_offload True
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 1.0
+)
+
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+)
+
+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/wan_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

fastvideo/configs/pipelines/registry.py

@@ -13,7 +13,7 @@ from fastvideo.configs.pipelines.wan import (
     FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
     Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config,
     WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig,
-    SelfForcingWanT2V480PConfig)
+    SelfForcingWanT2V480PConfig, SelfForcingMoEWanT2V480PConfig)
 # isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
@@ -36,6 +36,8 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
     "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
     "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+    "rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers": SelfForcingMoEWanT2V480PConfig,
+    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingMoEWanT2V480PConfig,
     "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config,
     "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
     "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,
@@ -137,4 +139,4 @@ def get_pipeline_config_cls_from_name(
             f"No match found for pipeline {pipeline_name_or_path}, please check the pipeline name or path."
         )
 
-    return pipeline_config_cls
+    return pipeline_config_cls

fastvideo/configs/pipelines/wan.py

@@ -165,3 +165,10 @@ class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 750, 500, 250])
     warp_denoising_step: bool = True
+
+@dataclass
+class SelfForcingMoEWanT2V480PConfig(SelfForcingWanT2V480PConfig):
+    boundary_ratio: float | None = 0.875
+
+    def __post_init__(self) -> None:
+        self.dit_config.boundary_ratio = self.boundary_ratio

fastvideo/configs/sample/registry.py

@@ -55,6 +55,8 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
 
     # Causal Self-Forcing Wan2.1
     "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+    "rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingWanT2V480PConfig,
     # Add other specific weight variants
 }
 

fastvideo/fastvideo_args.py

@@ -717,6 +717,7 @@ class TrainingArgs(FastVideoArgs):
     min_timestep_ratio: float = 0.2
     max_timestep_ratio: float = 0.98
     real_score_guidance_scale: float = 3.5
+    real_score_guidance_scale_2: float = 3.5
     fake_score_learning_rate: float = 0.0  # separate learning rate for fake_score_transformer, if 0.0, use learning_rate
     fake_score_lr_scheduler: str = "constant"  # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler
     fake_score_betas: str = "0.9,0.999"  # betas for fake score optimizer, format: "beta1,beta2"
@@ -1102,6 +1103,10 @@ class TrainingArgs(FastVideoArgs):
                             type=float,
                             default=TrainingArgs.real_score_guidance_scale,
                             help="Teacher guidance scale")
+        parser.add_argument("--real-score-guidance-scale-2",
+                            type=float,
+                            default=TrainingArgs.real_score_guidance_scale_2,
+                            help="Teacher guidance scale")
         parser.add_argument("--fake-score-learning-rate",
                             type=float,
                             default=TrainingArgs.fake_score_learning_rate,

fastvideo/models/loader/component_loader.py

@@ -445,12 +445,20 @@ class TransformerLoader(ComponentLoader):
             if 'transformer_2' in model_path:
                 custom_weights_path = getattr(fastvideo_args, 'init_weights_from_safetensors_2', None)
             assert custom_weights_path is not None, "Custom initialization weights must be provided"
+            
+            # Handle both directory and single file cases
             if os.path.isdir(custom_weights_path):
+                # Directory: collect all safetensors files
                 safetensors_list = glob.glob(
                     os.path.join(str(custom_weights_path), "*.safetensors"))
-            else:
+                if not safetensors_list:
+                    raise ValueError(f"No safetensors files found in directory {custom_weights_path}")
+            elif os.path.isfile(custom_weights_path):
+                # Single file: verify it's a safetensors file
                 assert custom_weights_path.endswith(".safetensors"), "Custom initialization weights must be a safetensors file"
                 safetensors_list = [custom_weights_path]
+            else:
+                raise ValueError(f"Custom weights path does not exist or is not accessible: {custom_weights_path}")
 
         logger.info("Loading model from %s safetensors files: %s",
                     len(safetensors_list), safetensors_list)

fastvideo/models/utils.py

@@ -180,3 +180,50 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
     sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
     pred_video = noise_input_latent - sigma_t * pred_noise
     return pred_video.to(dtype)
+
+def pred_noise_to_x_bound(pred_noise: torch.Tensor,
+                             noise_input_latent: torch.Tensor,
+                             timestep: torch.Tensor,
+                             boundary_timestep: torch.Tensor,
+                             scheduler: Any) -> torch.Tensor:
+    """
+    Convert predicted noise to clean latent.
+
+    Args:
+    pred_noise: the predicted noise with shape [B, C, H, W]
+        where B is batch_size or batch_size * num_frames
+    noise_input_latent: the noisy latent with shape [B, C, H, W],
+    timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames]
+    boundary_timestep: the boundary timestep with shape [1] or [bs * num_frames] or [bs, num_frames]
+    scheduler: the scheduler
+
+    Returns:
+        the predicted video with shape [B, C, H, W]
+    """
+    # If timestep is [bs, num_frames]
+    if timestep.ndim == 2:
+        timestep = timestep.flatten(0, 1)
+        assert timestep.numel() == noise_input_latent.shape[0]
+    elif timestep.ndim == 1:
+        # If timestep is [1]
+        if timestep.shape[0] == 1:
+            timestep = timestep.expand(noise_input_latent.shape[0])
+        else:
+            assert timestep.numel() == noise_input_latent.shape[0]
+    else:
+        raise ValueError(f"[pred_noise_to_pred_video] Invalid timestep shape: {timestep.shape}")
+    # timestep shape should be [B]
+    dtype = pred_noise.dtype
+    device = pred_noise.device
+    pred_noise = pred_noise.double().to(device)
+    noise_input_latent = noise_input_latent.double().to(device)
+    sigmas = scheduler.sigmas.double().to(device)
+    timesteps = scheduler.timesteps.double().to(device)
+    timestep_id = torch.argmin(
+        (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+    boundary_timestep_id = torch.argmin(
+        (timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1)
+    sigma_t_boundary = sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1)
+    pred_video = noise_input_latent - (sigma_t - sigma_t_boundary) * pred_noise
+    return pred_video.to(dtype)

fastvideo/pipelines/composed_pipeline_base.py

@@ -71,6 +71,7 @@ class ComposedPipelineBase(ABC):
 
         # Load modules directly in initialization
         logger.info("Loading pipeline modules...")
+        fastvideo_args.dit_cpu_offload = False
         self.modules = self.load_modules(fastvideo_args, loaded_modules)
 
     def set_trainable(self) -> None:
@@ -165,7 +166,7 @@ class ComposedPipelineBase(ABC):
             for key, value in kwargs.items():
                 setattr(fastvideo_args, key, value)
 
-            fastvideo_args.dit_cpu_offload = False
+            fastvideo_args.dit_cpu_offload = True
             # we hijack the precision to be the master weight type so that the
             # model is loaded with the correct precision. Subsequently we will
             # use FSDP2's MixedPrecisionPolicy to set the precision for the

fastvideo/pipelines/stages/causal_denoising.py

@@ -1,4 +1,5 @@
 import torch  # type: ignore
+import gc
 
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -39,8 +40,8 @@ class CausalDMDDenosingStage(DenoisingStage):
         # KV and cross-attention cache state (initialized on first forward)
         self.transformer = transformer
         self.transformer_2 = transformer_2
-        self.kv_cache1: list | None = None
-        self.crossattn_cache: list | None = None
+        # self.kv_cache1: list | None = None
+        # self.crossattn_cache: list | None = None
         # Model-dependent constants (aligned with causal_inference.py assumptions)
         self.num_transformer_blocks = len(self.transformer.blocks)
         self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block
@@ -103,34 +104,43 @@ class CausalDMDDenosingStage(DenoisingStage):
         assert torch.isnan(prompt_embeds[0]).sum() == 0
 
         # Initialize or reset caches
-        if self.kv_cache1 is None:
-            self._initialize_kv_cache(batch_size=latents.shape[0],
-                                      dtype=target_dtype,
-                                      device=latents.device)
-            self._initialize_crossattn_cache(
-                batch_size=latents.shape[0],
-                max_text_len=fastvideo_args.pipeline_config.
-                text_encoder_configs[0].arch_config.text_len,
-                dtype=target_dtype,
-                device=latents.device)
-        else:
-            assert self.crossattn_cache is not None
-            # reset cross-attention cache
-            for block_index in range(self.num_transformer_blocks):
-                self.crossattn_cache[block_index][
-                    "is_init"] = False  # type: ignore
-            # reset kv cache pointers
-            for block_index in range(len(self.kv_cache1)):
-                self.kv_cache1[block_index][
-                    "global_end_index"] = torch.tensor(  # type: ignore
-                        [0],
-                        dtype=torch.long,
-                        device=latents.device)
-                self.kv_cache1[block_index][
-                    "local_end_index"] = torch.tensor(  # type: ignore
-                        [0],
-                        dtype=torch.long,
-                        device=latents.device)
+        kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0],
+                                              dtype=target_dtype,
+                                              device=latents.device)
+        crossattn_cache = self._initialize_crossattn_cache(
+            batch_size=latents.shape[0],
+            max_text_len=fastvideo_args.pipeline_config.text_encoder_configs[0].
+            arch_config.text_len,
+            dtype=target_dtype,
+            device=latents.device)
+        # if self.kv_cache1 is None:
+        #     self._initialize_kv_cache(batch_size=latents.shape[0],
+        #                               dtype=target_dtype,
+        #                               device=latents.device)
+        #     self._initialize_crossattn_cache(
+        #         batch_size=latents.shape[0],
+        #         max_text_len=fastvideo_args.pipeline_config.
+        #         text_encoder_configs[0].arch_config.text_len,
+        #         dtype=target_dtype,
+        #         device=latents.device)
+        # else:
+        #     assert self.crossattn_cache is not None
+        #     # reset cross-attention cache
+        #     for block_index in range(self.num_transformer_blocks):
+        #         self.crossattn_cache[block_index][
+        #             "is_init"] = False  # type: ignore
+        #     # reset kv cache pointers
+        #     for block_index in range(len(self.kv_cache1)):
+        #         self.kv_cache1[block_index][
+        #             "global_end_index"] = torch.tensor(  # type: ignore
+        #                 [0],
+        #                 dtype=torch.long,
+        #                 device=latents.device)
+        #         self.kv_cache1[block_index][
+        #             "local_end_index"] = torch.tensor(  # type: ignore
+        #                 [0],
+        #                 dtype=torch.long,
+        #                 device=latents.device)
 
         # Optional: cache context features from provided image latents prior to generation
         current_start_frame = 0
@@ -153,8 +163,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                         image_first_btchw,
                         prompt_embeds,
                         t_zero,
-                        kv_cache=self.kv_cache1,
-                        crossattn_cache=self.crossattn_cache,
+                        kv_cache=kv_cache1,
+                        crossattn_cache=crossattn_cache,
                         current_start=current_start_frame *
                         self.frame_seq_length,
                         **image_kwargs,
@@ -179,8 +189,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                         ref_btchw,
                         prompt_embeds,
                         t_zero,
-                        kv_cache=self.kv_cache1,
-                        crossattn_cache=self.crossattn_cache,
+                        kv_cache=kv_cache1,
+                        crossattn_cache=crossattn_cache,
                         current_start=current_start_frame *
                         self.frame_seq_length,
                         **image_kwargs,
@@ -280,8 +290,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                             latent_model_input,
                             prompt_embeds,
                             t_expanded_noise,
-                            kv_cache=self.kv_cache1,
-                            crossattn_cache=self.crossattn_cache,
+                            kv_cache=kv_cache1,
+                            crossattn_cache=crossattn_cache,
                             current_start=(pos_start_base + start_index) *
                             self.frame_seq_length,
                             start_frame=start_index,
@@ -345,8 +355,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                         context_bcthw,
                         prompt_embeds,
                         t_expanded_context,
-                        kv_cache=self.kv_cache1,
-                        crossattn_cache=self.crossattn_cache,
+                        kv_cache=kv_cache1,
+                        crossattn_cache=crossattn_cache,
                         current_start=(pos_start_base + start_index) *
                         self.frame_seq_length,
                         start_frame=start_index,
@@ -355,6 +365,11 @@ class CausalDMDDenosingStage(DenoisingStage):
                     )
                 start_index += current_num_frames
 
+        del kv_cache1
+        del crossattn_cache
+        gc.collect()
+        torch.cuda.empty_cache()
+
         batch.latents = latents
         return batch
 
@@ -392,7 +407,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                 torch.tensor([0], dtype=torch.long, device=device),
             })
 
-        self.kv_cache1 = kv_cache1
+        # self.kv_cache1 = kv_cache1
+        return kv_cache1
 
     def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype,
                                     device) -> None:
@@ -421,7 +437,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                 "is_init":
                 False,
             })
-        self.crossattn_cache = crossattn_cache
+        # self.crossattn_cache = crossattn_cache
+        return crossattn_cache
 
     def verify_input(self, batch: ForwardBatch,
                      fastvideo_args: FastVideoArgs) -> VerificationResult:
@@ -445,4 +462,4 @@ class CausalDMDDenosingStage(DenoisingStage):
         result.add_check(
             "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x:
             not batch.do_classifier_free_guidance or V.list_not_empty(x))
-        return result
+        return result

fastvideo/training/distillation_pipeline.py

@@ -7,7 +7,7 @@ import time
 from abc import abstractmethod
 from collections import deque
 from collections.abc import Iterator
-from typing import Any
+from typing import Any, Tuple
 
 import imageio
 import numpy as np
@@ -260,6 +260,7 @@ class DistillationPipeline(TrainingPipeline):
                                 self.num_train_timestep)
 
         self.real_score_guidance_scale = self.training_args.real_score_guidance_scale
+        self.real_score_guidance_scale_2 = self.training_args.real_score_guidance_scale_2
 
         self.generator_ema: EMA_FSDP | None = None
         self.generator_ema_2: EMA_FSDP | None = None
@@ -541,6 +542,15 @@ class DistillationPipeline(TrainingPipeline):
             self.generator_ema_2.decay = original_decay_2
             logger.info("EMA_2 reset completed")
 
+    def _get_real_score_guidance_scale(self, timestep: torch.Tensor):
+        """
+        Get the appropriate real score guidance scale based on timestep and boundary logic.
+        """
+        if self.boundary_timestep is not None and self.real_score_transformer_2 is not None:
+            if timestep.item() < self.boundary_timestep:
+                return self.real_score_guidance_scale_2
+        return self.real_score_guidance_scale
+
     def _get_real_score_transformer(self, timestep: torch.Tensor):
         """
         Get the appropriate real score transformer based on timestep and boundary logic.
@@ -615,7 +625,11 @@ class DistillationPipeline(TrainingPipeline):
             noisy_latent, timestep, training_batch.conditional_dict,
             training_batch)
 
-        pred_noise = self.transformer(**training_batch.input_kwargs).permute(
+        if self.transformer_2 is not None and self.train_transformer_2:
+            current_transformer = self.transformer_2
+        else:
+            current_transformer = self.transformer
+        pred_noise = current_transformer(**training_batch.input_kwargs).permute(
             0, 2, 1, 3, 4)
         pred_video = pred_noise_to_pred_video(
             pred_noise=pred_noise.flatten(0, 1),
@@ -626,7 +640,9 @@ class DistillationPipeline(TrainingPipeline):
         return pred_video
 
     def _generator_multi_step_simulation_forward(
-            self, training_batch: TrainingBatch) -> torch.Tensor:
+            self,
+            training_batch: TrainingBatch,
+            exit_flags: list[int] | None = None) -> Tuple[torch.Tensor, float]:
         """Forward pass through student transformer matching inference procedure."""
         latents = training_batch.latents
         dtype = latents.dtype
@@ -662,13 +678,17 @@ class DistillationPipeline(TrainingPipeline):
             with torch.no_grad():
                 for step_idx in range(max_target_idx):
                     current_timestep = self.denoising_step_list[step_idx]
+                    if self.transformer_2 is not None and current_timestep < self.boundary_timestep:
+                        current_transformer = self.transformer_2
+                    else:
+                        current_transformer = self.transformer
                     current_timestep_tensor = current_timestep * torch.ones(
                         1, device=self.device, dtype=torch.long)
                     # Run student model to get flow prediction
                     training_batch_temp = self._build_distill_input_kwargs(
                         current_noise_latents, current_timestep_tensor,
                         training_batch.conditional_dict, training_batch)
-                    pred_flow = self.transformer(
+                    pred_flow = current_transformer(
                         **training_batch_temp.input_kwargs).permute(
                             0, 2, 1, 3, 4)
                     pred_clean = pred_noise_to_pred_video(
@@ -711,8 +731,12 @@ class DistillationPipeline(TrainingPipeline):
         training_batch = self._build_distill_input_kwargs(
             noisy_input, target_timestep, training_batch.conditional_dict,
             training_batch)
-        pred_noise = self.transformer(**training_batch.input_kwargs).permute(
-            0, 2, 1, 3, 4)
+        if self.transformer_2 is not None:
+            pred_noise = self.transformer_2(**training_batch.input_kwargs).permute(
+                0, 2, 1, 3, 4)
+        else:
+            pred_noise = self.transformer(**training_batch.input_kwargs).permute(
+                0, 2, 1, 3, 4)
         pred_video = pred_noise_to_pred_video(
             pred_noise=pred_noise.flatten(0, 1),
             noise_input_latent=noisy_input.flatten(0, 1),
@@ -720,15 +744,30 @@ class DistillationPipeline(TrainingPipeline):
             scheduler=self.noise_scheduler).unflatten(0, pred_noise.shape[:2])
         training_batch.dmd_latent_vis_dict[
             "generator_timestep"] = target_timestep.float().detach()
-        return pred_video
+        return pred_video, target_timestep.item()
 
     def _dmd_forward(self, generator_pred_video: torch.Tensor,
-                     training_batch: TrainingBatch) -> torch.Tensor:
+                     training_batch: TrainingBatch,
+                     exit_timestep: float) -> torch.Tensor:
         """Compute DMD (Diffusion Model Distillation) loss."""
         original_latent = generator_pred_video
         with torch.no_grad():
-            timestep = torch.randint(0,
-                                     self.num_train_timestep, [1],
+            # Detect whether we're doing MoE DMD
+            # if self.boundary_timestep is not None and self.transformer_2 is not None:
+            #     if exit_timestep < self.boundary_timestep:  # low noise
+            #         start_timestep = 0
+            #         end_timestep = int(self.boundary_timestep)
+            #     else:  # high noise
+            #         start_timestep = int(self.boundary_timestep)
+            #         end_timestep = self.num_train_timestep
+            # else:
+            #     start_timestep = 0
+            #     end_timestep = self.num_train_timestep
+            start_timestep = 0
+            end_timestep = self.num_train_timestep
+
+            timestep = torch.randint(start_timestep,
+                                     end_timestep, [1],
                                      device=self.device,
                                      dtype=torch.long)
             world_group = get_world_group()
@@ -738,7 +777,8 @@ class DistillationPipeline(TrainingPipeline):
             timestep = shift_timestep(
                 timestep,
                 self.timestep_shift,  # type: ignore
-                self.num_train_timestep)
+                min_timestep=start_timestep,
+                max_timestep=end_timestep)
 
             timestep = timestep.clamp(self.min_timestep, self.max_timestep)
 
@@ -805,7 +845,7 @@ class DistillationPipeline(TrainingPipeline):
 
             real_score_pred_video = pred_real_video_cond + (
                 pred_real_video_cond -
-                pred_real_video_uncond) * self.real_score_guidance_scale
+                pred_real_video_uncond) * self._get_real_score_guidance_scale(timestep)
 
             grad = (faker_score_pred_video - real_score_pred_video) / torch.abs(
                 original_latent - real_score_pred_video).mean()
@@ -837,13 +877,29 @@ class DistillationPipeline(TrainingPipeline):
                 current_timestep=training_batch.timesteps,
                 attn_metadata=training_batch.attn_metadata_vsa):
             if self.training_args.simulate_generator_forward:
-                generator_pred_video = self._generator_multi_step_simulation_forward(
+                generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward(
                     training_batch)
             else:
                 generator_pred_video = self._generator_forward(training_batch)
 
-        fake_score_timestep = torch.randint(0,
-                                            self.num_train_timestep, [1],
+        # Detect whether we're doing MoE DMD
+        # If yes, then we need to use the appropriate timestep range
+        # if self.boundary_timestep is not None and self.transformer_2 is not None:
+        #     if exit_timestep < self.boundary_timestep:
+        #         start_timestep = 0
+        #         end_timestep = self.num_train_timestep
+        #         # end_timestep = int(self.boundary_timestep)
+        #     else:
+        #         start_timestep = int(self.boundary_timestep)
+        #         end_timestep = self.num_train_timestep
+        # else:
+        #     start_timestep = 0
+        #     end_timestep = self.num_train_timestep
+        start_timestep = 0
+        end_timestep = self.num_train_timestep
+
+        fake_score_timestep = torch.randint(start_timestep,
+                                            end_timestep, [1],
                                             device=self.device,
                                             dtype=torch.long)
         world_group = get_world_group()
@@ -853,7 +909,8 @@ class DistillationPipeline(TrainingPipeline):
         fake_score_timestep = shift_timestep(
             fake_score_timestep,
             self.timestep_shift,  # type: ignore
-            self.num_train_timestep)
+            min_timestep=start_timestep,
+            max_timestep=end_timestep)
 
         fake_score_timestep = fake_score_timestep.clamp(self.min_timestep,
                                                         self.max_timestep)
@@ -918,7 +975,7 @@ class DistillationPipeline(TrainingPipeline):
 
     def _prepare_dit_inputs(self,
                             training_batch: TrainingBatch) -> TrainingBatch:
-        super()._prepare_dit_inputs(training_batch)
+        # super()._prepare_dit_inputs(training_batch, prepare_timesteps=prepare_timesteps)
         conditional_dict = {
             "encoder_hidden_states": training_batch.encoder_hidden_states,
             "encoder_attention_mask": training_batch.encoder_attention_mask,
@@ -929,10 +986,17 @@ class DistillationPipeline(TrainingPipeline):
                 "encoder_attention_mask": self.negative_prompt_attention_mask,
             }
             training_batch.unconditional_dict = unconditional_dict
+        else:
+            unconditional_dict = {
+                "encoder_hidden_states": training_batch.encoder_hidden_states,
+                "encoder_attention_mask": training_batch.encoder_attention_mask,
+            }
+            training_batch.unconditional_dict = unconditional_dict
 
         training_batch.dmd_latent_vis_dict = {}
         training_batch.fake_score_latent_vis_dict = {}
 
+        training_batch.timesteps = 0
         training_batch.conditional_dict = conditional_dict
         training_batch.raw_latent_shape = training_batch.latents.shape
         training_batch.latents = training_batch.latents.permute(0, 2, 1, 3, 4)
@@ -967,6 +1031,8 @@ class DistillationPipeline(TrainingPipeline):
             batches.append(batch)
 
         self.optimizer.zero_grad()
+        if self.transformer_2 is not None:
+            self.optimizer_2.zero_grad()
         total_dmd_loss = 0.0
         dmd_latent_vis_dict = {}
         fake_score_latent_vis_dict = {}
@@ -978,7 +1044,7 @@ class DistillationPipeline(TrainingPipeline):
                         current_timestep=batch_gen.timesteps,
                         attn_metadata=batch_gen.attn_metadata_vsa):
                     if self.training_args.simulate_generator_forward:
-                        generator_pred_video = self._generator_multi_step_simulation_forward(
+                        generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward(
                             batch_gen)
                     else:
                         generator_pred_video = self._generator_forward(
@@ -988,7 +1054,8 @@ class DistillationPipeline(TrainingPipeline):
                                          attn_metadata=batch_gen.attn_metadata):
                     dmd_loss = self._dmd_forward(
                         generator_pred_video=generator_pred_video,
-                        training_batch=batch_gen)
+                        training_batch=batch_gen,
+                        exit_timestep=exit_timestep)
 
                 with set_forward_context(
                         current_timestep=batch_gen.timesteps,
@@ -997,12 +1064,20 @@ class DistillationPipeline(TrainingPipeline):
                 total_dmd_loss += dmd_loss.detach().item()
 
             # Only clip gradients for the model that is currently training
-            self._clip_model_grad_norm_(batch_gen, self.transformer)
-            for param in self.transformer.parameters():
-                # check if the gradient is not None and not zero
-                assert param.grad is not None and param.grad.abs().sum() > 0
-            self.optimizer.step()
-            self.optimizer.zero_grad(set_to_none=True)
+            if self.transformer_2 is not None and self.train_transformer_2:
+                self._clip_model_grad_norm_(batch_gen, self.transformer_2)
+                for param in self.transformer_2.parameters():
+                    # check if the gradient is not None and not zero
+                    assert param.grad is not None and param.grad.abs().sum() > 0
+                self.optimizer_2.step()
+                self.optimizer_2.zero_grad(set_to_none=True)
+            else:
+                self._clip_model_grad_norm_(batch_gen, self.transformer)
+                for param in self.transformer.parameters():
+                    # check if the gradient is not None and not zero
+                    assert param.grad is not None and param.grad.abs().sum() > 0
+                self.optimizer.step()
+                self.optimizer.zero_grad(set_to_none=True)
 
             if self.generator_ema is not None:
                 self.generator_ema.update(self.transformer)
@@ -1395,7 +1470,7 @@ class DistillationPipeline(TrainingPipeline):
                 video = video.permute(0, 2, 1, 3, 4)
                 video = (video * 255).numpy().astype(np.uint8)
                 wandb_loss_dict[latent_key] = wandb.Video(
-                    video, fps=24, format="mp4")  # change to 16 for Wan2.1
+                    video, fps=16, format="mp4")  # change to 16 for Wan2.1 (change to 24 for wan 2.2)
                 # Clean up references
                 del video, latents
 
@@ -1426,7 +1501,7 @@ class DistillationPipeline(TrainingPipeline):
                 video = video.permute(0, 2, 1, 3, 4)
                 video = (video * 255).numpy().astype(np.uint8)
                 wandb_loss_dict[latent_key] = wandb.Video(
-                    video, fps=24, format="mp4")  # change to 16 for Wan2.1
+                    video, fps=16, format="mp4")  # change to 16 for Wan2.1
                 # Clean up references
                 del video, latents
 

fastvideo/training/self_forcing_distillation_pipeline.py

@@ -3,6 +3,7 @@ import copy
 import time
 from collections import deque
 from typing import Any
+import gc
 
 import numpy as np
 import torch
@@ -60,7 +61,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
         self.noise_scheduler = SelfForcingFlowMatchScheduler(
             num_inference_steps=1000,
-            shift=5.0,
+            shift=self.timestep_shift,
             sigma_min=0.0,
             extra_one_step=True,
             training=True)
@@ -76,9 +77,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         self.last_step_only = getattr(training_args, 'last_step_only', False)
         self.context_noise = getattr(training_args, 'context_noise', 0)
 
-        self.kv_cache1: list[dict[str, Any]] | None = None
-        self.crossattn_cache: list[dict[str, Any]] | None = None
-
         logger.info("Self-forcing generator update ratio: %s",
                     self.dfake_gen_update_ratio)
         logger.info("RANK: %s, exiting initialize_training_pipeline",
@@ -127,17 +125,45 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         Compute generator loss using DMD-style approach.
         The generator tries to fool the critic (fake_score_transformer).
         """
+        exit_flags = None
+        # If the generator is MoE, randomly sample an exit timestep, and turn on training for one of the dits
+        if self.boundary_timestep is not None and self.transformer_2 is not None:
+            assert self.same_step_across_blocks, "same_step_across_blocks must be True for MoE generator. Otherwise we might need to train both transformers which will cause OOM"
+            exit_flags = self.generate_and_sync_list(
+                1, len(self.denoising_step_list), training_batch.latents.device)
+            exit_timestep = self.denoising_step_list[exit_flags[0]].item()
+            # logger.info("Exit timestep in generator_loss(): %s", exit_timestep)
+            if exit_timestep < self.boundary_timestep:
+                self.train_transformer_2 = True
+                self._enable_training(self.transformer_2, self.optimizer_2)
+                self._disable_training(self.transformer, self.optimizer)
+            else:
+                self.train_transformer_2 = False
+                self._enable_training(self.transformer, self.optimizer)
+                self._disable_training(self.transformer_2, self.optimizer_2)
+        else:  # Default to normal single-dit generator
+            self.train_transformer_2 = False
+            self._enable_training(self.transformer, self.optimizer)
+
+        # Turns off training for fake score transformers
+        self._disable_training(self.fake_score_transformer,
+                               self.fake_score_optimizer)
+        if self.fake_score_transformer_2 is not None:
+            self._disable_training(self.fake_score_transformer_2,
+                                   self.fake_score_optimizer_2)
+
         with set_forward_context(
                 current_timestep=training_batch.timesteps,
                 attn_metadata=training_batch.attn_metadata_vsa):
-            generator_pred_video = self._generator_multi_step_simulation_forward(
-                training_batch)
+            generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward(
+                training_batch, exit_flags=exit_flags)
 
         with set_forward_context(current_timestep=training_batch.timesteps,
                                  attn_metadata=training_batch.attn_metadata):
             dmd_loss = self._dmd_forward(
                 generator_pred_video=generator_pred_video,
-                training_batch=training_batch)
+                training_batch=training_batch,
+                exit_timestep=exit_timestep)
 
         log_dict = {
             "dmdtrain_gradient_norm": torch.tensor(0.0, device=self.device)
@@ -152,6 +178,18 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         Compute critic loss using flow matching between noise and generator output.
         The critic learns to predict the flow from noise to the generator's output.
         """
+        # Turns off training for all generators
+        self._disable_training(self.transformer, self.optimizer)
+        if self.transformer_2 is not None:
+            self._disable_training(self.transformer_2, self.optimizer_2)
+
+        # Turns on training for fake score transformers
+        self._enable_training(self.fake_score_transformer,
+                              self.fake_score_optimizer)
+        if self.fake_score_transformer_2 is not None:
+            self._enable_training(self.fake_score_transformer_2,
+                                  self.fake_score_optimizer_2)
+
         updated_batch, flow_matching_loss = self.faker_score_forward(
             training_batch)
         training_batch.fake_score_latent_vis_dict = updated_batch.fake_score_latent_vis_dict
@@ -162,7 +200,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
     def _generator_multi_step_simulation_forward(
             self,
             training_batch: TrainingBatch,
-            return_sim_steps: bool = False) -> torch.Tensor:
+            exit_flags: list[int] | None = None) -> tuple[torch.Tensor, float]:
         """Forward pass through student transformer matching inference procedure with KV cache management.
         
         This function is adapted from the reference self-forcing implementation's inference_with_trajectory
@@ -244,7 +282,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
         # Step 1: Initialize KV cache to all zeros
         cache_frames = num_generated_frames + num_input_frames
-        self.kv_cache1, self.crossattn_cache = self._initialize_simulation_caches(
+        kv_cache1, crossattn_cache = self._initialize_simulation_caches(
             batch_size, dtype, self.device, max_num_frames=cache_frames)
 
         # Step 2: Cache context feature
@@ -268,8 +306,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                     timestep=training_batch_temp.input_kwargs['timestep'],
                     encoder_hidden_states_image=training_batch_temp.
                     input_kwargs.get('encoder_hidden_states_image'),
-                    kv_cache=self.kv_cache1,
-                    crossattn_cache=self.crossattn_cache,
+                    kv_cache=kv_cache1,
+                    crossattn_cache=crossattn_cache,
                     current_start=current_start_frame * self.frame_seq_length,
                     start_frame=current_start_frame)
             current_start_frame += 1
@@ -279,9 +317,13 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         if self.independent_first_frame and initial_latent is None:
             all_num_frames = [1] + all_num_frames
         num_denoising_steps = len(self.denoising_step_list)
-        exit_flags = self.generate_and_sync_list(len(all_num_frames),
-                                                 num_denoising_steps,
-                                                 device=noise.device)
+
+        if exit_flags is None:
+            exit_flags = self.generate_and_sync_list(len(all_num_frames),
+                                                     num_denoising_steps,
+                                                     device=noise.device)
+        exit_timestep = self.denoising_step_list[exit_flags[0]].item()
+
         start_gradient_frame_index = max(0, num_output_frames - 21)
 
         for block_index, current_num_frames in enumerate(all_num_frames):
@@ -302,14 +344,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
                 if self.boundary_timestep is not None and current_timestep < self.boundary_timestep and self.transformer_2 is not None:
                     current_model = self.transformer_2
-                    self._enable_training(self.transformer_2, self.optimizer_2)
-                    self._disable_training(self.transformer, self.optimizer)
                 else:
                     current_model = self.transformer
-                    self._enable_training(self.transformer, self.optimizer)
-                    if self.boundary_timestep is not None and self.transformer_2 is not None:
-                        self._disable_training(self.transformer_2,
-                                               self.optimizer_2)
 
                 if not exit_flag:
                     with torch.no_grad():
@@ -327,8 +363,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                             input_kwargs['timestep'],
                             encoder_hidden_states_image=training_batch_temp.
                             input_kwargs.get('encoder_hidden_states_image'),
-                            kv_cache=self.kv_cache1,
-                            crossattn_cache=self.crossattn_cache,
+                            kv_cache=kv_cache1,
+                            crossattn_cache=crossattn_cache,
                             current_start=current_start_frame *
                             self.frame_seq_length,
                             start_frame=current_start_frame).permute(
@@ -351,6 +387,9 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                        dtype=torch.long)).unflatten(
                                            0, denoised_pred.shape[:2])
                 else:
+                    logger.info(
+                        "Exit timestep in _generator_multi_step_simulation_forward(): %s",
+                        current_timestep)
                     # Final prediction with gradient control
                     if current_start_frame < start_gradient_frame_index:
                         with torch.no_grad():
@@ -367,8 +406,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                 input_kwargs['timestep'],
                                 encoder_hidden_states_image=training_batch_temp.
                                 input_kwargs.get('encoder_hidden_states_image'),
-                                kv_cache=self.kv_cache1,
-                                crossattn_cache=self.crossattn_cache,
+                                kv_cache=kv_cache1,
+                                crossattn_cache=crossattn_cache,
                                 current_start=current_start_frame *
                                 self.frame_seq_length,
                                 start_frame=current_start_frame).permute(
@@ -387,8 +426,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                             input_kwargs['timestep'],
                             encoder_hidden_states_image=training_batch_temp.
                             input_kwargs.get('encoder_hidden_states_image'),
-                            kv_cache=self.kv_cache1,
-                            crossattn_cache=self.crossattn_cache,
+                            kv_cache=kv_cache1,
+                            crossattn_cache=crossattn_cache,
                             current_start=current_start_frame *
                             self.frame_seq_length,
                             start_frame=current_start_frame).permute(
@@ -407,7 +446,22 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                    current_num_frames] = denoised_pred
 
             # Step 3.3: rerun with timestep zero to update the cache
-            context_timestep = torch.ones_like(timestep) * self.context_noise
+            if self.boundary_timestep is not None and self.transformer_2 is not None:
+                if exit_timestep < self.boundary_timestep:
+                    context_timestep = torch.ones_like(
+                        timestep) * self.context_noise
+                    current_model = self.transformer_2
+                else:
+                    # For high noise expert, the context timestep should be the boundary timestep
+                    # context_timestep = torch.ones_like(timestep) * self.boundary_timestep
+                    context_timestep = torch.ones_like(
+                        timestep) * self.boundary_timestep
+                    current_model = self.transformer
+            else:
+                context_timestep = torch.ones_like(
+                    timestep) * self.context_noise
+                current_model = self.transformer
+
             denoised_pred = self.noise_scheduler.add_noise(
                 denoised_pred.flatten(0, 1),
                 torch.randn_like(denoised_pred.flatten(0, 1)),
@@ -419,7 +473,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                     training_batch.conditional_dict, training_batch)
 
                 # context_timestep is 0 so we use transformer_2
-                current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer
+                # current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer
                 current_model(
                     hidden_states=training_batch_temp.
                     input_kwargs['hidden_states'],
@@ -428,10 +482,35 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                     timestep=training_batch_temp.input_kwargs['timestep'],
                     encoder_hidden_states_image=training_batch_temp.
                     input_kwargs.get('encoder_hidden_states_image'),
-                    kv_cache=self.kv_cache1,
-                    crossattn_cache=self.crossattn_cache,
+                    kv_cache=kv_cache1,
+                    crossattn_cache=crossattn_cache,
                     current_start=current_start_frame * self.frame_seq_length,
                     start_frame=current_start_frame)
+                # update the high noise model kv cache to match the low noise kv cache
+                # if current_model is self.transformer_2:
+                #     # TODO: which timestep to use?
+                #     context_timestep = torch.ones_like(timestep) * self.context_noise
+                #     denoised_pred = self.noise_scheduler.add_noise(
+                #         denoised_pred.flatten(0, 1),
+                #         torch.randn_like(denoised_pred.flatten(0, 1)),
+                #         context_timestep).unflatten(0, denoised_pred.shape[:2])
+
+                #     training_batch_temp = self._build_distill_input_kwargs(
+                #         denoised_pred, context_timestep,
+                #         training_batch.conditional_dict, training_batch)
+
+                #     self.transformer(
+                #         hidden_states=training_batch_temp.
+                #         input_kwargs['hidden_states'],
+                #         encoder_hidden_states=training_batch_temp.
+                #         input_kwargs['encoder_hidden_states'],
+                #         timestep=training_batch_temp.input_kwargs['timestep'],
+                #         encoder_hidden_states_image=training_batch_temp.
+                #         input_kwargs.get('encoder_hidden_states_image'),
+                #         kv_cache=kv_cache1,
+                #         crossattn_cache=crossattn_cache,
+                #         current_start=current_start_frame * self.frame_seq_length,
+                #         start_frame=current_start_frame)
 
             # Step 3.4: update the start and end frame indices
             current_start_frame += current_num_frames
@@ -520,12 +599,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             training_batch.dmd_latent_vis_dict["min_num_frames"] = torch.tensor(
                 min_num_frames, dtype=torch.float32, device=self.device)
 
-        # Clean up caches
-        assert self.kv_cache1 is not None
-        assert self.crossattn_cache is not None
-        self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache)
-
-        return final_output if gradient_mask is not None else pred_image_or_video
+        if gradient_mask is not None:
+            return final_output, exit_timestep
+        else:
+            return pred_image_or_video, exit_timestep
 
     def _initialize_simulation_caches(
         self,
@@ -606,22 +683,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
         return kv_cache, crossattn_cache
 
-    def _reset_simulation_caches(self, kv_cache: list[dict[str, Any]],
-                                 crossattn_cache: list[dict[str, Any]]) -> None:
-        """Reset KV cache and cross-attention cache to clean state."""
-        if kv_cache is not None:
-            for cache_dict in kv_cache:
-                cache_dict["global_end_index"].fill_(0)
-                cache_dict["local_end_index"].fill_(0)
-                cache_dict["k"].zero_()
-                cache_dict["v"].zero_()
-
-        if crossattn_cache is not None:
-            for cache_dict in crossattn_cache:
-                cache_dict["is_init"] = False
-                cache_dict["k"].zero_()
-                cache_dict["v"].zero_()
-
     def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
         batch = next(self.train_loader_iter, None)  # type: ignore
         if batch is None:
@@ -709,9 +770,37 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                         setattr(batch_gen, key, copy.deepcopy(value))
 
                 generator_loss, gen_log_dict = self.generator_loss(batch_gen)
+                logger.info("train_transformer_2: %s", self.train_transformer_2)
+
                 with set_forward_context(current_timestep=batch_gen.timesteps,
                                          attn_metadata=batch_gen.attn_metadata):
                     (generator_loss / gradient_accumulation_steps).backward()
+
+                # Ensure that only one of the two transformers have received gradients
+                if self.train_transformer_2:
+                    # Assert that all gradients are None for transformer
+                    assert all(p.grad is None
+                               for p in self.transformer.parameters())
+                    grad_sum = 0
+                    for n, p in self.transformer_2.named_parameters():
+                        if p.grad is not None:
+                            grad_sum += p.grad.sum().item()
+                        else:
+                            raise ValueError(
+                                "Transformer 2 param %s has no gradient", n)
+                    assert grad_sum != 0, "Transformer 2 did not receive gradients"
+                else:
+                    assert all(p.grad is None
+                               for p in self.transformer_2.parameters())
+                    grad_sum = 0
+                    for n, p in self.transformer.named_parameters():
+                        if p.grad is not None:
+                            grad_sum += p.grad.sum().item()
+                        else:
+                            raise ValueError(
+                                "Transformer param %s has no gradient", n)
+                    assert grad_sum != 0, "Transformer did not receive gradients"
+
                 total_generator_loss += generator_loss.detach().item()
                 generator_log_dict.update(gen_log_dict)
                 # Store visualization data from generator training
@@ -753,6 +842,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
         logger.debug("Training critic at step %s", self.current_trainstep)
         self.fake_score_optimizer.zero_grad()
+        if self.fake_score_transformer_2 is not None:
+            self.fake_score_optimizer_2.zero_grad()
         total_critic_loss = 0.0
         critic_log_dict = {}
 
@@ -777,6 +868,26 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             with set_forward_context(current_timestep=batch_critic.timesteps,
                                      attn_metadata=batch_critic.attn_metadata):
                 (critic_loss / gradient_accumulation_steps).backward()
+
+            # if self.train_fake_score_transformer_2:
+            #     assert all(p.grad is None for p in self.fake_score_transformer.parameters())
+            #     grad_sum = 0
+            #     for n, p in self.fake_score_transformer_2.named_parameters():
+            #         if p.grad is not None:
+            #             grad_sum += p.grad.sum().item()
+            #         else:
+            #             raise ValueError("Fake score transformer 2 param %s has no gradient", n)
+            #     assert grad_sum != 0, "Fake score transformer 2 did not receive gradients"
+            # else:
+            #     assert all(p.grad is None for p in self.fake_score_transformer_2.parameters())
+            #     grad_sum = 0
+            #     for n, p in self.fake_score_transformer.named_parameters():
+            #         if p.grad is not None:
+            #             grad_sum += p.grad.sum().item()
+            #         else:
+            #             raise ValueError("Fake score transformer param %s has no gradient", n)
+            #     assert grad_sum != 0, "Fake score transformer did not receive gradients"
+
             total_critic_loss += critic_loss.detach().item()
             critic_log_dict.update(crit_log_dict)
             # Store visualization data from critic training
@@ -868,7 +979,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                     video = video.permute(0, 2, 1, 3, 4)
                     video = (video * 255).numpy().astype(np.uint8)
                     wandb_loss_dict[f"dmd_{latent_key}"] = wandb.Video(
-                        video, fps=24, format="mp4")
+                        video, fps=16, format="mp4")
                     del video, latents
 
         # Process critic predictions
@@ -910,7 +1021,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                     video = video.permute(0, 2, 1, 3, 4)
                     video = (video * 255).numpy().astype(np.uint8)
                     wandb_loss_dict[f"critic_{latent_key}"] = wandb.Video(
-                        video, fps=24, format="mp4")
+                        video, fps=16, format="mp4")
                     del video, latents
 
         # Log metadata
@@ -976,6 +1087,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         self._log_training_info()
         self._log_validation(self.transformer, self.training_args,
                              self.init_steps)
+        gc.collect()
+        torch.cuda.empty_cache()
 
         progress_bar = tqdm(
             range(0, self.training_args.max_train_steps),
@@ -1184,6 +1297,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
                 self._log_validation(self.transformer, self.training_args, step)
+                gc.collect()
+                torch.cuda.empty_cache()
 
         wandb.finish()
 
@@ -1223,3 +1338,4 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
         if get_sp_group():
             cleanup_dist_env_and_memory()
+

fastvideo/training/training_pipeline.py

@@ -107,6 +107,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         assert self.seed is not None, "seed must be set"
         set_random_seed(self.seed)
         self.transformer.train()
+        self.transformer.requires_grad_(True)
         if training_args.enable_gradient_checkpointing_type is not None:
             self.transformer = apply_activation_checkpointing(
                 self.transformer,
@@ -235,14 +236,15 @@ class TrainingPipeline(LoRAPipeline, ABC):
         for param in model.parameters():
             param.requires_grad = True
         model.train()
-        optimizer.zero_grad()
+        # optimizer.zero_grad()
 
     def _disable_training(self, model: torch.nn.Module,
                           optimizer: torch.optim.Optimizer) -> None:
         """Disable training mode and gradients for the specified model."""
         for param in model.parameters():
             param.requires_grad = False
-        optimizer.zero_grad(set_to_none=True)
+        model.eval()
+        # optimizer.zero_grad(set_to_none=True)
 
     def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
         batch = next(self.train_loader_iter, None)  # type: ignore
@@ -357,6 +359,20 @@ class TrainingPipeline(LoRAPipeline, ABC):
         #     pass
 
         indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
+        # timestep = self.noise_scheduler.timesteps[indices].to(device=device)
+
+        # if timestep < self.training_args.boundary_ratio * self.noise_scheduler.config.num_train_timesteps:
+        #     self.train_transformer_2 = True
+        # else:
+        #     self.train_transformer_2 = False
+
+        # # Broadcast the decision to all processes
+        # decision = torch.tensor(1.0 if self.train_transformer_2 else 0.0,
+        #                         device=self.device)
+        # dist.broadcast(decision, src=0)
+        # dist.broadcast(timestep, src=0)
+        # self.train_transformer_2 = decision.item() == 1.0
+
         return self.noise_scheduler.timesteps[indices].to(device=device)
 
     def _build_attention_metadata(

fastvideo/training/training_utils.py

@@ -1152,14 +1152,14 @@ def custom_to_hf_state_dict(
 
     return new_state_dict
 
-
+# More generalized version of shift_timestep
 def shift_timestep(timestep: torch.Tensor, shift: float,
-                   num_train_timestep: float) -> torch.Tensor:
+                   min_timestep: int = 0, max_timestep: int = 1000) -> torch.Tensor:
     if shift == 1:
         return timestep
-    t = timestep / num_train_timestep
+    t = (timestep - min_timestep) / (max_timestep - min_timestep)
     denominator = 1 + (shift - 1) * t
-    return num_train_timestep * (shift * t / denominator)
+    return min_timestep + (max_timestep - min_timestep) * (shift * t / denominator)
 
 
 # coding=utf-8

fastvideo/training/wan_distillation_pipeline.py

@@ -42,7 +42,7 @@ class WanDistillationPipeline(DistillationPipeline):
             training_args.model_path,
             args=args_copy,  # type: ignore
             inference_mode=True,
-            loaded_modules={"transformer": self.get_module("transformer")},
+            loaded_modules={"transformer": self.get_module("transformer"), "transformer_2": self.get_module("transformer_2")},
             tp_size=training_args.tp_size,
             sp_size=training_args.sp_size,
             num_gpus=training_args.num_gpus,

fastvideo/training/wan_training_pipeline.py

@@ -42,6 +42,7 @@ class WanTrainingPipeline(TrainingPipeline):
             inference_mode=True,
             loaded_modules={
                 "transformer": self.get_module("transformer"),
+                "transformer_2": self.get_module("transformer_2"),
             },
             tp_size=training_args.tp_size,
             sp_size=training_args.sp_size,