mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
cc90926a1c
Device-memory fixes found bringing up the released checkpoint (16.45 GiB of weights per 24 GB core): - Engram: the Neuron EngramBuffer keeps only this rank's SF rows, so skip the all-gather of the replicated SF table (~3 GB per layer, allocated on the device -> "Failed to allocate resource"). - gather_tl: index_select / advanced indexing of an int8 table allocates several times the table on the device (+4.9 GiB for the 1.3 GB embedding); gather through F.embedding on an int32 / fp16 view instead (bit-exact, no extra memory). - deep_gemm FP8/FP4 GEMMs: dequantize the weight in N chunks of 32M elements and one expert at a time instead of whole projections / all local experts in fp32 (OOM on a fragmented pool); bit-identical results. - transfer_to_neuron: import torch.utils._device explicitly (missing when torch_neuronx is not auto-loaded). examples/neuron/offline_inference_neuron.py: offline inference in basic mode (text prompts, no MTP / EPLB, fixed KV pool, max_model_len >= 2048 * 8 for the candidate-block filter), one process per NeuronCore. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
99 lines
4.5 KiB
Python
99 lines
4.5 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=176, type=int, help="fixed KV pool per rank; must hold max_model_len tokens")
|
||
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")
|
||
setup_common_env(model_name=args.model_name, sampler_force_topk=4096, 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)
|