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: Parallel/distributed training of machine learning models is a big and complex task. This paper proposed a dynamic programming algorithm in searching optimal parallelization configurations for machine learning models, considering tensor parallelization, data parallelization, and pipeline parallelization. Designing such algorithm is important and interesting, and potentially relieve machine learning app writers from manually designing parallelization configurations. To provided an optimal (or close to optimal) parallelization configuration without breaching memory limitation is already good contributions in real-world scenarios.