From cecaa6344d79cb586d57e49971f50ce830d52a17 Mon Sep 17 00:00:00 2001 From: hsliuustc0106 Date: Thu, 3 Sep 2026 12:27:26 +0000 Subject: [PATCH] Fix Hunyuan Image-3 expert mapping contract Signed-off-by: hsliuustc0106 --- .../models/test_kv_cache_scale_mapper.py | 50 ++++++++++++++++++- .../hunyuan_image3_transformer.py | 42 +++++++++------- .../models/hunyuan_image3/hunyuan_image3.py | 44 ++++++++-------- 3 files changed, 95 insertions(+), 41 deletions(-) diff --git a/tests/model_executor/models/test_kv_cache_scale_mapper.py b/tests/model_executor/models/test_kv_cache_scale_mapper.py index 98f016bb4..bba85e69a 100644 --- a/tests/model_executor/models/test_kv_cache_scale_mapper.py +++ b/tests/model_executor/models/test_kv_cache_scale_mapper.py @@ -67,8 +67,11 @@ class _FakeHunyuanModel(_FakeModel): use_cla=False, ) - def get_expert_mapping(self) -> tuple[list[object], dict[str, object]]: - return [], {} + def get_expert_mapping(self) -> list[object]: + return [] + + def _get_expert_weights_remapping(self) -> dict[str, object]: + return {} def _split_qkv_weight(self, weight: torch.Tensor) -> torch.Tensor: return weight @@ -81,6 +84,49 @@ def test_model_loaders_do_not_call_removed_get_cache_scale(rel_path: str) -> Non assert ".get_cache_scale(" not in source +@pytest.mark.parametrize( + ("module_name", "class_name"), + [ + ( + "vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3", + "HunyuanModel", + ), + ( + "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer", + "HunyuanImage3Model", + ), + ], +) +def test_hunyuan_expert_mapping_uses_vllm_public_contract( + module_name: str, + class_name: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The public expert mapping must be a flat list, not a loader-side pair.""" + module = __import__(module_name, fromlist=[class_name]) + expected_mapping = [("param", "weight", 0, "gate_proj")] + fake_model = SimpleNamespace( + config=SimpleNamespace(num_experts=1), + num_redundant_experts=0, + ) + monkeypatch.setattr(module, "_is_moe", lambda config: True) + if hasattr(module, "fused_moe_make_expert_params_mapping"): + monkeypatch.setattr(module, "fused_moe_make_expert_params_mapping", lambda **kwargs: expected_mapping) + else: + monkeypatch.setattr( + module, + "FusedMoE", + SimpleNamespace(make_expert_params_mapping=lambda *args, **kwargs: expected_mapping), + ) + + model_cls = getattr(module, class_name) + assert model_cls.get_expert_mapping(fake_model) == expected_mapping + assert model_cls._get_expert_weights_remapping(fake_model) == { + "gate_proj": ("gate_and_up_proj", 1, 2), + "up_proj": ("gate_and_up_proj", 0, 2), + } + + @pytest.mark.parametrize( ("module_name", "class_name", "fake_model_cls"), [ diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 856e6f189..dd1d4af0a 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -2190,24 +2190,27 @@ class HunyuanImage3Model(nn.Module): return torch.concat((q, k, v)) def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - if _is_moe(self.config): - # Params for weights, fp8 weight scales, fp8 activation scales - # (param_name, weight_name, expert_id, shard_id) - fused_moe_expert_mapping = FusedMoE.make_expert_params_mapping( - self, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, - num_redundant_experts=self.num_redundant_experts, - ) - expert_weights_remapping = { - "gate_proj": ("gate_and_up_proj", 1, 2), - "up_proj": ("gate_and_up_proj", 0, 2), - } - return fused_moe_expert_mapping, expert_weights_remapping - else: - return [], {} + if not _is_moe(self.config): + return [] + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + return FusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + num_redundant_experts=self.num_redundant_experts, + ) + + def _get_expert_weights_remapping(self) -> dict[str, tuple[str, int, int]]: + if not _is_moe(self.config): + return {} + return { + "gate_proj": ("gate_and_up_proj", 1, 2), + "up_proj": ("gate_and_up_proj", 0, 2), + } # rename for delay load def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): @@ -2243,7 +2246,8 @@ class HunyuanImage3Model(nn.Module): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - expert_params_mapping, expert_weights_remapping = self.get_expert_mapping() + expert_params_mapping = self.get_expert_mapping() + expert_weights_remapping = self._get_expert_weights_remapping() # List of unexpected keywords in weight names unexpected_keywords = [ diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index a8fac7fa6..ead3c6eb3 100755 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -128,25 +128,28 @@ class HunyuanModel(HunYuanModel): v = v.reshape(-1, hidden_size) return torch.concat((q, k, v)) - def get_expert_mapping(self) -> tuple[list[tuple[str, str, int, str]], dict[str, tuple[str, int, int]]]: - if _is_moe(self.config): - # Params for weights, fp8 weight scales, fp8 activation scales - # (param_name, weight_name, expert_id, shard_id) - fused_moe_expert_mapping = fused_moe_make_expert_params_mapping( - model=self, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, - num_redundant_experts=self.num_redundant_experts, - ) - expert_weights_remapping = { - "gate_proj": ("gate_and_up_proj", 1, 2), - "up_proj": ("gate_and_up_proj", 0, 2), - } - return fused_moe_expert_mapping, expert_weights_remapping - else: - return [], {} + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + if not _is_moe(self.config): + return [] + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + return fused_moe_make_expert_params_mapping( + model=self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + num_redundant_experts=self.num_redundant_experts, + ) + + def _get_expert_weights_remapping(self) -> dict[str, tuple[str, int, int]]: + if not _is_moe(self.config): + return {} + return { + "gate_proj": ("gate_and_up_proj", 1, 2), + "up_proj": ("gate_and_up_proj", 0, 2), + } def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): cla_factor = _get_cla_factor(self.config) @@ -174,7 +177,8 @@ class HunyuanModel(HunYuanModel): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - expert_params_mapping, expert_weights_remapping = self.get_expert_mapping() + expert_params_mapping = self.get_expert_mapping() + expert_weights_remapping = self._get_expert_weights_remapping() # List of unexpected keywords in weight names unexpected_keywords = [