[example] add retool v2 example with multi-turn framework interfaces

examples/retool/generate_with_retool.py

@@ -96,12 +96,11 @@ def format_conversation_with_tools(
 
 def postprocess_predictions(prediction: str):
     """Extract action and content from prediction string"""
-    # Check for Answer: \boxed{...} format (only format we need for math_dapo)
-    # Use a more robust regex that handles nested braces
-    answer_pattern = r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}"
-    answer_match = re.search(answer_pattern, prediction, re.DOTALL)
-    if answer_match:
-        content = answer_match.group(1).strip()
+    # Check for bare \boxed{...} (model may omit "Answer:" prefix)
+    boxed_pattern = r"\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}"
+    boxed_match = re.search(boxed_pattern, prediction, re.DOTALL)
+    if boxed_match:
+        content = boxed_match.group(1).strip()
         return "answer", content
 
     # Then check for <tool_call> tags (new format from Jinja2 template)
@@ -168,14 +167,17 @@ def postprocess_responses(resp: str) -> str:
             last_match = matches[-1]
             return resp[: last_match.end()]
 
-    # Handle Answer: \boxed{...} format (only format we need for math_dapo)
-    if "Answer:" in resp and "\\boxed{" in resp:
-        # Find the last occurrence of Answer: \boxed{...} with nested braces support
-        answer_pattern = r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}"
-        matches = list(re.finditer(answer_pattern, resp, re.DOTALL))
-        if matches:
-            last_match = matches[-1]
-            return resp[: last_match.end()]
+    # Handle Answer: \boxed{...} or bare \boxed{...}
+    if "\\boxed{" in resp:
+        # Try "Answer: \boxed{...}" first, then bare "\boxed{...}"
+        for pattern in [
+            r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}",
+            r"\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}",
+        ]:
+            matches = list(re.finditer(pattern, resp, re.DOTALL))
+            if matches:
+                last_match = matches[-1]
+                return resp[: last_match.end()]
 
     return resp
 
@@ -203,7 +205,7 @@ async def execute_predictions(prediction: str) -> str:
         next_obs = (
             "\nMy previous action is invalid. "
             "If I want to execute code, I should put the code between "
-            "<code> and </code>. "
+            "<tool_call> and </tool_call>. "
             "If I want to give the final answer, I should use the format "
             "'Answer: \\boxed{answer}'. Let me try again.\n"
         )
@@ -221,7 +223,12 @@ async def generate(args, sample: Sample, sampling_params) -> Sample:
 
     # Set up the initial prompt with system prompt and tools (outside the loop)
     tool_specs = tool_registry.get_tool_specs()
-    prompt = format_conversation_with_tools(prompt=sample.prompt, tools=tool_specs)
+
+    if isinstance(sample.prompt, str):
+        # Already formatted (e.g., by --apply-chat-template), use as-is to avoid double templating
+        prompt = sample.prompt
+    else:
+        prompt = format_conversation_with_tools(prompt=sample.prompt, tools=tool_specs)
 
     prompt_tokens_ids = state.tokenizer(prompt, add_special_tokens=False)["input_ids"]
     response = ""
@@ -355,8 +362,7 @@ async def reward_func(args, sample, **kwargs):
     if not isinstance(sample, Sample):
         raise TypeError("Sample must be an instance of Sample class.")
 
-    # Build complete solution string
-    solution_str = sample.prompt + sample.response
+    solution_str = sample.response
 
     # Get ground truth answer - label is a string, not a dict
     ground_truth = sample.label if sample.label is not None else ""

examples/retool/retool_qwen3_4b_rl.sh

@@ -31,7 +31,7 @@ CKPT_ARGS=(
    --ref-load /root/font-info/qwen3-4b-sft_torch_dist
    # --load /root/Qwen3-4B_miles/
    --save /root/font-info/qwen3-4b-sft/qwen3-4b-sft-multi-turn/
-   --save-interval 20
+   --save-interval 200
    --rotary-base 5000000
 )
 
@@ -43,12 +43,12 @@ ROLLOUT_ARGS=(
    --rollout-shuffle
    --reward-key score
    --num-rollout 3000
-   --rollout-batch-size 32
+   --rollout-batch-size 16
    --n-samples-per-prompt 8
    --rollout-max-response-len 8192
    --rollout-temperature 1
 
-   --global-batch-size 256
+   --global-batch-size 128
    --balance-data
 )
 
