NVIDIA-NeMo / NVIDIA-NeMo/RL

Support precomputing reference log probabilities for DPO training

Open
#2,484 0 comments 0 reactions 1 assignee Claimed by @ashors1 View on GitHub
community-request enhancement waiting-on-maintainers
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Is your feature request related to a problem? Please describe.**
NeMo-RL currently computes reference model log probs online during DPO training and validation. This increases GPU memory usage and slows down training, and can cause OOMs.

TRL already supports precomputing reference log probs:
https://huggingface.co/docs/trl/main/trainer#trl.KTOConfig.precompute_ref_log_probs

**Describe the solution you'd like**
Add support for precomputing and caching reference log probs before training, similar to TRL. This would reduce memory usage and training/validation overhead.

**Describe alternatives you've considered**
Current workarounds are:
* reducing batch size,
* increasing EP/TP,
* lowering sequence length,

These all reduce throughput or training efficiency.

**Additional context**
This would be especially useful for long-context DPO training.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.