PyTorch 2.15 tackles the 'one bad GPU' problem - The Register

Machine learning administrators can now breathe easier, as a single faulty graphics processing unit (GPU) will no longer force them to restart an entire training job from scratch.

The latest release of PyTorch, version 2.15, includes a completely rewritten backend for Nvidia's Collective Communications Library (NCCL). This new backend is designed to improve fault diagnosis and recovery when multiple GPUs work together.

The previous PyTorch backend for NCCL, known as ProcessGroupNCCL, was described by core developer Tristan Rice as 'drowning in technical debt'.

NCCL, pronounced 'Nickel', provides PyTorch with the essential communication primitives that allow GPUs to exchange data within a single-node or across multi-node clusters.

This data exchange is vital for running inference on large neural networks and for other tasks in building and maintaining large language models (LLMs).

Simplified backend

The old NCCL interface was chaotic, spawning dozens of communicators per process group, which led 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.

Diagnostic tools

A 'flight recorder' has also been added to give admins a much 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 that the cleaned-up NCCL backend, also called ProcessGroupNCCL, will give developers a blank slate, making it easier to add new features.

Future and support

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.

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.

The update also includes 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.