PaperlayerView exact revision on arXiv

Source-linked AI summary

Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning

Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P. Xing, Joseph E. Gonzalez, Ion Stoica

arXiv:2201.12023v3cs.LGcs.DCcs.PL

TL;DR

Large-model training requires expertise to tune parallelism across models and clusters, while existing automation covers limited plan spaces. Alpa hierarchically optimizes intra- and inter-operator parallelism; it matches or exceeds specialized systems in reported evaluations.

  • Problem

    Large-model training requires expert tuning across data, operator, and pipeline parallelism, while automated approaches cover limited plan spaces or rely on strong model and cluster assumptions.

    Source: §1 · p. 1

  • Method

    Alpa builds a hierarchical execution-plan space and uses compilation passes to optimize intra-operator execution within device meshes and inter-operator execution across stages and meshes.

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

  • Results

    Alpa matches Megatron-LM on GPT models and beats DeepSpeed on GShard MoE models, with 3.5× speedup on 2 nodes and 9.7× on 4 nodes.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Alpa generates efficient model-parallel plans across specialized and heterogeneous models, including models without manual strategies.

    Source: §1 · p. 2

    The inter-operator pass uses a static linear pipeline schedule and does not optimize its micro-batch count.

    Source: §7 · p. 10

Abstract

from arXiv · show

Alpa automates model-parallel training of large deep learning (DL) models by generating execution plans that unify data, operator, and pipeline parallelism. Existing model-parallel training systems either require users to manually create a parallelization plan or automatically generate one from a limited space of model parallelism configurations. They do not suffice to scale out complex DL models on distributed compute devices. Alpa distributes the training of large DL models by viewing parallelisms as two hierarchical levels: inter-operator and intra-operator parallelisms. Based on it, Alpa constructs a new hierarchical space for massive model-parallel execution plans. Alpa designs a number of compilation passes to automatically derive efficient parallel execution plans at each parallelism level. Alpa implements an efficient runtime to orchestrate the two-level parallel execution on distributed compute devices. Our evaluation shows Alpa generates parallelization plans that match or outperform hand-tuned model-parallel training systems even on models they are designed for. Unlike specialized systems, Alpa also generalizes to models with heterogeneous architectures and models without manually-designed plans. Alpa's source code is publicly available at https://github.com/alpa-projects/alpa

1 Introduction

Large-model training requires expertise to coordinate data, operator, and pipeline parallelism across model and cluster configurations. Alpa organizes these choices hierarchically and reports results matching or exceeding specialized systems on several model types.

1 Introduction

  • Large-model training requires model- and cluster-specific tuning, while automated approaches cover only restricted parallelism spaces or rely on strong assumptions.Effective tuning combines data, operator, and pipeline approaches at tensor-operator granularity and requires ML and systems expertise.Source: §1 · p. 1
  • Alpa recategorizes parallelism as intra-operator and inter-operator, whose distinct communication needs align with cluster bandwidth structure.Intra-operator parallelism improves device utilization but communicates at operator splits and merges; inter-operator parallelism communicates between stages but can leave devices idle.Source: §1 · pp. 1–2
  • Alpa builds a hierarchical plan space and compiler that assigns model stages to device meshes and orchestrates parallelism within and between meshes.Its contributions also include tractable optimization algorithms, compilation passes, a runtime for inter-operator execution, and system optimizations for compilation and cross-mesh communication.Source: §1 · p. 2
  • Alpa matches Megatron-LM on GPT models, beats DeepSpeed by 3.5× on 2-node and 9.7× on 4-node GShard MoE, and reaches 80% linear scaling on 4-node Wide-ResNet.The evaluation used an Amazon EC2 cluster with 8 instances and 64 GPUs.Source: §1 · p. 2

2 Background: Distributed Deep Learning

Distributed deep learning parallelizes computation when a model or dataset is too large for one device. This section contrasts conventional parallelism categories with Alpa's intra-operator and inter-operator view.

2 Background: Distributed Deep Learning

  • Dataflow graphs represent operators and tensor edges; distributed parallelization is used when a model or dataset exceeds what one device can train efficiently.Training an iteration includes a forward pass, a backward pass, and parameter updates.Source: §2 · p. 2

2.1 Conventional View of ML Parallelism

  • Conventional approaches categorize parallelism as data, operator, and pipeline; combining them is interdependent and prior automation covers at most one model-parallel approach alongside data parallelism.Operator parallelism partitions computations across devices, while pipeline parallelism assigns groups of operators to stages and passes activations between workers; hand-designed combinations may assume repeated transformer layers.Source: §2.1 · p. 3

2.2 Intra- and Inter-Operator Parallelisms

  • Alpa recategorizes parallelism as intra-operator, which partitions operators, and inter-operator, which assigns whole operators to distributed devices.Data parallelism, operator partitioning, and weight-update sharding are intra-operator cases; pipeline parallelism is inter-operator. Intra-operator approaches use collective communication, while inter-operator approaches communicate between stages and can leave devices idle; Alpa uses this categorization in an end-to-end system for automatic plan generation.Source: §2.2 · pp. 3–4

3 Overview

Alpa hierarchically optimizes model-parallel execution plans across intra-operator and inter-operator levels. Its compiler selects stage and mesh assignments, then compiles stages and orchestrates their execution across meshes.

3 Overview

  • Alpa optimizes each stage’s intra-operator plan on a device mesh and the inter-operator plan’s stage and mesh assignments, using intra-op costs to guide the higher level.The resulting execution plan combines both levels and is locally near-optimal at each level.Source: §3 · p. 4
  • The inter-op pass uses dynamic programming to assign stages to meshes, querying an intra-op ILP optimizer for each stage-mesh pair’s execution cost.The passes slice the model IR and device cluster into stages and meshes, then select the slicing scheme that minimizes inter-op latency.Source: §3 · p. 4
  • Alpa compiles each stage for its assigned mesh, then generates schedule-specific runtime instructions to coordinate communication and execution across meshes.Developers annotate functions with @parallelize; on the first call, Alpa traces and compiles the function into a parallel version while leaving the rest of the program intact.Source: Figure 4 · p. 4, §3 · p. 5

4 Intra-Operator Parallelism

Alpa searches intra-operator plans within a device mesh, choosing operator algorithms and tensor layouts to reduce graph execution cost. It formulates this choice as an ILP, with modeled compute, communication, and resharding costs.

4 Intra-Operator Parallelism

  • SPMD-style intra-op parallelism partitions operators evenly across equivalent devices, covering data parallelism, ZeRO, Megatron-LM operator parallelism, and combinations.The ILP formulation is intended to handle computational graphs with tens of thousands of operators.Source: §4 · p. 5

4.1 The Space of Intra-Operator Parallelism

  • For each operator, the intra-op pass selects a parallel algorithm and tensor layout; the choice determines compute and communication costs, while incompatible input layouts require resharding.A tensor’s sharding spec marks each axis partitioned or replicated and maps partitioned axes to one or both dimensions of a two-dimensional device mesh.Source: §4.1 · pp. 5–6
  • For batched matrix multiplication, parallelizing a reduction loop requires all-reduce to combine partial sums, unlike an algorithm whose devices already hold their required input tiles.Possible algorithms also differ in input and output sharding specs.Source: §4.1 · p. 6

4.2 ILP Formulation

  • The ILP minimizes graph-wide compute and communication costs at nodes plus resharding costs on edges, choosing one parallel algorithm for each operator.One-hot algorithm choices and edge resharding decisions encode the optimization; the quadratic resharding term is linearized for the ILP solver.Source: §4.2 · p. 6
  • Alpa estimates communication costs from bytes and mesh bandwidth, sets compute costs to zero under its stated workload assumptions, then applies communication optimizations such as reduce-scatter plus all-gather.The zero-compute assumption reflects evenly divided work for heavy operators and negligible computation for lightweight operators.Source: §4.2 · p. 6

5 Inter-Operator Parallelism

Alpa’s inter-operator pass jointly chooses model stages and device meshes, accounting for intra-operator plans at each stage to minimize end-to-end pipeline latency. A dynamic program makes this search tractable through operator clustering and practical optimizations.

5 Inter-Operator Parallelism

  • Alpa jointly chooses stage boundaries, device-mesh assignments, and per-stage intra-operator plans to minimize end-to-end pipeline latency.Earlier work considered simplified problems with preassigned devices or fixed parallelism plans.Source: §5 · p. 7

5.1 The Space for Inter-Operator Parallelism

  • The latency objective sums stage latencies for the first microbatch and adds (B−1) times the slowest stage latency for the remaining microbatches.The formulation also colocates forward and backward operators on the same submesh and requires submeshes to cover the cluster.Source: §5.1 · p. 7

5.2 DP Formulation

  • The search restricts submeshes to one-dimensional shapes and shapes spanning the cluster’s second dimension, which can fully cover the cluster.The authors report that excluded shapes tend to have inferior results because an alternative can use more high-bandwidth-connected devices.Source: §5.2 · p. 7
  • The dynamic program enumerates a maximum stage latency and minimizes total stage latency subject to that bound.Its stage costs come from intra-operator plans that are compiled, profiled, and checked against device-memory requirements.Source: §5.2 · pp. 7–8
  • Early pruning stops the tmax search when B·tmax cannot beat the best solution and skips nearby values, bounding the optimality gap by B·ε.With ε = 10−6 s, the algorithm returned the exact optimum for all evaluated settings.Source: §5.2 · p. 8
  • Operator clustering merges low-compute operators and neighboring operators with costly cross-mesh communication, reducing the graph used by the stage-search objective.The clustering DP balances layer FLOPs within 1+δ times the average while minimizing communication, and runs in O(K^2L) time.Source: §5.2 · p. 9

