PyTorch 2.15 tackles the 'one bad GPU' problem

Technical debt shoveled aside to make large-scale AI training more resilient

ML admins will be happy to learn that one bad GPU may no longer mean starting an entire training job from scratch. The latest 2.15 release of PyTorch includes a rewritten backend for Nvidia's Collective Communications Library (NCCL), designed to improve fault diagnosis and recovery when GPUs work together.

This latest feature shows how its creators continue to advance the software from being a tool for researchers to being a cornerstone in production AI work. 

The interface to the networking library in PyTorch needed a rewrite – it was drowning in technical debt, said PyTorch core developer Tristan Rice, a Meta software engineer specializing in distributed computing, during a live YouTube livestream on Thursday. 

REG AD

NCCL is pronounced "Nickel" and stands for the Nvidia Collective Communications Library. It provides PyTorch with a set of inter-GPU communication primitives, allowing GPUs to communicate within a single-node or across multi-node clusters, which is vital for running inference on large neural networks and for other tasks in building and maintaining LLMs.

REG AD

The old NCCL interface was too chaotic

The previous PyTorch backend for NCCL, called ProcessGroupNCCL, could spawn dozens of communicators per process group, leading to intermittent deadlocks and unpredictable performance. 

The new version of the NCCL backend comes with a strict rule: One process group gets only one communicator during initialization. By default, this rule eliminates crosstalk between multiple communicators and all the complexities that entails. If a program requires parallelism, it will have to dispatch more process groups.  

The updated library also better handles stragglers by allowing for easier reconfiguration. If a GPU crashes or runs slowly, a third-party orchestrator can exclude the laggard and even assign a replacement worker. 

A “flight recorder” has also been added to give admins a better idea of what went wrong. Previously, if a cluster of GPUs stalled—perhaps from a timeout or deadlock—the user would just get a generic message. The new flight recorder captures all operations in a series of snapshots that can be inspected after a failure occurs.

Rice noted the cleaned-up NCCL backend, also called ProcessGroupNCCL, will give developers a blank slate, making it easier to add new features.

The legacy ProcessGroupNCCL is part of the 2.15 package as a rollback, but if the new NCCL backend works as expected, the older version will be removed in the 2.16 release, planned for early next year. 

Nvidia origins

REG AD

Why, you may ask, does a machine learning framework need a specialized communications library? Like NumPy, PyTorch supports numerical computing, but it also provides tools for building and training neural networks on CPUs and GPUs.

In these environments, workers must constantly exchange data such as tensors, gradients, activations, and weights. The standard TCP-based networking stack can introduce bottlenecks in large GPU clusters, so NCCL uses optimized communication paths to move data between GPUs while reducing CPU overhead.

Nvidia initially released NCCL at SC15 as a general-purpose accelerator, and it was adopted by the Facebook folks when they first built PyTorch in 2016. Their goal was to build a Python-based version of the Lua-based Torch ML library.

“Every major ML lab (OpenAI, DeepMind, Meta, Nvidia) uses PyTorch. It powers everything from diffusion models and LLMs to recommender systems and robotics. This widespread adoption isn’t accidental,” noted Will Stewart, co-founder of AI software provider Northflank, in a blog post.     

More support for Apple and Nvidia’s next GPUs

Other new features of PyTorch 2.15 include a new Python transport API for asynchronous one-way data movement – ideal for training jobs where weights are synchronized between trainers and generators – as well as parallel processing improvements for Apple gear and for supporting Nvidia’s next-generation Blackwell GPUs. 

Overall, the update consisted of 3,117 commits from 583 contributors. ®