feat: add moe auxloss and metrics

2026-08-05 18:12:28 +08:00
parent 602b5ce216
commit 9b7e6c205f
9 changed files with 583 additions and 3 deletions
@@ -97,6 +97,7 @@ class AutoRegressiveLMConfig(BaseModelConfig):
    norm_topk_prob: bool = True
    decoder_sparse_step: int = 1
    mlp_only_layers: Optional[list[int]] = None
+    moe_aux_loss_coef: float = 0.01

    @field_validator("attn_type")
    def _validate_attn_type(cls, v: str) -> str: