Files
xserv/docs/19-moe-gpt-oss.md
Gahow Wang 057a3c68a3 docs: Phase 19 MoE (gpt-oss-20b) design + progress snapshot
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>
2026-05-29 19:13:23 +08:00

4.8 KiB
Raw Blame History

Phase 19: MoE — gpt-oss-20b

目标:在 xserv 支持 MoE,用 openai/gpt-oss-20b 端到端跑通,并与 llama.cpp 在 AIME 2025 / GSM8K 上对比正确性与性能。MXFP4 expert 权重加载时反量化为 BF16整模型 ~40GB 单卡放不下 → 复用 Phase 18 的 PPPP=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=8GQA 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 MLPgate_up/down 的 _blocks+_scales attn/router/embed/lm_head 为 BF16。

2. 参考数学HF transformers modeling_gpt_oss.py,逐字核对)

RMSNorm — 标准fp32 算 varianceeps=1e-5

RouterGptOssTopKRoutersoftmax 在 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)

ExpertsGptOssExpertsfused gate_upinterleavedclamped(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

Attentioneager_attention_forward带 sinks

scaling = head_dim**-0.5 = 64**-0.5q/k/v/o 都有 bias
RoPE(theta=150000) on q,krepeat_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 windowconfig layer_types 标明哪些层 window=128其余全注意力。

与 Qwen3 的新增点MoE FFN、MXFP4 反量化、attention sinkssoftmax 多一列再丢)、 交替 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/_scalesbias 为 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.rsdenseattn+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 接 PPexperts 随层切),--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