mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
0d2cbcb112
The DP-wide prefill vote runs every DEEPINFER_PREFILL_ATTEMPT_INTERVAL (70) / ..._ANY_ (100) decode steps. Requests that become schedulable after step 0 (multimodal-processor round trip) then wait 70 placeholder steps; with eager execution a full-model step on the NeuronCore takes ~150 s, so the smoke test sat ~3 hours in placeholder steps. Default both intervals to 1 on Neuron. offline_inference_neuron: force-topk off (#22), 400-block KV pool. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
100 lines
4.6 KiB
Python
100 lines
4.6 KiB
Python
"""Offline inference with the full DeepSeek-V4.1-Flash on one trn2.48xlarge (Neuron basic mode).
|
||
|
||
Same flow as examples/offline_inference.py (launch_dp -> one engine process per DP rank, each owning one
|
||
NeuronCore), with the settings basic mode needs:
|
||
* text-only prompts (VL kernels are stubbed in basic mode);
|
||
* no DSpark/MTP drafting (DEEPINFER_USE_NEXTN_DECODE=0) and no EPLB;
|
||
* a fixed KV-block pool instead of the engine's memory-profiling prefill (a max_num_batched_tokens dummy
|
||
prefill does not fit next to the eager pure-torch kernels);
|
||
* small batches and few output tokens: eager execution compiles every new op/shape on the first run.
|
||
|
||
export DEEPINFER_ACCELERATOR=neuron WORLD_SIZE=1 RANK=0
|
||
python examples/neuron/offline_inference_neuron.py --model_name /mnt/nvme/models/DeepSeek-V4.1-Flash-1017 --ep 64
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import math
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
EXAMPLES = Path(__file__).resolve().parent.parent
|
||
sys.path.insert(0, str(EXAMPLES))
|
||
|
||
from _config import DEFAULT_MODEL_NAME, create_llm, setup_common_env, setup_cuda_visible_devices # noqa: E402
|
||
from _launch import launch_dp # noqa: E402
|
||
|
||
PROMPTS = [
|
||
"问题：法国的首都是哪里？\n答案：巴黎\n\n问题：一年有几个月？\n答案：12 个月\n\n问题：中国的首都是哪里？\n答案：",
|
||
"Question: What is the largest planet in the solar system?\nAnswer: Jupiter\n\n"
|
||
"Question: How many legs does a spider have?\nAnswer: Eight\n\n"
|
||
"Question: Who wrote Romeo and Juliet?\nAnswer:",
|
||
]
|
||
|
||
parser = argparse.ArgumentParser(description="V4.1 offline inference on Neuron (basic mode)")
|
||
parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME, help="HF model id or local HF checkpoint directory")
|
||
parser.add_argument("--ep", default=64, type=int)
|
||
parser.add_argument("--max_tokens", default=16, type=int)
|
||
parser.add_argument("--num_kv_blocks", default=400, type=int, help="fixed KV pool per rank; the scheduler admits a request only if max_model_len tokens fit, per running sequence")
|
||
parser.add_argument("--max_model_len", default=16384, type=int, help="the candidate-block filter needs >= candidate_topk_blocks * candidate_block_size (2048 * 8)")
|
||
parser.add_argument("--max_num_batched_tokens", default=1024, type=int)
|
||
args = parser.parse_args()
|
||
dp_size = args.ep
|
||
|
||
|
||
def _fixed_kv_pool(num_blocks: int) -> None:
|
||
"""Skip WorkerBase.profile_num_available_blocks (a max_num_batched_tokens dummy prefill)."""
|
||
from deepinfer.worker.base import WorkerBase
|
||
|
||
def profile_num_available_blocks(self): # noqa: ANN001, ANN202
|
||
self.model_runner.set_block_size(self.cache_config.block_size)
|
||
self.cache_config.num_gpu_blocks = num_blocks
|
||
return num_blocks
|
||
|
||
WorkerBase.profile_num_available_blocks = profile_num_available_blocks
|
||
|
||
|
||
def main(dp_size: int, dp_rank: int) -> None:
|
||
os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
|
||
os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron")
|
||
# sampler_force_topk=0: torch.topk with k=4096 over the vocabulary does not compile on the device (issue #22).
|
||
setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=False, enable_eplb=False)
|
||
# Pick this rank's NeuronCore before anything initializes deepinfer.PLATFORM (importing deepinfer.worker does):
|
||
# PlatformNeuron claims the core from CUDA_VISIBLE_DEVICES once, at init.
|
||
from deep_cluster.hardware import get_actual_num_gpus
|
||
|
||
setup_cuda_visible_devices(dp_rank, 1, get_actual_num_gpus())
|
||
_fixed_kv_pool(args.num_kv_blocks)
|
||
|
||
llm = create_llm(
|
||
args.model_name,
|
||
dp_rank,
|
||
ep_size=args.ep,
|
||
cp_size=1,
|
||
env={"DEEPINFER_LOGGING_INTERVAL_STEP": "1"},
|
||
data_parallel_size=dp_size,
|
||
max_num_batched_tokens=args.max_num_batched_tokens,
|
||
max_model_len=args.max_model_len,
|
||
max_num_seqs=len(PROMPTS),
|
||
gpu_memory_utilization=0.85,
|
||
lookahead_tokens=0,
|
||
)
|
||
|
||
from deepinfer.sampling_params import SamplingParams
|
||
|
||
params = [SamplingParams(temperature=0.0, max_tokens=args.max_tokens) for _ in PROMPTS]
|
||
outputs = llm.generate(prompts=PROMPTS, sampling_params=params, use_tqdm=False)
|
||
for i, output in enumerate(outputs):
|
||
text = llm.llm_engine.tokenizer.decode([x.token_id for x in output])
|
||
logps = [x.token_logprobs for x in output]
|
||
ppl = math.exp(-sum(logps) / len(logps)) if logps else float("nan")
|
||
if dp_rank == 0:
|
||
print(f"DP {dp_rank}; INPUT: {PROMPTS[i]!r}; OUTPUT: {text!r}; output_tokens={len(output)}; decode_ppl={ppl}", flush=True)
|
||
llm.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
launch_dp(main, dp_size)
|