Initial K3 snapshot: 0.5B KDA/MLA/MoE train path

Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
2026-08-25 14:43:17 +08:00
commit 584f7e9e73
140 changed files with 21592 additions and 0 deletions
@@ -0,0 +1,158 @@
+% 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
+
+\section{SiTU-GLU 与 Stable LatentMoE}
+\splabel{C5}
+
+\subsection{SiTU-GLU：带软上限的激活}
+
+SwiGLU 在低精度（fp16/bf16）训练时可能溢出：$\mathrm{silu}(x) \cdot x$ 没有上限。
+SiTU-GLU 用 $\tanh$ 给门控和上投影加软上限：
+
+\[
+  \mathrm{SiTU}(x) = W_o \big[\underbrace{\beta_1 \tanh\!\left(\frac{W_g x}{\beta_1}\right) \cdot \sigma(W_g x)}_{\text{gate}} \;\cdot\; \underbrace{\beta_2 \tanh\!\left(\frac{W_u x}{\beta_2}\right)}_{\text{up}}\big]
+\]
+
+\begin{center}
+\begin{tabular}{lp{8cm}}
+\toprule
+性质 & 说明 \\
+\midrule
+输出上限 & $\|\mathrm{SiTU}\|_\infty \leq \beta_1 \cdot \beta_2 = 4 \times 25 = 100$ \\
+原点附近 & $\tanh(x/\beta) \approx x/\beta$，所以 $\beta \cdot \tanh(x/\beta) \approx x$，退化为 SwiGLU \\
+远端 & 软饱和，防 fp16 溢出 \\
+\bottomrule
+\end{tabular}
+\end{center}
+
+\begin{codemathtop}{layers/latent\_moe.py — SiTU}
+\begin{lstlisting}
+class SiTU(nn.Module):
+    def __init__(self, dim_in, dim_ff, beta1=4.0, beta2=25.0):
+        self.w_g = nn.Linear(dim_in, dim_ff, bias=False)
+        self.w_u = nn.Linear(dim_in, dim_ff, bias=False)
+        self.w_o = nn.Linear(dim_ff, dim_in, bias=False)
+
+    def forward(self, x):                    # [*, dim_in]
+        wg = self.w_g(x)
+        g = self.beta1 * tanh(wg / self.beta1) * sigmoid(wg)   # gate
+        u = self.beta2 * tanh(self.w_u(x) / self.beta2)        # up
+        return self.w_o(g * u)               # [*, dim_in]
+\end{lstlisting}
+\end{codemathtop}
+
+\subsection{LatentMoE 架构}
+
+\begin{center}
+\begin{tabular}{rl}
+\toprule
+组件 & 说明 \\
+\midrule
+\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 归一化 \\
+\bottomrule
+\end{tabular}
+\end{center}
+
+\subsection{计算流（五步）}
+
+\begin{enumerate}[leftmargin=2em]
+\item \textbf{Latent 压缩}：
+\[
+  z = W_\downarrow \cdot x \qquad \shape{B, T, \ell}
+\]
+
+\item \textbf{Routing}：
+\[
+  \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}
+\]
+
+\item \textbf{Routed 专家}（在 latent 空间 $\ell$ 上）：
+\[
+  u = \sum_{i \in \mathrm{Top\text{-}k}} p_i \cdot E_i^{\mathrm{rt}}(z)
+  \qquad \shape{B, T, \ell}
+\]
+
+\item \textbf{Shared 专家}（全宽 $d$）：
+\[
+  s = \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}
+\]
+\end{enumerate}
+
+\begin{importantbox}{如果你只记一件事}
+Routed 专家只在 $\ell = d/2$ 的 latent 空间操作，
+参数量是全宽专家的 $1/4$（$\ell^2$ vs $d^2$）。
+Shared 专家保持全宽 $d$，提供基础表达能力。
+\end{importantbox}
+
+\subsection{代码对照}
+
+\begin{codemathtop}{layers/latent\_moe.py — LatentMoE.forward}
+\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]
+\end{lstlisting}
+\end{codemathtop}
+
+\begin{warningbox}{为什么 router 用 $x$（全宽）而不是 $z$（latent）？}
+路由需要看到 token 的完整表示才能做好选择。
+如果用 $z$ 路由，压缩过程可能丢失路由需要的信息。
+K3 论文里也是用全宽 $x$ 做 Top-k，然后在 $\ell$ 空间计算。
+\end{warningbox}
+
+\subsection{形状总览}
+
+\begin{center}
+\begin{tabular}{llll}
+\toprule
+变量 & 形状 & 说明 \\
+\midrule
+$x$ & \shape{B, T, d} & 输入 \\
+$z$ & \shape{B, T, \ell} & latent（$\ell = d/2$）\\
+logits & \shape{B, T, n_r} & router 输出 \\
+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 专家输出 \\
+$u$ & \shape{B, T, \ell} & 加权求和后的 routed 输出 \\
+\texttt{shared\_out} & \shape{B, T, d} & shared 专家求和 \\
+$y$ & \shape{B, T, d} & 最终输出 \\
+\bottomrule
+\end{tabular}
+\end{center}
+
+\subsection{本章小结}
+
+LatentMoE 把 routed 专家限制在 $\ell = d/2$ 的 latent 空间，省参数。
+SiTU-GLU 给 gate 和 up 加 $\tanh$ 软上限（$\beta_1=4, \beta_2=25$），
+防止低精度溢出。Shared 专家全宽，提供基础能力；routed 专家通过 Top-k 路由提供专业化能力。