[Feature] Get Token IDs with Engine.generate() #2636
Merged
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
zhaochenyang20
requested changes
Dec 29, 2024
merrymercy
requested changes
Dec 29, 2024
Comment on lines
1261
to
1264
| if self.skip_tokenizer_init: | ||
| output_ids.append(req.output_ids) | ||
| origin_input_ids.append(req.origin_input_ids) | ||
| output_ids.append(req.output_ids) |
There was a problem hiding this comment.
Suggested change
| if self.skip_tokenizer_init: | |
| output_ids.append(req.output_ids) | |
| origin_input_ids.append(req.origin_input_ids) | |
| output_ids.append(req.output_ids) | |
| origin_input_ids.append(req.origin_input_ids) | |
| output_ids.append(req.output_ids) |
| @@ -657,6 +657,8 @@ async def handle_loop(self): | |||
| out_dict = { | |||
| "text": recv_obj.output_strs[i], | |||
| "meta_info": meta_info, | |||
| "input_ids": recv_obj.origin_input_ids[i], | |||
| "output_ids": recv_obj.output_ids[i], | |||
There was a problem hiding this comment.
by default, we should not return this because this introduce extra overhead.
There was a problem hiding this comment.
Yes. We can keep them disabled by default to avoid extra overhead, and just add a startup flag to enable returning these IDs when needed (like in OpenRLHF scenarios).
zhaochenyang20
requested changes
Dec 29, 2024
test/srt/test_engine_token_ids.py
Outdated
Comment on lines
6
to
12
| 1. Whether the input token IDs from the SGLang Engine match those produced | ||
| by the Hugging Face tokenizer for the same input string. | ||
| 2. Whether the output token IDs from the SGLang Engine match those produced | ||
| by the Hugging Face tokenizer (excluding the start token, which is | ||
| typically added by Hugging Face but not returned by SGLang). | ||
| 3. Whether the meta information, such as prompt token count and completion | ||
| token count, aligns with the actual lengths of input and output token IDs. |
There was a problem hiding this comment.
State about the start_token, where we add it, where we not.
test/srt/test_engine_token_ids.py
Outdated
Comment on lines
25
to
29
| def get_test_engine(): | ||
| """Returns a test engine with the 'Meta-Llama-3.1-8B-Instruct' model.""" | ||
| return sgl.Engine( | ||
| model_path="meta-llama/Meta-Llama-3.1-8B-Instruct", return_token_ids=True | ||
| ) |
test/srt/test_engine_token_ids.py
Outdated
|
|
||
|
|
||
| class TestEngineTokenIds(unittest.TestCase): | ||
| """Tests SGLang's token IDs against Hugging Face tokenizer.""" |
There was a problem hiding this comment.
Tests SGLang's return token IDs against Hugging Face tokenizer.
test/srt/test_engine_token_ids.py
Outdated
Comment on lines
35
to
51
| def setUp(self): | ||
| """Creates engine, tokenizer, and prompts.""" | ||
| self.llm = get_test_engine() | ||
| self.tokenizer = AutoTokenizer.from_pretrained( | ||
| "meta-llama/Meta-Llama-3.1-8B-Instruct" | ||
| ) | ||
| self.prompts = [ | ||
| "Hello, my name is", | ||
| "The president of the United States is", | ||
| "The capital of France is", | ||
| "The future of AI is", | ||
| ] | ||
| self.sampling_params = {"temperature": 0.8, "top_p": 0.95} | ||
|
|
||
| def tearDown(self): | ||
| """Shuts down the engine.""" | ||
| self.llm.shutdown() |
There was a problem hiding this comment.
Merge into one function.
test/srt/test_engine_token_ids.py
Outdated
| f"Input token IDs mismatch for: {prompt}", | ||
| ) | ||
|
|
||
| # Remove start token from HuggingFace as SGLang output doesn't include it |
There was a problem hiding this comment.
State for the input about start token.
There was a problem hiding this comment.
Explain why input has a token but output does not. Make this concise.
test/srt/test_engine_token_ids.py
Outdated
Comment on lines
1
to
12
| """Test token ID alignment between SGLang and Hugging Face. | ||
|
|
||
| This test suite ensures that the token IDs generated by the SGLang Engine | ||
| are consistent with those of the Hugging Face tokenizer for a given set of | ||
| prompts. Specifically, it checks: | ||
| 1. Whether the input token IDs from the SGLang Engine match those produced | ||
| by the Hugging Face tokenizer for the same input string. | ||
| 2. Whether the output token IDs from the SGLang Engine match those produced | ||
| by the Hugging Face tokenizer (excluding the start token, which is | ||
| typically added by Hugging Face but not returned by SGLang). | ||
| 3. Whether the meta information, such as prompt token count and completion | ||
| token count, aligns with the actual lengths of input and output token IDs. |
|
|
||
| self.assertEqual( | ||
| len(output["input_ids"]), | ||
| output["meta_info"]["prompt_tokens"], |
There was a problem hiding this comment.
Make a new PR to change the name, prompt_tokens_length.
There was a problem hiding this comment.
In the next PR, change input_ids to prompt_tokens. And current prompt_tokens into prompt_tokens_length.
test/srt/test_engine_token_ids.py
Outdated
| ) | ||
| self.assertEqual( | ||
| output["meta_info"]["completion_tokens"], | ||
| 128, |
There was a problem hiding this comment.
len(sgl_output_ids)
python/sglang/srt/server_args.py
Outdated
| "--return-token-ids", | ||
| action="store_true", | ||
| default=ServerArgs.return_token_ids, | ||
| help="Whether to return token IDs in the output. Experimental feature.", |
There was a problem hiding this comment.
Whether to return token IDs in the output.
zhaochenyang20
requested changes
Dec 29, 2024
python/sglang/srt/server_args.py
Outdated
Comment on lines
288
to
289
| help="Whether to return token IDs in the output.", | ||
| ) |
There was a problem hiding this comment.
This may introduce additional overhead.
shuaills
added a commit
to shuaills/sglang
that referenced
this pull request
Jan 11, 2025
)" This reverts commit 35bdb48.
timethink
pushed a commit
to timethink/sglang
that referenced
this pull request
Mar 9, 2025
Co-authored-by: Chayenne <zhaochen20@outlook.com>
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.
Motivation
related to #2634
Checklist