Laguna: config-driven MoE router scoring (#35362)
This commit is contained in:
@@ -114,6 +114,7 @@ class LagunaConfig(PretrainedConfig):
|
|||||||
moe_routed_scaling_factor: float = 1.0,
|
moe_routed_scaling_factor: float = 1.0,
|
||||||
moe_router_logit_softcapping: float = 0.0,
|
moe_router_logit_softcapping: float = 0.0,
|
||||||
moe_apply_router_weight_on_input: bool = False,
|
moe_apply_router_weight_on_input: bool = False,
|
||||||
|
moe_router_score_func: str = "sigmoid",
|
||||||
# Per-layer-type rope dict; nested under "full_attention" / "sliding_attention".
|
# Per-layer-type rope dict; nested under "full_attention" / "sliding_attention".
|
||||||
rope_parameters: Optional[Dict[str, Any]] = None,
|
rope_parameters: Optional[Dict[str, Any]] = None,
|
||||||
partial_rotary_factor: Optional[float] = None,
|
partial_rotary_factor: Optional[float] = None,
|
||||||
@@ -155,6 +156,8 @@ class LagunaConfig(PretrainedConfig):
|
|||||||
self.moe_routed_scaling_factor = moe_routed_scaling_factor
|
self.moe_routed_scaling_factor = moe_routed_scaling_factor
|
||||||
self.moe_router_logit_softcapping = moe_router_logit_softcapping
|
self.moe_router_logit_softcapping = moe_router_logit_softcapping
|
||||||
self.moe_apply_router_weight_on_input = moe_apply_router_weight_on_input
|
self.moe_apply_router_weight_on_input = moe_apply_router_weight_on_input
|
||||||
|
# Router scoring nonlinearity: "sigmoid" (default) or "sqrtsoftplus".
|
||||||
|
self.moe_router_score_func = moe_router_score_func
|
||||||
|
|
||||||
# Synthesise per-layer schedules when the caller omits them so the model
|
# Synthesise per-layer schedules when the caller omits them so the model
|
||||||
# file can index by layer_id without per-call guards.
|
# file can index by layer_id without per-call guards.
|
||||||
|
|||||||
@@ -172,7 +172,8 @@ class LagunaMoE(nn.Module):
|
|||||||
layer_id=layer_id,
|
layer_id=layer_id,
|
||||||
renormalize=True,
|
renormalize=True,
|
||||||
use_grouped_topk=False,
|
use_grouped_topk=False,
|
||||||
scoring_func="sigmoid",
|
# "sigmoid" (default) or "sqrtsoftplus"; the branch lives in topk.py.
|
||||||
|
scoring_func=config.moe_router_score_func,
|
||||||
correction_bias=self.gate.e_score_correction_bias,
|
correction_bias=self.gate.e_score_correction_bias,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user