🔖 Fix: ensure user-provided labels are retained in self._signature_columns

trl/trainer/sft_trainer.py

@@ -147,7 +147,10 @@ def torch_call(self, examples: list[Union[list[int], Any, dict[str, Any]]]) -> d
                 position_ids = [torch.tensor(example["position_ids"]) for example in examples]
             else:
                 position_ids = [torch.arange(len(ids)) for ids in input_ids]
-        labels = [torch.tensor(example["input_ids"]) for example in examples]
+        if "labels" in examples[0]:
+            labels = [torch.tensor(example["labels"]) for example in examples]
+        else:
+            labels = [torch.tensor(example["input_ids"]) for example in examples]
         if self.completion_only_loss and "completion_mask" in examples[0]:
             completion_mask = [torch.tensor(example["completion_mask"]) for example in examples]
 
@@ -704,7 +707,7 @@ def _set_signature_columns_if_needed(self):
         # and "attention_mask"). When using `train_on_completion_only` we add a "completion_mask" column to the
         # dataset. So we need to override the default signature columns to include "completion_mask" as well.
         if self._signature_columns is None:
-            self._signature_columns = ["input_ids", "attention_mask", "position_ids", "completion_mask"]
+            self._signature_columns = ["input_ids", "labels", "attention_mask", "position_ids", "completion_mask"]
 
     def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
         """