From 137f61f2fee9dafb30a67bf71fdd4c0f5da8aaf4 Mon Sep 17 00:00:00 2001 From: Kevin Lin <42618777+kevin314@users.noreply.github.com> Date: Fri, 11 Apr 2025 03:40:01 -0500 Subject: [PATCH] Port tests to v1 (#333) --- .github/workflows/pr-test.yml | 122 +++++++- fastvideo/v1/tests/conftest.py | 20 +- .../v1/tests/encoders/test_clip_encoder.py | 6 +- .../v1/tests/encoders/test_llama_encoder.py | 6 +- .../v1/tests/old_tests/test_hunyuan_vae.py | 179 ------------ .../v1/tests/old_tests/test_hunyuanvideo.py | 260 ------------------ .../tests/old_tests/test_hunyuanvideo_load.py | 212 -------------- fastvideo/v1/tests/ssim/compute_ssim.py | 2 + .../tests/ssim/test_inference_similarity.py | 2 + .../tests/transformers/test_hunyuanvideo.py | 137 +++++++++ .../transformers/test_hunyuanvideo_load.py | 119 ++++++++ fastvideo/v1/tests/vaes/test_hunyuan_vae.py | 78 ++++++ 12 files changed, 474 insertions(+), 669 deletions(-) delete mode 100644 fastvideo/v1/tests/old_tests/test_hunyuan_vae.py delete mode 100644 fastvideo/v1/tests/old_tests/test_hunyuanvideo.py delete mode 100644 fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py create mode 100644 fastvideo/v1/tests/transformers/test_hunyuanvideo.py create mode 100644 fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py create mode 100644 fastvideo/v1/tests/vaes/test_hunyuan_vae.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 7928f2865..6aaa0abf9 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -19,6 +19,16 @@ on: required: false default: false type: boolean + run_vae_test: + description: "Run vae-test" + required: false + default: false + type: boolean + run_transformer_test: + description: "Run transformer-test" + required: false + default: false + type: boolean run_ssim_test: description: "Run ssim-test" required: false @@ -39,6 +49,8 @@ jobs: if: ${{ github.event.pull_request.draft == false || github.event_name == 'workflow_dispatch' }} outputs: encoder-test: ${{ steps.filter.outputs.encoder-test }} + vae-test: ${{ steps.filter.outputs.vae-test }} + transformer-test: ${{ steps.filter.outputs.transformer-test }} steps: - uses: actions/checkout@v4 - uses: dorny/paths-filter@v3 @@ -49,6 +61,14 @@ jobs: - 'fastvideo/v1/models/encoders/**' - 'fastvideo/v1/models/loaders/**' - 'fastvideo/v1/tests/encoders/**' + vae-test: + - 'fastvideo/v1/models/vaes/**' + - 'fastvideo/v1/models/loaders/**' + - 'fastvideo/v1/tests/vaes/**' + transformer-test: + - 'fastvideo/v1/models/dits/**' + - 'fastvideo/v1/models/loaders/**' + - 'fastvideo/v1/tests/transformers/**' encoder-test: needs: change-filter @@ -99,6 +119,104 @@ jobs: JOB_ID: "encoder-test" run: python .github/scripts/runpod_cleanup.py + vae-test: + needs: change-filter + if: >- + (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.vae-test == 'true') || + (github.event_name == 'workflow_dispatch' && github.event.inputs.run_vae_test == 'true') + runs-on: ubuntu-latest + environment: runpod-runners + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Set up SSH key + run: | + mkdir -p ~/.ssh + echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa + chmod 600 ~/.ssh/id_rsa + ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub + + - name: Install dependencies + run: pip install requests + + - name: Run tests on RunPod + env: + JOB_ID: "vae-test" + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + GITHUB_RUN_ID: ${{ github.run_id }} + timeout-minutes: 30 + run: >- + python .github/scripts/runpod_api.py + --gpu-type "NVIDIA A40" + --gpu-count 1 + --volume-size 100 + --test-command "pip install -e .[test] && + pip install flash-attn==2.7.0.post2 --no-build-isolation && + pytest ./fastvideo/v1/tests/vaes -s" + + - name: Terminate RunPod Instances + if: ${{ always() }} + env: + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + GITHUB_RUN_ID: ${{ github.run_id }} + JOB_ID: "vae-test" + run: python .github/scripts/runpod_cleanup.py + + transformer-test: + needs: change-filter + if: >- + (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.transformer-test == 'true') || + (github.event_name == 'workflow_dispatch' && github.event.inputs.run_transformer_test == 'true') + runs-on: ubuntu-latest + environment: runpod-runners + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Set up SSH key + run: | + mkdir -p ~/.ssh + echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa + chmod 600 ~/.ssh/id_rsa + ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub + + - name: Install dependencies + run: pip install requests + + - name: Run tests on RunPod + env: + JOB_ID: "transformer-test" + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + GITHUB_RUN_ID: ${{ github.run_id }} + timeout-minutes: 30 + run: >- + python .github/scripts/runpod_api.py + --gpu-type "NVIDIA L40S" + --gpu-count 1 + --volume-size 100 + --test-command "pip install -e .[test] && + pip install flash-attn==2.7.0.post2 --no-build-isolation && + pytest ./fastvideo/v1/tests/transformers -s" + + - name: Terminate RunPod Instances + if: ${{ always() }} + env: + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + GITHUB_RUN_ID: ${{ github.run_id }} + JOB_ID: "transformer-test" + run: python .github/scripts/runpod_cleanup.py + ssim-test: needs: change-filter if: >- @@ -150,7 +268,7 @@ jobs: run: python .github/scripts/runpod_cleanup.py runpod-cleanup: - needs: [encoder-test, ssim-test] # Add other jobs to this list as you create them + needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }} runs-on: ubuntu-latest steps: @@ -167,7 +285,7 @@ jobs: - name: Cleanup all RunPod instances env: - JOB_IDS: '["encoder-test", "ssim-test"]' # JSON array of job IDs + JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test"]' # JSON array of job IDs RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} GITHUB_RUN_ID: ${{ github.run_id }} run: python .github/scripts/runpod_cleanup.py diff --git a/fastvideo/v1/tests/conftest.py b/fastvideo/v1/tests/conftest.py index eaac2668e..075afa886 100644 --- a/fastvideo/v1/tests/conftest.py +++ b/fastvideo/v1/tests/conftest.py @@ -1,9 +1,12 @@ -import pytest -import torch.distributed as dist +# SPDX-License-Identifier: Apache-2.0 -from fastvideo.v1.distributed import (destroy_model_parallel, - init_distributed_environment, - initialize_model_parallel) +import pytest +import torch +import numpy as np + +from fastvideo.v1.distributed import (init_distributed_environment, + initialize_model_parallel, + cleanup_dist_env_and_memory) @pytest.fixture(scope="function") @@ -13,6 +16,9 @@ def distributed_setup(): This ensures proper cleanup even if tests fail. """ + torch.manual_seed(42) + np.random.seed(42) + init_distributed_environment(world_size=1, rank=0, distributed_init_method="env://", @@ -23,6 +29,4 @@ def distributed_setup(): backend="nccl") yield - if dist.is_initialized(): - destroy_model_parallel() - dist.destroy_process_group() + cleanup_dist_env_and_memory() diff --git a/fastvideo/v1/tests/encoders/test_clip_encoder.py b/fastvideo/v1/tests/encoders/test_clip_encoder.py index 0c8bf1be0..3c1df3296 100644 --- a/fastvideo/v1/tests/encoders/test_clip_encoder.py +++ b/fastvideo/v1/tests/encoders/test_clip_encoder.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + # TODO: check if correct import os @@ -19,10 +21,6 @@ logger = init_logger(__name__) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29503" -# Set fixed random seed for reproducibility -torch.manual_seed(42) -np.random.seed(42) - BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo" MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join( diff --git a/fastvideo/v1/tests/encoders/test_llama_encoder.py b/fastvideo/v1/tests/encoders/test_llama_encoder.py index f812ab0b7..84a0eff4e 100644 --- a/fastvideo/v1/tests/encoders/test_llama_encoder.py +++ b/fastvideo/v1/tests/encoders/test_llama_encoder.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + import os import numpy as np @@ -18,10 +20,6 @@ logger = init_logger(__name__) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29503" -# Set fixed random seed for reproducibility -torch.manual_seed(42) -np.random.seed(42) - BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo" MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join( diff --git a/fastvideo/v1/tests/old_tests/test_hunyuan_vae.py b/fastvideo/v1/tests/old_tests/test_hunyuan_vae.py deleted file mode 100644 index 99ac4151a..000000000 --- a/fastvideo/v1/tests/old_tests/test_hunyuan_vae.py +++ /dev/null @@ -1,179 +0,0 @@ -import argparse -import json -import os - -import numpy as np -import torch -import torch.nn as nn -from diffusers import AutoencoderKLHunyuanVideo as DiffusersHunyuanVAE -from safetensors.torch import load_file - -from fastvideo.v1.logger import init_logger -from fastvideo.v1.models.vaes.hunyuanvae import ( - AutoencoderKLHunyuanVideo as MyHunyuanVAE) - -logger = init_logger(__name__) - - -def initialize_identical_weights(model1, model2, seed=42): - """Initialize both models with identical weights using a fixed seed for reproducibility.""" - # Get all parameters from both models - params1 = dict(model1.named_parameters()) - params2 = dict(model2.named_parameters()) - - # Initialize each layer with identical values - with torch.no_grad(): - # Initialize weights - for name1, param1 in params1.items(): - if 'weight' in name1: - # Set seed before each weight initialization - torch.manual_seed(seed) - nn.init.normal_(param1, mean=0.0, std=0.05) - - for name2, param2 in params2.items(): - if 'weight' in name2: - # Reset seed to get same initialization - torch.manual_seed(seed) - nn.init.normal_(param2, mean=0.0, std=0.05) - - # Initialize biases - for name1, param1 in params1.items(): - if 'bias' in name1: - torch.manual_seed(seed) - nn.init.normal_(param1, mean=0.0, std=0.05) - param1.data = param1.data.to(torch.bfloat16) - - for name2, param2 in params2.items(): - if 'bias' in name2: - torch.manual_seed(seed) - nn.init.normal_(param2, mean=0.0, std=0.05) - param2.data = param2.data.to(torch.bfloat16) - - logger.info("Both models initialized with identical weights in bfloat16") - return model1, model2 - - -def setup_args(): - parser = argparse.ArgumentParser(description='HunyuanVAE Test') - parser.add_argument('--in-channels', - type=int, - default=4, - help='Number of input channels') - parser.add_argument('--out-channels', - type=int, - default=4, - help='Number of output channels') - parser.add_argument('--latent-channels', - type=int, - default=4, - help='Number of latent channels') - return parser.parse_args() - - -def test_hunyuan_vae(): - args = setup_args() - - # Set fixed random seed for reproducibility - torch.manual_seed(42) - np.random.seed(42) - - # Model parameters - in_channels = args.in_channels - out_channels = args.out_channels - latent_channels = args.latent_channels - - device = torch.device("cuda:0") - print - # Initialize the two model implementations - path = "data/hunyuanvideo-community/HunyuanVideo/vae" - config_path = os.path.join(path, "config.json") - config = json.load(open(config_path)) - config.pop("_class_name") - config.pop("_diffusers_version") - model1 = MyHunyuanVAE(**config).to(torch.bfloat16) - - model2 = DiffusersHunyuanVAE(**config).to(torch.bfloat16) - - loaded = load_file(os.path.join(path, - "diffusion_pytorch_model.safetensors")) - model1.load_state_dict(loaded) - model2.load_state_dict(loaded) - - # Set both models to eval mode - model1.eval() - model2.eval() - - # Move to GPU - model1 = model1.to(device) - model2 = model2.to(device) - - model1.enable_tiling(tile_sample_min_height=32, - tile_sample_min_width=32, - tile_sample_min_num_frames=8, - tile_sample_stride_height=16, - tile_sample_stride_width=16, - tile_sample_stride_num_frames=4) - model2.enable_tiling(tile_sample_min_height=32, - tile_sample_min_width=32, - tile_sample_min_num_frames=8, - tile_sample_stride_height=16, - tile_sample_stride_width=16, - tile_sample_stride_num_frames=4) - - # Create identical inputs for both models - batch_size = 1 - - # Video input [B, C, T, H, W] - input_tensor = torch.randn(batch_size, - 3, - 21, - 64, - 64, - device=device, - dtype=torch.bfloat16) - - # Disable gradients for inference - with torch.no_grad(): - # Test encoding - logger.info("Testing encoding...") - latent1 = model1.encode(input_tensor).mean - print("--------------------------------") - latent2 = model2.encode(input_tensor).latent_dist.mean - # Check if latents have the same shape - assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}" - assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}" - # Check if latents are similar - max_diff_encode = torch.max(torch.abs(latent1 - latent2)) - mean_diff_encode = torch.mean(torch.abs(latent1 - latent2)) - logger.info( - f"Maximum difference between encoded latents: {max_diff_encode.item()}" - ) - logger.info( - f"Mean difference between encoded latents: {mean_diff_encode.item()}" - ) - assert max_diff_encode < 1e-4, f"Encoded latents differ significantly: max diff = {max_diff_encode.item()}" - # Test decoding - logger.info("Testing decoding...") - output1 = model1.decode(latent1) - output2 = model2.decode(latent2).sample - # Check if outputs have the same shape - assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" - - # Check if outputs are similar - max_diff_decode = torch.max(torch.abs(output1 - output2)) - mean_diff_decode = torch.mean(torch.abs(output1 - output2)) - logger.info( - f"Maximum difference between decoded outputs: {max_diff_decode.item()}" - ) - logger.info( - f"Mean difference between decoded outputs: {mean_diff_decode.item()}" - ) - assert max_diff_decode < 1e-4, f"Decoded outputs differ significantly: max diff = {max_diff_decode.item()}" - - logger.info( - "Test passed! Both VAE implementations produce similar outputs.") - logger.info("Test completed successfully") - - -if __name__ == "__main__": - test_hunyuan_vae() diff --git a/fastvideo/v1/tests/old_tests/test_hunyuanvideo.py b/fastvideo/v1/tests/old_tests/test_hunyuanvideo.py deleted file mode 100644 index b8a882ed5..000000000 --- a/fastvideo/v1/tests/old_tests/test_hunyuanvideo.py +++ /dev/null @@ -1,260 +0,0 @@ -import argparse -import os -from itertools import chain - -import numpy as np -import torch -import torch.nn as nn -from torch.distributed.device_mesh import init_device_mesh - -from fastvideo.v1.distributed.parallel_state import ( - cleanup_dist_env_and_memory, destroy_distributed_environment, - destroy_model_parallel, get_sequence_model_parallel_rank, - get_sequence_model_parallel_world_size, init_distributed_environment, - initialize_model_parallel) -from fastvideo.v1.logger import init_logger -from fastvideo.v1.models.dits.hunyuanvideo import ( - HunyuanVideoTransformer3DModel as HunyuanVideoDit) -from fastvideo.v1.models.hunyuan.modules.models import ( - HYVideoDiffusionTransformer) -from fastvideo.v1.models.loader.fsdp_load import shard_model -from fastvideo.v1.utils.parallel_states import ( - initialize_sequence_parallel_state) - -logger = init_logger(__name__) - - -def initialize_identical_weights(model1, model2, seed=42): - """Initialize both models with identical weights using a fixed seed for reproducibility.""" - # Get all parameters from both models - params1 = dict(model1.named_parameters()) - params2 = dict(model2.named_parameters()) - - # Initialize each layer with identical values - with torch.no_grad(): - # Initialize weights - for name1, param1 in params1.items(): - if 'weight' in name1: - # Set seed before each weight initialization - torch.manual_seed(seed) - nn.init.normal_(param1, mean=0.0, std=0.05) - - for name2, param2 in params2.items(): - if 'weight' in name2: - # Reset seed to get same initialization - torch.manual_seed(seed) - nn.init.normal_(param2, mean=0.0, std=0.05) - - # Initialize biases - for name1, param1 in params1.items(): - if 'bias' in name1: - torch.manual_seed(seed) - nn.init.normal_(param1, mean=0.0, std=0.05) - param1.data = param1.data.to(torch.bfloat16) - - for name2, param2 in params2.items(): - if 'bias' in name2: - torch.manual_seed(seed) - nn.init.normal_(param2, mean=0.0, std=0.05) - param2.data = param2.data.to(torch.bfloat16) - - logger.info("Both models initialized with identical weights in bfloat16") - return model1, model2 - - -def setup_args(): - parser = argparse.ArgumentParser( - description='Distributed HunyuanVideo Test') - parser.add_argument('--sequence_model_parallel_size', - type=int, - default=1, - help='Degree of sequence model parallelism') - parser.add_argument('--hidden-size', - type=int, - default=128, - help='Hidden size for the model') - parser.add_argument('--heads-num', - type=int, - default=4, - help='Number of attention heads') - parser.add_argument('--double-blocks-depth', - type=int, - default=2, - help='Number of double stream blocks') - parser.add_argument('--single-blocks-depth', - type=int, - default=2, - help='Number of single stream blocks') - return parser.parse_args() - - -def test_hunyuanvideo_distributed(): - args = setup_args() - - # Initialize distributed environment - local_rank = int(os.environ.get("LOCAL_RANK", 0)) - rank = int(os.environ.get("RANK", 0)) - world_size = int(os.environ.get("WORLD_SIZE", 1)) - - logger.info( - f"Initializing process: rank={rank}, local_rank={local_rank}, world_size={world_size}" - ) - - # Initialize distributed environment - init_distributed_environment(world_size=world_size, - rank=rank, - local_rank=local_rank) - - # Initialize tensor model parallel groups - initialize_model_parallel( - sequence_model_parallel_size=args.sequence_model_parallel_size) - initialize_sequence_parallel_state(world_size) - # Get tensor parallel info - sp_rank = get_sequence_model_parallel_rank() - sp_world_size = get_sequence_model_parallel_world_size() - - logger.info( - f"Process rank {rank} initialized with SP rank {sp_rank} in SP world size {sp_world_size}" - ) - - # Set fixed random seed for reproducibility - torch.manual_seed(42) - np.random.seed(42) - - # Small model parameters for testing - hidden_size = args.hidden_size - heads_num = args.heads_num - mm_double_blocks_depth = args.double_blocks_depth - mm_single_blocks_depth = args.single_blocks_depth - patch_size = [1, 2, 2] - torch.cuda.set_device(f"cuda:{local_rank}") - # Initialize the two model implementations - model1 = HunyuanVideoDit( - patch_size=2, - patch_size_t=1, - in_channels=4, - out_channels=4, - attention_head_dim=hidden_size // heads_num, - num_attention_heads=heads_num, - num_layers=mm_double_blocks_depth, - num_single_layers=mm_single_blocks_depth, - rope_axes_dim=[8, 16, 8], # sum = hidden_size // heads_num = 32 - dtype=torch.bfloat16).to(torch.bfloat16) - model2 = HYVideoDiffusionTransformer( - patch_size=patch_size, - in_channels=4, - hidden_size=hidden_size, - heads_num=heads_num, - mm_double_blocks_depth=mm_double_blocks_depth, - mm_single_blocks_depth=mm_single_blocks_depth, - rope_dim_list=[8, 16, 8], # sum = hidden_size // heads_num = 32 - dtype=torch.bfloat16).to(torch.bfloat16) - - # print("--------------------------------") - # for name, param in model3.named_parameters(): - # print(name) - # import pdb; pdb.set_trace() - # # Initialize with identical weights - model1, model2 = initialize_identical_weights(model1, model2, seed=42) - device_mesh = init_device_mesh( - "cuda", - mesh_shape=(sp_world_size, ), - mesh_dim_names=("dp", ), - ) - shard_model(model1, cpu_offload=False, reshard_after_forward=True) - for n, p in chain(model1.named_parameters(), model1.named_buffers()): - if p.is_meta: - raise RuntimeError( - f"Unexpected param or buffer {n} on meta device.") - for p in model1.parameters(): - p.requires_grad = False - # Set both models to eval mode - model1.eval() - model2.eval() - - # Move to GPU based on local rank (0 or 1 for 2 GPUs) - device = torch.device(f"cuda:{local_rank}") - model1 = model1.to(device) - model2 = model2.to(device) - - # Create identical inputs for both models - batch_size = 1 - seq_len = 3 - - # Video latents [B, C, T, H, W] - hidden_states = torch.randn(batch_size, - 4, - 8, - 16, - 16, - device=device, - dtype=torch.bfloat16) - chunk_per_rank = hidden_states.shape[2] // sp_world_size - hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * - chunk_per_rank] - - # Text embeddings [B, L, D] (including global token) - encoder_hidden_states = torch.randn(batch_size, - seq_len + 1, - 4096, - device=device, - dtype=torch.bfloat16) - - # Timestep - timestep = torch.tensor([500], device=device, dtype=torch.bfloat16) - - # Attention mask for text - encoder_attention_mask = torch.ones(batch_size, - seq_len, - device=device, - dtype=torch.bfloat16) - guidance = torch.tensor([1.0], device=device, dtype=torch.bfloat16) - # Disable gradients for inference - with torch.no_grad(): - output1 = model1( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - ) - print("--------------------------------") - output2, _ = model2( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - encoder_attention_mask=encoder_attention_mask, - ) - - # Check if outputs have the same shape - assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" - - # Check if outputs are similar (allowing for small numerical differences) - max_diff = torch.max(torch.abs(output1 - output2)) - logger.info(f"Maximum difference between outputs: {max_diff.item()}") - # mean diff - mean_diff = torch.mean(torch.abs(output1 - output2)) - logger.info(f"Mean difference between outputs: {mean_diff.item()}") - # diff sum - diff_sum = torch.sum(torch.abs(output1 - output2)) - logger.info(f"Diff sum between outputs: {diff_sum.item()}") - # sum - sum_output1 = torch.sum(output1.float()) - sum_output2 = torch.sum(output2.float()) - logger.info(f"Rank {sp_rank} Sum of output1: {sum_output1.item()}") - logger.info(f"Rank {sp_rank} Sum of output2: {sum_output2.item()}") - # The outputs should be very close if not identical - assert max_diff < 1e-3, f"Outputs differ significantly: max diff = {max_diff.item()}" # Increased tolerance for bf16 - - logger.info( - "Test passed! Both model implementations produce the same outputs.") - - # Clean up - logger.info("Cleaning up distributed environment") - destroy_model_parallel() - destroy_distributed_environment() - cleanup_dist_env_and_memory() - - logger.info("Test completed successfully") - - -if __name__ == "__main__": - test_hunyuanvideo_distributed() diff --git a/fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py b/fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py deleted file mode 100644 index fa471e777..000000000 --- a/fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py +++ /dev/null @@ -1,212 +0,0 @@ -import argparse -import glob -import json -import os - -import torch - -from fastvideo.models.hunyuan.modules.models import ( - HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer) -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state -from fastvideo.v1.distributed.parallel_state import ( - cleanup_dist_env_and_memory, destroy_distributed_environment, - destroy_model_parallel, get_sequence_model_parallel_rank, - get_sequence_model_parallel_world_size, init_distributed_environment, - initialize_model_parallel) -from fastvideo.v1.logger import init_logger -from fastvideo.v1.models.dits.hunyuanvideo import ( - HunyuanVideoTransformer3DModel as HunyuanVideoDit) -from fastvideo.v1.models.loader.fsdp_load import load_fsdp_model - -logger = init_logger(__name__) - - -def setup_args(): - parser = argparse.ArgumentParser( - description='Distributed HunyuanVideo Test') - parser.add_argument('--sequence_model_parallel_size', - type=int, - default=1, - help='Degree of sequence model parallelism') - return parser.parse_args() - - -def test_hunyuanvideo_distributed(): - args = setup_args() - - # Initialize distributed environment - local_rank = int(os.environ.get("LOCAL_RANK", 0)) - rank = int(os.environ.get("RANK", 0)) - world_size = int(os.environ.get("WORLD_SIZE", 1)) - - logger.info( - f"Initializing process: rank={rank}, local_rank={local_rank}, world_size={world_size}" - ) - - # Initialize distributed environment - init_distributed_environment(world_size=world_size, - rank=rank, - local_rank=local_rank) - torch.cuda.set_device(f"cuda:{local_rank}") - # Initialize tensor model parallel groups - initialize_model_parallel( - sequence_model_parallel_size=args.sequence_model_parallel_size) - initialize_sequence_parallel_state(args.sequence_model_parallel_size) - # Get tensor parallel info - sp_rank = get_sequence_model_parallel_rank() - sp_world_size = get_sequence_model_parallel_world_size() - - logger.info( - f"Process rank {rank} initialized with SP rank {sp_rank} in SP world size {sp_world_size}" - ) - - # load data/hunyuanvideo_community/transformer/config.json - with open( - "data/hunyuanvideo-community/HunyuanVideo/transformer/config.json") as f: - config = json.load(f) - # remove "_class_name": "HunyuanVideoTransformer3DModel", "_diffusers_version": "0.32.0.dev0", - # TODO: write normalize config function - config.pop("_class_name") - config.pop("_diffusers_version") - # load data/hunyuanvideo_community/transformer/*.safetensors - weight_dir_list = glob.glob( - "data/hunyuanvideo-community/HunyuanVideo/transformer/*.safetensors") - # to str - weight_dir_list = [str(path) for path in weight_dir_list] - model1 = load_fsdp_model(HunyuanVideoDit, - init_params=config, - weight_dir_list=weight_dir_list, - device=torch.device(f"cuda:{local_rank}"), - cpu_offload=False) - - # successfully sharded the model (hunyuanvideo bf16 should take around 26GB in total) - total_params = sum(p.numel() for p in model1.parameters()) - logger.info(f"Total parameters: {total_params / 1e9}B") - - # Calculate weight sum for model1 (converting to float64 to avoid overflow) - weight_sum_model1 = sum( - p.to(torch.float64).sum().item() for p in model1.parameters()) - # Also calculate mean for more stable comparison - weight_mean_model1 = weight_sum_model1 / total_params - - model2 = HYVideoDiffusionTransformer( - in_channels=16, - out_channels=16, - **HUNYUAN_VIDEO_CONFIG["HYVideo-T/2-cfgdistill"], - device=torch.device(f"cuda:{local_rank}"), - dtype=torch.bfloat16).bfloat16() - # data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt - state_dict = torch.load( - "/mbz/users/hao.zhang/peiyuan/FastVideo/data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt", - map_location=lambda storage, loc: storage)["module"] - model2.load_state_dict(state_dict, strict=True) - model2.to(torch.device(f"cuda:{local_rank}")).bfloat16() - print("load state dict done") - - # Calculate weight sum for model2 (converting to float64 to avoid overflow) - total_params_model2 = sum(p.numel() for p in model2.parameters()) - weight_sum_model2 = sum( - p.to(torch.float64).sum().item() for p in model2.parameters()) - # Also calculate mean for more stable comparison - weight_mean_model2 = weight_sum_model2 / total_params_model2 - logger.info(f"Model 2 weight sum: {weight_sum_model2}") - logger.info(f"Model 2 weight mean: {weight_mean_model2}") - - # Set both models to eval mode - model1.eval() - model2.eval() - - # Create random inputs for testing - batch_size = 1 - seq_len = 3 - device = torch.device(f"cuda:{local_rank}") - - # Video latents [B, C, T, H, W] - hidden_states = torch.randn(batch_size, - 16, - 8, - 16, - 16, - device=device, - dtype=torch.bfloat16) - chunk_per_rank = hidden_states.shape[2] // sp_world_size - hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * - chunk_per_rank] - - # Text embeddings [B, L, D] (including global token) - encoder_hidden_states = torch.randn(batch_size, - seq_len + 1, - 4096, - device=device, - dtype=torch.bfloat16) - - # Timestep - timestep = torch.tensor([500], device=device, dtype=torch.bfloat16) - - # Attention mask for text - encoder_attention_mask = torch.ones(batch_size, - seq_len, - device=device, - dtype=torch.bfloat16) - guidance = torch.tensor([1.0], device=device, dtype=torch.bfloat16) - - # Disable gradients for inference - with torch.no_grad(): - # Run inference on model1 - with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): - logger.info("Running inference on model1") - output1 = model1( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - ) - logger.info("Model 1 inference completed") - - # Run inference on model2 - output2, _ = model2( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - encoder_attention_mask=encoder_attention_mask, - ) - logger.info("Model 2 inference completed") - - # Check if outputs have the same shape - assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" - - # Compare weight sums and means - logger.info(f"Model 1 weight sum: {weight_sum_model1}") - logger.info(f"Model 2 weight sum: {weight_sum_model2}") - weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2) - logger.info(f"Weight sum difference: {weight_sum_diff}") - - logger.info(f"Model 1 weight mean: {weight_mean_model1}") - logger.info(f"Model 2 weight mean: {weight_mean_model2}") - weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2) - logger.info(f"Weight mean difference: {weight_mean_diff}") - - # mean diff - mean_diff = torch.mean(torch.abs(output1 - output2)) - assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}" - - # diff sum - diff_sum = torch.sum(torch.abs(output1 - output2)) - logger.info(f"Diff sum between outputs: {diff_sum.item()}") - - # sum - sum_output1 = torch.sum(output1.float()) - sum_output2 = torch.sum(output2.float()) - logger.info(f"Rank {sp_rank} Sum of output1: {sum_output1.item()}") - logger.info(f"Rank {sp_rank} Sum of output2: {sum_output2.item()}") - - # Clean up - logger.info("Cleaning up distributed environment") - destroy_model_parallel() - destroy_distributed_environment() - cleanup_dist_env_and_memory() - - logger.info("Test completed successfully") - - -if __name__ == "__main__": - test_hunyuanvideo_distributed() diff --git a/fastvideo/v1/tests/ssim/compute_ssim.py b/fastvideo/v1/tests/ssim/compute_ssim.py index 8a6647903..943adf7f2 100644 --- a/fastvideo/v1/tests/ssim/compute_ssim.py +++ b/fastvideo/v1/tests/ssim/compute_ssim.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + import argparse import os diff --git a/fastvideo/v1/tests/ssim/test_inference_similarity.py b/fastvideo/v1/tests/ssim/test_inference_similarity.py index 39c29a794..495c31205 100644 --- a/fastvideo/v1/tests/ssim/test_inference_similarity.py +++ b/fastvideo/v1/tests/ssim/test_inference_similarity.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + import json import os diff --git a/fastvideo/v1/tests/transformers/test_hunyuanvideo.py b/fastvideo/v1/tests/transformers/test_hunyuanvideo.py new file mode 100644 index 000000000..b7f2def35 --- /dev/null +++ b/fastvideo/v1/tests/transformers/test_hunyuanvideo.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +from itertools import chain + +import pytest +import torch +import torch.nn as nn + +from fastvideo.v1.distributed.parallel_state import ( + get_sequence_model_parallel_rank, + get_sequence_model_parallel_world_size) +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.dits.hunyuanvideo import ( + HunyuanVideoTransformer3DModel as HunyuanVideoDit) + +from fastvideo.v1.models.loader.fsdp_load import shard_model +from fastvideo.v1.forward_context import set_forward_context + + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +# Latent generated on commit c021e8a27cf437ac22827f2bc58b7f006561317f with 1 x L40S +REFERENCE_LATENT = 1472.079828262329 + + +def initialize_identical_weights(model, seed=42): + """Initialize both models with identical weights using a fixed seed for reproducibility.""" + # Get all parameters from both models + params1 = dict(model.named_parameters()) + + # Initialize each layer with identical values + with torch.no_grad(): + # Initialize weights + for name1, param1 in params1.items(): + if 'weight' in name1: + # Set seed before each weight initialization + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + + # Initialize biases + for name1, param1 in params1.items(): + if 'bias' in name1: + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + param1.data = param1.data.to(torch.bfloat16) + + logger.info("Model initialized with identical weights in bfloat16") + return model + + +@pytest.mark.usefixtures("distributed_setup") +def test_hunyuanvideo_distributed(): + # Get tensor parallel info + sp_rank = get_sequence_model_parallel_rank() + sp_world_size = get_sequence_model_parallel_world_size() + + # Small model parameters for testing + hidden_size = 128 + heads_num = 4 + mm_double_blocks_depth = 2 + mm_single_blocks_depth = 2 + torch.cuda.set_device("cuda:0") + # Initialize the two model implementations + model = HunyuanVideoDit( + patch_size=2, + patch_size_t=1, + in_channels=4, + out_channels=4, + attention_head_dim=hidden_size // heads_num, + num_attention_heads=heads_num, + num_layers=mm_double_blocks_depth, + num_single_layers=mm_single_blocks_depth, + rope_axes_dim=[8, 16, 8], # sum = hidden_size // heads_num = 32 + dtype=torch.bfloat16).to(torch.bfloat16) + + # Initialize with identical weights + model = initialize_identical_weights(model, seed=42) + shard_model(model, cpu_offload=False, reshard_after_forward=True) + for n, p in chain(model.named_parameters(), model.named_buffers()): + if p.is_meta: + raise RuntimeError( + f"Unexpected param or buffer {n} on meta device.") + for p in model.parameters(): + p.requires_grad = False + + model.eval() + + # Move to GPU based on local rank (0 or 1 for 2 GPUs) + device = torch.device(f"cuda:0") + model = model.to(device) + + batch_size = 1 + seq_len = 3 + + # Video latents [B, C, T, H, W] + hidden_states = torch.randn(batch_size, + 4, + 8, + 16, + 16, + device=device, + dtype=torch.bfloat16) + chunk_per_rank = hidden_states.shape[2] // sp_world_size + hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * + chunk_per_rank] + + # Text embeddings [B, L, D] (including global token) + encoder_hidden_states = torch.randn(batch_size, + seq_len + 1, + 4096, + device=device, + dtype=torch.bfloat16) + + # Timestep + timestep = torch.tensor([500], device=device, dtype=torch.bfloat16) + + # Disable gradients for inference + with torch.no_grad(): + with set_forward_context(current_timestep=0, attn_metadata=None): + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + ) + + latent = output.double().sum().item() + + # Check if latents are similar + diff_output_latents = abs(REFERENCE_LATENT - latent) + logger.info( + f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}" + ) + assert diff_output_latents < 1e-4, f"Output latents differ significantly: max diff = {diff_output_latents}" diff --git a/fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py b/fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py new file mode 100644 index 000000000..1b4d9c7c0 --- /dev/null +++ b/fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 + +import glob +import json +import os +import pytest +import torch + +from fastvideo.v1.distributed.parallel_state import ( + get_sequence_model_parallel_rank, + get_sequence_model_parallel_world_size) +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.dits.hunyuanvideo import ( + HunyuanVideoTransformer3DModel as HunyuanVideoDit) +from fastvideo.v1.models.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.utils import maybe_download_model +from fastvideo.v1.forward_context import set_forward_context + + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + "data", BASE_MODEL_PATH)) +TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") +CONFIG_PATH = os.path.join(TRANSFORMER_PATH, "config.json") + +LOCAL_RANK = 0 +RANK = 0 +WORLD_SIZE = 1 + +# Latent generated on commit c021e8a27cf437ac22827f2bc58b7f006561317f with 1 x L40S +REFERENCE_LATENT = 89.7002067565918 + + +@pytest.mark.usefixtures("distributed_setup") +def test_hunyuanvideo_distributed(): + logger.info( + f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}" + ) + + torch.cuda.set_device(f"cuda:{LOCAL_RANK}") + + # Get tensor parallel info + sp_rank = get_sequence_model_parallel_rank() + sp_world_size = get_sequence_model_parallel_world_size() + + logger.info( + f"Process rank {RANK} initialized with SP rank {sp_rank} in SP world size {sp_world_size}" + ) + + config = json.load(open(CONFIG_PATH)) + # remove "_class_name": "HunyuanVideoTransformer3DModel", "_diffusers_version": "0.32.0.dev0", + # TODO: write normalize config function + config.pop("_class_name") + config.pop("_diffusers_version") + + weight_dir_list = glob.glob(os.path.join(TRANSFORMER_PATH, "*.safetensors")) + weight_dir_list = [str(path) for path in weight_dir_list] + model = load_fsdp_model(HunyuanVideoDit, + init_params=config, + weight_dir_list=weight_dir_list, + device=torch.device(f"cuda:{LOCAL_RANK}"), + cpu_offload=False) + + model.eval() + + # Create random inputs for testing + batch_size = 1 + seq_len = 3 + device = torch.device(f"cuda:{LOCAL_RANK}") + + # Video latents [B, C, T, H, W] + hidden_states = torch.randn(batch_size, + 16, + 8, + 16, + 16, + device=device, + dtype=torch.bfloat16) + chunk_per_rank = hidden_states.shape[2] // sp_world_size + hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * + chunk_per_rank] + + # Text embeddings [B, L, D] (including global token) + encoder_hidden_states = torch.randn(batch_size, + seq_len + 1, + 4096, + device=device, + dtype=torch.bfloat16) + + # Timestep + timestep = torch.tensor([500], device=device, dtype=torch.bfloat16) + + # Disable gradients for inference + with torch.no_grad(): + # Run inference on model + with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): + with set_forward_context(current_timestep=0, attn_metadata=None): + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + ) + + latent = output.double().sum().item() + + # Check if latents are similar + diff_output_latents = abs(REFERENCE_LATENT - latent) + logger.info( + f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}" + ) + assert diff_output_latents < 1e-4, f"Output latents differ significantly: max diff = {diff_output_latents}" + + diff --git a/fastvideo/v1/tests/vaes/test_hunyuan_vae.py b/fastvideo/v1/tests/vaes/test_hunyuan_vae.py new file mode 100644 index 000000000..be876ec83 --- /dev/null +++ b/fastvideo/v1/tests/vaes/test_hunyuan_vae.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json +import os + +import pytest +import torch +from safetensors.torch import load_file + +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.vaes.hunyuanvae import ( + AutoencoderKLHunyuanVideo as MyHunyuanVAE) +from fastvideo.v1.utils import maybe_download_model + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" + +BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + "data", BASE_MODEL_PATH)) +VAE_PATH = os.path.join(MODEL_PATH, "vae") +CONFIG_PATH = os.path.join(VAE_PATH, "config.json") + +# Latent generated on commit 250f0b916cebb18a1c15c4aae1a0b480604d066a with 1 x A40 +REFERENCE_LATENT = -105.51324462890625 + + +@pytest.mark.usefixtures("distributed_setup") +def test_hunyuan_vae(): + device = torch.device("cuda:0") + # Initialize the two model implementations + config = json.load(open(CONFIG_PATH)) + config.pop("_class_name") + config.pop("_diffusers_version") + model = MyHunyuanVAE(**config).to(torch.bfloat16) + + loaded = load_file(os.path.join(VAE_PATH, + "diffusion_pytorch_model.safetensors")) + model.load_state_dict(loaded) + + # Set model to eval mode + model.eval() + + # Move to GPU + model = model.to(device) + + model.enable_tiling(tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_min_num_frames=8, + tile_sample_stride_height=16, + tile_sample_stride_width=16, + tile_sample_stride_num_frames=4) + + batch_size = 1 + + # Video input [B, C, T, H, W] + input_tensor = torch.randn(batch_size, + 3, + 21, + 64, + 64, + device=device, + dtype=torch.bfloat16) + + # Disable gradients for inference + with torch.no_grad(): + latent = model.encode(input_tensor).mean.double().sum().item() + + # Check if latents are similar + diff_encoded_latents = abs(REFERENCE_LATENT - latent) + logger.info( + f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}" + ) + assert diff_encoded_latents < 1e-4, f"Encoded latents differ significantly: max diff = {diff_encoded_latents}" +