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

Hello everyone, uh, my name is Gabriel. I work at Krea, and I'll be talking about the infrastructure that allowed us to train K2, and also how we serve it. So what is K2? K2 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. They're quite, I don't know, saltless. They have no spice. And the whole idea was we want to give creatives tools 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. The model was trained from scratch, no base checkpoint, not anything. Everything done in-house. And also, as I said, built for exploration. Andright now, this is what you can get out of Krea 2: very different styles and, you know, like, pixel art, and, like, photo reel, and, like, some silly stuff.

And the whole idea of the model, as I said, let's explore this medium. Krea 2 is open source. Right now, you can go play with it. There are two checkpoints we also serve in production. There's a raw checkpoint, which is pre-trained, so people can post-train and do whatever they wish to do with it.

And there's also the post-trained version, which is the turbo one, which is very, very fast. You can get, like, an image in, like, I don't know, less than a second. And this is, as I said, like, this type of images you can get in less than a second.

On Hugging Face, GitHub, go play with it. Also, you can go in production on Krea.ai and go play with it. So let's talk about how we trained this model. As I said, it's going to be how we train and how we serve.

First, the model was trained from scratch on thousands of GPUs. We have a big cluster, one main cluster with a lot of GPUs, all InfiniBand connected. And you put those GPUs to work, and it trained. But, like, that's I wish it was that simple, but it's not.

So at the beginning, we did a bunch of, like, small ablations on, like, a small number of GPUs to see, like, how things go,right? So you want to test some hypotheses, and you do a small number of GPUs and let it train for a little bit.

Oh, this works, this doesn't work, let's scale. And as we were, like, training this model, the whole idea was to, like, kind of bridge the gap between, like, LLM research and diffusion transformers. So my AI researchers, they ported a lot of research from LLMs into DITs.

And the whole, like, architecture of the model was meant to be extremely, extremely simple. And so, like, it is very, very dumb, but, like, very effective. And so let's start talking about numbers. Incredibly, our maybe skill issue on our part, maybe our cluster, was very interesting as we were, like, scaling.

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. And as we, like, started scaling, like, getting, like, more and more and more GPUs, like 128, 256, 512, whatever number, and, like, you scale and scale, like, things start crashing more.

That's expected,right? Like, there's more surface area for things to break, and things are going to go wrong. And a lot of the times, things would go wrong in silent ways. I don't know, NICO timeout. Like, it just crashes, and, like, the metrics are all good.

And it is extremely annoying stuff. At the beginning, we were, like, paranoid. Swap node, change nodes, whatever, whatever, whatever. And we learned that, like, sometimes you just let it crash. It crashes, like, it runs for, like, an hour crash, runs for an hour crash, and then it runs again on the same set of machines, same code, same data for, like, 12 hours, 16 hours, 24 hours.

There is this paper from Meta that kind of gives you, like, a rough estimate of, like, how many failures you should expect. We kind of saw the same pattern, but, like, not the same numbers. Our runs would last way, way less than this.

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? You want to keep those GPUs fed, and if things are crashing, you are not doing progress and losing time, and the model is going to be delayed.

So for us, extremely important, and, like, at least from the infra side, was to get metrics. Metrics are everything. That's how I can support my researchers. That's how I have visibility in the system. And, like, if you're doing large-scale pre-training, I highly, highly recommend for you to invest heavily in metrics.

Don't go blind, because you're going to go crazy. So I'm going to share some of the metrics that, like, were important for us. And, like, they're quite silly, but, like, extremely effective. First one, like, GPU temperature. GPUs are very, very annoying.

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 the training is going to be unstable, and you have weird problems. So for us, it was, like, if there are any GPUs above, like, 78 degrees, you remove them.

Don't think about it. Don't try to fix. Don't try to be smart. Just remove the GPU. You're going to save time, and just ask your provider, like, this GPU is hot, please replace it.

And there are 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, there is GPU utilization, which is a lie. Don't trust this.

This is dumb. This tells you, oh, the GPU is doing work. And this is the amount of time the GPU is doing work, but, like, not good work, not how efficient the GPU is working. So, like, as you can see during our pre-training, the GPU is at 100%.

This is not true. We are not fully utilizing the GPU. This is 100% a lie. What we would use as a proxy was tensor core utilization. This is actually how much of your tensor cores you're using and, like, how effective they are being.

And it was also very interesting as we were doing pre-training, and then you go from pre-training to mid-training to post-training, you start, like, scaling the resolution of the images. So, 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 resolutions, because now you're doing more work on images. And, like, also very, very interesting was InfiniBand and NVLink metrics. These are by default not exported by the NVIDIA metrics, the DCGM stuff.

