Hugging Face Journal Club: AsyncOPD and How Stale Can On-Policy Distillation Be?

[00:00:02]This paper is doing something if I understand correctly that is very similar to what pipeline RL proposed for reinforcement learning which is where you want to basically have like really decoupled generators from learners from trainers um but you have the additional complexity in um distillation that you need to for um the rollouts of the student. [clears throat] And uh this is a bit different from like reinforcement learning because in conventional like GRPO for example um you just have to assign scalar scores to your trajectories. Um but here you need to also factor in that you need the the log probabilities um of the student and the teacher and once you go fully asynchronous um these can start to diverge. So for example the the rollouts from the student um they can be you know several steps off policy um and then the way that these are then graded by the teacher can potentially introduce uh some instability and I think after reading all the the math the most simple way to understand this paper is I think this figure.

[00:01:24]So they they basically compare um like three main ways of doing uh on policy distillation. You've got the synchronous case which looks very familiar um to the RL. So you you have your your roll out then you have a teacher doing the scoring then you um have I guess like the the the back prop on the on the policy you sync the weights and then you generate again. So this is like I think pretty familiar and um if I understand correctly this is also how we have it in taterial like in the distillation trainer um you have this kind of like you know GPU does this then you do that then you do this. So this is like the the reason that usually you get things like blocked GPUs because while you're doing back prop for example in the green blocks you're not generating.

[00:02:19]So they mentioned that um libraries like Verl um have actually got like a kind of kstep off policy uh version of distillation. And so here um the difference is that you you generate from your policy and while you're doing the scoring in red with a teacher and the uh the backdrop you're starting to generate the next um uh batch or rollouts uh from the existing uh policy and then at some point you have like a a weight sync and then you know the next set of rollouts use that new policy. And so here you get like a throughput advantage because you're not basically waiting um all the time. Um, but I suppose you're limited to how many steps off policy you can go before things go off the rails. And then their thing is this async OPD where you are just continuously generating uh from the policies and you've got all these kind of like mixed uh trajectories um where you've got potentially you know partially completed rollouts from one policy versus the next one and then you're just doing like uh you know uh uh scoring from the teacher and back prop uh all the time. So this is this is really like pipeline RL uh for distillation.

[00:03:49]Um and they they show here actually like uh if you study uh the quen models um using a bunch of like math uh tasks then there's you know pretty significant uh gains in throughput. Uh so here like for example it's like about 9 to 23. So a [clears throat] factor of 2 point something. Um here it's a bit less with a larger model but roughly you know I would say 1.5 to 2ishx speed up. I mean they even have it here I guess 2.7.

[00:04:26]Um and they also show that like the performance uh on these various math benchmarks is um you know comparable I would say within the variance the benchmarks.

[00:04:38]So it seems to me like it's a like no-brainer that one should go fully async. But the thing that I found a bit difficult to understand and maybe someone here can explain it is why they have to do all this like complicated uh Monte Carlo sampling in order to achieve their like desired results. So they they say here that okay when you do like uh conventional um knowledge distillation you you have like what they call dense KL and so here you've got like two ways of computing the K divergence one which is forward KL is weighted towards the teacher so you compute Q log Q minus Q log P from the student and then reverse KL is the inverse so now it's student weighted and now you have P log P minus P log cube and here like if you had all the logits uh from the student and the teacher you would just compute the dense kale depending on whether you're forward or backward and they mention you know this is something we know that this is expensive in practice so you have to do some sort of like specification or truncation and then they sort of say well if you just take the top k um logits then for forward to kl uh it's like kind of stable so you don't get any mis match. Um, but then for some reason here, I don't fully understand and maybe Carlos can explain better. They say when you do reverse KL, which is now student weighted, you get this kind of like mismatch uh in in the cache of the logits. Um, this then motivates this long discussion about Monte Carlo uh like sampling and it seems like pretty gnarly. But uh before we jump into the nollie bits, maybe just on this point here, Cuddles, did you understand why the reverse KL leads to these like cash misses?

[00:06:38]>> Yeah, it it has to do with how the top K and the cache interact. So basically what you do is when you score the the probabilities or the lock proofs for each of the tokens here, they call it actions, you store that in the cache.

[00:06:53]But if you're using top three, then you score, let's say, token one, two, three for both the student and the teacher.

[00:07:01]Uh, but then in the new policy that's in green, you may not have one of the log props that you need to actually compute the loss. So here you had three of those, but you're missing one that you didn't consider at the beginning because the policy changed and the top K is different than before.

[00:07:22]Ah, makes sense. Okay. And so in the forward kale, the reason it works is because the teacher is static.

[00:07:29]>> Exactly. Yeah.

[00:07:31]>> Nice. Okay. Cool. All right. So that that's nice. Does anyone understand the Monte Carlo bit?

