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/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..71920910c --- /dev/null +++ b/training/grpo/README.md @@ -0,0 +1,116 @@ +# Multi-Turn GRPO with Tool Calling (ReTool) + +This example trains Qwen3 with GRPO to solve math problems while calling a Python interpreter. In each trajectory the model can write code, read the interpreter's output and continue reasoning, for several turns, before giving its answer. + +Rollout and training share one DeepSpeed **hybrid engine** (ZeRO-2): the policy generates with its own training weights, so there is no separate inference server and no weight sync. The example is built on `deepspeed.runtime.rollout`, which provides the multi-turn rollout, the tool interface and the GRPO loss. + +## How it works + +Each training step has three phases: + +1. **Rollout.** Each rank takes its share of the step's prompts and samples `n_samples_per_prompt` trajectories per prompt with `HybridEngineRollout`. When the model writes a `` block, the rollout runs it in the sandboxed `code_interpreter` (`tools.py`) and appends what the code prints as a tool message, for up to `max_turns` turns. With continuous batching, each turn takes a free decode row and releases it when the turn ends, and tools run while the other rows keep decoding. +2. **Reward.** A DAPO strict-box check (`reward.py`) scores each trajectory +1 if the last `\boxed{...}` of the response equals the label, and −1 otherwise. `group_advantages` normalizes each reward within its prompt's group. +3. **Update.** `ResponseTokenBatch.from_rollout` turns the trajectories into training targets with a mask that keeps only the tokens the model generated, so prompts, padding and tool output are not trained on. `grpo_policy_loss` computes PPO-clip with clip-higher (0.2 / 0.28) and no KL term. The loss is averaged over every generated token across ranks, or per trajectory first with `algorithm.loss_aggregation: "sample_mean"`, and the step takes one optimizer update, accumulated over micro-batches. + +With one update per rollout, the old policy is the current one, so the old log-probs are the current ones detached and the importance ratio is 1. `rollout.return_log_probs` records them during the rollout and takes them from there instead, through `old_log_probs_from_rollout`. + +## Layout + +| Path | Contents | +|---|---| +| `main.py` | Entry point: builds the model, the hybrid engine, the rollout and the data, then trains | +| `trainer.py` | The training loop: rollout, reward, update, evaluation, logging and checkpoints | +| `config.py` | The configuration dataclasses and their checks | +| `policy.py` | The policy as DeepSpeed trains it: token log-probs from a chunked LM head | +| `tools.py` | The sandboxed Python interpreter the model calls | +| `reward.py` | The strict-box answer check | +| `data.py` | JSONL prompt loading | +| `configs/` | The three training configs | +| `scripts/` | Data preparation and one launch script per config | + +## Prerequisites + +- **DeepSpeed with the multi-turn rollout.** `deepspeed.runtime.rollout` with multi-turn tool calling, continuous batching and the GRPO loss is not in a DeepSpeed release yet; it is stage 1 of [RFC #8942](https://github.com/deepspeedai/DeepSpeed/issues/8942). Install DeepSpeed from source with that code. +- **Python packages.** `transformers` 5.0 or later, `numpy`, and `huggingface_hub` for the `hf` download tool: + + ```bash + pip install "transformers>=5.0.0" numpy huggingface_hub + ``` + +- **GPUs.** The smoke test runs on one GPU. The single-GPU config is sized for one 40 GB A100, and also needs more than 64 GB of host memory for the optimizer state it keeps on the CPU; 160 GB is enough. The distributed config is sized for 8 GPUs with 80 GB each. + +## Data + +The training set is `dapo-math-17k` and the evaluation set is `aime-2024`, in JSONL with a chat-message `prompt` and a string `label`. Download both into `data/`: + +```bash +cd training/grpo +bash scripts/prepare_data.sh +``` + +The script also wraps each AIME prompt in the instructions every `dapo-math-17k` prompt has. The downloaded problems don't ask for a `\boxed{}` answer, which is the only form the scorer accepts, so without them nearly every sample scores −1. + +## Run + +```bash +cd training/grpo + +# smoke test: Qwen3-0.6B, three steps on one GPU +bash scripts/run_smoke_test.sh + +# Qwen3-4B-Instruct-2507 on one GPU +bash scripts/run_single_gpu.sh + +# Qwen3-8B on 8 GPUs +bash scripts/run_8gpu.sh +``` + +Each script runs `deepspeed --num_gpus 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_instruct_1gpu.json` | `qwen3_8b_8gpu.json` | +|---|---|---|---| +| Model | `Qwen/Qwen3-0.6B` | `Qwen/Qwen3-4B-Instruct-2507` | `Qwen/Qwen3-8B` | +| GPUs | 1 | 1 | 8 | +| Steps | 3 | 500 | 500 | +| Trajectories per step | 8 (2 prompts × 4) | 64 (8 prompts × 8) | 256 (32 prompts × 8) | +| Trajectories per micro-batch, all ranks | 2 | 2 | 8 | +| Turns, context, tokens per turn | 4, 4096, 1024 | 16, 8192, 4096 | 16, 8192, 4096 | +| Continuous batching rows per GPU | 8 | 16 | 16 | +| ZeRO-2 | optimizer on the GPU | optimizer state on the CPU | optimizer on the GPU | +| AIME evaluation | off | every 100 steps | every 20 steps | +| Thinking mode | off | none (an instruct model) | off | + +All three turn on every rollout optimization that works together: CUDA-graph decode (`use_graph_capture`) and multi-turn continuous batching (`continuous_batch_size`). Training uses bf16, gradient checkpointing, the chunked LM head, SDPA attention, and ZeRO-2 with overlapped communication. + +Each config holds the DeepSpeed config inline under `deepspeed_config`, which can also be the path of a DeepSpeed JSON. The trainer fills in `train_micro_batch_size_per_gpu` and `gradient_accumulation_steps`, so that one step covers all of a rank's trajectories. + +The main settings: + +- `rollout.max_new_tokens`: generated tokens per trajectory, **summed over its turns**. +- `rollout.max_new_tokens_per_turn`: an additional cap on each turn. `eval` has the same two fields. +- `rollout.max_context_len`: total prompt, response and tool tokens per trajectory. A trajectory that reaches it stops as truncated. +- `rollout.continuous_batch_size`: decode rows per GPU for continuous batching; leave it unset for static batching. Rows are sized by `max_context_len`, so the KV cache holds `continuous_batch_size × max_context_len` tokens, about 144 KiB per token for Qwen3-4B and Qwen3-8B. Lower it if rollout runs out of memory. +- `rollout.return_log_probs`: record each generated token's log-prob during the rollout and use them as the PPO old log-probs, off by default. It applies to evaluation rollouts too. +- `rollout.generate_batch_size`: trajectories per `generate` call, in whole prompts; by default a rank's whole share of the step. +- `rollout.chat_template_kwargs`: extra arguments for the chat template, such as `{"enable_thinking": false}` for Qwen3. +- `rollout.repetition_penalty`: 1.0 by default; set it to null to use the model's `generation_config`. +- `training.global_batch_size`: trajectories per step across all ranks, a multiple of ranks × `n_samples_per_prompt` so that each rank rolls out whole prompts. +- `training.micro_batch_size`: trajectories per forward and backward pass across all ranks, split evenly between them. It must divide `global_batch_size`; gradient accumulation covers the rest, so it changes memory and speed, not the update. +- `tool.timeout`, `tool.max_output_chars`: limits on each interpreter call; long output is cut so that one large print cannot fill the context. +- `tool.show_last_value`: also return the value of a bare expression on the code's last line, as a notebook does, and say so in the tool description. Off by default, so only printed output comes back. Models that write notebook-style code, such as Qwen3-4B-Instruct-2507, which almost never ends its code with `print()`, otherwise get empty tool output. +- `algorithm.tool_call_bonus`: an optional reward bonus for calling tools on wrong answers, off by default. +- `algorithm.loss_aggregation`: `"token_mean"` (default) weighs every generated token of the step equally; `"sample_mean"` averages each trajectory over its own tokens first, then over trajectories. + +## Notes and limitations + +- **The sandbox only screens obvious misuse.** It applies an import allowlist, a pattern filter, a memory limit and a timeout. It is not a security boundary, so run untrusted training on isolated machines. +- **ZeRO-2 only.** Multi-turn continuous batching does not support ZeRO-3, whose parameter gathers cannot run inside a captured graph or across rows that finish at different steps. +- **Hybrid engine without injected kernels.** DeepSpeed has no inference policy for Qwen3, so generation runs the Hugging Face model on the training weights, with DeepSpeed's own static cache, graph capture and continuous batching around it. + +## Acknowledgement + +This example is adapted from the [miles retool_v2](https://github.com/radixark/miles/tree/main/examples/retool_v2) recipe (Copyright 2025-2026 RadixArk, Inc., Apache License 2.0). The code interpreter in `tools.py` and the answer check in `reward.py` are based on miles. The answer check is in turn adapted from verl and EleutherAI's lm-evaluation-harness, both under the Apache License 2.0. diff --git a/training/grpo/config.py b/training/grpo/config.py new file mode 100644 index 000000000..e28501c44 --- /dev/null +++ b/training/grpo/config.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""GRPO application configuration. + +``GRPOConfig`` is loaded from JSON. The ``rollout`` sub-config is consumed +by DeepSpeed's rollout engine; optimizer, ZeRO and hybrid-engine settings +live in ``deepspeed_config``, either inline or as the path of a DeepSpeed JSON. +""" + +import json +from dataclasses import asdict, dataclass, field +from typing import Optional, Union + +from deepspeed.runtime.rollout import RolloutConfig as _BaseRolloutConfig + + +@dataclass +class ModelConfig: + model_name_or_path: str + dtype: str = "bfloat16" + attn_implementation: str = "sdpa" + gradient_checkpointing: bool = True + trust_remote_code: bool = False + + +@dataclass +class RolloutConfig(_BaseRolloutConfig): + """Extends DeepSpeed's RolloutConfig with the sampling knobs of training rollouts. + + The base config holds the engine settings: ``continuous_batch_size`` (decode rows for continuous + batching, sized by ``max_context_len``; None keeps static batching) and ``return_log_probs`` (record + each generated token's log-prob, used as the PPO old log-probs when set). + """ + multi_turn: bool = True + max_turns: Optional[int] = 16 + max_context_len: Optional[int] = 16384 + # Generated tokens per trajectory, summed over its turns. + max_new_tokens: int = 8192 + # Generated tokens per turn, on top of max_new_tokens; None caps turns only by what is left of it. + max_new_tokens_per_turn: Optional[int] = None + temperature: float = 1.0 + # 1.0 disables the penalty; None takes the value in the model's generation_config. + repetition_penalty: Optional[float] = 1.0 + n_samples_per_prompt: int = 8 + # Trajectories per generate call, in whole prompts and at least one; None rolls out a rank's share of the + # step in one call. Bounds the KV cache and prefill memory of a call. + generate_batch_size: Optional[int] = None + + +@dataclass +class ToolConfig: + timeout: float = 120.0 + memory_limit_gb: int = 4 + max_concurrency: int = 32 + # Long tool outputs are cut so one print cannot overflow the context. + max_output_chars: int = 4096 + # Also return the value of a bare expression on the code's last line, as a notebook shows it; off returns only + # what the code prints. + show_last_value: bool = False + + +@dataclass +class AlgorithmConfig: + eps_clip: float = 0.2 + eps_clip_high: float = 0.28 + # Sequence positions per LM-head chunk when computing token log-probs. + logprob_chunk_size: int = 1024 + # Reward bonus for tool calls on wrong answers. + tool_call_bonus: bool = False + # "token_mean": every generated token of the step weighs the same. + # "sample_mean": each trajectory's loss is averaged over its own tokens, then over trajectories. + loss_aggregation: str = "token_mean" + + +@dataclass +class TrainingConfig: + # Trajectories per step across all ranks, in whole prompts per rank: rolled out together, then + # one optimizer step over all of them. + global_batch_size: int = 256 + # Trajectories per forward/backward pass across all ranks, split evenly between them; None means + # one pass per step, global_batch_size. + micro_batch_size: Optional[int] = None + num_rollouts: int = 3000 + logging_steps: int = 1 + save_steps: int = 1000 + save_dir: str = "./grpo_ckpt" + # Continue from the latest checkpoint in this directory (weights, optimizer state, step count and data + # position); empty starts from the model. It may differ from save_dir. + resume_from: str = "" + seed: int = 42 + + +@dataclass +class DataConfig: + path: str = "" + prompt_key: str = "prompt" + label_key: str = "label" + shuffle: bool = True + + +@dataclass +class EvalConfig: + # Empty path disables evaluation. + path: str = "" + interval: int = 20 + # Evaluate once more after the last step. + at_end: bool = True + n_samples_per_prompt: int = 16 + max_new_tokens: int = 16384 + max_new_tokens_per_turn: Optional[int] = None + temperature: float = 1.0 + + +@dataclass +class GRPOConfig: + model: ModelConfig + rollout: RolloutConfig = field(default_factory=RolloutConfig) + tool: ToolConfig = field(default_factory=ToolConfig) + algorithm: AlgorithmConfig = field(default_factory=AlgorithmConfig) + training: TrainingConfig = field(default_factory=TrainingConfig) + data: DataConfig = field(default_factory=DataConfig) + eval: EvalConfig = field(default_factory=EvalConfig) + # The DeepSpeed config: a dict, or the path of a JSON file. + deepspeed_config: Union[str, dict] = "" + + @classmethod + def from_json(cls, path: str) -> "GRPOConfig": + with open(path, "r") as f: + raw = json.load(f) + return cls( + model=ModelConfig(**raw["model"]), + rollout=RolloutConfig(**raw.get("rollout", {})), + tool=ToolConfig(**raw.get("tool", {})), + algorithm=AlgorithmConfig(**raw.get("algorithm", {})), + training=TrainingConfig(**raw.get("training", {})), + data=DataConfig(**raw.get("data", {})), + eval=EvalConfig(**raw.get("eval", {})), + deepspeed_config=raw.get("deepspeed_config", ""), + ) + + def to_dict(self) -> dict: + return asdict(self) + + def load_deepspeed_config(self) -> dict: + """A copy of the DeepSpeed config, read from its file when it is a path.""" + if isinstance(self.deepspeed_config, dict): + return json.loads(json.dumps(self.deepspeed_config)) + if not self.deepspeed_config: + raise ValueError("deepspeed_config is not set") + with open(self.deepspeed_config, "r") as f: + return json.load(f) + + def validate(self, world_size: int) -> None: + if self.rollout.engine != "hybrid_engine": + raise ValueError(f"Unknown rollout engine {self.rollout.engine!r}; expected 'hybrid_engine'") + if self.rollout.n_samples_per_prompt < 2: + raise ValueError("GRPO needs at least 2 samples per prompt to compare them") + if self.rollout.temperature <= 0: + raise ValueError("GRPO samples with temperature > 0") + if self.rollout.generate_batch_size is not None and self.rollout.generate_batch_size < 1: + raise ValueError("rollout.generate_batch_size must be positive") + if self.rollout.continuous_batch_size is not None and not self.rollout.max_context_len: + raise ValueError("continuous batching sizes its rows by rollout.max_context_len; set it") + if self.algorithm.loss_aggregation not in ("token_mean", "sample_mean"): + raise ValueError(f"Unknown loss_aggregation {self.algorithm.loss_aggregation!r}; " + "expected 'token_mean' or 'sample_mean'") + # Each rank rolls out whole prompts, every one with all of its samples. + trajectories_per_prompt_round = world_size * self.rollout.n_samples_per_prompt + if self.training.global_batch_size < 1 or self.training.global_batch_size % trajectories_per_prompt_round: + raise ValueError(f"global_batch_size {self.training.global_batch_size} must be a multiple of world size " + f"x n_samples_per_prompt = {trajectories_per_prompt_round}") + if self.training.micro_batch_size is None: + self.training.micro_batch_size = self.training.global_batch_size + micro_batch_size = self.training.micro_batch_size + if micro_batch_size < 1 or micro_batch_size % world_size: + raise ValueError(f"micro_batch_size {micro_batch_size} must split into whole trajectories across the " + f"world size {world_size}") + if self.training.global_batch_size % micro_batch_size: + raise ValueError(f"global_batch_size {self.training.global_batch_size} does not split into " + f"micro-batches of {micro_batch_size}") diff --git a/training/grpo/configs/qwen3_4b_instruct_1gpu.json b/training/grpo/configs/qwen3_4b_instruct_1gpu.json new file mode 100644 index 000000000..a5d979819 --- /dev/null +++ b/training/grpo/configs/qwen3_4b_instruct_1gpu.json @@ -0,0 +1,96 @@ +{ + "model": { + "model_name_or_path": "Qwen/Qwen3-4B-Instruct-2507", + "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, + "return_log_probs": false + }, + "tool": { + "timeout": 120, + "memory_limit_gb": 4, + "max_concurrency": 32, + "max_output_chars": 4096, + "show_last_value": false + }, + "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_instruct_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": 100, + "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..b171aa112 --- /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, + "return_log_probs": false, + "chat_template_kwargs": { + "enable_thinking": false + } + }, + "tool": { + "timeout": 120, + "memory_limit_gb": 4, + "max_concurrency": 32, + "max_output_chars": 4096, + "show_last_value": false + }, + "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..1aceb8968 --- /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, + "return_log_probs": false, + "chat_template_kwargs": { + "enable_thinking": false + } + }, + "tool": { + "timeout": 120, + "memory_limit_gb": 4, + "max_concurrency": 32, + "max_output_chars": 4096, + "show_last_value": false + }, + "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..59db37c33 --- /dev/null +++ b/training/grpo/main.py @@ -0,0 +1,124 @@ +# 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, + show_last_value=cfg.tool.show_last_value) + 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..8551e07f5 --- /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_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_instruct_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..e0590f0f4 --- /dev/null +++ b/training/grpo/tools.py @@ -0,0 +1,257 @@ +# 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 ast +import asyncio +import os +import re +import sys +import tempfile +import traceback + +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 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, 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}" + 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 + + +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" + # 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": { + "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, + 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 + + 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) + 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)" + 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(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")) + 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() 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 -