From 55074bda76c100df9a5bdb39320e6dd4c117f0f5 Mon Sep 17 00:00:00 2001 From: Yongqi Chen <144848849+BrianChen1129@users.noreply.github.com> Date: Wed, 18 Jun 2025 00:13:06 -0400 Subject: [PATCH] [CI] Add STA-inference/VSA-training test (#527) Co-authored-by: Will Lin --- .github/workflows/pr-test.yml | 60 +++++++- docker/Dockerfile.python3.12 | 4 +- fastvideo/v1/STA_configuration.py | 2 + .../tests/inference/STA/test_STA_inference.py | 52 +++++++ .../VSA/reference_wandb_summary_VSA.json | 1 + .../training/VSA/test_training_loss_VSA.py | 136 ++++++++++++++++++ .../reference_wandb_summary.json | 0 .../{ => Vanilla}/test_training_loss.py | 6 +- 8 files changed, 254 insertions(+), 7 deletions(-) create mode 100644 fastvideo/v1/tests/inference/STA/test_STA_inference.py create mode 100644 fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json create mode 100644 fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py rename fastvideo/v1/tests/training/{ => Vanilla}/reference_wandb_summary.json (100%) rename fastvideo/v1/tests/training/{ => Vanilla}/test_training_loss.py (95%) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index addff1089..c9a9532e6 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -44,6 +44,16 @@ on: required: false default: false type: boolean + run_training_test_VSA: + description: "Run training-test-VSA" + required: false + default: false + type: boolean + run_inference_test_STA: + description: "Run inference-test-STA" + required: false + default: false + type: boolean run_nightly_test: description: "Run nightly-test" required: false @@ -70,6 +80,8 @@ jobs: vae-test: ${{ steps.filter.outputs.vae-test }} transformer-test: ${{ steps.filter.outputs.transformer-test }} training-test: ${{ steps.filter.outputs.training-test }} + training-test-VSA: ${{ steps.filter.outputs.training-test-VSA }} + inference-test-STA: ${{ steps.filter.outputs.inference-test-STA }} steps: - uses: actions/checkout@v4 - uses: dorny/paths-filter@v3 @@ -92,6 +104,10 @@ jobs: - 'fastvideo/v1/attention/**' training-test: - 'fastvideo/v1/**' + training-test-VSA: + - 'fastvideo/v1/**' + inference-test-STA: + - 'fastvideo/v1/**' encoder-test: needs: change-filter @@ -187,7 +203,47 @@ jobs: volume_size: 100 disk_size: 100 image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" - test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training -srP" + test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/Vanilla -srP" + timeout_minutes: 30 + secrets: + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }} + WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} + + training-test-VSA: + needs: change-filter + if: >- + (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || + (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test_VSA == 'true') + uses: ./.github/workflows/runpod-test.yml + with: + job_id: "training-test-VSA" + gpu_type: "NVIDIA H100 NVL" + gpu_count: 1 + volume_size: 100 + disk_size: 100 + image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" + test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/VSA -srP" + timeout_minutes: 30 + secrets: + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }} + WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} + + inference-test-STA: + needs: change-filter + if: >- + (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || + (github.event_name == 'workflow_dispatch' && github.event.inputs.run_inference_test_STA == 'true') + uses: ./.github/workflows/runpod-test.yml + with: + job_id: "inference-test-STA" + gpu_type: "NVIDIA H100 NVL" + gpu_count: 1 + volume_size: 100 + disk_size: 100 + image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" + test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/inference/STA -srP" timeout_minutes: 30 secrets: RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} @@ -204,7 +260,7 @@ jobs: volume_size: 100 disk_size: 100 image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" - test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs" + test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs" timeout_minutes: 30 secrets: RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} diff --git a/docker/Dockerfile.python3.12 b/docker/Dockerfile.python3.12 index ccd3480dc..002e8ed3b 100644 --- a/docker/Dockerfile.python3.12 +++ b/docker/Dockerfile.python3.12 @@ -58,14 +58,14 @@ RUN source $HOME/.local/bin/env && \ # Install STA (Sliding Tile Attention) RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ - cd csrc/sliding_tile_attention/ && \ + cd csrc/attn/ && \ git submodule update --init --recursive && \ python setup_sta.py install # Install VSA RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ - cd csrc/sliding_tile_attention/ && \ + cd csrc/attn/ && \ git submodule update --init --recursive && \ python setup_vsa.py install diff --git a/fastvideo/v1/STA_configuration.py b/fastvideo/v1/STA_configuration.py index df9d21ac7..7c1f1ea4f 100644 --- a/fastvideo/v1/STA_configuration.py +++ b/fastvideo/v1/STA_configuration.py @@ -360,6 +360,8 @@ def dict_to_3d_list(mask_strategy: Optional[Dict[str, List[int]]], return result for key, value in mask_strategy.items(): t, layer_idx, h = map(int, key.split('_')) + if t >= t_max: + continue result[t][layer_idx][h] = value return result diff --git a/fastvideo/v1/tests/inference/STA/test_STA_inference.py b/fastvideo/v1/tests/inference/STA/test_STA_inference.py new file mode 100644 index 000000000..d980ef9a4 --- /dev/null +++ b/fastvideo/v1/tests/inference/STA/test_STA_inference.py @@ -0,0 +1,52 @@ +import os +import sys +import subprocess +from pathlib import Path +import pytest + +NUM_NODES = "1" +NUM_GPUS_PER_NODE = "1" + +# Set environment variables +os.environ["FASTVIDEO_ATTENTION_CONFIG"] = "assets/mask_strategy_wan.json" +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN" + +def test_inference(): + """Test the inference functionality""" + # Create command as in wan_14B-STA.sh + cmd = [ + "fastvideo", "generate", + "--model-path", "Wan-AI/Wan2.1-T2V-14B-Diffusers", + "--sp-size", "1", + "--tp-size", "1", + "--num-gpus", "1", + "--height", "768", + "--width", "1280", + "--num-frames", "69", + "--num-inference-steps", "2", + "--fps", "16", + "--guidance-scale", "5.0", + "--flow-shift", "5.0", + "--prompt", "A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic.", + "--negative-prompt", "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards", + "--seed", "1024", + "--output-path", "outputs_video/STA_1024/", + ] + + # Run the command + subprocess.run(cmd, check=True) + + # Verify output directory exists + output_dir = Path("outputs_video/STA_1024/") + assert output_dir.exists(), f"Output directory {output_dir} does not exist" + + # Verify that video files were generated + video_files = list(output_dir.glob("*.mp4")) + assert len(video_files) > 0, "No video files were generated" + + # Verify the video file properties + for video_file in video_files: + assert video_file.stat().st_size > 0, f"Video file {video_file} is empty" + +if __name__ == "__main__": + test_inference() diff --git a/fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json b/fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json new file mode 100644 index 000000000..0b682ad1b --- /dev/null +++ b/fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json @@ -0,0 +1 @@ +{"grad_norm":0.478515625,"_runtime":95.727033597,"_wandb":{"runtime":95},"_step":5,"validation_videos_50_steps":{"videos":[{"_type":"video-file","sha256":"42a1c311521a9d460db788713be1cbf2db767494e02619b43be5bf3eed8381d8","size":158632,"path":"media/videos/validation_videos_50_steps_0_42a1c311521a9d460db7.mp4"},{"path":"media/videos/validation_videos_50_steps_0_818505095b4b5e8b7f51.mp4","_type":"video-file","sha256":"818505095b4b5e8b7f511012d45f04d151ce3344bc058fc0f3225a414a851e4a","size":147825},{"sha256":"fc334ba9ed5e66c8527ee3b408e3be2d76167fef03588bf2840f4a0792f2fe34","size":136933,"path":"media/videos/validation_videos_50_steps_0_fc334ba9ed5e66c8527e.mp4","_type":"video-file"},{"size":201797,"path":"media/videos/validation_videos_50_steps_0_ccd98f6f907635d266a7.mp4","_type":"video-file","sha256":"ccd98f6f907635d266a74783688e7ecf1dac752d79d72d69eab9ef0e3f7413eb"},{"_type":"video-file","sha256":"ca79f40a0aed38f676f12779b349ce40e9e3fb7f36c578f49a20854c70508fb4","size":147114,"path":"media/videos/validation_videos_50_steps_0_ca79f40a0aed38f676f1.mp4"},{"size":175104,"path":"media/videos/validation_videos_50_steps_0_32c9b33ff920c17e5881.mp4","_type":"video-file","sha256":"32c9b33ff920c17e588133d7a27aa400ff3dc529b01ed4f16ac4d6bb2afa0f00"},{"sha256":"2cf520bfb93401c914e93c87ef791c2f12a4e043b95dfdc98115c930e11dfe67","size":139655,"path":"media/videos/validation_videos_50_steps_0_2cf520bfb93401c914e9.mp4","_type":"video-file"},{"_type":"video-file","sha256":"1d73aba17ce582c7aef4af4d64079e3e9d3df205634eff453446bdaf2340b214","size":149028,"path":"media/videos/validation_videos_50_steps_0_1d73aba17ce582c7aef4.mp4"}],"captions":false,"_type":"videos","count":8},"train_loss":0.08922439813613892,"_timestamp":1.750202051751466e+09,"avg_step_time":0.7536672964692116,"step_time":0.4742048177868128,"learning_rate":1e-05,"vsa_sparsity":0.05} \ No newline at end of file diff --git a/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py b/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py new file mode 100644 index 000000000..13b3413e4 --- /dev/null +++ b/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py @@ -0,0 +1,136 @@ +import os +import sys +import subprocess +from pathlib import Path +import json +from huggingface_hub import snapshot_download + +# Import the training pipeline +sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) +from fastvideo.v1.training.wan_training_pipeline import main +from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.v1.utils import FlexibleArgumentParser + +wandb_name = "test_training_loss_VSA" +reference_wandb_summary_file = "fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json" + +NUM_NODES = "1" +NUM_GPUS_PER_NODE = "1" + +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" + +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) + + # Set the arguments as they are in finetune_v1_test.sh + args = parser.parse_args([ + "--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", + "--cache_dir", "/home/.cache", + "--data_path", "data/mini_dataset_i2v_VSA/combined_parquet_dataset", + "--validation_prompt_dir", "data/mini_dataset_i2v_VSA/validation_parquet_dataset", + "--train_batch_size", "1", + "--num_latent_t", "4", + "--num_gpus", "1", + "--sp_size", "1", + "--tp_size", "1", + "--hsdp_replicate_dim", "1", + "--hsdp_shard_dim", "1", + "--train_sp_batch_size", "1", + "--dataloader_num_workers", "4", + "--gradient_accumulation_steps", "1", + "--max_train_steps", "5", + "--learning_rate", "1e-5", + "--mixed_precision", "bf16", + "--checkpointing_steps", "30", + "--validation_steps", "10", + "--validation_sampling_steps", "50", + "--log_validation", + "--checkpoints_total_limit", "3", + "--allow_tf32", + "--ema_start_step", "0", + "--cfg", "0.0", + "--output_dir", "data/wan_finetune_test_VSA", + "--tracker_project_name", "wan_finetune_ci_VSA", + "--wandb_run_name", wandb_name, + "--num_height", "384", + "--num_width", "512", + "--num_frames", "13", + "--flow_shift", "3", + "--validation_guidance_scale", "1.0", + "--num_euler_timesteps", "50", + "--weight_decay", "0.01", + "--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) + +def test_distributed_training(): + """Test the distributed training setup""" + os.environ["WANDB_MODE"] = "online" + + data_dir = Path("data/mini_dataset_i2v_VSA") + + if not data_dir.exists(): + print(f"Downloading test dataset to {data_dir}...") + snapshot_download( + repo_id="BrianChen1129/mini_dataset_i2v_VSA", + local_dir=str(data_dir), + repo_type="dataset", + local_dir_use_symlinks=False + ) + + # Get the current file path + current_file = Path(__file__).resolve() + + # Run torchrun command + cmd = [ + "torchrun", + "--nnodes", NUM_NODES, + "--nproc_per_node", NUM_GPUS_PER_NODE, + str(current_file) + ] + + process = subprocess.run(cmd, check=True) + + summary_file = 'wandb/latest-run/files/wandb-summary.json' + + reference_wandb_summary = json.load(open(reference_wandb_summary_file)) + wandb_summary = json.load(open(summary_file)) + + fields_and_thresholds = { + 'avg_step_time': 1.0, + 'grad_norm': 0.1, + 'step_time': 0.5, + 'train_loss': 0.001 + } + + failures = [] + for field, threshold in fields_and_thresholds.items(): + ref_value = reference_wandb_summary[field] + current_value = wandb_summary[field] + diff = abs(ref_value - current_value) + print(f"INFO: {field}, diff: {diff}, threshold: {threshold}, reference: {ref_value}, current: {current_value}") + if diff > threshold: + failures.append(f"FAILED: {field} difference {diff} exceeds threshold of {threshold} (reference: {ref_value}, current: {current_value})") + + if failures: + raise AssertionError("\n".join(failures)) + +if __name__ == "__main__": + if os.environ.get("LOCAL_RANK") is not None: + # We're being run by torchrun + run_worker() + else: + # We're being run directly + test_distributed_training() diff --git a/fastvideo/v1/tests/training/reference_wandb_summary.json b/fastvideo/v1/tests/training/Vanilla/reference_wandb_summary.json similarity index 100% rename from fastvideo/v1/tests/training/reference_wandb_summary.json rename to fastvideo/v1/tests/training/Vanilla/reference_wandb_summary.json diff --git a/fastvideo/v1/tests/training/test_training_loss.py b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py similarity index 95% rename from fastvideo/v1/tests/training/test_training_loss.py rename to fastvideo/v1/tests/training/Vanilla/test_training_loss.py index efab00595..1b75725af 100644 --- a/fastvideo/v1/tests/training/test_training_loss.py +++ b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py @@ -18,7 +18,7 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.v1.utils import FlexibleArgumentParser wandb_name = "test_training_loss" -reference_wandb_summary_file = "fastvideo/v1/tests/training/reference_wandb_summary.json" +reference_wandb_summary_file = "fastvideo/v1/tests/training/Vanilla/reference_wandb_summary.json" NUM_NODES = "1" NUM_GPUS_PER_NODE = "4" @@ -81,7 +81,7 @@ def run_worker(): def test_distributed_training(): """Test the distributed training setup""" - os.environ["WANDB_MODE"] = "offline" + os.environ["WANDB_MODE"] = "online" data_dir = Path("data/crush-smol_parq") @@ -107,7 +107,7 @@ def test_distributed_training(): process = subprocess.run(cmd, check=True) - summary_file = "fastvideo/v1/tests/training/reference_wandb_summary.json" + summary_file = 'wandb/latest-run/files/wandb-summary.json' reference_wandb_summary = json.load(open(reference_wandb_summary_file)) wandb_summary = json.load(open(summary_file))