[misc]: Refactor separate training pipeline runner

fastvideo/tests/training/VSA/test_training_loss_VSA.py

@@ -13,9 +13,11 @@
 
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-from fastvideo.training.wan_training_pipeline import main
+
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
+from fastvideo.training.runner import main
+from fastvideo.utils import build_parser
 
 wandb_name = "test_training_loss_VSA"
 reference_wandb_summary_file = "fastvideo/tests/training/VSA/h100_reference_wandb_summary_VSA.json"
@@ -27,12 +29,12 @@
 def run_worker():
     """Worker function that will be run on each GPU"""
     # Create and populate args
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
+    parser = build_parser()
 
     # Set the arguments as they are in finetune_v1_test.sh
     args = parser.parse_args([
+        "--pipeline_class", "WanTrainingPipeline",
+        "--pipeline_module", "fastvideo.training.wan_training_pipeline",
         "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--inference_mode", "False",
         "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--data_path",
         "data/mini_dataset_i2v_VSA/combined_parquet_dataset", "--validation_dataset_file",
@@ -49,7 +51,7 @@ def run_worker():
         "--dit_precision", "fp32", "--max_grad_norm", "1.0", "--VSA_decay_rate", "0.01", "--VSA_decay_interval_steps",
         "1", "--VSA_sparsity", "0.9"
     ])
-
+    
     # Call the main training function
     main(args)
 

fastvideo/tests/training/Vanilla/mfu_calculation.py

@@ -14,10 +14,11 @@
 from huggingface_hub import snapshot_download
 from fastvideo.utils import logger
 # Import the training pipeline
-from fastvideo.training.wan_training_pipeline import main
+
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
-from fastvideo.training.wan_training_pipeline import WanTrainingPipeline
+from fastvideo.training.runner import main
+from fastvideo.utils import build_parser
 
 MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_PATH = "data/crush-smol_processed_t2v/training_dataset/worker_1/worker_0/"
@@ -38,12 +39,12 @@
 def run_worker():
     """Worker function that will be run on each GPU"""
     # Create and populate args
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
+    parser = build_parser()
 
     # Set the arguments as they are in finetune_t2v.sh
     args = parser.parse_args([
+        "--pipeline_class", "WanTrainingPipeline",
+        "--pipeline_module", "fastvideo.training.wan_training_pipeline",
         "--model_path",
         MODEL_PATH,
         "--inference_mode",
@@ -122,11 +123,9 @@ def run_worker():
         "--hsdp_shard_dim",
         "1"
     ])
+
     # Call the main training function
-    pipeline = WanTrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
+    main(args)
 
 
 def test_distributed_training():

fastvideo/tests/training/Vanilla/test_training_loss.py

@@ -11,10 +11,12 @@
 from fastvideo.utils import logger
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-from fastvideo.training.wan_training_pipeline import main
+
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
-from fastvideo.training.wan_training_pipeline import WanTrainingPipeline
+from fastvideo.utils import FlexibleArgumentParser
+from fastvideo.training.runner import main
+from fastvideo.utils import build_parser
 
 wandb_name = "test_training_loss"
 a40_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/a40_reference_wandb_summary.json"
@@ -28,12 +30,12 @@
 def run_worker():
     """Worker function that will be run on each GPU"""
     # Create and populate args
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
+    parser = build_parser()
 
     # Set the arguments as they are in finetune_v1_test.sh
     args = parser.parse_args([
+        "--pipeline_class", "WanTrainingPipeline",
+        "--pipeline_module", "fastvideo.training.wan_training_pipeline",
         "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--inference_mode", "False",
         "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--data_path",
         "data/crush-smol_processed_t2v/combined_parquet_dataset", "--validation_dataset_file",
@@ -49,11 +51,9 @@ def run_worker():
         "--num_euler_timesteps", "50", "--multi_phased_distill_schedule", "4000-1", "--weight_decay", "0.01",
         "--not_apply_cfg_solver", "--dit_precision", "fp32", "--max_grad_norm", "1.0"
     ])
+
     # Call the main training function
-    pipeline = WanTrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
+    main(args)
 
 
 def test_distributed_training():

fastvideo/tests/training/distill/test_distill_dmd.py

@@ -11,10 +11,12 @@
 from fastvideo.utils import logger
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-from fastvideo.training.wan_training_pipeline import main
+
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
-from fastvideo.training.wan_distillation_pipeline import WanDistillationPipeline
+from fastvideo.utils import FlexibleArgumentParser
+from fastvideo.training.runner import main
+from fastvideo.utils import build_parser
 
 wandb_name = "test_distill_dmd"
 