6 Parallelism Orchestration

Alpa compiles each stage for its assigned mesh and orchestrates communication and execution across meshes. Its cross-mesh resharding plan exploits destination-mesh all-gather, while static instructions are dispatched before execution.

6 Parallelism Orchestration

  • Alpa compiles each stage for its assigned device mesh, following the intra-operator plan and inserting collective communication when needed.Compilation depends on XLA and GSPMD.Source: §6 · p. 9
  • For cross-mesh resharding, Alpa matches source and destination tensor tiles, then replaces repeated sends with one transfer followed by destination-mesh all-gather when the destination sharding replicates the tensor.This local all-gather approach is used because destination devices can exchange data over higher-bandwidth connections.Source: §6 · p. 9, Figure 6 · p. 9
  • Alpa generates distinct static execution instructions for each mesh and dispatches them to workers before execution, avoiding runtime driver-worker coordination overhead.The instructions cover memory allocation, tensor communication, synchronization, and computation.Source: §6 · pp. 9–10

7 Limitations and Discussion

Alpa’s hierarchical design allows heterogeneous stage boundaries, mesh shapes, and within-stage parallelism. The authors report near-optimal plans for many notable models, while identifying limits in cost modeling, scheduling, and supported graphs.

7 Limitations and Discussion

  • Alpa permits uneven stage sizes, different mesh shapes across stages, and operator-specific data and operator parallelism, supporting heterogeneous architectures and cluster setups.The authors contrast this flexibility with manually combined parallelism approaches such as 3D parallelism and PTD-P.Source: §7 · p. 10
  • Alpa does not model cross-stage communication costs, optimize microbatch count or computation-communication overlap, or support dynamic pipeline schedules and graphs without known tensor shapes.The authors note that modeling cross-stage costs would require exponentially more intra-operator passes and DP states.Source: §7 · p. 10
  • Weak-scaling results suggest Alpa generates near-optimal execution plans for many notable models.Source: §7 · p. 10

8 Evaluation

Alpa is evaluated on large models spanning homogeneous and heterogeneous architectures, then compared with specialized baselines and ablated across its optimization stages. Results show strong end-to-end scaling, effective intra- and inter-operator optimizations, and a plan for Wide-ResNet that is difficult to design manually.

8 Evaluation

  • The evaluation covers GPT-3, GShard MoE, and Wide-ResNet on an 8-node, 64-GPU cluster, with models scaled alongside GPU count.The models cover homogeneous and heterogeneous architectures; Megatron-LM is the GPT-3 baseline, while DeepSpeed is used for MoE.Source: §8 · p. 10, §8.1 · pp. 10–11, Table 4 · p. 10

8.1 End-to-End Performance

  • Alpa achieves 3.5× speedup over DeepSpeed on 2 nodes and 9.7× on 4 nodes, while maintaining linear scaling to 64 GPUs with inter-op only.On Wide-ResNet, Alpa reaches 80% scaling on 32 GPUs; the PP-DP and inter-op-only baselines run out of memory, while intra-only does not scale across nodes.Source: Figure 7 · p. 11, §8.1 · pp. 11–12

8.2 Intra-Op Parallelism Ablation Study

  • Auto-sharding performs best in all tested intra-op cases and maintains near-linear scaling, while data parallelism runs out of memory and ZeRO variants can suffer from communication.The comparison disables pipeline parallelism and gradient accumulation and runs on one 8-GPU instance.Source: Figure 8 · p. 12, §8.2 · p. 12

8.3 Inter-Op Parallelism Ablation Study

  • On Wide-ResNet with 32 GPUs, DP outperforms equal-operator and equal-layer clustering by 2.6× and 1.6×, respectively.DP clusters by communication cost and compute balance; equal-layer clustering matches DP on homogeneous GPT but performs worse on heterogeneous Wide-ResNet.Source: Figure 9 · p. 12, §8.3 · p. 13

8.4 Compilation Time

  • Compilation and search for the largest evaluated GPT model take 2,393.26 seconds, compared with more than 40 hours without the reported optimizations.Compilation time grows linearly with model size and cluster GPU count; stage compilation is parallelized and profiling uses an XLA-instruction-level cost model.Source: Table 5 · p. 13, §8.5 · p. 13, §8.4 · p. 13

8.5 Cross-Mesh Resharding

  • Local all-gather cross-mesh resharding brings 2.0× speedup on 32 GPUs by moving more communication from slow connections to fast local connections.The evaluation is on Wide-ResNet and compares the optimization with send/recv alone.Source: §8.5 · p. 13, Figure 11 · p. 13

