Use assert_close in tests (#429)
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
This commit is contained in:
Wenxuan Tan
co-authored by
SolitaryThinker
SolitaryThinker
parent
0082bc66fc
commit
754292c419
@@ -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__"]
|
||||
__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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()}"
|
||||
assert_close(output1, output2, atol=1e-1, rtol=1e-2)
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user