clean

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

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",
 })
 
 

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,

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)

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())

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

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):