[NPU] add tau bench example for npu

examples/tau-bench/README.md

@@ -0,0 +1,65 @@
+# Tau bench
+This example shows vime training in an agentic multi-turn tool use environment.
+
+
+## Environment Setup
+This example assumes a vime container image. Install tau-bench dependencies:
+
+```bash
+cd /root/
+git clone https://github.com/JD-ETH/tau-bench.git
+cd tau-bench
+git checkout feature/litellm-retry
+pip install -e . --no-deps
+pip install litellm
+```
+
+Use the following script to generate task index jsonl for training:
+
+```bash
+cd /root/vime/examples/tau-bench
+python tau1_mock.py --local_dir /root/tau-bench/
+```
+
+Initialize the Qwen3-4B-Instruct-2507 model needed for tool use:
+
+```bash
+# hf checkpoint
+hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir /root/Qwen3-4B-Instruct-2507
+
+# mcore checkpoint
+cd /root/vime
+source scripts/models/qwen3-4B-Instruct-2507.sh
+PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
+    ${MODEL_ARGS[@]} \
+    --hf-checkpoint /root/Qwen3-4B-Instruct-2507 \
+    --save /root/Qwen3-4B-Instruct-2507_torch_dist
+```
+
+## Running the Script
+
+You need to configure your litellm API in generate_with_tau.py for user simulation:
+
+TAU_CONFIGS = {
+    "env": "retail",  # Select between ["retail", "airline"]
+    "agent_strategy": "tool-calling",  # Select between ["tool-calling", "act", "react", "few-shot"], only tool-calling implemented for now
+    "user_model": "gemini-2.0-flash-lite",  # Cheap Model for user simulator
+    "user_model_provider": "gemini",
+    "task_split": "train",  # Select between ["train", "test", "dev"] for retail, ["test"] for airline
+    "user_strategy": "llm",  # Select between ["llm", "react", "verify", "reflection"]
+    "model_provider": "auto_router", # Unused, required
+    "model": "qwen3-4b", # Unused, reqired
+}
+# Replace with your actual API key for user sim    
+GEMINI_API_KEY = "YOUR KEY" 
+
+Multi-turn limit: set env `TAU_MAX_TURNS` (default 10) or pass `--max-turns` to train.py.
+
+Agent rollout always uses vLLM (`/inference/v1/generate`); only `TAU_CONFIGS` controls the user simulator.
+
+And run:
+
+```bash
+cd /root/vime
+bash examples/tau-bench/run_qwen3_4B.sh
+```

examples/tau-bench/generate_with_tau.py

