<aside> 💡
TL;DR
We introduce speculative decoding into the RL sampling process, achieving a significant improvement in sampling speed under appropriate batch sizes. Furthermore, the draft model is also updated during the training process. Compared to freezing the draft model, the accepted length is consistently maintained at a high level, generating long-term stable positive gains.
</aside>
slime/docs/zh/advanced/speculative-decoding.md at main · THUDM/slime
Speculative Decoding is well-known as a clever inference acceleration technique. Specifically, during inference, instead of having the expensive target model decode token by token, a lightweight draft model first decodes and generates multiple tokens, which are then verified in batch by the large model. Tokens that pass verification are directly used as the final inference result, while failed tokens are resampled using the large model. Ideally, if all tokens generated by the draft model pass verification, the system can accept these k tokens at once, significantly improving inference efficiency. However, if the discrepancy between the draft model and the target model is too large, resulting in too few verified tokens, it might conversely have a negative effect.
Speculative Decoding, this double-edged sword, has already seen significant application in industrial-grade inference engines and also holds enticing potential in On-Policy RL. On one hand, speculative decoding can significantly accelerate the rollout sampling process, and the sampled tokens are, in distribution, identical to those from the target model. In addition, long-tail trajectories significantly reduce the effective parallelism during RL rollouts, making it difficult to fully saturate hardware compute; this regime naturally matches the application scenario of speculative decoding. Of course, all this is premised on the sampling probability difference between the draft model and the target model being within a reasonable range—if the policy difference between the two is too large, the acceptance rate of tokens predicted by the draft model will plummet.
This is the problem solved in this article—we introduce speculative decoding into the RL sampling process and synchronously update the draft model as training progresses, stably improving sampling speed.
This feature has already been merged into the slime main branch and is available out of the box. See the documentation for usage details.
Megatron added support for EAGLE MTP SFT in v0.12.0rc3. Based on this, we consider performing online SFT on the draft model during the training process. Specifically, we add a new cross entropy loss (CE Loss) flow inside the Megatron backend, using the target model's hidden state and generated tokens as input to the MTP layer (i.e., the draft model), with the expectation that the MTP layer can accurately predict the next token that the target model will actually generate. When the target model's GPRO Loss calculation is complete and backward() is called, it triggers the backpropagation of both the target model's GPRO Loss and the MTP CE Loss simultaneously. The specific process is as follows: