From c972c4e09018be0646ce1abd110239cd03fc0b24 Mon Sep 17 00:00:00 2001
From: xiao-baia <76036940+xiao-baia@users.noreply.github.com>
Date: Wed, 7 Jan 2026 21:00:46 +0800
Subject: [PATCH] Fix DPO loss_mask boundary (include first assistant token)

Fix off-by-one in generate_loss_mask so the first assistant token contributes to loss in DPO.
---
 dataset/lm_dataset.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py
index 1fa0397..c5fc48e 100644
--- a/dataset/lm_dataset.py
+++ b/dataset/lm_dataset.py
@@ -159,7 +159,7 @@ class DPODataset(Dataset):
                     if input_ids[end:end + len(self.eos_id)] == self.eos_id:
                         break
                     end += 1
-                for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
+                for j in range(start, min(end + len(self.eos_id) + 1, self.max_length)):
                     loss_mask[j] = 1
                 i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
             else:
@@ -202,4 +202,4 @@ class RLAIFDataset(Dataset):
         }
 
 if __name__ == "__main__":
-    pass
\ No newline at end of file
+    pass