@@ -98,8 +98,8 @@ OPTIMIZER_ARGS=(
 
 WANDB_ARGS=(
    --use-wandb
-   --wandb-project miles-dapo
-   --wandb-group qwen3-4B-test-multi-turn
+   --wandb-project miles-dev-retool-v2
+   --wandb-group retool-v1-qwen3-4b-sft-new
    --wandb-key ${WANDB_KEY}
 )
 
@@ -117,6 +117,7 @@ MISC_ARGS=(
    --attention-softmax-in-fp32
    # need to comment this when using model with MLA
    --attention-backend flash
+   --log-passrate
 )
 
 CUSTOM_ARGS=(

examples/retool_v2/README.md

@@ -0,0 +1,31 @@
+# Retool v2
+
+This example is an upgraded version of [retool](../retool), using the updated interfaces provided by the miles framework to implement multi-turn RL training with tool calls in a cleaner way.
+
+## Key Differences from v1
+
+**v1 (retool)** requires manually implementing the full multi-turn conversation loop in `generate_with_retool.py`, directly depending on low-level `GenerateState` and `sglang_rollout` interfaces — resulting in verbose code tightly coupled to the framework internals.
+
+**v2 (retool_v2)** uses the framework's standard plugin interfaces. Users only need to implement three functions and mount them via command-line arguments:
+
+| Argument | Description |
+|----------|-------------|
+| `--custom-generate-function-path` | Uses the built-in `miles.rollout.generate_hub.multi_turn.generate` — no need to implement the multi-turn loop yourself |
+| `--generate-tool-specs-path` | Declare tool definitions (user-implemented) |
+| `--generate-execute-tool-function-path` | Implement tool execution logic (user-implemented) |
+| `--custom-rm-path` | Implement the reward function (user-implemented) |
+
+Users only need to focus on business logic (tool definitions, tool execution, reward calculation). Multi-turn scheduling, token concatenation, loss masking, etc. are all handled by the framework.
+
+## Files
+
+- `tool_sandbox.py`: Tool definitions (`tool_specs`), tool execution (`execute_tool`), reward function (`reward_func`), and sandboxed safe execution environment
+- `run_retool_multi_turn.py`: Training launch script
+
+## Quick Start
+
+```bash
+python examples/retool_v2/run_retool_multi_turn.py
+```
+
+For data and model preparation, refer to the [retool v1 README](../retool/README.md).

examples/retool_v2/run_retool_multi_turn.py

@@ -0,0 +1,208 @@
+import os
+from dataclasses import dataclass, field
+from typing import Literal
+
+import typer
+
+import miles.utils.external_utils.command_utils as U
+
+WANDB_PROJECT = "miles-dev-retool-v2"
+WANDB_GROUP = "sft-multi-turn-batch-32"
+
+
+@dataclass
+class ScriptArgs(U.ExecuteTrainConfig):
+    mode: Literal["normal", "debug_minimal"] = "normal"
+    run_id: str = field(default_factory=U.create_run_id)
+    hardware: Literal["H100", "GB200", "GB300"] = "H100"
+    num_gpus_per_node: int | None = None
+    use_sft_model: bool = True
+    save_path: str = "/root/Qwen3-4B_miles/retool_v2_multi_turn"
+    prompt_data: str = "/root/dapo-math-17k/dapo-math-17k.jsonl"
+    generate_max_turns: int = 16
+    rollout_num_gpus_per_engine: int = 2
+    extra_args: str = ""
+
+    # resolved in __post_init__, not set by user
+    hf_checkpoint: str = field(init=False)
+    ref_load: str = field(init=False)
+
+    def __post_init__(self):
+        self.num_gpus_per_node = self.num_gpus_per_node or U.NUM_GPUS_OF_HARDWARE[self.hardware]
+        if self.use_sft_model:
+            self.hf_checkpoint = "/root/font-info/qwen3-4b-sft"
+            self.ref_load = "/root/font-info/qwen3-4b-sft_torch_dist"
+        else:
+            self.hf_checkpoint = "/root/models/Qwen3-4B"
+            self.ref_load = "/root/models/Qwen3-4B_torch_dist"
+
+
+def _get_wandb_args() -> str:
+    WANDB_API_KEY = os.environ.get("WANDB_API_KEY")
+    return (
+        "--use-wandb "
+        f"--wandb-project {WANDB_PROJECT} "
+        f"--wandb-group {WANDB_GROUP} "
+        f"--wandb-key {WANDB_API_KEY} "
+    )
+
+
+def prepare(args: ScriptArgs):
+    U.exec_command("mkdir -p /root/dapo-math-17k /root/aime-2024")
+    U.exec_command("hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k")
+    U.exec_command("hf download --repo-type dataset zhuzilin/aime-2024 --local-dir /root/aime-2024")
+
+    if args.use_sft_model:
+        U.exec_command("mkdir -p /root/font-info")
+        U.exec_command(f"hf download font-info/qwen3-4b-sft-SGLang-RL --local-dir {args.hf_checkpoint}")
+        U.convert_checkpoint(
+            model_name="qwen3-4b-sft",
+            megatron_model_type="qwen3-4B",
+            num_gpus_per_node=args.num_gpus_per_node,
+            hf_checkpoint=args.hf_checkpoint,
+            dir_dst="/root/font-info",
+        )
+    else:
+        U.exec_command("mkdir -p /root/models")
+        U.exec_command("hf download Qwen/Qwen3-4B --local-dir /root/models/Qwen3-4B")
+        U.convert_checkpoint(
+            model_name="Qwen3-4B",
+            megatron_model_type="qwen3-4B",
+            num_gpus_per_node=args.num_gpus_per_node,
+            dir_dst="/root/models",
+        )
+
+
+def execute(args: ScriptArgs):
+    megatron_model_type = "qwen3-4B"
+
+    ckpt_args = (
+        f"--hf-checkpoint {args.hf_checkpoint} "
+        f"--ref-load {args.ref_load} "
+        f"--save {args.save_path} "
+        f"--save-interval {2 if args.mode == 'debug_minimal' else 1000} "
+        f"{'--rotary-base 5000000 ' if args.use_sft_model else ''}"
+    )
+
+    custom_args = (
+        "--custom-generate-function-path miles.rollout.generate_hub.multi_turn.generate "
+        "--generate-tool-specs-path examples.retool_v2.tool_sandbox.tool_specs "
+        "--generate-execute-tool-function-path examples.retool_v2.tool_sandbox.execute_tool "
+        "--generate-tool-call-parser qwen25 "
+        f"--generate-max-turns {args.generate_max_turns} "
+        "--log-multi-turn "
+    )
+
+    rollout_args = (
+        f"--prompt-data {args.prompt_data} "
+        "--input-key prompt "
+        "--label-key label "
+        "--apply-chat-template "
+        "--rollout-shuffle "
+        "--custom-rm-path examples.retool_v2.tool_sandbox.reward_func "
+        "--reward-key score "
+        "--num-rollout 3000 "
+        "--rollout-batch-size 32 "
+        "--n-samples-per-prompt 8 "
+        f"--rollout-max-response-len {100 if args.mode == 'debug_minimal' else 8192} "
+        "--rollout-temperature 1 "
+        "--global-batch-size 256 "
+        "--balance-data "
+    )
+
+    eval_args = ""
+    if args.mode != "debug_minimal":
+        eval_args = (
+            "--eval-interval 20 "
+            "--eval-prompt-data aime /root/aime-2024/aime-2024.jsonl "
+            "--n-samples-per-eval-prompt 16 "
+            "--eval-max-response-len 16384 "
+            "--eval-top-p 1 "
+        )
+
+    grpo_args = (
+        "--advantage-estimator grpo "
+        "--use-kl-loss "
+        "--kl-loss-coef 0.00 "
+        "--kl-loss-type low_var_kl "
+        "--entropy-coef 0.00 "
+        "--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 "
+    )
+
+    sglang_args = (
+        f"--rollout-num-gpus-per-engine {args.rollout_num_gpus_per_engine} " "--sglang-mem-fraction-static 0.7 "
+    )
+
+    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 "
+    )
+
+    misc_args = (
+        f"--actor-num-nodes {args.num_nodes} "
+        f"--actor-num-gpus-per-node {args.num_gpus_per_node} "
+        "--colocate "
+        # default dropout in megatron is 0.1
+        "--attention-dropout 0.0 "
+        "--hidden-dropout 0.0 "
+        # should be good for model performance
+        "--accumulate-allreduce-grads-in-fp32 "
+        "--attention-softmax-in-fp32 "
+        # need to comment this when using model with MLA
+        "--attention-backend flash "
+        "--log-passrate "
+    )
+
+    train_args = (
+        f"{ckpt_args} "
+        f"{rollout_args} "
+        f"{optimizer_args} "
+        f"{grpo_args} "
+        f"{_get_wandb_args()} "
+        f"{perf_args} "
+        f"{eval_args} "
+        f"{sglang_args} "
+        f"{misc_args} "
+        f"{custom_args} "
+        f"{args.extra_args} "
+    )
+
+    U.execute_train(
+        train_args=train_args,
+        config=args,
+        num_gpus_per_node=args.num_gpus_per_node,
+        megatron_model_type=megatron_model_type,
+        extra_env_vars={
+            "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1",
+            "PYTHONPATH": "/root/Megatron-LM/:/root/miles",
+        },
+    )
+
+
+@U.dataclass_cli
+def main(args: ScriptArgs):
+    prepare(args)
+    execute(args)
+
+
+if __name__ == "__main__":
+    typer.run(main)