PaperlayerView exact revision on arXiv

Source-linked AI summary

Oobleck: Resilient Distributed Training of Large Models Using Pipeline Templates

Insu Jang, Zhenning Yang, Zhen Zhang, Xin Jin, Mosharaf Chowdhury

arXiv:2309.08125v2cs.DCcs.LG

TL;DR

Hybrid-parallel training has limited systematic fault-tolerance support. Oobleck uses precomputed heterogeneous pipeline templates and replicated model states to recover from failures. It reports outperforming Bamboo and Varuna by up to 29.6×.

  • Problem

    Existing frameworks provide little systematic fault-tolerance support for hybrid-parallel training, while cited data-parallel approaches do not extend their benefits to hybrid parallelism.

    Source: §1 · p. 1

  • Method

    Oobleck plans heterogeneous pipeline templates, instantiates at least f+1 pipelines, and recovers by reinstantiating pipelines and copying missing model states from replicas.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    Oobleck outperforms Bamboo and Varuna by up to 29.6×.

    Source: §9 · p. 12

  • Takeaways & Limitations

    Oobleck provides guaranteed fault tolerance while maximizing throughput in resilient distributed training.

    Source: §1 · p. 2

    The cited state-of-the-art approaches perform poorly for larger models, especially when failures are frequent.

    Source: §2.3 · p. 3

Abstract

from arXiv · show

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$.

1 Introduction

Oobleck addresses limited fault-tolerance support for hybrid-parallel training, where existing strategies trade training overhead against recovery overhead. It uses precomputed heterogeneous pipeline templates and replicated model states for recovery, and reports throughput improvements over Bamboo and Varuna.

1 Introduction

  • Hybrid-parallel training lacks systematic fault-tolerance support, and existing data-parallel elasticity does not extend to it.Data-parallel approaches use workers’ existing model-state copies, while hybrid-parallel approaches add redundancy in other forms.Source: §1 · p. 1
  • Existing hybrid-parallel approaches trade training overhead against recovery overhead, favoring different strategies depending on failure frequency.The cited approaches use standby GPUs, checkpoint-based reconfiguration, or redundant computation.Source: §1 · pp. 1–2
  • Oobleck’s pipeline templates are logically equivalent but physically heterogeneous, enabling pipelines to be instantiated for a model using different node counts and configurations.Templates are generated during planning and reused by the execution engine; replicas can recover missing model states by copying layers rather than restarting from checkpoints.Source: §1 · p. 2
  • For a job tolerating f simultaneous failures, Oobleck instantiates at least f+1 heterogeneous pipelines and guarantees that a template combination can use all remaining nodes after f or fewer failures.The fixed global batch is distributed in proportion to the pipelines’ computing capability.Source: §1 · p. 2
  • Oobleck outperforms Bamboo and Varuna by up to 29.6× across different failure frequencies, model sizes, and computational complexities.The evaluation includes large models such as GPT-3 variants with billions of parameters.Source: §1 · p. 2

2 Background and Motivation

Large-model training commonly combines data and model parallelism, but failures complicate fault tolerance when models span multiple GPUs. Existing approaches have limitations in recovery overhead, added computation, or resource use, particularly for larger models and frequent failures.

2.1 Hybrid Parallelism

  • Hybrid parallelism combines data and model parallelism because either alone can be insufficient for efficient large-DNN training.Pipeline parallelism incurs bubble overheads, while tensor parallelism has communication costs that cannot be hidden in computation.Source: §2.1 · p. 2

2.2 Fault Tolerance in Distributed Training

  • Checkpointing-based recovery can take unsustainable time at high failure rates, while redundant computation adds fixed compute and memory overhead.Bamboo also requires restart after two adjacent node failures.Source: §2.2 · p. 3

2.3 Limitations of the State-of-the-Art

  • Varuna offers higher throughput than Bamboo but has restart overhead, which dominates as failures become more frequent.Figure 1 compares effective training time for BERT-Large and GPT-3 6.7B at failure intervals of six hours and ten minutes.Source: §2.3 · p. 3, Figure 1 · p. 3
  • Bamboo’s redundant computation adds performance overhead, and two adjacent node failures can still require full checkpoint reconfiguration.Forward redundant computation runs even without failures; backward redundant computation can create stragglers.Source: §2.3 · p. 3
  • Both Bamboo and Varuna perform poorly for larger models, especially when failures are frequent.Bamboo runs out of memory, while Varuna spends most of its time preparing to train.Source: §2.3 · p. 3

