Adaptra: Straggler-Resilient Hybrid-Parallel Training with Pipeline Adaptation
Tianyuan Wu, Lunxi Cao, Hanfeng Lu, Xiaoxiao Jiang, Yinghao Yu, Siran Yang, Guodong Yang, Jiamang Wang, Lin Qu, Liping Zhang, Wei Wang
TL;DR
Adaptra addresses the dual challenges of dependency bubbles and head-of-line blocking in pipeline-parallel DNN training under communication stragglers. It introduces a straggler-resilient pipeline adaptation algorithm that increases inter-stage slackness and a fully decoupled data plane that offloads communications to CPU delegates to eliminate HOL blocking, plus RNIC fault tolerance. The approach yields 1.2–3.5× faster iterations under various straggler conditions and maintains uninterrupted training during RNIC failures, demonstrated on GPT-2-scale models up to 140B parameters. By combining semantic pipeline optimization with hardware-oblivious delegation, Adaptra offers practical resilience for large-scale production training.
Abstract
Training large Deep Neural Network (DNN) models at scale often encounters straggler issues, mostly in communications due to network congestion, RNIC/switch defects, or topological asymmetry. Under advanced pipeline parallelism, even minor communication delays can induce significant training slowdowns. This occurs because (1) slow communication disrupts the pipeline schedule, creating cascading "bubbles" in a domino effect, and (2) current GPU kernel scheduling is susceptible to head-of-line blocking, where slow communication blocks subsequent computations, further adding to these bubbles. To address these challenges, we present ADAPTRA, a straggler-resilient training system with two key optimizations. First, it optimally adapts the pipeline schedule in the presence of stragglers to absorb communication delays without inducing cascading bubbles, using a simple yet effective algorithm guided by an analytical model. Second, upon detecting slow communication, ADAPTRA offloads communication operations from GPU to host memory and utilizes CPU-side RDMA for data transfer. This eliminates head-of-line blocking as subsequent computation kernels can be scheduled immediately on GPUs. Together, these optimizations effectively reduce pipeline stalls in the presence of communication stragglers, improving the training iteration time by 1.2-3.5x in our experiments under various settings.
