PaperlayerView exact revision on arXiv

Source-linked AI summary

Memory-Efficient Pipeline-Parallel DNN Training

Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, Matei Zaharia

arXiv:2006.09503v3cs.LGcs.DCstat.ML

TL;DR

Large DNNs often exceed single-accelerator memory, motivating distributed training that also avoids utilization, communication, and weight-version problems. PipeDream-2BW combines double-buffered, coalesced weight updates with hardware-aware model planning, achieving up to 20× acceleration with similar accuracy. Its main scope boundary is that large-scale training can require impractically large batches for additional gradient accumulation.

  • Problem

    Large DNN parameters and activations may not fit in one accelerator, while distributed training must address utilization, communication, weight-version, memory, and partitioning constraints.

  • Method

    PipeDream-2BW combines double-buffered weight updates with coalesced gradients and a planner that partitions operators under memory and hardware-interconnect constraints.

  • Results

    Up to 20× speedups over model-parallel baselines and up to 3.2× over GPipe are reported, with similar Adam loss trajectories and downstream accuracy.

  • Takeaways & Limitations

    Two weight versions provide memory-efficient pipeline parallelism with high throughput and data-parallelism-like weight-update semantics for models with billions of parameters.

  • Takeaways & Limitations

    For large worker counts, even minimal gradient accumulation can imply high batch sizes, making additional accumulation infeasible because batch size cannot scale indefinitely without affecting convergence.

Abstract

from arXiv · show

Many state-of-the-art ML results have been obtained by scaling up the number of parameters in existing models. However, parameters and activations for such large models often do not fit in the memory of a single accelerator device; this means that it is necessary to distribute training of large models over multiple accelerators. In this work, we propose PipeDream-2BW, a system that supports memory-efficient pipeline parallelism. PipeDream-2BW uses a novel pipelining and weight gradient coalescing strategy, combined with the double buffering of weights, to ensure high throughput, low memory footprint, and weight update semantics similar to data parallelism. In addition, PipeDream-2BW automatically partitions the model over the available hardware resources, while respecting hardware constraints such as memory capacities of accelerators and interconnect topologies. PipeDream-2BW can accelerate the training of large GPT and BERT language models by up to 20$\times$ with similar final model accuracy.

1. Introduction

Large DNNs exceed single-accelerator memory, while conventional model-parallel and pipelined approaches face utilization, communication, weight-version, and partitioning challenges. PipeDream-2BW addresses these constraints with double-buffered weight updates and hardware-aware planning, achieving high throughput and low memory footprint.

  • Motivation: Large models can exceed the memory available on a single accelerator, requiring training to be distributed across multiple devices.Model parallelism partitions the model over accelerator devices.
  • Motivation: Traditional model parallelism can underutilize resources or incur high communication overhead, increasing training time and dollar cost.Good scaling may be limited to devices within a multi-GPU server.
  • Challenges: Naive pipelining can harm convergence because an input’s forward and backward passes may use inconsistent weight versions.Existing techniques trade off memory footprint and throughput to avoid this problem.
  • Challenges: Operator partitioning must satisfy memory capacity and heterogeneous-interconnect constraints while searching an exponentially large placement space.These constraints make effective model partitioning challenging as model size increases.
  • Approach: PipeDream-2BW uses double-buffered, coalesced weight updates to avoid pipeline flushes while maintaining only two weight versions instead of up to d.It uses new weights for newly admitted inputs and a shadow version for in-flight inputs, with a constant weight delay of 1.
  • Approach: PipeDream-2BW’s planner partitions DNN operators across workers while accounting for accelerator memory and exploiting repetitive model structure.The planner also seeks to keep expensive communication on high-speed intra-server interconnects.
  • Results: 1.3× to 20× end-to-end speedups are achieved for GPT models versus an optimized model-parallel baseline, while performance is up to 3.2× higher than GPipe.The system also trains transformer models that vanilla PipeDream cannot fit in memory, with similar Adam loss trajectories and downstream accuracy.

