PyTorch/XLA 2.4
For deep learning academics and practitioners, the open-source PyTorch machine learning (ML) library and XLA ML compiler provide flexible, powerful model training, fine-tuning, and serving. The PyTorch/XLA team is happy to announce the release of PyTorch/XLA 2.4 today. This version includes several noteworthy enhancements to address issues raised by developers and builds on the previous release. Here, we go over a few of the most recent additions that facilitate using PyTorch/XLA:
- Pallas, a proprietary kernel language that supports GPUs and TPUs, has been improved.
- Fresh calls to the API
- The “eager mode” experiment is introduced.
- The TPU command line interface has been updated.
Pallas improvements
Although the XLA compiler can optimize your current models, there are situations in which bespoke kernel code can provide model authors with superior performance. Pallas is a bespoke kernel language that supports TPU and GPUs, so instead of requiring you to use a more complex and lower-level language like C++, you can write more performant code in Python that is closer to the hardware. Pallas is comparable to the Triton library, but it makes porting your model from one machine learning accelerator to another easier because it runs on both TPUs and GPUs.
The latest version of PyTorch/XLA 2.4 brings improvements to Pallas’ functionality and user experience.
- Flash Attention is now completely integrated with PyTorch autograd, allowing for automatic gradient calculation.
- Integrated assistance for Paged Focus on Inference.
- Support for group matrix multiplication using Mega blocks’ block sparse kernels as an Autograd function, eliminating the requirement for backpropagation to be done manually.
API modifications
A few new calls are included in PyTorch/XLA 2.4 to facilitate integration with your current PyTorch workflow, such as:
device = torch_xla.device()
And now you can call torch_xla.sync() in place of having to do xm.mark_step(). The developer workflow is enhanced and the process of converting your code to PyTorch/XLA is made simpler by these enhancements.
import torch_xla.core.xla_model as xm
device = xm.device()
Try out the eager mode
If you’ve worked with PyTorch/XLA for any length of time, you are aware of the term “lazily executed” models. This implies that before models are sent to be performed on the XLA device target hardware, PyTorch/XLA 2.4 builds the compute graph of operation. Operations are compiled and then instantly carried out on the target hardware with the new eager mode.
The drawback of this feature is that, because each instruction is not conveyed to the TPU immediately by default, TPUs themselves lack a real eager mode. In order to compel the compilation and execution, Google cloud add a “mark step” call to each PyTorch action on TPUs. As a result, eager mode functions, albeit as an emulator rather than a built-in feature.
With this release, Google cloud want for eager mode to be used in your local surroundings rather than in your production environment. Eager mode is intended to simplify local model debugging on your PCs without requiring you to deploy it to a broader fleet of devices, as is the case with most production systems.
CLI to view Cloud TPU information
The nvidia-smi tool, which you can use to troubleshoot your GPU workloads, determine which cores are being used, and check how much memory a particular workload is consuming, may be familiar to you if you’ve previously used Nvidia GPUs. Additionally, a comparable command line tool has been developed for Cloud TPUs that facilitates the retrieval of device and utilization data.
Start using PyTorch/XLA 2.4 right now
The best aspect is that your current code is still compatible with PyTorch/XLA 2.4, despite the fact that it has certain API changes. Additionally, the new API methods will make your future development processes easier. What are you waiting for? Try the most recent version.
The PyTorch logo: Key Features
Prepared for Production
Use TorchScript to switch between eager and graph modes with ease, and TorchServe to quicken the production process.
Dispersed Instruction
The torch.distributed backend enables scalable distributed training and performance optimisation in research and production.
Sturdy Ecosystem
A plethora of tools and frameworks complement PyTorch/XLA 2.4 Improved Pallas and developer experience, “eager mode” and facilitate its development in computer vision, natural language processing, and other fields.
Cloud Assistance
Major cloud platforms support PyTorch well, enabling easy scaling and frictionless development.
XLA Features
Accelerated Linear Algebra, or XLA, is an open-source machine learning compiler. Models from well-known frameworks like PyTorch, TensorFlow, and JAX are imported into the XLA compiler, which then optimises them for high-performance execution on a variety of hardware platforms, including GPUs, CPUs, and ML accelerators. For instance, employing XLA with 8 Volta V100 GPUs produced a ~7x performance boost and ~5x batch-size improvement over the same GPUs without XLA in a BERT MLPerf submission.
Leading ML hardware and software companies, including as Alibaba, Amazon Web Services, AMD, Apple, Arm, Google, Intel, Meta, and NVIDIA, are working together to develop XLA as part of the OpenXLA initiative.
Principal advantages
Construct anywhere: Prominent machine learning frameworks like TensorFlow, PyTorch, and JAX have already incorporated XLA.
Run anywhere: It has pluggable infrastructure to offer support for other backends, such as GPUs, CPUs, and ML accelerators, among other backends.
Optimise and scale performance: It makes use of automated partitioning for model parallelism and production-tested optimization stages to maximize a model’s performance.
Reduce complexity: By utilising MLIR, it combines the greatest features into a single compiler toolchain, saving you from having to handle a variety of domain-specific compilers.
Future-ready: XLA is an open-source project that was developed in conjunction with top ML software and hardware providers. Its goal is to be the industry leader in machine learning.
