[NPU] support multi-agent rollout pipeline#298
Conversation
Documentation build overview
42 files changed ·
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a compatibility helper for extracting token IDs and log probabilities from vLLM rollout responses, updates multi-agent rollouts to share a single rollout_id among sibling samples to prevent over-counting during loss reduction, and adds a training script for Qwen3-4B on NPU. The review feedback highlights several critical improvements: adding defensive checks to prevent TypeError when converting potentially None log probabilities, handling cases where the multi-agent function returns None or empty samples, and replacing hardcoded absolute paths in the shell script with dynamic path resolution to ensure portability.
| new_response_log_probs = [ | ||
| float(item.get("logprob", 0.0)) if isinstance(item, dict) else 0.0 | ||
| for item in content_items | ||
| ] |
There was a problem hiding this comment.
If item.get("logprob") is explicitly None or missing, calling float(None) will raise a TypeError and crash the rollout pipeline. We should defensively check if the logprob value is not None before converting it to a float.
| new_response_log_probs = [ | |
| float(item.get("logprob", 0.0)) if isinstance(item, dict) else 0.0 | |
| for item in content_items | |
| ] | |
| new_response_log_probs = [ | |
| float(item.get("logprob")) if isinstance(item, dict) and item.get("logprob") is not None else 0.0 | |
| for item in content_items | |
| ] |
| 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}" |
There was a problem hiding this comment.
The script contains hardcoded absolute paths specific to a single user's environment (e.g., /workspace/issue205/... and /home/w00893744/...). This makes the script non-portable and prone to immediate failure on other machines or for other users.
We should dynamically resolve the project root (VIME_ROOT) and workspace directory relative to the script's location, and avoid hardcoding user-specific home directories for DATA_FILE and LOG_FILE.
| 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}" | |
| SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" | |
| VIME_ROOT="${VIME_ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}" | |
| WORKSPACE_DIR="$(cd "${VIME_ROOT}/.." && pwd)" | |
| export PYTHONPATH="${WORKSPACE_DIR}/Megatron-Bridge/src:${WORKSPACE_DIR}/Megatron-Bridge:${WORKSPACE_DIR}/Megatron-LM:${VIME_ROOT}:${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 | |
| WEIGHT_DIR="${WEIGHT_DIR:-/home/data/weights/Qwen3-4B}" | |
| if [ -z "${DATA_FILE}" ]; then | |
| echo "Error: DATA_FILE environment variable is not set." >&2 | |
| exit 1 | |
| fi | |
| LOG_FILE="${LOG_FILE:-/tmp/train_qwen3_4b_multi_agent_vllm.log}" |
| for sibling in samples: | ||
| sibling.rollout_id = compact_rollout_id | ||
|
|
||
| random.shuffle(samples) |
There was a problem hiding this comment.
If the custom multi-agent function custom_multi_agent_func returns None or an empty list, iterating over samples or calling random.shuffle(samples) will raise a TypeError. We should add a defensive check to ensure samples is not None or empty before processing.
| for sibling in samples: | |
| sibling.rollout_id = compact_rollout_id | |
| random.shuffle(samples) | |
| if samples: | |
| for sibling in samples: | |
| if sibling is not None: | |
| sibling.rollout_id = compact_rollout_id | |
| random.shuffle(samples) | |
| else: | |
| samples = [] |
Windfeng8
changed the title
|
|
||
| source "${SCRIPT_DIR}/models/qwen3-4B.sh" | ||
|
|
||
| python "${VIME_ROOT}/train.py" \ |
There was a problem hiding this comment.
Please use ray to submit task instead of using python directly, just ref the gpu script.
There was a problem hiding this comment.
You should put this script under examples/multi_agent folder
| --balance-data | ||
| ) | ||
|
|
||
| EVAL_ARGS=( |
There was a problem hiding this comment.
EVAL_ARGS and WANDB_ARGS can be removed since they are not used.
There was a problem hiding this comment.
I think it can be reserved for future expansion.
|
Could you add a short PR description covering what changes were made to the multi-agent example and why? |
done |
| compact_rollout_id = ( | ||
| sample.rollout_id | ||
| if sample.rollout_id is not None | ||
| else sample.index |
There was a problem hiding this comment.
Which rule/design determines the priority selection for this sample?
There was a problem hiding this comment.
This is not a priority rule for selecting
the best candidate response. It only determines the shared rollout_id
for all sibling samples produced by one multi-agent rollout.
The intended precedence is:
- Use the existing
sample.rollout_idwhen available. - Otherwise use
sample.index, which is the stable per-sample ID
assigned by VIME.
All sibling samples share this ID so that the loss reducer treats them as
one rollout instead of over-counting them. group_index is a
prompt-level grouping key rather than a rollout identifier, and
id(sample) is process-local, so they should not be used as normal
fallbacks.
| 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" \ |
There was a problem hiding this comment.
Please launch via ray job submit (with runtime-env-json for PYTHONPATH / Ascend env), matching scripts/run-qwen3-4B-npu.sh and examples/multi_agent/run-qwen3-30B-A3B-multi-agent.sh. Direct python train.py after ray start is inconsistent with the Ascend launch convention and was already requested in review.
38186cc to
5098165
Compare
CalvinXKY
left a comment
CalvinXKY
left a comment
There was a problem hiding this comment.
LGTM. clear the DCO with sign off commit.
Signed-off-by: Windfeng8 <523758380@qq.com>
Signed-off-by: Windfeng8 <523758380@qq.com>
Signed-off-by: wky <wky@example.com>
Signed-off-by: wky <wky@example.com>
Signed-off-by: Windfeng8 <523758380@qq.com>
5098165 to
81fed0d
Compare
done |
|
Please rebase your branch onto the latest to address the NPU-CI failures. |
5 participants
Added local token/logprob parsing for compatibility with the latest VIME.Assigned the same to sibling agent samples.
[RFC] #299
Qwen3-4B 20-Rollout Pilot: Single vs. Multi-Agent Analysis
Technical Summary
The 20-rollout pilot completed for both conditions. Multi-Agent was higher in the observed AIME results:
pass@1increased from 0.333% to 0.500%,pass@8increased from 2.547% to 3.540%, and correct-sample efficiency per output token was 1.498× the Single baseline.Scope and Comparability
20260720_201425_pilot20_single_qwen3-4b_seed4220260720_213101_pilot20_multi_qwen3-4b_seed42The two conditions used the same model, data, seeds, rollout budget, and AIME evaluation budget. Multi-Agent training used solver / rewriter / selector collaboration. Evaluation deliberately used the standard single-generation path for both conditions; this comparison therefore measures the model-capability effect of Multi-Agent training, not the end-to-end benefit of Multi-Agent inference orchestration.
Main Results: Multi-Agent Shows a Small Observed Gain
pass@1(primary score)pass@8pass@11Multi - Singlemeans the Multi-Agent value minus the Single baseline value. A positive value indicates an observed increase; a negative value indicates a regression. For example,pass@1is0.00500 - 0.00333 = +0.00167, which is +0.167 percentage points, not +0.167 percent.pass@1estimates the probability that one randomly selected answer is correct.pass@8andpass@11estimate the probability of obtaining at least one correct answer with 8 or 11 samples, respectively. These are score metrics; “correct samples per million output tokens” is an efficiency metric and should not be interpreted as a replacement for correctness.