coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 23:07:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

offline_inference_neuron.py --cp

Browse Source 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-24 00:37:23 +08:00
parent 3ee74e8a4e
commit d44cf253c5
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
examples/neuron/offline_inference_neuron.py
View File

@ -36,6 +36,7 @@ PROMPTS = [
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")
@ -44,7 +45,7 @@ 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
dp_size = args.ep // args.cp
def _fixed_kv_pool(num_blocks: int) -> None:
@ -69,14 +70,14 @@ def main(dp_size: int, dp_rank: int) -> None:
# 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())
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=1,
cp_size=args.cp,
env={"DEEPINFER_LOGGING_INTERVAL_STEP": "1"},
data_parallel_size=dp_size,
max_num_batched_tokens=args.max_num_batched_tokens,