Gahow Wang
057a3c68a3
Architecture + exact HF reference math (router softmax-after-topk, interleaved clamped (up+1)*glu experts, attention sinks, alternating sliding window, head_dim 64, rope 150000), MXFP4 dequant plan, and the correctness-first -> PP -> llama.cpp roadmap. MOE_PROGRESS.md captures live state for resuming after a machine reboot (HF is firewalled here; download via proxy + hf-mirror; gpt-oss-20b not yet downloaded). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
94 lines
4.8 KiB
Markdown
94 lines
4.8 KiB
Markdown
# Phase 19: MoE — gpt-oss-20b
|
||
|
||
> 目标：在 xserv 支持 **MoE**，用 `openai/gpt-oss-20b` 端到端跑通，并与 llama.cpp 在
|
||
> AIME 2025 / GSM8K 上对比正确性与性能。MXFP4 expert 权重加载时反量化为 BF16；整模型
|
||
> ~40GB 单卡放不下 → 复用 Phase 18 的 **PP**（PP=2 ~20GB/卡，PP=4 ~10GB/卡）。
|
||
>
|
||
> 实时进度与重启续作指南见 `docs/MOE_PROGRESS.md`。
|
||
|
||
## 1. 架构（config.json，已核对）
|
||
|
||
num_hidden_layers=24, hidden=2880, **head_dim=64**（≠hidden/heads）, n_heads=64,
|
||
n_kv_heads=8（GQA n_rep=8）, expert intermediate=2880, **num_local_experts=32**,
|
||
**num_experts_per_tok=4**, vocab=201088, max_pos=131072, rope_theta=150000,
|
||
sliding_window=128（交替层，见 `layer_types`）, rms_norm_eps=1e-5, swiglu_limit=7.0,
|
||
alpha=1.702, tie_embeddings=false。
|
||
量化：**MXFP4**，仅 expert MLP（gate_up/down 的 `_blocks`+`_scales`）；
|
||
attn/router/embed/lm_head 为 BF16。
|
||
|
||
## 2. 参考数学（HF transformers `modeling_gpt_oss.py`，逐字核对）
|
||
|
||
### RMSNorm — 标准（fp32 算 variance，eps=1e-5）。
|
||
|
||
### Router（`GptOssTopKRouter`，softmax 在 topk **之后**，含 bias）
|
||
```
|
||
logits = x @ W_router^T + b_router # [T, 32]
|
||
top_val, idx = topk(logits, k=4, dim=-1) # [T, 4]
|
||
top_val = softmax(top_val, dim=-1) # 仅对选中的 4 个归一化
|
||
scores = zeros[T,32].scatter(1, idx, top_val)
|
||
```
|
||
|
||
### Experts（`GptOssExperts`，fused gate_up，**interleaved**；clamped；(up+1)·glu）
|
||
```
|
||
alpha=1.702; limit=7.0
|
||
gate_up = x @ gate_up_proj[e] + gate_up_proj_bias[e] # [.., 2*dim]
|
||
gate = gate_up[..., ::2]; up = gate_up[..., 1::2] # 偶/奇 交错
|
||
gate = clamp(gate, max=limit) # 仅上界
|
||
up = clamp(up, min=-limit, max=limit)
|
||
glu = gate * sigmoid(gate * alpha)
|
||
h = (up + 1) * glu # 注意 (up+1)
|
||
y_e = h @ down_proj[e] + down_proj_bias[e]
|
||
out = Σ_{e∈top4} scores[t,e] * y_e
|
||
```
|
||
|
||
### Attention（`eager_attention_forward`，**带 sinks**）
|
||
```
|
||
scaling = head_dim**-0.5 = 64**-0.5；q/k/v/o 都有 bias
|
||
RoPE(theta=150000) on q,k；repeat_kv(n_rep=8)
|
||
attn = (q @ k^T) * scaling + causal_mask # 滑窗层叠加 banded(window=128)
|
||
sinks = module.sinks[head] # 每 head 一个标量
|
||
combined = cat([attn, sinks broadcast], dim=-1) # 多一列
|
||
combined -= combined.max(-1, keepdim) # 数值稳定
|
||
probs = softmax(combined, -1)
|
||
scores = probs[..., :-1] # 丢掉 sink 列 => 概率不归一到 1
|
||
o = (scores @ v) -> merge heads -> @Wo + bo
|
||
```
|
||
> sinks 等价于 softmax 分母多了 `exp(sink)`——可学习的"不注意"通道。
|
||
> 交替 sliding window：config `layer_types` 标明哪些层 window=128，其余全注意力。
|
||
|
||
与 Qwen3 的新增点：MoE FFN、MXFP4 反量化、attention sinks（softmax 多一列再丢）、
|
||
交替 sliding window、q/k/v/o bias、head_dim=64、clamped `(up+1)*glu`、rope_theta=150000。
|
||
|
||
## 3. MXFP4 反量化（expert 权重）
|
||
|
||
expert 张量名：`model.layers.{i}.mlp.experts.gate_up_proj_blocks/_scales`、
|
||
`...down_proj_blocks/_scales`（bias 为 BF16）。MXFP4：每 32 元素一 block 共享一个
|
||
E8M0(8-bit 指数) scale，每元素 4-bit FP4(E2M1)。反量化
|
||
`val = fp4_lut[code] * 2^(e8m0 - 127)`。**P19.1 先用 Python(numpy) 反量化并与 HF 一层
|
||
数值对照**（block 方向 / LUT / gate_up interleave），再写进 Rust loader。
|
||
|
||
## 4. 路线（正确优先）
|
||
|
||
1. **P19.1** Python 侦查 + MXFP4 反量化验证（不依赖 GPU）。
|
||
2. **P19.2** `config.rs` 加 MoE 字段（Qwen3 路径不变）。
|
||
3. **P19.3** `gptoss.rs`：dense（attn+sinks+bias+滑窗 / norm / lm_head）+ MoE FFN
|
||
（正确优先：逐 token top-4 gather→clamped SwiGLU→加权和）；MXFP4 在 `from_weights`
|
||
反量化为 BF16。验收：prefill logits 与 HF BF16 容差内一致（top-1 一致）。
|
||
4. **P19.4** 接 PP（experts 随层切），`--pp` 端到端；PP=2/4 与 PP=1 等价。
|
||
5. **P19.5** llama.cpp 对比（升级 submodule 到支持 gpt-oss 的版本 + 取/转 GGUF），
|
||
跑 AIME 2025 + GSM8K，复用 `tools/bench` + `summarize_fullq.py`。
|
||
|
||
## 5. 风险
|
||
|
||
- MXFP4 格式细节必须逐字对 → Python 反量化兜底。
|
||
- attention sinks + 交替滑窗：现有 flash/paged kernel 未必支持 → 正确优先版本先走朴素
|
||
attention（显式 mask + sink 列）。
|
||
- llama.cpp pinned b9371 早于 gpt-oss（约 2025-08）→ 需升级 submodule，有连锁影响。
|
||
- 性能：MoE 正确优先版本（逐 expert gather/scatter）会慢；先对再快。
|
||
- **环境**：huggingface.co 被墙，需经代理 + hf-mirror 下载（见 `MOE_PROGRESS.md` §2）。
|
||
|
||
## 6. 不在本阶段范围
|
||
|
||
GPU 原生 MXFP4 + 按需反量化 kernel（先全 BF16）；高性能 grouped-GEMM / expert parallel；
|
||
TP×MoE；单卡运行（需 MXFP4-native）。
|