hao-ai-lab-FastVideo/fastvideo/tests/encoders/test_llama_encoder.py

# SPDX-License-Identifier: Apache-2.0
import os

import numpy as np
import pytest
import torch
from transformers import AutoConfig, AutoTokenizer, LlamaModel
import gc
from fastvideo.configs.pipelines import HunyuanConfig
from fastvideo.forward_context import set_forward_context
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.logger import init_logger
from fastvideo.models.loader.component_loader import TextEncoderLoader
from fastvideo.utils import maybe_download_model
from fastvideo.configs.models.encoders import LlamaConfig
from torch.distributed.tensor import DTensor
from torch.testing import assert_close

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))
TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder")
TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer")


@pytest.mark.usefixtures("distributed_setup")
def test_llama_encoder():
    """
    Tests compatibility between two different implementations for loading text encoders:
    1. load_text_encoder from fastvideo.models.hunyuan.text_encoder
    2. TextEncoderLoader from fastvideo.models.loader

    The test verifies that both implementations:
    - Load models with the same weights and parameters
    - Produce nearly identical outputs for the same input prompts
    """
    args = FastVideoArgs(model_path="meta-llama/Llama-2-7b-hf", pipeline_config=HunyuanConfig())

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # Initialize the two model implementations
    logger.info("Loading models from %s", args.model_path)
    hf_config = AutoConfig.from_pretrained(TEXT_ENCODER_PATH)
    print(hf_config)

    # Load HuggingFace implementation
    model1 = LlamaModel.from_pretrained(TEXT_ENCODER_PATH).to(torch.float16).to(device).eval()
    loader = TextEncoderLoader()
    device = torch.device("cuda:0")
    model2 = loader.load(TEXT_ENCODER_PATH, args)

    # Convert to float16 and move to device
    # model2 = model2.to(torch.float16)
    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 %d parameters", len(params1))
    logger.info("Model2 has %d parameters", len(params2))

    # check if embed_tokens are the same
    device = model1.embed_tokens.weight.device
    assert torch.allclose(
        model1.embed_tokens.weight,
        model2.embed_tokens.weight.to_local().to(device)
        if isinstance(model2.embed_tokens.weight, DTensor)
        else model2.embed_tokens.weight.to(device),
    )
    weights = ["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
        skip = False
        for param_name, weight_name, shard_id in model2.config.arch_config.stacked_params_mapping:
            if weight_name not in name1:
                skip = True
        # stacked params are more troublesome
        if skip:
            continue
        param2 = params2[name2]
        param2 = param2.to_local().to(device) if isinstance(param2, DTensor) else param2.to(device)
        assert_close(param1, param2, atol=1e-4, rtol=1e-4)
    gc.collect()
    torch.cuda.empty_cache()

    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)

    # 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 LLaMA 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 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, output_hidden_states=True)
            print("--------------------------------")
            logger.info("Testing model2")

            # Get outputs from HuggingFace implementation
            with set_forward_context(current_timestep=0, attn_metadata=None):
                outputs2 = model2(
                    input_ids=tokens.input_ids, attention_mask=tokens.attention_mask, output_hidden_states=True
                )

            # Compare last hidden states
            last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask == 1]
            last_hidden_state2 = outputs2.last_hidden_state[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-1, rtol=1e-4)