fix(quant): load offline-quantized compressed-tensors checkpoints directly on CPU under offload
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.
This commit is contained in:
1 parent
e778d21f91
commit
321e8bb47f
3 files changed
+52
-6
No files matched your search
@@ -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())
|
||||
Reference in new issue
Block a user