Add a multi-turn GRPO example with tool calling, and port OPSD to the per-sample rollout API

benchmarks/opsd/README.md

@@ -12,8 +12,7 @@ DeepSpeed PR #8295:
 
 https://github.com/deepspeedai/DeepSpeed/pull/8295
 
-Use a DeepSpeed checkout that contains that API and place it first on
-`PYTHONPATH`. The current validation scope is one process, one GPU, and ZeRO-0.
+It also builds its requests from per-sample `RolloutRecord`s, the rollout API of stage 1 of [RFC #8942](https://github.com/deepspeedai/DeepSpeed/issues/8942). Use a DeepSpeed checkout that contains both and place it first on `PYTHONPATH`. The current validation scope is one process, one GPU, and ZeRO-0.
 This is a HybridEngine rollout benchmark, not a complete OPSD training-step
 benchmark; it does not measure teacher inference, loss computation, backward,
 or optimizer work.

benchmarks/opsd/benchmark_hybrid_engine_rollout.py

@@ -119,13 +119,15 @@ def _build_engine(model, args):
 
 def _make_request(model, batch_size, prompt_length, device):
     import torch
-    from deepspeed.runtime.rollout.base import RolloutRequest
+    from deepspeed.runtime.rollout.base import RolloutRecord, RolloutRequest
 
     vocab_size = model.config.vocab_size
     first_token_id = min(3, vocab_size - 1)
     prompt_ids = torch.randint(first_token_id, vocab_size, (batch_size, prompt_length), device=device)
-    prompt_attention_mask = torch.ones_like(prompt_ids)
-    return RolloutRequest(prompt_ids=prompt_ids, prompt_attention_mask=prompt_attention_mask)
+    return RolloutRequest(samples=[
+        RolloutRecord(group_index=row, tokens=ids, attention_mask=[1] * prompt_length)
+        for row, ids in enumerate(prompt_ids.tolist())
+    ])
 
 
 def _run_case(rollout, model, args, batch_size, samples_per_prompt, prompt_length, response_length, device):

training/grpo/.gitignore

@@ -0,0 +1,3 @@
+checkpoints/
+data/
+__pycache__/

training/grpo/README.md

@@ -0,0 +1,116 @@
+# Multi-Turn GRPO with Tool Calling (ReTool)
+
+This example trains Qwen3 with GRPO to solve math problems while calling a Python interpreter. In each trajectory the model can write code, read the interpreter's output and continue reasoning, for several turns, before giving its answer.
+
+Rollout and training share one DeepSpeed **hybrid engine** (ZeRO-2): the policy generates with its own training weights, so there is no separate inference server and no weight sync. The example is built on `deepspeed.runtime.rollout`, which provides the multi-turn rollout, the tool interface and the GRPO loss.
+
+## How it works
+
+Each training step has three phases:
+
+1. **Rollout.** Each rank takes its share of the step's prompts and samples `n_samples_per_prompt` trajectories per prompt with `HybridEngineRollout`. When the model writes a `<tool_call>` block, the rollout runs it in the sandboxed `code_interpreter` (`tools.py`) and appends what the code prints as a tool message, for up to `max_turns` turns. With continuous batching, each turn takes a free decode row and releases it when the turn ends, and tools run while the other rows keep decoding.
+2. **Reward.** A DAPO strict-box check (`reward.py`) scores each trajectory +1 if the last `\boxed{...}` of the response equals the label, and −1 otherwise. `group_advantages` normalizes each reward within its prompt's group.
+3. **Update.** `ResponseTokenBatch.from_rollout` turns the trajectories into training targets with a mask that keeps only the tokens the model generated, so prompts, padding and tool output are not trained on. `grpo_policy_loss` computes PPO-clip with clip-higher (0.2 / 0.28) and no KL term. The loss is averaged over every generated token across ranks, or per trajectory first with `algorithm.loss_aggregation: "sample_mean"`, and the step takes one optimizer update, accumulated over micro-batches.
+
+With one update per rollout, the old policy is the current one, so the old log-probs are the current ones detached and the importance ratio is 1. `rollout.return_log_probs` records them during the rollout and takes them from there instead, through `old_log_probs_from_rollout`.
+
+## Layout
+
+| Path | Contents |
+|---|---|
+| `main.py` | Entry point: builds the model, the hybrid engine, the rollout and the data, then trains |
+| `trainer.py` | The training loop: rollout, reward, update, evaluation, logging and checkpoints |
+| `config.py` | The configuration dataclasses and their checks |
+| `policy.py` | The policy as DeepSpeed trains it: token log-probs from a chunked LM head |
+| `tools.py` | The sandboxed Python interpreter the model calls |
+| `reward.py` | The strict-box answer check |
+| `data.py` | JSONL prompt loading |
+| `configs/` | The three training configs |
+| `scripts/` | Data preparation and one launch script per config |
+
+## Prerequisites
+
+- **DeepSpeed with the multi-turn rollout.** `deepspeed.runtime.rollout` with multi-turn tool calling, continuous batching and the GRPO loss is not in a DeepSpeed release yet; it is stage 1 of [RFC #8942](https://github.com/deepspeedai/DeepSpeed/issues/8942). Install DeepSpeed from source with that code.
+- **Python packages.** `transformers` 5.0 or later, `numpy`, and `huggingface_hub` for the `hf` download tool:
+
+  ```bash
+  pip install "transformers>=5.0.0" numpy huggingface_hub
+  ```
+
+- **GPUs.** The smoke test runs on one GPU. The single-GPU config is sized for one 40 GB A100, and also needs more than 64 GB of host memory for the optimizer state it keeps on the CPU; 160 GB is enough. The distributed config is sized for 8 GPUs with 80 GB each.
+
+## Data
+
+The training set is `dapo-math-17k` and the evaluation set is `aime-2024`, in JSONL with a chat-message `prompt` and a string `label`. Download both into `data/`:
+
+```bash
+cd training/grpo
+bash scripts/prepare_data.sh
+```
+
+The script also wraps each AIME prompt in the instructions every `dapo-math-17k` prompt has. The downloaded problems don't ask for a `\boxed{}` answer, which is the only form the scorer accepts, so without them nearly every sample scores −1.
+
+## Run
+
+```bash
+cd training/grpo
+
+# smoke test: Qwen3-0.6B, three steps on one GPU
+bash scripts/run_smoke_test.sh
+
+# Qwen3-4B-Instruct-2507 on one GPU
+bash scripts/run_single_gpu.sh
+
+# Qwen3-8B on 8 GPUs
+bash scripts/run_8gpu.sh
+```
+
+Each script runs `deepspeed --num_gpus <N> main.py --config configs/<config>.json`.
+
+Rank 0 prints the loss and the mean reward after each step. When `eval.path` is set, it also prints the AIME accuracy, over every sample, before training, every `eval.interval` steps and at the end. Checkpoints go to `training.save_dir`: a resumable DeepSpeed checkpoint plus a Hugging Face copy in `hf_step_<N>`. To continue a run, set `training.resume_from` to its `save_dir`.
+
+## Configurations
+
+| | `smoke_qwen3_0.6b.json` | `qwen3_4b_instruct_1gpu.json` | `qwen3_8b_8gpu.json` |
+|---|---|---|---|
+| Model | `Qwen/Qwen3-0.6B` | `Qwen/Qwen3-4B-Instruct-2507` | `Qwen/Qwen3-8B` |
+| GPUs | 1 | 1 | 8 |
+| Steps | 3 | 500 | 500 |
+| Trajectories per step | 8 (2 prompts × 4) | 64 (8 prompts × 8) | 256 (32 prompts × 8) |
+| Trajectories per micro-batch, all ranks | 2 | 2 | 8 |
+| Turns, context, tokens per turn | 4, 4096, 1024 | 16, 8192, 4096 | 16, 8192, 4096 |
+| Continuous batching rows per GPU | 8 | 16 | 16 |
+| ZeRO-2 | optimizer on the GPU | optimizer state on the CPU | optimizer on the GPU |
+| AIME evaluation | off | every 100 steps | every 20 steps |
+| Thinking mode | off | none (an instruct model) | off |
+
+All three turn on every rollout optimization that works together: CUDA-graph decode (`use_graph_capture`) and multi-turn continuous batching (`continuous_batch_size`). Training uses bf16, gradient checkpointing, the chunked LM head, SDPA attention, and ZeRO-2 with overlapped communication.
+
+Each config holds the DeepSpeed config inline under `deepspeed_config`, which can also be the path of a DeepSpeed JSON. The trainer fills in `train_micro_batch_size_per_gpu` and `gradient_accumulation_steps`, so that one step covers all of a rank's trajectories.
+
+The main settings:
+
+- `rollout.max_new_tokens`: generated tokens per trajectory, **summed over its turns**.
+- `rollout.max_new_tokens_per_turn`: an additional cap on each turn. `eval` has the same two fields.
+- `rollout.max_context_len`: total prompt, response and tool tokens per trajectory. A trajectory that reaches it stops as truncated.
+- `rollout.continuous_batch_size`: decode rows per GPU for continuous batching; leave it unset for static batching. Rows are sized by `max_context_len`, so the KV cache holds `continuous_batch_size × max_context_len` tokens, about 144 KiB per token for Qwen3-4B and Qwen3-8B. Lower it if rollout runs out of memory.
+- `rollout.return_log_probs`: record each generated token's log-prob during the rollout and use them as the PPO old log-probs, off by default. It applies to evaluation rollouts too.
+- `rollout.generate_batch_size`: trajectories per `generate` call, in whole prompts; by default a rank's whole share of the step.
+- `rollout.chat_template_kwargs`: extra arguments for the chat template, such as `{"enable_thinking": false}` for Qwen3.
+- `rollout.repetition_penalty`: 1.0 by default; set it to null to use the model's `generation_config`.
+- `training.global_batch_size`: trajectories per step across all ranks, a multiple of ranks × `n_samples_per_prompt` so that each rank rolls out whole prompts.
+- `training.micro_batch_size`: trajectories per forward and backward pass across all ranks, split evenly between them. It must divide `global_batch_size`; gradient accumulation covers the rest, so it changes memory and speed, not the update.
+- `tool.timeout`, `tool.max_output_chars`: limits on each interpreter call; long output is cut so that one large print cannot fill the context.
+- `tool.show_last_value`: also return the value of a bare expression on the code's last line, as a notebook does, and say so in the tool description. Off by default, so only printed output comes back. Models that write notebook-style code, such as Qwen3-4B-Instruct-2507, which almost never ends its code with `print()`, otherwise get empty tool output.
+- `algorithm.tool_call_bonus`: an optional reward bonus for calling tools on wrong answers, off by default.
+- `algorithm.loss_aggregation`: `"token_mean"` (default) weighs every generated token of the step equally; `"sample_mean"` averages each trajectory over its own tokens first, then over trajectories.
+
+## Notes and limitations
+
+- **The sandbox only screens obvious misuse.** It applies an import allowlist, a pattern filter, a memory limit and a timeout. It is not a security boundary, so run untrusted training on isolated machines.
+- **ZeRO-2 only.** Multi-turn continuous batching does not support ZeRO-3, whose parameter gathers cannot run inside a captured graph or across rows that finish at different steps.
+- **Hybrid engine without injected kernels.** DeepSpeed has no inference policy for Qwen3, so generation runs the Hugging Face model on the training weights, with DeepSpeed's own static cache, graph capture and continuous batching around it.
+
+## Acknowledgement
+
+This example is adapted from the [miles retool_v2](https://github.com/radixark/miles/tree/main/examples/retool_v2) recipe (Copyright 2025-2026 RadixArk, Inc., Apache License 2.0). The code interpreter in `tools.py` and the answer check in `reward.py` are based on miles. The answer check is in turn adapted from verl and EleutherAI's lm-evaluation-harness, both under the Apache License 2.0.

training/grpo/config.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""GRPO application configuration.
+
+``GRPOConfig`` is loaded from JSON. The ``rollout`` sub-config is consumed
+by DeepSpeed's rollout engine; optimizer, ZeRO and hybrid-engine settings
+live in ``deepspeed_config``, either inline or as the path of a DeepSpeed JSON.
+"""
+
+import json
+from dataclasses import asdict, dataclass, field
+from typing import Optional, Union
+
+from deepspeed.runtime.rollout import RolloutConfig as _BaseRolloutConfig
+
+
+@dataclass
+class ModelConfig:
+    model_name_or_path: str
+    dtype: str = "bfloat16"
+    attn_implementation: str = "sdpa"
+    gradient_checkpointing: bool = True
+    trust_remote_code: bool = False
+
+
+@dataclass
+class RolloutConfig(_BaseRolloutConfig):
+    """Extends DeepSpeed's RolloutConfig with the sampling knobs of training rollouts.
+
+    The base config holds the engine settings: ``continuous_batch_size`` (decode rows for continuous
+    batching, sized by ``max_context_len``; None keeps static batching) and ``return_log_probs`` (record
+    each generated token's log-prob, used as the PPO old log-probs when set).
+    """
+    multi_turn: bool = True
+    max_turns: Optional[int] = 16
+    max_context_len: Optional[int] = 16384
+    # Generated tokens per trajectory, summed over its turns.
+    max_new_tokens: int = 8192
+    # Generated tokens per turn, on top of max_new_tokens; None caps turns only by what is left of it.
+    max_new_tokens_per_turn: Optional[int] = None
+    temperature: float = 1.0
+    # 1.0 disables the penalty; None takes the value in the model's generation_config.
+    repetition_penalty: Optional[float] = 1.0
+    n_samples_per_prompt: int = 8
+    # Trajectories per generate call, in whole prompts and at least one; None rolls out a rank's share of the
+    # step in one call. Bounds the KV cache and prefill memory of a call.
+    generate_batch_size: Optional[int] = None
+
+
+@dataclass
+class ToolConfig:
+    timeout: float = 120.0
+    memory_limit_gb: int = 4
+    max_concurrency: int = 32
+    # Long tool outputs are cut so one print cannot overflow the context.
+    max_output_chars: int = 4096
+    # Also return the value of a bare expression on the code's last line, as a notebook shows it; off returns only
+    # what the code prints.
+    show_last_value: bool = False
+
+
+@dataclass
+class AlgorithmConfig:
+    eps_clip: float = 0.2
+    eps_clip_high: float = 0.28
+    # Sequence positions per LM-head chunk when computing token log-probs.
+    logprob_chunk_size: int = 1024
+    # Reward bonus for tool calls on wrong answers.
+    tool_call_bonus: bool = False
+    # "token_mean": every generated token of the step weighs the same.
+    # "sample_mean": each trajectory's loss is averaged over its own tokens, then over trajectories.
+    loss_aggregation: str = "token_mean"
+
+
+@dataclass
+class TrainingConfig:
+    # Trajectories per step across all ranks, in whole prompts per rank: rolled out together, then
+    # one optimizer step over all of them.
+    global_batch_size: int = 256
+    # Trajectories per forward/backward pass across all ranks, split evenly between them; None means
+    # one pass per step, global_batch_size.
+    micro_batch_size: Optional[int] = None
+    num_rollouts: int = 3000
+    logging_steps: int = 1
+    save_steps: int = 1000
+    save_dir: str = "./grpo_ckpt"
+    # Continue from the latest checkpoint in this directory (weights, optimizer state, step count and data
+    # position); empty starts from the model. It may differ from save_dir.
+    resume_from: str = ""
+    seed: int = 42
+
+
+@dataclass
+class DataConfig:
+    path: str = ""
+    prompt_key: str = "prompt"
+    label_key: str = "label"
+    shuffle: bool = True
+
+
+@dataclass
+class EvalConfig:
+    # Empty path disables evaluation.
+    path: str = ""
+    interval: int = 20
+    # Evaluate once more after the last step.
+    at_end: bool = True
+    n_samples_per_prompt: int = 16
+    max_new_tokens: int = 16384
+    max_new_tokens_per_turn: Optional[int] = None
+    temperature: float = 1.0
+
+
+@dataclass
+class GRPOConfig:
+    model: ModelConfig
+    rollout: RolloutConfig = field(default_factory=RolloutConfig)
+    tool: ToolConfig = field(default_factory=ToolConfig)
+    algorithm: AlgorithmConfig = field(default_factory=AlgorithmConfig)
+    training: TrainingConfig = field(default_factory=TrainingConfig)
+    data: DataConfig = field(default_factory=DataConfig)
+    eval: EvalConfig = field(default_factory=EvalConfig)
+    # The DeepSpeed config: a dict, or the path of a JSON file.
+    deepspeed_config: Union[str, dict] = ""
+
+    @classmethod
+    def from_json(cls, path: str) -> "GRPOConfig":
+        with open(path, "r") as f:
+            raw = json.load(f)
+        return cls(
+            model=ModelConfig(**raw["model"]),
+            rollout=RolloutConfig(**raw.get("rollout", {})),
+            tool=ToolConfig(**raw.get("tool", {})),
+            algorithm=AlgorithmConfig(**raw.get("algorithm", {})),
+            training=TrainingConfig(**raw.get("training", {})),
+            data=DataConfig(**raw.get("data", {})),
+            eval=EvalConfig(**raw.get("eval", {})),
+            deepspeed_config=raw.get("deepspeed_config", ""),
+        )
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    def load_deepspeed_config(self) -> dict:
+        """A copy of the DeepSpeed config, read from its file when it is a path."""
+        if isinstance(self.deepspeed_config, dict):
+            return json.loads(json.dumps(self.deepspeed_config))
+        if not self.deepspeed_config:
+            raise ValueError("deepspeed_config is not set")
+        with open(self.deepspeed_config, "r") as f:
+            return json.load(f)
+
+    def validate(self, world_size: int) -> None:
+        if self.rollout.engine != "hybrid_engine":
+            raise ValueError(f"Unknown rollout engine {self.rollout.engine!r}; expected 'hybrid_engine'")
+        if self.rollout.n_samples_per_prompt < 2:
+            raise ValueError("GRPO needs at least 2 samples per prompt to compare them")
+        if self.rollout.temperature <= 0:
+            raise ValueError("GRPO samples with temperature > 0")
+        if self.rollout.generate_batch_size is not None and self.rollout.generate_batch_size < 1:
+            raise ValueError("rollout.generate_batch_size must be positive")
+        if self.rollout.continuous_batch_size is not None and not self.rollout.max_context_len:
+            raise ValueError("continuous batching sizes its rows by rollout.max_context_len; set it")
+        if self.algorithm.loss_aggregation not in ("token_mean", "sample_mean"):
+            raise ValueError(f"Unknown loss_aggregation {self.algorithm.loss_aggregation!r}; "
+                             "expected 'token_mean' or 'sample_mean'")
+        # Each rank rolls out whole prompts, every one with all of its samples.
+        trajectories_per_prompt_round = world_size * self.rollout.n_samples_per_prompt
+        if self.training.global_batch_size < 1 or self.training.global_batch_size % trajectories_per_prompt_round:
+            raise ValueError(f"global_batch_size {self.training.global_batch_size} must be a multiple of world size "
+                             f"x n_samples_per_prompt = {trajectories_per_prompt_round}")
+        if self.training.micro_batch_size is None:
+            self.training.micro_batch_size = self.training.global_batch_size
+        micro_batch_size = self.training.micro_batch_size
+        if micro_batch_size < 1 or micro_batch_size % world_size:
+            raise ValueError(f"micro_batch_size {micro_batch_size} must split into whole trajectories across the "
+                             f"world size {world_size}")
+        if self.training.global_batch_size % micro_batch_size:
+            raise ValueError(f"global_batch_size {self.training.global_batch_size} does not split into "
+                             f"micro-batches of {micro_batch_size}")