[misc] fix moe (#10334)

Co-authored-by: frozenleaves <frozen@Mac.local>
This commit is contained in:
浮梦
2026-03-31 23:04:45 +08:00
committed by GitHub
parent 97433c53b6
commit e5e8546493

View File

@@ -147,6 +147,11 @@ def add_z3_leaf_module(model: "PreTrainedModel") -> None:
_set_z3_leaf_modules(model, [Qwen3NextSparseMoeBlock])
if model_type == "qwen3_5_moe":
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeSparseMoeBlock
_set_z3_leaf_modules(model, [Qwen3_5MoeSparseMoeBlock])
def configure_moe(config: "PretrainedConfig", model_args: "ModelArguments", is_trainable: bool) -> None:
if not is_trainable or not model_args.moe_aux_loss_coef: