From 3c15952e765b91dcc15dd0004f28fe8723b39321 Mon Sep 17 00:00:00 2001 From: loopyt Date: Wed, 21 Oct 2026 16:16:53 +0800 Subject: [PATCH] debug: add layer-wise dump + reference comparison for H20 bring-up 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=`), last position only, reduced over `hc_mult`, written to `.npz` at exit. `DEEPINFER_DUMP_TOKENS=` 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. --- H20-DEBUG-PLAN.md | 42 +++++-- deepinfer/h20_ops/debug_dump.py | 67 +++++++++++ .../model_executor/models/decoder_layer.py | 10 ++ deepinfer/model_executor/models/deepseek.py | 3 + scripts/compare_layers.py | 85 ++++++++++++++ scripts/ref_layer_dump.py | 105 ++++++++++++++++++ 6 files changed, 302 insertions(+), 10 deletions(-) create mode 100644 deepinfer/h20_ops/debug_dump.py create mode 100644 scripts/compare_layers.py create mode 100644 scripts/ref_layer_dump.py diff --git a/H20-DEBUG-PLAN.md b/H20-DEBUG-PLAN.md index 8d80a65..f26afc6 100644 --- a/H20-DEBUG-PLAN.md +++ b/H20-DEBUG-PLAN.md @@ -55,21 +55,43 @@ length : max_new_tokens = 6 --- -## P2 分层对拍（下一步，进行中） +## P2 分层对拍 —— 已完成，定位到第 0 层的 FFN 阶段 -输入固定为 14-token 那条（官方 `6127` / 我方 `666`），这是最短的可复现分歧。 +**方法**：`scripts/ref_layer_dump.py`（官方，forward hook）+ `DEEPINFER_DUMP_LAYERS`（我方， +`h20_ops/debug_dump.py`），两侧都取最后位置、都按 hc_mult 求均值，用 +`scripts/compare_layers.py` 对比。输入固定为那条 14-token ids。 -- [ ] P2.1 官方 `model.py` 挂 forward hook，dump 每层 `h`（最后一行）到 `.npy` -- [ ] P2.2 DeepInfer 侧加同样 dump（`DEEPINFER_DEBUG_DUMP=`） -- [ ] P2.3 逐层算 `rel_err = ||x_mine - x_ref|| / ||x_ref||`，找**首个**超 1e-2 的层 +**踩过的两个坑（重要，否则会得出错误结论）**： -**判据**：拿到具体层号。若是第 0 层就已经偏，说明输入侧（embedding / engram）有问题； -若是中间层开始偏，则集中查该层的 attention 或 MoE。 +1. 我方 `decoder_layer` 返回的 `hidden_states` 只是**子层输出**（2D），完整残差流在 `residual` + （`HCLazyPostData`，hc_mult=4）里。必须 `hc_post(hidden_states, residual)` 之后再比， + 否则第一版对比得到的 cos≈0.02 是假的。 +2. 引擎在分配 KV cache 前会跑一次 **dummy profile forward**，token 数与真实 prompt 不同。 + 必须用 `DEEPINFER_DUMP_TOKENS=` 门控，否则 dump 到的是 dummy 输入。 -**注意**：engram 在第 1~2 层就有贡献，而 engram 表是我自己 gather 的，属于高嫌疑区， -所以「第 0/1 层就偏」的可能性不小。 +**结果（14-token 输入，cos / rel_err / 范数比）**： ---- +| 位置 | cos | rel_err | ‖mine‖/‖ref‖ | +|---|---|---|---| +| layer0 attention 子层输出 | **0.99897** | **4.71e-02** | 0.986 | +| layer0 整层输出 | 0.99190 | 1.27e-01 | 1.000 | +| layer1 整层 | 0.97707 | 2.29e-01 | 1.062 | +| layer1 attention 子层 | 0.99324 | 1.21e-01 | 0.959 | +| layer2 | 0.96411 | 2.76e-01 | 1.039 | +| layer5 | 0.93275 | 3.61e-01 | 0.943 | +| layer11 | 0.86515 | 5.28e-01 | 1.030 | + +**结论**： + +1. **稀疏 MLA attention 本身是对的**：第 0 层 attention 子层 cos=0.999，误差 4.7%， + 在 FP8/量化口径下属于合理量级。之前怀疑的 TileLang kernel 可以排除。 +2. **误差主要在第 0 层的 FFN/MoE 阶段引入**：attention 4.7% → 整层 12.7%， + 即 FFN 贡献了主要偏差。 +3. **误差单调累积**，到第 11 层 cos 已降到 0.865，最终把 top-1 从「北京」翻成「**」。 + +**下一步**：在 layer0 的 FFN 内再切一刀 —— 分别 dump `routed_experts_forward` 的输出、 +shared expert 输出、以及 swiglu/route 权重，看是路由（top-6 选中了不同专家）还是 +专家计算本身偏。`engram` 在第 1/14 层，本轮已可排除其为**首个**偏离点。 ## P3 层内细分（按嫌疑排序） diff --git a/deepinfer/h20_ops/debug_dump.py b/deepinfer/h20_ops/debug_dump.py new file mode 100644 index 0000000..53bf55f --- /dev/null +++ b/deepinfer/h20_ops/debug_dump.py @@ -0,0 +1,67 @@ +"""Opt-in tensor dumps for H20 bring-up comparisons. + +Enabled with `DEEPINFER_DUMP_LAYERS=`. Each call appends one row to +`/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) diff --git a/deepinfer/model_executor/models/decoder_layer.py b/deepinfer/model_executor/models/decoder_layer.py index 6773756..4b939a2 100644 --- a/deepinfer/model_executor/models/decoder_layer.py +++ b/deepinfer/model_executor/models/decoder_layer.py @@ -5,6 +5,7 @@ import torch from torch import nn from deepinfer import AttrDict, envs +from deepinfer.h20_ops import debug_dump as _debug_dump from deepinfer.logger import init_logger from deepinfer.model_executor.input_metadata import InputMetadata from deepinfer.model_executor.input_metadata.yoco import YocoMetadataMixin @@ -179,6 +180,7 @@ class DeepseekDecoderLayer(HCLayerMixin): **attn_kwargs, ) _check_finite(f"layer{self.layer_idx}.attn_out", hidden_states) + _debug_dump.record(f"layer{self.layer_idx}_attn", hidden_states) hidden_states, hidden_states_bf16, residual = self.ffn_norm_forward(hidden_states, residual) mlp_timer_scale = self.num_prefill_mlp_layers if input_metadata.is_prompt else self.num_mlp_layers @@ -193,6 +195,14 @@ class DeepseekDecoderLayer(HCLayerMixin): ) _check_finite(f"layer{self.layer_idx}.mlp_out", hidden_states) + if _debug_dump.enabled(): + # `hidden_states` is only this layer's sub-layer output; the residual stream lives in + # `residual` and is merged by `hc_post` ([n, hc_mult, hidden]). The reference dumps the + # merged stream, so merge it here too -- otherwise the two are not comparable. + _debug_dump.record( + f"layer{self.layer_idx}", + HCLayerMixin.hc_post(hidden_states, residual) if self.hc_enabled else hidden_states, + ) return hidden_states, residual diff --git a/deepinfer/model_executor/models/deepseek.py b/deepinfer/model_executor/models/deepseek.py index ca31b1c..36eb614 100644 --- a/deepinfer/model_executor/models/deepseek.py +++ b/deepinfer/model_executor/models/deepseek.py @@ -10,6 +10,7 @@ from torch import nn from deepinfer import 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 @@ -160,6 +161,7 @@ class DeepseekModel(nn.Module): ) -> CausalLMForwardOutput: if hidden_states is None: hidden_states = self.embed_tokens(input_metadata.input_ids) + _debug_dump.record("embed", hidden_states) residual = None # 跨层持有的运行时值 (shared_hidden / pre_rope_kv / topk): 逐层传同一引用，由各层原地写入。 @@ -237,6 +239,7 @@ class DeepseekModel(nn.Module): hidden_states = hidden_states.unflatten(-1, (self.hc_head.hc_mult, -1)) hidden_states = self.hc_head(hidden_states) hidden_states = self.output_norm(hidden_states) + _debug_dump.record("final_norm", hidden_states) if input_metadata.is_prompt and PG.context_parallel_group().size() > 1: with userbuffer.cp_session(): diff --git a/scripts/compare_layers.py b/scripts/compare_layers.py new file mode 100644 index 0000000..89d1bd7 --- /dev/null +++ b/scripts/compare_layers.py @@ -0,0 +1,85 @@ +"""Diff per-layer hidden states between DeepInfer and the reference implementation. + +Both dumps keep only the final position of the prompt, reduced over the `hc_mult` copies, so +the comparison is one vector per layer. Reports cosine similarity, relative L2 error and the +ratio of norms, which together distinguish "slightly off" from "different direction". + + python scripts/compare_layers.py --ref ref_layers.npz --mine deepinfer_layers.npz +""" + +import argparse + +import numpy as np + + +def rel_err(a: np.ndarray, b: np.ndarray) -> float: + return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-12)) + + +def cosine(a: np.ndarray, b: np.ndarray) -> float: + na, nb = np.linalg.norm(a), np.linalg.norm(b) + if na == 0 or nb == 0: + return float("nan") + return float(np.dot(a, b) / (na * nb)) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--ref", required=True) + parser.add_argument("--mine", required=True) + parser.add_argument("--tol", type=float, default=1e-2, help="first layer above this rel_err is the culprit") + args = parser.parse_args() + + ref = np.load(args.ref) + mine = np.load(args.mine) + + layers = sorted(int(k[5:]) for k in ref.files if k.startswith("layer") and k[5:].isdigit()) + mine_layers = {int(k[5:]) for k in mine.files if k.startswith("layer") and k[5:].isdigit()} + missing = [i for i in layers if i not in mine_layers] + if missing: + print(f"WARNING: no dump from DeepInfer for layers {missing}") + + 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) + if r.shape == m.shape: + print( + f"{key:>12} 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() + print(f"{'layer':>5} {'cos':>9} {'rel_err':>9} {'|mine|/|ref|':>12}") + first_bad = None + for i in layers: + if i not in mine_layers: + continue + r, m = ref[f"layer{i}"].astype(np.float64), mine[f"layer{i}"].astype(np.float64) + if r.shape != m.shape: + print(f"{i:>5} shape mismatch ref={r.shape} mine={m.shape}") + continue + err, cos = rel_err(m, r), cosine(m, r) + ratio = float(np.linalg.norm(m) / (np.linalg.norm(r) + 1e-12)) + flag = "" + if first_bad is None and err > args.tol: + first_bad = i + flag = " <-- first > tol" + print(f"{i:>5} {cos:9.5f} {err:9.2e} {ratio:12.4f}{flag}") + + ref_tok = int(ref["next_token"][0]) if "next_token" in ref.files else None + print(f"\nreference next token: {ref_tok}") + + if "final_row_logits" in ref.files: + rl = ref["final_row_logits"].astype(np.float64) + top = np.argsort(-rl)[:5] + 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() + if first_bad is None: + print(f"RESULT: every dumped layer is within rel_err <= {args.tol}; divergence is downstream") + else: + print(f"RESULT: first layer above rel_err {args.tol} is layer {first_bad}") + + +if __name__ == "__main__": + main() diff --git a/scripts/ref_layer_dump.py b/scripts/ref_layer_dump.py new file mode 100644 index 0000000..6bd875f --- /dev/null +++ b/scripts/ref_layer_dump.py @@ -0,0 +1,105 @@ +"""Dump the reference model's per-layer hidden state for one prompt. + +Counterpart to DeepInfer's `DEEPINFER_DUMP_LAYERS`, so the two can be diffed layer by layer. +The reference residual stream is `hc_mult` parallel copies, so each layer is reduced with +`mean(dim=1)` -- the same convention `debug_dump.record` uses -- and only the final position +is kept, because that is the only row generation depends on. + + torchrun --nproc-per-node 8 ref_layer_dump.py --ids 0,128803,... --out /data/deepinfer-h20/ref +""" + +import argparse +import json +import os +import sys +from pathlib import Path + +import numpy as np +import torch +import torch.distributed as dist + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from safetensors.torch import load_model # noqa: E402 +from transformers import AutoTokenizer # noqa: E402 + +from model import ModelArgs, Transformer # noqa: E402 + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--ckpt-path", required=True) + parser.add_argument("--config", required=True) + parser.add_argument("--ids", required=True) + parser.add_argument("--out", default="/data/deepinfer-h20/ref") + args = parser.parse_args() + + ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x] + rank = int(os.getenv("RANK", "0")) + local_rank = int(os.getenv("LOCAL_RANK", "0")) + world_size = int(os.getenv("WORLD_SIZE", "1")) + torch.cuda.set_device(local_rank) + if world_size > 1: + dist.init_process_group("nccl") + + torch.set_default_dtype(torch.bfloat16) + torch.set_num_threads(8) + with open(args.config) as f: + model_args = ModelArgs(**json.load(f)) + model_args.temperature = 0.0 + model_args.max_batch_size = 1 + model_args.max_seq_len = 4096 + + tokenizer = AutoTokenizer.from_pretrained(args.ckpt_path) + with torch.device("cuda"): + model = Transformer(model_args, tokenizer) + load_model(model, os.path.join(args.ckpt_path, f"model{rank}-mp{world_size}.safetensors")) + torch.set_default_device("cuda") + + captured: dict[str, np.ndarray] = {} + + def make_hook(layer_id: int): + def hook(_module, _inputs, output): + # Block.forward returns (x, pre_mix); x is [b, s, hc_mult, dim]. + x = output[0] if isinstance(output, tuple) else output + captured[f"layer{layer_id}"] = x.mean(dim=2)[0, -1].float().cpu().numpy() + + return hook + + def make_attn_hook(layer_id: int): + def hook(_module, _inputs, output): + # The attention sub-layer output is already collapsed: [b, s, dim]. + x = output[0] if isinstance(output, tuple) else output + captured[f"layer{layer_id}_attn"] = x[0, -1].float().cpu().numpy() + + return hook + + def embed_hook(_module, _inputs, output): + captured["embed"] = output[0, -1].float().cpu().numpy() + + 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)) + + tokens = torch.tensor([ids], dtype=torch.long) + with torch.inference_mode(): + output_ids, logits, _ = model.forward(tokens, 0) + + captured["final_row_logits"] = logits[0, -1].float().cpu().numpy() + captured["next_token"] = np.array([int(output_ids[-1])], dtype=np.int32) + + if rank == 0: + out = Path(args.out) + out.mkdir(parents=True, exist_ok=True) + path = out / "ref_layers.npz" + np.savez(path, **captured) + print(f"wrote {len(captured)} tensors to {path}", flush=True) + print(json.dumps({"next_token": int(output_ids[-1]), "layers": len(model.layers)}), flush=True) + + if world_size > 1: + dist.destroy_process_group() + + +if __name__ == "__main__": + main()