PaperlayerView exact revision on arXiv

Source-linked AI summary

PipeMare: Asynchronous Pipeline Parallel DNN Training

Bowen Yang, Jian Zhang, Jonathan Li, Christopher Ré, Christopher R. Aberger, Christopher De Sa

arXiv:1910.05124v2cs.DCcs.LGstat.ML

TL;DR

Synchronous pipeline-parallel methods trade hardware efficiency for statistical efficiency through pipeline bubbles or extra weight memory. PipeMare uses asynchronous updates with learning-rate rescheduling and discrepancy correction; on ResNet50 and Transformer models, it reports competitive accuracy with better hardware utilization than GPipe and PipeDream.

  • Problem

    Existing pipeline-parallel methods preserve synchronous execution by reducing pipeline utilization or incurring extra memory costs.

    Source: §1 · p. 2

  • Method

    PipeMare combines an asynchronous pipeline model with learning rate rescheduling and discrepancy correction to improve statistical efficiency.

    Source: §1 · p. 2

  • Results

    PipeMare achieves competitive model accuracy with better hardware utilization than GPipe and PipeDream on ResNet50 and Transformer models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors present PipeMare as a promising candidate algorithm for a new generation of DNN-training chips.

    Source: §5 · p. 12

    The learning-rate analysis assumes a fixed delay shared across weights and no delay discrepancy, which the paper considers separately.

    Source: §3 · p. 5, §3.2 · p. 7

Abstract

from arXiv · show

Pipeline parallelism (PP) when training neural networks enables larger models to be partitioned spatially, leading to both lower network communication and overall higher hardware utilization. Unfortunately, to preserve the statistical efficiency of sequential training, existing PP techniques sacrifice hardware efficiency by decreasing pipeline utilization or incurring extra memory costs. In this paper, we investigate to what extent these sacrifices are necessary. We devise PipeMare, a simple yet robust training method that tolerates asynchronous updates during PP execution without sacrificing utilization or memory, which allows efficient use of fine-grained pipeline parallelism. Concretely, when tested on ResNet and Transformer networks, asynchrony enables PipeMare to use up to $2.7\times$ less memory or get $4.3\times$ higher pipeline utilization, with similar model quality, when compared to state-of-the-art synchronous PP training techniques.

1 Introduction

Pipeline parallelism partitions neural-network layers across hardware, but synchronous execution can trade away utilization or memory as pipeline depth grows. The paper introduces asynchronous PipeMare and reports competitive accuracy with better hardware utilization than GPipe and PipeDream.

1 Introduction

  • Asynchronous pipeline parallelism partitions neural-network layers across hardware resources while pipelining computation between them.The paper also notes PP can reduce network bandwidth by up to 95%.Source: §1 · pp. 1–2
  • Existing pipeline-parallel methods preserve synchronous execution by adding pipeline bubbles or storing extra weight copies, and these costs increase with pipeline depth.This trade-off challenges high-depth fine-grained PP, which could otherwise improve throughput on new accelerators.Source: §1 · p. 2
  • PipeMare’s asynchronous model avoids pipeline bubbles and substantial increases in weight memory, and its two techniques aim to improve statistical efficiency.The techniques are learning rate rescheduling and discrepancy correction.Source: §1 · p. 2
  • PipeMare achieves competitive model accuracy with better hardware utilization than GPipe and PipeDream on ResNet50 and Transformer models.The paper identifies future high-depth fine-grained PP accelerators as a motivation for this work.Source: §1 · p. 2

1.1 Related Work

  • PipeDream uses weight stashing to keep forward and backward weights synchronized, while GPipe uses microbatching and extra activation memory to preserve synchronization.Hogwild! also computes gradients from delayed weights, but its delays vary randomly rather than remaining fixed as in PP.Source: §1.1 · pp. 2–3

2 Preliminaries

The preliminaries formalize pipeline stages, microbatches, utilization, and delayed weights. They then contrast synchronous PP methods with PipeMare’s asynchronous execution, which avoids weight copies and pipeline bubbles.

2.1 Model of Pipeline Parallelism

  • Pipeline parallelism partitions a network’s L layers into P stages assigned to parallel workers, with each stage processing microbatches of size M within a minibatch.Fine-grained PP has P approximately equal to L; coarse-grained PP has P much smaller than L.Source: §2.1 · p. 3
  • Pipeline utilization is the share of stages active at a given time, while pipeline delay counts optimizer steps between a weight read and its gradient update.Forward and backward delays can differ by stage; synchronous execution requires the same weights in both passes, corresponding to equal delays.Source: §2.1 · pp. 3–4

