[feat]: add FP8 vendor support for Wan weights

v2/_vendor/configs/pipeline_registry.py

@@ -24,8 +24,8 @@ from v2._vendor.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfi
 from v2._vendor.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from v2._vendor.configs.pipelines.gen3c import Gen3CConfig
 from v2._vendor.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig, Hunyuan15I2V480PStepDistilledConfig,
-                                                   Hunyuan15T2V720PConfig, Hunyuan15I2V720PConfig,
-                                                   Hunyuan15SR1080PConfig)
+                                                    Hunyuan15T2V720PConfig, Hunyuan15I2V720PConfig,
+                                                    Hunyuan15SR1080PConfig)
 from v2._vendor.configs.pipelines.hyworld import HYWorldConfig
 from v2._vendor.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
 from v2._vendor.configs.pipelines.longcat import LongCatT2V480PConfig
@@ -736,8 +736,9 @@ def _register_configs() -> None:
         hf_model_paths=[
             "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
             "FastVideo/FastWan2.1-T2V-14B-480P-Diffusers",
+            "FastVideo/FastWan-QAD-FP8-1.3B",
         ],
-        model_detectors=[lambda path: "wandmdpipeline" in path.lower()],
+        model_detectors=[lambda path: "wandmdpipeline" in path.lower() or "fastwan-qad-fp8" in path.lower()],
         model_family="wan",
         default_preset="fast_wan_t2v_480p",
     )

v2/_vendor/layers/quantization/__init__.py

@@ -2,7 +2,7 @@ from typing import Literal, get_args
 
 from v2._vendor.layers.quantization.base_config import QuantizationConfig
 
-QuantizationMethods = Literal[None, "AbsMaxFP8", "NVFP4", "nvfp4_qat"]
+QuantizationMethods = Literal[None, "AbsMaxFP8", "FP8", "NVFP4", "nvfp4_qat"]
 
 QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
 
@@ -49,20 +49,21 @@ 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 .nvfp4_config import NVFP4Config
-    from .nvfp4_qat_config import NVFP4QATConfig
-
-    method_to_config: dict[str, type[QuantizationConfig]] = {
-        "AbsMaxFP8": AbsMaxFP8Config,
-        "NVFP4": NVFP4Config,
-        "nvfp4_qat": NVFP4QATConfig,
-    }
-    # Update the `method_to_config` with customized quantization methods.
-    method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
-
-    return method_to_config[quantization]
+    if quantization in _CUSTOMIZED_METHOD_TO_QUANT_CONFIG:
+        return _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization]
+    if quantization == "AbsMaxFP8":
+        from .absmax_fp8 import AbsMaxFP8Config
+        return AbsMaxFP8Config
+    if quantization == "FP8":
+        from .fp8_config import FP8Config
+        return FP8Config
+    if quantization == "NVFP4":
+        from .nvfp4_config import NVFP4Config
+        return NVFP4Config
+    if quantization == "nvfp4_qat":
+        from .nvfp4_qat_config import NVFP4QATConfig
+        return NVFP4QATConfig
+    raise ValueError(f"Invalid quantization method: {quantization}")
 
 
 all = [

v2/_vendor/layers/quantization/fp8_config.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Generic FP8 quantization backed by ``torch._scaled_mm``.
+
+Matches linear layers by suffix (``to_q/k/v/to_out``, ``ffn.fc_in/fc_out``).
+Supports per-tensor (default, fast) and per-channel (higher accuracy) granularity.
+Falls back to bf16 dequant on GPUs older than sm89.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch.nn.parameter import Parameter
+
+from v2._vendor.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from v2._vendor.models.utils import set_weight_attrs
+
+logger = logging.getLogger(__name__)
+
+FP8_DTYPE = torch.float8_e4m3fn
+FP8_MAX = float(torch.finfo(FP8_DTYPE).max)
+FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0)
+
+_FP8_SUFFIXES = (
+    "ffn.fc_in",
+    "ffn.fc_out",
+    "to_q",
+    "to_k",
+    "to_v",
+    "to_out",
+)
+
+
+def _supports_fp8_compute() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    cap = torch.cuda.get_device_capability()
+    return cap[0] > 8 or (cap[0] == 8 and cap[1] >= 9)
+
+
+def _quantize_tensorwise(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    x_absmax = x_2d.abs().amax().float()
+    x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
+    x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
+    return x_fp8, x_scale.view(1)
+
+
+def _quantize_rowwise(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    x_absmax = x_2d.abs().amax(dim=-1, keepdim=True).float()
+    x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
+    x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
+    return x_fp8, x_scale
+
+
+class FP8QuantizeMethod(QuantizeMethodBase):
+    """FP8 linear method.
+
+    ``granularity='tensor'`` uses per-tensor weight and dynamic activation
+    scales. ``granularity='channel'`` uses per-output-channel weight and
+    per-token activation scales.
+    """
+
+    def __init__(self, granularity: str = "tensor"):
+        super().__init__()
+        self.granularity = granularity
+
+    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,
+    ):
+        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)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        out_dim = layer._fp8_weight.shape[0]
+        original_shape = x.shape
+
+        if not _supports_fp8_compute():
+            return self._apply_dequant(layer, x, bias)
+
+        if self.granularity == "channel":
+            x_fp8, x_scale = _quantize_rowwise(x.reshape(-1, x.shape[-1]))
+        else:
+            x_fp8, x_scale = _quantize_tensorwise(x.reshape(-1, x.shape[-1]))
+
+        w_fp8 = layer._fp8_weight
+        w_scale = layer._fp8_weight_scale
+        scale_a = x_scale.to(device=x_fp8.device, dtype=torch.float32).contiguous()
+        scale_b = w_scale.view(1, -1) if self.granularity == "channel" else w_scale
+        scale_b = scale_b.to(device=w_fp8.device, dtype=torch.float32).contiguous()
+
+        out = torch._scaled_mm(
+            x_fp8,
+            w_fp8.t(),
+            scale_a=scale_a,
+            scale_b=scale_b,
+            out_dtype=torch.bfloat16,
+        )
+        if isinstance(out, tuple):
+            out = out[0]
+        if bias is not None:
+            out = out + bias
+        return out.view(*original_shape[:-1], out_dim)
+
+    def _apply_dequant(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        out_dim = layer._fp8_weight.shape[0]
+        original_shape = x.shape
+        w_fp8 = layer._fp8_weight
+        w_scale = layer._fp8_weight_scale.to(x.dtype)
+        weight = w_fp8.to(x.dtype) * w_scale.unsqueeze(1)
+        out = F.linear(x, weight, bias)
+        return out.view(*original_shape[:-1], out_dim)
+
+
+class FP8Config(QuantizationConfig):
+    """FP8 e4m3 quantization via suffix matching on standard linear layer names."""
+
+    def __init__(self, granularity: str = "tensor"):
+        super().__init__()
+        if granularity not in ("tensor", "channel"):
+            raise ValueError(f"granularity must be 'tensor' or 'channel', got {granularity!r}")
+        self.granularity = granularity
+
+    def get_name(self) -> str:
+        return "FP8"
+
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 89
+
+    @staticmethod
+    def get_config_filenames() -> list[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> FP8Config:
+        return cls(granularity=config.get("granularity", "tensor"))
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
+        from v2._vendor.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase) and any(s in prefix for s in _FP8_SUFFIXES):
+            return FP8QuantizeMethod(granularity=self.granularity)
+        return None
+
+
+def convert_model_to_fp8(model: torch.nn.Module) -> None:
+    """Quantize all FP8-tagged linear layers in-place after weights are loaded."""
+    import gc
+    from torch.distributed.tensor import DTensor
+
+    with torch.no_grad():
+        for mod in model.modules():
+            qm = getattr(mod, "quant_method", None)
+            if not isinstance(qm, FP8QuantizeMethod):
+                continue
+            weight = getattr(mod, "weight", None)
+            if weight is None:
+                continue
+            weight_local = weight.to_local() if isinstance(weight, DTensor) else weight
+            if getattr(qm, "granularity", "tensor") == "channel":
+                w_absmax = weight_local.detach().abs().amax(dim=1).nan_to_num().float()
+                w_scale = (w_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
+                w_fp8 = (weight_local / w_scale.to(weight_local.dtype).unsqueeze(1)).clamp(-FP8_MAX,
+                                                                                           FP8_MAX).to(FP8_DTYPE)
+            else:
+                w_absmax = weight_local.detach().abs().amax().nan_to_num().to(torch.float32)
+                w_scale = (w_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE).view(1)
+                w_fp8 = (weight_local / w_scale.to(weight_local.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
+            mod.register_buffer("_fp8_weight", w_fp8.contiguous(), persistent=False)
+            mod.register_buffer("_fp8_weight_scale", w_scale.to(torch.float32), persistent=False)
+            removed_weight = mod._parameters.pop("weight", None)
+            if removed_weight is not None:
+                removed_weight.grad = None
+            del removed_weight, weight, weight_local, w_absmax, w_scale, w_fp8
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+__all__ = [
+    "FP8Config",
+    "FP8QuantizeMethod",
+    "convert_model_to_fp8",
+]

v2/_vendor/models/dits/wanvideo.py

@@ -7,22 +7,18 @@ import torch
 import torch.nn as nn
 
 import v2._vendor.envs as envs
-from v2._vendor.attention import (DistributedAttention, DistributedAttention_VSA,
-                                 LocalAttention)
+from v2._vendor.attention import (DistributedAttention, DistributedAttention_VSA, LocalAttention)
 from v2._vendor.configs.models.dits import WanVideoConfig
-from v2._vendor.distributed.communication_op import (
-    sequence_model_parallel_all_gather_with_unpad,
-    sequence_model_parallel_shard)
-from v2._vendor.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift,
-                                        RMSNorm, ScaleResidual,
-                                        ScaleResidualLayerNormScaleShift)
+from v2._vendor.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad,
+                                                     sequence_model_parallel_shard)
+from v2._vendor.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual,
+                                         ScaleResidualLayerNormScaleShift)
 from v2._vendor.layers.linear import ReplicatedLinear
 # from torch.nn import RMSNorm
 # TODO: RMSNorm ....
 from v2._vendor.layers.mlp import MLP
 from v2._vendor.layers.rotary_embedding import get_rotary_pos_embed
-from v2._vendor.layers.visual_embedding import (ModulateProjection, PatchEmbed,
-                                               TimestepEmbedder)
+from v2._vendor.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder)
 from v2._vendor.logger import init_logger
 from v2._vendor.models.dits.base import BaseDiT
 from v2._vendor.platforms import AttentionBackendEnum, current_platform
