[Dynamo] Refactor get_wrapper and pickling compiled graph #438
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
Simple model with one conv2d failed. - fix signature for conv* ops to corresponds torch.nn.functional] - add missed padding normalization After that the model works
Partial changes related to hidet-org#18
Add .vscode to .gitignore
Previously, if a performance regression fails due to an exception, the job that stops the runner VM instances will be skipped, leaving the instances on. This will make the stop_instances job run even when previous jobs failed. Not sure if always() will override the inputs.shutdown_instances flag, if it does we can move it into the step scope.
Module wrapper around groupnorm operator. Supports compiled app development.
destefy
force-pushed
the
stefan/refactor_get_wrapper
branch
from
9555421
to
4132253
Compare
Format Small typo
destefy
force-pushed
the
stefan/refactor_get_wrapper
branch
from
4132253
to
13f83e3
Compare
yaoyaoding
requested changes
Mar 25, 2024
There was a problem hiding this comment.
Hi @destefy,
I don't like the idea of wrapping what we have returned as a torch.fx.GraphModule. If you want to have something that can be saved, you can use CompiledGraph. Or make the CompiledForwardFunction serializable.
Comment on lines
134
to
141
| use_cuda_graph = dynamo_config['use_cuda_graph'] | ||
| if use_cuda_graph: | ||
| try: | ||
| runner = self.cgraph.cuda_graph() | ||
| except CudaGraphCreationError: | ||
| runner = self.cgraph | ||
| else: | ||
| runner = self.cgraph |
There was a problem hiding this comment.
We should not create the cuda graph for every call invoke.
| try: | ||
| runner = cgraph.cuda_graph() | ||
| except CudaGraphCreationError: | ||
| class CompiledForwardFunction(torch.nn.Module): |
There was a problem hiding this comment.
I prefer the name TorchCompiledModule
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.
The CentML compilation backend I am working on wants to wrap the
CompiledGraphs forward function (the one returned byget_wrapper) in atorch.fx.GraphModule. ThisGraphModulewould then be pickled and sent from a server to a client.However, it isn't possible to pickle the lambda/local function returned by
get_wrapper. Therefore, I am turningget_wrapperinto a classCompiledForwardFunctionwhoseforwardfunction behaves like thewrapperreturned byget_wrapper.Additionally, in order to pickle
CompiledForwardFunction, I have defined pickling and unpickling behaviour forCompiledGraphusing__getstate__and__setstate__respectively. These just callCompiledGraph's existingsaveandloadfunctions.