Infra behind Krea 2: How to train and serve at scale — Gabriel Jorge Menezes, Krea.ai

AI Engineer · 16 min · 279 sentences · from YouTube's caption track

Each timecode opens YouTube at the start of that sentence. Line anchors (#s42) are the cue ids in the WebVTT, and every line carries its start and end seconds. All transcripts has every talk, and the whole corpus as one file.

00:01[music]
00:12Hello everyone.
00:14My name is Gabriel.
00:15I work at Korea and I'll be talking about the infrastructure that allowed us to train K2 and also how we serve it.
00:23So, what is K2?
00:24K2 is our pre-trained from scratch model we just released like less than a month ago and the whole idea about training this model was because we were kind of bored of AI images.
00:35They're quite, you know, soulless.
00:38They have no spice.
00:40And the whole idea was we want to give creatives tools to to explore out of distribution extremely interesting images, do composition and like actually give tools to creatives and that was the whole idea of the model.
00:52The model was trained from scratch.
00:54No base checkpoint, not anything.
00:56Everything done in house and also as I said, built for exploration and right now this is what you can get out of Korea 2. Very different styles and, you know, like pixel art and and like
01:11like photo real and like some silly stuff.
01:14Uh, and the whole idea of the model, as I said, let's explore this medium.
01:18Uh, Korea 2 is open source.
01:20Right now, you can go play with it.
01:23There's two checkpoints.
01:24We also serve in production.
01:26There's a raw checkpoint we just pre-trained so people can post-train and do whatever they wish to do with it.
01:32And there's also the post-train version which is the turbo one which is very very fast.
01:35You can get like an image in like, I don't know, less than a second.
01:38And this is just like this type of images you can get in less than a second.
01:42Uh, on Hugging Face GitHub, go go play with it.
01:45Also, you can go in production on korea.ai and go play with it.
01:48Uh, so let's talk about how we trained this model.
01:52Uh, as I said, it's going to be how we train and how we serve.
01:55Uh, first the model I was trained from scratch on thousands of GPUs.
01:59Uh, we have a big cluster, one main cluster with a lot of GPUs, all Infiniband connected.
02:03And you put those GPUs work and it train.
02:06But like that that's I wish was that simple, but it's not.
02:09Uh, so at the beginning we did a bunch of like small ablations on like like a small number of GPUs to see like how things will work.
02:17So you want to test some hypotheses and you do a small number of GPUs and let it train for a little bit.
02:22Oh, this works, it doesn't work.
02:23Let's scale.
02:24Uh, and as we are like training this model, the whole idea was to like kind of bridge the gap between like LLM research and diffusion transformers.
02:33Uh, so my AI researchers they ported a lot of research from LLMs into into DiTs.
02:39And the whole like like the whole like architecture of the model was meant to be extremely extremely simple.
02:46And so like it is very very dumb, but like very effective.
02:51Uh, and so let's start talking about numbers.
02:54Uh, incredibly our maybe skill issue on our part, maybe our cluster uh, was very interesting as we were like scaling.
03:02Uh, when you did like small experiments, experiments would like run for days and like even less than we would like to, but like they would still run fine.
03:09Uh, and as we like start scaling, like getting like more and more and more GPUs, like 128, 256, 512, whatever number, and like you scale and scale, like things are crashing more.
03:19That's expected, right?
03:20Like there's more surface area for things to break and things going to go wrong.
03:25And a lot of the times things would go wrong in silent ways.
03:28I don't know nickel timeouts, like just crashes and you like the metrics are all good.
03:33Uh, and it is extremely annoying stuff.
03:35At the beginning we were like paranoid, swap node, change node, whatever, whatever, whatever.
03:40And we learned that like sometimes you just let it crash.
03:42It crashes, like it runs for like an hour, crash, runs for an hour, crash, and then runs again on the same set of machines, same code, same data for like 12 hours, 16 hours, 24 hours.
03:52Uh, there is this paper from Meta that kind of gives you a like rough estimate of like how many failures you should expect.
03:58Uh, we we kind of saw the same pattern but like not the same numbers.
04:02Uh, our runs would last way way less than this.
04:05Uh, so it was extremely annoying and you can imagine doing large-scale pre-training on runs that last less than 8 hours, it is a problem, right?
04:12You want kept those GPUs fed and if things are crashing, you are not doing progress and losing time and models are going to be late.
04:19Uh, so for us extremely important and like for at least from the infra side was to get metrics.
04:26Metrics are everything.
04:27That's how I can support my researchers.
04:28That's how I have visibility in the system.
04:30And like if you're doing large-scale pre-training, I highly highly recommend for you to invest heavily on metrics.
04:37Uh, don't don't go blind because you're going to go you're going to go crazy.
04:40So I'm going to share some of the metrics that like were important for us.
04:43And like they're quite silly but like extremely effective.
04:46Uh, first one like GPU temperature.
04:47Uh, GPUs are very very annoying.
04:49They if you're if you have a single GPU that is like a bit warmer than the others, it's going to start like throttling and slow down and and the training's going to be unstable
04:59and you have weird problems.
05:00So for us it was like if there is any GPUs above like 78°, you you remove them.
05:05Don't don't think about it.
05:06Don't try to fix.
05:07Don't don't try to be smart.
05:08You just remove the GPU.
05:10Uh, you're going to save time and just ask your provider and like this GPU's hot.
05:15Please replace it.
05:17Uh and this there there is two metrics that at the beginning we did not fully understand and as we are like getting more and more used to GPUs and like how they work,
05:26uh, there is GPU utilization which is a lie.
05:28Uh, don't trust this.
05:29Uh, this is dumb.
05:30Uh, this tells you, oh, the GPU's doing work.
05:33And this is amount of time GPU's doing work but like not good work.
05:36Not how efficient the GPU's working.
05:38So like as you can see during our pre-training, the GPUs are 100%.
05:41This is not true.
05:42We are not fully utilizing the GPU.
05:43Uh, this is 100% a lie.
05:46What we would use uh, as a proxy was tensor core utilization.
05:49This is actually how how much of a tensor core you're using and like how effective they are being.
05:53Uh and it was also very interesting as we're doing pre-training and then you go from pre-training to mid-training to post-training, you start like scaling the like the resolution of the images.
06:02So like for example, pre-training we do like I think 128, 256, 512 1024 pixels as they scale and you could see the tensor core utilization go up as we like would scale on these these resolutions
06:14cuz now you're doing more work on on on images.
06:17And like also very very interesting was InfiniBand and NVLink metrics.
06:21These by default not exported by uh the NVIDIA metrics that DCGM stuff.
06:26Some NVLink stuff, yes, but no InfiniBand.
06:28So if you don't have InfiniBand metrics, uh go go go get it.
06:32Uh I'm telling you right now, if you're doing large-scale pre-training with a bunch of GPUs talking to each other between machines and you have no InfiniBand metrics, you're doing something wrong.
06:41These was probably the most important stuff for us because most of our failures were like related to like cross-node communication.
06:47So like for example, here you just have throughput, but like on our like reference dashboard, we have a bunch of stuff like from like wait like when you when a message sent on on on the fabric like how much time the the the message is waiting
06:58or like the like number of errors and different types of errors and like number of packets and all of the things that like InfiniBand's exported uh we collect.
07:08We had to build custom stuff to get this.
07:10It was not hard.
07:11Uh you can figure it out.
07:12It's very very easy.
07:13Same thing for NVLink.
07:15Uh NVIDIA exports some stuff about NVLink, but like for example, NVLink NVLink errors, NVIDIA doesn't export this.
07:20Uh so you can collect this and like as I said, InfiniBand was extremely important and NVLink was a little bit less.
07:26Uh in some cases, these helped us catch some problems in like especially like cuz NVLink single node, right?
07:31Like communication inside a node.
07:33So sometimes a single node would have a weird failure where the GPUs seem to be fine, but like some weird error happens and then you can look at NVLink NVLink errors and like you see errors happening.
07:43And then replace that machine.
07:45So, go get this matrix.
07:47They they're extremely important and without this we would not be able to to train at all.
07:52Also, as I said, our our trains would crash constantly and a hacky way to do it to fix the problem just checkpoint.
08:02Use and abuse the file system that you have.
08:04At the beginning we use SEF.
08:05SEF didn't work well.
08:07Was very annoying.
08:08It broke.
08:09We lost trust in data.
08:11So, I recommend if if you have the money go go with something paid and cuz you can trust your data.
08:17You can see numbers.
08:18This is our our worker cluster.
08:19We can do like 1.8 terabytes of second of reads, almost a terabyte of writes.
08:25And the file system would not choke on the training.
08:27So, we could checkpoint every like 30 minutes, 20 minutes, produce like a terabyte of data in like less than 30 seconds.
08:34So, this would not delay trainings.
08:36That was like probably one of the like most important things that we need to to like recoup the loss and just checkpoint.
08:43Don't don't don't don't forget about it.
08:46And like how we serve.
08:47This goes in connection on how the trains are launched.
08:51Because at the beginning I don't want my researchers to think about GPUs.
08:55I just want them to launch stuff and this goes into a queue and if we have GPUs we have GPUs.
09:01If we don't have GPUs we don't have GPUs.
09:03So, queue.
09:05This is a open source project.
09:07You can look it up.
09:08It does gang scheduling.
09:09Gang scheduling is important for for for trainings in general.
09:14And this gives us a semantic of like two tiers of priority where you have a workload priority and this you can say like oh this training is more important than this one.
09:23So, it skips on the queue in front of the queue.
09:25And then after this we have the normal Kubernetes priority if you're used to Kubernetes.
09:31And the way this system works is like the training pods they always have like high priority for everything.
09:35So, like once they are submitted they going on a schedule.
09:38If there's inference running on those machines, the inference gets kicked out.
09:41And you'd say, "Oh, this is bad.
09:42Production is going to go down."
09:43No, you can build on top of that to to make production not go down, which is very, very cool.
09:47Uh the only one of the problems with Q, which is annoying, you can automate that.
09:51We have not.
09:52It's just that you specify the queues.
09:54The queues have like amount of resources, CPU, and like Nvidia GPUs, memory, whatever.
10:00Uh but this is manually like manually like specified.
10:05And at least our cluster is quite fluid.
10:07Nodes phasing in and out of existence.
10:09They go into to maintenance, whatever.
10:12You lose a few nodes here and there.
10:13Uh this number gets out of sync.
10:15And sometimes these these would break gang scheduling.
10:17Uh so, FYI, this this is a bit annoying.
10:20You're going to face this if you use Q. Uh but yeah, very good project.
10:24Kubernetes 1.15, uh we have not had the chance to play with it.
10:27Has gang scheduling uh out of the box.
10:29Something very similar to Q. Uh so, maybe you can use Kubernetes 1.15.
10:34Uh and as I said, this this is the system that we built that allowed us to train using the whole cluster.
10:39As I said, we have one big cluster that runs production and and trainings.
10:43Uh so, I don't want my researchers to think about GPUs, and I don't want make the trainings and research be delayed because production is running, right?
10:52Uh production is lower priority.
10:53The site still needs to work.
10:55People still need to be use the website.
10:56But like the GPUs, they're like the value that we get off the GPUs doing trainings is more like higher than than we get out of production.
11:04So, the whole system works by default where there is this magical system that I'm going to explain in in a bit that allows us to flip traffic between clusters magically.
11:12And not just clusters, and like external providers, GPU rentals, whatever.
11:16And you can see like the green like the the dark green is like inference running in cluster.
11:22Uh and then someone launches a train, and then suddenly start flipping to the other cluster, and then training is done or whatever happens, it flips back.
11:31So, we stop wasting money.
11:32And this is seamless.
11:33No one needs to think about it.
11:34Uh the whole system like like handles itself.
11:39And you get this this very nice pattern of like I'm going to use all the GPUs in my cluster uh for trainings.
11:45Production is going to run somewhere else.
11:46I don't need to think about it.
11:47My users on production they not going to feel anything.
11:49Uh research is going to be happy, and we can get values out of the GPUs.
11:54Uh so, how does this work?
11:56There's this very nice project called virtual kubelet, also open source.
11:59Uh you can build on top of it.
12:00It is a very nice code base.
12:02Uh and this works by creating a fake machine in Kubernetes.
12:05Uh Kubernetes has nodes.
12:06This creates a fake machine that is like up to you to control how it works.
12:10So, Kubernetes does normal scheduling as you would expect.
12:13Things would go into into these nodes.
12:15For example, here all the GPUs are in the cluster in use, right?
12:18So, this pod goes into the virtual kubelet node, and in there you can do whatever.
12:23This is the system that we built.
12:24There is like you receive the pods back, and then you find a provider.
12:28You can like This is up to you.
12:30Let's say you have a deal with some provider that gives nice prices.
12:32You integrate integrate into here.
12:34We built like some nice interfaces to be to to like not leak things.
12:38So, like we just implement a provider, and there is a a algo that decides which one it goes.
12:44Uh you translate this back of the pods into into the provider stuff, and you deploy, and then you have something that reconciles between the both sides.
12:51And what's extremely extremely nice, if you guys know about Kubernetes, Kubernetes has like the horizontal pod auto scaler, which like scales the number of of replicas.
12:59Uh number of replicas uh Could you stop being annoying?
13:04Thank you, sir.
13:05I appreciate it.
13:06There you go.
13:07Let's go back.
13:08Uh now we go back.
13:10Uh back back.
13:12There you go.
13:15Like Kubernetes has the HPA, and the HPA scales the pods.
13:17And so, if something fails, it is very interesting you don't need to handle the fail.
13:21Uh the only thing you need to handle is like, "Oh, something failed."
13:23You mark the pod as failed.
13:24Kubernetes is it's to detect that something has failed and create a new one.
13:27Uh you don't need to try to save the world.
13:29Let Kubernetes handle it for you, which is extremely nice way of handling stuff.
13:33If something breaks on your side, something breaks on the the other side, just mark it as failed, let Kubernetes handle it, create a new one, and things keep working.
13:40Very, very nice way to handle stuff.
13:41Uh and also very interesting, let's say you have GPUs on your cluster available, right?
13:45You don't want to waste money.
13:46This would be very, very bad.
13:47So, the system works about like works with like using taints, Kubernetes taints.
13:52They allow and disallow things to run.
13:55Pods have tolerations for the taints.
13:56And when we have GPUs in the cluster, uh if you look back, there is the taint system on the bottom.
14:02This taint system, it is what would like by itself decide what you like if you have GPUs or not GPUs in the cluster.
14:09Uh and this adds a taint into the into the node when we have a lot of GPUs.
14:14So, like a lot of GPUs in the cluster, we taint the node, nothing can schedule on it.
14:18Uh so, we stop wasting GPUs.
14:20Like the pods, they go into into the GPUs in the cluster, we don't waste money.
14:25Uh and then imagine someone launches a training, right?
14:27This training is going to take all the GPUs in the cluster, it's going to hog all of the GPUs, no GPUs in the cluster.
14:31The system detects this, removes the taint, new pod schedules there.
14:34Very, very nice.
14:35Uh you also don't think about it.
14:37Uh and for us, for example, we use just some Prometheus metrics.
14:40That's how we do it, very simple, uh but it works very, very, very well.
14:43Uh you let the system run by itself.
14:46Someone's going to launch stuff.
14:47You're going to like the training SLI is going to kick out the pods, it's going to schedule, it's going to take the GPUs, the system is going to detect that, remove the taint, pod schedules there.
14:55Very nice.
14:56Someone finish the training.
14:57Now we have pods running on the other side, you're wasting money.
15:00Uh how do we fix this?
15:02Uh same thing, you run something else that detects the system and removes things back.
15:06Uh in this case, a descheduler.
15:08Uh once the taint's added back, so like GPUs available, we add the taint, the descheduler says, "Oh, these pods, they don't tolerate the taints.
15:15I'm going to migrate them back."
15:17And you can ask, "Oh, why you don't use a no no no execute taint?"
15:20No execute in Kubernetes would kick everything out at the same time.
15:23So, as as like at moment you put the tanks, everything will be kicked out and that's bad.
15:27Production would go down.
15:28So, like this system my like slowly migrates the pods back, so production doesn't go down and we don't waste money.
15:34Um it is a like a very like self-healing system.
15:38You don't need to interfere with it.
15:40Uh it just runs.
15:41Uh yes, of course, there was bugs at the beginning.
15:43Uh nothing's perfect.
15:44But like once you calibrate it was was like very very well and like changed the way we do research cuz no one else needs to care about GPUs.
15:51They just launch stuff.
15:52If have GPUs, we have GPUs.
15:53If not, we don't have GPUs, go into the queue.
15:55And we fully utilize the cluster for trainings.
15:57Production runs somewhere else.
15:59And the GPUs are doing like useful work.
16:01And also like if you're doing like diffusion transformers, they're not huge like LLMs that need like multi-node like inference.
16:08Uh something that we learned like whatever GPU works.
16:11Uh the GPU can be hot, falling out of the bus.
16:13It can be exploding.
16:15Uh inference still going to run.
16:16It is very interesting.
16:17So, like you can have very very bad GPUs for inference.
16:19Uh and everyone's going to be happy.
16:22Uh we are hiring.
16:23Uh if you're interested in building sort of stuff, uh doing large-scale pre-training, uh build this sort of system for researchers, shoot me a message at gabriel plus hiring at cerebra.ai.
16:33There's also jobs listing and that's it.
16:36Thank you.
16:51[music]