Support precomputing reference log probabilities for DPO training
- 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
Assessment
This issue has not been assessed yet.