An Open Source Post-Training Stack: Kubernetes + Ray + PyTorch + vLLM - Robert Nishihara, Anyscale

Key Takeaways

Hey everyone, thanks for coming. I'm Robert. Um, spoke earlier today and wanted to go a little deeper technically here talking about uh the work we're doing with Ray, how it integrates with Kubernetes, PyTorch, VLM, uh, and really the other components of the open source uh, compute stack for AI to support post-training and just some of the more um, modern computensive AI workloads we've been seeing recently. So um wanted to share a little bit about my background. I want to talk about uh how AI has been changing over the last year, what new workloads and new requirements we've been seeing and then a little bit about how the infrastructure software stack has been evolving uh to support these new workloads. Um and earlier today we announced that Ry is joining the PyTorch Foundation. So super excited about this. Um, Ry is used has you know been used along with PyTorch along with VLM with Deep Speed other projects in the PyTorch Foundation for um for as long as people have been using Ray. And so it's super exciting for us to just you know be part of that uh uh foundation. Um so my background I started doing machine learning research I think probably around 2012 2013 and uh was a grad student at Berkeley working on uh deep learning algorithms reinforcement learning uh algorithms optimization algorithms and actually had no background in distributed systems or infrastructure uh we were really focused on algorithms but because um you know even at the time to really experiment with AI to innovate on algorithms. You had to AI is very empirical. You had to run experiments. You had to try things out and see if they worked. And to see if they worked, you had to do things at some scale, at some meaningful scale, right? Of course, the scale at the time was not um the scale that it is today. But even uh you know, as people focused on algorithms and research, we spent most of our time working on infrastructure. spent most of our time building systems to uh scale compute across clusters to manage instances to run stuff on on GPUs. And because that was what was taking up all of our time, we thought there was an opportunity to build better tools for distributed computing. Okay, so that led us to start Ry um back at UC Berkeley and um Ray adoption has really taken off over the last year really with generative AI with today's reasoning models with agents with reinforcement learning. These are the things that have really led to um just a need for Ray and and a really explosion in Ray adoption. Um, and [snorts] a lot of this is really just in the last year. And, uh, there are a ton of use cases I'm very excited about. But this is, of course, regular, you know, um, established tech companies, companies like Uber and Pinterest and Discord and so forth. It's AI startups like XAI or Physical Intelligence or Runway or Thinking Machines, you know, it's traditional enterprises as well. So, it's really um, Ray is is is uh, you know, widely used across the industry. Now um so what I want to talk a little bit about now is some of the changes in AI workloads and how AI workloads have grown more complex and I'll talk about data processing and I'll talk about um post- trainining so people have been doing big data processing for a very long time right you have companies like data bricks companies like snowflake tools like spark and flink and Hadoop um but the nature of data processing is kind of changing right and You have um a world of tabular data that's moving more and more toward multimodal data, right? You have today you're people working with video data, images, audio, text, sensor, robotic sensor data, all sorts of things. Um and the way you work with tabular data and with multimodal data is different, right? With tabular data, you're often doing regular processing, running SQL queries and things like that. With video, you're probably not running a SQL query on your video, right? You're generally running inference. along with other processing on this type of multimodal data and so data processing is becoming very inferenceheavy and as data processing becomes very inferenceheavy um it becomes more GPU heavy right so a lot of data processing is shifting to this heterogeneous regime where it's it's a a CPU plus GPU workload and that's very different because when people were building the previous generation of data processing systems like Apache Spark you know Snowflake Fling, Hadoop, um they were really designed for a homogeneous cluster of CPUs, right? They were designed for more uh SQL like queries, more simple analytics and often more for tabular data than and and so the shift now if we were to summarize it in one sentence is from doing SQL on CPUs with tabular data to doing inference on GPUs with multimodal data. And that's a pretty big shift if you think about uh and a lot of this is driven by the fact that multimodal models are getting better, right? If you think about um you know, you work at a company and there are a lot of meetings and you might record some of those meetings, but most companies still don't record most meetings. And why is that? Well, no one's going to go back and watch the the recordings and so you don't get much value out of recording it. Um, and you know, you could use AI to kind of uh, you know, certainly even today using AI to really understand video and get value out of video is uh, it's kind of video is kind of unwieldy, right? It's like hard to do. Uh, it's hard to work with. But as these models get better, as we're able to get more and more value from all different types of data, um, well, we're going to get more value from the data and we're going to start storing way more of the data. And so I think we'll reach that point where kind of overnight everyone's recording all of the the meetings and things like that. Um you're starting to see a little bit of that, but we're still not getting anywhere near the value that we could be. And I think um we're going to start to get far more value out of multimodal data. Uh and then just as a consequence, the amount of multimodal data we store is going to explode and the amount of uh inference we do on that multimodal data to get value out of it is going to explode. So there's a huge shift in the nature of data processing that's that's happening here. Uh the second thing I want to talk about is reinforcement learning right. Um of course we've you know all aware of reinforcement learning from Alph Go and Atari and and uh you know physics simulators that this was happening in 2015 2016. uh the whole field of AI was was really focused on this um and a and reinforcement learning has really made a comeback recently because it's with deepseek and with its use to build uh really powerful reasoning models and to and to build agents and this is very interesting from a systems perspective. So this is a diagram from uh researchers at UC Berkeley talking about Sky RL which is an RL uh open source RL framework for post-training large language models and the it illustrates the anatomy from a systems perspective of running a reinforcement learning algorithm. So if you think about how this compares to uh pre-training or just you know training regular training of a model the pink box there that says GRPO PO that's basically running your training algorithm. If you were just doing regular training it would be that pink box but to do reinforcement learning you are not only doing training with the model you are also um using the model to generate data. So you're also doing inference and the process and then you're you know you're so you're taking the latest model um shipping the model weights over to the the data generation side of things of generating data and then shipping the the data back to to train on. And it's more complex than that, right? Because the data generation side actually has multiple components. There's inference where you're running the model um and then you often have some environment that you're simulating. If you're trying to build a coding model, then maybe you have a bunch of dockerized code bases and you know paired with GitHub issues and the model is saying things like um you know clone this repository uh CD into this directory like um apply this patch right run you know compile the project run this test suite and and all those commands are getting executed in the docker container and so you need to scale the um the containers that are running the environments you need to scale the inference you need to scale the training and each of these could be a bottleneck. You might need different software dependencies, different compute resources for each of these different um sort of pools of of of compute. You might need and and different ones can be bottlenecks at different time, right? As you as these models get better and better at reasoning for longer amounts of time throughout the course of of training, um the inference can go for longer and longer and so the amount of compute you need to shift toward inference might increase um as you progress. Right? So there's all these different uh complexities not you know to give you an example of more kinds of complexities that can arise um as you're reasoning for longer periods of time you may want to have multiple you may want to kind of split fork the um the reasoning and go down explore multiple different um reasoning paths and then recombine later. Right? these different environments um that you're executing may all um have some shared rate limited resource that they're querying like like an LOM API or something like that and so you need coordination among all these different um processes that you're scaling to kind of handle that rate limiting. So there are many different systems challenges that arise with reinforcement learning that just simply don't come up with regular training. And so um and and one of the main things here is increased heterogeneity just more moving parts each with their own requirements. And so um what I and and um Ray was really designed for heterogeneity both hardware heterogeneity application level heterogeneity and this is something um we talked about earlier today but as a consequence nearly every open source RL framework is built on top of Ray. So Ray is the open source uh framework we've been developing and Ry here is used in conjunction Ray is handling this process orchestration and coordination and and data movement uh in conjunction with a training engine like a Megatron if you're using large mixture of expert models or torch titan or or FSTP um or and then on the serving side we see a lot of SG lang a lot of VLM and other engines as well. Okay, so this is these are some of the uh technology stacks that people use for doing reinforcement learning for uh post-training LLMs. Um but this is by no means converged right this is uh you know you have new RL frameworks emerging all the time um new and and uh depending on the type of model you're using the architecture of the RL framework make uh that makes sense could be different and so I expect just more and more um iteration on this space before things converge um and so now I want to talk a little bit about the software stack for actually supporting these types of use cases Right? We have growing complexity with AI workloads. Um we talked about reinforcement learning. We talked about data processing with AI. Um and we have um all of this powered by GPUs under the hood. And so I want to, you know, repeat and I um a few of the different layers of the stack that we see, right? Because you have growing complexity on the applications, growing complexity on the hardware. And so you need a software stack in between in between that kind of connects the applications to the hardware and helps um manage that complexity, right? And that software stack has to solve these types of challenges, challenges around scale, um resource management, failure handling, reliability, and and so forth. And so we see a few common layers in this stack. We all know PyTorch, right? This is the dominant deep learning framework. Uh this layer is responsible for running the model efficiently on the GPU, right? And there's a rich ecosystem of um around PyTorch for handling model parallelism, different parallelism strategies for handling uh you know optimizations for transformers, right? both inference and training engines and uh but fundamentally this layer is about squeezing the most performance out of the GPU when you're running running the model right and then ray which we've been building sits at this distributed compute engine layer and this is about solving the distributed systems challenges of scaling right so that includes all this process orchestration coordination um right going back to this picture for a second um you have all these different pools of processes that you need to manage to autoscale to handle failures to move data between them um to coordinate them and synchronize them um and uh you know they might fail you might need to uh autoscale them these kinds of things. So there's a just this is all about solving those distributed systems challenges um within a single workload and then this all runs on top of a container orchestrator like Kubernetes or slurm which is responsible for provisioning the compute handling u more coarse grained considerations like user multi-tenency job multi-tenency and um and really together all of these um uh provide a stack for enabling complex AI applications. And so one thing I want to talk about is that we're actually seeing a tremendous amount of co-evolution of these different layers. I want to give just one example of how these layers by being aware of each other can work better together and provide uh better support for AI workloads and that is uh the the integration of Ray and VLLM. So if you're trying to deploy large mixture of expert models, right, um the largest deepseek models or things like that, it's uh a multi-node thing, right? You're not just using a single node. And there's quite a bit of complexity around getting this right and making it really performant. And doing it well requires not only um optimizing the code that runs on the GPUs and really getting the inference engine right, but it also means doing a great job at the process orchestration and coordination layer. Um and so Ray and VLM the open source communities have worked closely together to really enable uh fine grain control and over placement of the different VLM workers. For example, if you're doing uh pre-fill decode disagregation, right, with LLM inference, you are there's the prefill stage where you process the inputs, which can be more computebound. Uh there's the decode stage where you're generating tokens, which can be more um uh like GPU memory bandwidth bound. it can make sense to separate those onto different pools of compute or different GPUs in order to um basically so they don't interrupt each other and you get more predictable performance and better uh tail latencies and things like that. And but when you do this, you may have um you know, you're sharding the pre-fill stage, you're sharting the decode stage, you have a bunch of different workers. You may want to say like the matching I have u certain ranks on the prefill side that match to certain ranks on the decode side. Those correspond to different uh processes that need to be maybe colllocated or things like that. And so expressing this kind of thing is very much a a process orchestration coordination thing. And and the complexity is only growing as the models get bigger and the uh the types of parallelism strategies grow more complex, right? You have different parallelism strategies for the like expert layers. You have different parallelism strategies for the attention layers. And um and so really getting this right requires doing a good job at multiple layers of the software stack here. So not just uh something that's handled at at a single layer. Um and so these are the kinds of things that we've been able to do by have you know having a collaboration among the uh the two open source communities and this same could be said of Ray and Kubernetes along with as well as PyTorch. Um but I just wanted to share the the VLM example. Okay. So this is a little bit about how we see the software stack evolving for supporting these emerging AI use cases. All of this is motivated again by the fact that you just have so much growing complexity with AI workloads. Um and also because you have so much growing complexity on the hardware side. Um and uh wanted to share a little bit about what we're building at any scale right so we started any scale uh to commercialize ray and continue developing ray and just recently um you know today donated ray to the the pietorch foundation. Um so at any scale we're investing heavily in developing Ray. Um and then we're building our platform which is basically um a managed optimized version of Ray uh with a ton of the observability tooling around Ray. And so it's basically uh for companies using Ray. It's just a uh the best way to run Ray. Um the last thing I want to share is that similar to PyTorch conference, we have a Ray summit every year. U this is our fourth one happening um in a couple weeks here in San Francisco. So if you're in town this is uh it's a deeply technical audience and conference. So think about if you're interested in reinforcement learning uh folks from perplexity from cursor from uh thinking machines will be talking about how they scale reinforcement learning a lot of uh folks working on robotics folks working on other in other fields. um this is it's uh you know audience of basically like AI and AI infrastructure people and so um would encourage you to check it out. All right, I will uh stop there. Thank you all so much. Happy to take a few questions uh or or chat with folks afterwards. So thank you. [applause] >> Um yeah, why don't you shout the question? I can repeat it. >> Yeah. Could you share a little bit about the story for the donation of [snorts] >> Yeah, so a little bit about the story of donating Ray to the PyTorch Foundation. Absolutely. Um, so first of all, I would say most successful open source projects end up in a foundation at some point, right? If you look at PyTorch, Kubernetes, uh, so this is a very normal outcome. Um and for us it's really kind of a signal to the community that we are uh trying to grow the open source community that we like welcome uh contributors from um you know from every company and and that we're um really trying to make this a community project and that's a big part of it. It's really about growing the contributor community, growing the open source community. I'll also say that because um you know Ry is used so closely with PyTorch with VLM and these projects it's just a very natural home um for Ray and you know because of all these projects focused on uh really supporting like AI infrastructure. Yeah. So it was a uh very like uh good experience with um you know working with PyTorch Foundation and Linux Foundation and it's not like um you know it's not much change in terms of day-to-day processes. It's not like all of a sudden you adopt all these like heavyweight processes. It's it's actually um um been been quite good. Yeah. Yes. Yeah. >> Oh. >> Um, take on my take on Monarch. Um, so I think the, uh, I think Monarch has a lot of of really good ideas. I think, um, you know, and I think draws some inspiration from Ry, right, with, um, like the single controller model, the actor framework, stuff like that. I think um it's uh uh you know obviously like more natively uh pietorch right whereas ray is designed to be framework agnostic and you have people using ray with um of course with pietorch and but any different engine really um and uh but it's an interesting project [snorts] yeah >> can you repeat that >> here hold >> I got a microphone now. Here >> uh can you share how um why why is it called Ray or if it's like >> why is it called Ray? >> It's like a Metal Gear reference or >> ah um we went through a number of names actually and I think the first name was Hermes because we wanted it's like fast you know it had a I think it's a Greek god that's really fast or something like that. Um but and then we changed it to orchestra actually John Schulman suggested orchestra where because it's like a bunch of distributed system and and a lot of things working together. Um I think there was some naming conflict some other project in the lab had the similar name and so we changed it I think to photon next because it's fast and like you know light and positive connotation. Um, and then I think there was a naming conflict with some other uh project and then we changed it to Halo I believe and because um but um I don't know people were concerned about it the the game Halo just um and so we changed it to Ray and like nobody complained since then and so um yeah and I think you know we're thinking about like a ray of light like kind of or like a u just something positive and and kind of fast or lightweight. Yeah. Uh yeah. >> Are we planning to uh donate Cubray to the PyTorch Foundation? Um so so we were discussing this um and we you know wanted to do this announcement here at this conference, right? And so we were confident that we wanted to donate uh Ray to the to the PyTorch Foundation and we were unsure about the best path for like so there's the Ray project GitHub or right and then there's a bunch of repos under that. Ray is the main one. Um and but there are a couple others that are interesting and um we weren't sure like about the best path for the others and so uh we decided to just start with the with with the with Ray and then over the next you know 6 to 12 months we will um have a little space to think about the what's the best path for the others. So uh TBD. Yeah. >> Okay. I think we'll stop there. Uh thank you guys so much. [applause]