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