@@ -0,0 +1,105 @@
+"""Tau-bench multi-turn custom rollout for vime (vLLM render + generate)."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from typing import Any
+
+from tau_bench.types import RunConfig
+from trainable_agents import TrainableTauBenchAgent, agent_factory, patch_tau_user_retries
+
+from vime.utils.types import Sample
+
+logger = logging.getLogger(__name__)
+
+_TAU_DEFAULT_MAX_TURNS = 10
+
+_inflight_sem: asyncio.Semaphore | None = None
+
+# Tau-bench user-simulator configuration (edit TAU_CONFIGS below).
+# Agent rollout uses vLLM; only user_model / user_model_provider affect the user simulator here.
+TAU_CONFIGS = {
+    "env": "retail",  # Select between ["retail", "airline"]
+    "agent_strategy": "tool-calling",  # Select between ["tool-calling", "act", "react", "few-shot"]
+    # Default: local vLLM user sim (no external API). For Gemini API user sim, switch to:
+    # "user_model": "gemini-2.5-flash-lite", "user_model_provider": "gemini",
+    "user_model": "openai/local-qwen3-4b",
+    "user_model_provider": "openai",
+    "task_split": "train",  # Select between ["train", "test", "dev"] for retail
+    "user_strategy": "llm",  # Select between ["llm", "react", "verify", "reflection"]
+    "model_provider": "auto_router",  # Unused, required
+    "model": "qwen3-4b",  # Unused, required
+}
+# Replace with your actual API key when user_model_provider is gemini.
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "NONE")
+os.environ["GEMINI_API_KEY"] = GEMINI_API_KEY
+tau_config = RunConfig(**TAU_CONFIGS)
+
+
+def _get_inflight_sem() -> asyncio.Semaphore:
+    global _inflight_sem
+    if _inflight_sem is None:
+        _inflight_sem = asyncio.Semaphore(int(os.environ.get("TAU_MAX_INFLIGHT", "8")))
+    return _inflight_sem
+
+
+patch_tau_user_retries()
+
+
+def _ensure_tau_args(args: Any) -> None:
+    if getattr(args, "max_turns", None) is None:
+        env_max = os.environ.get("TAU_MAX_TURNS")
+        args.max_turns = int(env_max) if env_max is not None else _TAU_DEFAULT_MAX_TURNS
+
+
+def resolve_tau_config(args: Any) -> RunConfig:
+    """Build RunConfig from TAU_CONFIGS, with optional local-vLLM user-sim routing."""
+    user_model = tau_config.user_model
+    user_model_provider = tau_config.user_model_provider
+
+    if user_model_provider == "openai" and "local" in user_model:
+        vllm_router_host = getattr(args, "vllm_router_ip", "127.0.0.1")
+        vllm_router_port = getattr(args, "vllm_router_port", 3250)
+        vllm_model_name = getattr(args, "vllm_model_name", getattr(args, "hf_checkpoint", ""))
+        os.environ["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY", "dummy")
+        os.environ["OPENAI_API_BASE"] = f"http://{vllm_router_host}:{vllm_router_port}/v1"
+        user_model = vllm_model_name
+
+    return RunConfig(
+        env=tau_config.env,
+        agent_strategy=tau_config.agent_strategy,
+        user_model=user_model,
+        user_model_provider=user_model_provider,
+        task_split=tau_config.task_split,
+        user_strategy=tau_config.user_strategy,
+        model_provider=tau_config.model_provider,
+        model=tau_config.model,
+    )
+
+
+async def batched_tau_bench_rm(args, samples, **kwargs) -> list[float] | float:
+    if isinstance(samples, Sample):
+        return samples.reward if samples.reward is not None else 0.0
+    rewards = [s.reward if s.reward is not None else 0.0 for s in samples]
+    max_r = max(rewards) if rewards else 1.0
+    if max_r > 0:
+        rewards = [r / max_r for r in rewards]
+    return rewards
+
+
+async def generate(args: Any, sample: Sample, sampling_params) -> Sample:
+    assert not args.partial_rollout, "Partial rollout is not supported for tau-bench interactions."
+    _ensure_tau_args(args)
+    args.tau_bench_config = resolve_tau_config(args)
+
+    task_index = sample.prompt
+    logger.info(f"Starting agent-environment interaction for task {task_index}")
+
+    async with _get_inflight_sem():
+        agent: TrainableTauBenchAgent = agent_factory()
+        result = await agent.asolve(args, sample, sampling_params)
+
+    logger.info(f"Finished agent-environment interaction for task {task_index}")
+    return result

examples/tau-bench/openai_tool_adapter.py

@@ -0,0 +1,67 @@
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+try:
+    from .vllm_tool_parser import parse_tools
+except ImportError:
+    from vllm_tool_parser import parse_tools
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OpenAIToolCall:
+    id: str
+    type: str = "function"
+    function: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class OpenAIAssistantMessage:
+    role: str = "assistant"
+    content: str | None = None
+    tool_calls: list[OpenAIToolCall] | None = None
+
+
+class OpenAICompatibleToolCallAdapter:
+    def __init__(self, tools_info: list[dict[str, Any]], parser_type: str = "qwen25"):
+        self.tools_info = tools_info
+        self.parser_type = parser_type
+
+    def parse_response_to_openai_format(self, response: str) -> dict[str, Any]:
+        try:
+            parsed = parse_tools(response, self.tools_info, self.parser_type)
+            normal_text = parsed["normal_text"]
+            calls = parsed["calls"]
+            openai_message = self._convert_to_openai_message(normal_text, calls)
+            return {"openai_message": openai_message, "parsed_result": parsed, "success": True}
+        except Exception as e:
+            logger.warning(f"Parsing failed with error: {e}")
+            return {"openai_message": None, "parsed_result": None, "success": False, "error": str(e)}
+
+    def _convert_to_openai_message(self, normal_text: str, calls: list[dict[str, Any]]) -> OpenAIAssistantMessage:
+        if not calls:
+            return OpenAIAssistantMessage(role="assistant", content=normal_text, tool_calls=None)
+
+        openai_tool_calls = []
+        for i, call in enumerate(calls):
+            openai_tool_calls.append(
+                OpenAIToolCall(
+                    id=f"call_{i}_{call.get('name', 'unknown')}",
+                    type="function",
+                    function={"name": call.get("name", ""), "arguments": call.get("parameters", "{}")},
+                )
+            )
+
+        return OpenAIAssistantMessage(
+            role="assistant",
+            content=normal_text if normal_text.strip() else None,
+            tool_calls=openai_tool_calls,
+        )
+
+
+def create_openai_adapter(
+    tools_info: list[dict[str, Any]], parser_type: str = "qwen25"
+) -> OpenAICompatibleToolCallAdapter:
+    return OpenAICompatibleToolCallAdapter(tools_info, parser_type)

examples/tau-bench/run_qwen3_4B_npu.sh

@@ -0,0 +1,148 @@
+#!/bin/bash
+
+if grep -q $'\r' "$0" 2>/dev/null; then
+  exec bash <(sed 's/\r$//' "$0") "$@"
+fi
+
+# 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
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=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
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/../../scripts/models/qwen3-4B-Instruct-2507.sh"
+
+export PYTHONPATH="${SCRIPT_DIR}:/root/Megatron-Bridge/src:/root/Megatron-LM:${PYTHONPATH:-}"
+
+DATA_ROOT="${DATA_ROOT:-/root}"
+TAU_BENCH_ROOT="${TAU_BENCH_ROOT:-/root/tau-bench}"
+
+CKPT_ARGS=(
+   --hf-checkpoint ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/
+   --load ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/
+   --ref-load ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data ${TAU_BENCH_ROOT}/retail_train_tasks.jsonl
+   --input-key index
+   --rollout-shuffle
+   --num-rollout 500
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 4096
+   --rollout-max-context-len 16384
+   --rollout-temperature 0.7
+   --global-batch-size 256
+   --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std
+   --balance-data
+)
+
+EVAL_ARGS=(
+   --eval-interval 5
+   --eval-prompt-data retail-dev ${TAU_BENCH_ROOT}/retail_dev_tasks.jsonl
+   --n-samples-per-eval-prompt 1
+   --eval-max-response-len 4096
+   --eval-top-k 1
+)
+
+PERF_ARGS=(
+   --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 9216
+)
+
+GRPO_ARGS=(
+   --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_ARGS=(
+   --optimizer adam
+   --lr 5e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 1
+   --vllm-gpu-memory-utilization 0.7
+   --vllm-max-model-len 16384
+)
+
+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
+   --no-gradient-accumulation-fusion
+)
+
+CUSTOM_ARGS=(
+   --custom-generate-function-path generate_with_tau.generate
+   --custom-rm-path generate_with_tau.batched_tau_bench_rm
+)
+
+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[@]}" \
+   "${EVAL_ARGS[@]}" \
+   "${OPTIMIZER_ARGS[@]}" \
+   "${GRPO_ARGS[@]}" \
+   "${PERF_ARGS[@]}" \
+   "${VLLM_ARGS[@]}" \
+   "${CUSTOM_ARGS[@]}" \
+   "${MISC_ARGS[@]}"
+