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
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user