2.2 Pipeline Parallel Training Methods

  • PipeDream maintains synchronous execution with weight stashing and full utilization, at the cost of extra weight memory that can grow with pipeline depth.GPipe avoids extra weight memory but has bubbles and additional activation memory; microbatching can reduce its bubbles.Source: §2.2 · p. 4
  • GPipe and PipeDream preserve synchronous execution with different costs, while PipeMare updates weights asynchronously to avoid extra weight copies and pipeline bubbles.GPipe instead fills and drains the pipeline at minibatch boundaries, reducing utilization and using additional activation memory.Source: PipeMare · p. 4, §2.2 · p. 4
  • PipeMare uses the weights in memory when gradients are computed, with forward delay matching PipeDream and no delay between backward pass and weight updates.Reducing microbatch size can lower activation memory and the discrepancy between forward and backward delays while keeping stages utilized.Source: PipeMare · pp. 4–5
  • Figure 2 describes how increasing delay can cause divergence at a fixed step size and plots cpusmall linear-regression loss against step size and delay.Its black curve is the upper bound from Lemma 1, using the objective’s largest curvature in place of λ.Source: Figure 2 · p. 5

3 PipeMare

PipeMare uses fixed-delay asynchronous SGD analysis on a quadratic objective to develop learning-rate rescheduling and forward-backward discrepancy correction. The techniques address stability risks associated with delay and delay discrepancy.

3 PipeMare

  • The analysis begins with fixed-delay asynchronous SGD on a one-dimensional convex quadratic, assuming bounded gradient-estimation noise and no delay discrepancy.The no-discrepancy update uses delayed weights in the gradient step.Source: §3 · p. 5

3.1 Learning rate rescheduling (T1)

  • A large pipeline delay can make asynchronous SGD diverge even at a step size where synchronous training converges.The quadratic-model example uses τ=10 and α=0.2; the passage also points to a ResNet50 example in Appendix B.1.Source: §3.1 · p. 6
  • The stable step size scales as O(τ^-1); the threshold also predicts divergence in pipeline-parallel SGD on cpusmall linear regression.The authors report that the same threshold extends to momentum SGD and motivates using the technique with optimizers such as Adam.Source: §3.1 · p. 6

3.2 Discrepancy correction (T2)

  • With ∆=5, D=0.1 discrepancy correction reduces the largest-magnitude companion-matrix eigenvalue toward its no-discrepancy value.The analysis reports that correction seems to increase the range of stable step sizes.Source: Figure 3 · p. 7, §3.2 · p. 8

3.1 Learning rate rescheduling (T1)

  • The rescheduling scheme is motivated by avoiding a uniform 1/τ step-size reduction, which can slow convergence early and persist after the base rate has fallen.The proposed schedule is intended to behave like inverse-delay scaling early, then return toward the baseline learning-rate scheme later.Source: §3.1 · pp. 6–7

3.2 Discrepancy correction (T2)

  • With delay discrepancy, the maximal stable step size remains inversely proportional to delay, but becomes smaller when ∆ is large.The passage relates this to discrepancy worsening divergence.Source: §3.2 · p. 8
  • Forward-backward delay discrepancy can amplify divergence; correction approximates forward-pass weights from their recent trajectory rather than storing their full history.The method uses extra memory for an estimate of weight velocity; the added accumulator tracks changes over time.Source: §3.2 · pp. 7–8

4 Experiments

The experiments evaluate PipeMare’s hardware utilization, memory use, and model quality on image-classification and translation benchmarks. They compare end-to-end performance with synchronous methods and test how learning-rate rescheduling, discrepancy correction, and warmup affect the tradeoffs.

4.2 End-to-End Comparison

  • PipeMare achieves higher pipeline utilization with comparable final model quality than synchronous methods across image-classification and translation tasks.On IWSLT14 and WMT17, PipeMareW achieves 3.2× and 1.7× higher pipeline utilization than GPipe, respectively.Source: Table 2 · p. 9, §4.2 · pp. 10–11

4.1 Experimental Setting

  • Experiments simulate asynchronous updates on ResNet50 and a 12-layer Transformer, partitioned into up to 107 and 93 stages, respectively.Benchmarks use CIFAR10, ImageNet, IWSLT14, and WMT17; evaluation reports pipeline utilization, weight and optimizer memory, and test accuracy or BLEU score.Source: §4.1 · pp. 9–10, Figure 4 · p. 10

