From 39705f9ee8a9fef6f8115ac12b7782035db3c931 Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Thu, 28 May 2026 10:31:58 +0800 Subject: [PATCH] [Feat] Support online quantization (fp8/int8) for DreamID-Omni (#3902) Signed-off-by: Yuanheng Zhao --- .../x_to_video_audio/x_to_video_audio.md | 2 + .../x_to_video_audio/x_to_video_audio.py | 44 +++- .../diffusion/models/dreamid_omni/fusion.py | 84 +++++-- .../dreamid_omni/pipeline_dreamid_omni.py | 81 +++++- .../diffusion/models/dreamid_omni/wan2_2.py | 238 ++++++++++++++++-- 5 files changed, 397 insertions(+), 52 deletions(-) diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md index 921fc1bd9..f5071ce6d 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md @@ -102,3 +102,5 @@ Key arguments: - `--audio-negative-prompt`: negative prompt for audio generation. - `--enable-cpu-offload`: enable CPU offload (defaults False). - `--cache-backend`: enable `cache_dit` for acceleration. +- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`. + (VAEs, the T5 text encoder, norms and modulation stay bf16.) diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py index 96e5b3592..30de8aa45 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py @@ -4,6 +4,7 @@ import argparse import re import time +from typing import Any import numpy as np from PIL import Image @@ -41,6 +42,12 @@ def parse_args() -> argparse.Namespace: choices=[1, 2, 3, 4], help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).", ) + parser.add_argument( + "--tensor-parallel-size", + type=int, + default=1, + help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."), + ) parser.add_argument( "--video-negative-prompt", default="jitter, bad hands, blur, distortion", @@ -52,6 +59,13 @@ def parse_args() -> argparse.Namespace: help="Negative prompt for audio.", ) parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.") + parser.add_argument( + "--quantization", + type=str, + default=None, + choices=["fp8", "int8"], + help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "), + ) parser.add_argument( "--enable-cpu-offload", action="store_true", @@ -100,6 +114,27 @@ def load_image_and_audio(image_paths, audio_paths): return image, audio +def _extract_peak_memory_mb(result: Any) -> float: + """Pull worker-reported peak VRAM (MiB) generation result. + + Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb. + """ + if isinstance(result, list): + result = result[0] if result else None + if result is None: + return 0.0 + val = getattr(result, "peak_memory_mb", 0.0) + if not val: + inner = getattr(result, "request_output", None) + if isinstance(inner, list): + inner = inner[0] if inner else None + val = getattr(inner, "peak_memory_mb", 0.0) + try: + return float(val or 0.0) + except (TypeError, ValueError): + return 0.0 + + def main() -> None: args = parse_args() if args.prompt is None and args.prompt_file is None: @@ -142,6 +177,7 @@ def main() -> None: parallel_config = DiffusionParallelConfig( cfg_parallel_size=args.cfg_parallel_size, + tensor_parallel_size=args.tensor_parallel_size, use_hsdp=args.use_hsdp, hsdp_shard_size=args.hsdp_shard_size, hsdp_replicate_size=args.hsdp_replicate_size, @@ -162,7 +198,7 @@ def main() -> None: "scm_steps_policy": "dynamic", } - omni = Omni( + omni_kwargs = dict( model=args.model, parallel_config=parallel_config, model_type=args.model_type, @@ -171,6 +207,9 @@ def main() -> None: cache_backend=args.cache_backend, cache_config=cache_config, ) + if args.quantization is not None: + omni_kwargs["quantization"] = args.quantization + omni = Omni(**omni_kwargs) start = time.perf_counter() outputs = omni.generate(prompt, sampling_params) elapsed = time.perf_counter() - start @@ -208,6 +247,9 @@ def main() -> None: f.write(video_bytes) print(f"Saved generated video to {output_path}") print(f"Total time: {elapsed:.2f}s") + peak_mb = _extract_peak_memory_mb(outputs) + if peak_mb: + print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)") if __name__ == "__main__": diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py index d0b00bc5b..423935427 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py +++ b/vllm_omni/diffusion/models/dreamid_omni/fusion.py @@ -1,18 +1,28 @@ import re +from typing import TYPE_CHECKING import torch import torch.nn as nn +from vllm.distributed import get_tensor_model_parallel_world_size from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm_omni.diffusion.attention.layer import Attention +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + try: - from dreamid_omni.modules.model import WanLayerNorm, WanRMSNorm + from dreamid_omni.modules.model import WanLayerNorm except ImportError: raise ImportError("Failed to import from dependency 'dreamid_omni'.") from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import WanModel, rope_apply +from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import ( + DistributedRMSNorm, + WanModel, + rope_apply, +) logger = init_logger(__name__) @@ -243,26 +253,36 @@ class FusedBlock(nn.Module): class FusionModel(nn.Module): _layerwise_offload_blocks_attrs = ["fused_blocks"] + packed_modules_mapping = { + "to_qkv": ["q", "k", "v"], + } + @staticmethod def _is_fused_block(name: str, module) -> bool: return "fused_blocks" in name and name.split(".")[-1].isdigit() _hsdp_shard_conditions = [_is_fused_block] - def __init__(self, video_config=None, audio_config=None): + def __init__( + self, + video_config=None, + audio_config=None, + *, + quant_config: "QuantizationConfig | None" = None, + ): super().__init__() has_video = True has_audio = True self.device = get_local_device() if video_config is not None: - self.video_model = WanModel(**video_config) + self.video_model = WanModel(quant_config=quant_config, prefix="video_model", **video_config) else: has_video = False self.video_model = None logger.warning("No video model is provided!") if audio_config is not None: - self.audio_model = WanModel(**audio_config) + self.audio_model = WanModel(quant_config=quant_config, prefix="audio_model", **audio_config) else: has_audio = False self.audio_model = None @@ -272,11 +292,12 @@ class FusionModel(nn.Module): assert len(self.video_model.blocks) == len(self.audio_model.blocks) self.num_blocks = len(self.video_model.blocks) - self.inject_cross_attention_kv_projections() + self.inject_cross_attention_kv_projections(quant_config=quant_config) - self.num_heads = self.video_model.num_heads + tp_size = get_tensor_model_parallel_world_size() + self.full_num_heads = self.video_model.num_heads + self.num_heads = self.full_num_heads // tp_size self.head_dim = self.video_model.dim // self.video_model.num_heads - # Make a single shared instance to pass in at forward time self.attn = Attention( num_heads=self.num_heads, head_size=self.head_dim, @@ -315,21 +336,42 @@ class FusionModel(nn.Module): return super().load_state_dict(state_dict, strict=strict, assign=assign) - def inject_cross_attention_kv_projections(self): - for vid_block in self.video_model.blocks: - vid_block.cross_attn.k_fusion = nn.Linear(vid_block.dim, vid_block.dim) - vid_block.cross_attn.v_fusion = nn.Linear(vid_block.dim, vid_block.dim) - vid_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True) - vid_block.cross_attn.norm_k_fusion = ( - WanRMSNorm(vid_block.dim, eps=1e-6) if vid_block.qk_norm else nn.Identity() + def inject_cross_attention_kv_projections( + self, + quant_config: "QuantizationConfig | None" = None, + ): + def _make_kv_proj(dim: int, prefix: str) -> ColumnParallelLinear: + # K/V projections for the cross-modal (a2v / v2a) attention path. + # Inputs are full latents, so we quantize matched to self_attn. + return ColumnParallelLinear( + dim, + dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=prefix, ) - for audio_block in self.audio_model.blocks: - audio_block.cross_attn.k_fusion = nn.Linear(audio_block.dim, audio_block.dim) - audio_block.cross_attn.v_fusion = nn.Linear(audio_block.dim, audio_block.dim) - audio_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True) - audio_block.cross_attn.norm_k_fusion = ( - WanRMSNorm(audio_block.dim, eps=1e-6) if audio_block.qk_norm else nn.Identity() + tp_size = get_tensor_model_parallel_world_size() + for i, vid_block in enumerate(self.video_model.blocks): + cross = vid_block.cross_attn + block_prefix = f"video_model.blocks.{i}.cross_attn" + cross.k_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.k_fusion") + cross.v_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.v_fusion") + cross.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True) + cross.norm_k_fusion = ( + DistributedRMSNorm(vid_block.dim // tp_size, eps=1e-6) if vid_block.qk_norm else nn.Identity() + ) + + for i, audio_block in enumerate(self.audio_model.blocks): + cross = audio_block.cross_attn + block_prefix = f"audio_model.blocks.{i}.cross_attn" + cross.k_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.k_fusion") + cross.v_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.v_fusion") + cross.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True) + cross.norm_k_fusion = ( + DistributedRMSNorm(audio_block.dim // tp_size, eps=1e-6) if audio_block.qk_norm else nn.Identity() ) def _detach_blocks_from_backbones(self) -> None: diff --git a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py index cc932f8c1..e223c516e 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py +++ b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py @@ -13,6 +13,11 @@ from PIL import Image, ImageOps from torch import nn from torchvision.transforms import Compose, Normalize from tqdm import tqdm +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin @@ -124,7 +129,8 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor self.text_encoder = self.text_model.model # Fusion model — weights are loaded later via load_weights() - self.model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG) + quant_config = getattr(self.od_config, "quantization_config", None) + self.model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG, quant_config=quant_config) self.transformer = self.model fusion_path = self.od_config.tf_model_config.get("fusion", None) @@ -236,10 +242,75 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor return ref_vae_latents, ref_audio_lengths def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - prefix = "model." - state_dict = {name[len(prefix) :]: tensor for name, tensor in weights if name.startswith(prefix)} - self.model.load_state_dict(state_dict, strict=True) - return {prefix + k for k in state_dict} + pipeline_prefix = "model." + + qkv_stacked_mapping = [ + (".self_attn.to_qkv", ".self_attn.q", "q"), + (".self_attn.to_qkv", ".self_attn.k", "k"), + (".self_attn.to_qkv", ".self_attn.v", "v"), + ] + + params_dict = dict(self.model.named_parameters()) + loaded_local: set[str] = set() + + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0 + + def _maybe_shard_weight(weight: torch.Tensor, param: torch.Tensor) -> torch.Tensor: + """Slice a checkpoint tensor to this rank's TP shard when the + param is a TP-local shape but the checkpoint stores the full one. + """ + if tp_size <= 1 or weight.shape == param.shape: + return weight + if weight.ndim == 1 and weight.numel() == param.numel() * tp_size: + return weight.chunk(tp_size, dim=0)[tp_rank] + if weight.ndim == 2: + if weight.shape[0] == param.shape[0] * tp_size: + return weight.chunk(tp_size, dim=0)[tp_rank] + if weight.shape[1] == param.shape[1] * tp_size: + return weight.chunk(tp_size, dim=1)[tp_rank] + return weight + + for full_name, loaded_weight in weights: + if not full_name.startswith(pipeline_prefix): + continue + name = full_name[len(pipeline_prefix) :] + + # Fused-QKV stacking: route checkpoint's separate + # - ...self_attn.{q,k,v}: weights into the corresponding + # - ...self_attn.to_qkv : via its sharded weight_loader. + stacked_hit = False + for fused_name, raw_name, shard_id in qkv_stacked_mapping: + if raw_name not in name: + continue + stacked = name.replace(raw_name, fused_name) + if stacked not in params_dict: + continue + param = params_dict[stacked] + param.weight_loader(param, loaded_weight, shard_id) + loaded_local.add(stacked) + stacked_hit = True + break + if stacked_hit: + continue + + if name not in params_dict: + logger.warning(f"Skipping DreamID-Omni weight {full_name} -- not found in model parameters") + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", None) + if weight_loader is not None: + # Parallel Linears handle TP sharding internally + weight_loader(param, loaded_weight) + else: + default_weight_loader(param, _maybe_shard_weight(loaded_weight, param)) + loaded_local.add(name) + + # Detach must happen before `DiffusersPipelineLoader._process_weights_after_loading` + # and before any forward pass that triggers layerwise-offload hooks. + self.model._detach_blocks_from_backbones() + + return {pipeline_prefix + k for k in loaded_local} def get_scheduler_time_steps(self, sampling_steps, solver_name="unipc", device=0, shift=5.0): torch.manual_seed(4) diff --git a/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py b/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py index db289daf9..c7e464cc8 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py +++ b/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py @@ -6,6 +6,16 @@ import torch.amp as amp import torch.nn as nn from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.modeling_utils import ModelMixin +from vllm.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion.attention.layer import Attention @@ -17,7 +27,6 @@ try: MLPProj, ModulationAdd, WanLayerNorm, - WanRMSNorm, rope_apply, # diff from wan2.2 rope_params, sinusoidal_embedding_1d, @@ -26,42 +35,99 @@ except ImportError: raise ImportError("Failed to import from dependency 'dreamid_omni'.") +class DistributedRMSNorm(nn.Module): + """RMSNorm that computes global RMS across tensor parallel ranks. + + Mirrors `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py::DistributedRMSNorm` + """ + + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + tp_size = get_tensor_model_parallel_world_size() + x_dtype = x.dtype + x_float = x.float() + local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True) + local_count = x.shape[-1] + + if tp_size > 1: + local_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq) + local_count = local_count * tp_size + + rms = torch.sqrt(local_sum_sq / local_count + self.eps) + return (x_float / rms * self.weight.float()).to(x_dtype) + + class WanSelfAttention(nn.Module): """Optimized self-attention module using vLLM layers.""" - def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6): + def __init__( + self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): assert dim % num_heads == 0 super().__init__() self.dim = dim - self.num_heads = num_heads self.head_dim = dim // num_heads self.window_size = window_size self.qk_norm = qk_norm self.eps = eps - # layers - self.q = nn.Linear(dim, dim) - self.k = nn.Linear(dim, dim) - self.v = nn.Linear(dim, dim) - self.o = nn.Linear(dim, dim) - self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=self.head_dim, + total_num_heads=num_heads, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_qkv" if prefix else "to_qkv", + ) + self.num_heads = self.to_qkv.num_heads + self.num_kv_heads = self.to_qkv.num_kv_heads + self.tp_inner_dim = self.num_heads * self.head_dim + self.tp_kv_dim = self.num_kv_heads * self.head_dim + + self.o = RowParallelLinear( + dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o" if prefix else "o", + ) + self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = DistributedRMSNorm(self.tp_kv_dim, eps=eps) if qk_norm else nn.Identity() # Unified attention layer self.attn = Attention( num_heads=self.num_heads, head_size=self.head_dim, - num_kv_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, softmax_scale=1.0 / (self.head_dim**0.5), causal=False, + prefix=prefix, ) # query, key, value function def qkv_fn(self, x): - b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + b, s = x.shape[:2] + qkv, _ = self.to_qkv(x) + q_size = self.num_heads * self.head_dim + kv_size = self.num_kv_heads * self.head_dim + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) - q = self.norm_q(self.q(x)).view(b, s, n, d) - k = self.norm_k(self.k(x)).view(b, s, n, d) - v = self.v(x).view(b, s, n, d) + q = self.norm_q(q).view(b, s, self.num_heads, self.head_dim) + k = self.norm_k(k).view(b, s, self.num_kv_heads, self.head_dim) + v = v.view(b, s, self.num_kv_heads, self.head_dim) return q, k, v def forward(self, x, seq_lens, grid_sizes, freqs, ref_lengths=None, freqs_scaling=None): @@ -87,14 +153,97 @@ class WanSelfAttention(nn.Module): return x -class WanT2VCrossAttention(WanSelfAttention): - def qkv_fn(self, x, context): - b, n, d = x.size(0), self.num_heads, self.head_dim +class WanT2VCrossAttention(nn.Module): + """Text -> latent cross-attention.""" - # compute query, key, value - q = self.norm_q(self.q(x)).view(b, -1, n, d) - k = self.norm_k(self.k(context)).view(b, -1, n, d) - v = self.v(context).view(b, -1, n, d) + def __init__( + self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + assert dim % num_heads == 0 + self.dim = dim + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + + # Separate Q/K/V — Q's input differs from K/V's input (text context) + self.q = ColumnParallelLinear( + dim, + dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q" if prefix else "q", + ) + self.k = ColumnParallelLinear( + dim, + dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.k" if prefix else "k", + ) + self.v = ColumnParallelLinear( + dim, + dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.v" if prefix else "v", + ) + + tp_size = get_tensor_model_parallel_world_size() + self.num_heads = num_heads // tp_size + self.num_kv_heads = self.num_heads + self.tp_inner_dim = self.num_heads * self.head_dim + + self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() + + self.o = RowParallelLinear( + dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o" if prefix else "o", + ) + + # Kept for interface parity with WanSelfAttention; + # the active cross-attention path in DreamID-Omni uses the shared FusionModel.attn. + # Disable kv quant to prevent affecting quality. + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + num_kv_heads=self.num_kv_heads, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + prefix=prefix, + disable_kv_quant=True, + ) + + def qkv_fn(self, x, context): + b = x.size(0) + + q = self.q(x) + k = self.k(context) + v = self.v(context) + q = self.norm_q(q).view(b, -1, self.num_heads, self.head_dim) + k = self.norm_k(k).view(b, -1, self.num_kv_heads, self.head_dim) + v = v.view(b, -1, self.num_kv_heads, self.head_dim) return q, k, v @@ -142,6 +291,9 @@ class WanAttentionBlock(nn.Module): cross_attn_norm=False, eps=1e-6, additional_emb_length=None, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ): super().__init__() self.dim = dim @@ -154,7 +306,15 @@ class WanAttentionBlock(nn.Module): # layers self.norm1 = WanLayerNorm(dim, eps) - self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps) + self.self_attn = WanSelfAttention( + dim, + num_heads, + window_size, + qk_norm, + eps, + quant_config=quant_config, + prefix=f"{prefix}.self_attn" if prefix else "self_attn", + ) self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() if cross_attn_type == "i2v_cross_attn": assert additional_emb_length is not None, "additional_emb_length should be specified for i2v_cross_attn" @@ -167,9 +327,32 @@ class WanAttentionBlock(nn.Module): (-1, -1), qk_norm, eps, + quant_config=quant_config, + prefix=f"{prefix}.cross_attn" if prefix else "cross_attn", ) self.norm2 = WanLayerNorm(dim, eps) - self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim)) + # FFN layout: Stay compatible with upstream DreamID-Omni checkpoints + self.ffn = nn.Sequential( + ColumnParallelLinear( + dim, + ffn_dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.ffn.0" if prefix else "ffn.0", + ), + nn.GELU(approximate="tanh"), + RowParallelLinear( + ffn_dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.ffn.2" if prefix else "ffn.2", + ), + ) # modulation self.modulation = ModulationAdd(dim, 6) @@ -243,6 +426,9 @@ class WanModel(ModelMixin, ConfigMixin): cross_attn_norm=True, temporal_rope_scaling_factor=1.0, eps=1e-6, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ): r""" Initialize the diffusion model backbone. @@ -357,8 +543,10 @@ class WanModel(ModelMixin, ConfigMixin): cross_attn_norm, eps, additional_emb_length, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{layer_idx}" if prefix else f"blocks.{layer_idx}", ) - for _ in range(num_layers) + for layer_idx in range(num_layers) ] )