PyTorch 2.15 tackles the 'one bad GPU' problem

ML administrators can breathe easier knowing that a faulty GPU will no longer require restarting an entire training job, thanks to PyTorch's 2.15 release. The updated backend for Nvidia's Collective Communications Library (NCCL) aims to enhance fault diagnosis and recovery when GPUs work together. This feature highlights PyTorch's evolution from a research tool to a critical component in production AI work.

PyTorch core developer Tristan Rice explained that the network library in PyTorch needed a rewrite due to technical debt. NCCL, pronounced "Nickel," is Nvidia's Collective Communications Library and provides PyTorch with inter-GPU communication primitives for running inference on large neural networks and building and maintaining large language models (LLMs).

The previous PyTorch backend for NCCL, ProcessGroupNCCL, could spawn numerous communicators per process group, leading to deadlocks and unpredictable performance. The new version enforces a strict rule: one process group gets only one communicator during initialization, which eliminates crosstalk and the complexities that come with it.

If more parallelism is required, additional process groups must be dispatched. The updated library also handles stragglers more effectively by allowing for easier reconfiguration and enabling a third-party orchestrator to exclude a slow GPU and assign a replacement worker. A "flight recorder" has been added to provide administrators with a better understanding of what caused a failure.

Previously, a generic message was displayed when a cluster of GPUs stalled. The flight recorder captures all operations in snapshots that can be inspected after a failure occurs. The cleaned-up NCCL backend, called ProcessGroupNCCL, gives developers a fresh start to add new features. The legacy ProcessGroupNCCL is included in the 2.15 package as a rollback, but it will be removed in the 2.16 release, scheduled for early next year.

NCCL's origins trace back to Nvidia, which initially released it at SC15 as a general-purpose accelerator. It was later adopted by Facebook when they first built PyTorch in 2016. The goal was to create a Python-based version of the Lua-based Torch ML library. PyTorch's widespread adoption is no coincidence, as it powers diffusion models, LLMs, recommender systems, and robotics.

PyTorch 2.15 also introduces a new Python transport API for asynchronous one-way data movement, parallel processing improvements for Apple devices, and support for Nvidia's upcoming Blackwell GPUs. The update comprises 3,117 commits from 583 contributors.