PyTorch 2.0 Released, Introduces Model Compiler For Free Speedups
PyTorch 2.0 has been released, introducing a new model compiler for model speedups at no extra trouble, and a major move away from C++ and towards a full Python backend.
Created on December 5|Last edited on December 5
Comment
Early into the PyTorch Conference 2022 presentation on Friday, PyTorch 2.0 was announced. Currently available in preview, this newest version of PyTorch jumps from 1.X to 2.X thanks to a major move from C++ code to Python code in the backend, and thanks to several new technologies, introduces a model compiler utility providing considerable speedups for the majority of existing models.
PyTorch 2.0 is currently available in preview/nightly builds, with plans for a stable release in early 2023. You can install it for yourself following the instructions on the PyTorch website like usual (make sure to install the preview/nightly build). Also check out a tutorial for using PyTorch 2.0's new features in this blog post.
Given that the new additions to PyTorch are entirely optional, PyTorch 2.0 is fully backwards compatible with earlier versions of PyTorch so upgrading should pose no issues to already-existing projects.
PyTorch 2.0's new compiler function
The main thing most of us will care about is the new compiler function, torch.compile(). The new compiler is possible thanks to four new key technologies developed for PyTorch 2.0 (TorchDynamo, AOTAutograd, PrimTorch, and TorchInductor) which each do their own part in taking standard PyTorch code and squishing it down into a significantly more optimized representation that can be run more efficiently. Full details on the new tech are available on PyTorch 2.0's introductory post.
Using the new compiler is as simple as adding a single line call to torch.compile() with your model, no need to change your code or model structure anywhere (it can even compute dynamic shapes, though some optimizations cannot be applied). Different modes let you optimize for compile speed vs memory cost, etc.
Performance improvements with torch.compile() were validated (against eager mode) on a variety of model types from Hugging Face Transformers, TIMM, and TorchBench. From the 163 tests, it worked 93% of the time, with average 43% training speedups on NVIDIA A100 GPUs. Though, with consumer-grade GPUS (such as the NVIDIA 3090), speed increases were less significant.
Making contribution easier
Because PyTorch is moving away from a C++ backend and towards a fully Python codebase, contributors will no longer need to worry about knowing both languages or dealing with the ways they might interact. Instead, developers for PyTorch can stay totally in the realm of Python, keeping a lot more consistency across the board.
Find out more
Learn more about PyTorch 2.0, the compiler, and all the new underlying technologies in the introductory post as well as the tutorial walkthrough.
Add a comment
Tags: ML News
Iterate on AI agents and models faster. Try Weights & Biases today.