Radius: Range-based Gradient Sparsity for Large Foundation Model Pre-training

Radius: Range-based Gradient Sparsity for Large Foundation Model Pre-training

Speakers: Mingkai Zheng, PhD Student, Rutgers University; Zhao Zhang

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

The pre-training of large language models (LLMs) has become a cornerstone of modern AI, yet it remains an immensely resource-intensive and time-consuming endeavor. As models scale to billions or even trillions of parameters, the computational and communication overheads associated with distributed training frameworks become dominant bottlenecks, particularly in data-parallel (DP) setups. This talk introduces Radius, a novel range-based gradient sparsity algorithm and system designed to significantly accelerate the pre-training of large foundation models by intelligently reducing the communication volume without sacrificing model quality.

Presented by Mingkai Zheng from Rutgers University, Radius addresses the critical challenge of communication overhead in data-parallel training, which often surpasses computation time on high-performance computing clusters. The core innovation lies in its ability to leverage the temporal stability of important gradient values, enabling periodic selection of top-k indices and subsequent gradient compression. By doing so, Radius offers a practical and effective solution to a pervasive problem in large-scale machine learning, promising substantial speedups and improved scalability for the next generation of foundation models.

The significance of Radius extends to anyone involved in the development and deployment of large AI models, from researchers pushing the boundaries of model scale to infrastructure teams managing vast GPU clusters. By mitigating the communication bottleneck, Radius not only reduces the wall-clock time required for pre-training but also enhances the overall efficiency and scalability of distributed training systems. Its demonstrated ability to achieve considerable speedups while maintaining or even improving model performance on downstream tasks positions it as a vital advancement in the quest for more efficient and accessible large-scale AI training.

Background

▶ Watch: Introduction and communication bottleneck problem (0:00)

The era of large foundation models, characterized by billions of parameters, has brought unprecedented capabilities but also formidable challenges, primarily concerning the sheer computational and communication resources required for their pre-training. Data parallelism (DP) is a widely adopted technique to distribute the training workload across multiple workers, each processing a partition of the mini-batch. In this setup, after each worker completes its forward and backward passes, the gradients computed on each local batch must be aggregated (typically averaged) across all workers. This aggregation step, usually performed via an all-reduce operation, ensures that all workers maintain a synchronized view of the model parameters for the next optimization step.

However, as the number of workers and model size increase, the communication overhead associated with gradient synchronization often becomes the dominant bottleneck, eclipsing the actual computation time. For instance, profiling the pre-training of a GPT-20B model with a global batch size of 512 on 128 A100 GPUs on the Perlmutter supercomputer (equipped with NVLink and Slingshot 11 inter-node communication) revealed that DP communication consumed more time than computation. This imbalance necessitates effective strategies to reduce the volume of data transmitted during gradient synchronization.

Several methods have been proposed to address this communication bottleneck. One common approach is top-k gradient sparsity, where only the gradient values with the largest magnitudes (top-k) are transmitted. While intuitively appealing, this method faces significant hurdles. Because each worker computes slightly different gradients, their respective top-k indices are generally not identical. Consequently, to correctly average the sparsified gradients, workers must exchange both the top-k values and their corresponding indices, typically using an all-gather operation. This all-gather approach introduces substantial memory overhead, as each worker must buffer results from all other workers, severely limiting scalability. Furthermore, when many workers are involved, the collection of top-k indices from all workers can lead to a "filling effect," where the union of all workers' top-k indices becomes nearly dense, negating the benefits of sparsity.

Another class of methods involves low-rank approximation, which seeks to represent the gradient with a lower-dimensional structure. Techniques like singular value decomposition (SVD) can be used, but in practice, PowerSGD is often favored due to its lower computational overhead. PowerSGD approximates the gradient with a low-rank representation that is structured, allowing for the direct use of an all-reduce operation for aggregation, thus avoiding the all-gather memory issue. However, compressing gradients to very high degrees of sparsity using low-rank methods can introduce significant computational overhead, potentially shifting the bottleneck from communication to computation. For example, using PowerSGD with rank 128 on 128 A100 GPUs was shown to incur computation overhead even higher than the baseline communication time, undermining its utility for aggressive compression.

