Repository navigation
Qwen3 autotp rollout - #1
Closed
OkoroAchraf wants to merge 18 commits into
Closed
OkoroAchraf wants to merge 18 commits into
OkoroAchraf wants to merge 18 commits into
Conversation
…e paths Graph capture applied only to greedy decoding, and continuous batching rejected temperature > 0. Only the decode forward is captured, so tokens can be picked from its output logits after each replay. - One token selector picks the first and every decoded token on graph decode and on continuous batching with and without graphs: argmax at temperature <= 0, otherwise Hugging Face's temperature, top-k and top-p warpers in generate()'s order, then multinomial sampling. - generate() routes sampled requests to the graph path. - Continuous batching accepts temperature > 0 and applies the model's repetition penalty before the warpers. - One helper builds the generate() sampling arguments for the main path and for continuous batching's legacy prefill. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
Under the same seed, the graph decode and continuous batching paths must generate exactly the tokens Hugging Face generate() samples. Each prompt gets its own slot, so every path draws all rows in one call per step, as generate() does. This replaces the log-prob and argmax z-score checks. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
- Sample-based RolloutRequest/RolloutBatch with Sample.Status. - HybridEngineRollout.generate_multi_turn: chat-template prompt rendering with tool specs, per-turn generation bounded by max_new_tokens (summed over turns), max_turns and max_context_len, Hermes tool-call parsing, tool execution, loss_mask and optional rollout log-probs. - Tool/ToolSet/ToolCall and parse_hermes_tool_calls (tool.py). - Prompt rendering, sample expansion and tool-response tokenization helpers (utils.py); chat_template_kwargs passed to every render. - Generation stops at EOS unless ZeRO-3 needs ranks in lockstep. - Multi-turn profile covers all turns, adds tool_ms and num_turns, and times the module whose generate() runs. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
…ple API HybridEngineRollout.generate() is now the single entry point and dispatches to _generate_multi_turn or _generate_single_turn. Both take and return Sample objects. - Single-turn rollout renders, expands and left-pads Samples like the multi-turn path, records rollout log-probs when return_log_probs is set, and marks each sample COMPLETED or TRUNCATED. - Shared helpers (_check_sampling, _pad_token_id, _generate_tokens) replace the duplicated validation and generate/graph dispatch code. The graph path samples with the token selector and records each sampled token's log-prob under softmax(logits / temperature). - Continuous batching, with and without CUDA graphs, accepts Sample requests, samples with the token selector, records log-probs for return_log_probs, and reads the repetition penalty from SamplingConfig.repetition_penalty instead of the checkpoint's generation_config. generate() receives the same penalty. - ResponseTokenBatch.from_rollout builds training targets from Samples: response tokens with loss_mask 0 are not targets. - Add a TODO on SamplingConfig.repetition_penalty about defaulting to the checkpoint value. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
max_new_tokens bounds a multi-turn trajectory's generated tokens summed over its turns. Recipes such as miles' rollout_max_response_len cap each turn instead, which max_new_tokens alone cannot express. - SamplingConfig.max_new_tokens_per_turn (default None) caps every turn on top of max_new_tokens; a turn may generate up to the smallest of the remaining budget, the per-turn cap and the room left in the context. - Single-turn rollout, which has one turn, folds the cap into max_new_tokens. - _check_sampling rejects a non-positive cap. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
_generate_graph prefilled into a Hugging Face StaticCache sized for the whole sequence, then allocated a DeepSpeedStaticCache of the same size and copied the prompt into it. Both stayed alive for the call, doubling its KV memory: for 16 Qwen3-4B rows with a 16k-token budget, 79 GB instead of 39 GB, which ran out of memory on a 95 GB GPU. - Prefill writes straight into the DeepSpeedStaticCache, from write position 0, and the HF cache and the copy are gone. - Prefill passes explicit position ids: the cache's get_seq_length also counts the slot being written, so positions derived from it would shift the prompt by one. - Prefill keeps only the last position's logits when the model accepts logits_to_keep (_last_logits_kwargs), instead of [batch, prompt_len, vocab]. - A CUDA test checks the graph path's recorded log-probs against a full forward pass over the sampled sequences, with and without left padding. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
…mple API Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
…arser
Qwen3.5 (and Qwen3-Coder) chat templates teach an XML-style call,
<tool_call><function=NAME><parameter=KEY>VALUE</parameter></function></tool_call>,
not the Hermes JSON block parse_hermes_tool_calls reads, so every call those
models made came back as "not valid JSON" and no tool ever ran.
- parse_qwen3_coder_tool_calls reads that format. The newline the template puts
around each value is removed, and values are converted to the tool's schema
type (integer, number, boolean, object/array, null), as SGLang's qwen3_coder
parser does; string and undeclared parameters stay text. Malformed blocks
yield a ToolCall with an error, like the Hermes parser.
- TOOL_CALL_PARSERS maps format names to parsers, and the new
tool_call_parser option ("hermes" by default) on RolloutConfig and
HybridEngineRolloutConfig picks the one multi-turn rollout uses.
Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
The GPU tests of continuous batching with CUDA graphs and of sampling against Hugging Face generate() build their requests from RolloutRecords. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
return_log_probs records log-probs under softmax(logits / temperature), the distribution a training forward pass reproduces. The rejection of the settings that sample from another distribution (top-k or top-p when sampling, or a repetition penalty) moves from the rollout's generation paths into SamplingConfig itself, together with the per-turn cap check. An invalid config now fails when it is built, on every path: continuous batching never ran the rollout-side check. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
…ph path The plain graph decode loop ignored the repetition penalty, which is now an explicit SamplingConfig field rather than a checkpoint generation_config setting. It applies Hugging Face's RepetitionPenaltyLogitsProcessor to the sequence so far, prompt included, before the token selector's warpers, as generate() and the continuous batching paths do. The GPU test against Hugging Face generate() adds a repetition-penalty case on all three decode paths. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
deepspeed/runtime/rollout/grpo.py adds the training-side primitives for GRPO on rollout records: - group_advantages: each reward minus the mean of its group_index group, optionally divided by the group's standard deviation. Single-record and equal-reward groups get advantage 0. - grpo_policy_loss: PPO-clip on per-token log-probs, with separate low and high clip ranges (clip-higher), an optional k3 KL penalty against reference log-probs, and token-mean or sample-mean aggregation over the response mask, which excludes tool output. It takes log-probs, not logits, so it works with plain log_softmax and vocabulary-parallel log-prob ops. - GRPOLossOutput: the loss with its unreduced sum and normalizer for aggregation across ranks, plus clip fraction and KL for logging. - old_log_probs_from_rollout: the records' rollout_log_probs laid out like ResponseTokenBatch's shifted targets, to use as the old policy. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
…GRPO inference-engine.rst describes RolloutRecord and return_log_probs, the multi-turn loop and the tool interface with an example, the GRPO primitives, the explicit repetition penalty on every decode path, and ResponseTokenBatch's loss_mask handling. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
The rollout places prompts on the accelerator's current device, where the policy lives in training. The test built its model on the CPU, so on a GPU node the prompts and the model were on different devices. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
_pad_after_eos treated every response token equal to pad_token_id as padding. Record-based rollouts drop padding from the response, so a sampled pad id disappeared: later tokens and their log-probs shifted, and the next turn continued from a different context than the one the model generated. Generation pads a row only after it ends, so only positions after a row's first EOS are padding now. Signed-off-by: Puyuan Yao <williamyao034@gmail.com>
AutoTPRollout generates with a separate AutoTP tensor-parallel copy of the training model that spans all ranks. sync_weights() refreshes that copy in memory from the ZeRO training engine: safe_get_full_fp32_param gathers each master weight, and _resolve_autotp_partition (the universal-checkpoint restore path) cuts this rank's shard from the layout metadata every AutoTP layer records. Ranks generate together: prompts are gathered, the union is generated with a shared seed and no EOS early stop, tokens are checked to match on every rank, and each rank keeps its own samples. Single-turn generate() only. build_rollout() accepts engine="autotp". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Achraf okoro <aokoro2003@gmail.com>
OkoroAchraf
force-pushed
the
qwen3-autotp-rollout
branch
from
November 2, 2026 16:05
441c89c to
00f06ab
Compare
anyadontfly
force-pushed
the
multi-turn-rollout
branch
from
November 5, 2026 06:20
24df6fe to
cd2694a
Compare
Owner
|
Opened in DeepSpeed #9023 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.