@@ -42,8 +38,7 @@ class WanImageEmbedding(torch.nn.Module):
         self.ff = MLP(in_features, in_features, out_features, act_type="gelu")
         self.norm2 = FP32LayerNorm(out_features)
 
-    def forward(self,
-                encoder_hidden_states_image: torch.Tensor) -> torch.Tensor:
+    def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor:
         dtype = encoder_hidden_states_image.dtype
         hidden_states = self.norm1(encoder_hidden_states_image)
         hidden_states = self.ff(hidden_states)
@@ -62,15 +57,9 @@ class WanTimeTextImageEmbedding(nn.Module):
     ):
         super().__init__()
 
-        self.time_embedder = TimestepEmbedder(
-            dim, frequency_embedding_size=time_freq_dim, act_layer="silu")
-        self.time_modulation = ModulateProjection(dim,
-                                                  factor=6,
-                                                  act_layer="silu")
-        self.text_embedder = MLP(text_embed_dim,
-                                 dim,
-                                 dim,
-                                 bias=True,
+        self.time_embedder = TimestepEmbedder(dim, frequency_embedding_size=time_freq_dim, act_layer="silu")
+        self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu")
+        self.text_embedder = MLP(text_embed_dim, dim, dim, bias=True,
                                  act_type="gelu_pytorch_tanh") if text_embed_dim > 0 else None
 
         self.image_embedder = None
@@ -90,11 +79,12 @@ class WanTimeTextImageEmbedding(nn.Module):
         if self.text_embedder is not None:
             encoder_hidden_states = self.text_embedder(encoder_hidden_states)
         else:
-            encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]), device=temb.device, dtype=temb.dtype)
+            encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]),
+                                                device=temb.device,
+                                                dtype=temb.dtype)
         if encoder_hidden_states_image is not None:
             assert self.image_embedder is not None
-            encoder_hidden_states_image = self.image_embedder(
-                encoder_hidden_states_image)
+            encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image)
 
         return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image
 
@@ -129,17 +119,16 @@ class WanSelfAttention(nn.Module):
         self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
         # Scaled dot product attention
-        self.attn = LocalAttention(
-            num_heads=num_heads,
-            head_size=self.head_dim,
-            dropout_rate=0,
-            softmax_scale=None,
-            causal=False,
-            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
-                                          AttentionBackendEnum.TORCH_SDPA))
+        self.attn = LocalAttention(num_heads=num_heads,
+                                   head_size=self.head_dim,
+                                   dropout_rate=0,
+                                   softmax_scale=None,
+                                   causal=False,
+                                   supported_attention_backends=(AttentionBackendEnum.SAGE_ATTN,
+                                                                 AttentionBackendEnum.FLASH_ATTN,
+                                                                 AttentionBackendEnum.TORCH_SDPA))
 
