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.

kda/layers/latent_moe.py

@@ -9,12 +9,14 @@ SiTU-GLU:  gate = β1·tanh(W_g x/β1)⊙σ(W_g x);  up = β2·tanh(W_u x/β2)
     ||SiTU-GLU||_∞ ≤ β1·β2 (=100), 原点附近≈SwiGLU, 远端软饱和防低精度溢出.
     E: R^in → R^in (内部中间维 d_ff).
 
-Router: Top-k logits 基于全宽 x (笔记 Topk(x)); 归一化权重取 softmax(topk).
+Router (K3 eq.13): s=σ(W_r x), Top-k(s+b), p_i = s_i / Σ_{j∈T} s_j.
+Routed 执行: permute-dispatch, pad 到 [R, C, ℓ], 三次 bmm（SiTU 参数结构不变）.
+训练: Switch/GShard aux + router z-loss, 由 train loop 加到 CE 上.
 """
+
 from __future__ import annotations
 
 import torch
-import torch.nn.functional as F
 from torch import nn
 
 from .rmsnorm import RMSNorm
@@ -23,7 +25,9 @@ from .rmsnorm import RMSNorm
 class SiTU(nn.Module):
     """SiTU-GLU expert: gate 支软上限 β1, up 支软上限 β2, 输出回到输入维."""
 
-    def __init__(self, dim_in: int, dim_ff: int, beta1: float = 4.0, beta2: float = 25.0):
+    def __init__(
+        self, dim_in: int, dim_ff: int, beta1: float = 4.0, beta2: float = 25.0
+    ):
         super().__init__()
         self.beta1, self.beta2 = beta1, beta2
         self.w_g = nn.Linear(dim_in, dim_ff, bias=False)
@@ -48,14 +52,19 @@ class LatentMoE(nn.Module):
         d_ff: int,
         beta1: float = 4.0,
         beta2: float = 25.0,
+        aux_loss_coef: float = 1e-2,
+        z_loss_coef: float = 1e-3,
     ):
         super().__init__()
         self.latent_size = latent_size
         self.n_routed = n_routed
         self.top_k = top_k
+        self.aux_loss_coef = aux_loss_coef
+        self.z_loss_coef = z_loss_coef
 
-        self.down = nn.Linear(hidden_size, latent_size, bias=False)   # W↓
-        self.router = nn.Linear(hidden_size, n_routed, bias=False)    # Top-k logits
+        self.down = nn.Linear(hidden_size, latent_size, bias=False)  # W↓
+        self.router = nn.Linear(hidden_size, n_routed, bias=False)  # logits → sigmoid
+        self.register_buffer("expert_bias", torch.zeros(n_routed), persistent=False)
         self.shared = nn.ModuleList(
             [SiTU(hidden_size, d_ff, beta1, beta2) for _ in range(n_shared)]
         )
@@ -63,8 +72,11 @@ class LatentMoE(nn.Module):
             [SiTU(latent_size, d_ff, beta1, beta2) for _ in range(n_routed)]
         )
         self.norm = RMSNorm(latent_size)
-        self.up = nn.Linear(latent_size, hidden_size, bias=False)     # W↑
+        self.up = nn.Linear(latent_size, hidden_size, bias=False)  # W↑
         self.last_route_ids: torch.Tensor | None = None
+        self.last_capacity: int = 0
+        self.last_aux_loss: torch.Tensor | None = None
+        self.last_z_loss: torch.Tensor | None = None
 
     @classmethod
     def from_config(cls, config) -> LatentMoE:
@@ -77,28 +89,102 @@ class LatentMoE(nn.Module):
             config.moe_d_ff,
             config.situ_beta1,
             config.situ_beta2,
+            float(getattr(config, "moe_aux_loss_coef", 1e-2)),
+            float(getattr(config, "moe_z_loss_coef", 1e-3)),
         )
 
+    def _routed_u(
+        self, z: torch.Tensor, ids: torch.Tensor, probs: torch.Tensor
+    ) -> torch.Tensor:
+        """Permute-dispatch + pad to [R, C, ℓ] + 3 bmm + scatter-add.
+
+        ``C = max(counts)``: FLOPs are ``R·C``, not ``sum(counts)``. Padding
+        slots are not gathered, so they contribute zero gradient. Empty
+        experts stay in the stacked weights (padded grouped GEMM).
+        """
+        B, T, ell = z.shape
+        N = B * T
+        R, k = self.n_routed, self.top_k
+        device = z.device
+
+        tok = torch.arange(N, device=device).unsqueeze(1).expand(N, k).reshape(-1)
+        eid = ids.reshape(-1)
+        pw = probs.reshape(-1)
+
+        order = eid.argsort(stable=True)
+        tok, eid, pw = tok[order], eid[order], pw[order]
+
+        counts = torch.bincount(eid, minlength=R)
+        offsets = counts.cumsum(0) - counts
+        local_pos = torch.arange(N * k, device=device) - offsets[eid]
+        C = int(counts.max().item()) if eid.numel() else 0
+        self.last_capacity = C
+
+        u_flat = z.new_zeros(N, ell)
+        if C == 0 or eid.numel() == 0:
+            return u_flat.view(B, T, ell)
+
+        gathered = z.reshape(N, ell)[tok]
+        padded = torch.index_put(
+            gathered.new_zeros(R, C, ell), (eid, local_pos), gathered
+        )
+
+        w_g = torch.stack([e.w_g.weight for e in self.experts])  # [R, ff, ℓ]
+        w_u = torch.stack([e.w_u.weight for e in self.experts])
+        w_o = torch.stack([e.w_o.weight for e in self.experts])  # [R, ℓ, ff]
+        beta1 = self.experts[0].beta1
+        beta2 = self.experts[0].beta2
+
+        wg = torch.bmm(padded, w_g.transpose(-1, -2))  # [R, C, ff]
+        g = beta1 * torch.tanh(wg / beta1) * torch.sigmoid(wg)
+        wu = torch.bmm(padded, w_u.transpose(-1, -2))
+        hidden = beta2 * torch.tanh(wu / beta2)
+        out = torch.bmm(g * hidden, w_o.transpose(-1, -2))  # [R, C, ℓ]
+
+        weighted = pw.unsqueeze(-1) * out[eid, local_pos]
+        u_flat = u_flat.index_add(0, tok, weighted)
+        return u_flat.view(B, T, ell)
+
+    def _route(self, logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """K3: s=σ(l), T=TopK(s+b), p_i = s_i / Σ_{j∈T} s_j. Bias does not enter p."""
+        scores = torch.sigmoid(logits)
+        ids = torch.topk(scores + self.expert_bias, self.top_k, dim=-1).indices
+        selected = scores.gather(-1, ids)
+        probs = selected / selected.sum(dim=-1, keepdim=True).clamp_min(1e-9)
+        return ids, probs
+
+    def _balancing_losses(
+        self, logits: torch.Tensor, ids: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Switch/GShard aux on sigmoid scores; z-loss on raw logits."""
+        routed = self.n_routed
+        flat_logits = logits.reshape(-1, routed).float()
+        scores = torch.sigmoid(flat_logits)
+        counts = torch.bincount(ids.reshape(-1), minlength=routed).to(
+            dtype=scores.dtype
+        )
+        frac = counts / counts.sum().clamp_min(1.0)
+        prob_mean = scores.mean(dim=0)
+        aux = routed * (frac * prob_mean).sum()
+        z_loss = torch.logsumexp(flat_logits, dim=-1).square().mean()
+        return self.aux_loss_coef * aux, self.z_loss_coef * z_loss
+
     def forward(self, x: torch.Tensor):
         B, T, _ = x.shape
