diff --git a/videox_fun/models/__init__.py b/videox_fun/models/__init__.py index 1b6c9d8..c6540de 100755 --- a/videox_fun/models/__init__.py +++ b/videox_fun/models/__init__.py @@ -91,7 +91,8 @@ from .qwenimage_transformer2d_control import QwenImageControlTransformer2DModel from .qwenimage_transformer2d_instantx import QwenImageInstantXControlNetModel from .qwenimage_vae import AutoencoderKLQwenImage from .spark_transformer3d import (SparkWan2_2Transformer3DModel, - SparkWanSelfAttention) + SparkWanSelfAttention, + SparkWanTransformer3DModel) from .turbowan_transformer3d import TurboWanTransformer3DModel from .wan_audio_encoder import WanAudioEncoder from .wan_image_encoder import CLIPModel diff --git a/videox_fun/models/spark_transformer3d.py b/videox_fun/models/spark_transformer3d.py index 8624feb..f6ec709 100644 --- a/videox_fun/models/spark_transformer3d.py +++ b/videox_fun/models/spark_transformer3d.py @@ -1,9 +1,10 @@ -# SparkDiffusion sparse student backbone for the Wan2.2 A14B MoE models. +# SparkDiffusion sparse student backbones for the Wan2.2 A14B MoE models and the +# Wan2.1 single-expert models. import torch from .attention_utils import (RoLASparseAttention, SparseLinearAttention, attention) -from .wan_transformer3d import (Wan2_2Transformer3DModel, WanSelfAttention, +from .wan_transformer3d import (Wan2_2Transformer3DModel, WanSelfAttention, WanTransformer3DModel, rope_apply, rope_apply_qk) @@ -261,3 +262,141 @@ class SparkWan2_2Transformer3DModel(Wan2_2Transformer3DModel): def sparse_trainable_params(self): """Names of the extra sparse-branch parameters, for logging / stage-A freezing.""" return [name for name, _ in self.named_parameters() if "attn_op." in name] + + +class SparkWanTransformer3DModel(WanTransformer3DModel): + """Wan2.1 transformer (T2V / I2V, single expert) with a sparse student + self-attention. Block type stays WanAttentionBlock, so the existing + `--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock` policy keeps working.""" + + # NOTE: diffusers from_pretrained filters config.json keys by __init__ signature, so + # every parent parameter must be declared explicitly (a *args/**kwargs passthrough would + # silently drop dim/ffn_dim/... and build a wrongly-sized model). + def __init__( + self, + model_type='t2v', + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6, + in_channels=16, + hidden_size=2048, + add_control_adapter=False, + in_dim_control_adapter=24, + downscale_factor_control_adapter=8, + add_ref_conv=False, + in_dim_ref_conv=16, + cross_attn_type=None, + student_topk=0.03, + attn_variant='rola', + rola_rank=64, + ): + super().__init__( + model_type=model_type, + patch_size=patch_size, + text_len=text_len, + in_dim=in_dim, + dim=dim, + ffn_dim=ffn_dim, + freq_dim=freq_dim, + text_dim=text_dim, + out_dim=out_dim, + num_heads=num_heads, + num_layers=num_layers, + window_size=window_size, + qk_norm=qk_norm, + cross_attn_norm=cross_attn_norm, + eps=eps, + in_channels=in_channels, + hidden_size=hidden_size, + add_control_adapter=add_control_adapter, + in_dim_control_adapter=in_dim_control_adapter, + downscale_factor_control_adapter=downscale_factor_control_adapter, + add_ref_conv=add_ref_conv, + in_dim_ref_conv=in_dim_ref_conv, + cross_attn_type=cross_attn_type, + ) + self.student_topk = student_topk + self.attn_variant = attn_variant + self.rola_rank = rola_rank + self._apply_sparse_self_attention(student_topk) + + def _apply_sparse_self_attention(self, topk): + for block in self.blocks: + dense_attn = block.self_attn + sparse_attn = SparkWanSelfAttention( + dense_attn.dim, dense_attn.num_heads, dense_attn.window_size, + dense_attn.qk_norm, dense_attn.eps, topk=topk, + attn_variant=self.attn_variant, rola_rank=self.rola_rank, + ) + # Under low_cpu_mem_usage from_pretrained the replaced params are meta tensors + # and get filled by the checkpoint load afterwards; only upgrade in place (copy + # the real tensors, the freshly-init RoLa params always survive) when weights are + # concrete. + if not next(dense_attn.parameters()).is_meta: + old_sd = {k: v for k, v in dense_attn.state_dict().items() if not v.is_meta} + # Match the live params' device AND dtype: freshly built RoLa params are fp32, + # and mixing them with bf16 backbone weights breaks FSDP flat-param flattening. + p0 = next(dense_attn.parameters()) + sparse_attn = sparse_attn.to(device=p0.device, dtype=p0.dtype) + sparse_attn.load_state_dict(old_sd, strict=False) + block.self_attn = sparse_attn + + def set_student_topk(self, topk): + """(Re)build the sparse kernels at the given top-k ratio. Always recreates + attn_op, so the RoLa params (proj_q/proj_k/gate_proj/gate_bias) are reset to their + init: call right after from_pretrained and BEFORE loading a stage-1 sparse checkpoint + (its trained RoLa params would be dropped otherwise).""" + self.student_topk = topk + self._apply_sparse_self_attention(topk) + + def enable_multi_gpus_inference(self,): + """Bind the sparse-aware USP forward instead of the parent's dense one. + + ``WanTransformer3DModel.enable_multi_gpus_inference`` monkey-patches every + ``blocks[i].self_attn.forward`` to ``usp_attn_forward`` (dense full-softmax via + ``xFuserLongContextAttention``), which never reaches ``SparkWanSelfAttention.attn_op`` + and would silently drop the trained RoLa params under multi-GPU. Rebind to + ``usp_attn_sparkwan_forward`` (ulysses head-parallel + RoLa kernel) so the sparse branch + stays live; a plain ``shard_model`` FSDP wrap does not help here -- it only shards + params and shares the same SP branch, so the patch still applies. + """ + import types + + from ..dist import (get_sequence_parallel_rank, + get_sequence_parallel_world_size, get_sp_group, + usp_attn_sparkwan_forward) + self.sp_world_size = get_sequence_parallel_world_size() + self.sp_world_rank = get_sequence_parallel_rank() + self.all_gather = get_sp_group().all_gather + + # Ulysses head-split needs num_heads divisible by the SP degree. + n_heads = self.blocks[0].self_attn.num_heads + if self.sp_world_size > 1 and n_heads % self.sp_world_size != 0: + raise ValueError( + f"sparse multi-GPU (ulysses) needs num_heads ({n_heads}) divisible by the " + f"sequence-parallel degree ({self.sp_world_size}); set ulysses_degree accordingly " + f"and keep ring_degree=1 (the SLA kernel has no ring-attention path).") + + for block in self.blocks: + block.self_attn.forward = types.MethodType( + usp_attn_sparkwan_forward, block.self_attn) + + if hasattr(self, 'vace_blocks'): + for block in self.vace_blocks: + block.self_attn.forward = types.MethodType( + usp_attn_sparkwan_forward, block.self_attn) + + def sparse_trainable_params(self): + """Names of the extra sparse-branch parameters, for logging / stage-A freezing.""" + return [name for name, _ in self.named_parameters() if "attn_op." in name]