transcribe

How Prime Intellect Builds Scalable Infrastructure for Agentic RL | Ray Summit 2025

Anyscale · 30m · transcribed 15d ago
More from Anyscale Business
𝕏 Share ▶ YouTube 📥 PDF 🤖 .md

Transcript

0:05 All right. yeah, thank you for having us. I'm Johannes, co-founder and CTO at Prime Intellect. And, yeah, I'm here with Will Brown, who is our research lead on reinforcement learning. And, yeah, I'm going to share a bit more about, scaling open distributed infrastructure at Prime Intellect and how we think about environments for Agentica is what Will is going to touch on afterwards. And yeah, I'll keep it very short here on the high level overview of what we are doing at prime to then get into the actual framework we've been building and our research and so on. but we like think of ourselves at prime intellect as an open source frontier lab in a sense that we build our own models.

0:43 but the differentiating factor is that we also want to make the whole infrastructure to build frontier models open source and not necessarily only like the model checkpoints and yeah for people to actually train their own models for their own application end to end from the pre-training stage to the post- training stage and we start from the beginning from like the compute layer in a sense and yeah want to offer there the end toend pipeline as an open source infrastructure and yeah I brought here the the slide from Jensen at Nvidia on like the the different scaling laws in a sense. I have I don't have to go into detail into this one because I think most of you guys got the memo at this point and in terms of like the different paradigms we are seeing right now. but I brought it to basically talk a little bit about the yeah different yeah areas we've also went through with prime intellect over the last two years of our existence. So we started a lot with the like pre-training paradigm and especially this topic of globally distributed training to make use of like all the different idle compute in the world and you have been doing fundamental work on yeah how to actually scale up this globally distributed low communication training on a technique called DLCO distributed low communication training by Google deep mind and yeah scaled this to the largest model sizes last year and I also ever more interested to scale that further and are doing yeah more research on that but there would be a talk for another other session. so I'm going to keep it are we going to keep it focused today on this whole topic of test time scaling and our recent efforts in this direction.

2:15 with the like release of 01 and this whole like reasoning paradigm kicking off and especially with Deepseeker 1 coming out which yeah pretty much yeah catalyzed the whole open source space and open source reinforcement learning space over the last year. we we going to focus on this one and yeah to to give a bit of an overview of our reinforcement learning stack we start at the compute layer as I've mentioned to aggregate like all the different data centers all the different providers out there and then we have our open source reinforcement learning libraries with primer which I'm going to touch on a little bit for our large scale async trainer as well as verifiers repository by real brown and how it basically ties into our whole ecosystem we've built around reinforcement learning with our environment hub which utilizing sandboxes, eval as well as a reinforcement fine-tuning offering to basically enable anyone to train and deploy their own open frontier agentic models.

3:11 And yeah, to give a bit of an history of what we've been doing with reinforcement learning over the last year, one of the releases we had earlier this year was intellect 2 which is a model similarly trained as like Deepsea 1 with reinforcement learning with verifiable rewards. but the unique thing about intellect 2 was that we trained it in a globally distributed fashion similar as our work in the pre-training paradigm. We applied similar concept to the whole post-training and all scaling paradigm and in generally like reinforcement learning is a process that is a lot more async a lot more like disagregated into training and inference of your model and yeah through a technique called async reinforcement learning where we don't necessarily use the latest model rates to actually generate our rollouts yeah this is possible to do actually do in a distributed setting.