Some NVLink stuff, yes, but no InfiniBand. So if you don't have InfiniBand metrics, go get it. I'm telling youright 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 are doing something wrong.

This was probably, like, the most important stuff for us, because most of our failures were, like, related to, like, cross-node communication. So, like, for example, here you just have throughput. But, like, on our, like, Grafana dashboard, we have a bunch of stuff.

Like, from, like, weight, like, when a message is sent on the fabric, like, how much time the message is waiting, or, like, the, like, number of errors and different types of errors and, like, number of packets. And all of the things that, like, InfiniBand is exported, we collect.

We had to build custom stuff to get this. It was not hard. You can figure it out. It's very, very easy. Same thing for NVLink. NVIDIA exports some stuff about NVLink, but, like, for example, NVLink errors. NVIDIA doesn't export this.

So you can collect this. And, like, as I said, InfiniBand was extremely important. NVLink was a little bit less. In some cases, this helped us catch some problems in, like, especially, like, because NVLink is a single node,right? Like, the communication inside the node.

So sometimes a single node would have a weird failure where the GPU seemed to be fine, but, like, some weird error happens. And then you can look at the NVLink errors and, like, you see errors happening. And then replace that machine.

So go get these metrics. They're extremely important. And without this, we would not be able to train at all. Also, as I said, our trainings would crash constantly. And a hacky way to do it, to fix the problem, just checkpoints.

Use and abuse the filesystem that you have. At the beginning, we used CEF. CEF did not work well. It was very annoying. It broke. We did not trust the data. So I recommend, if you have the money, go with something paid, because you can trust your data.

You can see numbers. This is our Wekker cluster. We can do, like, 1.2 terabytes a second of reads, almost a terabyte of writes. And the filesystem would not choke on the training. So we could checkpoints every, like, 30 minutes, 20 minutes, produce, like, a terabyte of data in, like, less than 30 seconds.

So this would not delay trainings. That was, like, probably one of the, like, most important things that we did to, like, recoup the loss. Like, just checkpoint. Don't think about it. And, like, how we serve. This goes in connection on how the trainings are launched.

Because at the beginning, I don't want my researchers to think about GPUs. I just want them to launch stuff. And this goes into a queue. And if we have GPUs, we have GPUs. If we don't have GPUs, we don't have GPUs.

So queue. This is an open-source project. You can look it up. It does gang scheduling. Gang scheduling, extremely important for trainings in general. And 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, so it keeps on the queue in front of the queue. And then after this, we have the normal Kubernetes priority, if you're used to Kubernetes.

And the way the system works is, like, the training pods, they always have, like, high priority for everything. So, like, once they are admitted, they're going to schedule. If there is inference running on those machines, the inference gets kicked out.

And you would say, oh, this is bad. Production is going to go down. No, you can build on top of that to make production not go down, which is very, very cool. One of the problems with queue, which is annoying, you can automate that.

We have not. It's just that you specify the queues. The queues have, like, a month of resources, CPU, like, NVIDIA GPUs, memory, whatever. But this is manually, like, specified. And at least our cluster is quite fluid. Nodes phasing in and out of existence.

They go into maintenance, whatever. You lose a few nodes here and there. This number gets out of sync. And sometimes this would break gang scheduling. So FYI, this is a bit annoying. You're going to face this if you use queue.

But yeah, very good project. Kubernetes 135. We have not had the chance to play with it. It has gang scheduling out of the box. Something very similar to queue. So maybe you can use Kubernetes 135. And as I said, this is the system that we built.

They allowed us to train using the whole cluster. As I said, we have one big cluster that runs production and trainings. So I don't want my researchers to think about GPUs, and I don't want to make the trainings and research be delayed because production is running,right?

Production is lower priority. The site still needs to work. People still need to use the website. But, like, the GPUs, the value that we get off the GPUs doing training is more, like, higher than we get out of production.

So the whole system works by default, where there is this magical system that I'm going to explain in a bit that allows us to flip traffic between clusters magically. And not just clusters, like, external providers, GPU rentals, whatever.

And you can see, like, the green, like, the dark green is, like, inference running in cluster. And then someone launches the train, and then suddenly it starts flipping to the other cluster. And then training is done or whatever happens, it flips back.

So we stop wasting money. And this is seamless. No one needs to think about it. The whole system, like, handles itself. And you get this very nice pattern of, like, I'm going to use all the GPUs in my cluster for trainings.

