[NPU] add tau bench example for npu

[floatlibai]

This pr mainly does the following:

1. mv the gpu tau-bench example from the main branch to the ascend branch [Example] Add tau-bench multi-turn tool-use example #142 fix(examples/tau-bench): use RunConfig.agent_strategy in TAU_CONFIGS (slime #2101) #280
2. add an npu tau-bench example script

Note: To use dynamic-sampling-filter, you need to cherry-pick #296 first, otherwise rollout abort may cause a hang.

vllm: 967c5c3bc38891f4465d3f4e99917ed837bb3833
vllm-ascend: 4fcffdae97aa13c2752350385bb04288951eb6b5

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #33369537 | 📁 Comparing 6e20fe6 against latest (d16d1dc)

  🔍 Preview build  

41 files changed · + 1 added · ± 27 modified · - 13 deleted

+ Added

• get_started/NPU.html

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/fault-tolerance.html
• advanced/megatron-config.html
• advanced/pd-disaggregation.html
• advanced/reproducibility.html
• advanced/speculative-decoding.html
• advanced/vllm-config.html
• and 17 more...

- Deleted

• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/low-precision.html
• advanced/observability.html
• examples/deepseek-r1.html
• examples/glm4-9B.html
• examples/glm4.7-30B-A3B.html
• examples/glm4.7-355B-A32B.html
• examples/glm5.2-744B-A40B.html
• examples/qwen3-4b-base-openhermes.html
• and 3 more...

[gemini-code-assist]

Code Review

This pull request introduces a multi-turn tau-bench training example with trainable agents, custom rollout logic, tool-calling adapters, and execution scripts. It also refactors the vLLM rollout abort mechanism to use a direct /abort_requests endpoint, avoiding worker pause/resume cycles and potential deadlocks. Feedback on the changes highlights a fragile lexicographical string comparison in trainable_agents.py that should be replaced with a length-based check, a specification violation in openai_tool_adapter.py where the arguments field must be JSON-serialized, and a redundant try-except block around build_env that can be simplified.

[floatlibai]

Smoke test script (click to expand)

#!/usr/bin/env bash

# for rerun the task
pkill -9 vllm 2>/dev/null || true
pkill -9 VLLM 2>/dev/null || true
sleep 3
ray stop --force 2>/dev/null || true
pkill -9 ray 2>/dev/null || true
pkill -9 -f 'python3 train.py' 2>/dev/null || true
sleep 3

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"

export VIME_ROOT="/root/vime"

ACTOR_GPUS="${ACTOR_GPUS:-4}"
ROLLOUT_GPUS="${ROLLOUT_GPUS:-4}"
NUM_ROLLOUT="${NUM_ROLLOUT:-200}"

TAU_DIR="/${VIME_ROOT}/../tau-bench"
TAU_EXAMPLE="${VIME_ROOT}/examples/tau-bench"
MODEL_ROOT="/home/data/weights/Qwen"
INSTRUCT="${MODEL_ROOT}/Qwen3-4B-Instruct-2507"
MODEL_SCRIPT="${VIME_ROOT}/scripts/models/qwen3-4B-Instruct-2507.sh"

# shellcheck source=/dev/null
source "${MODEL_SCRIPT}"
cd "${VIME_ROOT}"

export TAU_MAX_TURNS="${TAU_MAX_TURNS:-3}"
export TAU_SHORT_TOOLS="${TAU_SHORT_TOOLS:-0}"
ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-2048}"
MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-6144}"
VLLM_GPU_MEMORY_UTIL="${VLLM_GPU_MEMORY_UTIL:-0.60}"
TRAIN_MEMORY_MARGIN_BYTES="${TRAIN_MEMORY_MARGIN_BYTES:-4294967296}"
ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-8}"
N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-16}"
VLLM_SERVER_CONCURRENCY="${VLLM_SERVER_CONCURRENCY:-8}"
TAU_MAX_INFLIGHT="${TAU_MAX_INFLIGHT:-8}"
TAU_USER_LITELLM_RETRIES="${TAU_USER_LITELLM_RETRIES:-30}"
TAU_USER_LITELLM_RETRY_DELAY="${TAU_USER_LITELLM_RETRY_DELAY:-2}"
export TAU_MAX_INFLIGHT TAU_USER_LITELLM_RETRIES TAU_USER_LITELLM_RETRY_DELAY
export PYTHONPATH="${VIME_ROOT}/../Megatron-Bridge/src:${VIME_ROOT}/../Megatron-LM:${TAU_EXAMPLE}:${VIME_ROOT}:${PYTHONPATH:-}"

