[Pass] PruneRelaxFunc to remove Relax function based on target #1555
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
MasterJH5574
force-pushed
the
01-07-prune
branch
from
60998ab to
ce7fa2c
Compare
junrushao
reviewed
Jan 7, 2024
| @@ -75,7 +76,8 @@ def _mlc_llm_pipeline( # pylint: disable=too-many-arguments | |||
| def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule: | |||
| seq = tvm.transform.Sequential( | |||
| [ | |||
| # Phase 0. Add additional information for compilation | |||
| # Phase 0. Add additional information for compilation and remove unused Relax func | |||
| PruneRelaxFunc(), | |||
There was a problem hiding this comment.
how about this: let's add a boolean flag flashinfer_enabled in _mlc_llm_pipeline, which comes from OptimizationFlags, and base on this flag we may choose to prune flashinfer-related functions
There was a problem hiding this comment.
Just updated with a new flag here. Actually I think it might be more extensible to pass the entire compilation flag object to the pipeline, given right now we have already passed flashinfer and cublas-gemm. What do you think?
There was a problem hiding this comment.
The principle I'm trying to maintain here is that we wanted every nn.Module to be standalone-compilation-and-debuggable, it means:
SomeModule(...).jit(format="torch", pipeline=..., pipeline_args=...)and it would be ideal if pipeline_args are package-independent, i.e. one doesn't have to import anything extra to create a pipeline. In this sense, while there are indeed abusively huge amount of parameters to tweak, I believe we don't have to aggregate them into one dataclass that requires extra import logics - as long as we document them clearly
junrushao
reviewed
Jan 7, 2024
MasterJH5574
force-pushed
the
01-07-prune
branch
from
ce7fa2c to
dbb7672
Compare
We recently noticed that when FlashInfer is not built due to unsupported cuda architecture or platform, running single-sequence ChatModule will hit VM function initialization error, where the function is used in `create_flashinfer_paged_kv_cache`, which won't actually be invoked in single-sequence flow. This is due to relax VM eagerly initializes all used PackedFunc at initialization stage (instead of lazy load). Therefore, even when the `create_flashinfer_paged_kv_cache` is not invoked, the PackedFuncs will be looked up. So whenever FlashInfer is not available, the issue will happen. This PR adds a compiler pass which removes `create_flashinfer_paged_kv_cache` (and also other similar functions that may be introduced in the future) based on the target. This pass can effectively address the issue.
MasterJH5574
force-pushed
the
01-07-prune
branch
from
dbb7672 to
9d35c77
Compare
junrushao
approved these changes
Jan 8, 2024
MasterJH5574
added a commit
to MasterJH5574/mlc-llm
that referenced
this pull request
Jan 9, 2024
…i#1555) We recently noticed that when FlashInfer is not built due to unsupported cuda architecture or platform, running single-sequence ChatModule will hit VM function initialization error, where the function is used in `create_flashinfer_paged_kv_cache`, which won't actually be invoked in single-sequence flow. This is due to relax VM eagerly initializes all used PackedFunc at initialization stage (instead of lazy load). Therefore, even when the `create_flashinfer_paged_kv_cache` is not invoked, the PackedFuncs will be looked up. So whenever FlashInfer is not available, the issue will happen. This PR adds a compiler pass which removes `create_flashinfer_paged_kv_cache` (and also other similar functions that may be introduced in the future) based on the target. This pass can effectively address the issue.
MasterJH5574
added a commit
to MasterJH5574/mlc-llm
that referenced
this pull request
Jan 9, 2024
…i#1555) We recently noticed that when FlashInfer is not built due to unsupported cuda architecture or platform, running single-sequence ChatModule will hit VM function initialization error, where the function is used in `create_flashinfer_paged_kv_cache`, which won't actually be invoked in single-sequence flow. This is due to relax VM eagerly initializes all used PackedFunc at initialization stage (instead of lazy load). Therefore, even when the `create_flashinfer_paged_kv_cache` is not invoked, the PackedFuncs will be looked up. So whenever FlashInfer is not available, the issue will happen. This PR adds a compiler pass which removes `create_flashinfer_paged_kv_cache` (and also other similar functions that may be introduced in the future) based on the target. This pass can effectively address the issue.
MasterJH5574
added a commit
to MasterJH5574/mlc-llm
that referenced
this pull request
Jan 9, 2024
…i#1555) We recently noticed that when FlashInfer is not built due to unsupported cuda architecture or platform, running single-sequence ChatModule will hit VM function initialization error, where the function is used in `create_flashinfer_paged_kv_cache`, which won't actually be invoked in single-sequence flow. This is due to relax VM eagerly initializes all used PackedFunc at initialization stage (instead of lazy load). Therefore, even when the `create_flashinfer_paged_kv_cache` is not invoked, the PackedFuncs will be looked up. So whenever FlashInfer is not available, the issue will happen. This PR adds a compiler pass which removes `create_flashinfer_paged_kv_cache` (and also other similar functions that may be introduced in the future) based on the target. This pass can effectively address the issue.
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.
We recently noticed that when FlashInfer is not built due to unsupported cuda architecture or platform, running single-sequence ChatModule will hit VM function initialization error, where the function is used in
create_flashinfer_paged_kv_cache, which won't actually be invoked in single-sequence flow.This is due to relax VM eagerly initializes all used PackedFunc at initialization stage (instead of lazy load). Therefore, even when the
create_flashinfer_paged_kv_cacheis not invoked, the PackedFuncs will be looked up. So whenever FlashInfer is not available, the issue will happen.This PR adds a compiler pass which removes
create_flashinfer_paged_kv_cache(and also other similar functions that may be introduced in the future) based on the target. This pass can effectively address the issue.