From 1a1477dcbabd7fb97a33bc99e26c3c8689061300 Mon Sep 17 00:00:00 2001 From: Francesco Bertolotti Date: Thu, 6 Aug 2026 08:39:10 +0200 Subject: [PATCH] fix --- torchtitan/hf_datasets/text_datasets.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/torchtitan/hf_datasets/text_datasets.py b/torchtitan/hf_datasets/text_datasets.py index d3d09ff0b5..67e9773ce3 100644 --- a/torchtitan/hf_datasets/text_datasets.py +++ b/torchtitan/hf_datasets/text_datasets.py @@ -160,7 +160,10 @@ def __iter__(self): ] input = x[:-1] - label = x[1:] + label = x[1:].clone() + # Mask EOS tokens in the label to avoid predicting from + # next-document context. + label[input == self._tokenizer.eos_id] = IGNORE_INDEX positions = pos[:-1] yield {"input": input, "positions": positions}, label