Torchtitan: Large-Scale LLM Training Using Native PyTorch 3D Parallel... Wanchao Liang & Linsong Chu

Key Takeaways

uh my name is W Chan and I'm an engineer from the P torch cor team working on distributed training today we're going to talk about torch Titan a large scale large language model Library using the pytorch native 3D parm uh let's first discuss uh large language model pre-training challenges the industry invested heavily into pre-training what is special about it and why it requires foundational Investments the scaling law drivve the larg langage model inov with great model quality but it also impos lots of technical challenges first the model size grows exponentially over the last few years today many large langage models ranging from 7 billion parameter all the way to trillions of parameters second huge amount of data needed data size also increased crazily where trillions of tokens are needed given the large model size and huge amount of data the training would take millions of GPU hours to finish we need a robust large scale training system that could scale the model training to tens of thousands of gpus staying training efficient and keep converging as if we're training in a single device this requires many system level Innovations to implement a training system that could compose uh different training techniques together including but lot not limited to uh UND dimensional paradism Fusion compiler activation checkpoint flow Aid training fast checkpoint restart Etc however composing these techniques together itself is very challenging let's only look at the UND dimensional paradism part to train the large language model at scale we need to compose different type of paradism together starting from data parallel here we primarily focus on fly shed data parallel or fsdp fsdp scales pretty well for medium siiz models and could scale up to maybe like 12 52 gpus but when training at large scale it needs to compose with model parallelism together and address several composability problems including how do you represent a uh shed uh undimensioned state for for training saving and loading the checkpoint how to interact with tensor paradism and pipeline paradism how to perform things like sharded grading Norm clipping in a numerical correct way and combining with tensor or sequence parallel we're able to train larger models and improve the scalability to some somewhere between like 2K to 4K gpus but tensor paradism usually have drastic model changes and blocking Communications that will dra draw down the efficiency Beyond 2,000 gpus we often incorporate P plan paradism to further improve the scalability but pipeline paradism also introduced drastic changes to the model and training setup along with some of the complicated interaction with data paradism all of these composability challenges are overhead to researchers and Developers where many researchers ended up just using fsdp for majority of the use case due to its Simplicity to demonstrate how we as a PCH uh address the compossibility challenges with our undimensioned NN training with those techniques we're introducing torch Titan torch Titan is the open- source proof of concept for large scale training with Native pytorch it is buil with the pytorch design principles focusing on complexity and ease of use it implemented a composable and dimensional parallelism training for the Transformer models including fly sharded data parallel tensor or sequence parallel and pipeline parallel it showcased how to compose many efficient training techniques together in a simple Manner and Achieve great performance next let's dive deeper into torch data and look at how it enables 3D parallelism including a fully shed data parallel with two tensor or sequence parallel and pipeline parallel from pytorch all the relevant code for each paradism can be found by scanning the QR code here now fsdp we have redesigned fsdp to a new version called fsdp 2 to simplify the algorithm implementation and better tackle the compossibility challenges this is done by switching to per parameter sharting based on D tensor instead of a flat parameter and enhance the per parameter sharting performance after careful optimizations we found that the performance of the fsp2 is on par with flight parameter fsp and with even better memory usage it also offers simpler 3D paradism composability uh capabilities communication shed communication free shed DCT for checkpoint 7 load and could compose easier with other pator core features like things like meta initialization or features that needs to manipulate individual parameters for example flow a quantization we actually enable the flow a training with fsp2 and it works very well with great performance which I'll cover later the tensor or sequence parallel in torch tan is using the pytorch native tensor parallel API that is built on top of the tensor it gives the user the ability to implement sequence parallel easily and apply tensor sequence parallel to the model without touching the model code tensor parallel allow larger model training by lower the effective local batch size to be smaller than one to save the activation memory and it performs inode communication by leveraging the fast EnV link connection when using it together uh with fsdp this is also known as uh two 2D sharting it could effectively it could effectively reduce the fsdp communication latency in large cluster furthermore given that tensor sequence parallel introduced blocking communication we have also enabled float a all together with tensor parallel torch compile based optimizations to allow overlapping between compute and communication to unleash the four power of TP pipeline paradism pipeline parm inan utilize the pyto native pip planing module and it could be used together with the 2D sharding to form a 3D parm setup or if you are in a bandus limited cluster you may want to train mainly with P parm as incurs less communication for 3D paradism it allows scaling on large cluster by further reduce fsdp communication latency the py native pipelining module offers an easy to ous model splitting tool to perform a manual or Auto split and a set of popular pipeline schedules including gpip 1 f1b zero bubble and lo schedules furthermore Advan user could Implement their own schedules based on their specific use case now let's look at the checkpoint 7 load to support longer training and failure recovery it is critical to save the checkpoint periodically and reload them whenever there's a failure happen P offer a distributed checkpoint module that could allow efficient safe and low during training it offers a communication free shy state that could represent the 3D paradism from torch Satan where each parameter is represented by deter for example if we look at the oper projection layer inside the attention of the Llama model it could be represented as a sh zero Shard one placement on a two-dimensional device MH to represent the sharding on on the data parallel and tensor parallel Dimension correspondingly pipeline parallel here is identified by the submat uh from the 3D device mesh that involves the TP and DP Dimension we can see that pipeline parallels stage zero and Stage n they have different 2D device match DET tensors given the easy to understand uh short static or DCP offers efficient checkpoint s low solution including asynchronized checkpointing uh that brings over 10x performance Improvement which is sh by uh showed by our uh partner at IBM checkpoint load rearing to different War size this is very useful when you're uh training at a happens in a specific GPU scale but you're you want to continue the training with a smaller World size and easy to transform to a for for dat for model publish next let's look into various P torch native training techniques that have applied in torch Titan to further improve the training efficiency in today's talk we're mainly focusing on three techniques activation recomputation torch compel plus distribut training and mix Precision training with flow eight activation recomputation this is an important training technique that could help lower the memory consumption by treating with treating it with comput the idea is to drop the forward intermediate uh competition result that is supposed to be saved in memory for backward and recompute them again in the backward in my opinion there are three tires of activation checkpoint first is full activation checkpoint where you recompute every forward activation in your back cord this will give the most memory savings but it's also the most expensive Choice second is a selective activation checkpoint where it allows developer to pass in a custom policy to decide which activation to recompute or not this this will give medium memory savings and medium performance slow down third is auto activation checkpoint theoretically we should be able to figure out uh recomputation policy automatically based on the competition and communication happen inside the model this will yield the most efficient memory saving and still keep the training be be more performant inside torch Titan we have enabled four activation checkpoint and developed the uh selective activation checkpoint policy for Lama like models uh using the improved activation checkpoint API from pytorch we're also actively working on enabling something like Auto selective activation checkpoint policy next thing I want to highlight is that uh uh torch comp the pytorch compiler that could automatically capture the graph from pytorch Models to improve the performance with fusion and more optimization opportunities uh in torch Titan we adopt a regional torch compel approach where we apply torch compel for each Transformer block this allows a faster code restart by reusing the compel cach for the repeated model architecture this is a new feature that is being developed inside torch Dynamo it also allow us to only focusing on the per layer compel compatibility instead of navigating the full model compile for optimizing the parallelisms with torch compile we mainly focus on the intr layer optimizations when applying tensor sequence parallel the P torch D tensor is now fully compatible with torch compile so does the tensor parallelism apis the compiler could capture both the compute and communication operations and based on that we Implement two important optimizations the first optimization we call it async tensor parallel where we implemented a micro pip planing optimization to overlap the all gather or reduce GA communication together with the MMO compute uh we use approach based on uh semant symmetric memory the second optimization is that we have a compiler path that we could reorder the communication inside the tensor parm backward to overlap with with this compute because of the regional compell approach we could simply graph break around the uh fsp2 per layer boundary and we don't really need to worry about pipeline parm because it's layerwise partitioning enabling torch compile with those optimizations would greatly improve the training performance next let's go through the mixed Precision training B FL 16 has been widely adopted uh in training uh inside torch Titan user could config a mixed Precision training policies through fsp2 to specify the parameter and reduction D type for newer Hardware Beyond h100 where it offers the flow a training capability the pytorch team has been working very hard to get the pytorch native flow a training receipt to be ready we have developed the flow a quantization inside torch AO where it supports Dynamic delayed or even rowwise scaling and here's an API to in place mutate the model to start the flow a training flow a is kind of special compared to B FL B FL 16 it needs to have a finer control on specific Ops for example MMOs instead of casting all the activations or parameters inside our model uh this bring complications to paradism as casting now needs to happen uh inside linear layers only and uh do need to be paradism aware for example we need to cast indiv indidual parameters and ensure that it is casted before the communication happen in fsp or TP we similarly integrate the flow a dynamic scaling with fsp2 and tensor parallel and furthermore turning on the torch compel for the full F performance let's take a look at how uh 48 works with fsp2 since fscp keeps shed weight and all gather the weights during training we want to C the sharded weight inside the linear layer to flow eight and do the flow eight computation or fcb2 brings tons of flexibility here because of the per paramet sharting that replace the flight parameter with casting individual shed parameter it will perform an all reduce to compute the ax or scale and user could choose to all gather whether to allgather the flow in flow a to sell the communication volume we also put many efforts optimizing flow a training with fsp2 we found that the scale computation for all sharded parameters are expensive because each of them need to do a individual a ruce so we offer an API to explicitly precompute all the sharded model parameters with a single R ruse we can take a snap at the performance Improvement when using fsdp and flow8 uh quantization together flow a comput brings over 40% C curve gains and enabling flow it all together brings another uh 7% so in total 47% gains next paradism we need to work on is tensor and sequence parallel since our tensor or sequence parallel is implemented using p d tensor and the P native flow a is also using a tensor subclass called flow a tensor so mainly here we are we address the uh tensor subclass interoperability problem where when the flow a casting happens it converts a plan D tensor that may be holding a b46 d type uh local Shard to a d tensor that holds a FL 8 tensor local Shar since the scale computation is actually a reduction operation it will simply leverage the DOR reduction operator to perform the r Ru we also improve the torch compel stack to handle tensor subclass and Ne tensor subclass uh to support this case we offer dedicated tensor uh FL a tensor parallel sharting styles that encapsulates the flow a casting it ensure that flow a activation casting in terens or sequence parallel happen before it's all together therefore housing the communication volume and improve the training efficiency to to achieve the best performance we could turn on async tensor parallel once we turn on Torch compile that uh this is a recent feature that we develop that could f fully overlap the communication and computation here are some results that we have when training with 2D sharing and enables tensor parallel all together we can see that FL a compute plus tensor parallel all together gives over 35% gains and ASN TP you addition bring another uh 3% performance so in total uh 30 38% gains beyond all these training techniques we also enhance the many aspects of torch Tian uh to make the pre trining become easier first thing debugging we have been developing a flight recorder that could help debug uh help debug the stocking job uh from the collective information training at Large Scale expose more challenges that the flight recorder could help navigate a stu job or crash quickly developers could enable this debugging tool by setting simply setting this two environment variables the torch niik TR buffer size which controls the buffer that saves the n and the torch ni D on timeout which allows uh distributed to write out the diagnostic fails on Collective timeout this will greatly help uh the developers to navigate failures when training at Large Scale second analyzing the ongoing progress of a training job is important where we have implemented many things like metrix logging profiling and loss curve measuring uh this this ensures that we can observe every aspect of the TR training process while we're on training third we implement the metad device initialization for large models which empowers Ultra fast model initialization after sharting and leveraging detor meta initialization uh to ensure numerical correctness this matters a lot when your model become like super big and initializing on CPU take a very long time last but not least I want to highlight that torch Tian features into ability within the P torch ecosystem where users could training their models using torch Titan save a model checkpoint and then Lo that checkpoint later for fine tuning with torch tune or go inference with GPT fast or torch chat next I'll hand it off to linon or close partner from IBM research to share their experience as the early adopter of Tor Satan at IBM thank you uh thank you thank you wo for inviting me to this session uh so I think wo just shared a very detailed overview on Torch Titan and uh well we seen a lot of great features and how t Titan can help large scale RM training uh but I think the are still a couple questions remain one we want to know how easily we can turn the this three po into production because remember here we are talking about training and when we talk about training we mean scale and we mean production so no matter how fancy this playground looks like we want to know how easily we can turn this playground into production in our own environment uh maybe by either switching from our existing code base to Tor Titan or borrow some of the features from Tor Titan into our own code base uh the second question might be you know what kind of performance can we expect from this repo right because all of us already have some existing code base that is already running maybe some 100% tou native some using other code bases like uh magron deep speed so we want to know what kind of performance we can expect if we move to Tor Titan or if we use torch Titan so we can compare the performance and uh you know make better decision on you know what we should do if we should move if we should borrow the features or if we should stick with our existing code base so for the next two slides which are also the last two slides of this session I will just very brief briefly talk about our experience with t Titan so uh yeah the first question is what it really took us uh so it turns out it actually only took us like four four five hours with only two PRS to make torch Titan uh really production in our uh environment uh the first PR we did was to really uh replace the data loader inside Tor Titan with our own data loader I mean after all t Titan is still a demonstration repo so for many of the components like data loader they just use the vanilla data loader for Simplicity but the vanilla data loader is not suitable for you know real word real large scale LM training so the first thing we did was to replace the vanilla data loader with our own production data loader uh I saw many of you were there yesterday with our uh full session on data loader so you know what I will be talking about but for those of you for those of you who are not there in very short we had a pretty powerful production level data loader that is 100% touch native so really it's just to us a few you know clicks on copy paste and swap the interface and we were able to just drop our data loader into Tor Titan uh and we fully open this and currently it sits in one of the open PRS in the torch Titan repo so you know welcome to check it out and if you want to know a little bit more uh you know welcome to watch the replay from our yesterday session uh but in short that's really the first thing we did we replaced the data loader we replace the vanilla data loader with a production data loader uh the second thing is uh even minor we just you know dump our own model variance into Tor Titan you know we all have our own you know sort of In-House model variance uh but really all these two PRS were really what we need uh in order to make torch Titan production ready in our own environment so hopefully this give you a little bit answer on the first question you know how easily we can turn t Titan into you know your own production environment uh so next just couple numbers you know we want to know what can uh well t Titan post posted a a few numbers there but we really want to know what kind of number we can get so we run dozens of different benchmarking performance numbers uh this is far from being a comprehensive list but we handpick a few of this uh you know for some standard configuration that I know many of us have been using uh so just to give you some ideas on some of the numbers we run towards Titan as is so all these numbers are as is without our further tweak uh we did find we can you know apply some further tweaks to make these numbers higher uh but we tend to put you know all these n all all these as is numb so you get the idea on you know if you just move move t Titan or if you use t Titan in your environment as is with similar setup what kind of performance you can get on these different model variance uh yeah so these two slides are really what I have so hopefully could give you some sort of ideas on the previous questions that I mentioned thank you