48 lines
2.0 KiB
Python
48 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Omni's online FP8 option backed by vLLM's online quantization API."""
|
|
|
|
from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
|
|
from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
|
|
|
|
|
|
class DiffusionFp8Config(OnlineQuantizationConfig):
|
|
"""Keep the diffusion ``fp8`` option for quantizing BF16/FP16 weights.
|
|
|
|
Upstream Fp8Config now describes serialized checkpoints only. Keeping the
|
|
Omni method name also lets checkpoint metadata replace this online config
|
|
through the existing serialized-checkpoint reconciliation path.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
is_checkpoint_fp8_serialized: bool = False,
|
|
activation_scheme: str = "dynamic",
|
|
ignored_layers: list[str] | None = None,
|
|
weight_block_size: list[int] | None = None,
|
|
store_dtype: str | None = None,
|
|
) -> None:
|
|
if is_checkpoint_fp8_serialized:
|
|
raise ValueError("Serialized FP8 checkpoints require upstream Fp8Config")
|
|
if activation_scheme != "dynamic":
|
|
raise ValueError("Online FP8 quantization requires activation_scheme='dynamic'")
|
|
if weight_block_size is not None:
|
|
raise ValueError("Block-wise FP8 quantization requires a serialized checkpoint or fp8_per_block")
|
|
if store_dtype is not None:
|
|
raise ValueError("Online FP8 quantization does not support store_dtype")
|
|
super().__init__(
|
|
QuantizationConfigArgs(
|
|
linear=QuantSpec(weight="fp8_per_tensor_static"),
|
|
moe=QuantSpec(weight="fp8_per_tensor_static"),
|
|
ignore=list(ignored_layers or []),
|
|
)
|
|
)
|
|
self.is_checkpoint_fp8_serialized = False
|
|
self.activation_scheme = activation_scheme
|
|
self.weight_block_size = None
|
|
self.store_dtype = None
|
|
|
|
@classmethod
|
|
def get_name(cls) -> str:
|
|
return "fp8"
|