[00:07:39]>> I do get the Monte Carlo stuff, but maybe Carlos if you want. So you >> No, no, go ahead.

[00:07:44]>> Go ahead. [laughter] The way the way I understand it is that they they avoid the top K altogether by just sampling let's say top one repeatedly and that's very cheap in terms of the the cost because you are scaling for example if you're sampling m times you're just uh storing m logits for the student and the teacher and then at the uh When are you actually calculating the the loss with the new policy, you just correct the the the estimate using the the ones that you stored. So you're just using top one uh sampling instead of top K or the full V full V cup and you just do the correction of the important sampling.

[00:08:38]Okay. So you're saying when you say sampling you mean like selecting like from this uh let's say list of log props uh which ones I should choose from the teacher and the student or let's say actor actor and the teacher >> basically for each token like having the full distribution that like a 64 token that you would have chosen >> uh basically so you you can do estimates on like the the probability that the new uh student that has shifted through time would still output a token that is in the distribution of the MC64 is like uh conversions to one basically. So you can still do uh conversions.

[00:09:23]>> I think there's a nice figure that uh if you go to figure eight in page eight I think.

[00:09:32]>> Ah yeah that one.

[00:09:34]>> Sorry >> it's a small one. Yeah, I I think this is a is a good example. This is with M2.

[00:09:39]So by sampling I mean literally predicting the next token repeatedly. So here you've got like hello and you do twice the process of decoding the next token and you get what and how >> and you just store those two for and then when you and store on the cache the probability from the current policy uh and the teacher and then when you are going to actually calculate the loss for the back prop and everything you get the log props for those specific ones because you already know what you stored.

[00:10:15]And you also know the probabilities of the current policy and you do the corrections using important sampling.

[00:10:22]>> Okay. But does this mean that like you have to change the way you do your rollouts in some non-trivial way? Like if I was doing like vanilla distillation, right? I would take my prompt and then I would just roll out till the end of whatever end is like end of sequence or whatever and then that's now my trajectory and then I score it.

[00:10:42]Um whereas here it looks like uh a full roll out is decomposed into a series of steps where each step is presumably a token and then at each step you are like doing a sort of one token ahead uh prediction is is that correct?

[00:11:03]>> Yeah.

[00:11:04]Yeah. I don't know in terms of for example in inference engines the mechanics of how that work but it would be like generating multiple times from the same prefix >> like >> yeah I mean >> no it's basically like you don't do arg basically on on the sampling you sample multiple times like uh you do uh top p top n but you do it like three times for the same prefix so you get three different tokens basically Yeah. Yeah. I I get that. The thing I'm wondering though is like does this mean like you have to like for a given prefix, right? You you do this like m time sampling, right? That that that's fine. That's like just as you would it basically creates like a tree, right?

[00:11:52]you have like your your prefix, it branches and then you know maybe you you pick one and then you branch and branch and that's that makes sense like from a Monte Carlo perspective how you then select you know which of those tokens um are the most important ones. The question I have is like if you're trying to get like the throughput gains um that you want from like being fully async, does this introduce some overhead where you now have to create essentially this tree, right? Like so for a given imagine I have like a task or a prompt and I want the student to to to solve it. Do I have to create this like fully branched tree in order to get all of these like um you know intermediate log problems and if so what's the tradeoff compared to just doing like you know a full roll out >> in my head it's not a tree >> like okay >> from what I understand like they just do it like auto reggressively and then for each step sample like three times.

[00:12:59]So they still like on this picture you can see like the ST+1 token is what it's not like you don't have another with hello how >> ah okay >> okay but okay fine but let's suppose like in this case right let's suppose like we generated like let's suppose the full final trajectory was like three tokens and we're doing two local samples per time step that would effectively mean that we are now going generate like uh sixish tokens, right?

[00:13:35]>> Yes.

[00:13:36]>> Like so in a sense your your roll out like token count increases by a factor of m >> but my question is like when we do top K distribution like autogressively we choose one token but we can like choose two at the same time.

[00:13:57]Ah, you're right. I see. Yeah. Yeah, you're right. Okay. So, if we're sampling with like temperature like you could just say instead of Yeah. picking the AR. Okay. I now understand what you're saying. Yep. Okay. So, it is actually a free free computation.

[00:14:14]>> Yeah. Yeah. And they also mentioned that that they they are aware of at the beginning of section 6.1 in the the first paragraph I think that they don't roll out into additional trajectories that they don't do the branching thing they don't do the tree things just what you were mentioning about just one roll out and then additional tokens per step.

[00:14:41]>> Okay.

[00:14:43]So that's what they mean. This draws multiple local next token samples without rolling them out. Okay. Okay.

