PaperlayerView exact revision on arXiv

Source-linked AI summary

Varuna: Scalable, Low-cost Training of Massive Deep Learning Models

Sanjith Athlur, Nitika Saran, Muthian Sivathanu, Ramachandran Ramjee, Nipun Kwatra

arXiv:2111.04007v2cs.DC

TL;DR

Massive-model training relies on costly specialized hyperclusters, whose scale and fragmented resources constrain access to GPUs. Varuna combines pipeline and data parallelism to train on commodity networks and spot VMs; it reports up to 18x higher training performance than state-of-the-art approaches while reducing cost.

  • Problem

    Training massive models depends on expensive specialized hyperclusters, which constrain parallelism to cluster size and fragment GPU resources across clusters.

    Source: §1 · p. 1

  • Method

    Varuna combines pipeline and data parallelism and morphs job configurations to use available GPUs over commodity networking.

    Source: §1 · pp. 1–2

  • Results

    Up to 18x higher training performance than state-of-the-art approaches is reported for massive models trained on low-priority VMs.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Varuna’s use of commodity VMs reduces massive-model training cost and allows jobs to scale opportunistically across more GPUs.

    Source: §8 · p. 13

    Pipeline schemes including Varuna require mini-batches about 6x larger for efficiency.

    Source: §7.3 · p. 12

Abstract

from arXiv · show

Systems for training massive deep learning models (billions of parameters) today assume and require specialized "hyper-clusters": hundreds or thousands of GPUs wired with specialized high-bandwidth interconnects such as NV-Link and Infiniband. Besides being expensive, such dependence on hyper-clusters and custom high-speed inter-connects limits the size of such clusters, creating (a) scalability limits on job parallelism; (b) resource fragmentation across hyper-clusters. In this paper, we present Varuna, a new system that enables training massive deep learning models on commodity networking. Varuna makes thrifty use of networking resources and automatically configures the user's training job to efficiently use any given set of resources. Therefore, Varuna is able to leverage "low-priority" VMs that cost about 5x cheaper than dedicated GPUs, thus significantly reducing the cost of training massive models. We demonstrate the efficacy of Varuna by training massive models, including a 200 billion parameter model, on 5x cheaper "spot VMs", while maintaining high training throughput. Varuna improves end-to-end training time by up to 18x compared to other model-parallel approaches and up to 26% compared to other pipeline parallel approaches. The code for Varuna is available at https://github.com/microsoft/varuna.

1 Introduction

Varuna targets the cost, scale, and resource-utilization limits of training massive models on specialized hyperclusters. It combines pipeline and data parallelism, adapts jobs to changing resources, and reports large-model training on cheaper low-priority VMs.

1 Introduction

  • Specialized hyperclusters are expensive, constrain parallelism to cluster size, and fragment GPU resources across separate clusters.The passage estimates fully loaded BERT-large training costs at $200K and describes limits to scaling high-speed interconnects.Source: §1 · p. 1
  • Varuna trains massive models over commodity networking without specialized hyperclusters and uses discounted spot VMs to reduce cost by 4–5x without losing performance.Spot VMs use spare capacity, and Varuna can orchestrate jobs across available GPUs in a data center.Source: §1 · p. 1
  • Varuna combines pipeline and data parallelism, schedules many micro-batches to tolerate network jitter, and constrains stage sizing by bandwidth.It also demonstrates equivalent accuracy for a 2.5-billion-parameter GPT-2 model despite a 16x larger mini-batch.Source: §1 · p. 2
  • Varuna morphs jobs across pipeline depth and data parallelism as resources change, without changing hyperparameters, using profiling-driven simulation to select configurations.It also addresses failures and stragglers in spot-VM settings.Source: §1 · p. 2
  • Varuna eases model partitioning by automatically identifying cut-points and tracing cross-partition sharing so shared tensors can be synchronized.The tracer detects sharing in user code and libraries.Source: §1 · p. 2
  • Varuna trains models up to 200 billion parameters on low-priority VMs and improves training performance by up to 18x over state-of-the-art approaches.Evaluated models include BERT-large and GPT-2 models ranging from 2.5 billion to 200 billion parameters.Source: §1 · p. 2

2 Background and Related Work

