Refactor llama / mixtral / grok for shared features #267
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
rsuderman
requested review from
KyleHerndon,
archana-ramalingam and
dan-garvey
rsuderman
force-pushed
the
refactor_llm
branch
from
62aecb3 to
bb5be7e
Compare
| @@ -122,33 +123,32 @@ def prefill( | |||
| self._assert_device(seq_block_ids) | |||
| self._assert_device(*cache_state, dtype=self.activation_dtype) | |||
| h = self.token_embedding(tokens) | |||
| h *= 78.38367176906169 | |||
| h *= math.sqrt(h.shape[-1]) | |||
There was a problem hiding this comment.
This requires import math
There was a problem hiding this comment.
better as torch? I don't think math traces well
There was a problem hiding this comment.
Do we not have prefill here testing export? If not that needs to be added
There was a problem hiding this comment.
We should try to make faked versions of each model that is locally exportable. @KyleHerndon can you look into making faked theta parameters for the model? It can just be a single layer with smallers tensors.
There was a problem hiding this comment.
Here are the theta generators for attention/FFN/MOE blocks. We already have attention and prefill tests, might need to add export component in there.
IanNod
reviewed
Oct 14, 2024
| @@ -122,33 +123,32 @@ def prefill( | |||
| self._assert_device(seq_block_ids) | |||
| self._assert_device(*cache_state, dtype=self.activation_dtype) | |||
| h = self.token_embedding(tokens) | |||
| h *= 78.38367176906169 | |||
| h *= math.sqrt(h.shape[-1]) | |||
There was a problem hiding this comment.
Do we not have prefill here testing export? If not that needs to be added
rsuderman
force-pushed
the
refactor_llm
branch
from
bb5be7e to
770f83f
Compare
rsuderman
force-pushed
the
refactor_llm
branch
from
770f83f to
2bd1bbc
Compare
rsuderman
force-pushed
the
refactor_llm
branch
from
5beb357 to
4fb7663
Compare
Many of these features can toggle between depending on architecture. Replumbing the configurations separately allows better reuse and understanding of how models vary between eachother. grok uses a softcap, plumbing a value enables `sc * tanh( v / sc)` grok has some hardcoded values that have better representations, e.g. `sqrt(6144)` and `sqrt(3)`. output normalization is optional but used by mixtral. Presence of the tensor is sufficient for performing the normalization.
Many of these features can toggle between depending on architecture. Replumbing the configurations separately allows better reuse and understanding of how models vary between eachother. grok uses a softcap, plumbing a value enables `sc * tanh( v / sc)` grok has some hardcoded values that have better representations, e.g. `sqrt(6144)` and `sqrt(3)`. output normalization is optional but used by mixtral. Presence of the tensor is sufficient for performing the normalization. We remove the sparse moe block as we now know it will not be used due to poor performance.
rsuderman
force-pushed
the
refactor_llm
branch
from
4fb7663 to
80891fe
Compare
archana-ramalingam
approved these changes
Oct 16, 2024
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.
Many of these features can toggle between depending on architecture. Replumbing the configurations separately allows better reuse and understanding of how models vary between eachother.
grok uses a softcap, plumbing a value enables
sc * tanh( v / sc)grok has some hardcoded values that have better representations, e.g.sqrt(6144)andsqrt(3).output normalization is optional but used by mixtral. Presence of the tensor is sufficient for performing the normalization.