Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AB-Training: A Communication-Efficient Approach for Distributed Low-Rank Learning

Daniel Coquelin, Katherina Flügel, Marie Weiel, Nicholas Kiefer, Muhammed Öz, Charlotte Debus, Achim Streit, Markus Götz

TL;DR

AB-training presents a communication-efficient, data-parallel approach that leverages low-rank weight representations and independent training groups to reduce interconnect traffic in distributed neural network training. By decomposing weight matrices into AB components and alternating training across A- and B-groups, then reconciling via synchronization and a subsequent full-rank rebound, the method achieves substantial traffic reductions with competitive accuracy across ImageNet and CIFAR-10 benchmarks. The work demonstrates ~70% average traffic reduction and variable compression ratios (e.g., up to 44.14:1 in CIFAR-10 with VGG16) while maintaining training times comparable to traditional DP; it also reveals regularization benefits at smaller scales but notes challenges with large-scale batch effects. These results suggest AB-training as a practical strategy for scalable, communication-efficient distributed training in HPC environments, with future work needed to refine update mechanisms and hyperparameter schedules for extreme-scale deployments.

Abstract

Communication bottlenecks severely hinder the scalability of distributed neural network training, particularly in high-performance computing (HPC) environments. We introduce AB-training, a novel data-parallel method that leverages low-rank representations and independent training groups to significantly reduce communication overhead. Our experiments demonstrate an average reduction in network traffic of approximately 70.31\% across various scaling scenarios, increasing the training potential of communication-constrained systems and accelerating convergence at scale. AB-training also exhibits a pronounced regularization effect at smaller scales, leading to improved generalization while maintaining or even reducing training time. We achieve a remarkable 44.14 : 1 compression ratio on VGG16 trained on CIFAR-10 with minimal accuracy loss, and outperform traditional data parallel training by 1.55\% on ResNet-50 trained on ImageNet-2012. While AB-training is promising, our findings also reveal that large batch effects persist even in low-rank regimes, underscoring the need for further research into optimized update mechanisms for massively distributed training.

AB-Training: A Communication-Efficient Approach for Distributed Low-Rank Learning

TL;DR

AB-training presents a communication-efficient, data-parallel approach that leverages low-rank weight representations and independent training groups to reduce interconnect traffic in distributed neural network training. By decomposing weight matrices into AB components and alternating training across A- and B-groups, then reconciling via synchronization and a subsequent full-rank rebound, the method achieves substantial traffic reductions with competitive accuracy across ImageNet and CIFAR-10 benchmarks. The work demonstrates ~70% average traffic reduction and variable compression ratios (e.g., up to 44.14:1 in CIFAR-10 with VGG16) while maintaining training times comparable to traditional DP; it also reveals regularization benefits at smaller scales but notes challenges with large-scale batch effects. These results suggest AB-training as a practical strategy for scalable, communication-efficient distributed training in HPC environments, with future work needed to refine update mechanisms and hyperparameter schedules for extreme-scale deployments.

Abstract

Communication bottlenecks severely hinder the scalability of distributed neural network training, particularly in high-performance computing (HPC) environments. We introduce AB-training, a novel data-parallel method that leverages low-rank representations and independent training groups to significantly reduce communication overhead. Our experiments demonstrate an average reduction in network traffic of approximately 70.31\% across various scaling scenarios, increasing the training potential of communication-constrained systems and accelerating convergence at scale. AB-training also exhibits a pronounced regularization effect at smaller scales, leading to improved generalization while maintaining or even reducing training time. We achieve a remarkable 44.14 : 1 compression ratio on VGG16 trained on CIFAR-10 with minimal accuracy loss, and outperform traditional data parallel training by 1.55\% on ResNet-50 trained on ImageNet-2012. While AB-training is promising, our findings also reveal that large batch effects persist even in low-rank regimes, underscoring the need for further research into optimized update mechanisms for massively distributed training.
Paper Structure (14 sections, 4 equations, 5 figures, 2 tables, 1 algorithm)

This paper contains 14 sections, 4 equations, 5 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Distributed Training of Neural Networks
  4. Low-Rank Neural Network Training
  5. Pruning
  6. Methodology
  7. Experimental Evaluation
  8. Computational Environment
  9. Datasets and Models
  10. Hyperparameter Considerations
  11. Constant Local Batch Size Scaling
  12. Constant Global Batch Size Scaling
  13. Discussion
  14. Conclusion

Figures (5)

  • Figure 1: A UML diagram of the AB training procedure.
  • Figure 2: Highest top 1 accuracy measurements for each training run on ImageNet-2012 for two network architectures with a constant local batch size of 256. Global batch sizes range from 2,048 to 32,768 in powers of 2. Error are plotted, though not always visible.
  • Figure 3: Highest top 1 accuracy measurements for each training run on ImageNet-2012 for two network architectures with a constant global batch size of 4,096. Error are plotted, though not always visible.
  • Figure 4: Compression ratios for AB training with and without groups and a traditional DP baseline on ImageNet-2012 for two network architectures
  • Figure 5: Scaled interconnect traffic and job wall-clock time for the ViT B/16 trained on ImageNet-2012. Scaling is based on the average time required for the calculation and communication of gradients on 4 nodes.