[00:14:49]Okay. Good.

[00:14:52]>> Okay. So the idea is that like you you store all of these like local samples and then you do the like Monte Carlo.

[00:15:06]Um well, let me think about it. What do you do at that point?

[00:15:13]>> At at that point, you have the you have the loss function like the usual KL estimate, but then you need to correct it using the important sampling correction because you got the old Yeah, exactly. You got the old one and the new priority.

[00:15:31]>> Okay, cool.

[00:15:34]>> Okay. So then I think I take back my my claim that it seems overly complex.

[00:15:43]>> Yeah, it's it's actually pretty cool paper because they split the challenges very nicely and they come up with solutions to each of them.

[00:15:54]It's good. Yeah.

[00:15:56]What I still um like my question on this MC MCM Monte Carlo estimate stuff is um basically the bet is that the stainless is not far enough that when we do like m M equals 2 in the new policy we would end up like sampling a token that is in this uh m2 for each step if I understand correctly.

[00:16:22]Uh so we can like import and sample it and we still have the same support between the distributions. Um but what happens if we if we don't like let's say we we put stainless like to 100 do we have to like rescore using like the teacher if it's a new token we didn't keep the log distributions >> but I think you always do keep the distributions no because you know which tokens you sample at step zero and then at step 100 you got the full policy right so you need you can get the tokens that you need to calculate the loss.

[00:17:02]>> Uh so I I maybe I didn't explain like uh so the in my head like the the issue is like this cashmiss problem. So the MC part is like to solve the cash miss problems.

[00:17:14]>> That's right.

[00:17:15]>> So we don't cash miss with MC because we basically have more more um more distribution for each step that we sampled.

[00:17:27]But I think >> go ahead. Go ahead.

[00:17:30]>> Like I say, so I mean I think this plot here, figure seven, kind of gets to the core of your thing. So here they are showing the accuracy as a function of staleness and they're showing for each of these curves the the number of samples in the Monte Carlo estimate. So you can see that if you have larger MC, then it's like a bit more stable, but it it still goes down, right? like once you start getting to like 64 and 128 is down.

[00:17:59]>> So perhaps there's like some kind of trade-off, right? It's like I can push for larger staleness, but I have to increase the Monte Carlo sample size.

[00:18:10]Um, and then maybe there's some additional imputation that you have to worry about. Um but at least like you know for pretty large values of stalness it seems pretty robust at least on you know like just as a reference right like um I don't know if this has changed much since I last looked at it but I seem to recall in like the scale RL paper that you can only really go to like staleness of eight or something in in RL like at least in GPO before it starts to go off the rail. So yeah, this seems like roughly consistent like I mean correct me if I'm wrong. I mean maybe you've seen new papers but I seem to remember that even with like asynchl you can't really go crazy off of off policy.

[00:18:58]>> No, you're right. Hm. But then if you go to figure two, then we don't need to jump through these hoops, right? Because on staleness eight, like the vanilla reverse scale to get to pi h5. Yeah, this one.

[00:19:16]>> Yeah.

[00:19:18]>> Like you really see the dip starting on 16 or so, which is >> Yeah.

[00:19:27]You don't need all these corrections and MC implementations just the vanilla reverse scale works.

[00:19:36]>> Yeah, it's a really good point actually like on this aspect right when they compare the one step or twostep.

[00:19:44]>> Yes, exactly. So here they have this like two-step off right which is essentially this one here. Um so you're not even doing like the fully async thing but you're just doing the simpler you know generate k steps and then optimize uh on those steps right and this like you know okay it's debatable if it's like giving you totally comparable like speed up but it's still giving you some like let's say non-trivial speed up um compared to the fully synchronous case and is like way simpler.

[00:20:22]>> Yeah.

[00:20:25]>> Actually, is Kashif here or not?

[00:20:29]>> No.

[00:20:30]>> Okay. So maybe we should have a look at his like async implementation um to see what he did because um for me it's like if you can get like compar like some non-trivial improvement in throughput with like something simpler like you know kstep off policy this is perhaps the preferred thing over the fully async one. Um but you know to be decided.

[00:21:00]>> Why not do you think?

[00:21:02]>> Oh just because then you have to do this like Monte Carlo stuff right which again maybe is not that big a deal but it's like more annoying I guess in some small way.

[00:21:14]>> Okay. Well I checked the implementation this morning and uh it's mostly a sing with a different objective. So he went for fully a sync like the third figure.

[00:21:26]>> Okay, cool. So then let's try to make it work.

[00:21:31]>> Actually, I think these guys have a code base too. Um so could be kind of interesting to see um how actually hard it is to do the mod color thing.

