From 7c6fdc7c6ed2fac8e6ddcecdc38ef69733b1a4a6 Mon Sep 17 00:00:00 2001 From: loopyt Date: Wed, 21 Oct 2026 16:50:31 +0800 Subject: [PATCH] fix(sm90): expand expert scales on the contraction axis only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `requant_expert_fp4` expanded the UE8M0 scale with `_expand_blocks`, which repeats *both* axes. That is right for dense weights, whose scales are `(n_out/32, k/32)`, but expert scales are already per-output-channel: `(n_out, k/32)`. Expanding both axes therefore treated `(2304, 160)` as `(72, 5)` and handed every one of the 2304 output channels the scales of the first 72 rows, scrambling each expert's weights. This is what the layer comparison was pointing at: `shared_expert` matched the reference (cos 0.9987) because it goes through `requant_dense_fp8`, which uses `_expand_blocks` correctly, while `combined` -- routed experts plus shared -- was off by 3.6e-1. Verified against the authoritative implementation, `inference/convert.py`'s `cast_e2m1fn_to_e4m3fn`, on layer 0 / expert 370: w1 cos 0.8936 -> 1.0000 rel_err 5.35e-1 -> 2.84e-2 w3 cos 0.8960 -> 0.9997 rel_err 5.26e-1 -> 2.69e-2 w2 cos 0.8964 -> 0.9996 rel_err 5.20e-1 -> 2.65e-2 The residual ~2.7% is the per-32 -> per-128 requantization loss, not a defect; the reference is lossless because it folds the relative per-32 offset into the FP8 values. End to end on the 11-token chat prompt: before: [17607, 41175, 389, 666] "中国的首都是**" after: [17607, 2793, 3139, 666, 6127, 666, 320, 1] "中国的首都是**北京**。" ref: [17607, 2793, 3139, 6127, 320] "中国的首都是北京。" The answer is now correct and the first three tokens match the reference exactly. The remaining difference is markdown emphasis (`**`), so some numerical drift still moves the top-1 on the occasional position. Also adds two checks that would have caught this immediately: `scripts/h20_expert_weight_check.py` diffs our requantization against DeepSeek's own converter, and `scripts/h20_quant_check.py` pins the E2M1 table, nibble order and UE8M0 bias. The nibble order and table turn out to match the reference; only the scale broadcast was wrong. --- H20-DEBUG-PLAN.md | 66 ++++++++-- deepinfer/h20_ops/debug_dump.py | 32 +++++ deepinfer/h20_ops/quant.py | 7 +- .../model_executor/models/v4_modules/mlp.py | 11 ++ scripts/compare_layers.py | 27 ++++ scripts/h20_expert_weight_check.py | 119 ++++++++++++++++++ scripts/h20_quant_check.py | 105 ++++++++++++++++ scripts/ref_layer_dump.py | 27 ++++ 8 files changed, 380 insertions(+), 14 deletions(-) create mode 100644 scripts/h20_expert_weight_check.py create mode 100644 scripts/h20_quant_check.py diff --git a/H20-DEBUG-PLAN.md b/H20-DEBUG-PLAN.md index f26afc6..b29ad99 100644 --- a/H20-DEBUG-PLAN.md +++ b/H20-DEBUG-PLAN.md @@ -93,21 +93,61 @@ length : max_new_tokens = 6 shared expert 输出、以及 swiglu/route 权重，看是路由（top-6 选中了不同专家）还是 专家计算本身偏。`engram` 在第 1/14 层，本轮已可排除其为**首个**偏离点。 -## P3 层内细分（按嫌疑排序） +## P3 已完成：根因是专家权重量化的 scale 广播错误 —— 已修复，答案正确 -- [ ] P3.1 **indexer 选出的 top-k 索引**（离散量，最容易判断）：同一层、同一输入， - 比较索引集合的重合率。低于 ~90% 就说明 indexer 错了。 -- [ ] P3.2 **index KV 字节布局**：SM90 用 `FP8_FLOAT32_SF_INDEX_KV_FORMAT` - （128 fp8 + 1 f32 = 132B/token）。store kernel 是否按「整页 values 后接整页 scales」 - 写、`_decode_p_sum_sm90` 的读法是否一致 —— 用一个 round-trip 测试验证： - 写入已知向量 → 读出 → 比 bf16 原值。 -- [ ] P3.3 **主 KV cache 的 ue8m0 解量化**：SM90 用 fp32 scale，但 cache 仍存 ue8m0 - 字节，`store_kv_cache` / `load_paged_kv_cache` 两侧口径要对齐。 -- [ ] P3.4 **engram**：表 96M 行、per-32 scale，本地 gather 已单测过，但和官方 - 对拍还没做。 -- [ ] P3.5 **MoE**：路由（sqrtsoftplus + route_scale 1.5）+ swiglu_limit 10.0 是否一致。 +### 根因 ---- +`h20_ops/quant.py::requant_expert_fp4` 里用了 `_expand_blocks(sf, n_out, k, 32)` 去展开 +UE8M0 scale。但两者的布局前提不同： + +| | scale 形状 | 需要展开的维度 | +|---|---|---| +| dense 权重（`requant_dense_fp8`） | `(n_out/32, k/32)` | **两个**维度 | +| 专家权重（`requant_expert_fp4`） | `(n_out, k/32)` | **只有 k 维度** | + +`_expand_blocks` 对两个维度都做 `repeat_interleave(32)` 再截断，于是 +`(2304, 160)` 被当成 `(72, 5)` 来展开 —— **全部 2304 个输出通道都拿到了前 72 行的 scale**， +整个专家权重被打乱。这精确解释了 P2 观察到的「shared 对、routed 错」的非对称： +shared expert 走 `requant_dense_fp8`（用法正确），routed experts 走这条错误路径。 + +修复：只沿 contraction 轴展开，并加断言锁住布局。 + +```python +assert tuple(sf.shape) == (n_out, k // FP4_BLOCK) +x = _unpack_fp4(packed) * sf.to(torch.float32).repeat_interleave(FP4_BLOCK, dim=1) +``` + +### 验证（`scripts/h20_expert_weight_check.py`，与官方 `inference/convert.py` +### 的 `cast_e2m1fn_to_e4m3fn` 逐元素对拍） + +| 权重 | 修复前 cos / rel_err | 修复后 cos / rel_err | +|---|---|---| +| w1 | 0.8936 / 5.35e-1 | **1.0000 / 2.84e-2** | +| w3 | 0.8960 / 5.26e-1 | **0.9997 / 2.69e-2** | +| w2 | 0.8964 / 5.20e-1 | **0.9996 / 2.65e-2** | + +剩余 2.7% 是 per-32 → per-128 的量化损失（FP8 e4m3 三位尾数，理论量级 ~3%）， +不是 bug。官方 `cast_e2m1fn_to_e4m3fn` 是**无损**的（把 per-32 的相对 offset 折进 FP8 数值）， +如果要进一步压掉这 2.7% 可以照搬，见下方 P3.6。 + +### 端到端结果 + +``` +修复前: [17607, 41175, 389, 666] -> 中国的首都是** +修复后: [17607, 2793, 3139, 666, 6127, 666, 320, 1] -> 中国的首都是**北京**。 +官方: [17607, 2793, 3139, 6127, 320] -> 中国的首都是北京。 +``` + +**语义已正确**（「北京」），前 3 个 token 与官方逐 token 一致。差异只剩 markdown 加粗 +`**`（我方多出 666/666 两个 token），说明仍有残余数值差异让 top-1 在个别位置偏移。 + +### 下一步（可选，按收益排序） + +- [ ] **P3.6 改用官方无损重化**：照 `cast_e2m1fn_to_e4m3fn` 的做法，把 per-32 e8m0 相对 + offset 折进 FP8、scale 保持 e8m0 per-32，理论上可把 2.7% 降到 0。需要确认 SM90 + DeepGEMM 是否接受 e8m0 per-32 的 expert scale（当前配方是 fp32 per-128）。 +- [ ] **P3.7 重跑 P2 分层对拍**：确认 layer0 整层 rel_err 从 1.27e-1 降到 ~3e-2 量级。 +- [ ] **P3.8 `**` 的来源**：定位是哪个位置翻的 top-1，判断是否还有别的量化口径差异。 ## P4 decode 专项（仅当 prefill 对、decode 错） diff --git a/deepinfer/h20_ops/debug_dump.py b/deepinfer/h20_ops/debug_dump.py index 53bf55f..6629906 100644 --- a/deepinfer/h20_ops/debug_dump.py +++ b/deepinfer/h20_ops/debug_dump.py @@ -22,6 +22,9 @@ _DUMP_DIR = os.getenv("DEEPINFER_DUMP_LAYERS") # different token count. Gating on the real prompt length keeps the dump comparable. _DUMP_TOKENS = int(os.getenv("DEEPINFER_DUMP_TOKENS", "0") or 0) _CAPTURED: dict[str, np.ndarray] = {} +# Only the first N layers of MoE internals are kept: the routing and per-branch outputs are the +# interesting part for bring-up, and dumping all 40 layers would slow the forward down for no gain. +_MOE_LAYERS = int(os.getenv("DEEPINFER_DUMP_MOE_LAYERS", "2") or 0) def enabled() -> bool: @@ -38,6 +41,8 @@ def record(tag: str, tensor: torch.Tensor) -> None: return if _DUMP_TOKENS and tensor.shape[0] != _DUMP_TOKENS: return + if not _want_moe(tag): + return if tensor.dim() == 3: tensor = tensor.mean(dim=1) row = tensor[-1] if tensor.dim() == 2 else tensor.reshape(-1) @@ -49,6 +54,33 @@ def record(tag: str, tensor: torch.Tensor) -> None: _CAPTURED[tag] = row.detach().float().cpu().numpy() +def record_ids(tag: str, tensor: torch.Tensor) -> None: + """Store an integer tensor as-is (expert ids); reduction would destroy its meaning.""" + if not _DUMP_DIR: + return + if not _want_moe(tag): + return + if _DUMP_TOKENS and tensor.shape[0] != _DUMP_TOKENS: + return + from deepinfer.utils import allow_cuda_sync + + with allow_cuda_sync(): + _CAPTURED[tag] = tensor.detach().to(torch.int32).cpu().numpy() + + +def _want_moe(tag: str) -> bool: + """Keep MoE internals only for the first `_MOE_LAYERS` layers.""" + if "_" not in tag: + return True + head = tag.split("_", 1)[0] + if not head.startswith("layer"): + return True + try: + return int(head[5:]) < _MOE_LAYERS + except ValueError: + return True + + def flush() -> None: """Write everything captured so far. Called once at interpreter exit.""" if not _DUMP_DIR or not _CAPTURED: diff --git a/deepinfer/h20_ops/quant.py b/deepinfer/h20_ops/quant.py index ff3124a..6a1e0c3 100644 --- a/deepinfer/h20_ops/quant.py +++ b/deepinfer/h20_ops/quant.py @@ -81,7 +81,12 @@ def requant_expert_fp4( """ n_out, half_k = packed.shape k = half_k * 2 - x = _unpack_fp4(packed) * _expand_blocks(sf.to(torch.float32), n_out, k, FP4_BLOCK) + # Expert scales are per-output-channel already: shape `(n_out, k / 32)`. Only the contraction + # axis is blocked, so only that axis gets expanded -- running this through `_expand_blocks` + # (which repeats *both* axes) would hand every output channel the scales of the first + # `n_out / 32` rows, scrambling the whole expert. + assert tuple(sf.shape) == (n_out, k // FP4_BLOCK), f"unexpected expert scale layout: {tuple(sf.shape)}" + x = _unpack_fp4(packed) * sf.to(torch.float32).repeat_interleave(FP4_BLOCK, dim=1) q, scale = _requant(x, out_block) if contraction_major_sf: scale = scale.t().contiguous() diff --git a/deepinfer/model_executor/models/v4_modules/mlp.py b/deepinfer/model_executor/models/v4_modules/mlp.py index 3040aeb..2edcc4c 100644 --- a/deepinfer/model_executor/models/v4_modules/mlp.py +++ b/deepinfer/model_executor/models/v4_modules/mlp.py @@ -22,6 +22,7 @@ from torch import nn from deepinfer import PLATFORM, AttrDict, envs from deepinfer.engine.engine_type import EngineType +from deepinfer.h20_ops import debug_dump as _debug_dump from deepinfer.logger import init_logger from deepinfer.model_executor import eplb from deepinfer.model_executor.input_metadata import InputMetadata @@ -623,6 +624,14 @@ class DeepseekMoE(nn.Module): combine_pre_evt = self.ep_buffer.capture() with limit_num_sms(envs.DEEPINFER_USERBUFFER_COMM_SMS, enabled=PLATFORM.is_gpu): shared_out = self.shared_forward(hidden_states) + if _debug_dump.enabled(): + # Capture the routing decision as well as each branch's contribution: the ids are + # a discrete quantity, so a set mismatch proves routing is wrong while an equal + # set with different outputs points at the expert weights instead. + _debug_dump.record_ids(f"layer{self.layer_idx}_topk_idx", topk_idx) + _debug_dump.record(f"layer{self.layer_idx}_topk_w", topk_weights) + _debug_dump.record(f"layer{self.layer_idx}_shared_out", shared_out) + _debug_dump.record(f"layer{self.layer_idx}_moe_out", moe_out_hidden) event_overlap = self.combine( moe_out_hidden, shared_out, @@ -630,6 +639,8 @@ class DeepseekMoE(nn.Module): combine_pre_evt, ) moe_combined_hidden, _combined_topk_weights = event_overlap.current_stream_wait() + if _debug_dump.enabled(): + _debug_dump.record(f"layer{self.layer_idx}_combined", moe_combined_hidden) except Exception: logger.exception(f"[rank {PG.get_rank()}] Exception during MoE forward, topk_idx.shape = {topk_idx.shape}") diff --git a/scripts/compare_layers.py b/scripts/compare_layers.py index 89d1bd7..b3a01d3 100644 --- a/scripts/compare_layers.py +++ b/scripts/compare_layers.py @@ -39,6 +39,24 @@ def main() -> None: if missing: print(f"WARNING: no dump from DeepInfer for layers {missing}") + # Routing first: expert ids are discrete, so a set mismatch is a definitive verdict and the + # weights/branch comparisons below only make sense once routing is known to agree. + if args.ref and args.mine: + for layer in range(4): + key = f"layer{layer}_topk_idx" + if key not in ref.files or key not in mine.files: + continue + r_idx, m_idx = ref[key].astype(np.int64), mine[key].astype(np.int64) + r_set = set(r_idx.reshape(-1).tolist()) + m_set = set(m_idx.reshape(-1).tolist()) + inter = r_set & m_set + print( + f"{key:>18} ref={sorted(r_set)}\n{'':>18} mine={sorted(m_set)}" + f"\n{'':>18} overlap={len(inter)}/{len(r_set)}" + f" {'IDENTICAL' if r_idx.shape == m_idx.shape and (r_idx == m_idx).all() else 'DIFFERS'}" + ) + print() + for key in ("embed", "layer0_attn", "layer1_attn"): if key in ref.files and key in mine.files: r, m = ref[key].astype(np.float64), mine[key].astype(np.float64) @@ -74,6 +92,15 @@ def main() -> None: print("reference top-5 tokens:", [int(x) for x in top], "logits:", [round(float(rl[x]), 3) for x in top]) # DeepInfer's logits are not dumped; derive its choice from the harness output instead. + print() + for key in sorted(k for k in ref.files if k.endswith(("_shared_out", "_combined", "_moe_out"))): + if key in mine.files and ref[key].shape == mine[key].shape: + r, m = ref[key].astype(np.float64), mine[key].astype(np.float64) + print( + f"{key:>18} cos={cosine(m, r):9.5f} rel_err={rel_err(m, r):9.2e}" + f" |mine|/|ref|={np.linalg.norm(m) / (np.linalg.norm(r) + 1e-12):8.4f}" + ) + print() if first_bad is None: print(f"RESULT: every dumped layer is within rel_err <= {args.tol}; divergence is downstream") diff --git a/scripts/h20_expert_weight_check.py b/scripts/h20_expert_weight_check.py new file mode 100644 index 0000000..7077a27 --- /dev/null +++ b/scripts/h20_expert_weight_check.py @@ -0,0 +1,119 @@ +"""Compare our FP4->FP8 expert requantization against DeepSeek's own reference conversion. + +`inference/convert.py::cast_e2m1fn_to_e4m3fn` is the authoritative implementation: it is what +produced the TP8 checkpoint that the reference model actually runs on. Two things are worth +checking there, and both matter because a wrong answer scrambles every routed expert while +leaving the (already-FP8) shared expert intact -- exactly the asymmetry the layer comparison +showed. + +1. Equality of the *decoded* weights. Both paths should represent the same real values; any + difference is accuracy we are throwing away. +2. The scale grid. The reference keeps a per-32 e8m0 scale (relative offsets folded into the FP8 + values), while the SM90 path re-quantizes on a per-128 fp32 grid. That is a legitimate design + difference only if the magnitude error stays small. + + python scripts/h20_expert_weight_check.py --layer 0 --expert 370 +""" + +import argparse +import json +from pathlib import Path + +import torch +from safetensors import safe_open + +from deepinfer.h20_ops.quant import requant_expert_fp4, ue8m0_to_fp32 + +# Copied verbatim from inference/convert.py so the comparison uses the reference algorithm +# rather than a paraphrase of it. +FP4_TABLE = 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 cast_e2m1fn_to_e4m3fn(x: torch.Tensor, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """DeepSeek's lossless FP4 -> FP8 cast (inference/convert.py).""" + assert x.dtype == torch.int8 and x.ndim == 2 + out_dim, in_dim = x.size() + in_dim *= 2 + fp8_block_size = 32 + fp4_block_size = 32 + assert in_dim % fp8_block_size == 0 and out_dim % fp8_block_size == 0 + assert scale.size(0) == out_dim and scale.size(1) == in_dim // fp4_block_size + + x = x.view(torch.uint8) + low = x & 0x0F + high = (x >> 4) & 0x0F + x = torch.stack([FP4_TABLE[low.long()], FP4_TABLE[high.long()]], dim=-1).flatten(2) + + MAX_OFFSET_BITS = 6 + b_out = out_dim // fp8_block_size + b_in = in_dim // fp8_block_size + x = x.view(b_out, fp8_block_size, b_in, fp8_block_size).transpose(1, 2) + scale = scale.float().view(b_out, fp8_block_size, b_in, -1).transpose(1, 2).flatten(2) + scale_max_offset_bits = scale.amax(dim=-1, keepdim=True) / (2**MAX_OFFSET_BITS) + offset = scale / scale_max_offset_bits + offset = offset.unflatten(-1, (fp8_block_size, -1)).repeat_interleave(fp4_block_size, dim=-1) + x = (x * offset).transpose(1, 2).reshape(out_dim, in_dim) + return x.to(torch.float8_e4m3fn), scale_max_offset_bits.squeeze(-1).to(torch.float8_e8m0fnu) + + +def read_tensor(hf_dir: Path, index: dict, name: str) -> torch.Tensor: + shard = index["weight_map"][name] + with safe_open(str(hf_dir / shard), framework="pt", device="cpu") as f: + return f.get_tensor(name) + + +def dequant_ref(fp8: torch.Tensor, sf_e8m0: torch.Tensor, block: int = 32) -> torch.Tensor: + """Reference layout: (out, in) FP8, scale (out/32, in/32) e8m0 on a 32x32 grid.""" + s = ue8m0_to_fp32(sf_e8m0) + s = s.repeat_interleave(block, 0).repeat_interleave(block, 1)[: fp8.shape[0], : fp8.shape[1]] + return fp8.float() * s + + +def dequant_ours(fp8: torch.Tensor, sf_fp32: torch.Tensor, block: int = 128) -> torch.Tensor: + """SM90 layout: (out, in) FP8, scales (in/block, out/block) fp32, contraction-major.""" + s = sf_fp32.t().repeat_interleave(block, 0).repeat_interleave(block, 1)[: fp8.shape[0], : fp8.shape[1]] + return fp8.float() * s + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--hf-dir", default="/data/models/DeepSeek-V4.1-Flash-1017") + parser.add_argument("--layer", type=int, default=0) + parser.add_argument("--expert", type=int, default=370) + args = parser.parse_args() + + hf_dir = Path(args.hf_dir) + index = json.loads((hf_dir / "model.safetensors.index.json").read_text()) + prefix = f"layers.{args.layer}.ffn.experts.{args.expert}." + + for proj in ("w1", "w3", "w2"): + packed = read_tensor(hf_dir, index, prefix + proj + ".weight") + scale = read_tensor(hf_dir, index, prefix + proj + ".scale") + print(f"\n{proj}: packed {tuple(packed.shape)} {packed.dtype}, scale {tuple(scale.shape)} {scale.dtype}") + + ref_fp8, ref_sf = cast_e2m1fn_to_e4m3fn(packed, scale) + ours_fp8, ours_sf = requant_expert_fp4(packed, scale, out_block=128) + + # `requant_expert_fp4` returns (n_out, k) FP8 and a contraction-major scale; the reference + # returns (out, in) too, so the value grids line up directly. + ref_vals = dequant_ref(ref_fp8, ref_sf) + ours_vals = dequant_ours(ours_fp8, ours_sf) + + assert ref_vals.shape == ours_vals.shape, f"{ref_vals.shape} vs {ours_vals.shape}" + num = (ours_vals - ref_vals).norm().item() + den = ref_vals.norm().item() + cos = torch.nn.functional.cosine_similarity( + ours_vals.reshape(-1), ref_vals.reshape(-1), dim=0 + ).item() + print(f" ref scale: e8m0 {tuple(ref_sf.shape)} ours: fp32 {tuple(ours_sf.shape)}") + print(f" |ref|={den:.4f} |ours|={ours_vals.norm().item():.4f} ratio={ours_vals.norm().item() / den:.5f}") + print(f" rel_err={num / den:.5f} cos={cos:.7f}") + + print("\nEXPERT_WEIGHT_CHECK_DONE") + + +if __name__ == "__main__": + main() diff --git a/scripts/h20_quant_check.py b/scripts/h20_quant_check.py new file mode 100644 index 0000000..e9f2351 --- /dev/null +++ b/scripts/h20_quant_check.py @@ -0,0 +1,105 @@ +"""Check `h20_ops.quant` against PyTorch's native FP4/FP8 semantics. + +The SM90 path re-quantizes packed E2M1 experts into FP8 while loading, and that code makes three +assumptions the checkpoint does not spell out: which nibble holds the even element, the E2M1 code +table, and the UE8M0 bias. If any of them is wrong every expert weight is scrambled, which is +exactly the symptom the layer comparison shows (shared expert matches, routed experts do not). + + python scripts/h20_quant_check.py +""" + +import torch + +from deepinfer.h20_ops.quant import _E2M1_VALUES, _unpack_fp4, requant_expert_fp4, ue8m0_to_fp32 + +# PyTorch's FP4 cast kernel does not exist for SM90 (device-side assert on H20), and the decode +# under test is pure integer/table work, so everything here runs on the CPU. +DEV = "cpu" + + +def check_e2m1_table() -> None: + """Every 4-bit code must decode to the value PyTorch assigns it.""" + codes = torch.arange(16, dtype=torch.uint8, device=DEV) + # PyTorch's own decoder: pack two codes per byte and view as the FP4 dtype. + packed = torch.stack([codes, torch.zeros_like(codes)], dim=-1).reshape(-1).contiguous() + ref = packed.view(torch.float4_e2m1fn_x2).to(torch.float32)[0::2] + table = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=DEV) + ok = torch.equal(ref, table) + print(f"E2M1 table: {'OK' if ok else 'MISMATCH'}") + if not ok: + for code in range(16): + if ref[code] != table[code]: + print(f" code {code:04b}: torch={ref[code].item():+.1f} ours={table[code].item():+.1f}") + raise SystemExit("E2M1 table disagrees with PyTorch") + + +def check_nibble_order() -> None: + """`_unpack_fp4` must interleave in the same order as the packed dtype.""" + # Distinct, exactly representable values so an order swap is unambiguous. + values = torch.tensor([1.0, 2.0, 4.0, 6.0, 0.5, 1.5, 3.0, -1.0], dtype=torch.float32, device=DEV) + packed = values.to(torch.float4_e2m1fn_x2).view(torch.uint8) # (4,) bytes + ref = packed.view(torch.float4_e2m1fn_x2).to(torch.float32) + ours = _unpack_fp4(packed.view(torch.int8)).reshape(-1) + ok = torch.equal(ref, ours) + print(f"nibble order: {'OK' if ok else 'MISMATCH'}") + if not ok: + print(f" torch: {ref.tolist()}") + print(f" ours : {ours.tolist()}") + swapped = ours.view(-1, 2).flip(-1).reshape(-1) + print(f" ours (nibbles swapped): {swapped.tolist()}") + print(f" -> swapping would {'fix it' if torch.equal(ref, swapped) else 'NOT fix it'}") + raise SystemExit("nibble order disagrees with PyTorch") + + +def check_ue8m0() -> None: + """`ue8m0_to_fp32` must match the checkpoint's scale dtype.""" + bits = torch.tensor([127, 128, 126, 130, 124], dtype=torch.uint8, device=DEV) + scales = bits.view(torch.float8_e8m0fnu) + ours = ue8m0_to_fp32(scales) + ref = torch.pow(2.0, bits.to(torch.float32) - 127) + ok = torch.equal(ours, ref) + print(f"UE8M0 bias: {'OK' if ok else 'MISMATCH'} ours={ours.tolist()} ref={ref.tolist()}") + if not ok: + raise SystemExit("UE8M0 decode disagrees") + + +def check_requant_roundtrip() -> None: + """End to end: dequantize FP4 on the CPU, requantize, and measure the error.""" + torch.manual_seed(0) + n_out, k = 256, 512 + x = torch.randn(n_out, k, device=DEV, dtype=torch.float32) * 0.05 + + # Quantize to FP4 with the checkpoint's per-32 UE8M0 grid, so the round trip starts from a + # realistic representation rather than an exact float. FP8 casts also stay on the CPU here. + groups = x.view(n_out, k // 32, 32) + amax = groups.abs().amax(dim=-1).clamp_min(torch.finfo(torch.float32).tiny) + exp = torch.ceil(torch.log2(amax / 6.0)).clamp(-127, 127) + scale = torch.pow(2.0, exp) + q4 = (groups / scale.unsqueeze(-1)).to(torch.float4_e2m1fn_x2) + packed = q4.view(torch.int8) + sf_bits = (exp.to(torch.int32) + 127).to(torch.uint8) + sf = sf_bits.view(torch.float8_e8m0fnu) + + deq_ref = q4.to(torch.float32) * scale.repeat_interleave(32, dim=1) + + weight, sfb = requant_expert_fp4(packed, sf, out_block=128) + assert weight.dtype == torch.float8_e4m3fn, weight.dtype + assert sfb.dtype == torch.float32, sfb.dtype + assert sfb.shape == (k // 128, n_out // 128), f"scale layout {tuple(sfb.shape)}" + + deq_ours = weight.to(torch.float32) * sfb.t().repeat_interleave(128, 0).repeat_interleave(128, 1)[:n_out, :k] + err = (deq_ours - deq_ref).norm() / deq_ref.norm() + print(f"requant roundtrip: rel_err={err.item():.4f} (fp8 e4m3 floor ~0.01, fp4 source itself is ~0.1)") + assert err < 0.05, "requantization lost more precision than FP8 should" + + +def main() -> None: + check_e2m1_table() + check_nibble_order() + check_ue8m0() + check_requant_roundtrip() + print("QUANT_CHECK_DONE") + + +if __name__ == "__main__": + main() diff --git a/scripts/ref_layer_dump.py b/scripts/ref_layer_dump.py index 6bd875f..c4d4cc9 100644 --- a/scripts/ref_layer_dump.py +++ b/scripts/ref_layer_dump.py @@ -32,6 +32,7 @@ def main() -> None: parser.add_argument("--config", required=True) parser.add_argument("--ids", required=True) parser.add_argument("--out", default="/data/deepinfer-h20/ref") + parser.add_argument("--moe_layers", type=int, default=2, help="dump MoE internals for this many layers") args = parser.parse_args() ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x] @@ -77,10 +78,36 @@ def main() -> None: def embed_hook(_module, _inputs, output): captured["embed"] = output[0, -1].float().cpu().numpy() + def make_gate_hook(layer_id: int): + def hook(_module, _inputs, output): + # Gate.forward returns (weights [n, topk], indices [n, topk]). + weights, indices = output + captured[f"layer{layer_id}_topk_idx"] = indices[-1].to(torch.int32).cpu().numpy() + captured[f"layer{layer_id}_topk_w"] = weights[-1].float().cpu().numpy() + + return hook + + def make_shared_hook(layer_id: int): + def hook(_module, _inputs, output): + captured[f"layer{layer_id}_shared_out"] = output.view(-1, output.shape[-1])[-1].float().cpu().numpy() + + return hook + + def make_moe_hook(layer_id: int): + def hook(_module, _inputs, output): + # The MoE output already includes the shared expert, matching our `_combined`. + captured[f"layer{layer_id}_combined"] = output.view(-1, output.shape[-1])[-1].float().cpu().numpy() + + return hook + model.embed.register_forward_hook(embed_hook) for layer_id, layer in enumerate(model.layers): layer.register_forward_hook(make_hook(layer_id)) layer.attn.register_forward_hook(make_attn_hook(layer_id)) + if layer_id < args.moe_layers: + layer.ffn.gate.register_forward_hook(make_gate_hook(layer_id)) + layer.ffn.shared_experts.register_forward_hook(make_shared_hook(layer_id)) + layer.ffn.register_forward_hook(make_moe_hook(layer_id)) tokens = torch.tensor([ids], dtype=torch.long) with torch.inference_mode():