We are excited to announce the release of PyTorch® 2.3 (release note)! PyTorch 2.3 offers support for user-defined Triton kernels in torch.compile, allowing for users to migrate their own Triton kernels from eager without experiencing performance regressions or graph breaks. Tensor Parallelism improves the experience for training Large Language Models using native PyTorch functions, which has been validated on training runs for 100B parameter models. As well, semi-structured sparsity implements semi-structured sparsity as a Tensor subclass, with observed speedups of up to 1.6 over dense matrix multiplication.
This release is composed of 3393 commits and 426 contributors since PyTorch 2.2. We want to sincerely thank our dedicated community for your contributions. As always, we encourage you to try these out and report any issues as we improve 2.3. More information about how to get started with the PyTorch 2-series can be found at our Getting Started page.
| Beta | Prototype | Performance Improvements |
| User-defined Triton kernels in torch.compile | torch.export adds new API to specify dynamic_shapes | Weight-Only-Quantization introduced into Inductor CPU backend |
| Tensor parallelism within PyTorch Distributed | Asynchronous checkpoint generation | |
| Support for semi-structured sparsity |
*To see a full list of public feature submissions click here.
Beta Features
[Beta] Support for User-defined Triton kernels in torch.compile
Allows for PyTorch code that contains triton kernels to be executed natively using torch.compile. This enables users to migrate code containing triton kernels from eager PyTorch to torch.compile without running into performance regressions or graph breaks. Native support also creates an opportunity for Torch Inductor to precompile the user-defined Triton kernel as well as better organize code around the Triton kernel allowing for further optimizations.
You can find more information about how to utilize user defined Triton kernels in torch.compile within this tutorial.
[Beta] Tensor Parallelism introduces more efficient ways to train LLMs
The Tensor Parallel API facilitates various tensor manipulations across GPUs/hosts and integrates with FSDP for 2D Parallelism (Tensor parallelism across devices + Data Parallelism across hosts). It also offers a low-level API for constructing higher-level Tensor parallel APIs. This API has been validated to support the training of transformer models with over 100 billion parameters.
You can find more information on how to utilize this within your workflows within this tutorial.
[Beta] Semi-structured sparsity provides users with a way to take advantage of accelerated sparse inference and memory savings
torch.sparse.SparseSemiStructuredTensor implements semi-structured sparsity as a Tensor subclass, which have observed speedups of up to 1.6 over dense matrix multiplication.
In particular it adds:
- Additional support for quantization composability (mixed dtype, dequant fusion)
- Updated cuSPARSELt and CUTLASS kernels
- torch.compile support
You can find more information on how to take advantage of semi-structured sparsity here.
Prototype Features
[PROTOTYPE] torch.export adds new API to specify dynamic_shapes
You can now use torch.export.Dim to better represent dynamic shapes by enabling developers to specify ranges (min and max values) that can be reused across different input dimensions that are constrained to be equal.
To learn more about torch.export.Dim as well as how it can be used to express more interesting relationships (such as linear arithmetic expressions) check out the tutorial here.
[PROTOTYPE] Asynchronous checkpoint generation
Asynchronous checkpoint generation allows users to continue their training loops while checkpoints are being generated, essentially offloading much of the checkpointing cost.
You can find out how to utilize this within your own workflows with this example.
Performance Improvements
[PROTOTYPE] Weight-Only-Quantization introduced into Inductor CPU backend
PyTorch 2.3 enhances LLM inference performance on torch inductor CPU backend. The project gpt-fast offers a simple and efficient PyTorch native acceleration for transformer text generation with torch.compile. Prior to 2.3 only CUDA devices were supported and this feature enables the CPU counterpart by providing highly optimized kernels for the int4 and int8 weight only quantization Linear.
For more information / how to utilize this feature please refer to the gpt-fast README.