DeepInfer/scripts/ref_next_token.py

"""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 numpy as np
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)
    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]

    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")

    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()))

        # 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)

    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)
        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()