mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
28fb291db1
No behaviour change. `CONTRIBUTING.md` requires `ruff check --no-cache` and `ruff format --check --no-cache` to pass, and the SM90 branch predates that requirement: 72 findings, all in the files added here. Most of it is mechanical (`ruff format`), plus: - annotations for the helper signatures (`torch.Tensor` / `int` / `Generator` / return types); hooks are typed as `Callable[[nn.Module, tuple, object], None]` rather than left bare; - two `if`/`else` assignments become ternaries (SIM108), and the index read in `h20_real_requant_check.py` uses a context manager (SIM115). Also corrects the `sparse_mla_fwd` docstring, which still described the old `0 <= idx <= q_position` masking contract. Passing `is_causal=True` there is what produced the all-NaN layer 0: indices address the assembled `[swa cache | current swa | compressed]` buffer, so the row number is not a valid causal bound and the engine must pass False.
88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
"""Reference next-token prediction from the official V4.1 inference code.
|
|
|
|
`generate.py` always wraps text in a chat template, so its outputs are not directly comparable
|
|
with a raw-token harness. This script takes explicit token ids and prints the greedy next token
|
|
(plus the top-5), which makes the comparison exact regardless of prompt formatting.
|
|
|
|
torchrun --nproc-per-node 8 ref_next_token.py --ids 0,128803,2056,... --ckpt-path ...
|
|
|
|
Writes `<out>.json` with the ids, the chosen token and the top-5 so a harness can diff it.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
from model import ModelArgs, Transformer # noqa: E402
|
|
from safetensors.torch import load_model # 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, 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)
|
|
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
|
|
|
|
# The engram hash table is built from the tokenizer's compressed token map, so it is not
|
|
# optional even for a text-only forward.
|
|
from transformers import AutoTokenizer
|
|
|
|
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")
|
|
|
|
tokens = torch.tensor([ids], dtype=torch.long)
|
|
with torch.inference_mode():
|
|
# `Transformer.forward` returns (output_ids, logits, main_hidden).
|
|
output_ids, logits, _ = model.forward(tokens, 0)
|
|
|
|
row = logits[0, -1].float()
|
|
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,
|
|
"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],
|
|
}
|
|
if rank == 0:
|
|
print(json.dumps(payload, ensure_ascii=False), flush=True)
|
|
os.makedirs(os.path.dirname(args.out), exist_ok=True)
|
|
with open(args.out, "w") as f:
|
|
json.dump(payload, f, ensure_ascii=False, indent=2)
|
|
if world_size > 1:
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|