-        z = self.down(x)                                    # [B, T, ℓ]
+        z = self.down(x)  # [B, T, ℓ]
 
-        logits = self.router(x)                             # [B, T, n_routed]
-        topk = torch.topk(logits, self.top_k, dim=-1)
-        ids = topk.indices                                  # [B, T, k]
+        logits = self.router(x)  # [B, T, n_routed]
+        ids, probs = self._route(logits)
         self.last_route_ids = ids.detach()
-        probs = F.softmax(topk.values, dim=-1)              # [B, T, k]
+        if self.training and (self.aux_loss_coef != 0.0 or self.z_loss_coef != 0.0):
+            self.last_aux_loss, self.last_z_loss = self._balancing_losses(logits, ids)
+        else:
+            zero = logits.new_zeros(())
+            self.last_aux_loss = zero
+            self.last_z_loss = zero
 
-        # 向量化 routed: 预计算全部专家输出, 按 token 的 Top-k id 取
-        all_out = torch.stack([e(z) for e in self.experts])     # [R, B, T, ℓ]
-        all_out = all_out.permute(1, 2, 0, 3).reshape(B * T, self.n_routed, self.latent_size)
-        u = torch.zeros(B, T, self.latent_size, device=x.device, dtype=x.dtype)
-        for i in range(self.top_k):
-            idx = ids[:, :, i].reshape(B * T)               # [B*T]
-            sel = all_out[torch.arange(B * T, device=x.device), idx]  # [B*T, ℓ]
-            u += probs[:, :, i : i + 1] * sel.reshape(B, T, self.latent_size)
-
-        shared_out = torch.stack([e(x) for e in self.shared]).sum(0)   # [B, T, d]
+        u = self._routed_u(z, ids, probs)
+        shared_out = torch.stack([e(x) for e in self.shared]).sum(0)  # [B, T, d]
         return shared_out + self.up(self.norm(u))
 
 
@@ -116,3 +202,21 @@ def moe_route_frac(model: nn.Module) -> torch.Tensor | None:
         return None
     stacked = torch.stack(hists).sum(0)
     return stacked / stacked.sum().clamp_min(1.0)
+
+
+def moe_router_losses(model: nn.Module) -> tuple[torch.Tensor, torch.Tensor]:
+    """Sum Switch aux and z-loss over every LatentMoE layer (0 if none)."""
+    auxes: list[torch.Tensor] = []
+    z_losses: list[torch.Tensor] = []
+    for module in model.modules():
+        if not isinstance(module, LatentMoE):
+            continue
+        if module.last_aux_loss is None or module.last_z_loss is None:
+            continue
+        auxes.append(module.last_aux_loss)
+        z_losses.append(module.last_z_loss)
+    if not auxes:
+        param = next(model.parameters(), None)
+        zero = param.new_zeros(()) if param is not None else torch.zeros(())
+        return zero, zero
+    return torch.stack(auxes).sum(), torch.stack(z_losses).sum()

kda/models/k3_config.py

@@ -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"
 
@@ -73,7 +75,7 @@ class K3Config:
         if name == "toy":
             return cls()
         if name in {"0.5b", "500m"}:
