LatentMoE: K3 sigmoid routing and Switch aux/z-loss

Route with σ(W_r x), Top-k(s+b), then L1-normalize over the selected set.
Add Switch/GShard aux and router z-loss into train_k3 and train_sft.
Wiki parquet URLs honor HF_ENDPOINT for mirrored downloads.
This commit is contained in:
dela
2026-08-25 19:50:07 +08:00
parent d1da0816f2
commit 7a12f61de1
8 changed files with 381 additions and 82 deletions
+2
View File
@@ -53,6 +53,8 @@ class K3Config:
moe_d_ff: int = 96
situ_beta1: float = 4.0
situ_beta2: float = 25.0
moe_aux_loss_coef: float = 1e-2 # Switch/GShard N Σ f_e P_e
moe_z_loss_coef: float = 1e-3 # mean (logsumexp logits)^2
kda_backend: str = "reference"