diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index f4a6a6681..6cd5766e3 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -35,6 +35,7 @@ surfaces: prompt_txt: request.inputs.prompt_path override_text_encoder_safetensors: generator.pipeline.components.text_encoder_weights override_text_encoder_quant: generator.engine.quantization.text_encoder_quant + transformer_quant: generator.engine.quantization.transformer_quant override_transformer_cls_name: generator.pipeline.components.override_transformer_cls_name init_weights_from_safetensors: generator.pipeline.components.transformer_weights init_weights_from_safetensors_2: generator.pipeline.components.transformer_2_weights diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 7047e10bf..5c055f33e 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -65,6 +65,7 @@ _FROM_PRETRAINED_CONVENIENCE_KWARGS = frozenset({ "pin_cpu_memory", "enable_torch_compile", "torch_compile_kwargs", + "transformer_quant", }) diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 55e079a76..f30f88a7f 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -172,6 +172,7 @@ class FastVideoArgs: override_text_encoder_safetensors: str | None = None # path to safetensors file for text encoder override override_text_encoder_quant: QuantizationMethods = None + transformer_quant: QuantizationMethods = None override_transformer_cls_name: str | None = None init_weights_from_safetensors: str = "" # path to safetensors file for initial weight loading @@ -626,6 +627,13 @@ class FastVideoArgs: default=FastVideoArgs.override_text_encoder_quant, help="Quantization method for text encoder override", ) + parser.add_argument( + "--transformer-quant", + type=str, + choices=QUANTIZATION_METHODS, + default=FastVideoArgs.transformer_quant, + help="Quantization method for transformer loading", + ) parser.add_argument( "--override-transformer-cls-name", type=str, diff --git a/fastvideo/layers/quantization/__init__.py b/fastvideo/layers/quantization/__init__.py index 366fd463b..3067e4976 100644 --- a/fastvideo/layers/quantization/__init__.py +++ b/fastvideo/layers/quantization/__init__.py @@ -2,7 +2,7 @@ from typing import Literal, get_args from fastvideo.layers.quantization.base_config import QuantizationConfig -QuantizationMethods = Literal[None, "AbsMaxFP8"] +QuantizationMethods = Literal[None, "AbsMaxFP8", "fp4"] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) @@ -51,9 +51,11 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: # 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) diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 837faea44..a2b0418fe 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -873,6 +873,9 @@ class TransformerLoader(ComponentLoader): # Config from Diffusers supersedes fastvideo's model config dit_config = deepcopy(fastvideo_args.pipeline_config.dit_config) dit_config.update_model_arch(config) + if fastvideo_args.transformer_quant is not None: + quant_cls = get_quantization_config(fastvideo_args.transformer_quant) + dit_config.quant_config = quant_cls() model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) @@ -952,6 +955,7 @@ class TransformerLoader(ComponentLoader): training_mode=fastvideo_args.training_mode, enable_torch_compile=fastvideo_args.enable_torch_compile, torch_compile_kwargs=fastvideo_args.torch_compile_kwargs, + transformer_quant=fastvideo_args.transformer_quant, ) total_params = sum(p.numel() for p in model.parameters()) diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 53cb2fe10..cf190d030 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -20,10 +20,10 @@ from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule, from torch.nn.modules.module import _IncompatibleKeys from fastvideo.logger import init_logger +from fastvideo.layers.quantization import QuantizationMethods from fastvideo.models.loader.utils import (get_param_names_mapping, hf_to_custom_state_dict) from fastvideo.models.loader.weight_utils import safetensors_weights_iterator -from fastvideo.layers.fp4_config import convert_model_to_fp4 from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available logger = init_logger(__name__) @@ -57,7 +57,6 @@ def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]: torch.set_default_dtype(old_dtype) -USE_FP4 = False # TODO: move to FastVideoArgs # TODO(PY): add compile option def maybe_load_fsdp_model( model_cls: type[nn.Module], @@ -77,6 +76,7 @@ def maybe_load_fsdp_model( pin_cpu_memory: bool = True, enable_torch_compile: bool = False, torch_compile_kwargs: dict[str, Any] | None = None, + transformer_quant: QuantizationMethods = None, ) -> torch.nn.Module: """ Load the model with FSDP if is training, else load the model without FSDP. @@ -159,7 +159,10 @@ def maybe_load_fsdp_model( # Avoid unintended computation graph accumulation during inference if isinstance(p, torch.nn.Parameter): p.requires_grad = False - if USE_FP4: + if transformer_quant == "fp4": + from fastvideo.layers.quantization.fp4_config import convert_model_to_fp4 + + logger.info("Converting model to FP4 quantized runtime weights") convert_model_to_fp4(model) compile_in_loader = enable_torch_compile and training_mode diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index c0f79580b..f1f1de9d0 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -9,6 +9,7 @@ from fastvideo.api import ( GenerationResult, GeneratorConfig, InputConfig, + QuantizationConfig, SamplingConfig, ) from fastvideo.configs.sample import SamplingParam @@ -152,9 +153,10 @@ def test_prepare_output_path_empty_prompt_fallback(tmp_path): def test_from_config_normalizes_and_translates(monkeypatch): captured = _patch_from_fastvideo_args(monkeypatch) - _patch_fastvideo_args_from_kwargs(monkeypatch) + fastvideo_args_capture = _patch_fastvideo_args_from_kwargs(monkeypatch) config = GeneratorConfig(model_path="test-model") config.engine.num_gpus = 2 + config.engine.quantization = QuantizationConfig(transformer_quant="fp4") config.pipeline.workload_type = "t2v" generator = VideoGenerator.from_config(config) @@ -162,6 +164,7 @@ def test_from_config_normalizes_and_translates(monkeypatch): assert captured["fastvideo_args"].model_path == "test-model" assert captured["fastvideo_args"].num_gpus == 2 assert captured["fastvideo_args"].workload_type.value == "t2v" + assert fastvideo_args_capture["kwargs"]["transformer_quant"] == "fp4" assert generator.config == config @@ -195,6 +198,7 @@ def test_from_pretrained_convenience_kwargs_do_not_warn(monkeypatch): "test-model", num_gpus=4, use_fsdp_inference=False, + transformer_quant="fp4", text_encoder_cpu_offload=True, pin_cpu_memory=True, dit_cpu_offload=False, @@ -205,6 +209,7 @@ def test_from_pretrained_convenience_kwargs_do_not_warn(monkeypatch): assert captured["fastvideo_args"].model_path == "test-model" assert captured["fastvideo_args"].num_gpus == 4 assert fastvideo_args_capture["kwargs"]["use_fsdp_inference"] is False + assert fastvideo_args_capture["kwargs"]["transformer_quant"] == "fp4" assert fastvideo_args_capture["kwargs"]["text_encoder_cpu_offload"] is True assert fastvideo_args_capture["kwargs"]["pin_cpu_memory"] is True assert fastvideo_args_capture["kwargs"]["dit_cpu_offload"] is False @@ -212,6 +217,8 @@ def test_from_pretrained_convenience_kwargs_do_not_warn(monkeypatch): assert generator.config is not None assert generator.config.model_path == "test-model" assert generator.config.engine.num_gpus == 4 + assert generator.config.engine.quantization is not None + assert generator.config.engine.quantization.transformer_quant == "fp4" def test_from_pretrained_legacy_only_kwargs_warn(monkeypatch):