-            # H * head_dim == hidden. Routed 16: LatentMoE still runs every expert.
+            # H * head_dim == hidden. Routed 16 Top-2; LatentMoE padded bmm.
             # ~482M with tied Qwen3 embeddings. 6×(3 KDA + 1 MLA).
             return cls(
                 hidden_size=768,

kda/training/data.py

@@ -19,12 +19,15 @@ import torch
 from .prompts import instruction_prompt
 
 WIKI_SHARD_TOTAL = {"zh": 6, "en": 41}
-WIKI_BASE = (
-    "https://huggingface.co/datasets/wikimedia/wikipedia/resolve/main/20231101.{lang}"
-)
+WIKI_PATH = "datasets/wikimedia/wikipedia/resolve/main/20231101.{lang}"
 IGNORE_INDEX = -100
 
 
+def _hf_endpoint() -> str:
+    """Hub origin. OpenBayes/CN: export HF_ENDPOINT=https://hf-mirror.com"""
+    return os.environ.get("HF_ENDPOINT", "https://huggingface.co").rstrip("/")
+
+
 class Tokenizer(Protocol):
     vocab_size: int
 
@@ -91,7 +94,7 @@ def _wiki_files(lang: str, n_shards: int) -> list[str]:
         raise ValueError(f"unsupported wiki lang {lang!r}; expected zh or en")
     total = WIKI_SHARD_TOTAL[lang]
     n = min(max(n_shards, 1), total)
-    base = WIKI_BASE.format(lang=lang)
+    base = f"{_hf_endpoint()}/{WIKI_PATH.format(lang=lang)}"
     return [f"{base}/train-{i:05d}-of-{total:05d}.parquet" for i in range(n)]
 
 

notes/sections/sec-07.tex

@@ -1,8 +1,9 @@
 % teach:
-% gap: 读者知道 MoE 的 top-k 路由但不知道 LatentMoE 的 latent 接口和 SiTU-GLU
-% takeaway: LatentMoE 通过 latent 接口把 routed 专家限制在 ℓ=d/2 上算, SiTU-GLU 用软上限防溢出
-% jump: 为什么 routed 专家在 latent 空间而 shared 在全宽？省参数
-% omit: load balancing loss
+% gap: 读者知道 MoE 的 top-k 路由但不知道 LatentMoE 的 latent 接口、稀疏 dispatch 和负载均衡损失
+% takeaway: LatentMoE = shared 全宽 + routed 半宽 latent；路由用 K3 sigmoid-TopK + L1 归一化；
+%           执行用 permute-dispatch + padded bmm（每个 token 只算 k 个专家）；训练加 Switch/GShard aux + z-loss
+% jump: 为什么不用 dense stack 算全部专家？稀疏 dispatch 让算力只随 k 不随 R 涨
+% omit: none
 
 \section{SiTU-GLU 与 Stable LatentMoE}
 \splabel{C5}
@@ -54,7 +55,7 @@ class SiTU(nn.Module):
 \textbf{Shared 专家} & $n_{\mathrm{shared}}$ 个 SiTU，全宽 $d \to d$，所有 token 都经过 \\
 \textbf{Routed 专家} & $n_{\mathrm{routed}}$ 个 SiTU，半宽 $\ell \to \ell$（$\ell = d/2$） \\
 \textbf{Latent 接口} & $W_\downarrow: d \to \ell$, $W_\uparrow: \ell \to d$（压缩/还原） \\
-\textbf{Router} & $W_r: d \to n_{\mathrm{routed}}$，Top-k 选择 + softmax 归一化 \\
+\textbf{Router} & $W_r: d \to n_{\mathrm{routed}}$，K3：$s=\sigma(W_r x)$，Top-$k(s+b)$，$p_i=s_i/\sum_{j\in T}s_j$ \\
 \bottomrule
 \end{tabular}
 \end{center}
@@ -67,29 +68,31 @@ class SiTU(nn.Module):
   z = W_\downarrow \cdot x \qquad \shape{B, T, \ell}
 \]
 
-\item \textbf{Routing}：
+\item \textbf{Routing}（K3 eq.13）：
 \[
-  \mathrm{logits} = W_r \cdot x \qquad \shape{B, T, n_{\mathrm{routed}}}
-\]
-\[
-  \mathrm{ids}, \mathrm{probs} = \mathrm{TopK}(\mathrm{logits}, k)
-  \qquad \mathrm{ids}: \shape{B, T, k}, \;\; \mathrm{probs}: \shape{B, T, k}
+  s = \sigma(W_r \cdot x),\quad
+  T = \mathrm{TopK}(s+b, k),\quad
+  p_i = \frac{s_i}{\sum_{j\in T} s_j}
 \]
 
-\item \textbf{Routed 专家}（在 latent 空间 $\ell$ 上）：
+\item \textbf{Routed 专家}（在 latent 空间 $\ell$ 上，稀疏执行，见 \S7.4）：
 \[
   u = \sum_{i \in \mathrm{Top\text{-}k}} p_i \cdot E_i^{\mathrm{rt}}(z)
   \qquad \shape{B, T, \ell}
 \]
 
+  数学上是"对选中的 $k$ 个专家加权求和"，但\textbf{实现上不是}用
+  \texttt{stack([e(z) for e in experts])} 把所有专家都算一遍——
+  而是每个 token 只被送进它选中的 $k$ 个专家（permute-dispatch + padded bmm）。
+
 \item \textbf{Shared 专家}（全宽 $d$）：
 \[
-  s = \sum_j E_j^{\mathrm{sh}}(x) \qquad \shape{B, T, d}
+  s_{\mathrm{sh}} = \sum_j E_j^{\mathrm{sh}}(x) \qquad \shape{B, T, d}
 \]
 
 \item \textbf{合并}：
 \[
-  y = s + W_\uparrow \cdot \mathrm{RMSNorm}(u) \qquad \shape{B, T, d}
+  y = s_{\mathrm{sh}} + W_\uparrow \cdot \mathrm{RMSNorm}(u) \qquad \shape{B, T, d}
 \]
 \end{enumerate}
 
@@ -99,38 +102,132 @@ Routed 专家只在 $\ell = d/2$ 的 latent 空间操作，
 Shared 专家保持全宽 $d$，提供基础表达能力。
 \end{importantbox}
 
-\subsection{代码对照}
+\subsection{代码对照：路由}
 
-\begin{codemathtop}{layers/latent\_moe.py — LatentMoE.forward}
+\begin{codemathtop}{layers/latent\_moe.py — \_route（K3 eq.13）}
 \begin{lstlisting}
-def forward(self, x):                       # [B, T, d]
-    z = self.down(x)                        # [B, T, ell]
-
-    logits = self.router(x)                 # [B, T, n_routed]
-    topk = torch.topk(logits, self.top_k, dim=-1)
-    ids = topk.indices                      # [B, T, k]
-    probs = F.softmax(topk.values, dim=-1)  # [B, T, k]
-
-    # All expert outputs (vectorized)
-    all_out = stack([e(z) for e in self.experts])  # [R, B, T, ell]
-    # Gather top-k and weighted sum
-    u = zeros(B, T, ell)
-    for i in range(self.top_k):
-        idx = ids[:,:,i].reshape(B*T)
-        sel = all_out[arange, idx]
-        u += probs[:,:,i:i+1] * sel.reshape(B, T, ell)
-
-    shared_out = stack([e(x) for e in self.shared]).sum(0)  # [B, T, d]
-    return shared_out + self.up(self.norm(u))                # [B, T, d]
+def _route(self, logits):                       # logits: [B, T, n_routed]
+    scores = sigmoid(logits)                    # s = σ(W_r x)
+    ids = topk(scores + self.expert_bias, k).indices   # T = TopK(s+b)
+    selected = scores.gather(-1, ids)
+    probs = selected / selected.sum(-1).clamp_min(1e-9) # p_i = s_i / Σ_{j∈T} s_j
+    return ids, probs
 \end{lstlisting}
 \end{codemathtop}
 
