diff --git a/H20-DRIFT-ANALYSIS.md b/H20-DRIFT-ANALYSIS.md new file mode 100644 index 0000000..e81baec --- /dev/null +++ b/H20-DRIFT-ANALYSIS.md @@ -0,0 +1,152 @@ +# 漂移调查：2.7% 不是主因，问题在 decode 路径 + +承接 `H20-PORTING-JOURNAL.md` §6.3 的 P3.6（「改用官方无损重化把 2.7% 降到 0」）。 + +**两个结论**： + +1. **P3.6 在 SM90 上不可行** —— DeepGEMM 在 Hopper 上强制 per-128 scale，per-32 是 SM100 + 专属。2.7% 是硬件地板，不是实现缺陷。 +2. **但 2.7% 不是 `**` 的成因** —— prefill 路径的 top-1 与官方一致，翻转发生在 **decode**。 + +第 2 点是本次最有价值的发现，它把问题从「量化精度无解」重新定位到一个具体的、可能可修的 +bug 上。 + +--- + +## 一、P3.6 不可行：DeepGEMM 的 SM90 硬约束 + +### 证据 + +`csrc/apis/layout.hpp:100`（DeepGEMM 2.8.2），`check_sf_layout` 里唯一的架构断言： + +```cpp +DG_HOST_ASSERT((arch_major == 9 and gran_k == 128 and k_alignment == 128) or + (arch_major == 10 and (gran_k == 32 or gran_k == 128) and k_alignment % 128 == 0)); +``` + +- `arch_major == 9`（Hopper）→ **`gran_k` 必须 128** +- `arch_major == 10`（Blackwell）→ 32 或 128 都行 + +同文件 38-51 行的 `transform_sf_into_required_layout` 印证同一件事 —— SM90 只有两个分支， +都要求 `gran_k == 128` 且 `sf` 是 **fp32**： + +```cpp +// (FP32, 1, 128) on SM90 +if (sf.scalar_type() == kFloat and gran_mn == 1 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast)) + return get_mn_major_tma_aligned_tensor(sf); + +// (FP32, 128, 128) on SM90 +if (sf.scalar_type() == kFloat and gran_mn == 128 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast)) + return check_sf_layout(...); + +// (FP32, x, gran_k) on SM100 —— 只有这里才出现 packed ue8m0 +if (sf.scalar_type() == kFloat and (gran_k == 32 or gran_k == 128) and arch_major == 10) { ... } +``` + +官方 `cast_e2m1fn_to_e4m3fn` 产出的是 **e8m0 per-32**。SM90 要 **fp32 per-128**。 +两者不兼容，无中间等价物。 + +### 逐条排除的降损方案 + +| 想法 | 结论 | +|---|---| +| 换 per-32 recipe | ❌ 上述断言，SM90 不允许 | +| scale 用 e8m0 省位宽 | ❌ SM90 分支要求 fp32，e8m0 触发 `sf_dtype` 断言 | +| 专家权重用 BF16 绕开量化 | ❌ 显存 78.5 GiB/卡 → 约 157 GiB/卡 | +| 让 per-128 的 scale 更精确 | ❌ 已是 `amax/448`，per-128 网格上的最优表示 | + +**所以 2.7% 是「用 SM90 的 GEMM 跑一个按 MXFP8 per-32 导出的 checkpoint」的固有代价。** +参照量级：dense 权重（`requant_dense_fp8`，用法本来就对）也有 4.7e-2，同一量级。 + +--- + +## 二、修复确实有效：layer0 误差降了 4.4 倍 + +专家 scale 修复（提交 `7c6fdc7`）之后重跑逐层对拍，14-token 输入： + +| 层 | 修复前 cos / rel_err | **修复后** cos / rel_err | +|---|---|---| +| layer 0 | 0.99190 / 1.27e-1 | **0.99959 / 2.88e-2** | +| layer 1 | 0.97707 / 2.29e-1 | **0.99908 / 4.45e-2** | +| layer 2 | 0.96411 / 2.76e-1 | **0.99852 / 5.75e-2** | +| layer 5 | 0.93275 / 3.61e-1 | **0.99834 / 7.16e-2** | +| layer 10 | — | 0.99310 / 1.22e-1 | +| layer 20 | — | 0.98709 / 1.61e-1 | +| layer 39 | 0.86515 / 5.28e-1（旧 layer11 数据） | **0.97972 / 2.00e-1** | + +这是 P3.7 想验证的东西，结论明确：误差曲线整体压平了一个数量级。 + +--- + +## 三、关键发现：prefill 的 top-1 是对的 + +之前一直没能比较最终 logits，原因是 `ref_layer_dump.py` 写的是 `logits[0, -1]`， +而 `MainPredictHeadHC` 在 `full_logits=False` 时已经压掉了 seq 维（`x = x[:, -1]`）， +所以 logits 是 `[b, vocab]` —— 那个索引会得到一个 **0 维标量**。修成 `logits[0]` 之后 +才拿到可以比较的数据。 + +14-token 输入，两侧最终 logits： + +| | 官方 | 移植后 | +|---|---|---| +| cos | — | **0.98460** | +| top-1 | **6127（北京）** | **6127（北京）** ✅ | +| logit(6127) | 39.343 | 36.047 | +| logit(666 = `**`) | 36.232 | 35.791 | +| **margin** | **+3.112** | **+0.256** | + +**prefill 路径的 top-1 与官方一致。** 残余误差的表现为 margin 从 3.11 压到 0.26（约 12 倍）， +但方向没翻。 + +### 所以 `**` 从哪来 + +端到端（11-token 输入，greedy）实际输出： + +``` +[17607, 2793, 3139, 666, 6127, 666, 320, 1] → 中国的首都是**北京**。 + ↑ 前 3 个与官方逐 token 一致 ↑ 第 4 步翻成 666 +``` + +而同样这 14 个 token **一次性 prefill** 时，top-1 是 6127（正确）。 + +**同一个 14-token 上下文，prefill 给 6127，增量 decode 给 666。** + +这就是 `**` 的来源：不是量化精度，是 **decode 路径**引入了额外误差，把本来只有 0.26 的 +margin 彻底翻转。 + +--- + +## 四、下一步：P4 decode 专项 + +定位已经收窄到 decode，建议按这个顺序： + +- [ ] **P4.1 同位置 prefill vs decode 对拍**（最高优先） + 输入 14 token：① 一次 prefill，dump 最后位置 logits（已有：margin +0.256）； + ② 先喂 11 token，再用已确认正确的 3 个 token（17607/2793/3139）逐步 decode， + dump 第 4 步的 logits。两者不同即确认 decode 引入偏差，且可以直接量出偏差大小。 + +- [ ] **P4.2 decode 逐步误差曲线** + 从第 1 步到第 8 步，每步记录 top-1 和 margin，找出误差开始放大的步数。 + +- [ ] **P4.3 KV cache 写入/读回一致性** + 写入已知向量后读回，与 bf16 原值比。这是 P3.9 遗留项，也是最可能的嫌疑人： + decode 用 `_decode_p_sum_sm90` 走 torch 兜底算 indexer 分数，而 prefill 走 + TileLang kernel —— 两条路径从未对拍过。 + +- [ ] **P4.4 `_decode_p_sum_sm90` vs prefill 的 MQA logits** + 同一个 KV cache、同一个位置，两条 indexer 路径的分数差。 + +判据：如果 P4.1 的差距 < 1e-3，说明 decode 无误，`**` 另有来源（回到采样器/logits +processor 查）；如果差距明显，就从 P4.3/P4.4 里找。 + +--- + +## 五、两个工具修正（否则上面的数据拿不到） + +1. **`ref_layer_dump.py` 的 logits 索引**：`logits[0, -1]` → `logits[0]`。之前所有 + dump 里的 `final_row_logits` 都是 0 维标量，等于没有。 +2. **profile run 污染 dump**：`execute_model_profile_run` 会走完整的 forward + sampler， + 用的是合成 token。层 dump 靠 token 数门控侥幸过滤掉了，但 logits 的行数是「选中 token 数」 + 而不是 prompt 长度，门控不适用，于是抓到的是 profile 的垃圾（第一版得到 cos = 0.34）。 + 现在用 `debug_dump.profiling()` 上下文在 profile 期间整体挂起，并且把 logits 拆成 + `logits_prompt` / `logits_decode` 两个键 —— 只有前者能和 prefill-only 的参考 dump 比较。 diff --git a/deepinfer/h20_ops/debug_dump.py b/deepinfer/h20_ops/debug_dump.py index 6629906..14b96ee 100644 --- a/deepinfer/h20_ops/debug_dump.py +++ b/deepinfer/h20_ops/debug_dump.py @@ -27,19 +27,47 @@ _CAPTURED: dict[str, np.ndarray] = {} _MOE_LAYERS = int(os.getenv("DEEPINFER_DUMP_MOE_LAYERS", "2") or 0) +_SUSPENDED = False + + def enabled() -> bool: - return bool(_DUMP_DIR) + return bool(_DUMP_DIR) and not _SUSPENDED -def record(tag: str, tensor: torch.Tensor) -> None: +class profiling: + """Suspend capture while the engine runs its dummy profile forward. + + `execute_model_profile_run` goes through the same forward and sampler as a real step, with + synthetic tokens, so anything captured there is garbage that later masking cannot detect. The + layer dumps happen to be filtered by their token-count gate; the logits are not, because their + row count is the number of *selected* tokens rather than the prompt length. + """ + + def __enter__(self) -> None: + global _SUSPENDED + self._previous = _SUSPENDED + _SUSPENDED = True + + def __exit__(self, *_exc: object) -> None: + global _SUSPENDED + _SUSPENDED = self._previous + + +def record(tag: str, tensor: torch.Tensor, once: bool = False) -> 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. + + `once=True` keeps the first write only. Logits need this: the row count is the number of + selected tokens rather than the prompt length (so the token gate does not apply), and the + decode step that follows prefill would otherwise overwrite the value we care about. """ if not _DUMP_DIR: return - if _DUMP_TOKENS and tensor.shape[0] != _DUMP_TOKENS: + if _DUMP_TOKENS and not once and tensor.shape[0] != _DUMP_TOKENS: + return + if once and tag in _CAPTURED: return if not _want_moe(tag): return diff --git a/deepinfer/model_executor/layers/sampler.py b/deepinfer/model_executor/layers/sampler.py index 384a588..dc3adac 100644 --- a/deepinfer/model_executor/layers/sampler.py +++ b/deepinfer/model_executor/layers/sampler.py @@ -12,6 +12,7 @@ import torch.nn as nn from deep_select.interface import topk from deepinfer import 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.layers.mtp_info_pool import MTPInfoPool @@ -140,6 +141,12 @@ 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(): + # 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) del hidden_states if driver_inputs is None: diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index 9a03065..e08bf82 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -15,6 +15,7 @@ from deepinfer import envs from deepinfer._C.engine import prepare_sampler_inputs, update_prompt_logprobs # pyright: ignore[reportMissingModuleSource] from deepinfer.config import CacheConfig, ModelConfig, ParallelConfig, SchedulerConfig 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.input_metadata import InputMetadata from deepinfer.model_executor.layers.mtp_info_pool import get_nextn_pad_token_ids @@ -283,26 +284,28 @@ class ModelRunner: kv_caches: GPUCacheInterface, ) -> None: """Profile-run only: run one full step (forward + sample + post-process) synchronously.""" - main_model_output = self.execute_main_model( - seq_metadata_list, - kv_caches, - ) - input_metadata = main_model_output.input_metadata - model_output = Movable(main_model_output.model_output) - del main_model_output + # Synthetic tokens: keep tensor dumps off, or they capture this step instead of the real one. + with _debug_dump.profiling(): + main_model_output = self.execute_main_model( + seq_metadata_list, + kv_caches, + ) + input_metadata = main_model_output.input_metadata + model_output = Movable(main_model_output.model_output) + del main_model_output - sampler_output = self.execute_sample( - seq_metadata_list, - input_metadata, - model_output, - kv_caches, - ) - torch.cuda.current_stream().synchronize() - self.post_process( - seq_metadata_list, - sampler_output, - input_metadata.is_prompt, - ) + sampler_output = self.execute_sample( + seq_metadata_list, + input_metadata, + model_output, + kv_caches, + ) + torch.cuda.current_stream().synchronize() + self.post_process( + seq_metadata_list, + sampler_output, + input_metadata.is_prompt, + ) @torch.inference_mode() def execute_main_model( diff --git a/scripts/h20_recipe_probe.py b/scripts/h20_recipe_probe.py new file mode 100644 index 0000000..e8e4bef --- /dev/null +++ b/scripts/h20_recipe_probe.py @@ -0,0 +1,88 @@ +"""Probe which FP8 scale recipes DeepGEMM accepts on this GPU, via its own layout helper. + +The SM90 path re-quantizes everything to fp32 per-128 because `get_gemm_recipe()` assumes Hopper +has nothing finer. The reference converter instead keeps an e8m0 scale per 32x32 block and folds +the relative offset into the FP8 values, which is lossless -- worth adopting only if the grouped +GEMM actually accepts it. + +`transform_sf_into_required_layout` is the single place DeepGEMM validates a recipe (the engine +calls it for every MoE weight), so probing it answers the question directly and in milliseconds, +without materializing weights or launching a kernel. + + python scripts/h20_recipe_probe.py +""" + +import deep_gemm +import torch + +E8M0 = torch.float8_e8m0fnu + + +def make_sf(m: int, k: int, block_m: int, block_k: int, num_groups: int, e8m0: bool) -> torch.Tensor: + """Scales in the plainest layout: one value per (block_m x block_k) tile, per group.""" + shape = (num_groups * (m // block_m), k // block_k) + if e8m0: + return torch.full(shape, 127, dtype=torch.uint8, device="cuda").view(E8M0) + return torch.ones(shape, dtype=torch.float32, device="cuda") + + +def probe(recipe_b: tuple[int, int], block_a: int, e8m0: bool, n: int, k: int, num_groups: int) -> str: + block_n, block_k = recipe_b + sf = make_sf(n, k, block_n, block_k, num_groups, e8m0) + try: + out = deep_gemm.transform_sf_into_required_layout( + sf, + n, + k, + recipe=(1, block_n, block_k), + num_groups=num_groups, + is_sfa=False, + disable_ue8m0_cast=e8m0, + ) + got = tuple(out.shape) + # The helper silently re-lays-out whatever it is given, so acceptance alone is not proof; + # report the shape so an obviously wrong result is visible too. + return f"ACCEPTED out={got} dtype={out.dtype}" + except Exception as exc: # noqa: BLE001 - report whatever it raises + first = (str(exc).strip().splitlines() or [type(exc).__name__])[0] + return f"REJECTED {first[:100]}" + + +def main() -> None: + print(f"device: {torch.cuda.get_device_name(0)}") + print(f"deep_gemm: {deep_gemm.__version__}\n") + + n, k, num_groups = 2304, 2560, 8 + + print("B-operand recipes (what expert weights use, is_sfa=False):") + for recipe, block_a, e8m0, label in [ + ((128, 128), 128, False, "fp32 per-128 (current SM90)"), + ((128, 128), 128, True, "e8m0 per-128"), + ((32, 32), 128, False, "fp32 per-32 (ref-converter grid)"), + ((32, 32), 128, True, "e8m0 per-32 (ref-converter, lossless)"), + ((1, 32), 128, True, "e8m0 per-row/32 (MXFP8)"), + ((1, 128), 128, True, "e8m0 per-row/128"), + ]: + print(f" {label:34} recipe={recipe} e8m0={e8m0:<5} -> {probe(recipe, block_a, e8m0, n, k, num_groups)}") + + print("\nA-operand recipes (activations, is_sfa=True):") + for block_k, e8m0, label in [ + (128, False, "fp32 per-128 (current SM90)"), + (32, True, "e8m0 per-32 (MXFP8)"), + (32, False, "fp32 per-32"), + ]: + sf = make_sf(n, k, 1, block_k, 1, e8m0) + try: + out = deep_gemm.transform_sf_into_required_layout( + sf, n, k, recipe=(1, 1, block_k), num_groups=1, is_sfa=True, disable_ue8m0_cast=e8m0 + ) + print(f" {label:34} block_k={block_k:<4} e8m0={e8m0:<5} -> ACCEPTED out={tuple(out.shape)}") + except Exception as exc: # noqa: BLE001 + first = (str(exc).strip().splitlines() or [type(exc).__name__])[0] + print(f" {label:34} block_k={block_k:<4} e8m0={e8m0:<5} -> REJECTED {first[:100]}") + + print("\nPROBE_DONE") + + +if __name__ == "__main__": + main() diff --git a/scripts/ref_layer_dump.py b/scripts/ref_layer_dump.py index 76ba49a..3990a18 100644 --- a/scripts/ref_layer_dump.py +++ b/scripts/ref_layer_dump.py @@ -114,7 +114,9 @@ def main() -> None: with torch.inference_mode(): output_ids, logits, _ = model.forward(tokens, 0) - captured["final_row_logits"] = logits[0, -1].float().cpu().numpy() + # The head already collapsed the sequence axis (`x = x[:, -1]` when not full_logits), so the + # logits are [b, vocab] -- indexing [0, -1] here would silently store a scalar. + captured["final_row_logits"] = logits[0].float().cpu().numpy() captured["next_token"] = np.array([int(output_ids[-1])], dtype=np.int32) if rank == 0: