mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
106 lines
5.2 KiB
Python
106 lines
5.2 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("--cp", default=1, type=int, help="TP group size: attention heads sharded over cp ranks running one batch (issue #28)")
|
||
parser.add_argument("--max_tokens", default=16, type=int)
|
||
parser.add_argument("--num_prompts", default=len(PROMPTS), type=int, help="run only the first N prompts (eager prefill compiles per prompt length)")
|
||
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)
|
||
parser.add_argument("--max_num_seqs", default=1, type=int,
|
||
help="decode batch; every captured batch size adds one set of region NEFFs per layer (HBM: model code)")
|
||
args = parser.parse_args()
|
||
dp_size = args.ep // args.cp
|
||
|
||
|
||
def _fixed_kv_pool(num_blocks: int) -> None:
|
||
os.environ["DEEPINFER_NEURON_NUM_KV_BLOCKS"] = str(num_blocks) # CP followers read it too
|
||
"""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, args.cp, 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=args.cp,
|
||
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=args.max_num_seqs,
|
||
gpu_memory_utilization=0.85,
|
||
lookahead_tokens=0,
|
||
)
|
||
|
||
from deepinfer.sampling_params import SamplingParams
|
||
|
||
prompts = PROMPTS[: args.num_prompts]
|
||
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)
|