+注意三点：
+
+\begin{enumerate}[leftmargin=2em]
+\item \textbf{sigmoid 代替 softmax}：K3 的 router 对每个专家独立打分
+      $s_i = \sigma(w_i \cdot x)$，不再是 softmax 归一化。这样"专家之间"不互相竞争
+      归一化预算，便于用 bias 做负载调节。
+\item \textbf{bias 只进选择、不进权重}：$\texttt{expert\_bias}$ 是
+      \texttt{register\_buffer(..., persistent=False)} 的非持久 buffer（初始化为 $0$，
+      不进 \texttt{state\_dict}），只参与 $\mathrm{TopK}(s+b)$ 的选择，
+      归一化 $p_i$ 仍用原始 $s_i$。
+\item \textbf{L1 归一化}：$p_i = s_i / \sum_{j \in T} s_j$（在选中的 $k$ 个上做），
+      权重之和为 1，等价于"选中的 sigmoid 分数重新归一化"。
+\end{enumerate}
+
+\subsection{稀疏执行：permute-dispatch + padded bmm}\label{sec:sparse-dispatch}
+
+\begin{codemathtop}{layers/latent\_moe.py — \_routed\_u}
+\begin{lstlisting}
+def _routed_u(self, z, ids, probs):             # z: [B,T,ell]  ids/probs: [B,T,k]
+    N = B * T;  R, k = self.n_routed, self.top_k
+    tok = arange(N).unsqueeze(1).expand(N, k).reshape(-1)   # 每个 token 重复 k 次
+    eid, pw = ids.reshape(-1), probs.reshape(-1)
+
+    order = eid.argsort(stable=True)            # 按专家 id 排序 → 同专家连续
+    tok, eid, pw = tok[order], eid[order], pw[order]
+
+    counts  = bincount(eid, minlength=R)        # 每个专家的 token 数
+    offsets = counts.cumsum(0) - counts
+    local_pos = arange(N*k) - offsets[eid]      # 专家内局部位置
+    C = int(counts.max().item())                # capacity = 最大负载
+
+    gathered = z.reshape(N, ell)[tok]
+    padded = index_put(zeros(R, C, ell), (eid, local_pos), gathered)  # [R, C, ell]
+
+    w_g = stack([e.w_g.weight for e in self.experts])  # [R, ff, ell]
+    w_u = stack([e.w_u.weight for e in self.experts])
+    w_o = stack([e.w_o.weight for e in self.experts])  # [R, ell, ff]
+
+    wg = bmm(padded, w_g.transpose(-1,-2))      # [R, C, ff]  grouped GEMM
+    g  = beta1 * tanh(wg / beta1) * sigmoid(wg)
+    wu = bmm(padded, w_u.transpose(-1,-2))
+    h  = beta2 * tanh(wu / beta2)
+    out = bmm(g * h, w_o.transpose(-1,-2))      # [R, C, ell]
+
+    weighted = pw.unsqueeze(-1) * out[eid, local_pos]
+    return index_add(zeros(N, ell), 0, tok, weighted).view(B, T, ell)  # scatter-add
+\end{lstlisting}
+\end{codemathtop}
+
+三步走：\textbf{① permute-dispatch}（按专家排序 + pad 到 $[R, C, \ell]$）→
+\textbf{② padded bmm}（专家参数堆成 batch 维，三次 batched GEMM 一次算完 $R$ 个专家）→
+\textbf{③ scatter-add}（\texttt{index\_add} 把加权输出按 \texttt{tok} 累加回 $u$）。
+
+\begin{importantbox}{为什么不用 dense stack？}
+朴素写法 \texttt{stack([e(z) for e in experts])} 会让每个专家都算全部 $B\cdot T$ 个 token，
+FLOPs 是 $R \cdot N$，退化成 dense，失去 MoE 的加速。稀疏 dispatch 把每个 token 只送进
+它选中的 $k$ 个专家，FLOPs 是 $R \cdot C$（$C = \max_e \text{count}_e \approx k \cdot N / R$），
+当 $k \ll R$ 时远小于 $R \cdot N$。padding 槽位不被 \texttt{index\_add} 收集，
+贡献零梯度；空专家仍留在堆叠权重里（padded grouped GEMM）。
+\end{importantbox}
+
 \begin{warningbox}{为什么 router 用 $x$（全宽）而不是 $z$（latent）？}
 路由需要看到 token 的完整表示才能做好选择。
 如果用 $z$ 路由，压缩过程可能丢失路由需要的信息。
 K3 论文里也是用全宽 $x$ 做 Top-k，然后在 $\ell$ 空间计算。
 \end{warningbox}
 
