LLM Avalanche: Xingyang Gang : Building Simple and Scalable Training Framework for LLMs

LLM Avalanche: Xingyang Gang : Building Simple and Scalable Training Framework for LLMs

Recording: LLM Avalanche: Xingyang Gang : Building Simple and Scalable Training Framework for LLMs

thank you hi everyone my name is hin Young I'm a PhD student in UC Berkeley and as Shreya mentioned last time in last talk that training LMS is very difficult and talking with the model is also not easy to do and today I'm presenting eclm a simple and scalable framework precisely to do that and hopefully with this 15 minutes of talk I will convince you that it's actually not so difficult to customize and train your own large language models so if we look at existing Frameworks for training large language models I will borrow this picture from this famous picture from Angry kapasi from his most famous mean GPT now Nano GPT implementation most of the things Frameworks four into two categories the scalable Frameworks are very complex they're like a humongous Battleship that has all these spells and whistles and also the simple diabetic Frameworks that are very easy and minimal and very easy to use so let's dive into each of these Frameworks and see why that's the case so if you look can mostly these educational framework that's very easy to use and easy to learn we see that this is because essentially causal language models are very simple they're standardized Transformer architectures with only minor tweaks and the training methods for these language models are also quite standardized you just use supervised fine tuning you use atom Optimizer all these things have basic standard implementations in today's neural network libraries and only require minimal coding on the user side so why do we still have these complicated language model Frameworks the heavy lifting is mostly for the scalability the previous minimal language model framework if you just wanted to run it on a single GPU or on a few gpus on a single machine that might be sufficient but if you want to scale up your training in a cloud on many different machines and hundreds of gpus or gpus it's very difficult so these scalable language model Frameworks are often built to handle that case on the other hand these scaling mechanisms are often deeply integrated into the model and training Loop make it very difficult to customize for example if you want to build your own language model you want to take an existing complicated implementation it's very difficult to do and I was wondering can we actually bridge the gap between this minimal language model framework and the scalable language model Frameworks and I believe with today's machine learning framework we actually can and this is my analogy here it's like a futuristic Battleship that's very powerful but automatically control that means it doesn't need much input from the user to be able to do a very scalable heavy lifting so how does it work the key ideas is we can actually keep the model and the training Loop of the minimal language model framework to be minimal and simple and we leverage modern machine learning compilers to automatically scale up the computation so this is the framework devop called easyom and as you can see it gains quite a lot of stars in GitHub and you can also scan this QR code you see the framework itself easyim has a lot of support for existing popular open source language models such as Lama and gbtj and it also integrates deeply with hugging phase Transformer ecosystem that means that frame language model trained with ecim can be easily exported to the hugging phase format and it supports flexible parallelism and has a lot of building bells and whistles for model serving and evaluation so let's dive deeper into it so first of all when you write the language model in Easy RM training you mostly just write a model for a single GPU so look at the training code here it's essentially just Computing the largest and Computing the loss as if you're doing the computation on a single GPU this gives this makes it possible to directly expose all the essential model details such as the model definition and the loss function directly to the user without dealing with hundreds of lines of complicated code this makes the language model really customizable but how do we scale it up we'll scale it up automatically with modern machine learning compiler so since ezrm is building Jacks we are able to leverage Jack's paralleloget computation compiler to automatically transform the simple train Loop we defined before into scalable multi-host multi-accelerator programs so how to do that the first step is for user to Simply provide what's called a sharding rule and these charting rules are mostly standard for modern Transformers you can just copy it from somewhere else and you don't have to really write it yourself and once you have the sharding rule you can use this compiler to directly transform the single GPU program you defined before into a scalable program so after you transform it into scalable program you then provide simple specifications of parallelism by basically providing three numbers of the shapes of your accelerator configuration and easy item supports a flexible 3D mesh of parallelism including a mixture of tensor parallelism fully shorted their parallelism also known as zero stage three and also the normal replicated data parallelism and as the example here I'm showing you simply just provide three numbers then you can scale up your training program to 256 accelerators covering a group of four data parallelism groups 16 ways of fully shattering data parallelism and four ways of tensor parallelism now with all these simple program to set up how does it perform they actually performed decently compared to both existing open source implementations and also commercial Solutions here's a comparison of our throughput in training our 7 billion open Llama model we can see that uh a100 we actually achieves comparable throughput to mosaics MPT 7B model which is a very popular commercial solution these days acquired by databricks and also we're doing better on TPU than Google's t5x open source solution which is a much more complicated implementation which is better for larger models but not so good for these smaller models and running ecim is also very simple all the training configurations is basically exposed via command line arguments this means that you only need to write one single stroke for all configurations you don't have to deal with all these complicated levels of yaml configuration file just one single script you can launch your experiment and run now in today's talk I also give you two use cases of ezlm of models trained with Ecom oh one more thing to add here we also have built-in support with serving the model and also evaluating the model and having a chat UI so you can play with the directly play with the model you train today we'll also give you two examples of model training with ECM hopefully to convince you that it's actually not so difficult to do so the first model is what we call the koala a dialogue model we train here in UC Berkeley it's essentially very similar to all these various of styles of fine-tuned llama including things like alpaca and vicuna and basically we took the shared gbt data online and fine-tuned our llama to become an instruction following model the fine tuning as you can see here is very cheap it's fine tuning on a single agpu machine for only six hours this means that you can actually do the complete fine-tuning run with less than a hundred dollars on cloud with spot instances and I'm pretty sure any company can afford this cost so hopefully this will convince you this it's not so costly and so difficult to do now how does this fine tune model perform it actually on a recently released results of LM chat Arena which is comparing lots of different chat boxing user study we showed that koala actually performs very comparably to the top of the line fine-tune llama model which is uh by Kuna and in the further study we study ourself we actually found out that some of our data we used during our fine-tuning time actually hurts the performance of the model this really showcases that you know curating high quality data set is very important which is also another reason why you should fine tune your own language model instead of using an existing language model via API now this is a fine tuning this is for smaller scale study how does easy arm help us train large language models scale up on hundreds of accelerators we also did a project for pre-training language model which is open llama one of the one of the main bottleneck we see with llama style models is that llama is released not under a commercial licenses under academic research only license and many people who want to use llama for their startups for their business cannot do so because of a license and on the other hand people have built a lot of great ecosystems around the Llama models such as you know llama CPP such as these very fast on-device inference engines so seeing these opportunities we decided to take some of the computation resources we have also some resources sponsored from stability AI to basically replicate the pre-training of llama to create a completely permissive licensed llama reproduction that can serve as direct Replacements as the original Lama weights so we released the serious operation including the smallest model like a three billion Model A 7 billion and a 13 billion model on this red pajama data set which is a reproduction of a llama data set released with uh from a company called together and the model pre-trained on these representment data sets can directly serve as dropping remake placement it used the exact same llama architecture any code you written for llama can directly use this on our zero shot Benchmark evaluations we actually also see that our reproduction of lamba reaches comparable performance to the original llama models which is shown in a table here so this hopefully can convince you that actually training your own language model is not so difficult and you can leverage a lot of Open Source Frameworks in the open source model as your base model to start fine-tuning and customize it for your own business if you're interested in this work please scan this QR code to CR framework it up and also maybe talk to me after the talk all right [Applause]