[anyadontfly]

Depends on deepspeedai/DeepSpeed#9037.

Summary

Adds training/grpo/, a ReTool-style example that trains Qwen3 with GRPO to solve math problems while calling a Python interpreter across several turns. It is the example for stage 1 of the agentic RL proposal (deepspeedai/DeepSpeed#8942) and is built on the rollout API that deepspeedai/DeepSpeed#9037 adds. Rollout and training share one hybrid engine, so the policy generates with its live training weights, with no separate inference server and no weight sync.

That PR changes RolloutRequest and RolloutBatch from tensors to lists of per-sample RolloutRecords, which breaks the existing OPSD example. This PR ports OPSD to the new API as well.

The GRPO example

Each step has three phases:

1. Rollout. HybridEngineRollout samples n_samples_per_prompt trajectories per prompt. When the model writes a <tool_call>, the rollout runs the code in the interpreter and appends its output as a tool message, for up to max_turns turns. Multi-turn continuous batching and CUDA-graph decode are on in every config.
2. Reward. A strict \boxed{} check scores each trajectory +1 or −1, and group_advantages normalizes rewards within each prompt's group.
3. Update. ResponseTokenBatch.from_rollout builds the targets with a mask over generated tokens only, 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, aggregated per token (default) or per trajectory.

File	Contents
main.py	Builds the model, the hybrid engine (ZeRO-2), the rollout and the data, then trains
trainer.py	Rollout, reward, update, AIME evaluation, logging, and resumable checkpoints with a Hugging Face copy
config.py	JSON-loaded dataclasses and their checks
policy.py	Token log-probs from a chunked, checkpointed LM head, so full logits of long rollouts are never materialized
tools.py	The sandboxed Python interpreter
reward.py	The strict-box answer check
data.py, scripts/	JSONL loading, data preparation (dapo-math-17k, aime-2024) and one launch script per config

Configs:

	smoke_qwen3_0.6b.json	qwen3_4b_instruct_1gpu.json	qwen3_8b_8gpu.json
Model	Qwen3-0.6B	Qwen3-4B-Instruct-2507	Qwen3-8B
GPUs	1	1 (40 GB)	8 (80 GB)
Trajectories per step	8	64	256
Turns, context	4, 4096	16, 8192	16, 8192
ZeRO-2 optimizer	GPU	CPU	GPU

Interpreter. Each call runs in a subprocess with an import allowlist (checked on the parsed code, so from math import sqrt works), a pattern filter, a memory limit, a timeout and an output cap. It is not a security boundary, and the README says so. Code wrapped in Markdown fences is unwrapped, code that doesn't compile returns Python's own SyntaxError, and the tool description asks the model to print() what it needs. tool.show_last_value (off by default) also returns the value of a bare last-line expression, as a notebook does, for models such as Qwen3-4B-Instruct-2507 that rarely end their code with print().

OPSD port

• trainer.py turns each prompt row's real tokens into a RolloutRecord and gets input_ids, attention_mask and the response mask from ResponseTokenBatch.from_rollout, which replaces utils.build_response_mask. utils.py and its tests are removed.
• main.py passes the student engine to build_rollout as generation_engine.
• Both rollout benchmarks build their requests from records.
• The READMEs say which DeepSpeed the examples need.

Testing

All on one A100 40 GB, with DeepSpeed from deepspeedai/DeepSpeed#9037 and transformers 5.

Test	Result
GRPO smoke config, 3 steps	Passes. Each step's loss equals the loss recomputed independently from the dumped trajectories, and tool output replayed from the records matches the interpreter
Qwen3-4B-Instruct-2507 config, 4 steps, with and without show_last_value	Loss matches the recomputed loss; recorded log-probs agree with a training forward (mean absolute difference about 0.01). With show_last_value, 62 of 82 tool calls return output, against 2 of 108 without it
Checkpoint and resume	Resuming from step 4 restores the weights (all 398 tensors equal the saved Hugging Face copy) and the optimizer state, and step 5 trains on the next prompts in the data order with the expected loss
OPSD smoke config (ZeRO-0, 5 steps)	Losses identical to the unported example on DeepSpeed master at every step
OPSD smoke config with graph capture	Runs
OPSD decode benchmark, with and without graph capture	Runs
OPSD unit tests (CPU)	34 passed, 1 skipped

Notes

Decode attention. Hugging Face's SDPA used to repeat each KV head once per query head whenever a mask was passed, which made masked decode of grouped-query models such as Qwen3 about three times slower than needed. huggingface/transformers#49743 fixes this by folding the query heads onto their KV head when decoding. It is merged and should be in the next Transformers release, so this example uses plain sdpa and does not work around it.