Accelerating Mixture of Experts Training With Rail Optimized InfiniBand Networking in Crusoe Cloud

Thank you for coming over to our session today. Lots of really interesting stuff is happening on the AI world these days,right, with all the recent model developments and the GPU developments. It is really cool to see all the use cases.

However, here I want to talk now a little bit about the infrastructure and the way how we can support the newest models of the GPUs and the newest, um, the newest models, the machine learning models, and how we can help that everything is working smoothly and fast and, and productive.

So my name is Yevgen Bakulenko, I'm a product manager at Crusoe. My main responsibility is, is the infrastructure, and specifically GPU networking infrastructure, and we are always looking for a way how we can increase the performance of that network, because as we will see later in the presentation, it is really important to do that.

Now, a little bit about the Crusoe. Crusoe is an AI cloud platform, which has one, I think, very important mission for all of us: it's to align the future of computing with the future of climate. There is a really strong demandright now for the computing power.

The GPUs are really energy hungry. There is a lot of investments being done in the data center area, and of course that puts an additional pressure on the grid and on the energy sources. What we are trying to do here at Crusoe, we are trying to utilize these trended energy sources, wasted energy sources, and renewables to power our data centers.

We want really to be able to make sure that every time when you train your model, every time when you're using GPU for inference, you're not causing any negative impact on the climate. Now, whenever we are building the cloud,right, the AI cloud, we are building it based on three important pillars.

First of all, there is a high performance pillar, as the customers are buying our services and procuring, you know, the GPU times and training their models, we have to ensure that all the infrastructure is optimized for this training.

Every time when it's not optimized, every time when there is a delay or a glitch or any sort of outage or simply not that great performance, it causes the direct impact on the customer's bottom line, it causes a direct impact on the time to train, and kind of raises the cost to train the model.

Now, the second one, which I think is very important for everybody around here, is the easy-to-use. We want really to separate ourselves from the general-purpose clouds. We do know all of them. The hyperscalers are building the great infrastructure and are trying to support each and every use case the customers might have for the cloud computing.

However, in our case, we really want to focus on the experience of the AI engineers. So we want to make sure that we are providing a simpler user interface that allows developers to spin up the compute resources, to deploy the models, to train them, to use them for inference, and, and so on.

All the underlying complexity of the infrastructure is being hidden by us, and I believe that's our job to make sure that that is, that stays the case. And now, as I mentioned, we are, as I mentioned before, we are climate-aligned, which means we as a company really aiming to power 100% of our data centers with the renewable wasted energy sources, with some, some, some form of stranded energy sources to ensure that we are, we are being net zero emission, net new zero emissions from the carbon perspective.

We have a big story around that. Feel free to check it on our website or come over to our booth on the, on the show floor, and the team will help you to talk about that. Now, where are we presentright now?

We have a number of the data centers located across the US. As you see, three of them in the continental United States, and they are generally located close to the energy sources I was mentioning before. So we have the one in Texas, we have the one in the northern central part of the country, and on the east.

We are also buildingright now one big data center in Iceland that will be powered by the geothermal energy. I mean, again, a way, amazing way to use the constrained energy sources or the renewables to power the data center.

We are trying to follow that model, hence we are placing our data center strategically. The placement of the data center in Iceland, though, will be also very important for our EMEA customers given the latency and the, and the general connectivity to the Europe.

That is something I think might be helpful for them as well. Now, what is our platform,right? I say Crusoe Cloud, but generally whenever we are talking about any cloud, we are talking about three general types of the products.

First and foremost, we have the compute. We are offering the VMs with, with GPUs attached to them. So every time when customer wants to spin up, when customer wants to get access to the GPUs, they're able to get it through the VM.

They can get a bunch of VM connected together and use them as a one single training cluster. We also offer CPU instances for any potential data preprocessing or any general purpose compute tasks you might have for the data preparation, for the offload, whatever you have.

From the storage perspective, we are offering ephemeral and persistent disks on the node. So those are delivered from the NVMe on the local server where your VMs are being placed. We also have the persistent block storage solution available for our customers, and we are working on providing and delivering the managed file system, the network file systems for the customers.

On the networking side, of course, more traditional, more typical VPC networking, that's the network, sometimes we call it front-end network, that is used to deliver the customer traffic from the internet or from the customer environment, wherever the customer might have the data sources, to deliver that towards the VM.

So that's your kind of main connectivity path to the outside world. Now, we do offer a number of the additional services on that. That's not simply connectivity. We also have the firewalls. We will be offering the load balancers soon.

But generally, we are trying to follow more traditional paths for the VPC networking and, and the requirements the customers usually have there. Now, what is more interesting and what we will be talking a little bit later today in greater details is our rail-optimized InfiniBand cluster networking.

So for those of you who don't know, typically cust, typically providers, the GPU providers are separating their network. They have the front-end network, which is used for general purpose traffic. But then all the communication between the GPUs is happening on the standalone separate network that is really high performance, low latency, and how bandwidth, and the whole topology is optimized for the GPU to GPU communication.

Now, last but not least, the user experience. As I mentioned before, we are our main customers, our main persona, the people who are using Crusoe Cloud are the AI developers and machine learning engineers. So we want to make sure they have what they need in order, you know, to be successful and not to think too much about infrastructure.

We also offer CLI, we have APIs, we have GUIs, so everything can be automated, everything can be, can be consumed and configured in the way you like it more.

We do have a lot of customers already, and it was very fun for me to see on the floor that some of them are there and some of them are talking about their solutions. Probably this is the first time in my life, whenever I'm attending a conference and standing at the booth, I don't have to compete with all the people around us.

So we do see all the companies that are presenting their solutionsright now as our partner, partners. We do partner with a bunch of them already. We have the Together AI here, we have the Boson AI, and all of them are using our infrastructure for different purposes.

So Together AI, for example, they're really into using Crusoe infrastructure for the ML training, for the fine-tuning their models, and some, sometimes for inference. The C dot AI is, they're, they're using our compute infrastructure to train the new foundational models.

This is really great. I mean, if you are the customer of Together AI, for example, or Codium or whatnot, it is likely that you have been somehow exposed to the Crusoe infrastructure. Now, the distributed training has a very specific set of problems or issues,right?

There is a compute part of it when the computation is being done on the GPUs, but since we are talking about a distributed training, which means there are a lot of GPUs, at certain stages, whenever there is a, whenever there is a training step being completed, all the GPUs have to exchange the information, have to exchange the data that they calculated on their own.

This is typically done through the all-reduce or, or all, all, all get

through the all-reduce process and, and the protocols, and it contains a forward path, the backward path, but then the networking part takes, without any optimization, about 25, 30% of the network of the time of the training time. Now, this is the time when, when your GPUs are staying idle, they're not being able to compute anything because they have to wait for all the information to be gathered together.

This is kind of a bad thing for everybody,right? This is bad thing for the customers because they still pay for that infrastructure, they still have to wait, it delays the, the model training, but it also bad for us because we have the infrastructure that is not being performant enough.

There are a couple of tricks we can do. First of all, the computation and communication overlap allows you to start the network exchange or the data, data exchange when the computation is still ongoing. But even with that, when we were working with the customers, we saw just the reduction of about 10%, so about 25% of the training time was still spent on the networking.

Me, as a product manager on the infrastructure side, are constantly being asked like how can we reduce that, how can we use the network as much as possible and reduce that gap. So we, we have been looking into that and we were trying to figure out what would be theright cluster networking topology.

How can we make sure that our data fabric that is used for connecting the GPUs is being fully optimized and is being, is able to provide the bandwidth needed and the latency needed. The standard fed tree, those of you who have been working in the data center infrastructure before, that is something that we were traditionally doing for years.

That's a great way to build a scalable, maybe non-blocking fabric,right? But there are a bunch of issues with that. First of all, if we will be connecting our servers that are shown below to a single leaf, that introduces the single choke point as well as the single fault domain,right?

If we are losing the leaf, we are losing all the GPUs that are connected to that leaf. Now, what else we were thinking about is like, look, we have that switch. We, we have that switch. Sorry, what is it?

The time?

Yes.

Okay. So we have that switch that can be used for the backend traffic propagation. And why don't we use that switch for, from the bandwidth perspective and kind of, you know, have an additional path? Let me just use the simple two-node example to explain the topology and to explain how we are using it.

So first of all, whenever we have the GPUs that want to communicate within one server, they can use their embedded NVLink NVSwitch, and that provides a good communication. They don't have to go to the outside fabric anywhere and, and whatnot.

Now, whenever we have the data communication between the GPUs on the different nodes, if they collect, if they are connected to the single leaf, that's something we called one single rail. That means that the traffic communication will be passing through the, through the one single leaf, just one hope away, and you will get the, to the destination.

Now, what is interesting here is when we want to talk to the different rails,right? We have to go all the way to the spine, and that introduces the additional hope. Besides the bandwidth saturation problems, that might lead to the additional latency, which will be really important for all, for your all-reduce, all-reduce operations.

But luckily for us, NVIDIA with the recent version of Nikko introduced the feature called PXN, which allows you to use the internal NVSwitch inside the host to communicate across the rail. So whenever we want to have the GPU zero to communicate with the GPU eight on the another host, we can use an internal switch to do the traffic hop between the GPUs and then send it to the leaf where it is connected to.

So it still allows us to use one single hop and have access across the different rails of the GPUs. Now, we did some nickel test results, and we saw quite a significant improvement, 50% for the small messages and 50% of the large messages.

Now, those numbers here are of course for the smaller, for the smaller messages are about latency. For larger ones, we care more about bandwidth because latency tends to stand, to stay roughly the same. Those numbers are great,right? Everyone, everybody would love them, but not the customers.

And it does make sense because those numbers are synthetic and more are showing you the workload that is applied to your network. What customers care about is the time to train the particular model. So we use the sparse mixture of expert as an example.

And I mean, I'm not going to dive into the details how, how it works, but essentially the sparse network, the sparse mixture of experts shows you, gives you a different layers of the experts and routes the traffic between them.

Whenever you are deploying that on the really large GPU cluster, that makes, that creates a ton of traffic. Like all the GPUs have to send the traffic to each other, they have to extend the information. The workload on the network is pretty significant.

So we use the Mixtral model, the open source sparse mixture of experts, which is contained of eight feed forward blocks, eight, seven billions of parameters. And we use the fine-tuning to use this model to fine-tune it on 240 H100 GPUs.

And we did a quite significant, we saw a quite significant improvement when we had the PXN enabled and without it. The 14% of improvement is something that can be directly connected to the time to train the model, that can be directly connected to cost of train the model.

And that is something that everybody really, really got excited. I definitely got excited and, and our customers as well, because that shows them some real value numbers they can get with a model. Now, that was it from my side.

Sorry for going through that. It's so fast. It's a very, you know, large topic. It's, it's hard to talk about that, but I'm happy to answer all the additional questions, anything you guys might have.

Thank you.