dela
584f7e9e73
Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
159 lines
5.5 KiB
TeX
159 lines
5.5 KiB
TeX
% 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 路由提供专业化能力。
|