Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[RFC]PyTorchTVM #25

Merged
merged 25 commits into from
Sep 13, 2021
Merged

[RFC]PyTorchTVM #25

merged 25 commits into from
Sep 13, 2021

Conversation

Meteorix
Copy link
Contributor

This RFC add a PyTorchTVM module to support: compile TorchScript to TVM and use accelerated module in PyTorch.

Initial PR: apache/tvm#8777

Discuss: https://discuss.tvm.apache.org/t/rfc-pytorchtvm-compile-torchscript-to-tvm-and-use-accelerated-module-in-pytorch/10873

@masahi masahi self-assigned this Aug 24, 2021
@masahi
Copy link
Member

masahi commented Aug 24, 2021

cc @t-vi

@junrushao
Copy link
Member

Thanks for the nice RFC @Meteorix! It is definitely super important direction to embed TVM into PyTorch. I will read it carefully later this week.

@Meteorix Meteorix mentioned this pull request Aug 25, 2021
Merged
@masahi
Copy link
Member

masahi commented Aug 25, 2021

Is this supposed to be useful for accelerating PyTorch training?

@Meteorix
Copy link
Contributor Author

Is this supposed to be useful for accelerating PyTorch training?

From my perspective, we are a little far from supporting PyTorch training. I can think of 2 possible ways:

  1. use torchscript to train
  2. use relay to train

It seems both ways are in early stage. Maybe we can further discuss the possibility, but not in the scope of this RFC.

@t-vi
Copy link

t-vi commented Aug 27, 2021

I wonder whether this would make the torch fallback op (apache/tvm#7401) more or less useful (it would depend on what you (plan to) do with unsupported ops). I am still pondering whether to close it or dust it off.

I should note that as far as I know NVidia has a TensorRT front end serving a similar goal and there also is one for ONNXRuntime-as-a-module (recently featured in the PyTorch blog). There may be useful design insights in how they work (or maybe they're too different to what you have in mind).

junrushao
junrushao reviewed Aug 29, 2021
View reviewed changes
rfcs/0025-add-pytorch-tvm.md Outdated Show resolved Hide resolved
rfcs/0025-add-pytorch-tvm.md Outdated Show resolved Hide resolved
rfcs/0025-add-pytorch-tvm.md Outdated Show resolved Hide resolved
rfcs/0025-add-pytorch-tvm.md Outdated Show resolved Hide resolved
rfcs/0025-add-pytorch-tvm.md Outdated Show resolved Hide resolved
rfcs/0025-add-pytorch-tvm.md Outdated Show resolved Hide resolved
junrushao
junrushao reviewed Aug 30, 2021
View reviewed changes
rfcs/0025-add-pytorch-tvm.md
* Onnx cannot cover all models with dynamic control flow (e.g. for loop)
* TensorRT can only accelerate some standard networks

So we hope to use TVM to accelerate PyTorch model inference.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should discuss more about the benefit of PyTorchTVM here

junrushao
junrushao reviewed Aug 30, 2021
View reviewed changes
rfcs/0025-add-pytorch-tvm.md
Comment on lines +60 to +64
images: List[torch.Tensor] = []
for path in image_path:
img = read_image(path)
images.append(img)
x = torch.stack(images).cuda().half()
Copy link
Member

@junrushao junrushao Aug 30, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shall we conclude these lines to an abstract data loader to be consistent with our description? like, self.load_data()

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe just keep the class simple?

junrushao
junrushao reviewed Aug 30, 2021
View reviewed changes
rfcs/0025-add-pytorch-tvm.md Outdated
# Prior art
[prior-art]: #prior-art

Our implementation is inspired by this RFC: https://discuss.tvm.apache.org/t/rfc-add-tensorflow-custom-op-to-embed-tvm-runtime-in-tensorflow-graph-and-session/4601
Copy link
Member

@junrushao junrushao Aug 30, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Our implementation is inspired by this RFC: https://discuss.tvm.apache.org/t/rfc-add-tensorflow-custom-op-to-embed-tvm-runtime-in-tensorflow-graph-and-session/4601
Our implementation is inspired by this RFC which allows TVM operators to be embedded into TensorFlow: https://discuss.tvm.apache.org/t/rfc-add-tensorflow-custom-op-to-embed-tvm-runtime-in-tensorflow-graph-and-session/4601. Beyond a single operator, our RFC further extends the scope to embedding arbitrary TVM graphs to TorchScript, and enables smooth integration and serving with both `torch.jit.trace` and `torch.jit.script` APIs.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To clarify: a major difference is that, the TF RFC only support embedding a TVM op, but not a TVM graph. And we also have an internal implementation that embeds TVM graph into TF.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated with clarification

@junrushao
Copy link
Member

Thanks again for the contribution! I have no doubt that it is going to be hugely important piece of work. Just made some suggestions in terms of wording.

@Meteorix
Copy link
Contributor Author

Meteorix commented Aug 30, 2021 •
edited
Loading

I wonder whether this would make the torch fallback op (apache/tvm#7401) more or less useful (it would depend on what you (plan to) do with unsupported ops). I am still pondering whether to close it or dust it off.

@t-vi It will help a lot. If we have the torch fallback op support, users will less likely get stuck in frontend conversion phase. We are looking forward to this feature. Could you complete the PR?

I should note that as far as I know NVidia has a TensorRT front end serving a similar goal and there also is one for ONNXRuntime-as-a-module (recently featured in the PyTorch blog). There may be useful design insights in how they work (or maybe they're too different to what you have in mind).

And thanks for this information. Actually, we are inspired by another NVIDIA work: TRTorch.

@suluner
Copy link

suluner commented Sep 9, 2021

Can this method split the graph automaticly given a torchscript, and convert the subgraph that tvm Relay IR can support or tvm outperforming pytorch?

junrushao
junrushao approved these changes Sep 9, 2021
View reviewed changes
Copy link
Member

@junrushao junrushao left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@junrushao
Copy link
Member

@masahi @t-vi Please take another look :-)

hogepodge
hogepodge suggested changes Sep 9, 2021
View reviewed changes
Copy link
Contributor

@hogepodge hogepodge left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In general I support this effort. I'd like to see stronger integrations of TVM into adjacent projects like PyTorch. This looks like a good start.

One minor spelling point, but otherwise this looks good to me.

rfcs/0025-add-pytorch-tvm.md Outdated Show resolved Hide resolved
Laurawly
Laurawly approved these changes Sep 9, 2021
View reviewed changes
Copy link

@Laurawly Laurawly left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great work!

@tqchen
Copy link
Member

tqchen commented Sep 10, 2021

Thanks for the great work, a minor note on rendering. it would be great to add the codeblock type to the code so we have syntax highlights

@Meteorix
Copy link
Contributor Author

@hogepodge @tqchen thanks for your advices. I have updated the PR.

@junrushao
Copy link
Member

It seems that we all reach consensus and no more comments are raised for several days, let's get this RFC finally merged! 🔥 🎉

@junrushao junrushao merged commit 1ed4d70 into apache:main Sep 13, 2021
@zhiqwang zhiqwang mentioned this pull request Nov 17, 2021
Open
10 tasks
@masahi masahi mentioned this pull request Jan 10, 2022
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@hogepodge hogepodge hogepodge requested changes

@junrushao junrushao junrushao approved these changes

@Laurawly Laurawly Laurawly approved these changes

@masahi masahi masahi approved these changes

Assignees

@masahi masahi

@junrushao junrushao

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
8 participants
@Meteorix @masahi @junrushao @t-vi @suluner @tqchen @hogepodge @Laurawly