-    def forward(self, x: torch.Tensor, context: torch.Tensor,
-                context_lens: int):
+    def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int):
         r"""
         Args:
             x(Tensor): Shape [B, L, num_heads, C / num_heads]
@@ -201,8 +190,14 @@ class WanI2VCrossAttention(WanSelfAttention):
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
-        super().__init__(dim, num_heads, window_size, qk_norm, eps,
-                         supported_attention_backends, quant_config=quant_config, prefix=prefix)
+        super().__init__(dim,
+                         num_heads,
+                         window_size,
+                         qk_norm,
+                         eps,
+                         supported_attention_backends,
+                         quant_config=quant_config,
+                         prefix=prefix)
 
         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")
@@ -224,8 +219,7 @@ class WanI2VCrossAttention(WanSelfAttention):
         q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
         k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
         v = self.to_v(context)[0].view(b, -1, n, d)
-        k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(
-            b, -1, n, d)
+        k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(b, -1, n, d)
         v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d)
         img_x = self.attn(q, k_img, v_img)
         # compute attention
@@ -262,12 +256,11 @@ class WanTransformerBlock(nn.Module):
         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, quant_config=quant_config, prefix=f"{prefix}.to_out")
-        self.attn1 = DistributedAttention(
-            num_heads=num_heads,
-            head_size=dim // num_heads,
-            causal=False,
-            supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.attn1")
+        self.attn1 = DistributedAttention(num_heads=num_heads,
+                                          head_size=dim // num_heads,
+                                          causal=False,
+                                          supported_attention_backends=supported_attention_backends,
+                                          prefix=f"{prefix}.attn1")
         self.hidden_dim = dim
         self.num_attention_heads = num_heads
         dim_head = dim // num_heads
@@ -282,13 +275,12 @@ class WanTransformerBlock(nn.Module):
             print("QK Norm type not supported")
             raise Exception
         assert cross_attn_norm is True
-        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
-            dim,
-            norm_type="layer",
-            eps=eps,
-            elementwise_affine=True,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim,
+                                                                        norm_type="layer",
+                                                                        eps=eps,
+                                                                        elementwise_affine=True,
+                                                                        dtype=torch.float32,
+                                                                        compute_dtype=torch.float32)
 
         # 2. Cross-attention
         if added_kv_proj_dim is not None:
@@ -307,13 +299,12 @@ class WanTransformerBlock(nn.Module):
                                               eps=eps,
                                               quant_config=quant_config,
                                               prefix=f"{prefix}.attn2")
-        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
-            dim,
-            norm_type="layer",
-            eps=eps,
-            elementwise_affine=False,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim,
+                                                                         norm_type="layer",
+                                                                         eps=eps,
+                                                                         elementwise_affine=False,
+                                                                         dtype=torch.float32,
+                                                                         compute_dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn")
@@ -338,8 +329,7 @@ class WanTransformerBlock(nn.Module):
         if temb.dim() == 4:
             # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
             shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-                self.scale_shift_table.unsqueeze(0) + temb.float()
-            ).chunk(6, dim=2)
+                self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2)
             # batch_size, seq_len, 1, inner_dim
             shift_msa = shift_msa.squeeze(2)
             scale_msa = scale_msa.squeeze(2)
@@ -350,13 +340,11 @@ class WanTransformerBlock(nn.Module):
         else:
             # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
             e = self.scale_shift_table + temb.float()
-            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
-                6, dim=1)
+            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1)
         assert shift_msa.dtype == torch.float32
 
         # 1. Self-attention
-        norm_hidden_states = (self.norm1(hidden_states.float()) *
-                              (1 + scale_msa) + shift_msa).to(orig_dtype)
+        norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype)
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
@@ -382,19 +370,15 @@ class WanTransformerBlock(nn.Module):
         attn_output = attn_output.squeeze(1)
 
         null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
-        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
-            hidden_states, attn_output, gate_msa, null_shift, null_scale)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
+        norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa,
+                                                                         null_shift, null_scale)
+        norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)
 
         # 2. Cross-attention
-        attn_output = self.attn2(norm_hidden_states,
-                                 context=encoder_hidden_states,
-                                 context_lens=None)
-        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
-            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
+        attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None)
+        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa,
+                                                                          c_scale_msa)
+        norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)
 
         # 3. Feed-forward
         ff_output = self.ffn(norm_hidden_states)
@@ -425,14 +409,17 @@ class WanTransformerBlock_VSA(nn.Module):
         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_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,
-            causal=False,
-            supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.attn1")
+        self.attn1 = DistributedAttention_VSA(num_heads=num_heads,
+                                              head_size=dim // num_heads,
+                                              causal=False,
+                                              supported_attention_backends=supported_attention_backends,
+                                              prefix=f"{prefix}.attn1")
         self.hidden_dim = dim
         self.num_attention_heads = num_heads
         dim_head = dim // num_heads
@@ -447,13 +434,12 @@ class WanTransformerBlock_VSA(nn.Module):
             print("QK Norm type not supported")
             raise Exception
         assert cross_attn_norm is True
-        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
-            dim,
-            norm_type="layer",
-            eps=eps,
-            elementwise_affine=True,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim,
+                                                                        norm_type="layer",
+                                                                        eps=eps,
+                                                                        elementwise_affine=True,
+                                                                        dtype=torch.float32,
+                                                                        compute_dtype=torch.float32)
 
         # 2. Cross-attention
         if added_kv_proj_dim is not None:
@@ -472,13 +458,12 @@ class WanTransformerBlock_VSA(nn.Module):
                                               eps=eps,
                                               quant_config=quant_config,
                                               prefix=f"{prefix}.attn2")
-        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
-            dim,
-            norm_type="layer",
-            eps=eps,
-            elementwise_affine=False,
-            dtype=torch.float32,
-            compute_dtype=torch.float32)
+        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim,
+                                                                         norm_type="layer",
+                                                                         eps=eps,
+                                                                         elementwise_affine=False,
+                                                                         dtype=torch.float32,
+                                                                         compute_dtype=torch.float32)
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn")
@@ -500,13 +485,11 @@ class WanTransformerBlock_VSA(nn.Module):
         orig_dtype = hidden_states.dtype
         # assert orig_dtype != torch.float32
         e = self.scale_shift_table + temb.float()
-        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
-            6, dim=1)
+        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1)
         assert shift_msa.dtype == torch.float32
 
         # 1. Self-attention
-        norm_hidden_states = (self.norm1(hidden_states.float()) *
-                              (1 + scale_msa) + shift_msa).to(orig_dtype)
+        norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype)
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
@@ -520,8 +503,7 @@ class WanTransformerBlock_VSA(nn.Module):
         query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
-        gate_compress = gate_compress.squeeze(1).unflatten(
-            2, (self.num_attention_heads, -1))
+        gate_compress = gate_compress.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
 
         attn_output, _ = self.attn1(
             query,
@@ -536,19 +518,15 @@ class WanTransformerBlock_VSA(nn.Module):
         attn_output = attn_output.squeeze(1)
 
         null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
-        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
-            hidden_states, attn_output, gate_msa, null_shift, null_scale)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
+        norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa,
+                                                                         null_shift, null_scale)
+        norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)
 
         # 2. Cross-attention
-        attn_output = self.attn2(norm_hidden_states,
-                                 context=encoder_hidden_states,
-                                 context_lens=None)
-        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
-            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
+        attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None)
+        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa,
+                                                                          c_scale_msa)
+        norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)
 
         # 3. Feed-forward
         ff_output = self.ffn(norm_hidden_states)
@@ -561,14 +539,12 @@ class WanTransformerBlock_VSA(nn.Module):
 class WanTransformer3DModel(BaseDiT):
     _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
     _compile_conditions = WanVideoConfig()._compile_conditions
-    _supported_attention_backends = WanVideoConfig(
-    )._supported_attention_backends
+    _supported_attention_backends = WanVideoConfig()._supported_attention_backends
     param_names_mapping = WanVideoConfig().param_names_mapping
     reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
     lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping
 
-    def __init__(self, config: WanVideoConfig, hf_config: dict[str,
-                                                               Any]) -> None:
+    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
 
@@ -581,7 +557,8 @@ class WanTransformer3DModel(BaseDiT):
         self.patch_size = config.patch_size
         self.text_len = config.text_len
 
-        assert config.num_attention_heads % get_sp_world_size() == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})"
+        assert config.num_attention_heads % get_sp_world_size(
+        ) == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})"
 
         # 1. Patch & position embedding
         self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
@@ -610,8 +587,7 @@ class WanTransformer3DModel(BaseDiT):
                               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)
+                              prefix=f"{config.prefix}.blocks.{i}") for i in range(config.num_layers)
         ])
 
         # 4. Output norm & projection
@@ -621,10 +597,8 @@ class WanTransformer3DModel(BaseDiT):
                                             elementwise_affine=False,
                                             dtype=torch.float32,
                                             compute_dtype=torch.float32)
-        self.proj_out = nn.Linear(
-            inner_dim, config.out_channels * math.prod(config.patch_size))
-        self.scale_shift_table = nn.Parameter(
-            torch.randn(1, 2, inner_dim) / inner_dim**0.5)
+        self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size))
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)
 
         self.gradient_checkpointing = False
         self.__post_init__()
@@ -640,8 +614,7 @@ class WanTransformer3DModel(BaseDiT):
         orig_dtype = hidden_states.dtype
         if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
+        if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0:
             encoder_hidden_states_image = encoder_hidden_states_image[0]
         else:
             encoder_hidden_states_image = None
@@ -655,27 +628,19 @@ class WanTransformer3DModel(BaseDiT):
         # Get rotary embeddings
         d = self.hidden_size // self.num_attention_heads
         rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
-        freqs_cos, freqs_sin = get_rotary_pos_embed(
-            (post_patch_num_frames, post_patch_height,
-             post_patch_width),
-            self.hidden_size,
-            self.num_attention_heads,
-            rope_dim_list,
-            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
-            rope_theta=10000)
-        freqs_cis = (freqs_cos.to(hidden_states.device).float(),
-                     freqs_sin.to(hidden_states.device).float())
+        freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width),
+                                                    self.hidden_size,
+                                                    self.num_attention_heads,
+                                                    rope_dim_list,
+                                                    dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+                                                    rope_theta=10000)
+        freqs_cis = (freqs_cos.to(hidden_states.device).float(), freqs_sin.to(hidden_states.device).float())
 
         hidden_states = self.patch_embedding(hidden_states)
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
         # Shard with padding support - returns (sharded_tensor, original_seq_len)
         hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
-        
-        current_seq_len = hidden_states.shape[1]
-        sp_world_size = get_sp_world_size()
-        padded_seq_len = current_seq_len * sp_world_size
-        
 
         # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
         if timestep.dim() == 2:
@@ -695,28 +660,25 @@ class WanTransformer3DModel(BaseDiT):
 
         if encoder_hidden_states_image is not None:
             if encoder_hidden_states is not None:
-                encoder_hidden_states = torch.concat(
-                    [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+                encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
             else:
                 encoder_hidden_states = encoder_hidden_states_image
 
         if current_platform.is_mps() or current_platform.is_npu():
             encoder_hidden_states = encoder_hidden_states.to(orig_dtype)
         else:
-            encoder_hidden_states = encoder_hidden_states # cast to orig_dtype for MPS & NPU
+            encoder_hidden_states = encoder_hidden_states  # cast to orig_dtype for MPS & NPU
 
         assert encoder_hidden_states.dtype == orig_dtype
 
         # 4. Transformer blocks
         if torch.is_grad_enabled() and self.gradient_checkpointing:
             for block in self.blocks:
-                hidden_states = self._gradient_checkpointing_func(
-                    block, hidden_states, encoder_hidden_states,
-                    timestep_proj, freqs_cis, original_seq_len)
+                hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states,
+                                                                  timestep_proj, freqs_cis, original_seq_len)
         else:
             for block in self.blocks:
-                hidden_states = block(hidden_states, encoder_hidden_states,
-                                      timestep_proj, freqs_cis, original_seq_len)
+                hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, original_seq_len)
         # 5. Output norm, projection & unpatchify
         if temb.dim() == 3:
             # batch_size, seq_len, inner_dim (wan 2.2 ti2v)
@@ -726,23 +688,20 @@ class WanTransformer3DModel(BaseDiT):
         else:
             # batch_size, inner_dim
             shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
-            
 
         hidden_states = self.norm_out(hidden_states, shift, scale)
 
         # Gather and unpad in one operation
-        hidden_states = sequence_model_parallel_all_gather_with_unpad(
-            hidden_states, original_seq_len, dim=1)
+        hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
         hidden_states = self.proj_out(hidden_states)
 
-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
-                                              post_patch_height,
-                                              post_patch_width, p_t, p_h, p_w,
-                                              -1)
+        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width,
+                                              p_t, p_h, p_w, -1)
         hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
         output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
 
         return output
 
+
 # Entry point for model registry
 EntryClass = WanTransformer3DModel

v2/_vendor/models/loader/fsdp_load.py

@@ -15,55 +15,43 @@ import torch
 from torch import nn
 from torch.distributed import DeviceMesh, init_device_mesh
 from torch.distributed._tensor import distribute_tensor
-from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule,
-                                    MixedPrecisionPolicy, fully_shard)
+from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, fully_shard)
 from torch.nn.modules.module import _IncompatibleKeys
 
 from v2._vendor.logger import init_logger
-from v2._vendor.models.loader.utils import (get_param_names_mapping,
-                                           hf_to_custom_state_dict)
+from v2._vendor.models.loader.utils import (get_param_names_mapping, hf_to_custom_state_dict)
 from v2._vendor.models.loader.weight_utils import safetensors_weights_iterator
 from v2._vendor.utils import set_mixed_precision_policy, is_pin_memory_available
 
 logger = init_logger(__name__)
 
 
-def _maybe_convert_model_to_nvfp4(model: nn.Module) -> None:
-    """Quantize NVFP4-tagged linear layers in-place after weights are loaded.
+def _maybe_quantize_model(model: nn.Module) -> None:
+    """Quantize tagged linear layers in-place after weights are loaded.
 
-    Walks the module tree once, looking for layers whose ``quant_method``
-    is an :class:`NVFP4QuantizeMethod` (attached at construction time by
-    :meth:`NVFP4Config.get_quant_method`). When at least one such layer
-    exists, calls :func:`convert_model_to_nvfp4` to register the
-    ``_nvfp4_weight*`` / ``_nvfp4_alpha`` / ``_weight_global_sf`` buffers
-    on each targeted layer.
-
-    The walk returns on the first NVFP4 layer found so non-NVFP4 callers
-    pay only an ``isinstance`` check per module. flashinfer is imported
-    lazily inside :func:`convert_model_to_nvfp4` so this helper is a
-    no-op on hosts without the NVFP4 backend.
+    Quantization configs attach a ``quant_method`` to each targeted linear layer
+    at construction time. The loader materializes the matching quantized weight
+    buffers from freshly-loaded dense weights, then removes the dense parameter
+    where that quantizer owns storage.
     """