2. Background

Pipeline parallelism improves resource utilization by processing multiple inputs across inter-layer model partitions, but existing systems trade pipeline flushing and activation costs against weight-version memory. GPipe keeps one weight version with periodic flushes, whereas PipeDream avoids flushes by stashing multiple versions.

  • Related Training Strategies: Data parallelism replicates the entire model on every worker, so it cannot train models that do not fit on a single worker.It remains applicable to smaller model partitions.
  • Related Training Strategies: Traditional inter-layer model parallelism partitions weights across workers but can leave at most one worker active at a time.Intermediate activations and gradients are communicated across workers.
  • Pipeline Parallelism: Pipeline parallelism injects multiple inputs into inter-layer model partitions so different workers can process different inputs concurrently.Naive pipelining can still create forward-backward weight-version mismatches.
  • GPipe: GPipe maintains one weight version by accumulating gradients and periodically flushing the pipeline, but these flushes can limit throughput.Additional accumulation can mitigate flush overhead but may require large batches and activation stashes proportional to batch size.
  • Execution Timelines: Figure 1 compares pipeline timelines by showing forward passes in blue, backward passes in green, microbatch IDs as numbers, time on the x-axis, and per-worker utilization on the y-axis.The comparison highlights GPipe’s periodic flushes and PipeDream’s multiple weight versions without periodic flushes.
  • PipeDream: PipeDream stashes weight versions so each input uses the same version during forward and backward passes, requiring up to d versions in the worst case.Its default updates also have different delay terms across pipeline stages and perform no within-pipeline accumulation.

3. PipeDream-2BW System Design

PipeDream-2BW combines pipeline parallelism with double-buffered weight updates and gradient accumulation to reduce memory use without expensive pipeline stalls. It also offers PipeDream-Flush, which uses one weight version and periodic flushes for lower memory and vanilla optimizer semantics at lower throughput.

  • PipeDream-2BW Overview: PipeDream-2BW splits models into stages across workers and replicates each stage equally, combining model and data parallelism with input pipelining.Parallel pipelines suit models with repeated structures such as transformer layers.
  • Double-Buffered Weight Updates: 2BW coalesces gradients across microbatches, applies updates at batch granularity, and maintains only two weight versions while avoiding pipeline flushes.New weights serve newly admitted inputs, while in-flight inputs retain compatible versions for backward passes.
  • Double-Buffered Weight Updates: Each microbatch uses the same weight version for forward and backward passes, with version selection determined by its batch and the accumulation parameter m ≥ d.Larger m increases the global batch size.
  • Memory Footprint: With activation recomputation, PipeDream-2BW stores input activations for at most d in-flight microbatches instead of full intermediate activations.Its worst-case footprint is expressed as 2|W| plus d times the input-activation storage.
  • Weight Update Semantics: 2BW’s update semantics use a one-step delay across stages, while the paper states this delay does not significantly affect convergence.The same analysis extends to replication and momentum-based optimizers such as Adam without extra shadow variables.
  • PipeDream-Flush: PipeDream-Flush maintains one weight version and periodically flushes the pipeline, yielding vanilla optimizer semantics and lower memory footprint than 2BW at lower throughput.Its active input activations are bounded by pipeline depth, unlike GPipe’s accumulation-dependent activation storage.

4. Planner

PipeDream-2BW’s planner searches model and hardware configurations using compute and memory cost models. It selects feasible pipeline layouts, microbatch sizes, and accumulation settings while respecting device memory and safe global-batch constraints.

  • Planning Algorithm: The planner exhaustively searches reduced parallel-pipeline configurations that exploit repeated DNN blocks and equal stage replication.This reduces the search space while targeting balanced computation and communication.
  • Cost Modeling: Cost models estimate block-level compute times and memory footprints to evaluate pipeline width, depth, microbatch size, and memory-saving optimizations.The planner can model activation recomputation and the throughput effects of deeper configurations.
  • Activation Recomputation: Activation recomputation lowers memory use by storing only stage inputs and rerunning forward computations when backward-pass activations are needed.The trade-off is extra computation for lower memory footprint, potentially enabling larger per-GPU microbatches.
  • Hardware-Constrained Selection: For total device capacity M, the planner chooses the largest fitting per-GPU microbatch size, verifies the global batch remains below the maximum safe size, and discards over-capacity configurations.It then sweeps pipeline width and depth to select a configuration with a high compute-to-communication ratio.