4:00 So how Intellect 2 worked was that we had a single training cluster where we did our actual update steps and then we had yeah a bunch of different inference rollout workers. yeah Louis across the world that were doing actually the rollouts and yeah we were able to completely hide the communication of this broadcast from the new model weights to the inference workers by running a few steps async in this case two to basically yeah hide the internet communication of the broadcast to all of those different inference workers and yeah had a great infrastructure release with that as well as a technical report where we basically showed this whole topic of yeah async rein enforcement learning is actually scaling and can be stable at large scale and yeah compared it to like a synchronized baseline in a sense and yeah I've been building out this whole framework called primer all since then which yeah as I said fully disagregates the training and the inference of your model the training is just an FSTP2 based training framework yeah pretty much inspired by torch titan we have an orchestrator which is like yeah orchestrating the whole core data and scheduling logic And then we have like the inference workers which are just hosting an open API server in a sense with the additional feature of doing on the fly weight updates and yeah since then we've been adopting the best practices that have been coming out out of open research for async reinforcement learning. One of the things a lot of people have been adopting is a concept called aerial where you do continuous batching and on the-fly updates of your weights. So you don't necessarily wait until all the rollouts are done with this particular checkpoint. You can on the fly update your weights when the the training is actually done, the broadcast is done.

5:45 yeah, we use an FTP trainer and a VLM inference. as I've mentioned, we've been yeah doing a lot of work on actually trying to get a stable reinforcement learning algorithm and yeah, I've been validating Cisco which I believe is from Miniax and their V1 release. and yeah we've been scaling this to the 100b scale with async reinforcement learning at got stable running stable reinforcement learning running at that scale and yeah it has a multi-backend broadcast so as I've mentioned we've been working on this whole system of globally distributed sailing so one of the backends is called shortcast so you can actually broadcast your weights over the internet another one is nickel if you want to be do it in a in a single data center and yeah the whole trainer is based pretty much on torch titan are heavily inspired by it and its expert parallelism and yeah we've been using primal quite heavily so here you guys see a preview of what's coming next in a sense at prime where we've been building our own frontier open a magentaic model it's a 100 billion plus scale mixture of expert 100k sequence length and yeah over 512 H200 that we've been training it on we've done our whole own post train so the whole SFT stage with stage one and stage to with single turn examples for mass code science and then as well agentic SFT as well as reinforcement learning via our primal frameworks which also supports SFT.

7:12 we've been trying to maximize our training and inference throughput with AIAL and continuous batching. And yeah, another very unique thing about Entle Lake 3 is that we've been training and evaluating that model via yeah all those different environments that have been contributed to our environment hub. And yeah, for the environment hub, I'm going to hand it over to Will here too to tell you more about how that is all working and why in general agental needs those reinforcement learning environments.

7:46 >> thanks Johannes. I'm excited to talk to you guys a bit about how we've been thinking about environments and the ecosystem around environments. And I think it's definitely been a shift for a lot of people's perspectives where on one hand going from single turn to multi-turn might seem like just a tech technicality in your RL algorithms but on the other hand it totally changes the infrastructure and the way that we think about building and composition of people contributing to these pieces of the pipeline in that before it was all data. The way you kind of add new skills to a model is you need new data sets. in SFT or in RLHF or in kind of just single turn reinforcement learning in the the deepcar1 style. But once you move into the agentic paradigm, a lot of things become complicated very quickly in terms of like what you're actually training on, where are these coming from, how are you building these, how are you validating them, what are the toolkits you're giving to people, and what is your also your pitch to people as to why they should be building these?

8:39 How do you get these at scale? and so some of the labs will like have lots of these like contracts where they're going to like try to have lots of startups go build them for very like large contract values and have them kind of come inhouse and vetted through some like large scale pipeline that they're doing. We wanted to take a more open and kind of collaborative approach to this. And the framework that we've been working on over the past year is called verifiers which is a developer friendly toolkit SDK for building environments as well as eval. And we've really thought about the developer friendly angle as like one of the most essential things to prioritize where like to us the goal of the environment SDK is not to have it be the most like native to RL or have it be the most easy to use for the trainer but really to have it be the thing that makes it interesting and fun and useful for people to go build these environments.

9:30 and so a big part of this is we have decided that using the OpenAI standard is what makes the most sense for building out these environments in terms of what the inference logic should look like. how people should be prototyping these. And this means that people can prototype environments on CPU. They don't need to like have a tokenizer. They don't need to like think about the tensors or kind of applying a chat template. We kind of are abstracting this away to create a layer which allows developers to just write the agent logic for the task that they care about. And we're doing this to leverage community feedback from the open source world. Sorry, I'm just looking at time. and to ensure that we have a lot of people who are able to contribute even if they aren't like already doing lots of reinforcement learning or already doing large scale training. And the environments hub has been our platform that we've released two months back to kind of be a home for these environments to have a place where people can kind of come browse, explore, find examples, see eval fork environments, star them, favorite them, and ultimately like figure out how to both build them for themselves as well as browse other ones that they can use that other people have already created.

