feat: support streaming external rollouts - #2272
Open
Aphoh wants to merge 1 commit into
Open
Conversation
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.
Summary
This PR makes Slime's external SGLang rollout path work with streaming
/generateresponses.It adds:
No dependencies are added.
Why
The non-streaming rollout path receives generation state only after
/generatecompletes. If a rollout is interrupted early, it cannot retain the generated prefix without a separate polling or abort-response API.Streaming exposes that prefix as it is generated. Slime can apply each observed chunk directly to the
Sample, cancel the request when partial-rollout collection fires, and later resume from the retained tokens.Incremental streaming is also important for long contexts: Slime processes only the new suffix from each chunk and materializes response text once, avoiding repeated concatenation or rescanning of the accumulated response. Cumulative streams remain supported for compatibility.
Design
Stream normalization
SGLangStreamAccumulatornormalizes SGLang's two stream formats into updates that can be applied directly to a sample:The accumulator handles empty terminal chunks, incremental or terminal-full top-p snapshots, encoded metadata arrays, routed-expert metadata, Unicode and special-token text, and resumed samples. Text chunks are collected and joined once; when text is absent, accumulated token IDs are decoded once.
The stream must end with a
finish_reason. An unexpected EOF raises an error rather than allowing an incomplete sample into reward computation or training.Abort lifecycle
A generation function opts into request-level abort by setting an attribute:
Slime tracks active request-abortable tasks separately from active server-abort generations:
CancelledErroris converted into anABORTEDsample only when Slime deliberately canceled that task; unrelated cancellation still propagatesThis keeps the lifecycle generic: custom generation functions can opt in without adding function paths to a registry or introducing a global abort-mode CLI option.
Partial-rollout resumption
Slime already skips completed or truncated samples when a buffered group is reused. The request-abort path preserves that behavior: a group is buffered only when at least one aborted sample has a non-empty observed prefix, completed siblings remain unchanged, and only aborted siblings resume generation. Existing off-policy masking and reward behavior are also preserved.
Configuration
Select the streaming generator:
For cumulative SGLang output, no additional Slime flag is needed.
For incremental output, configure both sides:
The Slime flag must match the server's output mode. Length validation fails early if cumulative output is sent to the incremental consumer or vice versa.
Request cancellation can preserve only metadata received before disconnect. Metadata that SGLang emits only in a terminal chunk cannot be recovered from a canceled partial request; this PR does not attempt to reconstruct it.
Diff size
Measured against the current PR base:
The test contribution is 527 lines in the new focused streaming suite plus two fixture lines adapting the existing custom-generator contract tests. CI changes only register the new CPU test file in the generated and checked-in matrices.
Automated validation
The new streaming suite contains 26 parameterized cases covering:
The five existing custom-generator contract tests and fourteen argument-validation regression tests were also run, for 45 focused passing cases in total.
Additional checks:
python tests/...entrypoints passed for the modified test filesuvx prek run --all-filespassedLive external-SGLang E2E training
The PR at
4d39b5a606a410c80dd74475a2467b9a94a137fawas built into one image and that exact image was used for both the trainer and external SGLang engine in both runs.Shared setup:
/generate--stream-interval 20--enable-return-routed-experts--use-rollout-routing-replayThe only A/B difference was stream representation:
--incremental-streaming-output--sglang-incremental-streaming-outputResults:
All responses reached the configured length limit, exercising streamed length termination. Saved rollout artifacts were inspected directly. Every sample in both modes had:
rollout_routed_experts.dtype == torch.int32rollout_routed_experts.shape == (284, 48, 8), matching(len(tokens) - 1, num_moe_layers, experts_per_token)Both modes consumed those tensors with routing replay enabled and completed forward, backward, optimizer, and post-train weight synchronization.