update sfwan ssim

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA.slurm

@@ -91,7 +91,7 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 1e-6
   --mixed_precision "bf16"
   --weight_only_checkpointing_steps 1000
   --training_state_checkpointing_steps 1000

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm

@@ -91,7 +91,7 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 1e-6
   --mixed_precision "bf16"
   --weight_only_checkpointing_steps 1000
   --training_state_checkpointing_steps 1000

fastvideo/fastvideo_args.py

@@ -160,6 +160,8 @@ class FastVideoArgs:
         "vae": True,
     })
     override_transformer_cls_name: str | None = None
+    init_weights_from_safetensors: str = ""  # path to safetensors file for initial weight loading
+    init_weights_from_safetensors_2: str = ""  # path to safetensors file for initial weight loading for transformer_2
 
     # # DMD parameters
     # dmd_denoising_steps: List[int] | None = field(default=None)
@@ -411,6 +413,14 @@ class FastVideoArgs:
             default=FastVideoArgs.override_transformer_cls_name,
             help="Override transformer cls name",
         )
+        parser.add_argument(
+            "--init-weights-from-safetensors",
+            type=str,
+            help="Path to safetensors file for initial weight loading")
+        parser.add_argument(
+            "--init-weights-from-safetensors-2",
+            type=str,
+            help="Path to safetensors file for initial weight loading")
         # Add pipeline configuration arguments
         PipelineConfig.add_cli_args(parser)
 
@@ -659,8 +669,6 @@ class TrainingArgs(FastVideoArgs):
     output_dir: str = ""
     checkpoints_total_limit: int = 0
     resume_from_checkpoint: str = ""  # specify the checkpoint folder to resume from
-    init_weights_from_safetensors: str = ""  # path to safetensors file for initial weight loading
-    init_weights_from_safetensors_2: str = ""  # path to safetensors file for initial weight loading for transformer_2
 
     # optimizer & scheduler
     num_train_epochs: int = 0
@@ -917,14 +925,6 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--resume-from-checkpoint",
                             type=str,
                             help="Path to checkpoint to resume from")
-        parser.add_argument(
-            "--init-weights-from-safetensors",
-            type=str,
-            help="Path to safetensors file for initial weight loading")
-        parser.add_argument(
-            "--init-weights-from-safetensors-2",
-            type=str,
-            help="Path to safetensors file for initial weight loading")
         parser.add_argument("--logging-dir",
                             type=str,
                             help="Directory for logging")

fastvideo/models/loader/component_loader.py

@@ -438,7 +438,6 @@ class TransformerLoader(ComponentLoader):
         # Check if we should use custom initialization weights
         custom_weights_path = getattr(fastvideo_args, 'init_weights_from_safetensors', None)
         use_custom_weights = (custom_weights_path and os.path.exists(custom_weights_path) and 
-                            fastvideo_args.training_mode and 
                             not hasattr(fastvideo_args, '_loading_teacher_critic_model'))
 
         if use_custom_weights:

scripts/distill/v1_distill_dmd_wan.sh

@@ -30,7 +30,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
     --dataloader_num_workers 0 \
     --gradient_accumulation_steps 8 \
     --max_train_steps 30000 \
-    --learning_rate 1e-5 \
+    --learning_rate 2e-6 \
     --mixed_precision "bf16" \
     --checkpointing_steps 400 \
     --validation_steps 100 \

scripts/distill/v1_distill_dmd_wan_VSA.sh

@@ -31,7 +31,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
     --dataloader_num_workers 0 \
     --gradient_accumulation_steps 8 \
     --max_train_steps 30000 \
-    --learning_rate 1e-5 \
+    --learning_rate 2e-6 \
     --mixed_precision "bf16" \
     --checkpointing_steps 400 \
     --validation_steps 100 \

scripts/finetune/finetune_v1.sh

@@ -28,7 +28,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --dataloader_num_workers 10\
     --gradient_accumulation_steps=1 \
     --max_train_steps=5000 \
-    --learning_rate=1e-5\
+    --learning_rate=1e-6\
     --mixed_precision="bf16"\
     --checkpointing_steps=6000 \
     --validation_steps 200\

scripts/finetune/finetune_v1_VSA.sh

@@ -34,7 +34,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
     --dataloader_num_workers 4 \
     --gradient_accumulation_steps 8 \
     --max_train_steps 30000 \
-    --learning_rate 1e-5 \
+    --learning_rate 1e-6 \
     --mixed_precision "bf16" \
     --checkpointing_steps 6000 \
     --validation_steps 100 \