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

Neuron CP followers: keep the EPLB step when skipping the sampler (#28)

Browse Source 
Its periodic load-statistics all-gather spans the EP group; the driver waited
in it forever (allgather seq=10 timeout after the warmup request at cp=64).
Also bench_neuron.py --cp.

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

7
deepinfer/worker/model_runner.py
View File

@ -414,7 +414,12 @@ class ModelRunner:
"""
if PLATFORM.is_neuron and not self.is_driver_worker:
# Only the driver samples; a follower's sampler work (lm_head, prompt logprobs) has no collectives on
# Neuron, so it is skipped together with the sampler-input broadcast.
# Neuron, so it is skipped together with the sampler-input broadcast. The EPLB step is not: its periodic
# load-statistics all-gather spans the EP group.
from deepinfer.model_executor import eplb
if eplb.current() is not None:
eplb.current().step(input_metadata.is_prompt)
return SamplerOutput(post_funcs=[])
sampler_inputs = self.prepare_sampler_inputs(seq_metadata_list, input_metadata)
kwargs = {

7
examples/neuron/bench_neuron.py
View File

@ -36,6 +36,7 @@ from _launch import launch_dp # noqa: E402
parser = argparse.ArgumentParser(description="V4.1 TTFT/TPOT benchmark on Neuron (vllm-neuron#71 method)")
parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME)
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("--input_lens", default="512,1024,2048,8192", help="comma-separated prompt lengths (tokens)")
parser.add_argument("--output_len", default=128, type=int)
parser.add_argument("--repeats", default=3, type=int, help="requests per cell (medians are reported)")
@ -46,7 +47,7 @@ parser.add_argument("--max_num_batched_tokens", default=128, type=int)
parser.add_argument("--max_num_seqs", default=1, type=int)
parser.add_argument("--out", default="", help="write the results as JSON here (rank 0)")
args = parser.parse_args()
dp_size = args.ep
dp_size = args.ep // args.cp
def _fixed_kv_pool(num_blocks: int) -> None:
@ -97,13 +98,13 @@ def main(dp_size: int, dp_rank: int) -> None:
setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=False, enable_eplb=False)
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": "100"},
data_parallel_size=dp_size,
max_num_batched_tokens=args.max_num_batched_tokens,