[Feat][Training] Rename weight conversion function and update gradient checkpoint in scripts (#589)

Author: BrianChen1129

fastvideo/v1/fastvideo_args.py

@@ -413,8 +413,7 @@ class TrainingArgs(FastVideoArgs):
     lr_scheduler: str = "constant"
     lr_warmup_steps: int = 0
     max_grad_norm: float = 0.0
-    gradient_checkpointing: bool = False
-    gradient_checkpointing_type: str = "full"
+    enable_gradient_checkpointing_type: Optional[str] = None
     selective_checkpointing: float = 0.0
     allow_tf32: bool = False
     mixed_precision: str = ""
@@ -613,13 +612,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parser.add_argument("--max-grad-norm",
                             type=float,
                             help="Maximum gradient norm")
-        parser.add_argument("--gradient-checkpointing",
-                            action=StoreBoolean,
-                            help="Whether to use gradient checkpointing")
-        parser.add_argument("--gradient-checkpointing-type",
+        parser.add_argument("--enable-gradient-checkpointing-type",
                             type=str,
                             choices=["full", "ops", "block_skip"],
-                            default="full",
+                            default=None,
                             help="Gradient checkpointing type")
         parser.add_argument("--selective-checkpointing",
                             type=float,

fastvideo/v1/training/training_pipeline.py

@@ -84,11 +84,11 @@ def initialize_training_pipeline(self, training_args: TrainingArgs):
 
         self.transformer.requires_grad_(True)
         self.transformer.train()
-
-        if training_args.gradient_checkpointing:
+        if training_args.enable_gradient_checkpointing_type is not None:
             self.transformer = apply_activation_checkpointing(
                 self.transformer,
-                checkpointing_type=training_args.gradient_checkpointing_type)
+                checkpointing_type=training_args.
+                enable_gradient_checkpointing_type)
 
         noise_scheduler = self.modules["scheduler"]
         params_to_optimize = self.transformer.parameters()

fastvideo/v1/training/training_utils.py

@@ -162,8 +162,8 @@ def save_checkpoint(transformer,
                     weight_path,
                     local_main_process_only=False)
 
-        # Convert training format to diffusers format and save
-        diffusers_state_dict = convert_training_to_diffusers_format(
+        # Convert fastvideo custom format to diffusers format and save
+        diffusers_state_dict = convert_custom_format_to_diffusers_format(
             cpu_state, transformer)
         save_file(diffusers_state_dict, weight_path)
 
@@ -487,10 +487,10 @@ def _has_foreach_support(tensors: List[torch.Tensor],
         t is None or type(t) in [torch.Tensor] for t in tensors)
 
 
-def convert_training_to_diffusers_format(state_dict: Dict[str, Any],
-                                         transformer) -> Dict[str, Any]:
+def convert_custom_format_to_diffusers_format(state_dict: Dict[str, Any],
+                                              transformer) -> Dict[str, Any]:
     """
-    Convert training format state dict to diffusers format using reverse_param_names_mapping.
+    Convert fastvideo custom format state dict to diffusers format using reverse_param_names_mapping.
     
     Args:
         state_dict: State dict in training format

scripts/finetune/finetune_v1.sh

@@ -48,4 +48,5 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --weight_decay 0.01 \
     --not_apply_cfg_solver \
     --dit_precision "fp32" \
-    --max_grad_norm 1.0
+    --max_grad_norm 1.0 \
+    --enable_gradient_checkpointing_type "full"

scripts/finetune/finetune_v1_VSA.sh

@@ -58,5 +58,6 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \
     --VSA_decay_sparsity 0.9 \
     --VSA_decay_rate 0.03 \
     --VSA_decay_interval_steps 30 \
-    --VSA_val_sparsity 0.9
+    --VSA_val_sparsity 0.9 \
+    --enable_gradient_checkpointing_type "full"
 # --resume_from_checkpoint "$CHECKPOINT_PATH"