[Bugfix] Support K2 Horizon MoE without MoVA (#37825)

Co-authored-by: BBuf <1182563586@qq.com>
This commit is contained in:
Yash Akhauri
2026-09-03 23:53:16 +08:00
committed by GitHub
co-authored by BBuf
parent bf71035d39
commit 392841f47c
+28 -10
View File
@@ -188,6 +188,17 @@ def _normalize_k2_horizon_config(config: PretrainedConfig) -> None:
f"got {mova_num_experts!r}." f"got {mova_num_experts!r}."
) )
is_mova = mova_num_experts > 0 is_mova = mova_num_experts > 0
num_experts = getattr(config, "num_experts", 0)
if (
isinstance(num_experts, bool)
or not isinstance(num_experts, int)
or num_experts < 0
):
raise ValueError(
"K2Horizon num_experts must be a non-negative integer, "
f"got {num_experts!r}."
)
has_moe_ffn = num_experts > 0
if is_mova: if is_mova:
if _get_xllm_source_router_gemm_partitions(config) is None: if _get_xllm_source_router_gemm_partitions(config) is None:
@@ -247,14 +258,21 @@ def _normalize_k2_horizon_config(config: PretrainedConfig) -> None:
target_name="num_values_per_tok", target_name="num_values_per_tok",
value=0, value=0,
) )
for field in ("num_experts", "num_experts_per_tok", "num_shared_experts"): if not has_moe_ffn:
value = getattr(config, field, 0) for field in (
if isinstance(value, bool) or not isinstance(value, int) or value != 0: "num_experts",
raise ValueError(f"Dense K2Horizon requires {field}=0, got {value!r}") "num_experts_per_tok",
# Some dense exports omit the MoE-only fields. Downstream model "num_shared_experts",
# construction reads them directly, so materialize the validated ):
# dense defaults instead of relying on getattr fallbacks forever. value = getattr(config, field, 0)
setattr(config, field, 0) if isinstance(value, bool) or not isinstance(value, int) or value != 0:
raise ValueError(
f"Dense K2Horizon requires {field}=0, got {value!r}"
)
# Some dense exports omit the MoE-only fields. Downstream model
# construction reads them directly, so materialize the validated
# dense defaults instead of relying on getattr fallbacks forever.
setattr(config, field, 0)
if getattr(config, "query_key_norm", False): if getattr(config, "query_key_norm", False):
raise ValueError( raise ValueError(
"Dense K2Horizon native loading does not support query/key " "Dense K2Horizon native loading does not support query/key "
@@ -489,7 +507,7 @@ def _normalize_k2_horizon_config(config: PretrainedConfig) -> None:
"K2Horizon MoVA requires mlp_only_layers to be a contiguous " "K2Horizon MoVA requires mlp_only_layers to be a contiguous "
f"prefix starting at zero, got {list(mlp_only_layers)}" f"prefix starting at zero, got {list(mlp_only_layers)}"
) )
if not is_mova and list(mlp_only_layers) != list( if not has_moe_ffn and list(mlp_only_layers) != list(
range(config.num_hidden_layers) range(config.num_hidden_layers)
): ):
raise ValueError( raise ValueError(
@@ -501,7 +519,7 @@ def _normalize_k2_horizon_config(config: PretrainedConfig) -> None:
target_name="num_dense_layers", target_name="num_dense_layers",
value=len(mlp_only_layers), value=len(mlp_only_layers),
) )
elif not is_mova: elif not has_moe_ffn:
raise ValueError( raise ValueError(
"Dense K2Horizon native loading requires explicit mlp_only_layers" "Dense K2Horizon native loading requires explicit mlp_only_layers"
) )