nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Port tests to v1 (#333)

Browse Source
This commit is contained in:
Kevin Lin
2025-04-11 01:40:01 -07:00
committed by GitHub
parent ccb262974e
commit 137f61f2fe
12 changed files with 474 additions and 669 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pr-test.yml
+120 -2
View File
@@ -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
fastvideo/v1/tests/conftest.py
+12 -8
View File
@@ -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()
fastvideo/v1/tests/encoders/test_clip_encoder.py
+2 -4
View File
@@ -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(
fastvideo/v1/tests/encoders/test_llama_encoder.py
+2 -4
View File
@@ -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(
fastvideo/v1/tests/old_tests/test_hunyuan_vae.py
-179
View File
@@ -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()
fastvideo/v1/tests/old_tests/test_hunyuanvideo.py
-260
View File
@@ -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()
fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py
-212
View File
@@ -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()
fastvideo/v1/tests/ssim/compute_ssim.py
+2
View File
@@ -1,3 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
import argparse
import os
fastvideo/v1/tests/ssim/test_inference_similarity.py
+2
View File
@@ -1,3 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
import json
import os
fastvideo/v1/tests/transformers/test_hunyuanvideo.py
+137
View File
@@ -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}"
fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py
+119
View File
@@ -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}"
fastvideo/v1/tests/vaes/test_hunyuan_vae.py
+78
View File
@@ -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}"