moe(wip): gpt-oss-20b groundwork — config fields, arch doc, MXFP4 tools
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 <noreply@anthropic.com>
This commit is contained in:
79
tools/gptoss_dequant.py
Normal file
79
tools/gptoss_dequant.py
Normal file
@@ -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()
|
||||
82
tools/mxfp4_probe.py
Normal file
82
tools/mxfp4_probe.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user