DPOTrainer fails with FileNotFoundError when `precompute_ref_log_probs=True`

[saattrupdan]

Reproduction

When using DPOTrainer with precompute_ref_log_probs=True, training fails with a
FileNotFoundError on the Arrow cache file. The issue is in
DPOTrainer._precompute_ref_logps() which computes a cache file path but doesn't pass
it to dataset.map().

Minimal Reproduction Script

Only precompute_ref_log_probs=True is required to trigger the bug. Any causal LM
works—the bug is in TRL's dataset caching, not model-specific.

import os
import tempfile
from datasets import Dataset
from trl import DPOConfig, DPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer

with tempfile.TemporaryDirectory() as tmpdir:
    os.environ["TMPDIR"] = tmpdir

    # Minimal dataset (any causal LM format works)
    dataset = Dataset.from_dict({
        "prompt": ["x"] * 4,
        "chosen": ["a"] * 4,
        "rejected": ["b"] * 4,
    })

    # Any causal LM works—not model-specific
    model_name = "danish-foundation-models/munin-apertus-8b"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map="cpu",  # or "cuda" if GPU available
    )

    # Only precompute_ref_log_probs=True triggers the bug
    config = DPOConfig(
        output_dir="/tmp/test-dpo",
        precompute_ref_log_probs=True,
    )

    # Raises FileNotFoundError: cache file written to wrong path
    trainer = DPOTrainer(
        model=model,
        ref_model=None,
        args=config,
        train_dataset=dataset,
        processing_class=tokenizer,
    )

Error Output

FileNotFoundError: [Errno 2] Failed to open local file
'/tmp/hf_datasets-*/cache-*.arrow'. Detail: [errno 2] No such file or directory

Traceback (most recent call last):
  File "repro.py", line XX, in <module>
    trainer = DPOTrainer(...)
              ^^^^^^^^^^^^^
  File "/path/to/trl/trl/trainer/dpo_trainer.py", line 871, in __init__
    self.train_dataset = self._precompute_ref_logps(...)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/to/trl/trl/trainer/dpo_trainer.py", line 1092, in _precompute_ref_logps
    return concatenate_datasets([dataset, Dataset.from_file(cache_file)], axis=1)
                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/to/datasets/datasets/arrow_dataset.py", line 827, in from_file
    table = ArrowReader.read_table(filename, in_memory=in_memory)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  ...

Root Cause

In trl/trainer/dpo_trainer.py, the _precompute_ref_logps() method:

1. Line ~1047: Computes expected cache file path:

   cache_file = dataset._get_cache_file_path(fingerprint)

2. Lines ~1085-1091: Calls dataset.map() without passing cache_file_name:

   dataset.map(
       add_ref_logps,
       with_indices=True,
       batched=True,
       remove_columns=dataset.column_names,
       new_fingerprint=fingerprint,
       # cache_file_name=cache_file,  # <-- MISSING!
       desc=f"Caching reference log probs for {name} dataset",
   )

3. Line ~1092: Tries to read from cache_file which was never written:

   return concatenate_datasets([dataset, Dataset.from_file(cache_file)], axis=1)

Why it fails: When dataset.map() is called without cache_file_name and the
dataset has no cache_files, the datasets library generates a different random cache
path than what TRL computed. The cache is written to the random path, but TRL tries to
read from the original cache_file path.

Expected Behavior

The trainer should initialize successfully and precompute reference log probs.

Workaround

Patch DPOTrainer._precompute_ref_logps() to pass cache_file_name=cache_file to
dataset.map():

dataset.map(
    add_ref_logps,
    with_indices=True,
    batched=True,
    remove_columns=dataset.column_names,
    new_fingerprint=fingerprint,
    cache_file_name=cache_file,  # <-- THE FIX
    desc=f"Caching reference log probs for {name} dataset",
)

System Info

Platform: Linux (aarch64)
Python version: 3.12.13
PyTorch version: 2.8.0+cuda12.8
Transformers version: 4.53.0
TRL version: 1.7.0
CUDA version: 13.0
GPU: NVIDIA GB10 (128GB unified memory)

Checklist

• I have checked that my issue isn't already filed (see
  open issues)
• I have included my system information
• Any code provided is minimal, complete, and reproducible
• Any code provided is properly formatted in code blocks, (no screenshot)
• Any traceback provided is complete

[Functionhx]

Found the bug — _precompute_ref_logps computes a cache_file path but doesn't pass it to dataset.map(), so datasets writes to a different path and FileNotFoundError follows. Fixed in both DPO and KTO. PR ready.