`precompute_ref_log_probs=True` crashes under FSDP2 (device mismatch)
- Dominant language
- Python
- Stars
- 19.3k
- Forks
- 3k
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 194
Description
### Reproduction
With `ref_model is None`, the reference forward runs in `__init__` before `train` prepares the model. Under FSDP2 the params are still on CPU → forward hits CPU weights with CUDA inputs. This is the FSDP2 analogue of the DeepSpeed crashes fixed in #6403; `_precompute_ref_logps` has a DeepSpeed branch but no FSDP2 one, so it falls through to the raw un-prepared module.
**MRE** (2 GPUs):
```python
# mre.py
import tempfile
from datasets import load_dataset
from trl import DPOConfig, DPOTrainer
ds = load_dataset("trl-internal-testing/zen", "standard_preference")
with tempfile.TemporaryDirectory() as tmp:
args = DPOConfig(output_dir=tmp, precompute_ref_log_probs=True, max_steps=2, report_to="none")
DPOTrainer("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", args, train_dataset=ds["train"]).train()
```
```bash
accelerate launch --config_file tests/distributed/data/accelerate_configs/fsdp2.yaml mre.py
```
**Actual:** `RuntimeError: Expected all tensors to be on the same device, but got index is on cuda:0, different from other tensors on cpu`
**Fix:** mirror the DeepSpeed branch via `prepare_fsdp`, reusing one engine like #6403. Apply to DPO + KTO.
### System Info
- Platform: Linux-5.15.0-1048-aws-x86_64-with-glibc2.31
- Python version: 3.13.13
- TRL version: 1.7.0.dev0+da48072
- PyTorch version: 2.11.0+cu128
- accelerator(s): NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3
- Transformers version: 5.14.0.dev0
- Accelerate version: 1.15.0.dev0
- Accelerate config: not found
- Datasets version: 5.0.0
- HF Hub version: 1.23.0
- bitsandbytes version: 0.49.2
- DeepSpeed version: 0.19.2
- Liger-Kernel version: 0.8.0
- PEFT version: 0.19.1
- vLLM version: not installed
### Checklist
- [x] I have checked that my issue isn't already filed (see [open issues](https://github.com/huggingface/trl/issues?q=is%3Aissue))
- [x] I have included my system information
- [x] Any code provided is minimal, complete, and reproducible ([more on MREs](https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/creating-and-highlighting-code-blocks))
- [x] Any code provided is properly formatted in code blocks, (no screenshot, [more on code blocks](https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/creating-and-highlighting-code-blocks))
- [x] Any traceback provided is complete
Contributor guide
Research direction
Start at `_precompute_ref_logps` and compare its existing DeepSpeed branch with the FSDP2 preparation path using `prepare_fsdp`; apply the investigation to the DPO and KTO entry points. Run `mre.py` with `tests/distributed/data/accelerate_configs/fsdp2.yaml`, and consider the issue done when the two-GPU reproduction completes without the CPU/CUDA device mismatch.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 66/100