Fix FSDP error (huggingface#1196)

* Fix FSDP error Fixes error when `loss` field of model output is non-empty, and indexing as [0] returns loss instead of logits. Can happen with FSDP. * Apply suggestions from code review force return_dict Co-authored-by: Younes Belkada <[email protected]> --------- Co-authored-by: Younes Belkada <[email protected]>
lapp0 · May 10, 2024 · f30932e · f30932e
1 parent 29cf6c8
commit f30932e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py
@@ -220,11 +220,13 @@ def compute_loss(
         rewards_chosen = model(
             input_ids=inputs["input_ids_chosen"],
             attention_mask=inputs["attention_mask_chosen"],
-        )[0]
+            return_dict=True,
+        )["logits"]
         rewards_rejected = model(
             input_ids=inputs["input_ids_rejected"],
             attention_mask=inputs["attention_mask_rejected"],
-        )[0]
+            return_dict=True,
+        )["logits"]
         # calculate loss, optionally modulate with margin
         if "margin" in inputs:
             loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs["margin"]).mean()