-    # Defer the import: nvfp4_config imports heavy diffusers /
-    # torch.distributed symbols at module-load time, and unconditional
-    # import would penalize every loader call regardless of whether
-    # NVFP4 is wired.
-    from v2._vendor.layers.quantization.nvfp4_config import (
-        NVFP4QuantizeMethod, convert_model_to_nvfp4,
-    )
-    from v2._vendor.layers.quantization.nvfp4_qat_config import (
-        NVFP4QATQuantizeMethod, convert_model_to_fp4,
-    )
-
     for mod in model.modules():
         qm = getattr(mod, "quant_method", None)
-        if isinstance(qm, NVFP4QuantizeMethod):
+        quant_method_name = type(qm).__name__
+        if quant_method_name == "NVFP4QuantizeMethod":
+            from v2._vendor.layers.quantization.nvfp4_config import convert_model_to_nvfp4
             logger.info("Converting loaded model weights for NVFP4 linear layers")
             convert_model_to_nvfp4(model)
             return
-        if isinstance(qm, NVFP4QATQuantizeMethod):
+        if quant_method_name == "NVFP4QATQuantizeMethod":
+            from v2._vendor.layers.quantization.nvfp4_qat_config import convert_model_to_fp4
             logger.info("Converting loaded model weights for NVFP4-QAT linear layers")
             convert_model_to_fp4(model)
             return
