Fix generation using Jetstream Pytorch #94
Merged
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
tengomucho
force-pushed
the
investigate-generation
branch
from
28da563
to
6ec35ed
Compare
Text generation was not correct because the weights in the model were not correctly loaded. This is not something that it was easy to spot just looking at few tokens generated, and it was something that it was actually fixed already in the Jetstream/Pytorch code, but the fix hadn't been ported to optimum-tpu. This fix implement the necessary weights changes, aligning to Jetstream Pytorch, and tests expected output has been modified accordingly.
tengomucho
force-pushed
the
investigate-generation
branch
2 times, most recently
from
edb7ee3
to
a8cad0a
Compare
tengomucho
force-pushed
the
investigate-generation
branch
from
fa50f3b
to
23ab23b
Compare
The main workflow was failing due to an OS error. I suspect that being related to a problem of space. Separating the workflow will make it easier to analyse this issue.
I was previously referencing a given git revision and install from github, but since the Jetstream Pytorch package install its dependencies from its git submodels, these are installed in temporary directories, that can disappear afterwards. This happened on CI, making the installation fail. To work around that, a dedicated install script has been added, and it is now used to install that.
Since this is error-prone, a better solution is just to use this. This hadn't been done before mainly because in the model config we do not have some of the params anymore (ffn_dim_multiplier and multiple_of). We do have intermediate_size though, and that is enough to reconstruct parameters that end up producing the same calculation. This refactor should allow for future code to follow Jetstream/Pytorch changes in an easier way.
tengomucho
force-pushed
the
investigate-generation
branch
from
23ab23b
to
8519c81
Compare
mfuntowicz
approved these changes
Sep 23, 2024
| @@ -330,6 +331,9 @@ def warmup(self, batch: Batch) -> int: | |||
| # Counter-intuitively, now we ignore the input batch. Instead, we create dummy batches to cover all possible | |||
| # batch sizes and sequence lengths. | |||
| seq_len = self.model.config.sequence_length | |||
| if os.environ.get("SKIP_WARMUP", "0") == "1": | |||
| logger.debug("Skipping warmup") | |||
There was a problem hiding this comment.
Does this is used mostly for debug only? Or can it be turned on for other reasons? In the later case I would use logger.warning if not, logger.debug is fine
There was a problem hiding this comment.
I would say it's only for debugging. Warmup will check if model can fit in memory and prepare inference so prefill and decode is rather fast afterwards, but it can last around 4-5 minutes, so it can be annoying for debugging the container.
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?
Text generation was not correct because the weights in the model were not correctly loaded. This is not something that it was easy to spot just looking at few tokens generated, and it was something that it was actually fixed already in the Jetstream/Pytorch code, but the fix hadn't been ported to optimum-tpu.
This fix implement the necessary weights changes, aligning to Jetstream Pytorch, and tests expected output has been modified accordingly.
I plan to submit a change directly to the Jetstream Pytorch project, so to be able to reuse their modeling code and avoid surprises like this one in the future. EDIT: I added a commit that uses the original model, so there will be no need to request a change.Also, I changed the way we install Jetstream Pytorch, to make it more reliable, as the
pip installwith the git revision was failing sometimes.