3 Oobleck Overview

Oobleck uses heterogeneous pipeline templates to support high-throughput training and fast fault tolerance. Its overview covers guarantees for pipeline reconfiguration, system components, and the training lifecycle.

3 Oobleck Overview

  • Oobleck is a resilient training platform that uses pipeline templates to combine high throughput with fast fault tolerance.Its guarantees allow reconfiguration without restarts for up to f simultaneous failures in the worst case.Source: §3 · p. 3

3.1 Pipeline Templates

  • Pipeline templates specify node assignment, stage count, and layer-to-GPU mapping; heterogeneous templates let Oobleck use arbitrary node counts and replace failed pipelines without replanning.Planning creates templates that execution reuses, including templates requiring fewer nodes after a failure.Source: §3.1 · p. 3

3.2 Fault Tolerance Guarantees

  • Oobleck tolerates up to f simultaneous pipeline failures without restarting, and can tolerate more than f node failures if a complete copy of model states remains.In the illustrated case with f=2, eight node failures still leave one copy of each stage alive.Source: §3.2 · pp. 3–4, Figure 2 · p. 4

3.3 System Components

  • Oobleck’s generator creates templates once, while its execution engine instantiates at least f+1 replicas and redistributes microbatches to balance heterogeneous pipeline latency.The engine reconfigures from precomputed templates when nodes fail or are added.Source: §3.3 · p. 4

3.4 Training Lifecycle

  • Users provide f, a model and dataset, and batch information; after failures, Oobleck reconfigures and copies missing model states when a complete replica remains.If the cluster cannot maintain f+1 replicas, Oobleck saves progress and exits; users can restart from a recent checkpoint after enough nodes recover.Source: §3.4 · p. 4, Figure 3 · p. 4

4 Oobleck Planning Algorithm

Oobleck plans heterogeneous pipeline templates once, then instantiates combinations that use available nodes while meeting the replica requirement. It selects fast mappings and plans, distributing work across pipelines to balance their processing times.

4 Oobleck Planning Algorithm

  • Oobleck instantiates r ≥ f+1 heterogeneous pipeline replicas, and its template combinations can use all currently available nodes despite failures.The execution engine selects among precomputed templates rather than deriving a new optimal configuration after each failure.Source: §4 · p. 4

4.1 Generating Pipeline Templates

  • Each template specifies its node count and how GPUs and model layers map into pipeline stages; templates are generated once for the training job.The generator selects template node counts, then maps GPUs and model layers to minimize iteration time.Source: §4.1 · p. 5, §4.1.1 · p. 5

4.1.2 GPU–Stage Mapping.

  • A divide-and-conquer search partitions model layers and GPUs together, selecting the mapping and stage count with minimum iteration time.It evaluates a pipeline’s critical path using T1 for initialization and teardown, T2 for the slowest stage’s steady alternation, and T3 for the remaining steady-phase work.Source: §4.1.2 · pp. 5–6

4.2 Pipeline Instantiation

  • For available nodes, Oobleck enumerates feasible combinations of templates and selects the plan with the highest estimated throughput.Different combinations can cover the same node count, so it evaluates all feasible combinations rather than choosing a template in isolation.Source: §4.2 · p. 7, §4.2.1 · p. 7

4.2.1 Enumerating All Instantiation Options.

  • Dynamic programming enumerates pipeline-count vectors that use all N nodes and contain at least f+1 pipelines.The recurrence builds feasible sets in O(Np), after which Oobleck filters for the replica-count requirement.Source: §4.2.1 · p. 7

4.2.2 Calculating Throughput with Batch Distribu-

  • Batch distribution assigns microbatches across heterogeneous pipelines to minimize differences in processing times and avoid stragglers.The integer optimization may fail when a small global batch cannot be split into integers; Oobleck recommends a nearby distributable global batch instead of changing it arbitrarily.Source: §4.2.2 · pp. 7–8

5 Dynamic Reconfiguration