+\subsection{负载均衡损失：Switch/GShard aux + z-loss}\label{sec:load-balancing}
+
+Top-k 路由容易"塌缩"到少数专家（router 学出永远选某几个专家），导致负载不均衡、
+专家利用率低。训练时加两个损失，由 train loop 加到 CE 上：
+
+\[
+  f_e = \frac{\#\{\text{路由到 } e\}}{N \cdot k},\qquad
+  P_e = \frac{1}{N}\sum_{t} \sigma(W_r x_t)_e
+\]
+\[
+  \mathcal{L}_{\mathrm{aux}} = n_{\mathrm{routed}} \sum_e f_e \cdot P_e,\qquad
+  \mathcal{L}_{z} = \frac{1}{N}\sum_t \left(\log\!\sum_j e^{l_{tj}}\right)^2
+\]
+
+\begin{center}
+\begin{tabular}{ll}
+\toprule
+项 & 作用 \\
+\midrule
+$\mathcal{L}_{\mathrm{aux}}$ & Switch/GShard 风格：$f_e$ 是专家 $e$ 被路由到的 token 占比，$P_e$ 是它的平均 sigmoid 分数。塌缩时 $f$ 集中到单个专家、损失变大，逼着路由均匀化 \\
+$\mathcal{L}_{z}$ & z-loss：对 raw logits 的 logsumexp 求平方，压住 logits 幅度、防 router 分数爆炸 \\
+\bottomrule
+\end{tabular}
+\end{center}
+
+\begin{codemathtop}{layers/latent\_moe.py — \_balancing\_losses}
+\begin{lstlisting}
+def _balancing_losses(self, logits, ids):
+    flat = logits.reshape(-1, self.n_routed).float()
+    scores = sigmoid(flat)
+    counts = bincount(ids.reshape(-1), minlength=self.n_routed).float()
+    frac = counts / counts.sum().clamp_min(1.0)     # f_e
+    prob_mean = scores.mean(dim=0)                   # P_e
+    aux = self.n_routed * (frac * prob_mean).sum()   # N Σ f_e P_e
+    z_loss = logsumexp(flat, dim=-1).square().mean() # mean (logsumexp)^2
+    return self.aux_loss_coef * aux, self.z_loss_coef * z_loss
+\end{lstlisting}
+\end{codemathtop}
+
+两个损失只在 \texttt{self.training} 且系数非零时计算；系数默认
+$\alpha_{\mathrm{aux}} = 10^{-2}$、$\alpha_z = 10^{-3}$（\texttt{K3Config.moe\_aux\_loss\_coef} /
+\texttt{moe\_z\_loss\_coef}，可用 \texttt{--moe-aux-coef} / \texttt{--moe-z-coef} 覆盖）。
+train loop 里 \texttt{moe\_router\_losses(model)} 把所有 LatentMoE 层的损失求和，
+\texttt{loss = task + aux + z\_loss} 一起反传。aux/z 只更新 router 参数，
+不碰专家权重（\texttt{ids} 已 \texttt{.detach()}）。
+
 \subsection{形状总览}
 
 \begin{center}
@@ -140,12 +237,17 @@ K3 论文里也是用全宽 $x$ 做 Top-k，然后在 $\ell$ 空间计算。
 \midrule
 $x$ & \shape{B, T, d} & 输入 \\
 $z$ & \shape{B, T, \ell} & latent（$\ell = d/2$）\\
-logits & \shape{B, T, n_r} & router 输出 \\
+logits & \shape{B, T, n_r} & router logits \\
+$s$ & \shape{B, T, n_r} & $\sigma(\mathrm{logits})$ \\
+$b$ & \shape{n_r} & expert bias（非持久，只进 TopK 不进 $p$）\\
 ids & \shape{B, T, k} & Top-k 专家索引 \\
-probs & \shape{B, T, k} & Top-k softmax 权重 \\
-\texttt{all\_out} & \shape{n_r, B, T, \ell} & 所有 routed 专家输出 \\
+probs & \shape{B, T, k} & K3 sigmoid-L1 权重 \\
+padded & \shape{n_r, C, \ell} & dispatch 后 pad 到容量 $C$ 的张量 \\
+$C$ & 标量 & 最大专家负载（pad 宽度）\\
 $u$ & \shape{B, T, \ell} & 加权求和后的 routed 输出 \\
 \texttt{shared\_out} & \shape{B, T, d} & shared 专家求和 \\
+$f_e$, $P_e$ & 标量 & aux loss 的负载占比 / 平均分数 \\
+$\mathcal{L}_{\mathrm{aux}}$, $\mathcal{L}_z$ & 标量 & 负载均衡 / z-loss \\
 $y$ & \shape{B, T, d} & 最终输出 \\
 \bottomrule
 \end{tabular}
@@ -154,5 +256,7 @@ $y$ & \shape{B, T, d} & 最终输出 \\
 \subsection{本章小结}
 
 LatentMoE 把 routed 专家限制在 $\ell = d/2$ 的 latent 空间，省参数。
-SiTU-GLU 给 gate 和 up 加 $\tanh$ 软上限（$\beta_1=4, \beta_2=25$），
-防止低精度溢出。Shared 专家全宽，提供基础能力；routed 专家通过 Top-k 路由提供专业化能力。
+SiTU-GLU 给 gate 和 up 加 $\tanh$ 软上限（$\beta_1=4, \beta_2=25$），防止低精度溢出。
+路由走 K3 eq.13（sigmoid-TopK + L1 归一化），执行用稀疏 permute-dispatch + padded bmm
+（每个 token 只算 $k$ 个专家，FLOPs $R \cdot C$ 而非 $R \cdot N$），训练加
+Switch/GShard aux loss + router z-loss 防路由塌缩。

notes/sections/sec-11.tex

@@ -30,6 +30,9 @@ $n_r$ & routed 专家数 & 16 \\
 $k$ & Top-$k$ & 2 \\
 $n_s$ & shared 专家数 & 2 \\
 $d_{\mathrm{ff}}$ & 专家中间维度 & 96 \\
+$C$ & MoE 专家容量（pad 宽度） & 动态 \\
+$\alpha_{\mathrm{aux}}$ & Switch/GShard aux 系数 & $10^{-2}$ \\
+$\alpha_z$ & router z-loss 系数 & $10^{-3}$ \\
 $N$ & AttnRes 原子层数 ($= 2L$) & 8 \\
 $S$ & AttnRes 块大小（原子层） & 2--24 \\
 \bottomrule
@@ -105,12 +108,19 @@ gate & \shape{B, T, H \cdot d_v} & $\sigma(W_g x)$ \\
 \midrule
 $x$ & \shape{B, T, D} & 输入 \\
 $z$ & \shape{B, T, \ell} & latent ($\ell = D/2$) \\
-logits & \shape{B, T, n_r} & router logits \\
+logits & \shape{B, T, n_r} & router logits $W_r x$ \\
+$s$ & \shape{B, T, n_r} & sigmoid 分数 $\sigma(\mathrm{logits})$ \\
+$b$ & \shape{n_r} & expert bias（非持久，只进 TopK） \\
 ids & \shape{B, T, k} & Top-$k$ 专家索引 \\
-probs & \shape{B, T, k} & softmax 权重 \\
+$p_i$ & \shape{B, T, k} & sigmoid-L1 权重 $s_i/\sum_{j\in T}s_j$ \\
+padded & \shape{n_r, C, \ell} & dispatch 后 pad 到容量 $C$ \\
+$C$ & 标量 & 最大专家负载（pad 宽度） \\
 $u$ & \shape{B, T, \ell} & routed 加权输出 \\
-$s$ & \shape{B, T, D} & shared 专家求和 \\
-$y$ & \shape{B, T, D} & $s + W_\uparrow \mathrm{RMSNorm}(u)$ \\
+$s_{\mathrm{sh}}$ & \shape{B, T, D} & shared 专家求和 \\
+$y$ & \shape{B, T, D} & $s_{\mathrm{sh}} + W_\uparrow \mathrm{RMSNorm}(u)$ \\
+$f_e, P_e$ & 标量 & aux loss 负载占比 / 平均分数 \\
+$\mathcal{L}_{\mathrm{aux}}, \mathcal{L}_z$ & 标量 & 负载均衡 / z-loss \\
+$\alpha_{\mathrm{aux}}, \alpha_z$ & 标量 & 对应系数（$10^{-2}$ / $10^{-3}$） \\
 \bottomrule
 \end{tabular}
 \end{center}

tests/integration/test_k3_arch.py

@@ -4,7 +4,7 @@ import torch.nn.functional as F
 import pytest
 
 from kda.layers.kda_attn import KDAAttention
-from kda.layers.latent_moe import LatentMoE
+from kda.layers.latent_moe import LatentMoE, moe_router_losses
 from kda.layers.mla import GatedMLA
 from kda.models.causal_lm import CausalLM
 from kda.models.k3_config import K3Config
@@ -88,30 +88,176 @@ def test_preset_0_5b_schedule():
     assert cfg.layer_specs()[3] == ("mla", "moe")
 
 
+def _dense_moe_forward(moe: LatentMoE, x: torch.Tensor) -> torch.Tensor:
+    """Dense path: run every routed expert, then gather K3 sigmoid-norm top-k."""
+    z = moe.down(x)
+    ids, probs = moe._route(moe.router(x))
+    all_out = torch.stack([expert(z) for expert in moe.experts])
+    B, T, _ = x.shape
+    all_out = all_out.permute(1, 2, 0, 3).reshape(B * T, moe.n_routed, moe.latent_size)
+    u = z.new_zeros(B, T, moe.latent_size)
+    for i in range(moe.top_k):
+        idx = ids[:, :, i].reshape(B * T)
+        sel = all_out[torch.arange(B * T, device=x.device), idx]
+        u = u + probs[:, :, i : i + 1] * sel.reshape(B, T, moe.latent_size)
+    shared = torch.stack([expert(x) for expert in moe.shared]).sum(0)
+    return shared + moe.up(moe.norm(u))
+
+
 def test_moe_router_activates_topk_only():
-    from kda.layers.latent_moe import LatentMoE
     torch.manual_seed(3)
     moe = LatentMoE(hidden_size=32, latent_size=16, n_routed=8, top_k=2, n_shared=1, d_ff=24)
     x = torch.randn(2, 6, 32)
     with torch.no_grad():
         y = moe(x)
-        logits = moe.router(x)
-        topk = torch.topk(logits, moe.top_k, dim=-1)
+        scores = torch.sigmoid(moe.router(x))
+        ids, probs = moe._route(moe.router(x))
         z = moe.down(x)
-        # 手算: 只有 top-k 专家输出被加权, 再经 shared + up(norm(u))
         expected_u = torch.zeros(2, 6, moe.latent_size)
         all_out = torch.stack([e(z) for e in moe.experts])   # [R,B,T,ℓ]
-        probs = F.softmax(topk.values, dim=-1)
         for i in range(moe.top_k):
-            idx = topk.indices[:, :, i]
+            idx = ids[:, :, i]
             for b in range(2):
                 for t in range(6):
                     expected_u[b, t] += probs[b, t, i] * all_out[idx[b, t], b, t]
         shared = torch.stack([e(x) for e in moe.shared]).sum(0)
         expected_y = shared + moe.up(moe.norm(expected_u))
+        selected = scores.gather(-1, ids)
+        torch.testing.assert_close(
+            probs, selected / selected.sum(-1, keepdim=True).clamp_min(1e-9)
+        )
     torch.testing.assert_close(y, expected_y, atol=1e-5, rtol=1e-5)
     assert moe.last_route_ids is not None
     assert moe.last_route_ids.shape[-1] == moe.top_k
+    counts = torch.bincount(moe.last_route_ids.reshape(-1), minlength=moe.n_routed)
+    assert moe.last_capacity == int(counts.max())
+
+
+def test_moe_sparse_matches_dense_fwd_bwd():
+    torch.manual_seed(3)
+    moe = LatentMoE(hidden_size=32, latent_size=16, n_routed=8, top_k=2, n_shared=1, d_ff=24)
+    x = torch.randn(2, 6, 32)
+    y_sparse = moe(x)
+    y_dense = _dense_moe_forward(moe, x)
+    torch.testing.assert_close(y_sparse, y_dense, atol=1e-5, rtol=1e-5)
+
+    moe.zero_grad(set_to_none=True)
+    xs = x.detach().requires_grad_(True)
+    moe(xs).square().mean().backward()
+    grads_s = {
+        name: param.grad.detach().clone()
+        for name, param in moe.named_parameters()
+        if param.grad is not None
+    }
+    dx_s = xs.grad.detach().clone()
+
+    moe.zero_grad(set_to_none=True)
+    xd = x.detach().requires_grad_(True)
+    _dense_moe_forward(moe, xd).square().mean().backward()
+    grads_d = {
+        name: param.grad.detach().clone()
+        for name, param in moe.named_parameters()
+        if param.grad is not None
+    }
+    dx_d = xd.grad.detach().clone()
+
+    torch.testing.assert_close(dx_s, dx_d, atol=1e-5, rtol=1e-5)
+    assert grads_s.keys() == grads_d.keys()
+    for name in grads_s:
+        torch.testing.assert_close(grads_s[name], grads_d[name], atol=1e-5, rtol=1e-5)
+
+
+def test_moe_unselected_experts_have_zero_grad():
+    torch.manual_seed(0)
+    moe = LatentMoE(hidden_size=32, latent_size=16, n_routed=8, top_k=2, n_shared=1, d_ff=24)
+    with torch.no_grad():
+        moe.router.weight.zero_()
+        moe.router.weight[0] = 1.0
+        moe.router.weight[1] = 0.5
+    x = torch.ones(2, 4, 32)
+    moe(x).square().mean().backward()
+    selected = set(moe.last_route_ids.reshape(-1).tolist())
+    assert selected == {0, 1}
+    for idx, expert in enumerate(moe.experts):
+        for param in (expert.w_g.weight, expert.w_u.weight, expert.w_o.weight):
+            assert param.grad is not None
+            if idx in selected:
+                assert param.grad.abs().sum() > 0
+            else:
+                assert torch.equal(param.grad, torch.zeros_like(param.grad))
+
+
+def test_moe_aux_loss_penalizes_collapse():
+    torch.manual_seed(0)
+    moe = LatentMoE(
+        hidden_size=32, latent_size=16, n_routed=8, top_k=2, n_shared=1, d_ff=24,
+        aux_loss_coef=1.0, z_loss_coef=1.0,
+    )
+    moe(torch.randn(4, 16, 32))
+    spread = float(moe.last_aux_loss.detach())
+    with torch.no_grad():
+        moe.router.weight.zero_()
+        moe.router.weight[0] = 1.0
+        moe.router.weight[1] = 0.5
+    moe(torch.ones(4, 16, 32))
+    collapsed = float(moe.last_aux_loss.detach())
+    assert collapsed > spread
+    assert collapsed > 1.2
+    assert float(moe.last_z_loss.detach()) > 0
+
+
+def test_moe_aux_loss_updates_router_only():
+    torch.manual_seed(1)
+    moe = LatentMoE(
+        hidden_size=32, latent_size=16, n_routed=8, top_k=2, n_shared=1, d_ff=24,
+        aux_loss_coef=1.0, z_loss_coef=1.0,
+    )
+    moe(torch.randn(2, 8, 32))
+    aux, z_loss = moe_router_losses(moe)
+    (aux + z_loss).backward()
+    assert moe.router.weight.grad is not None
+    assert moe.router.weight.grad.abs().sum() > 0
+    for expert in moe.experts:
+        assert expert.w_g.weight.grad is None
+        assert expert.w_u.weight.grad is None
+        assert expert.w_o.weight.grad is None
+
+
+def test_moe_router_losses_sums_layers():
+    cfg = K3Config(
+        hidden_size=32, num_hidden_layers=2, num_heads=4, head_dim=8,
+        chunk_size=4, vocab_size=64, moe_latent_size=16, moe_d_ff=16,
+        n_routed=4, kv_lora_rank=16, q_lora_rank=32, qk_nope_head_dim=8, v_head_dim=8,
+        moe_aux_loss_coef=1.0, moe_z_loss_coef=1.0,
+    )
+    model = CausalLM(cfg)
+    model(torch.randint(0, cfg.vocab_size, (2, 8)))
+    aux, z_loss = moe_router_losses(model)
+    layers = [module for module in model.modules() if isinstance(module, LatentMoE)]
+    assert len(layers) == 2
+    torch.testing.assert_close(aux, layers[0].last_aux_loss + layers[1].last_aux_loss)
+    torch.testing.assert_close(z_loss, layers[0].last_z_loss + layers[1].last_z_loss)
+
+
+def test_moe_aux_backward_with_checkpoint():
+    cfg = K3Config(
+        hidden_size=32, num_hidden_layers=2, num_heads=4, head_dim=8,
+        chunk_size=4, vocab_size=64, moe_latent_size=16, moe_d_ff=16,
+        n_routed=4, kv_lora_rank=16, q_lora_rank=32, qk_nope_head_dim=8, v_head_dim=8,
+        gradient_checkpointing=True, moe_aux_loss_coef=1.0, moe_z_loss_coef=1.0,
+    )
+    model = CausalLM(cfg)
+    model.train()
+    tokens = torch.randint(0, cfg.vocab_size, (2, 8))
+    task = model(tokens, labels=tokens)
+    aux, z_loss = moe_router_losses(model)
+    (task + aux + z_loss).backward()
+    router_grad = sum(
+        param.grad.abs().sum().item()
+        for name, param in model.named_parameters()
+        if "router" in name and param.grad is not None
+    )
+    assert router_grad > 0
 
 
 def test_k3_causal_future_does_not_change_past_logits():

train_k3.py

@@ -17,7 +17,7 @@ from dataclasses import asdict
 
 import torch
 
-from kda.layers.latent_moe import moe_route_frac
+from kda.layers.latent_moe import LatentMoE, moe_route_frac, moe_router_losses
 from kda.models.causal_lm import CausalLM
 from kda.models.k3_config import K3Config
 from kda.training.data import iter_indexed, load_pretrain_chunks, load_tokenizer
@@ -91,6 +91,8 @@ def _init_swanlab(cfg: K3Config, args: argparse.Namespace):
                 "langs": args.langs,
                 "kda_backend": cfg.kda_backend,
                 "gradient_checkpointing": cfg.gradient_checkpointing,
+                "moe_aux_loss_coef": cfg.moe_aux_loss_coef,
+                "moe_z_loss_coef": cfg.moe_z_loss_coef,
             },
         )
     except Exception as exc:
