bug-fix: avoid 'Expected all tensors to be on the same device' error when doing multi-GPU training
Author: kallewoof
Pull request: yes

When doing DPO training, if the model has been split over multiple GPUs, the `tr_loss` and the `tr_loss_step` end up on different devices at some point, resulting in a 

```
Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1
```
error. ~This patch makes an explicit copy of the `tr_loss_step` value on the same device as the `tr_loss` value, when necessary.~

Update: This patch now explicitly moves the loss value to `self.args.device` to match `tr_loss` in the training step function. This is already done in the `is_sagemaker_mp_enabled()` case.

Ping @patrickvonplaten (git blame last touched), @younesbelkada


Comments:
@SunMarc Still working on that minimal reproducer. I did do some debugging and I am seeing this behavior which does not seem to be correct:

https://github.com/huggingface/trl/blob/4219cbfedcf27ef75754f87079c0d97b767f27f0/trl/trainer/dpo_trainer.py#L870-L872

(above does not render for some reason, so code is copied below)
```python
        pi_logratios = pi_logratios.to(self.accelerator.device)
        ref_logratios = ref_logratios.to(self.accelerator.device)
        logits = pi_logratios - ref_logratios
```

On entry, the policy and reference logps are on separate devices (policy ones on cuda:0, reference ones on cuda:1). After these lines, all of them are on cuda:1, which propagates down to the loss value which is now also on cuda:1.

Compare this to the `tr_loss` variable here:

https://github.com/huggingface/transformers/blob/9322576e2f49d1014fb0c00a7a7c8c34b6a5fd35/src/transformers/trainer.py#L1874

My conclusion (which could be wrong) is that for DPO training, the returned loss var can be on any device, since `self.accelerator.device` can be cuda:1, whereas `tr_loss` is set to cuda:0.

Am I missing some subtlety here?

Update: I tried to set the `tr_loss` device to  `.to(self.accelerator.device)`, which turned out to be `cuda:0`, so the problem remains.

So if we look at 

https://github.com/huggingface/transformers/blob/d45f47ab7f7c31991bb98a0302ded59ab6adac31/src/transformers/trainer.py#L2851-L2888

we see that in one case (L2874) we are transferring the loss to `self.args.device`, and in the other case we aren't. The right fix seems to be to transfer `.to(self.args.device)` in the latter case as well, i.e. on L2888.

Hey @younesbelkada, have you ever had this issue when running training with dpo + device_map on multi-gpu setup ? 

> In principle this looks ok if it fixes some issues ! however I never had this issue in the past, @kallewoof can you clarify a bit on how you're getting this issue?

That's very odd indeed. I will see if I can make a minimum reproducible example on my end. In the meantime, the details of my configuration are as follows:

I have been using Axolotl to do QLoRA DPO training on 2x3090 cards in linux, using `python -m` (not `accelerate launch`). Models were e.g. mixtral-8x7b or llama-70b, but I assume any model that is big enough that it has to be split across the devices should trigger this error, at least sporadically.

See https://github.com/OpenAccess-AI-Collective/axolotl/issues/1302 for details on the Axolotl setup.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

This is now fixed by https://github.com/huggingface/trl/pull/1439 - can we close this PR @kallewoof ? 🙏 

@younesbelkada It looks like it fixes the issue, and does so in a more explicit way (which would explain why not everyone is having this issue). Thanks for pointing it out.