Thank you guys. I'm going to talk a little bit about AI infrastructure, especially for reinforcement learning, and also touch a little bit on some of the multimodal data use cases that we see a lot. But just to some of the stuff I want to cover. We've been building Ray, an open-source project for scaling compute-intensive AI workloads. I've been working on that for a while. And want to talk about two things. Some of the most popular use cases that we see today: reinforcement learning as well as AI data pipelines. And I want to talk about how Ray fits into the overall AI infrastructure tech stack.
We actually started this open-source project back at UC Berkeley as grad students. We were doing research in reinforcement learning. If you remember the previous wave of excitement around AlphaGo and Atari, Mu Zero, these types of things. We were trying to design new reinforcement learning algorithms and optimization algorithms. And despite wanting to spend our time on the algorithms, we found ourselves spending all of our time managing clusters. So we thought there had to be a better way to build and program scalable Python applications and run them across clusters. So we started Ray as an open-source project to really achieve that. At the time, 2019 or so and 2020, when we were getting the company started, there was not that much demand for distributed computing for AI. It's really over the last couple of years that the adoption of Ray has exploded with the rise of multimodal data, of reinforcement learning, of reasoning models, custom models, every enterprise building custom models to take advantage of their proprietary data. And so we've seen really Ray bloom over the last year or two.
These are just an example of the types of companies that we work with, but these include companies building foundation models like xAI, Cursor, companies like Coinbase, Canva, Instacart, as well as traditional enterprise and financial institutions. To illustrate one of our design goals, with Ray has always been to support a huge range of scalable workloads. I want to give three examples of very different types of use cases to illustrate the breadth. Tencent, they build WeChat, right? Huge company. They use Ray to deploy all of their models powering WeChat. These include generative models, LLMs, traditional non-generative models, classifiers, content moderation, image models, audio models, text models, all sorts of things. That runs at the scale of roughly a million CPU cores and many tens of thousands of GPUs. That's on the inference side. On the data processing side, Amazon migrated an exabyte scale data processing workload on Ray. When they started using Ray, they actually saved over a hundred million dollars a year and they have written publicly about this. And then Cursor, as you all know, they do tons of reinforcement learning. They run all of their data infrastructure and reinforcement learning using Ray. So we're talking about reinforcement learning, data processing, inference, this is the kind of breadth of use cases that we are trying to support with Ray.
But I want to go a little bit deeper on reinforcement learning because this is one of the fastest-growing use cases, or just one of the fastest-growing workloads in the world today. It's the path that has really enabled tons of recent progress in AI. Why is reinforcement learning exploding? The rise of custom models, of agents, and reasoning. Reinforcement learning is quite different from regular training. You can think of training as a subroutine in reinforcement learning, but inference is also a subroutine in reinforcement learning. In doing reinforcement learning, you're simultaneously doing training and inference, running simulations, coordinating between all of these things. I'm going to walk through in a little more detail all these different components and why this is challenging from a systems perspective.
One more thing I want to share. Just a couple days ago, Microsoft AI released seven new models. They released MAI thinking one, which is a trillion-parameter model. Seems very impressive. They used Ray in many different parts of the frontier model creation. They used Ray for creating the RL environments. They used Ray for the pre-training. They used Ray for the post-training reinforcement learning. They have a 100-page technical report detailing all of this. There's so much complexity around actually implementing reinforcement learning algorithms. This is a depiction from their paper of some of the different components of the overall RL system.
I want to talk about two things. One is how you actually can implement a reinforcement learning algorithm in Python using Ray. This is intended to illustrate Ray's actor API, just to give you a sense of what Ray does. And then second, I want to walk through why it's actually hard from a systems perspective to implement these algorithms and how you have to change things as you scale. Let's say you want to implement an RL algorithm in Python. The first thing you need to do is have a trainer and a generator. There's two parts to RL. One is updating model weights with the trainer, and the other is generating new data to train on. You actually have to keep them in sync. As you're updating the model, you need to send the latest model weights to the generator so that you're generating data with the latest model. You might have a method for getting the weights and setting the weights. What Ray lets you do, you can add this annotation on top of the two classes and specify what kind of hardware they need. Maybe they need an H100 or a B200. Once you've done that, when you actually instantiate an instance of these classes, instead of just creating a local copy of the object in your Python process, it will spin up a new process in your cluster. You can specify. Here I'm creating one trainer process and one generator process, and they will interact with each other, talk to each other to actually run the algorithm. Of course, in reality, I'm creating thousands of these. I'm specifying more details using Ray to say, 'Hey, these have to be co-located on the same rack. These ones have to be co-located on a different rack.' Those are the kinds of things that Ray lets you do. The actual overall for loop in the application will alternate between doing a few things. It'll run some training by invoking a method on the training actor. It'll generate data by invoking a method on the generator actor. It will also synchronize the weights in between. This is vastly oversimplified and these systems are very complex and involve many more lines of code, but this is the backbone of what's going on. The way you coordinate between these different actors is by passing objects around. You can invoke the get weights method on one, pass the reference into the set weights method on the other, and that will transfer the weights in between. There are many more things you actually need to do. The way I wrote it here, this weight transfer happens over the network, copies into CPU memory, and over to CPU memory. You actually want this to go directly from GPU to GPU, point-to-point. You can use annotations like this to specify that the object should be transferred directly from GPU to GPU, using Ray's native RDMA support. This is just trying to illustrate the low-level Ray actor API that is used for spinning up, managing processes, coordinating processes, things like that.
Why is this hard? Reinforcement learning is far more complex to scale than regular model training. I want to walk through some of the different components visually. To implement a reinforcement learning algorithm, here I've got a bunch of different racks. You can imagine these are racks of GB200s. Each rack has 18 nodes, each node has four GPUs. I'm running a bunch of training processes on each of these GPUs. Because GPU health is a complex thing, there are often some number of GPUs that are broken. I'm going to run with spares that are just hot spares, not doing anything, but I may need to use them when something goes wrong. I may even have spare racks. If something goes wrong, it may be easier to swap out an entire rack. That's the training part. We also have inference, which has a bunch of GPUs and workers running on those GPUs, often running vLLM, SG Lang, things like that. Then there are environments running on CPU machines. You can think of these as dockerized coding environments which you can take actions in, like check out a certain branch, clone the repository, apply some patch to the code, compile the code, run some tests, that kind of thing. We also have a bunch of processes that are basically running the agent loop, coordinating between inference and the environments to actually generate the data. These are all the different pieces that you need. But how do you actually set this all up? You've got your driver script. This application driver is your Python program that you're writing. It is a single imperative Python program. It will be in charge of setting up all the processes. It's going to place all of these processes, specify what type of hardware they need, topology requirements, co-location requirements, what goes where. That's what you do with the Ray actor API.
But there's a lot more coordination after you set things up. I mentioned you have to sync weights from training to inference. That means sending the model weights from the training part to the inference components. That is more complex than it looks because first of all, these large models, the Microsoft models of trillion parameters, they're not fitting on a single GPU. They might be sharded over all the GPUs in the rack or even across multiple racks. The way they're sharded in training is different from how they're sharded in inference. So as you shift the weights from training to inference, you have to reshard them in a different way. You might even have multiple different inference deployments, each with different shardings. The reason is to get the best inference performance, you need to tune the inference deployment based on the input lengths and the output lengths of the sequences that you're expecting. For example, if you're doing prefill decode disaggregation in inference, you may tune the number of prefill workers and the number of decode workers. It's one kind of parameter you would tune. But when you're training these models, you're training on a mixture of different types of problems. You might have some coding environments, some math environments, some other type of thing. These different problems have different input and output length characteristics. So you may want different inference deployments tuned to different ones. You're managing multiple inference deployments, and that complicates your weight sync. There's also tight coordination between the agent loop, the inference, and the environments. You're going back and forth to generate the data. You need to use inference to determine an action to take, then use the environment to execute that action, then go back to inference to determine the next action, and so on. After you've generated data, you have to transfer that data, the rollouts, back to the training. That simple thing is more complex than it may look because the data you're generating, these rollouts, have variable lengths. They all have different lengths. When you're feeding them into training, you're batching them together because you need to do batching to get good efficiency. But when you're batching together a bunch of things that have different lengths, you could end up with a lot of padding that's just wasting your compute. You actually want to pack the rollout data into batches in a way that minimizes the amount of padding, because that can double your efficiency right there. In order to do that, you can imagine a lot of metadata being shipped around and used to shuffle the data, send the appropriate data to the appropriate rollout worker or appropriate training worker. That's the type of additional coordination complexity you have between all of these processes. There's a lot of moving parts you can see here. If you compare this to traditional or regular training, where you basically have your Python script, you replicate it across all of the GPUs, they're all running in lockstep executing the same code, aggregating gradients, and then executing the next gradient step. They're basically all doing the same thing. Here you have a lot of heterogeneity. Each different component has its own placement and topology requirements. Each different component might have a different failure handling strategy.
Just to go one level deeper, I want to talk about what happens in the event of a failure. There are many different components here. Things can fail in any of these different components, but I want to zoom in on training for a second to talk about what happens in the case of a training failure. The right strategy for handling failures varies depending on the scale and how quickly you need to recover from failures. If you're a company like Cursor or Microsoft AI running on, in their case they ran on 8,000 GB200s, you want to recover from failures as quickly as possible. At the next level scale, say 30,000 GPUs, you need to recover maybe in 3 minutes instead of 5 minutes. You just need to keep pushing that. Let's say something goes wrong. One of these processes crashes. The first step is failure detection. Before you even get around to handling the failure, you need to detect it. Failure detection itself is subtle. Failures can happen in all different places. It could be a hardware level failure, GPU, InfiniBand. It could be a rack level issue. It could be an application level issue. There are many different places failures can happen. The first thing you do: you can get an actual explicit error message, maybe some kind of NCCL timeout, or NCCL error. But you may not. Failures can be silent. You can get silent numerical corruption. You can get silent stragglers. The first thing you have to do is localize the failure within the rack. Just because you got a NCCL timeout in one place doesn't mean that place was the problematic thing. The actual root cause could be somewhere else in the rack. There's a combination of passive background health checks you can run on all of the nodes, but there are also a lot of application level logic that people run to detect failures. For example, when you're doing a forward backward pass in your model, there are certain times you may recompute certain activations in order to not have to keep them all around in memory. When you recompute activations, if you keep around a hash of all of your activations, you can compare the hashes, and that can help you detect if there might be numerical corruption. That's the kind of thing you might do at the application level. Let's say you've detected the error, localized it. You need to tear down processes that are affected. The simplest starting point is tear down everything and recreate everything. But that gets slow as you scale. So you tear down some processes. You have to keep a list of bad nodes. Companies will typically build a database of GPU health, tracking every GPU they have, every health check that ever failed on every GPU, every job that ever ran on every GPU, and what happened to that job. Over time you try to identify the healthiest subset of your fleet, and that's the subset you're going to run your large pre-training runs on. Within the job you're also tracking health, and you're saying, 'Okay, this node is bad. I'm going to schedule around it.' You have to take all of this application level knowledge of GPU health into account when you're scheduling and placing processes. The next thing you need to do is recreate relevant processes. You may do that on one of your spare nodes. That takes into account knowledge of GPU health and topology requirements. Recreate the processes. You have to reconstruct the state on all of those processes. That's the training engine state, the process state. You have to fetch the weights. You may fetch the weights from a checkpoint in blob storage, but you may also fetch the weights from another GPU in the cluster, from live replicas. That can avoid the need to replay certain gradient computations, but it's harder to do. You have to reconstruct NCCL groups, the communication groups. A lot of times, NCCL doesn't handle failures very well. So you may have to reconstruct the NCCL groups across all the nodes, but you don't want to tear down all of the training processes. What people will do is run NCCL in a separate subprocess on each of the nodes or each of the GPUs and then tear down just those subprocesses and reconstruct them. Those are the types of things you have to do to handle failures. The more you scale, the more you need to recover quickly because failures become more expensive, and the more you need to do these types of complex things to actually handle it efficiently.
I will stop there and leave you with this picture. These are the sum of the anatomy of a complex reinforcement learning system. This is only getting more complex and I will end there. Thank you so much.
