Elamur Research⌕
☰
home / papers / 2024 / PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation

PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation

Paper · 2024 · Jason Ansel, Edward Yang, Horace He et al.

Summary

Eager-mode frameworks such as PyTorch are popular because models are plain Python code that can be debugged and mixed freely with other libraries. The cost is that the framework only sees one operator at a time, so it cannot apply cross-operator optimizations like fusion. Earlier graph-capture tools for PyTorch each had a weakness. Tracing can silently give wrong results. AST-based scripting breaks on much real Python. Lazy evaluation adds runtime overhead and delays kernel launches.

The authors present the two components behind torch.compile in PyTorch 2. TorchDynamo is a JIT compiler that works on Python bytecode. It intercepts frames just before CPython runs them, pulls PyTorch operations out into FX graphs, and leaves everything it cannot handle to ordinary Python. TorchInductor is the default backend and generates Triton code for GPUs and C++/OpenMP for CPUs.

Tested on more than 180 real-world models, TorchDynamo captures graphs more reliably than earlier tools at low overhead. TorchInductor beats six other compiler backends. This shows that compiler optimizations can be added to eager frameworks without giving up usability.

Key points

  • On an NVIDIA A100, TorchInductor achieves geometric-mean speedups over eager PyTorch of 2.27× for inference and 1.41× for training. It outperforms the nvFuser, NNC, XLA, ONNX Runtime, TVM and Hidet backends.
  • The evaluation covers three benchmark suites (TorchBench, HuggingFace, timm), CPU and GPU, float32 and float16, and both training and inference.
  • On TorchBench models, torch.jit.script succeeds only about half the time. TorchDynamo usually captures a single whole-program graph and falls back to partial graphs when it cannot.
  • TorchInductor uses a define-by-run, loop-level IR and is written in Python. torch.compile was officially released in March 2023.
  • The code ships inside PyTorch under the BSD-3 license (torch/_dynamo, torch/_inductor). Benchmark scripts are in benchmarks/dynamo.
  • Ablation: if both fusion and inlining are turned off, TorchInductor makes models slower rather than faster.

Interesting ideas

  • I1. Partial capture with fallback instead of all-or-nothing. TorchDynamo compiles the parts it understands and hands the rest back to the normal interpreter. Any tool that optimizes or translates user code can adopt this to stay robust on messy real-world inputs.
  • I2. Hooking in at the bytecode level rather than parsing source. Working on bytecode right before execution sees the program as it actually runs. This avoids re-implementing a whole language's semantics statically, which is where AST-based approaches tend to fail.
  • I3. Writing the compiler in the users' own language. Because TorchInductor is in Python, the people who use it can also extend it. Internal tooling in data or NLP pipelines could benefit from the same choice.
  • I4. Evaluating on a large set of unmodified real models. Reporting geometric-mean speedups across 180+ models gives a more honest picture than a few hand-picked benchmarks. The artifact also records failures as zero speedup, so failures visibly lower the average.
  • I5. Ablations that disable passes jointly. Fusion and inlining overlap, so the authors also turn both off together. This reveals dependencies between passes that single-pass ablations would hide.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢MetaMenlo Park · United States183 works · 228,558 citations🏢OpenAISan Francisco · United States18 works · 27,146 citations🎓University of MichiganAnn Arbor · United States70 works · 13,693 citations🏢IntelSanta Clara · United States15 works · 41,257 citations🏢QuansightAustin · United States3 works · 62,887 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksComputer Vision and Pattern Recognition · 2,869 worksHardware and Architecture · 81 works

💡 Concepts 1Key ideas and methods explained on our concept pages.Machine learning · 51 works

Cite

Cite as Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard and Evgeni Burovski. 2024. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. ASPLOS 2024.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.1145/3620665.3640366

