Mesh-TensorFlow: Model Parallelism for Supercomputers (TF Dev Summit ‘19)

Key Takeaways

[Music] uh hi my name is no shazir um I'm going to tell you about mesh tensor flow which is a system we built uh uh for training giant models on TPU pods so so we're going to talk about uh data parallelism and model parallelism and uh that also about why I want to train giant models and why we need model parallelism to do it then then I'll tell you about the mesh tensor flow system so first uh data parallelism this is how roughly everybody trains neural networks if you have uh distributed Hardware this is what we use on TPU Pods at Google um and generally you put your entire model on every device split up a training batch into a lot of little trunks one on each device run it then add all the gradients on the parameters across all the devices and do your update um this this works really well the communication is fast on on uh on our kinds of networks on lots of different types of networks in fact um so this is roughly what everyone's doing uh the great things about it are it's Universal you can use any model architecture um fast to compile because you're writing spmd code you're writing code for what one device is doing and then uh you can compile it and send it to every device uh which is using the same code and you get roughly full utilization because every device is doing roughly the same thing so if uh they're all similar then uh nobody's waiting for anybody else and uh the communication is f the the only problem is that you cannot train giant models because your entire model has to fit on every device so why do I want to train giant models well I like uh working on language modeling there are lots of important applications machine translation question answering uh dialogue sentimental analysis lots of interesting things to do with language and we find that quality improves with model size so a bigger model tends to know more about the world understand things better and give you overall better results there's plenty of data out there to train a giant model just download the text to the web common crawl whatever and uh and you've got uh you've got billions to trillions of words of training data and in fact you can train one big model and then F tuna to do lots of different things uh there's been a lot of research open Ai and bird at Google on um on transfer uh learning in language so it's a great candidate for building giant models now uh as an example I uh trained a Transformer language model with the roughly a 100 million parameters on the text of Wikipedia the Abraham Lincoln article was in the dev set in the heldout set and I told it to generate a random Abraham Lincoln article and it looks roughly grammatical uh remembers that somehow that he's a politician American politician there's plenty it doesn't know about the world like who Abraham Lincoln was or that America doesn't have a prime minister and lots of other stuff um but if you make the similar model just bigger here's with five billion parameters instead of 100 million and now it seems to have picked up a lot more about Abraham Lincoln roughly half of that stuff is correct but you know mostly it's fake news but there are more important uh applications out there than than generating fake news um but but this is just a nice demonstration that that model size is important um what would a model look like with a trillion parameters uh we have not done that yet but we hope to do to do that soon um okay so if all the parameters will not fit on One Core we need to do something called Model parallelism which means that we're splitting the model itself between different devices and that should let us train really large models and it should also be very good for inference latency because now the computation for one example can be split across multiple devices the problem is it's very tricky to design these kinds of algorithms um how do people tend to do it now well you use device placement you say this operation's going on this device this operation's going on that device tensor flow makes it easy to do that still It's tricky to design an efficient algorithm and you end up with a giant graph if you're generating enough operations to go on 2,000 different cores um here's an example of some uh model parallelism by device placement from uh Google neural machine trans ation they had eight lstms which they distributed across eight different gpus the softmax layer they put somewhere else and you need some kind of uh interesting pipelining to keep all the gpus busy it works but uh you know a lot of work to to get this thing to work right now we're going to take a totally different approach which is going to be inspired by what works well about synchronous data parallelism so we will have every processor involved in every operation uh we're going to use spmd style programming where you'll have the same program on every device and it's going to use Collective communication like all reduce um just like uh data parallelism um and the way and our library for doing this is called mesh tensorflow we should be able to implement data parallelism model parallelism we should be able to split in different dimensions like uh split a an image or video spatially or any sorts of combinations of these things and we're uh targeting Hardware where you have a homogeneous set of similar processors ideally um well connected like a TPU pod we've got these two-dimensional supercomputers at Google that we that we've been using uh and you're going to view your set of processors as an n-dimensional mesh it doesn't have to correspond to a physical n-dimensional mesh you could view a two dimensional mesh of processors is a one-dimensional mesh um but of course performance will depend on uh on those considerations so how does this all work well in data parallelism you can view it as splitting the batch dimension of the computation across all your processors so any tensor that has a batch Dimension is going to be split across all the processors and any tensor that does not have a batch Dimension meaning the uh meaning the parameters gets fully replicated now we're going to do the same thing but for model parallelism we will choose different dimensions to split so maybe Dimensions representing the sizes of hidden layers and we will decide to split those Dimensions across the uh set of processors so and the communication will happen us usually an operation will not involve communication but some operations will involve Collective commun communication like all reduce particularly when you're reducing outs split dimensions and this is going to be somewhat similar to how uh how things work in synchronous data parallelism okay so let's do an example a simple three layer neural network input layer x hidden layer H output layer Y and we have two weight matrices W and V the data parallel way to do this is that we're going to split anything with a batch Dimension meaning the activations x h and y so all of those tensors are split evenly across processors and each tensor that does not have a batch Dimension W and V will be replicated across every processor so here it's showing what processor zero is doing what processor one is doing they're both doing something roughly similar except they have different halves of the activations you don't see any communication in the forwards paths but if you were to see the backwards pass where you're Computing the parameter gradients you would see some mat moles where the split Dimension b gets uh gets reduced out and there would be some all reduces in there so we're going to now instead of splitting B let's split the size of the Hidden layer Dimension H so we do that and now X and Y the input and output are fully replicated because they do not have an H Dimension but the hidden layer H is split because it does have an H Dimension and the parameter matrices W and V also have an H Dimension so they're split so again you have a parallel computation on on the two processors and you see an all reduce communication when you're Computing why because we're reducing out the split Dimension H we didn't have to split H instead we could have split D which is the uh dimension of of X and Y so uh in that case you would have a different pattern of which tensors are split and which ones are replicated and you'd have communication in different places and if you want to get really fancy let's split let's let's do data parallelism and model parallelism at once we're going to split Dimension B the batch across one axis of our two-dimensional supercomputer and we're going to split the hidden layer H across the other axis of our mesh of processor so now we have different tensors being either split in one dimension and replicated in the other or the or since tensor H has both of those Dimensions it ends up tiled among all the processors and there are going to be all reduced Communications in there um but not across all the processors they'll be partitioned all reduces just across rows or just across columns so the general case is give all of the tensor Dimensions names and the we Define the layout of the communication as a map from tensor Dimensions to mesh Dimensions saying which tensor Dimensions get split across which mesh dimensions for example in the previous slide we had the batch tensor Dimension split across processor rows and the hidden size dimension split across processor columns we did this to our Transformer machine translation language mod model and here are the layouts we use for data parallelism model parallelism and combined parallelism for that uh for the model parallelism it works to split the size of the vocabulary the size of the feedforward Hidden layer and the number of attention heads and if you do that you end up splitting up all of your communication very nicely and and get a nice model parallel algorithm and that can also be combined with data parallelism by splitting the batch across the other dimension of the supercomputer now picking a good layout is for now something that you need a well-trained human to do uh you need to make sure that all of your expensive operations are split you're not allowed to split the same two dimensions of the same tensor across different dimensions of the mesh uh and depending on what dimensions you chop up it may result in more or less communication so example model data parallelism you'd like the batch to be big otherwise you're going to get a lot of communication similarly if you're splitting up a hidden layer you might want that layer to be really big um so how do you use mesh tensor flow well download our open- source repository you build a graph in Python much like a regular tensorflow graph uh except that you're using named Dimensions um you define what your mesh is and how it maps to your physical processors uh you define your layout of what gets split across what and then mesh tensorflow turns your mesh tensor flow graph into part of a tensor flow graph and you still use tensorflow for anything else you want to use it for like the data pipelines and everything else um so so far we've trained Transformer models on with up to five billion parameters on entire TPU pods getting uh good performance out of the thing um and uh these giant models give state-of-the-art quality on on uh on some Benchmark tasks like in language modeling and machine translation not surprisingly uh bigger models are better lots of people are finding that out um and in the future we would like to try even bigger models uh I think with some uh well-placed uh sparsity we we would have the computation to train models with a trillion parameters we've tried up to a couple hundred billion for now and it runs um so next thing is uh see if we can get a trillion parameter model to run and give us great quality and um this should be useful for other things like low latency inference and situations where you have giant um uh inputs that you want to process and for now what works well we have uh we're emitting spmd code for uh for TPU including Cloud TPU so this runs nicely on TPU pods and um for CPU and GPU it's still emitting the oldfashioned device placement code so it runs but not as scalable um everything is out there on GitHub and uh runs with tensor flow one not yet with tensorflow two uh we're uh and then uh in the future we want to uh use this for different types of models it would be great to automate the process of choosing a distributed layout because then you wouldn't need to know about mesh tensor flow and it would just figure out how to distribute your computation for you uh and we welcome uh contributions to the open source code uh or um just contact us and um I'd like to thank uh all of my collaborators uh the the authors of uh of our uh paper um also like to thank the tensorflow teams and xlaa team for uh for a lot of technical support and help with all of this implementing what we needed to be implemented and uh everything's out there in our our open source repository uh thank you [Music]