5. Evaluation

The evaluation compares PipeDream-2BW with model-parallel and pipelined baselines on large BERT and GPT models, measuring convergence, throughput, memory, and scalability. It reports comparable model quality, substantial speedups, and support for larger models under accelerator memory constraints.

  • Throughput: 20× faster training is achieved for a 3.8-billion-parameter GPT model than non-pipelined model-parallel approaches.On a single node, the speedup over tensor model parallelism is 1.3×; across 64 GPUs, communication across servers increases the advantage.
  • Quality of Convergence: 2BW produces comparable downstream accuracy to vanilla Adam while maintaining nearly identical training and validation loss trajectories after the first 100k iterations.BERT finetuning on MNLI and RACE shows similar accuracy, while GPT Wikitext-103 perplexities are 19.28 versus 19.56.
  • Throughput: 6.1× higher throughput than inter-layer model parallelism is achieved, while pipelining with flushes provides up to 5.3× higher throughput on the 2.2-billion-parameter model.PipeDream-2BW is up to 1.6× faster than PipeDream-Flush at low batch size, but the advantage drops to 15% with batch size 2048.
  • Throughput: 3.2× higher throughput than GPipe is achieved at the same global batch size because PipeDream-2BW avoids periodic pipeline flushes.GPipe’s larger activation-stash footprint also limits admitted microbatches and enlarges the pipeline bubble.
  • Scalability and Memory: Almost 30 billion parameters can be trained with PipeDream-2BW using 64 16-GB GPUs, compared with 8.3 billion parameters trained by Megatron-LM using 8 32-GB GPUs.Deeper pipelines provide additional memory capacity, although increasing pipeline depth can reduce throughput by making execution more communication-bound.

6. Related Work and Discussion

The discussion situates PipeDream-2BW among model-parallel, pipeline-parallel, memory-saving, and planning approaches. It emphasizes speed and memory advantages while noting tradeoffs and opportunities for combining orthogonal optimizations.

  • Pipeline Parallelism: PipeDream-2BW trains large models up to 3.2× faster than GPipe because it avoids periodic pipeline flushes and uses less memory for activation stashes.Its lower memory footprint allows more input microbatches to enter the pipeline, whereas PipeDream cannot train these large models.
  • Pipeline Parallelism: PipeMare uses asynchronous weight-update semantics, whereas PipeDream-2BW and the evaluated baselines use synchronous execution.The comparison distinguishes throughput-oriented asynchronous execution from the synchronous semantics used in this paper.
  • Memory-Saving Optimizations: ZeRO partitions optimizer state across data-parallel replicas, and the paper states that such memory-saving optimizations can be combined with PipeDream-2BW.Other related techniques include activation compression and systematic activation recomputation.
  • Planning Algorithms: Existing planners such as PipeDream, DAPPLE, and FlexFlow do not exploit the repetitive structure of modern transformer models or account for memory-saving optimizations.PipeDream’s planner explores O(n^3m^2) configurations, motivating PipeDream-2BW’s more structured planning approach.

7. Conclusion

PipeDream-2BW combines memory-efficient pipeline parallelism with data parallelism-like weight-update semantics and hardware-aware model partitioning. It accelerates billion-parameter model training by up to 20× over model-parallel baselines and up to 3.2× over GPipe on commodity hardware.

  • PipeDream-2BW achieves high throughput, low memory footprint, and data parallelism-like semantics through its 2BW weight-update strategy.The system also uses a planner to partition the model's operator graph over training resources in a memory-aware way.
  • 20× is the maximum acceleration over model-parallel baselines for training models with billions of trainable parameters.
  • 3.2× is the maximum acceleration over GPipe on commodity hardware.

