Sponsored Session: NeMo-Aligner: A Scalable Toolkit for Model Alignment - Gerald Shen & Jimmy Zhang

Key Takeaways

so let's maybe get started um so yeah okay uh thanks for coming to our talk today Jimmy and I would talk about Nemo aligner a skillable toolkit that we built for model alignment uh I'm Gerald I'm the framework lead for Nemo liner and Jimmy is the perf lead for Nemo liner so to start I first want to talk about the nvidia's nwin LM software ecosystem that we built so everything here is open source and Nemo aligner Falls inside the Nemo framework uh more broadly so Nemo framework is essentially um a core component of our llm training stack where it focuses on very high performance pre-training and as part of that aligner builds on top of everything we're using such as Transformer engine Megatron core which Megatron LM uses and we build on top of that to do efficient alignment this involves using uh par Ms and high performance training kernels and as well we also take advantage of the INF stack that we have inidia um in particular trlm that Jimmy will talk about later on so to motivate why we need to build a framework um I want to go through the definition of alignment um alignment broadly speaking is getting Foundation models to follow instructions and generate human prefer responses and nowadays almost every model uh that comes out under go some alignment stage uh to demonstrate uh what the alignment stage does to a model I could let's say have a query that is such as uh how many units are in a dozen and the pre-train model let's say 3.1 AB base would sort of know how to answer this question but it's not generally easy to read and it sort of just blabs on the answer over and over again whereas if you look at the instruct model let's say 3.1 AP instruct um it response is short concise and just generally what um I like to see out of a model that comes out so what Nemo liner is doing is creating a toolkit to go from models on the left and turn that into models on the right um for Neo liner we want to implement state of the AR alignment techniques while also making sure that our performance and scalability are great so by state of the art techniques we sort of say um have rhf DPO RPO constitutional AI uh sft and while we implement this we sort of want to make sure that the performance and scalability are also great this involves uh using model parallelism doing accelerated inference which is what Jimmy is going to talk about next uh having FP support and we're currently working on Long sequence support and sequence packing um as part also we want to make sure that the open source Community can sort of uh use the toolkit easily so we open source the data sets that we've used for for some of our alignment pipelines let's say h steer and H steer 2 and during and eer uh we also open source and commercially available uh model checkpoints uh which include the 340b instruct and the 340b reward model that we was trained with Nemo liner so uh to sort of double click a little bit I want to talk about the performance optimizations that we've been doing and um in particular for the rhf pipeline so rhf is stands for reinforcement learning from Human feedback and this is sort of the algorithm that got us from uh to chbt essentially um this is uh prescribed very similar to the instruct PD paper and it involves generally three steps so the first step is the supervised fine truning Step and this is what happens when you take your pre-train model let's say 3.1 base and ask humans to let's say generate what the model should have said so you give it a question and the human would write out what exactly to say and then in this stage we sort of just like train on this annotated data and this could get you pretty far but if you want to have even more performant and a more aligned model you would go and collect this thing called the preference data set the preference data is essentially um um given a prompt we have two responses and sort of the human labeler will pick which response is better than the other and with this preference data in mind we could go and train a reward model and what this reward model that comes out of would do is it would assign high scores for responses that it deems as good and low scores for responses that it is deems as bad with the reward model in mind we could go and plug this into a very canonical RL process and which is uh what we talk about when we say uh po so in po we have let's say our policy Network and this is like the GPT that we want to turn to chat GPT and we have a bunch of prompts and we sort of go and use this against the reward model to do the whole R process one thing I want to keep in mind is um even though we're talking about just the R Performance generally these optimizations would apply to all online algorithms like rejection sampling and other algorithms but we just focus on RF because that's the most wellknown to talk more about step three um we sort of have this kind of like a messy diagram here but roughly what happens is that we have the policy in that Network which is a network that we want to turn to jbt we sample from this policy Network and we ask the reward model to say how good it is uh roughly speaking when we do RL we're just trying to increase the likelihood of the good response and decrease the likelihood of the bad response but in rhf it's slightly more complicated than that because we have this thing called the initial policy and the value Network so uh these two things are mostly here to stabilize training in in particular the initial policy is to prevent uh deviations from that are too far from the initial model and the value network is here to stabilize the RL training in here and Jimmy will talk about this later on we have two distinct stages in the r pipeline one is sort of the inference stage and that's when we generate samples and run inference on all of these models once we finish the inference we have this thing called a roll out buffer which caches uh the generation generator results and then we go and do training so one unique thing about rhf is that we have inference stage and training stage and this is part of the greater training pipeline uh in contrast with pre-training we just really do training and during evaluation we almost like just doing training without gradients uh but here we're actually doing autog sampling and um inference on many different models so um as you can see um ourf is complicated which means that per optimizing the performance of RF is also complicated first we need to orchestrate four different models on the previous Slide the policy the initial policy reward model and the value model and these four models could be of any size let's say for VB one thing that to note is that the policy Network and the initial policy are architecturally very similar and the reward model and the value Network are also architecturally very similar so for Aigner what we trying to do is we we want to make sure that uh um it's as fast as possible and we want to make sure that we're the fastest open source RF implementation um to do this we have three key performance optimizations one is distributed architecture two is trtm fast inference three is dynamic low balancing and I'll let Jimmy talk about these three hey everyone I'm Jimmy I was the perf lead for Nemo liner and yeah I'll be I'll be doing a sort of flyby of a the per optimization we do to make um RF fast so the first per feature we talking about is sort of our general architecture to how we Implement um rhf so in Nemo liner when we launch a job uh we launch two jobs per each rhf run so we'll have one job that holds the value and reward model and we'll have another job that holds the policy and initial policy and the reasoning for this is that it maximize performance uh generally when you do rhf you might have like very very differently sized um awarded policy models so for instance your REM more model might be like a really big model like llama 45b and your policy might be a smaller model so you want to be able you want to have the ability to precisely tune what gpus go to what job right because obviously the smaller model needs less gpus and a bigger one you want to give more so we have these two heter heterogeneous jobs um that run in parallel um these two jobs uh coordinate asyn asynchronously so this means that they're largely independent and when they do share communication for instance the generated de generation results and the loog props um that communication process happens asynchronously um so because these two jobs uh run independently we're able to maximize the GPU performance because they're asynchronous and because we can precisely tune um the number of gpus per model uh so to perform to communication we use this software called py Triton um which is a l which is a LM serving framework um that Nvidia has part of its Nemo um broader Nemo uh framework um another detail is that uh for each job we have two models right um we are we have the ability to dynamically switch between these two models so let's say you're doing some inference on the value model and you need to do some forward passes on the reward model um in order to like minimize GPU idling we on demand SWA rates so we push the value model to host and we load the reward model onto GPU um at runtime and this because it reduces like the amount of jobs it simplifies the orchestration of these models um it really improves uh overall the GP utilization the next feature I'll be talking about is our ability to use tensor tlm for inference so tensor rtlm is another uh software that's part of like the broader Neo framework and in particular it's Nemo's it's nvidia's uh solution to perform really fast uh production level quality uh inference serving um so as J mentioned there's this really big inference stage where you have to sample your model with a lot generation and in fact this is the bulk of the comp uh the computation of rhf so if you want really great performance you have to make sure that you have really great generation performance so the way we do this is that we use tensor rtlm um so we take our Nemo model which is doing training and on the first step we actually create a replica of the model inside trlm and because um we have this replica model we can do autogeneration inside trtm um so shown here in the diagram um the trlm engine will produce Generations the Nemo model will consume those generations and do its training and when it does training it'll call its Optimizer to step the weights and send those updated weights back to trlm uh so because trlm is you know part of the broader Neo framework we have really great interop ability between these two Frameworks in particular we're able to update the weights inside t tlm with really really really fast um at really fast speed with really low low overhead so this like process of switching between T tlm and Nemo and you know vice versa happens really fast if really low overhead um in particular the thing we use is called a tensor refit which consumes the weights that that the N model produces updates its weights and then has updated weights produces new generations and so on so forth uh another detail that I want to mention is that uh inference uses a lot less memory than training and this is because inference is just you know lots of forward passes so we don't need things like optimizers or gradients so when we transition from when we switch from the from training stage to inference stage what we'll do is we'll actually reshard our M parallelism um from training to inference uh to a low a lower modalism and the reason we want to do this is because typically in inference we have much higher computational costs because we are in that auto regressive Loop so we we want to minimize minimize those uh communication overheads and maximize performance in particular what we commonly do is that we have the ability to unfold to train pipeline parallelism so if we're training if we're using some pipeline parallel size uh when we when we transition into the uh trt uh inference stage will unfold the the pipeline parallelism and run that as data parallelism and this you know really improves our aurf um some other things we do is that as I mentioned we will unload the optimizer in gradients and for training you know uh because we're doing training we we don't we don't need the engine the TR tlm engine so we can also unload that as well during training so overall the memory footprint of incorporating trlm isn't much higher than training um and because we are able to you know have low memory Footprints um and because we inherit really good perf from tlm um this whole um utilizing TTM for inference um really uh gives us really great perf with really like not much downsides the last um really important part feature that I have is is um how we handle low balancing so typically when we you know run that scale um we'll run with some amount of data parallelism so this means taking the input like input batch and sharting the sharting it evenly across some data parallel workers uh for uh generation there uh because generation is like each uh generation might have a variable length of uh tokens generated right because you're just sampling the model um if you give each DP worker the same amount of prompts uh you'll get idoling and this is because you know commonly what will happen is one DP worker will you know get prompts that cause really long generations and you'll Force the other DP workers to wait and this is not good because when you wait you have lower performance so our solution is to dynamically allocate um our data loader from from from like a server so what we will do is on the First Data parallel rank we'll have a small server that dynamically pushes uh like requests from the uh from the DP workers it it'll receive requests for work and they will push prompts to those DB workers each worker will like consume those prompts du generations and then they'll push it push those uh Generations into like a global R buffer and when we finish inference we'll rebalance that Global rout buffer for a training and this is great because um generally we'll for rhf because when we um when we run at scale we'll have many many DB workers so this is a really um effective way to make sure that D work is constantly being utilized which is great for perf yeah so I've talked about three really important per optimizations um here in this table we have a sort of a small oblation where we where we turn off each per feature and you can see each of them gives really great perf um the most important in particular the most important per optimization is is indeed like tensorrt L inference and that optimization gives us a 7x speed up Alone um overall our perf uplift against a a naive implantation is 15x which is really big um 15x is huge it's like if previously your run took a week you can now finish it overnight so it it's really important um to have really great perf and you know it's it's really important to for productivity too because like now that your experence run fast um you can also iterate really fast um which which we found to be really important um this is the same data as the previous slide but we present it as a graph and on on a vertical axis we have step time and on the horizontal axis we have like the optimizations we've talked about and you can see that um each optimization sort of targets a different part of the ENT time so you can see that each each chunk of the like vertical bar is is is a different part of the rhf um Step time and you can see that each obervation targets like a different color and you can see that in our base line by turning on all our optimizations we've shrunk every chunk so this means that um our optimizations sort of Target the whole ARL process like holistically um to deliver like really great performance at every part of the uh process so perf is certainly very important but you have to make sure that perf is robust and one of the ways um you want per to be robust is that you want perf to scale so this means that if you have thousands of gpus you want to make sure that at thousands of gpus you you still have really great perve um overall NE L achiev this achieves this we have uh we we're able to observe a 1.8x weak scaling factor with data parallelism so this means that let's say you want your experiments to run faster what you can do is you can just double the amount of gpus and Nemo liner will almost will like efficiently be able to use that EXT those extra gpus and run about 1.8x times faster so this means that if you have thousands of gpus you can use you can use all of them and your your experience will accordingly get also run really fast and and the reason um we can get this really excellent scaling is for many reasons but in particular in particular it's because of the really good Dynamic low balancing we have that I talked about in in um the slides prior uh so we not only can scale to thousands of gpus but we can also scale to the very largest models so for Neo liner um this was like a top priority and we support scaling to the largest models by supporting 3D modelism so for both training and inference um we support 3D parm so we can do data parallelism tensor parallelism for inter no scaling and pipeline parallelism for inter no scaling so because we have like this really great scaling from theism we can make sure that even if your model is really large uh we can always fit it uh because we have um these ways to Shard the model into like over hundreds of gpus um as as an example here we've shown us training um llama 405b right so we have a 405b actor a 405b sorry 45b policy a 45b role model we run this on a thousand gpus and we still get really good perf yeah so all in all um I want to highlight that Nemo liner is really fast um we were a we Implement you know all these State alignment algorithms with excellent performance um in particular rhf in N liner is faster than any other open source arlet implementation to our knowledge uh neon allance performance is robust so we can scale to thousands of gpus to the largest models um we also provide many knobs to T performance to your particular workload and setup so for instance we have you know 3D modm which I've described recharting this distributive approach to rhf we expose many performance knobs so that even if your like workload is kind of unique chances are you'll be able to toune perf with these knobs so that it's also excellent for your case um as an anecdote um in our experiments we were able to finish llama 370b rhf in just under two hours and using only 12 h100 nodes and this is actually the fastest we've ever seen um RF performed yeah thanks so um to conclude things uh for future work we're obviously going to look into more alignment techniques um in particular one thing that's of interest is doing knowledge isolation during alignment and this proposes new challenges because we also have to run a large model inference while doing the alignment process we also want to apply trtm generally to other online like algorithms like reinforce leave on out rejection sampling and obviously we all want to uh further improve and extend performance this involves during infuence faster engine offloading faster pipeline parallelism using Dynamic batching uh more work on low precision and in training we want to do sequence packing context parallelism and also more work on low Precision we also want to EXP expore other modalities as well um this concludes our talk um our repo is linked here it's open sourced uh we have the docs and paper and thanks to all the contributors on the paper and as well wwin who's our uh product manager yeah than