diff --git a/fastvideo/__init__.py b/fastvideo/__init__.py index 45fd11fa5..44c7938e9 100644 --- a/fastvideo/__init__.py +++ b/fastvideo/__init__.py @@ -3,4 +3,4 @@ from fastvideo.configs.sample import SamplingParam from fastvideo.entrypoints.video_generator import VideoGenerator from fastvideo.version import __version__ -__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"] \ No newline at end of file +__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"] diff --git a/fastvideo/tests/encoders/test_clip_encoder.py b/fastvideo/tests/encoders/test_clip_encoder.py index 90d21b508..a6c4da380 100644 --- a/fastvideo/tests/encoders/test_clip_encoder.py +++ b/fastvideo/tests/encoders/test_clip_encoder.py @@ -22,8 +22,8 @@ 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)) + local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod + ) TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder_2") TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer_2") @@ -130,17 +130,6 @@ def test_clip_encoder(): assert last_hidden_state1.shape == last_hidden_state2.shape, \ f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" - - max_diff_hidden = torch.max( - torch.abs(last_hidden_state1 - last_hidden_state2)) - mean_diff_hidden = torch.mean( - torch.abs(last_hidden_state1 - last_hidden_state2)) - - logger.info("Maximum difference in last hidden states: %f", - max_diff_hidden.item()) - logger.info("Mean difference in last hidden states: %f", - mean_diff_hidden.item()) - # Compare pooler outputs pooler_output1 = outputs1.pooler_output pooler_output2 = outputs2.pooler_output @@ -148,22 +137,5 @@ def test_clip_encoder(): assert pooler_output1.shape == pooler_output2.shape, \ f"Pooler output shapes don't match: {pooler_output1.shape} vs {pooler_output2.shape}" - max_diff_pooler = torch.max( - torch.abs(pooler_output1 - pooler_output2)) - mean_diff_pooler = torch.mean( - torch.abs(pooler_output1 - pooler_output2)) - - logger.info("Maximum difference in pooler outputs: %f", - max_diff_pooler.item()) - logger.info("Mean difference in pooler outputs: %f", - mean_diff_pooler.item()) - - # Check if outputs are similar (allowing for small numerical differences) - assert mean_diff_hidden < 1e-2, \ - f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}" - assert mean_diff_pooler < 1e-2, \ - f"Pooler outputs differ significantly: mean diff = {mean_diff_pooler.item()}" - assert max_diff_hidden < 1e-1, \ - f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}" - assert max_diff_pooler < 2e-2, \ - f"Pooler outputs differ significantly: max diff = {max_diff_pooler.item()}" + assert_close(pooler_output1, pooler_output2, atol=1e-2, rtol=1e-3) + assert_close(last_hidden_state1, last_hidden_state2, atol=1e-2, rtol=1e-3) diff --git a/fastvideo/tests/encoders/test_llama_encoder.py b/fastvideo/tests/encoders/test_llama_encoder.py index 0a7d220ce..da50ba393 100644 --- a/fastvideo/tests/encoders/test_llama_encoder.py +++ b/fastvideo/tests/encoders/test_llama_encoder.py @@ -22,8 +22,8 @@ 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)) + local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod + ) TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder") TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer") @@ -68,8 +68,7 @@ def test_llama_encoder(): logger.info("Model1 has %d parameters", len(params1)) logger.info("Model2 has %d parameters", len(params2)) - # Compare a few key parameters - weight_diffs = [] + # check if embed_tokens are the same device = model1.embed_tokens.weight.device assert torch.allclose(model1.embed_tokens.weight, @@ -78,6 +77,18 @@ def test_llama_encoder(): "layers.{}.input_layernorm.weight", "layers.{}.post_attention_layernorm.weight" ] + for layer_idx in range(hf_config.num_hidden_layers): + for w in weights: + name1 = w.format(layer_idx) + name2 = w.format(layer_idx) + p1 = params1[name1] + p2 = params2[name2] + if "gate_up" in name2: + # print("skipping gate_up") + continue + p1 = p1.to_local().to(device) if isinstance(p1, DTensor) else p1.to(device) + p2 = p2.to_local().to(device) if isinstance(p2, DTensor) else p2.to(device) + assert_close(p1, p2, atol=1e-4, rtol=1e-4) for name1, param1 in sorted(params1.items()): name2 = name1 @@ -139,19 +150,4 @@ def test_llama_encoder(): assert last_hidden_state1.shape == last_hidden_state2.shape, \ f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" - - max_diff_hidden = torch.max( - torch.abs(last_hidden_state1 - last_hidden_state2)) - mean_diff_hidden = torch.mean( - torch.abs(last_hidden_state1 - last_hidden_state2)) - - logger.info("Maximum difference in last hidden states: %f", - max_diff_hidden.item()) - logger.info("Mean difference in last hidden states: %f", - mean_diff_hidden.item()) - - # Check if outputs are similar (allowing for small numerical differences) - assert mean_diff_hidden < 1e-2, \ - f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}" - assert max_diff_hidden < 1e-1, \ - f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}" + assert_close(last_hidden_state1, last_hidden_state2, atol=1e-1, rtol=1e-4) diff --git a/fastvideo/tests/encoders/test_t5_encoder.py b/fastvideo/tests/encoders/test_t5_encoder.py index 27e2bf76c..ef587b6a8 100644 --- a/fastvideo/tests/encoders/test_t5_encoder.py +++ b/fastvideo/tests/encoders/test_t5_encoder.py @@ -133,24 +133,7 @@ def test_t5_encoder(t5_model_paths): last_hidden_state1 = outputs1[tokens.attention_mask == 1] last_hidden_state2 = outputs2[tokens.attention_mask == 1] - assert last_hidden_state1.shape == last_hidden_state2.shape, \ - f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" - - max_diff_hidden = torch.max( - torch.abs(last_hidden_state1 - last_hidden_state2)) - mean_diff_hidden = torch.mean( - torch.abs(last_hidden_state1 - last_hidden_state2)) - - logger.info("Maximum difference in last hidden states: %s", - max_diff_hidden.item()) - logger.info("Mean difference in last hidden states: %s", - mean_diff_hidden.item()) - logger.info("Max memory allocated: %s GB", torch.cuda.max_memory_allocated() / 1024**3) - # Check if outputs are similar (allowing for small numerical differences) - assert mean_diff_hidden < 1e-4, \ - f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}" - assert max_diff_hidden < 1e-4, \ - f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}" + assert_close(last_hidden_state1, last_hidden_state2, atol=1e-4, rtol=1e-4) @pytest.mark.usefixtures("distributed_setup") @@ -252,18 +235,4 @@ def test_t5_large_encoder(t5_large_model_paths): assert last_hidden_state1.shape == last_hidden_state2.shape, \ f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" - max_diff_hidden = torch.max( - torch.abs(last_hidden_state1 - last_hidden_state2)) - mean_diff_hidden = torch.mean( - torch.abs(last_hidden_state1 - last_hidden_state2)) - - logger.info("Maximum difference in last hidden states: %s", - max_diff_hidden.item()) - logger.info("Mean difference in last hidden states: %s", - mean_diff_hidden.item()) - logger.info("Max memory allocated: %s GB", torch.cuda.max_memory_allocated() / 1024**3) - # Check if outputs are similar (allowing for small numerical differences) - assert mean_diff_hidden < 1e-4, \ - f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}" - assert max_diff_hidden < 1e-4, \ - f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}" + assert_close(last_hidden_state1, last_hidden_state2, atol=1e-4, rtol=1e-4) diff --git a/fastvideo/tests/old_tests/distributed_run.sh b/fastvideo/tests/old_tests/distributed_run.sh deleted file mode 100644 index 6c2ee410e..000000000 --- a/fastvideo/tests/old_tests/distributed_run.sh +++ /dev/null @@ -1,19 +0,0 @@ -#!/bin/bash - -num_gpus=8 -torchrun --standalone --nnodes=1 --nproc_per_node=$num_gpus \ - --master_port 29503 \ - tp_example.py - - - -num_gpus=2 -torchrun --standalone --nnodes=1 --nproc_per_node=$num_gpus \ - --master_port 29503 \ - fastvideo/tests/test_hunyuanvideo_load.py --sequence_model_parallel_size $num_gpus - -torchrun --nnodes=1 --nproc_per_node=1 --master_port 29503 fastvideo/tests/test_llama_encoder.py - - -export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN -torchrun --nnodes=1 --nproc_per_node=1 --master_port 29503 fastvideo/tests/test_clip_encoder.py \ No newline at end of file diff --git a/fastvideo/tests/old_tests/test_t5_encoder.py b/fastvideo/tests/old_tests/test_t5_encoder.py deleted file mode 100644 index 788010e2f..000000000 --- a/fastvideo/tests/old_tests/test_t5_encoder.py +++ /dev/null @@ -1,164 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -import argparse - -import numpy as np -import torch -from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel - -from fastvideo.distributed import (maybe_init_distributed_environment_and_model_parallel) -from fastvideo.logger import init_logger - -logger = init_logger(__name__) - - -def setup_args(): - parser = argparse.ArgumentParser(description='T5 Encoder Test') - parser.add_argument('--model_path', type=str, default="google/umt5-xxl") - parser.add_argument( - '--dit-precision', - type=str, - default="float32", - help='Precision to use for the model (float32, float16, bfloat16)') - return parser.parse_args() - - -def test_t5_encoder(): - maybe_init_distributed_environment_and_model_parallel(1, 1) - - # Set fixed random seed for reproducibility - torch.manual_seed(42) - np.random.seed(42) - - device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - - # Initialize the two model implementations - model_path = "/workspace/data/Wan2.1-T2V-1.3B-Diffusers/text_encoder" - tokenizer_path = "/workspace/data/Wan2.1-T2V-1.3B-Diffusers/tokenizer" - - hf_config = AutoConfig.from_pretrained(model_path) - print(hf_config) - precision = torch.float16 # It must be float16 because the weight loader is in float16 - # Load our implementation using the loader from text_encoder/__init__.py - model1 = UMT5EncoderModel.from_pretrained(model_path).to(precision).to( - device).eval() - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - - from fastvideo.models.loader.component_loader import TextEncoderLoader - loader = TextEncoderLoader() - model2 = loader.load_model(model_path, hf_config, device) - - # Convert to float16 and move to device - model2 = model2.to(precision) - model2 = model2.to(device) - model2.eval() - - # Sanity check weights between the two models - logger.info("Comparing model weights for sanity check...") - params1 = dict(model1.named_parameters()) - params2 = dict(model2.named_parameters()) - - # Check number of parameters - logger.info(f"Model1 has {len(params1)} parameters") - logger.info(f"Model2 has {len(params2)} parameters") - - weight_diffs = [] - # check if embed_tokens are the same - weights = ["encoder.block.{}.layer.0.layer_norm.weight", "encoder.block.{}.layer.0.SelfAttention.relative_attention_bias.weight", \ - "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\ - "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \ - "encoder.block.{}.layer.1.layer_norm.weight", "encoder.final_layer_norm.weight", "shared.weight"] - # for (name1, param1), (name2, param2) in zip( - # sorted(params1.items()), sorted(params2.items()) - # ): - for l in range(hf_config.num_hidden_layers): - for w in weights: - name1 = w.format(l) - name2 = w.format(l) - p1 = params1[name1] - p2 = params2[name2] - assert p1.dtype == p2.dtype - try: - logger.info(f"Parameter: {name1} vs {name2}") - max_diff = torch.max(torch.abs(p1 - p2)).item() - mean_diff = torch.mean(torch.abs(p1 - p2)).item() - weight_diffs.append((name1, name2, max_diff, mean_diff)) - logger.info(f" Max diff: {max_diff}, Mean diff: {mean_diff}") - except Exception as e: - logger.info(f"Error comparing {name1} and {name2}: {e}") - - total_params = sum(p.numel() for p in model1.parameters()) - weight_sum_model1 = sum( - p.to(torch.float64).sum().item() for p in model1.parameters()) - weight_mean_model1 = weight_sum_model1 / total_params - print("Model 1 Weight Sum: ", weight_sum_model1) - print("Model 1 Weight Mean: ", weight_mean_model1) - - total_params = 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 - print("Model 2 Weight Sum: ", weight_sum_model2) - print("Model 2 Weight Mean: ", weight_mean_model2) - - # Test with some sample prompts - prompts = [ - "Once upon a time", "The quick brown fox jumps over", - "In a galaxy far, far away" - ] - - logger.info("Testing T5 encoder with sample prompts") - - with torch.no_grad(): - for prompt in prompts: - logger.info(f"Testing prompt: '{prompt}'") - - # Tokenize the prompt - tokens = tokenizer(prompt, - padding="max_length", - max_length=512, - truncation=True, - return_tensors="pt").to(device) - - # Get outputs from our implementation - # filter out padding input_ids - # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1] - # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1] - outputs1 = model1(input_ids=tokens.input_ids, - attention_mask=tokens.attention_mask, - output_hidden_states=True).last_hidden_state - print("--------------------------------") - logger.info("Testing model2") - - # Get outputs from HuggingFace implementation - outputs2 = model2( - input_ids=tokens.input_ids, - attention_mask=tokens.attention_mask, - ) - - # Compare last hidden states - last_hidden_state1 = outputs1[tokens.attention_mask == 1] - last_hidden_state2 = outputs2[tokens.attention_mask == 1] - - assert last_hidden_state1.shape == last_hidden_state2.shape, \ - f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" - - max_diff_hidden = torch.max( - torch.abs(last_hidden_state1 - last_hidden_state2)) - mean_diff_hidden = torch.mean( - torch.abs(last_hidden_state1 - last_hidden_state2)) - - logger.info( - f"Maximum difference in last hidden states: {max_diff_hidden.item()}" - ) - logger.info( - f"Mean difference in last hidden states: {mean_diff_hidden.item()}" - ) - - logger.info( - "Test passed! Both T5 encoder implementations produce similar outputs.") - logger.info("Test completed successfully") - - -if __name__ == "__main__": - test_t5_encoder() diff --git a/fastvideo/tests/old_tests/test_wan_vae.py b/fastvideo/tests/old_tests/test_wan_vae.py deleted file mode 100644 index b35e24657..000000000 --- a/fastvideo/tests/old_tests/test_wan_vae.py +++ /dev/null @@ -1,123 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -import json -import os - -import numpy as np -import torch -from diffusers import AutoencoderKLWan -from safetensors.torch import load_file - -from fastvideo.logger import init_logger -from fastvideo.models.vaes.wanvae import AutoencoderKLWan as MyWanVAE - -logger = init_logger(__name__) - - -def test_wan_vae(): - # Set fixed random seed for reproducibility - torch.manual_seed(42) - np.random.seed(42) - - device = torch.device("cuda:0") - # Initialize the two model implementations - path = "/workspace/data/Wan2.1-T2V-1.3B-Diffusers/vae" - config_path = os.path.join(path, "config.json") - config = json.load(open(config_path)) - config.pop("_class_name") - config.pop("_diffusers_version") - model1 = MyWanVAE(**config).to(torch.bfloat16) - - model2 = AutoencoderKLWan(**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 - # ) - - # Create identical inputs for both models - batch_size = 1 - - # Video input [B, C, T, H, W] - input_tensor = torch.randn(batch_size, - 3, - 81, - 32, - 32, - device=device, - dtype=torch.bfloat16) - latent_tensor = torch.randn(batch_size, - 16, - 21, - 32, - 32, - device=device, - dtype=torch.bfloat16) - - # Disable gradients for inference - with torch.no_grad(): - # Test encoding - logger.info("Testing encoding...") - latent2 = model2.encode(input_tensor).latent_dist.mean - print("--------------------------------") - latent1 = model1.encode(input_tensor).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 mean_diff_encode < 5e-1, f"Encoded latents differ significantly: mean diff = {mean_diff_encode.item()}" - # Test decoding - logger.info("Testing decoding...") - latent1 = latent2 = latent_tensor - latents_mean = (torch.tensor(model2.config.latents_mean).view( - 1, model2.config.z_dim, 1, 1, 1).to(latent2.device, latent2.dtype)) - latents_std = 1.0 / torch.tensor(model2.config.latents_std).view( - 1, model2.config.z_dim, 1, 1, 1).to(latent2.device, latent2.dtype) - latent2 = latent2 / latents_std + latents_mean - 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 mean_diff_decode < 1e-1, f"Decoded outputs differ significantly: mean diff = {mean_diff_decode.item()}" - - logger.info( - "Test passed! Both VAE implementations produce similar outputs.") - logger.info("Test completed successfully") - - -if __name__ == "__main__": - test_wan_vae() diff --git a/fastvideo/tests/old_tests/tp_example.py b/fastvideo/tests/old_tests/tp_example.py deleted file mode 100644 index c66dc79eb..000000000 --- a/fastvideo/tests/old_tests/tp_example.py +++ /dev/null @@ -1,152 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -import argparse -import os - -import torch -import torch.nn as nn - -from fastvideo.distributed.parallel_state import ( - cleanup_dist_env_and_memory, destroy_distributed_environment, - destroy_model_parallel, get_tp_rank, - get_tp_world_size, maybe_init_distributed_environment_and_model_parallel, get_world_group) -from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear -from fastvideo.logger import init_logger - -logger = init_logger(__name__) - - -class SimpleTPModel(nn.Module): - """A simple model that uses tensor parallelism.""" - - def __init__(self, hidden_size=1024, intermediate_size=4096): - super().__init__() - # Column parallel linear layer (splits output dimension) - self.fc1 = ColumnParallelLinear( - input_size=hidden_size, - output_size=intermediate_size, - bias=True, - gather_output= - False, # Don't gather output since we're passing to row parallel - skip_bias_add=False) - - # Row parallel linear layer (splits input dimension) - self.fc2 = RowParallelLinear( - input_size=intermediate_size, - output_size=hidden_size, - bias=True, - input_is_parallel=True, # Input is already split from previous layer - skip_bias_add=False) - - self.activation = nn.GELU() - - def forward(self, x): - # Forward through column parallel layer - hidden_states, _ = self.fc1(x) - - # Apply activation - hidden_states = self.activation(hidden_states) - - # Forward through row parallel layer - output, _ = self.fc2(hidden_states) - - return output - - -def initialize_random_weights(model, seed=42): - """Initialize the model with random weights using a fixed seed for reproducibility.""" - # Set seed for reproducibility - torch.manual_seed(seed) - - # Initialize weights for each layer - with torch.no_grad(): - # For ColumnParallelLinear layers - if hasattr(model, 'fc1'): - nn.init.normal_(model.fc1.weight, mean=0.0, std=0.02) - if model.fc1.bias is not None: - nn.init.zeros_(model.fc1.bias) - - # For RowParallelLinear layers - if hasattr(model, 'fc2'): - nn.init.normal_(model.fc2.weight, mean=0.0, std=0.02) - if model.fc2.bias is not None: - nn.init.zeros_(model.fc2.bias) - - logger.info("Model initialized with random weights") - return model - - -def setup_args(): - parser = argparse.ArgumentParser( - description='Simple Tensor Parallelism Example') - parser.add_argument('--tensor-model-parallel-size', - type=int, - default=8, - help='Degree of tensor model parallelism') - parser.add_argument('--batch-size', - type=int, - default=8, - help='Batch size for the example') - parser.add_argument('--hidden-size', - type=int, - default=1024, - help='Hidden size for the model') - parser.add_argument('--intermediate-size', - type=int, - default=4096, - help='Intermediate size for the model') - return parser.parse_args() - - -def main(): - args = setup_args() - - maybe_init_distributed_environment_and_model_parallel(args.tensor_model_parallel_size, args.tensor_model_parallel_size) - rank = get_world_group().rank - local_rank = get_world_group().local_rank - # Get tensor parallel info - tp_rank = get_tp_rank() - tp_world_size = get_tp_world_size() - - logger.info( - f"Process rank {rank} initialized with TP rank {tp_rank} in TP world size {tp_world_size}" - ) - - # Create a simple model - model = SimpleTPModel(hidden_size=args.hidden_size, - intermediate_size=args.intermediate_size) - - # Initialize with random weights - model = initialize_random_weights(model) - - # Create a random input tensor - batch_size = args.batch_size - hidden_size = args.hidden_size - x = torch.randn(batch_size, hidden_size, dtype=torch.float) - - # Move to GPU if available - device = torch.device( - f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu") - model = model.to(device) - x = x.to(device) - - # Forward pass - logger.info(f"Running forward pass on TP rank {tp_rank}") - with torch.no_grad(): - output = model(x) - - # Print output shape and statistics - logger.info(f"Output shape: {output.shape}") - logger.info( - f"Output mean: {output.mean().item()}, std: {output.std().item()}") - - # Clean up - logger.info("Cleaning up distributed environment") - destroy_model_parallel() - destroy_distributed_environment() - cleanup_dist_env_and_memory() - - logger.info("Example completed successfully") - - -if __name__ == "__main__": - main() diff --git a/fastvideo/tests/transformers/test_hunyuanvideo_load.py b/fastvideo/tests/transformers/test_hunyuanvideo_load.py index 5d83032e9..b223dcc52 100644 --- a/fastvideo/tests/transformers/test_hunyuanvideo_load.py +++ b/fastvideo/tests/transformers/test_hunyuanvideo_load.py @@ -25,8 +25,8 @@ 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)) + local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod + ) TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") CONFIG_PATH = os.path.join(TRANSFORMER_PATH, "config.json") diff --git a/fastvideo/tests/transformers/test_wanvideo.py b/fastvideo/tests/transformers/test_wanvideo.py index 60b0a8626..1901a1811 100644 --- a/fastvideo/tests/transformers/test_wanvideo.py +++ b/fastvideo/tests/transformers/test_wanvideo.py @@ -5,6 +5,7 @@ import numpy as np import pytest import torch from diffusers import WanTransformer3DModel +from torch.testing import assert_close from fastvideo.configs.pipelines import PipelineConfig from fastvideo.forward_context import set_forward_context @@ -23,8 +24,8 @@ os.environ["MASTER_PORT"] = "29503" BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, - local_dir=os.path.join( - 'data', BASE_MODEL_PATH)) + local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod + ) TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") @@ -120,10 +121,4 @@ def test_wan_transformer(): assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}" # Check if outputs are similar (allowing for small numerical differences) - max_diff = torch.max(torch.abs(output1 - output2)) - mean_diff = torch.mean(torch.abs(output1 - output2)) - logger.info("Max Diff: %s", max_diff.item()) - logger.info("Mean Diff: %s", mean_diff.item()) - assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}" - # mean diff - assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}" \ No newline at end of file + assert_close(output1, output2, atol=1e-1, rtol=1e-2) \ No newline at end of file diff --git a/fastvideo/tests/vaes/test_hunyuan_vae.py b/fastvideo/tests/vaes/test_hunyuan_vae.py index 4e77a4f65..f37b78f83 100644 --- a/fastvideo/tests/vaes/test_hunyuan_vae.py +++ b/fastvideo/tests/vaes/test_hunyuan_vae.py @@ -23,8 +23,8 @@ 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)) + local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod + ) VAE_PATH = os.path.join(MODEL_PATH, "vae") CONFIG_PATH = os.path.join(VAE_PATH, "config.json") diff --git a/fastvideo/tests/vaes/test_wan_vae.py b/fastvideo/tests/vaes/test_wan_vae.py index c96b31c00..f72417b45 100644 --- a/fastvideo/tests/vaes/test_wan_vae.py +++ b/fastvideo/tests/vaes/test_wan_vae.py @@ -12,6 +12,7 @@ from fastvideo.logger import init_logger from fastvideo.models.loader.component_loader import VAELoader from fastvideo.configs.models.vaes import WanVAEConfig from fastvideo.utils import maybe_download_model +from torch.testing import assert_close logger = init_logger(__name__) @@ -20,16 +21,16 @@ os.environ["MASTER_PORT"] = "29503" BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, - local_dir=os.path.join( - 'data', BASE_MODEL_PATH)) + local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod + ) VAE_PATH = os.path.join(MODEL_PATH, "vae") @pytest.mark.usefixtures("distributed_setup") def test_wan_vae(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - precision = torch.bfloat16 - precision_str = "bf16" + precision = torch.float32 + precision_str = "fp32" args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str)) args.device = device args.vae_cpu_offload = False @@ -70,13 +71,7 @@ def test_wan_vae(): # Check if latents have the same shape assert latent1.mean.shape == latent2.mean.shape, f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}" # Check if latents are similar - max_diff_encode = torch.max(torch.abs(latent1.mean - latent2.mean)) - mean_diff_encode = torch.mean(torch.abs(latent1.mean - latent2.mean)) - logger.info("Maximum difference between encoded latents: %s", - max_diff_encode.item()) - logger.info("Mean difference between encoded latents: %s", - mean_diff_encode.item()) - assert max_diff_encode < 1e-5, f"Encoded latents differ significantly: max diff = {mean_diff_encode.item()}" + assert_close(latent1.mean, latent2.mean, atol=1e-4, rtol=1e-4) # Test decoding logger.info("Testing decoding...") latent1_tensor = latent1.mode() @@ -98,10 +93,4 @@ def test_wan_vae(): 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("Maximum difference between decoded outputs: %s", - max_diff_decode.item()) - logger.info("Mean difference between decoded outputs: %s", - mean_diff_decode.item()) - assert max_diff_decode < 1e-5, f"Decoded outputs differ significantly: max diff = {mean_diff_decode.item()}" + assert_close(output1, output2, atol=1e-5, rtol=1e-3)