diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index ed0536e07..bbea4bbc2 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -37,6 +37,7 @@ surfaces: torch_compile_kwargs_text_encoder: generator.engine.compile.text_encoder_kwargs torch_compile_kwargs_vae: generator.engine.compile.vae_kwargs torch_compile_kwargs_audio_vae: generator.engine.compile.audio_vae_kwargs + transformer_quant: generator.engine.quantization.transformer_quant disable_autocast: generator.engine.disable_autocast enable_stage_verification: generator.engine.enable_stage_verification prompt_txt: request.inputs.prompt_path diff --git a/fastvideo/api/compat.py b/fastvideo/api/compat.py index 5e88de870..a92f25c92 100644 --- a/fastvideo/api/compat.py +++ b/fastvideo/api/compat.py @@ -268,7 +268,15 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An if quantization is not None and quantization.text_encoder_quant is not None: kwargs["override_text_encoder_quant"] = quantization.text_encoder_quant if quantization is not None and quantization.transformer_quant is not None: - kwargs["transformer_quant"] = quantization.transformer_quant + # Resolve the typed quant name to a concrete ``QuantizationConfig`` + # instance and pin it on ``dit_config.quant_config``. The legacy + # path expected callers to do this themselves via + # ``pipeline_config.dit_config.quant_config = FP4Config()``; the + # typed surface accepts a string and does the wiring here so + # downstream code can rely on a single source of truth. + from fastvideo.layers.quantization import get_quantization_config + _resolved_quant_cls = get_quantization_config(quantization.transformer_quant) + kwargs["transformer_quant"] = _resolved_quant_cls() components = normalized.pipeline.components if components.pipeline_config_path is not None: diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 2c7fbccb0..51742e1e9 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -217,6 +217,15 @@ class FastVideoArgs: override_text_encoder_safetensors: str | None = None # path to safetensors file for text encoder override override_text_encoder_quant: QuantizationMethods = None + # Typed transformer quantization carrier. The typed inference API + # accepts ``engine.quantization.transformer_quant: "FP4"`` and the + # compat layer resolves the name to a concrete ``QuantizationConfig`` + # instance (e.g. ``FP4Config()``); ``__post_init__`` then pins it on + # ``pipeline_config.dit_config.quant_config`` so the loader can detect + # FP4 layers via the standard ``get_quant_method`` path. ``None`` + # leaves whatever value the caller already set on ``dit_config`` + # untouched. + transformer_quant: Any | None = None override_transformer_cls_name: str | None = None init_weights_from_safetensors: str = "" # path to safetensors file for initial weight loading @@ -245,8 +254,29 @@ class FastVideoArgs: raise self._apply_ltx2_vae_overrides() self._resolve_refine_args() + self._apply_transformer_quant() self.check_fastvideo_args() + def _apply_transformer_quant(self) -> None: + """Pin the typed ``transformer_quant`` instance onto ``dit_config``. + + ``transformer_quant`` is populated by the typed compat layer when + a caller writes ``engine.quantization.transformer_quant: "FP4"`` + in their config. We pin it here rather than at request time so + the model loader sees the quant_config when constructing the + DiT (linear layers attach their quant_method during ``__init__``). + """ + if self.transformer_quant is None or self.pipeline_config is None: + return + dit_config = getattr(self.pipeline_config, "dit_config", None) + if dit_config is None: + return + # Don't overwrite if the caller already set it explicitly on + # dit_config (e.g. via ``pipeline_config.dit_config.quant_config = FP4Config()``); + # the explicit setter wins. + if getattr(dit_config, "quant_config", None) is None: + dit_config.quant_config = self.transformer_quant + def _resolve_refine_args(self) -> None: """Map generic refine_* args to LTX-2-specific refine fields.""" if self.refine_enabled is not None: diff --git a/fastvideo/layers/linear.py b/fastvideo/layers/linear.py index e33a738ec..efbcc8e76 100644 --- a/fastvideo/layers/linear.py +++ b/fastvideo/layers/linear.py @@ -191,7 +191,15 @@ class LinearBase(torch.nn.Module): if quant_config is None: self.quant_method: QuantizeMethodBase | None = (UnquantizedLinearMethod()) else: + # ``get_quant_method`` returns ``None`` for layers the config + # has decided not to quantize (e.g. ``FP4Config`` only tags + # a curated subset of LTX-2 attention/FFN layers). Fall back + # to ``UnquantizedLinearMethod`` so untagged layers behave + # like a plain ``nn.Linear`` instead of breaking subclass + # asserts. self.quant_method = quant_config.get_quant_method(self, prefix=prefix) + if self.quant_method is None: + self.quant_method = UnquantizedLinearMethod() def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]: raise NotImplementedError @@ -230,8 +238,14 @@ class ReplicatedLinear(LinearBase): prefix=prefix, ) - # All the linear layer supports quant method. - assert self.quant_method is not None + # ``QuantizationConfig.get_quant_method`` may return ``None`` for + # layers it doesn't intend to quantize (e.g. ``FP4Config`` only + # tags a specific subset of LTX-2 attention/FFN layers). Fall + # back to ``UnquantizedLinearMethod`` so non-matched layers + # behave like a plain ``nn.Linear``. + if self.quant_method is None: + self.quant_method = UnquantizedLinearMethod() + self.quant_method.create_weights( self, self.input_size, diff --git a/fastvideo/layers/quantization/__init__.py b/fastvideo/layers/quantization/__init__.py index 366fd463b..8eba57e9d 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/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index d7051cb29..8760a9eea 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -28,6 +28,8 @@ from fastvideo.distributed.communication_op import ( ) from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size from fastvideo.forward_context import ForwardContext, get_forward_context, set_forward_context +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.quantization.base_config import QuantizationConfig from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum @@ -35,6 +37,50 @@ from fastvideo.platforms import AttentionBackendEnum logger = init_logger(__name__) +def _supports_prequantized_input(linear: ReplicatedLinear) -> bool: + """Whether ``linear``'s quant method is willing to accept a + pre-quantized ``(x_fp4, x_scale, x_global_sf)`` input tuple. + + Used by the LTX-2 attention forward path so that a single input + tensor can be quantized once and reused across q/k/v projections + when they share the same source (self-attention). + """ + quant_method = getattr(linear, "quant_method", None) + if not callable(getattr(quant_method, "quantize_input", None)): + return False + wants_prequant = getattr(quant_method, "wants_prequantized_input", None) + if callable(wants_prequant): + try: + return bool(wants_prequant()) + except Exception: + return False + return True + + +def _linear_project_with_optional_prequant( + linear: ReplicatedLinear, + x: torch.Tensor, + pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None, +) -> torch.Tensor: + """Project ``x`` through ``linear``, optionally bypassing the + in-method quantize step when ``pre_quantized`` is supplied. + + When the linear's quant method does not support pre-quantized + inputs (e.g. ``UnquantizedLinearMethod``), falls back to the + standard ``linear(x)`` call and discards the bias-pass-through + tuple element. + """ + if pre_quantized is None or not _supports_prequantized_input(linear): + return linear(x)[0] + bias = linear.bias if not linear.skip_bias_add else None + return linear.quant_method.apply( # type: ignore[union-attr] + linear, + x, + bias=bias, + pre_quantized=pre_quantized, + ) + + def get_timestep_embedding( timesteps: torch.Tensor, embedding_dim: int, @@ -172,25 +218,57 @@ class PixArtAlphaTextProjection(torch.nn.Module): class GELUApprox(nn.Module): """Linear + tanh-approximate GELU used by LTX-2 FFN.""" - def __init__(self, in_features: int, out_features: int): + def __init__( + self, + in_features: int, + out_features: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): super().__init__() - self.proj = nn.Linear(in_features, out_features) + self.proj = ReplicatedLinear( + in_features, + out_features, + quant_config=quant_config, + prefix=f"{prefix}.fc_in", + ) self.act = nn.GELU(approximate="tanh") def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.act(self.proj(x)) + return self.act(self.proj(x)[0]) class FeedForward(nn.Module): """LTX-2 FFN: GELUApprox -> Identity -> Linear.""" - def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None: + def __init__( + self, + dim: int, + dim_out: int, + mult: int = 4, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: super().__init__() inner_dim = int(dim * mult) - project_in = GELUApprox(dim, inner_dim) - self.net = nn.Sequential(project_in, nn.Identity(), nn.Linear(inner_dim, dim_out)) + project_in = GELUApprox( + dim, + inner_dim, + quant_config=quant_config, + prefix=f"{prefix}.ffn", + ) + project_out = ReplicatedLinear( + inner_dim, + dim_out, + quant_config=quant_config, + prefix=f"{prefix}.ffn.fc_out", + ) + self.net = nn.ModuleList([project_in, nn.Identity(), project_out]) def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.net(x) + x = self.net[0](x) + x = self.net[1](x) + x = self.net[2](x)[0] + return x class VideoLatentShape(tuple): @@ -1246,6 +1324,8 @@ class LTXSelfAttention(nn.Module): norm_eps: float, rope_type: LTXRopeType, supported_attention_backends: tuple[AttentionBackendEnum, ...], + quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: super().__init__() inner_dim = dim_head * heads @@ -1257,10 +1337,37 @@ class LTXSelfAttention(nn.Module): self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) - self.to_q = nn.Linear(query_dim, inner_dim, bias=True) - self.to_k = nn.Linear(context_dim, inner_dim, bias=True) - self.to_v = nn.Linear(context_dim, inner_dim, bias=True) - self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()) + self.to_q = ReplicatedLinear( + query_dim, + inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_q", + ) + self.to_k = ReplicatedLinear( + context_dim, + inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_k", + ) + self.to_v = ReplicatedLinear( + context_dim, + inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_v", + ) + self.to_out = nn.ModuleList([ + ReplicatedLinear( + inner_dim, + query_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_out", + ), + nn.Identity(), + ]) self.attn = LTXLocalAttention( num_heads=heads, @@ -1271,9 +1378,15 @@ class LTXSelfAttention(nn.Module): causal=False, supported_attention_backends=supported_attention_backends, ) - self.to_gate_compress: nn.Linear | None = None + self.to_gate_compress: ReplicatedLinear | None = None if self.attn.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: - self.to_gate_compress = nn.Linear(context_dim, inner_dim, bias=True) + self.to_gate_compress = ReplicatedLinear( + context_dim, + inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_gate_compress", + ) self.attn_masked = LTXLocalAttention( num_heads=heads, head_size=dim_head, @@ -1292,11 +1405,24 @@ class LTXSelfAttention(nn.Module): pe: tuple[torch.Tensor, torch.Tensor] | None = None, k_pe: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: - q = self.to_q(x) context = x if context is None else context - k = self.to_k(context) - v = self.to_v(context) - gate_compress = (self.to_gate_compress(context) + + q_pre_quantized = ( + self.to_q.quant_method.quantize_input(x) # type: ignore[union-attr] + if _supports_prequantized_input(self.to_q) else None) + kv_pre_quantized = None + if (_supports_prequantized_input(self.to_k) + and _supports_prequantized_input(self.to_v)): + if context is x and q_pre_quantized is not None: + kv_pre_quantized = q_pre_quantized + else: + kv_pre_quantized = self.to_k.quant_method.quantize_input( # type: ignore[union-attr] + context) + + q = _linear_project_with_optional_prequant(self.to_q, x, q_pre_quantized) + k = _linear_project_with_optional_prequant(self.to_k, context, kv_pre_quantized) + v = _linear_project_with_optional_prequant(self.to_v, context, kv_pre_quantized) + gate_compress = (self.to_gate_compress(context)[0] if self.to_gate_compress is not None else None) q = self.q_norm(q) @@ -1346,7 +1472,8 @@ class LTXSelfAttention(nn.Module): ltx_freqs_cis=pe, ltx_k_freqs_cis=k_pe) out = out.reshape(b, q_len, -1) - return self.to_out(out) + out = self.to_out[0](out)[0] + return self.to_out[1](out) class LTXDistributedSelfAttention(nn.Module): @@ -1361,6 +1488,7 @@ class LTXDistributedSelfAttention(nn.Module): norm_eps: float, rope_type: LTXRopeType, supported_attention_backends: tuple[AttentionBackendEnum, ...], + quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() @@ -1373,10 +1501,37 @@ class LTXDistributedSelfAttention(nn.Module): self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) - self.to_q = nn.Linear(query_dim, inner_dim, bias=True) - self.to_k = nn.Linear(context_dim, inner_dim, bias=True) - self.to_v = nn.Linear(context_dim, inner_dim, bias=True) - self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()) + self.to_q = ReplicatedLinear( + query_dim, + inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_q", + ) + self.to_k = ReplicatedLinear( + context_dim, + inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_k", + ) + self.to_v = ReplicatedLinear( + context_dim, + inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_v", + ) + self.to_out = nn.ModuleList([ + ReplicatedLinear( + inner_dim, + query_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_out", + ), + nn.Identity(), + ]) self.attn = LTXDistributedAttention( num_heads=heads, @@ -1386,9 +1541,15 @@ class LTXDistributedSelfAttention(nn.Module): supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn", ) - self.to_gate_compress: nn.Linear | None = None + self.to_gate_compress: ReplicatedLinear | None = None if self.attn.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: - self.to_gate_compress = nn.Linear(context_dim, inner_dim, bias=True) + self.to_gate_compress = ReplicatedLinear( + context_dim, + inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_gate_compress", + ) def forward( self, @@ -1408,11 +1569,24 @@ class LTXDistributedSelfAttention(nn.Module): k_pe: Rotary position embeddings for K (cos, sin), or None to use pe original_seq_len: Original (unpadded) full sequence length """ - q = self.to_q(x) context = x if context is None else context - k = self.to_k(context) - v = self.to_v(context) - gate_compress = (self.to_gate_compress(context) + + q_pre_quantized = ( + self.to_q.quant_method.quantize_input(x) # type: ignore[union-attr] + if _supports_prequantized_input(self.to_q) else None) + kv_pre_quantized = None + if (_supports_prequantized_input(self.to_k) + and _supports_prequantized_input(self.to_v)): + if context is x and q_pre_quantized is not None: + kv_pre_quantized = q_pre_quantized + else: + kv_pre_quantized = self.to_k.quant_method.quantize_input( # type: ignore[union-attr] + context) + + q = _linear_project_with_optional_prequant(self.to_q, x, q_pre_quantized) + k = _linear_project_with_optional_prequant(self.to_k, context, kv_pre_quantized) + v = _linear_project_with_optional_prequant(self.to_v, context, kv_pre_quantized) + gate_compress = (self.to_gate_compress(context)[0] if self.to_gate_compress is not None else None) q = self.q_norm(q) @@ -1443,7 +1617,8 @@ class LTXDistributedSelfAttention(nn.Module): ) out = out.reshape(b, q_len, -1) - return self.to_out(out) + out = self.to_out[0](out)[0] + return self.to_out[1](out) class BasicAVTransformerBlock(torch.nn.Module): @@ -1457,6 +1632,7 @@ class BasicAVTransformerBlock(torch.nn.Module): rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, norm_eps: float = 1e-6, use_distributed_attention: bool = False, + quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() @@ -1488,15 +1664,8 @@ class BasicAVTransformerBlock(torch.nn.Module): norm_eps=norm_eps, rope_type=rope_type, supported_attention_backends=video_self_attn_backends, - prefix=f"{prefix}.blocks.{idx}.attn1" if use_distributed_attention else "", - ) if use_distributed_attention else LTXSelfAttention( - query_dim=video.dim, - context_dim=None, - heads=video.heads, - dim_head=video.d_head, - norm_eps=norm_eps, - rope_type=rope_type, - supported_attention_backends=video_self_attn_backends, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{idx}.attn1", ) # Text cross-attention - always local (text is replicated) self.attn2 = CrossAttnCls( @@ -1507,8 +1676,15 @@ class BasicAVTransformerBlock(torch.nn.Module): norm_eps=norm_eps, rope_type=rope_type, supported_attention_backends=dense_attn_backends, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{idx}.attn2", + ) + self.ff = FeedForward( + video.dim, + dim_out=video.dim, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{idx}", ) - self.ff = FeedForward(video.dim, dim_out=video.dim) self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim)) if audio is not None: @@ -1521,15 +1697,8 @@ class BasicAVTransformerBlock(torch.nn.Module): norm_eps=norm_eps, rope_type=rope_type, supported_attention_backends=dense_attn_backends, - prefix=f"{prefix}.blocks.{idx}.audio_attn1" if use_distributed_attention else "", - ) if use_distributed_attention else LTXSelfAttention( - query_dim=audio.dim, - context_dim=None, - heads=audio.heads, - dim_head=audio.d_head, - norm_eps=norm_eps, - rope_type=rope_type, - supported_attention_backends=dense_attn_backends, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{idx}.audio_attn1", ) # Text cross-attention - always local (text is replicated) self.audio_attn2 = CrossAttnCls( @@ -1540,8 +1709,15 @@ class BasicAVTransformerBlock(torch.nn.Module): norm_eps=norm_eps, rope_type=rope_type, supported_attention_backends=dense_attn_backends, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{idx}.audio_attn2", + ) + self.audio_ff = FeedForward( + audio.dim, + dim_out=audio.dim, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{idx}.audio", ) - self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim) self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim)) if audio is not None and video is not None: @@ -1555,6 +1731,8 @@ class BasicAVTransformerBlock(torch.nn.Module): norm_eps=norm_eps, rope_type=rope_type, supported_attention_backends=dense_attn_backends, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{idx}.audio_to_video_attn", ) # Video-to-audio cross-attention # Uses local attention - context is gathered from all SP ranks in forward() @@ -1566,6 +1744,8 @@ class BasicAVTransformerBlock(torch.nn.Module): norm_eps=norm_eps, rope_type=rope_type, supported_attention_backends=dense_attn_backends, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{idx}.video_to_audio_attn", ) self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim)) self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim)) @@ -1889,6 +2069,7 @@ class LTXModel(torch.nn.Module): rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, double_precision_rope: bool = False, use_distributed_attention: bool = False, + quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() @@ -1942,6 +2123,7 @@ class LTXModel(torch.nn.Module): audio_cross_attention_dim=audio_cross_attention_dim, norm_eps=norm_eps, use_distributed_attention=use_distributed_attention, + quant_config=quant_config, prefix=prefix, ) @@ -2073,6 +2255,7 @@ class LTXModel(torch.nn.Module): audio_cross_attention_dim: int, norm_eps: float, use_distributed_attention: bool = False, + quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: video_config = ( @@ -2105,6 +2288,7 @@ class LTXModel(torch.nn.Module): rope_type=self.rope_type, norm_eps=norm_eps, use_distributed_attention=use_distributed_attention, + quant_config=quant_config, prefix=prefix, ) for idx in range(num_layers) @@ -2293,6 +2477,7 @@ class LTX2Transformer3DModel(BaseDiT): audio_positional_embedding_max_pos=arch.audio_positional_embedding_max_pos, av_ca_timestep_scale_multiplier=arch.av_ca_timestep_scale_multiplier, use_distributed_attention=use_distributed_attention, + quant_config=config.quant_config, prefix=config.prefix, ) diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index f0e449414..23e9c59ec 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -28,6 +28,35 @@ from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available logger = init_logger(__name__) +def _maybe_convert_model_to_fp4(model: nn.Module) -> None: + """Quantize FP4-tagged linear layers in-place after weights are loaded. + + Walks the module tree once, looking for layers whose ``quant_method`` + is an :class:`FP4QuantizeMethod` (attached at construction time by + :meth:`FP4Config.get_quant_method`). When at least one such layer + exists, calls :func:`convert_model_to_fp4` to register the + ``_fp4_weight*`` / ``_fp4_alpha`` / ``_weight_global_sf`` buffers + on each targeted layer. + + The walk is short-circuited on the first non-FP4 layer in the + ``UnquantizedLinearMethod`` case so non-FP4 callers pay only a + single ``isinstance`` check. flashinfer is imported lazily inside + :func:`convert_model_to_fp4` so this helper is a no-op on hosts + without the FP4 backend. + """ + # Defer the import: fp4_config imports heavy diffusers/torch.distributed + # symbols at module-load time, and unconditional import would penalize + # every loader call regardless of whether FP4 is wired. + from fastvideo.layers.quantization.fp4_config import (FP4QuantizeMethod, + convert_model_to_fp4) + + for mod in model.modules(): + if isinstance(getattr(mod, "quant_method", None), FP4QuantizeMethod): + logger.info("Converting loaded model weights for FP4 linear layers") + convert_model_to_fp4(model) + return + + # TODO(PY): move this to utils elsewhere @contextlib.contextmanager def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]: @@ -158,6 +187,15 @@ def maybe_load_fsdp_model( if isinstance(p, torch.nn.Parameter): p.requires_grad = False + # FP4 weight prequantization. We detect by the registered + # ``quant_method`` on linear layers rather than by a separate flag — + # construction-time ``FP4Config.get_quant_method`` already attached + # ``FP4QuantizeMethod`` to every targeted layer, so the loader's + # responsibility is just to materialize the per-layer fp4 weight / + # scale buffers from the freshly-loaded bf16 weights. No-op when + # ``flashinfer`` is not installed (lazy import inside the helper). + _maybe_convert_model_to_fp4(model) + compile_in_loader = enable_torch_compile and training_mode if compile_in_loader: compile_kwargs = torch_compile_kwargs or {}