Bay.Area.AI: Tensor and 2D Parallelism -- Junjie Wang, Xilun Wu,Iris Zhang

Recording: Bay.Area.AI: Tensor and 2D Parallelism -- Junjie Wang, Xilun Wu,Iris Zhang

uh engineer from pyto distributed today uh I Iris and then shilin we're going to talk about uh tensor parallel 2D parallel and also our update on the Pythor scalability efforts that we've been working on during the past uh 6 to 9 months yeah so uh first of all uh when we look at the like diagram from the like uh how how the model size uh scales it kind of go up craz uh so for example gpd2 only have like 1.5 billion parameters but gpt3 have uh 175 billion parameter right and also the data size also increase uh hugely for example for llama V uh llama V2 we're using two trillion tokens to train the model right and also what's more is that we are using way more gpus now to train a llm like 2K 4K even 16k you know anything above that so that brings some challenges and some composable requirement to us uh so for the um I guess to for for large language models originally we have this kind of data parallel training right like DDP so this really data parallel so that's simple you copy the model multiple like replicas and then you send different data to each gpus but this won't work when the model becomes so big that the model cannot fit in the memory of one single uh gpus then we have fsdp that we sh the way across uh all gpus or even across host and then during the forward and backward we call all together we get all the weight and then we do the forward backward uh but the thing is that when it comes to super large scale the communication across host becomes like a bottleneck for fsp so the trick for fsdp is that although it'll call all together but we kind of hide the communication with like basically overlap it with computation so that the all GA cost is is is is like a free lunch right but but with with the like super large scale uh the all GA becomes super super long and then it's like completely comound then there's nothing you can hide it so then um when we then have like hybrid shter data parallel and also um we we have like tensor parallel sequence parallel and then pipeline parallel so for tensor parallel and sequence parallel instead of um I mean we still sh the weight but instead of call all gather the model parameters we now do some like all gather or all reduce for the activations and for pipeline parallel this Al this will like reduce the Comm communication cost a lot because there's only like Point too communication and of course uh like uh Megatron is having this 3D parallel already uh so I'm going to also talk about our like native solution in the next slide but this come comes with like a uh some some challenge for us is that all these uh Solutions like pyto native Solutions are kind of like built individually and then we when we build first start building we don't think of like composability so that's why uh we we see some compos composable issues between each of these training paradigms and also uh when it comes to like like I said 2D 3D parallelism the statei uh and then the process group management becomes like a very complicated I think iris is going to talk about how uh how we kind of handle this static and then checkpointing in a more uniform way uh last but not least uh a lot of those kind of communication and computation optimization and overlapping are right now are kind of hand tuned for example fsdp has its prefetching logic uh so DDP has it bucketization logic and then tensor parallel has its own um like a a trick there but like when we combine them together there are way more uh optimization space for us so we want to leverage uh like new technique for example um p uh P toor to torch compile to um to further optimize the like in in the case of 2D and 3D parallelism okay all right right so during the past 6 to 9 months uh we we I mean we we implement the tensor parallel based on based on top of distribute tensor last year and then when we had a prototyping for fsdp plus TP so that's why that's why we call it 2D parallel and then we we now had the DDP working with TP as well um I mean for forward and backward and also we also had a prototyping for pipeline parm working with tensor parallel together uh yeah and then we're right now working on 3D parallel uh also uh another thing is that when we apply tensor pel on on Transformer model we see some uh ux pinpoint some like API uh are are not very good for user to use so we we made some changes to the uh tensor parallel apis so now they're all available in py toch nightly build so uh I mean feel free to try it out and and let us know how that goes uh and the third thing is that we have this activation checkpoint wrapper for tensor parallel so activation check pointing is that uh uh because activation size is so big so uh we just throw them away in the forward and then in the backward we just recompute the forward again and then do the do the backward that's activation checkpointing and then for tensor parallel because the input are the same across all the tensor parallel ranks then we we somehow just uh Shard the input as well and then all called all together in the four uh in in the back word again yeah so last but not least we're now uh working on enabling the this re inference for lar language model using native uh py to tens tensor parm another thing we are having that available in the nightly build is sequence parallel like I mentioned in tensor parallel uh all the input are the same across all TP ranks but for this one instead of sending same exact input across all TP ranks we Shard across sequence Dimension so that um we so that we first need to call one all together across like to to together all input together then it's like a very classical uh Megatron style uh uh two linear layer and then in the end we call one reduce scatter the tensor becomes um to becomes back to its original shape uh and then the the way it works is that uh when you first have a columnwise sharing and then you have a rowwise sharing uh in the end um in in the middle there's no no need to call any like Collective and then you just need to either in the in the end in the in the scenario of tensor parallel case you just call one Oru or in the scenario of sequence parallel we call a reduce scatter and then after reduce scatter uh you can call layer nor or soft Max uh which is in the last Dimension they they're all fing and then uh when it comes to like a next module then we'll repeat this uh uh this all gather reduce scatter pair thing again so that that's we already have enable it uh for sequence parallel in nightly build okay another thing uh I want to mention here is that we have noticed some like performance Gap in distribut tensor and also have some new features for distribute tensor uh so per for performance uh because distribut tensor is using tensor sub class so it comes with some like CPU overhead so we try to reduce the CPU overhead as much as possible uh this is going to benefit both small models and an inference because for large models uh it's CPU overhead can be well hidden by like fscp and also uh I mean because it's mostly CPU bound right so then all these CP cost can be well uh hidden but for small models and inference um we need to reduce the CPU overhead uh as much as possible and also like I mentioned we we are integrating tensor parallel with torch compile that torch compile can help us further do some reordering and then scheduling to better overlap the computation and communication last but not least uh we we enable random generators up like for example for drop out on a replicate tensor across like a different ranks we want them to get have exact result across all ranks so just give you an example all this kind of like random or some distribution this op will have correct behaviors on on both sharded and replicate detor all right so next I'm going to hand over to Iris to talk about uh distribute tensor and then all the static and the checkpointing stuff yeah so uh thanks J for going over like all the challenge for comp possible distri training and sharing updates to our features now I want to go uh into our ux for trainer initialization as well as like a backend powering our tensor parallel apis so to better support like composable distribut training initialization we're introducing device mesh so device mesh is a higher level abstraction that manage process group it allows the user to easily um create internal and internal process group without worrying about too much about how to set up the ranks correctly and how to manage um uh subprocess group as well so we can see from the example from the bottom left if you want to set up like two host with two device each host we can simply create a 2d mesh that uh connects devices within within each host and connect and it also would connect connect each device with its counterpart inside the other host in a homogeneous setup so um why is device mesh useful so because for multi-dimensional parallelism right in 2D 3D settings we need to set up the nio communicator to correctly before applying any sort of pism for example um here is an example of like how distributed developer could be manually creating process good for 2D hybrid sharting uh as you can see is fairly complicated and very aopr and with Device Mash things can become so much easier and you can see that we can see that user just need to tell any device Mash uh on the right hand side the number of Cuda device per Hol and the device Mash will automatic set up like Nick commuter cators uh within each Hol and across uh across holes so it it really like it becomes like into two lines uh it's also important to not note that like after creation of device Mash user uh can still get access to the underlying process group if needed um this significantly simplify the setup of like a multi-dimensional parallelism and next I want to uh go get into like DET tensor for a little bit more um so which DET tensor is like the backhand powering our pytor native tensor P apis so D tensor is our new like f FAL abstraction that perform tensor level sharding computations so it brings a lot of like a compossibility Advantage first it provides a uniform tensor sharing layout to represent different pism and their static right and secondly it also allows user to perform sharded computation easily in an spmd style um so without needing to worry about to actually needing to write back for for formula for colc collect Ives and potentially hitting uh the risk of that lock with some manual collectives um so deter allow so here you can see like DET tensor allows sharting on a tensor Dimension over a set of devices and replication over a set of device it automatically petition the data to sh it or replicate it or even more complicating sharting strategy um according to the device mash and placement user def fine so for example you can you can create a shed deter that only shed on a set of devices or and replicate on the rest of devices by just spec specifying a 2d device mesh and the corresponding placement so um in the uh third diagram so you will specify shart and replicate on um different uh mesh Dimension um so with Device mesh and deter as like the fundamental building block we can make large model training much easier so uh Jin went uh went through earlier um that we have our 2D um uh 2D Solutions and um a couple pain points that we have seen in the past for tensor parallel is usually tensor parallel is like intrusive to the modal and so when like the distribute C is intertwined with training C it really increases your maintenance B burden right and another paino is that compos composition makes checkpointing uh hard as well so for our native to the uh apis other than having like the competitive M2M uh performance we're also looking to solve the pain points that I just mentioned so we can see from this cool example that device management can be easily achieved by creating a 2d device mesh with the uh zero dimension of device mesh performing data parallel and the First Dimension performing tensor parallel and as you can see uh we you can easily slice this 2D device mesh uh into 1D device mesh with annotation telling which dimension of device mesh uh the Palm should be using and to apply tensor parallel to your model you can specify paralyzed plan for the Transformer block by composing columnwise parallel rowwise par style together and apply them each layer of the Transformer and adjusting the the number of your head um the head so uh by applying pism outside of the model code user can really minimize the code change um to uh allow efficient tensor parallel uh Plus data parallel solution um before going into our checkpointing solution I want to quickly talk about what makes checkpointing hard in the distributed settings right checkpointing would not be hard if you always um load with the same parm solution and Al and same W size you can simply set up the same number of trainer every time and you directly use torch. Save and torch. L without worrying too much about anything but however in reality that's not the case as in distributed there are different scenario where you need to um do data shuffling uh we call loading the safe tensor uh into like a different pism or different W siiz rearing so for example rearing is needed when you um say you're moving from training to inference say if you want to do 2D training um uh at in a large wall size and later you might only want to do inference with tensor parallel only in a smaller W size you would need rearing and also from offline training to online training uh which requires both recharting the model and optimizes St is and also you can see from pre-training and fine-tuning you need to um uh rear like the model as well so for that we have a pie torch native checkpoint solution distributed checkpoint uh this is beta ready in py 2.1 and together with um detener we're aiming to help user to solve the pain points for checkpointing under distributed settings so um distributed uh checkpoint supports efficient shed static safe and low for fully sharded data parallel and tensor parallel it avoids writing redundant copy to this and allow parallel to read and right also supports checkpoint low time recharting to like different wall size when the cluster environment changes and furthermore like with uh this uniforms they dig that contains detener um distribut checkpoint allow user to save in one type of pism and load into another type of pism very similarly so you can checkpoint in a 2d parallel setting for pre-training and later easily load it back uh using 1D pism for fine tuning or inference um yeah and next I'll pass it to sheilan to talk about control playe thanks Aris so uh my name is shulan I'm also from P page distributed and um the my focus is on the um on the Comm communication backhand of py of pyto so the communication backend has of P has proven to be stable and reliable in the modern distributed training jobs so by modern Distributing distributed training jobs I mean like 60 uh 64 gpus DDP fsdp but recently we have the example of llama which takes 2K gpus to train right which is crazy so we have seen this trend we are using more and more gpus and um this put the uh put the pressure on the or or question mark on the uh on the communication backand is it okay or is it still reliable for that number so this initiates the our effort to to Benchmark the Comm communication backand at this super large scale by super large scale I mean like 4K 8K even 16 16k 24K so this numbers are might be crazy right now but I I think it will not be that unreal that rare in the following like 5 years so for the commun communication backand there are two critical uh components the first is called TCP store and the second is called um so the second is called process group and a glue is one of the backand of process group you can Implement your your own backand for your custom device Hardware so process group is a uh abstra abstraction layer above the actual communication you we use the we use PG to to to represent collectives performed on like a certain number of uh gpus or ranks so the so the the glue is the library uh for uh for PG on CPU so it can be used for like distributed checkpointing it can be used for like um metadata you know uh communication so it's mainly for the data on on the host and we also have Nicko as the backhand for GPU which is responsible for the data on the device and a TCP store is a key value store it is centralized usually hosted on the on the first rank rank zero and it is responsible for uh for exchanging data uh which don't fall into the categories I mentioned above so so this is the roles of them in the P distributed back end so fir first of all we tested TCB store and I will show the result in the in the next slides so basically um we we add it we we added a non-blocking implementation for TCP for TCP store called libuv and uh it in the test we see it it scales to up to 32k ranks which is a super large number so that means temporary we are safe on the on this thing like until until 32k you okay you'll be okay and uh yeah let's see that first so here is the Benchmark so we see we see some lines so this vertical lines means the the the uh the old TCP store uh is likely to fail at 8K so this is the number 800 oh sorry 8,000 uh gpus or ranks and U with the LI UV backand we make it uh scale to 32k which is the right the right end and um we see that the start out time is is the same before like uh uh before like 4K but after uh between 4K and 8K the lip UV takes more time to uh to start up so this is because we add some uh some random randomized uh slowness to to each client so that they will not rush to the server at the same time which which which is like a DS attack right so at AK the old Li uh TCP store fail and uh the lip UV backand still still works and the time is linear linearly growing which means it is still scalable and another thing we test is the barrier so we so some in some cases we still use the TCP store to perform barrier to synchronize among the ranks so this happens where the process group has not been initialized and we see that uh this line is the uh execution time of uh barrier it is even at 32k one barrier just takes less than 1 second around like 700 milliseconds so that time is acceptable in most cases so we can say that the the TCP store uh the LI UV version is okay for 32k you don't need to worry about it there is there are no there are no issues yep and the next thing is the glue so what we do is we test uh this process group uh process group using glue backend on initialization and uh we see that it scales to akk but not Beyond yep I asked why it was limited to 32k you were um sorry what why was limited to 32k um TCB store or glue glue is limited at 8K yeah so um TCP story is we don't have more machines to perform Benchmark I don't I don't think in in near future you will have the use case of more than 32k so we didn't yeah we just don't test but but based on the on the trend it should still scale okay so then for for glue we see that the curve doesn't look good right it's it it's not linearly um it's not linearly uh growing so we can see the number here for AK ranks the the total initialization takes half uh takes like um like uh like 20 20 uh more than 20 minutes like 25 minutes so it's so it's like uh 1,500 seconds so this number is this number is is is okay for for a job startup because usually the the job fails around uh fails uh fails and it needs restart in like in a day or in like 20 hours so like 20 minutes is it start up is okay for every like 20 20 hours but with what will be the case Beyond so so um we have tested on like 12K 12K ranks the number would be 1 hour for initialization that that is a long time already especially with more ranks in involved the time to fail for for the job will be will be shorter because you have more gpus you have a uh you have a higher chance to to to fail on one of the device so it the job cannot last for like 20 hours it will be even shorter and uh you will have one hour to restart the job that is unacceptable so in this sense we think the current glue implementation does not scale Beyond 8K but at at 8K it is okay and uh we have we we also have the breakdown we see that the use of TCP store takes the the orange part and the glue full mesh connect takes the yellow part and they are they're mostly the same so both are the bottom neck so we we plan to optimize optimize the and we have some some ideas on how to how to do that so hopefully we will get to be able to scale Beyond 8K yeah so this should this should be able to happen in the in the next three months or six months so the last thing is about observability so this is this is a more urgent uh topic for us because because um this is about debugging like um now we have seen this issue in many like in many real jobs it is super hard for user to debug in like a 2K 2K GPU job because you have so many logs just like uh Omar has mentioned right the super large the the large logs you may even need like LM technique to help P that so the thing is the current debugging tool in py in py dist distribut is not suitable for this case so currently um we we have we have two things to help user to debug one is called um nickel Watch Dog which will get the thread time out if the collective is hand is hand there it doesn't if it doesn't proceed and this will kill kill the uh main threat and another thing another thing is called um uh nickco desync debug so this happens so this this feature will be used whenever a timeout happens it will send a report to the um to the TCB store which which is hosted on the first rank so that you can you can find a report on the on that host and uh that report should give you information about which rank joins the last Collective which one doesn't so this is very useful because if you see one rank doesn't doesn't uh join the collective we know that rank has some issue you just go to that Rank and see the log and figure out what happened so this is very useful in our um based on our past past experience however this one doesn't uh work well with like 2K gpus because um because um because of yeah SE several limitations so the first thing is when The Watch Dog kill the threat right suppose uh assume that it happens on rank zero which where the TCP store is hosted and the thread is killed and other thread will other ranks will try to send the report to that that rank however the store is gone so in that sense the TCP store or the DC report will not help because the store is no longer there you cannot see the report so we want so we add a um PG hook to handle to handle that thing so basically we separate um that you know um Watch Dog and the DC report uh uh from one uh we used to have them on the on one threat we separate them to two threats so that the time out will not kill the TCP store thread so in this case we will have the report even if the main thread was killed and also this um this hook should be uh customizable so that our custom customers can register their their own exception Handler to PG so whenever uh so um whenever a collect active failure happens the custom you know Handler will do the do do the job to help user to to debug so this is a very useful tool we think so we get feedback from um from any customers they all want this feature yeah so I think we believe that with those techniques or features we should be able to uh have U to make the error detection easier so so that we can find figure out stragglers like for example when rank join the collective very late and eventually it causes uh uh Collective Collective timeout or we can figure out like hands or crashes on certain ranks so this is this Clos closes the talk any questions is that is that going to continue to expand yeah so I Ste one of your mics with yeah oh okay it's right here oh oh yeah yeah so so um yeah we're we're working on like this device Mash support so naturally if you're seeing like the device Mash initialization it can uh it can actually support to ND right so because you you would have a 3D you can still slice like uh your your ond device mesh out of it for 3D composition as well so device and we're we're uh we're going to be moving device mmh to a standalone uh package and for distributor so it will be like this um genetic a generic package to for your to for you to manage your device layout yeah yeah so what are some of the differences with a defens ux compared to what was done before I know there are other shed cancer implementations in pytorch so what was the reason behind the switch and what you guys sort of think yeah I I can I can talk about that so when I first joined like a this through a team we were still working on sharded tenser but later on we found that sharded there's some limitations around sharded tenser for example Shard T only can represent sharding right however in the distributed World sometimes you want to Su for example for DDP DDP essentially just replicate right so we want the so dor itself has like a multiple uh scenario like partial replicate and shed tensor and then uh originally shed tensor is even not a subass of tensor so that actually took us a lot of effort to eventually convert to a tensor subass and then we found that oh okay like all these kind of things doesn't like we don't want to go down the you know wrong path for a very long and then we decide to redesign and then we recreate D tenser and then we deprecate a sh tensor already we we send the warning message there so I guess the recommendation here we have for our customer and for Community is that just don't use sh tensor and then use D tensor if you have any feature requirement or if you see any feature Gap just come to us say hey okay here's your requirement yeah and and to add on that for a little bit so like naturally like Shard tensor as sh say is not aware of any replication and it doesn't really support compos possible so it it it's not aware of your Global layout but with dor you can do 1D 2D and 3D which help naturally help checkpointing to be much easier we actually have a prototype uh for 2D uh uh 2D checkpointing with SH tensor is really like a pain as we learn a lot from like the uh the we don't want to get into that ux Health um so we just want to switch to detor which is like a better support for this yeah and also deter now has you know work with PCH compile already right so then you get the to compile wait so yes S one more quick so what are you guys thinking about in terms of torch compile and detor are you just going to like if I give you like a torch module will you just sort of figure out the best way to to parallelize that or is that is that sort of the vision uh I I think uh the other way around like you still like call like uh after you call TP API and then r with fstp and then you you call torch compile yeah so that's the current uh thought around that yeah um you mentioned somewhere on the slide that there's some by netive parallelism work going on for inference can you elaborate uh we are still working on that so currently I mean um people have applied D tensor for tensor parallel in the inference but I think we're still working on that so yeah and I think there's a separate talk yesterday talking about this and you can go back to the py forch conference schedule yesterday I think there will will be releasing the uh YouTube video soon right right right Horrors and then hame they will talk all yeah thank you