mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 16:56: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.
106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
"""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()
|