Data and Environment Curation for Post-Training LLMs — Mahesh Sathiamoorthy, Bespoke Labs

AI Engineer · 19 min · 118 sentences · from YouTube's caption track

Each timecode opens YouTube at the start of that sentence. Line anchors (#s42) are the cue ids in the WebVTT, and every line carries its start and end seconds. All transcripts has every talk, and the whole corpus as one file.

00:12Hey everyone, um today I'll be talking about data and uh environment curation for uh post- training LLMs.
00:21And I am Mahesh Satyimi.
00:23Um I'm co-founder and CEO of Bespoke Labs.
00:27And previously I was a researcher and uh engineer at uh Google deep mind.
00:32So very briefly I will tell you a little bit about uh bespoke and uh after that the talk will be mostly around uh opensource work we have done.
00:41So bespoke is an applied data research lab with a mission to help enterprises and frontier labs access high quality data and RL environments for their post training needs.
00:53So very briefly what we do and what we have done is that last year we put out something called curator which is a tool for curating uh synthetic data
01:03for post training with basically SFT and right after that actually deepse landed and we started an effort to curate reasoning data and that's how we started something called bespoke stratos
01:15which eventually formed something uh into the project called open thoughts which some of you hopefully know about and we have also been core contributors to terminal bench.
01:26Um you know these days we we do a lot of research and build and ship RL environments.
01:32So I was actually looking forward to the previous talk uh from Nick who is also you know uh doing something similar uh and and the other thing we do is we
01:43do a lot of post training and help enterprises uh to get their own custom models.
01:49Right?
01:49That's the name of that's how we ended up with the bespoke uh t title for the company.
01:56Uh the the other thing I want to kind of mention is there's there is this you know you know in in our industry there are a lot of people who create data
02:04uh create RL environments and then there are the uh researchers who consume this.
02:09But I feel like there is this slight mismatch and it's kind of beneficial for someone to kind of go do both at the same time.
02:18And in fact uh as you're curating data you want to put yourself in the shoes of the researcher to see what what does it take to you know uh actually move the metrics on the models.
02:30So that's one of the motivations of how we kind of think about the other thing I want to kind of talk about is you know how you know uh AI has evolved right.
02:40So early on we used to think about and evaluate models on what they know.
02:46Um for example this is a this was a very popular benchmark on uh testing LLMs on various kinds of STEM humanities and all that knowledge and these days we have all these benchmarks that test
02:59uh how how agents are able to do things we have moved on from knowing to doing right so that's the idea of agents obviously and the one of the key
03:10principles or one of the key things about agents is that they are autonomous And there are as I was saying there are many benchmarks including uh swb bench terminal bench and so on but ultimately for many people what they care about is
03:25are these agents autonomous for long durations of time uh Nick had uh sorry uh Ross had a great talk on long horizon right so that's the goal is eventually we make these agents
03:38autonomous for maybe few hours or you know few days or few weeks and what is it that's blocking the uh autonomy of agents.
03:48It's basically reliability, right?
03:50So, at some point something falls apart like either they called the wrong tool or they made a mistake and you know what what not, right?
03:58And what what's one lever to improve reliability there?
04:02There are of course many um obviously you can prompt your way to improving the agents u reliability or you can uh update the harness you know the tools and whatnot
04:14but postraining is a very powerful tool to improve reliability or or maybe even pre-train good models right so if you think of frontier labs this is one of their primary mechanisms of improving agents over to to
04:29uh get better [clears throat] u capability ities in uh various domains or you know or for better you know benchmark numbers or better u um autonomy for long longer and longer durations
04:46and for post training one of the popular techniques as you know is reinforcement learning and that's kind of um something you know a lot of you are excited about is the you know notion of RL environments
04:59but ultimately for post training be it SFT or or uh reinforcement learning data is the bottleneck right so when when I talk about data RLNs are also something I'm calling it as data it's just the data is now in a very different shape
05:16u again here you know compute is kind of well definfined models you know uh good sort of models exist and the uh infrastructure to post train for example u the There are various providers like fireworks,
05:33tinker or uh slime world and whatnot.
05:36So all all of those are somewhat well defined.
05:40Most of the places where people struggle especially enterprises is that they don't have access to good quality data and RLNs and this obviously also applies to frontier labs where they have all this infra setup and they are you know needing good quality RLMs
05:55right uh beyond so that that's one of the this is kind of how we are thinking about why to invest time in you know doing data research and RLN research
06:06and as a side note one of the um other there are many other benefits of post- training for example you can reduce latency or improve cost throughput and whatnot
06:18and I'll give one concrete example of a post- training work we did uh with one of the enterprises so in this talk I will mostly uh cover some of the work we have done in the open-source
06:31uh community so we did some work on curating reasoning data for reasoning models and for uh you curating trajectories and uh environments for agents and recently we had an engagement with post training which uh I I'll very briefly talk about and some tools on data curation
06:51so open thoughts um is a reasoning data set as well as a paper right so we we started this effort last year as I was saying this uh we we after DC came out we realized that there is uh lack lack of very high quality reasoning data in the uh community.
07:12Obviously the labs have access to good data but outside we didn't have access to data right so we we at bespoke started this effort called bespoke stratos and then we realized that this is actually quite useful so we joined
07:26um together with various folks in uh Stanford UC Berkeley Udub and so on to create this consortium called open thoughts and we did lot of work on basically identifying the curation recipe
07:40and we also published this as a paper in night of this year and this is the main figure of the paper.
07:46So what it shows is like we we figured out a curation recipe and it shows the scaling law right.
07:52So again this is last year when Amy uh and and uh live codebench and these these were some of the popular benchmarks.
08:01What we showed is that with this recipe if you keep you know scaling up the data set size the the you know the the it's a scalable recipe right the the metrics also improve
08:12uh it's actually very widely used as well for example this is um Microsoft cso tweeting about the work and this alle Alex is my uh co-founder he's a chief scientist
08:24and also a professor at UC Berkeley and this is John Schulman talking about open thoughts that he as he and his uh colleagues have been using it internally at uh thinking machines right and some of their blog posts also reference this
08:39so I'll uh talk about how we did the curation for open thoughts um this is the pipeline that we used so you start with curate you start with a bunch of source questions right so there are various data sets out there that have the
08:56uh prompt response and we choose with the prompt so we start with the prompts.
09:00These are various uh sources we have.
09:03And then uh if you look at the paper, so if you look at this graph for any given data point, say if there are 10,000 samples that you want, the question is then how do you choose
09:16uh the questions from all these different data set so that you have 10,000 uh for the data point.
09:24So the the then there is the aspect around how do you mix these questions.
09:28So uh you can use various methods.
09:32So the paper talks about for example using LLMs to check for whether this is a good good question hardness of a question and so on.
09:40And then you want to filter questions u and generate the answers.
09:44Again this is all like driven by LLMs right?
09:47So this is the curation recipe we did for uh creating this reasoning data set and the answer generation is using teacher models.
09:55So you can take other reasoning data uh reasoning models such as deepseek or quenbased models or even gemini and whatnot and then you can also filter the answers once you have the answers for uh these questions.
10:09Um and and then you can also you know given a question generate multiple answers or a single answer.
10:16So these are various knobs in the curation recipe and the systematic way of doing this is like you run ablations and figure out which uh you know in each of these stages what works and you kind of proceed to the next.
10:31So after doing all of this you get the final recipe right.
10:34So this uh you can read this paper it has lots and lots of uh you know information about how we did the curation but here are some of the learnings that you know some of them are quite
10:47uh counterintuitive and some of this was also covered in last year's uh AI uh AI engineer conference.
10:54For example sampling um multiple answers per question works pretty well.
11:00This is something that uh we it's it's kind of counterintuitive.
11:04So as an example, something else we could have done is we could have had more much many more questions and then just answered them exactly once versus taking one question and answering them 16 times.
11:18The I think the the reasoning is probably that it gives like a variety of how reasoning is done.
11:24So the the during finetuning we also use the the reasoning traces, right?
11:29So I think the diversity helps there.
11:32And the other thing we saw is like the stronger teachers are not always the best uh uh stronger models are not always the better teachers.
11:42And there were a few other counterintuitive aspects around like you know uh synthetic question generation or question answering working whereas answer filtering and other aspects not working very well.
11:54And after the open thoughts work which was around uh data curation for reasoning models such as you know uh deepse kind of models, we moved on to open thoughts agents which is
12:05um very similar but how do you curate these the the data and RL environments for uh training agents now right not models.
12:15We we have a very similar figure here.
12:17Again we want to establish scaling loss.
12:19Um so as you increase the data set size we want to make sure that the curation recipe actually works.
12:25Uh and again I'm I'm not going to go into details here but very similarly there are various ways of choosing different sources for example stack exchange and and whatnot.
12:38How do you mix the test?
12:39How do you filter?
12:40Generating the rollouts uh choosing the teacher and so on.
12:44And again these are some of the lessons learnings u as an example even here we saw that stronger models are not necessarily the uh best teachers right so we found out some some some of the I think uh
13:01um quen models were better than for example um um um claude models I think and sampling multiple answers again helped in this case synthetic rewriting and task augmentation um
13:16is something we thought will work but it didn't very work work very well and the other thing is like in in this whole process of building this open thoughts agent
13:26SFT still contributed a lot to the gains um RL was kind of you know it's very comput inensive and for for the last few few percentages it really helped uh but but you know in many of the situations for example in enterprises
13:42SFT actually works works pretty well, right?
13:46And here is one concrete example I wanted to share on um uh actually deploying something to production, right?
13:54By post training.
13:55So we have seen a lot of people talk about post- training but in enterprise settings we haven't seen lot of successes at least haven't seen uh that here is a very concrete example of uh with intude there is a this app called credit karma which if you install there is a page
14:11place where you can uh the the the app gives you a reasoning as to why a credit card has been recommended and this you can prompt a model to do this.
14:22But one of the reasons one of the places where it fails is that the you know it it it's not always compliant.
14:29So you have to have a long list of rules to make sure the responses are compliant and that actually blows up the latency.
14:37So answer here is like you want to curate data and post train right seems kind of straightforward but one of the things that we ran into is um the data set can be quite impa imbalance and lot lots and lots of places for example you will have 0% APR and the model after fine-tuning
14:57can kind of hallucinate the these uh numbers so this again kind of ties back to what Ross talked about some time back with respect to attacks And we kind uh we created this specific
15:11uh curation recipe where instead of just having these uh question the the prompt response pairs in plain language, we added these uh tags which helped the model to focus on you know uh the the kind of form rather than the specific numbers itself
15:28and that gave a big boost and uh we we saw that um the the overall the compliance metrics improved the latency improved, the throughput improved and eventually you know they they are able to own the model right as frontier models improve
15:44they don't need to kind of go and u um uh update it and also as we see now the uh frontier [clears throat] models are also getting more and more expensive
15:55and you know this kind of g gives them a very good way for owning the model and also um lowering the costs.
16:04I think with that I want to briefly touch upon um uh you know curator the tooling that we had built last year um which is for curating reasoning data.
16:16So uh what it does is you can basically uh you know um specify the you you can either go with say a hugging face data set where you have various prompts
16:29or uh in many situations you may have collected logs and you want to get the responses and fine-tune a model.
16:36So this curator kind of makes it pretty easy to do that and it comes with the integration with you know uh tinker and fireworks and this this is again the tool that we used um originally for curating open thoughts
16:52and here is a very very detailed diagram of what we are building today but um this again connects back to um what Ross was talking about where he was talking about algorithms
17:04uh environments and compute root right so it it feels like you know we are kind of converging on something very similar so if you think about uh the the stack that is needed to say not just curate these RL environments but to post train models
17:22one of the things you need is obviously handle on like how do you build these RL environments how do you measure the quality how do you track the different versions and so on so that's one of the layers and below that you want various infrastructure
17:35to uh um sand use sandboxes, right?
17:39to spin up the rollouts to to spin up the sandboxes to generate rollouts and especially if you have long horizon rollouts then maybe at some points you need to do a checkpointing and then you need to be able to
17:50snapshot or roll back to something else right so that's the other uh the lower level u you know compute and orchestration and at the top I have been giving examples on post training so there is all this uh layer around like how do you do SFT
18:05how do you do RL and so on but there is also this method called Japa which is around uh which is on prompt optimization.
18:12I don't know if you if you guys have heard of it but you can use LLMs itself to uh to to kind of optimize the prompts based on reflection.
18:22Um so that also works pretty well for updating the system prompts and also the harnesses.
18:28So this is kind of I feel like you know the the new architecture or the new reference uh stack for how how at least we are building and how many others are building um the the stack on how to build the RLMs
18:42and then also post train agents.
18:45I think with that uh I will uh end the talk and u you know happy to take questions offline.
18:53[applause]