@@ -25,12 +27,12 @@
 def run_worker():
     """Worker function that will be run on each GPU"""
     # Create and populate args
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
+    parser = build_parser()
 
     # Set the arguments as they are in finetune_v1_test.sh
     args = parser.parse_args([
+        "--pipeline_class", "WanDistillationPipeline",
+        "--pipeline_module", "fastvideo.training.wan_distillation_pipeline",
         "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--inference_mode", "False",
         "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--real_score_model_path",
         "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--fake_score_model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
@@ -48,11 +50,9 @@ def run_worker():
         "--min_timestep_ratio", "0.02", "--max_timestep_ratio", "0.98", "--real_score_guidance_scale", "3.5",
         "--enable_gradient_checkpointing_type", "full"
     ])
+
     # Call the main training function
-    pipeline = WanDistillationPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
+    main(args)
 
 
 def test_distributed_training():

fastvideo/tests/training/self-forcing/test_self_forcing.py

@@ -11,9 +11,12 @@
 from fastvideo.utils import logger
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-from fastvideo.training.wan_self_forcing_distillation_pipeline import WanSelfForcingDistillationPipeline
+
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.utils import FlexibleArgumentParser
+from fastvideo.utils import FlexibleArgumentParser
+from fastvideo.training.runner import main
+from fastvideo.utils import build_parser
 
 wandb_name = "test_self_forcing_distill"
 
@@ -24,12 +27,12 @@
 def run_worker():
     """Worker function that will be run on each GPU"""
     # Create and populate args
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
+    parser = build_parser()
 
     # Set the arguments based on the distill_dmd_t2v_1.3B.sh script
     args = parser.parse_args([
+        "--pipeline_class", "WanSelfForcingDistillationPipeline",
+        "--pipeline_module", "fastvideo.training.wan_self_forcing_distillation_pipeline",
         "--model_path",
         "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",
         "--inference_mode",
@@ -146,12 +149,9 @@ def run_worker():
         "--ema_start_step",
         "100",
     ])
-
+    
     # Call the main training function
-    pipeline = WanSelfForcingDistillationPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Self-forcing distillation training pipeline done")
+    main(args)
 
 
 def test_distributed_training():

fastvideo/tests/training/test_runner.py

@@ -0,0 +1,65 @@
+from unittest.mock import patch, MagicMock
+from fastvideo.training.runner import main
+import argparse
+from fastvideo.utils import build_parser
+
+def test_runner_invokes_correct_class():
+    with patch('fastvideo.training.runner.importlib.import_module') as mock_import_module:
+        mock_module = MagicMock()
+        mock_import_module.return_value = mock_module
+
+        mock_pipeline_class = MagicMock()
+        setattr(mock_module, "WanTrainingPipeline", mock_pipeline_class)
+
+        mock_pipeline_instance = MagicMock()
+        mock_pipeline_class.from_pretrained.return_value = mock_pipeline_instance
+
+        args = argparse.Namespace(
+            pipeline_module="fastvideo.training.wan_training_pipeline",
+            pipeline_class="WanTrainingPipeline",
+            pretrained_model_name_or_path="test_model"
+        )
+
+        main(args)
+
+        mock_import_module.assert_called_once_with("fastvideo.training.wan_training_pipeline")
+        mock_pipeline_class.from_pretrained.assert_called_once_with("test_model", args=args)
+        mock_pipeline_instance.train.assert_called_once()
+
+
+def test_runner_cli_argument_parsing():
+
+    base_required_args = [
+        "--data_path", "test_path",
+        "--dataloader_num_workers", "0",
+        "--num_height", "448",
+        "--num_width", "832",
+        "--num_frames", "61",
+        "--train_batch_size", "1",
+        "--num_latent_t", "16",
+        "--output_dir", "outputs/test",
+        "--learning_rate", "2e-6",
+    ]
+
+    # Test with underscores
+    parser1 = build_parser()
+    args1 = parser1.parse_args([
+        "--pipeline_class", "WanDistillationPipeline",
+        "--pipeline_module", "fastvideo.training.wan_distillation_pipeline",
+        "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        *base_required_args,
+    ])
+    assert args1.pipeline_class == "WanDistillationPipeline"
+    assert args1.pipeline_module == "fastvideo.training.wan_distillation_pipeline"
+
+    # Test with dashes
+    parser2 = build_parser()
+    args2 = parser2.parse_args([
+        "--pipeline-class", "WanDistillationPipeline",
+        "--pipeline-module", "fastvideo.training.wan_distillation_pipeline",
+        "--pretrained-model-name-or-path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        *base_required_args,
+    ])
+    assert args2.pipeline_class == "WanDistillationPipeline"
+    assert args2.pipeline_module == "fastvideo.training.wan_distillation_pipeline"
+

