[HybridEngineRollout] Add multi-turn tool-calling rollout, GRPO primitives and turn-level continuous batching by anyadontfly · Pull Request #9037 · deepspeedai/DeepSpeed

anyadontfly · 2026-11-08T23:55:47Z

Implements stage 1 of the agentic RL proposal (#8942): HybridEngineRollout can roll out multi-turn, tool-calling trajectories for RL training, with the primitives a GRPO trainer needs.

Rollout records. RolloutRequest and RolloutBatch carry RolloutRecords: one prompt and its rollout across all turns, with loss_mask (0 on tool output), recorded log-probs, reward and status.
Multi-turn tool calling. With multi_turn=True, each turn's response is parsed for tool calls, the tools run, and their output is appended as tool messages. Tools subclass Tool; parse errors, unknown tools and exceptions go back to the model as text. Parsers for the Hermes/Qwen3 and Qwen3-Coder formats.
GRPO primitives. group_advantages, grpo_policy_loss (PPO-clip with clip-higher, optional k3 KL penalty, token-mean or sample-mean aggregation) and old_log_probs_from_rollout.
Unified generate path. A single-turn call is a rollout of one turn without tool calls, so every call runs one driver. Under it are static batching, and continuous batching with or without graph capture.
Turn-level continuous batching. Built on master's continuous-batching loops. A multi-turn trajectory is a request that comes back: a row is freed when its turn ends, a turn that called a tool waits in a ToolQueue while its tools run, and its next turn returns to the front of the pending queue.
Sampling for RL. temperature, top_k, top_p and repetition_penalty apply in generate()'s order on every path, and return_log_probs records each token's log-prob under the sampling distribution.
Data-parallel and AutoTP rollout. Each data-parallel rank rolls out its own prompts. Under AutoTP, ranks of a tensor-parallel group share their random state and run tools once, so they sample and append the same tokens.

Changes

API

Change	Notes
`RolloutRequest` / `RolloutBatch` hold `RolloutRecord`s instead of tensors	Breaking. `RolloutRecord` holds a series of states to manage trajectories across turns
`continuous_batch_size` and `return_log_probs` move from `SamplingConfig` to `HybridEngineRolloutConfig`	Breaking. The row count sizes the KV cache and the captured graph, so it is an engine setting rather than a per-call one
The rollout caps each trajectory (prompt, responses and tool output) at `hybrid_engine.max_out_tokens`, which must be set explicitly in the DeepSpeed config	Breaking. The hybrid engine's default would silently cap every trajectory
`HybridEngineRolloutConfig` gains `multi_turn`, `max_turns`, `chat_template_kwargs`, `tool_call_parser`; `SamplingConfig` gains `max_new_tokens_per_turn`	New
`SamplingConfig.max_tool_output_chars` cuts each tool output longer than that many characters before it is added, keeping its `"start"`, `"end"` or `"both"` ends (`tool_output_keep`) with a `...(truncated)` marker, as verl's `max_tool_response_length` does; `RolloutRecord` counts the cuts in `num_tool_outputs_truncated` and the turns in `num_turns`	New, off by default. Without it, a tool output that would pass `max_out_tokens` is dropped and the trajectory ends `TRUNCATED`
`build_rollout(..., generation_engine=...)` replaces `student_engine`, which still works as a deprecated alias	Backward compatible
A turn ends on any EOS id from the tokenizer or the model's `generation_config`	Matches the stop set of HF `generate()`, for example Llama 3.1's `<
`SamplingConfig.repetition_penalty` is optional; `None` uses the model's `generation_config`	Same contract as `generate()`
`SamplingConfig` validates its own options when created	Invalid values now raise at construction instead of at `generate()`
`deepspeed.runtime.rollout.tool` and `deepspeed.runtime.rollout.grpo`	New modules

Generation paths

Static batching runs each turn for every active trajectory with HF generate, or on a captured decode graph.
Continuous batching keeps master's _generate_continuous and _generate_continuous_graph loops and ContinuousBatchScheduler, with three additions for multi-turn: each request carries its own token budget, a request returning from tools goes to the front of the pending queue, and a ToolQueue runs tool calls on a background event loop while the other rows decode; the rollout ends each turn itself and resubmits a turn once its tools return. Both loops read EOS once per step for all rows, and record each step for the profile.
Without graph capture, all rows share master's write column, with compaction, align_decode_fronts and enable_cache_trimming. A returning turn whose context is longer than the space behind the column moves the column forward, and live rows mask the skipped columns. A single-turn call keeps master's sizing; a multi-turn rollout gets the longest context a turn can be admitted with (max_out_tokens, plus any left padding of pre-tokenized samples) plus one turn's budget, which trimming at full always leaves room for. continuous_cache_capacity stays a memory cap for single-turn calls, and is raised to that bound, with a warning, for multi-turn rollouts.
With graph capture, each row keeps its own write position. An admission is prefilled in one padded batch, through the same prefill as the eager path, and only its real tokens are stored, from slot 0, so a returning turn's width does not matter. The captured graph is reused across calls as on master since Reuse continuous rollout graphs and release captured resources on shutdown #9009: an in-place weight update no longer captures it again, and shutdown() releases it.

Parallelism

Setup	Behavior
Data parallel, ZeRO 0–2, with or without optimizer offload	Supported on every path; ranks run independently
ZeRO-3	Single-turn static batching on HF `generate`. Ranks decode up to the same configured cap and stop together, through master's group-aware criterion from #8390, once every row on every rank has reached EOS or its own budget, which `max_out_tokens` can make differ between ranks. A rank with no room joins with a placeholder row whose budget is 0, so it does not hold the others back. Multi-turn rollout, graph capture and continuous batching raise, since ranks would run different forwards
AutoTP	Every generate path, with and without graph capture. Every KV cache is sized from each rank's own prefill, as in #9008, ranks share their group's random state once per `generate()` call, and tool calls run on the group's first rank, which broadcasts the output. A tool that draws from the accelerator RNG therefore advances only that rank's state, and the group samples different tokens from the next turn on; such tools should use their own `torch.Generator`. In multi-turn continuous batching the first rank also decides when a turn returns from its tools: at each step where a row frees, it sends the finished tools' output to its group over a host-side Gloo group, so every rank admits the same turns at the same step. With a tied output head that AutoTP shards by vocabulary (`VocabParallelLinear`, which master enables automatically for tied-embedding models under Transformers 5.x), each rank's logits cover only its shard, so ranks picked different, wrong tokens (R-1 in RFC #8790); the head now gathers the full logits for the duration of each rollout call and returns rank-local logits to training afterwards. Risk: the gather covers whatever positions the forward returns. Every rollout path, and HF `generate()`, requests last-position logits only (`logits_to_keep=1`), so it gathers `[B, V]` per step, about 10 MB at 32 rows in bf16. A model whose forward does not accept `logits_to_keep` would return logits for every prompt position, and the gather would then build the full `[B, T, V]` tensor at prefill (about 20 GB at 32 rows x 2048 tokens with a 152k vocabulary); current Hugging Face causal LMs accept it
Hybrid engine `inference_tp_size > 1`	Raises: the rollout calls the module directly and skips the hybrid engine's input gather

Validation

Hardware and software: NVIDIA A100-SXM4-40GB (1 or 2 per run, driver 595.71.05) and NVIDIA A40, torch 2.14.1+cu130, CUDA 13.0, Python 3.12.15, Transformers 5.19.0 and 4.57.6.

Check	Command	Result
Rollout, static cache and v1 rollout suites, including the 2- and 4-GPU tests, Transformers 5.19.0, 4× A100	`pytest tests/unit/runtime/rollout tests/unit/utils/test_static_cache.py tests/unit/v1/test_hybrid_engine_rollout_v1.py`	272 passed, 1 skipped (the opt-in pretrained-model test), including the tied-head test on all four paths
Existing AutoTP tests, 4× A100	`pytest tests/unit/v1/autotp/test_tp_plan_real_models.py tests/unit/v1/autotp/test_tp_plan_e2e.py`	17 passed
AutoTP on a tied-embedding model: Qwen3-4B-Instruct-2507, TP2, fp32, greedy, single- and multi-turn, 2× A100	static, static graph, continuous and continuous graph, against static and against one GPU	Ranks identical and every mode equal to static in all 8 cases (log-prob gap at most 7.4e-5), with a tool call in all 8 multi-turn trajectories; TP2 static equal to one GPU, 5/5 single-turn and 8/8 multi-turn rows. Before the fix, ranks differed in all 8 cases
Cost of the gather: Qwen3-4B, bf16, TP2	tok/s with the fix vs `tensor_parallel.vocab_parallel_lm_head: false`	Within ±1.1% on all four paths
CPU suites	`DS_ACCELERATOR=cpu pytest tests/unit/runtime/rollout tests/unit/utils/test_static_cache.py`	192 passed, 4 skipped; the multi-rank classes run in the GPU suites
`pre-commit`	`pre-commit run --files <changed files>`	Passes
Docs build	`sphinx-build -b html docs/code-docs/source`	Succeeds with the same 21 warnings as master
Same GPU suites, Transformers 4.57.6, 2× A40	as above	230 passed, 7 skipped, on the revision before the continuous-batching rewrite; to rerun

Behavior checks with a real model (Qwen3-4B SFT checkpoint, stored in bf16; the equivalence checks load it in float32, so that rounding cannot flip near-tied tokens, and bf16 results are noted where run):

All modes give the same records. For single-turn and multi-turn rollouts with a calculator tool, all modes (static and continuous batching, with and without graph capture, aligned decode fronts, cache trimming) give greedy records identical to static batching. Recorded log-probs differ by at most 3.4e-5.
Graph decode matches HF. Greedy graph decode matches HF generate on 205 of 205 tokens for a padded batch of 8 (bf16: 190 of 207; rows split on near-tied logits).
Recorded log-probs match a training forward. On tool-calling rollouts on every decode path, recorded log-probs differ from a full forward by at most 0.0002. In bf16, the largest gap is 0.49–0.68, in line with HF's own incremental decode, which differs from a full forward by 0.47–0.72 in earlier runs.
Tool output cap. On Qwen3-4B-Instruct-2507 with a calculator that also returns about 6,000 characters of log, every trajectory ends TRUNCATED after its first turn without the cap, since the output would pass max_out_tokens. With max_tool_output_chars=512, every trajectory completes its second turn on static, continuous and continuous graph batching, num_tool_outputs_truncated counts each cut output, and the chat template stays intact, with one assistant header per turn.
Response text. In every record of these runs, response equals a decode of all its response tokens at once, although it is built one turn and tool output at a time.

Distributed tests added (2 GPUs):

Data parallel. Each rank rolls out its own prompts, of different counts and lengths, on static and continuous batching, with and without graph capture. Its greedy records match HF generate on its own module, before and after an optimizer step, and the step changes them. Covered: ZeRO-2 on every path, ZeRO-2 with CPU optimizer offload and ZeRO-2 fp16 on graph capture, and ZeRO-3 on single-turn static batching.
AutoTP. Ranks seeded differently produce identical records, greedy and sampled, with and without graph capture, and across a sampled tool turn. Greedy output equals the unsharded model's, and a tool that returns its own rank gives both ranks the first rank's output.

Generation speed against master, single-turn, path by path. Setup: Qwen3-4B in bf16 on one A100-SXM4-40GB; 32 chat prompts of varied length per call; 256 new tokens; 32 rows for continuous batching, as many as static batching decodes per call. Each cell is 2 timed calls after a warm-up, one round per tree, in tokens per second, master → branch. "Fixed length" disables EOS, including the model's generation_config EOS ids, so both trees generate exactly the same tokens.

Path	Greedy, fixed length	Sampled, fixed length	Greedy, real EOS
Static (HF `generate`)	774 → 770 (−0.6%)	772 → 770 (−0.3%)	497 → 492 (−0.9%)
Static, graph capture	1060 → 1054 (−0.6%)	1055 → 1048 (−0.7%)	675 → 669 (−0.8%)
Continuous	674 → 667 (−1.0%)	675 → 666 (−1.4%)	416 → 425 (+2.2%)
Continuous, aligned fronts	607 → 605 (−0.4%)	608 → 603 (−0.9%)	378 → 385 (+2.0%)
Continuous, cache trimming	675 → 667 (−1.3%)	674 → 664 (−1.5%)	416 → 424 (+1.8%)
Continuous, graph capture	998 → 1056 (+5.8%)	1002 → 1056 (+5.4%)	633 → 677 (+6.9%)
Continuous, graph capture, weights updated in place between calls	1009 → 1071 (+6.1%)	1002 → 1056 (+5.4%)	633 → 677 (+6.9%)

The weight-update row is the RL case. Neither tree captures the graph again after an in-place weight update (0 captures in the timed calls), so it runs at the speed of plain graph capture. In the real-EOS column, continuous graph capture on the branch generates 5,174 tokens per call against 5,242 on master and on every other continuous path, likely because its prefill rounds differently in bf16 and a few greedy rows split at near-tied logits; in float32 all modes give identical records. Peak extra memory matches master on every path except static graph decode, which needs about half as much (1.45 vs 2.72 GiB), because the prefill writes straight into the decode cache. Differences within about 2% are inside what one round with 2 timed calls can resolve.

Compatibility matrix

Axis	Validated	Not validated
ZeRO stage	0, 2, 3 (single-turn static only)	1
Offload	CPU optimizer offload, ZeRO-2	NVMe, parameter offload
Precision	fp32, bf16, fp16	fp8 / quantized
Parallelism	Data parallel (2 GPUs), AutoTP (2 GPUs)	Pipeline, sequence parallel, MoE, more than 2 GPUs
Accelerator	CUDA (A100, A40), CPU	ROCm, XPU, HPU and the rest
Transformers	5.19.0, 4.57.6	Other versions

Companion PR

A DeepSpeedExamples PR adds a Qwen3 GRPO ReTool training example on this API. deepspeedai/DeepSpeedExamples#1023.

AI assistance

This PR was written with AI assistance (Claude Code). The submitter reviewed every changed line and ran the tests above.

Stage 1 of the agentic RL RFC (deepspeedai#8942). HybridEngineRollout generated one response per prompt as a padded tensor, which cannot carry interleaved model and tool spans or per-trajectory metadata. - Requests and outputs are lists of RolloutRecord: tokens across all turns, loss mask, rollout log-probs, reward, status and group index. ResponseTokenBatch.from_rollout honors the loss mask, so tool output is never a training target. - Multi-turn generation: generate, parse tool calls, run tools, append their output, within turn, token and context budgets. Model tokens are kept exactly as sampled, so recorded log-probs match the trained sequence. - Tool, ToolSet and ToolCall, with hermes and qwen3_coder parsers. Tool errors go back to the model as text. - Temperature, top-k, top-p and repetition penalty on the graph and continuous batching paths, in generate()'s order. The graph decode forward gets a pre-built 4-D mask, so capture works on Transformers 4.57.x with SDPA. Includes the changes in deepspeedai#8939 and deepspeedai#9004. - SamplingConfig gains return_log_probs, max_new_tokens_per_turn and an optional repetition_penalty; None uses the checkpoint's value. - Turn-level continuous batching: a row is freed when a turn ends and serves another trajectory while tools run. It rejects ZeRO-3. - group_advantages, grpo_policy_loss and old_log_probs_from_rollout. Tested on 1x A100, torch 2.14.1+cu130: rollout, static cache and v1 suites pass on Transformers 5.19.0 and 4.57.6. On a Qwen3-4B checkpoint with real tool calls in float32, recorded log-probs match a full forward within 1e-4 on continuous and static batching, with and without graph capture. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…gine settings continuous_batch_size sizes the KV cache and the captured decode graph, so it cannot change between generate() calls. It moves with return_log_probs from SamplingConfig to RolloutConfig and HybridEngineRolloutConfig, and a value below 1 is rejected when the config is built. Idle-row reset is removed. It only shortens KV reads for a kernel that reads up to the write position; SDPA reads the full cache width either way. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

Multi-turn continuous batching read the rows' state back to the host every continuous_poll_steps steps (16 by default), so a turn that ended in between kept its row until the next check, and its tool call and the next queued turn waited with it. The static graph path skipped its EOS check on the same schedule. Both now check after every step, and continuous_poll_steps is removed. The profile's per-poll trace becomes a per-step trace, continuous_step_trace, read in the same host sync as the row state. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…rojects The tool and parser docstrings now say what each piece does in a line or two. References to SGLang's parser names and to miles' behavior are replaced by the behavior itself. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…peedai#8602 The test belongs to deepspeedai#8602, which adds the same file. Keeping it here would make both PRs add one path. Graph decode against native generation stays covered by test_static_graph_matches_native_llama_generation. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…on_2 test The test checks a rejection on the static graph path, but on an accelerator without graph capture generate() falls back to HF generate and never reaches it, so the test failed on CPU. The graph_capture_accelerator fixture turns the opt-in on, as the other graph-path tests do. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

Single-turn generation is now a rollout of one turn without tool calls, so generate() has one driver and three schedulers instead of four drivers: static batching, continuous batching on a captured graph, and continuous batching without graph capture. - Continuous batching without graph capture keeps the shared write column with align_decode_fronts and enable_cache_trimming, and now also runs multi-turn rollouts. A returning turn longer than the column moves the column forward. The cache width is continuous_cache_capacity, or the longest admissible context plus one turn's budget. - Continuous batching on a captured graph gives each row its own write position, prefills each admission in one padded forward and stores only real tokens. Its decode mask is additive, since eager attention adds the mask to the attention weights. - Both continuous runs live in continuous_batching.py, and ContinuousBatchScheduler gains per-request budgets and front-of-queue submission for returning turns. - AutoTP is supported with static batching: the graph cache is sized per TP rank, ranks share their group's random state, and tools run on the group's first rank. AutoTP with continuous batching, hybrid engine inference tensor parallelism, and ZeRO-3 with multi-turn rollout or graph capture are rejected. - AutoTP's input consistency check treats DeepSpeedStaticCache as a cache, and the static graph path no longer fails without an EOS id. - Shared prefill applies to each trajectory's first turn and is rejected at construction where it can't run. - SamplingConfig checks its own options when created. - Graph reuse across calls, master's cache-size estimate and the single-turn continuous drivers are removed. Greedy float32 records from Qwen3-4B are identical to the previous tree and to static batching in all modes. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…g step On two ranks with ZeRO-2, each rank rolls out its own prompts on static and continuous batching, with and without graph capture. Its greedy records must match HF generate on its own module, before and after an optimizer step, and the step must change them. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

Master gained deepspeedai#8939 (sampling on the graph decode path) and deepspeedai#9004 (a pre-built 4-D mask for the static graph decode forward). This branch already carries reworked versions of both, so the rollout module, its tests and the docs keep the branch's versions. The docs keep master's graph decode section, merged with the branch's sampling notes. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…llel rollout test ZeRO-3 runs single-turn static batching on HF generate, the only path it supports, with the same number of prompts on every rank. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

chatgpt-codex-connector · 2026-11-08T23:56:02Z

Codex Review Summary: ✅ Completed

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-14T09:13:14.357430Z	`ad555e4`	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: bb5660f598

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-11-09T00:03:36Z

+        if sampling.max_new_tokens_per_turn is not None:
+            room = min(room, sampling.max_new_tokens_per_turn)
+        if self.max_context_len is not None:
+            room = min(room, self.max_context_len - sum(sample.attention_mask))


Keep ZeRO-3 decode counts equal under context caps

When ZeRO-3 data-parallel ranks receive prompts with different effective lengths and max_context_len is set, this rank-local subtraction can produce different turn budgets, or no active row at all on one rank. _generate() then either skips generation or calls module.generate() with a rank-local max(budgets), even though _generate_tokens() relies on every ZeRO-3 rank executing the same number of parameter-gathering forwards, so one rank can exit while another enters a collective and hang; reject this combination or preserve a synchronized decode count.

AGENTS.md reference: AGENTS.md:L19-L19

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-09T00:03:37Z

-            attention_mask=attention_mask,
-            response_start_idx=torch.full((total, ), response_start, dtype=torch.long, device=device),
-        )
+            prompt_ids, prompt_attn = left_pad_samples(turn_samples, pad_token_id)


Restore the EOS fallback for padless tokenizers

When a tokenizer defines eos_token_id but not pad_token_id (an explicitly documented and previously supported setup), pad_token_id remains None here and left_pad_samples() passes it to torch.full, causing every static rollout to fail before generation, even for a single unpadded prompt. Restore the EOS fallback before prompt rendering/padding and retain the clear error when neither token is defined.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-09T00:03:37Z

+            for row, i in enumerate(active):
+                tool_response = []
+                if row in turn_calls:
+                    tool_response = self._tool_response(turn_calls[row], tool_outputs[row])


Enforce max_context_len on appended tool output

When a tool call is generated near the context limit, _turn_room() reserves space only for model-generated tokens, but the entire unbounded tool response is appended here without checking the remaining max_context_len. The next turn merely marks the record truncated, leaving a returned training sequence longer than the configured/model context despite the contract in docs/code-docs/source/inference-engine.rst:241-243; a long tool result can therefore make the subsequent training forward exceed positional limits or memory, so settle or truncate the trajectory before appending over-cap output.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-09T00:03:37Z

+    return_log_probs: bool = False
+
+    def __post_init__(self) -> None:
+        if self.multi_turn and not self.max_turns:


Reject negative max_turns values

With multi_turn=True, max_turns=-1, this truthiness check accepts the value; static rollout feeds it to range(-1) and returns untouched PENDING records, while continuous rollout never reaches its terminal-turn comparison and instead runs until another cap. Validate that max_turns is positive here and in HybridEngineRolloutConfig.__post_init__, which duplicates the same check.

AGENTS.md reference: AGENTS.md:L19-L19

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-09T00:03:37Z

+            looks_integral = "." not in value and "e" not in value.lower()
+            return int(number) if looks_integral and number.is_integer() else number
+        if param_type in ("boolean", "bool"):
+            return value.lower() == "true"


Preserve malformed boolean tool arguments

When a Qwen3-Coder response supplies an invalid boolean value such as <parameter=flag>yes</parameter>, this expression silently converts it to False instead of preserving the unconvertible text as _convert_parameter() promises or reporting a parse error. The tool can then execute with a plausible opposite argument and corrupt the trajectory/reward; only map valid true/false text and retain or reject other values.

Useful? React with 👍 / 👎.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: bb5660f598

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-11-09T01:14:37Z

+                forward_profiler = _ForwardProfiler(accelerator)
+                handles.extend(forward_profiler.register(module))
+
+            turn_max_new_tokens = max(budgets)


Stop each static row at its context budget

When static batching on ZeRO 0–2 receives rows with unequal effective prompt lengths and max_context_len is set, this chooses the largest room for one HF generate() call, so a row whose own room is smaller still executes the extra decode steps until it happens to emit EOS; slicing the result later cannot prevent those forwards. If the cap is the model's positional limit, a non-EOS row is driven past that limit and can error or OOM despite the configured cap. The continuous runners stop per request, so the static path needs to group by budget or otherwise stop rows before the model forward.

AGENTS.md reference: AGENTS.md:L19-L19

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-09T01:14:37Z

+            tokens = prefill_output.sequences[:, -1:]
+            # output_logits holds the raw logits; return_log_probs rules out every other setting that changes them.
+            scale = sampling.temperature if sampling.temperature > 0.0 else 1.0
+            log_probs = _chosen_log_probs(prefill_output.logits[-1] / scale, tokens)


Score fallback prefill tokens after repetition penalty

When continuous eager batching takes this compatibility branch for a model whose _supports_cache_class is None and the repetition penalty is not 1, generate() samples the first token from penalized scores but output_logits is explicitly raw, so this line records a probability from the unpenalized distribution. The first token of every admitted turn then has an incorrect old log-probability, silently distorting GRPO ratios; apply the same penalty used by generation or consume processed scores, as the cache-class sibling path does.

AGENTS.md reference: AGENTS.md:L19-L19

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-09T01:14:37Z



-def build_rollout(rollout_cfg, student_engine=None, tokenizer=None, **kwargs):
+def build_rollout(rollout_cfg, generation_engine=None, tokenizer=None, tools=None, **kwargs):


Preserve the public student_engine keyword

For callers using the prior public form build_rollout(cfg, student_engine=engine, tokenizer=tok), the renamed argument is now swallowed by **kwargs, leaving generation_engine as None and raising the misleading "needs both" error; positional callers alone continue to work. Retain student_engine as a deprecated alias (or explicitly translate it from kwargs) so this factory signature remains backward compatible.

AGENTS.md reference: AGENTS.md:L40-L40

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-09T01:14:37Z

+
+    def _end_turn(self, sample, response, turn, generated, max_new_tokens):
+        """Settle a trajectory after a turn, and return the tool calls it should run next, if any."""
+        if not response or response[-1] != self.tokenizer.eos_token_id:


Treat EOS ID lists as membership sets

When a tokenizer exposes multiple EOS IDs, which _pad_after_eos explicitly supports, response[-1] is an integer and can never equal the list, so every response is marked TRUNCATED and multi-turn tool calls are never parsed or run. The continuous graph sibling also assigns the list to a scalar stop_token and compares it row-wise; normalize both paths to test membership against all EOS IDs.

AGENTS.md reference: AGENTS.md:L19-L19

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-09T01:14:37Z

+    if value.lower() == "null":
+        return None


Preserve null text for string tool parameters

When a Qwen3-Coder call supplies the literal text null for a schema-declared string or an undeclared parameter, this unconditional branch runs before the type check and passes Python None to the tool instead of the promised original text. Only convert null when the schema actually admits a null value; string and unknown parameters must retain the literal.

Useful? React with 👍 / 👎.

…ol arguments - ZeRO-3 with max_context_len: each rank's turn budget depended on its own prompts, so ranks ran different numbers of forwards and the parameter gathers hung. Ranks now agree on the longest budget, and a rank with no room runs the same forwards on a placeholder prompt. - A tokenizer without a pad token pads with its EOS id again, as on master. - A tool output that would pass max_context_len ends the trajectory as truncated instead of overrunning the cap, on every scheduler. - multi_turn requires a positive integer max_turns. - The Qwen3-Coder parser converts only "true" and "false" to booleans and keeps other text unchanged. The ZeRO-3 test setup sets stage3_param_persistence_threshold to 0, so the test model's parameters are partitioned and the lockstep is exercised. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

delock · 2026-11-10T01:39:12Z

Hi @anyadontfly thanks for your PR. I left a quick inline comments and may have other comments as I review. Thanks!

anyadontfly · 2026-11-10T09:11:26Z

Hi @anyadontfly thanks for your PR. I left a quick inline comments and may have other comments as I review. Thanks!

Thanks for viewing! I'll check the commented lines.

anyadontfly · 2026-11-10T20:21:19Z

Known limitation (AutoTP): ranks of a tensor-parallel group share their random state once per generate() call (moving from once per turn in the next push), and stay in step because tools run only on the group's first rank. A tool that draws from the accelerator RNG advances only that rank's state, so the group samples different tokens from the next turn on. Such tools should use their own torch.Generator.

…g in render_prompts - get_pad_token_id(tokenizer) replaces HybridEngineRollout._pad_token_id(), so the continuous-batching runs read it without going through the rollout. - get_repetition_penalty moves to utils next to it. - render_prompts only renders and tokenizes. Every generation path pads its own batch, so padding the whole request to one width was redundant. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

- index: nothing read it; group_index already ties samples to their prompt. - reward: rollout never sets it, and the trainer keeps rewards itself. - Status.ABORTED: no path produces it. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…an up names - The rollout reads ZeRO-3 partitioning, the AutoTP size and the inference TP size once in __init__, through zero_optimization_partition_weights(), autotp_size() and _config.hybrid_engine, as the rest of the repo does, instead of probing the engine with getattr fallbacks. The ZeRO-3 step all-reduce uses seq_data_parallel_group directly. - A turn's token allowance is its budget: _turn_room becomes _turn_budget. - Clearer names on the static path (static_batching, generated_length, prompt_attn_mask, the static graph's *_attn_mask buffers), and the profile reports max_prompt_length. _turn_limit is inlined, and some comments that restated the code are dropped. - Test engine stand-ins provide the accessors through _engine_for. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…ap length with hybrid_engine.max_out_tokens - RolloutConfig only selects the backend (engine). The hybrid engine's settings live on HybridEngineRolloutConfig, which subclasses it and is exported from deepspeed.runtime.rollout. HybridEngineRollout reads them directly and raises TypeError for another backend's config, which would otherwise leave every setting at its default. - max_context_len is removed. A trajectory's prompt, responses and tool output are capped by the DeepSpeed config's hybrid_engine.max_out_tokens, the hybrid engine's own total-length setting, which HybridEngineRollout requires to be set explicitly. Continuous batching sizes its rows from it, so multi-turn continuous batching no longer needs a separate capacity. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

The tokenizer has a single eos_token_id, but generation configs often list more, such as Llama 3.1's <|eom_id|> after a tool call. Passing the tokenizer's id to generate() replaced that list, and _end_turn and the graph paths compared tokens against a scalar, so a list EOS marked every turn TRUNCATED. get_eos_token_ids merges the tokenizer's and the generation config's ids. Every path stops on, trims at and ends turns on that set, and tokenize_tool_responses joins tool output after whichever of them the chat template ends the assistant turn with. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…ecated alias Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…any type Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…ool queue Replace the _ContinuousRun classes with master's _generate_continuous and _generate_continuous_graph, plus what multi-turn needs: per-request budgets, requeueing at the front of the pending queue, and a ToolQueue that runs tool calls on a background event loop while the other rows decode. Read EOS once per step, write the graph inputs with indexed writes, keep the step trace on both paths, record the fallback prefill's log-probs after the repetition penalty, and size the multi-turn shared column for the longest context plus one turn. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…te tests Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…rofile fields Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

Combine master's ZeRO-3 synchronized early stopping (deepspeedai#8390) and KV caches sized from local prefill shapes under AutoTP (deepspeedai#9008) with the branch's rollout: the static path uses the group-aware stopping criterion, with a rank decoding only a placeholder prompt counted as idle; every decode path sizes its cache lazily from the prefill; single-turn continuous batching supports AutoTP again, and only multi-turn continuous batching rejects it. Master's new tests are ported to the records API. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…ens, use the accelerator API in tests Remove _continuous_prefill's target_slots, which no caller passes since the graph path prefills with _prefill_rows. Document hybrid_engine.max_out_tokens, which HybridEngineRollout requires as the cap on each trajectory, and route the on-device rollout tests' device checks through get_accelerator(). Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…edai#9058) ## what `TestZero3SynchronizedEarlyStop` (added in deepspeedai#8390) builds its tensors on `torch.device(f"cuda:{rank}")`. on the cpu multi-rank nightly (`cpu-torch-multi-latest`, `LOCAL_SIZE=4`) the test is admitted by the world-size gate, and the first `torch.tensor(..., device=device)` raises `AssertionError: Torch not compiled with CUDA enabled`. that is the 2026-11-11 recurrence in deepspeedai#9033 (run 42957495892). this takes the device from `get_accelerator().current_device_name()`, which is `cpu` on the cpu accelerator and `cuda:<local rank>` on cuda, so the gpu path is unchanged. ## why this is not a duplicate no open pr touches this test; deepspeedai#9009 and deepspeedai#9037 edit the same file but not this class. ## tests on cpu (8 vcpu x86_64, torch 2.15.0+cpu, python 3.12), with `LOCAL_SIZE=4`: - before: `pytest unit/runtime/rollout/test_hybrid_engine_rollout.py::TestZero3SynchronizedEarlyStop` fails with `Torch not compiled with CUDA enabled`, same traceback as the nightly - after: `1 passed` - `pytest unit/runtime/rollout/`: 94 passed, 2 skipped; `TestZero3SubgroupGeneration` failed locally only because `accelerate` is not installed in my env - flake8 and yapf clean on the changed file not run on gpu; the cuda device string is the same as before apart from using the local rank instead of the global rank. fixes deepspeedai#9033 (the hybrid engine rollout part; the earlier `test_tp_compile.py` failure looked flaky per the thread) Signed-off-by: Dundy Pasupuleti <dundysm@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

Tools finish at wall-clock times, so tensor-parallel ranks admitted returning turns at different steps and ran different forwards. The first rank of each tensor-parallel group now runs the tools and, at each step where a row frees, sends the finished tools' output to its group over a host-side Gloo group, so every rank appends the same output and admits the same turns. deepspeed.comm.new_group takes an optional backend for that group. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

… tool group explicitly ZeRO-3 ranks now decode up to the same configured cap, and master's _GlobalEosStoppingCriteria counts a row as finished once it reaches EOS or its own budget. This replaces the per-turn all-reduce of the largest budget over seq_data_parallel_group, so the criterion's optimizer.dp_process_group is the only group, and a row past its budget no longer keeps every rank decoding. A rank with nothing left joins with a placeholder row whose budget is 0, in place of the idle flag. ToolQueue takes the AutoTP Gloo group and its first rank as arguments instead of reading the rollout's fields. Redundant local imports of deepspeed.comm are removed, and _chosen_log_probs is inlined into _TokenSelector. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

Static batching now retires each row of a turn into ToolQueue and drains it before the next turn, so both batching modes share one implementation of the turn bookkeeping, the tool-output size check and the AutoTP tool sharing. The new ToolQueue.drain waits for every running tool and makes one pop_ready call, one collective per turn as before. _run_tools and its NCCL broadcast_object_list are removed, and the Gloo group is created for any multi-turn rollout under AutoTP. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

… group When the AutoTP size equals the world size, the tensor-parallel group is the default group, for which get_global_rank is the identity and never raises. get_all_ranks_from_group probes ranks until an error, so building the Gloo tool group looped forever in HybridEngineRollout.__init__ and exhausted host memory. Count the group's ranks from its size instead. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…toTP rollout Under Transformers 5.x, AutoTP turns a tied-embedding model's output head into VocabParallelLinear(gather_output=False), so each rank's logits cover only its vocabulary shard. Every rollout path selected tokens from that shard: ranks picked different, wrong tokens and normalized log-probs over one shard. generate() now sets gather_output=True on every vocabulary-parallel head for the duration of the call and restores it afterwards, so all decode paths see the full vocabulary and training keeps rank-local logits for the vocabulary-parallel loss. The fallback prefill sizes the repetition-penalty mask from the head's vocab_size instead of the local weight shard. Every path requests last-position logits only, so the gather is [B, V] per step. A model whose forward does not accept logits_to_keep would gather [B, T, V] at prefill. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

ToolQueue.requeue pops the turns whose tools finished, computes each next turn's budget (capped by the cache row on the graph path, from the sample's unpadded context), truncates trajectories with no budget left, and submits the rest to the front of the scheduler's pending queue. The continuous loops now set a request's context length and open its response buffers when it is admitted, and pop the buffers when it retires, so _requeue_tool_turns and its loop-state arguments are gone. Also drops an unused import. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…t helpers The class decides where each trajectory goes after a turn (end, run its tools, or back to the pending queue), not only tool execution. Move it next to the rollout that drives it, so continuous_batching.py again holds only scheduling primitives. Also move the sampled log-prob recorder and the host tensor-parallel group to rollout utils, run a turn's tool calls through ToolSet.run_all, and rename _continuous_turn to _pop_response. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…ous prefill Fold the cache-class check for continuous batching into _check_model (was _check_attention), which generate() runs before dispatch. Graph continuous batching now prefills through the same helper as the eager path, which uses HF generate() for models without _supports_cache_class, as master's graph path did; it then stores each row's tokens from column 0 of its cache row. Rename the two prefills after their cache layouts: _continuous_prefill_shared_column (eager) and _continuous_prefill_per_row (graph). Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…l execution Ending a turn is synchronous: every retired turn is recorded, and the trajectory either ends or calls tools, before anything else happens to it. Tool execution is asynchronous: calls are submitted and collected later at a boundary, after the whole batch's turn for static batching and at steps where a row frees for continuous batching. PostTurnRouter mixed the two. ToolQueue (now in tool.py) only runs tool calls: submit, pop_ready, wait, drain, and the AutoTP broadcast of the first rank's output. The rollout ends each turn (_finish_turn), which records the response and decides whether the trajectory ends or which tools it calls; appends tool output when it returns (_append_tool_output); and in continuous batching queues each turn with its budget (_submit_turn), for first and returning turns alike. All four decode paths share these. RolloutRecord gains num_turns, and the tokens a trajectory has generated are read from its loss mask, so the rollout keeps no per-trajectory counters of its own. RolloutRequest rejects records that already hold a rollout. tokenize_tool_responses takes the calls and their outputs and builds the tool messages itself. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

A tool output that would take a trajectory past max_out_tokens is dropped and the trajectory ends truncated, so one long output, such as a stack trace, ends a rollout that could have gone on. SamplingConfig.max_tool_output_chars cuts each tool output longer than that many characters before it is templated, keeping its "start", "end" or "both" ends (tool_output_keep) with a "...(truncated)" marker at the cut, as verl's max_tool_response_length does. The cap is per output and off by default; the check on the whole conversation stays as the backstop. Each record counts its cut outputs in num_tool_outputs_truncated. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

…ey serve Move only the methods this PR adds: the checks, the turn lifecycle and the static loop after generate(), each continuous-batching helper beside its loop or the profile code, and the EOS helpers where master's _is_eos was. Master's methods return to master's order. No code changes. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

Resolve conflicts with deepspeedai#9009 (reuse continuous rollout graphs). Take its graph signature, which keys on each parameter's and buffer's address and layout instead of its version, so an in-place optimizer step no longer recaptures, and its shutdown(). Keep this branch's graph loop, which already dropped the prompt width from the signature and keeps its own single-turn capacity check. Port the new GPU test to rollout records, and describe reuse and shutdown() in the graph-capture part of the docs. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

Remove the static graph tests that test_decode_paths_match_hf_generate and test_graph_capture_decodes_left_padded_rows_as_if_alone subsume, the AutoTP tool-output test that the continuous AutoTP test covers on both schedules, and two repetition-penalty tests that other tests cover. Trim the parametrizations that add no path: the multi-turn as-if-alone cases at temperature 1.0, the second model in the reused-row repetition-penalty test, the static cases of the GPU over-budget tool-output test (covered on CPU), and the data-parallel matrix, which keeps every decode path on ZeRO-2 and one or two for offload, fp16 and ZeRO-3. Merge the config and sampling validation tests into one test each. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ad555e4710

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-11-14T09:13:10Z

+        if not samples:
+            return RolloutBatch(samples=[])


Keep empty ZeRO-3 ranks in the rollout collectives

When ZeRO-3 data-parallel ranks have uneven local work and one receives RolloutRequest(samples=[]), this rank-local return skips every parameter-gathering forward and GlobalWorkRemaining, while nonempty ranks enter module.generate() and block on ZeRO collectives. The zero-row rank needs to participate with placeholder synchronization just like over-cap rows, or empty local requests must be rejected collectively.

AGENTS.md reference: AGENTS.md:L19-L19

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-14T09:13:10Z

+        token_kl = torch.exp(log_ratio) - log_ratio - 1.0
+        token_loss = token_loss + kl_coef * token_kl


Mask excluded KL positions before exponentiation

When ref_log_probs is supplied and a masked prompt/tool/pad position has a large positive ref_log_probs - log_probs (which is easy for a very unlikely fixed token), this exponentiation becomes inf; the later mask then evaluates inf * 0 as NaN, poisoning both the loss and KL metric, and even kl_coef=0 computes 0 * inf. Zero the masked log-ratios before exp or compute KL only on valid positions.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-14T09:13:10Z

-            ],
-                               dim=0)
-        return select_next_tokens(logits)
+        tokens, log_probs = selector.first(prefill_output.logits[:, -1, :], prompt_ids, prompt_attention, rows)


Align padded repetition history across decode paths

When a batch is left-padded and repetition_penalty != 1, the cache-class continuous prefill removes masked pad IDs from the selector history here, while static graph prefill calls selector.first without a mask and Hugging Face's repetition processor on the static generate() path scans all input IDs. If the pad ID is EOS (a common fallback) or otherwise has a competitive logit, continuous rollout can sample or stop differently from static rollout for the same record; all paths need one padding policy.

AGENTS.md reference: AGENTS.md:L19-L19

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-14T09:13:10Z

-    return cdb.new_group(ranks)
+    if backend is None:
+        return cdb.new_group(ranks)
+    return cdb.new_group(ranks, backend=backend)


Forward backend through the CCL new_group override

When cdb is CCLBackend and multi-turn AutoTP initializes host_tensor_parallel_group(), this dispatch passes backend="gloo", but deepspeed/comm/ccl.py still overrides new_group(self, ranks) without accepting that keyword. The rollout therefore raises TypeError before generation on the CCL path; update that override to accept and forward the backend while preserving its subgroup bookkeeping.

AGENTS.md reference: AGENTS.md:L40-L40

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-11-14T09:13:10Z

+                    profile["tool_ms"] += tool_ms
+                    profile["post_processing_ms"] += (turn_end - generation_end) * 1000.0 - tool_ms
+                    profile["num_generated_tokens"] += int(output_ids.shape[0] * (output_ids.shape[1] - prompt_len))
+                    profile["num_turns"] += 1


Count trajectory turns in static profiles

With static batching of multiple records, each batched generation round increments num_turns only once, while both the records themselves and the continuous profiler count one turn per trajectory. Thus a one-turn batch of 32 reports num_turns=1 in static mode and 32 in continuous mode, making cross-mode tool and throughput profiles misleading; increment by the number of turn_samples processed in the round.

AGENTS.md reference: AGENTS.md:L19-L19

Useful? React with 👍 / 👎.

anyadontfly added 10 commits November 5, 2026 00:20

test(rollout): cover ZeRO-2 offload, fp16 and ZeRO-3 in the data-para…

bb5660f

…llel rollout test ZeRO-3 runs single-turn static batching on HF generate, the only path it supports, with the same number of prompts on every rank. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

anyadontfly requested review from loadams, tjruwase and tohtana as code owners November 8, 2026 23:55

anyadontfly marked this pull request as draft November 8, 2026 23:56

chatgpt-codex-connector Bot reviewed Nov 9, 2026

View reviewed changes

anyadontfly marked this pull request as ready for review November 9, 2026 01:06

chatgpt-codex-connector Bot reviewed Nov 9, 2026

View reviewed changes

anyadontfly mentioned this pull request Nov 9, 2026

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

Open

delock self-requested a review November 9, 2026 12:51

delock reviewed Nov 10, 2026

View reviewed changes

Comment thread deepspeed/runtime/rollout/hybrid_engine_rollout.py Resolved

tohtana added the needs-gpu-test GPU testing is recommended to validate the changes in this PR. label Nov 10, 2026

anyadontfly added 4 commits November 10, 2026 15:43

anyadontfly added 3 commits November 10, 2026 21:35

fix(rollout): accept build_rollout's student_engine keyword as a depr…

2d57b6f

…ecated alias Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

docs(rollout): note that Qwen3-Coder null parameters become None for …

cde4cb7

…any type Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

dundysm mentioned this pull request Nov 11, 2026

Use the accelerator device in TestZero3SynchronizedEarlyStop #9058

Merged

anyadontfly added 6 commits November 11, 2026 17:46

test(rollout): restore master's graph-state, capture and cache-estima…

178f41b

…te tests Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

test(rollout): move the on-device rollout tests to the v1 suite

b3f2a6b

Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

docs(rollout): describe continuous-batching sizing, graph reuse and p…

8858de9

…rofile fields Signed-off-by: Puyuan Yao <williamyao034@gmail.com>

anyadontfly added 8 commits November 13, 2026 12:23

jinyouzhi mentioned this pull request Nov 14, 2026

[RFC] Cleanup: move process-wide parallel-group globals to per-engine groups — follow-up to #8241 / #8851 #9081

Open

anyadontfly added 5 commits November 14, 2026 00:04

tohtana removed the needs-gpu-test GPU testing is recommended to validate the changes in this PR. label Nov 14, 2026

anyadontfly marked this pull request as ready for review November 14, 2026 09:02

anyadontfly requested a review from delock November 14, 2026 09:02

anyadontfly mentioned this pull request Nov 14, 2026

[RFC] Vocabulary-parallel per-token ops for RL and on-policy distillation #8790

Open

chatgpt-codex-connector Bot reviewed Nov 14, 2026

View reviewed changes



		def build_rollout(rollout_cfg, student_engine=None, tokenizer=None, **kwargs):
		def build_rollout(rollout_cfg, generation_engine=None, tokenizer=None, tools=None, **kwargs):

		token_kl = torch.exp(log_ratio) - log_ratio - 1.0
		token_loss = token_loss + kl_coef * token_kl

Conversation

anyadontfly commented Nov 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

API

Generation paths

Parallelism

Validation

Compatibility matrix

Companion PR

AI assistance

Uh oh!

chatgpt-codex-connector Bot commented Nov 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Nov 9, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 9, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 9, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 9, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 9, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Nov 9, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 9, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 9, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 9, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 9, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

delock commented Nov 10, 2026

Uh oh!

anyadontfly commented Nov 10, 2026

Uh oh!

anyadontfly commented Nov 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Nov 14, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 14, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 14, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 14, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Nov 14, 2026

Choose a reason for hiding this comment

Uh oh!

anyadontfly commented Nov 8, 2026 •

edited

Loading

chatgpt-codex-connector Bot commented Nov 8, 2026 •

edited

Loading

anyadontfly commented Nov 10, 2026 •

edited

Loading