A. Planner, Additional Details

The planner evaluates alternative parallel configurations using throughput and memory models, then selects the fastest configuration that fits accelerator memory. Its models account for computation, communication, activation recomputation, device bandwidth, and pipeline memory.

  • The planner compares non-pipelined, pipelined, and pipelined-with-recomputation executions for each width and depth configuration.These alternatives differ in whether they use pipelining and whether activation recomputation reduces memory use.
  • The planner estimates each execution's throughput and memory footprint, selecting the highest-throughput configuration that fits accelerator memory.Memory capacity is supplied as an input to the planner.
  • The partitioning algorithm searches widths and depths, filters configurations exceeding memory capacity, and returns the best width, depth, microbatch size, recomputation choice, and gradient-accumulation degree.The planner can use profile-based cost functions or closed-form estimates based on finer-grained profile information.
  • The throughput model uses profiled forward and backward times, communication costs, pipeline width and depth, and microbatch size.With pipelining, computation in different stages can overlap, allowing a microbatch to be processed every t seconds.
  • Communication modeling includes activation and gradient exchange between stages, weight-gradient all-reduce across replicas, and bandwidth differences from hierarchical network topologies.For example, communication can remain within a server when the pipeline width is smaller than the number of workers in that server.
  • Activation recomputation increases floating-point operations because forward passes are repeated to reconstruct activation stashes for backpropagation.The model represents this overhead with cextra; cextra = 4/3 is described as a reasonable value when the backward pass takes twice as long as the forward pass.
  • The memory model accounts for partitioned weights and activations, while 2BW maintains two weight versions and either d or one activation version without or with recomputation, respectively.Each pipeline stage has weight parameters of size |W|/d and activations of size |Atotal(b)|/d.

A.2. Partitioning Algorithm

The full partitioning algorithm is presented in pseudocode in Algorithm 1.

  • Algorithm 1 provides pseudocode for the full partitioning algorithm.

B. Evaluation, Additional Graphs

This section presents additional evaluation results that were omitted from the main paper because of space constraints.

  • The section contains additional results that could not fit in the main paper due to space.

B.1. Throughput and Memory Footprint with BERT Models

PipeDream-2BW was evaluated on BERT models with 2.2 billion and 3.8 billion parameters, producing results similar to those observed for GPT models. GPipe fit batch size 64 for BERT but still used more memory than every other baseline.

  • PipeDream-2BW was evaluated on BERT models with 2.2 billion and 3.8 billion parameters.
  • Results for BERT were similar to those for GPT across throughput and memory-footprint comparisons.
  • GPipe fit batch size 64 for BERT, whereas only batch size 32 fit for GPT, producing a larger pipeline bubble for GPT.
  • GPipe still had a higher memory footprint than all other baselines for the BERT models.

B.2. Impact of Activation Recomputation

Activation recomputation trades extra computation for lower peak memory, enabling larger per-GPU microbatches in some GPT configurations. Its throughput effect depends on the model and configuration: it helps in some cases but not others.

  • Activation recomputation can reduce peak training memory, making it possible to train larger models in extreme cases.This benefit comes at the cost of extra computation.
  • Activation recomputation increases throughput in Figure 12b but not in Figure 12a.
  • Figures 10 and 11 report BERT throughput across batch sizes and worst-case memory footprints for an 8-GPU, 2.2B-parameter model.
  • Figure 12 compares throughput against per-GPU microbatch size for PipeDream-2BW configurations using GPT models and 8 16-GB V100 GPUs.The comparison uses a maximum sequence length of 512 and includes configurations with and without activation recomputation.
  • Recomputation enables larger per-GPU microbatches to fit on workers, sometimes producing higher throughput.
  • 33% is the cap on recomputation overhead because backward passes take twice as long as forward passes for most operators.
Loading 2006.09503v3…