From c7d0750c32acf44efdf1192f1a41d9bdfbf4a3a6 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Fri, 29 May 2026 21:01:53 +0800 Subject: [PATCH] =?UTF-8?q?moe(wip):=20gpt-oss-20b=20groundwork=20?= =?UTF-8?q?=E2=80=94=20config=20fields,=20arch=20doc,=20MXFP4=20tools?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 19 start. config.rs: explicit head_dim (gpt-oss=64) + MoE fields (num_local_experts, num_experts_per_tok, swiglu_limit, sliding_window, layer_types) with accessors; Qwen3/GPT-2 paths unchanged (fall back to hidden/num_heads when head_dim absent). docs/19-moe-gpt-oss.md: architecture + exact HF reference math (router softmax-after-topk, interleaved clamped (up+1)*glu experts, attention sinks, alternating sliding window, rotate_half RoPE theta=150000, head_dim 64), verified tensor layout, MXFP4 dequant plan. docs/MOE_PROGRESS.md: resume/handoff snapshot. tools/mxfp4_probe.py: inspect safetensors + validate MXFP4 decode (done). tools/gptoss_dequant.py: MXFP4 experts -> plain BF16 safetensors dir so the existing loader reads it (no MXFP4 in Rust for the first pass). Verified: llama.cpp (dash5, LLM_ARCH_OPENAI_MOE) runs the gpt-oss-20b MXFP4 GGUF correctly (17*24 -> 408) = the correctness oracle. MXFP4 decode validated in numpy. Model + GGUF staged on dash5. Co-Authored-By: Claude Opus 4.8 --- crates/xserv-model/src/config.rs | 65 ++++++++++++++++++++++++- docs/19-moe-gpt-oss.md | 35 ++++++++++++++ tools/gptoss_dequant.py | 79 ++++++++++++++++++++++++++++++ tools/mxfp4_probe.py | 82 ++++++++++++++++++++++++++++++++ 4 files changed, 260 insertions(+), 1 deletion(-) create mode 100644 tools/gptoss_dequant.py create mode 100644 tools/mxfp4_probe.py diff --git a/crates/xserv-model/src/config.rs b/crates/xserv-model/src/config.rs index 5c88358..9d485b9 100644 --- a/crates/xserv-model/src/config.rs +++ b/crates/xserv-model/src/config.rs @@ -46,6 +46,28 @@ pub struct ModelConfig { pub rope_theta: Option, #[serde(default)] pub tie_word_embeddings: Option, + + // Explicit head_dim (gpt-oss: 64, which is NOT hidden/num_heads). When + // absent, head_dim() falls back to hidden/num_heads (Qwen3, GPT-2). + #[serde(default)] + pub head_dim: Option, + + // MoE (gpt-oss). Absent for dense models. + #[serde(default)] + pub num_local_experts: Option, + #[serde(default)] + pub num_experts_per_tok: Option, + // gpt-oss clamped-SwiGLU limit (config: swiglu_limit, default 7.0). + #[serde(default)] + pub swiglu_limit: Option, + + // Sliding-window attention (gpt-oss: 128 on alternating layers). The + // pattern is given by `layer_types` (e.g. "sliding_attention" / + // "full_attention" per layer); absent for dense models. + #[serde(default)] + pub sliding_window: Option, + #[serde(default)] + pub layer_types: Option>, } impl ModelConfig { @@ -81,7 +103,48 @@ impl ModelConfig { } pub fn head_dim(&self) -> usize { - self.hidden() / self.num_heads() + // gpt-oss sets head_dim explicitly (64 != 2880/64). Dense models omit it. + self.head_dim.unwrap_or_else(|| self.hidden() / self.num_heads()) + } + + // ----- MoE (gpt-oss) ----- + + /// True for MoE models (have an expert count in config). + pub fn is_moe(&self) -> bool { + self.num_local_experts.is_some() + } + + pub fn num_experts(&self) -> usize { + self.num_local_experts.unwrap_or(0) + } + + pub fn experts_per_tok(&self) -> usize { + self.num_experts_per_tok.unwrap_or(0) + } + + /// Clamp bound for gpt-oss SwiGLU (config `swiglu_limit`, default 7.0). + pub fn swiglu_limit(&self) -> f32 { + self.swiglu_limit.unwrap_or(7.0) as f32 + } + + /// Whether layer `i` uses sliding-window attention. gpt-oss alternates per + /// `layer_types`; if that's absent but `sliding_window` is set, fall back to + /// the common "every other layer" pattern (even = sliding). Dense → false. + pub fn layer_uses_sliding_window(&self, layer: usize) -> bool { + if self.sliding_window.is_none() { + return false; + } + match &self.layer_types { + Some(types) => types + .get(layer) + .map(|t| t.contains("sliding")) + .unwrap_or(false), + None => layer % 2 == 0, + } + } + + pub fn sliding_window(&self) -> Option { + self.sliding_window } pub fn ln_eps(&self) -> f32 { diff --git a/docs/19-moe-gpt-oss.md b/docs/19-moe-gpt-oss.md index 898ea11..2ee0d92 100644 --- a/docs/19-moe-gpt-oss.md +++ b/docs/19-moe-gpt-oss.md @@ -59,6 +59,41 @@ o = (scores @ v) -> merge heads -> @Wo + bo 与 Qwen3 的新增点：MoE FFN、MXFP4 反量化、attention sinks（softmax 多一列再丢）、 交替 sliding window、q/k/v/o bias、head_dim=64、clamped `(up+1)*glu`、rope_theta=150000。 +### 实测张量布局（layer 0，已用 `tools/mxfp4_probe.py` 核对） +``` +self_attn.q_proj.weight [4096,2880] +bias[4096] # 64 heads*64 +self_attn.k_proj.weight [512,2880] +bias[512] # 8 kv*64 +self_attn.v_proj.weight [512,2880] +bias[512] +self_attn.o_proj.weight [2880,4096] +bias[2880] +self_attn.sinks [64] # 每 q-head 一个标量（BF16） +input_layernorm.weight [2880]; post_attention_layernorm.weight [2880] +mlp.router.weight [32,2880] +bias[32] +mlp.experts.gate_up_proj_blocks [32,5760,90,16] U8 + _scales [32,5760,90] U8 + _bias[32,5760] BF16 +mlp.experts.down_proj_blocks [32,2880,90,16] U8 + _scales [32,2880,90] U8 + _bias[32,2880] BF16 +# 全局: model.embed_tokens.weight, model.norm.weight, lm_head.weight (BF16) +``` +MXFP4 打包：`[..., nblk=90, 16]` U8，每 16 字节 = 32 个 FP4 码（低 nibble=偶 idx，高 nibble=奇 idx）， +每 block 一个 E8M0 scale；`90*32 = 2880 = 输入(hidden)维`。即 gate_up 每 expert 权重逻辑 shape +`[5760 out, 2880 in]`（**已转置存储**：行=out，列=in，与 HF `nn.Linear` 一致 `y=x·Wᵀ`）。 + +### RoPE（**rotate_half，非 interleave**） +``` +dim = head_dim = 64; base = rope_theta = 150000 +inv_freq = 1 / base^(arange(0,64,2)/64) # 32 项 +freqs = pos ⊗ inv_freq # [S, 32]；cos/sin = cos(freqs)/sin(freqs) (不 doubling) +# 应用: x=[.., 64], first=x[:32], second=x[32:] +# out_first = first*cos - second*sin +# out_second = second*cos + first*sin +``` +> ⚠️ 与 Qwen3 的 RoPE kernel（interleave）不同 —— gptoss 走 rotate_half。需单独处理。 + +### Decoder layer（pre-norm 残差，结构同 Qwen3） +``` +h = x + attn(input_norm(x)) # attn 含 sinks/bias/滑窗 +out = h + moe(post_norm(h)) # moe = router + top4 experts 加权和 +``` +最终：`logits = lm_head(norm(h_last))`。无 q_norm/k_norm（与 Qwen3 不同，gptoss 没有）。 + ## 3. MXFP4 反量化（expert 权重） expert 张量名：`model.layers.{i}.mlp.experts.gate_up_proj_blocks/_scales`、 diff --git a/tools/gptoss_dequant.py b/tools/gptoss_dequant.py new file mode 100644 index 0000000..ac2e8b6 --- /dev/null +++ b/tools/gptoss_dequant.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +"""Dequantize gpt-oss-20b MXFP4 expert weights -> a plain BF16 safetensors dir. + +Only the expert MLPs are MXFP4 (`*_blocks` uint8 packed 4-bit + `*_scales` uint8 +E8M0, block=32); everything else is already BF16. We decode experts to BF16 and +re-emit a standard HF-format dir so xserv's normal safetensors loader reads it +(keeps the first MoE pass free of any MXFP4 code in Rust). + +Fused expert outputs (per layer i), matching HF `GptOssExperts` param shapes: + model.layers.{i}.mlp.experts.gate_up_proj [E, hidden, 2*inter] bf16 + model.layers.{i}.mlp.experts.down_proj [E, inter, hidden] bf16 + (*_bias tensors pass through unchanged) + +NOTE on transpose: the MXFP4 `_blocks` decode to [E, OUT, IN] (out-major, the +contraction dim = nblk*32 last). HF's nn.Parameter for these is [E, IN, OUT] +(it does `x @ gate_up_proj`). We emit [E, IN, OUT] (transpose last two dims) so +the names/shapes match HF exactly and xserv can treat them uniformly. + +Run on the GPU host (torch + the model + disk): + python3 tools/gptoss_dequant.py /opt/wjh/models/gpt-oss-20b /opt/wjh/models/gpt-oss-20b-bf16 +""" +import sys, os, json, glob +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +# FP4 E2M1 code -> value (OCP MX). 16 entries. +FP4 = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float32) + +def dequant(blocks: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """blocks uint8 [..., nblk, 16], scales uint8 [..., nblk] -> bf16 [..., nblk*32].""" + blocks = blocks.to(torch.int64) + lo = blocks & 0xF + hi = (blocks >> 4) & 0xF + codes = torch.stack([lo, hi], dim=-1).reshape(*blocks.shape[:-1], blocks.shape[-1] * 2) + vals = FP4[codes] # [..., nblk, 32] f32 + scale = torch.exp2(scales.to(torch.float32) - 127.0) # [..., nblk] + out = vals * scale[..., None] + out = out.reshape(*out.shape[:-2], out.shape[-2] * out.shape[-1]) + return out.to(torch.bfloat16) + +def main(): + src, dst = sys.argv[1], sys.argv[2] + os.makedirs(dst, exist_ok=True) + wm = json.load(open(os.path.join(src, "model.safetensors.index.json")))["weight_map"] + + shards = {} + for name, shard in wm.items(): + shards.setdefault(shard, []).append(name) + + out = {} + for shard in sorted(shards): + h = safe_open(os.path.join(src, shard), framework="pt") + keys = set(h.keys()) + for name in shards[shard]: + if name.endswith("_scales"): + continue + if name.endswith("_blocks"): + base = name[:-len("_blocks")] # ...gate_up_proj / ...down_proj + sc = base + "_scales" + sc_h = h if sc in keys else safe_open(os.path.join(src, wm[sc]), framework="pt") + deq = dequant(h.get_tensor(name), sc_h.get_tensor(sc)) # [E, OUT, IN] + out[base] = deq.transpose(1, 2).contiguous() # [E, IN, OUT] (HF param layout) + print("dequant", base, tuple(out[base].shape), flush=True) + else: + out[name] = h.get_tensor(name) # already bf16/other; pass through + + save_file(out, os.path.join(dst, "model.safetensors"), metadata={"format": "pt"}) + for f in glob.glob(os.path.join(src, "*.json")) + glob.glob(os.path.join(src, "*.jinja")): + b = os.path.basename(f) + if b == "model.safetensors.index.json": + continue + open(os.path.join(dst, b), "wb").write(open(f, "rb").read()) + print("DEQUANT_DONE ->", dst, flush=True) + +if __name__ == "__main__": + main() diff --git a/tools/mxfp4_probe.py b/tools/mxfp4_probe.py new file mode 100644 index 0000000..14766e4 --- /dev/null +++ b/tools/mxfp4_probe.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""P19.1 — inspect gpt-oss-20b safetensors + validate MXFP4 dequant (CPU only). + +Run: python3 tools/mxfp4_probe.py /path/to/gpt-oss-20b + +Does three things: + 1. List the layer-0 tensor names, shapes, dtypes (esp. expert _blocks/_scales). + 2. Dequantize one expert's gate_up_proj from MXFP4 -> fp32 with our own LUT/decode. + 3. Print stats so we can eyeball sanity (range, mean, a few values). + +MXFP4 (OCP microscaling) as used by gpt-oss: + - elements are FP4 E2M1 (4-bit), packed 2-per-byte. + - every block of 32 consecutive elements shares one E8M0 scale (8-bit exponent). + - value = fp4_e2m1_lut[code] * 2**(scale_e8m0 - 127) +The `_blocks` tensor holds the packed 4-bit codes; `_scales` holds the per-block +E8M0 exponents (uint8). We confirm shapes line up (last dim of blocks * 2 / ... ) +and that decoded values are in a sane range. +""" +import sys, json, os +import numpy as np + +# FP4 E2M1 code -> value lookup (OCP MX spec). 16 codes: sign(1) exp(2) mant(1). +# Values: 0, 0.5, 1, 1.5, 2, 3, 4, 6 and their negatives. +FP4_E2M1 = np.array( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=np.float32, +) + +def dequant_mxfp4(blocks: np.ndarray, scales: np.ndarray) -> np.ndarray: + """gpt-oss layout: + blocks: uint8 [..., nblk, 16] -> each 16-byte row = 32 FP4 codes (2/byte) + scales: uint8 [..., nblk] -> one E8M0 exponent per block + Returns fp32 [..., nblk*32] (the input/contraction dim).""" + lo = blocks & 0x0F + hi = (blocks >> 4) & 0x0F + # within a byte, low nibble is element 2i, high nibble is element 2i+1 + codes = np.empty(blocks.shape[:-1] + (blocks.shape[-1] * 2,), dtype=np.uint8) + codes[..., 0::2] = lo + codes[..., 1::2] = hi # [..., nblk, 32] + vals = FP4_E2M1[codes] # [..., nblk, 32] + scale_f = np.power(2.0, scales.astype(np.float32) - 127.0) # [..., nblk] + out = vals * scale_f[..., None] # [..., nblk, 32] + return out.reshape(out.shape[:-2] + (out.shape[-2] * out.shape[-1],)) # [..., nblk*32] + +def main(): + d = sys.argv[1] if len(sys.argv) > 1 else "/home/gahow/models/gpt-oss-20b" + from safetensors import safe_open + idx = json.load(open(os.path.join(d, "model.safetensors.index.json"))) + wm = idx["weight_map"] + l0 = {k: v for k, v in wm.items() if "layers.0." in k} + print("=== layer 0 tensors ===") + # open each shard lazily + handles = {} + def get(name): + shard = wm[name] + if shard not in handles: + handles[shard] = safe_open(os.path.join(d, shard), framework="numpy") + return handles[shard] + for k in sorted(l0): + h = get(k) + t = h.get_slice(k) + print(f" {k.replace('model.layers.0.','L0.')} shape={t.get_shape()} dtype={t.get_dtype()}") + + # find expert gate_up blocks/scales + gu_b = next((k for k in l0 if "gate_up_proj_blocks" in k), None) + gu_s = next((k for k in l0 if "gate_up_proj_scales" in k), None) + if gu_b and gu_s: + print(f"\n=== dequant {gu_b.split('.')[-1]} (expert 0) ===") + blocks = get(gu_b).get_tensor(gu_b) + scales = get(gu_s).get_tensor(gu_s) + print(" blocks", blocks.shape, blocks.dtype, " scales", scales.shape, scales.dtype) + b0 = blocks[0]; s0 = scales[0] # expert 0: blocks [5760,90,16], scales [5760,90] + deq = dequant_mxfp4(b0.astype(np.uint8), s0.astype(np.uint8)) # -> [5760, 2880] + print(" expert0 dequant shape", deq.shape, "(expect [5760, 2880] = [2*inter, hidden])", + "\n min %.4f max %.4f mean %.4f std %.4f" % (deq.min(), deq.max(), deq.mean(), deq.std())) + print(" row0 first 8 vals:", np.round(deq[0, :8], 4)) + else: + print("\n(no gate_up_proj_blocks/_scales found — check tensor names above)") + +if __name__ == "__main__": + main()