From 321e8bb47ff43d18cc81990fb076d0f6f042f89b Mon Sep 17 00:00:00 2001 From: Kyle Merritt Date: Thu, 8 Oct 2026 07:36:01 -0400 Subject: [PATCH] fix(quant): load offline-quantized compressed-tensors checkpoints directly on CPU under offload MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The diffusion loader treated any quant config without the MXFP8/ is_checkpoint_quantized markers as online quantization: it loaded every component on the accelerator before offloading, so serving the int8 weight-only compressed-tensors artifact OOM'd at model construction (encoder 17 GiB + VAE + packed DiT skeleton > 24 GiB). compressed-tensors is always a serialized format with no on-load quantization pass — detect it (through component wrappers too) and load on the host, round-tripping each quantized module through the accelerator only for the post-load Marlin repack, which has no CPU path. --- tests/quantization/test_quant_factory.py | 16 ++++++++++++++ .../model_loader/diffusers_loader.py | 22 ++++++++++++++----- vllm_omni/quantization/component_config.py | 20 +++++++++++++++++ 3 files changed, 52 insertions(+), 6 deletions(-) diff --git a/tests/quantization/test_quant_factory.py b/tests/quantization/test_quant_factory.py index 09921acd5..15636b71a 100644 --- a/tests/quantization/test_quant_factory.py +++ b/tests/quantization/test_quant_factory.py @@ -12,6 +12,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso from vllm_omni.quantization import build_quant_config, resolve_component_quant_config from vllm_omni.quantization.component_config import ( ComponentQuantizationConfig, + is_prequantized_config, resolve_encoder_quant_config, ) from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS, _build_single @@ -75,3 +76,18 @@ def test_unknown_method_still_raises() -> None: def test_supported_list_advertises_hyphenated_registry_names() -> None: assert "compressed-tensors" in SUPPORTED_QUANTIZATION_METHODS + + +def test_is_prequantized_config_recognizes_serialized_formats() -> None: + # compressed-tensors is always a serialized format: its linear methods + # consume packed checkpoint tensors and there is no on-load quantization. + assert is_prequantized_config(build_quant_config(_int8_weight_only_disk_config())) + component = build_quant_config({"transformer": _int8_weight_only_disk_config()}) + assert isinstance(component, ComponentQuantizationConfig) + assert is_prequantized_config(component) + assert not is_prequantized_config(None) + # Online (non-serialized) methods stay unrecognized so the loader keeps + # the accelerator-side load-quantize-offload path for them. + from vllm_omni.quantization.fp8_config import DiffusionFp8Config + + assert not is_prequantized_config(DiffusionFp8Config()) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index c74808efb..a142072f5 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -63,6 +63,7 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery, Pipe from vllm_omni.diffusion.offloader.offload_plan import get_offload_plan from vllm_omni.diffusion.registry import initialize_model from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.quantization.component_config import is_prequantized_config from vllm_omni.transformers_utils.repo_utils import hf_api @@ -695,16 +696,17 @@ class DiffusersPipelineLoader(HWRLoaderMixin): } if load_format is None: load_format = "default" - # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8), - # weights are already quantized in the checkpoint — load directly on CPU. - # For online quantization, load on device so quantization can run on accelerator, - # then move back to CPU afterward. + # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8, + # compressed-tensors), weights are already quantized in the checkpoint — load + # directly on CPU. For online quantization, load on device so quantization can + # run on accelerator, then move back to CPU afterward. offload_after_quant = False + offline_quant_direct_cpu = False if load_device == "cpu" and self.quant_config is not None and device is not None: quant_cfg = self.quant_config is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr( quant_cfg, "is_checkpoint_quantized", False - ) + ) or is_prequantized_config(quant_cfg) if not is_offline: load_device = device.type offload_after_quant = True @@ -713,6 +715,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin): load_device, ) else: + offline_quant_direct_cpu = True logger.info("Offline-quantized model with CPU offload, loading weights directly on CPU") target_device = torch.device(load_device) @@ -874,7 +877,14 @@ class DiffusersPipelineLoader(HWRLoaderMixin): else: self.load_weights(model) self._maybe_fuse_distilled_lora(model) - self._process_weights_after_loading(model, target_device) + # Post-load transforms (e.g. the Marlin repack for + # compressed-tensors) only run on the accelerator, so an + # offline-quantized CPU-offload load round-trips each + # quantized module through the device and back to the host. + process_device = ( + device if offline_quant_direct_cpu and device is not None else target_device + ) + self._process_weights_after_loading(model, process_device) # A warm final-layout hit has already completed all # byte-changing work through the restorer. Shared runtime diff --git a/vllm_omni/quantization/component_config.py b/vllm_omni/quantization/component_config.py index bf6f3b6b2..5ade0fea2 100644 --- a/vllm_omni/quantization/component_config.py +++ b/vllm_omni/quantization/component_config.py @@ -33,6 +33,26 @@ PRE_QUANTIZED_METHODS: frozenset[str] = frozenset( ) +def is_prequantized_config(quant_config: QuantizationConfig | None) -> bool: + """True when the config can only describe pre-quantized (serialized) weights. + + Such checkpoints need no accelerator-side quantization pass at load, so a + CPU-offload load can read weights directly on the host. compressed-tensors + is the canonical case: its linear methods consume packed checkpoint tensors + (weight_packed/weight_scale) directly and vLLM has no on-load quantization + path for the format. A component wrapper is pre-quantized when every routed + component (and the default) is itself pre-quantized or unquantized. + """ + if quant_config is None: + return False + if isinstance(quant_config, ComponentQuantizationConfig): + routed = [*quant_config.component_configs.values(), quant_config.default_config] + return any(c is not None for c in routed) and all( + c is None or is_prequantized_config(c) for c in routed + ) + return quant_config.get_name() == "compressed-tensors" + + def resolve_component_quant_config( quant_config: QuantizationConfig | None, component: str,