+        if quant_method_name == "FP8QuantizeMethod":
+            from v2._vendor.layers.quantization.fp8_config import convert_model_to_fp8
+            logger.info("Converting loaded model weights for FP8 linear layers")
+            convert_model_to_fp8(model)
+            return
 
 
 # TODO(PY): move this to utils elsewhere
@@ -119,10 +107,7 @@ def maybe_load_fsdp_model(
     """
     # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
     # manually casting the inputs to the model
-    mp_policy = MixedPrecisionPolicy(param_dtype,
-                                     reduce_dtype,
-                                     output_dtype,
-                                     cast_forward_inputs=False)
+    mp_policy = MixedPrecisionPolicy(param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False)
 
     set_mixed_precision_policy(
         param_dtype=param_dtype,
@@ -150,7 +135,7 @@ def maybe_load_fsdp_model(
         if not training_mode and not fsdp_inference:
             hsdp_replicate_dim = world_size
             hsdp_shard_dim = 1
-        
+
         if current_platform.is_npu():
             with torch.device("cpu"):
                 device_mesh = init_device_mesh(
@@ -161,11 +146,11 @@ def maybe_load_fsdp_model(
                 )
         else:
             device_mesh = init_device_mesh(
-            "cuda",
-            # (Replicate(), Shard(dim=0))
-            mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
-            mesh_dim_names=("replicate", "shard"),
-        )
+                "cuda",
+                # (Replicate(), Shard(dim=0))
+                mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
+                mesh_dim_names=("replicate", "shard"),
+            )
         shard_model(model,
                     cpu_offload=cpu_offload,
                     reshard_after_forward=True,
@@ -186,30 +171,23 @@ def maybe_load_fsdp_model(
         param_names_mapping=param_names_mapping_fn,
     )
     if hasattr(model, "materialize_non_persistent_buffers"):
-        model.materialize_non_persistent_buffers(
-            device=device, dtype=default_dtype)
+        model.materialize_non_persistent_buffers(device=device, dtype=default_dtype)
     for n, p in chain(model.named_parameters(), model.named_buffers()):
         if p.is_meta:
-            raise RuntimeError(
-                f"Unexpected param or buffer {n} on meta device.")
+            raise RuntimeError(f"Unexpected param or buffer {n} on meta device.")
         # Avoid unintended computation graph accumulation during inference
         if isinstance(p, torch.nn.Parameter):
             p.requires_grad = False
 
-    # NVFP4 weight prequantization. We detect by the registered
-    # ``quant_method`` on linear layers rather than by a separate flag —
-    # construction-time ``NVFP4Config.get_quant_method`` already attached
-    # ``NVFP4QuantizeMethod`` to every targeted layer, so the loader's
-    # responsibility is just to materialize the per-layer nvfp4 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_nvfp4(model)
+    # Post-load weight quantization. We detect by the registered
+    # ``quant_method`` on linear layers rather than by a separate flag.
+    _maybe_quantize_model(model)
 
-    compile_in_loader = enable_torch_compile and training_mode
+    compile_in_loader = enable_torch_compile
     if compile_in_loader:
         compile_kwargs = torch_compile_kwargs or {}
-        logger.info("Enabling torch.compile for FSDP training module with kwargs=%s",
-                    compile_kwargs)
+        profile = "training" if training_mode else "inference"
+        logger.info("Enabling torch.compile for %s module with kwargs=%s", profile, compile_kwargs)
         model = torch.compile(model, **compile_kwargs)
         logger.info("torch.compile enabled for %s", type(model).__name__)
     return model
@@ -253,7 +231,7 @@ def shard_model(
     """
     # Check if we should use size-based filtering
     use_size_filtering = os.environ.get("FASTVIDEO_FSDP2_AUTOWRAP", "0") == "1"
-    
+
     if not fsdp_shard_conditions:
         logger.warning("No FSDP shard conditions provided; nothing will be sharded.")
         return
@@ -264,44 +242,41 @@ def shard_model(
         "mp_policy": mp_policy,
     }
     if cpu_offload:
-        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(
-            pin_memory=pin_cpu_memory)
+        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(pin_memory=pin_cpu_memory)
 
     # iterating in reverse to start with
     # lowest-level modules first
     num_layers_sharded = 0
-    
+
     if use_size_filtering:
         # Size-based filtering mode
         min_params = int(os.environ.get("FASTVIDEO_FSDP2_MIN_PARAMS", "10000000"))
         logger.info("Using size-based filtering with threshold: %.2fM", min_params / 1e6)
-        
+
         for n, m in reversed(list(model.named_modules())):
             if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]):
                 # Count all parameters
                 param_count = sum(p.numel() for p in m.parameters(recurse=True))
