diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 87fe89ecd..2986fc139 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -264,7 +264,7 @@ jobs: with: job_id: "training-test-VSA" gpu_type: "NVIDIA H100 NVL" - gpu_count: 1 + gpu_count: 2 volume_size: 100 disk_size: 100 image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" @@ -284,7 +284,7 @@ jobs: with: job_id: "inference-test-STA" gpu_type: "NVIDIA H100 NVL" - gpu_count: 1 + gpu_count: 2 volume_size: 100 disk_size: 100 image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" diff --git a/fastvideo/v1/layers/layernorm.py b/fastvideo/v1/layers/layernorm.py index 66032ecd0..85c8f44d7 100644 --- a/fastvideo/v1/layers/layernorm.py +++ b/fastvideo/v1/layers/layernorm.py @@ -38,6 +38,12 @@ class RMSNorm(CustomOp): if self.has_weight: self.weight = nn.Parameter(self.weight) + # if we do fully_shard(model.layer_norm), and we call layer_form.forward_native(input) instead of layer_norm(input), + # we need to call model.layer_norm.register_fsdp_forward_method(model, "forward_native") to make sure fsdp2 hooks are triggered + # for mixed precision and cpu offloading + + # the even better way might be fully_shard(model.layer_norm, mp_policy=, cpu_offloading=), and call model.layer_norm(input). everything should work out of the box + # because fsdp2 hooks will be triggered with model.layer_norm.__call__ def forward_native( self, x: torch.Tensor, diff --git a/fastvideo/v1/tests/inference/STA/test_STA_inference.py b/fastvideo/v1/tests/inference/STA/test_STA_inference.py index d980ef9a4..b08378f3d 100644 --- a/fastvideo/v1/tests/inference/STA/test_STA_inference.py +++ b/fastvideo/v1/tests/inference/STA/test_STA_inference.py @@ -5,7 +5,7 @@ from pathlib import Path import pytest NUM_NODES = "1" -NUM_GPUS_PER_NODE = "1" +NUM_GPUS_PER_NODE = "2" # Set environment variables os.environ["FASTVIDEO_ATTENTION_CONFIG"] = "assets/mask_strategy_wan.json" @@ -17,9 +17,9 @@ def test_inference(): cmd = [ "fastvideo", "generate", "--model-path", "Wan-AI/Wan2.1-T2V-14B-Diffusers", - "--sp-size", "1", - "--tp-size", "1", - "--num-gpus", "1", + "--sp-size", "2", + "--tp-size", "2", + "--num-gpus", "2", "--height", "768", "--width", "1280", "--num-frames", "69", diff --git a/fastvideo/v1/tests/modal/pr_test.py b/fastvideo/v1/tests/modal/pr_test.py index 2c542b7d7..587e4ba78 100644 --- a/fastvideo/v1/tests/modal/pr_test.py +++ b/fastvideo/v1/tests/modal/pr_test.py @@ -69,11 +69,11 @@ def run_ssim_tests(): def run_training_tests(): run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/Vanilla -srP") -@app.function(gpu="H100:1", image=image, timeout=1800, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})]) +@app.function(gpu="H100:2", image=image, timeout=1800, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})]) def run_training_tests_VSA(): run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/VSA -srP") -@app.function(gpu="H100:1", image=image, timeout=1800) +@app.function(gpu="H100:2", image=image, timeout=1800) def run_inference_tests_STA(): run_test("pytest ./fastvideo/v1/tests/inference/STA -srP") diff --git a/fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json b/fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json index 0b682ad1b..2136ee17c 100644 --- a/fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json +++ b/fastvideo/v1/tests/training/VSA/reference_wandb_summary_VSA.json @@ -1 +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 +{"step_time":0.6983645600266755,"_wandb":{"runtime":107},"grad_norm":0.50390625,"avg_step_time":1.002151239803061,"_step":5,"validation_videos_50_steps":{"captions":false,"_type":"videos","count":8,"videos":[{"size":159131,"path":"media/videos/validation_videos_50_steps_0_dc447599dbe48350e9c9.mp4","_type":"video-file","sha256":"dc447599dbe48350e9c920f4971e1786bde580dd20d52b4aa147ae8d3dc564d6"},{"_type":"video-file","sha256":"4e283876ddfbf5a2cb6f8aca07a39f832b5f806fbefde8854bc73ed904ff20ee","size":160315,"path":"media/videos/validation_videos_50_steps_0_4e283876ddfbf5a2cb6f.mp4"},{"size":135225,"path":"media/videos/validation_videos_50_steps_0_78185c41e1935306e93c.mp4","_type":"video-file","sha256":"78185c41e1935306e93c2d416ee40b31d038abffb029cb5bfb11c2a634eb2fcf"},{"_type":"video-file","sha256":"27e9819d002d3f63c8918bbdc5bf2857b5effe0caf5d5b8374b9c590fc6432eb","size":197873,"path":"media/videos/validation_videos_50_steps_0_27e9819d002d3f63c891.mp4"},{"_type":"video-file","sha256":"46fe548e86144ca60a9396fcd15e8788d3a05c066c6819ccdd6a9041feaaec8f","size":170601,"path":"media/videos/validation_videos_50_steps_0_46fe548e86144ca60a93.mp4"},{"sha256":"91ec338774bec870b9c5c81be4330a8f0f2535124cb372e881c3703e6b65ed77","size":164462,"path":"media/videos/validation_videos_50_steps_0_91ec338774bec870b9c5.mp4","_type":"video-file"},{"_type":"video-file","sha256":"ee4e811080a619215fd7541b39203dfb69c2db4cdadbfdd7a234a2cff684f6f3","size":139435,"path":"media/videos/validation_videos_50_steps_0_ee4e811080a619215fd7.mp4"},{"sha256":"22e31e048ba5e5b9d6587d7306904d6edc6c618b8f77c3ceb05ba2d602309274","size":147072,"path":"media/videos/validation_videos_50_steps_0_22e31e048ba5e5b9d658.mp4","_type":"video-file"}]},"_timestamp":1.75118195270901e+09,"vsa_sparsity":0.05,"learning_rate":1e-05,"train_loss":0.19960195198655128,"_runtime":107.325113071} \ 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 index 7a21c5986..5e98f7247 100644 --- a/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py +++ b/fastvideo/v1/tests/training/VSA/test_training_loss_VSA.py @@ -15,7 +15,7 @@ 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" +NUM_GPUS_PER_NODE = "2" os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" @@ -35,14 +35,14 @@ def run_worker(): "--validation_preprocessed_path", "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", + "--num_gpus", "2", + "--sp_size", "2", + "--tp_size", "2", "--hsdp_replicate_dim", "1", - "--hsdp_shard_dim", "1", + "--hsdp_shard_dim", "2", "--train_sp_batch_size", "1", "--dataloader_num_workers", "4", - "--gradient_accumulation_steps", "1", + "--gradient_accumulation_steps", "2", "--max_train_steps", "5", "--learning_rate", "1e-5", "--mixed_precision", "bf16", @@ -110,7 +110,7 @@ def test_distributed_training(): fields_and_thresholds = { 'avg_step_time': 1.0, 'grad_norm': 0.1, - 'step_time': 0.5, + 'step_time': 1.0, 'train_loss': 0.001 } diff --git a/fastvideo/v1/tests/training/Vanilla/test_training_loss.py b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py index ac2b1032a..b49ba603b 100644 --- a/fastvideo/v1/tests/training/Vanilla/test_training_loss.py +++ b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py @@ -121,9 +121,9 @@ def test_distributed_training(): wandb_summary = json.load(open(summary_file)) fields_and_thresholds = { - 'avg_step_time': 1.0, + 'avg_step_time': 6.0, 'grad_norm': 0.3, - 'step_time': 0.5, + 'step_time': 6.0, 'train_loss': 0.0025 } diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index b9a79ad31..72b749167 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -245,9 +245,8 @@ class TrainingPipeline(ComposedPipelineBase, ABC): current_vsa_sparsity = training_batch.current_vsa_sparsity if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": - dit_seq_shape = [ - latents.shape[2] // patch_size[0], + latents.shape[2] * self.sp_world_size // patch_size[0], latents.shape[3] // patch_size[1], latents.shape[4] // patch_size[2] ]