diff --git a/.gitmodules b/.gitmodules index bdb04686d..8d090c699 100644 --- a/.gitmodules +++ b/.gitmodules @@ -5,3 +5,6 @@ [submodule "csrc/attn/sliding_tile_attn/tk"] path = csrc/attn/sliding_tile_attn/tk url = https://github.com/HazyResearch/ThunderKittens.git +[submodule "DeepGEMM"] + path = DeepGEMM + url = https://github.com/deepseek-ai/DeepGEMM diff --git a/DeepGEMM b/DeepGEMM new file mode 160000 index 000000000..c9f8b34dc --- /dev/null +++ b/DeepGEMM @@ -0,0 +1 @@ +Subproject commit c9f8b34dcdacc20aa746b786f983492c51072870 diff --git a/examples/inference/optimizations/quantization_example.py b/examples/inference/optimizations/quantization_example.py new file mode 100644 index 000000000..680cd8fc5 --- /dev/null +++ b/examples/inference/optimizations/quantization_example.py @@ -0,0 +1,105 @@ +import torch +from fastvideo import VideoGenerator +from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.layers.quantization.fp4_config import FP4Config +# from fastvideo.layers.quantization.fp8_config import FP8Config +from fastvideo.layers.linear import ReplicatedLinear + +OUTPUT_PATH = "fp4_video_samples" +# OUTPUT_PATH = "video_samples" + + +def main(): + print("=== FP4 Quantization Video Generation Example ===") + + if not torch.cuda.is_available(): + print("Warning: CUDA not available. FP4 quantization requires GPU.") + return + + gpu_capability = torch.cuda.get_device_capability() + if gpu_capability[0] < 9: # H100 and newer + print(f"Warning: GPU capability {gpu_capability} may not support FP4. Recommended: 9.0+") + + print(f"GPU: {torch.cuda.get_device_name()}") + print(f"GPU Capability: {gpu_capability}") + + model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + # model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers" + pipeline_config = PipelineConfig.from_pretrained(model_id) + pipeline_config.dit_precision = "bf16" + pipeline_config.dit_config.quant_config = FP4Config() + # pipeline_config.dit_config.quant_config = FP8Config() + + ReplicatedLinear.print_shape_summary() + + print("\nLoading model with FP4 quantization...") + generator = VideoGenerator.from_pretrained( + model_id, + pipeline_config=pipeline_config, + num_gpus=1, + use_fsdp_inference=True, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + pin_cpu_memory=False, + ) + + print("FP4 configuration applied. Generating videos...") + + print("\n=== Generating Video with FP4 Quantization ===") + + prompt1 = ( + "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones." + ) + + print(f"Prompt: {prompt1}") + print("Generating video...") + + try: + video1 = generator.generate_video( + prompt1, + output_path=OUTPUT_PATH, + save_video=True, + ) + print("✓ First video generated successfully with FP4 quantization!") + + # # Generate a second video to show the model can be reused + # prompt2 = ( + # "A majestic lion strides across the golden savanna, its powerful frame " + # "glistening under the warm afternoon sun. The tall grass ripples gently in " + # "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + # "embodying the raw energy of the wild. Low angle, steady tracking shot, " + # "cinematic." + # ) + + # print(f"\nGenerating second video...") + # print(f"Prompt: {prompt2}") + + # video2 = generator.generate_video( + # prompt2, + # output_path=OUTPUT_PATH, + # save_video=True, + # ) + # print("✓ Second video generated successfully with FP4 quantization!") + + except Exception as e: + print(f"Error during video generation: {e}") + print("Make sure you have the required deep_gemm library installed for FP4 support") + return + + print(f"\n=== FP4 Quantization Summary ===") + print(f"Videos saved to: {OUTPUT_PATH}") + print("FP4 quantization benefits:") + print("- ✓ Reduced memory usage compared to FP16/BF16") + print("- ✓ Faster inference on supported hardware (H100+)") + print("- ✓ Maintains model quality with proper scaling") + print("\nFor optimal FP4 performance:") + print("1. Use H100 or newer GPUs") + print("2. Ensure deep_gemm library is installed") + print("3. Keep models on GPU (avoid CPU offloading for quantized layers)") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/fastvideo/DeepGEMM b/fastvideo/DeepGEMM new file mode 160000 index 000000000..f85ec649d --- /dev/null +++ b/fastvideo/DeepGEMM @@ -0,0 +1 @@ +Subproject commit f85ec649d76846552cfd637b6c99fb9c985fd9eb diff --git a/fastvideo/layers/linear.py b/fastvideo/layers/linear.py index d711a1ad4..a90854398 100644 --- a/fastvideo/layers/linear.py +++ b/fastvideo/layers/linear.py @@ -180,6 +180,10 @@ class ReplicatedLinear(LinearBase): prefix: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj) """ + + # Class-level set to track unique shapes and dictionary to map shapes to layer types + _unique_shapes = set() + _shape_to_layer_types = {} def __init__(self, input_size: int, @@ -197,7 +201,9 @@ class ReplicatedLinear(LinearBase): prefix=prefix) # All the linear layer supports quant method. - assert self.quant_method is not None + if self.quant_method is None: + self.quant_method = UnquantizedLinearMethod() + self.quant_method.create_weights(self, self.input_size, [self.output_size], self.input_size, @@ -232,8 +238,25 @@ class ReplicatedLinear(LinearBase): def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]: bias = self.bias if not self.skip_bias_add else None - assert self.quant_method is not None + if self.quant_method is None: + self.quant_method = UnquantizedLinearMethod() + output = self.quant_method.apply(self, x, bias) + + # Create a shape key for uniqueness checking + shape_key = (x.shape, output.shape) + + # Track unique shapes and map them to layer types + if shape_key not in self._unique_shapes: + self._unique_shapes.add(shape_key) + self._shape_to_layer_types[shape_key] = [] + print(f"Layer: {self.prefix} | input shape: {x.shape} --> output shape: {output.shape}, Quant Method: {self.quant_method.__class__.__name__}") + + # Add this layer type to the list for this shape if not already present + layer_type = self.__class__.__name__ + if layer_type not in self._shape_to_layer_types[shape_key]: + self._shape_to_layer_types[shape_key].append(layer_type) + output_bias = self.bias if self.skip_bias_add else None return output, output_bias @@ -242,6 +265,28 @@ class ReplicatedLinear(LinearBase): s += f", output_features={self.output_size}" s += f", bias={self.bias is not None}" return s + + @classmethod + def get_shape_mapping(cls) -> dict: + """Get the mapping from (input_shape, output_shape) to layer types.""" + return cls._shape_to_layer_types.copy() + + @classmethod + def print_shape_summary(cls): + """Print a summary of all unique shapes and their layer types.""" + if not cls._shape_to_layer_types: + print("No shapes have been processed yet.") + return + + print(f"\n=== Matrix Multiplication Shape Summary ===") + print(f"Total unique shapes: {len(cls._shape_to_layer_types)}") + print() + + for i, (shape_key, layer_types) in enumerate(cls._shape_to_layer_types.items(), 1): + input_shape, output_shape = shape_key + print(f"{i}. Input: {input_shape} → Output: {output_shape}") + print(f" Layer types: {', '.join(layer_types)}") + print() class ColumnParallelLinear(LinearBase): @@ -298,7 +343,9 @@ class ColumnParallelLinear(LinearBase): if output_sizes is None: output_sizes = [output_size] - assert self.quant_method is not None + if self.quant_method is None: + self.quant_method = UnquantizedLinearMethod() + self.quant_method.create_weights( layer=self, input_size_per_partition=self.input_size_per_partition, @@ -359,7 +406,8 @@ class ColumnParallelLinear(LinearBase): bias = self.bias if not self.skip_bias_add else None # Matrix multiply. - assert self.quant_method is not None + if self.quant_method is None: + self.quant_method = UnquantizedLinearMethod() output_parallel = self.quant_method.apply(self, input_, bias) if self.gather_output: # All-gather across the partitions. @@ -866,7 +914,9 @@ class RowParallelLinear(LinearBase): self.input_is_parallel = input_is_parallel self.reduce_results = reduce_results - assert self.quant_method is not None + if self.quant_method is None: + self.quant_method = UnquantizedLinearMethod() + self.quant_method.create_weights( layer=self, input_size_per_partition=self.input_size_per_partition, @@ -935,7 +985,8 @@ class RowParallelLinear(LinearBase): input_parallel = splitted_input[tp_rank].contiguous() # Matrix multiply. - assert self.quant_method is not None + if self.quant_method is None: + self.quant_method = UnquantizedLinearMethod() # Only fuse bias add into GEMM for rank 0 (this ensures that # bias will not get added more than once in TP>1 case) bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias diff --git a/fastvideo/layers/mlp.py b/fastvideo/layers/mlp.py index 2d5ab2a34..89b89c29e 100644 --- a/fastvideo/layers/mlp.py +++ b/fastvideo/layers/mlp.py @@ -5,6 +5,7 @@ import torch.nn as nn from fastvideo.layers.activation import get_act_fn from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.quantization import QuantizationConfig class MLP(nn.Module): @@ -21,13 +22,21 @@ class MLP(nn.Module): act_type: str = "gelu_pytorch_tanh", dtype: torch.dtype | None = None, prefix: str = "", + quant_config: QuantizationConfig | None = None, ): super().__init__() self.fc_in = ReplicatedLinear( input_dim, mlp_hidden_dim, # For activation func like SiLU that need 2x width bias=bias, - params_dtype=dtype) + params_dtype=dtype, + quant_config=quant_config, + prefix=f"{prefix}.fc_in", + ) + if quant_config is not None: + quant_method = self.fc_in.quant_config.get_quant_method(self.fc_in, f"{prefix}.fc_in") + if quant_method is not None: + quant_method.process_weights_after_loading(self.fc_in) self.act = get_act_fn(act_type) if output_dim is None: @@ -35,10 +44,16 @@ class MLP(nn.Module): self.fc_out = ReplicatedLinear(mlp_hidden_dim, output_dim, bias=bias, - params_dtype=dtype) + params_dtype=dtype, + quant_config=quant_config, + prefix=f"{prefix}.fc_out") + if quant_config is not None: + quant_method = self.fc_out.quant_config.get_quant_method(self.fc_out, f"{prefix}.fc_out") + if quant_method is not None: + quant_method.process_weights_after_loading(self.fc_out) def forward(self, x: torch.Tensor) -> torch.Tensor: x, _ = self.fc_in(x) x = self.act(x) x, _ = self.fc_out(x) - return x + return x \ No newline at end of file diff --git a/fastvideo/layers/quantization/fp4_config.py b/fastvideo/layers/quantization/fp4_config.py new file mode 100644 index 000000000..72f938043 --- /dev/null +++ b/fastvideo/layers/quantization/fp4_config.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase +import torch +from torch.nn.parameter import Parameter + +from typing import Any, Tuple +from flashinfer import nvfp4_quantize, mm_fp4, SfLayout +from fastvideo.models.utils import set_weight_attrs +import logging +logger = logging.getLogger(__name__) + +class FP4QuantizeMethod(QuantizeMethodBase): + def __init__(self): + super().__init__() + self.weight_fp8 = None + self.weight_scale = None + + def create_weights(self, layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], input_size: int, + output_size: int, params_dtype: torch.dtype, + **extra_weight_attrs): + """Create weights for a linear layer. Note the corrected signature to match LinearMethodBase.""" + weight = Parameter(torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + @torch.compile + def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + """Apply FP4 quantized computation.""" + out_dim = layer.weight.shape[0] + original_shape = x.shape + assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}" + + x = x.view(-1, x.shape[-1]) + + x_global_sf = (448 * 6) / x.float().abs().nan_to_num().max(); + x_fp4, x_scale = nvfp4_quantize(x, x_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False) + weight_fp4 = layer._fp4_weight + weight_scale = layer._fp4_weight_scale + weight_global_sf = layer._weight_global_sf.item() + + out = mm_fp4(x_fp4, weight_fp4.T, x_scale, weight_scale.T, 1.0/(x_global_sf * weight_global_sf), torch.bfloat16, None, backend='cutlass') + + if bias is not None: + logger.info(f"Bias shape: {bias.shape}") + logger.info(f"Out shape: {out.shape}") + if bias.device != out.device or bias.dtype != out.dtype: + bias = bias.to(device=out.device, dtype=out.dtype) + out = out + bias + + if len(original_shape) == 3: + out = out.view(original_shape[0], original_shape[1], out_dim) + + return out + + +class FP4Config(QuantizationConfig): + def __init__(self): + super().__init__() + + def get_name(self): + return "fp4" + + def get_supported_act_dtypes(self): + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls): + return 100 + + @staticmethod + def get_config_filenames(): + return [] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "FP4Config": + return cls() + + def get_quant_method(self, layer: torch.nn.Module, prefix: str): + from fastvideo.layers.linear import LinearBase + fp8_layers = ["ffn.fc_in", "ffn.fc_out", "to_q", "to_k", "to_v", "to_out"] + if isinstance(layer, LinearBase) and any(layer_name in prefix for layer_name in fp8_layers): + return FP4QuantizeMethod() + return None + +@torch.compile +def convert_model_to_fp4(model: torch.nn.Module): + from torch.distributed.tensor import DTensor # type: ignore + for mod in model.modules(): + qm = getattr(mod, "quant_method", None) + if isinstance(qm, FP4QuantizeMethod): + weight = getattr(mod, "weight", None) + if weight is None: + continue + if isinstance(weight, DTensor): # type: ignore + weight_local = weight.to_local() + else: + weight_local = weight + weight_global_sf = (448 * 6) / weight_local.float().abs().nan_to_num().max(); + fp4_w, fp4_s = nvfp4_quantize(weight_local, weight_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False) + mod.register_buffer("_fp4_weight", fp4_w, persistent=False) + mod.register_buffer("_fp4_weight_scale", fp4_s, persistent=False) + mod.register_buffer("_weight_global_sf", torch.tensor(weight_global_sf, dtype=torch.bfloat16), persistent=False) \ No newline at end of file diff --git a/fastvideo/layers/quantization/fp8_config.py b/fastvideo/layers/quantization/fp8_config.py new file mode 100644 index 000000000..12f0f73ff --- /dev/null +++ b/fastvideo/layers/quantization/fp8_config.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase +import torch +from torch.nn.parameter import Parameter + +from typing import Any, Tuple +import deep_gemm +from deep_gemm import ceil_div, get_mn_major_tma_aligned_tensor +from deep_gemm.utils import per_token_cast_to_fp8, per_block_cast_to_fp8 +from fastvideo.models.utils import set_weight_attrs + +class FP8QuantizeMethod(QuantizeMethodBase): + def __init__(self): + super().__init__() + self.weight_fp8 = None + self.weight_scale = None + + def create_weights(self, layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], input_size: int, + output_size: int, params_dtype: torch.dtype, + **extra_weight_attrs): + """Create weights for a linear layer. Note the corrected signature to match LinearMethodBase.""" + weight = Parameter(torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + @torch.compile + def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + """Apply FP8 quantized computation.""" + out_dim = layer.weight.shape[0] + # Need contiguous tensors for collectives. + assert x.dtype == torch.bfloat16 or x.dtype == torch.float16 or x.dtype == torch.float32, f"only allow bf16/fp16/fp32 inputs to fp8 linear, got {x.dtype}" + x_fp8, x_scale = per_token_cast_to_fp8(x.view(-1, x.shape[-1]), use_ue8m0=False) + x_scale = get_mn_major_tma_aligned_tensor(x_scale) + weight_fp8 = layer._fp8_weight + weight_scale = layer._fp8_weight_scale + + original_shape = x.shape + out = torch.zeros((x_fp8.shape[0], out_dim), device=x.device, dtype=x.dtype) + deep_gemm.fp8_gemm_nt( + (x_fp8, x_scale), + (weight_fp8, weight_scale), + out, + disable_ue8m0_cast=False + ) + + if bias is not None: + if bias.device != out.device or bias.dtype != out.dtype: + bias = bias.to(device=out.device, dtype=out.dtype) + out = out + bias + + if len(original_shape) == 3: + out = out.view(original_shape[0], original_shape[1], out_dim) + + return out + + +class FP8Config(QuantizationConfig): + def __init__(self): + super().__init__() + + def get_name(self): + return "fp8" + + def get_supported_act_dtypes(self): + return [torch.bfloat16, torch.float16, torch.float32] + + @classmethod + def get_min_capability(cls): + return 90 + + @staticmethod + def get_config_filenames(): + return [] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "FP8Config": + return cls() + + def get_quant_method(self, layer: torch.nn.Module, prefix: str): + from fastvideo.layers.linear import LinearBase + fp8_layers = ["ffn.fc_in", "ffn.fc_out", "to_q", "to_k", "to_v", "to_out"] + if isinstance(layer, LinearBase) and any(layer_name in prefix for layer_name in fp8_layers): + return FP8QuantizeMethod() + return None + +@torch.compile +def convert_model_to_fp8(model: torch.nn.Module): + from torch.distributed.tensor import DTensor # type: ignore + for mod in model.modules(): + qm = getattr(mod, "quant_method", None) + if isinstance(qm, FP8QuantizeMethod): + weight = getattr(mod, "weight", None) + if weight is None: + continue + if isinstance(weight, DTensor): # type: ignore + weight_local = weight.to_local() + else: + weight_local = weight + fp8_w, fp8_s = per_block_cast_to_fp8(weight_local, use_ue8m0=False) + mod.register_buffer("_fp8_weight", fp8_w, persistent=False) + mod.register_buffer("_fp8_weight_scale", fp8_s, persistent=False) \ No newline at end of file diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 3d0f80975..6004f8fe1 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -30,6 +30,7 @@ from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, from fastvideo.logger import init_logger from fastvideo.models.dits.base import CachableDiT from fastvideo.platforms import AttentionBackendEnum, current_platform +from fastvideo.layers.quantization import QuantizationConfig from fastvideo.distributed.parallel_state import get_sp_world_size from fastvideo.distributed.utils import create_attention_mask_for_padding @@ -108,7 +109,9 @@ class WanSelfAttention(nn.Module): window_size=(-1, -1), qk_norm=True, eps=1e-6, - parallel_attention=False) -> None: + parallel_attention=False, + quant_config: QuantizationConfig | None = None, + prefix: str = "") -> None: assert dim % num_heads == 0 super().__init__() self.dim = dim @@ -120,10 +123,10 @@ class WanSelfAttention(nn.Module): self.parallel_attention = parallel_attention # layers - self.to_q = ReplicatedLinear(dim, dim) - self.to_k = ReplicatedLinear(dim, dim) - self.to_v = ReplicatedLinear(dim, dim) - self.to_out = ReplicatedLinear(dim, dim) + self.to_q = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_q") + self.to_k = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_k") + self.to_v = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_v") + self.to_out = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_out") self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() @@ -196,13 +199,15 @@ class WanI2VCrossAttention(WanSelfAttention): qk_norm=True, eps=1e-6, supported_attention_backends: tuple[AttentionBackendEnum, ...] - | None = None + | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: super().__init__(dim, num_heads, window_size, qk_norm, eps, - supported_attention_backends) + supported_attention_backends, quant_config=quant_config, prefix=prefix) - self.add_k_proj = ReplicatedLinear(dim, dim) - self.add_v_proj = ReplicatedLinear(dim, dim) + self.add_k_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_k_proj") + self.add_v_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_v_proj") self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() @@ -248,16 +253,17 @@ class WanTransformerBlock(nn.Module): added_kv_proj_dim: int | None = None, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + quant_config: QuantizationConfig | None = None, prefix: str = ""): super().__init__() # 1. Self-attention self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) - self.to_q = ReplicatedLinear(dim, dim, bias=True) - self.to_k = ReplicatedLinear(dim, dim, bias=True) - self.to_v = ReplicatedLinear(dim, dim, bias=True) + self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q") + self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k") + self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v") - self.to_out = ReplicatedLinear(dim, dim, bias=True) + self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out") self.attn1 = DistributedAttention( num_heads=num_heads, head_size=dim // num_heads, @@ -292,13 +298,17 @@ class WanTransformerBlock(nn.Module): self.attn2 = WanI2VCrossAttention(dim, num_heads, qk_norm=qk_norm, - eps=eps) + eps=eps, + quant_config=quant_config, + prefix=f"{prefix}.attn2") else: # T2V self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, - eps=eps) + eps=eps, + quant_config=quant_config, + prefix=f"{prefix}.attn2") self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, norm_type="layer", @@ -308,7 +318,7 @@ class WanTransformerBlock(nn.Module): compute_dtype=torch.float32) # 3. Feed-forward - self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn") self.mlp_residual = ScaleResidual() self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) @@ -402,17 +412,17 @@ class WanTransformerBlock_VSA(nn.Module): added_kv_proj_dim: int | None = None, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + quant_config: QuantizationConfig | None = None, prefix: str = ""): super().__init__() # 1. Self-attention self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) - self.to_q = ReplicatedLinear(dim, dim, bias=True) - self.to_k = ReplicatedLinear(dim, dim, bias=True) - self.to_v = ReplicatedLinear(dim, dim, bias=True) - self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True) - - self.to_out = ReplicatedLinear(dim, dim, bias=True) + self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q") + self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k") + self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v") + self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_gate_compress") + self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out") self.attn1 = DistributedAttention_VSA( num_heads=num_heads, head_size=dim // num_heads, @@ -447,13 +457,17 @@ class WanTransformerBlock_VSA(nn.Module): self.attn2 = WanI2VCrossAttention(dim, num_heads, qk_norm=qk_norm, - eps=eps) + eps=eps, + quant_config=quant_config, + prefix=f"{prefix}.attn2") else: # T2V self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, - eps=eps) + eps=eps, + quant_config=quant_config, + prefix=f"{prefix}.attn2") self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, norm_type="layer", @@ -463,7 +477,7 @@ class WanTransformerBlock_VSA(nn.Module): compute_dtype=torch.float32) # 3. Feed-forward - self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn") self.mlp_residual = ScaleResidual() self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) @@ -549,6 +563,7 @@ class WanTransformer3DModel(CachableDiT): def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None: super().__init__(config=config, hf_config=hf_config) + self.quant_config = config.quant_config inner_dim = config.num_attention_heads * config.attention_head_dim self.hidden_size = config.hidden_size @@ -587,6 +602,7 @@ class WanTransformer3DModel(CachableDiT): config.eps, config.added_kv_proj_dim, self._supported_attention_backends, + quant_config=config.quant_config, prefix=f"{config.prefix}.blocks.{i}") for i in range(config.num_layers) ]) diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 2fca474a2..f6bc3e5a6 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -24,6 +24,8 @@ 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.utils import set_mixed_precision_policy +from fastvideo.layers.quantization.fp8_config import convert_model_to_fp8 +from fastvideo.layers.quantization.fp4_config import convert_model_to_fp4 logger = init_logger(__name__) @@ -56,7 +58,9 @@ def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]: torch.set_default_dtype(old_dtype) -# Supports optional torch.compile for FSDP-wrapped models during training +USE_FP8 = False +USE_FP4 = False +# TODO(PY): add compile option def maybe_load_fsdp_model( model_cls: type[nn.Module], init_params: dict[str, Any], @@ -152,6 +156,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: + convert_model_to_fp4(model) + elif USE_FP8: + convert_model_to_fp8(model) compile_in_loader = enable_torch_compile and training_mode if compile_in_loader: