PyTorch 1.9: Powering the Next Generation of Distributed and Scientific Workloads

PyTorch has rapidly become one of the most popular deep learning frameworks, loved by researchers and engineers for its simplicity, flexibility and performance. With the release of PyTorch 1.9, the PyTorch team at Facebook has delivered a massive set of new features and improvements to expand its capabilities for distributed training and scientific computing. Let‘s dive into what‘s new and how it can supercharge your PyTorch projects.

Distributed Training Gets Elastic

One of the headline features of PyTorch 1.9 is the incorporation of TorchElastic into the core library. TorchElastic is a framework to enable fault-tolerant and elastic distributed training on platforms like Kubernetes. It provides primitives and interfaces for jobs to be dynamically scaled up or down based on resource availability.

This is a game changer for anyone running large-scale PyTorch training jobs in production environments. Instead of needing to handle interruptions or hardware failures manually, TorchElastic can automatically recover and continue training from the latest checkpoint. It also allows you to efficiently utilize resources by scaling jobs to fill available capacity.

TorchElastic was previously used by libraries like PyTorch Lightning and Kubernetes CRD. By bringing it into core PyTorch, it will be easier than ever to build robust distributed training pipelines. The PyTorch team is also investing in improved documentation and tutorials for distributed training best practices.

Along with TorchElastic, PyTorch RPC has gained support for CUDA-accelerated communication. PyTorch RPC enables you to easily build distributed models by providing primitives like remote procedure calls and distributed autograd. With CUDA RPC, tensors can be directly communicated peer-to-peer between GPUs, leading to major performance gains over CPU-based communication.

Here‘s an example of using CUDA RPC to run a distributed model:

import torch
import torch.distributed.rpc as rpc

@torch.jit.script
def my_rpc_function(t1, t2):
  return torch.add(t1, t2)

# On worker 0
t1 = torch.rand(3, 3).cuda(0)
t2 = torch.rand(3, 3).cuda(1)
res = rpc.rpc_sync("worker1", my_rpc_function, (t1, t2))

# On worker 1 
res = res.cuda(1)

In this example, my_rpc_function is executed on a remote worker, taking CUDA tensors located on different devices as arguments. The result is returned as a CUDA tensor on the invoking worker. With CUDA RPC, the tensors are communicated directly between GPUs without costly CPU transfers.

Accelerating Scientific Computing

While PyTorch is most well known as a deep learning library, it has also gained traction with researchers using it for scientific and numerical computing. PyTorch 1.9 doubles down on this use case in several ways.

First, the torch.linalg module for linear algebra has moved from beta to stable. torch.linalg provides GPU-accelerated and differentiable versions of common linear algebra operations like matrix decompositions and solvers. It aims for close compatibility with NumPy‘s linalg module, while extending it to support autograd and hardware acceleration.

Here‘s an example of using torch.linalg to compute the Cholesky decomposition of a matrix:

import torch

A = torch.randn(3, 3, dtype=torch.float64)
A = torch.matmul(A, A.t())  # make symmetric positive definite
L = torch.linalg.cholesky(A)
print(L)

This code generates a random symmetric positive definite matrix A and computes its Cholesky decomposition using torch.linalg.cholesky. The result L is the lower-triangular matrix satisfying A = LL^T. By computing this on the GPU with autograd support, it becomes much easier to use PyTorch for applications like Bayesian inference and physical simulations.

Another major update for scientific computing is expanded support for complex numbers and complex autograd. PyTorch 1.9 extends support for complex-valued tensors and complex-variable autodiff to over 98% of operators. This means you can perform forward and backward passes on complex-valued functions and optimize complex parameters. Complex numbers are essential for domains like signal processing, electromagnetism, and quantum machine learning.

As an example, here‘s how you can compute the gradient of a complex-valued function:

import torch

def f(z):
  return z.conj() * z

z = torch.tensor([3+4j, 2-5j], requires_grad=True)
w = f(z)
print(w)  # tensor([25, 29])

w.sum().backward()
print(z.grad)  # tensor([6+8j, 4-10j])