fastvideo/training/cosmos2_5_training_pipeline.py

@@ -143,6 +143,11 @@ def main(args) -> None:
 
 
 if __name__ == "__main__":
+    logger.warning("\n"
+                   "================================================================================\n"
+                   "[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
+                   "Please use `fastvideo/training/runner.py` instead.\n"
+                   "================================================================================")
     from fastvideo.fastvideo_args import TrainingArgs
     from fastvideo.utils import FlexibleArgumentParser
 

fastvideo/training/ltx2_training_pipeline.py

@@ -484,6 +484,11 @@ def main(args) -> None:
 
 
 if __name__ == "__main__":
+    logger.warning("\n"
+                   "================================================================================\n"
+                   "[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
+                   "Please use `fastvideo/training/runner.py` instead.\n"
+                   "================================================================================")
     argv = sys.argv
     from fastvideo.fastvideo_args import TrainingArgs
     from fastvideo.utils import FlexibleArgumentParser

fastvideo/training/matrixgame2_ar_diffusion_pipeline.py

@@ -438,6 +438,11 @@ def main(args) -> None:
 
 
 if __name__ == "__main__":
+    logger.warning("\n"
+                   "================================================================================\n"
+                   "[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
+                   "Please use `fastvideo/training/runner.py` instead.\n"
+                   "================================================================================")
     argv = sys.argv
     from fastvideo.fastvideo_args import TrainingArgs
     from fastvideo.utils import FlexibleArgumentParser

fastvideo/training/matrixgame2_ode_causal_pipeline.py

@@ -404,6 +404,11 @@ def main(args) -> None:
 
 
 if __name__ == "__main__":
+    logger.warning("\n"
+                   "================================================================================\n"
+                   "[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
+                   "Please use `fastvideo/training/runner.py` instead.\n"
+                   "================================================================================")
     argv = sys.argv
     from fastvideo.fastvideo_args import TrainingArgs
     from fastvideo.utils import FlexibleArgumentParser

fastvideo/training/matrixgame2_self_forcing_distillation_pipeline.py

@@ -887,6 +887,11 @@ def main(args) -> None:
 
 
 if __name__ == "__main__":
+    logger.warning("\n"
+                   "================================================================================\n"
+                   "[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
+                   "Please use `fastvideo/training/runner.py` instead.\n"
+                   "================================================================================")
     argv = sys.argv
     from fastvideo.fastvideo_args import TrainingArgs
     from fastvideo.utils import FlexibleArgumentParser

fastvideo/training/matrixgame2_training_pipeline.py

@@ -205,6 +205,11 @@ def main(args) -> None:
 
 
 if __name__ == "__main__":
+    logger.warning("\n"
+                   "================================================================================\n"
+                   "[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
+                   "Please use `fastvideo/training/runner.py` instead.\n"
+                   "================================================================================")
     argv = sys.argv
     from fastvideo.fastvideo_args import TrainingArgs
     from fastvideo.utils import FlexibleArgumentParser

fastvideo/training/ode_causal_pipeline.py

@@ -319,6 +319,11 @@ def main(args) -> None:
 
 
 if __name__ == "__main__":
+    logger.warning("\n"
+                   "================================================================================\n"
+                   "[DEPRECATED]: Direct execution of this pipeline is deprecated!\n"
+                   "Please use `fastvideo/training/runner.py` instead.\n"
+                   "================================================================================")
     argv = sys.argv
     from fastvideo.fastvideo_args import TrainingArgs
     from fastvideo.utils import FlexibleArgumentParser

fastvideo/training/runner.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+import importlib
+
+from fastvideo.logger import init_logger
+from fastvideo.utils import build_parser
+
+logger = init_logger(__name__)
+
+
+def main(args) -> None:
+    logger.info("Starting training pipeline %s...", args.pipeline_class)
+
+    module = importlib.import_module(args.pipeline_module)
+    pipeline_class = getattr(module, args.pipeline_class)
+
+    pipeline = pipeline_class.from_pretrained(args.pretrained_model_name_or_path, args=args)
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    parser = build_parser()
+
+    args = parser.parse_args()
+    args.dit_cpu_offload = False
+
+    main(args)