Correct extra token, start preparing docker image for TGI/Jetstream Pt #93
Conversation
This file contains 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
Jetstream's `generate` function returns input token as result token. The next token is instead available in the decode_state, so this change uses this instead.
Before, we could have an error is the seed was bigger than a 64 bit number.
This allows testing TGI images with Jetstream Pytorch.
This is required when there are no more tokens generated after prefill.
A new slot is created at each prefill request, and its selector is passed as argument to a jitted function. The problem is that each new slot has a new signature, even if the contents are the same. The solution is to wrap that in a singleton slot object for the prefill, so the compiler will always see the same object and stop recompiling.
mfuntowicz
approved these changes
Sep 16, 2024
Comment on lines
41
to
43
| 16, | ||
| 32, | ||
| 64, |
There was a problem hiding this comment.
umh, it's actually aligned with Jetstream's buckets. But I just realized that's been extracted from the function, so I can import the constant now, I'll make the change.
| @@ -55,6 +55,8 @@ def __init__( | |||
| self.eos_token_ids = eos_token_ids | |||
| self.pad_token_id = pad_token_id | |||
| self.logits_warper = logits_warper | |||
| # Seed needs to fit a 64-bit integer, so we modulo it in case is bigger (that can happen!) | |||
| seed = seed % jnp.iinfo(jnp.int64).max | |||
There was a problem hiding this comment.
Can't we just store it as int64? Modulo op sounds very expensive
There was a problem hiding this comment.
The seed comes from the TGI router and sometimes it comes as a crazy big number (bigger than int64). numpy and torch seem to silently do the operation, but jax requires the seed to be an int64 and it raises an error if the number is higher. Let me know if you think about a "cheaper" alternative.
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.
What does this PR do?
There was an issue related to the fact that decode used to return the token passed in the decode state as result, making TGI see as if the first token was returned twice. This is resolved by using the data in the decode state (the API should be stable, according to internal discussions with the Jetstream Pytorch team).
Also, the docker image is now capable of serving TGI using Jetstream Pytorch on the supported models.
Before submitting