4.3 Ablation study

  • Learning-rate rescheduling alone achieves 7.6× higher pipeline utilization than GPipe while reaching 95.0% test accuracy and 34.1 test BLEU.On CIFAR10, its accuracy matches synchronous training, while asynchronous training without rescheduling diverges.Source: Learning rate rescheduling (T1) · p. 11, Table 3 · p. 11
  • Discrepancy correction alone reaches 94.5% test accuracy on ResNet50 but 0.0 test BLEU on Transformer; adding rescheduling raises Transformer BLEU to 34.1.The combined techniques improve convergence speed, especially on IWSLT14, at a weight-memory cost of 33% for SGD with momentum or 25% for Adam.Source: Discrepancy Correction (T2) · pp. 11–12, Table 3 · p. 11
  • On IWSLT14, 10 synchronous warmup epochs raise BLEU from 34.1 to 34.5 while reducing pipeline utilization from 100% to 55%.Warmup was not needed on CIFAR10 to recover GPipe performance.Source: Table 3 · p. 11, Discrepancy Correction (T2) · p. 12

5 Conclusion

PipeMare combines bubble-free asynchronous pipeline parallelism with learning-rate rescheduling and discrepancy correction, improving hardware efficiency while maintaining statistical efficiency. The supplementary analysis compares memory and throughput across pipeline methods and explains when activation-memory reductions are effective.

5 Conclusion

  • PipeMare combines bubble-free pipeline execution with learning-rate rescheduling and discrepancy correction, improving hardware efficiency in utilization and memory over competing algorithms.The conclusion presents these techniques as improving statistical efficiency as well.Source: §5 · p. 12

A Supplementary material for Section 2

  • The appendix compares memory footprint and throughput across methods, treating normalized throughput as linearly proportional to pipeline utilization.Its memory analysis assumes equally partitioned layers and equal per-layer activation memory.Source: A Supplementary material for Section 2 · p. 14

A.1 Activation Memory

  • PipeMare and PipeDream require the same activation memory because both cache activations under the same bubble-free pipeline behavior.The cached activation count at stage i is proportional to O(2(P −i) + 1).Source: §A.1 · p. 14
  • When stages approach the number of layers, checkpointing within each stage yields negligible memory savings; PipeMare Recompute is proposed to reduce memory in this fine-grained setting.The appendix describes the technique as applicable to both synchronous GPipe and asynchronous PipeMare.Source: §A.1 · p. 14

A.2 Trade compute for memory via PipeMare Recompute

PipeMare Recompute trades compute for lower activation memory by recomputing activations across segments of pipeline stages. The surrounding analysis examines asynchronous-training divergence and stability, including the effects of delay discrepancy and momentum.

A.2 Trade compute for memory via PipeMare Recompute

  • PipeMare Recompute recomputes activations across multiple-stage segments, overlapping recomputation with forward and backward operations to avoid stalling backpropagation.The overlap consumes approximately 25% of total compute resources.Source: §A.2 · p. 15
  • For fine-grained pipelines with P = L, segmenting stages into groups of S gives segment memory O(2(P −i) + S2).The first term represents cached activations at a segment's first stage; the second represents buffers for recomputed activations used by backpropagation.Source: §A.2 · pp. 14–15
  • PipeMare Recompute significantly reduces activation memory for fine-grained pipeline parallelism with many stages.The appendix summarizes activation-memory requirements for GPipe, PipeDream, and PipeMare; its table caption identifies task-specific savings for PipeMare.Source: §A.2 · p. 16, Table 5 · p. 16

B.1 Motivating examples in deep learning

  • Large fixed delays can cause pipeline-parallel SGD to diverge, and forward-backward delay discrepancy can exacerbate that divergence.The ResNet50 discussion reports divergence with equal forward and backward delays, and further divergence when discrepancy is added.Source: §B.1 · pp. 16–17

B.2 Proof of Lemma 1

  • In the stability proof, roots of p start inside the unit disk for small α and cross the unit circle as α increases, after which they cannot all remain inside.The argument counts τ + 1 crossing opportunities for τ + 1 roots.Source: §B.2 · pp. 18–19

B.3 An extension to SGD with momentum.

  • Adding momentum does not remove the O(τ −1) step-size requirement under the analysis, motivating Technique 1 for algorithms beyond plain SGD.The extension assumes constant τ and quadratic loss.Source: §B.3 · pp. 19–20

B.4 Proof of Lemma 2

The appendix gives stability arguments for delayed asynchronous updates and justifies the discrepancy correction used in PipeMare. Its numerical experiments report that T2 generally widens the stable step-size range for nonnegative discrepancy sensitivity, with occasional negative effects outside that range.

B.4 Proof of Lemma 2

  • For every ∆≥0, the proof establishes an α at which a polynomial root lies on the unit circle, while α has no absolute lower bound as ∆ varies.The argument maps nonnegative ∆ values to θ and notes that α approaches zero as θ approaches zero.Source: §B.4 · pp. 24–25

B.5 Justification for Claims in Section 3.2

  • The second-order expansion around ω = 1 becomes invariant to ∆ when γ is set appropriately, motivating the discrepancy-correction parameter choice.The appendix also justifies focusing on ω = 1 as delays grow.Source: §B.5 · pp. 25–27
  • In quadratic-model experiments, T2 consistently enlarges the stable α range for ∆>0, but can occasionally reduce it for ∆≤0.Figure 7 illustrates a case with τfwd = 40 and τbkwd = 10.Source: §B.5 · p. 27, Figure 7 · p. 27

C Supplementary material for Section 4

  • The supplementary section introduces additional ImageNet and WMT results alongside further PipeMare ablation results.It directs readers to Appendix C.2 for those additional results.Source: C Supplementary material for Section 4 · p. 28

C.1 Experiment setup

The appendix documents model, dataset, and hyperparameter choices for ResNet and Transformer experiments, including PipeMare tuning. It also identifies additional ImageNet and WMT results and supplementary ablations.

C.1 Experiment setup

  • ResNet experiments use CIFAR10 and ImageNet, while 12-layer Transformer experiments use IWSLT14 and WMT17 German-to-English.The setup inherits most model and training hyperparameters from the respective implementations, with specified learning-rate adjustments.Source: §C.1 · p. 28
  • PipeMare tunes annealing epochs, discrepancy-correction decay, and warmup epochs sequentially, then re-sweeps annealing epochs for CIFAR10 and IWSLT14.Warmup epochs are set to zero for CIFAR10 because they did not improve statistical efficiency.Source: §C.1 · pp. 28, 30, Table 8 · p. 29
  • For ImageNet and WMT, PipeMare hyperparameters are transferred from CIFAR10 and IWSLT with minimal search centered around those settings.The appendix also identifies additional ImageNet and WMT results and supplementary ablation results.Source: §C.1 · p. 30, §C.2 · p. 30

C.2.1 ImageNet and WMT results

On ImageNet, PipeMare achieves higher test accuracy than PipeDream; on WMT, it reaches competitive BLEU to GPipe while PipeDream fails to converge. PipeMare and PipeMareW also achieve time-to-accuracy speedups over GPipe.

C.2.1 ImageNet and WMT results

  • PipeMare achieves higher ImageNet test accuracy than PipeDream and competitive WMT BLEU to GPipe, while PipeDream attains 0.0 BLEU on WMT.Figure 9 compares model accuracy and normalized time on ImageNet and WMT.Source: §C.2.1 · p. 30, Figure 9 · p. 31
  • PipeMare and PipeMareW achieve time-to-accuracy speedups over GPipe on ImageNet and WMT.The comparison uses normalized time.Source: Figure 9 · p. 31
  • PipeMareW attains higher model accuracy on ImageNet and WMT while remaining competitive with GPipe in the same number of epochs.Figure 9 presents accuracy and normalized-time comparisons for both datasets.Source: Figure 9 · p. 31

C.2.2 PipeMare ablation study

The ablations show discrepancy correction can prevent ResNet 152 divergence, and PipeMare’s techniques contribute more to statistical and hardware efficiency at twice the pipeline stages.

C.2.2 PipeMare ablation study

  • At 214 CIFAR10 and 186 IWSLT stages, learning-rate rescheduling, discrepancy correction, and warmup show larger statistical- and hardware-efficiency contributions than at half as many stages.The comparison is against 107 and 93 stages, respectively.Source: §C.2.2 · p. 31
  • With 150 pipeline stages on CIFAR10, PipeMare T1 diverges, whereas adding discrepancy correction lets it converge and match GPipe test accuracy.The match is reported after the first learning-rate drop, following 80 epochs.Source: §C.2.2 · p. 31
  • Discrepancy correction prevents divergence for ResNet 152 on CIFAR10, where learning-rate rescheduling alone is insufficient.The result extends the reported correction benefit to a model with more pipeline stages.Source: §C.2.2 · p. 31

C.2.3 Hyperparameter sensitivity studies

Hyperparameter choices affect convergence, accuracy, and throughput, while recomputation introduces additional delays that PipeMare’s discrepancy correction can address. The appendix also reports recompute accuracy and memory benefits.

