mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
8a5209ee65
- examples/neuron: tiny random V4.1 checkpoint in the real HF layout, pure-torch replacements for the official reference kernels, golden generation, a teacher-forced DeepInfer runner (LLM offline API) and a comparison tool. CPU mode: 25/27 greedy tokens match, worst logits cos 0.9985 / rel 0.056 against a reference self-noise floor of 0.032-0.060. - model_loader: check parameters against the platform device, not is_cuda. - deep_gemm paged MQA logits: do not dereference block-table slots past the context. - bind_cpu: skip in Neuron CPU mode (1 core per NeuronCore share would serialize host compute). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
352 lines
16 KiB
Python
352 lines
16 KiB
Python
"""Run the tiny V4.1 checkpoint through the real DeepInfer engine, teacher-forced on the reference golden.
|
|
|
|
Engine path: deepinfer.entrypoints.llm.LLM (scheduler, KV cache, model runner, sampler), EP=DP=CP=1, Neuron basic
|
|
mode defaults (eager, DEEPINFER_USE_NEXTN_DECODE=0). Teacher forcing uses the engine's own mechanism,
|
|
SamplingParams(golden_token_ids=...), so decode step i is fed the golden token i-1 exactly like the reference
|
|
(reference_golden.py --force-tokens-dir). Logits are captured by wrapping the sampler's `get_logits` (the fp32
|
|
[num_reqs, vocab] lm-head output before any sampling transform); per-layer block outputs by forward hooks on the
|
|
decoder layers (best effort, see --hidden).
|
|
|
|
Output: <out>/<prompt>.pt in the reference_golden.py format (prompt_ids, tokens = argmax per step, fed_tokens,
|
|
logits [steps+1, vocab] fp32, optional hidden), ready for compare_golden.py.
|
|
|
|
cd ~/di/DeepInfer && ~/di/venv/bin/python examples/neuron/run_deepinfer_tiny.py \
|
|
--hf-ckpt ~/di/tiny_ckpt --golden-dir ~/di/golden --out ~/di/deepinfer_out [--device cpu|neuron] \
|
|
[--hidden] [--compare] [--threads 32]
|
|
|
|
Placeholder engine steps (is_placeholder=True, no real sequence) also run the model and the sampler; they are
|
|
filtered out of the recorded logits/hidden. `install_workarounds` patches, at runtime and only in this process,
|
|
the bugs found in other modules while bringing this up (BUG-1/2/4, NOTE-3/5 in the code); the list of the ones
|
|
that actually triggered is printed at the end. Set RUN_TINY_DEBUG=1 to print every sampler call.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
def parse_args():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--hf-ckpt", default="~/di/tiny_ckpt")
|
|
ap.add_argument("--golden-dir", default="~/di/golden")
|
|
ap.add_argument("--out", default="~/di/deepinfer_out")
|
|
ap.add_argument("--prompts", default="short,medium,long")
|
|
ap.add_argument("--device", choices=["cpu", "neuron"], default="cpu", help="sets DEEPINFER_NEURON_DEVICE")
|
|
ap.add_argument("--max-model-len", type=int, default=2048)
|
|
ap.add_argument("--max-num-batched-tokens", type=int, default=1024)
|
|
ap.add_argument("--threads", type=int, default=16)
|
|
ap.add_argument("--num-kv-blocks", type=int, default=32, help="fixed KV block pool (128 tokens each); 0 = engine profiling")
|
|
ap.add_argument("--hidden", action="store_true", help="also record per-layer block outputs (hc residual stream)")
|
|
ap.add_argument("--compare", action="store_true", help="run compare_golden on every prompt afterwards")
|
|
return ap.parse_args()
|
|
|
|
|
|
def setup_env(device: str) -> None:
|
|
# Must happen before deepinfer is imported (see examples/_config.py).
|
|
os.environ["DEEPINFER_ACCELERATOR"] = "neuron"
|
|
os.environ["DEEPINFER_NEURON_DEVICE"] = device
|
|
os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
|
|
os.environ.setdefault("DEEPINFER_ENFORCE_EAGER", "1")
|
|
os.environ.setdefault("DEEPINFER_ZERO_OVERHEAD_SCHEDULING", "0")
|
|
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0") # one engine rank; mapped to the target device
|
|
os.environ.setdefault("RANK", "0")
|
|
os.environ.setdefault("WORLD_SIZE", "1")
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
|
os.environ.setdefault("DEEPINFER_USE_MULTIMODAL_POOL", "0") # text only; avoids an 8 GB pinned host pool
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
here = str(Path(__file__).resolve().parent)
|
|
if here not in sys.path:
|
|
sys.path.insert(0, here)
|
|
|
|
|
|
WORKAROUNDS: list[str] = []
|
|
|
|
|
|
def install_workarounds(num_kv_blocks: int) -> None:
|
|
"""Work around bugs in other modules without editing them; each one is reported (see the module docstring)."""
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
# BUG-1: PlatformNeuron(cpu) uses gloo and fp8_sf_type=torch.int16; gloo collectives reject int16
|
|
# ("RuntimeError: Invalid scalar type") in v4_modules/engram.py NgramEmbedding.gather_sf. Byte-reinterpret
|
|
# dtypes gloo lacks as uint8 for the gather-type collectives (lossless; concatenation is along dim 0).
|
|
orig_agit = dist.all_gather_into_tensor
|
|
unsupported = {torch.int16, torch.uint16, torch.float8_e4m3fn, torch.float8_e8m0fnu}
|
|
|
|
def all_gather_into_tensor(output_tensor, input_tensor, group=None, async_op=False):
|
|
if input_tensor.dtype in unsupported and dist.get_backend(group) == "gloo":
|
|
if "BUG-1" not in WORKAROUNDS:
|
|
WORKAROUNDS.append("BUG-1")
|
|
return orig_agit(output_tensor.view(torch.uint8), input_tensor.contiguous().view(torch.uint8), group=group, async_op=async_op)
|
|
return orig_agit(output_tensor, input_tensor, group=group, async_op=async_op)
|
|
|
|
dist.all_gather_into_tensor = all_gather_into_tensor
|
|
|
|
# BUG-2: model_executor/model_loader.py get_model asserts `p.is_cuda` for every parameter; on Neuron the
|
|
# parameters live on the transfer_to_neuron target (neuron:0 or cpu). Same function, device-agnostic check.
|
|
import deepinfer.model_executor.model_loader as ml
|
|
import deepinfer.worker.model_runner as mr
|
|
from deepinfer import PLATFORM
|
|
|
|
def get_model(model_config, skip_load=False):
|
|
model_class = ml._get_model_architecture(model_config)
|
|
with ml._set_default_torch_dtype(model_config.dtype), torch.device("cuda"):
|
|
model = model_class.load_model(model_config.model, model_config.model_cfg, skip_load)
|
|
target = PLATFORM.device.type
|
|
for n, p in model.named_parameters():
|
|
if p.device.type != target:
|
|
raise AssertionError(f"{n} on {p.device}, expected {target}")
|
|
if hasattr(p, "sf") and p.sf.device.type != target:
|
|
raise AssertionError(f"{n}.sf on {p.sf.device}, expected {target}")
|
|
if not any("BUG-2" == w for w in WORKAROUNDS):
|
|
WORKAROUNDS.append("BUG-2")
|
|
return model.eval().requires_grad_(False)
|
|
|
|
ml.get_model = get_model
|
|
mr.get_model = get_model
|
|
|
|
# NOTE-5: WorkerBase pins the worker to phys_cpus // get_actual_num_gpus() CPUs; on Trn2 that is 96 // 64 = 1
|
|
# CPU, so CPU-mode runs are single-threaded (~10x slower). Keep the process affinity instead.
|
|
import deepinfer.worker.base as wb
|
|
|
|
def bind_cpu(local_rank):
|
|
WORKAROUNDS.append("NOTE-5: skipped bind_cpu (1 CPU per rank on Trn2)")
|
|
|
|
wb.bind_cpu = bind_cpu
|
|
|
|
# BUG-4: vendor deep_gemm fp8_fp4_paged_mqa_logits gathers keys for every column < max_context_len and only
|
|
# masks afterwards, so block-table slots past a request's last page (unused entries, not valid page ids) are
|
|
# dereferenced -> IndexError. Sanitize out-of-range table entries (their columns are masked to -inf anyway) and
|
|
# flag it loudly if an out-of-range entry is ever needed by a valid position.
|
|
import deep_gemm
|
|
|
|
orig_paged = deep_gemm.fp8_fp4_paged_mqa_logits
|
|
|
|
def fp8_fp4_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, *a, **k):
|
|
num_pages, page_kv = kv_cache.shape[0], kv_cache.shape[1]
|
|
bad = (block_table < 0) | (block_table >= num_pages)
|
|
if bool(bad.any()):
|
|
lens = context_lens.reshape(-1).to(torch.int64)
|
|
needed_pages = (lens + page_kv - 1) // page_kv
|
|
col = torch.arange(block_table.shape[1], device=block_table.device)[None, :]
|
|
if bool((bad & (col < needed_pages[:, None])).any()):
|
|
WORKAROUNDS.append(f"BUG-4b: out-of-range page id inside the valid context: table={block_table.tolist()} lens={lens.tolist()}")
|
|
if "BUG-4" not in WORKAROUNDS:
|
|
WORKAROUNDS.append("BUG-4")
|
|
block_table = torch.where(bad, torch.zeros_like(block_table), block_table)
|
|
return orig_paged(q, kv_cache, weights, context_lens, block_table, *a, **k)
|
|
|
|
deep_gemm.fp8_fp4_paged_mqa_logits = fp8_fp4_paged_mqa_logits
|
|
|
|
# NOTE-3 (not a bug, a CPU-mode hazard): WorkerBase.profile_num_available_blocks runs a dummy prefill of
|
|
# max_num_batched_tokens tokens (903 s on the host for 2048 tokens) and then sizes the KV cache from
|
|
# mem_get_info(), which on cpu is host RAM -> 0.5 * 2 TB = 559 GiB of KV blocks. Use a fixed small pool instead.
|
|
if num_kv_blocks > 0:
|
|
import torch.distributed as _dist # noqa: F401
|
|
from deepinfer.worker.base import WorkerBase
|
|
|
|
def profile_num_available_blocks(self):
|
|
self.model_runner.set_block_size(self.cache_config.block_size)
|
|
self.cache_config.num_gpu_blocks = num_kv_blocks
|
|
WORKAROUNDS.append(f"NOTE-3: fixed {num_kv_blocks} KV blocks, no profile run")
|
|
return num_kv_blocks
|
|
|
|
WorkerBase.profile_num_available_blocks = profile_num_available_blocks
|
|
|
|
|
|
class Recorder:
|
|
"""Collects the lm-head logits of every sampler call (one request in flight -> one row per call)."""
|
|
|
|
def __init__(self):
|
|
self.step_kind = None # (is_prompt, is_placeholder) of the engine step being executed
|
|
self.n_placeholder = 0
|
|
self.logits = []
|
|
self.hidden_calls = []
|
|
self._cur_hidden = {}
|
|
|
|
def reset(self):
|
|
self.n_placeholder = 0
|
|
self.logits.clear()
|
|
self.hidden_calls.clear()
|
|
self._cur_hidden = {}
|
|
|
|
|
|
def install_logits_hook(rec: Recorder) -> None:
|
|
import deepinfer.model_executor.layers.sampler as sampler_mod
|
|
|
|
orig = sampler_mod.get_logits
|
|
|
|
def get_logits(hidden_states, embedding, vocab_size):
|
|
logits = orig(hidden_states, embedding, vocab_size)
|
|
if rec.step_kind is not None and rec.step_kind[1]:
|
|
rec.n_placeholder += 1 # placeholder step (no real sequence): not part of the request
|
|
return logits
|
|
rec.logits.append(logits.detach().float().cpu().clone())
|
|
if os.environ.get("RUN_TINY_DEBUG"):
|
|
import traceback as _tb
|
|
|
|
st = " <- ".join(f"{f.name}:{f.lineno}" for f in _tb.extract_stack()[-6:-1][::-1])
|
|
print(f"[run] sampler call {len(rec.logits)}: logits {tuple(logits.shape)} argmax {int(logits[0].argmax())} "
|
|
f"hs {tuple(hidden_states.shape)} at {time.strftime('%H:%M:%S')} via {st}", flush=True)
|
|
return logits
|
|
|
|
sampler_mod.get_logits = get_logits
|
|
|
|
from deepinfer.worker.worker import Worker
|
|
|
|
orig_step = Worker.execute_model_step
|
|
|
|
def execute_model_step(self, is_prompt, is_placeholder, *a, **k):
|
|
rec.step_kind = (bool(is_prompt), bool(is_placeholder))
|
|
try:
|
|
return orig_step(self, is_prompt, is_placeholder, *a, **k)
|
|
finally:
|
|
rec.step_kind = None
|
|
|
|
Worker.execute_model_step = execute_model_step
|
|
|
|
|
|
def install_hidden_hooks(rec: Recorder, llm) -> str:
|
|
"""Hook the decoder layers of the driver worker's model. Returns a status string."""
|
|
try:
|
|
worker = llm.llm_engine.driver_worker
|
|
runner = getattr(worker, "model_runner", None)
|
|
model = getattr(runner, "model", None) if runner is not None else getattr(worker, "model", None)
|
|
layers = model.model.layers
|
|
except Exception as exc: # noqa: BLE001
|
|
return f"hidden hooks not installed: {type(exc).__name__}: {exc}"
|
|
|
|
from deepinfer.model_executor.models.v4_modules.hc import HCLayerMixin
|
|
|
|
n = len(layers)
|
|
|
|
def hook(i):
|
|
def fn(mod, inp, out):
|
|
if rec.step_kind is None or rec.step_kind[1]:
|
|
return None # placeholder / non-request forward
|
|
h, residual = out
|
|
h = h.move() if hasattr(h, "move") else h
|
|
try:
|
|
# 2D output = before hc_post; the reference block output is the expanded hc stream
|
|
x = HCLayerMixin.hc_post(h, residual) if h.dim() == 2 else h
|
|
except Exception: # noqa: BLE001
|
|
x = h
|
|
rec._cur_hidden[i] = x.detach().float().cpu().clone()
|
|
if i == n - 1:
|
|
rec.hidden_calls.append(dict(rec._cur_hidden))
|
|
rec._cur_hidden = {}
|
|
return None
|
|
|
|
return fn
|
|
|
|
for i, layer in enumerate(layers):
|
|
layer.register_forward_hook(hook(i))
|
|
return f"hidden hooks on {n} layers of {type(model).__name__}"
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
for k in ("hf_ckpt", "golden_dir", "out"):
|
|
setattr(args, k, os.path.abspath(os.path.expanduser(getattr(args, k))))
|
|
setup_env(args.device)
|
|
os.makedirs(args.out, exist_ok=True)
|
|
|
|
import torch
|
|
|
|
torch.set_num_threads(args.threads)
|
|
import deepinfer # noqa: F401
|
|
from deepinfer import PLATFORM
|
|
|
|
print(f"[run] platform device = {PLATFORM.device}", flush=True)
|
|
if args.device == "cpu":
|
|
assert str(PLATFORM.device) == "cpu", "refusing to run: target device is not the host CPU"
|
|
|
|
install_workarounds(args.num_kv_blocks)
|
|
rec = Recorder()
|
|
install_logits_hook(rec)
|
|
|
|
from deepinfer.entrypoints.llm import LLM
|
|
from deepinfer.sampling_params import SamplingParams
|
|
|
|
t0 = time.time()
|
|
llm = LLM(
|
|
model=args.hf_ckpt,
|
|
trust_remote_code=True,
|
|
expert_parallel_size=1,
|
|
context_parallel_size=1,
|
|
data_parallel_size=1,
|
|
engine_rank=0,
|
|
chunked_prefill=True,
|
|
decode_sequence=False,
|
|
max_model_len=args.max_model_len,
|
|
max_num_batched_tokens=args.max_num_batched_tokens,
|
|
max_num_seqs=1,
|
|
gpu_memory_utilization=0.5,
|
|
)
|
|
print(f"[run] engine up in {time.time() - t0:.1f}s", flush=True)
|
|
status = install_hidden_hooks(rec, llm) if args.hidden else "hidden not requested"
|
|
print(f"[run] {status}", flush=True)
|
|
|
|
for name in args.prompts.split(","):
|
|
gold = torch.load(os.path.join(args.golden_dir, f"{name}.pt"), weights_only=False)
|
|
ids, gtok = list(gold["prompt_ids"]), list(gold["tokens"])
|
|
steps = len(gtok)
|
|
rec.reset()
|
|
sp = SamplingParams(temperature=0.0, max_tokens=steps, ignore_eos=True, golden_token_ids=gtok)
|
|
t1 = time.time()
|
|
outputs = llm.generate(prompt_token_ids=[ids], sampling_params=sp, use_tqdm=False)
|
|
dt = time.time() - t1
|
|
out_tokens = [o.token_id for o in outputs[0]]
|
|
rows = [r for r in rec.logits]
|
|
assert all(r.size(0) == 1 for r in rows), [tuple(r.shape) for r in rows]
|
|
logits = torch.cat(rows)[-steps:] if len(rows) >= steps else torch.cat(rows)
|
|
print(f"[run] {name}: {len(rows)} request sampler calls, {rec.n_placeholder} placeholder steps", flush=True)
|
|
if len(rows) != steps:
|
|
print(f"[run] WARNING {name}: {len(rows)} sampler calls for {steps} steps (chunked prefill?)", flush=True)
|
|
payload = {
|
|
"name": name,
|
|
"prompt_ids": ids,
|
|
"tokens": logits.argmax(-1).tolist(),
|
|
"engine_output_tokens": out_tokens,
|
|
"fed_tokens": gtok[: steps - 1],
|
|
"logits": logits,
|
|
"source": "deepinfer LLM engine, teacher-forced via SamplingParams.golden_token_ids",
|
|
"device": args.device,
|
|
}
|
|
if rec.hidden_calls:
|
|
hid = []
|
|
for call in rec.hidden_calls[-steps:]:
|
|
hid.append(torch.stack([call[i].reshape(call[i].shape[0], -1, gold["hidden"][0].shape[-1]) for i in sorted(call)]).bfloat16())
|
|
payload["hidden"] = hid
|
|
path = os.path.join(args.out, f"{name}.pt")
|
|
torch.save(payload, path)
|
|
match = sum(a == b for a, b in zip(payload["tokens"], gtok))
|
|
print(f"[run] {name}: {len(ids)} prompt tokens, engine output {out_tokens} (golden-forced: {out_tokens == gtok}), "
|
|
f"argmax==golden {match}/{steps}, {dt:.1f}s -> {path}", flush=True)
|
|
if args.compare:
|
|
import compare_golden
|
|
|
|
res = compare_golden.compare(gold, payload)
|
|
for s in res["steps"]:
|
|
print(f" step {s['step']}: match {s['match']} margin {s['margin_ref']:.3f} rel {s['rel']:.2e} cos {s['cos']:.6f} kl {s['kl']:.2e}")
|
|
for L in res["layers"][:1]:
|
|
print(" prefill hidden rel per layer:", L.get("error") or " ".join(f"{v:.1e}" for v in L["rel"]))
|
|
print(f"[run] workarounds triggered: {WORKAROUNDS}", flush=True)
|
|
llm.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except Exception:
|
|
traceback.print_exc()
|
|
sys.exit(1)
|