@@ -148,6 +150,13 @@ def _heldout_loss(
     return sum(losses) / max(len(losses), 1)
 
 
+def _apply_moe_coefs(model, cfg: K3Config) -> None:
+    for module in model.modules():
+        if isinstance(module, LatentMoE):
+            module.aux_loss_coef = cfg.moe_aux_loss_coef
+            module.z_loss_coef = cfg.moe_z_loss_coef
+
+
 def _moe_log(model) -> dict:
     frac = moe_route_frac(model)
     if frac is None:
@@ -225,6 +234,18 @@ def main() -> None:
         dest="grad_checkpoint",
         action="store_false",
     )
+    p.add_argument(
+        "--moe-aux-coef",
+        type=float,
+        default=None,
+        help="Switch/GShard aux loss weight (default 0.01; 0 disables)",
+    )
+    p.add_argument(
+        "--moe-z-coef",
+        type=float,
+        default=None,
+        help="router z-loss weight (default 0.001; 0 disables)",
+    )
     args = p.parse_args()
     if args.gen_prefix is None:
         args.gen_prefix = ["人工智能的发展", "The history of computing"]
@@ -252,6 +273,10 @@ def main() -> None:
     cfg.attnres_block_size = args.attnres_block_size
     if args.grad_checkpoint is not None:
         cfg.gradient_checkpointing = args.grad_checkpoint
