huggingface / huggingface/trl

`precompute_ref_log_probs=True` crashes under FSDP2 (device mismatch)

Open
#6,470 2 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.