The speakers' critical observation, which underpins Radius, stems from an analysis of GPT-3 355M model pre-training. They found that after 15-20% of the total pre-training steps, the top-k gradient values exhibit temporal stability. Specifically, examining the distribution of top 1% gradient values across thousands of iterations (e.g., 80,000 to 84,000) revealed that certain indices consistently held high-magnitude gradient values. This stability suggests that the "important" gradient components, i.e., those that contribute most significantly to the model update, do not change drastically from one iteration to the next. This insight sparked the question: can this temporal stability be leveraged to reduce the computation and memory overheads inherent in traditional top-k gradient sparsity methods? Radius is the answer to this question, designed to capitalize on this inherent property of gradients in large foundation models.

Key Findings

▶ Watch: Key observation: temporal stability of top-k gradients (3:00)

Radius presents several pivotal findings that collectively demonstrate its effectiveness in accelerating large foundation model pre-training while preserving or enhancing model quality:

1. Amortized Communication Overhead: Radius successfully overcomes the limitations of traditional top-k gradient sparsity by leveraging the temporal stability of important gradient values. By performing top-k index selection periodically (every T iterations) rather than every iteration, it significantly amortizes the computational overhead associated with identifying the most salient gradients. This approach drastically reduces the overall cost of sparsity.

1. Elimination of All-Gather and Filling Effect: A critical innovation in Radius is the selection of top-k indices after the all-reduce operation. This ensures that all workers operate on the same set of indices, eliminating the need for expensive all-gather operations that cause high memory overhead and prevent scalability. Furthermore, this design inherently avoids the "filling effect," where the union of different workers' top-k indices becomes dense, thereby preserving the benefits of sparsity even with a large number of workers.

1. Significant Wall-Clock Time Speedup: When pre-training a GPT-20B model, Radius achieved a 90% speedup on wall-clock time compared to a dense gradient baseline, using a gradient density level of 40% (i.e., transmitting 40% of gradient values) and a resampling period of 200 iterations. This substantial acceleration was accomplished without sacrificing the final training and validation perplexity, which remained approximately the same as the baseline.

1. Improved Iteration Efficiency: A detailed breakdown of computation and communication time showed that Radius, with a 40% density and 200-iteration resampling period on 64 A100 GPUs, compressed communication time by 47% and accelerated each training iteration by 21%. This highlights Radius's effectiveness in directly tackling the communication bottleneck.

1. Enhanced Strong Scaling Efficiency: Radius demonstrated superior strong scaling efficiency. When scaling from 8 to 128 A100 GPUs for GPT-20B pre-training, Radius (40% density, T=200) maintained a strong scaling efficiency of 43% at 128 GPUs, whereas the baseline method only achieved 33%. This indicates Radius's ability to better utilize increased hardware resources.

1. Comparable or Superior Downstream Task Performance: Models pre-trained with Radius (40% density, T=200) not only matched the perplexity of the baseline but also performed better on several zero-shot downstream tasks, including LAMBADA, MathQA, and PIQA. This validates that the gradient sparsity introduced by Radius does not degrade the model's generalization capabilities; in some cases, it may even act as a regularization, leading to improved performance.

1. Robustness Against Naive Top-K: Comparative evaluations against a "naive top-k" method (which sparsifies at every iteration and uses all-gather) revealed that Radius consistently followed the baseline validation perplexity curve. In contrast, the naive top-k method showed a clear trend of divergence, especially at higher sparsity rates (e.g., 40% or 50%), underscoring Radius's stability and reliability.

These key findings establish Radius as a robust and highly effective solution for accelerating the pre-training of large foundation models, offering significant performance gains and improved scalability without compromising the quality of the resulting models.

Technical Deep Dive

▶ Watch: Radius's core design principles explained (4:20)

Radius's design is predicated on two core technical ideas that directly address the limitations of prior gradient sparsity methods: leveraging temporal stability for periodic index selection and performing index selection after gradient aggregation to ensure global consistency.

1. Periodic Top-K Indices Selection

The first key idea is to exploit the observed temporal stability of top-k gradient values. Instead of re-selecting the top-k indices in every training iteration, Radius performs this computationally intensive operation only once every T iterations. During the T-1 intermediate iterations, all workers utilize the same fixed set of top-k indices identified in the most recent selection step. Only the gradient values corresponding to these pre-selected indices are then transmitted and aggregated.