export PYTHONUNBUFFERED=1
export HYDRA_FULL_ERROR=1
export CUDA_DEVICE_MAX_CONNECTIONS=1
export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
export ASCEND_RT_VISIBLE_DEVICES=8,9,10,11,12,13,14,15
export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
export DISABLE_L2_CACHE=1
export VLLM_ASCEND_ENABLE_NZ=0
export VLLM_USE_AOT_COMPILE=0
export VIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC=900

unset PYTORCH_CUDA_ALLOC_CONF PYTORCH_ALLOC_CONF http_proxy https_proxy

python -c "import generate_with_tau"

set +e
python train.py \
  --train-backend megatron \
  --actor-num-nodes 1 \
  --actor-num-gpus-per-node "${ACTOR_GPUS}" \
  --rollout-num-gpus "${ROLLOUT_GPUS}" \
  --rollout-num-gpus-per-engine 1 \
  "${MODEL_ARGS[@]}" \
  --hf-checkpoint "${INSTRUCT}" \
  --ref-load "${INSTRUCT}" \
  --load "${INSTRUCT}" \
  --megatron-to-hf-mode bridge \
  --custom-generate-function-path generate_with_tau.generate \
  --custom-rm-path generate_with_tau.batched_tau_bench_rm \
  --prompt-data "${TAU_DIR}/retail_train_tasks.jsonl" \
  --input-key index \
  --rollout-shuffle \
  --num-rollout "${NUM_ROLLOUT}" \
  --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" \
  --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" \
  --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN}" \
  --rollout-max-context-len 16384 \
  --rollout-temperature 0.7 \
  --seq-length 16384 \
  --global-batch-size "${GLOBAL_BATCH_SIZE}" \
  --disable-rewards-normalization \
  --advantage-estimator grpo \
  --use-kl-loss \
  --kl-loss-coef 0.001 \
  --kl-loss-type low_var_kl \
  --entropy-coef 0.01 \
  --eps-clip 0.2 \
  --eps-clip-high 0.28 \
  --optimizer adam \
  --lr "${LR:-1e-6}" \
  --lr-decay-style constant \
  --weight-decay 0.1 \
  --adam-beta1 0.9 \
  --adam-beta2 0.98 \
  --tensor-model-parallel-size 2 \
  --sequence-parallel \
  --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 "${MAX_TOKENS_PER_GPU}" \
  --vllm-gpu-memory-utilization "${VLLM_GPU_MEMORY_UTIL}" \
  --vllm-server-concurrency "${VLLM_SERVER_CONCURRENCY}" \
  --router-disable-circuit-breaker \
  --vllm-max-model-len 16384 \
  --train-memory-margin-bytes "${TRAIN_MEMORY_MARGIN_BYTES}" \
  --attention-dropout 0.0 \
  --hidden-dropout 0.0 \
  --accumulate-allreduce-grads-in-fp32 \
  --attention-softmax-in-fp32 \
  --attention-backend flash \
  --use-flash-attn

Test Results

Environment: A3, Qwen3-4B-Instruct-2507, GRPO, bs 16, full 200 steps

Pipeline Verification:

• Ray cluster startup + vLLM rollout engine initialization (4 NPU train + 4 NPUrollout)
• Full cycle: multi-turn rollout → tau-bench env interaction → reward computation → weight update
• Model correctly generates -formatted tool calls (e.g., find_user_id_by_email, get_order_details, etc.)
• vLLM render route multi-turn token budget tracking works correctly
• Prefix stability validation passes across turns
• Training throughput: ~9525 tok/s, 200 steps 2h21min

[momo609]

/lgtm

[CalvinXKY]

If the .sh files are example scripts, I’d suggest keeping only one. The two scripts differ in many ways, so multiple variants could be derived from them. For example, the GPU version uses the --colocate flag, while the NPU one does not adopt that mode.

[floatlibai]

Do you mean removing the GPU script and keeping only the NPU one?

[CalvinXKY]

I would suggest keeping the NPU version.

[floatlibai]

I would suggest keeping the NPU version.

done

[CalvinXKY]

I noticed that the test uses 8 GPUs, but here it is marked as 16. Could you please confirm whether the scripts are consistent (i.e., the same for both)?

[floatlibai]

This example script also uses 8 NPUs. ASCEND_RT_VISIBLE_DEVICES is set to 0–15 just as a generic/common environment variable (not indicative of actual device usage).

[CalvinXKY]

LGTM