Table of Contents
Fetching ...

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.

Adaptra: Straggler-Resilient Hybrid-Parallel Training with Pipeline Adaptation

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.
Paper Structure (25 sections, 2 theorems, 4 equations, 21 figures, 3 tables, 4 algorithms)

This paper contains 25 sections, 2 theorems, 4 equations, 21 figures, 3 tables, 4 algorithms.

Key Result

Lemma 1

For any pipeline schedule, the warm-up forward count is non-increasing over stages, i.e., $x_i \geq x_{i+1}$ for all $i = 0, 1, \dots, S-1$.

Figures (21)

  • Figure 1: GPT-2 14B training performance on 8 nodes (one H800 GPU per node) with 8-stage PP, where minor communication delays between PP stages trigger dependency bubbles and blocking stalls, causing significant slowdowns in 1F1B and ZeroBubble (ZB) pipelines.
  • Figure 2: Head-of-line blocking due to sequential kernel scheduling: slow comm. blocks subsequent comp.
  • Figure 3: Iteration time growth under different per packet delays, where DP is more communication-sensitive to PP.
  • Figure 4: An ideal straggler-free ZeroBubble qi2023zero pipeline with $4$ stages and $12$ microbatches, completing in 390 ms.
  • Figure 5: ZeroBubble schedule under $c_0=10/20$ ms delay between stage 0 and 1. Increasing $c_0$ from 0 to 10 ms only prolongs iteration time $T$ by 10 ms, while an additional 10 ms delay introduces a 40 ms growth in $T$.
  • ...and 16 more figures

Theorems & Definitions (2)

  • Lemma 1: Monotonic warm-up
  • Theorem 1: Delay resilience