test_vllm.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD 3-Clause license found in the
# LICENSE file in the root directory of this source tree.

import importlib.util
import os
import random
import shutil
from pathlib import Path
from typing import List

import numpy as np
import pytest
import torch

from torchao.utils import TORCH_VERSION_AT_LEAST_2_7

if not TORCH_VERSION_AT_LEAST_2_7:
    pytest.skip("Requires PyTorch 2.7 or higher", allow_module_level=True)


VLLM_AVAILABLE = importlib.util.find_spec("vllm") is not None
TRANSFORMERS_AVAILABLE = importlib.util.find_spec("transformers") is not None

if not VLLM_AVAILABLE:
    pytest.skip("vLLM not installed", allow_module_level=True)

if not TRANSFORMERS_AVAILABLE:
    pytest.skip("transformers not installed", allow_module_level=True)

from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
from vllm import LLM, SamplingParams

from torchao.quantization.granularity import PerRow, PerTensor
from torchao.quantization.quant_api import (
    CutlassInt4PackedLayout,
    Float8DynamicActivationFloat8WeightConfig,
    Int8DynamicActivationInt4WeightConfig,
    Int8WeightOnlyConfig,
)


def get_tests() -> List[TorchAoConfig]:
    """Get all the tests based off of device info"""

    # Helper objects for granularity
    per_tensor = PerTensor()
    per_row = PerRow()

    BASE_TESTS = [TorchAoConfig(Int8WeightOnlyConfig())]
    SM89_TESTS = [
        TorchAoConfig(
            Float8DynamicActivationFloat8WeightConfig(granularity=per_tensor)
        ),
        TorchAoConfig(Float8DynamicActivationFloat8WeightConfig(granularity=per_row)),
    ]
    SM90_ONLY_TESTS = [
        TorchAoConfig(
            Int8DynamicActivationInt4WeightConfig(layout=CutlassInt4PackedLayout())
        )
    ]
    SM100_TESTS = [
        # TorchAoConfig(MXFPInferenceConfig())
    ]  # Failing for : https://github.com/pytorch/ao/issues/2239

    # Check CUDA availability first
    if not torch.cuda.is_available():
        return []  # No CUDA, no tests

    major, minor = torch.cuda.get_device_capability()

    # Build test list based on compute capability
    all_tests = []

    # Always include base tests if we have CUDA
    all_tests.extend(BASE_TESTS)

    # Add SM89+ tests
    if major > 8 or (major == 8 and minor >= 9):
        all_tests.extend(SM89_TESTS)

    # Add SM100+ tests
    if major >= 10:
        all_tests.extend(SM100_TESTS)

    # Only work for sm 90
    if major == 9:
        all_tests.extend(SM90_ONLY_TESTS)

    return all_tests


class TestVLLMIntegration:
    """Integration tests for vLLM with quantized models."""

    @classmethod
    def setup_class(cls):
        """Set up test environment."""
        # Set seeds for reproducibility
        cls.set_seed(42)

        # See https://github.com/pytorch/ao/issues/2239 for details
        os.environ["VLLM_TEST_STANDALONE_COMPILE"] = "1"
        # For Small testing this makes it faster
        os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"

    @classmethod
    def teardown_class(cls):
        """Clean up after all tests."""
        torch.cuda.empty_cache()
        import gc

        gc.collect()

    def setup_method(self, method):
        """Clean up before each test method."""
        torch.cuda.empty_cache()
        import gc

        gc.collect()

    def teardown_method(self, method):
        """Clean up after each test method."""
        torch.cuda.empty_cache()
        import gc

        gc.collect()

    @staticmethod
    def set_seed(seed):
        """Set random seeds for reproducibility."""
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

    def quantize_and_save_model(
        self,
        model_name: str,
        quantization_config: TorchAoConfig,
        output_dir: Path,
    ):
        """Quantize a model and save it to disk."""
        # Load and quantize model
        quantized_model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype="bfloat16",
            device_map="cuda",
            quantization_config=quantization_config,
        )
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        # Save quantized model
        quantized_model.save_pretrained(output_dir, safe_serialization=False)
        tokenizer.save_pretrained(output_dir)

        # Clean up to free memory
        del quantized_model
        torch.cuda.empty_cache()

        return output_dir

    def cleanup_model_directory(self, model_path: Path):
        """Clean up the model directory safely."""
        try:
            if model_path.exists() and model_path.is_dir():
                shutil.rmtree(model_path)
        except (OSError, PermissionError) as e:
            # Log the error but don't fail the test
            print(f"Warning: Failed to clean up {model_path}: {e}")

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
    @pytest.mark.skipif(not VLLM_AVAILABLE, reason="vLLM not installed")
    @pytest.mark.parametrize(
        "quantization_config", get_tests(), ids=lambda config: f"{config.quant_type}"
    )
    @pytest.mark.parametrize("compile", [True, False])
    @pytest.mark.parametrize(
        "tp_size", [1, 2] if torch.cuda.device_count() > 1 else [1]
    )
    def test_vllm_smoke_test(self, tmp_path, quantization_config, compile, tp_size):
        """Test vLLM generation with quantized models."""
        # Skip per_row tests if not supported
        torch._dynamo.reset()

        # Use a small model for testing
        base_model = "facebook/opt-125m"

        # Create a descriptive name for the output directory
        config_name = str(quantization_config).replace("/", "_").replace(" ", "_")[:50]
        output_dir = tmp_path / f"{config_name}-opt-125m"

        llm = None
        quantized_model_path = None

        try:
            # Quantize the model
            quantized_model_path = self.quantize_and_save_model(
                base_model, quantization_config, output_dir
            )

            # Test generation with vLLM
            sampling_params = SamplingParams(
                temperature=0.8,
                top_p=0.95,
                seed=42,
                max_tokens=16,  # Small for testing
            )

            # Create LLM instance
            llm = LLM(
                model=str(quantized_model_path),
                tensor_parallel_size=tp_size,
                enforce_eager=not compile,
                dtype="bfloat16",
                num_gpu_blocks_override=128,
            )

            # Test prompts
            prompts = [
                "Hello, my name is",
                "The capital of France is",
            ]

            # Generate outputs
            outputs = llm.generate(prompts, sampling_params)

            # Verify outputs
            assert len(outputs) == len(prompts)
            for output in outputs:
                assert output.prompt in prompts
                assert len(output.outputs) > 0
                generated_text = output.outputs[0].text
                assert isinstance(generated_text, str)
                assert len(generated_text) > 0

        finally:
            # Clean up resources
            if llm is not None:
                del llm

            # Clean up CUDA memory
            torch.cuda.empty_cache()

            # Clean up the saved model directory
            if quantized_model_path is not None:
                self.cleanup_model_directory(quantized_model_path)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])