From d878cd78992748af26deb80eb3e40254c0e1d6f4 Mon Sep 17 00:00:00 2001 From: Windfeng8 <523758380@qq.com> Date: Tue, 30 Jun 2026 11:00:12 +0800 Subject: [PATCH 1/5] fix: resolve multi-agent Signed-off-by: Windfeng8 <523758380@qq.com> --- examples/multi_agent/agent_system.py | 18 +++- .../multi_agent/rollout_with_multi_agents.py | 15 ++++ scripts/run-qwen3-4B-multi-agent-npu.sh | 86 +++++++++++++++++++ 3 files changed, 118 insertions(+), 1 deletion(-) create mode 100755 scripts/run-qwen3-4B-multi-agent-npu.sh diff --git a/examples/multi_agent/agent_system.py b/examples/multi_agent/agent_system.py index b8fe390c5..22e50906b 100644 --- a/examples/multi_agent/agent_system.py +++ b/examples/multi_agent/agent_system.py @@ -5,13 +5,29 @@ 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 +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): try: sampling_params = args.sampling_params diff --git a/examples/multi_agent/rollout_with_multi_agents.py b/examples/multi_agent/rollout_with_multi_agents.py index d8a843c37..421e93d95 100644 --- a/examples/multi_agent/rollout_with_multi_agents.py +++ b/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 diff --git a/scripts/run-qwen3-4B-multi-agent-npu.sh b/scripts/run-qwen3-4B-multi-agent-npu.sh new file mode 100755 index 000000000..692798c7d --- /dev/null +++ b/scripts/run-qwen3-4B-multi-agent-npu.sh @@ -0,0 +1,86 @@ +#!/usr/bin/env bash +set -o pipefail + +export SLIME_SCRIPT_TRAIN_BACKEND=megatron +export PYTHONPATH="/workspace/issue205/Megatron-Bridge/src:/workspace/issue205/Megatron-Bridge:/workspace/issue205/Megatron-LM:/workspace/issue205/vime-ascend:${PYTHONPATH}" +export ASCEND_RT_VISIBLE_DEVICES="${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 MASTER_PORT="${MASTER_PORT:-$(shuf -i 20000-65000 -n 1)}" +export DISABLE_L2_CACHE=1 +export VLLM_ASCEND_ENABLE_NZ=0 + +VIME_ROOT="${VIME_ROOT:-/workspace/issue205/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}" +LOG_FILE="${LOG_FILE:-/home/w00893744/train_qwen3_4b_multi_agent_vllm.log}" + +source "${SCRIPT_DIR}/models/qwen3-4B.sh" + +python "${VIME_ROOT}/train.py" \ + --train-backend megatron \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --rollout-num-gpus 4 \ + --rollout-num-gpus-per-engine 4 \ + "${MODEL_ARGS[@]}" \ + --hf-checkpoint "${WEIGHT_DIR}" \ + --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 \ + --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 adam \ + --lr 1e-6 \ + --lr-decay-style constant \ + --weight-decay 0.1 \ + --adam-beta1 0.9 \ + --adam-beta2 0.98 \ + --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 \ + --load "${WEIGHT_DIR}" \ + --megatron-to-hf-mode bridge \ + --attention-dropout 0.0 \ + --hidden-dropout 0.0 \ + --accumulate-allreduce-grads-in-fp32 \ + --attention-softmax-in-fp32 \ + --attention-backend flash \ + --micro-batch-size 1 \ + --use-flash-attn \ + --train-memory-margin-bytes 2147483648 \ + 2>&1 | tee -a "${LOG_FILE}" From f5be233fc4606b71b2297871302ad1e2d3b690e1 Mon Sep 17 00:00:00 2001 From: Windfeng8 <523758380@qq.com> Date: Wed, 1 Jul 2026 15:53:42 +0800 Subject: [PATCH 2/5] fix: run-qwen3-4B-multi-agent-npu.sh Signed-off-by: Windfeng8 <523758380@qq.com> --- .../run-qwen3-4B-multi-agent-npu.sh | 148 ++++++++++++++++++ scripts/run-qwen3-4B-multi-agent-npu.sh | 86 ---------- 2 files changed, 148 insertions(+), 86 deletions(-) create mode 100755 examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh delete mode 100755 scripts/run-qwen3-4B-multi-agent-npu.sh diff --git a/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh b/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh new file mode 100755 index 000000000..5792849ef --- /dev/null +++ b/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh @@ -0,0 +1,148 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f "vllm serve" +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +export SLIME_SCRIPT_TRAIN_BACKEND=megatron +export PYTHONPATH="/workspace/issue205/Megatron-Bridge/src:/workspace/issue205/Megatron-Bridge:/workspace/issue205/Megatron-LM:/workspace/issue205/vime-ascend:${PYTHONPATH}" +export ASCEND_RT_VISIBLE_DEVICES="${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 MASTER_PORT="${MASTER_PORT:-$(shuf -i 20000-65000 -n 1)}" +export DISABLE_L2_CACHE=1 +export VLLM_ASCEND_ENABLE_NZ=0 + +VIME_ROOT="${VIME_ROOT:-/workspace/issue205/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 +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +python "${VIME_ROOT}/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[@]} \ + 2>&1 | tee -a "${LOG_FILE}" diff --git a/scripts/run-qwen3-4B-multi-agent-npu.sh b/scripts/run-qwen3-4B-multi-agent-npu.sh deleted file mode 100755 index 692798c7d..000000000 --- a/scripts/run-qwen3-4B-multi-agent-npu.sh +++ /dev/null @@ -1,86 +0,0 @@ -#!/usr/bin/env bash -set -o pipefail - -export SLIME_SCRIPT_TRAIN_BACKEND=megatron -export PYTHONPATH="/workspace/issue205/Megatron-Bridge/src:/workspace/issue205/Megatron-Bridge:/workspace/issue205/Megatron-LM:/workspace/issue205/vime-ascend:${PYTHONPATH}" -export ASCEND_RT_VISIBLE_DEVICES="${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 MASTER_PORT="${MASTER_PORT:-$(shuf -i 20000-65000 -n 1)}" -export DISABLE_L2_CACHE=1 -export VLLM_ASCEND_ENABLE_NZ=0 - -VIME_ROOT="${VIME_ROOT:-/workspace/issue205/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}" -LOG_FILE="${LOG_FILE:-/home/w00893744/train_qwen3_4b_multi_agent_vllm.log}" - -source "${SCRIPT_DIR}/models/qwen3-4B.sh" - -python "${VIME_ROOT}/train.py" \ - --train-backend megatron \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 4 \ - --rollout-num-gpus 4 \ - --rollout-num-gpus-per-engine 4 \ - "${MODEL_ARGS[@]}" \ - --hf-checkpoint "${WEIGHT_DIR}" \ - --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 \ - --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 adam \ - --lr 1e-6 \ - --lr-decay-style constant \ - --weight-decay 0.1 \ - --adam-beta1 0.9 \ - --adam-beta2 0.98 \ - --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 \ - --load "${WEIGHT_DIR}" \ - --megatron-to-hf-mode bridge \ - --attention-dropout 0.0 \ - --hidden-dropout 0.0 \ - --accumulate-allreduce-grads-in-fp32 \ - --attention-softmax-in-fp32 \ - --attention-backend flash \ - --micro-batch-size 1 \ - --use-flash-attn \ - --train-memory-margin-bytes 2147483648 \ - 2>&1 | tee -a "${LOG_FILE}" From 9632acace5cbf1aa9ae0b5e1f6b74ca508728644 Mon Sep 17 00:00:00 2001 From: wky Date: Wed, 22 Jul 2026 10:32:18 +0800 Subject: [PATCH 3/5] fix: submit Qwen3-4B multi-agent training via Ray Job Signed-off-by: wky --- .../run-qwen3-4B-multi-agent-npu.sh | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh b/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh index 5792849ef..9b8424be3 100755 --- a/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh +++ b/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh @@ -1,7 +1,8 @@ #!/bin/bash # for rerun the task -pkill -9 -f "vllm serve" +pkill -9 -f '[v]llm serve|VLL[M]::' +pkill -9 -f VLLM sleep 3 ray stop --force pkill -9 ray @@ -9,25 +10,27 @@ pkill -9 python sleep 3 pkill -9 ray pkill -9 python +pkill -9 redis set -ex -# will prevent ray from buffering stdout/stderr export PYTHONUNBUFFERED=1 export SLIME_SCRIPT_TRAIN_BACKEND=megatron -export PYTHONPATH="/workspace/issue205/Megatron-Bridge/src:/workspace/issue205/Megatron-Bridge:/workspace/issue205/Megatron-LM:/workspace/issue205/vime-ascend:${PYTHONPATH}" -export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}" +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 MASTER_PORT="${MASTER_PORT:-$(shuf -i 20000-65000 -n 1)}" export DISABLE_L2_CACHE=1 export VLLM_ASCEND_ENABLE_NZ=0 +export VLLM_USE_AOT_COMPILE=0 -VIME_ROOT="${VIME_ROOT:-/workspace/issue205/vime-ascend}" +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}" @@ -126,11 +129,11 @@ MISC_ARGS=( --train-memory-margin-bytes 2147483648 ) -# launch the master node of ray in container export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +ray start --head --node-ip-address ${MASTER_ADDR} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 -python "${VIME_ROOT}/train.py" \ +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 \ @@ -144,5 +147,4 @@ python "${VIME_ROOT}/train.py" \ ${PERF_ARGS[@]} \ ${EVAL_ARGS[@]} \ ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ - 2>&1 | tee -a "${LOG_FILE}" + ${MISC_ARGS[@]} From e7a5ff8eb5bd5f4c484e87058df5c6eff769529d Mon Sep 17 00:00:00 2001 From: wky Date: Wed, 22 Jul 2026 10:47:06 +0800 Subject: [PATCH 4/5] style: apply black formatting Signed-off-by: wky --- examples/multi_agent/agent_system.py | 3 +-- examples/multi_agent/rollout_with_multi_agents.py | 10 +++++----- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/examples/multi_agent/agent_system.py b/examples/multi_agent/agent_system.py index 22e50906b..51b3fa9cf 100644 --- a/examples/multi_agent/agent_system.py +++ b/examples/multi_agent/agent_system.py @@ -20,8 +20,7 @@ def _inference_generate_tokens_and_logprobs(choice): 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 + 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) diff --git a/examples/multi_agent/rollout_with_multi_agents.py b/examples/multi_agent/rollout_with_multi_agents.py index 421e93d95..b3e4725cb 100644 --- a/examples/multi_agent/rollout_with_multi_agents.py +++ b/examples/multi_agent/rollout_with_multi_agents.py @@ -34,11 +34,11 @@ async def generate_with_multi_agents(args, sample: Sample, sampling_params, eval 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) + 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 From c9f828020cd461594b9028db587d228eec48be2e Mon Sep 17 00:00:00 2001 From: Windfeng8 <523758380@qq.com> Date: Wed, 22 Jul 2026 11:45:24 +0800 Subject: [PATCH 5/5] feat: preserve multi-agent worker metadata Signed-off-by: Windfeng8 <523758380@qq.com> --- examples/multi_agent/agent_system.py | 39 ++++++++++++++++++---------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/examples/multi_agent/agent_system.py b/examples/multi_agent/agent_system.py index 51b3fa9cf..ec026860e 100644 --- a/examples/multi_agent/agent_system.py +++ b/examples/multi_agent/agent_system.py @@ -27,12 +27,18 @@ def _inference_generate_tokens_and_logprobs(choice): return new_response_tokens, new_response_log_probs -async def generate_response(args, prompt, key): +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" @@ -105,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}") @@ -124,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): @@ -136,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. @@ -148,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): @@ -157,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. @@ -169,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.""" @@ -188,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 @@ -199,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: @@ -260,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) @@ -269,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: