`use_reference_model` is much slower than expected

While testing DPO, I am finding that the weight swap when getting the reference logprobs is a big bottleneck. On a single GPU with Llama 1B, it takes ~11s to get the reference logprobs with the weight swap but under 1s without the weight swap.