116 lines
4.2 KiB
Python
116 lines
4.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""E2E tests for FLUX AutoRound W4A16 quantized inference.
|
|
|
|
These tests require:
|
|
- A CUDA GPU
|
|
- The quantized model checkpoint (vllm-project-org/FLUX.1-dev-AutoRound-w4a16)
|
|
"""
|
|
|
|
import gc
|
|
import os as _os
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
|
|
|
|
from tests.conftest import OmniRunner
|
|
from tests.utils import DeviceMemoryMonitor, hardware_test
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
QUANTIZED_MODEL = "vllm-project-org/FLUX.1-dev-AutoRound-w4a16"
|
|
BASELINE_MODEL = "black-forest-labs/FLUX.1-dev"
|
|
|
|
QUANTIZED_MODEL = _os.environ.get("FLUX_AUTOROUND_MODEL", QUANTIZED_MODEL)
|
|
BASELINE_MODEL = _os.environ.get("FLUX_BASELINE_MODEL", BASELINE_MODEL)
|
|
|
|
# Small resolution to keep GPU memory & time manageable
|
|
HEIGHT = 256
|
|
WIDTH = 256
|
|
NUM_STEPS = 2 # minimal for smoke-test
|
|
|
|
|
|
def _generate_image(model_name: str, **extra_kwargs) -> tuple[list, float]:
|
|
"""Load a FLUX model, generate one image, return (images, peak_memory_mb)."""
|
|
gc.collect()
|
|
current_omni_platform.empty_cache()
|
|
device_index = current_omni_platform.current_device()
|
|
current_omni_platform.reset_peak_memory_stats()
|
|
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
|
|
monitor.start()
|
|
|
|
with OmniRunner(model_name, enforce_eager=True, **extra_kwargs) as runner:
|
|
current_omni_platform.reset_peak_memory_stats()
|
|
outputs = runner.omni.generate(
|
|
"a photo of a cat sitting on a laptop keyboard",
|
|
OmniDiffusionSamplingParams(
|
|
height=HEIGHT,
|
|
width=WIDTH,
|
|
num_inference_steps=NUM_STEPS,
|
|
guidance_scale=0.0,
|
|
generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
|
|
),
|
|
)
|
|
|
|
peak = monitor.peak_used_mb
|
|
monitor.stop()
|
|
|
|
first_output = outputs[0]
|
|
assert first_output.final_output_type == "image"
|
|
req_out = first_output.request_output
|
|
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
|
|
images = req_out.images
|
|
|
|
gc.collect()
|
|
current_omni_platform.empty_cache()
|
|
|
|
return images, peak
|
|
|
|
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.diffusion
|
|
@hardware_test(res={"cuda": "L4"})
|
|
def test_flux_autoround_w4a16_generates_image():
|
|
"""Load the W4A16 quantized FLUX model and verify it produces a valid image."""
|
|
images, _ = _generate_image(QUANTIZED_MODEL)
|
|
|
|
assert len(images) >= 1, "Expected at least one generated image"
|
|
img = images[0]
|
|
assert img.width == WIDTH, f"Expected width {WIDTH}, got {img.width}"
|
|
assert img.height == HEIGHT, f"Expected height {HEIGHT}, got {img.height}"
|
|
|
|
# Sanity: image should not be blank (all-zero)
|
|
import numpy as np
|
|
|
|
arr = np.array(img)
|
|
assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"
|
|
|
|
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.diffusion
|
|
@hardware_test(res={"cuda": "L4"})
|
|
def test_flux_autoround_w4a16_memory_savings():
|
|
"""Compare peak GPU memory of quantized vs FP16 baseline.
|
|
|
|
The W4A16 model should use meaningfully less memory than the
|
|
BF16/FP16 baseline since weights are 4-bit instead of 16-bit.
|
|
"""
|
|
quant_images, quant_peak = _generate_image(QUANTIZED_MODEL)
|
|
cleanup_dist_env_and_memory()
|
|
_, baseline_peak = _generate_image(BASELINE_MODEL)
|
|
|
|
print(f"Quantized (W4A16) peak memory: {quant_peak:.0f} MB")
|
|
print(f"Baseline (BF16) peak memory: {baseline_peak:.0f} MB")
|
|
print(f"Savings: {baseline_peak - quant_peak:.0f} MB")
|
|
|
|
# W4A16 weights are 4x smaller than BF16/FP16. FLUX.1-dev transformer
|
|
# is ~12 GB in BF16, so we expect ~9 GB savings on weights alone.
|
|
# Use a conservative threshold to account for activations and overhead.
|
|
min_savings_mb = 2000
|
|
assert quant_peak + min_savings_mb < baseline_peak, (
|
|
f"Quantized model ({quant_peak:.0f} MB) should use at least "
|
|
f"{min_savings_mb} MB less than baseline ({baseline_peak:.0f} MB)"
|
|
)
|