Oobleck replaces failure-affected pipelines using precomputed templates, then redistributes batches to keep the global batch size constant. Reinstantiation may borrow nodes or merge pipelines when no suitable template is immediately available.

5 Dynamic Reconfiguration

  • Oobleck replaces a failure-affected pipeline with one instantiated from a precomputed template, then copies missing layers from unaffected replicas.It can recover while at least (f+1)n0 nodes remain.Source: §5 · p. 8

5.1 Pipeline Reinstantiation

  • Pipeline reinstantiation uses three steps: simple reinstantiation, borrowing nodes, and merging pipelines.These steps address cases where remaining nodes do not match a suitable template.Source: §5.1 · p. 8
  • If remaining nodes fit a template, Oobleck simply reinstantiates it; otherwise, it borrows nodes from other pipelines to reach the smallest template size.Pipelines that yield nodes are themselves reinstantiated with fewer nodes.Source: §5.1 · p. 8
  • When pipelines can no longer yield nodes, Oobleck merges them; a merged pipeline is guaranteed to have an instantiable template if it has at least n0 nodes.n0 is the minimum number of nodes to maintain one pipeline.Source: §5.1 · p. 8

5.2 Batch Redistribution

  • After reinstantiation, Oobleck recalculates batch sizes for the changed configuration while keeping the global batch size constant.Each pipeline may have more batches to compute.Source: §5.2 · p. 8

6 Implementation

Oobleck implements heterogeneous pipelines with layer-level gradient synchronization, failure detection through CPU-process TCP connections, and components including PyTorch, Merak, and DeepSpeed.

6 Implementation

  • Oobleck integrates PyTorch FSDP, pipeline and data parallelism, Merak, and DeepSpeed; Figure 9 illustrates execution with heterogeneous pipeline stages.It uses Pyomo and MindtPy for nonlinear integer optimization.Source: §6 · pp. 8–9, Figure 9 · p. 9

6.1 Model Synchronization Between Heterogeneous Pipelines

  • Oobleck synchronizes gradients layer by layer because heterogeneous pipelines have different stage configurations, with layers potentially communicating across different peer nodes.In the example, one three-layer stage in a two-stage pipeline corresponds to stages B and C in a three-stage pipeline.Source: §6.1 · p. 9

6.2 Detecting Node Failures

  • Oobleck detects node failure through a CPU process on each node connected by TCP to a centralized CPU process, which broadcasts socket disconnections for reconfiguration.NCCL can hang while communicating with a failed node until a timer expires.Source: §6.2 · p. 9

7 Evaluation

Oobleck is evaluated against Bamboo and Varuna on large models across controlled failures and spot-instance traces. The evaluation reports strong throughput, low planning overhead, and differences in how the systems spend time on fault tolerance.

7.1 Experimental Setup

  • Evaluation uses 30 NVIDIA A40 GPUs and compares Oobleck with Varuna and Bamboo, which do not support 3D parallelism and are tested with one GPU per node.The spot-instance experiments leverage preemption warnings, although such notifications are not generally available.Source: §7.1 · p. 9

7.2 Throughput Under Controlled Failures

  • 29.6× for GPT3-6.7B; Oobleck outperforms or matches other frameworks for every model and controlled failure scenario.Table 2 reports throughput in samples/s across failure frequencies.Source: §7.2 · p. 10, Table 2 · p. 10

7.3 Throughput in Spot Instance Traces

  • In spot-instance traces, Varuna's checkpointing grows increasingly costly with model size and it made no progress on GPT-3 6.7B; Bamboo could not run models larger than GPT-2.The traces include node additions as well as preemptions; Figure 10 shows throughput changes across models and cloud environments.Source: §7.3 · p. 10, Figure 10 · p. 11

7.4.1 Overhead of Oobleck Planning.

  • Planning overhead is less than 0.1% of estimated end-to-end training time for large models using hundreds of GPUs, and memoization and intermediate caches make later template creation negligible.The planning algorithm's latency is reported in Table 3.Source: §7.4.1 · p. 11, Table 3 · p. 11

7.4.2 Throughput Breakdown.

  • Oobleck has a small copying overhead for missing layers, while Varuna's checkpoint-related waste occupies up to 95% of wall-clock time and Oobleck reaches at least 75% of no-failure effective throughput.Figure 11 presents the time-occupation breakdown; all frameworks experience fallback overhead after mid-iteration failures.Source: Figure 11 · p. 11, §7.4.2 · p. 12