10:38 And so verifiers, the framework is it's out on pipe. We've kind of seen a big surge in adoption since the environments hub launch a couple months back and it's been really exciting to kind of like get a lot of feedback from the community about what are they building and why and what is the goal of the people who are doing reinforcement learning. how can we make their jobs easier? How can we make their research process more fun and more collaborative?

11:00 and so there's a lot of design choices that I want to kind of talk through today about what the library is, what it isn't, why we did things the way that we did. and kind of one of the high level goals that we had from the start was really taking inspiration from the eval ecosystem. and so I think the eval eosystem for agentic benchmarks on one hand there's a lot of like really cool projects that are like very becoming like standard evals in terms of the scores that people report when they release a frontier model. But there's also a huge range of diversity there. So some of them it's like a single turn math benchmark. Some of them it requires a a loop with some game environment. some of them it requires simulation of a database. Some of them requires like a a full Docker container. Some of them require executing code on GPUs. and so there's a pretty wide range of stuff that we're doing to evaluate these kinds of tasks. And to us, this is really like the kind of baseline of what you want to support is that if it's not possible to very easily and directly port some existing eval or environment into your framework in a way that feels very natural where it feels like this could have been the initial implementation, then you've done something wrong. And so our goal was to have it be first like a first class eval framework too. So we wanted to be something where you can evaluate like any model endpoint with the same framework. and develop with API models. we also wanted it to be very portable. So we want environments to be the sort of thing where you don't need one big overflowing folder of like all the environments people have ever created that is like your hub. But we wanted environments themselves to be projects, packages. and so the environments hub is a package registry under the hood where you can install environments from it. And when you upload a project, it is uploaded as a Python project, a a package that you can then install. we've definitely been very heavily adopting kind of object-oriented design principles, especially when it comes to compositionality. So to kind of cover the range of sorts of things people have wanted out of evals or people want to do with environments, we think a lot about like what are the kind of baseline assumptions. So like you need to have a roll out, you need to have scoring, but then also what are the different flavors these can take? What are the sorts of things that people might want that will poke at the edges of these or that some paper did in a very bespoke way that maybe there should be a better abstraction for? And so thinking about like what is the total landscape of the sorts of things people might care about in evaluating these patterns and then having a hierarchy of features and abstractions that really support building these in a very direct manner. and so verifiers as a library is really opinionated about the abstractions it chooses. is not an off-the-shelf eval framework in terms of like a lot of the existing ones. I think there's a lot of eval frameworks that are like really good for single turn problems, but a lot of them start to kind of hit rough edges where if you want agentic loops in there, it's just pretty hard to hack them in. it's also not gymnasium. gymnasium is like very popular in the old days of RL with kind of this rollout loop of state action reward MDPs which is really useful for kind of studying classical reinforcement learning and doing reinforcement learning research but it's also not the most it was designed pre-LM it was designed before the era we're in now and it also kind of the the kind of direct way to do it forces people to think about the tokens and requires you to kind of repeat a lot of logic around like tool call parsing and chat templates that in our view is just not what you want developers to be working with. It's not how you kind of scale these because there's lots of kind of foot guns and things that can go wrong very quickly. And if you are building things in a way where like this is kind of offloaded to every developer to kind of get it right the first try without having it supported by native abstractions, this becomes a mess very quickly. and so we've thought about the idea of environments as really this encapsulation of a data set potentially an eval data set as well.

