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>
This commit is contained in:
93
docs/19-moe-gpt-oss.md
Normal file
93
docs/19-moe-gpt-oss.md
Normal file
@@ -0,0 +1,93 @@
|
||||
# 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）。
|
||||
177
docs/MOE_PROGRESS.md
Normal file
177
docs/MOE_PROGRESS.md
Normal file
@@ -0,0 +1,177 @@
|
||||
# MoE (gpt-oss-20b) — 工作进度与续作指南
|
||||
|
||||
> **中断原因**：用户要重启 dash5 机器（IP 等可能变），让我先把当前 MoE 支持工作的状态
|
||||
> 完整记录到本文件，重启后据此继续。本文件是"重启后从这里接着干"的唯一入口。
|
||||
|
||||
最后更新：Phase 18 (PP) 已完成并 push；Phase 19 (MoE/gpt-oss-20b) 刚起步（下载受阻，
|
||||
架构与参考数学已侦查清楚）。
|
||||
|
||||
---
|
||||
|
||||
## 0. 一句话现状
|
||||
|
||||
- ✅ **Phase 18 流水线并行 (PP)** 全部完成、验证、benchmark，已 commit 并 push 到
|
||||
`origin/phase18-pipeline-parallelism`（gitea）。
|
||||
- 🚧 **Phase 19 MoE (gpt-oss-20b)** 刚开始：架构 + HF 参考数学已核对（见
|
||||
`docs/19-moe-gpt-oss.md`），**模型还没下载完**（HF 被墙，正在解决下载路径），代码未动。
|
||||
|
||||
---
|
||||
|
||||
## 1. 环境关键事实（重启后很可能变 / 需重新确认）
|
||||
|
||||
- **本机**（开发机，非 GPU）：`/home/gahow/projects/xserv`，有公网（走代理）。
|
||||
- **dash5**（GPU 机，8×RTX 5090，无 NVLink，0-3/4-7 分组）：通过 `ssh dash5` 访问。
|
||||
- 远端仓库目录：`/opt/wjh/projects/xserv`，模型目录：`/opt/wjh/models/`。
|
||||
- **dash5 无外网、无 rsync**；同步用 `./tools/sync-and-build.sh`（tar over ssh）。
|
||||
- cargo 在 `$HOME/.cargo/bin`；CUDA 12.9 在 `/usr/local/cuda-12.9`。
|
||||
- ⚠️ **重启后 `ssh dash5` 的 IP/可达性可能变** —— 先 `ssh dash5 hostname` 确认；
|
||||
若连不上，检查 `~/.ssh/config` 里 `dash*` 配置 / 让用户给新地址。
|
||||
- **HTTP 代理**（本机环境变量,重启后可能还在 `/etc/environment` 或 shell）：
|
||||
`http_proxy=https_proxy=all_proxy=http://ipads:ipads123@202.120.40.82:11235`
|
||||
- **huggingface.co 被墙**（`SSL_ERROR_SYSCALL`，即使过代理）。pypi 可过代理。
|
||||
- **`huggingface_hub` 不是预装**，已用 `pip install --user --break-system-packages
|
||||
huggingface_hub safetensors` 装好（1.17.0）；venv 不可用（无 ensurepip）。
|
||||
|
||||
---
|
||||
|
||||
## 2. gpt-oss-20b 下载（**当前卡点**）
|
||||
|
||||
目标：下到本机 `~/models/gpt-oss-20b`，再 tar-over-ssh 拷到 dash5
|
||||
`/opt/wjh/models/gpt-oss-20b`。
|
||||
|
||||
**已验证可行的下载路径**（重启后照此做）：
|
||||
- huggingface.co 直连/经代理都失败。
|
||||
- hf-mirror.com 的 `/resolve/` 会 **308 跳回 huggingface.co**（也被墙）——所以不能用
|
||||
`curl -L` 跟跳转，`huggingface_hub` 设 `HF_ENDPOINT` 在新版(1.17)上 HEAD 也失败。
|
||||
- ✅ **能用的办法**：直接走 **hf-mirror 的 `/raw/`（小文件）和实际 CDN，经代理 curl**。
|
||||
已成功取到 `config.json`（200, 1799 bytes）：
|
||||
```bash
|
||||
curl -s -x "http://ipads:ipads123@202.120.40.82:11235" \
|
||||
"https://hf-mirror.com/openai/gpt-oss-20b/raw/main/config.json" -o config.json
|
||||
```
|
||||
大文件（safetensors）要用 `/resolve/main/<file>` 且 **指定 `-x` 代理、不要 `-L`**，
|
||||
若仍 308 跳回 hf.co，则改用 hf-mirror 的 LFS 直链或 `huggingface_hub` 配合
|
||||
`HF_ENDPOINT=https://hf-mirror.com` + 代理（库内部不跟 308）。**下载脚本草稿在
|
||||
`/tmp/dl_shards.sh`（重启后 /tmp 会清空，需重建）。**
|
||||
|
||||
**待下载文件**（3 个分片 + 元数据，总 ~13.5GB MXFP4）：
|
||||
- `model-00000-of-00002.safetensors`、`model-00001-of-00002.safetensors`、
|
||||
`model-00002-of-00002.safetensors`（注意是 0/1/2 三个，命名 of-00002）
|
||||
- `model.safetensors.index.json`、`config.json`、`tokenizer.json`、
|
||||
`tokenizer_config.json`、`special_tokens_map.json`、`generation_config.json`、
|
||||
`chat_template.jinja`
|
||||
|
||||
**重启后第一步**：`ls -la ~/models/gpt-oss-20b/` 看已下了哪些、`wc -c` 校验分片大小，
|
||||
断点续传用 `curl -C -`。
|
||||
|
||||
---
|
||||
|
||||
## 3. gpt-oss-20b 架构（config.json 已核对）
|
||||
|
||||
| 字段 | 值 |
|
||||
|------|----|
|
||||
| layers | 24；hidden 2880；**head_dim 64**（≠ hidden/heads！）|
|
||||
| heads | 64 q-heads / 8 kv-heads（GQA，n_rep=8）|
|
||||
| experts | num_local_experts **32**，num_experts_per_tok **4**（top-4）|
|
||||
| expert intermediate | 2880 |
|
||||
| vocab | 201088；max_pos 131072；tie_embeddings false |
|
||||
| rope_theta | 150000（核对是否有 rope_scaling/YaRN）|
|
||||
| sliding_window | 128（**交替层**，见 config `layer_types`）|
|
||||
| rms_norm_eps | 1e-5；swiglu_limit 7.0；alpha 1.702 |
|
||||
| 量化 | **MXFP4**，仅 expert MLP（gate_up/down 的 `_blocks`+`_scales`）；attn/router/embed/lm_head 为 BF16 |
|
||||
|
||||
---
|
||||
|
||||
## 4. HF 参考数学（已从 transformers `modeling_gpt_oss.py` 逐字核对，务必照抄）
|
||||
|
||||
完整版见 `docs/19-moe-gpt-oss.md` §2。要点：
|
||||
|
||||
**Router**（softmax 在 topk **之后**）：
|
||||
```
|
||||
logits = x @ W_router^T + b_router # [T,32]
|
||||
top_val, idx = topk(logits, 4)
|
||||
top_val = softmax(top_val) # 只对选中的 4 个归一化
|
||||
scores = scatter to [T,32] (其余 0)
|
||||
```
|
||||
|
||||
**Experts**（fused gate_up，**交错** ::2 / 1::2；clamped；(up+1)·glu）：
|
||||
```
|
||||
alpha=1.702, limit=7.0
|
||||
gate_up = x @ gate_up_proj[e] + bias # [.., 2*2880]
|
||||
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] + bias
|
||||
out = Σ_{e∈top4} scores[t,e] * y_e
|
||||
```
|
||||
|
||||
**Attention（带 sinks）**：
|
||||
```
|
||||
scaling = 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(+ 滑窗层叠加 banded window=128)
|
||||
combined = cat([attn, sinks_per_head], dim=-1) # 每 head 一个标量 sink，多一列
|
||||
combined -= combined.max(-1, keepdim) # 数值稳定
|
||||
probs = softmax(combined, -1)
|
||||
scores = probs[..., :-1] # 丢掉 sink 列（概率不归一到 1！）
|
||||
o = scores @ v -> merge heads -> @Wo + bo
|
||||
```
|
||||
|
||||
**RMSNorm**：标准（fp32 算 variance，eps=1e-5）。
|
||||
|
||||
参考源码已存（重启后 /tmp 清空需重取）：`pip download transformers --no-deps`
|
||||
解 wheel 取 `transformers/models/gpt_oss/modeling_gpt_oss.py`（967 行）。
|
||||
|
||||
---
|
||||
|
||||
## 5. MXFP4 反量化（expert 权重）
|
||||
|
||||
- expert 张量名：`model.layers.{i}.mlp.experts.gate_up_proj_blocks` + `..._scales`，
|
||||
`...down_proj_blocks` + `..._scales`（bias 是 BF16 的 `gate_up_proj_bias`/`down_proj_bias`）。
|
||||
- MXFP4：每 **32** 元素一 block，共享一个 **E8M0**（8-bit 指数）scale，每元素 4-bit
|
||||
FP4(E2M1，16 码字)。反量化 `val = fp4_lut[code] * 2^(e8m0 - 127)`。
|
||||
- **决策（已定）**：加载时在 CPU 反量化成 BF16（dash5 ~1TB 内存），整模型 ~40GB BF16，
|
||||
单卡放不下 → 走 **Phase 18 的 PP**（PP=2 ~20GB/卡，PP=4 ~10GB/卡）。不写 GPU 原生
|
||||
MXFP4 kernel（风险高、慢），先正确跑通+对比，后续再优化。
|
||||
|
||||
---
|
||||
|
||||
## 6. 实施路线（Phase 19，逐步可验证）
|
||||
|
||||
1. **P19.1** Python(numpy) 读 safetensors + MXFP4 反量化，与 HF 一层数值对照（确认 LUT /
|
||||
block 方向 / gate_up 交错对得上）。**不依赖 GPU，重启后可先做。**
|
||||
2. **P19.2** `crates/xserv-model/src/config.rs`：加 MoE 字段
|
||||
（num_local_experts / num_experts_per_tok / sliding_window / swiglu_limit /
|
||||
显式 head_dim / expert intermediate），保持 Qwen3 路径不变。
|
||||
3. **P19.3** 新文件 `crates/xserv-model/src/gptoss.rs`：dense（attn+sinks+bias+滑窗 /
|
||||
RMSNorm / lm_head）+ MoE FFN（正确优先：逐 token top-4 gather→clamped SwiGLU→加权和）。
|
||||
MXFP4 在 `from_weights` 反量化为 BF16。验收：prefill logits 与 HF BF16 容差内一致。
|
||||
4. **P19.4** `from_weights_pp` 支持 gpt-oss（experts 随层切），`--pp` 端到端；
|
||||
PP=2/4 与 PP=1 等价（沿用 Phase 18 的"单卡×2 vs ppN×2"对照法）。注：~40GB 需 PP≥2。
|
||||
5. **P19.5** llama.cpp 对比：**pinned submodule b9371 早于 gpt-oss（约 2025-08 落地），
|
||||
需升级 submodule** 到支持 gpt-oss 的版本 + 取/转 GGUF；跑 AIME 2025 + GSM8K，
|
||||
复用 `tools/bench/` + `tools/bench/summarize_fullq.py`（已有，PP 阶段写的）。
|
||||
|
||||
---
|
||||
|
||||
## 7. 复用 Phase 18 的资产
|
||||
|
||||
- 多卡：`--pp N`（已验证），`crates/xserv-distributed`（NCCL P2P + AllReduce）。
|
||||
- bench：`tools/bench/runner.py`（支持 `--pp`/`--tp`）、`summarize_fullq.py`、
|
||||
`tools/pp_quality_full.sh`（xserv 0-3 ‖ llama 4-7 并行跑 AIME+GSM8K 的范式可直接改用）。
|
||||
- 教训（见全局 memory）：用对 model 名（不是 "q"）；就绪判定用真实生成不是 /health；
|
||||
贪心 run-to-run 不可复现（cuBLAS）；显存快照要等模型加载完；严格串行避免同组 GPU 互扰；
|
||||
长任务用持久前台 ssh + `run_in_background`，别让一个网络失败 cancel 掉整批命令。
|
||||
|
||||
---
|
||||
|
||||
## 8. 重启后立即要做（checklist）
|
||||
|
||||
1. `ssh dash5 hostname` 确认 GPU 机可达（不行就问用户新地址 / 改 ~/.ssh/config）。
|
||||
2. `git -C ~/projects/xserv log --oneline -6` 确认 PP 5 个 commit 还在
|
||||
（`859c0cc..` 那串，分支 `phase18-pipeline-parallelism`）。
|
||||
3. `ls -la ~/models/gpt-oss-20b/` 看下载进度，续传缺的分片（§2）。
|
||||
4. 重新 `pip download transformers` 取参考源码（/tmp 已清）。
|
||||
5. 从 §6 的 P19.1 接着干。
|
||||
Reference in New Issue
Block a user