Two Dimensional Parallelism Using Distributed Tensors at PyTorch Conference 2022

Key Takeaways

hi my name is Juan Chao Liang and I'm from the pi torch distributed team at meta today we're going to print down our recent work about two-dimensional parallelism using distributed tensor we will briefly go through the background and motivation then introduce the distributed tensor concept and the tensor parallelism that builds on top of distributed tensor then we will go through the 2D parallelism that we build for the vision Transformer model let's briefly talk about the background and motivation currently the research Community trained large-scale models mainly in three ways data parallel within layer model parallel and pipeline parallel in pytorch for data parallel we are offering distributed data parallel and fully shorted data parallel to three large-scale models for pipeline parallel we are offering PP or pipeline parallel solution which is currently in the Prototype stage so what should be a general purpose within layer model parallel look like our answer to this is distributed tensor and the tensor parallelism that's built on top of it so what is distributed tensor distributed tensor also known as detensor is a common abstraction we proposed to build the bridge between different parallelism strategies it offers fundamental tensor level Primitives to describe data distribution and computation across holes details are consists of two important Concepts the first concept is device match an abstraction that describes the topology of devices within a cluster for example we could construct the U.S mesh that describes Two Hosts each with two gpus the second concept is the details replacement types there are three placement types chart where we Shard on Twitter Dimension across devices replicate where we replicate tensor across devices partial where we partition values across devices with these two concepts we can distribute the tensor in different ways for example we can chart a tensor with each piece on separate devices we can replicate a tensor across all devices or we can even shorter tensor on set of devices then replicate this shortly tensor on another side of devices even the detensor concepts now let's look at the details API data sir is fundamentally a torture tensor subclass and they really look like a normal tensor except that it's a distributed across devices with the device mesh we can create a detensor with your white starting and column wise charting and we can convert from our two local torture tensor on each rank we can read chart or redistribute the detailsor from one type of sharding to another we also offer a high level API to Shard and Scatter an existing tensor from one rank to all other ranks with the power of detailser we can build tensor parallelism apis on top of it to express model parallelism next I will hand it off to my teammate jinji to talk about tensor parallelism and 2D parallel hi I'm Jin jio Wang and I will talk about tensor parallelism and its user API first of all let's see how tensor parallelism works there are three parallelism Styles we support column wise row wise and pairwise parallelisms the first two are pretty intuitive as we just partition one dimension of a given tensor the third one pairwise parallelism is to combine the first two and Achieve calculation efficiency let's take a look at the diagram below as you can see here after multiplying with the First Column wise partition linear we don't need to do any Collective and leave the activation partitioned and this partition Matrix can then do another multiplication with a row wise partition linear directly to get the final result we just need to do one or reduce here when it comes to API users don't need to change their module or code they just need to call this high level API parallelize module with parallelized style specified per module to get the model partitioned we will handle all the hassles during model training and inference next I want to talk about 2D parallelism where we combine data parallel and tensor parallel together the way it works is that we first partition each model parameter across all gpus on each host and then the fully sharded data parallel which is also called fsdp will restart the local tensor of tensor parallelism this way we can get the benefits of both tensor parallel and data parallel the user API is also pretty simple that user just needs to create a 2d device mesh before calling the parallelized module API and user can then wrap the model with fsdp last but not least I want to talk about What preliminary result we have got so far from the experiment on the vit model we see 2.6 times faster in training a 60 billion vit model and enable the training of a 120 billion vit model on 64 gpus here we can do some backup envelope calculations if we simplify each layer of Transformer encoder as one input tensor and four linears we can calculate the peak memory ratio for 2D parallelism and fsdp only scenarios the result shows that when the batch size is 1024 fsdp is using 1.3 times more Peak memory than 2D parallelism and upon measurement this ratio is 1.1 meanwhile if we reduce the batch size to 512 this ratio increases to 2.4 the result shows that when the batch size is 1024 fsdp is using 1.3 times more Peak memory than 2D parallelism and the Palm measurement this ratio is 1.1 meanwhile if we reduce the batch size to 512 the ratio increases to 2.4 this calculation shows that when the number of gpus is limited and if we have to use small batch size for really large model 2D parallelism can indeed have some advantages