[Feature][Training] Add cfg rate for dataset loader (#556)

Author: BrianChen1129

fastvideo/v1/dataset/parquet_dataset_map_style.py

@@ -200,10 +200,6 @@ def __init__(
         super().__init__()
         self.path = path
         self.cfg_rate = cfg_rate
-        if cfg_rate > 0.0:
-            raise ValueError(
-                "cfg_rate > 0.0 is not supported for now because it will trigger bug when num_data_workers > 0"
-            )
         logger.info("Initializing LatentsParquetMapStyleDataset with path: %s",
                     path)
         self.parquet_files, self.lengths = get_parquet_files_and_length(path)
@@ -247,7 +243,7 @@ def get_validation_negative_prompt(
                                               [self.lengths[0]])
 
         all_latents_list, all_embs_list, all_masks_list, caption_text_list = collate_latents_embs_masks(
-            [row_dict], self.text_padding_length, self.keys)
+            [row_dict], self.text_padding_length, self.keys, cfg_rate=0.0)
         all_latents, all_embs, all_masks, caption_text = all_latents_list[
             0], all_embs_list[0], all_masks_list[0], caption_text_list[0]
         # add batch dimension
@@ -268,7 +264,7 @@ def __getitems__(self, indices: List[int]):
         ]
 
         all_latents, all_embs, all_masks, caption_text = collate_latents_embs_masks(
-            rows, self.text_padding_length, self.keys)
+            rows, self.text_padding_length, self.keys, self.cfg_rate)
         return all_latents, all_embs, all_masks, caption_text
 
     def __len__(self):

fastvideo/v1/dataset/preprocessing_datasets.py

@@ -461,7 +461,7 @@ def _init_stages(self, args, transform, transform_topcrop,
         self.text_encoding_stage = TextEncodingStage(
             tokenizer=tokenizer,
             text_max_length=args.text_max_length,
-            cfg_rate=args.cfg)
+            cfg_rate=args.training_cfg_rate)
 
     def _load_raw_data(self) -> List[Dict]:
         """Load raw data from JSON files."""

fastvideo/v1/dataset/utils.py

@@ -1,3 +1,4 @@
+import random
 from typing import Any, Dict, List
 
 import numpy as np
@@ -20,7 +21,7 @@ def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
         return t[:padding_length], torch.ones(padding_length)
 
 
-def get_torch_tensors_from_row_dict(row_dict, keys) -> Dict[str, Any]:
+def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> Dict[str, Any]:
     """
     Get the latents and prompts from a row dictionary.
     """
@@ -42,7 +43,10 @@ def get_torch_tensors_from_row_dict(row_dict, keys) -> Dict[str, Any]:
             bytes = row_dict[f"{key}_bytes"]
 
         # TODO (peiyuan): read precision
-        data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy()
+        if key == 'text_embedding' and random.random() < cfg_rate:
+            data = np.zeros((512, 4096), dtype=np.float32)
+        else:
+            data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy()
         data = torch.from_numpy(data)
         if len(data.shape) == 3:
             B, L, D = data.shape
@@ -53,8 +57,11 @@ def get_torch_tensors_from_row_dict(row_dict, keys) -> Dict[str, Any]:
 
 
 def collate_latents_embs_masks(
-        batch_to_process, text_padding_length,
-        keys) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str]]:
+        batch_to_process,
+        text_padding_length,
+        keys,
+        cfg_rate=0.0
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str]]:
     # Initialize tensors to hold padded embeddings and masks
     all_latents = []
     all_embs = []