14:56 a parser which is kind of an optional thing to kind of extract logic from whatever the LM gives you into the format for your reward function as well as reward functions themselves. And so we this rubric abstraction has been kind of one of the hearts of the verifiers library where we want it to be really easy to add new reward functions take a reward function from one environment and stick it to another have metrics that don't actually influence your reward. So here you can have a format reward function that has weight zero meaning you're just logging it. So we want it to be really easy to kind of get information out of the environment of what's happening in your rollout and like plot your metrics for it. and so these are the sorts of things that we really kind of thought came up with as we were designing the kind of initial versions and iterating on okay what are the actual problems we want to solve with reinforcement learning. What are the sorts of things that we want to be able to see while we're doing this? What are the kind of papers that we might think are interesting or we've read about or heard about or paper ideas we have? and what would it take to actually have all these supported initially in a way that doesn't require kind of throwing everything away? And so, for example, like MCP environments, I think there's a lot of things where like people want an MCP tool to be used as a backend for some like agent agent. And so, in the case where like these MCP servers are stateless, this is a different problem from when the MCP servers are stateful. if you have a like a stateless tool, you kind of want to do your hosting once if you can have the tool autoscale itself rather than having it hosted once per roll out. if it's stateful, then you need to host it once per roll out. And so like these design considerations are the sorts of things where like depending on what your environment is, there's like different things you'll want to do.

16:32 If you're using an LM judge, you want kind of a default way to bring your judge and like have it plug in to your environment. And so like the judge rubric abstraction is one of the things that we've found was like really useful because you can just there's a default prompt but you can also just bring your own prompt for the judge and then have a reward function whose job is just like to take whatever the judge returns and parse this into a score. and we can have ways of like bring these MCP environments that are then going to be can be launched in containers can be like managed for you by the environment abstractions so that the things that the users have to think about when they're building a new environment are kind of as minimal as possible without having to repeat the effort of kind of prior efforts. And this is very much like an ongoing process. So it's the sort of thing where like we are under no pretenses that like it's done. It's going to keep going because the things people want to do are going to keep changing and growing. but as an example, if we're talking about like environments with state in them. So to get to an environment where let's say we have an agent who is working with essentially a Jupyter notebook. So in this example, we have a an agent who's going to be given math questions and has a Jupyter environment, a Python notebook style environment to be able to use as a calculator to solve math problems.

17:41 So the way that we think about it is this environment layer is the base. Then we want multi-turn environment which is essentially our like react pattern. So this is anything that is going to be like a a while loop of multiple tool calls in a loop that all kind of are a growing roll out. tool environment is okay it's one of those and the interaction loop is tool calls or function calls. Okay, it's also a stateful tool environment which means that we need to be able to inject other things to these tools for routing to backend information in terms of like if there's a sandbox like what's my sandbox ID. Then we have sandbox environment which is kind of our managed layer for having like a docker container that you can interact with in these loops so that each rollout is going to have its own container declared and kind of grabbed from a pool to be used alongside the rollout and then torn down when the rollout is done. and then on top of this this is that's just kind of like a base Docker layer. You can then build a Python layer where the Python environment is just actually the last mile of determining what is going to be in the container that exposes a Jupyter style interface to the model itself. And the goal of all of this is to have everything be able to be more than sum of its parts and to grow over time and to have each subsequent environment be easier to build. we want this to really be easy for people. We want to be fun for people. We want it to be educational for people, but we also want it to be useful. We want it to be the sort of thing where once someone has built a really cool environment, we can then take what they've done in the open source community and make this easier for other people to build on top of. And this is kind of why we thought about this hierarchy of abstractions as being the most robust way to do this.