This strategy effectively amortizes the computation overhead of top-k index selection over T iterations. The speaker's experiments suggest that a default value of T=200 works well for large language model pre-training, providing a significant reduction in overhead without negatively impacting convergence or quality. This hyperparameter T allows for a tunable trade-off between the freshness of selected indices and the computational savings.

2. Top-K Indices Selection After All-Reduce

The second, and perhaps most crucial, innovation addresses the scalability and memory issues of traditional top-k sparsity. In conventional approaches, each worker selects its top-k gradients locally before aggregation. Since local gradients differ, their top-k indices also differ, necessitating an all-gather operation to collect all indices and values. This leads to high memory consumption and the "filling effect" (where the union of all top-k indices becomes dense).

Radius circumvents this by performing the top-k indices selection after the all-reduce operation. Because the average of the gradients (after all-reduce) is identical across all workers, the top-k indices derived from this averaged gradient will also be identical on every worker. This eliminates the need for all-gather; workers can independently compute the same set of top-k indices post-all-reduce and then apply them consistently. This design choice directly enables the use of efficient all-reduce for subsequent gradient aggregation, avoiding the memory overhead and filling effect.

High-Level Design and Algorithm Flow

The Radius algorithm operates in two distinct phases, alternating between a "resampling" step (when current_step % T == 0) and "intermediate" steps.

A. Resampling Step (Every T iterations)

1. Residual Compensation: The process begins by adding the residual buffer to the current local gradient. The residual buffer is a well-known error feedback mechanism that accumulates the non-top-k gradient values from the previous T-1 intermediate steps. This ensures that the information from the "pruned" gradients is not lost but is gradually incorporated into subsequent updates, maintaining convergence stability and accuracy.
2. All-Reduce Compensated Gradient: The compensated gradient (local gradient + residual) is then subjected to an all-reduce operation to compute its average across all workers. This results in a globally consistent, averaged gradient.
3. Gradient Clipping and Correction: Standard optimizer operations, specifically gradient clipping and gradient correction, are applied to the averaged gradient. These operations are crucial for the AdamW optimizer (which is explicitly mentioned as the underlying optimizer) to update its momentum and variance buffers. Critically, these steps are performed before top-k selection because the intermediate states from AdamW are required to accurately perform the selection.
4. Top-K Indices Selection: After clipping and correction, the globally averaged and processed gradient is used to perform the top-k indices selection. Each worker independently identifies the same set of top-k and non-top-k indices based on the magnitude of the elements in this averaged gradient. These indices are then stored locally and will be used for the next T-1 iterations.
5. Model Parameter Update: The clipped and corrected gradient (using the selected top-k values) is used to update the model parameters.
6. Residual Buffer Reset: The residual buffer is reset to all zeros, ready to accumulate new non-top-k values in the upcoming intermediate steps.

B. Intermediate Steps (For T-1 iterations after resampling)

1. Top-K Value Extraction: Using the pre-selected top-k indices from the most recent resampling step, each worker extracts only the gradient values corresponding to these indices from its current local gradient. This forms a much smaller gradient vector.
2. All-Reduce Top-K Values: This smaller, sparse gradient vector (containing only the top-k values) is then aggregated across all workers using an all-reduce operation. This is efficient because the indices are known and consistent, and the data volume is significantly reduced.
3. Gradient Clipping and Correction: Similar to the resampling step, gradient clipping and correction are applied to the averaged top-k gradient values.
4. Model Parameter Update: The processed top-k gradient values are used to update the model parameters.
5. Residual Accumulation: For the non-top-k values (those not selected in the previous resampling step), they are accumulated into the residual buffer. This buffer will hold the "error" or unapplied gradient components, which will be re-introduced in the next resampling step, ensuring that no gradient information is permanently discarded and maintaining convergence.

This dual-mode operation allows Radius to achieve significant communication savings by transmitting only a fraction of the gradients most of the time, while the periodic resampling and error feedback mechanism ensure that model quality and convergence are preserved. The integration with standard optimizers like AdamW is seamless, requiring intermediate states for accurate index selection.

Experimental Setup & Results

▶ Watch: Experimental results: speedup and perplexity trade-offs (8:00)

The evaluation of Radius was comprehensive, focusing on its impact on training efficiency and model quality during large foundation model pre-training.

Experimental Setup

