Accelerating Direct Preference Optimization with Prefix Sharing

TL;DR

This paper introduces prefix sharing with a custom attention mask to improve the training efficiency of preference tuning algorithms, achieving significant speedups without affecting convergence.

Contribution

It presents a novel prefix sharing technique with a block-sparse attention mask for preference tuning, boosting training throughput significantly.

Findings

Achieves 1.1-1.5x training speedup on DPO datasets

Combining with sequence packing yields 1.3-1.6x speedups

Method is applicable to various paired preference tuning approaches

Abstract

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.