LIBRA: Enabling Workload-aware Multi-dimensional Network Topology Optimization for Distributed Training of Large AI Models

Abstract

As model sizes in machine learning continue to scale, distributed training is necessary to accommodate model weights within each device and to reduce training time. However, this comes with the expense of increased communication overhead due to the exchange of gradients and activations, which become the critical bottleneck of the end-to-end training process. In this work, we motivate the design of multi-dimensional networks within machine learning systems as a cost-efficient mechanism to enhance overall network bandwidth. We also identify that optimal bandwidth allocation is pivotal for multi-dimensional networks to ensure efficient resource utilization. We introduce LIBRA, a framework specifically focused on optimizing multi-dimensional fabric architectures. Through case studies, we demonstrate the value of LIBRA, both in architecting optimized fabrics under diverse constraints and in enabling co-optimization opportunities.

Contains 10 main pages, 21 figures, 3 tables

References in corpus (11)

• Scaling Laws for Neural Language Models
• A Survey on Distributed Machine Learning
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Deep Learning Recommendation Model for Personalization and Recommendation Systems
• Hardware for Machine Learning: Challenges and Opportunities
• Scaling Distributed Machine Learning with In-Network Aggregation
• ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale
• Massively Distributed SGD: ImageNet/ResNet-50 Training in a Flash
• DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale
• Enabling Compute-Communication Overlap in Distributed Deep Learning Training Platforms
• Themis: A Network Bandwidth-Aware Collective Scheduling Policy for Distributed Training of DL Models

Cited by in corpus (2)

• TACOS: Topology-Aware Collective Algorithm Synthesizer for Distributed Machine Learning
• Towards a Standardized Representation for Deep Learning Collective Algorithms