• Models:
• GPT-20B: The primary model used for performance and quality evaluation.
• GPT-3 355M: Used for the initial observation of temporal stability of gradient values (specifically, top 1% gradient values between iterations 80,000 and 84,000).
• Hardware:
• Perlmutter Supercomputer: Used for initial profiling and large-scale experiments. Each compute node featured four A100 GPUs connected by NVLink, with inter-node communication handled by Slingshot 11 (100 GB/s bandwidth).
• NVIDIA A100 GPUs: General hardware platform for performance measurements and comparisons. Experiments were conducted using configurations from 8 to 128 GPUs.
• Framework: The Megatron framework was used for pre-training, which, as noted by the speaker, did not employ techniques to hide communication latency with computation.
• Baseline: The main baseline was a standard data-parallel training setup using dense gradients (no compression).
• Comparison Methods:
• PowerSGD: A low-rank approximation method, specifically tested with rank 128 on 128 GPUs.
• Naive Top-K: A variant of top-k gradient sparsity that sparsifies gradients at every iteration and uses an all-gather operation to collect values and indices from all workers.
• Hyperparameters for Radius:
• Density Level (Sparsity Rate): Explored various levels, with 40% being a key sweet spot (meaning 40% of gradient values are transmitted). Also tested 10% and 50%.
• Resampling Period (T): Defaulted to 200 iterations for large language models.
• Optimization: AdamW optimizer was used.
• Global Batch Size: 512 for GPT-20B pre-training.

Headline Numbers & Key Results

1. Initial Communication Bottleneck: For GPT-20B pre-training with a global batch size of 512 on 128 A100 GPUs on Perlmutter, data parallelism communication time exceeded computation time. This underscored the necessity of gradient compression.
2. PowerSGD Limitations: When using PowerSGD with rank 128 to compress gradients on 128 A100 GPUs, its computation overhead was higher than the baseline communication time, indicating that for aggressive compression, low-rank methods can introduce new bottlenecks.
3. Wall-Clock Time Speedup:

• Radius, configured with a 40% density level and a resampling period (T) of 200, achieved a remarkable 90% speedup on wall-clock time for GPT-20B pre-training. Crucially, the final training and validation perplexity remained approximately the same as the dense gradient baseline.
• Pushing the density to 10% yielded higher speedup but resulted in significantly higher perplexity, indicating a trade-off.

1. Per-Iteration Acceleration:

• On 64 A100 GPUs, Radius (40% density, T=200) compressed communication time by 47%.
• This translated to a 21% acceleration for each training iteration.

1. Strong Scaling Efficiency:

• Radius demonstrated superior strong scaling efficiency, achieving 43% efficiency when using 128 A100 GPUs for GPT-20B pre-training.
• The baseline method, in contrast, only managed 33% efficiency at 128 GPUs, highlighting Radius's ability to scale more effectively.

1. Training Loss and Perplexity:

• When the density rate was set to 40% or higher, the final training loss error and perplexity between Radius and the dense baseline were very small, confirming that Radius maintains model quality.
• At a very high sparsity rate (e.g., 10% density), the error between Radius and the baseline became significant, indicating that there are limits to compression without further methodological refinements.

1. Downstream Task Performance (Zero-Shot):

• Models pre-trained with Radius (40% density, T=200) performed better on several downstream tasks compared to the baseline. These included LAMBADA, MathQA, and PIQA. This suggests that Radius not only preserves but can sometimes enhance generalization, possibly due to a regularization effect.

1. Comparison with Naive Top-K:

• Radius consistently tracked the baseline validation perplexity curve, even with 50% density.
• In stark contrast, the naive top-k method showed a clear and increasing trend of divergence from the baseline perplexity curve, particularly when sparsity rates were pushed higher (e.g., 40% or 1%). This validated Radius's stability and robustness over simpler top-k approaches.

These results unequivocally demonstrate that Radius effectively mitigates the communication bottleneck in large-scale LLM pre-training, delivering substantial speedups and improved scalability while maintaining or even improving the quality of the resulting models.

Practical Implications

▶ Watch: Breakdown of computation and communication time savings (9:00)

Radius presents a compelling set of practical implications for various stakeholders involved in the development, deployment, and operation of large-scale AI systems. Its innovations directly address some of the most pressing challenges in foundation model pre-training.

