Hugging Face Journal Club: AsyncOPD and How Stale Can On-Policy Distillation Be?
The discussion details Asynchronous On-Policy Distillation (AsyncOPD), a method designed to significantly boost training throughput by making the distillation process fully asynchronous. While conventional methods are synchronous and suffer from GPU blocking during backpropagation, AsyncOPD continuously generates rollouts from policies while simultaneously scoring them with a teacher model. This approach achieves substantial speedups (1.5x to 2.7x) compared to synchronous methods, though it introduces complexity related to maintaining stability when student rollouts become significantly off-policy.
Key takeaways
-
AsyncOPD for Throughput Gains
23:23
By decoupling the generator (student policy), scorer (teacher model), and backpropagator, AsyncOPD eliminates GPU blocking inherent in synchronous distillation. This allows continuous operation, leading to throughput improvements of 1.5x to 2.7x on various math benchmarks.
-
Addressing Cache Misses via Monte Carlo Sampling
30:07
When calculating Reverse KL divergence using Top-K logits, cache misses can occur because the required log probabilities for the loss calculation may not have been stored during the initial sampling phase. MC sampling is proposed as a solution to estimate the loss accurately by storing and correcting estimates using important sampling.
-
Trade-offs in Off-Policy Distillation
20:50
While fully asynchronous methods offer high throughput, they require careful handling of off-policy rollouts. The stability and accuracy are dependent on the degree of staleness allowed (the difference between the current policy and the teacher's distribution).