C.2.3 Hyperparameter sensitivity studies

  • ResNet prefers fewer annealing epochs than Transformer for optimal test performance.Different models may require different annealing-epoch settings.Source: §C.2.3 · p. 31
  • For Transformer, correction decay D ≤0.2 can speed convergence, while an improper D can perform worse than no correction.The correction also requires an extra memory buffer and accumulation, adding 25–33% to total weight memory.Source: §C.2.3 · p. 33
  • More synchronous warmup epochs lower pipeline utilization and throughput but often help statistical convergence.Figure 13 examines the warmup-epoch tradeoff on IWSLT.Source: §C.2.3 · p. 33, Figure 13 · p. 33

D Statistical Efficiency and Recompute

  • Activation recomputation can reduce memory footprint while retaining matching or competitive accuracy relative to PipeMare without recomputation.The appendix frames recomputation as a way to reduce the memory required to store activations.Source: D Statistical Efficiency and Recompute · p. 34

D.1 Asynchronous pipeline-parallel recomputation

  • Recomputation adds delayed recomputed activations with delays distinct from backward-pass weights, and T2 extends correction to those activations using the same trajectory accumulator.The recomputation weights have a fixed delay, and the correction adjusts their assignment.Source: §D.1 · p. 34
  • The recompute stability model separates forward, recomputation, and backward delays, with Φ measuring gradient sensitivity to recomputed-versus-backward weight discrepancy.The model assumes τfwd > τrecomp > τbkwd.Source: §D.1 · p. 34
  • With recomputation, discrepancy correction increases the quadratic model’s stable step-size range and brings its behavior closer to the no-delay-discrepancy case.Figure 14 shows correction reducing the largest-magnitude eigenvalue for the illustrated model.Source: §D.1 · p. 35, Figure 14 · p. 35

D.2 Statistical efficiency and recompute

With discrepancy correction, PipeMare with recomputation generally matches or competes with its accuracy without recomputation across CIFAR10 and IWSLT. Correction is important for stable recompute training, which can reduce activation memory.

D.2 Statistical efficiency and recompute

  • On IWSLT, recompute without discrepancy correction can be unstable, while applying correction allows different checkpoint counts to match accuracy without recomputation.Without correction, 2 checkpoints failed to exceed BLEU 10.0 and 12 checkpoints diverged mid-training.Source: §D.2 · pp. 35, 37, Figure 16 · p. 36
  • Across CIFAR10 and IWSLT, PipeMare with recomputation generally matches or competes with the accuracy of PipeMare without recomputation.The comparison covers different numbers of gradient checkpoints.Source: §D.2 · p. 35, Figure 15 · p. 36, Figure 16 · p. 36
  • With discrepancy correction, asynchronous PipeMare can retain strong CIFAR10 and IWSLT accuracy while recomputation reduces activation-memory footprint.The reported accuracy comparison is against PipeMare without recomputation.Source: §D.2 · p. 37

E Hogwild! asynchrony

PipeMare’s learning-rate rescheduling heuristic also improves test performance under Hogwild!-style stochastic asynchrony, extending beyond the fixed-delay pipeline setting. The evaluation varies gradient delays by stage and tests ResNet50 on CIFAR10 and Transformer on IWSLT14.

Stochastic asynchrony model

  • Hogwild!-style asynchrony updates the model using a stochastic gradient computed from a stale iterate, with random delay modeling network transmission or shared-memory updates.The section asks whether the rescheduling heuristic can help beyond pipeline execution with a fixed gradient-delay pattern.Source: E Hogwild! asynchrony · p. 37
  • The evaluated variant assigns each stage its own stochastic gradient delay, with each stage’s update using the gradient dimensions corresponding to that stage.This stage-specific delay setting is also studied in the paper’s fixed-delay asynchronous pipeline training.Source: E Hogwild! asynchrony · p. 37
  • Stage delays are sampled from truncated exponential distributions bounded at τmax, using different expectations to model different delay levels across stages.The passage identifies the truncated exponential as a maximum-entropy distribution and says the same truncation is used uniformly across stages.Source: E Hogwild! asynchrony · p. 37

Evaluation results

  • The evaluation uses 107 ResNet stages and 93 Transformer stages, and inherits the optimal annealing-epoch configuration from PipeMare T1.The models are ResNet50 on CIFAR10 and Transformer on IWSLT14 German-to-English translation.Source: E Hogwild! asynchrony · p. 37
  • With learning-rate rescheduling, ResNet50 test accuracy rises from 94.51% to 94.80% and Transformer test BLEU from 3.6 to 33.8 under Hogwild!-style asynchrony.Figure 17 shows rescheduled training matching synchronous test performance and outperforming asynchronous training without rescheduling after the same number of epochs.Source: E Hogwild! asynchrony · pp. 37–38, Figure 17 · p. 38
Loading 1910.05124v2…