For practitioners and model builders, Radius offers a tangible path to significantly reduce the wall-clock time required for pre-training massive models like GPT-20B. A 90% speedup means that experiments that once took days could potentially complete in hours, drastically accelerating research cycles and enabling more iterative development. This directly translates to lower computational costs, as fewer GPU hours are consumed for the same amount of training progress. The ability to maintain or even improve model quality on downstream tasks (as seen with LAMBADA, MathQA, and PIQA) ensures that these speedups do not come at the expense of performance, making Radius a highly attractive solution for production-grade model development. Practitioners would need to carefully tune the density level and resampling period (T) to find the optimal balance between speedup and desired model quality, especially when pushing to very high sparsity rates.

For infrastructure teams managing large GPU clusters and supercomputers (like Perlmutter), Radius provides a crucial mechanism to improve the utilization and strong scaling efficiency of their expensive hardware. By reducing the communication bottleneck, Radius allows for more effective scaling to a larger number of GPUs, extracting more computational power from the distributed system. The observed improvement in strong scaling efficiency from 33% (baseline) to 43% (Radius) at 128 GPUs signifies better resource utilization, which is paramount in environments where GPU time is a premium commodity. This also alleviates pressure on network infrastructure, as the total data volume transmitted across inter-node connections is substantially reduced.

For deployers and researchers pushing the boundaries of model scale, Radius helps overcome a fundamental limitation of data parallelism. The "filling effect" and memory overhead associated with traditional all-gather-based top-k methods have historically limited the practical scalability of gradient sparsity. By selecting indices after all-reduce, Radius removes these barriers, making advanced gradient compression techniques viable for truly massive distributed training setups. This opens up possibilities for training even larger models on more distributed hardware than previously feasible, or for training existing models faster on current infrastructure.

Tradeoffs and Limitations: While highly effective, Radius is not without its considerations. The primary trade-off lies between the degree of gradient sparsity (density level) and the final model quality. While a 40% density level yielded excellent results, pushing sparsity to 10% resulted in a noticeable increase in perplexity. This means practitioners must carefully calibrate their desired speedup against acceptable quality degradation. The speaker noted that reducing this error for very high sparsity rates is an ongoing area of research. Additionally, the resampling period T is a hyperparameter that needs to be chosen; while a default of 200 iterations was suggested for LLMs, optimal values might vary for different models or tasks. The talk also clarified that the Megatron framework, as used, did not hide communication with computation, suggesting that even greater speedups might be possible if Radius were integrated into frameworks that do utilize such overlap.

In summary, Radius offers a robust, practical, and highly impactful solution for accelerating the pre-training of large foundation models. Its ability to achieve significant speedups and improved scalability without compromising model quality positions it as a key enabling technology for the continued advancement and broader accessibility of large-scale AI.

Key Takeaways

• Addresses Communication Bottleneck: Radius effectively mitigates the dominant communication bottleneck in data-parallel pre-training of large foundation models, which often exceeds computation time on high-performance clusters.
• Leverages Gradient Stability: The core innovation of Radius is its exploitation of the temporal stability of high-magnitude gradient values, enabling periodic (every T iterations, e.g., 200) selection of top-k indices to amortize computation overhead.
• Scalable Gradient Sparsity: By selecting top-k indices after the all-reduce operation, Radius ensures global consistency, eliminates the need for expensive all-gather operations, avoids memory overhead, and prevents the "filling effect," thereby improving scalability.
• Significant Performance Gains: Radius delivers a 90% speedup on wall-clock time for GPT-20B pre-training (with 40% gradient density) and a 21% acceleration per iteration on 64 A100 GPUs, while maintaining or improving model quality.
• Enhanced Scaling Efficiency & Quality: It improves strong scaling efficiency (43% vs. 33% for baseline at 128 GPUs) and leads to better performance on several zero-shot downstream tasks (e.g., LAMBADA, MathQA, PIQA) compared to dense baselines.
• Robustness and Practicality: Radius provides a stable and reliable approach to gradient sparsity, outperforming naive top-k methods that diverge in perplexity, making it a practical solution for real-world large-scale LLM training.

About the Speaker(s)

Mingkai Zheng is a second-year PhD student from Rutgers University. He is the primary speaker for this talk, presenting the research work on Radius. His work focuses on improving the efficiency and scalability of large foundation model pre-training.

Zhao Zhang is also credited as a speaker for this work. No further biographical details were provided in the transcript for either speaker beyond Mingkai Zheng's academic affiliation.