Add context forcing, ode_init to sf training

examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm

@@ -47,7 +47,7 @@ FAKE_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2
 # 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="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/vidprom_16k_text_embed"
-DATA_DIR_2="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-distilled-121"
+DATA_DIR_2="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-121"
 # 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"
@@ -59,7 +59,7 @@ VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/
 training_args=(
   --tracker_project_name SFhy1.5_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
   --output_dir "/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_3333_1e-5_1e-5_cfg6_corrected_scheduler"
-  --wandb_run_name "self_forcing_3333_1e-5_1e-5_cfg3.5_corrected_scheduler"
+  --wandb_run_name "self_forcing_3333_context_forcing"
   --max_train_steps 4000
   --train_batch_size 1
   --train_sp_batch_size 1
@@ -73,7 +73,8 @@ training_args=(
   --num_frames 81
   --num_frame_per_block 3  # Frame generation block size for self-forcing
   # --resume-from-checkpoint "/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_1333/checkpoint-300"
-  --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_distilled_ode_init_3333/checkpoint-2400/transformer/diffusion_pytorch_model.safetensors
+  --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_worldplay_df_init_3333/checkpoint-2000/transformer/diffusion_pytorch_model.safetensors
+  # --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_distilled_ode_init_3333/checkpoint-2400/transformer/diffusion_pytorch_model.safetensors
 )
 
 parallel_args=(
@@ -93,7 +94,7 @@ model_args=(
 
 dataset_args=(
   --data_path "$DATA_DIR"
-  --data_path_2 "$DATA_DIR_2"
+  # --data_path_2 "$DATA_DIR_2"
   --dataloader_num_workers 4
 )
 
@@ -108,7 +109,7 @@ validation_args=(
 )
 
 optimizer_args=(
-  --learning_rate 1e-6
+  --learning_rate 1e-5
   --mixed_precision "bf16"
   --training_state_checkpointing_steps 100
   --weight_only_checkpointing_steps 500
@@ -122,7 +123,7 @@ miscellaneous_args=(
   --checkpoints_total_limit 3
   --training_cfg_rate 0.0
   --dit_precision "fp32"
-  --flow_shift 7
+  --flow_shift 5
   --seed 1000
   --use_ema True
   --ema_decay 0.99
@@ -133,11 +134,12 @@ dmd_args=(
   --warp_denoising_step
   --dmd_denoising_steps '1000,875,750,625,500,375,250,125'
   # --dmd_denoising_steps '1000,760,520,280'
+  # --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.5
-  --fake_score_learning_rate 2e-7
+  --fake_score_learning_rate 5e-6
   --fake_score_betas '0.0,0.999'
 )
 
@@ -146,6 +148,7 @@ self_forcing_args=(
   --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)
+  # --use-context-forcing True
 )
 
 srun torchrun \