From 8abfe234af3b9aed882b93d71cc27d2f7bbae6dc Mon Sep 17 00:00:00 2001
From: Peiyuan Zhang <a1286225768@gmail.com>
Date: Mon, 29 Sep 2025 20:56:25 +0000
Subject: [PATCH] update

---
 .../wan_t2v_1.3B/crush_smol/finetune_t2v.sh   | 12 +--
 .../crush_smol/finetune_t2v_nvfp4.sh          | 94 +++++++++++++++++++
 fastvideo/fastvideo_args.py                   |  4 +
 3 files changed, 104 insertions(+), 6 deletions(-)
 create mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4.sh

diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
index 3dd2f77ec..a4fb457c9 100644
--- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
+++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
@@ -6,9 +6,9 @@ export TOKENIZERS_PARALLELISM=false
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-NUM_GPUS=4
+DATA_DIR="data/Wan-Syn_77x448x832_600k/train/"
+VALIDATION_DATASET_FILE="data/Wan-Syn_77x448x832_600k/validation_6.json"
+NUM_GPUS=6
 # export CUDA_VISIBLE_DEVICES=4,5
 
 
@@ -19,7 +19,7 @@ training_args=(
   --max_train_steps 5000
   --train_batch_size 1
   --train_sp_batch_size 1
-  --gradient_accumulation_steps 8
+  --gradient_accumulation_steps 2
   --num_latent_t 20
   --num_height 480
   --num_width 832
@@ -30,7 +30,7 @@ training_args=(
 # Parallel arguments
 parallel_args=(
   --num_gpus $NUM_GPUS 
-  --sp_size $NUM_GPUS 
+  --sp_size 1
   --tp_size 1
   --hsdp_replicate_dim 1
   --hsdp_shard_dim $NUM_GPUS
@@ -52,7 +52,7 @@ dataset_args=(
 validation_args=(
   --log_validation 
   --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 200
+  --validation_steps 100
   --validation_sampling_steps "50" 
   --validation_guidance_scale "3.0"
 )
diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4.sh
new file mode 100644
index 000000000..65ff6c4df
--- /dev/null
+++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4.sh
@@ -0,0 +1,94 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/Wan-Syn_77x448x832_600k/train/"
+VALIDATION_DATASET_FILE="data/Wan-Syn_77x448x832_600k/validation_6.json"
+NUM_GPUS=6
+# export CUDA_VISIBLE_DEVICES=4,5
+
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_t2v_finetune"
+  --output_dir "checkpoints/wan_t2v_finetune"
+  --max_train_steps 5000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 2
+  --num_latent_t 20
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --enable_gradient_checkpointing_type "full"
+  --generator_4bit_linear True
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS 
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim $NUM_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path $DATA_DIR
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation 
+  --validation_dataset_file $VALIDATION_DATASET_FILE
+  --validation_steps 100
+  --validation_sampling_steps "50" 
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 5e-5
+  --mixed_precision "bf16"
+  --checkpointing_steps 2000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+  --enable_gradient_checkpointing_type "full"
+  # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500"
+)
+
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py
index 418b72b35..fb822792f 100644
--- a/fastvideo/fastvideo_args.py
+++ b/fastvideo/fastvideo_args.py
@@ -1150,6 +1150,10 @@ class TrainingArgs(FastVideoArgs):
             "--generator-4bit-attn",
             action=StoreBoolean,
             help="Whether to use 4bit attention in generator in DMD distillation")
+        parser.add_argument(
+            "--generator-4bit-linear",
+            action=StoreBoolean,
+            help="Whether to use nvfp4 in generator in DMD distillation")
 
         # Self-forcing specific arguments
         parser.add_argument(