215 lines
9.5 KiB
Python
215 lines
9.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
import os
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from torch.distributed.tensor import DTensor
|
|
from torch.testing import assert_close
|
|
from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel, T5EncoderModel
|
|
|
|
from fastvideo.configs.pipelines import CosmosConfig, PipelineConfig, WanT2V480PConfig
|
|
from fastvideo.forward_context import set_forward_context
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.models.loader.component_loader import TextEncoderLoader
|
|
from fastvideo.tests.utils import skip_if_gated_repo_inaccessible
|
|
from fastvideo.utils import maybe_download_model, PRECISION_TO_TYPE
|
|
from fastvideo.fastvideo_args import FastVideoArgs
|
|
from fastvideo.configs.models.encoders import T5Config, T5LargeConfig
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
os.environ["MASTER_ADDR"] = "localhost"
|
|
os.environ["MASTER_PORT"] = "29503"
|
|
|
|
|
|
@pytest.fixture
|
|
def t5_model_paths_and_config():
|
|
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))
|
|
text_encoder_path = os.path.join(model_path, "text_encoder")
|
|
tokenizer_path = os.path.join(model_path, "tokenizer")
|
|
return text_encoder_path, tokenizer_path, WanT2V480PConfig()
|
|
|
|
|
|
@pytest.fixture
|
|
def t5_large_model_paths_and_config():
|
|
base_model_path = "nvidia/Cosmos-Predict2-2B-Video2World"
|
|
local_dir = os.path.join('data', base_model_path)
|
|
skip_if_gated_repo_inaccessible(base_model_path, local_path=local_dir, test_name="Cosmos T5-large encoder test")
|
|
model_path = maybe_download_model(base_model_path, local_dir=local_dir)
|
|
text_encoder_path = os.path.join(model_path, "text_encoder")
|
|
tokenizer_path = os.path.join(model_path, "tokenizer")
|
|
return text_encoder_path, tokenizer_path, CosmosConfig()
|
|
|
|
|
|
@pytest.mark.usefixtures("distributed_setup")
|
|
def test_t5_encoder(t5_model_paths_and_config):
|
|
# Initialize the two model implementations
|
|
text_encoder_path, tokenizer_path, pipeline_config = t5_model_paths_and_config
|
|
hf_config = AutoConfig.from_pretrained(text_encoder_path)
|
|
print(hf_config)
|
|
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
precision_str = "fp32"
|
|
precision = PRECISION_TO_TYPE[precision_str]
|
|
model1 = UMT5EncoderModel.from_pretrained(text_encoder_path).to(precision).to(device).eval()
|
|
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
|
|
|
args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False)
|
|
loader = TextEncoderLoader()
|
|
model2 = loader.load(text_encoder_path, args)
|
|
model2 = model2.to(precision)
|
|
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("Model1 has %s parameters", len(params1))
|
|
logger.info("Model2 has %s parameters", len(params2))
|
|
|
|
# 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"]
|
|
|
|
for idx in range(hf_config.num_hidden_layers):
|
|
for w in weights:
|
|
name1 = w.format(idx)
|
|
name2 = w.format(idx)
|
|
p1 = params1[name1]
|
|
p2 = params2[name2]
|
|
p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1)
|
|
assert_close(p1, p2, atol=1e-4, rtol=1e-4)
|
|
|
|
# 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("Testing prompt: %s", prompt)
|
|
|
|
# Tokenize the prompt
|
|
tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True,
|
|
return_tensors="pt").to(device)
|
|
|
|
# Get outputs from HuggingFace 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 our implementation
|
|
with set_forward_context(current_timestep=0, attn_metadata=None):
|
|
outputs2 = model2(
|
|
input_ids=tokens.input_ids,
|
|
attention_mask=tokens.attention_mask,
|
|
).last_hidden_state
|
|
|
|
# Compare last hidden states
|
|
last_hidden_state1 = outputs1[tokens.attention_mask == 1]
|
|
last_hidden_state2 = outputs2[tokens.attention_mask == 1]
|
|
|
|
assert_close(last_hidden_state1, last_hidden_state2, atol=1e-4, rtol=1e-4)
|
|
|
|
|
|
@pytest.mark.usefixtures("distributed_setup")
|
|
def test_t5_large_encoder(t5_large_model_paths_and_config):
|
|
# Initialize the two model implementations
|
|
text_encoder_path, tokenizer_path, pipeline_config = t5_large_model_paths_and_config
|
|
hf_config = AutoConfig.from_pretrained(text_encoder_path)
|
|
print(hf_config)
|
|
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
precision_str = "fp32"
|
|
precision = PRECISION_TO_TYPE[precision_str]
|
|
model1 = T5EncoderModel.from_pretrained(text_encoder_path).to(precision).to(device).eval()
|
|
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
|
|
|
args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False)
|
|
loader = TextEncoderLoader()
|
|
model2 = loader.load(text_encoder_path, args)
|
|
model2 = model2.to(precision)
|
|
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("Model1 has %s parameters", len(params1))
|
|
logger.info("Model2 has %s parameters", len(params2))
|
|
|
|
# Print parameter names for comparison
|
|
logger.info("Model1 parameters:")
|
|
for name in sorted(params1.keys()):
|
|
logger.info(" %s: %s", name, params1[name].shape)
|
|
|
|
logger.info("Model2 parameters:")
|
|
for name in sorted(params2.keys()):
|
|
logger.info(" %s: %s", name, params2[name].shape)
|
|
|
|
#check if embed_tokens are the same
|
|
weights = ["encoder.block.{}.layer.0.layer_norm.weight", "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight", \
|
|
"encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi.weight", \
|
|
"encoder.block.{}.layer.1.DenseReluDense.wo.weight", "encoder.final_layer_norm.weight"]
|
|
|
|
for idx in range(hf_config.num_hidden_layers):
|
|
for w in weights:
|
|
name1 = w.format(idx)
|
|
name2 = w.format(idx)
|
|
p1 = params1[name1]
|
|
p2 = params2[name2]
|
|
p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1)
|
|
assert_close(p1, p2, atol=1e-4, rtol=1e-4)
|
|
|
|
# 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 Large encoder with sample prompts")
|
|
|
|
with torch.no_grad():
|
|
for prompt in prompts:
|
|
logger.info("Testing prompt: %s", prompt)
|
|
|
|
# Tokenize the prompt
|
|
tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True,
|
|
return_tensors="pt").to(device)
|
|
|
|
# Get outputs from HuggingFace 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 with T5LargeConfig")
|
|
|
|
# Get outputs from our implementation
|
|
with set_forward_context(current_timestep=0, attn_metadata=None):
|
|
outputs2 = model2(
|
|
input_ids=tokens.input_ids,
|
|
attention_mask=tokens.attention_mask,
|
|
).last_hidden_state
|
|
|
|
# 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}"
|
|
|
|
assert_close(last_hidden_state1, last_hidden_state2, atol=1e-4, rtol=1e-4)
|