mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
3c15952e76
Narrows the SM90 numerical divergence to layer 0's FFN. The sparse MLA attention kernel is not the problem: at layer 0 its output matches the reference with cos=0.99897 / rel_err=4.7e-2, while the full layer output is already at 12.7% and the error compounds to cos=0.865 by layer 11, which finally flips the top-1 token. Tooling - `h20_ops/debug_dump.py`: opt-in capture (`DEEPINFER_DUMP_LAYERS=<dir>`), last position only, reduced over `hc_mult`, written to `.npz` at exit. `DEEPINFER_DUMP_TOKENS=<n>` gates on the real prompt length, because the engine runs a dummy profile forward with a different token count before it allocates the KV cache and that dump would otherwise be mistaken for the real one. - `scripts/ref_layer_dump.py`: the same anchors from the reference model via forward hooks (block output, attention sub-layer output, embedding), plus the final logits and greedy token. - `scripts/compare_layers.py`: cosine similarity, relative L2 error and norm ratio per layer, reporting the first layer above a tolerance. Two traps this had to fix before the numbers meant anything - `decoder_layer` returns the *sub-layer* output; the residual stream lives in `residual` (`HCLazyPostData`, hc_mult=4). Comparing without merging them through `hc_post` produced a bogus cos of 0.02 at every layer. - `DEEPINFER_DEBUG_NAN` and the dump both need a host copy inside the forward, so both take `allow_cuda_sync()` rather than fighting the engine's sync guard. Also documents the reproduction case and the findings in H20-DEBUG-PLAN.md.
68 lines
2.2 KiB
Python
68 lines
2.2 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] = {}
|
|
|
|
|
|
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 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 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)
|