Accelerating Direct Preference Optimization with Prefix Sharing
This work incrementally improves computational efficiency for preference-based fine-tuning, making it more accessible for various applications and model sizes.
The paper tackles computational inefficiency in Direct Preference Optimization by introducing prefix sharing, which processes chosen and rejected responses as one sequence with a shared prefix and uses a custom block-sparse attention mask to prevent contamination. This method achieves 1.1-1.5× training throughput improvements on DPO datasets without affecting convergence, with further speedups to 1.3-1.6× when combined with sequence packing.
Offline paired preference optimization algorithms have become a popular approach for fine-tuning on preference data, outperforming traditional supervised fine-tuning in various tasks. However, traditional implementations often involve redundant computations, especially for tasks with long shared prompts. We introduce prefix sharing for preference tuning, a novel technique that processes chosen and rejected responses as one sequence with a shared prefix. To prevent cross-response contamination, we use a custom block-sparse attention mask. Our method achieves $1.1$-$1.5\times$ improvement in training throughput on popular DPO datasets, without any effect on convergence. When combined with sequence packing, we observe consistent $1.3$-$1.6\times$ speedups, benefiting even datasets with smaller sequence lengths. While we focus on Direct Preference Optimization (DPO), our approach is applicable to other paired preference tuning methods. By enhancing computational efficiency, our work contributes to making preference-based fine-tuning more accessible for a wider range of applications and model sizes. We open-source our code at https://github.com/frankxwang/dpo-prefix-sharing.