From 11445db0893de1b9d643c1286a99e8eaed209e7b Mon Sep 17 00:00:00 2001 From: Puyuan Yao Date: Sat, 7 Nov 2026 03:28:32 -0600 Subject: [PATCH 1/5] Add multi-turn GRPO example with tool calling Trains Qwen3 with GRPO on dapo-math-17k while it calls a sandboxed Python interpreter, on the DeepSpeed hybrid engine rollout (deepspeed.runtime.rollout): multi-turn rollout, CUDA-graph decode and multi-turn continuous batching, then a PPO-clip update normalized over all ranks. Includes a Qwen3-0.6B smoke test, a Qwen3-4B single-GPU config and a Qwen3-8B 8-GPU config with launch scripts. Signed-off-by: Puyuan Yao --- training/grpo/.gitignore | 3 + training/grpo/README.md | 115 ++++++++ training/grpo/config.py | 179 +++++++++++++ training/grpo/configs/qwen3_4b_sft_1gpu.json | 96 +++++++ training/grpo/configs/qwen3_8b_8gpu.json | 96 +++++++ training/grpo/configs/smoke_qwen3_0.6b.json | 90 +++++++ training/grpo/data.py | 44 +++ training/grpo/main.py | 123 +++++++++ training/grpo/policy.py | 74 +++++ training/grpo/reward.py | 55 ++++ training/grpo/scripts/prepare_data.sh | 33 +++ training/grpo/scripts/run_8gpu.sh | 9 + training/grpo/scripts/run_single_gpu.sh | 9 + training/grpo/scripts/run_smoke_test.sh | 9 + training/grpo/tools.py | 193 ++++++++++++++ training/grpo/trainer.py | 267 +++++++++++++++++++ 16 files changed, 1395 insertions(+) create mode 100644 training/grpo/.gitignore create mode 100644 training/grpo/README.md create mode 100644 training/grpo/config.py create mode 100644 training/grpo/configs/qwen3_4b_sft_1gpu.json create mode 100644 training/grpo/configs/qwen3_8b_8gpu.json create mode 100644 training/grpo/configs/smoke_qwen3_0.6b.json create mode 100644 training/grpo/data.py create mode 100644 training/grpo/main.py create mode 100644 training/grpo/policy.py create mode 100644 training/grpo/reward.py create mode 100755 training/grpo/scripts/prepare_data.sh create mode 100755 training/grpo/scripts/run_8gpu.sh create mode 100755 training/grpo/scripts/run_single_gpu.sh create mode 100755 training/grpo/scripts/run_smoke_test.sh create mode 100644 training/grpo/tools.py create mode 100644 training/grpo/trainer.py diff --git a/training/grpo/.gitignore b/training/grpo/.gitignore new file mode 100644 index 000000000..758165ce9 --- /dev/null +++ b/training/grpo/.gitignore @@ -0,0 +1,3 @@ +checkpoints/ +data/ +__pycache__/ diff --git a/training/grpo/README.md b/training/grpo/README.md new file mode 100644 index 000000000..cf45c804a --- /dev/null +++ b/training/grpo/README.md @@ -0,0 +1,115 @@ +# 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 `` block, the rollout runs it in the sandboxed `code_interpreter` (`tools.py`) and appends the output 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 80 GB GPU, and also needs about 64 GB of host memory for the optimizer state it keeps on the CPU. 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 + +# a fine-tuned Qwen3-4B 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 main.py --config configs/.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_`. To continue a run, set `training.resume_from` to its `save_dir`. + +## Configurations + +| | `smoke_qwen3_0.6b.json` | `qwen3_4b_sft_1gpu.json` | `qwen3_8b_8gpu.json` | +|---|---|---|---| +| Model | `Qwen/Qwen3-0.6B` | `font-info/qwen3-4b-sft-SGLang-RL` | `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 | 32 | 16 | +| ZeRO-2 | optimizer on the GPU | optimizer state on the CPU | optimizer on the GPU | +| AIME evaluation | off | every 20 steps | every 20 steps | +| Thinking mode | off | the model's template | 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. +- `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. diff --git a/training/grpo/config.py b/training/grpo/config.py new file mode 100644 index 000000000..db066d931 --- /dev/null +++ b/training/grpo/config.py @@ -0,0 +1,179 @@ +# 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 + # Read by HybridEngineRollout: decode steps between checks for finished turns under continuous batching. + continuous_poll_steps: int = 16 + + +@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 + + +@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}") diff --git a/training/grpo/configs/qwen3_4b_sft_1gpu.json b/training/grpo/configs/qwen3_4b_sft_1gpu.json new file mode 100644 index 000000000..57890a193 --- /dev/null +++ b/training/grpo/configs/qwen3_4b_sft_1gpu.json @@ -0,0 +1,96 @@ +{ + "model": { + "model_name_or_path": "font-info/qwen3-4b-sft-SGLang-RL", + "dtype": "bfloat16", + "attn_implementation": "sdpa", + "gradient_checkpointing": true + }, + "rollout": { + "engine": "hybrid_engine", + "multi_turn": true, + "max_turns": 16, + "max_context_len": 8192, + "max_new_tokens": 8192, + "max_new_tokens_per_turn": 4096, + "temperature": 1.0, + "repetition_penalty": 1.0, + "n_samples_per_prompt": 8, + "tool_call_parser": "hermes", + "use_graph_capture": true, + "continuous_batch_size": 32, + "continuous_poll_steps": 16, + "return_log_probs": false + }, + "tool": { + "timeout": 120, + "memory_limit_gb": 4, + "max_concurrency": 32, + "max_output_chars": 4096 + }, + "algorithm": { + "eps_clip": 0.2, + "eps_clip_high": 0.28, + "logprob_chunk_size": 1024, + "tool_call_bonus": false, + "loss_aggregation": "token_mean" + }, + "training": { + "global_batch_size": 64, + "micro_batch_size": 2, + "num_rollouts": 500, + "logging_steps": 1, + "save_steps": 50, + "save_dir": "./checkpoints/qwen3_4b_sft_1gpu", + "resume_from": "", + "seed": 42 + }, + "data": { + "path": "data/dapo-math-17k/dapo-math-17k.jsonl", + "prompt_key": "prompt", + "label_key": "label", + "shuffle": true + }, + "eval": { + "path": "data/aime-2024/aime-2024.jsonl", + "interval": 20, + "at_end": true, + "n_samples_per_prompt": 16, + "max_new_tokens": 8192, + "max_new_tokens_per_turn": 4096, + "temperature": 1.0 + }, + "deepspeed_config": { + "bf16": { + "enabled": true + }, + "zero_optimization": { + "stage": 2, + "offload_optimizer": { + "device": "cpu", + "pin_memory": true + }, + "overlap_comm": true, + "contiguous_gradients": true, + "reduce_bucket_size": 200000000.0, + "allgather_bucket_size": 200000000.0 + }, + "optimizer": { + "type": "AdamW", + "params": { + "lr": 1e-06, + "betas": [ + 0.9, + 0.98 + ], + "eps": 1e-08, + "weight_decay": 0.1 + } + }, + "gradient_clipping": 1.0, + "hybrid_engine": { + "enabled": true + }, + "steps_per_print": 1000000, + "wall_clock_breakdown": false + } +} diff --git a/training/grpo/configs/qwen3_8b_8gpu.json b/training/grpo/configs/qwen3_8b_8gpu.json new file mode 100644 index 000000000..ec34f1ce0 --- /dev/null +++ b/training/grpo/configs/qwen3_8b_8gpu.json @@ -0,0 +1,96 @@ +{ + "model": { + "model_name_or_path": "Qwen/Qwen3-8B", + "dtype": "bfloat16", + "attn_implementation": "sdpa", + "gradient_checkpointing": true + }, + "rollout": { + "engine": "hybrid_engine", + "multi_turn": true, + "max_turns": 16, + "max_context_len": 8192, + "max_new_tokens": 8192, + "max_new_tokens_per_turn": 4096, + "temperature": 1.0, + "repetition_penalty": 1.0, + "n_samples_per_prompt": 8, + "tool_call_parser": "hermes", + "use_graph_capture": true, + "continuous_batch_size": 16, + "continuous_poll_steps": 16, + "return_log_probs": false, + "chat_template_kwargs": { + "enable_thinking": false + } + }, + "tool": { + "timeout": 120, + "memory_limit_gb": 4, + "max_concurrency": 32, + "max_output_chars": 4096 + }, + "algorithm": { + "eps_clip": 0.2, + "eps_clip_high": 0.28, + "logprob_chunk_size": 1024, + "tool_call_bonus": false, + "loss_aggregation": "token_mean" + }, + "training": { + "global_batch_size": 256, + "micro_batch_size": 8, + "num_rollouts": 500, + "logging_steps": 1, + "save_steps": 50, + "save_dir": "./checkpoints/qwen3_8b_8gpu", + "resume_from": "", + "seed": 42 + }, + "data": { + "path": "data/dapo-math-17k/dapo-math-17k.jsonl", + "prompt_key": "prompt", + "label_key": "label", + "shuffle": true + }, + "eval": { + "path": "data/aime-2024/aime-2024.jsonl", + "interval": 20, + "at_end": true, + "n_samples_per_prompt": 16, + "max_new_tokens": 8192, + "max_new_tokens_per_turn": 4096, + "temperature": 1.0 + }, + "deepspeed_config": { + "bf16": { + "enabled": true + }, + "zero_optimization": { + "stage": 2, + "overlap_comm": true, + "contiguous_gradients": true, + "reduce_bucket_size": 500000000.0, + "allgather_bucket_size": 500000000.0 + }, + "optimizer": { + "type": "AdamW", + "params": { + "lr": 1e-06, + "betas": [ + 0.9, + 0.98 + ], + "eps": 1e-08, + "weight_decay": 0.1, + "torch_adam": true + } + }, + "gradient_clipping": 1.0, + "hybrid_engine": { + "enabled": true + }, + "steps_per_print": 1000000, + "wall_clock_breakdown": false + } +} diff --git a/training/grpo/configs/smoke_qwen3_0.6b.json b/training/grpo/configs/smoke_qwen3_0.6b.json new file mode 100644 index 000000000..9ea941aa4 --- /dev/null +++ b/training/grpo/configs/smoke_qwen3_0.6b.json @@ -0,0 +1,90 @@ +{ + "model": { + "model_name_or_path": "Qwen/Qwen3-0.6B", + "dtype": "bfloat16", + "attn_implementation": "sdpa", + "gradient_checkpointing": true + }, + "rollout": { + "engine": "hybrid_engine", + "multi_turn": true, + "max_turns": 4, + "max_context_len": 4096, + "max_new_tokens": 2048, + "max_new_tokens_per_turn": 1024, + "temperature": 1.0, + "repetition_penalty": 1.0, + "n_samples_per_prompt": 4, + "tool_call_parser": "hermes", + "use_graph_capture": true, + "continuous_batch_size": 8, + "continuous_poll_steps": 16, + "return_log_probs": false, + "chat_template_kwargs": { + "enable_thinking": false + } + }, + "tool": { + "timeout": 120, + "memory_limit_gb": 4, + "max_concurrency": 32, + "max_output_chars": 4096 + }, + "algorithm": { + "eps_clip": 0.2, + "eps_clip_high": 0.28, + "logprob_chunk_size": 1024, + "tool_call_bonus": false, + "loss_aggregation": "token_mean" + }, + "training": { + "global_batch_size": 8, + "micro_batch_size": 2, + "num_rollouts": 3, + "logging_steps": 1, + "save_steps": 1000000, + "save_dir": "./checkpoints/smoke_qwen3_0.6b", + "resume_from": "", + "seed": 42 + }, + "data": { + "path": "data/dapo-math-17k/dapo-math-17k.jsonl", + "prompt_key": "prompt", + "label_key": "label", + "shuffle": true + }, + "eval": { + "path": "" + }, + "deepspeed_config": { + "bf16": { + "enabled": true + }, + "zero_optimization": { + "stage": 2, + "overlap_comm": true, + "contiguous_gradients": true, + "reduce_bucket_size": 500000000.0, + "allgather_bucket_size": 500000000.0 + }, + "optimizer": { + "type": "AdamW", + "params": { + "lr": 1e-06, + "betas": [ + 0.9, + 0.98 + ], + "eps": 1e-08, + "weight_decay": 0.1, + "torch_adam": true + } + }, + "gradient_clipping": 1.0, + "hybrid_engine": { + "enabled": true + }, + "steps_per_print": 1000000, + "wall_clock_breakdown": false + } +} diff --git a/training/grpo/data.py b/training/grpo/data.py new file mode 100644 index 000000000..972c04013 --- /dev/null +++ b/training/grpo/data.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Math prompts with reference answers, read from JSONL. + +Each line holds a prompt (a chat message list, or a string taken as the +user message) and its label, as in the ``dapo-math-17k`` and +``aime-2024`` files. Rows keep their line number as ``index``, which the +trainer uses as the GRPO group index. +""" + +import json + +from torch.utils.data import Dataset + + +def load_prompts(path: str, prompt_key: str = "prompt", label_key: str = "label") -> list: + rows = [] + with open(path, "r") as f: + for line in f: + if not line.strip(): + continue + record = json.loads(line) + prompt = record[prompt_key] + if isinstance(prompt, str): + prompt = [{"role": "user", "content": prompt}] + rows.append({"index": len(rows), "prompt": prompt, "label": str(record[label_key])}) + return rows + + +class PromptDataset(Dataset): + + def __init__(self, path: str, prompt_key: str = "prompt", label_key: str = "label"): + self.rows = load_prompts(path, prompt_key, label_key) + + def __len__(self) -> int: + return len(self.rows) + + def __getitem__(self, index: int) -> dict: + return self.rows[index] + + +def collate_rows(rows: list) -> list: + """Keep a batch as a list of rows; prompts are rendered by the rollout engine.""" + return rows diff --git a/training/grpo/main.py b/training/grpo/main.py new file mode 100644 index 000000000..d7efbb732 --- /dev/null +++ b/training/grpo/main.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""GRPO retool training entry point. + +Launch with the DeepSpeed launcher, or with a script in ``scripts/``:: + + deepspeed --num_gpus 8 main.py --config configs/qwen3_8b_8gpu.json + +The DeepSpeed launcher sets ``LOCAL_RANK``, ``RANK``, and ``WORLD_SIZE`` in +the environment; we call :func:`deepspeed.init_distributed` to take that over. +""" + +import argparse +import os +import random + +import deepspeed +import numpy as np +import torch +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from transformers import AutoModelForCausalLM, AutoTokenizer + +from config import GRPOConfig +from data import PromptDataset, collate_rows, load_prompts +from deepspeed.runtime.rollout import build_rollout +from policy import PolicyModel, RolloutEngineView +from tools import CodeInterpreter +from trainer import GRPOTrainer + + +def _seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if get_accelerator().is_available(): + get_accelerator().manual_seed_all(seed) + + +def _resolve_dtype(name: str) -> torch.dtype: + return {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}[name] + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config", required=True, help="Path to GRPOConfig JSON") + parser.add_argument("--local_rank", type=int, default=int(os.environ.get("LOCAL_RANK", 0))) + args = parser.parse_args() + + deepspeed.init_distributed() + world_size = dist.get_world_size() + + cfg = GRPOConfig.from_json(args.config) + cfg.validate(world_size) + _seed_everything(cfg.training.seed) + + # --- tokenizer (shared between rollout and training) ------------------- + tokenizer = AutoTokenizer.from_pretrained(cfg.model.model_name_or_path, + trust_remote_code=cfg.model.trust_remote_code, + padding_side="left") + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + # --- policy + DeepSpeed hybrid engine ---------------------------------- + model = AutoModelForCausalLM.from_pretrained( + cfg.model.model_name_or_path, + dtype=_resolve_dtype(cfg.model.dtype), + attn_implementation=cfg.model.attn_implementation, + trust_remote_code=cfg.model.trust_remote_code, + ) + if cfg.model.gradient_checkpointing: + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + policy = PolicyModel(model, chunk_size=cfg.algorithm.logprob_chunk_size) + + # Every step is one optimizer update over all of a rank's trajectories. + samples_per_rank = cfg.training.global_batch_size // world_size + prompts_per_rank = samples_per_rank // cfg.rollout.n_samples_per_prompt + ds_config = cfg.load_deepspeed_config() + ds_config["train_micro_batch_size_per_gpu"] = cfg.training.micro_batch_size // world_size + ds_config["gradient_accumulation_steps"] = cfg.training.global_batch_size // cfg.training.micro_batch_size + ds_config.pop("train_batch_size", None) + if not ds_config.get("hybrid_engine", {}).get("enabled", False): + raise ValueError("This example rolls out with the hybrid engine; enable it in the DeepSpeed config") + + engine, *_ = deepspeed.initialize(model=policy, model_parameters=policy.parameters(), config=ds_config) + + # --- rollout engine with the code interpreter -------------------------- + tool = CodeInterpreter(timeout=cfg.tool.timeout, + memory_limit_gb=cfg.tool.memory_limit_gb, + max_concurrency=cfg.tool.max_concurrency, + max_output_chars=cfg.tool.max_output_chars) + rollout = build_rollout(cfg.rollout, + generation_engine=RolloutEngineView(engine), + tokenizer=tokenizer, + tools=[tool]) + + # --- data: each rank rolls out its own share of every step's prompts ---- + dataset = PromptDataset(cfg.data.path, cfg.data.prompt_key, cfg.data.label_key) + sampler = DistributedSampler(dataset, + num_replicas=world_size, + rank=dist.get_rank(), + shuffle=cfg.data.shuffle, + seed=cfg.training.seed) + loader = DataLoader(dataset, batch_size=prompts_per_rank, sampler=sampler, collate_fn=collate_rows, drop_last=True) + eval_rows = load_prompts(cfg.eval.path, cfg.data.prompt_key, cfg.data.label_key) if cfg.eval.path else [] + + trainer = GRPOTrainer( + cfg=cfg, + engine=engine, + tokenizer=tokenizer, + rollout=rollout, + dataloader=loader, + eval_rows=eval_rows, + ) + if cfg.training.resume_from: + trainer.resume(cfg.training.resume_from) + trainer.train() + + +if __name__ == "__main__": + main() diff --git a/training/grpo/policy.py b/training/grpo/policy.py new file mode 100644 index 000000000..ff9473e05 --- /dev/null +++ b/training/grpo/policy.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""The policy as DeepSpeed trains it: a causal LM whose forward returns token log-probs. + +Full logits of a long rollout are too large to keep (16k tokens x 152k +vocab is 9.7 GB in fp32), so the LM head runs over the sequence in chunks +under activation checkpointing and only each target token's log-prob is +kept. Training goes through ``engine(...)`` rather than a +submodule, because DeepSpeed hooks the engine's forward outputs to prepare +ZeRO for the backward pass. + +Rollout instead drives the causal LM itself: ``generate``, CUDA graph capture +and the forward profiler all call ``engine.module`` as a Hugging Face model, so +the rollout gets a :class:`RolloutEngineView` whose ``module`` is that model. +""" + +import torch +import torch.nn.functional as F +from torch.utils.checkpoint import checkpoint + + +def _chunk_log_probs(hidden, lm_head, targets, temperature): + logits = lm_head(hidden).float() / temperature + return F.log_softmax(logits, dim=-1).gather(-1, targets.unsqueeze(-1)).squeeze(-1) + + +class PolicyModel(torch.nn.Module): + """Wraps a Hugging Face causal LM; its forward returns token log-probs for training.""" + + def __init__(self, model, chunk_size: int = 1024): + super().__init__() + self.model = model + self.chunk_size = chunk_size + + def forward(self, input_ids, attention_mask, temperature: float = 1.0): + """Log-prob of each next token under ``softmax(logits / temperature)``. + + Rollout samples from ``softmax(logits / temperature)``, so the + policy being trained is that tempered distribution. + + Returns: + ``[B, T - 1]`` log-probs; position ``t`` scores ``input_ids[:, t + 1]``. + """ + hidden = self.model.get_decoder()(input_ids=input_ids, attention_mask=attention_mask, + use_cache=False).last_hidden_state + hidden = hidden[:, :-1] + targets = input_ids[:, 1:] + lm_head = self.model.get_output_embeddings() + log_probs = [] + for start in range(0, hidden.shape[1], self.chunk_size): + end = start + self.chunk_size + log_probs.append( + checkpoint(_chunk_log_probs, + hidden[:, start:end], + lm_head, + targets[:, start:end], + temperature, + use_reentrant=False)) + return torch.cat(log_probs, dim=1) + + +class RolloutEngineView: + """The DeepSpeed engine as the rollout sees it: ``module`` is the wrapped causal LM. + + Everything else (the hybrid engine's ZeRO and shared-prefill methods) is the + engine's own, so rollout and training still share one set of weights. + """ + + def __init__(self, engine): + self._engine = engine + self.module = engine.module.model + + def __getattr__(self, name): + return getattr(self._engine, name) diff --git a/training/grpo/reward.py b/training/grpo/reward.py new file mode 100644 index 000000000..609ef202b --- /dev/null +++ b/training/grpo/reward.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Retool reward: DAPO strict-box answer check. + +The answer check is adapted from verl and EleutherAI's +lm-evaluation-harness under the Apache License 2.0: +Copyright 2024 Bytedance Ltd. and/or its affiliates, Copyright 2022 +EleutherAI and the HuggingFace Inc. team. +""" + +from typing import Optional + + +def last_boxed_only_string(text: str) -> Optional[str]: + """The last ``\\boxed{...}`` expression in ``text``, braces balanced, or None.""" + start = text.rfind("\\boxed{") + if start < 0: + return None + depth = 0 + for end in range(start, len(text)): + if text[end] == "{": + depth += 1 + elif text[end] == "}": + depth -= 1 + if depth == 0: + return text[start:end + 1] + return None + + +def remove_boxed(boxed: str) -> str: + return boxed[len("\\boxed{"):-1] + + +def compute_score(solution: str, label: str) -> dict: + """+1 if the last boxed answer near the end of ``solution`` equals ``label``, else -1. + + Only the last 100 characters are searched, so the answer must close the + response, as the DAPO prompt asks. + """ + boxed = last_boxed_only_string(solution[-100:]) + pred = remove_boxed(boxed) if boxed is not None else None + correct = pred == label + return {"score": 1.0 if correct else -1.0, "acc": correct, "pred": pred or ""} + + +def retool_reward(solution: str, label: str, num_tool_calls: int, tool_call_bonus: bool = False) -> dict: + """Retool reward: the strict-box score, plus an optional tool bonus. + + The bonus lifts a wrong answer by 0.05 per tool call, up to -0.6, so it + never reaches a correct one. + """ + result = compute_score(solution, label) + if tool_call_bonus and result["score"] < 0: + result["score"] = min(-0.6, result["score"] + num_tool_calls / 2 * 0.1) + return result diff --git a/training/grpo/scripts/prepare_data.sh b/training/grpo/scripts/prepare_data.sh new file mode 100755 index 000000000..25b752432 --- /dev/null +++ b/training/grpo/scripts/prepare_data.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +# +# Downloads the retool training and AIME eval data into data/. Assumes you're cd'd into training/grpo/. +set -euo pipefail + +hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir data/dapo-math-17k +hf download --repo-type dataset zhuzilin/aime-2024 --local-dir data/aime-2024 + +# aime-2024 has bare problems, but the scorer only accepts \boxed{} answers. Ask for one with the prefix and suffix +# that every dapo-math-17k prompt has. Prompts that already have them are left as they are. +python3 - data/aime-2024/aime-2024.jsonl <<'PY' +import json +import os +import sys + +PREFIX = ("Solve the following math problem step by step. The last line of your response should be of the form " + "Answer: \\boxed{$Answer} where $Answer is the answer to the problem.\n\n") +SUFFIX = '\n\nRemember to put your answer on its own line after "Answer:".' + +path = sys.argv[1] +with open(path) as f: + rows = [json.loads(line) for line in f if line.strip()] +for row in rows: + message = row["prompt"][0] + if not message["content"].startswith(PREFIX): + message["content"] = PREFIX + message["content"] + SUFFIX +with open(path + ".tmp", "w") as f: + for row in rows: + f.write(json.dumps(row) + "\n") +os.replace(path + ".tmp", path) +PY diff --git a/training/grpo/scripts/run_8gpu.sh b/training/grpo/scripts/run_8gpu.sh new file mode 100755 index 000000000..e702448bb --- /dev/null +++ b/training/grpo/scripts/run_8gpu.sh @@ -0,0 +1,9 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +# +# Trains with configs/qwen3_8b_8gpu.json on 8 GPUs. Run scripts/prepare_data.sh first. +set -euo pipefail +cd "$(dirname "$0")/.." + +deepspeed --num_gpus 8 main.py --config configs/qwen3_8b_8gpu.json diff --git a/training/grpo/scripts/run_single_gpu.sh b/training/grpo/scripts/run_single_gpu.sh new file mode 100755 index 000000000..413b9ab59 --- /dev/null +++ b/training/grpo/scripts/run_single_gpu.sh @@ -0,0 +1,9 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +# +# Trains with configs/qwen3_4b_sft_1gpu.json on 1 GPU. Run scripts/prepare_data.sh first. +set -euo pipefail +cd "$(dirname "$0")/.." + +deepspeed --num_gpus 1 main.py --config configs/qwen3_4b_sft_1gpu.json diff --git a/training/grpo/scripts/run_smoke_test.sh b/training/grpo/scripts/run_smoke_test.sh new file mode 100755 index 000000000..014517e14 --- /dev/null +++ b/training/grpo/scripts/run_smoke_test.sh @@ -0,0 +1,9 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +# +# Trains with configs/smoke_qwen3_0.6b.json on 1 GPU. Run scripts/prepare_data.sh first. +set -euo pipefail +cd "$(dirname "$0")/.." + +deepspeed --num_gpus 1 main.py --config configs/smoke_qwen3_0.6b.json diff --git a/training/grpo/tools.py b/training/grpo/tools.py new file mode 100644 index 000000000..53dd1821e --- /dev/null +++ b/training/grpo/tools.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""A Python code interpreter the model calls during rollout. + +Code passes an import allowlist and a pattern filter, then runs under a +memory limit and a timeout, as an async subprocess so concurrent calls do not +block each other. The filter only screens obvious misuse; it is not a security +boundary. +""" + +import asyncio +import os +import re +import sys +import tempfile + +from deepspeed.runtime.rollout import Tool + +ALLOWED_MODULES = { + "math", + "random", + "datetime", + "collections", + "itertools", + "functools", + "operator", + "statistics", + "decimal", + "fractions", +} + +DANGEROUS_PATTERNS = [ + r"import\s+os", + r"import\s+sys", + r"import\s+subprocess", + r"import\s+shutil", + r"import\s+glob", + r"import\s+pathlib", + r"__import__", + r"eval\s*\(", + r"exec\s*\(", + r"open\s*\(", + r"file\s*\(", + r"input\s*\(", + r"raw_input\s*\(", + r"compile\s*\(", + r"execfile\s*\(", + r"getattr\s*\(", + r"setattr\s*\(", + r"delattr\s*\(", + r"hasattr\s*\(", + r"globals\s*\(", + r"locals\s*\(", + r"vars\s*\(", + r"dir\s*\(", + r"type\s*\(", + r"isinstance\s*\(", + r"issubclass\s*\(", + r"super\s*\(", + r"property\s*\(", + r"staticmethod\s*\(", + r"classmethod\s*\(", + r"__\w+__", +] + +# The user code runs inside a try block so its output and errors come back in +# the format the model trains with. +WRAPPER = """import sys +import traceback +from io import StringIO +import resource + +try: + resource.setrlimit(resource.RLIMIT_AS, (MEMORY_LIMIT, -1)) +except Exception: + pass + +old_stdout = sys.stdout +old_stderr = sys.stderr +stdout_capture = StringIO() +stderr_capture = StringIO() +sys.stdout = stdout_capture +sys.stderr = stderr_capture + +try: +USER_CODE + + stdout_output = stdout_capture.getvalue() + stderr_output = stderr_capture.getvalue() + sys.stdout = old_stdout + sys.stderr = old_stderr + result = "" + if stdout_output: + result += f"Output:\\n{stdout_output}" + if stderr_output: + result += f"\\nErrors:\\n{stderr_output}" + print(result) +except Exception as e: + sys.stdout = old_stdout + sys.stderr = old_stderr + print(f"Error: {str(e)}\\nTraceback:\\n{traceback.format_exc()}") +""" + + +def check_code_safety(code: str): + """Return an error message if ``code`` matches a blocked pattern or import, else None.""" + for pattern in DANGEROUS_PATTERNS: + if re.search(pattern, code, re.IGNORECASE): + return f"Code contains dangerous pattern: {pattern}" + imports = set(re.findall(r"import\s+(\w+)", code) + re.findall(r"from\s+(\w+)", code)) + for module in sorted(imports): + if module not in ALLOWED_MODULES: + return f"Import of '{module}' is not allowed" + return None + + +class CodeInterpreter(Tool): + name = "code_interpreter" + description = "A tool for executing Python code in a safe sandbox environment." + parameters = { + "type": "object", + "properties": { + "code": { + "type": "string", + "description": "The Python code to execute" + } + }, + "required": ["code"], + } + + def __init__(self, + timeout: float = 120.0, + memory_limit_gb: int = 4, + max_concurrency: int = 32, + max_output_chars: int = 4096): + self.timeout = timeout + self.memory_limit = memory_limit_gb * 1024**3 + self.max_concurrency = max_concurrency + self.max_output_chars = max_output_chars + self._loop = None + self._semaphore = None + + def _limiter(self) -> asyncio.Semaphore: + # The rollout runs each turn's calls in a new event loop, and a + # semaphore only works in the loop it was first used in. + loop = asyncio.get_running_loop() + if self._loop is not loop: + self._loop = loop + self._semaphore = asyncio.Semaphore(self.max_concurrency) + return self._semaphore + + async def execute(self, arguments) -> str: + code = arguments.get("code", "") + if isinstance(code, list): + code = "\n".join(str(line) for line in code) + output = await self._execute(code) + if len(output) > self.max_output_chars: + output = output[:self.max_output_chars] + "\n... (output truncated)" + return output + + async def _execute(self, code) -> str: + """The sandbox's output for ``code``, before truncation.""" + if not isinstance(code, str) or not code.strip(): + return "Error: No code provided" + error = check_code_safety(code) + if error is not None: + return f"Error: {error}" + async with self._limiter(): + return await self._run(code) + + async def _run(self, code: str) -> str: + indented = "\n".join(" " + line for line in code.split("\n")) + script = WRAPPER.replace("MEMORY_LIMIT", str(self.memory_limit)).replace("USER_CODE", indented) + with tempfile.TemporaryDirectory(prefix="python_sandbox_") as workdir: + path = os.path.join(workdir, "code.py") + with open(path, "w") as f: + f.write(script) + env = dict(os.environ, PYTHONPATH=workdir, PYTHONUNBUFFERED="1") + process = await asyncio.create_subprocess_exec(sys.executable, + path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=workdir, + env=env) + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=self.timeout) + except asyncio.TimeoutError: + process.kill() + await process.wait() + return f"Error: Code execution timed out after {self.timeout:g} seconds" + if process.returncode != 0: + return f"Error: Process exited with code {process.returncode}\n{stderr.decode(errors='replace')}" + return stdout.decode(errors="replace").strip() diff --git a/training/grpo/trainer.py b/training/grpo/trainer.py new file mode 100644 index 000000000..992cddf67 --- /dev/null +++ b/training/grpo/trainer.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""GRPO training loop for multi-turn tool-calling rollouts. + +Each step: + + 1. **Rollout.** The hybrid engine switches to inference mode and the policy + solves each rank's prompts, calling the code interpreter between turns + (:class:`~deepspeed.runtime.rollout.HybridEngineRollout`). + 2. **Reward.** Every trajectory gets the strict-box answer score, and its + advantage is the score normalized within its prompt's group + (:func:`~deepspeed.runtime.rollout.group_advantages`). + 3. **Update.** One optimizer step over all trajectories, accumulated over + micro-batches: PPO-clip with clip-higher on the generated tokens + (:func:`~deepspeed.runtime.rollout.grpo_policy_loss`), averaged over every + generated token across ranks, or per trajectory first with + ``algorithm.loss_aggregation = "sample_mean"``. Tool output is masked out + (:class:`~deepspeed.runtime.rollout.ResponseTokenBatch`). + +With one update per rollout, the old policy is the current one, so by default +the old log-probs are the current ones detached; the importance ratio is 1 and +the loss reduces to the policy gradient. ``rollout.return_log_probs`` +takes them from the rollout instead +(:func:`~deepspeed.runtime.rollout.old_log_probs_from_rollout`). +""" + +import os +import re + +import torch +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator + +from config import GRPOConfig +from reward import retool_reward +from deepspeed.runtime.rollout import (ResponseTokenBatch, RolloutBatch, RolloutEngine, RolloutRequest, SamplingConfig, + group_advantages, grpo_policy_loss, old_log_probs_from_rollout) +from deepspeed.runtime.rollout.base import RolloutRecord + + +def _is_rank_zero() -> bool: + return (not dist.is_initialized()) or dist.get_rank() == 0 + + +def _mean_across_ranks(values: dict) -> dict: + """Average each metric over ranks with one all-reduce.""" + names = sorted(values) + tensor = torch.tensor([float(values[name]) for name in names], device=get_accelerator().current_device_name()) + dist.all_reduce(tensor) + tensor /= dist.get_world_size() + return dict(zip(names, tensor.tolist())) + + +class GRPOTrainer: + + def __init__( + self, + cfg: GRPOConfig, + engine, + tokenizer, + rollout: RolloutEngine, + dataloader, + eval_rows=None, + ): + self.cfg = cfg + self.engine = engine + self.tokenizer = tokenizer + self.rollout = rollout + self.dataloader = dataloader + self.eval_rows = eval_rows or [] + + self.device = get_accelerator().current_device_name() + self.pad_token_id = tokenizer.pad_token_id + self.step = 0 + self.sampling = SamplingConfig( + max_new_tokens=cfg.rollout.max_new_tokens, + max_new_tokens_per_turn=cfg.rollout.max_new_tokens_per_turn, + temperature=cfg.rollout.temperature, + n_samples_per_prompt=cfg.rollout.n_samples_per_prompt, + repetition_penalty=cfg.rollout.repetition_penalty, + ) + self.eval_sampling = SamplingConfig( + max_new_tokens=cfg.eval.max_new_tokens, + max_new_tokens_per_turn=cfg.eval.max_new_tokens_per_turn, + temperature=cfg.eval.temperature, + n_samples_per_prompt=cfg.eval.n_samples_per_prompt, + repetition_penalty=cfg.rollout.repetition_penalty, + ) + + # ------------------------------------------------------------------ + # Driver + # ------------------------------------------------------------------ + + def resume(self, checkpoint_dir: str) -> None: + """Load the latest checkpoint in ``checkpoint_dir``: weights, optimizer state and the step it was saved at.""" + load_path, _ = self.engine.load_checkpoint(checkpoint_dir) + if load_path is None: + raise FileNotFoundError(f"No checkpoint to resume from in {checkpoint_dir}") + # load_path is a file inside the checkpoint's tag directory, and tags are step_ (see _save). + tag = re.search(r"(?:^|/)step_(\d+)(?:/|$)", load_path) + if tag is None: + raise ValueError(f"Cannot read the step from checkpoint path {load_path}") + self.step = int(tag.group(1)) + if _is_rank_zero(): + print(f"[grpo] resumed from {load_path} at step {self.step}", flush=True) + + def train(self) -> None: + # A resumed run picks up the data order where its checkpoint left off. + epoch, skip = divmod(self.step, len(self.dataloader)) + while self.step < self.cfg.training.num_rollouts: + # Reseed the DistributedSampler so each epoch reshuffles each rank's shard. + self.dataloader.sampler.set_epoch(epoch) + for rows in self.dataloader: + if skip: + skip -= 1 + continue + if self.step >= self.cfg.training.num_rollouts: + break + if self._eval_due(): + self._log(self.evaluate(), prefix="eval") + metrics = self._train_step(rows) + self.step += 1 + if self.step % self.cfg.training.logging_steps == 0: + self._log(metrics, prefix="train") + if self.step % self.cfg.training.save_steps == 0: + self._save() + epoch += 1 + if self.cfg.eval.path and self.cfg.eval.at_end: + self._log(self.evaluate(), prefix="eval") + if self.step % self.cfg.training.save_steps: + self._save() + + def _eval_due(self) -> bool: + """Evaluate before training starts and every ``eval.interval`` steps.""" + return bool(self.cfg.eval.path) and self.cfg.eval.interval > 0 and self.step % self.cfg.eval.interval == 0 + + # ------------------------------------------------------------------ + # Rollout and reward + # ------------------------------------------------------------------ + + def _rollout(self, rows, sampling: SamplingConfig) -> list: + """Generate trajectories for ``rows`` in chunks, the hybrid engine in inference mode.""" + self.engine.eval() + samples = [] + # A call decodes rollout.generate_batch_size trajectories, by default as many as a training step does + # on this rank, so a step's rows roll out in one call. Evaluation's larger n gets fewer prompts per call. + trajectories_per_call = self.cfg.rollout.generate_batch_size + if trajectories_per_call is None: + trajectories_per_call = self.cfg.training.global_batch_size // dist.get_world_size() + prompts_per_chunk = max(1, trajectories_per_call // sampling.n_samples_per_prompt) + for start in range(0, len(rows), prompts_per_chunk): + request = RolloutRequest(samples=[ + RolloutRecord(group_index=row["index"], prompt=row["prompt"]) for row in rows[start:start + prompts_per_chunk] + ]) + samples.extend(self.rollout.generate(request, sampling).samples) + return samples + + def _score(self, samples, rows) -> list: + """Each sample's reward, also stored on the sample.""" + labels = {row["index"]: row["label"] for row in rows} + rewards = [] + for sample in samples: + num_tool_calls = sample.response.count("") + result = retool_reward(sample.response, labels[sample.group_index], num_tool_calls, + self.cfg.algorithm.tool_call_bonus) + sample.reward = result["score"] + rewards.append(result["score"]) + return rewards + + @torch.no_grad() + def evaluate(self) -> dict: + """Accuracy over the eval set, averaged over every sample; its prompts are split across ranks.""" + rank, world_size = dist.get_rank(), dist.get_world_size() + rows = self.eval_rows[rank::world_size] + rewards = self._score(self._rollout(rows, self.eval_sampling), rows) + # A correct answer scores +1; any wrong one, even with the tool bonus, stays below 0. + num_correct = sum(reward > 0 for reward in rewards) + totals = torch.tensor([num_correct, len(rewards)], dtype=torch.float, device=self.device) + dist.all_reduce(totals) + return {"accuracy": (totals[0] / totals[1].clamp(min=1)).item()} + + # ------------------------------------------------------------------ + # One step + # ------------------------------------------------------------------ + + def _train_step(self, rows) -> dict: + samples = self._rollout(rows, self.sampling) + rewards = self._score(samples, rows) + advantages = group_advantages(samples, rewards).to(self.device) + + # Rollout and training take turns on the GPU: return each phase's cached memory before the other runs. + get_accelerator().empty_cache() + self.engine.train() + loss = self._update(samples, advantages) + get_accelerator().empty_cache() + # Averaged over ranks, the per-rank losses are the step loss. + return _mean_across_ranks({"loss": loss, "reward": sum(rewards) / len(rewards)}) + + def _update(self, samples, advantages) -> float: + """One optimizer step over ``samples``: PPO-clip, averaged over all ranks. + + Returns this rank's loss; averaged over ranks it is the step loss. + """ + if self.cfg.algorithm.loss_aggregation == "sample_mean": + # Each trajectory with a generated token weighs the same, however many it generated; one cut + # before its first turn adds nothing, as grpo_policy_loss counts them. + num_units = float(sum(1 for sample in samples if any(sample.loss_mask or []))) + else: + # Every generated token across ranks weighs the same. + num_units = float(sum(sum(sample.loss_mask or []) for sample in samples)) + total_units = torch.tensor(num_units, device=self.device) + dist.all_reduce(total_units) + # grpo_policy_loss returns each micro-batch's unreduced loss_sum. DeepSpeed averages gradients over + # ranks, so each rank's sum is scaled by world_size / total to give the step's mean. + loss_scale = dist.get_world_size() / total_units.clamp(min=1).item() + + step_loss = 0.0 + # micro_batch_size counts trajectories across ranks. + micro_batch_size = self.cfg.training.micro_batch_size // dist.get_world_size() + for start in range(0, len(samples), micro_batch_size): + rollout = RolloutBatch(samples=samples[start:start + micro_batch_size]) + batch = ResponseTokenBatch.from_rollout(rollout, pad_token_id=self.pad_token_id, device=self.device) + # [B, T - 1] log-probs of batch.target_ids, aligned with batch.response_mask. + log_probs = self.engine(input_ids=batch.input_ids, + attention_mask=batch.attention_mask, + temperature=self.cfg.rollout.temperature) + if self.cfg.rollout.return_log_probs: + old_log_probs = old_log_probs_from_rollout(rollout, device=self.device) + else: + old_log_probs = log_probs.detach() + output = grpo_policy_loss(log_probs, + old_log_probs, + advantages[start:start + micro_batch_size], + batch.response_mask, + clip_low=self.cfg.algorithm.eps_clip, + clip_high=self.cfg.algorithm.eps_clip_high, + aggregation=self.cfg.algorithm.loss_aggregation) + loss = output.loss_sum * loss_scale + # The loss is already normalized over the whole step, so it must not + # be divided again by the gradient accumulation steps. + self.engine.backward(loss, scale_wrt_gas=False) + self.engine.step() + step_loss += loss.item() + return step_loss + + # ------------------------------------------------------------------ + # Logging / checkpointing + # ------------------------------------------------------------------ + + def _log(self, metrics: dict, prefix: str) -> None: + if not _is_rank_zero(): + return + values = " ".join(f"{name}={value:.4g}" for name, value in sorted(metrics.items())) + print(f"[grpo][{prefix}][step {self.step}] {values}", flush=True) + + def _save(self) -> None: + """Save a resumable DeepSpeed checkpoint and a Hugging Face copy of the weights.""" + tag = f"step_{self.step}" + save_dir = self.cfg.training.save_dir + self.engine.save_checkpoint(save_dir, tag=tag) + # ZeRO-2 keeps full weights on every rank, so rank 0 can export them alone. + if _is_rank_zero(): + hf_dir = os.path.join(save_dir, f"hf_{tag}") + self.engine.module.model.save_pretrained(hf_dir) + self.tokenizer.save_pretrained(hf_dir) + print(f"[grpo] saved checkpoint to {save_dir}/{tag} and {hf_dir}", flush=True) + dist.barrier() From 21a1ccfd51c367c8047b5d9d4cfcdab46b8da772 Mon Sep 17 00:00:00 2001 From: Puyuan Yao Date: Sun, 8 Nov 2026 03:16:17 -0600 Subject: [PATCH 2/5] Train Qwen3-4B-Instruct-2507 on one GPU and fix the interpreter's checks - Switch the single-GPU config from the SFT checkpoint to Qwen3-4B-Instruct-2507 (renamed to qwen3_4b_instruct_1gpu.json), with 16 decode rows to fit a 40 GB A100 and AIME evaluation every 100 steps. - Register an sdpa_gqa_decode attention implementation that reads grouped KV heads in place on decode steps. Hugging Face's SDPA copies them whenever there is a mask, which made Qwen3-4B decode about three times slower. - Strip a Markdown fence around the whole code argument, as Qwen3-2507 models write it, and tell the model that only printed output is returned. - Check imports with ast, so `from math import sqrt` is allowed, and return Python's SyntaxError for code that does not parse instead of guessing its imports. - Drop continuous_poll_steps, which the rollout no longer reads, and correct the host memory the single-GPU config needs. Signed-off-by: Puyuan Yao --- training/grpo/README.md | 20 ++++---- training/grpo/config.py | 5 +- ..._1gpu.json => qwen3_4b_instruct_1gpu.json} | 11 ++--- training/grpo/configs/qwen3_8b_8gpu.json | 3 +- training/grpo/configs/smoke_qwen3_0.6b.json | 3 +- training/grpo/main.py | 3 +- training/grpo/policy.py | 34 ++++++++++++++ training/grpo/scripts/run_single_gpu.sh | 4 +- training/grpo/tools.py | 46 +++++++++++++++++-- 9 files changed, 100 insertions(+), 29 deletions(-) rename training/grpo/configs/{qwen3_4b_sft_1gpu.json => qwen3_4b_instruct_1gpu.json} (90%) diff --git a/training/grpo/README.md b/training/grpo/README.md index cf45c804a..b4f7f5427 100644 --- a/training/grpo/README.md +++ b/training/grpo/README.md @@ -21,7 +21,7 @@ With one update per rollout, the old policy is the current one, so the old log-p | `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 | +| `policy.py` | The policy as DeepSpeed trains it: token log-probs from a chunked LM head, and the decode attention | | `tools.py` | The sandboxed Python interpreter the model calls | | `reward.py` | The strict-box answer check | | `data.py` | JSONL prompt loading | @@ -37,7 +37,7 @@ With one update per rollout, the old policy is the current one, so the old log-p 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 80 GB GPU, and also needs about 64 GB of host memory for the optimizer state it keeps on the CPU. The distributed config is sized for 8 GPUs with 80 GB each. +- **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 @@ -58,7 +58,7 @@ cd training/grpo # smoke test: Qwen3-0.6B, three steps on one GPU bash scripts/run_smoke_test.sh -# a fine-tuned Qwen3-4B on one GPU +# Qwen3-4B-Instruct-2507 on one GPU bash scripts/run_single_gpu.sh # Qwen3-8B on 8 GPUs @@ -71,20 +71,22 @@ Rank 0 prints the loss and the mean reward after each step. When `eval.path` is ## Configurations -| | `smoke_qwen3_0.6b.json` | `qwen3_4b_sft_1gpu.json` | `qwen3_8b_8gpu.json` | +| | `smoke_qwen3_0.6b.json` | `qwen3_4b_instruct_1gpu.json` | `qwen3_8b_8gpu.json` | |---|---|---|---| -| Model | `Qwen/Qwen3-0.6B` | `font-info/qwen3-4b-sft-SGLang-RL` | `Qwen/Qwen3-8B` | +| 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 | 32 | 16 | +| 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 20 steps | every 20 steps | -| Thinking mode | off | the model's template | off | +| 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. +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. + +The configs set `model.attn_implementation` to `sdpa_gqa_decode`, which `policy.py` registers with Hugging Face: SDPA, except that decode steps read Qwen3's shared KV heads in place. Hugging Face's SDPA copies each KV head once per query head it serves whenever there is an attention mask, and rollout decodes with a mask over each row's whole KV cache, so the copy dominates decode time; skipping it makes Qwen3-4B decode about three times faster on an A100 with the same outputs. Prefill and training use Hugging Face's SDPA unchanged. 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. diff --git a/training/grpo/config.py b/training/grpo/config.py index db066d931..5f776e633 100644 --- a/training/grpo/config.py +++ b/training/grpo/config.py @@ -18,7 +18,8 @@ class ModelConfig: model_name_or_path: str dtype: str = "bfloat16" - attn_implementation: str = "sdpa" + # "sdpa_gqa_decode" (policy.py) is SDPA with faster decode steps for models with grouped KV heads. + attn_implementation: str = "sdpa_gqa_decode" gradient_checkpointing: bool = True trust_remote_code: bool = False @@ -45,8 +46,6 @@ class RolloutConfig(_BaseRolloutConfig): # 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 - # Read by HybridEngineRollout: decode steps between checks for finished turns under continuous batching. - continuous_poll_steps: int = 16 @dataclass diff --git a/training/grpo/configs/qwen3_4b_sft_1gpu.json b/training/grpo/configs/qwen3_4b_instruct_1gpu.json similarity index 90% rename from training/grpo/configs/qwen3_4b_sft_1gpu.json rename to training/grpo/configs/qwen3_4b_instruct_1gpu.json index 57890a193..2da0f1042 100644 --- a/training/grpo/configs/qwen3_4b_sft_1gpu.json +++ b/training/grpo/configs/qwen3_4b_instruct_1gpu.json @@ -1,8 +1,8 @@ { "model": { - "model_name_or_path": "font-info/qwen3-4b-sft-SGLang-RL", + "model_name_or_path": "Qwen/Qwen3-4B-Instruct-2507", "dtype": "bfloat16", - "attn_implementation": "sdpa", + "attn_implementation": "sdpa_gqa_decode", "gradient_checkpointing": true }, "rollout": { @@ -17,8 +17,7 @@ "n_samples_per_prompt": 8, "tool_call_parser": "hermes", "use_graph_capture": true, - "continuous_batch_size": 32, - "continuous_poll_steps": 16, + "continuous_batch_size": 16, "return_log_probs": false }, "tool": { @@ -40,7 +39,7 @@ "num_rollouts": 500, "logging_steps": 1, "save_steps": 50, - "save_dir": "./checkpoints/qwen3_4b_sft_1gpu", + "save_dir": "./checkpoints/qwen3_4b_instruct_1gpu", "resume_from": "", "seed": 42 }, @@ -52,7 +51,7 @@ }, "eval": { "path": "data/aime-2024/aime-2024.jsonl", - "interval": 20, + "interval": 100, "at_end": true, "n_samples_per_prompt": 16, "max_new_tokens": 8192, diff --git a/training/grpo/configs/qwen3_8b_8gpu.json b/training/grpo/configs/qwen3_8b_8gpu.json index ec34f1ce0..0a938523c 100644 --- a/training/grpo/configs/qwen3_8b_8gpu.json +++ b/training/grpo/configs/qwen3_8b_8gpu.json @@ -2,7 +2,7 @@ "model": { "model_name_or_path": "Qwen/Qwen3-8B", "dtype": "bfloat16", - "attn_implementation": "sdpa", + "attn_implementation": "sdpa_gqa_decode", "gradient_checkpointing": true }, "rollout": { @@ -18,7 +18,6 @@ "tool_call_parser": "hermes", "use_graph_capture": true, "continuous_batch_size": 16, - "continuous_poll_steps": 16, "return_log_probs": false, "chat_template_kwargs": { "enable_thinking": false diff --git a/training/grpo/configs/smoke_qwen3_0.6b.json b/training/grpo/configs/smoke_qwen3_0.6b.json index 9ea941aa4..04ffe1732 100644 --- a/training/grpo/configs/smoke_qwen3_0.6b.json +++ b/training/grpo/configs/smoke_qwen3_0.6b.json @@ -2,7 +2,7 @@ "model": { "model_name_or_path": "Qwen/Qwen3-0.6B", "dtype": "bfloat16", - "attn_implementation": "sdpa", + "attn_implementation": "sdpa_gqa_decode", "gradient_checkpointing": true }, "rollout": { @@ -18,7 +18,6 @@ "tool_call_parser": "hermes", "use_graph_capture": true, "continuous_batch_size": 8, - "continuous_poll_steps": 16, "return_log_probs": false, "chat_template_kwargs": { "enable_thinking": false diff --git a/training/grpo/main.py b/training/grpo/main.py index d7efbb732..c38bfd709 100644 --- a/training/grpo/main.py +++ b/training/grpo/main.py @@ -26,7 +26,7 @@ from config import GRPOConfig from data import PromptDataset, collate_rows, load_prompts from deepspeed.runtime.rollout import build_rollout -from policy import PolicyModel, RolloutEngineView +from policy import PolicyModel, RolloutEngineView, register_attention from tools import CodeInterpreter from trainer import GRPOTrainer @@ -64,6 +64,7 @@ def main() -> None: tokenizer.pad_token = tokenizer.eos_token # --- policy + DeepSpeed hybrid engine ---------------------------------- + register_attention() model = AutoModelForCausalLM.from_pretrained( cfg.model.model_name_or_path, dtype=_resolve_dtype(cfg.model.dtype), diff --git a/training/grpo/policy.py b/training/grpo/policy.py index ff9473e05..04332162f 100644 --- a/training/grpo/policy.py +++ b/training/grpo/policy.py @@ -17,6 +17,40 @@ import torch import torch.nn.functional as F from torch.utils.checkpoint import checkpoint +from transformers import AttentionInterface +from transformers.integrations.sdpa_attention import sdpa_attention_forward +from transformers.masking_utils import AttentionMaskInterface, sdpa_mask + +# SDPA that reads grouped KV heads in place on decode steps; see _sdpa_gqa_decode. +SDPA_GQA_DECODE = "sdpa_gqa_decode" + + +def _sdpa_gqa_decode(module, query, key, value, attention_mask, dropout=0.0, scaling=None, **kwargs): + """Hugging Face SDPA attention, except that one-token decode steps read grouped KV heads in place. + + With an attention mask, Hugging Face's SDPA copies each KV head once per query head it serves before + attending. Rollout decodes with a mask over each row's whole KV cache, so for Qwen3 (8 KV heads for 32 + query heads) that copy reads and writes the full cache four times over at every step, and dominates + decode time. ``enable_gqa`` attends to the shared heads directly. Prefill and training, where the query + is longer than one token, keep Hugging Face's SDPA path. + """ + if query.shape[2] != 1 or getattr(module, "num_key_value_groups", 1) == 1: + return sdpa_attention_forward(module, query, key, value, attention_mask, dropout=dropout, scaling=scaling, + **kwargs) + output = F.scaled_dot_product_attention(query, + key, + value, + attn_mask=attention_mask, + dropout_p=dropout, + scale=scaling, + enable_gqa=True) + return output.transpose(1, 2).contiguous(), None + + +def register_attention() -> None: + """Make ``SDPA_GQA_DECODE`` available as an ``attn_implementation``, with SDPA's attention masks.""" + AttentionInterface.register(SDPA_GQA_DECODE, _sdpa_gqa_decode) + AttentionMaskInterface.register(SDPA_GQA_DECODE, sdpa_mask) def _chunk_log_probs(hidden, lm_head, targets, temperature): diff --git a/training/grpo/scripts/run_single_gpu.sh b/training/grpo/scripts/run_single_gpu.sh index 413b9ab59..8551e07f5 100755 --- a/training/grpo/scripts/run_single_gpu.sh +++ b/training/grpo/scripts/run_single_gpu.sh @@ -2,8 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team # -# Trains with configs/qwen3_4b_sft_1gpu.json on 1 GPU. Run scripts/prepare_data.sh first. +# Trains with configs/qwen3_4b_instruct_1gpu.json on 1 GPU. Run scripts/prepare_data.sh first. set -euo pipefail cd "$(dirname "$0")/.." -deepspeed --num_gpus 1 main.py --config configs/qwen3_4b_sft_1gpu.json +deepspeed --num_gpus 1 main.py --config configs/qwen3_4b_instruct_1gpu.json diff --git a/training/grpo/tools.py b/training/grpo/tools.py index 53dd1821e..55b5a8df5 100644 --- a/training/grpo/tools.py +++ b/training/grpo/tools.py @@ -8,11 +8,13 @@ boundary. """ +import ast import asyncio import os import re import sys import tempfile +import traceback from deepspeed.runtime.rollout import Tool @@ -102,13 +104,44 @@ """ +def imported_modules(tree: ast.AST) -> set: + """The top-level modules a parsed program imports: ``a`` for ``import a.b`` and for ``from a.b import c``.""" + modules = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + modules.update(alias.name.split(".")[0] for alias in node.names) + elif isinstance(node, ast.ImportFrom): + # A relative import has no module to check against the allowlist. + modules.add(node.module.split(".")[0] if node.module and not node.level else ".") + return modules + + +# A Markdown code fence around the whole program: ```py / ```python / ``` ... ```. +_CODE_FENCE = re.compile(r"\A\s*```[\w+-]*[ \t]*\n(.*?)\n?[ \t]*(?:```\s*)?\Z", re.DOTALL) + + +def strip_code_fence(code: str) -> str: + """``code`` without a Markdown fence around it, as Qwen3-2507 models write their ``code_interpreter`` calls. + + Only a fence that wraps the whole argument is removed (its closing line may be missing); code that merely + contains backticks is left as it is. + """ + match = _CODE_FENCE.match(code) + return match.group(1) if match else code + + def check_code_safety(code: str): - """Return an error message if ``code`` matches a blocked pattern or import, else None.""" + """Return an error message if ``code`` matches a blocked pattern, does not parse or imports a blocked module, else None.""" for pattern in DANGEROUS_PATTERNS: if re.search(pattern, code, re.IGNORECASE): return f"Code contains dangerous pattern: {pattern}" - imports = set(re.findall(r"import\s+(\w+)", code) + re.findall(r"from\s+(\w+)", code)) - for module in sorted(imports): + try: + tree = ast.parse(code, filename="") + except SyntaxError as e: + # Python compiles the whole program before running any of it, so code that does not parse cannot run; + # return the error Python would raise. + return "Code does not compile\n" + "".join(traceback.format_exception_only(type(e), e)).rstrip() + for module in sorted(imported_modules(tree)): if module not in ALLOWED_MODULES: return f"Import of '{module}' is not allowed" return None @@ -116,7 +149,10 @@ def check_code_safety(code: str): class CodeInterpreter(Tool): name = "code_interpreter" - description = "A tool for executing Python code in a safe sandbox environment." + # The tool returns only what the code prints, as a script does: a bare expression on the last line, as in a + # notebook, gives no output. + description = ("A tool for executing Python code in a safe sandbox environment. Only printed output is returned, " + "so print() every value you need.") parameters = { "type": "object", "properties": { @@ -153,6 +189,8 @@ async def execute(self, arguments) -> str: code = arguments.get("code", "") if isinstance(code, list): code = "\n".join(str(line) for line in code) + if isinstance(code, str): + code = strip_code_fence(code) output = await self._execute(code) if len(output) > self.max_output_chars: output = output[:self.max_output_chars] + "\n... (output truncated)" From 51c87a23bf90cd3065cc7367f7ac66ab92b6b298 Mon Sep 17 00:00:00 2001 From: Puyuan Yao Date: Sun, 8 Nov 2026 16:19:09 -0600 Subject: [PATCH 3/5] Add a notebook-style option to the code interpreter tool.show_last_value (off by default) makes the interpreter also return the value of a bare expression on the code's last line, as a notebook cell shows it, and says so in the tool description. Qwen3-4B-Instruct-2507 almost never ends its code with print(), so without it nearly every call returns empty output. Code that ends with print(...) runs as before. Signed-off-by: Puyuan Yao --- training/grpo/README.md | 3 +- training/grpo/config.py | 3 ++ .../grpo/configs/qwen3_4b_instruct_1gpu.json | 3 +- training/grpo/configs/qwen3_8b_8gpu.json | 3 +- training/grpo/configs/smoke_qwen3_0.6b.json | 3 +- training/grpo/main.py | 3 +- training/grpo/tools.py | 34 ++++++++++++++++--- 7 files changed, 43 insertions(+), 9 deletions(-) diff --git a/training/grpo/README.md b/training/grpo/README.md index b4f7f5427..634bec073 100644 --- a/training/grpo/README.md +++ b/training/grpo/README.md @@ -8,7 +8,7 @@ Rollout and training share one DeepSpeed **hybrid engine** (ZeRO-2): the policy 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 `` block, the rollout runs it in the sandboxed `code_interpreter` (`tools.py`) and appends the output 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. +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 `` 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. @@ -103,6 +103,7 @@ The main settings: - `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. diff --git a/training/grpo/config.py b/training/grpo/config.py index 5f776e633..c9e9bfa7b 100644 --- a/training/grpo/config.py +++ b/training/grpo/config.py @@ -55,6 +55,9 @@ class ToolConfig: 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 diff --git a/training/grpo/configs/qwen3_4b_instruct_1gpu.json b/training/grpo/configs/qwen3_4b_instruct_1gpu.json index 2da0f1042..d08b2a887 100644 --- a/training/grpo/configs/qwen3_4b_instruct_1gpu.json +++ b/training/grpo/configs/qwen3_4b_instruct_1gpu.json @@ -24,7 +24,8 @@ "timeout": 120, "memory_limit_gb": 4, "max_concurrency": 32, - "max_output_chars": 4096 + "max_output_chars": 4096, + "show_last_value": false }, "algorithm": { "eps_clip": 0.2, diff --git a/training/grpo/configs/qwen3_8b_8gpu.json b/training/grpo/configs/qwen3_8b_8gpu.json index 0a938523c..f0b1eb668 100644 --- a/training/grpo/configs/qwen3_8b_8gpu.json +++ b/training/grpo/configs/qwen3_8b_8gpu.json @@ -27,7 +27,8 @@ "timeout": 120, "memory_limit_gb": 4, "max_concurrency": 32, - "max_output_chars": 4096 + "max_output_chars": 4096, + "show_last_value": false }, "algorithm": { "eps_clip": 0.2, diff --git a/training/grpo/configs/smoke_qwen3_0.6b.json b/training/grpo/configs/smoke_qwen3_0.6b.json index 04ffe1732..24173e6b2 100644 --- a/training/grpo/configs/smoke_qwen3_0.6b.json +++ b/training/grpo/configs/smoke_qwen3_0.6b.json @@ -27,7 +27,8 @@ "timeout": 120, "memory_limit_gb": 4, "max_concurrency": 32, - "max_output_chars": 4096 + "max_output_chars": 4096, + "show_last_value": false }, "algorithm": { "eps_clip": 0.2, diff --git a/training/grpo/main.py b/training/grpo/main.py index c38bfd709..42a4693ea 100644 --- a/training/grpo/main.py +++ b/training/grpo/main.py @@ -91,7 +91,8 @@ def main() -> None: tool = CodeInterpreter(timeout=cfg.tool.timeout, memory_limit_gb=cfg.tool.memory_limit_gb, max_concurrency=cfg.tool.max_concurrency, - max_output_chars=cfg.tool.max_output_chars) + max_output_chars=cfg.tool.max_output_chars, + show_last_value=cfg.tool.show_last_value) rollout = build_rollout(cfg.rollout, generation_engine=RolloutEngineView(engine), tokenizer=tokenizer, diff --git a/training/grpo/tools.py b/training/grpo/tools.py index 55b5a8df5..e0590f0f4 100644 --- a/training/grpo/tools.py +++ b/training/grpo/tools.py @@ -147,12 +147,34 @@ def check_code_safety(code: str): return None +def show_last_value(code: str) -> str: + """``code`` changed to print the value of its last statement when that is an expression, as a notebook shows it. + + The value is printed with ``repr``, after anything the code printed itself, and not at all when it is None, so + code that ends with ``print(...)`` runs as it is. ``code`` must parse. + """ + last = ast.parse(code).body[-1:] + if not last or not isinstance(last[0], ast.Expr): + return code + expr = last[0] + # ast offsets count UTF-8 bytes from the start of a line. + lines = code.encode().split(b"\n") + before = b"\n".join(lines[:expr.lineno - 1] + [lines[expr.lineno - 1][:expr.col_offset]]) + after = b"\n".join([lines[expr.end_lineno - 1][expr.end_col_offset:]] + lines[expr.end_lineno:]) + value = ast.get_source_segment(code, expr.value) + # The last statement is at the top level, so the print runs right after it at no indent. + return (before.decode() + f"_last_value = ({value})" + after.decode() + + "\nif _last_value is not None:\n print(repr(_last_value))\n") + + class CodeInterpreter(Tool): name = "code_interpreter" - # The tool returns only what the code prints, as a script does: a bare expression on the last line, as in a - # notebook, gives no output. + # Without show_last_value the tool returns only what the code prints, as a script does. description = ("A tool for executing Python code in a safe sandbox environment. Only printed output is returned, " "so print() every value you need.") + # With show_last_value, like a notebook cell. + notebook_description = ("A tool for executing Python code in a safe sandbox environment. It returns the printed " + "output and, as in a notebook, the value of the last line when that is an expression.") parameters = { "type": "object", "properties": { @@ -168,11 +190,15 @@ def __init__(self, timeout: float = 120.0, memory_limit_gb: int = 4, max_concurrency: int = 32, - max_output_chars: int = 4096): + max_output_chars: int = 4096, + show_last_value: bool = False): self.timeout = timeout self.memory_limit = memory_limit_gb * 1024**3 self.max_concurrency = max_concurrency self.max_output_chars = max_output_chars + self.show_last_value = show_last_value + if show_last_value: + self.description = self.notebook_description self._loop = None self._semaphore = None @@ -204,7 +230,7 @@ async def _execute(self, code) -> str: if error is not None: return f"Error: {error}" async with self._limiter(): - return await self._run(code) + return await self._run(show_last_value(code) if self.show_last_value else code) async def _run(self, code: str) -> str: indented = "\n".join(" " + line for line in code.split("\n")) From 47b4f2e1988b851f9a1922fda7a3dc07e99e67b2 Mon Sep 17 00:00:00 2001 From: Puyuan Yao Date: Sun, 8 Nov 2026 20:44:47 -0600 Subject: [PATCH 4/5] Port the OPSD example to the per-sample rollout API The rollout in deepspeed.runtime.rollout now takes and returns lists of per-sample RolloutRecords (stage 1 of RFC #8942). The OPSD trainer builds one record per prompt row from its real tokens and gets the training tensors and the response mask from ResponseTokenBatch.from_rollout, which replaces utils.build_response_mask. main.py passes the student engine as generation_engine. Both rollout benchmarks build their requests from records, and the READMEs say which DeepSpeed the example needs. Signed-off-by: Puyuan Yao --- benchmarks/opsd/README.md | 3 +- .../opsd/benchmark_hybrid_engine_rollout.py | 8 ++-- training/opsd/README.md | 10 +---- training/opsd/benchmarks/bench_decode_1p1r.py | 10 ++--- training/opsd/data.py | 6 +-- training/opsd/main.py | 7 +--- training/opsd/tests/test_losses.py | 19 --------- training/opsd/trainer.py | 36 +++++++++------- training/opsd/utils.py | 41 ------------------- 9 files changed, 37 insertions(+), 103 deletions(-) delete mode 100644 training/opsd/utils.py diff --git a/benchmarks/opsd/README.md b/benchmarks/opsd/README.md index e58be3800..c867f7f32 100644 --- a/benchmarks/opsd/README.md +++ b/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. diff --git a/benchmarks/opsd/benchmark_hybrid_engine_rollout.py b/benchmarks/opsd/benchmark_hybrid_engine_rollout.py index 3d4b610ff..0079e13a8 100644 --- a/benchmarks/opsd/benchmark_hybrid_engine_rollout.py +++ b/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): diff --git a/training/opsd/README.md b/training/opsd/README.md index e79e32b01..ca3bdd6a1 100644 --- a/training/opsd/README.md +++ b/training/opsd/README.md @@ -37,7 +37,6 @@ training/opsd/ ├── teacher.py # frozen teacher + TeacherLogitCache ├── losses.py # streamed distillation loss (fwd/rev KL, JSD) ├── data.py # PromptDataset + left-padded collator -├── utils.py # response-mask helpers ├── configs/ │ ├── ds_zero3.json # base DeepSpeed ZeRO-3 + hybrid engine │ ├── opsd_hybrid_engine.json # production-ish hybrid-engine OPSD config @@ -57,19 +56,12 @@ training/opsd/ ### Install -This example uses DeepSpeed's `deepspeed.runtime.rollout` module (hybrid -engine + rollout API), which landed on `master` after the 0.19.2 release and -is not yet on PyPI. Until DeepSpeed 0.19.3 is released, install it from -source: +This example uses DeepSpeed's `deepspeed.runtime.rollout` module (hybrid engine + rollout API) in the form that stage 1 of [RFC #8942](https://github.com/deepspeedai/DeepSpeed/issues/8942) gives it: requests and results are lists of per-sample `RolloutRecord`s, and `ResponseTokenBatch` turns a result into training tensors. That API is not in a DeepSpeed release yet, so install DeepSpeed from source with that code, then: ``` -pip install git+https://github.com/deepspeedai/DeepSpeed.git@master pip install transformers>=5.0.0 datasets accelerate ``` -Once DeepSpeed 0.19.3 ships, a plain `pip install deepspeed` works and the -source line above can be dropped. - ### Hybrid-engine training ``` diff --git a/training/opsd/benchmarks/bench_decode_1p1r.py b/training/opsd/benchmarks/bench_decode_1p1r.py index e50c3cced..5909e9a2b 100644 --- a/training/opsd/benchmarks/bench_decode_1p1r.py +++ b/training/opsd/benchmarks/bench_decode_1p1r.py @@ -20,7 +20,7 @@ from deepspeed.accelerator import get_accelerator from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRollout -from deepspeed.runtime.rollout.base import RolloutRequest, SamplingConfig +from deepspeed.runtime.rollout.base import RolloutRecord, RolloutRequest, SamplingConfig def bench_decode_raw(model, tokenizer, device, prompt_len=64, max_new_tokens=64, num_warmup=3, num_iters=10): @@ -113,10 +113,9 @@ def avg_of_avg(list_of_lists, n): def bench_hybrid_rollout(rollout, tokenizer, device, prompt_len=64, max_new_tokens=64, num_warmup=3, num_iters=10): """Benchmark the full HybridEngineRollout.generate() path.""" - input_ids = torch.randint(10, 1000, (1, prompt_len), device=device) - attn_mask = torch.ones(1, prompt_len, dtype=torch.long, device=device) + input_ids = torch.randint(10, 1000, (prompt_len, )).tolist() sampling = SamplingConfig(max_new_tokens=max_new_tokens, temperature=1.0, top_p=1.0) - request = RolloutRequest(prompt_ids=input_ids, prompt_attention_mask=attn_mask) + request = RolloutRequest(samples=[RolloutRecord(tokens=input_ids, attention_mask=[1] * prompt_len)]) times = [] for _ in range(num_warmup + num_iters): @@ -162,7 +161,8 @@ def main(): print(f" Total: {raw['total_ms']:.1f} ms") print(f"\n=== HybridEngineRollout benchmark (graph_capture={args.graph_capture}) ===") - engine = type('Engine', (), {'module': model})() # lightweight wrapper + # A lightweight stand-in for a DeepSpeed engine: the model, with its weights unpartitioned. + engine = type('Engine', (), {'module': model, 'zero_optimization_partition_weights': lambda self: False})() from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRolloutConfig cfg = HybridEngineRolloutConfig(use_graph_capture=args.graph_capture) rollout = HybridEngineRollout(engine, tokenizer, cfg=cfg) diff --git a/training/opsd/data.py b/training/opsd/data.py index a7c8780fc..ff2d1a5b7 100644 --- a/training/opsd/data.py +++ b/training/opsd/data.py @@ -8,9 +8,9 @@ wrapped in a user-role message with ``add_generation_prompt=True`` so the student generates the assistant turn. -Batches are **left-padded** because causal generation requires real tokens at - the right edge — :class:`deepspeed.runtime.rollout.RolloutRequest` and the hybrid-engine -backend both assume this layout. +Batches are **left-padded** so the real tokens sit at the right edge; the trainer +turns each row's real tokens into one :class:`deepspeed.runtime.rollout.RolloutRecord` +of the rollout request. """ import json diff --git a/training/opsd/main.py b/training/opsd/main.py index dde605f30..d251ccba4 100644 --- a/training/opsd/main.py +++ b/training/opsd/main.py @@ -90,12 +90,7 @@ def main() -> None: teacher = TeacherWrapper(cfg.teacher, world_size=dist_world_size()) # --- rollout engine ---------------------------------------------------- - rollout = build_rollout( - cfg.rollout, - student_engine=student_engine, - tokenizer=tokenizer, - student_model_path=cfg.student.model_name_or_path, - ) + rollout = build_rollout(cfg.rollout, generation_engine=student_engine, tokenizer=tokenizer) # --- dataloader -------------------------------------------------------- dataset = PromptDataset( diff --git a/training/opsd/tests/test_losses.py b/training/opsd/tests/test_losses.py index 652cccc42..b744b44da 100644 --- a/training/opsd/tests/test_losses.py +++ b/training/opsd/tests/test_losses.py @@ -12,7 +12,6 @@ import torch from losses import chunked_distillation_loss -from utils import build_response_mask @pytest.mark.parametrize("loss_type", ["forward_kl", "reverse_kl", "jsd"]) @@ -127,21 +126,3 @@ def test_all_zero_mask_returns_zero(): mask = torch.zeros(2, 8) loss = chunked_distillation_loss(s, t, mask, loss_type="reverse_kl") assert loss.item() == pytest.approx(0.0, abs=1e-6) - - -def test_build_response_mask_basic(): - attention_mask = torch.tensor([[1, 1, 1, 1, 0], [1, 1, 1, 1, 1]]) - response_start_idx = torch.tensor([2, 3]) - resp = build_response_mask(response_start_idx, attention_mask) - expected = torch.tensor([[0, 0, 1, 1, 0], [0, 0, 0, 1, 1]]) - assert torch.equal(resp, expected) - - -def test_build_response_mask_validates_shapes(): - with pytest.raises(ValueError, match="response_start_idx must be 1-D"): - build_response_mask(torch.zeros(2, 2), torch.ones(2, 4)) - with pytest.raises(ValueError, match="attention_mask must be 2-D"): - build_response_mask(torch.zeros(2), torch.ones(4)) - with pytest.raises(ValueError, match="batch"): - build_response_mask(torch.zeros(3), torch.ones(2, 4)) - diff --git a/training/opsd/trainer.py b/training/opsd/trainer.py index 173f0d2a1..e2b55b312 100644 --- a/training/opsd/trainer.py +++ b/training/opsd/trainer.py @@ -33,8 +33,17 @@ from config import OPSDConfig from losses import streamed_distillation_loss -from utils import build_response_mask -from deepspeed.runtime.rollout import RolloutEngine, RolloutRequest, SamplingConfig +from deepspeed.runtime.rollout import ResponseTokenBatch, RolloutEngine, RolloutRequest, SamplingConfig +from deepspeed.runtime.rollout.base import RolloutRecord + + +def rollout_request(prompt_ids: torch.Tensor, prompt_attention_mask: torch.Tensor) -> RolloutRequest: + """A rollout request with one record per row of a left-padded prompt batch, holding the row's real tokens.""" + samples = [] + for row, (ids, mask) in enumerate(zip(prompt_ids.tolist(), prompt_attention_mask.tolist())): + tokens = [token for token, keep in zip(ids, mask) if keep] + samples.append(RolloutRecord(group_index=row, tokens=tokens, attention_mask=[1] * len(tokens))) + return RolloutRequest(samples=samples) def _is_rank_zero() -> bool: @@ -101,9 +110,6 @@ def train(self) -> None: def _train_step(self, batch) -> dict: t_start = time.time() - prompt_ids = batch["prompt_ids"].to(self.device, non_blocking=True) - prompt_attn = batch["prompt_attention_mask"].to(self.device, non_blocking=True) - # Sync student weights into the rollout backend. self.rollout.sync_weights(self.step) @@ -117,14 +123,10 @@ def _train_step(self, batch) -> dict: top_k=self.cfg.rollout.top_k, n_samples_per_prompt=self.cfg.rollout.n_samples_per_prompt, ) - roll = self.rollout.generate( - RolloutRequest(prompt_ids=prompt_ids, prompt_attention_mask=prompt_attn), - sampling, - ) - input_ids = roll.input_ids.to(self.device, non_blocking=True) - attention_mask = roll.attention_mask.to(self.device, non_blocking=True) - response_start_idx = roll.response_start_idx.to(self.device, non_blocking=True) - response_mask = build_response_mask(response_start_idx, attention_mask) + roll = self.rollout.generate(rollout_request(batch["prompt_ids"], batch["prompt_attention_mask"]), sampling) + # Prompt + response per sample, left-padded to one width, with the response targets shifted by one. + tokens = ResponseTokenBatch.from_rollout(roll, pad_token_id=self._pad_token_id(), device=self.device) + input_ids, attention_mask = tokens.input_ids, tokens.attention_mask t_rollout = time.time() - t_start # --- Phase 1: teacher forward → host-cached logits ---------------- @@ -141,8 +143,8 @@ def _train_step(self, batch) -> dict: # Shift for next-token prediction: logits at position t predict token # at t+1, so the loss aligns student_logits[:, :-1] with the position # t+1 entries of the response mask. - student_logits_shifted = student_logits[:, :-1, :] - mask_shifted = response_mask[:, 1:].contiguous() + student_logits_shifted = tokens.select_causal_logits(student_logits) + mask_shifted = tokens.response_mask def _fetch(start: int, end: int) -> torch.Tensor: # The cache holds *unshifted* teacher logits; for the next-token @@ -184,6 +186,10 @@ def _fetch(start: int, end: int) -> torch.Tensor: "response_tokens": int(mask_shifted.sum().item()), } + def _pad_token_id(self) -> int: + pad = self.tokenizer.pad_token_id + return pad if pad is not None else self.tokenizer.eos_token_id + # ------------------------------------------------------------------ # Logging / checkpointing # ------------------------------------------------------------------ diff --git a/training/opsd/utils.py b/training/opsd/utils.py deleted file mode 100644 index 4180da5f2..000000000 --- a/training/opsd/utils.py +++ /dev/null @@ -1,41 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# DeepSpeed Team -"""Small tensor/masking helpers shared by trainer, losses, and tests. - -These intentionally stay free of DeepSpeed / distributed imports so the -non-distributed unit tests can exercise them on CPU without a torchrun -launcher. -""" - -import torch - - -def build_response_mask(response_start_idx: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - """Mark positions belonging to the response (not prompt, not padding). - - Args: - response_start_idx: ``[B]`` int tensor — the first column index that is - part of the response, per sample. For *right-padded* prompts this - equals the prompt's token count; for the more common *left-padded* - convention used by causal generation it equals the prompt section - length (i.e. the column where prompt ends and response begins). - attention_mask: ``[B, T]`` — 1 on real tokens (prompt + response), 0 on - padding. - - Returns: - ``[B, T]`` 0/1 mask with the same dtype as ``attention_mask``. 1 only - at positions ``t >= response_start_idx[b]`` that are also attended. - """ - if response_start_idx.dim() != 1: - raise ValueError(f"response_start_idx must be 1-D, got shape {tuple(response_start_idx.shape)}") - if attention_mask.dim() != 2: - raise ValueError(f"attention_mask must be 2-D, got shape {tuple(attention_mask.shape)}") - B, T = attention_mask.shape - if response_start_idx.shape[0] != B: - raise ValueError(f"response_start_idx batch ({response_start_idx.shape[0]}) != " - f"attention_mask batch ({B})") - - pos = torch.arange(T, device=attention_mask.device).unsqueeze(0).expand(B, T) - is_response = pos >= response_start_idx.to(pos.dtype).unsqueeze(1) - return is_response.to(attention_mask.dtype) * attention_mask - From 5c1bc1a5ba6e6b37b718d044043896b1e4fe6c4b Mon Sep 17 00:00:00 2001 From: Puyuan Yao Date: Sun, 8 Nov 2026 21:16:48 -0600 Subject: [PATCH 5/5] Leave grouped-query decode attention to the framework Drop the sdpa_gqa_decode attention that policy.py registered with Hugging Face, and go back to plain SDPA in the configs. Reading grouped KV heads in place during decode is a rollout concern for DeepSpeed (or Transformers) to handle for every model, not something each example should register. Until then, decode pays Hugging Face's KV-head copy. Signed-off-by: Puyuan Yao --- training/grpo/README.md | 4 +-- training/grpo/config.py | 3 +- .../grpo/configs/qwen3_4b_instruct_1gpu.json | 2 +- training/grpo/configs/qwen3_8b_8gpu.json | 2 +- training/grpo/configs/smoke_qwen3_0.6b.json | 2 +- training/grpo/main.py | 3 +- training/grpo/policy.py | 34 ------------------- 7 files changed, 6 insertions(+), 44 deletions(-) diff --git a/training/grpo/README.md b/training/grpo/README.md index 634bec073..71920910c 100644 --- a/training/grpo/README.md +++ b/training/grpo/README.md @@ -21,7 +21,7 @@ With one update per rollout, the old policy is the current one, so the old log-p | `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, and the decode attention | +| `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 | @@ -86,8 +86,6 @@ Rank 0 prints the loss and the mean reward after each step. When `eval.path` is 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. -The configs set `model.attn_implementation` to `sdpa_gqa_decode`, which `policy.py` registers with Hugging Face: SDPA, except that decode steps read Qwen3's shared KV heads in place. Hugging Face's SDPA copies each KV head once per query head it serves whenever there is an attention mask, and rollout decodes with a mask over each row's whole KV cache, so the copy dominates decode time; skipping it makes Qwen3-4B decode about three times faster on an A100 with the same outputs. Prefill and training use Hugging Face's SDPA unchanged. - 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: diff --git a/training/grpo/config.py b/training/grpo/config.py index c9e9bfa7b..e28501c44 100644 --- a/training/grpo/config.py +++ b/training/grpo/config.py @@ -18,8 +18,7 @@ class ModelConfig: model_name_or_path: str dtype: str = "bfloat16" - # "sdpa_gqa_decode" (policy.py) is SDPA with faster decode steps for models with grouped KV heads. - attn_implementation: str = "sdpa_gqa_decode" + attn_implementation: str = "sdpa" gradient_checkpointing: bool = True trust_remote_code: bool = False diff --git a/training/grpo/configs/qwen3_4b_instruct_1gpu.json b/training/grpo/configs/qwen3_4b_instruct_1gpu.json index d08b2a887..a5d979819 100644 --- a/training/grpo/configs/qwen3_4b_instruct_1gpu.json +++ b/training/grpo/configs/qwen3_4b_instruct_1gpu.json @@ -2,7 +2,7 @@ "model": { "model_name_or_path": "Qwen/Qwen3-4B-Instruct-2507", "dtype": "bfloat16", - "attn_implementation": "sdpa_gqa_decode", + "attn_implementation": "sdpa", "gradient_checkpointing": true }, "rollout": { diff --git a/training/grpo/configs/qwen3_8b_8gpu.json b/training/grpo/configs/qwen3_8b_8gpu.json index f0b1eb668..b171aa112 100644 --- a/training/grpo/configs/qwen3_8b_8gpu.json +++ b/training/grpo/configs/qwen3_8b_8gpu.json @@ -2,7 +2,7 @@ "model": { "model_name_or_path": "Qwen/Qwen3-8B", "dtype": "bfloat16", - "attn_implementation": "sdpa_gqa_decode", + "attn_implementation": "sdpa", "gradient_checkpointing": true }, "rollout": { diff --git a/training/grpo/configs/smoke_qwen3_0.6b.json b/training/grpo/configs/smoke_qwen3_0.6b.json index 24173e6b2..1aceb8968 100644 --- a/training/grpo/configs/smoke_qwen3_0.6b.json +++ b/training/grpo/configs/smoke_qwen3_0.6b.json @@ -2,7 +2,7 @@ "model": { "model_name_or_path": "Qwen/Qwen3-0.6B", "dtype": "bfloat16", - "attn_implementation": "sdpa_gqa_decode", + "attn_implementation": "sdpa", "gradient_checkpointing": true }, "rollout": { diff --git a/training/grpo/main.py b/training/grpo/main.py index 42a4693ea..59db37c33 100644 --- a/training/grpo/main.py +++ b/training/grpo/main.py @@ -26,7 +26,7 @@ from config import GRPOConfig from data import PromptDataset, collate_rows, load_prompts from deepspeed.runtime.rollout import build_rollout -from policy import PolicyModel, RolloutEngineView, register_attention +from policy import PolicyModel, RolloutEngineView from tools import CodeInterpreter from trainer import GRPOTrainer @@ -64,7 +64,6 @@ def main() -> None: tokenizer.pad_token = tokenizer.eos_token # --- policy + DeepSpeed hybrid engine ---------------------------------- - register_attention() model = AutoModelForCausalLM.from_pretrained( cfg.model.model_name_or_path, dtype=_resolve_dtype(cfg.model.dtype), diff --git a/training/grpo/policy.py b/training/grpo/policy.py index 04332162f..ff9473e05 100644 --- a/training/grpo/policy.py +++ b/training/grpo/policy.py @@ -17,40 +17,6 @@ import torch import torch.nn.functional as F from torch.utils.checkpoint import checkpoint -from transformers import AttentionInterface -from transformers.integrations.sdpa_attention import sdpa_attention_forward -from transformers.masking_utils import AttentionMaskInterface, sdpa_mask - -# SDPA that reads grouped KV heads in place on decode steps; see _sdpa_gqa_decode. -SDPA_GQA_DECODE = "sdpa_gqa_decode" - - -def _sdpa_gqa_decode(module, query, key, value, attention_mask, dropout=0.0, scaling=None, **kwargs): - """Hugging Face SDPA attention, except that one-token decode steps read grouped KV heads in place. - - With an attention mask, Hugging Face's SDPA copies each KV head once per query head it serves before - attending. Rollout decodes with a mask over each row's whole KV cache, so for Qwen3 (8 KV heads for 32 - query heads) that copy reads and writes the full cache four times over at every step, and dominates - decode time. ``enable_gqa`` attends to the shared heads directly. Prefill and training, where the query - is longer than one token, keep Hugging Face's SDPA path. - """ - if query.shape[2] != 1 or getattr(module, "num_key_value_groups", 1) == 1: - return sdpa_attention_forward(module, query, key, value, attention_mask, dropout=dropout, scaling=scaling, - **kwargs) - output = F.scaled_dot_product_attention(query, - key, - value, - attn_mask=attention_mask, - dropout_p=dropout, - scale=scaling, - enable_gqa=True) - return output.transpose(1, 2).contiguous(), None - - -def register_attention() -> None: - """Make ``SDPA_GQA_DECODE`` available as an ``attn_implementation``, with SDPA's attention masks.""" - AttentionInterface.register(SDPA_GQA_DECODE, _sdpa_gqa_decode) - AttentionMaskInterface.register(SDPA_GQA_DECODE, sdpa_mask) def _chunk_log_probs(hidden, lm_head, targets, temperature):