diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index 2892e16a02e..de4d6966a61 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -295,6 +295,7 @@ def _wan_inputs() -> dict[str, torch.Tensor]: # Requires: hidden_states, encoder_hidden_states, timestep in_channels = getattr(cfg, "in_channels", 16) text_dim = getattr(cfg, "text_dim", 4096) + image_dim = getattr(cfg, "image_dim", None) max_seq_len = getattr(cfg, "rope_max_seq_len", 512) patch_dtype = getattr(getattr(model, "patch_embedding", None), "weight", None) @@ -308,13 +309,14 @@ def _wan_inputs() -> dict[str, torch.Tensor]: # Wan expects num_frames = 4 * n + 1; keep n small for dummy forward num_frames = 5 - text_seq_len = min(max_seq_len, 512) + # I2V attention splits image context from the final 512 text tokens. + text_seq_len = 512 if image_dim is not None else min(max_seq_len, 512) # Keep spatial dims small and divisible by patch size (default 2x2) height = 8 width = 8 - return { + dummy_inputs = { "hidden_states": torch.randn( batch_size, in_channels, num_frames, height, width, device=device, dtype=patch_dtype ), @@ -324,6 +326,21 @@ def _wan_inputs() -> dict[str, torch.Tensor]: "timestep": torch.randint(0, 1000, (batch_size,), device=device), "return_dict": False, } + if image_dim is not None: + image_embedder = model.condition_embedder.image_embedder + # The projection bias retains its compute dtype when weights are packed. + image_dtype = image_embedder.ff.net[0].proj.bias.dtype + image_batch_size = batch_size + image_seq_len = 257 + pos_embed_seq_len = getattr(cfg, "pos_embed_seq_len", None) + if pos_embed_seq_len is not None: + # FLF2V concatenates first/last-frame features inside the image embedder. + image_batch_size *= 2 + image_seq_len = pos_embed_seq_len // 2 + dummy_inputs["encoder_hidden_states_image"] = torch.randn( + image_batch_size, image_seq_len, image_dim, device=device, dtype=image_dtype + ) + return dummy_inputs def _qwen_inputs() -> dict[str, Any]: # QwenImageTransformer2DModel does NOT take the standard diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 37187dd3397..4a2a0ae95e1 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1328,7 +1328,13 @@ def _export_diffusers_checkpoint( # dummy forward fails for it, fail loudly rather than silently skipping # fusion (which would export un-unified amax values). is_qwen_component = "qwen" in type(component).__name__.lower() - _fuse_qkv_linears_diffusion(component, strict=is_qwen_component) + # Image-conditioned Wan also requires a successful dummy forward for fusion. + is_wan_i2v_component = ( + HAS_DIFFUSERS + and isinstance(component, getattr(diffusers, "WanTransformer3DModel", ())) + and getattr(getattr(component, "config", None), "image_dim", None) is not None + ) + _fuse_qkv_linears_diffusion(component, strict=is_qwen_component or is_wan_i2v_component) # Process quantized modules (convert weights, register scales) _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False) diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index eedeae677e1..ec7787687e5 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -286,6 +286,50 @@ def test_flux2_dummy_inputs_shape(): assert "guidance" in inputs +@pytest.mark.parametrize("pos_embed_seq_len", [None, 514], ids=["i2v", "flf2v"]) +@torch.no_grad() +def test_wan_dummy_inputs_support_qkv_fusion(pos_embed_seq_len): + torch.manual_seed(0) + model = get_tiny_wan22_transformer( + image_dim=32, added_kv_proj_dim=24, pos_embed_seq_len=pos_embed_seq_len + ).eval() + inputs = generate_diffusion_dummy_inputs(model, torch.device("cpu"), torch.float32) + assert inputs is not None + image_batch_size = 1 if pos_embed_seq_len is None else 2 + assert inputs["encoder_hidden_states_image"].shape == (image_batch_size, 257, 32) + assert inputs["encoder_hidden_states"].shape == (1, 512, 32) + + # Per-tensor INT8 exercises both input and weight amax fusion on CPU. + quant_cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG) + quant_cfg["quant_cfg"].append( + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": None}} + ) + mtq.quantize(model, quant_cfg, lambda m: m(**inputs)) + + groups = [] + for block in model.blocks: + groups.extend( + [ + (block.attn1.to_q, block.attn1.to_k, block.attn1.to_v), + (block.attn2.to_k, block.attn2.to_v), + (block.attn2.add_k_proj, block.attn2.add_v_proj), + ] + ) + # Distinct values ensure a skipped fusion cannot pass on already-equal scales. + for group in groups: + for index, module in enumerate(group, start=1): + module.input_quantizer.amax = torch.tensor(float(index)) + module.weight_quantizer.amax = torch.tensor(float(index)) + + unified_export_hf._fuse_qkv_linears_diffusion(model, strict=True) + + for group in groups: + expected_amax = torch.tensor(float(len(group))) + for module in group: + torch.testing.assert_close(module.input_quantizer.amax, expected_amax) + torch.testing.assert_close(module.weight_quantizer.amax, expected_amax) + + @pytest.mark.parametrize( ("model_factory", "export_attention_quantizers", "expect_attention_quantizers"), [