diff --git a/H20-DEBUG-PLAN.md b/H20-DEBUG-PLAN.md index c307a69..dfbd227 100644 --- a/H20-DEBUG-PLAN.md +++ b/H20-DEBUG-PLAN.md @@ -152,13 +152,109 @@ x = _unpack_fp4(packed) * sf.to(torch.float32).repeat_interleave(FP4_BLOCK, dim= > 后续跟踪（P3.6 起）已迁移到 `H20-PORTING-JOURNAL.md` 第 6 节，那里有完整的 > 漂移量化数据、建议顺序和前置条件。本文件的 P4 / P5 保留为方法清单。 -## P4 decode 专项（仅当 prefill 对、decode 错） +## P4 decode 专项 —— 已确认 decode 是 `**` 的来源 -- [ ] P4.1 KV cache 写入后立刻读回，比 bf16 -- [ ] P4.2 同一位置，decode 路径 vs prefill 路径重算的 logits 差异 -- [ ] P4.3 `_decode_p_sum_sm90` 的数值 vs prefill 的 `fp8_fp4_mqa_logits` +### 为什么锁定 decode ---- +P3 之后重跑对拍（`H20-DRIFT-ANALYSIS.md` 有完整数据）： + +| | top-1 | margin(6127 北京 − 666 `**`) | +|---|---|---| +| 官方，14 token prefill | 6127 | **+3.112** | +| 我方，14 token prefill | **6127** ✅ | **+0.256** | + +prefill 选中了正确的 token，误差只把 margin 压缩了 12 倍。但端到端跑 +（11 token 起，增量 decode 到第 14 个位置）给的是 666。 + +**同一个 14-token 上下文，一次 prefill 给 6127，增量 decode 给 666。** +所以剩余漂移在 decode 路径，而且 margin 只剩 0.256，decode 稍微再加一点误差就会翻。 + +### 判据（先定死，避免事后解释） + +设第 14 个位置上： + +- `A` = 我方 prefill 路径的 logits +- `B` = 我方 decode 路径的 logits +- `R` = 官方 decode 路径的 logits + +| 量 | 含义 | 阈值 | +|---|---|---| +| `rel(A, B)` | 我方两条路径自身是否一致 | > 3e-2 即为 decode 自身缺陷 | +| `rel(A, R)` | prefill 残余误差 | 已知约 3e-2（量化地板） | +| `rel(B, R)` | decode 总误差 | 明显大于 `rel(A,R)` 说明 decode 额外引入误差 | + +- 若 `rel(A,B)` 很小而 `rel(B,R)` 大 → decode 没有额外缺陷，是量化误差累积，`**` 属可接受 +- 若 `rel(A,B)` 大 → **decode 有 bug**，按 P4.3 → P4.4 往下查 + +### 步骤 + +- [x] **P4.1 同位置 prefill vs decode 对拍** —— **完成，确认 decode 有 bug** + + 四条路径走到同一个 14-token 上下文 + `[0,128803,2056,768,17607,2793,3139,12108,1148,128804,128822,17607,2793,3139]`： + + | 路径 | 下一个 token | 结果 | + |---|---|---| + | 官方 prefill（一次喂 14 个） | 6127 = 北京 | ✅ | + | 官方 decode（喂 11 + 逐 token forward） | 6127 = 北京 | ✅ | + | **我方 prefill（一次喂 14 个）** | **6127 = 北京** | ✅ | + | **我方 decode（喂 11 + 增量）** | **666 = `**`** | ❌ | + + 前两条来自 `ref_next_token.py`（prefill 一次 / `--force_ids` 增量），后两条来自 + `h20_ask.py`：`--ids <14个> --max_tokens 1` 给 `6127`，`--ids <11个> --max_tokens 7` + 的第 4 个 token 是 `666`。 + + **结论：prefill 路径正确，decode 路径错误，差额只由 decode 引入。** + 这把 P3 之后「margin 被压缩到 0.256，decode 再推一把就翻」的推断变成了实证。 + + 补充：官方自己 prefill 与 decode 的 margin 也不同（3.112 vs 1.638），所以 + 「两条路径不完全相等」是框架固有特性（KV cache 量化往返），但**方向不该翻**。 + - 路径 1：直接 14 token prefill（`--ids <14个>`，dump `logits_prompt`）—— 已有 + - 路径 2：11 token prefill + 3 步贪心 decode（贪心本来就给对了 + 17607/2793/3139，不需要强制喂 token），dump 第 3 步 decode 的 logits + - 参考侧同样跑增量：`ref_next_token.py` 逐步喂 11→12→13→14 个 token， + 每一步的记录就是官方 decode 路径的 logits + 产出：三条曲线的 margin 与 top-1 逐步对照。 + +- [ ] **P4.2 decode 逐步误差曲线** + 从第 1 步到第 8 步，每步记 `margin`、`top-1 是否与官方一致`、`rel(B_t, R_t)`。 + 用来判断误差是「一开始就有」还是「某一步突然放大」——后者意味着踩到了某个 + batch/位置相关的边界（例如 page 边界、SWA window 回绕）。 + +- [ ] **P4.3 KV cache 写入/读回一致性**（最高嫌疑之一） + 写入已知向量后立刻读回，与 bf16 原值比。decode 依赖 cache 能正确回读， + 而 prefill 的 KV 是当场算的、几乎不读 cache。若这一步就错，P4.1 的差距可解释。 + 同时验证 index KV 的字节布局（`FP8_FLOAT32_SF_INDEX_KV_FORMAT`：128 fp8 + 1 f32 + = 132 B/token），这是 P3.9 遗留项，**从未做过 round-trip 验证**。 + +- [ ] **P4.4 `_decode_p_sum_sm90` vs prefill 的 indexer 分数**（最高嫌疑之二） + decode 的 indexer 走 torch 兜底（SM90 paged kernel 要求 64-token page，cache 是 128）， + prefill 走 TileLang 的 `fp8_fp4_mqa_logits`。**两条路径从未对拍过。** + 给同一个 KV cache、同一个位置，比较两边算出的 top-k 索引集合与分数。 + P3 已证明「路由 top-k 一致」时误差仍存在，所以这里要查的是**稀疏 attention 选哪些 + KV**，而不是 MoE 选哪些专家。 + +- [ ] **P4.5 采样器 / logits processor 排除** + 若 P4.3、P4.4 都干净，则查采样侧：`DEEPINFER_SAMPLER_FORCE_TOPK=4096` + 会不会在 topk=4096 时改变数值，以及 logits 是否被额外 mask（blacklist、EOS 抑制）。 + +### P4.1 的工具坑（都已修，但 dump 仍建议配合黑盒交叉验证） + +| 坑 | 现象 | 修法 | +|---|---|---| +| 多 DP 进程写同一文件 | 8 个 engine 都是单进程组，`RANK` 全为 0，互相覆盖 | 文件名带 `_rank_pid

` | +| decode 的 batch 有 padding 行 | 取 `logits[-1]` 读到占位 token 的垃圾（恒定 `top1=69`） | 取第 0 行（第一个真实请求） | +| follower rank 的 logits 是占位值 | 同一向量被重复记录几十次，淹没真实步 | 判据 `driver_inputs is not None` | +| profile forward 用合成 token | 抓到 profile 的数据 | `debug_dump.profiling()` 挂起 | + +即便都修了，逐 step dump 与端到端 token 序列仍无法可靠对齐（`max_tokens=7` 的 run 里 +出现 75 次 sampler 调用）。所以 P4.1 的结论**以黑盒四格对照为准**，dump 只作辅助。 + +### 已备好的工具 + +- `debug_dump.record(..., once=True)` + `logits_prompt` / `logits_decode` 两个键 +- `debug_dump.profiling()`：profile forward 期间整体挂起，避免抓到合成 token +- `ref_next_token.py`：吃 token ids，给官方 greedy next token + top5 ## P5 logprob 异常 diff --git a/deepinfer/h20_ops/debug_dump.py b/deepinfer/h20_ops/debug_dump.py index 14b96ee..12f96a9 100644 --- a/deepinfer/h20_ops/debug_dump.py +++ b/deepinfer/h20_ops/debug_dump.py @@ -25,6 +25,12 @@ _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) +# Per-tag call counter for `record_step`. +_STEPS: dict[str, int] = {} +# Number of rows in the most recent logits tensor, used to skip warmup passes. +_ROWS = 0 +# A decode step that serves real requests has at most this many rows. Anything wider is warmup. +_REAL_ROWS = int(os.getenv("DEEPINFER_DUMP_MAX_ROWS", "1") or 1) _SUSPENDED = False @@ -82,6 +88,35 @@ def record(tag: str, tensor: torch.Tensor, once: bool = False) -> None: _CAPTURED[tag] = row.detach().float().cpu().numpy() +def set_rows(rows: int) -> None: + """Tell the dump how many rows the sampler just produced, so warmup can be filtered out.""" + global _ROWS + _ROWS = rows + + +def record_step(tag: str, tensor: torch.Tensor) -> None: + """Append `tensor`'s last row under `tag_`, numbering the calls. + + Decode needs every step kept, not just the first: the point of the decode comparison is to + watch the margin evolve and find the step where it collapses. + """ + if not _DUMP_DIR: + return + # Decode with one request produces exactly `_REAL_ROWS` rows; warmup batches are larger. + if _ROWS > _REAL_ROWS: + return + index = _STEPS.get(tag, 0) + _STEPS[tag] = index + 1 + # Row 0, not the last row: during decode the batch is padded (CUDA graph slots, or just + # `max_num_seqs`), and the padding rows hold logits for placeholder tokens. Row 0 is the + # first real request's current token, which is the one to compare against a batch-1 reference. + row = tensor[0] if tensor.dim() >= 2 else tensor + from deepinfer.utils import allow_cuda_sync + + with allow_cuda_sync(): + _CAPTURED[f"{tag}_{index}"] = 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: @@ -109,13 +144,34 @@ def _want_moe(tag: str) -> bool: return True +def _rank_suffix() -> str: + """`_rank` when running distributed, so ranks do not overwrite each other's dump. + + With data parallelism every rank registers `atexit(flush)` and writes the same filename; the + last one to exit wins, which silently mixes one rank's tensors with another's. Only rank 0 + holds the real prompt anyway, so the suffix makes that explicit instead of accidental. + """ + rank = os.getenv("RANK") + if rank is None: + try: + import torch.distributed as dist + + if dist.is_initialized(): + rank = str(dist.get_rank()) + except Exception: # noqa: BLE001 - torch may not be usable at exit + rank = None + # With data parallelism each engine is its own single-process group, so every rank reports + # RANK=0 and they all target the same file. The pid disambiguates them. + return f"_rank{rank if rank is not None else 'x'}_pid{os.getpid()}" + + 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" + path = out / f"deepinfer_layers{_rank_suffix()}.npz" np.savez(path, **_CAPTURED) print(f"[dump] wrote {len(_CAPTURED)} tensors to {path}", flush=True) _CAPTURED.clear() diff --git a/deepinfer/model_executor/layers/sampler.py b/deepinfer/model_executor/layers/sampler.py index dc3adac..d44026f 100644 --- a/deepinfer/model_executor/layers/sampler.py +++ b/deepinfer/model_executor/layers/sampler.py @@ -141,12 +141,19 @@ class Sampler(nn.Module): # CUDA Graph model_output retains graph-batch padding; gather trims hidden_states. hidden_states = gather(hidden_states, selected_token_indices) logits = get_logits(hidden_states, embedding, self.vocab_size) - if _debug_dump.enabled(): + # `driver_inputs is None` means this rank only participates in the embedding all-gather and + # never sees the real result, so its logits are placeholders -- identical on every call. + # Recording those buries the real steps under dozens of copies of the same vector. + if _debug_dump.enabled() and driver_inputs is not None: # Capture before any masking/penalty so the numbers are comparable with the reference, # which dumps the raw head output as well. Named per phase: with `max_tokens=1` both a # prompt step and a decode step run, and only the former is comparable with a # prefill-only reference dump. - _debug_dump.record("logits_prompt" if is_prompt else "logits_decode", logits, once=True) + if is_prompt: + _debug_dump.record("logits_prompt", logits, once=True) + else: + # Every decode step, so the comparison can watch the margin collapse step by step. + _debug_dump.record_step("logits_decode", logits) del hidden_states if driver_inputs is None: diff --git a/scripts/ref_next_token.py b/scripts/ref_next_token.py index 25ee053..1d93de9 100644 --- a/scripts/ref_next_token.py +++ b/scripts/ref_next_token.py @@ -14,6 +14,7 @@ import json import os import sys +import numpy as np import torch import torch.distributed as dist @@ -30,6 +31,15 @@ def main() -> None: parser.add_argument("--ids", required=True, help="comma-separated prompt token ids") parser.add_argument("--out", default="/data/deepinfer-h20/ref/next_token.json") parser.add_argument("--topk", type=int, default=5) + parser.add_argument( + "--force_ids", + default="", + help="Comma-separated ids fed one at a time after the prompt, i.e. incremental decode. " + "This mirrors `generate.py`: prefill the prompt, then forward a single token per step with " + "prev_pos advancing, so the resulting logits come from the decode path rather than from " + "recomputing the whole prefix.", + ) + parser.add_argument("--logits_out", default="", help="optional .npz to store every step's full logits") args = parser.parse_args() ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x] @@ -59,21 +69,45 @@ def main() -> None: load_model(model, os.path.join(args.ckpt_path, f"model{rank}-mp{world_size}.safetensors")) torch.set_default_device("cuda") + forced = [int(x) for x in args.force_ids.replace(" ", "").split(",") if x] + tokens = torch.tensor([ids], dtype=torch.long) + steps: list[tuple[str, torch.Tensor]] = [] with torch.inference_mode(): # `Transformer.forward` returns (output_ids, logits, main_hidden). output_ids, logits, _ = model.forward(tokens, 0) + steps.append(("prefill", logits[0].float().clone())) - row = logits[0, -1].float() + # Same loop shape as `generate.py`: one token per forward, prev_pos advancing. The token + # fed in is the ground truth, which is exactly what the engine does when it accepts a + # greedy token it already agrees with. + prev_pos = tokens.shape[1] + for token in forced: + output_ids, logits, _ = model.forward(torch.tensor([[token]], dtype=torch.long), prev_pos) + steps.append((f"decode_{prev_pos}", logits[0].float().clone())) + prev_pos += 1 + + row = steps[-1][1] top = torch.topk(row, args.topk) - assert int(top.indices[0]) == int(output_ids[0, -1]), "argmax disagrees with the sampled id" payload = { "prompt_ids": ids, + "force_ids": forced, "next_token": int(top.indices[0]), "top_tokens": [int(x) for x in top.indices], "top_logits": [round(float(x), 4) for x in top.values], + "steps": [ + { + "name": name, + "top_tokens": [int(x) for x in torch.topk(v, args.topk).indices], + "top_logits": [round(float(x), 4) for x in torch.topk(v, args.topk).values], + } + for name, v in steps + ], } + if args.logits_out and rank == 0: + os.makedirs(os.path.dirname(args.logits_out), exist_ok=True) + np.savez(args.logits_out, **{f"logits_{name}": v.cpu().numpy() for name, v in steps}) if rank == 0: print(json.dumps(payload, ensure_ascii=False), flush=True) os.makedirs(os.path.dirname(args.out), exist_ok=True)