[Mixtral] Use col major for MoE gemm and use small batch specialization #171
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
masahi
approved these changes
Jan 24, 2024
Lunderberg
pushed a commit
to Lunderberg/mlc-llm
that referenced
this pull request
Jan 30, 2024
This PR introduces the model reload functionality to the chat module and the CLI. * In CLI, the usage is `/reload [model_id]`, for example `model_id` can be `vicuna-v1-7b-q3f16_0`. When `model_id` is specified, CLI will search the corresponding model for `model_id` in candidate paths, and load the specified model. When it is not specified, CLI will reload the currently running model again. * The effect of the reload operation is to load the (specified) model as if loading it at the very beginning. This means the chat module will completely release and forget everything happened before reload. * Regarding the implementation, this PR unifies the previous `init_chat` function with the `Init` function into the `Reload` function. * Verified on Mac Studio and Linux that the reload operation will not cause memory leak. * Given the iOS/Android app does not support reading json config string yet, the refactor of this PR keeps some existing functions that required by the iOS/Android app as legacy functions. TODO items are marked where legacy functions are. I have verified that the legacy functions work.
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.
Changed MoE gemm to use more performant col major layouts. Also cleaned up param loading hack for mixtral that previously used PyTorch for fast weight transposition.
Benchmark: https://docs.google.com/spreadsheets/d/1w1uYvBK9bZluue4uyZjgT_4e8tl0rQDfMZ2XQfg43f4/edit#gid=0
need the commit https://github.com/octoml/tvm/commit/7fb6b704879f5accaf3ff08c96409123a02b8f58
cc @sunggg @masahi