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>
4.8 KiB
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:configlayer_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. 路线(正确优先)
- P19.1 Python 侦查 + MXFP4 反量化验证(不依赖 GPU)。
- P19.2
config.rs加 MoE 字段(Qwen3 路径不变)。 - 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 一致)。 - P19.4 接 PP(experts 随层切),
--pp端到端;PP=2/4 与 PP=1 等价。 - 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)。