-                
+
                 # Skip small modules
                 if param_count < min_params:
-                    logger.info("Skipping module %s (%.2fM params < %.2fM threshold)", 
-                               n, param_count / 1e6, min_params / 1e6)
+                    logger.info("Skipping module %s (%.2fM params < %.2fM threshold)", n, param_count / 1e6,
+                                min_params / 1e6)
                     continue
-                
+
                 # Shard this module
                 logger.info("Sharding module %s (%.2fM params)", n, param_count / 1e6)
                 fully_shard(m, **fsdp_kwargs)
                 num_layers_sharded += 1
     else:
-        # Shard all modules matching conditions        
+        # Shard all modules matching conditions
         for n, m in reversed(list(model.named_modules())):
             if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]):
                 fully_shard(m, **fsdp_kwargs)
                 num_layers_sharded += 1
-        
+
         if num_layers_sharded == 0:
-            raise ValueError(
-                "No layer modules were sharded. Please check if shard conditions are working as expected."
-            )
+            raise ValueError("No layer modules were sharded. Please check if shard conditions are working as expected.")
 
     # Finally shard the entire model to account for any stragglers
     fully_shard(model, **fsdp_kwargs)
@@ -341,16 +316,15 @@ def load_model_from_full_model_state_dict(
     meta_sd = model.state_dict()
     named_parameters = dict(model.named_parameters())
     sharded_sd = {}
-    custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
-        full_sd_iterator, param_names_mapping)  # type: ignore
+    custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(full_sd_iterator,
+                                                                           param_names_mapping)  # type: ignore
     for target_param_name, full_tensor in custom_param_sd.items():
         meta_sharded_param = meta_sd.get(target_param_name)
         if meta_sharded_param is None:
             # Some checkpoints include extra entries that are not part of the
             # instantiated model's state_dict (e.g. `_extra_state` keys from
             # some FSDP checkpoint formats). These can be safely skipped.
