Enable Shared Function in LiftTransformParam Pass#16717
Merged
vinx13 merged 12 commits intoapache:mainfrom Mar 19, 2024
Merged
Conversation
Currently, the `relax.transform.LiftTransformParams` pass produces a separate `transform_params` function for every function in the `IRModule`. In most cases, the functions in an `IRModule` all accept the same set of model weights (e.g. `"prefill"` and `"decode"` in a transformer model). However, the lifted `*_transform_params` functions may be different for each inference function. The goal is to introduce a new optional parameter `shared_transform` for `LiftTransformParams`. If set, a single parameter transformation function should be generated for the entire `IRModule`, rather than one parameter transformation function for each original function. Because the shared parameter transformation function must be compatible with all existing functions, it should only contain parameter transformation steps that are common across all input functions.
vinx13
reviewed
Mar 14, 2024
Lunderberg
reviewed
Mar 16, 2024
Contributor
Lunderberg
left a comment
Lunderberg
left a comment
There was a problem hiding this comment.
Overall, looks good, and thank you for the improvement! A couple of comments on it, then it should be good to go!
Member
Author
|
@Lunderberg thank you for the detailed review, comments are all addressed, please take another look when you got time, thanks a lot! |
Lunderberg
approved these changes
Mar 19, 2024
Contributor
Lunderberg
left a comment
Lunderberg
left a comment
There was a problem hiding this comment.
Thank you for making the changes, and LGTM!
thaisacs
pushed a commit
to thaisacs/tvm
that referenced
this pull request
Apr 3, 2024
* [WIP] LiftTransformParams for multiple functions * pass test * [In-Progress] Define desired behavior for shared LiftTransformParams Currently, the `relax.transform.LiftTransformParams` pass produces a separate `transform_params` function for every function in the `IRModule`. In most cases, the functions in an `IRModule` all accept the same set of model weights (e.g. `"prefill"` and `"decode"` in a transformer model). However, the lifted `*_transform_params` functions may be different for each inference function. The goal is to introduce a new optional parameter `shared_transform` for `LiftTransformParams`. If set, a single parameter transformation function should be generated for the entire `IRModule`, rather than one parameter transformation function for each original function. Because the shared parameter transformation function must be compatible with all existing functions, it should only contain parameter transformation steps that are common across all input functions. * [TIR] Implemented shared lift transform params * Comments & skip test. * Linting. * Avoid c++20 feature to pass CI. * Remove unused code. * Fix interface as suggested. * Fix docs. * Fix interface as suggested. * Move code for readability. --------- Co-authored-by: Wuwei Lin <wuwei@apache.org> Co-authored-by: Eric Lunderberg <elunderberg@octoml.ai>
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
This PR enables specifying a list of function names to extract shared transform parameters. A single parameter transformation function will be produced, containing the preprocessing steps common across each function whose name is in a given function name list.
Unit tests are passing except the one that has no shared prepocessing (transpose), skipping for now, will follow-up in another PR.
Cherry-picked from @vinx13's working branch.