[NPU] support multi-agent rollout pipeline

examples/multi_agent/agent_system.py

@@ -5,19 +5,40 @@
 from copy import deepcopy
 
 from vime.rollout.rm_hub import batched_async_rm
-from vime.rollout.vllm_rollout import _build_inference_sampling_params, _inference_generate_tokens_and_logprobs
+from vime.rollout.vllm_rollout import _build_inference_sampling_params
 from vime.utils.http_utils import post
 from vime.utils.types import Sample
 
 from .prompts import SOLVER_PROMPT_TEMPLATE, generate_rewriter_template, generate_select_template
 
 
-async def generate_response(args, prompt, key):
+def _inference_generate_tokens_and_logprobs(choice):
+    """Compatibility helper for current VIME vLLM rollout responses."""
+    new_response_tokens = choice.get("token_ids") or []
+    new_response_log_probs = []
+    lp = choice.get("logprobs")
+    if isinstance(lp, dict):
+        content_items = lp.get("content") or []
+        new_response_log_probs = [
+            float(item.get("logprob", 0.0)) if isinstance(item, dict) else 0.0 for item in content_items
+        ]
+    if not new_response_log_probs:
+        new_response_log_probs = [0.0] * len(new_response_tokens)
+    return new_response_tokens, new_response_log_probs
+
+
+async def generate_response(args, prompt, key, worker_id: int | None = None):
     try:
         sampling_params = args.sampling_params
         tokenizer = args.tokenizer
         max_context_length = args.rollout_max_context_len
         sample = deepcopy(args.sample)
+        sample.metadata = dict(sample.metadata or {})
+        sample.metadata["multi_agent_role"] = key
+        sample.metadata["multi_agent_parent_group_index"] = sample.group_index
+        sample.metadata["multi_agent_parent_index"] = sample.index
+        if worker_id is not None:
+            sample.metadata["multi_agent_worker_id"] = worker_id
 
         url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}/inference/v1/generate"
 
@@ -90,11 +111,11 @@ class Agent:
     def __init__(self):
         pass
 
-    async def run(self, args, prompt, max_retries: int = 1, key: str = None) -> str:
+    async def run(self, args, prompt, max_retries: int = 1, key: str = None, worker_id: int | None = None) -> str:
         """Runs the agent by sending a prompt to the LLM."""
         for _i in range(max_retries):
             try:
-                response = await generate_response(args, prompt, key=key)
+                response = await generate_response(args, prompt, key=key, worker_id=worker_id)
                 return response
             except Exception as e:
                 print(f"Error querying LLM: {e}")
@@ -109,10 +130,10 @@ class SolverAgent(Agent):
     def __init__(self):
         super().__init__()
 
-    async def generate_initial_solution(self, args, problem_statement) -> str:
+    async def generate_initial_solution(self, args, problem_statement, worker_id: int) -> str:
         """Generates the first solution attempt."""
         prompt = SOLVER_PROMPT_TEMPLATE.format(problem_statement=problem_statement)
-        return await self.run(args, prompt, max_retries=3, key="solver")
+        return await self.run(args, prompt, max_retries=3, key="solver", worker_id=worker_id)
 
 
 class RewriterAgent(Agent):
@@ -121,7 +142,7 @@ class RewriterAgent(Agent):
     def __init__(self):
         super().__init__()
 
-    async def rewrite(self, args, problem_statement, previous_solutions: list[str]) -> str:
+    async def rewrite(self, args, problem_statement, previous_solutions: list[str], worker_id: int) -> str:
         """Generates the rewrited solution."""
 
         # Build the prompt template dynamically.
@@ -133,7 +154,7 @@ async def rewrite(self, args, problem_statement, previous_solutions: list[str])
             format_params[f"solution{i+1}"] = solution
 
         prompt = template.format(**format_params)
-        return await self.run(args, prompt, max_retries=1, key="rewriter")
+        return await self.run(args, prompt, max_retries=1, key="rewriter", worker_id=worker_id)
 
 
 class SelectorAgent(Agent):
@@ -142,7 +163,7 @@ class SelectorAgent(Agent):
     def __init__(self):
         super().__init__()
 
-    async def select(self, args, problem_statement, candidate_solutions: list[str]) -> str:
+    async def select(self, args, problem_statement, candidate_solutions: list[str], worker_id: int = 0) -> str:
         """Generates the rewrited solution."""
 
         # Build the prompt template dynamically.
@@ -154,7 +175,7 @@ async def select(self, args, problem_statement, candidate_solutions: list[str])
             format_params[f"solution{i+1}"] = solution
 
         prompt = template.format(**format_params)
-        return await self.run(args, prompt, max_retries=10, key="selector")
+        return await self.run(args, prompt, max_retries=10, key="selector", worker_id=worker_id)
 
     def extract_selected_solution_idx(self, response: str, candidate_solutions: list[str]) -> int:
         """Extracts the selected solution ID from the response."""
