[Example][Bugfix]: align custom rollout paths with vLLM router API - #75
[Example][Bugfix]: align custom rollout paths with vLLM router API#75CalvinXKY wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request transitions the On-Policy Distillation (OPD) implementation from SGLang to vLLM, updating documentation, example scripts, command-line arguments, and reward processing logic to use the vLLM /v1/completions endpoint. Feedback on these changes highlights several robustness improvements in slime/rollout/on_policy_distillation.py: preventing potential alignment shifts when teacher logprobs are truncated by padding missing entries, adding defensive type-checking in _logprob_for_token to avoid AttributeErrors, and directly accessing sample.reward instead of calling sample.get_reward_value(args) to prevent KeyErrors when a global reward key is configured.
| choice = reward["choices"][0] | ||
| plp = choice.get("prompt_logprobs") or [] | ||
| per_pos = [ | ||
| _logprob_for_token(plp[i], sample.tokens[i]) for i in range(1, min(len(plp), len(sample.tokens))) | ||
| ] |
There was a problem hiding this comment.
If the teacher server truncates the input or returns fewer prompt logprobs than expected (len(plp) < len(sample.tokens)), the current list comprehension will truncate per_pos. When t_log_prob[-response_length:] is subsequently called, it will slice the last response_length elements of this truncated list, causing a silent alignment shift where logprobs are matched with the wrong tokens.
To prevent this, we should pad any missing logprobs with None (which _logprob_for_token safely handles as 0.0) so that per_pos always maintains a 1:1 index alignment with sample.tokens.
Additionally, we should defensively extract choice to avoid raising a KeyError or IndexError if the teacher response is empty or malformed.
| choice = reward["choices"][0] | |
| plp = choice.get("prompt_logprobs") or [] | |
| per_pos = [ | |
| _logprob_for_token(plp[i], sample.tokens[i]) for i in range(1, min(len(plp), len(sample.tokens))) | |
| ] | |
| choice = reward.get("choices", [{}])[0] if isinstance(reward, dict) and reward.get("choices") else {} | |
| plp = choice.get("prompt_logprobs") or [] | |
| per_pos = [ | |
| _logprob_for_token(plp[i] if i < len(plp) else None, sample.tokens[i]) | |
| for i in range(1, len(sample.tokens)) | |
| ] |
| def _logprob_for_token(pos_entry: dict | None, token_id: int) -> float: | ||
| """Pull the teacher's logprob for `token_id` out of one position's logprob dict.""" | ||
| if pos_entry is None: | ||
| return 0.0 |
There was a problem hiding this comment.
To prevent potential AttributeErrors if the teacher server returns an unexpected format where pos_entry is not a dictionary (e.g., a list or string), we should explicitly verify that pos_entry is an instance of dict before calling .get().
| def _logprob_for_token(pos_entry: dict | None, token_id: int) -> float: | |
| """Pull the teacher's logprob for `token_id` out of one position's logprob dict.""" | |
| if pos_entry is None: | |
| return 0.0 | |
| def _logprob_for_token(pos_entry: dict | None, token_id: int) -> float: | |
| """Pull the teacher's logprob for token_id out of one position's logprob dict.""" | |
| if not isinstance(pos_entry, dict): | |
| return 0.0 |
| The actual learning signal comes from the OPD KL penalty applied in compute_advantages_and_returns. | ||
| """ | ||
| """Process rewards from teacher model and extract teacher log probabilities.""" | ||
| raw_rewards = [sample.get_reward_value(args) for sample in samples] |
There was a problem hiding this comment.
Using sample.get_reward_value(args) here is risky. If args.reward_key is configured globally for other tasks, get_reward_value will attempt to look up that key in the teacher's JSON response dict. Since the teacher response only contains standard OpenAI-compatible fields (like choices), this will raise a KeyError.
Since we specifically need the raw teacher response dictionary here, we should access sample.reward directly.
| raw_rewards = [sample.get_reward_value(args) for sample in samples] | |
| raw_rewards = [sample.reward for sample in samples] |
CalvinXKY
changed the title
A100 smoke validation (vime_v22 @ A100-server)Re-ran
Log: geo3k_vlm_multi_turn fixPrevious failure was Fixed by aligning with
Retest log: Out of scope (unchanged)
|
|
I checked the corresponding References:
So the change to build initial train-side features from One thing I would still adjust before merging: please keep the multi-turn behavior closer to slime/miles instead of dropping it entirely. Concrete suggestions:
In short: I would keep this PR's initial-prompt fix, but mirror slime/miles by preserving the observation multimodal path and full multi-turn token/loss-mask trajectory. That gives us the bug fix without losing the behavior the original slime/miles rollout was maintaining. |
186a740 to
5cb8927
Compare
…me/miles Use processor(sample.prompt) for the initial turn, preserve observation multimodal encoding and full token/loss_mask trajectory on the vLLM render path.
5cb8927 to
5daa5c9
Compare
|
Related to #120. |
2 participants
Summary
Align remaining vime examples and related rollout hooks with the vLLM-first API surface after the SGLang removal on
main. Custom generate paths now useargs.router_ip/args.router_portinstead of the deprecatedvllm_router_*attributes, fixing smoke failures in examples such asmulti_agent.Changes
multi_agent/agent_system.pyandgeo3k_vlm_multi_turn/rollout.pyto call the router viarouter_ip/router_port--opd-type vllmwith example scripts and teacher scoring via/v1/completions+prompt_logprobsValidation
Smoke-tested on A100-server / vime_v22 using external scripts under
run_script/validate_examples/(not included in this PR):Removed on
mainand not covered here:retool,search-r1(SGLang-only).Test plan
multi_agentexample completes at least one rolloutgeo3k_vlm_multi_turnandgeo3k_vlmsmoke on available VLM checkpoints