fix: fix data type of num_gpus & generate_topk_per_token of vllmwrapper #117
Conversation
Summary of ChangesHello @ChenZiHong-Gavin, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request focuses on refining the integration with vLLM by addressing data type inconsistencies related to GPU resource allocation and vLLM engine parameters. Additionally, it significantly enhances the robustness and reliability of the token generation process, particularly when retrieving top-k probabilities. These changes contribute to more flexible resource management and more stable interactions with the LLM. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![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 correctly addresses data type issues for fractional GPU allocation and adds null handling for improved reliability in the vLLM wrapper. However, I've identified a significant breaking change in the generate_topk_per_token method's return structure. My review includes a suggestion to revert this structural change to maintain API compatibility while retaining the beneficial reliability improvements.
| candidate_tokens = [] | ||
| for _, logprob_obj in top_logprobs.items(): | ||
| tok_str = logprob_obj.decoded_token | ||
| tok_str = logprob_obj.decoded_token.strip() if logprob_obj.decoded_token else "" | ||
| prob = float(math.exp(logprob_obj.logprob)) | ||
| tokens.append(Token(tok_str, prob)) | ||
|
|
||
| tokens.sort(key=lambda x: -x.prob) | ||
| return tokens | ||
| candidate_tokens.append(Token(tok_str, prob)) | ||
|
|
||
| candidate_tokens.sort(key=lambda x: -x.prob) | ||
|
|
||
| if candidate_tokens: | ||
| main_token = Token( | ||
| text=candidate_tokens[0].text, | ||
| prob=candidate_tokens[0].prob, | ||
| top_candidates=candidate_tokens | ||
| ) | ||
| return [main_token] | ||
| else: | ||
| return [] |
There was a problem hiding this comment.
The structure of the returned value from this function has been changed significantly. Previously, it returned a flat list of candidate Token objects, sorted by probability. The new implementation returns a list containing a single Token object, which in turn contains all candidate tokens in its top_candidates attribute.
This is a breaking change for any consumer of this function that expects a list of all candidates. For example, code that uses len() on the result or iterates over it will now behave differently. This also introduces a potential inconsistency with other BaseLLMWrapper implementations, which might still return a flat list of tokens.
To maintain API compatibility and consistency, I suggest reverting to returning a flat list of candidate tokens, while keeping the beneficial null-safety check for decoded_token.
| candidate_tokens = [] | |
| for _, logprob_obj in top_logprobs.items(): | |
| tok_str = logprob_obj.decoded_token | |
| tok_str = logprob_obj.decoded_token.strip() if logprob_obj.decoded_token else "" | |
| prob = float(math.exp(logprob_obj.logprob)) | |
| tokens.append(Token(tok_str, prob)) | |
| tokens.sort(key=lambda x: -x.prob) | |
| return tokens | |
| candidate_tokens.append(Token(tok_str, prob)) | |
| candidate_tokens.sort(key=lambda x: -x.prob) | |
| if candidate_tokens: | |
| main_token = Token( | |
| text=candidate_tokens[0].text, | |
| prob=candidate_tokens[0].prob, | |
| top_candidates=candidate_tokens | |
| ) | |
| return [main_token] | |
| else: | |
| return [] | |
| candidate_tokens = [] | |
| for _, logprob_obj in top_logprobs.items(): | |
| tok_str = logprob_obj.decoded_token.strip() if logprob_obj.decoded_token else "" | |
| prob = float(math.exp(logprob_obj.logprob)) | |
| candidate_tokens.append(Token(tok_str, prob)) | |
| candidate_tokens.sort(key=lambda x: -x.prob) | |
| return candidate_tokens |
|
Ray makes it possible to use fractional GPUs, like: |
…er (InternScience#117) * fix: fix data type of num_gpus & generate_topk_per_token of vllmwrapper * fix: fix lint error
This PR fixes data type issues in the vLLM wrapper to support fractional GPU allocation and improves token generation reliability with better null handling.