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

[onnx.export] Avoid linear look up in env for exist_in_env #124909

Closed
wants to merge 1 commit into from
Closed

[onnx.export] Avoid linear look up in env for exist_in_env #124909

wants to merge 1 commit into from

Conversation

gustavla
Copy link
Contributor

This PR is part of a series of PRs to significantly speed up torch.onnx.export for models with many nodes (e.g. LLM). See #121422 for more analysis.

  • As part of torch.onnx.export, a reverse look-up is made in env. This is done for each node, and this look-up costs in proportional to the graph size, which incurs and overall O(N^2) time complexity.
  • A pragmatic solution is simply to keep a separate data structure to make this de facto constant time. So, this introduces a set containing all the values of env. Open to other ideas. Ideally exist_in_env wouldn't be needed at all, but to preserve current behavior exactly I'm not sure how that can be done.
  • Resolves (4) in ONNX export is unnecessarily slow (O(N^2)) #121422.
  • This code change and the choice of py::set looks a bit more natural on top of [onnx.export] Avoid dict <-> unordered_map implicit copies #123063, where the env is changed from a std::unordered_map to a py::dict.

Partially fixes #121422

@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Apr 25, 2024 •
edited

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/124909

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

✅ No Failures

As of commit 0726246 with merge base 6ea226b (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@linux-foundation-easycla Linux Foundation: EasyCLA
Copy link

linux-foundation-easycla bot commented Apr 25, 2024 •
edited

CLA Signed

The committers listed above are authorized under a signed CLA.

  • ✅ login: gustavla / name: Gustav Larsson (0726246)

@pytorch-bot pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Apr 25, 2024
@gustavla gustavla force-pushed the dev/gustav/onnx_export_speedup_121422_D_4_real branch from 6e80e7f to ab070ac Compare April 25, 2024 04:20
@gustavla gustavla mentioned this pull request Apr 25, 2024
Closed
9 tasks
@gustavla gustavla changed the title Avoid linear look up in env for exist_in_env [onnx.export] Avoid linear look up in env for exist_in_env Apr 25, 2024
@gustavla gustavla force-pushed the dev/gustav/onnx_export_speedup_121422_D_4_real branch from ab070ac to b59ab82 Compare April 25, 2024 15:57
@ezyang ezyang added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Apr 25, 2024
justinchuby
justinchuby reviewed May 9, 2024
View reviewed changes
torch/csrc/jit/passes/onnx.cpp
@@ -169,8 +169,14 @@ std::shared_ptr<Graph> ToONNX(
ConstantValueMap::ClearMaps();
auto new_graph = std::make_shared<Graph>(graph->current_scope());
std::unordered_map<Value*, Value*> env;
py::set values_in_env;
Copy link
Collaborator

Choose a reason for hiding this comment

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

Could you add the rationale as a comment here for future readers? Thanks!

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Added a comment, thanks! I also pushed a rebase since #123063 merged, that touches many of the same lines.

@gustavla
Avoid linear look up in env for exist_in_env
0726246 
- As part of torch.onnx.export, a reverse look-up is made in env. This
  is done for each node, and this look-up costs in proportion to the
  graph size, which incurs and overall O(N^2) time complexity.
- A pragmatic solution is simply to keep a separate data structure to
  make this de facto constant time. So, this introduces a set containing
  all the values of env.
- Resolves (4) in pytorch#121422.
@gustavla gustavla force-pushed the dev/gustav/onnx_export_speedup_121422_D_4_real branch from b59ab82 to 0726246 Compare May 9, 2024 18:28
@justinchuby
Copy link
Collaborator

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label May 9, 2024
@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@linshokaku linshokaku mentioned this pull request May 13, 2024
Open
@titaiwangms titaiwangms added the topic: improvements topic category label Jul 10, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@justinchuby justinchuby justinchuby approved these changes

@srikris-sridhar srikris-sridhar srikris-sridhar approved these changes

@BowenBao BowenBao Awaiting requested review from BowenBao

@thiagocrepaldi thiagocrepaldi Awaiting requested review from thiagocrepaldi

@wschin wschin Awaiting requested review from wschin wschin is a code owner

Assignees
No one assigned
Labels
ciflow/trunk Trigger trunk jobs on your pull request Merged open source release notes: onnx torch.onnx related changes that should show up in the release notes topic: improvements topic category triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

ONNX export is unnecessarily slow (O(N^2))
7 participants
@gustavla @justinchuby @pytorchmergebot @srikris-sridhar @ezyang @titaiwangms @pytorchbot