From 3ee74e8a4e3cfe7b274e712ea7c7ae3c936a5519 Mon Sep 17 00:00:00 2001 From: whn09 Date: Sat, 24 Oct 2026 00:04:02 +0800 Subject: [PATCH] Neuron CP followers: keep the EPLB step when skipping the sampler (#28) 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 --- deepinfer/worker/model_runner.py | 7 ++++++- examples/neuron/bench_neuron.py | 7 ++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index b93c4c1..9ac9676 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -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 = { diff --git a/examples/neuron/bench_neuron.py b/examples/neuron/bench_neuron.py index 56899e5..4de6244 100644 --- a/examples/neuron/bench_neuron.py +++ b/examples/neuron/bench_neuron.py @@ -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,