diff --git a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA.slurm b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA.slurm
index d64a988e6..aad328234 100644
--- a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA.slurm
+++ b/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
diff --git a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm
index de97f63a3..8befc289c 100644
--- a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm
+++ b/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
diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py
index 1f43b232a..5d31abb3f 100644
--- a/fastvideo/fastvideo_args.py
+++ b/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")
diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py
index 44eb9c06a..5f0186c78 100644
--- a/fastvideo/models/loader/component_loader.py
+++ b/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:
diff --git a/scripts/distill/v1_distill_dmd_wan.sh b/scripts/distill/v1_distill_dmd_wan.sh
index 962931aec..c0306df39 100644
--- a/scripts/distill/v1_distill_dmd_wan.sh
+++ b/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 \
diff --git a/scripts/distill/v1_distill_dmd_wan_VSA.sh b/scripts/distill/v1_distill_dmd_wan_VSA.sh
index e23c2c952..8f81749c3 100644
--- a/scripts/distill/v1_distill_dmd_wan_VSA.sh
+++ b/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 \
diff --git a/scripts/finetune/finetune_v1.sh b/scripts/finetune/finetune_v1.sh
index d25ed27ab..a94d18a23 100644
--- a/scripts/finetune/finetune_v1.sh
+++ b/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\
diff --git a/scripts/finetune/finetune_v1_VSA.sh b/scripts/finetune/finetune_v1_VSA.sh
index 7f444497f..d62770d86 100644
--- a/scripts/finetune/finetune_v1_VSA.sh
+++ b/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 \