7.4.3 Impact of Checkpointing Overhead.

  • Checkpointing overhead in Varuna largely comes from serialized checkpointing and full restart, and loading cannot be overlapped with computation.Full restart includes framework initialization and loading the last checkpoint; Table 4 compares Varuna variants with Oobleck.Source: §7.4.3 · p. 12, Table 4 · p. 12

8 Related Works

Related work addresses elasticity, spot-instance fault tolerance, and large-model training through distinct approaches, with the cited large-model optimization proposals described as complementary to Oobleck.

Elastic training

  • Horovod Elastic and TorchElastic restart after failure and recovery, while CoDDL, Aryl, and Pollux address elastic resource sharing for data-parallel training of models fitting within one GPU.Source: §8 · p. 12

Distributed training with spot instances

  • Varuna uses hybrid parallelism with checkpoint-based reconfiguration, while Bamboo uses redundant computation for resilience to frequent spot-instance preemptions; Oobleck matches or outperforms them across model sizes and failure frequencies.Source: §8 · p. 12

Large model training

  • Other large-model training proposals optimize training through diverse mechanisms but do not provide fault tolerance out of the box and are orthogonal to Oobleck.Source: §8 · p. 12

9 Conclusion

Oobleck co-designs planning and execution around pipeline templates to support fast recovery and high throughput. It reinstantiates pipelines and restores missing model states from replicas, outperforming Bamboo and Varuna by up to 29.6×.

9 Conclusion

  • Oobleck co-designs planning and execution around pipeline templates to support resilient large-model training.Templates are designed during planning and reused during training execution.Source: §9 · p. 12
  • Recovery reinstantiates pipelines from templates and copies missing model states from replicas without a full checkpoint restart.Source: §9 · p. 12
  • Oobleck outperforms Bamboo and Varuna by up to 29.6×.Source: §9 · p. 12

A Proof of Nodes Specification Covering All Nodes

The proof establishes that a finite set of pipeline templates can represent any feasible remaining node count. The guarantee holds under stated conditions on the number and sizes of templates.

A Proof of Nodes Specification Covering All Nodes

  • A finite set of p pipeline templates can represent any feasible N′ from (f+1)n0 through N, provided the theorem’s conditions hold.The stated conditions are p > n0 − 1 and consecutive template sizes.Source: A Proof of Nodes Specification… · p. 15
  • The proof represents N′ as an integer combination of template sizes, with x_i counting pipelines instantiated from the template of size n_i.Source: A Proof of Nodes Specification… · p. 15
  • With the stated conditions, the Frobenius number is n0 − 1, below the minimum feasible N′ of n0, so every feasible N′ is representable.The proof invokes the Frobenius number for an arithmetic sequence of template sizes.Source: A Proof of Nodes Specification… · p. 15

B Proof of Guarantee for Pipeline Template Availability When Merging Pipelines

The proof shows that templates are available for rebuilding pipelines after node failures. It extends the single-pipeline case to failures spread across several pipelines by repeatedly merging them.

B Proof of Guarantee for Pipeline Template Availability When Merging Pipelines

  • After losing k nodes from an n0-node pipeline, merging it with another n0-node pipeline can form a 2n0 − k-node pipeline, for which a template is guaranteed to exist.The stated template size satisfies 2n0 − k ≥ n0.Source: B Proof of Guarantee for Pipeline… · p. 15
  • The single-pipeline guarantee follows from maintaining at least f+1 model replicas after a merge and the resulting contradiction if the required template were absent.Source: B Proof of Guarantee for Pipeline… · p. 15
  • When failures leave several pipelines below n0 nodes, Oobleck repeatedly merges pipelines until the combined node count reaches n0 or more.The proof uses the single-pipeline guarantee to show a template exists for the resulting node count.Source: B Proof of Guarantee for Pipeline… · p. 15

C Throughput of All Models in Spot Instances

Spot-instance results vary with model size: Varuna matches Oobleck on average for small models when it checkpoints before preemption, while Oobleck outperforms it on large models.

C Throughput of All Models in Spot Instances

Loading 2309.08125v2…