Enhancing Goodput in Large-Scale LLM Training with Nonuniform Tensor Parallelism

Training LLMs at massive scale brings unique infrastructure challenges, especially as jobs span thousands of GPUs and run for extended periods. The longer these jobs run, the greater the likelihood of encountering unscheduled interruptions or resource fluctuations. Even infrequent device unavailability can have outsized effects on tightly interconnected clusters, resulting in slowdowns for a given training run.

For large-scale training jobs, elastically adapting the job to the number of available GPUs is a powerful method to improve Goodput. In the context of AI training, Goodput represents the critical measure of useful, convergence-driving work completed, rather than just raw hardware throughput.

Effective methods for elastic scaling today include dropping a data replica, utilizing fast checkpoint-restarts, or swapping to hot spares. These methods enable LLM jobs to adapt to GPU availability changes while maintaining balance across the system. However, they also incur some amount of lost throughput and higher cost during the period that the training job is running in a reduced availability condition.

A recent paper on Nonuniform Tensor Parallelism (NTP) introduces a forward-looking, experimental framework that builds on these methods in a way that minimizes throughput overheads. Combined with potential dynamic power boosting to offset any performance loss, throughput remains steady, transforming interruptions into manageable and recoverable events.

NTP’s core contribution is its ability to sustain high Goodput by preventing transient device issues from stalling large, interconnected training jobs. By dynamically adjusting the tensor parallelism degree and intelligently overlapping the necessary data resharding, NTP minimizes lost time and computational effort.

This ensures the cluster spends a maximal amount of its operating time performing useful work, preserving the overall integrity and efficiency of the training run even as hardware conditions fluctuate.

Challenges with large-scale training

AI model training is a parallel undertaking, spanning thousands of GPUs. A common technique to parallelize these workloads is tensor parallelism (TP), where the layers of the neural network are split across a tightly-coupled group of GPUs. The number of GPUs in this group coincides with the scale-up domain, which is interconnected with high-speed interconnects like NVIDIA NVLink. On NVIDIA Blackwell and NVIDIA Blackwell Ultra systems, NVLink connects up to 72 GPUs at 1,800 GB/s, supporting all-to-all communication within a single hop. 

In a typical scenario, a frontier LLM is trained across a cluster of racks, each housing multiple servers. A single rack of servers forms the scale-up domain and makes up a single TP group. Data parallelism (DP) then replicates the model across multiple such scale-up domains, each processing a different batch of data. A change in GPU status within a scale-up domain can affect the efficiency of that TP group. Because GPUs in the same TP group share tightly coupled computations, an issue with one device may reduce training performance or require temporary rebalancing to maintain progress. 

As data center architectures evolve to support larger scale-up domains, going from eight to 72 GPUs and beyond, maximizing the productive uptime of every healthy device becomes the key to achieving high Goodput.

The paper notes that rather than allowing localized, transient interruptions to dictate overall training throughput, systems can be designed to keep the vast majority of active GPUs continuously processing. By adapting to these hardware fluctuations, clusters can maintain the highly efficient resource use required to optimize Goodput during large-model training at scale.

Training usually involves a pipeline where each stage depends on the timely completion of prior steps. If one GPU experiences delays within a TP group, it can slow synchronization or processing across that group, causing temporary stalls or reduced throughput until the system recovers or redistributes the workload.

How NTP maintains training 

The fundamental idea behind NTP is to ensure a DP replica remains productive even during transient hardware interruptions. Upon resuming from the latest checkpoint, the model automatically adapts its configuration to the available hardware, maintaining partial functionality to keep the pipeline moving without completely losing that replica’s contribution.

The following dives deeper into how NTP achieves this resilience.

Dynamic TP degree adaptation

When a GPU within a scale-up domain experiences an interruption, the system identifies the affected group and automatically reconfigures its tensor parallelism to utilize only the remaining functional GPUs. For example, if a TP group of eight GPUs experiences one drop-out, it can dynamically switch to a TP degree of seven. That model shard continues its computations, preventing a complete loss of its contribution. The remaining GPUs within the group increase their individual workload, enabling the training job to maintain high Goodput and availability even when a subset of resources is experiencing issues.

Power boosting for performance compensation

Reducing the TP degree isn’t enough to maintain global throughput. A DP replica with fewer GPUs will inherently run slower, causing the entire DP system to stall, waiting for the slowest replica. To counteract this, the study proposes a rack design that incorporates improved electrical and thermal capabilities. This design enables power-boosting of the scale-up domains with reduced availability. By dynamically increasing the power supplied to the active GPUs in an affected domain, clock frequencies and computational throughput can be temporarily elevated.

The DP replica with the reduced TP degree can effectively catch up and keep pace with the other, fully functional replicas, preventing global synchronization bottlenecks and ensuring that the cluster’s Goodput remains highly optimized despite localized hardware fluctuations.

Efficient resharding

The dynamic adjustment of the TP degree requires an efficient mechanism for redistributing the model’s tensor shards among the remaining GPUs. NTP employs a clever resharding technique overlapped with other computational stages.

By performing this resharding during the backward computation and parameter synchronization phases, the overhead introduced to healthy replicas is minimized, often to less than 1%. This careful scheduling maximizes compute efficiency, seamlessly sustaining optimal Goodput without the adaptation mechanism itself becoming a performance bottleneck.

NTP builds a more resilient path to scaled AI training

This work underscores the critical importance of co-designing hardware and software to address the challenges inherent in large-scale AI training. The integration of NTP with advanced rack designs, which provide the electrical and thermal headroom needed for dynamic power-boosting, serves as a prime example of how thoughtful hardware innovations can profoundly complement sophisticated software solutions.

This symbiotic relationship between hardware and software is essential for unlocking higher levels of performance, efficiency, steady Goodput and resilience in the next generation of AI systems. As a forward-looking, experimental feature, NTP demonstrates what’s possible when resilience is baked directly into the parallelism strategy.

Building on this foundation, research is already underway to extend these concepts to Nonuniform Expert Parallelism (NEP), optimizing resilience for Mixture-of-Experts (MoE) models where standard tensor parallelism is less ideal.

Check out production-ready fault tolerance and resiliency features available in NVIDIA Resiliency Extension (NVRx) or the Nonuniform Tensor Parallelism Readme to learn more about its recent addition to the developer branch of  NVIDIA Megatron Core.