Piper: Multidimensional Planner for DNN Parallelization

1 Introduction . Deep Neural Network ( DNN ) models have grown exponentially in size over the past two decades ; consequently , modern DNNs are extremely computationally expensive . For example , a state-of-theart language model , GPT-3 [ 2 ] , has an astounding 175 billion parameters , requiring 314 ZettaFLOPS ( 3.14× 1023 FLOPS ) of compute . This trend of computationally expensive models with working-set sizes for training and inference that far exceed the memory capacities of individual DNN accelerators ( e.g. , GPUs , TPUs , FPGAs , ASICs ) has reinvigorated the study of parallel training techniques that split DNN model state across different devices . In the “ modern era ” , such model-parallel training techniques trace their roots back to AlexNet [ 14 ] and early influential systems such as DistBelief [ 6 ] and Project Adam [ 3 ] . Recent advances in model-parallel training include tensor ( intra-layer ) model parallelism [ 22 ] , where individual operators of the model are partitioned over multiple workers , and pipeline model parallelism that combines inter-layer model parallelism with pipelining [ 17 , 18 , 11 ] . Various hybrid parallelism techniques that combine tensor , pipeline , and data parallelism across operators [ 13 , 19 ] have also been developed to more efficiently train a wider range of models . Additionally , memorysaving optimizations like activation recomputation , which trade off throughput for memory , are also being used [ 18 ] to scale training to larger models such as GPT-3 . Combining these dimensions , however , is non-trivial [ 19 ] , since each dimension has trade-offs with respect to computational efficiency , amount of communication , and memory footprint . Given the importance of efficient model-parallel training ( and inference [ 7 , 4 ] ) , partitioning a model across 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . devices has thus received recent interest , evolving from a manual process driven by human experts to using algorithms for automated device placement of DNN operators . Focusing on algorithms , some take a black-box approach using learning-based techniques [ 16 , 15 , 8 , 1 , 26 , 21 , 25 ] ; others build a cost model and algorithmically solve an optimization problem [ 12 , 17 , 18 , 13 , 23 ] . However , each of the prior algorithmic approaches tackle a slightly different subset of parallelization techniques , system constraints , and DNN optimizations ; for example , FlexFlow [ 13 ] does not consider pipelining , PipeDream [ 17 ] does not consider memory-saving optimizations that are useful for training large models , and PipeDream and PipeDream-2BW [ 18 ] both do not consider tensor parallelism . Our contribution . In a first , this paper covers a broad search space of parallelization and optimization techniques while determining placement of DNN layers on the provided hardware deployment ( a given number of accelerators of a specific type with certain performance characteristics , memory capacity , and interconnect bandwidth ) . Specifically , our optimization algorithm , called Piper , considers three different dimensions of parallelism ( data , tensor model , and pipeline model parallelism ) , and a general framework of memory-saving optimizations that trade off throughput for lower memory footprint . In this work , we consider pipeline parallelism schedules without pipeline flushes [ 17 , 18 ] . We leave extensions to pipeline parallelism schedules with flushes [ 11 , 19 ] to future work . The focus of our work is on algorithms for DNN partitioning across devices to maximize throughput ( for training and large-batch inference of DNN models ) . Piper uses non-trivial combinations of data , tensor , and pipeline parallelism , as well as memory-saving optimizations such as activation recomputation , carefully chosen for each stage of the partition . Piper supports general DAG ( Directed Acyclic Graph ) model topologies and is agnostic of the type of DNN accelerator used . Piper is given a target “ safe ” batch size ( e.g. , 1920 for large GPT-style models [ 20 ] ) and returns a training configuration with this batch size . It is efficient both in theory and in practice . In Section 5 , we evaluate Piper on real-world DNN profiles , and study the effects of combining the various parallelism modes and memory-saving optimizations on performance . We find that Piper ’ s large search space combined with its principled algorithmic approach allows it to find non-trivial high-quality parallelization configurations for each given problem instance . Piper compares favorably to planners from prior work ( PipeDream , PipeDream-2BW ) . For example , we find ( as expected ) that in many settings , pipeline model parallelism combined with per-shard data parallelism is enough to obtain high throughput ; however , tensor parallelism is also required when one needs a batch size smaller than the number of devices , as well as in highly memory-constrained settings . Piper ’ s ability to form heterogeneous stages also gives it an advantage over prior work . 2 Background . In this section , we briefly describe methods to train models at scale . 2.1 Pipeline model parallelism . Pipeline model parallelism is a recent technique [ 17 , 11 ] used to accelerate the distributed training of both small and large models . The operators of a model are partitioned over the available devices ( i.e. , each device hosts a subset of the operators of the model ; typically these subsets are contiguous ) . A batch is then split into smaller microbatches , and execution is pipelined across these microbatches : while the second device starts executing the first microbatch , the first device starts executing the second microbatch , and so on . Pipeline model parallelism features cheap point-to-point communication . A key problem with deploying pipeline parallelism is determining how to schedule forward and backward microbatches on the available devices , and how to stash state to ensure a final high-quality model . Different schedules have different throughput , memory footprint , and semantics trade-offs . For example , GPipe [ 11 ] maintains just a single weight version , but introduces periodic pipeline flushes where microbatches in the pipeline are allowed to drain . This leads to idle resources and lower throughput , but also lower memory footprint . PipeDream [ 17 ] maintains a weight version for every in-flight input microbatch , but does not use pipeline flushes , leading to both high throughput and memory footprint . PipeDream-2BW [ 18 ] tries to find a middle ground , where gradients are coalesced across microbatches , reducing the number of weight versions , while still keeping throughput high . In this work , we assume that a non-flushing pipelining scheme is used when considering pipelines with multiple stages ( i.e. , we do not use inter-layer model parallelism without pipelining ) . 2.2 Tensor ( intra-layer ) model parallelism . A model can be partitioned in other ways as well . Megatron-LM [ 22 ] slices transformer blocks in NLP models like BERT and GPT across available devices , in effect performing distributed matrix multiplications . All-to-all communication in the form of scatters , gathers and all-reductions needs to be performed to ensure that these distributed operators retain their semantics . Megatron-LM uses a specific , hand-tuned strategy to distribute transformer layers across multiple GPUs . Tensor parallelism reduces per-device memory usage , as model parameters are split across devices , and shows good scaling performance when using high-bandwidth interconnects like NVLink [ 19 ] . Trying to find an optimal tensor parallelization strategy for the entire model by deciding how to split every operator is a highly challenging problem with a large search space , and requires understanding the semantics of each operator . Modern models consist of a large number of operators , often with considerable branching . Moreover , each operator can be split in many different ways . For example , a matrix multiplication can be split over the rows or columns of the left ( input ) matrix or right ( weight ) matrix , or a combination of these . Some splits will require inserting an all-reduction operator to recover the result , and some will require an all-gather . Some will divide weights across devices , which helps with memory usage . The optimal split in terms of latency depends on the matrix dimensions and also downstream operators ; a matrix multiplication split over some dimension would accept input tensors split over the same dimension , but a dropout operator requires the entire input tensor . 2.3 Data parallelism . Data parallelism is the de-facto way of parallel training today . With data parallelism , every worker has a copy of the full model . Inputs are sharded , and each worker computes weight gradients independently . Periodically , these weight gradients are aggregated using all-to-all communication or a parameter server . Vanilla data parallelism assumes that model parameters fit on a single device . For large models , data parallelism can be used on model shards ; gradients need to only be exchanged among devices responsible for the relevant model shard . Data parallelism can be used with both pipeline ( as shown in PipeDream [ 17 ] ) and tensor model parallelism ( as shown in Megatron-LM [ 22 ] ) . 2.4 Memory-saving optimizations . When training large models , it is common to use memory-saving optimizations that trade off a reduction in throughput for a lower memory footprint . One example is activation recomputation , where the intermediate activations produced during the forward pass are not stored . Instead , the forward pass is performed a second time , just before the backward pass , from a stashed input activation , which is much smaller than the full set of intermediate activations . This drastically reduces the footprint of pipelined training , at the cost of executing the forward pass twice . 2.5 Related work on algorithms for partitioning and placement of DNN graphs . A number of systems have looked at the problem of partitioning DNN models across devices . However , each of these systems considers a subset of the parallelism dimensions considered in this paper . FlexFlow [ 13 ] considers many dimensions ( data , tensor , and inter-layer model parallelism ) , but not pipelining ; moreover , it uses a heuristic MCMC-based approach that requires domain knowledge to be encoded into the “ parallel plan ” generator . PipeDream [ 17 ] and Tarnawski et al . [ 23 ] use dynamic programming and consider a combination of pipeline model parallelism and data parallelism ; however , they support different model topology classes . PipeDream supports only linear ( path ) graphs while [ 23 ] supports general DAGs . While PipeDream is oblivious to memory usage , its enhancement , PipeDream-2BW [ 18 ] , targets large models that do not necessarily fit on a single accelerator . Exploiting the repetitive structure of some of these large models , such as transformerbased language models , PipeDream-2BW ’ s planner only considers configurations where every stage in the pipeline is replicated an equal number of times ( equi-replicated stages ) , thus simplifying the search for parallelization configurations . PipeDream-2BW ’ s planner also explicitly considers per-device memory capacity constraints and includes activation recomputation in its configuration search space . None of these three approaches consider tensor parallelism . Learning-based approaches ( Mirhoseini et al . [ 16 , 15 ] , Spotlight [ 8 ] , Placeto [ 1 ] , GDP [ 26 ] , REGAL [ 21 ] ) treat the objective ( latency or throughput ) as a black box to be optimized using RL ; runtimes are measured in an online fashion on the target system , which is very accurate but computationally expensive . Moreover , some of them optimize for latency or peak memory usage rather than throughput , and they do not explicitly consider splitting individual operators or pipelining . All existing algorithms , except PipeDream-2BW , also do not consider memory-saving optimizations , that trade off throughput for reduction in memory footprint , in combination with different modes of parallelism ; activation recomputation is one such important optimization . These optimizations are commonly needed for training models that exceed the memory capacities of modern accelerators .

Summary: This paper proposes a novel algorithm to find a pipeline parallel configuration where each pipeline stage also includes data and tensor model parallelism. The method assumes that the input computational graph consists of multiple “layers” that form a DAG and each layer has a  set of predefined tensor model parallelism configurations. The algorithm is a two-stage dynamic programming algorithm and is mainly evaluated in a simulated environment.