@@ -173,7 +194,7 @@ def extract_selected_solution_idx(self, response: str, candidate_solutions: list
 
 async def rewrite_worker(args, previous_solutions, problem_statement, worker_id):
     rewriter = RewriterAgent()
-    new_solution = await rewriter.rewrite(args, problem_statement, previous_solutions)
+    new_solution = await rewriter.rewrite(args, problem_statement, previous_solutions, worker_id)
     return new_solution
 
 
@@ -184,7 +205,7 @@ async def solver_worker(args, problem_statement, worker_id):
 
     try:
         solver = SolverAgent()
-        current_solution = await solver.generate_initial_solution(args, problem_statement)
+        current_solution = await solver.generate_initial_solution(args, problem_statement, worker_id)
         return current_solution
 
     except Exception as e:
@@ -245,7 +266,7 @@ def reward_adjustment(samples, reward_weight):
 
     # Selection
     selector = SelectorAgent()
-    response = await selector.select(args, problem_statement, rewrited_solutions)
+    response = await selector.select(args, problem_statement, rewrited_solutions, worker_id=0)
     if len(args.results_dict["selector"]) == 0:
         reward_adjustment(args.results_dict["solver"], args.incorrect_reward_weight)
         reward_adjustment(args.results_dict["rewriter"], args.incorrect_reward_weight)
@@ -254,13 +275,18 @@ def reward_adjustment(samples, reward_weight):
     assert (
         len(args.results_dict["selector"]) == 1
     ), f"selector should only return one solution, but got {len(args.results_dict['selector'])}"
+    selector_sample = args.results_dict["selector"][0]
     if response is None:
-        args.results_dict["selector"][0].reward = 0
+        selector_sample.reward = 0
+        selector_sample.metadata["selector_parse_success"] = False
     else:
         selected_solution_idx = selector.extract_selected_solution_idx(response, rewrited_solutions)
         if selected_solution_idx is None:
-            args.results_dict["selector"][0].reward = 0
+            selector_sample.reward = 0
+            selector_sample.metadata["selector_parse_success"] = False
         else:
+            selector_sample.metadata["selector_parse_success"] = True
+            selector_sample.metadata["selector_choice"] = selected_solution_idx + 1
             selected_solution = rewrited_solutions[selected_solution_idx]
             for sample in args.results_dict["rewriter"]:
                 if sample.response_content is not None and selected_solution in sample.response_content:

examples/multi_agent/rollout_with_multi_agents.py

@@ -28,6 +28,21 @@ async def generate_with_multi_agents(args, sample: Sample, sampling_params, eval
     custom_multi_agent_func = load_function(args.custom_multi_agent_function_path)
     samples = await custom_multi_agent_func(args, sample)
 
+    # VIME compact rollouts return multiple training samples from one source
+    # sample. Newer VIME requires all siblings to share a rollout_id so loss
+    # reduction counts the source rollout once instead of over-counting agents.
+    compact_rollout_id = (
+        sample.rollout_id
+        if sample.rollout_id is not None
+        else (
+            sample.index
+            if sample.index is not None
+            else sample.group_index if sample.group_index is not None else id(sample)
+        )
+    )
+    for sibling in samples:
+        sibling.rollout_id = compact_rollout_id
+
     random.shuffle(samples)
 
     return samples

examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh

@@ -0,0 +1,150 @@
+#!/bin/bash
+
+# for rerun the task
+pkill -9 -f '[v]llm serve|VLL[M]::'
+pkill -9 -f VLLM
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+pkill -9 redis
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+
+export SLIME_SCRIPT_TRAIN_BACKEND=megatron
+export PYTHONPATH="/workspace/wky/Megatron-Bridge/src:/workspace/wky/Megatron-LM/:${PYTHONPATH:-}"
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=1
+export DISABLE_L2_CACHE=1
+export VLLM_ASCEND_ENABLE_NZ=0
+export VLLM_USE_AOT_COMPILE=0
+
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
+
+VIME_ROOT="${VIME_ROOT:-/workspace/wky/vime-ascend}"
+SCRIPT_DIR="${VIME_ROOT}/scripts"
+WEIGHT_DIR="${WEIGHT_DIR:-/home/data/weights/Qwen3-4B}"
+DATA_FILE="${DATA_FILE:-/home/w00893744/dataset/dapo-math-17k.jsonl}"
+RUN_TS="${RUN_TS:-$(date +%Y%m%d_%H%M%S)}"
+LOG_FILE="${LOG_FILE:-/home/w00893744/train_qwen3_4b_multi_agent_vllm_${RUN_TS}.log}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+
+cd "${VIME_ROOT}"
+source "${SCRIPT_DIR}/models/qwen3-4B.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint "${WEIGHT_DIR}"
+   --load "${WEIGHT_DIR}"
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --custom-generate-function-path examples.multi_agent.rollout_with_multi_agents.generate_with_multi_agents
+   --prompt-data "${DATA_FILE}"
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type math
+
+   --rollout-backend vllm
+   --vllm-weight-sync-mode native
+   --vllm-gpu-memory-utilization 0.6
+   --vllm-enable-sleep-mode
+   --vllm-max-model-len 4096
+   --vllm-enforce-eager
+
+   --num-rollout 200
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-context-len 4096
+   --rollout-max-response-len 2048
+   --rollout-temperature 1.0
+
+   --global-batch-size 256
+   --balance-data
+)
+
+EVAL_ARGS=(
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 4
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 8192
+   --micro-batch-size 1
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --kl-loss-coef 0.0
+   --kl-loss-type low_var_kl
+   --kl-coef 0.00
+   --entropy-coef 0.0
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+WANDB_ARGS=(
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 4
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --use-flash-attn
+   --train-memory-margin-bytes 2147483648
+)
+
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+ray job submit --address="http://127.0.0.1:8265" \
+   -- python3 train.py \
+   --train-backend megatron \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 4 \
+   --rollout-num-gpus 4 \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${VLLM_ARGS[@]} \
+   ${MISC_ARGS[@]}