1dc991809b
The public LTX-2 path was previously running full bf16 even when callers set ``pipeline_config.dit_config.quant_config = FP4Config()``, because: * the DiT used plain ``torch.nn.Linear`` for FP4-eligible projections (``attn1/attn2.to_q/k/v/out``, ``audio_to_video_attn``, ``video_to_audio_attn``, FFN ``fc_in``/``fc_out``), so ``LinearBase.quant_method.apply`` was never reached; * the loader did not call ``convert_model_to_fp4`` after weights loaded, so even matched layers had no ``_fp4_weight*`` buffers; * ``QuantizationMethods`` did not register ``"FP4"``, so the typed ``engine.quantization.transformer_quant`` surface had no way to select it. This commit completes the inference wire-up through the existing ``fastvideo.layers.quantization`` registry — no parallel pathway: * swap ``nn.Linear`` → ``ReplicatedLinear`` for the LTX-2 attention ``to_q``/``to_k``/``to_v``/``to_out``/``to_gate_compress`` and the FFN ``fc_in``/``fc_out`` (``GELUApprox.proj`` and ``FeedForward`` ``project_out``). Other linears (timestep MLP, caption proj, patchify, AdaLN scale/shift) stay ``nn.Linear`` to match internal. * port ``_supports_prequantized_input`` and ``_linear_project_with_optional_prequant`` so the attention forward quantizes input once and reuses the ``(x_fp4, x_scale, x_global_sf)`` tuple across q/k/v projections — matches internal exactly. * plumb ``quant_config`` and ``prefix`` through ``BasicAVTransformerBlock`` → ``_init_transformer_blocks`` → ``LTXModel`` → ``LTX2Transformer3DModel`` so each ``ReplicatedLinear`` gets the correct ``ltx2.blocks.<i>.<sub>.<proj>`` prefix that ``FP4Config.get_quant_method`` matches against. * register ``"FP4"`` in ``QuantizationMethods`` literal + ``get_quantization_config`` registry so typed ``engine.quantization.transformer_quant: "FP4"`` resolves to a concrete ``FP4Config()`` instance. * add ``transformer_quant`` carrier on ``FastVideoArgs`` and ``__post_init__._apply_transformer_quant`` to pin the resolved quant config onto ``dit_config.quant_config`` (without overwriting an explicit setter on the dit_config). * add ``_maybe_convert_model_to_fp4`` in ``fsdp_load.py`` that walks the loaded model once and registers ``_fp4_weight*``/``_fp4_alpha``/``_weight_global_sf`` buffers on layers whose ``quant_method`` is ``FP4QuantizeMethod``. flashinfer is imported lazily inside ``convert_model_to_fp4`` so this is a no-op on hosts without the FP4 kernels. * restore ``LinearBase`` fallback to ``UnquantizedLinearMethod`` when ``quant_config.get_quant_method`` returns ``None`` — ``FP4Config`` only tags a curated subset of layers, and the previous ``assert quant_method is not None`` would crash any non-tagged layer that received a quant_config. Non-FP4 callers are unaffected: ``UnquantizedLinearMethod.apply`` runs the standard ``F.linear`` path, and parameter names (``to_q.weight``, ``to_out.0.weight`` from ``nn.ModuleList``) are identical to the previous ``nn.Sequential`` layout. Co-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com> Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com> Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com> Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com>
72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
from typing import Literal, get_args
|
|
|
|
from fastvideo.layers.quantization.base_config import QuantizationConfig
|
|
|
|
QuantizationMethods = Literal[None, "AbsMaxFP8", "FP4"]
|
|
|
|
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
|
|
|
|
# The customized quantization methods which will be added to this dict.
|
|
_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {}
|
|
|
|
|
|
def register_quantization_config(quantization: str):
|
|
"""Register a customized vllm quantization config.
|
|
|
|
When a quantization method is not supported by vllm, you can register a customized
|
|
quantization config to support it.
|
|
|
|
Args:
|
|
quantization (str): The quantization method name.
|
|
|
|
Examples:
|
|
>>> from fastvideo.layers.quantization import register_quantization_config
|
|
>>> from fastvideo.layers.quantization import get_quantization_config
|
|
>>> from fastvideo.layers.quantization.base_config import QuantizationConfig
|
|
>>>
|
|
>>> @register_quantization_config("my_quant")
|
|
... class MyQuantConfig(QuantizationConfig):
|
|
... pass
|
|
>>>
|
|
>>> get_quantization_config("my_quant")
|
|
<class 'MyQuantConfig'>
|
|
""" # noqa: E501
|
|
|
|
def _wrapper(quant_config_cls):
|
|
if quantization in QUANTIZATION_METHODS:
|
|
raise ValueError(f"The quantization method `{quantization}` is already exists.")
|
|
if not issubclass(quant_config_cls, QuantizationConfig):
|
|
raise ValueError("The quantization config must be a subclass of "
|
|
"`QuantizationConfig`.")
|
|
_CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls
|
|
QUANTIZATION_METHODS.append(quantization)
|
|
return quant_config_cls
|
|
|
|
return _wrapper
|
|
|
|
|
|
def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
|
|
if quantization not in QUANTIZATION_METHODS:
|
|
raise ValueError(f"Invalid quantization method: {quantization}")
|
|
|
|
# lazy import to avoid triggering `torch.compile` too early
|
|
from .absmax_fp8 import AbsMaxFP8Config
|
|
from .fp4_config import FP4Config
|
|
|
|
method_to_config: dict[str, type[QuantizationConfig]] = {
|
|
"AbsMaxFP8": AbsMaxFP8Config,
|
|
"FP4": FP4Config,
|
|
}
|
|
# Update the `method_to_config` with customized quantization methods.
|
|
method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
|
|
|
|
return method_to_config[quantization]
|
|
|
|
|
|
all = [
|
|
"QuantizationMethods",
|
|
"QuantizationConfig",
|
|
"get_quantization_config",
|
|
"QUANTIZATION_METHODS",
|
|
]
|