Production is going to run somewhere else. I don't need to think about it. My users on production, they're not going to feel anything. Researchers are going to be happy, and we can get values out of the GPUs. So how does this work?

There is this very nice project called Virtual Qubelet, also open source. You can build on top of it. It is a very nice code base. And this works by creating a fake machine in Kubernetes. Kubernetes has nodes. This creates a fake machine that is, like, up to you to control how it works.

So Kubernetes does normal scheduling as you would expect. Things would go into these nodes. For example, here, all the GPUs in the cluster are in use,right? So this pod goes into the Virtual Qubelet node. And in there, you can do whatever.

This is the system that we built. There is, like, you receive the pod spec, and then you find a provider. You can, like, this is up to you. Let's say you have a deal with some provider that gives you nice prices.

You integrate into here. We built, like, some nice interfaces to, like, not leak things. So, like, we just implement a provider, and there is an algo that decides which one it goes. You translate the spec of the pod into the provider stuff, and you deploy.

And then you have something that reconciles between both sides. And it was extremely, extremely nice. If you guys know about Kubernetes, Kubernetes has, like, the horizontal pod autoscaler, which, like, scales the number of replicas. Number of replicas. Could you stop being annoying?

Thank you, sir. I appreciate it. There we go. Let's go back. Go back. Back, back. There we go.

Like, Kubernetes has the HPA, and the HPA scales the pods. And so if something fails, it is very interesting. You don't need to handle the fail. The only thing you need to handle is, like, oh, something failed. You mark the pod as failed.

Kubernetes is going to detect that something has failed and create a new one. You don't need to try to save the world. Let Kubernetes handle for you, which is an extremely nice way of handling stuff. If something breaks on your side, something breaks on the other side, just mark as failed.

Let Kubernetes handle. Create a new one, and things keep working. Very, very nice way to handle stuff. And also very interesting. Let's say you have GPUs on your cluster available,right? You don't want to waste money. This would be very, very bad.

So the system works with, like, using tanks. Kubernetes tanks. They allow and disallow things to run. Pods have tolerations for the tanks. And when we have GPUs in the cluster, if you look back, there is the tank system on the bottom.

This tank system, it is what would, like, by itself decide what, like, if you have GPUs or not GPUs in the cluster. And this adds a tank into the node when we have a lot of GPUs. So, like, a lot of GPUs in the cluster.

We tank the node. Nothing can schedule on it. So we stop wasting GPUs. Like, the pods, they go into the GPUs in the cluster. We don't waste money. And then imagine someone launches a train,right? This train is going to take all the GPUs in the cluster.

It's going to hog all of the GPUs. No GPUs in the cluster. The system detects this, removes the tanks. New pods schedule there. Very, very nice. You also don't think about it. And for us, for example, we use just some Prometheus metrics.

That's how we do it. Very simple, but it works very, very, very well. You let the system run by itself. Someone's going to launch stuff. You're going to, like, the train SLA 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 tanks. Pods schedule there. Very nice. Someone finished the train. Now we have pods running on the other side. You're wasting money.

How do we fix this? Same thing. You run something else that detects the system and removes things back. In this case, a descheduler. Once the tanks added back, so, like, GPUs are available, we add the tanks. The descheduler sees, oh, these pods, they don't tolerate the tanks.

I'm going to migrate them back. And you can ask, oh, why you don't use a no-execute tanks? No-execute in Kubernetes would kick everything out at the same time. So, like, the moment you put the tanks, everything would be kicked out, and that's bad.

Production would go down. So, like, this system, like, slowly migrates the pods back, so production doesn't go down, and we don't waste money. It is, like, a very, like, self-healing system. You don't need to interfere with it. It just runs.

Yes, of course, there were bugs at the beginning. Nothing's perfect. But, like, once you calibrate it, it was, like, very, very well and, like, changed the way we do research. Because no one else needs to care about GPUs.

They just launch stuff. If we have GPUs, we have GPUs. If not, we don't have GPUs, we go into the queue. And we fully utilize the cluster for trainings. Production runs somewhere else. And the GPUs are doing, like, useful work.

And also, like, if you're doing, like, diffusion transformers, they're not huge, like, LLMs. They need, like, multi-node, like, inference. Something that we learn, like, whatever GPU works. The GPU can be hot, falling out of the bus. It can be exploding.

Inference is still going to run. It is very interesting. So, like, you can have very, very bad GPUs for inference, and everyone is going to be happy. We are hiring. If you're interested in building sort of stuff, doing large-scale op training, build this sort of system for researchers, shoot me a message at gabriel+hiring@krea.ai.

There's also jobs listing, and that's it. Thank you.