The background distinguishes data-parallel training from model-parallel approaches and reviews systems with different trade-offs in synchronization, partitioning, memory, and resource elasticity.

2 Background and Related Work

  • A DLT iteration runs a forward pass, computes parameter gradients in a backward pass, and updates parameters; data parallelism replicates the model and synchronously averages gradients.Model parallelism instead spreads one model’s parameters across GPUs, using inter-layer or intra-layer partitioning.Source: §2 · pp. 2–3
  • Gpipe preserves synchronous SGD and uses micro-batches to reduce pipeline overhead, but can incur stalled-stage bubbles and requires users to provide a partitioned model.Source: §2 · p. 3
  • PipeDream reaches 100% steady-state pipeline utilization but uses stale updates that can hinder convergence and retains parameter copies that limit model size.PipeDream-2BW reduces copies to two but still uses stale updates; PipeMare introduces approximate corrections for stale updates and uses 25–33% more memory than Gpipe.Source: §2 · p. 3
  • DAPPLE preserves synchronous SGD but stores intermediate activations across micro-batches, while intra-layer partitioning depends on fast, low-jitter links and gradient checkpointing adds about 33% compute overhead.The largest model for which DAPPLE reports speedup is BERT-48 with 600 million parameters.Source: §2 · p. 3
  • Existing elasticity systems often leave resource adaptation to users or focus on single-GPU jobs, while Varuna automatically scales massive jobs using both model and data parallelism.Proteus is described as specific to parameter-server frameworks, and Spotnik as focused on smaller image models.Source: §2 · p. 4

3 Design Overview

Varuna combines pipeline model parallelism with data-parallel replicas and tunes the design for commodity-network constraints. Its schedule reduces stalls while retaining synchronous SGD and accommodating network jitter.

3.1 High-level architecture

  • Varuna combines pipeline stages across model partitions with multiple data-parallel replicas per partition, passing activations forward and gradients backward between stages.It recomputes activations for backpropagation rather than retaining them, and pipeline communication is about 300x smaller than intra-layer communication for the cited GPT-2 2.5B example.Source: §3.1 · p. 4
  • Varuna balances pipeline depth against the increased all-reduce bandwidth and network load that come with fewer stages, while using micro-batches and profiling to handle latency and jitter.The design also exploits the greater availability of single-GPU spot VMs; the passage reports only a 2% performance difference versus four-GPU VMs.Source: §3.1 · p. 5

3.2 Tuning pipeline efficiency

  • For a four-stage, five-micro-batch schedule, Varuna uses one fewer time unit than Gpipe, spreads idle slots through the schedule, and skips last-stage recomputation.Avoiding last-stage recomputation lets the final stage accommodate less-compute-intensive language-model embedding layers without upsetting pipeline balance.Source: Figure 4 · p. 5, §3.2 · p. 5
  • The schedule combines a static rule-based plan with opportunistic task selection when expected gradients have not arrived, keeping stages work-conserving during network jitter.Its rules prioritize backward work and avoid recomputation choices that would double activation memory.Source: §3.2 · p. 5
  • Varuna’s schedule is reported as more efficient than Gpipe, DeepSpeed, and Megatron pipeline schedules in commodity-network settings.Source: §3.2 · p. 6

4 Handling pre-emptions: Job Morphing

Varuna morphs training jobs to use changing sets of pre-emptible GPUs while keeping the user’s total mini-batch fixed. It combines one-time calibration and simulation with checkpointing and manager-driven monitoring to select and sustain configurations.

4 Handling pre-emptions: Job Morphing

  • Job morphing dynamically reconfigures training for available resources, using scale-invariant calibration and parametrized simulation to find a configuration.The profiling state space is kept small and calibration is done once at startup rather than repeated when GPU counts change.Source: §4 · p. 6

4.1 Problem setting and constraints

  • The configuration selects pipeline depth P and data-parallel replicas D within the available G GPUs, subject to cut-point and mini-batch constraints.The total mini-batch is the per-GPU micro-batch size multiplied by the number of micro-batches and GPUs; the micro-batch must balance operation efficiency against GPU memory.Source: §4.1 · p. 6

