Oobleck: Resilient Distributed Training of Large Models Using Pipeline Templates
This addresses the challenge of efficient and reliable training of billion-parameter models for AI researchers and practitioners, representing a novel method rather than an incremental improvement.
The paper tackles the problem of fault tolerance in distributed training of large DNN models by proposing Oobleck, which uses pipeline templates and replicas to guarantee resilience and avoid resource idling, resulting in up to 29.6x higher throughput compared to state-of-the-art solutions.
Oobleck enables resilient distributed training of large DNN models with guaranteed fault tolerance. It takes a planning-execution co-design approach, where it first generates a set of heterogeneous pipeline templates and instantiates at least $f+1$ logically equivalent pipeline replicas to tolerate any $f$ simultaneous failures. During execution, it relies on already-replicated model states across the replicas to provide fast recovery. Oobleck provably guarantees that some combination of the initially created pipeline templates can be used to cover all available resources after $f$ or fewer simultaneous failures, thereby avoiding resource idling at all times. Evaluation on large DNN models with billions of parameters shows that Oobleck provides consistently high throughput, and it outperforms state-of-the-art fault tolerance solutions like Bamboo and Varuna by up to $29.6x$.