Demystifying Workload Imbalances in Large Transformer Model Training over Variable-length Sequences

2024-12-10

arXiv.org

AI Summary

To address computational load imbalance in large Transformer model training—caused by skewed sequence-length distributions and the mismatch between linear memory complexity and quadratic computation complexity of attention mechanisms—this paper proposes a dynamic heterogeneous parallel training paradigm. Our method features a two-stage sequence-aware data allocation mechanism that jointly optimizes intra- and inter-group load balancing, integrated with cross-iteration workload modeling, dynamic scheduling, and heterogeneous GPU-aware resource allocation. To the best of our knowledge, this is the first approach to achieve joint optimization of computation, memory, and hardware resources for variable-length sequence training. Evaluated on real-world large-model training workloads, it delivers 1.32–2.66× speedup over baseline methods, significantly improving GPU utilization and training throughput.

Technology Category

• Machine Learning: Hardware-aware ML
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

To optimize large Transformer model training, efficient parallel computing and advanced data management are essential. However, current methods often assume a stable and uniform training workload, neglecting imbalances in data sampling and packing that can impede performance. Specifically, data sampling imbalance arises from uneven sequence length distribution of the training data, while data packing imbalance stems from the discrepancy between the linear memory complexity and quadratic time complexity of the attention mechanism. To address these imbalance issues, we develop Hydraulis, which jointly optimizes the parallel strategies and data assignment. For one thing, we introduce large model training with dynamic heterogeneous parallel strategies in response to the sequence length variations within and across training iterations. For another, we devise a two-stage data assignment approach, which strikes a good balance in terms of the training workloads both within and across model replicas. Empirical results demonstrate that Hydraulis outperforms existing systems by 1.32-2.66 times.

Problem

• Addressing data-induced imbalances in Transformer training
• Optimizing parallel strategies for varying sequence lengths
• Balancing workloads across model replicas via data assignment

Innovation

• Dynamic heterogeneous parallel strategies for sequence variations
• Two-stage data assignment balancing training workloads
• Joint optimization of parallel computing and data management