4.2 Correctness-preserving morphing

  • Varuna keeps the user-specified total mini-batch fixed across configurations and uses gradient accumulation to fit it when resources shrink.When the resulting micro-batch would exceed GPU memory, Varuna increases the number of micro-batches.Source: §4.2 · p. 6

4.3 Scale-invariant Calibration

  • Scale-invariant calibration measures orthogonal, configuration-independent primitive parameters once, in parallel where possible, rather than profiling each full configuration.Calibration takes under a minute even for a 10-billion-parameter model and is independent of the total GPU count.Source: §4.3 · pp. 6–7

4.4 Parametrized Simulation

  • The event-driven simulator estimates time per mini-batch for each configuration, with simulated times within 5% of measured times.The search limits exploration to O(G) configurations by sweeping pipeline depths and balancing stage assignments.Source: §4.4 · p. 7
  • Table 3 compares pipeline depths for 2.5B GPT-2; deeper pipelines can outperform shallower ones, and the optimal depth varies with G.With 100 GPUs, 6-way pipelining uses 96 GPUs while 9-way uses 99, making the total-throughput difference larger than the per-GPU difference.Source: Table 3 · p. 7, §4.4 · p. 7

4.5 Continuous Checkpointing

  • Varuna checkpoints each layer independently at mini-batch boundaries, enabling a resumed job to use a new layer-to-stage mapping after reconfiguration.Checkpoints are written to local SSD and copied to cloud storage in the background.Source: §4.5 · p. 7

4.6 Varuna Manager

  • The manager uses task heartbeats to omit slow GPU replicas, detect pre-emptions and trigger morphing, and periodically attempts to grow the cluster.A slow GPU can delay the synchronous job, so the manager tracks forward- and backward-pass compute times per micro-batch.Source: §4.6 · p. 7

5 Enabling Ease of Programming

Varuna aims to make pipeline model parallelism usable without requiring developers to rewrite their models. It automatically identifies partition points and tracks dependencies that cross those boundaries.

5 Enabling Ease of Programming

  • Existing pipeline systems can require model rewrites or library-specific implementations, making automatic partitioning challenging for complex models.The paper notes that PipeDream relies on a whitelisted set of PyTorch operations and that a standard BERT implementation crashes on it.Source: §5 · p. 8

5.1 Auto-partitioning

  • Varuna profiles execution times and activation sizes to identify fine-grained cut-points, then groups them at runtime into pipeline partitions suited to available resources.It favors endpoints with low activation sizes and checks for parameter sharing across cut-point boundaries.Source: §5.1 · p. 8

5.2 Tracking cross-partition dependencies

  • To preserve cross-partition dependencies such as tied embeddings or global optimizer state, Varuna traces a dry run and synchronizes tensors flagged as shared.The tracer flags functions that use tensors from multiple partitions and presents potential violations for the user to mark as shared.Source: §5.2 · p. 8

6 Implementation

Varuna implements automatic model partitioning and pipelined execution in PyTorch. Its implementation also supports communication overlap, synchronization, checkpointing, and recovery from pre-emption.

Implementation

  • At each activated cut-point, Varuna sends activations to the next rank in the forward pass and gradients to the previous rank in the backward pass.Separate send and receive threads use asynchronous communication so modules can compute when tensors are available.Source: §6 · pp. 8–9
  • Varuna uses one process group for gradient all-reduce across data-parallel replicas and another to synchronize state shared across pipeline partitions.The implementation also handles random-number state, scheduling, checkpoint recovery, fail-stutter monitoring, and cluster growth in about 2,400 Python and 550 C++ lines.Source: §6 · p. 9

7 Evaluation

Varuna is evaluated against model-parallel and pipeline-parallel systems on low-priority commodity VMs and hyperclusters, with tests of scaling, resource changes, and convergence. It delivers strong throughput, adapts to changing VM availability, and reaches GPT-2 2.5B’s reported accuracy in less time.

