AutoHAAP: Automated Heterogeneity-Aware Asymmetric Partitioning for LLM Training

Heterogeneous clusters with diverse devices mitigate computational and memory burdens in large language model (LLM) training, yet their inherent resource heterogeneity, characterized by divergent computation, memory, and bandwidth capabilities, renders manual parallelization strategy optimization both challenging and time-intensive. Automatic parallelization is critical for scaling complex workloads across heterogeneous architectures. However, previous methodologies face significant inefficiencies. First, insufficient pruning of the parameter initialization space results in impractically large search spaces. Second, the prevailing automatic parallel search strategies exhibit suboptimal performance in load balancing and resource constraint adaptation. Third, dynamic parallel strategy tuning incurs substantial overhead due to redundant latency calculations for operators with unchanged configurations, leading to unnecessary computational costs. Therefore, insufficient search space pruning, suboptimal load/resource adaptation, and redundant latency computation are identified as the major bottlenecks in our research. To address these challenges, we propose AutoHAAP (Automated Heterogeneity-Aware Asymmetric Partitioning), a novel framework incorporating three core innovations: (1) memory-aware initialization to drastically reduce viable search spaces; (2) a heterogeneity-aware load-balancing estimator that guides resource-efficient configuration search; and (3) state caching mechanisms eliminating redundant latency calculations. Evaluations across GPT3 and Llama3 models of varying scales on both homogeneous and heterogeneous clusters demonstrate that AutoHAAP achieves$\mathbf{0. 6 8}-\mathbf{9 8} \times$search efficiency gains,$\mathbf{6. 5 7 \%} \boldsymbol{-} \mathbf{1 0 6. 9 \%} \boldsymbol{\times}$throughput improvements in homogeneous environments, and$\mathbf{1 0. 1 \%} \boldsymbol{-} 22.28 \% \times$throughput enhancements in heterogeneous setups. These results validate AutoHAAP's effectiveness in distributed LLM training on diverse hardware.