19:15 We want it to be useful and we also want it to be communal. We want people to be able to share these and fork them and modify them. And I think there was a time when like every new paper was kind of talking about how it was a new unified RL framework. That to me always felt a little silly. Like not every paper needs to be like a new trainer. If you're not changing the trainer, you shouldn't have to hack in a judge to do LM judging or like different tool call parsing. Like this should have there should be a home for this in the stack that doesn't require pretending like you built everything from scratch. That doesn't require totally reinventing the wheel on things that are becoming common patterns at this point. and so like I think we think think about ecosystems like MCP N8N or like the old eval benchmarks like LMAL or Helm where people are doing these single turn benchmarks in a way that are composable and we've really been trying to think like what is the sustainable like reliable way to do this in the open for the agentic RL era and that brings us to the environment hub and so the environments hub is our platform for making these and sharing them we have a CLI tool just called primes you can do UV tool will install prime and you can start playing with it. You can download them. You can use models to run eval with them. You can then upload and share them. You can download them and star them and fork them. and you can really browse them. And I think the goal here is to start having a place where people can contribute environments. They can contribute evals. can use these evals in their research projects and their training runs to start having a an ecosystem that can really allow alternatives beyond just what is happening inside of the big labs to start becoming more of a part of people's daily workflows in terms of optimizing models for the tasks that they care about because reinforcement learning is hard. There's a lot of like pieces of it that you don't really want to worry about. Like not everyone who's using RL needs to change the underlying RL algorithm. Not everyone who's like doing computer use needs to like reinvent a whole computer use framework if what they're doing is able to stand on the shoulders of what's happened before. And so that is really at the heart of what we are trying to grow as an ecosystem is a way that this can all be build growing over time to something that is useful and powerful and allowing end users the people who are accustomed to agent frameworks and Python libraries and API calls to actually do large scale reinforcement learning. and so these environments just a little bit more about like the structure is there is essentially a a pi project file which declares what is in them. There's source code for the environment itself. You get a read me.

21:44 then we have some opinionated kind of defaults of how we think people should structure these projects but you can also break away from those as needed with kind of some warnings. but really the goal here is that these can also be the ways that you just measure if your system is working. If you want an eval running evals currently is quite annoying. Like if you want to run five popular benchmarks, you have to clone all the repos. You have to install them.

22:06 They all have a different setup instruction set. you have to set some environment variables. And our our goal with this is to really be able to like have this be as immediate as possible. So we have like continuous integration. We can test that these environments work out of the box. We can kind of highlight in green that like okay the builds are passing, the eval are running correctly. and then we can also like look at the scores and validate these versus reported benchmarks. So we can kind of also get a signal as to like if an environment is matching kind of a a source implementation this is a a good sign that it's been essentially implemented correctly. beyond this like eval are the sort of thing that we want this to be a first class platform for. So we want the environments here to also be the sorts of thing that are useful for people as a way to like report the eval of their models to share a new eval created rather than just it needing to be a kind of standalone git repo. but something that is more kind of native to the research practices of what people are actually doing when they are kind of evaluating and building models. as well as being able to just kind of understand performance like is my EV implemented correct correctly? how well do different models do? and so we've been kind of sourcing implementations from the broader community via a sponsorship program for kind of what we see as the most popular and most notable environments in eval. we also have launched inference as a kind of first-class feature of the platform where like we are able to kind of proxy to a number of like closed models as well as our own API models as well as kind of coming soon doing inference on the models you've trained with environments to be able to have kind of a unified experience layer for doing this kind of research for like prototyping and optimizing agents in RL environments. and this also includes things like building, we are a comput platform initially, and so we have GPUs, but we also have CPUs. And so if you want code sandboxes, we'll we'll happily run all of your Docker containers for you. If you want to like have a bunch of agents all running in different containers that each are going to only be alive for 45 seconds. And maybe you maybe you want to do this on your cluster, but maybe you don't, and you want to just kind of offload this.

24:08 And so we have this as a kind of feature of the platform as well. we have compute if you want to do some training. and we're also like just very excited about kind of community collaborations. And so we're actively integrating with lots of popular open source frameworks. Prime RL is our own main trainer, but we also have support for Tinker from Thinking Machines. We have a in the verifiers repo, there's a transformersbased trainer that looks a lot like TRL.