The function f computes the squared magnitude of each element of the complex vector z. By calling backward() on the sum of the result, we compute the gradient of f with respect to z. The autograd engine correctly handles the conjugate terms and accumulates the complex gradient in z.grad.

Deeper Insights with PyTorch Profiler

To get the best performance out of your PyTorch models, it‘s essential to be able to measure and analyze their runtime behavior. PyTorch 1.9 significantly expands the capabilities of the new PyTorch Profiler tool, which provides deep introspection of the performance of PyTorch operators.

Some of the key new features of the PyTorch Profiler include:

• Support for profiling distributed workloads, including data collected from all workers in a distributed job
• Improved trace analysis, with the ability to compare traces across workers or iterations
• Visualization of GPU kernel activity and memory usage over time
• New performance insights, like the top operators by FLOP count and GPU utilization statistics

Here‘s an example of using the PyTorch Profiler to analyze the performance of a simple model:

import torch
from torch.profiler import profile, record_function, ProfilerActivity

model = torch.nn.Linear(1024, 1024)

with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    with record_function("model_forward"):
        model(torch.randn(512, 1024))

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

This code profiles the forward pass of a linear layer, recording CPU activity and tensor shapes. The printed output shows a table of the top 10 operators by total CPU time, along with their input shapes. This makes it easy to identify performance bottlenecks at the operator level.

With the new distributed profiling support, you can get this level of insight across all machines in a distributed job. The profiler can aggregate traces from each worker, enabling you to spot communication bottlenecks, stragglers, and other distributed performance issues.

Deploying to the Edge with Mobile Interpreter

As deep learning moves to the edge, it‘s becoming increasingly important to be able to efficiently deploy models to mobile and embedded devices. PyTorch 1.9 makes this easier with the new mobile interpreter feature.

The mobile interpreter is a lightweight version of the PyTorch runtime designed for edge deployment. It executes models in a high-performance eager mode, with a focus on minimizing binary size and memory usage. This enables you to deploy complex models to resource-constrained devices without needing to first convert them to a special mobile format.

Along with the mobile interpreter, PyTorch 1.9 adds support for the TorchVision library on mobile. You can now easily integrate pre-trained TorchVision models into iOS and Android apps. This opens up new possibilities for on-device tasks like object detection, segmentation, and pose estimation.

Here‘s an example of using a pre-trained MobileNetV3 model for image classification in an Android app:

Module module = Module.load(assetFilePath(this, "mobilenet_v3.pt"));
float[] mean = {0.485f, 0.456f, 0.406f};
float[] std = {0.229f, 0.224f, 0.225f};

Bitmap bitmap = BitmapFactory.decodeStream(getAssets().open("image.jpg"));
Tensor input = TensorImageUtils.bitmapToFloat32Tensor(bitmap,
    NO_MEAN_RGB, mean, std);

Tensor output = module.forward(IValue.from(input)).toTensor();
float[] scores = output.getDataAsFloatArray();

This code snippet loads a pre-trained MobileNetV3 model, preprocesses an input image, and runs the model to get predicted class scores. With the new mobile interpreter, you can execute this model directly on an Android device without any additional conversion steps.

Embracing the Future of PyTorch

PyTorch 1.9 represents a major step forward for the library, with new features and improvements that touch almost every aspect of the PyTorch ecosystem. From elastic distributed training with TorchElastic to effortless mobile deployment with the mobile interpreter, this release enables PyTorch users to tackle new challenges and scale up their workloads like never before.

At the same time, the PyTorch team has remained true to the core values of simplicity and usability that have made PyTorch so beloved. Despite the wealth of updates, the library feels more polished and stable than ever, with clear documentation and intuitive APIs.

Whether you‘re a researcher pushing the boundaries of scientific computing, an engineer deploying models to production, or a student learning the ropes of deep learning, PyTorch 1.9 has something to offer you. As the PyTorch community continues to grow and thrive, I‘m excited to see what new breakthroughs and applications will be powered by this incredible tool.