Add MoE layer example #303
Draft
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
for more information, see https://pre-commit.ci
This was referenced May 1, 2024
|
Super exciting! Really looking forward to discuss this in more detail at a design review! |
|
Do we have any broader ideas for how this fits into the strategy for handling dynamic and data dependent shapes? I was under the impression that this was just something we were completely incapable of doing with the way that we're modeling traces. |
6 tasks
IvanYashchuk
commented
May 8, 2024
| @@ -1048,6 +1053,7 @@ def find_producer_symbols(trace: TraceCtx, proxies: Sequence[Proxy], stop_proxie | |||
| (__b = ltorch.sub(x, y) | |||
| # __b = prims.sub(x, y),) | |||
| """ | |||
| stop_proxies = filter(lambda x: isinstance(x, Proxy), stop_proxies) | |||
There was a problem hiding this comment.
Suggested change
| stop_proxies = filter(lambda x: isinstance(x, Proxy), stop_proxies) | |
| stop_proxies = tuple(filter(lambda x: isinstance(x, Proxy), stop_proxies)) |
Vaguely remember that I have run into |
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.
DRAFT MODE TO PREVENT MERGES The approach and code is ready for experimentation and review.
The main result of this PR is that Thunder can run a variant of the MoE layer from LitGPT. There are three modifications
None(need to create an issue). The workaround is to use unsqueeze instead ofNonewhen indexing.+=) is replaced withindex_add.The main missing operator is
nonzero(x, as_tuple=True). The problem with this operator is that the output shape is unknown at compile time and it's dynamic at runtime. I tried using NumberProxy with None, NumberProxy with a custom int subclass as value, using a custom int subclass directly. But simple-1in the shape worked best.The forward pass worked just with 14ce097. The backward pass required more of
-1-special handling.Currently,
index_add,index_select,topkare not fused with any of Thunder's fusing executors.