Compare commits

...
Author SHA1 Message Date
hsliuustc0106 ade3bb27ea [Docs] Clarify Hunyuan expert remapping scope
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-03 23:02:46 +00:00
hsliuustc0106 cecaa6344d Fix Hunyuan Image-3 expert mapping contract
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-03 17:05:10 +00:00
3 changed files with 97 additions and 41 deletions

No files matched your search

@@ -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_separates_public_and_local_contracts(
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"),
[
@@ -2190,24 +2190,28 @@ 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]]:
"""Return the vLLM-Omni-local remapping for Hunyuan checkpoint weights."""
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 +2247,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 = [
@@ -128,25 +128,29 @@ 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]]:
"""Return the vLLM-Omni-local remapping for Hunyuan checkpoint weights."""
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 +178,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 = [