mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
7c6fdc7c6e
`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.
100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
"""Opt-in tensor dumps for H20 bring-up comparisons.
|
|
|
|
Enabled with `DEEPINFER_DUMP_LAYERS=<dir>`. Each call appends one row to
|
|
`<dir>/deepinfer_layers.npz`-style storage under a per-tag key, always as the *last* row of
|
|
the batch: the reference implementation only generates from the final position, so that is
|
|
the only row the two sides can be compared on.
|
|
|
|
The dump is deliberately simple (host copy + `numpy.savez` at exit) because it exists to
|
|
answer "which layer diverges first", not to be a profiler.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
_DUMP_DIR = os.getenv("DEEPINFER_DUMP_LAYERS")
|
|
# The engine runs a dummy profile forward before it allocates the KV cache, and that one has a
|
|
# 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:
|
|
return bool(_DUMP_DIR)
|
|
|
|
|
|
def record(tag: str, tensor: torch.Tensor) -> None:
|
|
"""Store the last row of `tensor` under `tag`.
|
|
|
|
A 3D tensor is treated as `[tokens, hc_mult, dim]` and reduced with `mean(dim=1)`, which is
|
|
the closest match to the reference's collapsed residual stream.
|
|
"""
|
|
if not _DUMP_DIR:
|
|
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)
|
|
# Copying to the host is the whole point here, and it happens inside the forward where the
|
|
# engine forbids syncs, so take the sanctioned escape hatch for this one transfer.
|
|
from deepinfer.utils import allow_cuda_sync
|
|
|
|
with allow_cuda_sync():
|
|
_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:
|
|
return
|
|
out = Path(_DUMP_DIR)
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
path = out / "deepinfer_layers.npz"
|
|
np.savez(path, **_CAPTURED)
|
|
print(f"[dump] wrote {len(_CAPTURED)} tensors to {path}", flush=True)
|
|
_CAPTURED.clear()
|
|
|
|
|
|
if _DUMP_DIR:
|
|
import atexit
|
|
|
|
atexit.register(flush)
|