24:31 we've integrated with Sky RL. We're working with Nemo RL and others in the works because we we don't see this as like everyone uses RStack for everything. like people are going to have other considerations, other goals, and we want to ensure that the environments that are created here are not just for Primal Act, but they're for anybody who wants to make use of them. And kind of the main goal to be able to use them is we ask that trainers implement an OpenAI endpoint. And if trainers can do that, it's not too hard to plug in these environments. and so that is kind of how we thought about what community collaboration can look and our broader goal here is to kind of just foster more open research. So we have hundreds of environments so far on the platform. We have a large scale bounty program with lots of money committed from ourselves as well as other partners to reimplement papers to kind of sponsor projects around building on the platform. We want to be able to ensure that for any like high important benchmark sweeb bench or towbench or any of these things that we are committing to paying someone to reimplement it in our framework as well as validate the implementation is correct. and we have a kind of research residency program for people who want to kind of create new novel benchmarks and environments on our platform. and really our goal is just to grow the community and to grow the ecosystem around RL environments and evals to be something that is very developer friendly, very useful for people and the sort of thing that allows this research to not just happen behind the walls of big labs, but to become more of a first class feature of the developer life cycle around AI applications for researchers, institutions and companies around the world. and throw it ahead, we've been using environments hub for our intellect 3 run. We are growing the community with more integrations and partnerships. and we are learning from the community like iteratively by kind of engaging directly with people to understand what do they need? What are the things that we don't support yet that they really want us to support and how can we add this to a roadmap? How do we make sure that the sorts of things people care about in within this kind of domain of RL and RL environments are supported by our platform and really using the stack day-to-day and dog fooding it is like the the models we train are built with the tools we're building for everybody.

26:36 And to us that is kind of the one of the most important ways we can demonstrate our commitment to doing it right because we're building the things that we want. We don't want it to be hard. We want it to be easy and fun friendly. And so that's what we're building and that's a kind of sharing with everybody else as well because we see super intelligence as a meta goal of how does the world advance the capabilities of AI systems to solve problems for us to make our lives easier to innovate to discover to advance research and we think that the stack for this the the tools should be open and accessible and available to everybody.

27:08 Awesome. I think I have a couple minutes two minutes for questions. Johannes, Johannes can come back up as well. >> I'll get you to the microphone. >> hi, thanks for the talk. I just want to check whether you have plans for implementing the older style of state machine and and RL frameworks in your newer one. Not just considering LLMs, but probably something towards for robotic where we don't have data but want to use RL. That's a great question.

27:44 So, we definitely have like talked a lot about like robotics. I think it's it's the sort of thing that we will want to do some version of eventually. Initially, it'll look like VLMs. so we do support like multimodal environments. and so kind of anything that can be done on top of VLMs is the kind of most immediate version of this. I think to us it's we would want to like learn more from the broader community of like where does that break?

28:08 What do you need for robotics that we don't support? but I think our goal is not to have it be like like we are generally thinking of this as like things that are LLMish. So maybe this is VAS, but maybe it's not like the sorts of models people are doing for like small model Atari PO. I think that that is already a very like active developer community and has a lot of infrastructure built around it and our goal is not to kind of like compete with that. It's more to have a version of that for the current agentic LLM RL paradigm.

28:45 >> Well, just we'll just take one more questions. >> Mhm. >> just moving beyond RL, what do you see as other applications for the kind of async training protocols and reinforcement protocols you're working on? >> Yeah, so there's definitely a lot of interesting stuff around distillation and as well as just SFT and synthetic data generation. So like I think RL is kind of an instance of the sort of pattern and so these environments are like people are using them for RL now but they're not just for RL they are also kind of first class synthetic data engines where you can take any model you want run it through and kind of collect these rollouts and then use reward functions just as like a filtering mechanism where then you now have data you can train on you can do SFT you can do like true distillation via logits there's a lot of ways you can use environment abstractions and other models to kind of have more data. But really, I think a lot of it will be cases where you want to like have an environment mirroring some real world goal of an agent. some of this could be like collecting traces from your real world deployment and identifying patterns in these, identifying like ways to like have this as a starting point for like resampling a new completion.

29:53 Like if an agent messed up the last turn of a user conversation, how do you like replay this and try again and train on the better version? sorts of things like that I think is a pretty wide open space. Okay. if there are more questions you can catch them right outside in hallway. Thank you Will and Johannes. >> Thank you.

© transcribe · For agents Built with care and craft by Gokul Rajaram