7.1 Performance & Cost

  • 18x better training speed than Megatron on the same commodity low-priority VMs; Varuna on spot VMs also runs 17% faster than Megatron on a hypercluster.For GPT-2 8.3B, Varuna on spot VMs achieves 0.56 ex/s/GPU versus Megatron’s 0.48 on a hypercluster, with 5.85x better cost-performance.Source: §7.1.1 · p. 10, Figure 5 · p. 10
  • Varuna on commodity VMs is 78% faster than Megatron on a hypercluster for GPT-2 20B, while forcing Megatron to 20B with 18-way partitioning drops its performance 10x.Megatron on the hypercluster fits only a 19.2B model with 16-way model parallelism; Varuna’s 20B configuration uses 294 commodity-VM GPUs.Source: §7.1.1 · pp. 10–11, Table 4 · p. 10

7.1.2 Comparison with other Pipelining architectures.

  • Varuna delivers 15–70% better performance than GPipe and 20–26% better performance than DeepSpeed, plus 13–14% over Megatron-1F1B.For GPipe, the comparison uses a 4-stage pipeline; the single-GPU multi-node comparisons disable intra-layer parallelism and other orthogonal optimizations.Source: §7.1.2 · p. 11, Table 5 · p. 11, Table 6 · p. 11

7.1 Performance & Cost

  • With 5.1x more GPUs, from 56 to 288, Varuna’s per-GPU performance drops only about 7.5%; its TFlops/s/GPU remains roughly the same from a 2.5B to a 200B model.The authors note that opportunistic spot VMs can provide more GPUs for short periods, enabling faster completion at similar dollar cost.Source: Figure 5 · p. 10, §7.1.3 · pp. 11–12

7.2 Auto configuration and Job Morphing

  • Simulator estimates of mini-batch time are within 5% of actual values, and simulation takes 376–660 ms for the reported 128-GPU configurations.The simulator is profile-driven, and its sub-second runtime is described as quick enough to react to changes in spot-VM availability.Source: §7.2 · p. 12, Table 7 · p. 12
  • On 72 GPUs, Varuna reaches 1.77 ex/s/GPU on 1-GPU VMs versus 1.81 on 4-GPU VMs; over 60 hours, per-GPU performance varies only 15% as throughput changes fivefold.Varuna updates its configuration as available VMs change; periodic throughput spikes are checkpointing events.Source: §7.2 · p. 12, Figure 8 · p. 12

7.3 Accuracy and Time-to-Convergence

  • Varuna reaches the same reported validation perplexity of 10.81 as Megatron on GPT-2 2.5B while improving time-to-convergence by 4.1x on commodity VMs.Varuna uses batch size 8192 and 18.75K iterations, processing the same number of examples as Megatron’s 300K-iteration, batch-size-512 run.Source: Figure 9 · p. 12, §7.3 · p. 13

8 Conclusion

Varuna makes training deep learning models with billions of parameters possible on commodity VMs, reducing training cost. Without reliance on specialized hyperclusters, jobs can also scale opportunistically to more GPUs.

8 Conclusion

  • Varuna enables training deep learning models with billions of parameters on commodity VMs, reducing training cost by a significant factor.The authors argue that low cost, high performance, and higher scale can accelerate innovation in large-scale AI models.Source: §8 · p. 13
  • Removing dependence on specialized hyperclusters allows training jobs to scale opportunistically to much larger numbers of GPUs.Source: §8 · p. 13
  • The authors expect Varuna’s combination of low cost, high performance, and higher scale to significantly accelerate innovation in large-scale AI models.Source: §8 · p. 13

9 Appendix

The appendix examines PipeDream-2BW convergence on GPT2, reporting its WikiText perplexity and a training run that diverged. It also notes the trade-off between pipeline throughput and synchronous SGD semantics.

9.1 PipeDream-2BW Convergence

  • PipeDream-2BW boosts pipeline training throughput by sacrificing synchronous SGD semantics, which can lead to sub-optimal accuracy.The appendix notes that stale backward updates leave final accuracy effects unclear.Source: §9.1 · p. 14
  • PipeDream-2BW’s reported WikiText PPL was 19.56, versus Nvidia’s 19.28 and Varuna’s 19.16 at 8k batch size (lower is better).Source: §9.1 · p. 14
  • After 16k iterations, PipeDream-2BW training of a 355M GPT2 model diverged and its loss rose, as shown in Figure 10.The run used 48 GPUs and a batch size of 512.Source: Figure 10 · p. 14, §9.1 · p. 14

Related papers

Loading 2111.04007v2…