diff --git a/v2/_vendor/configs/pipeline_registry.py b/v2/_vendor/configs/pipeline_registry.py index d051eb7c4..e791ebafb 100644 --- a/v2/_vendor/configs/pipeline_registry.py +++ b/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", ) diff --git a/v2/_vendor/layers/quantization/__init__.py b/v2/_vendor/layers/quantization/__init__.py index d650153fa..ee27ccd53 100644 --- a/v2/_vendor/layers/quantization/__init__.py +++ b/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 = [ diff --git a/v2/_vendor/layers/quantization/fp8_config.py b/v2/_vendor/layers/quantization/fp8_config.py new file mode 100644 index 000000000..5250a68b1 --- /dev/null +++ b/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", +] diff --git a/v2/_vendor/models/dits/wanvideo.py b/v2/_vendor/models/dits/wanvideo.py index fb14f4c91..a461790d3 100644 --- a/v2/_vendor/models/dits/wanvideo.py +++ b/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 diff --git a/v2/_vendor/models/loader/fsdp_load.py b/v2/_vendor/models/loader/fsdp_load.py index 36cf9539a..3231b7e37 100644 --- a/v2/_vendor/models/loader/fsdp_load.py +++ b/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,