+    if args.moe_aux_coef is not None:
+        cfg.moe_aux_loss_coef = args.moe_aux_coef
+    if args.moe_z_coef is not None:
+        cfg.moe_z_loss_coef = args.moe_z_coef
 
     langs = [part.strip() for part in args.langs.split(",") if part.strip()]
     tpm = tokens_per_micro(args.batch, args.seq_len)
@@ -283,6 +308,10 @@ def main() -> None:
         cfg.attnres_block_size = args.attnres_block_size
         if args.grad_checkpoint is not None:
             cfg.gradient_checkpointing = args.grad_checkpoint
+        if args.moe_aux_coef is not None:
+            cfg.moe_aux_loss_coef = args.moe_aux_coef
+        if args.moe_z_coef is not None:
+            cfg.moe_z_loss_coef = args.moe_z_coef
         model.gradient_checkpointing = cfg.gradient_checkpointing
         model.to(device)
         payload = torch.load(args.resume, map_location="cpu", weights_only=False)
@@ -296,6 +325,7 @@ def main() -> None:
     else:
         model = CausalLM(cfg).to(device)
 
+    _apply_moe_coefs(model, cfg)
     tracker = _init_swanlab(cfg, args)
     n = sum(p.numel() for p in model.parameters())
     print(
@@ -304,7 +334,8 @@ def main() -> None:
     )
     print(
         f"vocab={cfg.vocab_size}  tied={cfg.tie_word_embeddings}  "
-        f"layers={cfg.layer_types()}  attnres={cfg.attnres}  langs={langs}"
+        f"layers={cfg.layer_types()}  attnres={cfg.attnres}  langs={langs}  "
+        f"moe_aux={cfg.moe_aux_loss_coef:g}  moe_z={cfg.moe_z_loss_coef:g}"
     )
     if args.max_tokens is None:
         print(
@@ -364,7 +395,9 @@ def main() -> None:
         scale = lr_scale(opt_step, args.warmup, horizon)
         _set_lr(optim, args.lr * scale)
         with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_bf16):