8.6 Case Study: Wide-ResNet

  • On 16 GPUs, Alpa assigns Wide-ResNet’s three stages 4, 4, and 8 GPUs, using data parallelism in the first two and nontrivial convolution partitioning in the third.The passage notes that manually creating such a strategy for this heterogeneous model can be opaque even for domain experts.Source: §8.6 · p. 13

9 Related Work

Related systems provide manual intra-operator plans, specialized pipeline or data-parallel approaches, or automatic searches with narrower supported spaces. Alpa is positioned as a compiler for a broader distributed-training plan space, while incorporating complementary techniques such as rematerialization and mixed precision.

Systems for model-parallel training

  • TensorFlow, GSPMD, and OneFlow expose APIs for users to manually specify intra-op plans, while pipeline systems such as GPipe and PipeDream use distinct scheduling approaches.PipeDream is asynchronous, unlike Alpa’s synchronous training; Pathway’s SPMD/MPMD runtime split is similar to Alpa’s intra-/inter-op runtime.Source: §9 · p. 14

Automatic search for model-parallel plans

  • Automatic plan-search systems cover narrower settings: Tofu targets single-node linear graphs, FlexFlow lacks pipeline parallelism, and Piper relies on manually designed intra-op strategies.The passage also characterizes FlexFlow’s search as unable to scale to large graphs or clusters and without optimality guarantees.Source: §9 · p. 14

Complementary techniques and compilers

  • Alpa can incorporate memory optimization and low-precision techniques, and differs from single-device-focused DL compilers by supporting distributed-training execution plans.The passage gives rematerialization and mixed-precision training as examples.Source: §9 · p. 14

10 Conclusion

Alpa presents automated model-parallel distributed training through a hierarchical view of intra- and inter-operator parallelism. The authors argue that this approach can reduce the labor involved in planning distributed training and support adoption of emerging large models.

10 Conclusion

  • Alpa builds a hierarchical execution-plan space and compilation passes to derive plans at intra- and inter-operator levels, then orchestrates execution across distributed devices.The conclusion frames automation as a way to address the historically labor-intensive work of planning distributed model-parallel training.Source: §10 · p. 14

A Proof of Submesh Shape Covering

The theorem gives conditions under which selected submesh shapes can cover the full cluster mesh. Its proof first places full-width submeshes, then reduces the remaining coverage problem to power-of-two one-dimensional submeshes.

Theorem and conditions

  • A list of submeshes covers the full (N,M) mesh when its total area is N·M and each shape is either (1,2^p) or has width M.The theorem assumes M is a power of two.Source: A Proof of Submesh Shape Covering · p. 19

Proof

  • The proof places submeshes of width M first, reducing the remaining problem to fitting one-dimensional shapes (1,2^p) into the mesh.This leaves submesh exponents in the range 0 through m−1.Source: A Proof of Submesh Shape Covering · p. 19
  • Induction pairs the remaining (1,1) submeshes into (1,2) meshes, then reduces each exponent and m by one.The proof uses the area equation to show the count of (1,1) submeshes is even in the induction step.Source: A Proof of Submesh Shape Covering · p. 19

B Model Specifications

The specifications fix shared sequence and vocabulary settings for GPT-3 and GShard MoE models, and shared image and class settings for Wide-ResNet models.

B Model Specifications

  • GPT-3 models use sequence length 1024 and vocabulary size 51200; other parameters and GPU counts are listed in Table 6.The table’s last column gives the number of GPUs used to train each model.Source: B Model Specifications · p. 19
  • GShard MoE models use sequence length 1024 and vocabulary size 32000; other parameters and GPU counts are listed in Table 7.The table’s last column gives the number of GPUs used to train each model.Source: B Model Specifications · p. 19
  • Wide-ResNet models use input image size (224, 224, 3) and 1024 classes; other parameters and GPU counts are listed in Table 8.The table’s last column gives the number of GPUs used to train each model.Source: B Model Specifications · p. 19

C Extra Case Study

Alpa’s Wide-ResNet case study visualizes parallelization strategies for 4 and 8 GPUs, showing how data, activations, and weights can be partitioned.

C Extra Case Study

  • Alpa visualizes the parallelization strategies it finds for Wide-ResNet on 4 and 8 GPUs.The appendix also identifies model-specification tables for GPT-3, GShard MoE, and Wide-ResNet.Source: C Extra Case Study · p. 19, Table 6 · p. 19, Table 7 · p. 19
  • In Figure 13, colors mark devices carrying distributed tensors and grey blocks mark tensors replicated across all devices.Inputs and resulting activations can be partitioned along batch and hidden axes; weights along input and output channel axes.Source: Figure 13 · p. 20

Related papers

Loading 2201.12023v3…