Inference: Bring chat completions API inline with vllm/official openAI spec - #5276
Conversation
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # Small RL/debug scalars (a few bytes each); harmless to keep for | ||
| # NeMo-RL compatibility. | ||
| message["generation_log_probs"] = result.get("generated_log_probs", []) |
There was a problem hiding this comment.
[SUGGESTION Naming] The comment "Small RL/debug scalars (a few bytes each)" accurately describes policy_epoch, kv_cache_epoch, and num_evictions, but it is placed on a block whose first member is generation_log_probs — a per-generated-token float list, not a scalar. For a long completion this is one of the larger fields in the payload, which contradicts both the "a few bytes each" wording and the PR's stated goal of cutting payload size.
Why it matters: a wrong comment is worse than none — a future reader trimming the payload may skip generation_log_probs believing it is negligible.
Suggested fix: move the comment below generation_log_probs so it only covers the three scalar fields, and note that generation_log_probs is always included because RL clients (e.g. megatron/rl/inference/megatron.py) read choice.message.generation_log_probs. If it should also be gated behind an opt-in flag like the token ids, consider doing so for consistency.
| # Choice-level prompt/generation_token_ids, generation_log_probs and | ||
| # raw_text were duplicates of message-level data (or reconstructable); | ||
| # dropped to match vLLM's response shape and cut payload size. | ||
| choice_data = { | ||
| "index": request_idx, | ||
| "message": message, | ||
| "prompt_token_ids": result["prompt_tokens"], | ||
| "generation_token_ids": result["generated_tokens"], | ||
| "generation_log_probs": result.get("generated_log_probs", []), | ||
| "raw_text": result["prompt"] + result["generated_text"], | ||
| # 'logprobs' in chat API is an object containing 'content' | ||
| # "logprobs": {"content": logprobs_content} if logprobs_content else None, | ||
| "logprobs": {"content": logprobs_content} if return_log_probs else None, | ||
| "finish_reason": finish_reason, | ||
| } |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This removes prompt_token_ids, generation_token_ids, generation_log_probs, and raw_text from the choice object. The in-repo consumer (megatron/rl/inference/megatron.py) is updated in this PR to read them from choice.message.* instead, so the internal path is fine. But this is a response-shape change on a public OpenAI-compatible endpoint: any external client that read these at the choice level (e.g. a NeMo-RL client pinned to the previous shape) will silently get None/AttributeError after this change, with no deprecation window.
Why it matters: silent breakage of an external client contract is hard to diagnose from the server side.
Suggestion: confirm NeMo-RL (and any other downstream) reads the message-level fields, and if this endpoint is externally consumed, call out the shape change in the PR description / release notes. No code change required if downstreams are already aligned.
There was a problem hiding this comment.
Is this deviating from the official openAI spec?
There was a problem hiding this comment.
No, these are not part of the spec.
|
Strict Review Summary Findings: CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 1 Small, well-scoped inference-server change: (1) reasoning-token retention delegated to the chat template (removing What I verified (correct)
Findings
Risk: Low. Logic changes are minimal and mirror existing tested patterns. The only real risk is the external response-shape compatibility item above. |
| # Small RL/debug scalars (a few bytes each); harmless to keep for | ||
| # NeMo-RL compatibility. | ||
| message["generation_log_probs"] = result.get("generated_log_probs", []) |
There was a problem hiding this comment.
[SUGGESTION Naming] The comment "Small RL/debug scalars (a few bytes each)" accurately describes policy_epoch, kv_cache_epoch, and num_evictions, but it is applied to a block whose first member is generation_log_probs — a per-generated-token float list, not a scalar. For a long completion this is one of the larger fields in the payload, which directly contradicts the "a few bytes each" characterization and the PR's stated goal of cutting payload size.
Why it matters: a wrong comment is worse than none — a future reader trimming the payload may skip generation_log_probs believing it is negligible.
Suggested fix: move the comment below generation_log_probs so it only covers the three scalar fields, and note separately that generation_log_probs is always included (it is required by RL clients such as megatron/rl/inference/megatron.py, which reads choice.message.generation_log_probs). If it should also be gated behind an opt-in flag like the token ids, consider doing so for consistency.
| # Choice-level prompt/generation_token_ids, generation_log_probs and | ||
| # raw_text were duplicates of message-level data (or reconstructable); | ||
| # dropped to match vLLM's response shape and cut payload size. | ||
| choice_data = { | ||
| "index": request_idx, | ||
| "message": message, | ||
| "prompt_token_ids": result["prompt_tokens"], | ||
| "generation_token_ids": result["generated_tokens"], | ||
| "generation_log_probs": result.get("generated_log_probs", []), | ||
| "raw_text": result["prompt"] + result["generated_text"], | ||
| # 'logprobs' in chat API is an object containing 'content' | ||
| # "logprobs": {"content": logprobs_content} if logprobs_content else None, | ||
| "logprobs": {"content": logprobs_content} if return_log_probs else None, | ||
| "finish_reason": finish_reason, | ||
| } |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This removes prompt_token_ids, generation_token_ids, generation_log_probs, and raw_text from the choice object. The in-repo consumer (megatron/rl/inference/megatron.py) is updated in this PR to read them from choice.message.* instead, so the internal path is fine. However, this is a response-shape change on a public OpenAI-compatible endpoint: any external client that read these at the choice level (e.g. a NeMo-RL client pinned to the previous shape) will silently get None/AttributeError after this change, with no deprecation window.
Why it matters: silent breakage of an external client contract is hard to diagnose from the server side.
Suggestion: confirm the NeMo-RL (and any other downstream) client reads the message-level fields, and if this endpoint is externally consumed, call out the shape change in the PR description / release notes so consumers can migrate. No code change required if downstreams are already aligned.
| # Choice-level prompt/generation_token_ids, generation_log_probs and | ||
| # raw_text were duplicates of message-level data (or reconstructable); | ||
| # dropped to match vLLM's response shape and cut payload size. | ||
| choice_data = { | ||
| "index": request_idx, | ||
| "message": message, | ||
| "prompt_token_ids": result["prompt_tokens"], | ||
| "generation_token_ids": result["generated_tokens"], | ||
| "generation_log_probs": result.get("generated_log_probs", []), | ||
| "raw_text": result["prompt"] + result["generated_text"], | ||
| # 'logprobs' in chat API is an object containing 'content' | ||
| # "logprobs": {"content": logprobs_content} if logprobs_content else None, | ||
| "logprobs": {"content": logprobs_content} if return_log_probs else None, | ||
| "finish_reason": finish_reason, | ||
| } |
There was a problem hiding this comment.
No, these are not part of the spec.
| choice_data = { | ||
| "index": request_idx, | ||
| "message": message, | ||
| "prompt_token_ids": result["prompt_tokens"], |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29343662639 |
sidsingh-nvidia
changed the title
… conversations (NVIDIA#5276) Signed-off-by: mchochowski <mchochowski@nvidia.com>
… conversations (NVIDIA#5276) (cherry picked from commit a79f49d)
… conversations (NVIDIA#5276) Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
7 participants
What does this PR do ?
ignore_eos.Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.