-            loss = model(x, labels=y) / args.grad_acc
+            task = model(x, labels=y)
+            aux, z_loss = moe_router_losses(model)
+            loss = (task + aux + z_loss) / args.grad_acc
         loss.backward()
         do_step = (micro_step + 1) % args.grad_acc == 0
         grad_norm = None
@@ -374,14 +407,20 @@ def main() -> None:
             optim.zero_grad(set_to_none=True)
             opt_step += 1
 
-        raw_loss = loss.item() * args.grad_acc
+        raw_loss = float(task.detach())
         tokens += tpm
         micro_step += 1
         lr_now = optim.param_groups[0]["lr"]
         if raw_loss < best_train:
             best_train = raw_loss
 
-        metrics = {"train/loss": raw_loss, "train/lr": lr_now, "train/tokens": tokens}
+        metrics = {
+            "train/loss": raw_loss,
+            "train/lr": lr_now,
+            "train/tokens": tokens,
+            "moe/aux": float(aux.detach()),
+            "moe/z": float(z_loss.detach()),
+        }
         if grad_norm is not None:
             metrics["train/grad_norm"] = grad_norm
         elapsed = time.perf_counter() - t0
@@ -403,6 +442,7 @@ def main() -> None:
                 f"micro {micro_step:6d}  opt {opt_step:6d}  tok {tokens:,}  "
                 f"loss {raw_loss:.4f}  lr {lr_now:.2e}"
                 + (f"  held {held:.4f}" if held is not None else "")
+                + f"  aux {metrics['moe/aux']:.4f}  z {metrics['moe/z']:.4f}"
             )
             if micro_step % (args.eval_every * 2) == 0 or micro_step <= args.eval_every:
                 for prefix in args.gen_prefix:

train_sft.py

@@ -14,6 +14,7 @@ from dataclasses import asdict
 
 import torch
 
+from kda.layers.latent_moe import moe_router_losses
 from kda.training.data import (
     IGNORE_INDEX,
     iter_sft_batches,
@@ -131,23 +132,32 @@ def main() -> None:
         x, y = x.to(device), y.to(device)
         _set_lr(optim, args.lr * lr_scale(opt_step, args.warmup, horizon))
         with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_bf16):
-            loss = model(x, labels=y, ignore_index=IGNORE_INDEX) / args.grad_acc
+            task = model(x, labels=y, ignore_index=IGNORE_INDEX)
+            aux, z_loss = moe_router_losses(model)
+            loss = (task + aux + z_loss) / args.grad_acc
         loss.backward()
         if (step + 1) % args.grad_acc == 0:
             torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
             optim.step()
             optim.zero_grad(set_to_none=True)
             opt_step += 1
-        raw = loss.item() * args.grad_acc
+        raw = float(task.detach())
         if raw < best:
             best = raw
         if tracker is not None:
             tracker.log(
-                {"sft/loss": raw, "sft/lr": optim.param_groups[0]["lr"]},
+                {
+                    "sft/loss": raw,
+                    "sft/lr": optim.param_groups[0]["lr"],
+                    "moe/aux": float(aux.detach()),
+                    "moe/z": float(z_loss.detach()),
+                },
                 step=step,
             )
         if step % args.eval_every == 0 or step == max_micro - 1:
-            print(f"step {step:4d}  sft loss {raw:.4f}  lr {optim.param_groups[0]['lr']:.2e}")
+            print(
+                f"step {step:4d}  sft loss {raw:.4f}  lr {optim.param_groups[0]['lr']:.2e}"
+            )
             if args.src and args.ref:
                 model.eval()
                 srcs, refs = _read_lines(args.src), _read_lines(args.ref)