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

TL;DR

LIBRA tackles the gradient and activation communication bottleneck in distributed training of large AI models by introducing a workload-aware, multi-dimensional network optimization framework. It models end-to-end training time as a function of bandwidth per dimension and uses a Quadratic Programming solver to allocate BW across dimensions for maximum performance or perf-per-cost, under design constraints. The framework supports multiple workloads, provides a cost model, and demonstrates notable speedups and cost reductions across large language models and vision/MLP workloads, including co-design opportunities with runtime schedulers. The work enables practical design-time co-optimization of topology, parallelization, and runtime strategies, and is poised to influence next-generation AI cluster architectures.

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.

Figures (21)

• Figure 1: Communication sizes (FP16 as the datatype) for ML model training across 1,024 NPUs. The parallelization strategy of Turing-NLG and smaller workloads are data parallel (minibatch size of 32), while GPT-3 and MSFT-1T use both tensor and data parallelism.
• Figure 2: (a) An abstract view of multi-dimensional networks. (b) An example of physical connotation assigned to 2--4D networks.
• Figure 3: The architecture of Libra, a workload-aware multi-dimensional network BW optimization framework. Inputs to the Libra framework are represented in obrounds, with corresponding example values shown in square brackets. Libra estimates the end-to-end training time based on the provided network shape, training loop, and target workloads. Additionally, it calculates the cost of a network using a specified cost model and network shape. Subsequently, Libra searches for the optimal network BW configuration that maximizes the given objective while adhering to designated design constraints.
• Figure 4: (a) An example of a switch-based 2D network. (b) Example 1D topology, which utilizes a 2-level hierarchical switch in Dim 1.
• Figure 5: Examples of training loops. (a) HP-(2,2) parallelization strategy (b) A training loop without any overlap (No Overlap) (c) A training loop with TP and DP running concurrently (TP-DP Overlap).