From 75ec62d1a788a26d4a3ca84c2ace75f661e6e432 Mon Sep 17 00:00:00 2001 From: flb Date: Mon, 29 Jun 2026 21:10:07 +0800 Subject: [PATCH 1/3] add gpu tau-bench example Signed-off-by: flb --- examples/tau-bench/README.md | 65 +++ examples/tau-bench/generate_with_tau.py | 105 ++++ examples/tau-bench/openai_tool_adapter.py | 67 +++ examples/tau-bench/run_qwen3_4B.sh | 151 ++++++ examples/tau-bench/tau1_mock.py | 39 ++ examples/tau-bench/trainable_agents.py | 581 ++++++++++++++++++++++ examples/tau-bench/vllm_tool_parser.py | 95 ++++ 7 files changed, 1103 insertions(+) create mode 100644 examples/tau-bench/README.md create mode 100644 examples/tau-bench/generate_with_tau.py create mode 100644 examples/tau-bench/openai_tool_adapter.py create mode 100644 examples/tau-bench/run_qwen3_4B.sh create mode 100644 examples/tau-bench/tau1_mock.py create mode 100644 examples/tau-bench/trainable_agents.py create mode 100644 examples/tau-bench/vllm_tool_parser.py diff --git a/examples/tau-bench/README.md b/examples/tau-bench/README.md new file mode 100644 index 000000000..85611d128 --- /dev/null +++ b/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 +``` diff --git a/examples/tau-bench/generate_with_tau.py b/examples/tau-bench/generate_with_tau.py new file mode 100644 index 000000000..3c331794e --- /dev/null +++ b/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 diff --git a/examples/tau-bench/openai_tool_adapter.py b/examples/tau-bench/openai_tool_adapter.py new file mode 100644 index 000000000..53c744f08 --- /dev/null +++ b/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) diff --git a/examples/tau-bench/run_qwen3_4B.sh b/examples/tau-bench/run_qwen3_4B.sh new file mode 100644 index 000000000..1603c4f3f --- /dev/null +++ b/examples/tau-bench/run_qwen3_4B.sh @@ -0,0 +1,151 @@ +#!/bin/bash + +if grep -q $'\r' "$0" 2>/dev/null; then + exec bash <(sed 's/\r$//' "$0") "$@" +fi + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' 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 + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +unset PYTORCH_CUDA_ALLOC_CONF PYTORCH_ALLOC_CONF + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../../scripts/models/qwen3-4B-Instruct-2507.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-4B-Instruct-2507/ + --ref-load /root/Qwen3-4B-Instruct-2507_torch_dist/ + --save /root/Qwen3-4B-Instruct-2507_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/tau-bench/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 /root/tau-bench/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 +) + +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"} +NUM_GPUS=2 + +ray start --head \ + --node-ip-address "${MASTER_ADDR}" \ + --num-gpus "${NUM_GPUS}" \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"VIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC\": \"900\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --train-backend megatron \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node "${NUM_GPUS}" \ + --rollout-num-gpus "${NUM_GPUS}" \ + --colocate \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${VLLM_ARGS[@]}" \ + "${CUSTOM_ARGS[@]}" \ + "${MISC_ARGS[@]}" diff --git a/examples/tau-bench/tau1_mock.py b/examples/tau-bench/tau1_mock.py new file mode 100644 index 000000000..4be62d629 --- /dev/null +++ b/examples/tau-bench/tau1_mock.py @@ -0,0 +1,39 @@ +import argparse +import json +import os + +from tau_bench.envs import get_env +from tau_bench.types import RunConfig + +ALL_DATA_MAPPINGS = {"retail": ["train", "test", "dev"], "airline": ["test"]} + + +def main(): + parser = argparse.ArgumentParser(description="Tau1 Mock Script") + parser.add_argument("--local_dir", required=True, help="Path to the local directory") + args = parser.parse_args() + + local_dir = args.local_dir + if not os.path.isdir(local_dir): + os.makedirs(local_dir) + config = RunConfig(model_provider="mock", user_model_provider="mock", user_strategy="human", model="mock") + for env, split in ALL_DATA_MAPPINGS.items(): + for s in split: + config.env = env + config.task_split = s + env_instance = get_env( + env_name=config.env, + user_strategy=config.user_strategy, + user_model=config.user_model, + task_split=config.task_split, + ) + output_path = os.path.join(local_dir, f"{env}_{s}_tasks.jsonl") + with open(output_path, "w") as f: + for i, task in enumerate(env_instance.tasks): + row = {"index": i, "metadata": task.model_dump()} + f.write(json.dumps(row) + "\n") + print(f"Saved preprocessed task indices for {env} ({s}) to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/tau-bench/trainable_agents.py b/examples/tau-bench/trainable_agents.py new file mode 100644 index 000000000..6e2043f9b --- /dev/null +++ b/examples/tau-bench/trainable_agents.py @@ -0,0 +1,581 @@ +"""Trainable tau-bench agent for vime vLLM rollout.""" + +from __future__ import annotations + +import base64 +import io +import json +import logging +import os +import uuid +from typing import Any + +import numpy as np +from openai_tool_adapter import create_openai_adapter +from tau_bench.agents.tool_calling_agent import RESPOND_ACTION_NAME +from tau_bench.envs import get_env +from tau_bench.types import Action, RunConfig + +from vime.rollout.vllm_rollout import ( + GenerateState, + _build_inference_sampling_params, + _coerce_flat_int_token_ids, + _mm_render_response_to_generate_body, +) +from vime.utils.http_utils import post +from vime.utils.types import Sample + +logger = logging.getLogger(__name__) + + +def patch_tau_user_retries() -> None: + """Reduce circuit-breaker fatality: more LiteLLM retries with backoff.""" + try: + import tau_bench.envs.user as user_mod + + user_mod.MAX_RETRIES = int(os.environ.get("TAU_USER_LITELLM_RETRIES", "30")) + user_mod.RETRY_DELAY_SECONDS = float(os.environ.get("TAU_USER_LITELLM_RETRY_DELAY", "2")) + except Exception: + pass + + +def _parse_choice_tokens_and_logprobs(choice: dict[str, Any]) -> tuple[list[int], list[float]]: + """Parse token_ids + logprobs from vLLM /inference/v1/generate choice.""" + tids_raw = choice.get("token_ids") + if not (isinstance(tids_raw, list) and tids_raw and all(isinstance(x, int) for x in tids_raw)): + return [], [] + tids = [int(x) for x in tids_raw] + lp = choice.get("logprobs") + if not isinstance(lp, dict): + return tids, [0.0] * len(tids) + content = lp.get("content") + if isinstance(content, list) and content: + log_probs = [ + float(content[i].get("logprob", 0.0)) if i < len(content) and isinstance(content[i], dict) else 0.0 + for i in range(len(tids)) + ] + return tids, log_probs + return tids, [0.0] * len(tids) + + +def _maybe_apply_routed_experts(args: Any, sample: Sample, choice: dict[str, Any]) -> None: + if choice.get("routed_experts") is None: + return + raw = base64.b64decode(choice["routed_experts"].encode("ascii"), validate=True) + arr = np.load(io.BytesIO(raw), allow_pickle=False) + sample.rollout_routed_experts = np.ascontiguousarray(arr.astype(np.int32, copy=True)).reshape( + len(sample.tokens) - 1, + args.num_layers, + args.moe_router_topk, + ) + + +class TauBenchEnv: + def __init__( + self, + *, + tau_config: RunConfig, + task_index: int | None = None, + max_turns: int = 30, + ): + self.tau_config = tau_config + self.task_index = task_index + self.max_turns = max_turns + self.turn = 0 + self.total_reward = 0.0 + self.info: dict[str, Any] = {} + self.env = None + self.openai_adapter = None + self.successful_tool_calls = 0 + self.total_tool_calls = 0 + self.format_correct_calls = 0 + + def reset(self): + self.turn = 0 + self.total_reward = 0.0 + self.info = {} + self.successful_tool_calls = 0 + self.total_tool_calls = 0 + self.format_correct_calls = 0 + + self.env = get_env( + env_name=self.tau_config.env, + user_strategy=self.tau_config.user_strategy, + user_model=self.tau_config.user_model, + user_provider=self.tau_config.user_model_provider, + task_split=self.tau_config.task_split, + task_index=self.task_index, + ) + + self.openai_adapter = create_openai_adapter( + tools_info=self.env.tools_info, + parser_type="qwen25", + ) + + env_reset_res = self.env.reset(task_index=self.task_index) if self.task_index is not None else self.env.reset() + observation = env_reset_res.observation + self.info = self._to_dict(env_reset_res.info) + + return { + "obs_str": observation, + "role": "user", + "wiki": self.env.wiki, + "tools_info": self.env.tools_info, + } + + def step(self, response_text: str): + self.turn += 1 + is_final_turn = self.turn >= self.max_turns + + openai_result = self.openai_adapter.parse_response_to_openai_format(response_text) + + if not openai_result["success"]: + logger.warning(f"Tool parsing failed: {openai_result.get('error')}") + return ( + { + "obs_str": "Failed to parse tool call. Please try again.", + "role": "tool", + }, + is_final_turn, + {"tool_executed": False, "parse_error": openai_result.get("error")}, + ) + + parsed = openai_result["parsed_result"] + agent_content, calls = parsed["normal_text"], parsed["calls"] + + if calls: + self.format_correct_calls += 1 + + action = self._call_to_action(calls, agent_content) + + is_tool_call = action.name != RESPOND_ACTION_NAME + if is_tool_call: + self.total_tool_calls += 1 + + try: + env_response = self.env.step(action) + except Exception as e: + logger.warning(f"Environment step failed: {e}") + return ( + { + "obs_str": f"Environment error: {e}", + "role": "tool", + }, + True, + {"tool_executed": False, "env_error": str(e)}, + ) + + self.total_reward = env_response.reward + self.info.update(self._to_dict(env_response.info)) + + obs_lower = env_response.observation.lower() if env_response.observation else "" + if is_tool_call and not obs_lower.startswith(("error", "failed", "invalid", "not found")): + self.successful_tool_calls += 1 + + if action.name != RESPOND_ACTION_NAME: + obs_role = "tool" + else: + obs_role = "user" + obs_content = env_response.observation + + done = env_response.done or is_final_turn + + return ( + { + "obs_str": obs_content, + "role": obs_role, + "reward": env_response.reward, + }, + done, + {"tool_executed": True, "action": action.name}, + ) + + def _call_to_action(self, calls: list[Any], text_response: str) -> Action: + action = Action(name=RESPOND_ACTION_NAME, kwargs={"content": text_response}) + if calls: + if len(calls) > 1: + logger.debug("Multiple tool calls identified, only taking first.") + tool_call = calls[0] + try: + params = ( + json.loads(tool_call["parameters"]) + if isinstance(tool_call["parameters"], str) + else tool_call["parameters"] + ) + if not isinstance(params, dict): + logger.warning(f"{params} does not follow dict structure for action") + else: + action = Action(name=tool_call["name"], kwargs=params) + except json.JSONDecodeError as e: + logger.warning(f"Failed to parse parameters as JSON: {e}") + return action + + def close(self): + pass + + def format_observation(self, observation: dict) -> dict: + observation = observation or {} + content = observation.get("obs_str", "") + return { + "role": observation.get("role", "user"), + "content": content + "\n/no_think" if observation.get("role") == "user" else content, + } + + @staticmethod + def _to_dict(info: Any) -> dict: + if hasattr(info, "model_dump"): + return info.model_dump() + if isinstance(info, dict): + return info + return {} + + +def build_env(sample: Sample | None = None, args: Any | None = None, **_: Any) -> TauBenchEnv: + tau_bench_config = getattr(args, "tau_bench_config", None) + if tau_bench_config is None: + raise RuntimeError("args.tau_bench_config is missing; generate_with_tau.generate must set it from TAU_CONFIGS") + + task_index = None + if sample is not None and sample.prompt is not None: + try: + task_index = int(sample.prompt) + except (ValueError, TypeError): + pass + + max_turns = getattr(args, "max_turns", 30) + if max_turns is None: + max_turns = 30 + + return TauBenchEnv( + tau_config=tau_bench_config, + task_index=task_index, + max_turns=max_turns, + ) + + +def compute_process_reward(env: TauBenchEnv, base_reward: float) -> float: + reward = 0.0 + if base_reward > 0: + reward += 1.0 + if env.successful_tool_calls > 0: + reward += 0.1 * env.successful_tool_calls + if env.format_correct_calls > 0: + reward += 0.05 * env.format_correct_calls + reward = min(reward, 1.5) + return reward + + +def _build_tools_section(tools_info: list[dict]) -> str: + if not tools_info: + return "" + tools_json = json.dumps(tools_info, ensure_ascii=False) + parts = [ + "", + "", + "# Tools", + "", + "You may call one or more functions to assist with the user query.", + "", + "You are provided with function signatures within XML tags:", + "", + tools_json, + "", + "", + "For each function call, return a json object with function name and arguments within", + "", + '{"name": , "arguments": }', + "", + "XML tags.", + ] + return "\n".join(parts) + + +def _messages_for_render(messages: list[dict]) -> list[dict]: + out: list[dict] = [] + for msg in messages: + out.append({"role": msg["role"], "content": msg.get("content", "")}) + return out + + +class TrainableTauBenchAgent: + """Trainable tau-bench agent using vLLM render + /inference/v1/generate.""" + + async def asolve(self, args: Any, sample: Sample, sampling_params) -> Sample: + assert not args.partial_rollout, "Partial rollout is not supported for tau-bench interactions." + + state = GenerateState(args) + base_url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}" + + sample.metadata = sample.metadata or {} + + headers = None + if getattr(args, "router_policy", None) == "consistent_hash": + sample.session_id = sample.session_id or str(uuid.uuid4()) + headers = {"x-session-id": sample.session_id} + + try: + env = build_env(sample=sample, args=args) + except TypeError: + env = build_env(sample, args) + + initial_obs = env.reset() + wiki = initial_obs.get("wiki", "") + short_wiki_chars = os.environ.get("TAU_SHORT_WIKI") + if short_wiki_chars is not None: + wiki = wiki[: int(short_wiki_chars)] + tools_info = initial_obs.get("tools_info", []) + if os.environ.get("TAU_SHORT_TOOLS") == "1": + keep = {"find_user_id_by_email", "find_user_id_by_name_zip", "get_user_details", "get_order_details"} + tools_info = [tool for tool in tools_info if tool.get("function", {}).get("name") in keep] + tools_section = _build_tools_section(tools_info) + + messages: list[dict] = [ + {"role": "system", "content": wiki + tools_section + "\n/no_think"}, + {"role": "user", "content": initial_obs.get("obs_str", "")}, + ] + + response_tokens: list[int] = [] + sample.loss_mask = sample.loss_mask or [] + sample.rollout_log_probs = sample.rollout_log_probs or [] + sample.tokens = list(sample.tokens) if sample.tokens else [] + + sampling_params = sampling_params.copy() + sampling_params["repetition_penalty"] = 1.1 + sampling_params.setdefault("stop", [""]) + inference_sampling_params = _build_inference_sampling_params(sampling_params) + max_response_budget = sampling_params.get("max_new_tokens") + + def remaining_budget() -> int | None: + return None if max_response_budget is None else max_response_budget - sample.response_length + + def _ensure_trainable_skeleton() -> None: + """Megatron padding needs total_length >= 1 (F.pad uses prompt_length - 1).""" + if not sample.tokens: + eos_id = getattr(state.tokenizer, "eos_token_id", None) + pad_id = getattr(state.tokenizer, "pad_token_id", None) + token_id = eos_id if eos_id is not None else pad_id if pad_id is not None else 0 + sample.tokens = [int(token_id)] + sample.loss_mask = sample.loss_mask or [] + sample.rollout_log_probs = sample.rollout_log_probs or [] + if len(sample.rollout_log_probs) < len(sample.loss_mask): + sample.rollout_log_probs.extend([0.0] * (len(sample.loss_mask) - len(sample.rollout_log_probs))) + sample.response_length = len(sample.loss_mask) + + def _mark_truncated(reward_base: float = 0.0, *, remove: bool = True) -> Sample: + _ensure_trainable_skeleton() + if remove: + sample.remove_sample = True + sample.reward = compute_process_reward(env, reward_base) + sample.status = Sample.Status.TRUNCATED + return sample + + async def safe_render() -> dict | None: + try: + render_messages = _messages_for_render(messages) + payload = {"model": args.hf_checkpoint, "messages": render_messages} + render_data = await post( + f"{base_url}/v1/chat/completions/render", payload, headers=headers, max_retries=3 + ) + return _mm_render_response_to_generate_body(render_data, args.hf_checkpoint) + except Exception as exc: + logger.warning("render failed, skipping task: %s", exc) + return None + + def append_response_window( + token_ids: list[int], + loss_mask: list[int], + log_probs: list[float] | None = None, + ) -> None: + if not token_ids: + return + if len(loss_mask) != len(token_ids): + raise ValueError(f"loss_mask length {len(loss_mask)} != token_ids length {len(token_ids)}") + sample.tokens.extend(token_ids) + sample.loss_mask.extend(loss_mask) + sample.rollout_log_probs.extend(log_probs if log_probs is not None else [0.0] * len(token_ids)) + sample.response_length += len(token_ids) + + def sampling_params_for_turn() -> dict | None: + params = dict(inference_sampling_params) + max_tokens = remaining_budget() + if max_tokens is None: + return params + if max_tokens <= 0: + return None + params["max_tokens"] = max_tokens + return params + + try: + pending_obs_offset: int | None = None + rendered_body = await safe_render() + if rendered_body is None: + return _mark_truncated() + prompt_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids")) + if not sample.tokens: + sample.tokens = list(prompt_ids) + if args.rollout_max_context_len is not None: + context_budget = max(0, args.rollout_max_context_len - len(sample.tokens)) + if max_response_budget is None: + max_response_budget = context_budget + else: + max_response_budget = min(max_response_budget, context_budget) + + vllm_max_len = getattr(args, "vllm_max_model_len", 16384) or 16384 + if len(prompt_ids) >= vllm_max_len - 64: + logger.info(f"prompt too long ({len(prompt_ids)} tokens >= {vllm_max_len - 64}), skipping task") + return _mark_truncated() + + for turn_idx in range(args.max_turns): + input_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids")) + + if pending_obs_offset is not None: + obs_tokens = input_ids[pending_obs_offset:] + remaining = remaining_budget() + if remaining is not None and len(obs_tokens) > remaining: + append_response_window(obs_tokens[: max(remaining, 0)], [0] * max(remaining, 0)) + sample.status = Sample.Status.TRUNCATED + break + append_response_window(obs_tokens, [0] * len(obs_tokens)) + pending_obs_offset = None + + current_sampling_params = sampling_params_for_turn() + if current_sampling_params is None: + sample.status = Sample.Status.TRUNCATED + break + + body = dict(rendered_body) + body["sampling_params"] = current_sampling_params + output = await post(f"{base_url}/inference/v1/generate", body, headers=headers) + choice = output["choices"][0] + finish_reason = choice.get("finish_reason") or "stop" + new_tokens, new_logprobs = _parse_choice_tokens_and_logprobs(choice) + + if not new_tokens: + if finish_reason in ("abort", "cancelled"): + sample.status = Sample.Status.ABORTED + break + + response_text = state.tokenizer.decode(new_tokens, skip_special_tokens=False) if new_tokens else "" + train_tokens = list(new_tokens) + train_logprobs = list(new_logprobs) + train_loss_mask = [1] * len(train_tokens) + + stop = current_sampling_params.get("stop") + if not stop: + stop = [""] + stop_strings = (stop,) if isinstance(stop, str) else tuple(stop) if stop else () + hit_stop_str = None + hit_stop_pos = len(response_text) + if stop_strings: + for ss in stop_strings: + pos = response_text.find(ss) + if pos != -1 and pos < hit_stop_pos: + hit_stop_str = ss + hit_stop_pos = pos + if hit_stop_str is not None: + truncated_text = response_text[: hit_stop_pos + len(hit_stop_str)] + trunc_token_count = 0 + for t in range(1, len(train_tokens) + 1): + partial = state.tokenizer.decode(train_tokens[:t], skip_special_tokens=False) + if partial >= truncated_text: + trunc_token_count = t + break + if trunc_token_count == 0: + trunc_token_count = len(train_tokens) + train_tokens = train_tokens[:trunc_token_count] + train_logprobs = train_logprobs[:trunc_token_count] + train_loss_mask = train_loss_mask[:trunc_token_count] + response_text = truncated_text + finish_reason = "stop" + + eos_token_id = getattr(state.tokenizer, "eos_token_id", None) + append_stop_eos = ( + stop + and eos_token_id is not None + and getattr(args, "append_eos_token_after_stop_str_in_multi_turn", True) + ) + if append_stop_eos: + already_has_eos = bool(train_tokens and train_tokens[-1] == eos_token_id) + if stop_strings and response_text.endswith(stop_strings) and not already_has_eos: + if getattr(args, "use_rollout_routing_replay", False): + raise RuntimeError( + "Routing replay is not supported when appending an artificial EOS after a stop string, " + "because vLLM does not return routed experts for that extra token." + ) + train_tokens.append(int(eos_token_id)) + train_logprobs.append(0.0) + train_loss_mask.append(0) + + response_tokens.extend(new_tokens) + append_response_window(train_tokens, train_loss_mask, train_logprobs) + _maybe_apply_routed_experts(args, sample, choice) + + messages.append({"role": "assistant", "content": response_text}) + + if finish_reason == "length": + sample.status = Sample.Status.TRUNCATED + break + if finish_reason in ("abort", "cancelled"): + sample.status = Sample.Status.ABORTED + break + + observation, done, step_info = env.step(response_text) + + if done: + base_reward = observation.get("reward", 0.0) + sample.reward = compute_process_reward(env, base_reward) + sample.status = Sample.Status.COMPLETED + break + + next_user_message = env.format_observation(observation) + messages.append(next_user_message) + + if turn_idx + 1 >= args.max_turns: + sample.reward = compute_process_reward(env, 0.0) + sample.status = Sample.Status.TRUNCATED + break + + pending_obs_offset = len(input_ids) + len(train_tokens) + max_ctx = args.rollout_max_context_len or 8192 + if len(sample.tokens) >= max_ctx - 64: + logger.info( + f"[turn={turn_idx}] context overflow: {len(sample.tokens)} tokens >= {max_ctx - 64}, truncating" + ) + sample.reward = compute_process_reward(env, 0.0) + sample.status = Sample.Status.TRUNCATED + break + rendered_body = await safe_render() + if rendered_body is None: + return _mark_truncated() + rendered_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids")) + is_prefix_stable = rendered_ids[:pending_obs_offset] == sample.tokens[:pending_obs_offset] + sample.metadata["multiturn_render"] = { + "prefix_stable": is_prefix_stable, + "prefix_len": pending_obs_offset, + "sample_len": len(sample.tokens), + "rendered_len": len(rendered_ids), + "turn": turn_idx + 1, + } + if getattr(args, "strict_multiturn_render_token_match", False) and not is_prefix_stable: + raise RuntimeError( + "Full conversation render is not prefix-stable with the generated token stream: " + f"{sample.metadata['multiturn_render']}" + ) + + sample.response = state.tokenizer.decode(response_tokens, skip_special_tokens=False) + sample.response_length = len(sample.loss_mask) + _ensure_trainable_skeleton() + if sample.status == Sample.Status.PENDING: + sample.status = Sample.Status.COMPLETED + if sample.reward is None or sample.reward == 0.0: + sample.reward = compute_process_reward(env, getattr(env, "total_reward", 0.0)) + return sample + finally: + try: + env.close() + except Exception: + pass + + +def agent_factory() -> TrainableTauBenchAgent: + return TrainableTauBenchAgent() diff --git a/examples/tau-bench/vllm_tool_parser.py b/examples/tau-bench/vllm_tool_parser.py new file mode 100644 index 000000000..603c0f360 --- /dev/null +++ b/examples/tau-bench/vllm_tool_parser.py @@ -0,0 +1,95 @@ +"""Local tool-call parser for vLLM rollout.""" + +import json +import re +from typing import Any + + +def parse_tools(response: str, tools: list[dict[str, Any]], parser: str = "qwen25") -> dict[str, Any]: + if parser == "qwen25": + return _parse_qwen25_tools(response) + return _parse_qwen25_tools(response) + + +def _try_parse_json_tool_call(text: str) -> dict[str, Any] | None: + try: + parsed = json.loads(text) + if isinstance(parsed, dict) and "name" in parsed: + name = parsed.get("name", "") + parameters = parsed.get("arguments", parsed.get("parameters", {})) + if isinstance(parameters, str): + try: + parameters = json.loads(parameters) + except json.JSONDecodeError: + pass + return {"name": name, "parameters": parameters} + except (json.JSONDecodeError, TypeError): + pass + return None + + +def _parse_qwen25_tools(response: str) -> dict[str, Any]: + call_open = chr(60) + "tool_call" + chr(62) + call_close = chr(60) + "/tool_call" + chr(62) + call_open_alt = chr(60) + "call" + chr(62) + call_close_alt = chr(60) + "/call" + chr(62) + pattern = r"(?:" + call_open + "|" + call_open_alt + r")\s*(.*?)\s*(?:" + call_close + "|" + call_close_alt + ")" + tool_call_pattern = re.compile(pattern, re.DOTALL) + matches = tool_call_pattern.findall(response) + + if matches: + parts = tool_call_pattern.split(response) + normal_text = parts[0].strip() if parts else "" + calls = [] + for match in matches: + match = match.strip() + parsed_call = _try_parse_json_tool_call(match) + if parsed_call: + calls.append(parsed_call) + else: + try: + json_match = re.search(r"\{.*\}", match, re.DOTALL) + if json_match: + parsed_call = _try_parse_json_tool_call(json_match.group()) + if parsed_call: + calls.append(parsed_call) + else: + calls.append({"name": match, "parameters": {}}) + else: + calls.append({"name": match, "parameters": {}}) + except (json.JSONDecodeError, AttributeError): + calls.append({"name": match, "parameters": {}}) + + return { + "normal_text": normal_text, + "calls": calls, + } + + cleaned = re.sub(r"<\|im_end\|>", "", response).strip() + parsed_call = _try_parse_json_tool_call(cleaned) + if parsed_call: + return { + "normal_text": "", + "calls": [parsed_call], + } + + json_pattern = re.compile(r'\{[^{}]*"name"\s*:\s*"[^"]+?"[^{}]*\}', re.DOTALL) + json_matches = json_pattern.findall(response) + if json_matches: + calls = [] + for jm in json_matches: + parsed_call = _try_parse_json_tool_call(jm) + if parsed_call: + calls.append(parsed_call) + if calls: + normal_text = json_pattern.sub("", response).strip() + normal_text = re.sub(r"<\|im_end\|>", "", normal_text).strip() + return { + "normal_text": normal_text, + "calls": calls, + } + + return { + "normal_text": response, + "calls": [], + } From 66cecbe0e699a1726f31e732af94c7a21911710e Mon Sep 17 00:00:00 2001 From: flb Date: Tue, 30 Jun 2026 09:49:41 +0800 Subject: [PATCH 2/3] add npu tau-bench example script Signed-off-by: flb --- examples/tau-bench/run_qwen3_4B_npu.sh | 148 +++++++++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 examples/tau-bench/run_qwen3_4B_npu.sh diff --git a/examples/tau-bench/run_qwen3_4B_npu.sh b/examples/tau-bench/run_qwen3_4B_npu.sh new file mode 100644 index 000000000..c5bad23fc --- /dev/null +++ b/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[@]}" + From 6e20fe6357a397d2426de5265085faced41da9dc Mon Sep 17 00:00:00 2001 From: flb Date: Tue, 30 Jun 2026 15:55:42 +0800 Subject: [PATCH 3/3] drop GPU tau-bench script from ascend branch Signed-off-by: flb --- examples/tau-bench/run_qwen3_4B.sh | 151 ----------------------------- 1 file changed, 151 deletions(-) delete mode 100644 examples/tau-bench/run_qwen3_4B.sh diff --git a/examples/tau-bench/run_qwen3_4B.sh b/examples/tau-bench/run_qwen3_4B.sh deleted file mode 100644 index 1603c4f3f..000000000 --- a/examples/tau-bench/run_qwen3_4B.sh +++ /dev/null @@ -1,151 +0,0 @@ -#!/bin/bash - -if grep -q $'\r' "$0" 2>/dev/null; then - exec bash <(sed 's/\r$//' "$0") "$@" -fi - -# for rerun the task -pkill -9 -f '[v]llm serve|VLL[M]::' 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 - -# will prevent ray from buffering stdout/stderr -export PYTHONUNBUFFERED=1 - -unset PYTORCH_CUDA_ALLOC_CONF PYTORCH_ALLOC_CONF - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/../../scripts/models/qwen3-4B-Instruct-2507.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-4B-Instruct-2507/ - --ref-load /root/Qwen3-4B-Instruct-2507_torch_dist/ - --save /root/Qwen3-4B-Instruct-2507_vime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data /root/tau-bench/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 /root/tau-bench/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 -) - -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"} -NUM_GPUS=2 - -ray start --head \ - --node-ip-address "${MASTER_ADDR}" \ - --num-gpus "${NUM_GPUS}" \ - --disable-usage-stats \ - --dashboard-host=0.0.0.0 \ - --dashboard-port=8265 - -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"VIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC\": \"900\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --train-backend megatron \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node "${NUM_GPUS}" \ - --rollout-num-gpus "${NUM_GPUS}" \ - --colocate \ - "${MODEL_ARGS[@]}" \ - "${CKPT_ARGS[@]}" \ - "${ROLLOUT_ARGS[@]}" \ - "${EVAL_ARGS[@]}" \ - "${OPTIMIZER_ARGS[@]}" \ - "${GRPO_ARGS[@]}" \ - "${PERF_ARGS[@]}" \ - "${VLLM_ARGS[@]}" \ - "${CUSTOM_ARGS[@]}" \ - "${MISC_ARGS[@]}"