NVIDIA NeMo Megatron Bridge 支持离线参考打分的 DPO 训练
DPO 训练不用再挂两份模型了,NVIDIA 把参考打分挪到离线做一次,显存和速度跟 SFT 一个样,调参还不用重算。
NVIDIA 在其训练库 NeMo Megatron Bridge 中实现了 DPO 训练,与 SFT 一样只用一份模型,显存占用和步进耗时不变。传统 DPO 需要在每个训练步额外运行一个冻结的参考模型,该实现改为离线打分一次,把每对 chosen/rejected 的 logprob 总和存到磁盘,训练时直接读取。文件保存的是未加 beta 的原始总和,beta 和学习率在损失计算时才引入,一次打分即可覆盖后续所有超参数搜索。
1/2 DPO in NeMo™ Megatron Bridge, @NVIDIA’s training library for Megatron-based models, now trains like SFT: one model on the GPUs, same memory, same step time.
DPO's reference model is frozen, so its logprobs over your dataset never change. Yet the usual setup keeps a second copy of the model and runs it at every step.
Our implementation scores the reference once (offline scoring) and saves chosen/rejected logprob sums per pair to disk. Training reads them back instead.
The best part: the file stores raw sums, with no beta applied. Beta and learning rate only enter at loss time, so one scoring run covers every hyperparameter sweep you do afterwards. >>