-            if (target_param_name.endswith("._extra_state")
-                    or target_param_name.endswith("_extra_state")):
+            if (target_param_name.endswith("._extra_state") or target_param_name.endswith("_extra_state")):
                 logger.warning(
                     "Skipping non-parameter checkpoint key: %s",
                     target_param_name,
@@ -379,9 +353,7 @@ def load_model_from_full_model_state_dict(
             # fall through to the original `sharded_tensor = full_tensor` below.
             if target_param is not None and callable(weight_loader) and tuple(target_param.shape) != tuple(
                     full_tensor.shape):
-                loaded_param = nn.Parameter(torch.empty(tuple(target_param.shape),
-                                                        device=device,
-                                                        dtype=param_dtype),
+                loaded_param = nn.Parameter(torch.empty(tuple(target_param.shape), device=device, dtype=param_dtype),
                                             requires_grad=False)
                 for attr_name, attr_value in vars(target_param).items():
                     setattr(loaded_param, attr_name, attr_value)
@@ -404,31 +376,23 @@ def load_model_from_full_model_state_dict(
     model.reverse_param_names_mapping = reverse_param_names_mapping
     unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys())
     if unused_keys:
-        logger.warning("Found unloaded parameters in meta state dict: %s",
-                       unused_keys)
+        logger.warning("Found unloaded parameters in meta state dict: %s", unused_keys)
 
     # List of allowed parameter name patterns
     ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress", "proj_l"]  # Can be extended as needed
     for new_param_name in unused_keys:
-        if not any(pattern in new_param_name
-                   for pattern in ALLOWED_NEW_PARAM_PATTERNS):
-            logger.error("Unsupported new parameter: %s. Allowed patterns: %s",
-                         new_param_name, ALLOWED_NEW_PARAM_PATTERNS)
-            raise ValueError(
-                f"New parameter '{new_param_name}' is not supported. "
-                f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed."
-            )
+        if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS):
+            logger.error("Unsupported new parameter: %s. Allowed patterns: %s", new_param_name,
+                         ALLOWED_NEW_PARAM_PATTERNS)
+            raise ValueError(f"New parameter '{new_param_name}' is not supported. "
+                             f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed.")
         meta_sharded_param = meta_sd.get(new_param_name)
         if not hasattr(meta_sharded_param, "device_mesh"):
             # Initialize with zeros
-            sharded_tensor = torch.zeros_like(meta_sharded_param,
-                                              device=device,
-                                              dtype=param_dtype)
+            sharded_tensor = torch.zeros_like(meta_sharded_param, device=device, dtype=param_dtype)
         else:
             # Initialize with zeros and distribute
-            full_tensor = torch.zeros_like(meta_sharded_param,
-                                           device=device,
-                                           dtype=param_dtype)
+            full_tensor = torch.zeros_like(meta_sharded_param, device=device, dtype=param_dtype)
             sharded_tensor = distribute_tensor(
                 full_tensor,
                 meta_sharded_param.device_mesh,