DPO 计算图随机剪枝:训练提速 1.4 倍,峰值显存更低
Stochastic Pruning of the DPO Computation Graph: 1.4x Faster Training and Lower Peak Memory
一篇讲怎么把 DPO 训练提速 30% 的论文,分支剪枝拿到 1.426 倍单步加速,代价是准确率略降,做偏好优化的可以看看。
Direct preference optimization (DPO) 的训练开销来自输出打分、序列执行和反向传播,仅减少监督数据并不能省算力。这项工作提出对 DPO 计算图做随机剪枝的统一框架,分别对打分位置、执行前缀和反向传播的偏好分支独立采样,并在对应操作前完成选择,从而在每个阶段省去部分计算。分析部分刻画了位置剪枝和序列剪枝对目标函数的影响,并用逆概率加权得到条件无偏的单分支梯度估计器。测量结果显示,位置和序列剪枝主要降低峰值显存,分支剪枝相对拼接式 DPO 基线取得 1.426 倍单步加速,在两个模型家族上训练时间比原始 DPO 减少约 30%,平均准确率略降,开放式评分随基准而异。
Stochastic Pruning of the DPO Computation Graph: 1.4x Faster Training and Lower Peak Memory
Direct preference optimization (DPO) incurs training costs in output scoring, sequence execution, and backpropagation, so reducing supervision alone does not necessarily save computation. We introduce a unified framework for stochastic pruning of the DPO computation graph that independently samples response positions for scoring, response prefixes for execution, and preference branches for backpropagation. By making these selections before the corresponding operations, the framework omits computation at each stage of a training step. Our analysis characterizes the objective changes induced by position and sequence pruning and uses inverse-probability weighting to obtain a conditionally unbiased single-branch gradient estimator for the sampled objective. Controlled measurements show that position and sequence pruning primarily reduce peak memory, with most output-memory savings attributable to an implementation that retains full response supervision. Branch pruning yields a $1.426\times$ step-time speedup over a concatenated DPO baseline and reduces training time by about 30\% relative to stock DPO across two model families, with slightly lower mean accuracy and benchmark-dependent changes in open-ended scores. These findings relate the resource benefits of stochastic graph pruning to the computation omitted and the resulting changes in the training objective.