@inproceedings{ansel2024pytorch,
  title = {{PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation}},
  author = {Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard and Evgeni Burovski},
  booktitle = {ASPLOS 2024},
  year = {2024},
  doi = {10.1145/3620665.3640366},
  url = {https://doi.org/10.1145/3620665.3640366}
}

References

60 references, 9 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Martin Abadi et al. 2016. TensorFlow: A system for large-scale machine learning. In 12th USENIX Symposium on Operating Systems Design and Implementation (OSDI 16), 265–283. https://www.usenix.org/system/files/conference/osdi16/osdi16-abadi.pdf. ↗
  2. [SW] Martín Abadi et al., TensorFlow, Large-scale machine learning on heterogeneous systems Nov. 2015. doi: 10.5281/zenodo.4724125. — [doi:10.5281/zenodo.4724125](https://doi.org/10.5281/zenodo.4724125) find
  3. Hameer Abbasi, Edward Z Yang, and Ralf Gommers. 2020. Improving subclassing Tensor by propagating subclass instances. https://githu b.com/pytorch/rfcs/blob/master/RFC-0001-torch-function-for-me thods.md. (Aug. 2020). find
  4. Akshay Agrawal et al. 2019. TensorFlow Eager: A Multi-Stage, Python-Embedded DSL for Machine Learning. CoRR, abs/1903.01855. http://arxiv.org/abs/1903.01855 arXiv: 1903.01855. — [arXiv:1903.01855](https://arxiv.org/abs/1903.01855) ↗
  5. Rami Al-Rfou et al. 2016. Theano: A Python framework for fast computation of mathematical expressions. CoRR, abs/1605.02688. http://arxiv.org/abs/1605.02688 arXiv: 1605.02688. ↗
  6. Jason Ansel, Shoaib Kamil, Kalyan Veeramachaneni, Jonathan Ragan-Kelley, Jeffrey Bosboom, Una-May O’Reilly, and Saman Amarasinghe. 2014. Opentuner: an extensible framework for program autotuning. In Proceedings of the 23rd International Conference on Parallel Architectures and Compilation (PACT ’14)… — [doi:10.1145/2628071.2628092](https://doi.org/10.1145/2628071.2628092) find
  7. Riyadh Baghdadi, Jessica Ray, Malek Ben Romdhane, Emanuele Del Sozzo, Abdurrahman Akkas, Yunming Zhang, Patricia Suriana, Shoaib Kamil, and Saman Amarasinghe. 2019. Tiramisu: a polyhedral compiler for expressing fast and portable code. In Proceedings of the 2019 IEEE/ACM International Symposium on C… find
  8. [SW] James Bradbury et al., JAX: composable transformations of Python+NumPy programs version 0.3.13, 2018. url: http://github.com/google/jax. ↗
  9. Dino Viehland Brett Cannon. 2016. PEP 523: adding a frame evaluation API to CPython. https://peps.python.org/pep-0523/. (2016). ↗
  10. Jack Cao. 2022. PyTorch/XLA 2022 Q4 dev update. https://dev-discu ss.pytorch.org/t/pytorch-xla-2022-q4-dev-update/961. (2022). find
  11. Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018. Learning to optimize tensor programs. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18). Curran Associates Inc., Montr… find
  12. Tianqi Chen et al. 2018. TVM: an automated End-to-End optimizing compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18). USENIX Association, Carlsbad, CA, (Oct. 2018), 578–594. isbn: 978-1-939133-08-3. https://www.usenix.org/conference/osdi18/p… ↗
  13. Sharan Chetlur, Cliff Woolley, Philippe Vandermersch, Jonathan Cohen, John Tran, Bryan Catanzaro, and Evan Shelhamer. 2014. cuDNN: efficient primitives for deep learning. (2014). arXiv: 1410.07 59 [cs.NE]. find
  14. Will Constable et al. 2020. TorchBench: a collection of open source benchmarks for PyTorch performance and usability evaluation. https://github.com/pytorch/benchmark. (Sept. 2020). ↗
  15. Leonardo Dagum and Ramesh Menon. 1998. OpenMP: an industry standard API for shared-memory programming. Computational Science & Engineering, IEEE, 5, 1, 46–55. find
  16. ONNX Runtime developers. 2021. ONNX runtime. https://www.onn xruntime.ai. (2021). ↗
  17. Zachary DeVito et al. 2018. TorchScript. https://pytorch.org/docs/1 .9.0/jit.html. (Sept. 2018). ↗
  18. Yaoyao Ding, Cody Hao Yu, Bojian Zheng, Yizhi Liu, Yida Wang, and Gennady Pekhimenko. 2023. Hidet: task-mapping programming paradigm for deep learning tensor programs. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Vo… find
  19. Siyuan Feng et al. 2022. TensorIR: an abstraction for automatic tensorized program optimization. (2022). arXiv: 2207.04296 [cs.LG]. find
  20. Alan Gray. 2019. Getting started with CUDA graphs. https://develo per.nvidia.com/blog/cuda-graphs/. (2019). find
  21. Charles R. Harris et al. 2020. Array programming with NumPy. Nature, 585, 357–362. doi: 10.1038/s41586-020-2649-2.
  22. Horace He. 2019. The state of machine learning frameworks in 2019. https://thegradient.pub/state-of-ml-frameworks-2019-pytorch-d ominates-research-tensorflow-dominates-industry/. (2019). ↗
  23. Mike Innes et al. 2017. On machine learning and programming languages. https://julialang.org/blog/2017/12/ml-pl/. (Dec. 2017). ↗
  24. ISO. 1998. ISO/IEC 14882:1998: Programming languages — C++. (Sept. 1998), 732. http://webstore.ansi.org/ansidocstore/product.asp?sku =ISO%2FIEC+14882%2D1998. ↗
  25. Yangqing Jia, Evan Shelhamer, Jeff Donahue, Sergey Karayev, Jonathan Long, Ross B. Girshick, Sergio Guadarrama, and Trevor Darrell. 2014. Caffe: Convolutional Architecture for Fast Feature Embedding. CoRR, abs/1408.5093. http://arxiv.org/abs/1408.5093 arXiv: 1408.5093. ↗
  26. Norman P. Jouppi et al. 2017. In-datacenter performance analysis of a tensor processing unit. In Proceedings of the 44th Annual International Symposium on Computer Architecture (ISCA ’17). Association for Computing Machinery, Toronto, ON, Canada, 1–12. isbn: 9781450348928. doi: 10.1145/3079856.30802… find
  27. Chris Lattner et al. 2021. MLIR: scaling compiler infrastructure for domain specific computation. In 2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), 2–14. doi: 10.1 109/CGO51591.2021.9370308. find
  28. Aaron Meurer et al. 2017. SymPy: symbolic computing in Python. PeerJ Computer Science, 3, (Jan. 2017). doi: 10.7717/peerj-cs.103. — [doi:10.7717/peerj-cs.103](https://doi.org/10.7717/peerj-cs.103) find
  29. Adrian Mönnich, Armin Ronacher, David Lord, Grey Li, Joshua Bronson, Markus Unterwaditzer, and Philip Jones. 2023. Jinja project. https://github.com/pallets/jinja. (2023). ↗
  30. NVIDIA, Péter Vingelmann, and Frank H.P. Fitzek. 2023. CUDA. https://developer.nvidia.com/cuda-toolkit. (2023). ↗
  31. 2023. ONNX. https://onnx.ai/. (2023). ↗
  32. 2019. Pytorch: an imperative style, high-performance deep learning library. Proceedings of the 33rd International Conference on Neural Information Processing Systems. Curran Associates Inc., Red Hook, NY, USA, 12 pages.
  33. Jonathan Ragan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. 2013. Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines. In Proceedings of the 34th ACM SIGPLAN Conference on Programming Lang… find
  34. James Reed, Zachary DeVito, Horace He, Ansley Ussery, and Jason Ansel. 2022. Torch.fx: practical program capture and transformation for deep learning in python. In Proceedings of Machine Learning and Systems. D. Marculescu, Y. Chi, and C. Wu, (Eds.) Vol. 4, 638–651. https://proceedings.mlsys.org/pap… ↗
  35. Elvis Saravia. 2021. Papers with Code 2021: a year in review. https://medium.com/paperswithcode/papers-with-code-2021-a-year-in-review-de75d5a77b8b. (2021). ↗
  36. Christian Sarofeen, Piotr Bialecki, Jie Jiang, Kevin Stephano, Masaki Kozuki, Neal Vaidya, and Stas Bekman. 2022. Introducing nvFuser, a deep learning compiler for PyTorch. https://pytorch.org/blog/introd ucing-nvfuser-a-deep-learning-compiler-for-pytorch/. (2022). ↗
  37. Frank Seide and Amit Agarwal. 2016. CNTK: microsoft’s open-source deep-learning toolkit. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD ’16). Association for Computing Machinery, San Francisco, California, USA, 2135. isbn: 9781450342322. do… — [doi:10.1145/2939672.2945397](https://doi.org/10.1145/2939672.2945397) find
  38. Junru Shao et al. 2022. Tensor program optimization with probabilistic programs. (2022). arXiv: 2205.13603 [cs.LG]. find
  39. Alex Suhan, Davide Libenzi, Ailing Zhang, Parker Schuh, Brennan Saeta, Jie Young Sohn, and Denys Shabalin. 2021. LazyTensor: combining eager execution with domain-specific compilers. arXiv preprint arXiv:2102.13267. find
  40. PyTorch Team. 2023. TorchDynamo Benchmarking Code. https://git hub.com/pytorch/pytorch/tree/main/benchmarks/dynamo. (2023). find
  41. PyTorch Team. 2023. TorchInductor Performance Dashboard. https://hud.pytorch.org/benchmark/compilers. (2023). ↗
  42. PyTorch XLA Team. 2023. PyTorch/XLA. https://github.com/pytorc h/xla. (2023). ↗
  43. Vijay Thakkar et al. 2023. CUTLASS. https://github.com/NVIDIA/cu tlass. Version 3.0.0. (Jan. 2023). ↗
  44. [SW] The IREE Authors, IREE Sept. 2019. url: https://github.com/o penxla/iree. ↗
  45. The XLA Team. 2017. XLA - Tensorflow, compiled. https://develop ers.googleblog.com/2017/03/xla-tensorflow-compiled.html. (Mar. 2017). find
  46. Philippe Tillet, H. T. Kung, and David Cox. 2019. Triton: an intermediate language and compiler for tiled neural network computations. In (MAPL 2019). Association for Computing Machinery, Phoenix, AZ, USA, 10–19. isbn: 9781450367196. doi: 10.1145/3315508.3329973.
  47. Seiya Tokui et al. 2019. Chainer: A Deep Learning Framework for Accelerating the Research Cycle. CoRR, abs/1908.00213. http://arxiv.org/abs/1908.00213 arXiv: 1908.00213. ↗
  48. Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S. Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. 2018. Tensor comprehensions: framework-agnostic high-performance machine learning abstractions. (2018). arXiv: 1802.04730 [cs.PL]. find
  49. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Proceedings of the 31st International Conference on Neural Information Processing Systems (NIPS’17). Curran Associates Inc., Long Beach, C…
  50. B. P. Welford. 1962. Note on a method for calculating corrected sums of squares and products. Technometrics, 4, 3, 419–420. doi: 10.1080/00401706.1962.10490022. — [doi:10.1080/00401706.1962.10490022](https://doi.org/10.1080/00401706.1962.10490022) find
  51. Jian Weng, Animesh Jain, Jie Wang, Leyuan Wang, Yida Wang, and Tony Nowatzki. 2021. Unit: unifying tensorized instruction compilation. In 2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), 77–89. doi: 10.1109/CGO51591.2021.93703 30. find
  52. Ross Wightman. 2019. PyTorch image models. https://github.com/r wightman/pytorch-image-models. (2019). doi: 10.5281/zenodo.4414 861. — [doi:10.5281/zenodo.4414](https://doi.org/10.5281/zenodo.4414) ↗
  53. Thomas Wolf et al. 2020. Transformers: State-of-the-Art Natural Language Processing. In Association for Computational Linguistics, (Oct. 2020), 38–45. https://www.aclweb.org/anthology/2020.emnlp-demos.6. ↗
  54. Jiarong Xing, Leyuan Wang, Shang Zhang, Jack Chen, Ang Chen, and Yibo Zhu. 2022. Bolt: bridging the gap between auto-tuners and hardware-native performance. In Proceedings of Machine Learning and Systems. D. Marculescu, Y. Chi, and C. Wu, (Eds.) Vol. 4, 204–216. https://proceedings.mlsys.org/paper_f… ↗
  55. Shangdi Yu and Horace He. 2023. Transcending runtime-memory tradeoffs in checkpointing by being fusion aware. In Proceedings of Machine Learning and Systems. find
  56. Bojian Zheng et al. 2022. Dietcode: automatic optimization for dynamic tensor programs. In Proceedings of Machine Learning and Systems. D. Marculescu, Y. Chi, and C. Wu, (Eds.) Vol. 4, 848–863. https://proceedings.mlsys.org/paper_files/paper/2022/file/fa7cdfa d1a5aaf8370ebeda47a1ff1c3-Paper.pdf. ↗
  57. Lianmin Zheng et al. 2020. Ansor: generating high-performance tensor programs for deep learning. In Proceedings of the 14th USENIX Conference on Operating Systems Design and Implementation (OSDI’20) Article 49. USENIX Association, USA, 17 pages. isbn: 978-1-939133-19-9. find
  58. Size Zheng et al. 2022. AMOS: enabling automatic mapping for tensor computations on spatial accelerators with hardware abstraction. In Proceedings of the 49th Annual International Symposium on Computer Architecture (ISCA ’22). Association for Computing Machinery, New York, New York, 874–887. isbn: 9… — [doi:10.1145/347049](https://doi.org/10.1145/347049) find
  59. Hongyu Zhu et al. 2022. ROLLER: fast and efficient tensor compilation for deep learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22). USENIX Association, Carlsbad, CA, (July 2022), 233–248. isbn: 978-1-939133-28-1. https://www.usenix.org/conference/osdi22/prese… ↗
  60. Mikhail Zolotukhin. 2021. NNC walkthrough: how PyTorch ops get fused. https://dev-discuss.pytorch.org/t/nnc-walkthrough-how-pyt orch-ops-get-fused/125. (2021). ↗

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/ansel-2024-pytorch-2-faster-machine-learning/