[00:21:47]I I do have a question for uh Louis and Carl like and and also like this might be stupid but like how like how costly it is to cash like like this this cash miss like the whole problem is because we only keep topk of lobs of teacher but like really how like how costly it is like if we engineer around this. H you mean if we did like the dense scale?

[00:22:14]I I think for for short rollouts it's fine but then when you start 1k 2k onwards rollouts it starts getting expensive >> because you're >> sorry >> like to store is it expensive to store >> to store or if you have an external teacher server that you're getting the laptops from uh it it scales really poorly the the the communication between the the teacher server and the GPUs you're using for training.

[00:22:47]>> Yeah. Okay. Like [clears throat] what what uh I'm sorry like I'm I never did like distillation. So what's case like serializing and sending these weights like this? Uh >> okay.

[00:23:01]>> Um okay. And can we like reduce the size or like speed up the >> do you have I mean one of obvious one is top K but it has a bunch of problems like this one especially for reverse KO uh what people are doing is what deepseek 4 before did it's a trick that where you just store the I think it was the matrix of the attention layer something like that that >> uh prevents you from from storing the full like uh logits and instead you rebuild them at runtime. If you store some certain weights of the of the of the model and and why do we need like in reverse scale we only need like some tokens to to compute like the the the teacher. So why why not select these tokens from the like store everything but then when you query the teacher uh tokens you just query some that you you have sampled from the student >> um because you you need to do that at at the at the moment you're calculating the loss and like it's nothing it's nothing that you it's nothing that you can cash like preemptively because I I But but let's say like you do this at loss time like you know the tokens that you need from the teacher >> and you just query it like how fast can you just get these tokens?

[00:24:32]>> Uh depends on how many tokens in in the experiments we run it usually scales like linearly like >> okay >> K more tokens it's K lower.

[00:24:45]>> Okay thank you sorry for the stupid question. Yeah, it's a good question.

[00:24:50]Um, since I've forgotten, can you remind me, Carlos, like what the trade-offs are between forward and reverse KO?

[00:25:00]Um, it's mostly like a specialization.

[00:25:06]Uh, so in in in in a nutshell, for KO is mean covering. So, it tries to absorb everything from the teacher. Whereas reverse scale is mode covering. So it's it can like very specifically target the capability from the teacher. So in terms of distillation, it's really useful. For example, in math domain, you're just really targeting the math domain present in the prompts and not trying to cover everything that the teacher is trying to do.

[00:25:34]>> Ah okay. So if you want to become if you want to train like an expert model that you then later use as like a teacher or something, you would be doing reverse kale.

[00:25:44]>> Exactly.

[00:25:49]And um in the distillation trainer that you you implemented um is it doing like dense or is it sparse?

[00:25:59]Uh depends on the setup and that's actually something that Quantan and I were discussing on on the on the refactor. If it's local, if the teacher fits into the training GPUs, uh it's fully uh the full vocabulary because you have access to the teacher. But uh when the teacher is outside in in a different server, we do just uh top one >> for reverse and top K for forward.

[00:26:31]Gotcha.

[00:26:38]for the reverse um no for forward KL. Um do you see like I would expect the signal to vanish at some point uh when the well the token sampled by the students starts to become like very unlikely in the teacher distribution.

[00:26:58]Um, is it something that uh that we can observe or not really?

[00:27:05]>> Yeah. Yeah. I I think and also to to to lose this question, I think 4K really starts struggling when the distribution between the teacher and the student are very far away.

[00:27:17]>> Okay.

[00:27:18]>> Because you're trying to target the teacher distribution waiting by the teacher distribution. So it's like you never ending loop where you can't get out of there.

[00:27:29]>> Yeah, it makes sense.

[00:27:36]And one question for the MC um I think I understand but just to be 100% sure when you sample the the the token you do it uh from from the distribution right you're not sampling uniformly on the full vocabulary.

[00:27:58]Yeah, I think yeah from the distribution I think from the policy.

[00:28:04]Yeah, >> I know that with VLM like you can get uh I think it's probably like the top end um return along with the the sample token. So technically it's not it's not exactly the same because you get uh is it the same? No, I don't think it's the same because you take the distribution, you get like the top k of the distribution. But here in this process of MC, you you would want this the uh server to just reample from this distribution which is slightly different, right?

[00:28:42]>> Yeah.

[00:28:44]>> Okay.

[00:28:52]>> Yeah. I suppose it would be quite interesting to look at how they did the inference here, right? Um I wonder what they're using.

[00:29:09]>> Yeah.

[00:29:36]Yeah, maybe with some monkey patching.

[00:29:50]Okay.

[00:29:54]Any final questions?

[00:30:04]>> This really cool work.

[00:30:12]Then with that, I'll catch you guys later. Thanks a lot.

[00:30:16]>> Thanks. Bye.

[00:30:17]>> Bye.