@@ -63,7 +70,7 @@ def collate_latents_embs_masks(
     # Process each row individually
     for i, row in enumerate(batch_to_process):
         # Get tensors from row
-        data = get_torch_tensors_from_row_dict(row, keys)
+        data = get_torch_tensors_from_row_dict(row, keys, cfg_rate)
         latents, emb = data["vae_latent"], data["text_embedding"]
 
         padded_emb, mask = pad(emb, text_padding_length)

fastvideo/v1/fastvideo_args.py

@@ -384,7 +384,7 @@ class TrainingArgs(FastVideoArgs):
     # diffusion setting
     ema_decay: float = 0.0
     ema_start_step: int = 0
-    cfg: float = 0.0
+    training_cfg_rate: float = 0.0
     precondition_outputs: bool = False
 
     # validation & logs
@@ -528,7 +528,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
                             type=int,
                             default=0,
                             help="Step to start EMA")
-        parser.add_argument("--cfg",
+        parser.add_argument("--training-cfg-rate",
                             type=float,
                             help="Classifier-free guidance scale")
         parser.add_argument(

fastvideo/v1/pipelines/preprocess/v1_preprocess.py

@@ -91,7 +91,7 @@ def main(args) -> None:
                         type=str,
                         default="google/t5-v1_1-xxl")
     parser.add_argument("--cache_dir", type=str, default="./cache_dir")
-    parser.add_argument("--cfg", type=float, default=0.0)
+    parser.add_argument("--training_cfg_rate", type=float, default=0.0)
     parser.add_argument(
         "--output_dir",
         type=str,

fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py

@@ -122,7 +122,7 @@ def run_training():
         "--checkpoints_total_limit", "3",
         "--allow_tf32",
         "--ema_start_step", "0",
-        "--cfg", "0.0",
+        "--training_cfg_rate", "0.0",
         "--output_dir", LOCAL_OUTPUT_DIR,
         "--tracker_project_name", "wan_finetune_overfit_ci",
         "--num_height", "480",

fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py

@@ -54,7 +54,7 @@ def run_worker():
         "--checkpoints_total_limit", "3",
         "--allow_tf32",
         "--ema_start_step", "0",
-        "--cfg", "0.0",
+        "--training_cfg_rate", "0.0",
         "--output_dir", "data/wan_finetune_test_VSA",
         "--tracker_project_name", "wan_finetune_ci_VSA",
         "--wandb_run_name", wandb_name,

fastvideo/v1/tests/training/Vanilla/test_training_loss.py

@@ -59,7 +59,7 @@ def run_worker():
         "--checkpoints_total_limit", "3",
         "--allow_tf32",
         "--ema_start_step", "0",
-        "--cfg", "0.0",
+        "--training_cfg_rate", "0.0",
         "--output_dir", "data/wan_finetune_test",
         "--tracker_project_name", "wan_finetune_ci",
         "--wandb_run_name", wandb_name,

fastvideo/v1/training/training_pipeline.py

@@ -105,6 +105,7 @@ def initialize_training_pipeline(self, training_args: TrainingArgs):
             training_args.data_path,
             training_args.train_batch_size,
             num_data_workers=training_args.dataloader_num_workers,
+            cfg_rate=training_args.training_cfg_rate,
             drop_last=True,
             text_padding_length=training_args.pipeline_config.
             text_encoder_configs[0].arch_config.
@@ -534,9 +535,9 @@ def _log_validation(self, transformer, training_args, global_step) -> None:
             training_args.validation_preprocessed_path,
             batch_size=1,
             num_data_workers=0,
+            cfg_rate=0.0,
             drop_last=False,
-            drop_first_row=sampling_param.negative_prompt is not None,
-            cfg_rate=training_args.cfg)
+            drop_first_row=sampling_param.negative_prompt is not None)
         if sampling_param.negative_prompt:
             _, negative_prompt_embeds, negative_prompt_attention_mask, _ = validation_dataset.get_validation_negative_prompt(
             )

scripts/finetune/finetune_v1.sh

@@ -36,7 +36,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --checkpoints_total_limit 3\
     --allow_tf32\
     --ema_start_step 0\
-    --cfg 0.0\
+    --training_cfg_rate 0.0\
     --output_dir="$DATA_DIR/outputs/wan_finetune"\
     --tracker_project_name wan_finetune \
     --num_height 480 \